Call warmup_speculative at startup to pre-compile LpB kernels

The warmup_speculative() function was defined but never called. Custom Metal kernels (LpB) require first-call compilation (~200ms). Without warmup, the first speculative cycle is slow, dragging down average TPS by 10-20% on short generations. In mlx_bench testing: cold 48 TPS → warm 60 TPS for DFlash, cold 39 TPS → warm 44 TPS for MTP. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Add EXO_DISABLE_LOGPROBS=1 to skip per-token logprobs extraction
2026-04-01 23:15:05 +01:00 · 2026-03-30 23:41:47 +01:00 · 2026-03-30 22:39:34 +01:00 · 2026-03-30 18:44:01 +01:00 · 2026-03-30 18:24:25 +01:00 · 2026-03-30 18:23:05 +01:00
124 changed files with 6801 additions and 3221 deletions
@@ -159,7 +159,7 @@ jobs:
          fi

      - name: Install Homebrew packages
-        run: brew install just awscli
+        run: brew install just awscli macmon

      - name: Install UV
        uses: astral-sh/setup-uv@v6
@@ -243,14 +243,6 @@ jobs:
      # Build the bundle
      # ============================================================

-      - name: Add pinned macmon to PATH
-        run: |
-          MACMON_DIR=$(nix develop --command sh -c 'dirname $(which macmon)')
-          echo "Using macmon from: $MACMON_DIR"
-          echo "$MACMON_DIR" >> $GITHUB_PATH
-          # Remove any Homebrew macmon so PyInstaller can't accidentally pick it up
-          brew uninstall macmon 2>/dev/null || true
-
      - name: Build PyInstaller bundle
        run: uv run pyinstaller packaging/pyinstaller/exo.spec

@@ -11,18 +11,9 @@ To run EXO from source:
  ```bash
  brew install uv
  ```
- [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now)
-  ```bash
-  curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
-  rustup toolchain install nightly
-  ```
 - [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
-  Use the pinned fork revision used by this repo instead of Homebrew `macmon`.
  ```bash
-  cargo install --git https://github.com/swiftraccoon/macmon \
-    --rev 9154d234f763fbeffdcb4135d0bbbaf80609699b \
-    macmon \
-    --force
+  brew install macmon
  ```

 ```bash
@@ -95,10 +95,11 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
  /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
  ```
 - [uv](https://github.com/astral-sh/uv) (for Python dependency management)
+- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
 - [node](https://github.com/nodejs/node) (for building the dashboard)

  ```bash
-  brew install uv node
+  brew install uv macmon node
  ```
 - [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now)

@@ -106,17 +107,6 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
  curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
  rustup toolchain install nightly
  ```
- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
-
-  Install the pinned fork revision used by this repo instead of Homebrew `macmon`.
-  Homebrew `macmon 0.6.1` still crashes on Apple M5.
-
-  ```bash
-  cargo install --git https://github.com/swiftraccoon/macmon \
-    --rev 9154d234f763fbeffdcb4135d0bbbaf80609699b \
-    macmon \
-    --force
-  ```

 Clone the repo, build the dashboard, and run exo:

@@ -295,9 +285,8 @@ exo supports several environment variables for configuration:

 | Variable | Description | Default |
 |----------|-------------|---------|
-| `EXO_DEFAULT_MODELS_DIR` | Default directory for model downloads and caches. Always first in the writable dirs list. | `~/.local/share/exo/models` (Linux) or `~/.exo/models` (macOS) |
-| `EXO_MODELS_DIRS` | Colon-separated additional writable directories for model downloads. Checked in order after the default; first with enough free space is used. | None |
-| `EXO_MODELS_READ_ONLY_DIRS` | Colon-separated read-only directories to search for pre-downloaded models (e.g., NFS mounts, shared storage). Models here cannot be deleted. | None |
+| `EXO_MODELS_PATH` | Colon-separated paths to search for pre-downloaded models (e.g., on NFS mounts or shared storage) | None |
+| `EXO_MODELS_DIR` | Directory where exo downloads and stores models | `~/.local/share/exo/models` (Linux) or `~/.exo/models` (macOS) |
 | `EXO_OFFLINE` | Run without internet connection (uses only local models) | `false` |
 | `EXO_ENABLE_IMAGE_MODELS` | Enable image model support | `false` |
 | `EXO_LIBP2P_NAMESPACE` | Custom namespace for cluster isolation | None |
@@ -307,11 +296,8 @@ exo supports several environment variables for configuration:
 **Example usage:**

 ```bash
-# Use pre-downloaded models from NFS mount (read-only)
-EXO_MODELS_READ_ONLY_DIRS=/mnt/nfs/models:/opt/ai-models uv run exo
-
-# Download models to an external SSD (falls back to default dir if full)
-EXO_MODELS_DIRS=/Volumes/ExternalSSD/exo-models uv run exo
+# Use pre-downloaded models from NFS mount
+EXO_MODELS_PATH=/mnt/nfs/models:/opt/ai-models uv run exo

 # Run in offline mode
 EXO_OFFLINE=true uv run exo
@@ -4,7 +4,6 @@ import Foundation

 private let customNamespaceKey = "EXOCustomNamespace"
 private let hfTokenKey = "EXOHFToken"
-private let hfEndpointKey = "EXOHFEndpoint"
 private let enableImageModelsKey = "EXOEnableImageModels"
 private let offlineModeKey = "EXOOfflineMode"
 private let onboardingCompletedKey = "EXOOnboardingCompleted"
@@ -54,14 +53,6 @@ final class ExoProcessController: ObservableObject {
            UserDefaults.standard.set(hfToken, forKey: hfTokenKey)
        }
    }
-    @Published var hfEndpoint: String = {
-        return UserDefaults.standard.string(forKey: hfEndpointKey) ?? ""
-    }()
-    {
-        didSet {
-            UserDefaults.standard.set(hfEndpoint, forKey: hfEndpointKey)
-        }
-    }
    @Published var enableImageModels: Bool = {
        return UserDefaults.standard.bool(forKey: enableImageModelsKey)
    }()
@@ -282,9 +273,6 @@ final class ExoProcessController: ObservableObject {
        if !hfToken.isEmpty {
            environment["HF_TOKEN"] = hfToken
        }
-        if !hfEndpoint.isEmpty {
-            environment["HF_ENDPOINT"] = hfEndpoint
-        }
        if enableImageModels {
            environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
        }
@@ -12,7 +12,6 @@ struct SettingsView: View {

    @State private var pendingNamespace: String = ""
    @State private var pendingHFToken: String = ""
-    @State private var pendingHFEndpoint: String = ""
    @State private var pendingEnableImageModels = false
    @State private var pendingOfflineMode = false
    @State private var needsRestart = false
@@ -43,7 +42,6 @@ struct SettingsView: View {
        .onAppear {
            pendingNamespace = controller.customNamespace
            pendingHFToken = controller.hfToken
-            pendingHFEndpoint = controller.hfEndpoint
            pendingEnableImageModels = controller.enableImageModels
            pendingOfflineMode = controller.offlineMode
            needsRestart = false
@@ -76,17 +74,6 @@ struct SettingsView: View {
                    .foregroundColor(.secondary)
            }

-            Section {
-                LabeledContent("HuggingFace Endpoint") {
-                    TextField("default", text: $pendingHFEndpoint)
-                        .textFieldStyle(.roundedBorder)
-                        .frame(width: 200)
-                }
-                Text("Defaults to huggingface.co. Use a mirror (e.g. hf-mirror.com) for China.")
-                    .font(.caption)
-                    .foregroundColor(.secondary)
-            }
-
            Section {
                Toggle("Offline Mode", isOn: $pendingOfflineMode)
                Text("Skip internet checks and use only locally available models.")
@@ -467,7 +454,6 @@ struct SettingsView: View {

    private var hasGeneralChanges: Bool {
        pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
-            || pendingHFEndpoint != controller.hfEndpoint
            || pendingOfflineMode != controller.offlineMode
    }

@@ -478,7 +464,6 @@ struct SettingsView: View {
    private func applyGeneralSettings() {
        controller.customNamespace = pendingNamespace
        controller.hfToken = pendingHFToken
-        controller.hfEndpoint = pendingHFEndpoint
        controller.offlineMode = pendingOfflineMode
        restartIfRunning()
    }
@@ -377,7 +377,7 @@ def run_planning_phase(
                f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
            )

-        # Delete from smallest to largest (skip read-only models)
+        # Delete from smallest to largest (skip read-only models from EXO_MODELS_PATH)
        completed = [
            (
                unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
@@ -0,0 +1,91 @@
+#!/usr/bin/env bash
+#
+# Run exo_bench.py for each model/mode from bench_params.json.
+#
+# For each entry, runs with:
+#   --pp 800 (fixed, representative LCB prompt length)
+#   --tg <mean completion tokens from vLLM>
+#   --sharding tensor --instance-meta jaccl
+#   --min-nodes 1 --max-nodes 4
+#   --repeat 1
+#   --danger-delete-downloads
+#   --settle-timeout 300
+#
+# Results go to bench/eval_results/<model_dir>/tps_<mode>.json
+#
+# Usage:
+#   bash bench/run_lcb_tps_bench.sh              # run all
+#   bash bench/run_lcb_tps_bench.sh --dry-run    # show what would run
+
+set -euo pipefail
+cd "$(dirname "$0")"
+
+PARAMS_FILE="eval_results/bench_params.json"
+PP=800
+HOST="${EXO_HOST:-s9}"
+DRY_RUN=false
+
+if [[ "${1:-}" == "--dry-run" ]]; then
+    DRY_RUN=true
+fi
+
+if [[ ! -f "$PARAMS_FILE" ]]; then
+    echo "ERROR: $PARAMS_FILE not found. Run compute_bench_params.py first."
+    exit 1
+fi
+
+# Parse bench_params.json and run each entry
+python3 -c "
+import json, sys
+data = json.load(open('$PARAMS_FILE'))
+for entry in data:
+    mlx_id = entry['mlx_model_id']
+    mode = entry['mode']
+    tg = entry['bench_params']['tg']
+    vllm_name = entry['vllm_name']
+    # Output dir: replace / with _
+    out_dir = 'eval_results/' + mlx_id.replace('/', '_')
+    out_file = out_dir + '/tps_' + mode + '.json'
+    print(f'{mlx_id}\t{mode}\t{tg}\t{out_file}\t{vllm_name}')
+" | while IFS=$'\t' read -r model mode tg out_file vllm_name; do
+    out_dir="$(dirname "$out_file")"
+    mkdir -p "$out_dir"
+
+    echo ""
+    echo "============================================================"
+    echo "Model:  $model"
+    echo "Mode:   $mode"
+    echo "vLLM:   $vllm_name"
+    echo "PP:     $PP"
+    echo "TG:     $tg"
+    echo "Output: $out_file"
+    echo "============================================================"
+
+    if [[ -f "$out_file" ]]; then
+        echo "SKIP: $out_file already exists"
+        continue
+    fi
+
+    if [[ "$DRY_RUN" == "true" ]]; then
+        echo "DRY-RUN: would run exo_bench.py"
+        continue
+    fi
+
+    uv run python exo_bench.py \
+        --host "$HOST" \
+        --model "$model" \
+        --pp "$PP" \
+        --tg "$tg" \
+        --repeat 1 \
+        --sharding tensor \
+        --instance-meta jaccl \
+        --min-nodes 1 \
+        --max-nodes 4 \
+        --settle-timeout 300 \
+        --force-download \
+        --danger-delete-downloads \
+        --json-out "$out_file" || echo "FAILED: $model ($mode)"
+done
+
+echo ""
+echo "All benchmarks complete."
@@ -88,12 +88,6 @@
      d="M22.012 0h1.032v.927H24v.968h-.956V3.78h-1.032V1.896h-1.878v-.97h1.878V0zM2.6 12.371V1.87h.969v10.502h-.97zm10.423.66h10.95v.918h-6.208v9.579h-4.742V13.03zM5.629 3.333v12.356H0v4.51h10.386V8L20.859 8l-.003-4.668-15.227.001z"
    />
  </svg>
-{:else if family === "nemotron"}
-  <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
-    <path
-      d="M8.948 8.798v-1.43a6.7 6.7 0 0 1 .424-.018c3.922-.124 6.493 3.374 6.493 3.374s-2.774 3.851-5.75 3.851c-.398 0-.787-.062-1.158-.185v-4.346c1.528.185 1.837.857 2.747 2.385l2.04-1.714s-1.492-1.952-4-1.952a6.016 6.016 0 0 0-.796.035m0-4.735v2.138l.424-.027c5.45-.185 9.01 4.47 9.01 4.47s-4.08 4.964-8.33 4.964c-.37 0-.733-.035-1.095-.097v1.325c.3.035.61.062.91.062 3.957 0 6.82-2.023 9.593-4.408.459.371 2.34 1.263 2.73 1.652-2.633 2.208-8.772 3.984-12.253 3.984-.335 0-.653-.018-.971-.053v1.864H24V4.063zm0 10.326v1.131c-3.657-.654-4.673-4.46-4.673-4.46s1.758-1.944 4.673-2.262v1.237H8.94c-1.528-.186-2.73 1.245-2.73 1.245s.68 2.412 2.739 3.11M2.456 10.9s2.164-3.197 6.5-3.533V6.201C4.153 6.59 0 10.653 0 10.653s2.35 6.802 8.948 7.42v-1.237c-4.84-.6-6.492-5.936-6.492-5.936z"
-    />
-  </svg>
 {:else}
  <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
    <path
@@ -31,7 +31,6 @@
    kimi: "Kimi",
    flux: "FLUX",
    "qwen-image": "Qwen Img",
-    nemotron: "NVIDIA",
  };

  function getFamilyName(family: string): string {
@@ -42,20 +41,31 @@
 </script>

 <div
-  class="flex flex-col gap-1 py-2 px-1 border-r border-exo-yellow/10 bg-exo-medium-gray/30 min-w-[80px] sm:min-w-[72px] overflow-y-auto scrollbar-hide"
+  class="flex flex-col gap-1 py-2 px-1 border-r border-exo-yellow/10 bg-exo-medium-gray/30 min-w-[72px] sm:min-w-[64px] overflow-y-auto scrollbar-hide"
 >
  <!-- All models (no filter) -->
  <button
    type="button"
    onclick={() => onSelect(null)}
-    class="group flex items-center justify-center px-3 py-2.5 rounded transition-all duration-200 cursor-pointer min-h-[44px] sm:min-h-0 {selectedFamily ===
+    class="group flex flex-col items-center justify-center p-2 sm:p-2 rounded transition-all duration-200 cursor-pointer min-h-[44px] sm:min-h-0 {selectedFamily ===
    null
      ? 'bg-exo-yellow/20 border-l-2 border-exo-yellow'
      : 'hover:bg-white/5 border-l-2 border-transparent'}"
    title="All models"
  >
+    <svg
+      class="w-5 h-5 {selectedFamily === null
+        ? 'text-exo-yellow'
+        : 'text-white/50 group-hover:text-white/70'}"
+      viewBox="0 0 24 24"
+      fill="currentColor"
+    >
+      <path
+        d="M4 8h4V4H4v4zm6 12h4v-4h-4v4zm-6 0h4v-4H4v4zm0-6h4v-4H4v4zm6 0h4v-4h-4v4zm6-10v4h4V4h-4zm-6 4h4V4h-4v4zm6 6h4v-4h-4v4zm0 6h4v-4h-4v4z"
+      />
+    </svg>
    <span
-      class="text-[12px] font-mono font-medium {selectedFamily === null
+      class="text-[9px] font-mono mt-0.5 {selectedFamily === null
        ? 'text-exo-yellow'
        : 'text-white/40 group-hover:text-white/60'}">All</span
    >
@@ -79,7 +89,7 @@
          : "text-white/50 group-hover:text-amber-400/70"}
      />
      <span
-        class="text-[11px] font-mono mt-0.5 {selectedFamily === 'favorites'
+        class="text-[9px] font-mono mt-0.5 {selectedFamily === 'favorites'
          ? 'text-amber-400'
          : 'text-white/40 group-hover:text-white/60'}">Faves</span
      >
@@ -104,7 +114,7 @@
          : "text-white/50 group-hover:text-white/70"}
      />
      <span
-        class="text-[11px] font-mono mt-0.5 {selectedFamily === 'recents'
+        class="text-[9px] font-mono mt-0.5 {selectedFamily === 'recents'
          ? 'text-exo-yellow'
          : 'text-white/40 group-hover:text-white/60'}">Recent</span
      >
@@ -128,7 +138,7 @@
        : "text-white/50 group-hover:text-orange-400/70"}
    />
    <span
-      class="text-[11px] font-mono mt-0.5 {selectedFamily === 'huggingface'
+      class="text-[9px] font-mono mt-0.5 {selectedFamily === 'huggingface'
        ? 'text-orange-400'
        : 'text-white/40 group-hover:text-white/60'}">Hub</span
    >
@@ -154,7 +164,7 @@
          : "text-white/50 group-hover:text-white/70"}
      />
      <span
-        class="text-[11px] font-mono mt-0.5 truncate max-w-full {selectedFamily ===
+        class="text-[9px] font-mono mt-0.5 truncate max-w-full {selectedFamily ===
        family
          ? 'text-exo-yellow'
          : 'text-white/40 group-hover:text-white/60'}"
@@ -16,9 +16,7 @@
      perNode?: Array<{
        nodeId: string;
        nodeName: string;
-        status: "completed" | "partial" | "pending" | "downloading";
-        percentage: number;
-        progress: DownloadProgress | null;
+        progress: DownloadProgress;
      }>;
    } | null;
    nodes?: Record<string, NodeInfo>;
@@ -147,7 +145,10 @@
    return `${s}s`;
  }

-  const perNode = $derived(downloadStatus?.perNode ?? []);
+  const isDownloading = $derived(downloadStatus?.isDownloading ?? false);
+  const progress = $derived(downloadStatus?.progress);
+  const percentage = $derived(progress?.percentage ?? 0);
+  let expandedNodes = $state<Set<string>>(new Set());

  function toggleNodeDetails(nodeId: string): void {
    const next = new Set(expandedNodes);
@@ -586,49 +587,23 @@
      </span>
    </div>

-    <!-- Download Status (per-node) -->
-    {#if perNode.length > 0}
+    <!-- Download Status -->
+    {#if isDownloading && progress}
      <div class="mb-2 space-y-1">
-        <div
-          class="text-[10px] font-mono text-white/20 tracking-widest uppercase"
-        >
-          Download progress
+        <div class="flex items-center justify-between text-xs font-mono">
+          <span class="text-blue-400 tracking-wider uppercase">Downloading</span
+          >
+          <span class="text-white/60"
+            >{percentage.toFixed(1)}% &middot; {formatSpeed(progress.speed)}
+            &middot; {formatEta(progress.etaMs)}</span
+          >
+        </div>
+        <div class="h-1 bg-exo-medium-gray/30 rounded overflow-hidden">
+          <div
+            class="h-full bg-blue-500/70 transition-all duration-300"
+            style="width: {percentage}%"
+          ></div>
        </div>
-        {#each perNode as node}
-          <div class="flex items-center gap-2 text-xs font-mono">
-            <span class="text-white/40 w-20 truncate" title={node.nodeId}
-              >{node.nodeName}</span
-            >
-            <div
-              class="flex-1 h-1 bg-exo-medium-gray/30 rounded overflow-hidden"
-            >
-              <div
-                class="h-full transition-all duration-300 {node.status ===
-                'downloading'
-                  ? 'bg-blue-500/70'
-                  : node.status === 'completed'
-                    ? 'bg-exo-yellow/40'
-                    : 'bg-white/20'}"
-                style="width: {node.percentage}%"
-              ></div>
-            </div>
-            <span
-              class="text-right {node.status === 'completed'
-                ? 'text-exo-yellow/60'
-                : node.status === 'downloading'
-                  ? 'text-blue-400/60'
-                  : 'text-white/30'}"
-            >
-              {#if node.status === "downloading" && node.progress}
-                {Math.round(node.percentage)}% {formatSpeed(
-                  node.progress.speed,
-                )}
-              {:else}
-                {node.percentage > 0 ? `${Math.round(node.percentage)}%` : "0%"}
-              {/if}
-            </span>
-          </div>
-        {/each}
      </div>
    {/if}

@@ -687,7 +662,15 @@
            {@const allConnections =
              isDebugMode && usedNodes.length > 1
                ? (() => {
-                    const conns: Array = [];
+                    const conns: Array<{
+                      ip: string;
+                      iface: string | null;
+                      from: string;
+                      to: string;
+                      midX: number;
+                      midY: number;
+                      arrow: string;
+                    }> = [];
                    for (let i = 0; i < usedNodes.length; i++) {
                      for (let j = i + 1; j < usedNodes.length; j++) {
                        const n1 = usedNodes[i];
@@ -699,12 +682,7 @@
                          const toPos = nodePositions[c.to];
                          const arrow =
                            fromPos && toPos ? getArrow(fromPos, toPos) : "→";
-                          conns.push({
-                            ...c,
-                            midX,
-                            midY,
-                            arrow,
-                          });
+                          conns.push({ ...c, midX, midY, arrow });
                        }
                      }
                    }
@@ -73,7 +73,7 @@

 <!-- svelte-ignore a11y_no_static_element_interactions -->
 <div
-  class="filter-popover absolute right-0 top-full mt-2 w-64 bg-exo-dark-gray border border-exo-yellow/10 rounded-lg shadow-xl z-20"
+  class="filter-popover absolute right-0 top-full mt-2 w-64 bg-exo-dark-gray border border-exo-yellow/10 rounded-lg shadow-xl z-10"
  transition:fly={{ y: -10, duration: 200, easing: cubicOut }}
  onclick={(e) => e.stopPropagation()}
  role="dialog"
@@ -459,7 +459,6 @@
      "llama",
      "flux",
      "qwen-image",
-      "nemotron",
    ];
    return Array.from(families).sort((a, b) => {
      const aIdx = familyOrder.indexOf(a);
@@ -2650,9 +2650,6 @@ class AppStore {
    this.syncActiveMessagesIfNeeded(targetConversationId);
    this.saveConversationsToStorage();

-    const abortController = new AbortController();
-    this.currentAbortController = abortController;
-
    try {
      // Determine the model to use
      const model = this.getModelForRequest(modelId);
@@ -2707,7 +2704,6 @@ class AppStore {
          "Content-Type": "application/json",
        },
        body: JSON.stringify(requestBody),
-        signal: abortController.signal,
      });

      if (!response.ok) {
@@ -2847,27 +2843,14 @@ class AppStore {
        );
      }
    } catch (error) {
-      if (abortController.signal.aborted) {
-        this.updateConversationMessage(
-          targetConversationId,
-          assistantMessage.id,
-          (msg) => {
-            msg.content = "Cancelled";
-            msg.attachments = [];
-          },
-        );
-        this.syncActiveMessagesIfNeeded(targetConversationId);
-      } else {
-        console.error("Error generating image:", error);
-        this.handleStreamingError(
-          error,
-          targetConversationId,
-          assistantMessage.id,
-          "Failed to generate image",
-        );
-      }
+      console.error("Error generating image:", error);
+      this.handleStreamingError(
+        error,
+        targetConversationId,
+        assistantMessage.id,
+        "Failed to generate image",
+      );
    } finally {
-      this.currentAbortController = null;
      this.isLoading = false;
      this.saveConversationsToStorage();
    }
@@ -2931,9 +2914,6 @@ class AppStore {
    // Clear editing state
    this.editingImage = null;

-    const abortController = new AbortController();
-    this.currentAbortController = abortController;
-
    try {
      // Determine the model to use
      const model = this.getModelForRequest(modelId);
@@ -2995,7 +2975,6 @@ class AppStore {
      const apiResponse = await fetch("/v1/images/edits", {
        method: "POST",
        body: formData,
-        signal: abortController.signal,
      });

      if (!apiResponse.ok) {
@@ -3096,27 +3075,14 @@ class AppStore {
        );
      }
    } catch (error) {
-      if (abortController.signal.aborted) {
-        this.updateConversationMessage(
-          targetConversationId,
-          assistantMessage.id,
-          (msg) => {
-            msg.content = "Cancelled";
-            msg.attachments = [];
-          },
-        );
-        this.syncActiveMessagesIfNeeded(targetConversationId);
-      } else {
-        console.error("Error editing image:", error);
-        this.handleStreamingError(
-          error,
-          targetConversationId,
-          assistantMessage.id,
-          "Failed to edit image",
-        );
-      }
+      console.error("Error editing image:", error);
+      this.handleStreamingError(
+        error,
+        targetConversationId,
+        assistantMessage.id,
+        "Failed to edit image",
+      );
    } finally {
-      this.currentAbortController = null;
      this.isLoading = false;
      this.saveConversationsToStorage();
    }
@@ -42,7 +42,6 @@
    setSelectedChatModel,
    selectedChatModel,
    sendMessage,
-    thinkingEnabled,
    generateImage,
    editImage,
    editingImage,
@@ -853,7 +852,7 @@
  ) {
    const model = selectedChatModel();
    if (!model) {
-      sendMessage(content, files, thinkingEnabled());
+      sendMessage(content, files, null);
      return;
    }

@@ -881,7 +880,7 @@
    }

    // Default: text chat
-    sendMessage(content, files, thinkingEnabled());
+    sendMessage(content, files, null);
  }

  let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
@@ -1536,44 +1535,34 @@
  }

  // Helper to get download status for a model (checks all downloads for matching model ID)
-  type NodeDownloadStatus = {
-    nodeId: string;
-    nodeName: string;
-    status: "completed" | "partial" | "pending" | "downloading";
-    percentage: number;
-    progress: DownloadProgress | null;
-  };
-
-  // Shared helper: collect per-node download status for a model across a set of nodes.
-  // Handles deduplication, entry parsing, and aggregation in one place.
-  function collectDownloadStatus(
-    modelId: string,
-    nodeIds?: string[],
-  ): {
+  function getModelDownloadStatus(modelId: string): {
    isDownloading: boolean;
    progress: DownloadProgress | null;
-    perNode: NodeDownloadStatus[];
-    failedError: string | null;
+    perNode: Array<{
+      nodeId: string;
+      nodeName: string;
+      progress: DownloadProgress;
+    }>;
  } {
-    const empty = {
-      isDownloading: false,
-      progress: null,
-      perNode: [] as NodeDownloadStatus[],
-      failedError: null,
-    };
-
    if (!downloadsData || Object.keys(downloadsData).length === 0) {
-      return empty;
+      return { isDownloading: false, progress: null, perNode: [] };
    }

-    // Deduplicate by nodeId — a node can have multiple entries for the same model
-    // (e.g. PipelineShardMetadata + TensorShardMetadata). Keep the last entry,
-    // which is the most recently applied event.
-    const perNodeMap = new Map<string, NodeDownloadStatus>();
+    let totalBytes = 0;
+    let downloadedBytes = 0;
+    let totalSpeed = 0;
+    let completedFiles = 0;
+    let totalFiles = 0;
+    let isDownloading = false;
+    const allFiles: DownloadProgress["files"] = [];
+    const perNode: Array<{
+      nodeId: string;
+      nodeName: string;
+      progress: DownloadProgress;
+    }> = [];

-    const nodeIdSet = nodeIds ? new Set(nodeIds) : null;
+    // Check all nodes for downloads matching this model
    for (const [nodeId, nodeDownloads] of Object.entries(downloadsData)) {
-      if (nodeIdSet && !nodeIdSet.has(nodeId)) continue;
      if (!Array.isArray(nodeDownloads)) continue;

      for (const downloadWrapped of nodeDownloads) {
@@ -1586,45 +1575,29 @@
        const downloadPayload = (downloadWrapped as Record<string, unknown>)[
          downloadKind
        ] as Record<string, unknown>;
-        if (!downloadPayload) continue;
-
-        const downloadModelId = extractModelIdFromDownload(downloadPayload);
-        if (!downloadModelId || downloadModelId !== modelId) continue;
-
-        // DownloadFailed — return with any data collected so far
-        if (downloadKind === "DownloadFailed") {
-          return {
-            isDownloading: false,
-            progress: null,
-            perNode: Array.from(perNodeMap.values()),
-            failedError:
-              (downloadPayload.errorMessage as string) ||
-              (downloadPayload.error_message as string) ||
-              "Download failed",
-          };
-        }

        if (
          downloadKind !== "DownloadOngoing" &&
-          downloadKind !== "DownloadPending" &&
-          downloadKind !== "DownloadCompleted"
+          downloadKind !== "DownloadPending"
        )
          continue;
+        if (!downloadPayload) continue;

-        const nodeName =
-          data?.nodes?.[nodeId]?.friendly_name ?? nodeId.slice(0, 8);
+        const downloadModelId = extractModelIdFromDownload(downloadPayload);

-        if (downloadKind === "DownloadCompleted") {
-          perNodeMap.set(nodeId, {
-            nodeId,
-            nodeName,
-            status: "completed",
-            percentage: 100,
-            progress: null,
-          });
-          continue;
+        // Match if the model ID contains or equals the requested model
+        // (handles cases like "mlx-community/Meta-Llama..." matching)
+        if (
+          !downloadModelId ||
+          !downloadModelId.includes(modelId.split("/").pop() || modelId)
+        ) {
+          // Try exact match or partial match
+          if (downloadModelId !== modelId) continue;
        }

+        // For DownloadPending with partial bytes (paused/resumed downloads),
+        // synthesize a progress object from the top-level downloaded/total fields
+        let progress: DownloadProgress | null;
        if (downloadKind === "DownloadPending") {
          const pendingDownloaded = getBytes(
            downloadPayload.downloaded ??
@@ -1637,67 +1610,44 @@
              downloadPayload.totalBytes,
          );
          if (pendingDownloaded <= 0 && pendingTotal <= 0) continue;
-          const pct =
-            pendingTotal > 0 ? (pendingDownloaded / pendingTotal) * 100 : 0;
-          perNodeMap.set(nodeId, {
-            nodeId,
-            nodeName,
-            status: pendingDownloaded > 0 ? "partial" : "pending",
-            percentage: pct,
-            progress: null,
-          });
-          continue;
+          isDownloading = true;
+          progress = {
+            totalBytes: pendingTotal,
+            downloadedBytes: pendingDownloaded,
+            speed: 0,
+            etaMs: 0,
+            percentage:
+              pendingTotal > 0 ? (pendingDownloaded / pendingTotal) * 100 : 0,
+            completedFiles: 0,
+            totalFiles: 0,
+            files: [],
+          };
+        } else {
+          isDownloading = true;
+          progress = parseDownloadProgress(downloadPayload);
        }

-        // DownloadOngoing
-        const progress = parseDownloadProgress(downloadPayload);
-        if (
-          !progress ||
-          (progress.downloadedBytes <= 0 && progress.totalBytes <= 0)
-        )
-          continue;
+        if (progress) {
+          // Sum all values across nodes - each node downloads independently
+          totalBytes += progress.totalBytes;
+          downloadedBytes += progress.downloadedBytes;
+          totalSpeed += progress.speed;
+          completedFiles += progress.completedFiles;
+          totalFiles += progress.totalFiles;
+          allFiles.push(...progress.files);

-        perNodeMap.set(nodeId, {
-          nodeId,
-          nodeName,
-          status: "downloading",
-          percentage: progress.percentage,
-          progress,
-        });
-      }
-    }
-
-    // Aggregate from deduplicated per-node entries
-    const perNode = Array.from(perNodeMap.values());
-    let totalBytes = 0;
-    let downloadedBytes = 0;
-    let totalSpeed = 0;
-    let completedFiles = 0;
-    let totalFiles = 0;
-    let isDownloading = false;
-    const allFiles: DownloadProgress["files"] = [];
-
-    for (const node of perNode) {
-      if (node.status === "downloading" && node.progress) {
-        isDownloading = true;
-        totalBytes += node.progress.totalBytes;
-        downloadedBytes += node.progress.downloadedBytes;
-        totalSpeed += node.progress.speed;
-        completedFiles += node.progress.completedFiles;
-        totalFiles += node.progress.totalFiles;
-        allFiles.push(...node.progress.files);
+          const nodeName =
+            data?.nodes?.[nodeId]?.friendly_name ?? nodeId.slice(0, 8);
+          perNode.push({ nodeId, nodeName, progress });
+        }
      }
    }

    if (!isDownloading) {
-      return {
-        isDownloading: false,
-        progress: null,
-        perNode,
-        failedError: null,
-      };
+      return { isDownloading: false, progress: null, perNode: [] };
    }

+    // ETA = total remaining bytes / total speed across all nodes
    const remainingBytes = totalBytes - downloadedBytes;
    const etaMs = totalSpeed > 0 ? (remainingBytes / totalSpeed) * 1000 : 0;

@@ -1714,21 +1664,9 @@
        files: allFiles,
      },
      perNode,
-      failedError: null,
    };
  }

-  function getModelDownloadStatus(
-    modelId: string,
-    nodeIds?: string[],
-  ): {
-    isDownloading: boolean;
-    progress: DownloadProgress | null;
-    perNode: NodeDownloadStatus[];
-  } {
-    return collectDownloadStatus(modelId, nodeIds);
-  }
-
  // Helper to get download status for an instance
  function getInstanceDownloadStatus(
    instanceId: string,
@@ -1739,9 +1677,26 @@
    errorMessage: string | null;
    progress: DownloadProgress | null;
    statusText: string;
-    perNode: NodeDownloadStatus[];
+    perNode: Array<{
+      nodeId: string;
+      nodeName: string;
+      progress: DownloadProgress;
+    }>;
  } {
-    // Unwrap the instance to get shard assignments
+    if (!downloadsData || Object.keys(downloadsData).length === 0) {
+      // No download data yet — defer to runner status instead of assuming RUNNING
+      const statusInfo = deriveInstanceStatus(instanceWrapped);
+      return {
+        isDownloading: false,
+        isFailed: false,
+        errorMessage: null,
+        progress: null,
+        statusText: statusInfo.statusText,
+        perNode: [],
+      };
+    }
+
+    // Unwrap the instance
    const [instanceTag, instance] = getTagged(instanceWrapped);
    if (!instance || typeof instance !== "object") {
      return {
@@ -1761,45 +1716,132 @@
        modelId?: string;
      };
    };
-    const instanceModelId = inst.shardAssignments?.modelId;
-
-    if (!instanceModelId) {
-      const statusInfo = deriveInstanceStatus(instanceWrapped);
-      return {
-        isDownloading: false,
-        isFailed: statusInfo.statusText === "FAILED",
-        errorMessage: null,
-        progress: null,
-        statusText: statusInfo.statusText,
-        perNode: [],
-      };
-    }
-
-    // Get node IDs assigned to this instance
    const nodeToRunner = inst.shardAssignments?.nodeToRunner || {};
    const runnerToShard = inst.shardAssignments?.runnerToShard || {};
+    const instanceModelId = inst.shardAssignments?.modelId;
+
+    // Build reverse mapping: runnerId -> nodeId
    const runnerToNode: Record<string, string> = {};
    for (const [nodeId, runnerId] of Object.entries(nodeToRunner)) {
      runnerToNode[runnerId] = nodeId;
    }
-    const instanceNodeIds = Object.keys(runnerToShard)
-      .map((runnerId) => runnerToNode[runnerId])
-      .filter(Boolean);

-    const result = collectDownloadStatus(instanceModelId, instanceNodeIds);
+    let totalBytes = 0;
+    let downloadedBytes = 0;
+    let totalSpeed = 0;
+    let completedFiles = 0;
+    let totalFiles = 0;
+    let isDownloading = false;
+    const allFiles: DownloadProgress["files"] = [];
+    const perNode: Array<{
+      nodeId: string;
+      nodeName: string;
+      progress: DownloadProgress;
+    }> = [];

-    if (result.failedError) {
-      return {
-        isDownloading: false,
-        isFailed: true,
-        errorMessage: result.failedError,
-        progress: null,
-        statusText: "FAILED",
-        perNode: [],
-      };
+    // Check downloads for nodes that are part of this instance
+    for (const runnerId of Object.keys(runnerToShard)) {
+      const nodeId = runnerToNode[runnerId];
+      if (!nodeId) continue;
+
+      const nodeDownloads = downloadsData[nodeId];
+      if (!Array.isArray(nodeDownloads)) continue;
+
+      for (const downloadWrapped of nodeDownloads) {
+        if (!downloadWrapped || typeof downloadWrapped !== "object") continue;
+
+        const keys = Object.keys(downloadWrapped as Record<string, unknown>);
+        if (keys.length !== 1) continue;
+
+        const downloadKind = keys[0];
+        const downloadPayload = (downloadWrapped as Record<string, unknown>)[
+          downloadKind
+        ] as Record<string, unknown>;
+
+        // Handle DownloadFailed - return immediately with error info
+        if (downloadKind === "DownloadFailed") {
+          const downloadModelId = extractModelIdFromDownload(downloadPayload);
+          if (
+            instanceModelId &&
+            downloadModelId &&
+            downloadModelId === instanceModelId
+          ) {
+            return {
+              isDownloading: false,
+              isFailed: true,
+              errorMessage:
+                (downloadPayload.errorMessage as string) || "Download failed",
+              progress: null,
+              statusText: "FAILED",
+              perNode: [],
+            };
+          }
+        }
+
+        if (
+          downloadKind !== "DownloadOngoing" &&
+          downloadKind !== "DownloadPending"
+        )
+          continue;
+        if (!downloadPayload) continue;
+
+        // Check if this download is for this instance's model
+        const downloadModelId = extractModelIdFromDownload(downloadPayload);
+        if (
+          instanceModelId &&
+          downloadModelId &&
+          downloadModelId === instanceModelId
+        ) {
+          // For DownloadPending with partial bytes, synthesize progress
+          let progress: DownloadProgress | null;
+          if (downloadKind === "DownloadPending") {
+            const pendingDownloaded = getBytes(
+              downloadPayload.downloaded ??
+                downloadPayload.downloaded_bytes ??
+                downloadPayload.downloadedBytes,
+            );
+            const pendingTotal = getBytes(
+              downloadPayload.total ??
+                downloadPayload.total_bytes ??
+                downloadPayload.totalBytes,
+            );
+            if (pendingDownloaded <= 0 && pendingTotal <= 0) continue;
+            isDownloading = true;
+            progress = {
+              totalBytes: pendingTotal,
+              downloadedBytes: pendingDownloaded,
+              speed: 0,
+              etaMs: 0,
+              percentage:
+                pendingTotal > 0 ? (pendingDownloaded / pendingTotal) * 100 : 0,
+              completedFiles: 0,
+              totalFiles: 0,
+              files: [],
+            };
+          } else {
+            isDownloading = true;
+            progress = parseDownloadProgress(downloadPayload);
+          }
+
+          if (progress) {
+            // Sum all values across nodes - each node downloads independently
+            totalBytes += progress.totalBytes;
+            downloadedBytes += progress.downloadedBytes;
+            totalSpeed += progress.speed;
+            completedFiles += progress.completedFiles;
+            totalFiles += progress.totalFiles;
+            allFiles.push(...progress.files);
+
+            const nodeName =
+              data?.nodes?.[nodeId]?.friendly_name ?? nodeId.slice(0, 8);
+            perNode.push({ nodeId, nodeName, progress });
+          }
+        }
+      }
    }

-    if (!result.isDownloading) {
+    if (!isDownloading) {
+      // Check runner status for other states
      const statusInfo = deriveInstanceStatus(instanceWrapped);
      return {
        isDownloading: false,
@@ -1807,17 +1849,30 @@
        errorMessage: null,
        progress: null,
        statusText: statusInfo.statusText,
-        perNode: result.perNode,
+        perNode: [],
      };
    }

+    // ETA = total remaining bytes / total speed across all nodes
+    const remainingBytes = totalBytes - downloadedBytes;
+    const etaMs = totalSpeed > 0 ? (remainingBytes / totalSpeed) * 1000 : 0;
+
    return {
      isDownloading: true,
      isFailed: false,
      errorMessage: null,
-      progress: result.progress,
+      progress: {
+        totalBytes,
+        downloadedBytes,
+        speed: totalSpeed,
+        etaMs,
+        percentage: totalBytes > 0 ? (downloadedBytes / totalBytes) * 100 : 0,
+        completedFiles,
+        totalFiles,
+        files: allFiles,
+      },
      statusText: "DOWNLOADING",
-      perNode: result.perNode,
+      perNode,
    };
  }

@@ -4575,7 +4630,7 @@
                type="button"
                onclick={() => {
                  completeOnboarding();
-                  sendMessage(chip, undefined, thinkingEnabled());
+                  sendMessage(chip);
                }}
                class="px-4 py-2 rounded-full border border-white/10 bg-white/5 text-sm text-white/60 hover:bg-white/10 hover:text-white/80 hover:border-white/20 transition-all duration-200 cursor-pointer"
              >
@@ -5314,10 +5369,10 @@
                            <div
                              class="mt-2 space-y-2 max-h-48 overflow-y-auto pr-1"
                            >
-                              {#each downloadInfo.perNode.filter((n) => n.status === "downloading" && n.progress) as nodeProg}
+                              {#each downloadInfo.perNode as nodeProg}
                                {@const nodePercent = Math.min(
                                  100,
-                                  Math.max(0, nodeProg.percentage),
+                                  Math.max(0, nodeProg.progress.percentage),
                                )}
                                {@const isExpanded =
                                  instanceDownloadExpandedNodes.has(
@@ -5373,17 +5428,15 @@
                                    >
                                      <span
                                        >{formatBytes(
-                                          nodeProg.progress?.downloadedBytes ??
-                                            0,
+                                          nodeProg.progress.downloadedBytes,
                                        )} / {formatBytes(
-                                          nodeProg.progress?.totalBytes ?? 0,
+                                          nodeProg.progress.totalBytes,
                                        )}</span
                                      >
                                      <span
-                                        >{formatSpeed(
-                                          nodeProg.progress?.speed ?? 0,
-                                        )} • ETA {formatEta(
-                                          nodeProg.progress?.etaMs ?? 0,
+                                        >{formatSpeed(nodeProg.progress.speed)} •
+                                        ETA {formatEta(
+                                          nodeProg.progress.etaMs,
                                        )}</span
                                      >
                                    </div>
@@ -5391,14 +5444,14 @@

                                  {#if isExpanded}
                                    <div class="mt-2 space-y-1.5">
-                                      {#if nodeProg.progress?.files ?? [].length === 0}
+                                      {#if nodeProg.progress.files.length === 0}
                                        <div
                                          class="text-[11px] font-mono text-exo-light-gray/70"
                                        >
                                          No file details reported.
                                        </div>
                                      {:else}
-                                        {#each nodeProg.progress?.files ?? [] as f}
+                                        {#each nodeProg.progress.files as f}
                                          {@const filePercent = Math.min(
                                            100,
                                            Math.max(0, f.percentage ?? 0),
@@ -5874,15 +5927,12 @@
                )}
                {@const allPreviews = filteredPreviews()}
                {#if selectedModel && allPreviews.length > 0}
+                  {@const downloadStatus = getModelDownloadStatus(
+                    selectedModel.id,
+                  )}
                  {@const tags = modelTags()[selectedModel.id] || []}
                  <div class="space-y-3">
                    {#each allPreviews as apiPreview, i}
-                      {@const downloadStatus = getModelDownloadStatus(
-                        selectedModel.id,
-                        apiPreview.memory_delta_by_node
-                          ? Object.keys(apiPreview.memory_delta_by_node)
-                          : undefined,
-                      )}
                      <div
                        role="group"
                        onmouseenter={() => {
@@ -6070,7 +6120,7 @@
                          onclick={() => {
                            chatLaunchState = "idle";
                            selectedChatCategory = null;
-                            sendMessage(prompt, undefined, thinkingEnabled());
+                            sendMessage(prompt);
                          }}
                          class="text-left px-3 py-2.5 text-xs text-exo-light-gray hover:text-white font-mono rounded-lg border border-exo-medium-gray/30 hover:border-exo-yellow/30 bg-exo-dark-gray/30 hover:bg-exo-dark-gray/60 transition-all duration-200 cursor-pointer"
                        >
@@ -6453,10 +6503,10 @@
                              <div
                                class="mt-2 space-y-2 max-h-48 overflow-y-auto pr-1"
                              >
-                                {#each downloadInfo.perNode.filter((n) => n.status === "downloading" && n.progress) as nodeProg}
+                                {#each downloadInfo.perNode as nodeProg}
                                  {@const nodePercent = Math.min(
                                    100,
-                                    Math.max(0, nodeProg.percentage),
+                                    Math.max(0, nodeProg.progress.percentage),
                                  )}
                                  {@const isExpanded =
                                    instanceDownloadExpandedNodes.has(
@@ -6515,17 +6565,16 @@
                                      >
                                        <span
                                          >{formatBytes(
-                                            nodeProg.progress
-                                              ?.downloadedBytes ?? 0,
+                                            nodeProg.progress.downloadedBytes,
                                          )} / {formatBytes(
-                                            nodeProg.progress?.totalBytes ?? 0,
+                                            nodeProg.progress.totalBytes,
                                          )}</span
                                        >
                                        <span
                                          >{formatSpeed(
-                                            nodeProg.progress?.speed ?? 0,
+                                            nodeProg.progress.speed,
                                          )} • ETA {formatEta(
-                                            nodeProg.progress?.etaMs ?? 0,
+                                            nodeProg.progress.etaMs,
                                          )}</span
                                        >
                                      </div>
@@ -6533,14 +6582,14 @@

                                    {#if isExpanded}
                                      <div class="mt-2 space-y-1.5">
-                                        {#if nodeProg.progress?.files ?? [].length === 0}
+                                        {#if nodeProg.progress.files.length === 0}
                                          <div
                                            class="text-[11px] font-mono text-exo-light-gray/70"
                                          >
                                            No file details reported.
                                          </div>
                                        {:else}
-                                          {#each nodeProg.progress?.files ?? [] as f}
+                                          {#each nodeProg.progress.files as f}
                                            {@const filePercent = Math.min(
                                              100,
                                              Math.max(0, f.percentage ?? 0),
@@ -72,7 +72,7 @@
      ];

      perSystem =
-        { config, self', pkgs, lib, system, ... }:
+        { config, self', inputs', pkgs, lib, system, ... }:
        let
          # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
          pkgsSwift = import inputs.nixpkgs-swift { inherit system; };
@@ -84,17 +84,6 @@
            config.allowUnfreePredicate = pkg: (pkg.pname or "") == "metal-toolchain";
            overlays = [
              (import ./nix/apple-sdk-overlay.nix)
-              (final: prev: {
-                macmon = prev.macmon.overrideAttrs (_: {
-                  version = "git";
-                  src = final.fetchFromGitHub {
-                    owner = "swiftraccoon";
-                    repo = "macmon";
-                    rev = "9154d234f763fbeffdcb4135d0bbbaf80609699b";
-                    hash = "sha256-CwhilKNbs5XL9/tF5DMwyPBlE/hpmjGNTuxQ36sM50M=";
-                  };
-                });
-              })
            ];
          };
          treefmt = {
@@ -1,8 +1,5 @@
 export NIX_CONFIG := "extra-experimental-features = nix-command flakes"

-default: lint fmt
-all: lint fmt check
-
 fmt:
    treefmt || nix fmt

@@ -34,10 +31,6 @@ build-dashboard:
 package:
    uv run pyinstaller packaging/pyinstaller/exo.spec

-build-app: package
-    xcodebuild build -project app/EXO/EXO.xcodeproj -scheme EXO -configuration Debug -derivedDataPath app/EXO/build
-    @echo "\nBuild complete. Run with:\n  open {{justfile_directory()}}/app/EXO/build/Build/Products/Debug/EXO.app"
-
 clean:
    rm -rf **/__pycache__
    rm -rf target/
@@ -71,9 +71,7 @@ MACMON_PATH = shutil.which("macmon")
 if MACMON_PATH is None:
    raise SystemExit(
        "macmon binary not found in PATH. "
-        "Install the pinned fork used by exo via: "
-        "cargo install --git https://github.com/swiftraccoon/macmon "
-        "--rev 9154d234f763fbeffdcb4135d0bbbaf80609699b macmon --force"
+        "Install it via: brew install macmon"
    )

 BINARIES: list[tuple[str, str]] = [
@@ -122,3 +120,4 @@ coll = COLLECT(
    upx_exclude=[],
    name="exo",
 )
+
@@ -25,7 +25,7 @@ dependencies = [
  "openai-harmony>=0.0.8",
  "httpx>=0.28.1",
  "tomlkit>=0.14.0",
-  "mflux==0.17.2",
+  "mflux==0.16.9",
  "python-multipart>=0.0.21",
  "msgspec>=0.19.0",
  "zstandard>=0.23.0",
@@ -61,7 +61,7 @@ members = ["rust/exo_pyo3_bindings", "bench"]
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
-mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/fix-deepseek-v32-indexer" }
+mlx-lm = { git = "https://github.com/ml-explore/mlx-lm", branch = "main" }
 # Uncomment to use local mlx/mlx-lm development versions:
 # mlx = { path = "/Users/Shared/mlx", editable=true }
 # mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
@@ -1,13 +0,0 @@
-model_id = "mlx-community/DeepSeek-V3.2-4bit"
-n_layers = 61
-hidden_size = 7168
-num_key_value_heads = 128
-supports_tensor = true
-tasks = ["TextGeneration"]
-family = "deepseek"
-quantization = "4bit"
-base_model = "DeepSeek V3.2"
-capabilities = ["text", "thinking", "thinking_toggle"]
-
-[storage_size]
-in_bytes = 378086226621
@@ -1,13 +0,0 @@
-model_id = "mlx-community/DeepSeek-V3.2-8bit"
-n_layers = 61
-hidden_size = 7168
-num_key_value_heads = 128
-supports_tensor = true
-tasks = ["TextGeneration"]
-family = "deepseek"
-quantization = "8bit"
-base_model = "DeepSeek V3.2"
-capabilities = ["text", "thinking", "thinking_toggle"]
-
-[storage_size]
-in_bytes = 755957120916
@@ -42,7 +42,7 @@ class MessageTooLargeError(builtins.Exception):

@typing.final
 class NetworkingHandle:
-    def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
+    def __new__(cls, identity: Keypair) -> NetworkingHandle: ...
    async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
        r"""
        Subscribe to a `GossipSub` topic.
@@ -180,12 +180,7 @@ impl PyNetworkingHandle {
    // ---- Lifecycle management methods ----

    #[new]
-    #[pyo3(signature = (identity, bootstrap_peers, listen_port))]
-    fn py_new(
-        identity: Bound<'_, PyKeypair>,
-        bootstrap_peers: Vec<String>,
-        listen_port: u16,
-    ) -> PyResult<Self> {
+    fn py_new(identity: Bound<'_, PyKeypair>) -> PyResult<Self> {
        // create communication channels
        let (to_swarm, from_client) = mpsc::channel(MPSC_CHANNEL_SIZE);

@@ -194,9 +189,7 @@ impl PyNetworkingHandle {

        // create networking swarm (within tokio context!! or it crashes)
        let _guard = pyo3_async_runtimes::tokio::get_runtime().enter();
-        let swarm = create_swarm(identity, from_client, bootstrap_peers, listen_port)
-            .pyerr()?
-            .into_stream();
+        let swarm = create_swarm(identity, from_client).pyerr()?.into_stream();

        Ok(Self {
            swarm: Arc::new(Mutex::new(swarm)),
@@ -16,14 +16,9 @@ async fn main() {
    let (to_swarm, from_client) = mpsc::channel(20);

    // Configure swarm
-    let mut swarm = swarm::create_swarm(
-        identity::Keypair::generate_ed25519(),
-        from_client,
-        vec![],
-        0,
-    )
-    .expect("Swarm creation failed")
-    .into_stream();
+    let mut swarm = swarm::create_swarm(identity::Keypair::generate_ed25519(), from_client)
+        .expect("Swarm creation failed")
+        .into_stream();

    // Create a Gossipsub topic & subscribe
    let (tx, rx) = oneshot::channel();
@@ -104,7 +104,6 @@ pub struct Behaviour {
    // state-tracking for managed behaviors & mDNS-discovered peers
    managed: managed::Behaviour,
    mdns_discovered: HashMap<PeerId, BTreeSet<Multiaddr>>,
-    bootstrap_peers: Vec<Multiaddr>,

    retry_delay: Delay, // retry interval

@@ -113,11 +112,10 @@ pub struct Behaviour {
 }

 impl Behaviour {
-    pub fn new(keypair: &identity::Keypair, bootstrap_peers: Vec<Multiaddr>) -> io::Result<Self> {
+    pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
        Ok(Self {
            managed: managed::Behaviour::new(keypair)?,
            mdns_discovered: HashMap::new(),
-            bootstrap_peers,
            retry_delay: Delay::new(RETRY_CONNECT_INTERVAL),
            pending_events: WakerDeque::new(),
        })
@@ -370,12 +368,6 @@ impl NetworkBehaviour for Behaviour {
                    self.dial(p, ma)
                }
            }
-            // dial bootstrap peers (for environments where mDNS is unavailable)
-            for addr in &self.bootstrap_peers {
-                self.pending_events.push_back(ToSwarm::Dial {
-                    opts: DialOpts::unknown_peer_id().address(addr.clone()).build(),
-                })
-            }
            self.retry_delay.reset(RETRY_CONNECT_INTERVAL) // reset timeout
        }

@@ -142,29 +142,19 @@ fn filter_swarm_event(event: SwarmEvent<BehaviourEvent>) -> Option<FromSwarm> {
    }
 }

-/// Create and configure a swarm.
-///
-/// - `listen_port`: TCP port to listen on. `0` lets the OS assign one.
-/// - `bootstrap_peers`: multiaddrs to dial for environments without mDNS.
+/// Create and configure a swarm which listens to all ports on OS
 pub fn create_swarm(
    keypair: identity::Keypair,
    from_client: mpsc::Receiver<ToSwarm>,
-    bootstrap_peers: Vec<String>,
-    listen_port: u16,
 ) -> alias::AnyResult<Swarm> {
-    let parsed_bootstrap_peers: Vec<libp2p::Multiaddr> = bootstrap_peers
-        .iter()
-        .filter(|s| !s.is_empty())
-        .filter_map(|s| s.parse().ok())
-        .collect();
-
    let mut swarm = SwarmBuilder::with_existing_identity(keypair)
        .with_tokio()
        .with_other_transport(tcp_transport)?
-        .with_behaviour(|keypair| Behaviour::new(keypair, parsed_bootstrap_peers))?
+        .with_behaviour(Behaviour::new)?
        .build();

-    swarm.listen_on(format!("/ip4/0.0.0.0/tcp/{listen_port}").parse()?)?;
+    // Listen on all interfaces and whatever port the OS assigns
+    swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
    Ok(Swarm { swarm, from_client })
 }

@@ -256,12 +246,9 @@ mod behaviour {
    }

    impl Behaviour {
-        pub fn new(
-            keypair: &identity::Keypair,
-            bootstrap_peers: Vec<libp2p::Multiaddr>,
-        ) -> alias::AnyResult<Self> {
+        pub fn new(keypair: &identity::Keypair) -> alias::AnyResult<Self> {
            Ok(Self {
-                discovery: discovery::Behaviour::new(keypair, bootstrap_peers)?,
+                discovery: discovery::Behaviour::new(keypair)?,
                gossipsub: gossipsub_behaviour(keypair),
            })
        }
@@ -1,107 +0,0 @@
-use futures_lite::StreamExt;
-use networking::swarm::{FromSwarm, create_swarm};
-use std::time::Duration;
-use tokio::sync::mpsc;
-use tokio::time::timeout;
-
-/// Helper: find a free TCP port.
-fn free_port() -> u16 {
-    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
-    listener.local_addr().unwrap().port()
-}
-
-/// Two nodes connect via bootstrap peers — no mDNS needed.
-///
-/// Node A listens on a fixed port. Node B bootstraps to A's address.
-/// We verify that B emits `FromSwarm::Discovered` for A's peer ID.
-#[tokio::test]
-async fn two_nodes_connect_via_bootstrap_peers() {
-    let port_a = free_port();
-
-    // Node A: listens on a known port, no bootstrap peers
-    let keypair_a = libp2p::identity::Keypair::generate_ed25519();
-    let peer_id_a = keypair_a.public().to_peer_id();
-    let (_tx_a, rx_a) = mpsc::channel(16);
-    let swarm_a = create_swarm(keypair_a, rx_a, vec![], port_a).expect("create swarm A");
-    let mut stream_a = swarm_a.into_stream();
-
-    // Node B: bootstraps to A's address
-    let keypair_b = libp2p::identity::Keypair::generate_ed25519();
-    let (_tx_b, rx_b) = mpsc::channel(16);
-    let swarm_b = create_swarm(
-        keypair_b,
-        rx_b,
-        vec![format!("/ip4/127.0.0.1/tcp/{port_a}")],
-        0,
-    )
-    .expect("create swarm B");
-    let mut stream_b = swarm_b.into_stream();
-
-    // Wait for B to discover A (connection established)
-    let connected = timeout(Duration::from_secs(10), async {
-        loop {
-            tokio::select! {
-                Some(event) = stream_a.next() => {
-                    // A will also see B connect, but we check from B's perspective
-                    let _ = event;
-                }
-                Some(event) = stream_b.next() => {
-                    if let FromSwarm::Discovered { peer_id } = event {
-                        if peer_id == peer_id_a {
-                            return true;
-                        }
-                    }
-                }
-            }
-        }
-    })
-    .await;
-
-    assert!(
-        connected.is_ok() && connected.unwrap(),
-        "Node B should discover Node A via bootstrap peer"
-    );
-}
-
-/// Empty bootstrap peers should work (backward compatible).
-#[tokio::test]
-async fn create_swarm_with_empty_bootstrap_peers() {
-    let keypair = libp2p::identity::Keypair::generate_ed25519();
-    let (_tx, rx) = mpsc::channel(16);
-    let swarm = create_swarm(keypair, rx, vec![], 0);
-    assert!(
-        swarm.is_ok(),
-        "create_swarm with no bootstrap peers should succeed"
-    );
-}
-
-/// Invalid multiaddr strings are silently filtered out.
-#[tokio::test]
-async fn create_swarm_ignores_invalid_bootstrap_addrs() {
-    let keypair = libp2p::identity::Keypair::generate_ed25519();
-    let (_tx, rx) = mpsc::channel(16);
-    let swarm = create_swarm(
-        keypair,
-        rx,
-        vec![
-            "not-a-valid-multiaddr".to_string(),
-            "".to_string(),
-            "/ip4/10.0.0.1/tcp/30000".to_string(), // valid
-        ],
-        0,
-    );
-    assert!(
-        swarm.is_ok(),
-        "create_swarm should succeed even with invalid bootstrap addrs"
-    );
-}
-
-/// Fixed listen port works correctly.
-#[tokio::test]
-async fn create_swarm_with_fixed_port() {
-    let port = free_port();
-    let keypair = libp2p::identity::Keypair::generate_ed25519();
-    let (_tx, rx) = mpsc::channel(16);
-    let swarm = create_swarm(keypair, rx, vec![], port);
-    assert!(swarm.is_ok(), "create_swarm with fixed port should succeed");
-}
@@ -1,57 +0,0 @@
-from .api import AddCustomModelParams as AddCustomModelParams
-from .api import AdvancedImageParams as AdvancedImageParams
-from .api import BenchChatCompletionRequest as BenchChatCompletionRequest
-from .api import BenchChatCompletionResponse as BenchChatCompletionResponse
-from .api import BenchImageGenerationResponse as BenchImageGenerationResponse
-from .api import BenchImageGenerationTaskParams as BenchImageGenerationTaskParams
-from .api import CancelCommandResponse as CancelCommandResponse
-from .api import ChatCompletionChoice as ChatCompletionChoice
-from .api import ChatCompletionMessage as ChatCompletionMessage
-from .api import ChatCompletionMessageText as ChatCompletionMessageText
-from .api import ChatCompletionRequest as ChatCompletionRequest
-from .api import ChatCompletionResponse as ChatCompletionResponse
-from .api import CompletionTokensDetails as CompletionTokensDetails
-from .api import CreateInstanceParams as CreateInstanceParams
-from .api import CreateInstanceResponse as CreateInstanceResponse
-from .api import DeleteDownloadResponse as DeleteDownloadResponse
-from .api import DeleteInstanceResponse as DeleteInstanceResponse
-from .api import DeleteTracesRequest as DeleteTracesRequest
-from .api import DeleteTracesResponse as DeleteTracesResponse
-from .api import ErrorInfo as ErrorInfo
-from .api import ErrorResponse as ErrorResponse
-from .api import FinishReason as FinishReason
-from .api import GenerationStats as GenerationStats
-from .api import HuggingFaceSearchResult as HuggingFaceSearchResult
-from .api import ImageData as ImageData
-from .api import ImageEditsTaskParams as ImageEditsTaskParams
-from .api import ImageGenerationResponse as ImageGenerationResponse
-from .api import ImageGenerationStats as ImageGenerationStats
-from .api import ImageGenerationTaskParams as ImageGenerationTaskParams
-from .api import ImageListItem as ImageListItem
-from .api import ImageListResponse as ImageListResponse
-from .api import ImageSize as ImageSize
-from .api import Logprobs as Logprobs
-from .api import LogprobsContentItem as LogprobsContentItem
-from .api import ModelList as ModelList
-from .api import ModelListModel as ModelListModel
-from .api import NodePowerStats as NodePowerStats
-from .api import PlaceInstanceParams as PlaceInstanceParams
-from .api import PlacementPreview as PlacementPreview
-from .api import PlacementPreviewResponse as PlacementPreviewResponse
-from .api import PowerUsage as PowerUsage
-from .api import PromptTokensDetails as PromptTokensDetails
-from .api import StartDownloadParams as StartDownloadParams
-from .api import StartDownloadResponse as StartDownloadResponse
-from .api import StreamingChoiceResponse as StreamingChoiceResponse
-from .api import ToolCall as ToolCall
-from .api import ToolCallItem as ToolCallItem
-from .api import TopLogprobItem as TopLogprobItem
-from .api import TraceCategoryStats as TraceCategoryStats
-from .api import TraceEventResponse as TraceEventResponse
-from .api import TraceListItem as TraceListItem
-from .api import TraceListResponse as TraceListResponse
-from .api import TraceRankStats as TraceRankStats
-from .api import TraceResponse as TraceResponse
-from .api import TraceStatsResponse as TraceStatsResponse
-from .api import Usage as Usage
-from .api import normalize_image_size as normalize_image_size
@@ -1,21 +1,17 @@
-from __future__ import annotations
-
 from dataclasses import dataclass, field
-from pathlib import Path

 import anyio
-from anyio import current_time, to_thread
+from anyio import current_time
 from loguru import logger

 from exo.download.download_utils import (
    RepoDownloadProgress,
    delete_model,
-    is_read_only_model_dir,
    map_repo_download_progress_to_download_progress_data,
-    resolve_existing_model,
+    resolve_model_in_path,
 )
 from exo.download.shard_downloader import ShardDownloader
-from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_MODELS_READ_ONLY_DIRS
+from exo.shared.constants import EXO_MODELS_DIR, EXO_MODELS_PATH
 from exo.shared.models.model_cards import ModelId, get_model_cards
 from exo.shared.types.commands import (
    CancelDownload,
@@ -28,7 +24,6 @@ from exo.shared.types.events import (
    Event,
    NodeDownloadProgress,
 )
-from exo.shared.types.memory import Memory
 from exo.shared.types.worker.downloads import (
    DownloadCompleted,
    DownloadFailed,
@@ -54,7 +49,6 @@ class DownloadCoordinator:
    active_downloads: dict[ModelId, anyio.CancelScope] = field(default_factory=dict)

    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
-    _stopped: anyio.Event = field(init=False, default_factory=anyio.Event)

    # Per-model throttle for download progress events
    _last_progress_time: dict[ModelId, float] = field(default_factory=dict)
@@ -62,23 +56,8 @@ class DownloadCoordinator:
    def __post_init__(self) -> None:
        self.shard_downloader.on_progress(self._download_progress_callback)

-    @staticmethod
-    def _default_model_dir(model_id: ModelId) -> str:
-        return str(EXO_DEFAULT_MODELS_DIR / model_id.normalize())
-
-    def _completed_from_path(
-        self,
-        shard: ShardMetadata,
-        found: Path,
-        total: Memory,
-    ) -> DownloadCompleted:
-        return DownloadCompleted(
-            shard_metadata=shard,
-            node_id=self.node_id,
-            total=total,
-            model_directory=str(found),
-            read_only=is_read_only_model_dir(found),
-        )
+    def _model_dir(self, model_id: ModelId) -> str:
+        return str(EXO_MODELS_DIR / model_id.normalize())

    async def _download_progress_callback(
        self, callback_shard: ShardMetadata, progress: RepoDownloadProgress
@@ -87,18 +66,12 @@ class DownloadCoordinator:
        throttle_interval_secs = 1.0

        if progress.status == "complete":
-            found = await to_thread.run_sync(resolve_existing_model, model_id)
-            if found is not None:
-                completed = self._completed_from_path(
-                    callback_shard, found, progress.total
-                )
-            else:
-                completed = DownloadCompleted(
-                    shard_metadata=callback_shard,
-                    node_id=self.node_id,
-                    total=progress.total,
-                    model_directory=self._default_model_dir(model_id),
-                )
+            completed = DownloadCompleted(
+                shard_metadata=callback_shard,
+                node_id=self.node_id,
+                total=progress.total,
+                model_directory=self._model_dir(model_id),
+            )
            self.download_status[model_id] = completed
            await self.event_sender.send(
                NodeDownloadProgress(download_progress=completed)
@@ -115,7 +88,7 @@ class DownloadCoordinator:
                download_progress=map_repo_download_progress_to_download_progress_data(
                    progress
                ),
-                model_directory=self._default_model_dir(model_id),
+                model_directory=self._model_dir(model_id),
            )
            self.download_status[model_id] = ongoing
            await self.event_sender.send(
@@ -127,16 +100,12 @@ class DownloadCoordinator:
        logger.info(
            f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
        )
-        try:
-            async with self._tg as tg:
-                tg.start_soon(self._command_processor)
-                tg.start_soon(self._emit_existing_download_progress)
-        finally:
-            self._stopped.set()
+        async with self._tg as tg:
+            tg.start_soon(self._command_processor)
+            tg.start_soon(self._emit_existing_download_progress)

-    async def shutdown(self) -> None:
+    def shutdown(self) -> None:
        self._tg.cancel_tasks()
-        await self._stopped.wait()

    async def _command_processor(self) -> None:
        with self.download_command_receiver as commands:
@@ -161,7 +130,7 @@ class DownloadCoordinator:
            pending = DownloadPending(
                shard_metadata=current_status.shard_metadata,
                node_id=self.node_id,
-                model_directory=self._default_model_dir(model_id),
+                model_directory=self._model_dir(model_id),
            )
            self.download_status[model_id] = pending
            await self.event_sender.send(
@@ -180,12 +149,18 @@ class DownloadCoordinator:
                )
                return

-        # Check all model directories for pre-existing complete models
-        found_path = await to_thread.run_sync(resolve_existing_model, model_id)
+        # Check EXO_MODELS_PATH for pre-downloaded models
+        found_path = resolve_model_in_path(model_id)
        if found_path is not None:
-            logger.info(f"DownloadCoordinator: Model {model_id} found at {found_path}")
-            completed = self._completed_from_path(
-                shard, found_path, shard.model_card.storage_size
+            logger.info(
+                f"DownloadCoordinator: Model {model_id} found in EXO_MODELS_PATH at {found_path}"
+            )
+            completed = DownloadCompleted(
+                shard_metadata=shard,
+                node_id=self.node_id,
+                total=shard.model_card.storage_size,
+                model_directory=str(found_path),
+                read_only=True,
            )
            self.download_status[model_id] = completed
            await self.event_sender.send(
@@ -197,7 +172,7 @@ class DownloadCoordinator:
        progress = DownloadPending(
            shard_metadata=shard,
            node_id=self.node_id,
-            model_directory=self._default_model_dir(model_id),
+            model_directory=self._model_dir(model_id),
        )
        self.download_status[model_id] = progress
        await self.event_sender.send(NodeDownloadProgress(download_progress=progress))
@@ -208,18 +183,12 @@ class DownloadCoordinator:
        )

        if initial_progress.status == "complete":
-            found = await to_thread.run_sync(resolve_existing_model, model_id)
-            if found is not None:
-                completed = self._completed_from_path(
-                    shard, found, initial_progress.total
-                )
-            else:
-                completed = DownloadCompleted(
-                    shard_metadata=shard,
-                    node_id=self.node_id,
-                    total=initial_progress.total,
-                    model_directory=self._default_model_dir(model_id),
-                )
+            completed = DownloadCompleted(
+                shard_metadata=shard,
+                node_id=self.node_id,
+                total=initial_progress.total,
+                model_directory=self._model_dir(model_id),
+            )
            self.download_status[model_id] = completed
            await self.event_sender.send(
                NodeDownloadProgress(download_progress=completed)
@@ -234,7 +203,7 @@ class DownloadCoordinator:
                shard_metadata=shard,
                node_id=self.node_id,
                error_message=f"Model files not found locally in offline mode: {model_id}",
-                model_directory=self._default_model_dir(model_id),
+                model_directory=self._model_dir(model_id),
            )
            self.download_status[model_id] = failed
            await self.event_sender.send(NodeDownloadProgress(download_progress=failed))
@@ -255,7 +224,7 @@ class DownloadCoordinator:
            download_progress=map_repo_download_progress_to_download_progress_data(
                initial_progress
            ),
-            model_directory=self._default_model_dir(model_id),
+            model_directory=self._model_dir(model_id),
        )
        self.download_status[model_id] = status
        self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
@@ -270,7 +239,7 @@ class DownloadCoordinator:
                    shard_metadata=shard,
                    node_id=self.node_id,
                    error_message=str(e),
-                    model_directory=self._default_model_dir(model_id),
+                    model_directory=self._model_dir(model_id),
                )
                self.download_status[model_id] = failed
                await self.event_sender.send(
@@ -287,11 +256,13 @@ class DownloadCoordinator:
        self.active_downloads[model_id] = scope

    async def _delete_download(self, model_id: ModelId) -> None:
-        # Protect read-only models from deletion
+        # Protect read-only models (from EXO_MODELS_PATH) from deletion
        if model_id in self.download_status:
            current = self.download_status[model_id]
            if isinstance(current, DownloadCompleted) and current.read_only:
-                logger.warning(f"Refusing to delete read-only model {model_id}")
+                logger.warning(
+                    f"Refusing to delete read-only model {model_id} (from EXO_MODELS_PATH)"
+                )
                return

        # Cancel if active
@@ -314,7 +285,7 @@ class DownloadCoordinator:
            pending = DownloadPending(
                shard_metadata=current_status.shard_metadata,
                node_id=self.node_id,
-                model_directory=self._default_model_dir(model_id),
+                model_directory=self._model_dir(model_id),
            )
            await self.event_sender.send(
                NodeDownloadProgress(download_progress=pending)
@@ -338,26 +309,22 @@ class DownloadCoordinator:
                        continue

                    if progress.status == "complete":
-                        found = await to_thread.run_sync(
-                            resolve_existing_model, model_id
+                        status: DownloadProgress = DownloadCompleted(
+                            node_id=self.node_id,
+                            shard_metadata=progress.shard,
+                            total=progress.total,
+                            model_directory=self._model_dir(
+                                progress.shard.model_card.model_id
+                            ),
                        )
-                        if found is not None:
-                            status: DownloadProgress = self._completed_from_path(
-                                progress.shard, found, progress.total
-                            )
-                        else:
-                            status = DownloadCompleted(
-                                node_id=self.node_id,
-                                shard_metadata=progress.shard,
-                                total=progress.total,
-                                model_directory=self._default_model_dir(model_id),
-                            )
                    elif progress.status in ["in_progress", "not_started"]:
                        if progress.downloaded_this_session.in_bytes == 0:
                            status = DownloadPending(
                                node_id=self.node_id,
                                shard_metadata=progress.shard,
-                                model_directory=self._default_model_dir(model_id),
+                                model_directory=self._model_dir(
+                                    progress.shard.model_card.model_id
+                                ),
                                downloaded=progress.downloaded,
                                total=progress.total,
                            )
@@ -368,7 +335,9 @@ class DownloadCoordinator:
                                download_progress=map_repo_download_progress_to_download_progress_data(
                                    progress
                                ),
-                                model_directory=self._default_model_dir(model_id),
+                                model_directory=self._model_dir(
+                                    progress.shard.model_card.model_id
+                                ),
                            )
                    else:
                        continue
@@ -377,8 +346,8 @@ class DownloadCoordinator:
                    await self.event_sender.send(
                        NodeDownloadProgress(download_progress=status)
                    )
-                # Scan read-only directories for pre-downloaded models
-                if EXO_MODELS_READ_ONLY_DIRS:
+                # Scan EXO_MODELS_PATH for pre-downloaded models
+                if EXO_MODELS_PATH is not None:
                    for card in await get_model_cards():
                        mid = card.model_id
                        if mid in self.active_downloads:
@@ -388,8 +357,8 @@ class DownloadCoordinator:
                            (DownloadCompleted, DownloadOngoing, DownloadFailed),
                        ):
                            continue
-                        found = await to_thread.run_sync(resolve_existing_model, mid)
-                        if found is not None and is_read_only_model_dir(found):
+                        found = resolve_model_in_path(mid)
+                        if found is not None:
                            path_shard = PipelineShardMetadata(
                                model_card=card,
                                device_rank=0,
@@ -398,10 +367,12 @@ class DownloadCoordinator:
                                end_layer=card.n_layers,
                                n_layers=card.n_layers,
                            )
-                            path_completed: DownloadProgress = (
-                                self._completed_from_path(
-                                    path_shard, found, card.storage_size
-                                )
+                            path_completed: DownloadProgress = DownloadCompleted(
+                                node_id=self.node_id,
+                                shard_metadata=path_shard,
+                                total=card.storage_size,
+                                model_directory=str(found),
+                                read_only=True,
                            )
                            self.download_status[mid] = path_completed
                            await self.event_sender.send(
@@ -30,11 +30,7 @@ from exo.download.huggingface_utils import (
    get_hf_endpoint,
    get_hf_token,
 )
-from exo.shared.constants import (
-    EXO_DEFAULT_MODELS_DIR,
-    EXO_MODELS_DIRS,
-    EXO_MODELS_READ_ONLY_DIRS,
-)
+from exo.shared.constants import EXO_MODELS_DIR, EXO_MODELS_PATH
 from exo.shared.models.model_cards import ModelTask
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
@@ -114,87 +110,50 @@ def map_repo_download_progress_to_download_progress_data(
    )


-class InsufficientDiskSpaceError(Exception):
-    """Raised when no writable model directory has enough free space."""
+def resolve_model_in_path(model_id: ModelId) -> Path | None:
+    """Search EXO_MODELS_PATH directories for a pre-existing model.

-
-def resolve_existing_model(model_id: ModelId) -> Path | None:
-    """Search all model directories for a complete, pre-existing model.
-
-    Checks read-only directories first, then writable directories.
-    A candidate is only returned if ``is_model_directory_complete`` confirms
-    all weight files are present.
+    Checks each directory for the normalized name (org--model).  A candidate
+    is only returned if ``is_model_directory_complete`` confirms all weight
+    files are present.
    """
+    if EXO_MODELS_PATH is None:
+        return None
    normalized = model_id.normalize()
-    for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
+    for search_dir in EXO_MODELS_PATH:
        candidate = search_dir / normalized
        if candidate.is_dir() and is_model_directory_complete(candidate):
            return candidate
    return None


-def is_read_only_model_dir(model_dir: Path) -> bool:
-    """Check if a model directory lives under a read-only models root."""
-    return any(model_dir.is_relative_to(d) for d in EXO_MODELS_READ_ONLY_DIRS)
-
-
 def build_model_path(model_id: ModelId) -> Path:
-    found = resolve_existing_model(model_id)
+    found = resolve_model_in_path(model_id)
    if found is not None:
        return found
-    return EXO_DEFAULT_MODELS_DIR / model_id.normalize()
+    return EXO_MODELS_DIR / model_id.normalize()


-def select_download_dir(required_bytes: int) -> Path:
-    """Pick the first writable model directory with enough free space.
-
-    Raises ``InsufficientDiskSpaceError`` if none have enough space.
-    """
-    for candidate_dir in EXO_MODELS_DIRS:
-        if not candidate_dir.exists():
-            continue
-        try:
-            usage = shutil.disk_usage(candidate_dir)
-            if usage.free >= required_bytes:
-                return candidate_dir
-        except OSError:
-            continue
-    raise InsufficientDiskSpaceError(
-        f"No writable model directory has {required_bytes / (1024**3):.1f} GiB free. "
-        f"Checked: {[str(d) for d in EXO_MODELS_DIRS]}"
-    )
+async def resolve_model_path_for_repo(model_id: ModelId) -> Path:
+    return (await ensure_models_dir()) / model_id.normalize()


-async def resolve_model_dir(model_id: ModelId) -> Path:
-    """Return the directory for a model's files, creating it if needed.
-
-    Checks all model directories for an existing complete model first,
-    then falls back to the default models directory.
-    """
-    target = await asyncio.to_thread(build_model_path, model_id)
-    await aios.makedirs(target, exist_ok=True)
-    return target
-
-
-async def ensure_cache_dir(model_id: ModelId) -> Path:
-    """Return the cache directory for a model's metadata, creating it if needed."""
-    target = EXO_DEFAULT_MODELS_DIR / "caches" / model_id.normalize()
-    await aios.makedirs(target, exist_ok=True)
-    return target
+async def ensure_models_dir() -> Path:
+    await aios.makedirs(EXO_MODELS_DIR, exist_ok=True)
+    return EXO_MODELS_DIR


 async def delete_model(model_id: ModelId) -> bool:
-    """Delete a model from writable directories. Skips read-only dirs."""
-    normalized = model_id.normalize()
-    deleted = False
-    for models_dir in EXO_MODELS_DIRS:
-        model_dir = models_dir / normalized
-        if await aios.path.exists(model_dir):
-            await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
-            deleted = True
+    models_dir = await ensure_models_dir()
+    model_dir = models_dir / model_id.normalize()
+    cache_dir = models_dir / "caches" / model_id.normalize()

-    # Clear cache from default dir
-    cache_dir = EXO_DEFAULT_MODELS_DIR / "caches" / normalized
+    deleted = False
+    if await aios.path.exists(model_dir):
+        await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
+        deleted = True
+
+    # Also clear cache
    if await aios.path.exists(cache_dir):
        await asyncio.to_thread(shutil.rmtree, cache_dir, ignore_errors=False)

@@ -202,10 +161,9 @@ async def delete_model(model_id: ModelId) -> bool:


 async def seed_models(seed_dir: str | Path):
-    """Move models from resources folder to the default models directory."""
+    """Move models from resources folder to EXO_MODELS_DIR."""
    source_dir = Path(seed_dir)
-    await aios.makedirs(EXO_DEFAULT_MODELS_DIR, exist_ok=True)
-    dest_dir = EXO_DEFAULT_MODELS_DIR
+    dest_dir = await ensure_models_dir()
    for path in source_dir.iterdir():
        if path.is_dir() and path.name.startswith("models--"):
            dest_path = dest_dir / path.name
@@ -295,16 +253,14 @@ async def _build_file_list_from_local_directory(
    a local directory must contain a *.safetensors.index.json and
    safetensors listed there.
    """
-    normalized = model_id.normalize()
-    for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
-        model_dir = search_dir / normalized
-        if await aios.path.exists(model_dir):
-            file_list = await asyncio.to_thread(
-                _scan_model_directory, model_dir, recursive
-            )
-            if file_list:
-                return file_list
-    return None
+    model_dir = (await ensure_models_dir()) / model_id.normalize()
+    if not await aios.path.exists(model_dir):
+        return None
+
+    file_list = await asyncio.to_thread(_scan_model_directory, model_dir, recursive)
+    if not file_list:
+        return None
+    return file_list


 _fetched_file_lists_this_session: set[str] = set()
@@ -317,7 +273,8 @@ async def fetch_file_list_with_cache(
    skip_internet: bool = False,
    on_connection_lost: Callable[[], None] = lambda: None,
 ) -> list[FileListEntry]:
-    target_dir = await ensure_cache_dir(model_id)
+    target_dir = (await ensure_models_dir()) / "caches" / model_id.normalize()
+    await aios.makedirs(target_dir, exist_ok=True)
    cache_file = target_dir / f"{model_id.normalize()}--{revision}--file_list.json"
    cache_key = f"{model_id.normalize()}--{revision}"

@@ -372,7 +329,7 @@ async def fetch_file_list_with_cache(
        )
        if local_file_list is not None:
            logger.warning(
-                f"Failed to fetch file list for {model_id} and no cache exists, using local file list"
+                f"Failed to fetch file list for {model_id} and no cache exists, "
            )
            return local_file_list
        raise FileNotFoundError(f"Failed to fetch file list for {model_id}: {e}") from e
@@ -701,7 +658,8 @@ def calculate_repo_progress(


 async def get_weight_map(model_id: ModelId, revision: str = "main") -> dict[str, str]:
-    target_dir = await resolve_model_dir(model_id)
+    target_dir = (await ensure_models_dir()) / model_id.normalize()
+    await aios.makedirs(target_dir, exist_ok=True)

    index_files_dir = snapshot_download(
        repo_id=model_id,
@@ -772,46 +730,30 @@ async def download_shard(
    if not skip_download:
        logger.debug(f"Downloading {shard.model_card.model_id=}")

-    model_id = shard.model_card.model_id
    revision = "main"
+    target_dir = await ensure_models_dir() / str(shard.model_card.model_id).replace(
+        "/", "--"
+    )
+    if not skip_download:
+        await aios.makedirs(target_dir, exist_ok=True)

    if not allow_patterns:
        allow_patterns = await resolve_allow_patterns(shard)

    if not skip_download:
-        logger.debug(f"Downloading {model_id=} with {allow_patterns=}")
+        logger.debug(f"Downloading {shard.model_card.model_id=} with {allow_patterns=}")

    all_start_time = time.time()
-    try:
-        file_list = await fetch_file_list_with_cache(
-            model_id,
-            revision,
-            recursive=True,
-            skip_internet=skip_internet,
-            on_connection_lost=on_connection_lost,
-        )
-    except FileNotFoundError:
-        not_started_progress = RepoDownloadProgress(
-            repo_id=str(model_id),
-            repo_revision=revision,
-            shard=shard,
-            completed_files=0,
-            total_files=0,
-            downloaded=Memory.from_bytes(0),
-            downloaded_this_session=Memory.from_bytes(0),
-            total=Memory.from_bytes(0),
-            overall_speed=0.0,
-            overall_eta=timedelta(0),
-            status="not_started",
-            file_progress={},
-        )
-        return EXO_DEFAULT_MODELS_DIR / model_id.normalize(), not_started_progress
+    file_list = await fetch_file_list_with_cache(
+        shard.model_card.model_id,
+        revision,
+        recursive=True,
+        skip_internet=skip_internet,
+        on_connection_lost=on_connection_lost,
+    )
    filtered_file_list = list(
        filter_repo_objects(
-            file_list,
-            allow_patterns=allow_patterns,
-            ignore_patterns=["original/*", "metal/*"],
-            key=lambda x: x.path,
+            file_list, allow_patterns=allow_patterns, key=lambda x: x.path
        )
    )

@@ -823,15 +765,6 @@ async def download_shard(
            for f in filtered_file_list
            if "/" in f.path or not f.path.endswith(".safetensors")
        ]
-
-    # Pick a writable directory with enough free space
-    total_size = sum(f.size or 0 for f in filtered_file_list)
-    models_dir = (
-        select_download_dir(total_size) if not skip_download else EXO_DEFAULT_MODELS_DIR
-    )
-    target_dir = models_dir / model_id.normalize()
-    if not skip_download:
-        await aios.makedirs(target_dir, exist_ok=True)
    file_progress: dict[str, RepoFileDownloadProgress] = {}

    async def on_progress_wrapper(
@@ -868,7 +801,7 @@ async def download_shard(
            else timedelta(seconds=0)
        )
        file_progress[file.path] = RepoFileDownloadProgress(
-            repo_id=model_id,
+            repo_id=shard.model_card.model_id,
            repo_revision=revision,
            file_path=file.path,
            downloaded=Memory.from_bytes(curr_bytes),
@@ -896,7 +829,7 @@ async def download_shard(
        downloaded_bytes = await get_downloaded_size(target_dir / file.path)
        final_file_exists = await aios.path.exists(target_dir / file.path)
        file_progress[file.path] = RepoFileDownloadProgress(
-            repo_id=model_id,
+            repo_id=shard.model_card.model_id,
            repo_revision=revision,
            file_path=file.path,
            downloaded=Memory.from_bytes(downloaded_bytes),
@@ -922,7 +855,7 @@ async def download_shard(
    async def download_with_semaphore(file: FileListEntry) -> None:
        async with semaphore:
            await download_file_with_retry(
-                model_id,
+                shard.model_card.model_id,
                revision,
                file.path,
                target_dir,
@@ -938,7 +871,7 @@ async def download_shard(
            *[download_with_semaphore(file) for file in filtered_file_list]
        )
    final_repo_progress = calculate_repo_progress(
-        shard, model_id, revision, file_progress, all_start_time
+        shard, shard.model_card.model_id, revision, file_progress, all_start_time
    )
    await on_progress(shard, final_repo_progress)
    if gguf := next((f for f in filtered_file_list if f.path.endswith(".gguf")), None):
@@ -1,6 +1,7 @@
 """Tests for download verification and cache behavior."""

 import time
+from collections.abc import AsyncIterator
 from datetime import timedelta
 from pathlib import Path
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -24,6 +25,15 @@ def model_id() -> ModelId:
    return ModelId("test-org/test-model")


+@pytest.fixture
+async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]:
+    """Set up a temporary models directory for testing."""
+    models_dir = tmp_path / "models"
+    await aios.makedirs(models_dir, exist_ok=True)
+    with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+        yield models_dir
+
+
 class TestFileVerification:
    """Tests for file size verification in _download_file."""

@@ -178,8 +188,7 @@ class TestFileListCache:
        ]

        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
-            patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
+            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
            patch(
                "exo.download.download_utils.fetch_file_list_with_retry",
                new_callable=AsyncMock,
@@ -225,8 +234,7 @@ class TestFileListCache:
            )

        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
-            patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
+            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
            patch(
                "exo.download.download_utils.fetch_file_list_with_retry",
                new_callable=AsyncMock,
@@ -244,8 +252,7 @@ class TestFileListCache:
        models_dir = tmp_path / "models"

        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
-            patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
+            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
            patch(
                "exo.download.download_utils.fetch_file_list_with_retry",
                new_callable=AsyncMock,
@@ -277,10 +284,7 @@ class TestModelDeletion:
        async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
            await f.write("[]")

-        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
-            patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
-        ):
+        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
            result = await delete_model(model_id)

            assert result is True
@@ -299,10 +303,7 @@ class TestModelDeletion:
        async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
            await f.write("[]")

-        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
-            patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
-        ):
+        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
            result = await delete_model(model_id)

            # Returns False because model dir didn't exist
@@ -317,10 +318,7 @@ class TestModelDeletion:
        models_dir = tmp_path / "models"
        await aios.makedirs(models_dir, exist_ok=True)

-        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
-            patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
-        ):
+        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
            result = await delete_model(model_id)

            assert result is False
@@ -1,297 +0,0 @@
-"""Tests for multi-directory model resolution, download target selection, and deletion."""
-
-import json
-import shutil
-from collections.abc import AsyncIterator
-from pathlib import Path
-from unittest.mock import patch
-
-import aiofiles
-import aiofiles.os as aios
-import pytest
-
-from exo.download.download_utils import (
-    InsufficientDiskSpaceError,
-    delete_model,
-    is_read_only_model_dir,
-    resolve_existing_model,
-    select_download_dir,
-)
-from exo.shared.types.common import ModelId
-
-MODEL_ID = ModelId("test-org/test-model")
-NORMALIZED = MODEL_ID.normalize()
-
-
-def _create_complete_model(model_dir: Path) -> None:
-    """Create a minimal complete model directory on disk."""
-    model_dir.mkdir(parents=True, exist_ok=True)
-    weight_map = {"layer.weight": "model.safetensors"}
-    index = {"metadata": {"total_size": 1024}, "weight_map": weight_map}
-    (model_dir / "model.safetensors.index.json").write_text(json.dumps(index))
-    (model_dir / "model.safetensors").write_bytes(b"weights")
-    (model_dir / "config.json").write_text('{"model_type": "test"}')
-
-
-def _create_incomplete_model(model_dir: Path) -> None:
-    """Create a model directory missing weight files."""
-    model_dir.mkdir(parents=True, exist_ok=True)
-    weight_map = {"layer.weight": "model.safetensors"}
-    index = {"metadata": {"total_size": 1024}, "weight_map": weight_map}
-    (model_dir / "model.safetensors.index.json").write_text(json.dumps(index))
-    # model.safetensors is missing
-
-
-# ---------------------------------------------------------------------------
-# resolve_existing_model
-# ---------------------------------------------------------------------------
-
-
-class TestResolveExistingModel:
-    def test_returns_none_when_no_dirs_have_model(self, tmp_path: Path) -> None:
-        writable = tmp_path / "writable"
-        writable.mkdir()
-        with (
-            patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", ()),
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
-        ):
-            assert resolve_existing_model(MODEL_ID) is None
-
-    def test_finds_model_in_writable_dir(self, tmp_path: Path) -> None:
-        writable = tmp_path / "writable"
-        _create_complete_model(writable / NORMALIZED)
-        with (
-            patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", ()),
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
-        ):
-            assert resolve_existing_model(MODEL_ID) == writable / NORMALIZED
-
-    def test_finds_model_in_read_only_dir(self, tmp_path: Path) -> None:
-        read_only = tmp_path / "readonly"
-        _create_complete_model(read_only / NORMALIZED)
-        writable = tmp_path / "writable"
-        writable.mkdir()
-        with (
-            patch(
-                "exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (read_only,)
-            ),
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
-        ):
-            assert resolve_existing_model(MODEL_ID) == read_only / NORMALIZED
-
-    def test_read_only_takes_priority_over_writable(self, tmp_path: Path) -> None:
-        read_only = tmp_path / "readonly"
-        _create_complete_model(read_only / NORMALIZED)
-        writable = tmp_path / "writable"
-        _create_complete_model(writable / NORMALIZED)
-        with (
-            patch(
-                "exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (read_only,)
-            ),
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
-        ):
-            result = resolve_existing_model(MODEL_ID)
-            assert result == read_only / NORMALIZED
-
-    def test_skips_incomplete_model(self, tmp_path: Path) -> None:
-        incomplete = tmp_path / "incomplete"
-        _create_incomplete_model(incomplete / NORMALIZED)
-        complete = tmp_path / "complete"
-        _create_complete_model(complete / NORMALIZED)
-        with (
-            patch(
-                "exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (incomplete,)
-            ),
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (complete,)),
-        ):
-            result = resolve_existing_model(MODEL_ID)
-            assert result == complete / NORMALIZED
-
-    def test_searches_multiple_read_only_dirs_in_order(self, tmp_path: Path) -> None:
-        ro1 = tmp_path / "ro1"
-        ro1.mkdir()
-        ro2 = tmp_path / "ro2"
-        _create_complete_model(ro2 / NORMALIZED)
-        writable = tmp_path / "writable"
-        writable.mkdir()
-        with (
-            patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (ro1, ro2)),
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
-        ):
-            assert resolve_existing_model(MODEL_ID) == ro2 / NORMALIZED
-
-
-# ---------------------------------------------------------------------------
-# is_read_only_model_dir
-# ---------------------------------------------------------------------------
-
-
-class TestIsReadOnlyModelDir:
-    def test_path_under_read_only_dir(self, tmp_path: Path) -> None:
-        ro = tmp_path / "readonly"
-        with patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (ro,)):
-            assert is_read_only_model_dir(ro / NORMALIZED) is True
-
-    def test_path_under_writable_dir(self, tmp_path: Path) -> None:
-        writable = tmp_path / "writable"
-        with patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", ()):
-            assert is_read_only_model_dir(writable / NORMALIZED) is False
-
-    def test_path_not_under_any_read_only_dir(self, tmp_path: Path) -> None:
-        ro = tmp_path / "readonly"
-        other = tmp_path / "other"
-        with patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (ro,)):
-            assert is_read_only_model_dir(other / NORMALIZED) is False
-
-
-# ---------------------------------------------------------------------------
-# select_download_dir
-# ---------------------------------------------------------------------------
-
-
-class TestSelectDownloadDir:
-    def test_picks_first_dir_with_enough_space(self, tmp_path: Path) -> None:
-        dir1 = tmp_path / "dir1"
-        dir2 = tmp_path / "dir2"
-        dir1.mkdir()
-        dir2.mkdir()
-        # Both exist on same filesystem so both have space; first wins
-        with patch("exo.download.download_utils.EXO_MODELS_DIRS", (dir1, dir2)):
-            assert select_download_dir(1) == dir1
-
-    def test_skips_dir_without_enough_space(self, tmp_path: Path) -> None:
-        dir1 = tmp_path / "dir1"
-        dir2 = tmp_path / "dir2"
-        dir1.mkdir()
-        dir2.mkdir()
-
-        real_disk_usage = shutil.disk_usage
-
-        def mock_disk_usage(path: str | Path) -> object:
-            if Path(path).is_relative_to(dir1):
-                real = real_disk_usage(path)
-                return shutil._ntuple_diskusage(real.total, real.total, 0)  # pyright: ignore[reportPrivateUsage]
-            return real_disk_usage(path)
-
-        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (dir1, dir2)),
-            patch("shutil.disk_usage", side_effect=mock_disk_usage),
-        ):
-            assert select_download_dir(1024) == dir2
-
-    def test_raises_when_no_dir_has_space(self, tmp_path: Path) -> None:
-        dir1 = tmp_path / "dir1"
-        dir1.mkdir()
-
-        real_disk_usage = shutil.disk_usage
-
-        def mock_disk_usage(path: str | Path) -> object:
-            real = real_disk_usage(path)
-            return shutil._ntuple_diskusage(real.total, real.total, 0)  # pyright: ignore[reportPrivateUsage]
-
-        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (dir1,)),
-            patch("shutil.disk_usage", side_effect=mock_disk_usage),
-            pytest.raises(InsufficientDiskSpaceError),
-        ):
-            select_download_dir(1024)
-
-    def test_skips_nonexistent_dir(self, tmp_path: Path) -> None:
-        nonexistent = tmp_path / "does-not-exist"
-        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (nonexistent,)),
-            pytest.raises(InsufficientDiskSpaceError),
-        ):
-            select_download_dir(1)
-
-    def test_skips_dir_raising_oserror(self, tmp_path: Path) -> None:
-        dir1 = tmp_path / "unmounted"
-        dir2 = tmp_path / "ok"
-        dir1.mkdir()
-        dir2.mkdir()
-
-        real_disk_usage = shutil.disk_usage
-
-        def mock_disk_usage(path: str | Path) -> object:
-            if Path(path).is_relative_to(dir1):
-                raise OSError("device not mounted")
-            return real_disk_usage(path)
-
-        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIRS", (dir1, dir2)),
-            patch("shutil.disk_usage", side_effect=mock_disk_usage),
-        ):
-            assert select_download_dir(1) == dir2
-
-
-# ---------------------------------------------------------------------------
-# delete_model
-# ---------------------------------------------------------------------------
-
-
-class TestDeleteModel:
-    @pytest.fixture
-    async def dirs(self, tmp_path: Path) -> AsyncIterator[tuple[Path, Path, Path]]:
-        writable1 = tmp_path / "w1"
-        writable2 = tmp_path / "w2"
-        default = tmp_path / "default"
-        await aios.makedirs(writable1, exist_ok=True)
-        await aios.makedirs(writable2, exist_ok=True)
-        await aios.makedirs(default, exist_ok=True)
-        with (
-            patch(
-                "exo.download.download_utils.EXO_MODELS_DIRS",
-                (writable1, writable2, default),
-            ),
-            patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", default),
-        ):
-            yield writable1, writable2, default
-
-    async def test_deletes_from_writable_dir(
-        self, dirs: tuple[Path, Path, Path]
-    ) -> None:
-        w1, _, _ = dirs
-        model_dir = w1 / NORMALIZED
-        await aios.makedirs(model_dir, exist_ok=True)
-        async with aiofiles.open(model_dir / "weights.safetensors", "w") as f:
-            await f.write("data")
-
-        result = await delete_model(MODEL_ID)
-        assert result is True
-        assert not await aios.path.exists(model_dir)
-
-    async def test_deletes_from_multiple_writable_dirs(
-        self, dirs: tuple[Path, Path, Path]
-    ) -> None:
-        w1, w2, _ = dirs
-        model_dir1 = w1 / NORMALIZED
-        model_dir2 = w2 / NORMALIZED
-        await aios.makedirs(model_dir1, exist_ok=True)
-        await aios.makedirs(model_dir2, exist_ok=True)
-        async with aiofiles.open(model_dir1 / "w.safetensors", "w") as f:
-            await f.write("data")
-        async with aiofiles.open(model_dir2 / "w.safetensors", "w") as f:
-            await f.write("data")
-
-        result = await delete_model(MODEL_ID)
-        assert result is True
-        assert not await aios.path.exists(model_dir1)
-        assert not await aios.path.exists(model_dir2)
-
-    async def test_cleans_cache_from_default_dir(
-        self, dirs: tuple[Path, Path, Path]
-    ) -> None:
-        _, _, default = dirs
-        cache_dir = default / "caches" / NORMALIZED
-        await aios.makedirs(cache_dir, exist_ok=True)
-        async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
-            await f.write("[]")
-
-        await delete_model(MODEL_ID)
-        assert not await aios.path.exists(cache_dir)
-
-    async def test_returns_false_when_model_not_found(
-        self, dirs: tuple[Path, Path, Path]
-    ) -> None:
-        result = await delete_model(MODEL_ID)
-        assert result is False
@@ -26,10 +26,7 @@ def model_id() -> ModelId:
 async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]:
    models_dir = tmp_path / "models"
    await aios.makedirs(models_dir, exist_ok=True)
-    with (
-        patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
-        patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
-    ):
+    with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
        yield models_dir


@@ -186,7 +186,7 @@ async def test_re_download_after_delete_completes() -> None:
                "Re-download after deletion should complete"
            )
        finally:
-            await coordinator.shutdown()
+            coordinator.shutdown()
            coordinator_task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await coordinator_task
@@ -11,9 +11,9 @@ from loguru import logger
 from pydantic import PositiveInt

 import exo.routing.topics as topics
-from exo.api.main import API
 from exo.download.coordinator import DownloadCoordinator
 from exo.download.impl_shard_downloader import exo_shard_downloader
+from exo.master.api import API  # TODO: should API be in master?
 from exo.master.main import Master
 from exo.routing.event_router import EventRouter
 from exo.routing.router import Router, get_node_id_keypair
@@ -47,11 +47,7 @@ class Node:
        keypair = get_node_id_keypair()
        node_id = NodeId(keypair.to_node_id())
        session_id = SessionId(master_node_id=node_id, election_clock=0)
-        router = Router.create(
-            keypair,
-            bootstrap_peers=args.bootstrap_peers,
-            listen_port=args.libp2p_port,
-        )
+        router = Router.create(keypair)
        await router.register_topic(topics.GLOBAL_EVENTS)
        await router.register_topic(topics.LOCAL_EVENTS)
        await router.register_topic(topics.COMMANDS)
@@ -228,7 +224,7 @@ class Node:
                    )
                if result.is_new_master:
                    if self.download_coordinator:
-                        await self.download_coordinator.shutdown()
+                        self.download_coordinator.shutdown()
                        self.download_coordinator = DownloadCoordinator(
                            self.node_id,
                            exo_shard_downloader(offline=self.offline),
@@ -240,7 +236,7 @@ class Node:
                        )
                        self._tg.start_soon(self.download_coordinator.run)
                    if self.worker:
-                        await self.worker.shutdown()
+                        self.worker.shutdown()
                        # TODO: add profiling etc to resource monitor
                        self.worker = Worker(
                            self.node_id,
@@ -268,17 +264,12 @@ def main():
    mp.set_start_method("spawn", force=True)
    # TODO: Refactor the current verbosity system
    logger_setup(EXO_LOG, args.verbosity)
-    logger.info(f"{'=' * 40}")
-    logger.info(f"Starting EXO | pid={os.getpid()}")
-    logger.info(f"{'=' * 40}")
+    logger.info("Starting EXO")
    logger.info(f"EXO_LIBP2P_NAMESPACE: {os.getenv('EXO_LIBP2P_NAMESPACE')}")

    if args.offline:
        logger.info("Running in OFFLINE mode — no internet checks, local models only")

-    if args.bootstrap_peers:
-        logger.info(f"Bootstrap peers: {args.bootstrap_peers}")
-
    if args.no_batch:
        os.environ["EXO_NO_BATCH"] = "1"
        logger.info("Continuous batching disabled (--no-batch)")
@@ -315,8 +306,6 @@ class Args(CamelCaseModel):
    offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
    no_batch: bool = False
    fast_synch: bool | None = None  # None = auto, True = force on, False = force off
-    bootstrap_peers: list[str] = []
-    libp2p_port: int

    @classmethod
    def parse(cls) -> Self:
@@ -374,22 +363,6 @@ class Args(CamelCaseModel):
            action="store_true",
            help="Disable continuous batching, use sequential generation",
        )
-        parser.add_argument(
-            "--bootstrap-peers",
-            type=lambda s: [p for p in s.split(",") if p],
-            default=os.getenv("EXO_BOOTSTRAP_PEERS", "").split(",")
-            if os.getenv("EXO_BOOTSTRAP_PEERS")
-            else [],
-            dest="bootstrap_peers",
-            help="Comma-separated libp2p multiaddrs to dial on startup (env: EXO_BOOTSTRAP_PEERS)",
-        )
-        parser.add_argument(
-            "--libp2p-port",
-            type=int,
-            default=0,
-            dest="libp2p_port",
-            help="Fixed TCP port for libp2p to listen on (0 = OS-assigned).",
-        )
        fast_synch_group = parser.add_mutually_exclusive_group()
        fast_synch_group.add_argument(
            "--fast-synch",
@@ -4,7 +4,7 @@ import time
 from collections.abc import AsyncGenerator
 from typing import Any

-from exo.api.types import (
+from exo.shared.types.api import (
    ChatCompletionChoice,
    ChatCompletionMessage,
    ChatCompletionMessageText,
@@ -221,7 +221,6 @@ async def generate_chat_stream(
                    if chunk.stats is not None:
                        yield f": generation_stats {chunk.stats.model_dump_json()}\n\n"
                    yield "data: [DONE]\n\n"
-                    return


 async def collect_chat_response(
@@ -5,8 +5,14 @@ import re
 from collections.abc import AsyncGenerator
 from typing import Any

-from exo.api.types import FinishReason, Usage
-from exo.api.types.claude_api import (
+from exo.shared.types.api import FinishReason, Usage
+from exo.shared.types.chunks import (
+    ErrorChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
+from exo.shared.types.claude_api import (
    ClaudeContentBlock,
    ClaudeContentBlockDeltaEvent,
    ClaudeContentBlockStartEvent,
@@ -29,12 +35,6 @@ from exo.api.types.claude_api import (
    ClaudeToolUseBlock,
    ClaudeUsage,
 )
-from exo.shared.types.chunks import (
-    ErrorChunk,
-    PrefillProgressChunk,
-    TokenChunk,
-    ToolCallChunk,
-)
 from exo.shared.types.common import CommandId
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams

@@ -4,7 +4,14 @@ import json
 from collections.abc import AsyncGenerator
 from typing import Any

-from exo.api.types.ollama_api import (
+from exo.shared.types.chunks import (
+    ErrorChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
+from exo.shared.types.common import CommandId
+from exo.shared.types.ollama_api import (
    OllamaChatRequest,
    OllamaChatResponse,
    OllamaDoneReason,
@@ -14,13 +21,6 @@ from exo.api.types.ollama_api import (
    OllamaToolCall,
    OllamaToolFunction,
 )
-from exo.shared.types.chunks import (
-    ErrorChunk,
-    PrefillProgressChunk,
-    TokenChunk,
-    ToolCallChunk,
-)
-from exo.shared.types.common import CommandId
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams


@@ -4,8 +4,15 @@ from collections.abc import AsyncGenerator
 from itertools import count
 from typing import Any

-from exo.api.types import Usage
-from exo.api.types.openai_responses import (
+from exo.shared.types.api import Usage
+from exo.shared.types.chunks import (
+    ErrorChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
+from exo.shared.types.common import CommandId
+from exo.shared.types.openai_responses import (
    FunctionCallInputItem,
    ResponseCompletedEvent,
    ResponseContentPart,
@@ -35,13 +42,6 @@ from exo.api.types.openai_responses import (
    ResponseTextDoneEvent,
    ResponseUsage,
 )
-from exo.shared.types.chunks import (
-    ErrorChunk,
-    PrefillProgressChunk,
-    TokenChunk,
-    ToolCallChunk,
-)
-from exo.shared.types.common import CommandId
 from exo.shared.types.text_generation import (
    InputMessage,
    TextGenerationTaskParams,
@@ -21,17 +21,17 @@ from hypercorn.config import Config
 from hypercorn.typing import ASGIFramework
 from loguru import logger

-from exo.api.adapters.chat_completions import (
+from exo.master.adapters.chat_completions import (
    chat_request_to_text_generation,
    collect_chat_response,
    generate_chat_stream,
 )
-from exo.api.adapters.claude import (
+from exo.master.adapters.claude import (
    claude_request_to_text_generation,
    collect_claude_response,
    generate_claude_stream,
 )
-from exo.api.adapters.ollama import (
+from exo.master.adapters.ollama import (
    collect_ollama_chat_response,
    collect_ollama_generate_response,
    generate_ollama_chat_stream,
@@ -39,12 +39,34 @@ from exo.api.adapters.ollama import (
    ollama_generate_request_to_text_generation,
    ollama_request_to_text_generation,
 )
-from exo.api.adapters.responses import (
+from exo.master.adapters.responses import (
    collect_responses_response,
    generate_responses_stream,
    responses_request_to_text_generation,
 )
-from exo.api.types import (
+from exo.master.event_log import DiskEventLog
+from exo.master.image_store import ImageStore
+from exo.master.placement import place_instance as get_instance_placements
+from exo.shared.apply import apply
+from exo.shared.constants import (
+    DASHBOARD_DIR,
+    EXO_CACHE_HOME,
+    EXO_EVENT_LOG_DIR,
+    EXO_IMAGE_CACHE_DIR,
+    EXO_MAX_CHUNK_SIZE,
+    EXO_TRACING_CACHE_DIR,
+)
+from exo.shared.election import ElectionMessage
+from exo.shared.logging import InterceptLogger
+from exo.shared.models.model_cards import (
+    ModelCard,
+    ModelId,
+    delete_custom_card,
+    get_model_cards,
+    is_custom_card,
+)
+from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file
+from exo.shared.types.api import (
    AddCustomModelParams,
    AdvancedImageParams,
    BenchChatCompletionRequest,
@@ -92,48 +114,6 @@ from exo.api.types import (
    TraceStatsResponse,
    normalize_image_size,
 )
-from exo.api.types.claude_api import (
-    ClaudeMessagesRequest,
-    ClaudeMessagesResponse,
-)
-from exo.api.types.ollama_api import (
-    OllamaChatRequest,
-    OllamaChatResponse,
-    OllamaGenerateRequest,
-    OllamaGenerateResponse,
-    OllamaModelDetails,
-    OllamaModelTag,
-    OllamaPsModel,
-    OllamaPsResponse,
-    OllamaShowRequest,
-    OllamaShowResponse,
-    OllamaTagsResponse,
-)
-from exo.api.types.openai_responses import (
-    ResponsesRequest,
-    ResponsesResponse,
-)
-from exo.master.image_store import ImageStore
-from exo.master.placement import place_instance as get_instance_placements
-from exo.shared.apply import apply
-from exo.shared.constants import (
-    DASHBOARD_DIR,
-    EXO_CACHE_HOME,
-    EXO_EVENT_LOG_DIR,
-    EXO_IMAGE_CACHE_DIR,
-    EXO_MAX_CHUNK_SIZE,
-    EXO_TRACING_CACHE_DIR,
-)
-from exo.shared.election import ElectionMessage
-from exo.shared.logging import InterceptLogger
-from exo.shared.models.model_cards import (
-    ModelCard,
-    ModelId,
-    add_to_card_cache,
-    get_card,
-    get_model_cards,
-)
-from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file
 from exo.shared.types.chunks import (
    ErrorChunk,
    ImageChunk,
@@ -142,11 +122,13 @@ from exo.shared.types.chunks import (
    TokenChunk,
    ToolCallChunk,
 )
+from exo.shared.types.claude_api import (
+    ClaudeMessagesRequest,
+    ClaudeMessagesResponse,
+)
 from exo.shared.types.commands import (
-    AddCustomModelCard,
    Command,
    CreateInstance,
-    DeleteCustomModelCard,
    DeleteDownload,
    DeleteInstance,
    DownloadCommand,
@@ -169,13 +151,29 @@ from exo.shared.types.events import (
    TracesMerged,
 )
 from exo.shared.types.memory import Memory
+from exo.shared.types.ollama_api import (
+    OllamaChatRequest,
+    OllamaChatResponse,
+    OllamaGenerateRequest,
+    OllamaGenerateResponse,
+    OllamaModelDetails,
+    OllamaModelTag,
+    OllamaPsModel,
+    OllamaPsResponse,
+    OllamaShowRequest,
+    OllamaShowResponse,
+    OllamaTagsResponse,
+)
+from exo.shared.types.openai_responses import (
+    ResponsesRequest,
+    ResponsesResponse,
+)
 from exo.shared.types.state import State
 from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.banner import print_startup_banner
 from exo.utils.channels import Receiver, Sender, channel
-from exo.utils.disk_event_log import DiskEventLog
 from exo.utils.power_sampler import PowerSampler
 from exo.utils.task_group import TaskGroup

@@ -415,7 +413,6 @@ class API:
                node_network=self.state.node_network,
                topology=self.state.topology,
                current_instances=self.state.instances,
-                download_status=self.state.downloads,
            )
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc
@@ -478,7 +475,6 @@ class API:
                    topology=self.state.topology,
                    current_instances=self.state.instances,
                    required_nodes=required_nodes,
-                    download_status=self.state.downloads,
                )
            except ValueError as exc:
                if (model_card.model_id, sharding, instance_meta, 0) not in seen:
@@ -1562,7 +1558,7 @@ class API:
                    storage_size_megabytes=card.storage_size.in_mb,
                    supports_tensor=card.supports_tensor,
                    tasks=[task.value for task in card.tasks],
-                    is_custom=card.is_custom,
+                    is_custom=is_custom_card(card.model_id),
                    family=card.family,
                    quantization=card.quantization,
                    base_model=card.base_model,
@@ -1573,7 +1569,7 @@ class API:
        )

    async def add_custom_model(self, payload: AddCustomModelParams) -> ModelListModel:
-        """Fetch a model from HuggingFace and save as a custom model card, then sync across the cluster."""
+        """Fetch a model from HuggingFace and save as a custom model card."""
        try:
            card = await ModelCard.fetch_from_hf(payload.model_id)
        except Exception as exc:
@@ -1581,17 +1577,6 @@ class API:
                status_code=400, detail=f"Failed to fetch model: {exc}"
            ) from exc

-        await self.command_sender.send(
-            ForwarderCommand(
-                origin=self._system_id,
-                command=AddCustomModelCard(model_card=card),
-            )
-        )
-
-        # Immediately update the local cache so the subsequent GET /models
-        # returns the new model without waiting for the event round-trip.
-        add_to_card_cache(card)
-
        return ModelListModel(
            id=card.model_id,
            hugging_face_id=card.model_id,
@@ -1605,18 +1590,10 @@ class API:
        )

    async def delete_custom_model(self, model_id: ModelId) -> JSONResponse:
-        """Delete a user-added custom model card and sync deletion across the cluster."""
-        card = get_card(model_id)
-        if card is None or not card.is_custom:
+        """Delete a user-added custom model card."""
+        deleted = await delete_custom_card(model_id)
+        if not deleted:
            raise HTTPException(status_code=404, detail="Custom model card not found")
-
-        await self.command_sender.send(
-            ForwarderCommand(
-                origin=self._system_id,
-                command=DeleteCustomModelCard(model_id=model_id),
-            )
-        )
-
        return JSONResponse(
            {"message": "Model card deleted", "model_id": str(model_id)}
        )
@@ -3,6 +3,7 @@ from datetime import datetime, timedelta, timezone
 import anyio
 from loguru import logger

+from exo.master.event_log import DiskEventLog
 from exo.master.placement import (
    add_instance_to_placements,
    cancel_unnecessary_downloads,
@@ -13,9 +14,7 @@ from exo.master.placement import (
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
 from exo.shared.types.commands import (
-    AddCustomModelCard,
    CreateInstance,
-    DeleteCustomModelCard,
    DeleteInstance,
    ForwarderCommand,
    ForwarderDownloadCommand,
@@ -31,8 +30,6 @@ from exo.shared.types.commands import (
 )
 from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
-    CustomModelCardAdded,
-    CustomModelCardDeleted,
    Event,
    GlobalForwarderEvent,
    IndexedEvent,
@@ -64,7 +61,6 @@ from exo.shared.types.tasks import (
 )
 from exo.shared.types.worker.instances import InstanceId
 from exo.utils.channels import Receiver, Sender
-from exo.utils.disk_event_log import DiskEventLog
 from exo.utils.event_buffer import MultiSourceBuffer
 from exo.utils.task_group import TaskGroup

@@ -298,7 +294,6 @@ class Master:
                                self.state.instances,
                                self.state.node_memory,
                                self.state.node_network,
-                                download_status=self.state.downloads,
                            )
                            transition_events = get_transition_events(
                                self.state.instances, placement, self.state.tasks
@@ -349,14 +344,6 @@ class Master:
                                    f"Finished command {command.finished_command_id} finished"
                                )

-                        case AddCustomModelCard():
-                            generated_events.append(
-                                CustomModelCardAdded(model_card=command.model_card)
-                            )
-                        case DeleteCustomModelCard():
-                            generated_events.append(
-                                CustomModelCardDeleted(model_id=command.model_id)
-                            )
                        case RequestEventLog():
                            # We should just be able to send everything, since other buffers will ignore old messages
                            # rate limit to 1000 at a time
@@ -32,10 +32,7 @@ from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import (
-    DownloadCompleted,
-    DownloadFailed,
    DownloadOngoing,
-    DownloadPending,
    DownloadProgress,
 )
 from exo.shared.types.worker.instances import (
@@ -63,45 +60,6 @@ def add_instance_to_placements(
    return {**current_instances, command.instance.instance_id: command.instance}


-def _get_node_download_fraction(
-    node_id: NodeId,
-    model_id: ModelId,
-    download_status: Mapping[NodeId, Sequence[DownloadProgress]],
-) -> float:
-    """Return the download fraction (0.0–1.0) for a model on a given node."""
-    for progress in download_status.get(node_id, []):
-        if progress.shard_metadata.model_card.model_id != model_id:
-            continue
-        match progress:
-            case DownloadCompleted():
-                return 1.0
-            case DownloadOngoing():
-                total = progress.download_progress.total.in_bytes
-                return (
-                    progress.download_progress.downloaded.in_bytes / total
-                    if total > 0
-                    else 0.0
-                )
-            case DownloadPending():
-                total = progress.total.in_bytes
-                return progress.downloaded.in_bytes / total if total > 0 else 0.0
-            case DownloadFailed():
-                return 0.0
-    return 0.0
-
-
-def _cycle_download_score(
-    cycle: Cycle,
-    model_id: ModelId,
-    download_status: Mapping[NodeId, Sequence[DownloadProgress]],
-) -> float:
-    """Sum of download fractions across all nodes in a cycle."""
-    return sum(
-        _get_node_download_fraction(node_id, model_id, download_status)
-        for node_id in cycle
-    )
-
-
 def place_instance(
    command: PlaceInstance,
    topology: Topology,
@@ -109,7 +67,6 @@ def place_instance(
    node_memory: Mapping[NodeId, MemoryUsage],
    node_network: Mapping[NodeId, NodeNetworkInfo],
    required_nodes: set[NodeId] | None = None,
-    download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None,
 ) -> dict[InstanceId, Instance]:
    cycles = topology.get_cycles()
    candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
@@ -173,21 +130,11 @@ def place_instance(
        if any(topology.node_is_leaf(node_id) for node_id in cycle)
    ]

-    resolved_download_status = download_status or {}
-    candidate_cycles = (
-        cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else smallest_cycles
-    )
-
    selected_cycle = max(
-        candidate_cycles,
-        key=lambda cycle: (
-            _cycle_download_score(
-                cycle, command.model_card.model_id, resolved_download_status
-            ),
-            sum(
-                (node_memory[node_id].ram_available for node_id in cycle),
-                start=Memory(),
-            ),
+        cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else smallest_cycles,
+        key=lambda cycle: sum(
+            (node_memory[node_id].ram_available for node_id in cycle),
+            start=Memory(),
        ),
    )

@@ -4,11 +4,10 @@ from typing import Any
 from fastapi import FastAPI, HTTPException
 from fastapi.testclient import TestClient

-from exo.api.main import API
-

 def test_http_exception_handler_formats_openai_style() -> None:
    """Test that HTTPException is converted to OpenAI-style error format."""
+    from exo.master.api import API

    app = FastAPI()

@@ -5,12 +5,12 @@ from unittest.mock import AsyncMock, MagicMock
 from fastapi import FastAPI
 from fastapi.testclient import TestClient

-from exo.api.main import API
 from exo.shared.types.common import CommandId


 def _make_api() -> Any:
    """Create a minimal API instance with cancel route and error handler."""
+    from exo.master.api import API

    app = FastAPI()
    api = object.__new__(API)
@@ -3,11 +3,11 @@
 import pydantic
 import pytest

-from exo.api.adapters.claude import (
+from exo.master.adapters.claude import (
    claude_request_to_text_generation,
    finish_reason_to_claude_stop_reason,
 )
-from exo.api.types.claude_api import (
+from exo.shared.types.claude_api import (
    ClaudeMessage,
    ClaudeMessagesRequest,
    ClaudeTextBlock,
@@ -4,12 +4,12 @@ import json
 from collections.abc import AsyncGenerator
 from typing import Any, cast

-from exo.api.adapters.claude import (
+from exo.master.adapters.claude import (
    ClaudeMessagesResponse,
    collect_claude_response,
    generate_claude_stream,
 )
-from exo.api.types import ToolCallItem
+from exo.shared.types.api import ToolCallItem
 from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
 from exo.shared.types.common import CommandId, ModelId

@@ -2,8 +2,8 @@ from pathlib import Path

 import pytest

+from exo.master.event_log import DiskEventLog
 from exo.shared.types.events import TestEvent
-from exo.utils.disk_event_log import DiskEventLog


@pytest.fixture
@@ -7,11 +7,11 @@ The responses adapter converts it to TextGenerationTaskParams for the pipeline.
 import pydantic
 import pytest

-from exo.api.types.openai_responses import (
+from exo.shared.types.common import ModelId
+from exo.shared.types.openai_responses import (
    ResponseInputMessage,
    ResponsesRequest,
 )
-from exo.shared.types.common import ModelId


 class TestResponsesRequestValidation:
@@ -25,12 +25,6 @@ from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
 from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.shared.types.topology import Connection, SocketConnection
-from exo.shared.types.worker.downloads import (
-    DownloadCompleted,
-    DownloadFailed,
-    DownloadOngoing,
-    DownloadProgressData,
-)
 from exo.shared.types.worker.instances import (
    Instance,
    InstanceId,
@@ -39,7 +33,7 @@ from exo.shared.types.worker.instances import (
    MlxRingInstance,
 )
 from exo.shared.types.worker.runners import ShardAssignments
-from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
+from exo.shared.types.worker.shards import Sharding


@pytest.fixture
@@ -582,183 +576,3 @@ def test_get_transition_events_delete_instance_cancels_only_matching_tasks(
    assert cancel_events[0].task_status == TaskStatus.Cancelled
    assert len(delete_events) == 1
    assert delete_events[0].instance_id == instance_id_a
-
-
-def _make_shard_metadata(model_card: ModelCard) -> PipelineShardMetadata:
-    return PipelineShardMetadata(
-        model_card=model_card,
-        device_rank=0,
-        world_size=1,
-        start_layer=0,
-        end_layer=model_card.n_layers,
-        n_layers=model_card.n_layers,
-    )
-
-
-def test_placement_prefers_cycle_with_downloaded_model(
-    model_card: ModelCard,
-) -> None:
-    """When two cycles are otherwise equal, prefer the one with the model already downloaded."""
-    topology = Topology()
-
-    model_card.storage_size = Memory.from_bytes(500)
-
-    node_a = NodeId()
-    node_b = NodeId()
-
-    node_memory = {
-        node_a: create_node_memory(1000),
-        node_b: create_node_memory(1000),
-    }
-    node_network = {
-        node_a: create_node_network(),
-        node_b: create_node_network(),
-    }
-
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    # No connections between them — two single-node cycles
-
-    shard_meta = _make_shard_metadata(model_card)
-
-    # node_b has the model fully downloaded, node_a does not
-    download_status = {
-        node_b: [
-            DownloadCompleted(
-                node_id=node_b,
-                shard_metadata=shard_meta,
-                total=model_card.storage_size,
-            ),
-        ],
-    }
-
-    cic = place_instance_command(model_card)
-    placements = place_instance(
-        cic, topology, {}, node_memory, node_network, download_status=download_status
-    )
-
-    assert len(placements) == 1
-    instance = list(placements.values())[0]
-    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
-    assert assigned_nodes == {node_b}
-
-
-def test_placement_prefers_cycle_with_higher_download_progress(
-    model_card: ModelCard,
-) -> None:
-    """When two cycles are otherwise equal, prefer the one with more download progress."""
-    topology = Topology()
-
-    model_card.storage_size = Memory.from_bytes(1000)
-
-    node_a = NodeId()
-    node_b = NodeId()
-
-    node_memory = {
-        node_a: create_node_memory(1000),
-        node_b: create_node_memory(1000),
-    }
-    node_network = {
-        node_a: create_node_network(),
-        node_b: create_node_network(),
-    }
-
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-
-    shard_meta = _make_shard_metadata(model_card)
-
-    # node_a: 30% downloaded, node_b: 80% downloaded
-    download_status = {
-        node_a: [
-            DownloadOngoing(
-                node_id=node_a,
-                shard_metadata=shard_meta,
-                download_progress=DownloadProgressData(
-                    total=Memory.from_bytes(1000),
-                    downloaded=Memory.from_bytes(300),
-                    downloaded_this_session=Memory.from_bytes(300),
-                    completed_files=0,
-                    total_files=1,
-                    speed=0.0,
-                    eta_ms=0,
-                    files={},
-                ),
-            ),
-        ],
-        node_b: [
-            DownloadOngoing(
-                node_id=node_b,
-                shard_metadata=shard_meta,
-                download_progress=DownloadProgressData(
-                    total=Memory.from_bytes(1000),
-                    downloaded=Memory.from_bytes(800),
-                    downloaded_this_session=Memory.from_bytes(800),
-                    completed_files=0,
-                    total_files=1,
-                    speed=0.0,
-                    eta_ms=0,
-                    files={},
-                ),
-            ),
-        ],
-    }
-
-    cic = place_instance_command(model_card)
-    placements = place_instance(
-        cic, topology, {}, node_memory, node_network, download_status=download_status
-    )
-
-    assert len(placements) == 1
-    instance = list(placements.values())[0]
-    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
-    assert assigned_nodes == {node_b}
-
-
-def test_placement_does_not_prefer_cycle_with_failed_download(
-    model_card: ModelCard,
-) -> None:
-    """A failed download should count as 0% — not preferred over a node with no download history."""
-    topology = Topology()
-
-    model_card.storage_size = Memory.from_bytes(500)
-
-    node_a = NodeId()
-    node_b = NodeId()
-
-    # node_a has slightly more RAM so it would win on the RAM tiebreaker
-    node_memory = {
-        node_a: create_node_memory(1001),
-        node_b: create_node_memory(1000),
-    }
-    node_network = {
-        node_a: create_node_network(),
-        node_b: create_node_network(),
-    }
-
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-
-    shard_meta = _make_shard_metadata(model_card)
-
-    # node_b has a failed download — should not be preferred
-    download_status = {
-        node_b: [
-            DownloadFailed(
-                node_id=node_b,
-                shard_metadata=shard_meta,
-                error_message="connection reset",
-            ),
-        ],
-    }
-
-    cic = place_instance_command(model_card)
-    placements = place_instance(
-        cic, topology, {}, node_memory, node_network, download_status=download_status
-    )
-
-    assert len(placements) == 1
-    instance = list(placements.values())[0]
-    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
-    # node_a should win on RAM tiebreaker since failed download scores 0.0
-    assert assigned_nodes == {node_a}
@@ -1,4 +1,3 @@
-from collections.abc import Sequence
 from copy import copy
 from itertools import count
 from math import inf
@@ -103,15 +102,8 @@ class TopicRouter[T: CamelCaseModel]:

 class Router:
    @classmethod
-    def create(
-        cls,
-        identity: Keypair,
-        bootstrap_peers: Sequence[str] = (),
-        listen_port: int = 0,
-    ) -> "Router":
-        return cls(
-            handle=NetworkingHandle(identity, list(bootstrap_peers), listen_port)
-        )
+    def create(cls, identity: Keypair) -> "Router":
+        return cls(handle=NetworkingHandle(identity))

    def __init__(self, handle: NetworkingHandle):
        self.topic_routers: dict[str, TopicRouter[CamelCaseModel]] = {}
@@ -7,8 +7,6 @@ from loguru import logger
 from exo.shared.types.common import NodeId
 from exo.shared.types.events import (
    ChunkGenerated,
-    CustomModelCardAdded,
-    CustomModelCardDeleted,
    Event,
    IndexedEvent,
    InputChunkReceived,
@@ -67,8 +65,6 @@ def event_apply(event: Event, state: State) -> State:
            | InputChunkReceived()
            | TracesCollected()
            | TracesMerged()
-            | CustomModelCardAdded()
-            | CustomModelCardDeleted()
        ):  # Pass-through events that don't modify state
            return state
        case InstanceCreated():
@@ -119,13 +115,7 @@ def apply_node_download_progress(event: NodeDownloadProgress, state: State) -> S

    replaced = False
    for i, existing_dp in enumerate(current):
-        # TODO(ciaran): deduplicate by model_id for now. Will need to use
-        # shard_metadata again when pipeline and tensor downloads differ.
-        # For now this is fine
-        if (
-            existing_dp.shard_metadata.model_card.model_id
-            == dp.shard_metadata.model_card.model_id
-        ):
+        if existing_dp.shard_metadata == dp.shard_metadata:
            current[i] = dp
            replaced = True
            break
@@ -26,35 +26,21 @@ EXO_CONFIG_HOME = _get_xdg_dir("XDG_CONFIG_HOME", ".config")
 EXO_DATA_HOME = _get_xdg_dir("XDG_DATA_HOME", ".local/share")
 EXO_CACHE_HOME = _get_xdg_dir("XDG_CACHE_HOME", ".cache")

-# Default models directory (always included as first entry in writable dirs)
-_EXO_DEFAULT_MODELS_DIR_ENV = os.environ.get("EXO_DEFAULT_MODELS_DIR", None)
-EXO_DEFAULT_MODELS_DIR = (
-    Path(_EXO_DEFAULT_MODELS_DIR_ENV).expanduser()
-    if _EXO_DEFAULT_MODELS_DIR_ENV is not None
-    else EXO_DATA_HOME / "models"
+# Models directory (data)
+_EXO_MODELS_DIR_ENV = os.environ.get("EXO_MODELS_DIR", None)
+EXO_MODELS_DIR = (
+    EXO_DATA_HOME / "models"
+    if _EXO_MODELS_DIR_ENV is None
+    else Path.home() / _EXO_MODELS_DIR_ENV
 )

-
-def _parse_colon_dirs(env_var: str) -> tuple[Path, ...]:
-    raw = os.environ.get(env_var, None)
-    if raw is None:
-        return ()
-    return tuple(Path(p).expanduser() for p in raw.split(":") if p)
-
-
-# Read-only model directories (colon-separated). Never written to or deleted from.
-_EXO_MODELS_READ_ONLY_DIRS_ENV = _parse_colon_dirs("EXO_MODELS_READ_ONLY_DIRS")
-# Writable model directories (colon-separated). Default dir is always prepended.
-_EXO_MODELS_DIRS_ENV = _parse_colon_dirs("EXO_MODELS_DIRS")
-
-# If a directory appears in both lists, treat it as read-only.
-_read_only_set = frozenset(_EXO_MODELS_READ_ONLY_DIRS_ENV)
-EXO_MODELS_DIRS: tuple[Path, ...] = tuple(
-    d
-    for d in (EXO_DEFAULT_MODELS_DIR, *_EXO_MODELS_DIRS_ENV)
-    if d not in _read_only_set
+# Read-only search path for pre-downloaded models (colon-separated directories)
+_EXO_MODELS_PATH_ENV = os.environ.get("EXO_MODELS_PATH", None)
+EXO_MODELS_PATH: tuple[Path, ...] | None = (
+    tuple(Path(p).expanduser() for p in _EXO_MODELS_PATH_ENV.split(":") if p)
+    if _EXO_MODELS_PATH_ENV is not None
+    else None
 )
-EXO_MODELS_READ_ONLY_DIRS: tuple[Path, ...] = _EXO_MODELS_READ_ONLY_DIRS_ENV

 _RESOURCES_DIR_ENV = os.environ.get("EXO_RESOURCES_DIR", None)
 RESOURCES_DIR = (
@@ -66,7 +66,7 @@ def logger_setup(log_file: Path | None, verbosity: int = 0):
    else:
        logger.add(
            sys.__stderr__,  # type: ignore
-            format="[ {time:YYYY-MM-DD HH:mm:ss.SSS} | <level>{level: <8}</level> | {name}:{function}:{line} ] <level>{message}</level>",
+            format="[ {time:HH:mm:ss.SSS} | <level>{level: <8}</level> | {name}:{function}:{line} ] <level>{message}</level>",
            level="DEBUG",
            colorize=True,
            enqueue=True,
@@ -76,7 +76,7 @@ def logger_setup(log_file: Path | None, verbosity: int = 0):
        logger.add(
            log_file,
            format="[ {time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} ] {message}",
-            level="DEBUG" if verbosity > 0 else "INFO",
+            level="INFO",
            colorize=False,
            enqueue=True,
            rotation=lambda _, __: next(rotate_once),
@@ -30,42 +30,30 @@ from exo.utils.pydantic_ext import CamelCaseModel
 # kinda ugly...
 # TODO: load search path from config.toml
 _custom_cards_dir = Path(str(EXO_CUSTOM_MODEL_CARDS_DIR))
-_BUILTIN_CARD_DIRS = [
+CARD_SEARCH_PATH = [
    Path(RESOURCES_DIR) / "inference_model_cards",
    Path(RESOURCES_DIR) / "image_model_cards",
+    _custom_cards_dir,
 ]

 _card_cache: dict[ModelId, "ModelCard"] = {}


-async def _load_cards_from_dir(directory: Path, *, is_custom: bool) -> None:
-    """Load all TOML model cards from a directory into the cache."""
-    async for toml_file in directory.rglob("*.toml"):
-        try:
-            card = await ModelCard.load_from_path(toml_file)
-            if is_custom:
-                card = card.model_copy(update={"is_custom": True})
-            if card.model_id not in _card_cache:
-                _card_cache[card.model_id] = card
-        except (ValidationError, TOMLKitError):
-            pass
-
-
-async def _refresh_card_cache() -> None:
-    for path in _BUILTIN_CARD_DIRS:
-        await _load_cards_from_dir(path, is_custom=False)
-    await _load_cards_from_dir(_custom_cards_dir, is_custom=True)
+async def _refresh_card_cache():
+    for path in CARD_SEARCH_PATH:
+        async for toml_file in path.rglob("*.toml"):
+            try:
+                card = await ModelCard.load_from_path(toml_file)
+                if card.model_id not in _card_cache:
+                    _card_cache[card.model_id] = card
+            except (ValidationError, TOMLKitError):
+                pass


 def _is_image_card(card: "ModelCard") -> bool:
    return any(t in (ModelTask.TextToImage, ModelTask.ImageToImage) for t in card.tasks)


-def get_card(model_id: ModelId) -> "ModelCard | None":
-    """Look up a single model card from the cache by ID."""
-    return _card_cache.get(model_id)
-
-
 async def get_model_cards() -> list["ModelCard"]:
    if len(_card_cache) == 0:
        await _refresh_card_cache()
@@ -104,7 +92,6 @@ class ModelCard(CamelCaseModel):
    capabilities: list[str] = []
    uses_cfg: bool = False
    trust_remote_code: bool = True
-    is_custom: bool = False

    @field_validator("tasks", mode="before")
    @classmethod
@@ -113,7 +100,7 @@ class ModelCard(CamelCaseModel):

    async def save(self, path: Path) -> None:
        async with await open_file(path, "w") as f:
-            py = self.model_dump(exclude_none=True, exclude={"is_custom"})
+            py = self.model_dump(exclude_none=True)
            data = tomlkit.dumps(py)  # pyright: ignore[reportUnknownMemberType]
            await f.write(data)

@@ -135,24 +122,17 @@ class ModelCard(CamelCaseModel):
        if (mc := _card_cache.get(model_id)) is not None:
            return mc

-        mc = await ModelCard.fetch_from_hf(model_id)
-        await mc.save_to_custom_dir()
-        _card_cache[model_id] = mc
-        return mc
+        return await ModelCard.fetch_from_hf(model_id)

    @staticmethod
    async def fetch_from_hf(model_id: ModelId) -> "ModelCard":
-        """Fetches storage size and number of layers for a Hugging Face model, returns Pydantic ModelMeta.
-
-        This is a pure fetch — it does NOT save to disk or update the cache.
-        Persistence is handled by the event-sourcing layer (worker event handler).
-        """
+        """Fetches storage size and number of layers for a Hugging Face model, returns Pydantic ModelMeta."""
        # TODO: failure if files do not exist
        config_data = await fetch_config_data(model_id)
        num_layers = config_data.layer_count
        mem_size_bytes = await fetch_safetensors_size(model_id)

-        return ModelCard(
+        mc = ModelCard(
            model_id=ModelId(model_id),
            storage_size=mem_size_bytes,
            n_layers=num_layers,
@@ -161,13 +141,10 @@ class ModelCard(CamelCaseModel):
            num_key_value_heads=config_data.num_key_value_heads,
            tasks=[ModelTask.TextGeneration],
            trust_remote_code=False,
-            is_custom=True,
        )
-
-
-def add_to_card_cache(card: "ModelCard") -> None:
-    """Add or update a model card in the in-memory cache."""
-    _card_cache[card.model_id] = card
+        await mc.save_to_custom_dir()
+        _card_cache[model_id] = mc
+        return mc


 async def delete_custom_card(model_id: ModelId) -> bool:
@@ -180,6 +157,16 @@ async def delete_custom_card(model_id: ModelId) -> bool:
    return False


+def is_custom_card(model_id: ModelId) -> bool:
+    """Check if a model card exists in the custom cards directory."""
+    import os
+
+    card_path = Path(str(EXO_CUSTOM_MODEL_CARDS_DIR)) / (
+        ModelId(model_id).normalize() + ".toml"
+    )
+    return os.path.isfile(str(card_path))
+
+
 class ConfigData(BaseModel):
    model_config = {"extra": "ignore"}  # Allow unknown fields

@@ -243,10 +230,11 @@ async def fetch_config_data(model_id: ModelId) -> ConfigData:
    """Downloads and parses config.json for a model."""
    from exo.download.download_utils import (
        download_file_with_retry,
-        resolve_model_dir,
+        ensure_models_dir,
    )

-    target_dir = await resolve_model_dir(model_id)
+    target_dir = (await ensure_models_dir()) / model_id.normalize()
+    await aios.makedirs(target_dir, exist_ok=True)
    config_path = await download_file_with_retry(
        model_id,
        "main",
@@ -264,11 +252,12 @@ async def fetch_safetensors_size(model_id: ModelId) -> Memory:
    """Gets model size from safetensors index or falls back to HF API."""
    from exo.download.download_utils import (
        download_file_with_retry,
-        resolve_model_dir,
+        ensure_models_dir,
    )
    from exo.shared.types.worker.downloads import ModelSafetensorsIndex

-    target_dir = await resolve_model_dir(model_id)
+    target_dir = (await ensure_models_dir()) / model_id.normalize()
+    await aios.makedirs(target_dir, exist_ok=True)
    index_path = await download_file_with_retry(
        model_id,
        "main",
@@ -105,9 +105,9 @@ def test_node_id_in_config_dir():


 def test_models_in_data_dir():
-    """Test that default models directory is in the data directory."""
-    # Clear EXO_MODELS_DIRS to test default behavior
-    env = {k: v for k, v in os.environ.items() if k != "EXO_MODELS_DIRS"}
+    """Test that models directory is in the data directory."""
+    # Clear EXO_MODELS_DIR to test default behavior
+    env = {k: v for k, v in os.environ.items() if k != "EXO_MODELS_DIR"}
    with mock.patch.dict(os.environ, env, clear=True):
        import importlib

@@ -115,106 +115,4 @@ def test_models_in_data_dir():

        importlib.reload(constants)

-        assert constants.EXO_DEFAULT_MODELS_DIR.parent == constants.EXO_DATA_HOME
-
-
-def test_default_dir_always_prepended_to_models_dirs():
-    """Test that the default models dir is always the first entry in EXO_MODELS_DIRS."""
-    env = {
-        k: v
-        for k, v in os.environ.items()
-        if k not in ("EXO_MODELS_DIRS", "EXO_MODELS_READ_ONLY_DIRS", "EXO_HOME")
-    }
-    env["EXO_MODELS_DIRS"] = "/tmp/custom-models"
-    with mock.patch.dict(os.environ, env, clear=True):
-        import importlib
-
-        import exo.shared.constants as constants
-
-        importlib.reload(constants)
-
-        assert constants.EXO_MODELS_DIRS[0] == constants.EXO_DEFAULT_MODELS_DIR
-        assert Path("/tmp/custom-models") in constants.EXO_MODELS_DIRS
-
-
-def test_default_models_dir_override():
-    """Test that EXO_DEFAULT_MODELS_DIR can be overridden via env var."""
-    env = {
-        k: v
-        for k, v in os.environ.items()
-        if k
-        not in (
-            "EXO_MODELS_DIRS",
-            "EXO_MODELS_READ_ONLY_DIRS",
-            "EXO_HOME",
-            "EXO_DEFAULT_MODELS_DIR",
-        )
-    }
-    env["EXO_DEFAULT_MODELS_DIR"] = "/Volumes/FastSSD/exo-models"
-    with mock.patch.dict(os.environ, env, clear=True):
-        import importlib
-
-        import exo.shared.constants as constants
-
-        importlib.reload(constants)
-
-        assert Path("/Volumes/FastSSD/exo-models") == constants.EXO_DEFAULT_MODELS_DIR
-        assert constants.EXO_MODELS_DIRS[0] == constants.EXO_DEFAULT_MODELS_DIR
-
-
-def test_default_dir_only_entry_when_env_unset():
-    """Test that EXO_MODELS_DIRS contains only the default when env var is not set."""
-    env = {
-        k: v
-        for k, v in os.environ.items()
-        if k not in ("EXO_MODELS_DIRS", "EXO_MODELS_READ_ONLY_DIRS", "EXO_HOME")
-    }
-    with mock.patch.dict(os.environ, env, clear=True):
-        import importlib
-
-        import exo.shared.constants as constants
-
-        importlib.reload(constants)
-
-        assert constants.EXO_MODELS_DIRS == (constants.EXO_DEFAULT_MODELS_DIR,)
-
-
-def test_overlap_between_dirs_and_read_only_dirs():
-    """Test that a directory in both lists is excluded from writable dirs."""
-    env = {
-        k: v
-        for k, v in os.environ.items()
-        if k not in ("EXO_MODELS_DIRS", "EXO_MODELS_READ_ONLY_DIRS", "EXO_HOME")
-    }
-    env["EXO_MODELS_DIRS"] = "/tmp/shared:/tmp/writable-only"
-    env["EXO_MODELS_READ_ONLY_DIRS"] = "/tmp/shared:/tmp/ro-only"
-    with mock.patch.dict(os.environ, env, clear=True):
-        import importlib
-
-        import exo.shared.constants as constants
-
-        importlib.reload(constants)
-
-        # /tmp/shared should be excluded from writable dirs
-        assert Path("/tmp/shared") not in constants.EXO_MODELS_DIRS
-        assert Path("/tmp/writable-only") in constants.EXO_MODELS_DIRS
-        # /tmp/shared should still be in read-only dirs
-        assert Path("/tmp/shared") in constants.EXO_MODELS_READ_ONLY_DIRS
-        assert Path("/tmp/ro-only") in constants.EXO_MODELS_READ_ONLY_DIRS
-
-
-def test_empty_read_only_dirs_when_unset():
-    """Test that EXO_MODELS_READ_ONLY_DIRS is empty when env var is not set."""
-    env = {
-        k: v
-        for k, v in os.environ.items()
-        if k not in ("EXO_MODELS_DIRS", "EXO_MODELS_READ_ONLY_DIRS", "EXO_HOME")
-    }
-    with mock.patch.dict(os.environ, env, clear=True):
-        import importlib
-
-        import exo.shared.constants as constants
-
-        importlib.reload(constants)
-
-        assert constants.EXO_MODELS_READ_ONLY_DIRS == ()
+        assert constants.EXO_MODELS_DIR.parent == constants.EXO_DATA_HOME
@@ -1,18 +1,18 @@
 from collections.abc import Generator
 from typing import Any, Literal

-from exo.api.types import (
-    FinishReason,
+from exo.shared.models.model_cards import ModelId
+from exo.shared.types.api import (
    GenerationStats,
    ImageGenerationStats,
-    ToolCallItem,
    TopLogprobItem,
    Usage,
 )
-from exo.shared.models.model_cards import ModelId
 from exo.utils.pydantic_ext import TaggedModel

+from .api import FinishReason
 from .common import CommandId
+from .worker.runner_response import ToolCallItem


 class BaseChunk(TaggedModel):
@@ -1,10 +1,10 @@
 from pydantic import Field

-from exo.api.types import (
+from exo.shared.models.model_cards import ModelCard, ModelId
+from exo.shared.types.api import (
    ImageEditsTaskParams,
    ImageGenerationTaskParams,
 )
-from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.common import CommandId, NodeId, SystemId
 from exo.shared.types.text_generation import TextGenerationTaskParams
@@ -81,14 +81,6 @@ class CancelDownload(BaseCommand):
    model_id: ModelId


-class AddCustomModelCard(BaseCommand):
-    model_card: ModelCard
-
-
-class DeleteCustomModelCard(BaseCommand):
-    model_id: ModelId
-
-
 DownloadCommand = StartDownload | DeleteDownload | CancelDownload


@@ -104,8 +96,6 @@ Command = (
    | TaskCancelled
    | TaskFinished
    | SendInputChunk
-    | AddCustomModelCard
-    | DeleteCustomModelCard
 )


@@ -3,10 +3,9 @@ from typing import final

 from pydantic import Field

-from exo.shared.models.model_cards import ModelCard
 from exo.shared.topology import Connection
 from exo.shared.types.chunks import GenerationChunk, InputImageChunk
-from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId, SystemId
+from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -107,14 +106,6 @@ class TopologyEdgeDeleted(BaseEvent):
    conn: Connection


-class CustomModelCardAdded(BaseEvent):
-    model_card: ModelCard
-
-
-class CustomModelCardDeleted(BaseEvent):
-    model_id: ModelId
-
-
@final
 class TraceEventData(FrozenModel):
    name: str
@@ -156,8 +147,6 @@ Event = (
    | TopologyEdgeDeleted
    | TracesCollected
    | TracesMerged
-    | CustomModelCardAdded
-    | CustomModelCardDeleted
 )


@@ -2,7 +2,7 @@ from enum import Enum

 from pydantic import Field

-from exo.api.types import (
+from exo.shared.types.api import (
    ImageEditsTaskParams,
    ImageGenerationTaskParams,
 )
@@ -1,7 +1,7 @@
 from collections.abc import Generator
 from typing import Any, Literal

-from exo.api.types import (
+from exo.shared.types.api import (
    FinishReason,
    GenerationStats,
    ImageGenerationStats,
@@ -10,10 +10,11 @@ from typing import Self, cast
 import anyio
 from anyio import fail_after, open_process, to_thread
 from anyio.streams.buffered import BufferedByteReceiveStream
+from anyio.streams.text import TextReceiveStream
 from loguru import logger
 from pydantic import ValidationError

-from exo.shared.constants import EXO_CONFIG_FILE, EXO_DEFAULT_MODELS_DIR
+from exo.shared.constants import EXO_CONFIG_FILE, EXO_MODELS_DIR
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
    DiskUsage,
@@ -287,7 +288,7 @@ class ThunderboltBridgeInfo(TaggedModel):
                )
            )
        except Exception as e:
-            logger.opt(exception=e).warning("Failed to gather Thunderbolt Bridge info")
+            logger.warning(f"Failed to gather Thunderbolt Bridge info: {e}")
            return None


@@ -328,7 +329,7 @@ class NodeDiskUsage(TaggedModel):
    async def gather(cls) -> Self:
        return cls(
            disk_usage=await to_thread.run_sync(
-                DiskUsage.from_path, EXO_DEFAULT_MODELS_DIR
+                lambda: DiskUsage.from_path(EXO_MODELS_DIR)
            )
        )

@@ -371,58 +372,36 @@ GatheredInfo = (
@dataclass
 class InfoGatherer:
    info_sender: Sender[GatheredInfo]
+    interface_watcher_interval: float | None = 10
+    misc_poll_interval: float | None = 60
+    system_profiler_interval: float | None = 5 if IS_DARWIN else None
+    memory_poll_rate: float | None = None if IS_DARWIN else 1
+    macmon_interval: float | None = 1 if IS_DARWIN else None
+    thunderbolt_bridge_poll_interval: float | None = 10 if IS_DARWIN else None
+    static_info_poll_interval: float | None = 60
+    rdma_ctl_poll_interval: float | None = 10 if IS_DARWIN else None
+    disk_poll_interval: float | None = 30
    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
-    _psutil_enabled: bool = field(init=False, default=False)
-
-    async def _can_read_macmon_metrics(self, macmon_path: str) -> bool:
-        try:
-            with fail_after(5):
-                proc = await anyio.run_process(
-                    [macmon_path, "pipe", "--samples", "1", "--interval", "100"],
-                    check=False,
-                )
-        except Exception as e:
-            logger.opt(exception=e).warning(
-                f"Failed to validate macmon at {macmon_path}"
-            )
-            return False
-
-        if proc.returncode != 0:
-            stderr = proc.stderr.decode("utf-8", errors="replace").strip()
-            logger.warning(
-                f"macmon preflight failed with return code {proc.returncode}: "
-                f"{stderr or 'no stderr'}"
-            )
-            return False
-
-        stdout = proc.stdout.decode("utf-8", errors="replace").strip()
-        if not stdout:
-            logger.warning("macmon preflight returned no metrics")
-            return False
-
-        try:
-            MacmonMetrics.from_raw_json(stdout.splitlines()[0])
-        except ValidationError as e:
-            logger.opt(exception=e).warning(
-                "macmon preflight returned unexpected metrics JSON"
-            )
-            return False
-
-        return True

    async def run(self):
        async with self._tg as tg:
            if IS_DARWIN:
-                tg.start_soon(self._monitor_macmon, 1)
-                tg.start_soon(self._monitor_system_profiler_thunderbolt_data, 5)
-                tg.start_soon(self._monitor_thunderbolt_bridge_status, 10)
-                tg.start_soon(self._monitor_rdma_ctl_status, 10)
-            if not IS_DARWIN:
-                tg.start_soon(self._monitor_memory_usage, 1)
-            tg.start_soon(self._watch_system_info, 10)
-            tg.start_soon(self._monitor_misc, 60)
-            tg.start_soon(self._monitor_static_info, 60)
-            tg.start_soon(self._monitor_disk_usage, 30)
+                if (macmon_path := shutil.which("macmon")) is not None:
+                    tg.start_soon(self._monitor_macmon, macmon_path)
+                else:
+                    # macmon not installed — fall back to psutil for memory
+                    logger.warning(
+                        "macmon not found, falling back to psutil for memory monitoring"
+                    )
+                    self.memory_poll_rate = 1
+                tg.start_soon(self._monitor_system_profiler_thunderbolt_data)
+                tg.start_soon(self._monitor_thunderbolt_bridge_status)
+                tg.start_soon(self._monitor_rdma_ctl_status)
+            tg.start_soon(self._watch_system_info)
+            tg.start_soon(self._monitor_memory_usage)
+            tg.start_soon(self._monitor_misc)
+            tg.start_soon(self._monitor_static_info)
+            tg.start_soon(self._monitor_disk_usage)

            nc = await NodeConfig.gather()
            if nc is not None:
@@ -431,27 +410,32 @@ class InfoGatherer:
    def shutdown(self):
        self._tg.cancel_tasks()

-    async def _monitor_static_info(self, static_info_poll_interval: float):
+    async def _monitor_static_info(self):
+        if self.static_info_poll_interval is None:
+            return
        while True:
            try:
                with fail_after(30):
                    await self.info_sender.send(await StaticNodeInformation.gather())
            except Exception as e:
-                logger.opt(exception=e).warning("Error gathering static node info")
-            await anyio.sleep(static_info_poll_interval)
+                logger.warning(f"Error gathering static node info: {e}")
+            await anyio.sleep(self.static_info_poll_interval)

-    async def _monitor_misc(self, misc_poll_interval: float):
+    async def _monitor_misc(self):
+        if self.misc_poll_interval is None:
+            return
        while True:
            try:
                with fail_after(10):
                    await self.info_sender.send(await MiscData.gather())
            except Exception as e:
-                logger.opt(exception=e).warning("Error gathering misc data")
-            await anyio.sleep(misc_poll_interval)
+                logger.warning(f"Error gathering misc data: {e}")
+            await anyio.sleep(self.misc_poll_interval)
+
+    async def _monitor_system_profiler_thunderbolt_data(self):
+        if self.system_profiler_interval is None:
+            return

-    async def _monitor_system_profiler_thunderbolt_data(
-        self, system_profiler_interval: float
-    ):
        while True:
            try:
                with fail_after(30):
@@ -472,41 +456,42 @@ class InfoGatherer:
                    conns = [it for i in data if (it := i.conn()) is not None]
                    await self.info_sender.send(MacThunderboltConnections(conns=conns))
            except Exception as e:
-                logger.opt(exception=e).warning("Error gathering Thunderbolt data")
-            await anyio.sleep(system_profiler_interval)
+                logger.warning(f"Error gathering Thunderbolt data: {e}")
+            await anyio.sleep(self.system_profiler_interval)

-    async def _monitor_memory_usage(self, memory_poll_rate: float):
-        if self._psutil_enabled:
-            return
-        self._psutil_enabled = True
+    async def _monitor_memory_usage(self):
        override_memory_env = os.getenv("OVERRIDE_MEMORY_MB")
        override_memory: int | None = (
            Memory.from_mb(int(override_memory_env)).in_bytes
            if override_memory_env
            else None
        )
+        if self.memory_poll_rate is None:
+            return
        while True:
            try:
                await self.info_sender.send(
                    MemoryUsage.from_psutil(override_memory=override_memory)
                )
            except Exception as e:
-                logger.opt(exception=e).warning("Error gathering memory usage")
-            await anyio.sleep(memory_poll_rate)
+                logger.warning(f"Error gathering memory usage: {e}")
+            await anyio.sleep(self.memory_poll_rate)

-    async def _watch_system_info(self, interface_watcher_interval: float):
+    async def _watch_system_info(self):
+        if self.interface_watcher_interval is None:
+            return
        while True:
            try:
                with fail_after(10):
                    nics = await get_network_interfaces()
                    await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
            except Exception as e:
-                logger.opt(exception=e).warning("Error gathering network interfaces")
-            await anyio.sleep(interface_watcher_interval)
+                logger.warning(f"Error gathering network interfaces: {e}")
+            await anyio.sleep(self.interface_watcher_interval)

-    async def _monitor_thunderbolt_bridge_status(
-        self, thunderbolt_bridge_poll_interval: float
-    ):
+    async def _monitor_thunderbolt_bridge_status(self):
+        if self.thunderbolt_bridge_poll_interval is None:
+            return
        while True:
            try:
                with fail_after(30):
@@ -514,49 +499,39 @@ class InfoGatherer:
                    if curr is not None:
                        await self.info_sender.send(curr)
            except Exception as e:
-                logger.opt(exception=e).warning(
-                    "Error gathering Thunderbolt Bridge status"
-                )
-            await anyio.sleep(thunderbolt_bridge_poll_interval)
+                logger.warning(f"Error gathering Thunderbolt Bridge status: {e}")
+            await anyio.sleep(self.thunderbolt_bridge_poll_interval)

-    async def _monitor_rdma_ctl_status(self, rdma_ctl_poll_interval: float):
+    async def _monitor_rdma_ctl_status(self):
+        if self.rdma_ctl_poll_interval is None:
+            return
        while True:
            try:
                curr = await RdmaCtlStatus.gather()
                if curr is not None:
                    await self.info_sender.send(curr)
            except Exception as e:
-                logger.opt(exception=e).warning("Error gathering RDMA ctl status")
-            await anyio.sleep(rdma_ctl_poll_interval)
+                logger.warning(f"Error gathering RDMA ctl status: {e}")
+            await anyio.sleep(self.rdma_ctl_poll_interval)

-    async def _monitor_disk_usage(self, disk_poll_interval: float):
+    async def _monitor_disk_usage(self):
+        if self.disk_poll_interval is None:
+            return
        while True:
            try:
                with fail_after(5):
                    await self.info_sender.send(await NodeDiskUsage.gather())
            except Exception as e:
-                logger.opt(exception=e).warning("Error gathering disk usage")
-            await anyio.sleep(disk_poll_interval)
+                logger.warning(f"Error gathering disk usage: {e}")
+            await anyio.sleep(self.disk_poll_interval)

-    async def _monitor_macmon(self, macmon_interval: float):
-        if (
-            macmon_path := os.getenv("EXO_MACMON_PATH") or shutil.which("macmon")
-        ) is None:
-            logger.warning(
-                "macmon not found, falling back to psutil for memory monitoring"
-            )
-            self._tg.start_soon(self._monitor_memory_usage, 1)
-            return
-        if not await self._can_read_macmon_metrics(macmon_path):
-            logger.warning(
-                f"macmon at {macmon_path} is unusable, falling back to psutil memory monitoring"
-            )
-            self._tg.start_soon(self._monitor_memory_usage, 1)
+    async def _monitor_macmon(self, macmon_path: str):
+        if self.macmon_interval is None:
            return
        # macmon pipe --interval [interval in ms]
        # Timeout: if macmon produces no output for this many seconds, restart it.
        # macmon writes every macmon_interval seconds, so 10x that is generous.
-        read_timeout = max(macmon_interval * 10, 30)
+        read_timeout = max(self.macmon_interval * 10, 30)
        while True:
            try:
                async with await open_process(
@@ -564,26 +539,21 @@ class InfoGatherer:
                        macmon_path,
                        "pipe",
                        "--interval",
-                        str(macmon_interval * 1000),
+                        str(self.macmon_interval * 1000),
                    ]
                ) as p:
                    if not p.stdout:
                        logger.critical("MacMon closed stdout")
                        return
-                    stream = BufferedByteReceiveStream(p.stdout)
+                    stream = TextReceiveStream(BufferedByteReceiveStream(p.stdout))
                    while True:
                        with fail_after(read_timeout):
-                            data = await stream.receive_until(
-                                delimiter=b"\n", max_bytes=8 * 1024
-                            )
-                            text = data.decode("utf-8", errors="replace").strip()
-                            metrics = MacmonMetrics.from_raw_json(text)
-                        await self.info_sender.send(metrics)
+                            text = await stream.receive()
+                        await self.info_sender.send(MacmonMetrics.from_raw_json(text))
            except TimeoutError:
                logger.warning(
                    f"MacMon produced no output for {read_timeout}s, restarting"
                )
-                self._tg.start_soon(self._monitor_memory_usage, 1)
            except CalledProcessError as e:
                stderr_msg = "no stderr"
                stderr_output = cast(bytes | str | None, e.stderr)
@@ -596,8 +566,6 @@ class InfoGatherer:
                logger.warning(
                    f"MacMon failed with return code {e.returncode}: {stderr_msg}"
                )
-                self._tg.start_soon(self._monitor_memory_usage, 1)
            except Exception as e:
-                logger.opt(exception=e).warning("Error in macmon monitor")
-                self._tg.start_soon(self._monitor_memory_usage, 1)
-            await anyio.sleep(macmon_interval)
+                logger.warning(f"Error in macmon monitor: {e}")
+            await anyio.sleep(self.macmon_interval)
@@ -5,7 +5,7 @@ from typing import final

 import anyio

-from exo.api.types import NodePowerStats, PowerUsage
+from exo.shared.types.api import NodePowerStats, PowerUsage
 from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import SystemPerformanceProfile

@@ -3,7 +3,7 @@ from collections.abc import Mapping
 import anyio
 import pytest

-from exo.api.types import PowerUsage
+from exo.shared.types.api import PowerUsage
 from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import SystemPerformanceProfile
 from exo.utils.power_sampler import PowerSampler
@@ -1,4 +1,4 @@
-from collections.abc import Callable, Generator
+from collections.abc import Generator
 from pathlib import Path
 from typing import Any, Literal, Optional

@@ -6,8 +6,8 @@ import mlx.core as mx
 from mflux.models.common.config.config import Config
 from PIL import Image

-from exo.api.types import AdvancedImageParams
 from exo.download.download_utils import build_model_path
+from exo.shared.types.api import AdvancedImageParams
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.shards import CfgShardMetadata, PipelineShardMetadata
 from exo.worker.engines.image.config import ImageModelConfig
@@ -116,7 +116,6 @@ class DistributedImageModel:
        image_path: Path | None = None,
        partial_images: int = 0,
        advanced_params: AdvancedImageParams | None = None,
-        cancel_checker: Callable[[], bool] | None = None,
    ) -> Generator[Image.Image | tuple[Image.Image, int, int], None, None]:
        if (
            advanced_params is not None
@@ -164,7 +163,6 @@ class DistributedImageModel:
            guidance_override=guidance_override,
            negative_prompt=negative_prompt,
            num_sync_steps=num_sync_steps,
-            cancel_checker=cancel_checker,
        ):
            if isinstance(result, tuple):
                # Partial image: (GeneratedImage, partial_index, total_partials)
@@ -3,14 +3,13 @@ import io
 import random
 import tempfile
 import time
-from collections.abc import Callable
 from pathlib import Path
 from typing import Generator, Literal

 import mlx.core as mx
 from PIL import Image

-from exo.api.types import (
+from exo.shared.types.api import (
    AdvancedImageParams,
    ImageEditsTaskParams,
    ImageGenerationStats,
@@ -70,7 +69,6 @@ def warmup_image_generator(model: DistributedImageModel) -> Image.Image | None:
 def generate_image(
    model: DistributedImageModel,
    task: ImageGenerationTaskParams | ImageEditsTaskParams,
-    cancel_checker: Callable[[], bool] | None = None,
 ) -> Generator[ImageGenerationResponse | PartialImageResponse, None, None]:
    """Generate image(s), optionally yielding partial results.

@@ -129,7 +127,6 @@ def generate_image(
                image_path=image_path,
                partial_images=partial_images,
                advanced_params=advanced_params,
-                cancel_checker=cancel_checker,
            ):
                if isinstance(result, tuple):
                    # Partial image: (Image, partial_index, total_partials)
@@ -1,4 +1,4 @@
-from collections.abc import Callable, Iterator
+from collections.abc import Iterator
 from dataclasses import dataclass
 from math import ceil
 from typing import Any, Optional, final
@@ -100,8 +100,6 @@ class DiffusionRunner:
        self.total_layers = config.total_blocks

        self._guidance_override: float | None = None
-        self._cancel_checker: Callable[[], bool] | None = None
-        self._cancelling: bool = False

        self._compute_assigned_blocks()

@@ -242,43 +240,6 @@ class DiffusionRunner:
    def is_distributed(self) -> bool:
        return self.group is not None

-    def _is_sentinel(self, tensor: mx.array) -> bool:
-        return bool(mx.all(mx.isnan(tensor)).item())
-
-    def _check_cancellation(self) -> None:
-        if self._cancelling:
-            return
-        if (
-            self.is_first_stage
-            and self._cancel_checker is not None
-            and self._cancel_checker()
-        ):
-            self._cancelling = True
-
-    def _send(self, data: mx.array, dst: int) -> mx.array:
-        assert self.group is not None
-        if self._cancelling:
-            data = mx.full(data.shape, float("nan"), dtype=data.dtype)
-        return mx.distributed.send(data, dst, group=self.group)
-
-    def _recv_and_check(self, result: mx.array) -> mx.array:
-        mx.eval(result)
-        if self._is_sentinel(result):
-            self._cancelling = True
-        return result
-
-    def _recv(self, shape: tuple[int, ...], dtype: mx.Dtype, src: int) -> mx.array:
-        assert self.group is not None
-        return self._recv_and_check(
-            mx.distributed.recv(shape, dtype, src, group=self.group)
-        )
-
-    def _recv_like(self, template: mx.array, src: int) -> mx.array:
-        assert self.group is not None
-        return self._recv_and_check(
-            mx.distributed.recv_like(template, src, group=self.group)
-        )
-
    def _get_effective_guidance_scale(self) -> float | None:
        if self._guidance_override is not None:
            return self._guidance_override
@@ -352,13 +313,19 @@ class DiffusionRunner:
        assert self.cfg_peer_rank is not None

        if is_positive:
-            noise = self._send(noise, self.cfg_peer_rank)
+            noise = mx.distributed.send(noise, self.cfg_peer_rank, group=self.group)
            mx.async_eval(noise)
-            noise_neg = self._recv_like(noise, src=self.cfg_peer_rank)
+            noise_neg = mx.distributed.recv_like(
+                noise, self.cfg_peer_rank, group=self.group
+            )
+            mx.eval(noise_neg)
            noise_pos = noise
        else:
-            noise_pos = self._recv_like(noise, src=self.cfg_peer_rank)
-            noise = self._send(noise, self.cfg_peer_rank)
+            noise_pos = mx.distributed.recv_like(
+                noise, self.cfg_peer_rank, group=self.group
+            )
+            mx.eval(noise_pos)
+            noise = mx.distributed.send(noise, self.cfg_peer_rank, group=self.group)
            mx.async_eval(noise)
            noise_neg = noise

@@ -465,7 +432,6 @@ class DiffusionRunner:
        guidance_override: float | None = None,
        negative_prompt: str | None = None,
        num_sync_steps: int = 1,
-        cancel_checker: Callable[[], bool] | None = None,
    ):
        """Primary entry point for image generation.

@@ -488,8 +454,6 @@ class DiffusionRunner:
            Final GeneratedImage
        """
        self._guidance_override = guidance_override
-        self._cancel_checker = cancel_checker
-        self._cancelling = False
        latents = self.adapter.create_latents(seed, runtime_config)
        prompt_data = self.adapter.encode_prompt(prompt, negative_prompt)

@@ -531,7 +495,7 @@ class DiffusionRunner:
            except StopIteration as e:
                latents = e.value  # pyright: ignore[reportAny]

-        if self.is_last_stage and not self._cancelling:
+        if self.is_last_stage:
            yield self.adapter.decode_latents(latents, runtime_config, seed, prompt)  # pyright: ignore[reportAny]

    def _run_diffusion_loop(
@@ -560,12 +524,7 @@ class DiffusionRunner:
            latents=latents,
        )

-        t = -1  # default if time_steps is empty; drain condition uses t
        for t in time_steps:
-            self._check_cancellation()
-            if self._cancelling and self.group is None:
-                break
-
            try:
                latents = self._diffusion_step(
                    t=t,
@@ -583,7 +542,7 @@ class DiffusionRunner:

                mx.eval(latents)

-                if t in capture_steps and self.is_last_stage and not self._cancelling:
+                if t in capture_steps and self.is_last_stage:
                    yield (latents, t)

            except KeyboardInterrupt:  # noqa: PERF203
@@ -592,24 +551,6 @@ class DiffusionRunner:
                    f"Stopping image generation at step {t + 1}/{len(time_steps)}"
                ) from None

-            if self._cancelling:
-                break
-
-        # Drain pending ring recvs after cancellation during async steps.
-        # The last stage sent patches during the final completed step, but
-        # the first stage will never enter the next step to recv them.
-        if (
-            self._cancelling
-            and self.is_first_stage
-            and not self.is_last_stage
-            and self.group is not None
-            and t >= runtime_config.init_time_step + num_sync_steps
-            and t != runtime_config.num_inference_steps - 1
-        ):
-            patch_latents_drain, _ = self._create_patches(latents, runtime_config)
-            for patch in patch_latents_drain:
-                self._recv_like(patch, src=self.last_pipeline_rank)
-
        ctx.after_loop(latents=latents)  # pyright: ignore[reportAny]

        return latents
@@ -836,16 +777,19 @@ class DiffusionRunner:
                    rank=self.rank,
                    category="comms",
                ):
-                    hidden_states = self._recv(
+                    hidden_states = mx.distributed.recv(
                        (batch_size, num_img_tokens, hidden_dim),
                        dtype,
                        self.prev_pipeline_rank,
+                        group=self.group,
                    )
-                    encoder_hidden_states = self._recv(
+                    encoder_hidden_states = mx.distributed.recv(
                        (batch_size, text_seq_len, hidden_dim),
                        dtype,
                        self.prev_pipeline_rank,
+                        group=self.group,
                    )
+                    mx.eval(hidden_states, encoder_hidden_states)

            assert self.joint_block_wrappers is not None
            assert encoder_hidden_states is not None
@@ -881,7 +825,9 @@ class DiffusionRunner:
                    rank=self.rank,
                    category="comms",
                ):
-                    concatenated = self._send(concatenated, self.next_pipeline_rank)
+                    concatenated = mx.distributed.send(
+                        concatenated, self.next_pipeline_rank, group=self.group
+                    )
                    mx.async_eval(concatenated)

        elif self.has_joint_blocks and not self.is_last_stage:
@@ -892,9 +838,11 @@ class DiffusionRunner:
                rank=self.rank,
                category="comms",
            ):
-                hidden_states = self._send(hidden_states, self.next_pipeline_rank)
-                encoder_hidden_states = self._send(
-                    encoder_hidden_states, self.next_pipeline_rank
+                hidden_states = mx.distributed.send(
+                    hidden_states, self.next_pipeline_rank, group=self.group
+                )
+                encoder_hidden_states = mx.distributed.send(
+                    encoder_hidden_states, self.next_pipeline_rank, group=self.group
                )
                mx.async_eval(hidden_states, encoder_hidden_states)

@@ -906,11 +854,13 @@ class DiffusionRunner:
                    rank=self.rank,
                    category="comms",
                ):
-                    hidden_states = self._recv(
+                    hidden_states = mx.distributed.recv(
                        (batch_size, text_seq_len + num_img_tokens, hidden_dim),
                        dtype,
                        self.prev_pipeline_rank,
+                        group=self.group,
                    )
+                    mx.eval(hidden_states)

            assert self.single_block_wrappers is not None
            with trace(
@@ -936,7 +886,9 @@ class DiffusionRunner:
                    rank=self.rank,
                    category="comms",
                ):
-                    hidden_states = self._send(hidden_states, self.next_pipeline_rank)
+                    hidden_states = mx.distributed.send(
+                        hidden_states, self.next_pipeline_rank, group=self.group
+                    )
                    mx.async_eval(hidden_states)

        hidden_states = hidden_states[:, text_seq_len:, ...]
@@ -1009,11 +961,16 @@ class DiffusionRunner:
            )

            if not self.is_first_stage:
-                hidden_states = self._send(hidden_states, self.first_pipeline_rank)
+                hidden_states = mx.distributed.send(
+                    hidden_states, self.first_pipeline_rank, group=self.group
+                )
                mx.async_eval(hidden_states)

        elif self.is_first_stage:
-            hidden_states = self._recv_like(prev_latents, src=self.last_pipeline_rank)
+            hidden_states = mx.distributed.recv_like(
+                prev_latents, src=self.last_pipeline_rank, group=self.group
+            )
+            mx.eval(hidden_states)

        else:
            hidden_states = prev_latents
@@ -1049,7 +1006,10 @@ class DiffusionRunner:
                    rank=self.rank,
                    category="comms",
                ):
-                    patch = self._recv_like(patch, src=self.last_pipeline_rank)
+                    patch = mx.distributed.recv_like(
+                        patch, src=self.last_pipeline_rank, group=self.group
+                    )
+                    mx.eval(patch)

            results: list[tuple[bool, mx.array]] = []

@@ -1106,9 +1066,10 @@ class DiffusionRunner:
                        rank=self.rank,
                        category="comms",
                    ):
-                        patch_latents[patch_idx] = self._send(
+                        patch_latents[patch_idx] = mx.distributed.send(
                            patch_latents[patch_idx],
                            self.first_pipeline_rank,
+                            group=self.group,
                        )
                        mx.async_eval(patch_latents[patch_idx])

@@ -1155,11 +1116,13 @@ class DiffusionRunner:
                    rank=self.rank,
                    category="comms",
                ):
-                    patch = self._recv(
+                    patch = mx.distributed.recv(
                        (batch_size, patch_len, hidden_dim),
                        patch.dtype,
                        self.prev_pipeline_rank,
+                        group=self.group,
                    )
+                    mx.eval(patch)

                if patch_idx == 0:
                    with trace(
@@ -1167,11 +1130,13 @@ class DiffusionRunner:
                        rank=self.rank,
                        category="comms",
                    ):
-                        encoder_hidden_states = self._recv(
+                        encoder_hidden_states = mx.distributed.recv(
                            (batch_size, text_seq_len, hidden_dim),
                            patch.dtype,
                            self.prev_pipeline_rank,
+                            group=self.group,
                        )
+                        mx.eval(encoder_hidden_states)

            if self.is_first_stage:
                patch, encoder_hidden_states = self.adapter.compute_embeddings(
@@ -1210,7 +1175,9 @@ class DiffusionRunner:
                    rank=self.rank,
                    category="comms",
                ):
-                    patch_concat = self._send(patch_concat, self.next_pipeline_rank)
+                    patch_concat = mx.distributed.send(
+                        patch_concat, self.next_pipeline_rank, group=self.group
+                    )
                    mx.async_eval(patch_concat)

        elif self.has_joint_blocks and not self.is_last_stage:
@@ -1220,7 +1187,9 @@ class DiffusionRunner:
                rank=self.rank,
                category="comms",
            ):
-                patch = self._send(patch, self.next_pipeline_rank)
+                patch = mx.distributed.send(
+                    patch, self.next_pipeline_rank, group=self.group
+                )
                mx.async_eval(patch)

            if patch_idx == 0:
@@ -1230,8 +1199,8 @@ class DiffusionRunner:
                    rank=self.rank,
                    category="comms",
                ):
-                    encoder_hidden_states = self._send(
-                        encoder_hidden_states, self.next_pipeline_rank
+                    encoder_hidden_states = mx.distributed.send(
+                        encoder_hidden_states, self.next_pipeline_rank, group=self.group
                    )
                    mx.async_eval(encoder_hidden_states)

@@ -1244,11 +1213,13 @@ class DiffusionRunner:
                    rank=self.rank,
                    category="comms",
                ):
-                    patch = self._recv(
+                    patch = mx.distributed.recv(
                        (batch_size, text_seq_len + patch_len, hidden_dim),
                        patch.dtype,
                        self.prev_pipeline_rank,
+                        group=self.group,
                    )
+                    mx.eval(patch)

            assert self.single_block_wrappers is not None
            with trace(
@@ -1274,7 +1245,9 @@ class DiffusionRunner:
                    rank=self.rank,
                    category="comms",
                ):
-                    patch = self._send(patch, self.next_pipeline_rank)
+                    patch = mx.distributed.send(
+                        patch, self.next_pipeline_rank, group=self.group
+                    )
                    mx.async_eval(patch)

        noise: mx.array | None = None
@@ -57,8 +57,8 @@ from mlx_lm.models.step3p5 import Model as Step35Model
 from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP
 from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel

+from exo.shared.logging import logger
 from exo.shared.types.worker.shards import PipelineShardMetadata
-from exo.worker.runner.bootstrap import logger

 if TYPE_CHECKING:
    from mlx_lm.models.cache import Cache
@@ -4,7 +4,7 @@ from typing import Any

 from mlx_lm.chat_templates import deepseek_v32

-from exo.api.types import ToolCallItem
+from exo.shared.types.api import ToolCallItem

 BOS_TOKEN: str = deepseek_v32.bos_token
 EOS_TOKEN: str = deepseek_v32.eos_token
@@ -15,28 +15,7 @@ USER_TOKEN = "<\uff5cUser\uff5c>"
 ASSISTANT_TOKEN = "<\uff5cAssistant\uff5c>"
 TOOL_CALLS_START = f"<{DSML_TOKEN}function_calls>"
 TOOL_CALLS_END = f"</{DSML_TOKEN}function_calls>"
-_ORPHAN_THINK_END = ASSISTANT_TOKEN + THINKING_END
-_FIXED_THINK_BLOCK = ASSISTANT_TOKEN + THINKING_START + "\n" + THINKING_END
-
-
-def encode_messages(
-    messages: list[dict[str, Any]],
-    thinking_mode: str = "thinking",
-    context: list[dict[str, Any]] | None = None,
-    drop_thinking: bool = True,
-    add_default_bos_token: bool = True,
-    tools: Any = None,  # pyright: ignore[reportAny]
-) -> str:
-    prompt: str = deepseek_v32.encode_messages(
-        messages,
-        thinking_mode=thinking_mode,
-        context=context,
-        drop_thinking=drop_thinking,
-        add_default_bos_token=add_default_bos_token,
-        tools=tools,
-    )
-    return prompt.replace(_ORPHAN_THINK_END, _FIXED_THINK_BLOCK)
-
+encode_messages = deepseek_v32.encode_messages

 _INVOKE_PATTERN = re.compile(
    rf"<{re.escape(DSML_TOKEN)}invoke\s+name=\"([^\"]+)\">"
@@ -1,3 +1,4 @@
+import os
 import time
 from dataclasses import dataclass, field
 from typing import Callable, cast
@@ -6,14 +7,11 @@ import mlx.core as mx
 from mlx_lm.generate import (
    BatchGenerator as MlxBatchGenerator,
 )
-from mlx_lm.generate import (
-    generation_stream,
-)
 from mlx_lm.models.cache import RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper

-from exo.api.types import (
+from exo.shared.types.api import (
    CompletionTokensDetails,
    FinishReason,
    GenerationStats,
@@ -66,7 +64,6 @@ class _EngineTask:
    potential_stop_sequence_text: str = ""
    completion_tokens: int = 0
    generation_start_time: float = 0.0
-    generation_time_at_start: float = 0.0
    in_thinking: bool = False
    reasoning_tokens: int = 0
    prefill_tps: float = 0.0
@@ -79,23 +76,183 @@ class ExoBatchGenerator:
    group: mx.distributed.Group | None
    kv_prefix_cache: KVPrefixCache | None

-    _mlx_gen: MlxBatchGenerator = field(init=False)
+    _exo_gen: MlxBatchGenerator = field(init=False)
    _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False)

    def __post_init__(self) -> None:
-        self._mlx_gen = MlxBatchGenerator(
-            model=self.model,
-            stop_tokens=set(eos_ids_from_tokenizer(self.tokenizer)),
-            prefill_step_size=4096,
+        use_speculative = os.environ.get("EXO_SPECULATIVE", "0") == "1"
+        stop_tokens = set(eos_ids_from_tokenizer(self.tokenizer))
+
+        if use_speculative:
+            try:
+                from exo.worker.engines.mlx.speculative.mtp_module import MTPPredictor
+                from exo.worker.engines.mlx.speculative.mtp_batch_generator import MTPBatchGenerator
+
+                mtp_weights = self._resolve_mtp_weights()
+                gamma = int(os.environ.get("EXO_SPECULATIVE_GAMMA", "2"))
+
+                if mtp_weights:
+                    mtp = MTPPredictor(self.model, mtp_weights, quantize=False)
+                    temp = float(os.environ.get("EXO_SPECULATIVE_TEMP", "0.7"))
+                    alpha = float(os.environ.get("EXO_SPECULATIVE_ALPHA", "1.0"))
+                    self._exo_gen = MTPBatchGenerator(
+                        model=self.model,
+                        mtp_predictor=mtp,
+                        gamma=gamma,
+                        temp=temp,
+                        alpha=alpha,
+                        stop_tokens=stop_tokens,
+                        prefill_step_size=4096,
+                    )
+                    logger.info(f"MTP speculative decoding enabled (γ={gamma}, T={temp})")
+                    self.warmup_speculative(self.model, self.tokenizer)
+                else:
+                    logger.warning("EXO_SPECULATIVE=1 but could not find MTP weights. Falling back to standard generation.")
+                    self._exo_gen = MlxBatchGenerator(
+                        model=self.model,
+                        stop_tokens=stop_tokens,
+                        prefill_step_size=4096,
+                    )
+            except Exception as e:
+                logger.warning(f"Failed to initialize MTP speculative decoding: {e}. Falling back to standard generation.")
+                self._exo_gen = MlxBatchGenerator(
+                    model=self.model,
+                    stop_tokens=stop_tokens,
+                    prefill_step_size=4096,
+                )
+        else:
+            self._exo_gen = MlxBatchGenerator(
+                model=self.model,
+                stop_tokens=stop_tokens,
+                prefill_step_size=4096,
+            )
+
+    def _resolve_mtp_weights(self) -> str | None:
+        """Find MTP weights: explicit path, explicit HF model, or auto-extract."""
+        # 1. Explicit path
+        explicit_path = os.environ.get("EXO_MTP_WEIGHTS", "")
+        if explicit_path and os.path.exists(explicit_path):
+            return explicit_path
+
+        # 2. Explicit HF model repo containing MTP weights
+        mtp_model = os.environ.get("EXO_MTP_MODEL", "")
+
+        # 3. Auto-detect: if no EXO_MTP_MODEL set, try to infer from model config
+        if not mtp_model:
+            try:
+                inner = getattr(self.model, 'model', None) or self.model.language_model.model
+                args = getattr(inner, 'args', None)
+                if args and getattr(args, 'mtp_num_hidden_layers', 0) > 0:
+                    model_type = getattr(args, 'model_type', '')
+                    if 'qwen3_5' in model_type or 'qwen3.5' in str(type(self.model).__module__):
+                        # Default pairing for Qwen3.5-27B
+                        mtp_model = "Qwen/Qwen3.5-27B"
+                        logger.info(f"Auto-detected MTP model: {mtp_model}")
+            except Exception:
+                pass
+
+        if not mtp_model:
+            return None
+
+        # Download and extract MTP weights from HF repo
+        try:
+            return self._extract_mtp_from_hf(mtp_model)
+        except Exception as e:
+            logger.warning(f"Failed to extract MTP weights from {mtp_model}: {e}")
+            return None
+
+    def _extract_mtp_from_hf(self, repo_id: str) -> str:
+        """Download MTP tensors from HF repo and cache as a single safetensors file."""
+        import hashlib
+        from pathlib import Path
+        from huggingface_hub import snapshot_download
+        from safetensors.torch import load_file, save_file
+
+        cache_dir = Path.home() / ".cache" / "exo" / "mtp_weights"
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        cache_key = hashlib.md5(repo_id.encode()).hexdigest()[:12]
+        cached_path = cache_dir / f"mtp_{cache_key}.safetensors"
+
+        if cached_path.exists():
+            logger.info(f"Using cached MTP weights: {cached_path}")
+            return str(cached_path)
+
+        logger.info(f"Downloading MTP weights from {repo_id}...")
+        model_dir = snapshot_download(
+            repo_id,
+            allow_patterns=["*.safetensors", "*.json"],
        )
-        self._mlx_gen._needs_topk = False  # pyright: ignore[reportAttributeAccessIssue]
+
+        # Extract MTP tensors from all safetensors files
+        mtp_tensors = {}
+        model_path = Path(model_dir)
+        for sf_file in sorted(model_path.glob("*.safetensors")):
+            tensors = load_file(str(sf_file))
+            for k, v in tensors.items():
+                if k.startswith("model.mtp."):
+                    # Strip "model." prefix to match our MTPPredictor format
+                    clean_key = k[len("model."):]
+                    mtp_tensors[clean_key] = v
+
+        if not mtp_tensors:
+            raise ValueError(f"No MTP tensors found in {repo_id}")
+
+        save_file(mtp_tensors, str(cached_path))
+        logger.info(f"Extracted {len(mtp_tensors)} MTP tensors → {cached_path} ({cached_path.stat().st_size / 1e6:.0f}MB)")
+        return str(cached_path)
+
+    def warmup_speculative(self, model, tokenizer) -> None:
+        """Warm up the speculative decoding path (MTP draft + verify kernels)."""
+        if not hasattr(self._exo_gen, 'mtp'):
+            return
+
+        from mlx_lm.models import cache as cache_mod
+        from exo.worker.engines.mlx.speculative.mtp_module import speculative_forward, draft_tokens
+
+        logger.info("Warming up speculative decoding kernels...")
+        mtp = self._exo_gen.mtp
+        gamma = self._exo_gen.gamma
+
+        # Small warmup: prefill a short prompt, run a few speculative cycles
+        warmup_prompt = tokenizer.encode("Warm up speculative decoding.")
+        cache = cache_mod.make_prompt_cache(model)
+        mtp.reset_cache()
+
+        # Prefill
+        pre_norm, logits = speculative_forward(model, mx.array([warmup_prompt]), cache)
+        mx.eval(pre_norm, logits)
+        next_token = mx.argmax(logits[0, -1], axis=-1).item()
+
+        # MTP prefill
+        if pre_norm.shape[1] > 1:
+            _ = mtp.predict(pre_norm[:, :-1, :], mx.array([warmup_prompt[1:]]))
+            mx.eval(_)
+
+        # Run a few speculative cycles to compile kernels
+        last_pn = pre_norm[:, -1:, :]
+        next_arr = mx.array([[next_token]])
+        for _ in range(3):
+            draft_ids, _ = draft_tokens(mtp, last_pn, next_arr, gamma, 0.0)
+            draft_concat = mx.concatenate([d.reshape(1, 1) for d in draft_ids], axis=1)
+            verify_input = mx.concatenate([next_arr, draft_concat], axis=1)
+            vpn, vl = speculative_forward(model, verify_input, cache, speculative=True)
+            all_next = mx.argmax(vl[0], axis=-1)
+            mx.eval(vpn, all_next)
+            # Accept all for warmup (don't care about correctness)
+            next_arr = all_next[0].reshape(1, 1)
+            last_pn = vpn[:, 0:1, :]
+            for i, c in enumerate(cache):
+                if hasattr(c, 'base'):
+                    cache[i] = c.base
+
+        logger.info("Speculative warmup complete")

    @property
    def has_work(self) -> bool:
        return (
            bool(self._active_tasks)
-            or bool(self._mlx_gen.unprocessed_prompts)
-            or self._mlx_gen.active_batch is not None
+            or bool(self._exo_gen.unprocessed_prompts)
+            or self._exo_gen.active_batch is not None
        )

    def submit(
@@ -136,10 +293,16 @@ class ExoBatchGenerator:
        seed = task_params.seed if task_params.seed is not None else 42
        mx.random.seed(seed)

+        spec_temp_override = os.environ.get("EXO_SPECULATIVE_TEMP")
+        if spec_temp_override is not None:
+            sampling_temp = float(spec_temp_override)
+        elif task_params.temperature is not None:
+            sampling_temp = task_params.temperature
+        else:
+            sampling_temp = 0.7
+
        sampler = make_sampler(
-            temp=task_params.temperature
-            if task_params.temperature is not None
-            else 0.7,
+            temp=sampling_temp,
            top_p=task_params.top_p if task_params.top_p is not None else 1.0,
            min_p=task_params.min_p if task_params.min_p is not None else 0.05,
            top_k=task_params.top_k if task_params.top_k is not None else 0,
@@ -156,6 +319,23 @@ class ExoBatchGenerator:
            distributed_prompt_progress_callback,
        )

+        # MTP prefill: build MTP KV cache from prompt hidden states
+        # Pair position i with token i+1 (MTP predicts token t+2 from hidden[t] + embed[t+1])
+        if hasattr(self._exo_gen, 'mtp'):
+            prompt_pre_norm = self._exo_gen._captured.get('prompt_pre_norm')
+            if prompt_pre_norm is not None:
+                mx.eval(prompt_pre_norm)
+                self._exo_gen.mtp.reset_cache()
+                S_pre = prompt_pre_norm.shape[1]
+                if S_pre > 0 and len(all_prompt_tokens) > S_pre:
+                    mtp_toks = all_prompt_tokens[1:S_pre + 1].tolist()
+                    _ = self._exo_gen.mtp.predict(
+                        prompt_pre_norm,
+                        mx.array([mtp_toks])
+                    )
+                    mx.eval(_)
+                logger.info(f"MTP cache prefilled ({S_pre} positions)")
+
        # We need to clamp rotating kv caches to max size so that mlx lm's _merge_caches behaves
        for c in cache:
            if (
@@ -193,7 +373,7 @@ class ExoBatchGenerator:

        max_tokens = task_params.max_output_tokens or MAX_TOKENS

-        uids = self._mlx_gen.insert(
+        uids = self._exo_gen.insert(
            prompts=[last_tokens.tolist()],
            max_tokens=[max_tokens],
            caches=[list(cache)],
@@ -205,6 +385,16 @@ class ExoBatchGenerator:

        uid = uids[0]

+        # Pass request temperature to speculative cycle
+        # EXO_SPECULATIVE_TEMP overrides if set; otherwise use request temp
+        if hasattr(self._exo_gen, '_request_temp'):
+            env_temp = os.environ.get("EXO_SPECULATIVE_TEMP")
+            if env_temp is not None:
+                self._exo_gen._request_temp[uid] = float(env_temp)
+            else:
+                request_temp = task_params.temperature if task_params.temperature is not None else 0.7
+                self._exo_gen._request_temp[uid] = request_temp
+
        self._active_tasks[uid] = _EngineTask(
            uid=uid,
            task_params=task_params,
@@ -216,7 +406,6 @@ class ExoBatchGenerator:
            on_generation_token=on_generation_token,
            generation_start_time=time.perf_counter(),
            prefill_tps=_prefill_tps,
-            generation_time_at_start=self._mlx_gen._stats.generation_time,
        )

        return uid
@@ -225,12 +414,7 @@ class ExoBatchGenerator:
        if not self.has_work:
            return []

-        self._mlx_gen._needs_topk = any(  # pyright: ignore[reportAttributeAccessIssue]
-            t.task_params.logprobs for t in self._active_tasks.values()
-        )
-        _step_tic = time.perf_counter()
-        responses = self._mlx_gen.next()
-        _next_elapsed = time.perf_counter() - _step_tic
+        responses = self._exo_gen.next()

        results: list[tuple[int, GenerationResponse]] = []

@@ -287,32 +471,29 @@ class ExoBatchGenerator:

            logprob: float | None = None
            top_logprobs: list[TopLogprobItem] | None = None
-            if task_params.logprobs:
-                with mx.stream(generation_stream):
-                    logprob, top_logprobs = extract_top_logprobs(
-                        logprobs=response.logprobs,
-                        tokenizer=self.tokenizer,
-                        top_logprobs=task_params.top_logprobs or DEFAULT_TOP_LOGPROBS,
-                        selected_token=response.token,
-                        precomputed_indices=getattr(response, "_topk_indices", None),
-                        precomputed_values=getattr(response, "_topk_values", None),
-                        precomputed_selected=getattr(
-                            response, "_selected_logprob", None
-                        ),
-                    )
+            if task_params.logprobs and os.environ.get("EXO_DISABLE_LOGPROBS") != "1":
+                logprob, top_logprobs = extract_top_logprobs(
+                    logprobs=response.logprobs,
+                    tokenizer=self.tokenizer,
+                    top_logprobs=task_params.top_logprobs or DEFAULT_TOP_LOGPROBS,
+                    selected_token=response.token,
+                )

            stats: GenerationStats | None = None
            usage: Usage | None = None
            if is_done:
-                gen_time_delta = (
-                    self._mlx_gen._stats.generation_time
-                    - state.generation_time_at_start
-                )
-                generation_tps = (
-                    state.completion_tokens / gen_time_delta
-                    if gen_time_delta > 0
-                    else 0.0
-                )
+                try:
+                    mlx_stats = self._exo_gen.stats()
+                    generation_tps = mlx_stats.generation_tps
+                except ZeroDivisionError:
+                    generation_elapsed = (
+                        time.perf_counter() - state.generation_start_time
+                    )
+                    generation_tps = (
+                        state.completion_tokens / generation_elapsed
+                        if generation_elapsed > 0
+                        else 0.0
+                    )

                stats = GenerationStats(
                    prompt_tps=state.prefill_tps,
@@ -359,22 +540,15 @@ class ExoBatchGenerator:
                    -max_stop_len:
                ]

-        _step_elapsed = time.perf_counter() - _step_tic
-        _overhead = _step_elapsed - _next_elapsed
-        if self._mlx_gen._next_count % 64 == 0 and responses:
-            logger.debug(
-                f"step overhead: {_overhead * 1000:.2f}ms (next={_next_elapsed * 1000:.2f}ms total={_step_elapsed * 1000:.2f}ms)"
-            )
-
        return results

    def cancel(self, uids: list[int]) -> None:
-        self._mlx_gen.remove(uids)
+        self._exo_gen.remove(uids)
        for uid in uids:
            self._active_tasks.pop(uid, None)

    def close(self) -> None:
-        self._mlx_gen.close()
+        self._exo_gen.close()

    def _save_prefix_cache(
        self,
@@ -393,8 +567,9 @@ class ExoBatchGenerator:
                if len(all_prompt_tokens) > 0
                else 0.0
            )
-            if matched_index is not None and (
-                prefix_hit_length > 1000 or hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
+            if (
+                matched_index is not None
+                and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
            ):
                self.kv_prefix_cache.update_kv_cache(
                    matched_index,
@@ -13,7 +13,7 @@ from mlx_lm.models.cache import ArraysCache, RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper

-from exo.api.types import (
+from exo.shared.types.api import (
    CompletionTokensDetails,
    FinishReason,
    GenerationStats,
@@ -313,46 +313,52 @@ def warmup_inference(
    model_id: ModelId,
 ) -> int:
    logger.info(f"warming up inference for instance: {model_id}")
+    t = time.monotonic()

    content = "Prompt to warm up the inference engine. Repeat this."

-    warmup_task_params = TextGenerationTaskParams(
-        model=model_id,
-        input=[InputMessage(role="user", content=content)],
-        max_output_tokens=50,
-        temperature=0.0,
-    )
-
    warmup_prompt = apply_chat_template(
        tokenizer=tokenizer,
-        task_params=warmup_task_params,
+        task_params=TextGenerationTaskParams(
+            model=ModelId(""),
+            input=[InputMessage(role="user", content=content)],
+        ),
    )

    tokens_generated = 0

+    cache = make_kv_cache(
+        model=model,
+    )
+
+    # Use a default sampler for warmup
+    sampler = make_sampler(temp=0.0)
+
    mx_barrier(group)

    logger.info("Generating warmup tokens")
-
-    t = time.monotonic()
-
-    for _r in mlx_generate(
+    for _r in stream_generate(
        model=model,
        tokenizer=tokenizer,
-        task=warmup_task_params,
        prompt=warmup_prompt,
-        kv_prefix_cache=None,
-        group=group,
+        max_tokens=50,
+        sampler=sampler,
+        prompt_cache=cache,
+        prefill_step_size=2048,
+        kv_group_size=KV_GROUP_SIZE,
+        kv_bits=KV_BITS,
    ):
+        logger.info("Generated warmup token: " + str(_r.text))
        tokens_generated += 1

-    check_for_cancel_every = min(
-        math.ceil(tokens_generated / min(time.monotonic() - t, 0.001)), 100
-    )
+    logger.info("Generated ALL warmup tokens")

    mx_barrier(group)

    logger.info(f"warmed up by generating {tokens_generated} tokens")
+    check_for_cancel_every = min(
+        math.ceil(tokens_generated / min(time.monotonic() - t, 0.001)), 100
+    )
    if group is not None:
        check_for_cancel_every = int(
            mx.max(
@@ -645,9 +651,9 @@ def mlx_generate(
                    if len(all_prompt_tokens) > 0
                    else 0.0
                )
-                if matched_index is not None and (
-                    prefix_hit_length > 1000
-                    or hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
+                if (
+                    matched_index is not None
+                    and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
                ):
                    kv_prefix_cache.update_kv_cache(
                        matched_index,
@@ -1,14 +1,45 @@
-from exo.worker.engines.mlx.patches.opt_batch_gen import apply_batch_gen_patch
-from exo.worker.engines.mlx.patches.standard_yarn_rope import patch_yarn_rope
+"""Model-specific kernel fusion patches for MLX inference.

-_applied = False
+Detects model type after loading and applies optimized kernel patches.
+Currently supports:
+- Qwen3.5 MoE (model_type: qwen3_5_moe): batched fused oproj (GDN + GQA + MoE)
+
+Set EXO_FUSED_KERNELS=0 to disable all patches (vanilla mode).
+Default: EXO_FUSED_KERNELS=1 (enabled).
+"""
+
+import json
+import os
+from pathlib import Path
+
+import mlx.nn as nn
+from loguru import logger


-def apply_mlx_patches() -> None:
-    global _applied
-    if _applied:
+def maybe_apply_patches(model: nn.Module, model_path: Path) -> None:
+    """Detect model type and apply kernel fusion patches if available."""
+    fused_mode = os.environ.get("EXO_FUSED_KERNELS", "1")
+    if fused_mode == "0":
+        logger.info("Kernel fusion patches disabled (EXO_FUSED_KERNELS=0)")
        return
-    _applied = True
-    patch_yarn_rope()
-    # patch_gdn_softplus()
-    apply_batch_gen_patch()
+
+    config_path = model_path / "config.json"
+    if not config_path.exists():
+        return
+
+    with open(config_path) as f:
+        config = json.load(f)
+
+    model_type = config.get("model_type", "")
+
+    if model_type == "qwen3_5_moe":
+        from .qwen3_5_moe.apply import apply_qwen35_batched_fused_patches
+
+        logger.info("Detected Qwen3.5 MoE model, applying batched fused kernel patches")
+        apply_qwen35_batched_fused_patches(model)
+
+    elif model_type == "qwen3_5":
+        from .qwen3_5.lpb_patch import apply_lpb_patches
+
+        logger.info("Detected Qwen3.5 dense model, applying LpB kernel patches")
+        apply_lpb_patches(model, batch_size=4)
@@ -0,0 +1,27 @@
+import sys
+
+import mlx.core as mx
+from mlx_lm.models.gated_delta import compute_g
+
+
+def _compute_g_f32(a_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array:
+    return mx.exp(
+        -mx.exp(a_log.astype(mx.float32))
+        * mx.where(
+            (a + dt_bias).astype(mx.float32) > 20,
+            (a + dt_bias).astype(mx.float32),
+            mx.log1p(mx.exp((a + dt_bias).astype(mx.float32))),
+        )
+    ).astype(a.dtype)
+
+
+def patch_gdn_softplus() -> None:
+    from mlx_lm.models import gated_delta
+
+    gated_delta.compute_g = _compute_g_f32
+
+    for mod in list(sys.modules.values()):
+        if mod is gated_delta:
+            continue
+        if getattr(mod, "compute_g", None) is compute_g:
+            object.__setattr__(mod, "compute_g", _compute_g_f32)
@@ -0,0 +1,135 @@
+#!/usr/bin/env python3
+"""Isolated loop-over-B GEMV kernel for quantized matmul.
+
+Extracts the loop-over-B pattern from batched_fused_gdn_projections_8bit
+but without any epilogues — pure Y = X @ dequant(W)^T output.
+
+For comparing our GEMV approach against MLX's affine_qmv_fast on
+an isolated QuantizedLinear operation (e.g., in_proj_qkv: N=8192, K=2048).
+
+TG: (32, 2, 1) = 64 threads = 2 SGs.
+Each SG: 4 output rows.
+B loop inside row loop for low register pressure (R = 4B + 5).
+
+Usage:
+    from custom_qmv_loop_over_b import custom_qmv_loop_over_b
+    y = custom_qmv_loop_over_b(x, w, scales, biases, M=8, N=8192, K=2048)
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_custom_qmv_source(M_val, N_val, K_val, group_size=64):
+    gs = group_size
+    sc_stride = 256 // gs
+    slid_div = gs // 8
+    K_groups = K_val // gs
+    B = M_val  # batch size = M
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int K = {K_val};
+    const int N = {N_val};
+    const int M = {M_val};
+    const int K_groups = {K_groups};
+    const int SC_STRIDE = {sc_stride};
+    const int SLID_DIV = {slid_div};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+    int tg = tgid.y;
+
+    int out_row = tg * 8 + sgid * RESULTS_PER_SG;
+    if (out_row >= N) return;
+
+    // Weight pointers
+    const device uint8_t* ws = (const device uint8_t*)w + (long)out_row * K + slid * VALUES_PER_THREAD;
+    const device bfloat16_t* sc = (const device bfloat16_t*)scales + (long)out_row * K_groups + slid / SLID_DIV;
+    const device bfloat16_t* bi = (const device bfloat16_t*)biases + (long)out_row * K_groups + slid / SLID_DIV;
+
+    // Result accumulators: 4 rows × B batches
+    float result[{4 * B}];
+    for (int i = 0; i < {4 * B}; i++) result[i] = 0;
+
+    int x_base = slid * VALUES_PER_THREAD;
+
+    // K-loop: loop over B inside row loop
+    for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{
+
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            const device uint8_t* wl = ws + row * K;
+            float s_val = float(sc[row * K_groups]);
+            float b_val = float(bi[row * K_groups]);
+
+            for (int b = 0; b < {B}; b++) {{
+                float accum = 0, xsum = 0;
+                for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+                    float xi = float(((const device bfloat16_t*)x)[b * K + x_base + i]);
+                    accum += xi * float(wl[i]);
+                    xsum += xi;
+                }}
+                result[b * 4 + row] += s_val * accum + xsum * b_val;
+            }}
+        }}
+
+        ws += BLOCK_SIZE; sc += SC_STRIDE; bi += SC_STRIDE; x_base += BLOCK_SIZE;
+    }}
+
+    // Reduction
+    for (int i = 0; i < {4 * B}; i++) result[i] = simd_sum(result[i]);
+
+    // Write output (bf16)
+    if (slid < 4u) {{
+        for (int b = 0; b < {B}; b++) {{
+            int r = out_row + (int)slid;
+            if (r < N) {{
+                y[b * N + r] = static_cast<bfloat16_t>(result[b * 4 + slid]);
+            }}
+        }}
+    }}
+"""
+
+
+_custom_qmv_cache = {}
+
+
+def custom_qmv_loop_over_b(x, w, scales, biases, M, N, K, group_size=64):
+    """Loop-over-B GEMV for quantized matmul.
+
+    Args:
+        x: (M, K) bfloat16 input
+        w: (N, K/4) uint32 packed 8-bit weights
+        scales: (N, K/gs) bfloat16
+        biases: (N, K/gs) bfloat16
+        M, N, K: dimensions
+    Returns:
+        y: (M, N) bfloat16
+    """
+    key = (M, N, K, group_size)
+    if key not in _custom_qmv_cache:
+        _custom_qmv_cache[key] = mx.fast.metal_kernel(
+            name=f"custom_qmv_loop_b_M{M}_N{N}_K{K}",
+            input_names=["x", "w", "scales", "biases"],
+            output_names=["y"],
+            source=_gen_custom_qmv_source(M, N, K, group_size),
+        )
+    kern = _custom_qmv_cache[key]
+
+    n_tg = ceil_div(N, 8)
+
+    result = kern(
+        inputs=[x, w, scales, biases],
+        output_shapes=[(M * N,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(32, n_tg * 2, 1),
+        threadgroup=(32, 2, 1),
+    )
+
+    return result[0].reshape(M, N)
@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+"""Loop-over-B patches for Qwen3.5-27B dense model.
+
+Replaces vanilla QuantizedLinear calls with custom loop-over-B GEMV
+for projections where N > K (expanding projections). Falls back to
+vanilla for N <= K (contracting projections like down_proj, o_proj).
+
+Usage:
+    from lpb_patch import apply_lpb_patches
+    apply_lpb_patches(model, batch_size=4)
+"""
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .custom_qmv_loop_over_b import custom_qmv_loop_over_b
+
+
+def _make_lpb_forward(original_module, N, K, BS, GS=64):
+    """Create a patched forward that uses loop-over-B."""
+    w = original_module.weight
+    s = original_module.scales
+    b = original_module.biases
+
+    MAX_M = 16  # Max total tokens (B*S) for custom kernel; above this use vanilla
+
+    def forward(self_unused, x):
+        # Use LpB for small M=B*S. Large prefill falls back to vanilla.
+        M_total = 1
+        for d in x.shape[:-1]:
+            M_total *= d
+        if M_total > MAX_M:
+            return original_module(x)
+        orig_shape = x.shape
+        x_2d = x.reshape(-1, K)
+        M = x_2d.shape[0]
+        y = custom_qmv_loop_over_b(x_2d, w, s, b, M, N, K, GS)
+        return y.reshape(*orig_shape[:-1], N)
+
+    return forward
+
+
+def apply_lpb_patches(model, batch_size=4):
+    """Patch all expanding QuantizedLinear projections with loop-over-B.
+
+    Only patches projections where N > K (expanding):
+    - gate_proj, up_proj (17408 > 5120)
+    - in_proj_qkv (10240 > 5120)
+    - in_proj_z (6144 > 5120)
+    - q_proj (12288 > 5120)
+
+    Skips N <= K projections (down_proj, o_proj, k_proj, v_proj)
+    where vanilla is already efficient.
+    """
+    inner = getattr(model, 'model', None) or model.language_model.model
+    patched = 0
+
+    for li, layer in enumerate(inner.layers):
+        # MLP: gate_proj, up_proj (N=17408, K=5120)
+        mlp = layer.mlp
+        for proj_name in ['gate_proj', 'up_proj', 'down_proj']:
+            proj = getattr(mlp, proj_name)
+            if isinstance(proj, nn.QuantizedLinear):
+                N = proj.weight.shape[0]  # output dim
+                K_packed = proj.weight.shape[1]
+                K = K_packed * 4  # 8-bit: 4 values per uint32
+                setattr(mlp, proj_name, type('LpBLinear', (), {
+                    '__call__': _make_lpb_forward(proj, N, K, batch_size),
+                    'weight': proj.weight,
+                    'scales': proj.scales,
+                    'biases': proj.biases,
+                })())
+                patched += 1
+
+        # Attention projections
+        if layer.is_linear:
+            attn = layer.linear_attn
+            for proj_name in ['in_proj_qkv', 'in_proj_z', 'out_proj']:
+                if hasattr(attn, proj_name):
+                    proj = getattr(attn, proj_name)
+                    if isinstance(proj, nn.QuantizedLinear):
+                        N = proj.weight.shape[0]
+                        K = proj.weight.shape[1] * 4
+                        setattr(attn, proj_name, type('LpBLinear', (), {
+                            '__call__': _make_lpb_forward(proj, N, K, batch_size),
+                            'weight': proj.weight,
+                            'scales': proj.scales,
+                            'biases': proj.biases,
+                        })())
+                        patched += 1
+        else:
+            attn = layer.self_attn
+            for proj_name in ['q_proj', 'o_proj']:
+                if hasattr(attn, proj_name):
+                    proj = getattr(attn, proj_name)
+                    if isinstance(proj, nn.QuantizedLinear):
+                        N = proj.weight.shape[0]
+                        K = proj.weight.shape[1] * 4
+                        setattr(attn, proj_name, type('LpBLinear', (), {
+                            '__call__': _make_lpb_forward(proj, N, K, batch_size),
+                            'weight': proj.weight,
+                            'scales': proj.scales,
+                            'biases': proj.biases,
+                        })())
+                        patched += 1
+
+    print(f"  Patched {patched} projections with loop-over-B")
+    return patched
@@ -0,0 +1,74 @@
+"""Apply batched fused kernel patches to Qwen3.5 MoE models.
+
+Entry point called from patches/__init__.py after model type detection.
+"""
+
+import time
+
+import mlx.nn as nn
+from loguru import logger
+
+from .common import (
+    _patch_swiglu_weights,
+    _patch_shared_expert,
+    _patch_down_proj,
+    _patch_oproj_gate_rms,
+    _patch_gdn_proj_weights,
+    _patch_gqa_proj_weights,
+)
+from mlx_lm.models.qwen3_5 import DecoderLayer
+from mlx_lm.models.qwen3_next import Qwen3NextAttention, Qwen3NextSparseMoeBlock
+from mlx_lm.models.qwen3_5 import GatedDeltaNet
+
+
+def apply_qwen35_batched_fused_patches(model: nn.Module) -> None:
+    """Apply batched fused patches (GDN + GQA attention + oproj MoE) to all layers.
+
+    Fused GDN attention (3/4 layers) + fused GQA projections (1/4 layers)
+    + batched oproj MoE (4 custom dispatches). Works with BatchGenerator for
+    any batch size 1..8. Falls back to vanilla for B>8 or S>1.
+    """
+    layers = model.layers  # type: ignore[attr-defined]
+    n_layers = len(layers)
+
+    t0 = time.time()
+
+    n_gdn = 0
+    n_gqa = 0
+    for li, layer in enumerate(layers):
+        moe = layer.mlp
+        if isinstance(moe, Qwen3NextSparseMoeBlock):
+            # MoE weight prep
+            _patch_swiglu_weights(moe)
+            _patch_shared_expert(moe)
+            _patch_down_proj(moe)
+            _patch_oproj_gate_rms(layer, gate_bm=8)
+
+            # Attention weight prep
+            if layer.is_linear:
+                _patch_gdn_proj_weights(layer.linear_attn)
+                n_gdn += 1
+            else:
+                _patch_gqa_proj_weights(layer.self_attn)
+                n_gqa += 1
+
+            if (li + 1) % 10 == 0 or li == 0:
+                logger.info(f"  Patched layer {li+1}/{n_layers}")
+
+    # Import patched __call__ methods
+    from .fused_gdn_attention import _fused_gdn_call
+    from .batched_fused_gqa_attention import _batched_fused_gqa_call
+    from .batched_moe import _batched_oproj_moe_call
+    from .decoder import _fused_gdn_decoder_call
+
+    # Class-level method replacement
+    GatedDeltaNet.__call__ = _fused_gdn_call
+    Qwen3NextAttention.__call__ = _batched_fused_gqa_call
+    Qwen3NextSparseMoeBlock.__call__ = _batched_oproj_moe_call
+    DecoderLayer.__call__ = _fused_gdn_decoder_call
+
+    t_patch = time.time() - t0
+    logger.info(
+        f"Qwen3.5 batched fused: {n_gdn} GDN + {n_gqa} GQA layers, "
+        f"{n_layers} total in {t_patch:.1f}s"
+    )
@@ -0,0 +1,102 @@
+"""Batched fused GQA attention for Qwen3.5 (projections + norm+rope fused, vanilla SDPA).
+
+Dispatches:
+  1. batched_fused_gqa_projections — merged q+gate+k+v GEMV with register weight sharing
+  2. fused_qk_norm_rope — per-head RMSNorm + RoPE (already supports B>1 via grid z)
+  3. Vanilla cache update (BatchKVCache)
+  4. Vanilla SDPA (MLX built-in, handles batching natively)
+  5. Vanilla gate multiply
+
+Returns pre-out_proj output for the oproj MoE block.
+Falls back to vanilla (with o_proj) for B>8 or S>1.
+"""
+
+from typing import Any, Optional
+
+import mlx.core as mx
+
+from .kernels.batched_fused_gqa_projections_8bit import batched_fused_gqa_projections
+
+
+def _batched_fused_gqa_call(
+    self,
+    x: mx.array,
+    mask: Optional[mx.array] = None,
+    cache: Optional[Any] = None,
+) -> mx.array:
+    """Batched fused GQA attention with custom projection + norm/rope kernels.
+
+    For 1<=B<=8, S=1: fused projections + fused norm+rope + vanilla SDPA.
+    For B>8 or S>1: vanilla fallback.
+    Returns pre-out_proj output [B, S, H_q*D].
+    """
+    B, S, _ = x.shape
+
+    if S > 1 or B > 8:
+        # Vanilla fallback
+        q_proj_output = self.q_proj(x)
+        queries, gate = mx.split(
+            q_proj_output.reshape(B, S, self.num_attention_heads, -1), 2, axis=-1
+        )
+        gate = gate.reshape(B, S, -1)
+        keys, values = self.k_proj(x), self.v_proj(x)
+        queries = self.q_norm(queries).transpose(0, 2, 1, 3)
+        keys = self.k_norm(
+            keys.reshape(B, S, self.num_key_value_heads, -1)
+        ).transpose(0, 2, 1, 3)
+        values = values.reshape(B, S, self.num_key_value_heads, -1).transpose(
+            0, 2, 1, 3
+        )
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        from mlx_lm.models.qwen3_next import scaled_dot_product_attention
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, S, -1)
+        return self.o_proj(output * mx.sigmoid(gate))
+
+    H_q = self.num_attention_heads
+    H_kv = self.num_key_value_heads
+    D = self.head_dim
+
+    # ── Dispatch 1: batched fused projections ──
+    queries, gate_sigmoid, keys, values = batched_fused_gqa_projections(
+        x,
+        self._merged_proj_w, self._merged_proj_s, self._merged_proj_b,
+        self._merged_proj_dims,
+        batch_size=B,
+        total_tg=getattr(self, '_d1_total_tg', None),
+    )
+
+    # ── Dispatch 2+: vanilla norm + rope (avoids mx.eval sync on BatchKVCache offset) ──
+    queries = self.q_norm(queries.reshape(B, 1, H_q, D)).transpose(0, 2, 1, 3)
+    keys = self.k_norm(keys.reshape(B, 1, H_kv, D)).transpose(0, 2, 1, 3)
+    values = values.reshape(B, 1, H_kv, D).transpose(0, 2, 1, 3)
+
+    if cache is not None:
+        queries = self.rope(queries, offset=cache.offset)
+        keys = self.rope(keys, offset=cache.offset)
+    else:
+        queries = self.rope(queries)
+        keys = self.rope(keys)
+
+    # ── Dispatch 3: KV cache update ──
+    if cache is not None:
+        keys, values = cache.update_and_fetch(keys, values)
+
+    # ── Dispatch 4: vanilla SDPA ──
+    from mlx_lm.models.qwen3_next import scaled_dot_product_attention
+    output = scaled_dot_product_attention(
+        queries, keys, values, cache=cache, scale=self.scale, mask=mask
+    )
+    output = output.transpose(0, 2, 1, 3).reshape(B, S, -1)
+
+    # ── Gate multiply ──
+    return output * gate_sigmoid.astype(output.dtype)
@@ -0,0 +1,94 @@
+"""Batched oproj MoE with 4 custom Metal kernel dispatches.
+
+Fuses o_proj + RMSNorm + gate GEMV + softmax + topk + SwiGLU + down_proj + epilogue
+into 4 dispatches with register-level weight sharing for the shared expert.
+Falls back to vanilla MoE when called without _residual (from vanilla decoder path).
+"""
+
+import mlx.core as mx
+
+from .kernels.batched_merged_down_proj_8bit import batched_merged_down_proj_8bit
+from .kernels.batched_oproj_gate_gemv_8bit import batched_oproj_gate_gemv
+from .kernels.batched_softmax_topk_swiglu_8bit import batched_softmax_topk_swiglu_8bit
+from .kernels.batched_moe_epilogue import batched_moe_epilogue
+
+
+def _batched_oproj_moe_call(self, attn_out_3d, _residual=None):
+    """Batched MoE with full oproj fusion (4 custom dispatches).
+
+    Receives raw attention output (pre-o_proj) and residual for B tokens.
+    All 4 dispatches use register-level weight sharing for shared weights.
+
+    When _residual is None, called from vanilla decoder — do vanilla MoE.
+    """
+    if _residual is None:
+        # Vanilla MoE path (called from vanilla decoder for B>8 or S>1)
+        x = attn_out_3d
+        gates = self.gate(x)
+        gates = mx.softmax(gates, axis=-1, precise=True)
+        k = self.top_k
+        inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]
+        scores = mx.take_along_axis(gates, inds, axis=-1)
+        if self.norm_topk_prob:
+            scores = scores / scores.sum(axis=-1, keepdims=True)
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2)
+        shared_y = self.shared_expert(x)
+        shared_y = mx.sigmoid(self.shared_expert_gate(x)) * shared_y
+        return y + shared_y
+
+    B_dim = attn_out_3d.shape[0]
+
+    K = self._oproj_M
+    K_attn = self._oproj_K_attn
+    n_active = self.top_k
+    E = self._oproj_n_experts
+
+    attn_out = attn_out_3d.reshape(B_dim, K_attn).astype(mx.bfloat16)
+    residual = _residual.reshape(B_dim, K).astype(mx.bfloat16)
+
+    # ── Dispatch 1: batched o_proj + gate GEMVs ──
+    h_scaled, h_out, x2_partials, gate_part_a, gate_part_b = \
+        batched_oproj_gate_gemv(
+            self._oproj_w, self._oproj_s, self._oproj_b,
+            attn_out, residual, self._oproj_rms_weight,
+            self._oproj_M1, self._oproj_W_fused,
+            M=K, K_attn=K_attn, batch_size=B_dim,
+            n_experts=E, gate_bm=self._oproj_gate_bm,
+            K_hidden=self._oproj_K_hidden,
+        )
+
+    n_oproj_tg = (K + 31) // 32
+    N_INTER = self.switch_mlp._fused_n_inter
+    SHARED_INTER = self._shared_inter
+
+    # ── Dispatch 2: batched softmax + topk + SwiGLU ──
+    y_routed, y_shared, out_inds, norm_scores, gate_raw = \
+        batched_softmax_topk_swiglu_8bit(
+            self.switch_mlp._fused_w_gu, self.switch_mlp._fused_s_gu,
+            self.switch_mlp._fused_b_gu,
+            self._shared_w_gu, self._shared_s_gu, self._shared_b_gu,
+            self._seg_w, self._seg_s, self._seg_b,
+            h_scaled, gate_part_a, gate_part_b, x2_partials,
+            n_inter=N_INTER, k_hidden=K, batch_size=B_dim,
+            n_active=n_active, n_oproj_tg=n_oproj_tg,
+            n_experts=E, shared_inter=SHARED_INTER,
+        )
+
+    # ── Dispatch 3: batched merged down_proj ──
+    d_routed, d_shared = batched_merged_down_proj_8bit(
+        self._down_w, self._down_s, self._down_b,
+        self._shared_down_w, self._shared_down_s, self._shared_down_b,
+        y_routed, y_shared.reshape(B_dim * SHARED_INTER), out_inds,
+        k_out=K, n_in=self._down_N, batch_size=B_dim,
+        n_active=n_active, shared_n_in=SHARED_INTER,
+    )
+
+    # ── Dispatch 4: batched epilogue ──
+    Y = batched_moe_epilogue(
+        d_routed, d_shared, norm_scores,
+        h_out, gate_raw,
+        k_val=K, batch_size=B_dim, n_active=n_active,
+    )
+
+    return Y.reshape(B_dim, 1, K).astype(attn_out_3d.dtype)
@@ -0,0 +1,500 @@
+"""Common weight preparation functions for Qwen3.5 fused kernel patches.
+
+Functions:
+  ceil_div               — integer ceiling division
+  _patch_swiglu_weights  — stack gate+up weights for fused SwiGLU kernel
+  _patch_down_proj       — extract down_proj weights for merged kernel dispatch
+  _patch_shared_expert   — prepare shared expert weights (8-bit)
+  dequantize_shared_expert — convert shared expert from 8-bit to bf16
+  _patch_oproj_gate_rms  — precompute M1/W_fused for fused o_proj + gate GEMV
+  _patch_gdn_proj_weights — merge GDN projection weights for fused GEMV
+  _patch_gqa_proj_weights — merge GQA q/k/v weights with q_proj permutation
+  make_qwen_random_cache — create pre-filled cache for testing
+  build_model            — build Qwen3.5 MoE layers with 8-bit quantization
+"""
+
+from types import SimpleNamespace
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx_lm.models.qwen3_5 import (
+    DecoderLayer,
+    TextModelArgs,
+)
+from mlx_lm.models.qwen3_next import Qwen3NextSparseMoeBlock
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _patch_swiglu_weights(moe):
+    """Stack gate+up weights for fused 8-bit SwiGLU kernel.
+
+    Creates concatenated (E, 2*N_INTER, K/4) weight, (E, 2*N_INTER, K/gs) scales/biases
+    from the separate gate_proj and up_proj QuantizedSwitchLinear layers.
+    """
+    gate_proj = moe.switch_mlp.gate_proj
+    up_proj = moe.switch_mlp.up_proj
+
+    moe.switch_mlp._fused_w_gu = mx.concatenate(
+        [gate_proj.weight, up_proj.weight], axis=1)
+    moe.switch_mlp._fused_s_gu = mx.concatenate(
+        [gate_proj.scales, up_proj.scales], axis=1)
+    moe.switch_mlp._fused_b_gu = mx.concatenate(
+        [gate_proj.biases, up_proj.biases], axis=1)
+    moe.switch_mlp._fused_n_inter = gate_proj.output_dims
+    moe.switch_mlp._fused_k_hidden = gate_proj.input_dims
+    moe.switch_mlp._fused_group_size = gate_proj.group_size
+
+    mx.eval(moe.switch_mlp._fused_w_gu,
+            moe.switch_mlp._fused_s_gu,
+            moe.switch_mlp._fused_b_gu)
+
+
+def _patch_shared_expert(moe):
+    """Prepare shared expert quantized weights for fused 8-bit path.
+
+    Stacks shared gate+up quantized weights (weight, scales, biases).
+    Stores down_proj quantized weights separately.
+    Shared expert stays in 8-bit — same as vanilla MLX dispatch.
+    """
+    shared = moe.shared_expert
+    gp = shared.gate_proj
+    up = shared.up_proj
+    dp = shared.down_proj
+
+    # Gate+up stacked: (2*SHARED_INTER, K/4) uint32, (2*SHARED_INTER, K/gs) bf16
+    moe._shared_w_gu = mx.concatenate([gp.weight, up.weight], axis=0)
+    moe._shared_s_gu = mx.concatenate([gp.scales, up.scales], axis=0)
+    moe._shared_b_gu = mx.concatenate([gp.biases, up.biases], axis=0)
+
+    # Down_proj: (K, SHARED_INTER/4) uint32, (K, SHARED_INTER/gs) bf16
+    moe._shared_down_w = dp.weight
+    moe._shared_down_s = dp.scales
+    moe._shared_down_b = dp.biases
+
+    # QuantizedLinear: weight is (out_features, in_features/pack_factor) uint32
+    # For 8-bit: pack_factor = 4, so in_features = weight.shape[1] * 4
+    moe._shared_inter = gp.weight.shape[0]  # SHARED_INTER (= out_features)
+    moe._shared_gs = gp.group_size           # gs (64)
+
+    mx.eval(moe._shared_w_gu, moe._shared_s_gu, moe._shared_b_gu,
+            moe._shared_down_w, moe._shared_down_s, moe._shared_down_b)
+
+
+def _patch_down_proj(moe):
+    """Extract down_proj weights for merged 8-bit kernel dispatch."""
+    dp = moe.switch_mlp.down_proj
+    moe._down_w = dp.weight       # (E, K_OUT, N_IN/4) uint32
+    moe._down_s = dp.scales       # (E, K_OUT, N_IN/gs) bf16
+    moe._down_b = dp.biases       # (E, K_OUT, N_IN/gs) bf16
+    moe._down_K = dp.output_dims  # K = 4096
+    moe._down_N = dp.input_dims   # N = 1024
+    moe._down_gs = dp.group_size  # gs = 64
+    mx.eval(moe._down_w, moe._down_s, moe._down_b)
+
+
+def dequantize_shared_expert(moe):
+    """Convert shared expert from 8-bit QuantizedLinear to bf16 weight wrappers.
+
+    The fused kernels expect bf16 shared expert weights. The real model (and our
+    random model) has shared expert quantized to 8-bit. This dequantizes in-place.
+    """
+    shared = moe.shared_expert
+    for proj_name in ["gate_proj", "up_proj", "down_proj"]:
+        proj = getattr(shared, proj_name)
+        if hasattr(proj, 'scales') and hasattr(proj, 'biases'):
+            w_bf16 = mx.dequantize(
+                proj.weight, proj.scales, proj.biases,
+                group_size=proj.group_size, bits=proj.bits,
+            ).astype(mx.bfloat16)
+            mx.eval(w_bf16)
+            setattr(shared, proj_name, SimpleNamespace(weight=w_bf16))
+
+
+def _patch_oproj_gate_rms(layer, gate_bm=8):
+    """Precompute M1/W_fused for fused o_proj + gate GEMV (oproj 4-dispatch mode).
+
+    Gate decomposition:
+      gate_score[e] = W_gate[e,:] @ rms_norm(h)
+      where h = residual + W_oproj @ attn_out
+      rms_norm(h) = h * w_rms * inv_rms
+
+    Expanding:
+      gate_score[e] = (W_fused @ residual + M1 @ attn_out) * inv_rms
+
+    Precomputed offline (per layer, stored on moe block):
+      W_fused = dequant(W_gate) · diag(w_rms)    — (E, K) bf16
+      M1 = W_fused @ dequant(W_oproj)            — (E, K_attn) bf16
+
+    Also stores o_proj quantized weights and shared_expert_gate weights
+    on the moe block for use by Dispatch 1 and Dispatch 2.
+
+    Args:
+        layer: DecoderLayer instance
+        gate_bm: SGs per gate TG in Dispatch 1 (1,2,4,8)
+    """
+    moe = layer.mlp
+
+    # ── Get attention output projection (works for both attention types) ──
+    if layer.is_linear:
+        oproj = layer.linear_attn.out_proj
+    else:
+        oproj = layer.self_attn.o_proj
+
+    # ── Dequantize gate and o_proj (temporary, for M1 computation) ──
+    # Eval incrementally to limit peak memory (E=512: dequant temps are ~140 MB)
+    gate = moe.gate
+    W_gate_f32 = mx.dequantize(
+        gate.weight, gate.scales, gate.biases,
+        group_size=gate.group_size, bits=gate.bits,
+    ).astype(mx.float32)
+
+    W_oproj_f32 = mx.dequantize(
+        oproj.weight, oproj.scales, oproj.biases,
+        group_size=oproj.group_size, bits=oproj.bits,
+    ).astype(mx.float32)
+    mx.eval(W_gate_f32, W_oproj_f32)
+
+    # ── RMSNorm weight ──
+    rms_weight = layer.post_attention_layernorm.weight.astype(mx.bfloat16)
+
+    # ── W_fused = dequant(W_gate) · diag(w_rms) ──
+    w_rms_f32 = rms_weight.astype(mx.float32)
+    W_fused = (W_gate_f32 * w_rms_f32).astype(mx.bfloat16)
+    mx.eval(W_fused)
+    del W_gate_f32  # free ~8 MB (E=512) or ~1 MB (E=64)
+
+    # ── M1 = W_fused @ W_oproj — precomputed in f32, stored bf16 ──
+    M1 = (W_fused.astype(mx.float32) @ W_oproj_f32).astype(mx.bfloat16)
+    mx.eval(M1)
+    del W_oproj_f32  # free ~128 MB
+
+    # Store on moe block
+    moe._oproj_M1 = M1              # (E, K_attn) bf16
+    moe._oproj_W_fused = W_fused    # (E, K) bf16
+    moe._oproj_rms_weight = rms_weight  # (K,) bf16
+
+    # ── O_proj quantized weights (for 8-bit GEMV in Dispatch 1) ──
+    moe._oproj_w = oproj.weight      # (K, K_attn/4) uint32
+    moe._oproj_s = oproj.scales      # (K, K_attn/gs) bf16
+    moe._oproj_b = oproj.biases      # (K, K_attn/gs) bf16
+    moe._oproj_K_attn = oproj.weight.shape[1] * 4  # 8192 (8-bit: pack_factor=4)
+
+    # ── Shared expert gate weights (for TG(0,0,0) fusion in Dispatch 2) ──
+    seg = moe.shared_expert_gate
+    moe._seg_w = seg.weight    # (1, K/4) uint32
+    moe._seg_s = seg.scales    # (1, K/gs) bf16
+    moe._seg_b = seg.biases    # (1, K/gs) bf16
+
+    # ── Dimensions ──
+    M = oproj.weight.shape[0]                    # 4096 (hidden_size)
+    K_hidden = W_fused.shape[1]                   # 4096 (same as M for Qwen)
+    n_experts = W_fused.shape[0]                  # E
+    moe._oproj_M = M
+    moe._oproj_K_hidden = K_hidden
+    moe._oproj_n_experts = n_experts
+    moe._oproj_n_tg = ceil_div(M, 32)            # 128 for M=4096
+    moe._oproj_gate_bm = gate_bm
+
+    mx.eval(moe._oproj_rms_weight)
+
+
+
+def _patch_gdn_proj_weights(attn):
+    """Merge all 4 GDN projection weights into contiguous buffers.
+
+    Concatenates in_proj_qkv/z/b/a weights, scales, biases into single
+    contiguous arrays for better memory locality in the fused GEMV kernel.
+    Stored on the GatedDeltaNet module as _merged_proj_*.
+    """
+    W_merged = mx.concatenate([
+        attn.in_proj_qkv.weight,
+        attn.in_proj_z.weight,
+        attn.in_proj_b.weight,
+        attn.in_proj_a.weight,
+    ], axis=0)
+    S_merged = mx.concatenate([
+        attn.in_proj_qkv.scales,
+        attn.in_proj_z.scales,
+        attn.in_proj_b.scales,
+        attn.in_proj_a.scales,
+    ], axis=0)
+    B_merged = mx.concatenate([
+        attn.in_proj_qkv.biases,
+        attn.in_proj_z.biases,
+        attn.in_proj_b.biases,
+        attn.in_proj_a.biases,
+    ], axis=0)
+    attn._merged_proj_w = W_merged
+    attn._merged_proj_s = S_merged
+    attn._merged_proj_b = B_merged
+    attn._merged_proj_dims = (
+        attn.in_proj_qkv.weight.shape[0],  # N_QKV = 8192
+        attn.in_proj_z.weight.shape[0],     # N_Z = 4096
+        attn.in_proj_b.weight.shape[0],     # N_B = 32
+        attn.in_proj_a.weight.shape[0],     # N_A = 32
+    )
+    mx.eval(W_merged, S_merged, B_merged)
+
+
+def _patch_gqa_proj_weights(attn):
+    """Merge GQA q_proj, k_proj, v_proj weights into contiguous buffers.
+
+    q_proj outputs (H_q * 2 * D) = interleaved [queries, gate] per head.
+    We permute rows so queries (H_q * D) come first, then gate (H_q * D),
+    then k_proj, then v_proj. This gives clean contiguous regions for
+    the fused GEMV kernel's TG routing.
+
+    Permutation for q_proj:
+        Original row layout: [head0_q[0:D], head0_gate[0:D], head1_q[0:D], ...]
+        After permutation:   [head0_q, head1_q, ..., head0_gate, head1_gate, ...]
+
+    Stored on Qwen3NextAttention as _merged_proj_*.
+    """
+    q = attn.q_proj
+    k = attn.k_proj
+    v = attn.v_proj
+
+    H_q = attn.num_attention_heads
+    D = attn.head_dim
+
+    # Permute q_proj weights: separate queries and gate rows
+    # q_proj.weight shape: (H_q * 2 * D, K / pack_factor) for 8-bit
+    # Reshape to (H_q, 2*D, ...), split into queries[:, :D, :] and gate[:, D:, :]
+    W_q = q.weight.reshape(H_q, 2 * D, -1)
+    S_q = q.scales.reshape(H_q, 2 * D, -1)
+    B_q = q.biases.reshape(H_q, 2 * D, -1)
+
+    W_queries = W_q[:, :D, :].reshape(H_q * D, -1)
+    W_gate = W_q[:, D:, :].reshape(H_q * D, -1)
+    S_queries = S_q[:, :D, :].reshape(H_q * D, -1)
+    S_gate = S_q[:, D:, :].reshape(H_q * D, -1)
+    B_queries = B_q[:, :D, :].reshape(H_q * D, -1)
+    B_gate = B_q[:, D:, :].reshape(H_q * D, -1)
+
+    # Merge: [queries, gate, keys, values]
+    W_merged = mx.contiguous(mx.concatenate([W_queries, W_gate, k.weight, v.weight], axis=0))
+    S_merged = mx.contiguous(mx.concatenate([S_queries, S_gate, k.scales, v.scales], axis=0))
+    B_merged = mx.contiguous(mx.concatenate([B_queries, B_gate, k.biases, v.biases], axis=0))
+
+    attn._merged_proj_w = W_merged
+    attn._merged_proj_s = S_merged
+    attn._merged_proj_b = B_merged
+    attn._merged_proj_dims = (
+        H_q * D,           # N_Q = 4096 (queries)
+        H_q * D,           # N_GATE = 4096 (gate)
+        k.weight.shape[0], # N_K = 512
+        v.weight.shape[0], # N_V = 512
+    )
+    mx.eval(W_merged, S_merged, B_merged)
+
+    # Pre-cache constant scalar arrays for kernel dispatch (avoid per-call creation)
+    N_Q, N_GATE, N_K, N_V = attn._merged_proj_dims
+    N_TOTAL = N_Q + N_GATE + N_K + N_V
+    K = q.weight.shape[1] * 4  # 8-bit: pack_factor=4
+    attn._kernel_scalars = {
+        # Dispatch 1: fused_gqa_projections
+        'K': mx.array(K, dtype=mx.int32),
+        'N_Q': mx.array(N_Q, dtype=mx.int32),
+        'N_GATE': mx.array(N_GATE, dtype=mx.int32),
+        'N_K': mx.array(N_K, dtype=mx.int32),
+        'N_TOTAL': mx.array(N_TOTAL, dtype=mx.int32),
+        'N_Q_TG': mx.array(ceil_div(N_Q, 8), dtype=mx.int32),
+        'N_GATE_TG': mx.array(ceil_div(N_GATE, 8), dtype=mx.int32),
+        'N_K_TG': mx.array(ceil_div(N_K, 8), dtype=mx.int32),
+        # Dispatch 4-5: custom SDPA
+        'scale': mx.array(attn.head_dim ** -0.5, dtype=mx.float32),
+        'H_Q': mx.array(attn.num_attention_heads, dtype=mx.int32),
+        'H_KV': mx.array(attn.num_key_value_heads, dtype=mx.int32),
+        'N_blocks': mx.array(128, dtype=mx.int32),
+    }
+    mx.eval(*attn._kernel_scalars.values())
+
+    # Precompute grid/TG dims for Dispatch 1
+    N_V_TG = ceil_div(N_V, 8)
+    attn._d1_total_tg = ceil_div(N_Q, 8) + ceil_div(N_GATE, 8) + ceil_div(N_K, 8) + N_V_TG
+
+    # Precompute RoPE inv_freq for fused norm+rope kernel (Dispatch 2)
+    # inv_freq[d] = theta^(-d / half_dims) for d in {0, ..., half_dims-1}
+    rope_dims = attn.rope.dims           # 64 (partial_rotary_factor * head_dim)
+    half_dims = rope_dims // 2            # 32
+    theta = attn.rope.base                # 10000000
+    d_indices = mx.arange(half_dims, dtype=mx.float32)
+    attn._rope_inv_freq = theta ** (-d_indices / half_dims)
+    mx.eval(attn._rope_inv_freq)
+
+
+
+def make_qwen_random_cache(layer, config, prefill_len):
+    """Create a pre-filled cache for a single Qwen3.5 decoder layer.
+
+    GatedDeltaNet layers get ArraysCache(size=2) with fixed-size state:
+      cache[0] = conv state: (B, conv_kernel_size-1, conv_dim) bf16
+      cache[1] = SSM state:  (B, num_v_heads, head_k_dim, head_v_dim) bf16
+
+    GQA layers get KVCache with prefill_len tokens:
+      keys:   (B, n_kv_heads, alloc_len, head_dim) bf16
+      values: (B, n_kv_heads, alloc_len, head_dim) bf16
+    """
+    if layer.is_linear:
+        from mlx_lm.models.cache import ArraysCache
+        cache = ArraysCache(size=2)
+        attn = layer.linear_attn
+        cache[0] = mx.random.normal(
+            (1, attn.conv_kernel_size - 1, attn.conv_dim)
+        ).astype(mx.bfloat16)
+        cache[1] = mx.random.normal(
+            (1, attn.num_v_heads, attn.head_k_dim, attn.head_v_dim)
+        ).astype(mx.bfloat16)
+        return cache
+    else:
+        from mlx_lm.models.cache import KVCache
+        cache = KVCache()
+        n_steps = (prefill_len + KVCache.step - 1) // KVCache.step
+        alloc_len = n_steps * KVCache.step
+        n_kv = config.num_key_value_heads
+        hd = config.head_dim
+        cache.keys = mx.random.normal((1, n_kv, alloc_len, hd)).astype(mx.bfloat16)
+        cache.values = mx.random.normal((1, n_kv, alloc_len, hd)).astype(mx.bfloat16)
+        cache.offset = prefill_len
+        return cache
+
+
+def build_model(n_experts=16, n_layers=1, top_k=4,
+                hidden_size=4096, moe_intermediate_size=1024,
+                shared_expert_intermediate_size=2048, tp=1,
+                n_attn_heads=32, n_kv_heads=2,
+                lin_v_heads=64, lin_k_heads=16,
+                head_dim=256):
+    """Build Qwen3.5 MoE decoder layers with 8-bit gs=64 quantization.
+
+    Matches real mlx-community Qwen3.5 quantization:
+      Everything 8-bit gs=64 (gate, experts, shared expert, shared_expert_gate, attention/SSM).
+      RMSNorm weights: bf16.
+
+    Default dimensions are for Qwen3.5-397B-A17B. For 35B-A3B, pass:
+      n_attn_heads=16, lin_v_heads=32, hidden_size=2048
+
+    Uses qwen3_5.DecoderLayer (same class as real model) with hybrid attention:
+      3/4 GatedDeltaNet (SSM-like), 1/4 full attention (full_attention_interval=4).
+
+    TP=2 halves all column-parallel dimensions.
+
+    Returns:
+        layers: list of DecoderLayer instances
+        config: TextModelArgs
+        GROUP_SIZE: int (64)
+    """
+    GROUP_SIZE = 64
+    BITS = 8
+
+    # Apply TP sharding
+    moe_inter = moe_intermediate_size // tp
+    shared_inter = shared_expert_intermediate_size // tp
+    n_attn_heads_tp = n_attn_heads // tp
+    n_kv_heads_tp = max(1, n_kv_heads // tp)
+    lin_v_heads_tp = lin_v_heads // tp
+    lin_k_heads_tp = lin_k_heads // tp
+
+    config = TextModelArgs(
+        model_type="qwen3_5_moe",
+        hidden_size=hidden_size,
+        num_hidden_layers=n_layers,
+        intermediate_size=moe_inter,
+        num_attention_heads=n_attn_heads_tp,
+        num_key_value_heads=n_kv_heads_tp,
+        linear_num_value_heads=lin_v_heads_tp,
+        linear_num_key_heads=lin_k_heads_tp,
+        linear_key_head_dim=128,
+        linear_value_head_dim=128,
+        linear_conv_kernel_dim=4,
+        num_experts=n_experts,
+        num_experts_per_tok=top_k,
+        decoder_sparse_step=1,
+        shared_expert_intermediate_size=shared_inter,
+        moe_intermediate_size=moe_inter,
+        norm_topk_prob=True,
+        rms_norm_eps=1e-6,
+        vocab_size=248320,
+        head_dim=head_dim,
+        full_attention_interval=4,
+        max_position_embeddings=262144,
+        rope_theta=10000000,
+        partial_rotary_factor=0.25,
+        rope_parameters={
+            "type": "default",
+            "rope_theta": 10000000,
+            "partial_rotary_factor": 0.25,
+        },
+    )
+
+    tp_str = f" (TP={tp})" if tp > 1 else ""
+    print(f"  Config: {n_layers} layer(s), {n_experts} experts, top_k={top_k}, "
+          f"hidden={hidden_size}, inter={moe_inter}, shared={shared_inter}{tp_str}")
+    print(f"  Quant: {BITS}-bit gs={GROUP_SIZE} (all weights)")
+
+    layers = [DecoderLayer(config, idx) for idx in range(n_layers)]
+
+    for li, layer in enumerate(layers):
+        # Cast all attention/SSM params to bf16 before quantizing, matching
+        # real safetensors model where all non-quantized params are bf16.
+        # nn.quantize only touches nn.Linear; other params (conv1d, dt_bias,
+        # norm, A_log) must be cast manually.
+        attn_mod = layer.linear_attn if layer.is_linear else layer.self_attn
+        for name, mod in attn_mod.named_modules():
+            if isinstance(mod, nn.Linear):
+                mod.weight = mod.weight.astype(mx.bfloat16)
+            elif isinstance(mod, nn.Conv1d):
+                mod.weight = mod.weight.astype(mx.bfloat16)
+        # Cast leaf parameters (dt_bias, norm.weight, q/k_norm) to bf16
+        # A_log stays f32 (matches real model)
+        if layer.is_linear:
+            gdn = layer.linear_attn
+            gdn.dt_bias = gdn.dt_bias.astype(mx.bfloat16)
+            gdn.norm.weight = gdn.norm.weight.astype(mx.bfloat16)
+        else:
+            gqa = layer.self_attn
+            gqa.q_norm.weight = gqa.q_norm.weight.astype(mx.bfloat16)
+            gqa.k_norm.weight = gqa.k_norm.weight.astype(mx.bfloat16)
+        nn.quantize(attn_mod, bits=BITS, group_size=GROUP_SIZE)
+        mx.eval(attn_mod.parameters())
+
+        # RMSNorm to bf16 (norms are never quantized)
+        layer.input_layernorm.weight = layer.input_layernorm.weight.astype(mx.bfloat16)
+        layer.post_attention_layernorm.weight = layer.post_attention_layernorm.weight.astype(mx.bfloat16)
+        mx.eval(layer.input_layernorm.weight, layer.post_attention_layernorm.weight)
+
+        # MoE block: quantize everything to 8-bit gs=64
+        moe = layer.mlp
+        if isinstance(moe, Qwen3NextSparseMoeBlock):
+            # Gate: random init (zeros get optimized away), then quantize.
+            # nn.quantize on a leaf nn.Linear is a no-op (walks children, finds none).
+            # Use QuantizedLinear.from_linear directly.
+            moe.gate.weight = (
+                mx.random.normal(moe.gate.weight.shape) * 0.01
+            ).astype(mx.float32)
+            moe.gate = nn.QuantizedLinear.from_linear(
+                moe.gate, group_size=GROUP_SIZE, bits=BITS)
+            mx.eval(moe.gate.parameters())
+
+            # Routed experts: quantize per-projection to limit peak memory
+            nn.quantize(moe.switch_mlp, bits=BITS, group_size=GROUP_SIZE)
+            mx.eval(moe.switch_mlp.gate_proj.parameters())
+            mx.eval(moe.switch_mlp.up_proj.parameters())
+            mx.eval(moe.switch_mlp.down_proj.parameters())
+
+            # Shared expert: quantize to 8-bit (matching real model)
+            nn.quantize(moe.shared_expert, bits=BITS, group_size=GROUP_SIZE)
+            mx.eval(moe.shared_expert.parameters())
+
+            # shared_expert_gate: quantize to 8-bit gs=64 (leaf nn.Linear fix)
+            moe.shared_expert_gate = nn.QuantizedLinear.from_linear(
+                moe.shared_expert_gate, group_size=GROUP_SIZE, bits=BITS)
+            mx.eval(moe.shared_expert_gate.parameters())
+
+        if (li + 1) % 10 == 0 or li == 0:
+            print(f"  Layer {li+1}/{n_layers} ready")
+
+    return layers, config, GROUP_SIZE
@@ -0,0 +1,199 @@
+"""Decoder layer __call__ variants for Qwen3.5.
+
+Two modes:
+  _fused_decoder_call: passes residual to fused MoE epilogue (~15 dispatches)
+  _oproj_decoder_call: fuses o_proj + RMSNorm + gate GEMV (4 dispatches)
+
+Attention patches for oproj mode:
+  _pre_oproj_attention_call: Qwen3NextAttention.__call__ that skips o_proj
+  _pre_oproj_qwen35_linear_attn_call: qwen3_5.GatedDeltaNet.__call__ that skips out_proj
+
+Note: qwen3_5.GatedDeltaNet (used by DecoderLayer) is a DIFFERENT class from
+qwen3_next.Qwen3NextGatedDeltaNet. They have different projection layouts:
+  - qwen3_5.GatedDeltaNet: separate in_proj_qkv, in_proj_z, in_proj_b, in_proj_a
+  - qwen3_next.Qwen3NextGatedDeltaNet: merged in_proj_qkvz, in_proj_ba
+The patch must match qwen3_5.GatedDeltaNet's __call__ structure.
+"""
+
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.nn.layers.activations import silu as nn_silu
+
+# Map moe block id → parent decoder layer (avoids circular refs in model tree)
+_parent_layer_map = {}
+
+
+def _fused_decoder_call(self, x, mask=None, cache=None):
+    """Decoder layer with residual passed to fused MoE epilogue.
+
+    Replaces:
+      h = x + attn(norm(x))
+      out = h + mlp(norm(h))        # mlp returns MoE output, then adds h
+
+    With:
+      h = x + attn(norm(x))
+      out = mlp(norm(h), _residual=h)  # epilogue fuses: moe_out + h
+    """
+    if self.is_linear:
+        r = self.linear_attn(self.input_layernorm(x), mask, cache)
+    else:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+    h = x + r
+    out = self.mlp(self.post_attention_layernorm(h), _residual=h)
+    return out  # already includes residual add from epilogue
+
+
+def _oproj_decoder_call(self, x, mask=None, cache=None):
+    """Decoder with fused o_proj + RMSNorm + gate GEMV (oproj 4-dispatch mode).
+
+    Skips o_proj, addmm, and post_attention_layernorm — all fused into Dispatch 1.
+    Attention __call__ is patched to return pre-o_proj output.
+
+    Flow:
+      pre_oproj = attn(input_layernorm(x))   # returns BEFORE o_proj
+      MoE receives (pre_oproj, residual=x) and handles o_proj + RMSNorm + gate internally
+    """
+    if self.is_linear:
+        pre_oproj = self.linear_attn(self.input_layernorm(x), mask, cache)
+    else:
+        pre_oproj = self.self_attn(self.input_layernorm(x), mask, cache)
+    _parent_layer_map[id(self.mlp)] = self
+    return self.mlp(pre_oproj, _residual=x)
+
+
+def _vanilla_decoder_call(self, x, mask=None, cache=None):
+    """Original vanilla DecoderLayer.__call__ (fallback for B>8 or S>1)."""
+    if self.is_linear:
+        r = self.linear_attn(self.input_layernorm(x), mask, cache)
+    else:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+    h = x + r
+    out = self.mlp(self.post_attention_layernorm(h))
+    return h + out
+
+
+def _fused_gdn_decoder_call(self, x, mask=None, cache=None):
+    """Decoder with batched fused kernels. Falls back to vanilla for B>8 or S>1.
+
+    When fused: attention returns pre-out_proj output, MoE handles oproj + gate + experts.
+    When vanilla: original DecoderLayer flow (attention + residual + layernorm + MoE).
+    """
+    B = x.shape[0]
+    S = x.shape[1]
+
+    # Full vanilla fallback for large batch or prefill
+    if B > 8 or S > 1:
+        return _vanilla_decoder_call(self, x, mask, cache)
+
+    # Fused path: attention returns pre-oproj, MoE handles the rest
+    if self.is_linear:
+        pre_oproj = self.linear_attn(self.input_layernorm(x), mask, cache)
+    else:
+        pre_oproj = self.self_attn(self.input_layernorm(x), mask, cache)
+    _parent_layer_map[id(self.mlp)] = self
+    return self.mlp(pre_oproj, _residual=x)
+
+
+def _pre_oproj_attention_call(self, x, mask=None, cache=None):
+    """Qwen3NextAttention.__call__ that returns pre-o_proj output.
+
+    Identical to original except final line returns output*sigmoid(gate)
+    instead of self.o_proj(output*sigmoid(gate)).
+    """
+    B, L, D = x.shape
+    q_proj_output = self.q_proj(x)
+    queries, gate = mx.split(
+        q_proj_output.reshape(B, L, self.num_attention_heads, -1), 2, axis=-1
+    )
+    gate = gate.reshape(B, L, -1)
+    keys, values = self.k_proj(x), self.v_proj(x)
+    queries = self.q_norm(queries).transpose(0, 2, 1, 3)
+    keys = self.k_norm(
+        keys.reshape(B, L, self.num_key_value_heads, -1)
+    ).transpose(0, 2, 1, 3)
+    values = values.reshape(B, L, self.num_key_value_heads, -1).transpose(
+        0, 2, 1, 3
+    )
+    if cache is not None:
+        queries = self.rope(queries, offset=cache.offset)
+        keys = self.rope(keys, offset=cache.offset)
+        keys, values = cache.update_and_fetch(keys, values)
+    else:
+        queries = self.rope(queries)
+        keys = self.rope(keys)
+
+    from mlx_lm.models.qwen3_next import scaled_dot_product_attention
+    output = scaled_dot_product_attention(
+        queries, keys, values, cache=cache, scale=self.scale, mask=mask
+    )
+    output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+    return output * mx.sigmoid(gate)  # skip o_proj
+
+
+def _pre_oproj_qwen35_linear_attn_call(
+    self,
+    inputs: mx.array,
+    mask: Optional[mx.array] = None,
+    cache: Optional[Any] = None,
+) -> mx.array:
+    """qwen3_5.GatedDeltaNet.__call__ that returns pre-out_proj output.
+
+    Identical to qwen3_5.GatedDeltaNet.__call__ except final line returns
+    out.reshape(B,S,-1) instead of self.out_proj(out.reshape(B,S,-1)).
+
+    Note: this targets qwen3_5.GatedDeltaNet (separate projections), NOT
+    qwen3_next.Qwen3NextGatedDeltaNet (merged projections). They are
+    different classes with different __call__ bodies.
+    """
+    from mlx_lm.models.gated_delta import gated_delta_update
+
+    B, S, _ = inputs.shape
+
+    qkv = self.in_proj_qkv(inputs)
+    z = self.in_proj_z(inputs).reshape(B, S, self.num_v_heads, self.head_v_dim)
+    b = self.in_proj_b(inputs)
+    a = self.in_proj_a(inputs)
+
+    if cache is not None and cache[0] is not None:
+        conv_state = cache[0]
+    else:
+        conv_state = mx.zeros(
+            (B, self.conv_kernel_size - 1, self.conv_dim),
+            dtype=inputs.dtype,
+        )
+
+    if mask is not None:
+        qkv = mx.where(mask[..., None], qkv, 0)
+    conv_input = mx.concatenate([conv_state, qkv], axis=1)
+    if cache is not None:
+        cache[0] = conv_input[:, -(self.conv_kernel_size - 1) :]
+    conv_out = nn.silu(self.conv1d(conv_input))
+
+    q, k, v = [
+        t.reshape(B, S, h, d)
+        for t, h, d in zip(
+            mx.split(conv_out, [self.key_dim, 2 * self.key_dim], -1),
+            [self.num_k_heads, self.num_k_heads, self.num_v_heads],
+            [self.head_k_dim, self.head_k_dim, self.head_v_dim],
+        )
+    ]
+
+    state = cache[1] if cache else None
+    inv_scale = k.shape[-1] ** -0.5
+    q = (inv_scale**2) * mx.fast.rms_norm(q, None, 1e-6)
+    k = inv_scale * mx.fast.rms_norm(k, None, 1e-6)
+
+    out, state = gated_delta_update(
+        q, k, v, a, b,
+        self.A_log, self.dt_bias,
+        state, mask,
+        use_kernel=True,
+    )
+
+    if cache is not None:
+        cache[1] = state
+
+    out = self.norm(out, z)
+    return out.reshape(B, S, -1)  # skip out_proj
@@ -0,0 +1,161 @@
+"""Fused GDN attention __call__ for qwen3_5.GatedDeltaNet (Dispatches 2-5).
+
+Replaces the vanilla GatedDeltaNet.__call__ with fused kernel dispatches:
+  Dispatch 2: fused_gdn_projections — merged 8-bit GEMV + conv1d + SiLU(qkv) + SiLU(z)
+              + sigmoid(b)→beta + g=exp(-exp(A_log)*softplus(a+dt_bias))
+  Dispatch 3: fused_qk_rmsnorm — per-head L2-norm on q (×Dk^(-½)) and k
+  Dispatch 4: gated_delta_kernel — GDN recurrence (receives pre-computed g, beta)
+  Dispatch 5: fused_rms_norm_gated — RMSNorm(out, weight) × z_silu
+
+All 4 projection weights are pre-merged into contiguous buffers at patch time
+(_patch_gdn_proj_weights) for better memory locality.
+
+g/beta computation is fused into Dispatch 2 epilogues, eliminating ~8 micro-
+dispatches that gated_delta_update would otherwise generate.
+
+Fused path is decode-only (S=1). For prefill (S>1), falls back to vanilla ops.
+
+Returns pre-out_proj output (same interface as _pre_oproj_qwen35_linear_attn_call).
+Dispatch 1 (input_layernorm) is handled by the decoder.
+Dispatch 6 (oproj_gate_gemv) is handled by the MoE __call__.
+"""
+
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .kernels.batched_fused_gdn_projections_8bit import batched_fused_gdn_projections as fused_gdn_projections
+from .kernels.fused_qk_rmsnorm import fused_qk_rmsnorm
+from .kernels.fused_rms_norm_gated import fused_rms_norm_gated
+
+
+def _vanilla_gdn_call(self, inputs, mask, cache):
+    """Vanilla GDN path for prefill (S>1). Returns pre-out_proj output."""
+    from mlx_lm.models.gated_delta import gated_delta_update
+
+    B, S, _ = inputs.shape
+
+    qkv = self.in_proj_qkv(inputs)
+    z = self.in_proj_z(inputs).reshape(B, S, self.num_v_heads, self.head_v_dim)
+    b = self.in_proj_b(inputs)
+    a = self.in_proj_a(inputs)
+
+    if cache is not None and cache[0] is not None:
+        conv_state = cache[0]
+    else:
+        conv_state = mx.zeros(
+            (B, self.conv_kernel_size - 1, self.conv_dim),
+            dtype=inputs.dtype,
+        )
+
+    if mask is not None:
+        qkv = mx.where(mask[..., None], qkv, 0)
+    conv_input = mx.concatenate([conv_state, qkv], axis=1)
+    if cache is not None:
+        cache[0] = conv_input[:, -(self.conv_kernel_size - 1):]
+    conv_out = nn.silu(self.conv1d(conv_input))
+
+    q, k, v = [
+        t.reshape(B, S, h, d)
+        for t, h, d in zip(
+            mx.split(conv_out, [self.key_dim, 2 * self.key_dim], -1),
+            [self.num_k_heads, self.num_k_heads, self.num_v_heads],
+            [self.head_k_dim, self.head_k_dim, self.head_v_dim],
+        )
+    ]
+
+    state = cache[1] if cache else None
+    inv_scale = k.shape[-1] ** -0.5
+    q = inv_scale * q * mx.rsqrt(
+        (q * q).sum(axis=-1, keepdims=True) + 1e-6
+    )
+    k = k * mx.rsqrt(
+        (k * k).sum(axis=-1, keepdims=True) + 1e-6
+    )
+
+    out, state = gated_delta_update(
+        q, k, v, a, b,
+        self.A_log, self.dt_bias,
+        state, mask,
+        use_kernel=True,
+    )
+
+    if cache is not None:
+        cache[1] = state
+
+    out = self.norm(out, z)
+    return self.out_proj(out.reshape(B, S, -1))  # include out_proj for vanilla decoder
+
+
+def _fused_gdn_call(
+    self,
+    inputs: mx.array,
+    mask: Optional[mx.array] = None,
+    cache: Optional[Any] = None,
+) -> mx.array:
+    """Fused GDN attention: merged projections + existing GDN kernel.
+
+    Decode (S=1): uses fused kernels with merged weight buffers.
+    Prefill (S>1): falls back to vanilla ops.
+
+    Returns pre-out_proj output [B, S, value_dim] for Dispatch 6.
+    """
+    B, S, _ = inputs.shape
+
+    # Vanilla fallback: fused kernels are decode-only (S=1, B<=8)
+    if S > 1 or B > 8:
+        return _vanilla_gdn_call(self, inputs, mask, cache)
+
+    from mlx_lm.models.gated_delta import gated_delta_kernel
+
+    # ── Cache: conv state ──
+    if cache is not None and cache[0] is not None:
+        conv_state = cache[0]
+    else:
+        conv_state = mx.zeros(
+            (B, self.conv_kernel_size - 1, self.conv_dim),
+            dtype=inputs.dtype,
+        )
+
+    # ── Dispatch 2: fused projections (merged GEMV + conv + SiLU + g/beta) ──
+    qkv_conv_silu, z_silu, beta, g, conv_state_out = fused_gdn_projections(
+        inputs,
+        self._merged_proj_w, self._merged_proj_s, self._merged_proj_b,
+        self._merged_proj_dims,
+        conv_state, self.conv1d.weight,
+        self.A_log, self.dt_bias,
+        batch_size=B,
+    )
+
+    if cache is not None:
+        cache[0] = conv_state_out
+
+    # ── Dispatch 3: fused Q/K L2-norm ──
+    qk_normed = fused_qk_rmsnorm(qkv_conv_silu, batch_size=B)
+
+    # ── Split q, k from normed output; v from conv output ──
+    q = qk_normed[:, :, :self.key_dim].reshape(B, S, self.num_k_heads, self.head_k_dim)
+    k = qk_normed[:, :, self.key_dim:].reshape(B, S, self.num_k_heads, self.head_k_dim)
+    v = qkv_conv_silu[:, :, 2 * self.key_dim:].reshape(B, S, self.num_v_heads, self.head_v_dim)
+
+    # ── Dispatch 4: GDN recurrence with pre-computed g/beta ──
+    state = cache[1] if cache else None
+    if state is None:
+        state = mx.zeros(
+            (B, self.num_v_heads, self.head_v_dim, self.head_k_dim),
+            dtype=inputs.dtype,
+        )
+
+    out, state_new = gated_delta_kernel(
+        q, k, v, g, beta, state, mask,
+    )
+
+    if cache is not None:
+        cache[1] = state_new
+
+    # ── Dispatch 5: fused RMSNorm × z_silu ──
+    norm_weight = self.norm.weight
+    result = fused_rms_norm_gated(out, z_silu, norm_weight, batch_size=B)
+
+    return result  # [B, S, value_dim] — skip out_proj (handled by Dispatch 6)
@@ -0,0 +1,265 @@
+"""Batched fused GDN projections for Qwen3.5-35B-A3B, batch_size 1..8.
+
+Register-level weight sharing: each TG loads weights once, computes B outputs.
+Adapts fused_gdn_projections_8bit with the same pattern as
+batched_fused_gqa_projections_8bit.
+
+4 regions with different epilogues:
+  - QKV: GEMV → conv1d(4-tap) → SiLU → bf16 + cache update
+  - Z:   GEMV → SiLU → f32
+  - B:   GEMV → sigmoid → f32 (beta for GDN kernel)
+  - A:   GEMV → g=exp(-exp(A_log)*softplus(a+dt_bias)) → f32
+
+All constants baked into Metal source. B unrolled at code-generation time.
+
+Grid: (32, total_tg * 2, 1), TG: (32, 2, 1)
+No grid z for batch — batch is handled in registers.
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_batched_fused_gdn_proj_source(K, N_QKV, N_Z, N_B, N_A, B, group_size=64):
+    gs = int(group_size)
+    sc_stride = 256 // gs
+    slid_div = gs // 8
+    N_TOTAL = N_QKV + N_Z + N_B + N_A
+    K_groups = K // gs
+    N_QKV_TG = ceil_div(N_QKV, 8)
+    N_Z_TG = ceil_div(N_Z, 8)
+    N_B_TG = ceil_div(N_B, 8)
+
+    # Per-batch x loading (B unrolled)
+    x_load = "\n".join(f"""
+        float x{b}_thread[VALUES_PER_THREAD]; float xsum{b} = 0;
+        for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+            float xi = float(x[{b} * K + x_base + i]); x{b}_thread[i] = xi; xsum{b} += xi;
+        }}""" for b in range(B))
+
+    # Per-batch dot product with weights in registers
+    qdot = "\n".join(f"""
+            float accum{b} = 0;
+            for (int i = 0; i < VALUES_PER_THREAD; i++) accum{b} += x{b}_thread[i] * w_vals[i];
+            result{b}[row] += s_val * accum{b} + xsum{b} * b_val;""" for b in range(B))
+
+    result_decls = " ".join(f"float result{b}[4] = {{0,0,0,0}};" for b in range(B))
+    simd_reduce = "\n    ".join(
+        f"for (int row = 0; row < 4; row++) result{b}[row] = simd_sum(result{b}[row]);" for b in range(B))
+
+    # QKV epilogue: conv1d + SiLU + cache update, per batch
+    qkv_write = "\n".join(f"""
+            if (slid < 4u && c < N_QKV) {{
+                float qkv_val = result{b}[slid];
+                long cs_base = (long){b} * 3 * conv_dim;
+                float s0 = float(conv_state[cs_base + 0 * conv_dim + c]);
+                float s1 = float(conv_state[cs_base + 1 * conv_dim + c]);
+                float s2 = float(conv_state[cs_base + 2 * conv_dim + c]);
+                float conv_out = float(conv_w[c * 4 + 0]) * s0
+                               + float(conv_w[c * 4 + 1]) * s1
+                               + float(conv_w[c * 4 + 2]) * s2
+                               + float(conv_w[c * 4 + 3]) * qkv_val;
+                float silu_out = conv_out / (1.0f + metal::exp(-conv_out));
+                conv_state_out[cs_base + 0 * conv_dim + c] = static_cast<bfloat16_t>(s1);
+                conv_state_out[cs_base + 1 * conv_dim + c] = static_cast<bfloat16_t>(s2);
+                conv_state_out[cs_base + 2 * conv_dim + c] = static_cast<bfloat16_t>(qkv_val);
+                qkv_out[{b} * conv_dim + c] = static_cast<bfloat16_t>(silu_out);
+            }}""" for b in range(B))
+
+    # Z epilogue: SiLU per batch
+    z_write = "\n".join(f"""
+            if (slid < 4u && z_row < N_Z) {{
+                float val = result{b}[slid];
+                z_silu_out[{b} * N_Z + z_row] = val / (1.0f + metal::exp(-val));
+            }}""" for b in range(B))
+
+    # B epilogue: sigmoid per batch
+    b_write = "\n".join(f"""
+            if (slid < 4u && b_row < N_B) {{
+                b_out[{b} * N_B + b_row] = 1.0f / (1.0f + metal::exp(-result{b}[slid]));
+            }}""" for b in range(B))
+
+    # A epilogue: g computation per batch
+    a_write = "\n".join(f"""
+            if (slid < 4u && a_row < N_A_val) {{
+                float a_val = result{b}[slid];
+                float dt = float(dt_bias_arr[a_row]);
+                float x_g = a_val + dt;
+                float sp = (x_g > 20.0f) ? x_g : metal::log(1.0f + metal::exp(x_g));
+                float g_val = metal::exp(-metal::exp(float(A_log_arr[a_row])) * sp);
+                a_out[{b} * N_A_val + a_row] = g_val;
+            }}""" for b in range(B))
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int GROUP_SIZE = {gs};
+    const int SC_STRIDE = {sc_stride};
+    const int SLID_DIV = {slid_div};
+    const int K = {K};
+    const int K_groups = {K_groups};
+    const int N_QKV = {N_QKV};
+    const int N_Z = {N_Z};
+    const int N_B = {N_B};
+    const int N_TOTAL = {N_TOTAL};
+    const int N_QKV_TG = {N_QKV_TG};
+    const int N_Z_TG = {N_Z_TG};
+    const int N_B_TG = {N_B_TG};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+    int tg = tgid.y;
+
+    int out_row, region;
+    if (tg < N_QKV_TG) {{
+        region = 0; out_row = tg * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_QKV_TG + N_Z_TG) {{
+        region = 1; out_row = N_QKV + (tg - N_QKV_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_QKV_TG + N_Z_TG + N_B_TG) {{
+        region = 2; out_row = N_QKV + N_Z + (tg - N_QKV_TG - N_Z_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else {{
+        region = 3; out_row = N_QKV + N_Z + N_B + (tg - N_QKV_TG - N_Z_TG - N_B_TG) * 8 + sgid * RESULTS_PER_SG;
+    }}
+    if (out_row >= N_TOTAL) return;
+
+    // Weight pointers (shared across all batch elements)
+    const device uint8_t* ws = (const device uint8_t*)W_merged + (long)out_row * K + slid * VALUES_PER_THREAD;
+    const device bfloat16_t* sc = (const device bfloat16_t*)S_merged + (long)out_row * K_groups + slid / SLID_DIV;
+    const device bfloat16_t* bi = (const device bfloat16_t*)B_merged + (long)out_row * K_groups + slid / SLID_DIV;
+
+    {result_decls}
+    int x_base = slid * VALUES_PER_THREAD;
+
+    // K-loop: load weights into registers once, compute {B} batch elements
+    for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{
+{x_load}
+
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            const device uint8_t* wl = ws + row * K;
+            float s_val = float(sc[row * K_groups]);
+            float b_val = float(bi[row * K_groups]);
+            float w_vals[VALUES_PER_THREAD];
+            for (int i = 0; i < VALUES_PER_THREAD; i++) w_vals[i] = float(wl[i]);
+{qdot}
+        }}
+
+        ws += BLOCK_SIZE; sc += SC_STRIDE; bi += SC_STRIDE; x_base += BLOCK_SIZE;
+    }}
+
+    {simd_reduce}
+
+    // Region-specific epilogues for all {B} batches
+    if (region == 0) {{
+        int c = out_row + (int)slid;
+        int conv_dim = N_QKV;
+{qkv_write}
+    }} else if (region == 1) {{
+        int z_row = out_row - N_QKV + (int)slid;
+{z_write}
+    }} else if (region == 2) {{
+        int b_row = out_row - N_QKV - N_Z + (int)slid;
+{b_write}
+    }} else {{
+        int a_row = out_row - N_QKV - N_Z - N_B + (int)slid;
+        int N_A_val = N_TOTAL - N_QKV - N_Z - N_B;
+{a_write}
+    }}
+"""
+
+
+_batched_gdn_proj_cache = {}
+
+
+def _get_batched_gdn_proj_kernel(K, N_QKV, N_Z, N_B, N_A, B, group_size=64):
+    key = (K, N_QKV, N_Z, N_B, N_A, B, group_size)
+    if key not in _batched_gdn_proj_cache:
+        _batched_gdn_proj_cache[key] = mx.fast.metal_kernel(
+            name=f"batched_fused_gdn_proj_K{K}_NQKV{N_QKV}_B{B}",
+            input_names=[
+                "x",
+                "W_merged", "S_merged", "B_merged",
+                "conv_state", "conv_w",
+                "A_log_arr", "dt_bias_arr",
+            ],
+            output_names=["qkv_out", "z_silu_out", "b_out", "a_out", "conv_state_out"],
+            source=_gen_batched_fused_gdn_proj_source(K, N_QKV, N_Z, N_B, N_A, B, group_size),
+        )
+    return _batched_gdn_proj_cache[key]
+
+
+def batched_fused_gdn_projections(
+    x,
+    W_merged, S_merged, B_merged,
+    proj_dims,
+    conv_state, conv_weights,
+    A_log, dt_bias,
+    batch_size=1,
+):
+    """Batched fused GDN projections with register-level weight sharing.
+
+    Same as fused_gdn_projections but loads weights once per TG and computes
+    B outputs from registers. No grid z for batch.
+
+    Args:
+        x: [B, 1, K] bf16 — post-RMSNorm hidden state
+        W_merged, S_merged, B_merged: merged quantized weights
+        proj_dims: (N_QKV, N_Z, N_B, N_A)
+        conv_state: [B, 3, conv_dim] bf16
+        conv_weights: [conv_dim, 4, 1] or [conv_dim, 4] bf16
+        A_log: [Hv] f32, dt_bias: [Hv] f32
+        batch_size: int (1..8)
+
+    Returns:
+        qkv_conv_silu: [B, 1, N_QKV] bf16
+        z_silu:        [B, 1, N_Z] f32
+        beta:          [B, 1, N_B] f32
+        g:             [B, 1, N_A] f32
+        conv_state_out: [B, 3, N_QKV] bf16
+    """
+    B = batch_size
+
+    N_QKV, N_Z, N_B, N_A = proj_dims
+    K = x.shape[-1]
+
+    kern = _get_batched_gdn_proj_kernel(K, N_QKV, N_Z, N_B, N_A, B)
+
+    N_QKV_TG = ceil_div(N_QKV, 8)
+    N_Z_TG = ceil_div(N_Z, 8)
+    N_B_TG = ceil_div(N_B, 8)
+    N_A_TG = ceil_div(N_A, 8)
+    total_tg = N_QKV_TG + N_Z_TG + N_B_TG + N_A_TG
+
+    conv_w_flat = conv_weights.reshape(-1, 4) if conv_weights.ndim == 3 else conv_weights
+    x_flat = x.reshape(B, K)
+
+    results = kern(
+        inputs=[
+            x_flat,
+            W_merged, S_merged, B_merged,
+            conv_state, conv_w_flat,
+            A_log, dt_bias,
+        ],
+        output_shapes=[
+            (B * N_QKV,),
+            (B * N_Z,),
+            (B * N_B,),
+            (B * N_A,),
+            (B * 3 * N_QKV,),
+        ],
+        output_dtypes=[mx.bfloat16, mx.float32, mx.float32, mx.float32, mx.bfloat16],
+        grid=(32, total_tg * 2, 1),   # No grid z — batch in registers
+        threadgroup=(32, 2, 1),
+    )
+
+    qkv_out = results[0].reshape(B, 1, N_QKV)
+    z_silu = results[1].reshape(B, 1, N_Z)
+    beta = results[2].reshape(B, 1, N_B)
+    g = results[3].reshape(B, 1, N_A)
+    conv_state_out = results[4].reshape(B, 3, N_QKV)
+
+    return qkv_out, z_silu, beta, g, conv_state_out
@@ -0,0 +1,205 @@
+"""Batched fused GQA projections (Dispatch 1) for batch_size 1..8.
+
+Adapts fused_gqa_projections_8bit for B>1 with register-level weight sharing.
+Each TG loads weights once, computes B outputs from registers.
+
+4 regions with different epilogues (same as B=1):
+  - Queries: GEMV → raw bf16
+  - Gate: GEMV → sigmoid → f32
+  - Keys: GEMV → raw bf16
+  - Values: GEMV → raw bf16
+
+All constants baked into Metal source. B unrolled at code-generation time.
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_batched_fused_gqa_proj_source(K, N_Q, N_GATE, N_K, N_V, B, group_size=64):
+    gs = int(group_size)
+    sc_stride = 256 // gs
+    slid_div = gs // 8
+    N_TOTAL = N_Q + N_GATE + N_K + N_V
+    K_groups = K // gs
+    N_Q_TG = ceil_div(N_Q, 8)
+    N_GATE_TG = ceil_div(N_GATE, 8)
+    N_K_TG = ceil_div(N_K, 8)
+
+    # Per-batch x loading
+    x_load = "\n".join(f"""
+        float x{b}_thread[VALUES_PER_THREAD]; float xsum{b} = 0;
+        for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+            float xi = float(x[{b} * K + x_base + i]); x{b}_thread[i] = xi; xsum{b} += xi;
+        }}""" for b in range(B))
+
+    # Per-batch qdot (weights in registers)
+    qdot = "\n".join(f"""
+            float accum{b} = 0;
+            for (int i = 0; i < VALUES_PER_THREAD; i++) accum{b} += x{b}_thread[i] * w_vals[i];
+            result{b}[row] += s_val * accum{b} + xsum{b} * b_val;""" for b in range(B))
+
+    result_decls = " ".join(f"float result{b}[4] = {{0,0,0,0}};" for b in range(B))
+    simd_reduce = "\n    ".join(
+        f"for (int row = 0; row < 4; row++) result{b}[row] = simd_sum(result{b}[row]);" for b in range(B))
+
+    # Queries epilogue (bf16 write per batch)
+    q_write = "\n".join(f"""
+            if (slid < 4u && q_row < N_Q) q_out[{b} * N_Q + q_row] = static_cast<bfloat16_t>(result{b}[slid]);"""
+        for b in range(B))
+
+    # Gate epilogue (sigmoid → f32 per batch)
+    gate_write = "\n".join(f"""
+            if (slid < 4u && g_row < N_GATE) {{
+                float sig{b} = 1.0f / (1.0f + metal::exp(-result{b}[slid]));
+                gate_out[{b} * N_GATE + g_row] = sig{b};
+            }}""" for b in range(B))
+
+    # Keys epilogue
+    k_write = "\n".join(f"""
+            if (slid < 4u && k_row < N_K) k_out[{b} * N_K + k_row] = static_cast<bfloat16_t>(result{b}[slid]);"""
+        for b in range(B))
+
+    # Values epilogue
+    v_write = "\n".join(f"""
+            if (slid < 4u && v_row < N_V) v_out[{b} * N_V + v_row] = static_cast<bfloat16_t>(result{b}[slid]);"""
+        for b in range(B))
+
+    N_V_val = N_TOTAL - N_Q - N_GATE - N_K
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int GROUP_SIZE = {gs};
+    const int SC_STRIDE = {sc_stride};
+    const int SLID_DIV = {slid_div};
+    const int K = {K};
+    const int K_groups = {K_groups};
+    const int N_Q = {N_Q};
+    const int N_GATE = {N_GATE};
+    const int N_K = {N_K};
+    const int N_V = {N_V_val};
+    const int N_TOTAL = {N_TOTAL};
+    const int N_Q_TG = {N_Q_TG};
+    const int N_GATE_TG = {N_GATE_TG};
+    const int N_K_TG = {N_K_TG};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+    int b_idx = tgid.z;
+    int tg = tgid.y;
+
+    int out_row, region;
+    if (tg < N_Q_TG) {{
+        region = 0; out_row = tg * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_Q_TG + N_GATE_TG) {{
+        region = 1; out_row = N_Q + (tg - N_Q_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_Q_TG + N_GATE_TG + N_K_TG) {{
+        region = 2; out_row = N_Q + N_GATE + (tg - N_Q_TG - N_GATE_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else {{
+        region = 3; out_row = N_Q + N_GATE + N_K + (tg - N_Q_TG - N_GATE_TG - N_K_TG) * 8 + sgid * RESULTS_PER_SG;
+    }}
+    if (out_row >= N_TOTAL) return;
+
+    // Weight pointers (shared across all batch elements)
+    const device uint8_t* ws = (const device uint8_t*)W_merged + (long)out_row * K + slid * VALUES_PER_THREAD;
+    const device bfloat16_t* sc = (const device bfloat16_t*)S_merged + (long)out_row * K_groups + slid / SLID_DIV;
+    const device bfloat16_t* bi = (const device bfloat16_t*)B_merged + (long)out_row * K_groups + slid / SLID_DIV;
+
+    {result_decls}
+    int x_base = slid * VALUES_PER_THREAD;
+
+    // K-loop: load weights once, compute {B} batch elements
+    for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{
+{x_load}
+
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            const device uint8_t* wl = ws + row * K;
+            float s_val = float(sc[row * K_groups]);
+            float b_val = float(bi[row * K_groups]);
+            float w_vals[VALUES_PER_THREAD];
+            for (int i = 0; i < VALUES_PER_THREAD; i++) w_vals[i] = float(wl[i]);
+{qdot}
+        }}
+
+        ws += BLOCK_SIZE; sc += SC_STRIDE; bi += SC_STRIDE; x_base += BLOCK_SIZE;
+    }}
+
+    {simd_reduce}
+
+    // Region-specific epilogues for all {B} batches
+    if (region == 0) {{
+        int q_row = out_row + (int)slid;
+{q_write}
+    }} else if (region == 1) {{
+        int g_row = out_row - N_Q + (int)slid;
+{gate_write}
+    }} else if (region == 2) {{
+        int k_row = out_row - N_Q - N_GATE + (int)slid;
+{k_write}
+    }} else {{
+        int v_row = out_row - N_Q - N_GATE - N_K + (int)slid;
+{v_write}
+    }}
+"""
+
+
+_batched_proj_cache = {}
+
+
+def _get_batched_proj_kernel(K, N_Q, N_GATE, N_K, N_V, B, group_size=64):
+    key = (K, N_Q, N_GATE, N_K, N_V, B, group_size)
+    if key not in _batched_proj_cache:
+        _batched_proj_cache[key] = mx.fast.metal_kernel(
+            name=f"batched_fused_gqa_proj_K{K}_NQ{N_Q}_B{B}",
+            input_names=["x", "W_merged", "S_merged", "B_merged"],
+            output_names=["q_out", "gate_out", "k_out", "v_out"],
+            source=_gen_batched_fused_gqa_proj_source(K, N_Q, N_GATE, N_K, N_V, B, group_size),
+        )
+    return _batched_proj_cache[key]
+
+
+def batched_fused_gqa_projections(x, W_merged, S_merged, B_merged, proj_dims,
+                                   batch_size, total_tg=None):
+    """Batched fused GQA projections with register weight sharing.
+
+    Args:
+        x: [B, 1, K] bf16
+        W_merged, S_merged, B_merged: merged q+gate+k+v weights
+        proj_dims: (N_Q, N_GATE, N_K, N_V)
+        batch_size: B (1..8)
+
+    Returns:
+        queries (B, 1, N_Q) bf16, gate_sigmoid (B, 1, N_GATE) f32,
+        keys (B, 1, N_K) bf16, values (B, 1, N_V) bf16
+    """
+    B = batch_size
+    N_Q, N_GATE, N_K, N_V = proj_dims
+    K = x.shape[-1]
+
+    kern = _get_batched_proj_kernel(K, N_Q, N_GATE, N_K, N_V, B)
+
+    if total_tg is None:
+        total_tg = ceil_div(N_Q, 8) + ceil_div(N_GATE, 8) + ceil_div(N_K, 8) + ceil_div(N_V, 8)
+
+    x_flat = x.reshape(B, K)
+
+    results = kern(
+        inputs=[x_flat, W_merged, S_merged, B_merged],
+        output_shapes=[
+            (B * N_Q,), (B * N_GATE,), (B * N_K,), (B * N_V,),
+        ],
+        output_dtypes=[mx.bfloat16, mx.float32, mx.bfloat16, mx.bfloat16],
+        grid=(32, total_tg * 2, 1),
+        threadgroup=(32, 2, 1),
+    )
+
+    return (results[0].reshape(B, 1, N_Q),
+            results[1].reshape(B, 1, N_GATE),
+            results[2].reshape(B, 1, N_K),
+            results[3].reshape(B, 1, N_V))
@@ -0,0 +1,252 @@
+"""Batched merged 8-bit down_proj GEMV for Qwen3.5 routed + shared experts.
+
+Adapts merged_down_proj_8bit for batch_size B (1..8).
+
+Grid z-dimension: B * n_active + 1
+  - tgid.z < B * n_active: routed experts (one TG per batch×expert pair)
+    Same structure as affine_gather_qmv — each TG independently indexes into
+    expert weights via inds[flat_idx].
+  - tgid.z == B * n_active: shared expert (ONE TG handles ALL B batch elements
+    with register-level weight sharing — loads weights once, computes B outputs)
+
+All constants baked into Metal source (no scalar kernel inputs).
+"""
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_batched_merged_down_8bit_source(K_OUT, N_IN, SHARED_N_IN, n_active, B, group_size=64):
+    gs = int(group_size)
+    sc_stride = 256 // gs
+    slid_divisor = gs // 8
+    N_groups = N_IN // gs
+    SHARED_N_groups = SHARED_N_IN // gs
+    total_routed = B * n_active
+
+    # Shared expert: generate unrolled batch loops
+    shared_x_load_lines = []
+    for b in range(B):
+        shared_x_load_lines.append(f"""
+            float x{b}_thread[VALUES_PER_THREAD];
+            float xsum{b} = 0;
+            for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+                float xi = X_shared[{b} * SHARED_N_IN + x_base + i];
+                x{b}_thread[i] = xi;
+                xsum{b} += xi;
+            }}""")
+    shared_x_load = "\n".join(shared_x_load_lines)
+
+    shared_qdot_lines = []
+    for b in range(B):
+        shared_qdot_lines.append(f"""
+                float accum{b} = 0;
+                for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+                    accum{b} += x{b}_thread[i] * w_vals[i];
+                }}
+                result{b}[row] += s_val * accum{b} + xsum{b} * b_val;""")
+    shared_qdot = "\n".join(shared_qdot_lines)
+
+    shared_result_decls = "\n        ".join(
+        f"float result{b}[RESULTS_PER_SG] = {{0, 0, 0, 0}};"
+        for b in range(B)
+    )
+
+    shared_write_lines = []
+    for b in range(B):
+        shared_write_lines.append(f"""
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            float r{b} = simd_sum(result{b}[row]);
+            if (slid == 0) {{
+                Y_shared[{b} * K_OUT + out_row + row] = r{b};
+            }}
+        }}""")
+    shared_write = "\n".join(shared_write_lines)
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int K_OUT = {K_OUT};
+    const int N_IN = {N_IN};
+    const int SHARED_N_IN = {SHARED_N_IN};
+    const int N_GROUPS = {N_groups};
+    const int SHARED_N_GROUPS = {SHARED_N_groups};
+    const int N_ACTIVE = {n_active};
+    const int TOTAL_ROUTED = {total_routed};
+    const int BATCH_SIZE = {B};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+
+    int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG;
+    if (out_row >= K_OUT) return;
+
+    if (tgid.z < (uint)TOTAL_ROUTED) {{
+        // ═══════ ROUTED EXPERT PATH (same as gather_qmv) ═══════
+        // tgid.z indexes flat (batch, expert) pairs
+        int flat_idx = (int)tgid.z;
+        int expert = inds[flat_idx];
+
+        const device uint8_t* ws = (const device uint8_t*)W
+            + (long)expert * K_OUT * N_IN + out_row * N_IN + slid * VALUES_PER_THREAD;
+        const device bfloat16_t* sc = (const device bfloat16_t*)S
+            + (long)expert * K_OUT * N_GROUPS + out_row * N_GROUPS + slid / {slid_divisor};
+        const device bfloat16_t* bi = (const device bfloat16_t*)B_q
+            + (long)expert * K_OUT * N_GROUPS + out_row * N_GROUPS + slid / {slid_divisor};
+
+        const device float* x_ptr = (const device float*)X_routed
+            + flat_idx * N_IN;
+
+        int x_base = slid * VALUES_PER_THREAD;
+        float result[4] = {{0, 0, 0, 0}};
+
+        for (int k = 0; k < N_IN; k += BLOCK_SIZE) {{
+            float x_thread[8];
+            float xsum = 0;
+            for (int i = 0; i < 8; i++) {{
+                float xi = x_ptr[x_base + i];
+                x_thread[i] = xi;
+                xsum += xi;
+            }}
+
+            for (int row = 0; row < RESULTS_PER_SG; row++) {{
+                const device uint8_t* wl = ws + row * N_IN;
+                float s = float(sc[row * N_GROUPS]);
+                float b = float(bi[row * N_GROUPS]);
+                float accum = 0;
+                for (int i = 0; i < 8; i++) {{
+                    accum += x_thread[i] * float(wl[i]);
+                }}
+                result[row] += s * accum + xsum * b;
+            }}
+
+            ws += BLOCK_SIZE;
+            sc += {sc_stride};
+            bi += {sc_stride};
+            x_base += BLOCK_SIZE;
+        }}
+
+        device float* yp = Y_routed + flat_idx * K_OUT + out_row;
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            float r = simd_sum(result[row]);
+            if (slid == 0) {{
+                yp[row] = r;
+            }}
+        }}
+
+    }} else {{
+        // ═══════ SHARED EXPERT PATH (register-level weight sharing) ═══════
+        // ONE TG handles ALL {B} batch elements.
+        // Load shared expert weights once, compute {B} outputs from registers.
+
+        const device uint8_t* ws = (const device uint8_t*)W_shared_down
+            + (long)out_row * SHARED_N_IN + slid * VALUES_PER_THREAD;
+        const device bfloat16_t* sc = (const device bfloat16_t*)S_shared_down
+            + (long)out_row * SHARED_N_GROUPS + slid / {slid_divisor};
+        const device bfloat16_t* bi = (const device bfloat16_t*)B_shared_down
+            + (long)out_row * SHARED_N_GROUPS + slid / {slid_divisor};
+
+        int x_base = slid * VALUES_PER_THREAD;
+        {shared_result_decls}
+
+        for (int k = 0; k < SHARED_N_IN; k += BLOCK_SIZE) {{
+            // Load x for all {B} batch elements
+{shared_x_load}
+
+            // Load weights once into registers, compute all {B} batches
+            for (int row = 0; row < RESULTS_PER_SG; row++) {{
+                const device uint8_t* wl = ws + row * SHARED_N_IN;
+                float s_val = float(sc[row * SHARED_N_GROUPS]);
+                float b_val = float(bi[row * SHARED_N_GROUPS]);
+
+                float w_vals[VALUES_PER_THREAD];
+                for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+                    w_vals[i] = float(wl[i]);
+                }}
+
+{shared_qdot}
+            }}
+
+            ws += BLOCK_SIZE;
+            sc += {sc_stride};
+            bi += {sc_stride};
+            x_base += BLOCK_SIZE;
+        }}
+
+        // Write all {B} outputs
+{shared_write}
+    }}
+"""
+
+
+_batched_down_cache = {}
+
+
+def _get_batched_down_kernel(K_OUT, N_IN, SHARED_N_IN, n_active, B, group_size=64):
+    key = (K_OUT, N_IN, SHARED_N_IN, n_active, B, group_size)
+    if key not in _batched_down_cache:
+        _batched_down_cache[key] = mx.fast.metal_kernel(
+            name=f"batched_down_K{K_OUT}_N{N_IN}_SN{SHARED_N_IN}_na{n_active}_B{B}",
+            input_names=["W", "S", "B_q",
+                         "W_shared_down", "S_shared_down", "B_shared_down",
+                         "X_routed", "X_shared", "inds"],
+            output_names=["Y_routed", "Y_shared"],
+            source=_gen_batched_merged_down_8bit_source(
+                K_OUT, N_IN, SHARED_N_IN, n_active, B, group_size),
+        )
+    return _batched_down_cache[key]
+
+
+def batched_merged_down_proj_8bit(w_q, s, b_q,
+                                   w_shared_down, s_shared_down, b_shared_down,
+                                   x_routed, x_shared, inds,
+                                   k_out, n_in, batch_size,
+                                   n_active, group_size=64,
+                                   shared_n_in=None):
+    """Batched merged down_proj for 8-bit routed + shared experts.
+
+    Args:
+        w_q: routed weights (E, K_OUT, N_IN/4) uint32
+        s: routed scales (E, K_OUT, N_IN/gs) bf16
+        b_q: routed biases (E, K_OUT, N_IN/gs) bf16
+        w_shared_down: shared weight (K_OUT, SHARED_N_IN/4) uint32
+        s_shared_down: shared scales (K_OUT, SHARED_N_IN/gs) bf16
+        b_shared_down: shared biases (K_OUT, SHARED_N_IN/gs) bf16
+        x_routed: (B * n_active, N_IN) f32
+        x_shared: (B * SHARED_N_IN,) f32
+        inds: (B * n_active,) uint32
+        k_out: output dimension
+        n_in: routed input dimension
+        batch_size: B
+        n_active: experts per token (top_k)
+        shared_n_in: shared input dim (defaults to n_in)
+
+    Returns:
+        Y_routed: (B * n_active, k_out) f32
+        Y_shared: (B, k_out) f32
+    """
+    B = batch_size
+    k_out_val = int(k_out)
+    n_in_val = int(n_in)
+    shared_n_in_val = int(shared_n_in) if shared_n_in is not None else n_in_val
+
+    kern = _get_batched_down_kernel(k_out_val, n_in_val, shared_n_in_val, n_active, B)
+
+    y_groups = ceil_div(k_out_val, 8)
+    total_routed = B * n_active
+
+    Y = kern(
+        inputs=[w_q, s, b_q,
+                w_shared_down, s_shared_down, b_shared_down,
+                x_routed, x_shared, inds],
+        output_shapes=[(total_routed * k_out_val,), (B * k_out_val,)],
+        output_dtypes=[mx.float32, mx.float32],
+        grid=(32, y_groups * 2, total_routed + 1),
+        threadgroup=(32, 2, 1),
+    )
+
+    return Y[0].reshape(total_routed, k_out_val), Y[1].reshape(B, k_out_val)
@@ -0,0 +1,92 @@
+"""Batched MoE epilogue for Qwen3.5: weighted sum + shared expert gate + residual.
+
+Computes per batch element:
+  Y[b, j] = bf16( Σ_a(scores[b,a] * D_routed[b*n_active+a, j])
+                 + sigmoid(gate_raw[b]) * D_shared[b, j]
+                 + H[b, j] )
+
+Grid z = B (one set of threads per batch element).
+All constants baked into Metal source.
+"""
+import mlx.core as mx
+
+
+def _gen_batched_epilogue_source(K, n_active, B):
+    return f"""
+    const int K_const = {K};
+    const int n_active_const = {n_active};
+    const int B_const = {B};
+
+    uint tid = thread_position_in_grid.x;
+    uint batch_id = thread_position_in_grid.z;
+
+    if (tid >= K_const || batch_id >= B_const) return;
+
+    // Weighted sum of routed expert outputs for this batch element
+    float acc = 0.0f;
+    int routed_base = (int)batch_id * n_active_const * K_const;
+    int score_base = (int)batch_id * n_active_const;
+    for (int a = 0; a < n_active_const; a++) {{
+        acc += scores[score_base + a] * D_routed[routed_base + a * K_const + tid];
+    }}
+
+    // Shared expert: sigmoid(gate_raw) * D_shared
+    float gate_raw_val = gate_raw[(int)batch_id];
+    float gate = 1.0f / (1.0f + metal::exp(-gate_raw_val));
+    float shared_val = D_shared[(int)batch_id * K_const + tid] * gate;
+
+    // Add residual and write
+    Y[(int)batch_id * K_const + tid] = static_cast<bfloat16_t>(
+        acc + shared_val + float(H[(int)batch_id * K_const + tid])
+    );
+"""
+
+
+_batched_epilogue_cache = {}
+
+
+def _get_batched_epilogue_kernel(K, n_active, B):
+    key = (K, n_active, B)
+    if key not in _batched_epilogue_cache:
+        _batched_epilogue_cache[key] = mx.fast.metal_kernel(
+            name=f"batched_epilogue_K{K}_na{n_active}_B{B}",
+            input_names=["D_routed", "D_shared", "scores", "H", "gate_raw"],
+            output_names=["Y"],
+            source=_gen_batched_epilogue_source(K, n_active, B),
+        )
+    return _batched_epilogue_cache[key]
+
+
+def batched_moe_epilogue(d_routed, d_shared, scores, h, gate_raw,
+                          k_val, batch_size, n_active):
+    """Batched MoE epilogue with fused sigmoid.
+
+    Args:
+        d_routed: (B * n_active, K) f32
+        d_shared: (B, K) f32
+        scores: (B * n_active,) f32
+        h: (B, K) bf16 — residual
+        gate_raw: (B,) f32 — raw shared expert gate (pre-sigmoid)
+        k_val: hidden dimension
+        batch_size: B
+        n_active: experts per token
+
+    Returns:
+        Y: (B, K) bf16
+    """
+    K = int(k_val)
+    B = batch_size
+    kern = _get_batched_epilogue_kernel(K, n_active, B)
+
+    tg_size = min(K, 1024)
+    n_tg = (K + tg_size - 1) // tg_size
+
+    Y = kern(
+        inputs=[d_routed, d_shared.reshape(B * K), scores, h.reshape(B * K), gate_raw],
+        output_shapes=[(B * K,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(n_tg * tg_size, 1, B),
+        threadgroup=(tg_size, 1, 1),
+    )
+
+    return Y[0].reshape(B, K)
@@ -0,0 +1,332 @@
+"""Batched Dispatch 1: Fused o_proj (8-bit) + gate GEMV parts + x² partials.
+
+Adapts custom_oproj_gate_gemv_8bit for batch_size B (1..8).
+Register-level weight sharing: each TG loads weights once, computes B outputs.
+
+Three GEMV regions (same as B=1):
+  TGs 0..N_OPROJ_TG-1: o_proj GEMV (8-bit) + residual + h_scaled + x²
+  TGs N_OPROJ_TG..+N_M1_TG-1: M1 × attn_out → gate_part_a (bf16 GEMV)
+  TGs +N_M1_TG..end: W_fused × residual → gate_part_b (bf16 GEMV)
+
+All constants baked into Metal source. B is unrolled at code-generation time.
+"""
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_batched_oproj_source(n_experts, M, K_attn, K_hidden, B, group_size=64, gate_bm=8):
+    E = int(n_experts)
+    gs = group_size
+    oproj_slid_divisor = gs // 8
+    oproj_sc_stride = 256 // gs
+    blockM_gate = gate_bm * 4
+    n_m1_tg = ceil_div(E, blockM_gate)
+
+    # Generate unrolled per-batch code for o_proj epilogue
+    oproj_epilogue = []
+    for b in range(B):
+        oproj_epilogue.append(f"""
+            float x2_acc{b} = 0.0f;
+            for (int tm = 0; tm < TM; tm++) {{
+                int k = out_row + tm;
+                float h{b} = result{b}[tm] + float(residual[{b} * M_DIM + k]);
+                x2_acc{b} += h{b} * h{b};
+                h_scaled[{b} * M_DIM + k] = static_cast<bfloat16_t>(h{b} * float(w_rms[k]));
+                h_out[{b} * M_DIM + k] = static_cast<bfloat16_t>(h{b});
+            }}""")
+    oproj_epilogue_code = "\n".join(oproj_epilogue)
+
+    oproj_x2_write = []
+    for b in range(B):
+        oproj_x2_write.append(f"""
+                total{b} += tgp_x2[s * {B} + {b}];""")
+    oproj_x2_sum = "\n".join(oproj_x2_write)
+
+    oproj_x2_final = []
+    for b in range(B):
+        oproj_x2_final.append(f"""
+                x2_partials[{b} * N_OPROJ_TG_DIM + tg_x] = total{b};""")
+    oproj_x2_final_code = "\n".join(oproj_x2_final)
+
+    # o_proj K-loop: load weights once, compute B batch elements
+    oproj_x_load = "\n".join(f"""
+            float xv{b}[VPT]; float xsum{b} = 0.0f;
+            for (int i = 0; i < VPT; i++) {{ xv{b}[i] = float(attn_out[{b} * K_ATTN_DIM + xb + i]); xsum{b} += xv{b}[i]; }}""" for b in range(B))
+
+    oproj_qdot = "\n".join(f"""
+                acc{b}[row] += s_val * wdot(xv{b}, w_vals) + xsum{b} * b_val;""" for b in range(B))
+
+    oproj_result_decls = " ".join(f"float acc{b}[TM] = {{0,0,0,0}};" for b in range(B))
+    oproj_simd_reduce = "\n".join(f"            float result{b}[TM]; for (int tm=0;tm<TM;tm++) result{b}[tm] = simd_sum(acc{b}[tm]);" for b in range(B))
+
+    oproj_tgp_write = "\n".join(f"            tgp_x2[sgid * {B} + {b}] = x2_acc{b};" for b in range(B))
+    oproj_total_decls = " ".join(f"float total{b} = 0.0f;" for b in range(B))
+
+    # Gate M1 GEMV: load M1 weights once, compute B dot products with B attn_outs
+    gate_a_x_load = "\n".join(f"""
+            float v{b}[TN];
+            for (int tn = 0; tn < TN; tn++) v{b}[tn] = float(attn_out[{b} * K_ATTN_DIM + bn + tn]);""" for b in range(B))
+
+    gate_a_dot = "\n".join(f"""
+                float gacc{b} = 0.0f;
+                for (int tn = 0; tn < TN; tn++) gacc{b} += w_row[tn] * v{b}[tn];
+                gresult{b}[tm] += gacc{b};""" for b in range(B))
+
+    gate_a_decls = " ".join(f"float gresult{b}[TM] = {{0,0,0,0}};" for b in range(B))
+    gate_a_reduce = "\n".join(f"            gresult{b}[tm] = simd_sum(gresult{b}[tm]);" for b in range(B))
+    gate_a_write = "\n".join(f"""
+                    gate_part_a[{b} * E_CONST + e] = gresult{b}[tm];""" for b in range(B))
+
+    # Gate W_fused GEMV: same pattern but with residual input
+    gate_b_x_load = "\n".join(f"""
+            float rv{b}[TN];
+            for (int tn = 0; tn < TN; tn++) rv{b}[tn] = float(residual[{b} * K_HIDDEN_DIM + bn + tn]);""" for b in range(B))
+
+    gate_b_dot = "\n".join(f"""
+                float wdot{b} = 0.0f;
+                for (int tn = 0; tn < TN; tn++) wdot{b} += w_row[tn] * rv{b}[tn];
+                bresult{b}[tm] += wdot{b};""" for b in range(B))
+
+    gate_b_decls = " ".join(f"float bresult{b}[TM] = {{0,0,0,0}};" for b in range(B))
+    gate_b_reduce = "\n".join(f"            bresult{b}[tm] = simd_sum(bresult{b}[tm]);" for b in range(B))
+    gate_b_write = "\n".join(f"""
+                    gate_part_b[{b} * E_CONST + e] = bresult{b}[tm];""" for b in range(B))
+
+    return f"""
+    const int TM = 4;
+    const int TN = 4;
+    const int blockN = 128;
+    const int E_CONST = {E};
+    const int M_DIM = {M};
+    const int K_ATTN_DIM = {K_attn};
+    const int K_HIDDEN_DIM = {K_hidden};
+    const int N_OPROJ_TG_DIM = {ceil_div(M, 32)};
+    const int BATCH_SIZE = {B};
+
+    // Helper: dot product of x_thread and w_vals (8 elements)
+    auto wdot = [](thread float* x, thread float* w) -> float {{
+        float a = 0;
+        for (int i = 0; i < 8; i++) a += x[i] * w[i];
+        return a;
+    }};
+
+    const int N_OPROJ_TG = N_OPROJ_TG_DIM;
+    const int N_M1_TG = {n_m1_tg};
+    const int blockM_gate = {blockM_gate};
+
+    uint tg_x = threadgroup_position_in_grid.x;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+
+    if (tg_x < (uint)N_OPROJ_TG) {{
+        // ══════ O_PROJ GEMV (8-bit, register-sharing for {B} batches) ══════
+        const int blockM = 32;
+        const int VPT = 8;
+        const int BLOCK_SIZE = 256;
+
+        int out_row = int(tg_x) * blockM + int(sgid) * TM;
+        if (out_row >= M_DIM) return;
+        out_row = (out_row + TM <= M_DIM) ? out_row : (M_DIM - TM);
+
+        threadgroup float tgp_x2[8 * {B}];
+
+        {oproj_result_decls}
+        int K_groups = K_ATTN_DIM / {gs};
+
+        // Weight pointers (shared across batch)
+        const device uint8_t* ws = (const device uint8_t*)W_oproj
+            + (long)out_row * K_ATTN_DIM + slid * VPT;
+        const device bfloat16_t* sc = (const device bfloat16_t*)S_oproj
+            + (long)out_row * K_groups + slid / {oproj_slid_divisor};
+        const device bfloat16_t* bi = (const device bfloat16_t*)B_oproj
+            + (long)out_row * K_groups + slid / {oproj_slid_divisor};
+
+        int xb = slid * VPT;
+
+        for (int k = 0; k < K_ATTN_DIM; k += BLOCK_SIZE) {{
+            // Load x for all {B} batches
+{oproj_x_load}
+
+            // Load weights once, compute all batches
+            for (int row = 0; row < TM; row++) {{
+                const device uint8_t* wl = ws + row * K_ATTN_DIM;
+                float s_val = float(sc[row * K_groups]);
+                float b_val = float(bi[row * K_groups]);
+                float w_vals[VPT];
+                for (int i = 0; i < VPT; i++) w_vals[i] = float(wl[i]);
+{oproj_qdot}
+            }}
+
+            ws += BLOCK_SIZE; sc += {oproj_sc_stride}; bi += {oproj_sc_stride};
+            xb += BLOCK_SIZE;
+        }}
+
+        // simd_sum for all batches
+{oproj_simd_reduce}
+
+        // Epilogue: residual add + x² + h_scaled + h_out
+        if (slid == 0) {{
+{oproj_epilogue_code}
+{oproj_tgp_write}
+        }}
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (sgid == 0 && slid == 0) {{
+            {oproj_total_decls}
+            for (int s = 0; s < 8; s++) {{
+{oproj_x2_sum}
+            }}
+{oproj_x2_final_code}
+        }}
+
+    }} else if (tg_x < (uint)(N_OPROJ_TG + N_M1_TG)) {{
+        // ══════ M1 GEMV (bf16, register-sharing for {B} batches) ══════
+        int local_tg = int(tg_x) - N_OPROJ_TG;
+        int out_row = local_tg * blockM_gate + int(sgid) * TM;
+        if (out_row >= E_CONST) return;
+        out_row = (out_row + TM <= E_CONST) ? out_row : (E_CONST - TM);
+
+        {gate_a_decls}
+        int bn = int(slid) * TN;
+        int n_iter = K_ATTN_DIM / blockN;
+
+        for (int i = 0; i < n_iter; i++) {{
+{gate_a_x_load}
+            for (int tm = 0; tm < TM; tm++) {{
+                float w_row[TN];
+                for (int tn = 0; tn < TN; tn++) w_row[tn] = float(M1[(out_row + tm) * K_ATTN_DIM + bn + tn]);
+{gate_a_dot}
+            }}
+            bn += blockN;
+        }}
+
+        for (int tm = 0; tm < TM; tm++) {{
+{gate_a_reduce}
+        }}
+
+        if (slid == 0) {{
+            for (int tm = 0; tm < TM; tm++) {{
+                int e = out_row + tm;
+                if (e < E_CONST) {{
+{gate_a_write}
+                }}
+            }}
+        }}
+
+    }} else {{
+        // ══════ W_FUSED GEMV (bf16, register-sharing for {B} batches) ══════
+        int local_tg = int(tg_x) - N_OPROJ_TG - N_M1_TG;
+        int out_row = local_tg * blockM_gate + int(sgid) * TM;
+        if (out_row >= E_CONST) return;
+        out_row = (out_row + TM <= E_CONST) ? out_row : (E_CONST - TM);
+
+        {gate_b_decls}
+        int bn = int(slid) * TN;
+        int n_iter = K_HIDDEN_DIM / blockN;
+
+        for (int i = 0; i < n_iter; i++) {{
+{gate_b_x_load}
+            for (int tm = 0; tm < TM; tm++) {{
+                float w_row[TN];
+                for (int tn = 0; tn < TN; tn++) w_row[tn] = float(W_fused[(out_row + tm) * K_HIDDEN_DIM + bn + tn]);
+{gate_b_dot}
+            }}
+            bn += blockN;
+        }}
+
+        for (int tm = 0; tm < TM; tm++) {{
+{gate_b_reduce}
+        }}
+
+        if (slid == 0) {{
+            for (int tm = 0; tm < TM; tm++) {{
+                int e = out_row + tm;
+                if (e < E_CONST) {{
+{gate_b_write}
+                }}
+            }}
+        }}
+    }}
+"""
+
+
+_batched_oproj_cache = {}
+
+
+def _get_batched_oproj_kernel(n_experts, M, K_attn, K_hidden, B, group_size=64, gate_bm=8):
+    key = (n_experts, M, K_attn, K_hidden, B, group_size, gate_bm)
+    if key not in _batched_oproj_cache:
+        _batched_oproj_cache[key] = mx.fast.metal_kernel(
+            name=f"batched_oproj_E{n_experts}_M{M}_Ka{K_attn}_Kh{K_hidden}_B{B}",
+            input_names=[
+                "W_oproj", "S_oproj", "B_oproj",
+                "attn_out", "residual", "w_rms",
+                "M1", "W_fused",
+            ],
+            output_names=["h_scaled", "h_out", "x2_partials",
+                          "gate_part_a", "gate_part_b"],
+            source=_gen_batched_oproj_source(n_experts, M, K_attn, K_hidden, B, group_size, gate_bm),
+        )
+    return _batched_oproj_cache[key]
+
+
+def batched_oproj_gate_gemv(W_oproj, S_oproj, B_oproj,
+                             attn_out, residual, w_rms,
+                             M1, W_fused,
+                             M, K_attn, batch_size,
+                             n_experts=256, gate_bm=8,
+                             K_hidden=None, group_size=64):
+    """Batched fused 8-bit o_proj + bf16 gate GEMVs.
+
+    Args:
+        W_oproj/S_oproj/B_oproj: 8-bit o_proj weights
+        attn_out: (B, K_attn) bf16
+        residual: (B, K) bf16
+        w_rms: (K,) bf16 — RMSNorm weight (shared)
+        M1: (E, K_attn) bf16 (shared)
+        W_fused: (E, K) bf16 (shared)
+        M: hidden size
+        K_attn: attention output dim
+        batch_size: B
+        n_experts: E
+        gate_bm: SGs per gate TG
+
+    Returns:
+        h_scaled (B, M) bf16, h_out (B, M) bf16,
+        x2_partials (B, N_TG) f32, gate_part_a (B, E) f32, gate_part_b (B, E) f32
+    """
+    B = batch_size
+    M_val = int(M)
+    K_attn_val = int(K_attn)
+    K_hidden_val = int(K_hidden) if K_hidden is not None else M_val
+
+    kern = _get_batched_oproj_kernel(n_experts, M_val, K_attn_val, K_hidden_val, B, group_size, gate_bm)
+
+    n_oproj_tg = ceil_div(M_val, 32)
+    blockM_gate = gate_bm * 4
+    n_m1_tg = ceil_div(n_experts, blockM_gate)
+    n_wf_tg = ceil_div(n_experts, blockM_gate)
+    total_tg = n_oproj_tg + n_m1_tg + n_wf_tg
+
+    results = kern(
+        inputs=[W_oproj, S_oproj, B_oproj,
+                attn_out.reshape(B * K_attn_val), residual.reshape(B * M_val), w_rms,
+                M1, W_fused],
+        output_shapes=[
+            (B * M_val,), (B * M_val,),
+            (B * n_oproj_tg,),
+            (B * n_experts,), (B * n_experts,),
+        ],
+        output_dtypes=[mx.bfloat16, mx.bfloat16, mx.float32, mx.float32, mx.float32],
+        grid=(total_tg * 32, 8, 1),
+        threadgroup=(32, 8, 1),
+    )
+
+    return (results[0].reshape(B, M_val),
+            results[1].reshape(B, M_val),
+            results[2].reshape(B, n_oproj_tg),
+            results[3].reshape(B, n_experts),
+            results[4].reshape(B, n_experts))
@@ -0,0 +1,578 @@
+"""Dispatch 2 (batched): Softmax prologue + 8-bit SwiGLU for Qwen3.5 MoE.
+
+Combines the prologue from oproj_softmax_topk_swiglu_8bit (B=1) with the
+batched body from batched_merged_swiglu_8bit. All constants are baked into
+the Metal source at Python code-generation time.
+
+Grid z-dimension: B * n_active + 1
+  - tgid.z < B * n_active: routed expert TGs
+    batch_id = tgid.z / n_active, local_z = tgid.z % n_active
+  - tgid.z == B * n_active: shared expert TG (register-level weight sharing
+    for B batch elements, including shared_expert_gate GEMV)
+
+Prologue (all TGs):
+  Phase 1: distributed x2 partial sum -> inv_rms (per batch_id)
+  Phase 2 (routed TGs only): gate scores -> softmax -> parallel top-k ->
+           norm_topk_prob -> write out_inds / norm_scores
+  Phase 3 (shared TG, SG 0): shared_expert_gate 8-bit GEMV with
+           register-level weight sharing -> gate_raw[B]
+
+Body (after TG barrier):
+  Routed: 8-bit gate+up+SwiGLU with h_scaled[batch_id] input
+  Shared: register-level weight sharing for B batch elements
+"""
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_batched_softmax_topk_swiglu_source(
+    N_INTER, SHARED_INTER, K, n_active, B,
+    n_experts=256, top_k=10, norm_topk=True, group_size=64,
+    n_oproj_tg=64,
+):
+    """Generate Metal source for batched softmax + top-k + SwiGLU.
+
+    All routed TGs compute the full softmax+topk for their batch_id into TG
+    memory. Only the TG with local_z==0 writes out_inds/norm_scores to device
+    memory. Each routed TG reads its own expert from tg_inds[local_z].
+    """
+    gs = group_size
+    sc_stride = 256 // gs
+    slid_divisor = gs // 8
+    N_TOTAL = 2 * N_INTER
+    K_groups = K // gs
+    SHARED_K_groups = K // gs
+    total_routed = B * n_active
+    E = int(n_experts)
+    K_TOP = int(top_k)
+    SPT = (E + 63) // 64
+
+    # ── Shared expert body: unrolled per-batch code ──
+    shared_x_load = "\n".join(f"""
+            float x{b}_thread[VALUES_PER_THREAD];
+            float xsum{b} = 0;
+            for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+                float xi = float(X[{b} * K_DIM + x_base + i]);
+                x{b}_thread[i] = xi;
+                xsum{b} += xi;
+            }}""" for b in range(B))
+
+    shared_gate_qdot = "\n".join(f"""
+                float accum_g{b} = 0;
+                for (int i = 0; i < VALUES_PER_THREAD; i++) accum_g{b} += x{b}_thread[i] * wg_vals[i];
+                gate{b}[row] += sg * accum_g{b} + xsum{b} * bg;""" for b in range(B))
+
+    shared_up_qdot = "\n".join(f"""
+                float accum_u{b} = 0;
+                for (int i = 0; i < VALUES_PER_THREAD; i++) accum_u{b} += x{b}_thread[i] * wu_vals[i];
+                up{b}[row] += su * accum_u{b} + xsum{b} * bu;""" for b in range(B))
+
+    shared_result_decls = "\n        ".join(
+        f"float gate{b}[RESULTS_PER_SG] = {{0,0,0,0}}; float up{b}[RESULTS_PER_SG] = {{0,0,0,0}};"
+        for b in range(B))
+
+    shared_write_lines = []
+    for b in range(B):
+        shared_write_lines.append(f"""
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            float g{b} = simd_sum(gate{b}[row]) * inv_rms_{b};
+            float u{b} = simd_sum(up{b}[row]) * inv_rms_{b};
+            if (slid == 0) {{
+                float silu_g{b} = g{b} / (1.0f + metal::exp(-g{b}));
+                Y_shared[{b} * SHARED_INTER_DIM + out_row + row] = silu_g{b} * u{b};
+            }}
+        }}""")
+    shared_write = "\n".join(shared_write_lines)
+
+    # inv_rms for all B batches in the shared TG
+    shared_inv_rms_lines = []
+    for b in range(B):
+        shared_inv_rms_lines.append(f"""
+        float local_x2_{b} = 0.0f;
+        for (int i = x2_start; i < x2_end; i++) local_x2_{b} += x2_partials[{b} * N_OPROJ_TG_DIM + i];
+        float sg_x2_{b} = simd_sum(local_x2_{b});
+        if (slid == 0) tg_x2_sg[sgid] = sg_x2_{b};
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        float inv_rms_{b} = metal::precise::rsqrt((tg_x2_sg[0] + tg_x2_sg[1]) / (float)K_DIM + 1e-6f);""")
+    shared_inv_rms_block = "\n".join(shared_inv_rms_lines)
+
+    # Shared expert gate GEMV accumulator declarations
+    seg_acc_decls = "\n        ".join(
+        f"float seg_gate_acc{b} = 0.0f;" for b in range(B))
+
+    seg_write = "\n".join(
+        f"            gate_raw[{b}] = seg_gate_acc{b} * inv_rms_{b};"
+        for b in range(B))
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int N_INTER_DIM = {N_INTER};
+    const int SHARED_INTER_DIM = {SHARED_INTER};
+    const int K_DIM = {K};
+    const int K_GROUPS = {K_groups};
+    const int N_TOTAL = {N_TOTAL};
+    const int N_ACTIVE = {n_active};
+    const int TOTAL_ROUTED = {total_routed};
+    const int BATCH_SIZE = {B};
+    const int E_CONST = {E};
+    const int K_TOP_CONST = {K_TOP};
+    const int SPT = {SPT};
+    const int N_OPROJ_TG_DIM = {n_oproj_tg};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;  // 0 or 1
+    uint slid = thread_index_in_simdgroup;       // 0..31
+    int tid = int(sgid) * 32 + int(slid);        // 0..63
+
+    threadgroup float tg_x2_sg[2];
+    threadgroup int tg_inds[{K_TOP}];
+    threadgroup float tg_selected_scores[{K_TOP}];
+
+    if (tgid.z < (uint)TOTAL_ROUTED) {{
+        // ═══════════════════════════════════════════════════════════════
+        // ROUTED TG PROLOGUE
+        // ═══════════════════════════════════════════════════════════════
+        int flat_idx = (int)tgid.z;
+        int batch_id = flat_idx / N_ACTIVE;
+        int local_z = flat_idx % N_ACTIVE;
+
+        // Phase 1: distributed x2 sum -> inv_rms for batch_id
+        int chunk = (N_OPROJ_TG_DIM + 63) / 64;
+        int x2_start = tid * chunk;
+        int x2_end = min(x2_start + chunk, N_OPROJ_TG_DIM);
+        float local_x2 = 0.0f;
+        for (int i = x2_start; i < x2_end; i++)
+            local_x2 += x2_partials[batch_id * N_OPROJ_TG_DIM + i];
+        float sg_x2_sum = simd_sum(local_x2);
+        if (slid == 0) tg_x2_sg[sgid] = sg_x2_sum;
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        float total_x2 = tg_x2_sg[0] + tg_x2_sg[1];
+        float inv_rms = metal::precise::rsqrt(total_x2 / (float)K_DIM + 1e-6f);
+
+        // Phase 2: ALL routed TGs compute full softmax + top-k for their
+        // batch_id. Each TG gets its own copy in TG memory (tg_inds,
+        // tg_selected_scores). This avoids cross-TG communication.
+        {{
+            float my_scores[SPT];
+            for (int j = 0; j < SPT; j++) {{
+                int e = tid * SPT + j;
+                if (e < E_CONST)
+                    my_scores[j] = (gate_part_a[batch_id * E_CONST + e]
+                                  + gate_part_b[batch_id * E_CONST + e]) * inv_rms;
+                else
+                    my_scores[j] = -1e30f;
+            }}
+
+            // Softmax: distributed max
+            float local_max = -1e30f;
+            for (int j = 0; j < SPT; j++)
+                local_max = max(local_max, my_scores[j]);
+            float sg_max_val = simd_max(local_max);
+            threadgroup float tg_softmax_sg[2];
+            if (slid == 0) tg_softmax_sg[sgid] = sg_max_val;
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+            float tg_max = max(tg_softmax_sg[0], tg_softmax_sg[1]);
+
+            // Softmax: exp + distributed sum
+            float local_sum = 0.0f;
+            for (int j = 0; j < SPT; j++) {{
+                float e_val = metal::exp(my_scores[j] - tg_max);
+                my_scores[j] = e_val;
+                local_sum += e_val;
+            }}
+            float sg_sum_val = simd_sum(local_sum);
+            if (slid == 0) tg_softmax_sg[sgid] = sg_sum_val;
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+            float tg_sum = tg_softmax_sg[0] + tg_softmax_sg[1];
+
+            // Softmax: normalize
+            float inv_sum = 1.0f / tg_sum;
+            for (int j = 0; j < SPT; j++)
+                my_scores[j] *= inv_sum;
+
+            // Parallel top-k: K_TOP rounds
+            threadgroup float tg_tk_val[2];
+            threadgroup int tg_tk_info[2];
+
+            for (int round = 0; round < K_TOP_CONST; round++) {{
+                float best = -1.0f;
+                int best_e = -1;
+                for (int j = 0; j < SPT; j++) {{
+                    int e = tid * SPT + j;
+                    if (e < E_CONST && my_scores[j] > best) {{
+                        best = my_scores[j];
+                        best_e = e;
+                    }}
+                }}
+
+                float sg_best = simd_max(best);
+                int candidate = (best == sg_best && best > 0.0f) ? int(slid) : 999;
+                int sg_winner = simd_min(candidate);
+
+                if (slid == 0) {{
+                    tg_tk_val[sgid] = sg_best;
+                    tg_tk_info[sgid] = sg_winner;
+                }}
+                threadgroup_barrier(mem_flags::mem_threadgroup);
+
+                int winner_sg = (tg_tk_val[0] >= tg_tk_val[1]) ? 0 : 1;
+                int winner_lane = tg_tk_info[winner_sg];
+                int winner_tid = winner_sg * 32 + winner_lane;
+
+                if (tid == winner_tid) {{
+                    tg_inds[round] = best_e;
+                    tg_selected_scores[round] = best;
+                    for (int j = 0; j < SPT; j++) {{
+                        if (tid * SPT + j == best_e) {{
+                            my_scores[j] = -1.0f;
+                            break;
+                        }}
+                    }}
+                }}
+                threadgroup_barrier(mem_flags::mem_threadgroup);
+            }}
+
+            // Only local_z==0 writes norm_scores + out_inds to device memory
+            if (local_z == 0 && tid == 0) {{
+                float total_score = 0.0f;
+                for (int a = 0; a < {K_TOP}; a++) total_score += tg_selected_scores[a];
+                float inv_total = {"1.0f / total_score" if norm_topk else "1.0f"};
+                for (int a = 0; a < {K_TOP}; a++) {{
+                    norm_scores[batch_id * {K_TOP} + a] = tg_selected_scores[a] * inv_total;
+                    out_inds[batch_id * {K_TOP} + a] = (uint)tg_inds[a];
+                }}
+            }}
+        }}
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // ═══════════════════════════════════════════════════════════════
+        // ROUTED BODY: 8-bit gate+up+SwiGLU
+        // ═══════════════════════════════════════════════════════════════
+        int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG;
+        if (out_row >= N_INTER_DIM) return;
+
+        int expert = tg_inds[local_z];
+
+        const device uint8_t* ws_gate = (const device uint8_t*)W
+            + (long)expert * N_TOTAL * K_DIM + out_row * K_DIM + slid * VALUES_PER_THREAD;
+        const device bfloat16_t* sc_gate = (const device bfloat16_t*)S
+            + (long)expert * N_TOTAL * K_GROUPS + out_row * K_GROUPS + slid / {slid_divisor};
+        const device bfloat16_t* bi_gate = (const device bfloat16_t*)B_q
+            + (long)expert * N_TOTAL * K_GROUPS + out_row * K_GROUPS + slid / {slid_divisor};
+
+        const device uint8_t* ws_up = (const device uint8_t*)W
+            + (long)expert * N_TOTAL * K_DIM + (out_row + N_INTER_DIM) * K_DIM + slid * VALUES_PER_THREAD;
+        const device bfloat16_t* sc_up = (const device bfloat16_t*)S
+            + (long)expert * N_TOTAL * K_GROUPS + (out_row + N_INTER_DIM) * K_GROUPS + slid / {slid_divisor};
+        const device bfloat16_t* bi_up = (const device bfloat16_t*)B_q
+            + (long)expert * N_TOTAL * K_GROUPS + (out_row + N_INTER_DIM) * K_GROUPS + slid / {slid_divisor};
+
+        int x_base = batch_id * K_DIM + slid * VALUES_PER_THREAD;
+
+        float gate_result[4] = {{0, 0, 0, 0}};
+        float up_result[4] = {{0, 0, 0, 0}};
+
+        for (int k = 0; k < K_DIM; k += BLOCK_SIZE) {{
+            float x_thread[8];
+            float xsum = 0;
+            for (int i = 0; i < 8; i++) {{
+                float xi = float(X[x_base + i]);
+                x_thread[i] = xi;
+                xsum += xi;
+            }}
+
+            for (int row = 0; row < RESULTS_PER_SG; row++) {{
+                const device uint8_t* wg = ws_gate + row * K_DIM;
+                float sg = float(sc_gate[row * K_GROUPS]);
+                float bg = float(bi_gate[row * K_GROUPS]);
+                float accum_g = 0;
+                for (int i = 0; i < 8; i++) accum_g += x_thread[i] * float(wg[i]);
+                gate_result[row] += sg * accum_g + xsum * bg;
+
+                const device uint8_t* wu = ws_up + row * K_DIM;
+                float su = float(sc_up[row * K_GROUPS]);
+                float bu = float(bi_up[row * K_GROUPS]);
+                float accum_u = 0;
+                for (int i = 0; i < 8; i++) accum_u += x_thread[i] * float(wu[i]);
+                up_result[row] += su * accum_u + xsum * bu;
+            }}
+
+            ws_gate += BLOCK_SIZE; ws_up += BLOCK_SIZE;
+            sc_gate += {sc_stride}; sc_up += {sc_stride};
+            bi_gate += {sc_stride}; bi_up += {sc_stride};
+            x_base += BLOCK_SIZE;
+        }}
+
+        // Epilogue: apply inv_rms (factored), SwiGLU, write f32
+        device float* yp = Y_routed + flat_idx * N_INTER_DIM + out_row;
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            float g = simd_sum(gate_result[row]) * inv_rms;
+            float u = simd_sum(up_result[row]) * inv_rms;
+            if (slid == 0) {{
+                float silu_g = g / (1.0f + metal::exp(-g));
+                yp[row] = silu_g * u;
+            }}
+        }}
+
+    }} else {{
+        // ═══════════════════════════════════════════════════════════════
+        // SHARED TG (tgid.z == TOTAL_ROUTED)
+        // ═══════════════════════════════════════════════════════════════
+
+        // Phase 1: compute inv_rms for ALL B batch elements
+        int chunk = (N_OPROJ_TG_DIM + 63) / 64;
+        int x2_start = tid * chunk;
+        int x2_end = min(x2_start + chunk, N_OPROJ_TG_DIM);
+{shared_inv_rms_block}
+
+        // Phase 3: shared_expert_gate 8-bit GEMV (SG 0 only)
+        // Load W_seg once, compute B dot products via register-level weight sharing
+        if (sgid == 0) {{
+            const int VPT = 8;
+            const int SEG_BLOCK = 256;  // 32 * VPT
+
+            const device uint8_t* seg_w_ptr = (const device uint8_t*)W_seg
+                + slid * VPT;
+            const device bfloat16_t* seg_sc = (const device bfloat16_t*)S_seg
+                + slid / {slid_divisor};
+            const device bfloat16_t* seg_bi = (const device bfloat16_t*)B_seg
+                + slid / {slid_divisor};
+            int seg_xb = slid * VPT;
+
+            {seg_acc_decls}
+
+            for (int k = 0; k < K_DIM; k += SEG_BLOCK) {{
+                // Load weight block once into registers
+                float seg_w_regs[VPT];
+                for (int i = 0; i < VPT; i++) seg_w_regs[i] = float(seg_w_ptr[i]);
+                float seg_sc_val = float(*seg_sc);
+                float seg_bi_val = float(*seg_bi);
+
+                // Compute B dot products from the same weight registers
+{chr(10).join(f'''                {{
+                    float xsum{b} = 0.0f, wacc{b} = 0.0f;
+                    for (int i = 0; i < VPT; i++) {{
+                        float xi = float(X[{b} * K_DIM + seg_xb + i]);
+                        xsum{b} += xi;
+                        wacc{b} += xi * seg_w_regs[i];
+                    }}
+                    seg_gate_acc{b} += seg_sc_val * wacc{b} + xsum{b} * seg_bi_val;
+                }}''' for b in range(B))}
+
+                seg_w_ptr += SEG_BLOCK;
+                seg_sc += {sc_stride};
+                seg_bi += {sc_stride};
+                seg_xb += SEG_BLOCK;
+            }}
+
+            // Reduce across SG and write gate_raw[B]
+{chr(10).join(f"            seg_gate_acc{b} = simd_sum(seg_gate_acc{b});" for b in range(B))}
+            if (slid == 0) {{
+{seg_write}
+            }}
+        }}
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // ═══════════════════════════════════════════════════════════════
+        // SHARED BODY: register-level weight sharing for B batch elements
+        // ═══════════════════════════════════════════════════════════════
+        int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG;
+        if (out_row >= SHARED_INTER_DIM) return;
+
+        const device uint8_t* ws_gate = (const device uint8_t*)W_shared
+            + (long)out_row * K_DIM + slid * VALUES_PER_THREAD;
+        const device bfloat16_t* sc_gate = (const device bfloat16_t*)S_shared
+            + (long)out_row * {SHARED_K_groups} + slid / {slid_divisor};
+        const device bfloat16_t* bi_gate = (const device bfloat16_t*)B_shared
+            + (long)out_row * {SHARED_K_groups} + slid / {slid_divisor};
+
+        const device uint8_t* ws_up = (const device uint8_t*)W_shared
+            + (long)(out_row + SHARED_INTER_DIM) * K_DIM + slid * VALUES_PER_THREAD;
+        const device bfloat16_t* sc_up = (const device bfloat16_t*)S_shared
+            + (long)(out_row + SHARED_INTER_DIM) * {SHARED_K_groups} + slid / {slid_divisor};
+        const device bfloat16_t* bi_up = (const device bfloat16_t*)B_shared
+            + (long)(out_row + SHARED_INTER_DIM) * {SHARED_K_groups} + slid / {slid_divisor};
+
+        int x_base = slid * VALUES_PER_THREAD;
+        {shared_result_decls}
+
+        for (int k = 0; k < K_DIM; k += BLOCK_SIZE) {{
+            // Load x for all {B} batch elements
+{shared_x_load}
+
+            for (int row = 0; row < RESULTS_PER_SG; row++) {{
+                // Load gate weights once into registers
+                const device uint8_t* wg = ws_gate + row * K_DIM;
+                float sg = float(sc_gate[row * {SHARED_K_groups}]);
+                float bg = float(bi_gate[row * {SHARED_K_groups}]);
+                float wg_vals[VALUES_PER_THREAD];
+                for (int i = 0; i < VALUES_PER_THREAD; i++) wg_vals[i] = float(wg[i]);
+
+                // Compute gate for all {B} batches from registers
+{shared_gate_qdot}
+
+                // Load up weights once into registers
+                const device uint8_t* wu = ws_up + row * K_DIM;
+                float su = float(sc_up[row * {SHARED_K_groups}]);
+                float bu = float(bi_up[row * {SHARED_K_groups}]);
+                float wu_vals[VALUES_PER_THREAD];
+                for (int i = 0; i < VALUES_PER_THREAD; i++) wu_vals[i] = float(wu[i]);
+
+                // Compute up for all {B} batches from registers
+{shared_up_qdot}
+            }}
+
+            ws_gate += BLOCK_SIZE; ws_up += BLOCK_SIZE;
+            sc_gate += {sc_stride}; sc_up += {sc_stride};
+            bi_gate += {sc_stride}; bi_up += {sc_stride};
+            x_base += BLOCK_SIZE;
+        }}
+
+        // SwiGLU epilogue + write for all {B} batches (with inv_rms)
+{shared_write}
+    }}
+"""
+
+_batched_softmax_topk_swiglu_cache = {}
+
+
+def _get_batched_softmax_topk_swiglu_kernel(
+    N_INTER, SHARED_INTER, K, n_active, B,
+    n_experts=256, top_k=10, norm_topk=True, group_size=64,
+    n_oproj_tg=64,
+):
+    key = (N_INTER, SHARED_INTER, K, n_active, B, n_experts, top_k, norm_topk, group_size, n_oproj_tg)
+    if key not in _batched_softmax_topk_swiglu_cache:
+        nt_tag = "_nt" if norm_topk else ""
+        _batched_softmax_topk_swiglu_cache[key] = mx.fast.metal_kernel(
+            name=(f"batched_softmax_topk_swiglu_8bit"
+                  f"_NI{N_INTER}_SI{SHARED_INTER}_K{K}"
+                  f"_na{n_active}_B{B}_E{n_experts}_k{top_k}"
+                  f"_gs{group_size}{nt_tag}"),
+            input_names=[
+                "W", "S", "B_q",                    # routed expert weights
+                "W_shared", "S_shared", "B_shared",  # shared expert weights
+                "W_seg", "S_seg", "B_seg",            # shared_expert_gate weights
+                "X",                                  # h_scaled (B, K) bf16
+                "gate_part_a", "gate_part_b",         # (B, E) f32
+                "x2_partials",                        # (B, N_OPROJ_TG) f32
+            ],
+            output_names=["Y_routed", "Y_shared", "out_inds",
+                          "norm_scores", "gate_raw"],
+            source=_gen_batched_softmax_topk_swiglu_source(
+                N_INTER, SHARED_INTER, K, n_active, B,
+                n_experts, top_k, norm_topk, group_size, n_oproj_tg),
+        )
+    return _batched_softmax_topk_swiglu_cache[key]
+
+
+def batched_softmax_topk_swiglu_8bit(
+    w_gu, s_gu, b_gu,               # routed gate+up weights (E, 2*N_INTER, K/4)
+    w_shared, s_shared, b_shared,    # shared gate+up weights (2*SHARED_INTER, K/4)
+    w_seg, s_seg, b_seg,             # shared_expert_gate weights (1, K/4)
+    h_scaled,                        # (B, K) bf16 — from Dispatch 1
+    gate_part_a,                     # (B, E) f32 — from Dispatch 1
+    gate_part_b,                     # (B, E) f32 — from Dispatch 1
+    x2_partials,                     # (B, N_OPROJ_TG) f32 — from Dispatch 1
+    n_inter, k_hidden, batch_size, n_active,
+    n_oproj_tg, n_experts=256,
+    shared_inter=None, group_size=64,
+):
+    """Batched Dispatch 2: softmax + top-k + merged 8-bit SwiGLU with oproj prologue.
+
+    Prologue (per-batch):
+      Phase 1: distributed x2 -> inv_rms (all TGs, indexed by batch_id)
+      Phase 2: softmax(gate_scores) -> top-k -> norm_topk_prob (all routed TGs)
+      Phase 3: shared_expert_gate 8-bit GEMV -> gate_raw[B] (shared TG, SG 0)
+
+    Body:
+      Routed: 8-bit gate+up+SwiGLU with h_scaled[batch_id] input, inv_rms factored
+      Shared: register-level weight sharing for B batch elements
+
+    Args:
+        w_gu: stacked routed weights (E, 2*N_INTER, K/4) uint32
+        s_gu: routed scales (E, 2*N_INTER, K/gs) bf16
+        b_gu: routed biases (E, 2*N_INTER, K/gs) bf16
+        w_shared: shared gate+up stacked (2*SHARED_INTER, K/4) uint32
+        s_shared: shared scales (2*SHARED_INTER, K/gs) bf16
+        b_shared: shared biases (2*SHARED_INTER, K/gs) bf16
+        w_seg/s_seg/b_seg: shared_expert_gate 8-bit weights (1, K/4) uint32
+        h_scaled: (B, K) bf16 — h * w_rms from Dispatch 1
+        gate_part_a: (B, E) f32 — partial gate scores from Dispatch 1
+        gate_part_b: (B, E) f32 — partial gate scores from Dispatch 1
+        x2_partials: (B, N_OPROJ_TG) f32 — per-TG x2 sums from Dispatch 1
+        n_inter: routed intermediate size
+        k_hidden: hidden size K
+        batch_size: B (1..8)
+        n_active: experts per token (top_k)
+        n_oproj_tg: number of o_proj TGs (for x2 partial sum reduction)
+        n_experts: total number of experts E
+        shared_inter: shared intermediate size (defaults to n_inter)
+        group_size: quantization group size (default 64)
+
+    Returns:
+        (Y_routed, Y_shared, out_inds, norm_scores, gate_raw):
+            Y_routed: (B * n_active, n_inter) f32
+            Y_shared: (B, shared_inter) f32
+            out_inds: (B * n_active,) uint32
+            norm_scores: (B * n_active,) f32
+            gate_raw: (B,) f32 — raw shared expert gate values (sigmoid in epilogue)
+    """
+    B = int(batch_size)
+    n_inter_val = int(n_inter)
+    shared_inter_val = int(shared_inter) if shared_inter is not None else n_inter_val
+    k_val = int(k_hidden)
+    n_active_val = int(n_active)
+    top_k = n_active_val
+    E = int(n_experts)
+    n_oproj_tg_val = int(n_oproj_tg)
+
+    kern = _get_batched_softmax_topk_swiglu_kernel(
+        n_inter_val, shared_inter_val, k_val, n_active_val, B,
+        E, top_k, True, int(group_size), n_oproj_tg_val,
+    )
+
+    max_inter = max(n_inter_val, shared_inter_val)
+    total_routed = B * n_active_val
+
+    results = kern(
+        inputs=[
+            w_gu, s_gu, b_gu,
+            w_shared, s_shared, b_shared,
+            w_seg, s_seg, b_seg,
+            h_scaled,
+            gate_part_a, gate_part_b,
+            x2_partials,
+        ],
+        output_shapes=[
+            (total_routed * n_inter_val,),        # Y_routed flat
+            (B * shared_inter_val,),               # Y_shared flat
+            (total_routed,),                       # out_inds
+            (total_routed,),                       # norm_scores
+            (B,),                                  # gate_raw
+        ],
+        output_dtypes=[
+            mx.float32,    # Y_routed
+            mx.float32,    # Y_shared
+            mx.uint32,     # out_inds
+            mx.float32,    # norm_scores
+            mx.float32,    # gate_raw
+        ],
+        grid=(32, ceil_div(max_inter, 8) * 2, total_routed + 1),
+        threadgroup=(32, 2, 1),
+    )
+
+    Y_routed = results[0].reshape(total_routed, n_inter_val)
+    Y_shared = results[1].reshape(B, shared_inter_val)
+    out_inds = results[2]
+    norm_scores = results[3]
+    gate_raw = results[4]
+
+    return Y_routed, Y_shared, out_inds, norm_scores, gate_raw
@@ -0,0 +1,288 @@
+"""Fused GDN projections for Qwen3.5-35B-A3B (Dispatch 2).
+
+Single dispatch fuses 4 quantized 8-bit GEMVs + depthwise conv1d + activations:
+  - in_proj_qkv (8192×2048): GEMV → conv1d(4-tap) → SiLU → write bf16 + cache update
+  - in_proj_z   (4096×2048): GEMV → SiLU → write f32
+  - in_proj_b   (32×2048):   GEMV → sigmoid → write f32 (beta for GDN kernel)
+  - in_proj_a   (32×2048):   GEMV → g=exp(-exp(A_log)*softplus(a+dt_bias)) → write f32
+
+All 4 projection weight matrices are pre-merged into one contiguous buffer
+(W_merged, S_merged, B_merged) for better memory locality and cache behavior.
+Merging is done offline at patch time by _patch_gdn_proj_weights().
+
+B/A epilogues compute g and beta in-kernel, eliminating ~8 micro-dispatches
+that gated_delta_update would otherwise generate (sigmoid, exp, log, etc.).
+The caller can pass g/beta directly to gated_delta_kernel.
+
+TG-level multiplexing: tgid.y routes to different epilogues.
+Each TG: 64 threads = 2 SGs of 32, produces 8 output rows (4 per SG).
+Standard 8-bit affine GEMV: result = scale * Σ(x[i]*w[i]) + bias * Σ(x[i])
+
+Grid: (32, total_tg * 2, B), TG: (32, 2, 1)
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_fused_gdn_projections_source(K, N_QKV, N_Z, N_B, N_A, group_size=64):
+    """Generate Metal source for fused GDN projections with merged weights.
+
+    All constants baked into Metal source (no scalar kernel inputs).
+    """
+    gs = int(group_size)
+    sc_stride = 256 // gs
+    slid_div = gs // 8
+    N_TOTAL = N_QKV + N_Z + N_B + N_A
+    K_groups = K // gs
+    N_QKV_TG = ceil_div(N_QKV, 8)
+    N_Z_TG = ceil_div(N_Z, 8)
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int GROUP_SIZE = {gs};
+    const int SC_STRIDE = {sc_stride};
+    const int SLID_DIV = {slid_div};
+    const int K = {K};
+    const int K_groups = {K_groups};
+    const int N_QKV = {N_QKV};
+    const int N_Z = {N_Z};
+    const int N_B = {N_B};
+    const int N_TOTAL = {N_TOTAL};
+    const int N_QKV_TG = {N_QKV_TG};
+    const int N_Z_TG = {N_Z_TG};
+    const int N_B_TG = {ceil_div(N_B, 8)};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;  // 0 or 1
+    uint slid = thread_index_in_simdgroup;       // 0..31
+    int b_idx = tgid.z;
+
+    int tg = tgid.y;
+
+    // ─── Determine region and absolute out_row in merged matrix ───
+    int out_row;
+    int region;  // 0=QKV, 1=Z, 2=B, 3=A
+
+    if (tg < N_QKV_TG) {{
+        region = 0;
+        out_row = tg * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_QKV_TG + N_Z_TG) {{
+        region = 1;
+        out_row = N_QKV + (tg - N_QKV_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_QKV_TG + N_Z_TG + N_B_TG) {{
+        region = 2;
+        out_row = N_QKV + N_Z + (tg - N_QKV_TG - N_Z_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else {{
+        region = 3;
+        out_row = N_QKV + N_Z + N_B + (tg - N_QKV_TG - N_Z_TG - N_B_TG) * 8 + sgid * RESULTS_PER_SG;
+    }}
+
+    if (out_row >= N_TOTAL) return;
+
+    // ─── Single pointer into merged weight buffer ───
+    const device uint8_t* ws = (const device uint8_t*)W_merged + (long)out_row * K + slid * VALUES_PER_THREAD;
+    const device bfloat16_t* sc = (const device bfloat16_t*)S_merged + (long)out_row * K_groups + slid / SLID_DIV;
+    const device bfloat16_t* bi = (const device bfloat16_t*)B_merged + (long)out_row * K_groups + slid / SLID_DIV;
+
+    // ─── 8-bit GEMV K-loop (unified for all regions) ───
+    float result[4] = {{0, 0, 0, 0}};
+    int x_base = b_idx * K + slid * VALUES_PER_THREAD;
+
+    for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{
+        float x_thread[8];
+        float xsum = 0;
+        for (int i = 0; i < 8; i++) {{
+            float xi = float(x[x_base + i]);
+            x_thread[i] = xi;
+            xsum += xi;
+        }}
+
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            const device uint8_t* w = ws + row * K;
+            float s_val = float(sc[row * K_groups]);
+            float b_val = float(bi[row * K_groups]);
+            float accum = 0;
+            for (int i = 0; i < 8; i++) {{
+                accum += x_thread[i] * float(w[i]);
+            }}
+            result[row] += s_val * accum + xsum * b_val;
+        }}
+
+        ws += BLOCK_SIZE;
+        sc += SC_STRIDE;
+        bi += SC_STRIDE;
+        x_base += BLOCK_SIZE;
+    }}
+
+    // ─── Reduction ───
+    for (int row = 0; row < RESULTS_PER_SG; row++) {{
+        result[row] = simd_sum(result[row]);
+    }}
+
+    // ─── Region-specific epilogues ───
+    // After simd_sum, all 32 threads have result[0..3].
+    // Threads 0-3 each handle one output row.
+
+    if (region == 0) {{
+        // ═══ QKV: conv1d(4-tap) + SiLU + cache update ═══
+        int c = out_row + (int)slid;  // channel index (= absolute row for QKV)
+        if (slid < (uint)RESULTS_PER_SG && c < N_QKV) {{
+            float qkv_val = result[slid];
+
+            int conv_dim = N_QKV;
+            long cs_base = (long)b_idx * 3 * conv_dim;
+            float s0 = float(conv_state[cs_base + 0 * conv_dim + c]);
+            float s1 = float(conv_state[cs_base + 1 * conv_dim + c]);
+            float s2 = float(conv_state[cs_base + 2 * conv_dim + c]);
+
+            float conv_out = float(conv_w[c * 4 + 0]) * s0
+                           + float(conv_w[c * 4 + 1]) * s1
+                           + float(conv_w[c * 4 + 2]) * s2
+                           + float(conv_w[c * 4 + 3]) * qkv_val;
+
+            float silu_out = conv_out / (1.0f + metal::exp(-conv_out));
+
+            conv_state_out[cs_base + 0 * conv_dim + c] = static_cast<bfloat16_t>(s1);
+            conv_state_out[cs_base + 1 * conv_dim + c] = static_cast<bfloat16_t>(s2);
+            conv_state_out[cs_base + 2 * conv_dim + c] = static_cast<bfloat16_t>(qkv_val);
+
+            qkv_out[b_idx * conv_dim + c] = static_cast<bfloat16_t>(silu_out);
+        }}
+
+    }} else if (region == 1) {{
+        // ═══ Z: SiLU → write f32 ═══
+        int z_row = out_row - N_QKV + (int)slid;
+        if (slid < (uint)RESULTS_PER_SG && z_row < N_Z) {{
+            float val = result[slid];
+            float silu_val = val / (1.0f + metal::exp(-val));
+            z_silu_out[b_idx * N_Z + z_row] = silu_val;
+        }}
+
+    }} else if (region == 2) {{
+        // ═══ B: sigmoid(result) → beta (f32) ═══
+        int b_row = out_row - N_QKV - N_Z + (int)slid;
+        if (slid < (uint)RESULTS_PER_SG && b_row < N_B) {{
+            float val = result[slid];
+            float beta = 1.0f / (1.0f + metal::exp(-val));
+            b_out[b_idx * N_B + b_row] = beta;
+        }}
+
+    }} else {{
+        // ═══ A: g = exp(-exp(A_log) * softplus(a + dt_bias)) → f32 ═══
+        int a_row = out_row - N_QKV - N_Z - N_B + (int)slid;
+        int N_A = N_TOTAL - N_QKV - N_Z - N_B;
+        if (slid < (uint)RESULTS_PER_SG && a_row < N_A) {{
+            float a_val = result[slid];
+            float dt = float(dt_bias_arr[a_row]);
+            float x_g = a_val + dt;
+            // softplus(x) = log(1 + exp(x)), with x>20 shortcut for numerical stability
+            float sp = (x_g > 20.0f) ? x_g : metal::log(1.0f + metal::exp(x_g));
+            float g_val = metal::exp(-metal::exp(float(A_log_arr[a_row])) * sp);
+            a_out[b_idx * N_A + a_row] = g_val;
+        }}
+    }}
+"""
+
+
+_fused_gdn_proj_cache = {}
+
+
+def _get_fused_gdn_proj_kernel(K, N_QKV, N_Z, N_B, N_A, group_size=64):
+    key = (K, N_QKV, N_Z, N_B, N_A, group_size)
+    if key not in _fused_gdn_proj_cache:
+        _fused_gdn_proj_cache[key] = mx.fast.metal_kernel(
+            name=f"fused_gdn_proj_K{K}_NQKV{N_QKV}_NZ{N_Z}_NB{N_B}_NA{N_A}",
+            input_names=[
+                "x",
+                "W_merged", "S_merged", "B_merged",
+                "conv_state", "conv_w",
+                "A_log_arr", "dt_bias_arr",
+            ],
+            output_names=["qkv_out", "z_silu_out", "b_out", "a_out", "conv_state_out"],
+            source=_gen_fused_gdn_projections_source(K, N_QKV, N_Z, N_B, N_A, group_size),
+        )
+    return _fused_gdn_proj_cache[key]
+
+
+def fused_gdn_projections(
+    x,
+    W_merged, S_merged, B_merged,
+    proj_dims,
+    conv_state, conv_weights,
+    A_log, dt_bias,
+    batch_size=1,
+):
+    """Fused GDN projections: 4 GEMVs + conv1d + activations + g/beta.
+
+    Uses pre-merged contiguous weight buffers for all 4 projections.
+    B epilogue computes beta = sigmoid(b) in f32.
+    A epilogue computes g = exp(-exp(A_log) * softplus(a + dt_bias)) in f32.
+    Caller passes g/beta directly to gated_delta_kernel (no micro-dispatches).
+
+    Args:
+        x: [B, 1, K] bf16 — post-RMSNorm hidden state
+        W_merged: [N_TOTAL, K/4] uint32 — merged quantized weights
+        S_merged: [N_TOTAL, K/gs] bf16 — merged scales
+        B_merged: [N_TOTAL, K/gs] bf16 — merged biases
+        proj_dims: (N_QKV, N_Z, N_B, N_A) — per-projection output dims
+        conv_state:  [B, 3, conv_dim] bf16 — previous 3 timesteps
+        conv_weights: [conv_dim, 4, 1] or [conv_dim, 4] bf16 — depthwise conv filters
+        A_log: [Hv] f32 — GDN decay log-parameter
+        dt_bias: [Hv] f32 — GDN time constant bias
+        batch_size: int
+
+    Returns:
+        qkv_conv_silu: [B, 1, N_QKV] bf16 — post-conv, post-SiLU
+        z_silu:        [B, 1, N_Z] f32  — post-SiLU
+        beta:          [B, 1, N_B] f32  — sigmoid(b), ready for GDN kernel
+        g:             [B, 1, N_A] f32  — gating, ready for GDN kernel
+        conv_state_out: [B, 3, N_QKV] bf16
+    """
+    B = batch_size
+
+    N_QKV, N_Z, N_B, N_A = proj_dims
+    K = x.shape[-1]
+
+    kern = _get_fused_gdn_proj_kernel(K, N_QKV, N_Z, N_B, N_A)
+
+    N_QKV_TG = ceil_div(N_QKV, 8)
+    N_Z_TG = ceil_div(N_Z, 8)
+    N_B_TG = ceil_div(N_B, 8)
+    N_A_TG = ceil_div(N_A, 8)
+    total_tg = N_QKV_TG + N_Z_TG + N_B_TG + N_A_TG
+
+    conv_w_flat = conv_weights.reshape(-1, 4) if conv_weights.ndim == 3 else conv_weights
+    x_flat = x.reshape(B, K)
+
+    results = kern(
+        inputs=[
+            x_flat,
+            W_merged, S_merged, B_merged,
+            conv_state, conv_w_flat,
+            A_log, dt_bias,
+        ],
+        output_shapes=[
+            (B * N_QKV,),              # qkv_out
+            (B * N_Z,),                # z_silu_out
+            (B * N_B,),                # beta_out (f32)
+            (B * N_A,),                # g_out (f32)
+            (B * 3 * N_QKV,),          # conv_state_out
+        ],
+        output_dtypes=[mx.bfloat16, mx.float32, mx.float32, mx.float32, mx.bfloat16],
+        grid=(32, total_tg * 2, B),
+        threadgroup=(32, 2, 1),
+    )
+
+    qkv_out = results[0].reshape(B, 1, N_QKV)
+    z_silu = results[1].reshape(B, 1, N_Z)
+    beta = results[2].reshape(B, 1, N_B)
+    g = results[3].reshape(B, 1, N_A)
+    conv_state_out = results[4].reshape(B, 3, N_QKV)
+
+    return qkv_out, z_silu, beta, g, conv_state_out
@@ -0,0 +1,128 @@
+"""Fused Q/K per-head L2-norm for GDN attention (Dispatch 3).
+
+Performs per-head L2 normalization on q and k vectors with different scaling.
+Matches vLLM and latest mlx-lm (qwen3_5.py) which use rsqrt(sum(x²) + eps),
+NOT rms_norm which uses rsqrt(mean(x²) + eps).
+
+From qwen3_5.py (updated to match vLLM):
+    inv_scale = Dk^(-0.5) = 128^(-0.5)
+    q = inv_scale * q * rsqrt(sum(q²) + 1e-6)   → L2-normalize then scale by 1/√Dk
+    k = k * rsqrt(sum(k²) + 1e-6)                → L2-normalize only (no extra scale)
+
+Grid: (32 heads × 32 threads, 1, B).
+Each TG = 32 threads = 1 SG, handles one 128-dim head.
+Dk=128 = 32 threads × 4 elements → exactly 1 SG, no cross-SG reduction.
+"""
+
+import mlx.core as mx
+
+
+def _gen_fused_qk_rmsnorm_source():
+    """Generate Metal source for fused Q/K per-head L2-norm.
+
+    Input: qkv [B, 8192] bf16 (flattened from [B, 1, 8192])
+      - [0, 2048): q = 16 heads × 128
+      - [2048, 4096): k = 16 heads × 128
+      - [4096, 8192): v (untouched)
+
+    Output: qk_out [B, 4096] bf16
+      - [0, 2048): q L2-normalized then scaled by 1/√Dk
+      - [2048, 4096): k L2-normalized (no extra scale)
+
+    Grid: (32 * 32, 1, B), TG: (32, 1, 1)
+      tgid.x 0..15:  q heads → scale = 1/√128
+      tgid.x 16..31: k heads → scale = 1.0
+      tgid.z: batch index
+    """
+    return """
+    const int N_READS = 4;
+    const int DK = 128;
+    const int HK = 16;
+    const float EPS = 1e-6f;
+    const float Q_SCALE = rsqrt(128.0f);        // inv_scale = Dk^(-0.5)
+    const float K_SCALE = 1.0f;                  // no extra scale for k
+
+    uint head_idx = threadgroup_position_in_grid.x;
+    uint slid = thread_index_in_simdgroup;
+    uint b_idx = thread_position_in_grid.z;
+
+    bool is_q = (head_idx < (uint)HK);
+
+    // Input offset: q heads at [0, 2048), k heads at [2048, 4096)
+    int in_base = is_q
+        ? (b_idx * 8192 + head_idx * DK)
+        : (b_idx * 8192 + 2048 + (head_idx - HK) * DK);
+
+    // Output offset: q at [0, 2048), k at [2048, 4096)
+    int out_base = b_idx * 4096 + head_idx * DK;
+
+    // ── Phase 1: Load 4 elements + sum of squares ──
+    float vals[4];
+    float partial_sq = 0.0f;
+    int elem_base = slid * N_READS;
+
+    for (int i = 0; i < N_READS; i++) {
+        float xi = float(qkv[in_base + elem_base + i]);
+        vals[i] = xi;
+        partial_sq += xi * xi;
+    }
+
+    // ── Phase 2: simd reduction (32 threads → full sum of 128 elements) ──
+    float sum_sq = simd_sum(partial_sq);
+
+    // ── Phase 3: compute L2 inv-norm (NOT rms_norm — no /Dk) ──
+    float inv_rms = metal::precise::rsqrt(sum_sq + EPS);
+
+    // ── Phase 4: scale and write ──
+    float scale = is_q ? Q_SCALE : K_SCALE;
+    float combined = inv_rms * scale;
+
+    for (int i = 0; i < N_READS; i++) {
+        qk_out[out_base + elem_base + i] = static_cast<bfloat16_t>(vals[i] * combined);
+    }
+"""
+
+
+_fused_qk_rmsnorm_kernel = None
+
+
+def _get_fused_qk_rmsnorm_kernel():
+    """Get or compile the fused Q/K RMSNorm kernel."""
+    global _fused_qk_rmsnorm_kernel
+    if _fused_qk_rmsnorm_kernel is None:
+        _fused_qk_rmsnorm_kernel = mx.fast.metal_kernel(
+            name="fused_qk_rmsnorm",
+            input_names=["qkv"],
+            output_names=["qk_out"],
+            source=_gen_fused_qk_rmsnorm_source(),
+        )
+    return _fused_qk_rmsnorm_kernel
+
+
+def fused_qk_rmsnorm(qkv_conv_silu, batch_size=1):
+    """Fused Q/K per-head RMSNorm for GDN attention.
+
+    Args:
+        qkv_conv_silu: [B, 1, 8192] bf16 — post-conv, post-SiLU output from Dispatch 2.
+            First 2048 = q (16 heads × 128), next 2048 = k, last 4096 = v.
+        batch_size: int — batch dimension.
+
+    Returns:
+        qk_normed: [B, 1, 4096] bf16 — normalized q (first 2048) and k (next 2048).
+            v is NOT copied; Dispatch 4 reads v directly from qkv_conv_silu[:, :, 4096:].
+    """
+    B = batch_size
+    kern = _get_fused_qk_rmsnorm_kernel()
+
+    # Flatten to [B, 8192] for kernel
+    qkv_flat = qkv_conv_silu.reshape(B, 8192)
+
+    n_heads = 32  # 16 q + 16 k
+    results = kern(
+        inputs=[qkv_flat],
+        output_shapes=[(B * 4096,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(n_heads * 32, 1, B),
+        threadgroup=(32, 1, 1),
+    )
+    return results[0].reshape(B, 1, 4096)
@@ -0,0 +1,117 @@
+"""Fused RMSNormGated for GDN attention (Dispatch 5).
+
+Fuses RMSNorm(out, weight) × z_silu into one kernel.
+SiLU on z was already applied in Dispatch 2, so z_silu arrives as f32.
+
+From qwen3_next.py (Qwen3NextRMSNormGated):
+    x = rms_norm(hidden_states, weight, eps)     # weight: [Dv=128]
+    gate = silu(z.float())                        # already done in Dispatch 2
+    return (gate * x).to(hidden_states.dtype)
+
+Grid: (32 heads × 32 threads, 1, B).
+Each TG = 32 threads = 1 SG, handles one 128-dim head.
+Dv=128 = 32 threads × 4 elements → exactly 1 SG.
+"""
+
+import mlx.core as mx
+
+
+def _gen_fused_rms_norm_gated_source():
+    """Generate Metal source for fused RMSNormGated.
+
+    Inputs:
+      gdn_out: [B, Hv*Dv] bf16  — GDN output, flattened (Hv=32, Dv=128)
+      z_silu:  [B, Hv*Dv] f32   — post-SiLU z from Dispatch 2
+      weight:  [Dv] f32         — RMSNormGated learned weight (128 elements)
+
+    Output:
+      out: [B, Hv*Dv] bf16      — result = z_silu * rms_norm(gdn_out, weight)
+
+    Grid: (32 * 32, 1, B), TG: (32, 1, 1)
+      tgid.x: head index (0..31)
+      tgid.z: batch index
+    """
+    return """
+    const int N_READS = 4;
+    const int DV = 128;
+    const int HV = 32;
+    const float EPS = 1e-6f;
+
+    uint head_idx = threadgroup_position_in_grid.x;
+    uint slid = thread_index_in_simdgroup;
+    uint b_idx = thread_position_in_grid.z;
+
+    int base = b_idx * HV * DV + head_idx * DV;
+    int elem_base = slid * N_READS;
+
+    // ── Phase 1: Load gdn_out elements + sum of squares ──
+    float gdn_vals[4];
+    float partial_sq = 0.0f;
+
+    for (int i = 0; i < N_READS; i++) {
+        float xi = float(gdn_out[base + elem_base + i]);
+        gdn_vals[i] = xi;
+        partial_sq += xi * xi;
+    }
+
+    // ── Phase 2: simd reduction (32 threads → full sum of 128 elements) ──
+    float sum_sq = simd_sum(partial_sq);
+
+    // ── Phase 3: compute inv_rms ──
+    float inv_rms = metal::precise::rsqrt(sum_sq / float(DV) + EPS);
+
+    // ── Phase 4: RMSNorm × z_silu, write bf16 ──
+    for (int i = 0; i < N_READS; i++) {
+        int idx = elem_base + i;
+        float w = float(weight[idx]);                           // learned weight[Dv]
+        float normed = gdn_vals[i] * inv_rms * w;              // RMSNorm
+        float z_val = z_silu[base + idx];                       // already f32, post-SiLU
+        out[base + idx] = static_cast<bfloat16_t>(z_val * normed);
+    }
+"""
+
+
+_fused_rms_norm_gated_kernel = None
+
+
+def _get_fused_rms_norm_gated_kernel():
+    """Get or compile the fused RMSNormGated kernel."""
+    global _fused_rms_norm_gated_kernel
+    if _fused_rms_norm_gated_kernel is None:
+        _fused_rms_norm_gated_kernel = mx.fast.metal_kernel(
+            name="fused_rms_norm_gated",
+            input_names=["gdn_out", "z_silu", "weight"],
+            output_names=["out"],
+            source=_gen_fused_rms_norm_gated_source(),
+        )
+    return _fused_rms_norm_gated_kernel
+
+
+def fused_rms_norm_gated(gdn_out, z_silu, weight, batch_size=1):
+    """Fused RMSNormGated: RMSNorm(out, weight) × z_silu.
+
+    Args:
+        gdn_out: [B, 1, Hv, Dv] bf16 — GDN recurrence output (Hv=32, Dv=128).
+        z_silu:  [B, 1, 4096] f32 — post-SiLU z from Dispatch 2.
+        weight:  [128] f32 — RMSNormGated learned weight (Dv elements).
+        batch_size: int.
+
+    Returns:
+        out: [B, 1, 4096] bf16 — ready for out_proj in Dispatch 6.
+    """
+    B = batch_size
+    kern = _get_fused_rms_norm_gated_kernel()
+
+    # Flatten to [B, 4096]
+    gdn_flat = gdn_out.reshape(B, 4096)
+    z_flat = z_silu.reshape(B, 4096)
+
+    n_heads = 32  # Hv
+    results = kern(
+        inputs=[gdn_flat, z_flat, weight],
+        output_shapes=[(B * 4096,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(n_heads * 32, 1, B),
+        threadgroup=(32, 1, 1),
+    )
+    return results[0].reshape(B, 1, 4096)
@@ -0,0 +1,177 @@
+"""GDN recurrence with pre-computed g and beta (Dispatch 4).
+
+Modified version of gated_delta_step from mlx-lm-fork/mlx_lm/models/gated_delta.py.
+Instead of computing g = exp(-exp(A_log) * softplus(a + dt_bias)) and beta = sigmoid(b)
+inside the kernel, accepts them as pre-computed f32 inputs from Dispatch 2.
+
+Non-vectorized only (Qwen3.5-35B-A3B uses scalar gating per head).
+
+Grid: (32, Dv, B*Hv) = (32, 128, B*32), TG: (32, 4, 1)
+"""
+
+from typing import Optional, Tuple
+
+import mlx.core as mx
+
+
+def _make_gdn_precomputed_kernel(has_mask=False):
+    """Build the GDN kernel with pre-computed g and beta."""
+    if not mx.metal.is_available():
+        return None
+
+    mask_source = "mask[b_idx * T + t]" if has_mask else "true"
+
+    source = f"""
+        auto n = thread_position_in_grid.z;
+        auto b_idx = n / Hv;
+        auto hv_idx = n % Hv;
+        auto hk_idx = hv_idx / (Hv / Hk);
+        constexpr int n_per_t = Dk / 32;
+
+        // q, k: [B, T, Hk, Dk]
+        auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
+        auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
+
+        // v, y: [B, T, Hv, Dv]
+        auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
+        y += b_idx * T * Hv * Dv + hv_idx * Dv;
+
+        auto dk_idx = thread_position_in_threadgroup.x;
+        auto dv_idx = thread_position_in_grid.y;
+
+        // state_in, state_out: [B, Hv, Dv, Dk]
+        auto i_state = state_in + (n * Dv + dv_idx) * Dk;
+        auto o_state = state_out + (n * Dv + dv_idx) * Dk;
+
+        float state[n_per_t];
+        for (int i = 0; i < n_per_t; ++i) {{
+          auto s_idx = n_per_t * dk_idx + i;
+          state[i] = static_cast<float>(i_state[s_idx]);
+        }}
+
+        // g: [B, T, Hv] f32 — pre-computed decay gate
+        auto g_ = g + b_idx * T * Hv;
+        // beta: [B, T, Hv] f32 — pre-computed sigmoid(b)
+        auto beta_ = beta + b_idx * T * Hv;
+
+        for (int t = 0; t < T; ++t) {{
+          if ({mask_source}) {{
+            // Pre-computed g and beta (no softplus/exp/sigmoid needed)
+            float g_val = g_[hv_idx];
+            float beta_val = beta_[hv_idx];
+
+            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_val;
+              kv_mem += state[i] * k_[s_idx];
+            }}
+            kv_mem = simd_sum(kv_mem);
+
+            auto delta = (v_[dv_idx] - kv_mem) * beta_val;
+
+            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]);
+        }}
+    """
+
+    inputs = ["q", "k", "v", "g", "beta", "state_in", "T"]
+    if has_mask:
+        inputs.append("mask")
+
+    suffix = "_precomputed"
+    if has_mask:
+        suffix += "_mask"
+
+    return mx.fast.metal_kernel(
+        name=f"gated_delta_step{suffix}",
+        input_names=inputs,
+        output_names=["y", "state_out"],
+        source=source,
+    )
+
+
+_gdn_precomputed_kernel = None
+_gdn_precomputed_kernel_masked = None
+
+
+def _get_gdn_precomputed_kernel(has_mask=False):
+    """Get or compile the pre-computed GDN kernel."""
+    global _gdn_precomputed_kernel, _gdn_precomputed_kernel_masked
+    if has_mask:
+        if _gdn_precomputed_kernel_masked is None:
+            _gdn_precomputed_kernel_masked = _make_gdn_precomputed_kernel(has_mask=True)
+        return _gdn_precomputed_kernel_masked
+    else:
+        if _gdn_precomputed_kernel is None:
+            _gdn_precomputed_kernel = _make_gdn_precomputed_kernel(has_mask=False)
+        return _gdn_precomputed_kernel
+
+
+def gated_delta_update_precomputed(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: mx.array,
+    mask: Optional[mx.array] = None,
+) -> Tuple[mx.array, mx.array]:
+    """GDN recurrence with pre-computed g and beta.
+
+    Args:
+        q: [B, T, Hk, Dk] bf16 — normalized q from Dispatch 3
+        k: [B, T, Hk, Dk] bf16 — normalized k from Dispatch 3
+        v: [B, T, Hv, Dv] bf16 — v from Dispatch 2 (qkv_conv_silu[:, :, 4096:])
+        g: [B, T, Hv] f32 — pre-computed decay gate from Dispatch 2
+        beta: [B, T, Hv] f32 — pre-computed sigmoid(b) from Dispatch 2
+        state: [B, Hv, Dv, Dk] bf16 — recurrent state from cache
+        mask: [B, T] optional
+
+    Returns:
+        y: [B, T, Hv, Dv] bf16
+        new_state: [B, Hv, Dv, Dk] bf16
+    """
+    B, T, Hk, Dk = k.shape
+    Hv, Dv = v.shape[2:]
+    input_type = q.dtype
+
+    kernel = _get_gdn_precomputed_kernel(has_mask=mask is not None)
+    inputs = [q, k, v, g, beta, state, T]
+    if mask is not None:
+        inputs.append(mask)
+
+    return kernel(
+        inputs=inputs,
+        template=[
+            ("InT", input_type),
+            ("Dk", Dk),
+            ("Dv", Dv),
+            ("Hk", Hk),
+            ("Hv", Hv),
+        ],
+        grid=(32, Dv, B * Hv),
+        threadgroup=(32, 4, 1),
+        output_shapes=[(B, T, Hv, Dv), state.shape],
+        output_dtypes=[input_type, input_type],
+    )
@@ -0,0 +1,324 @@
+#!/usr/bin/env python3
+"""MTP Speculative Decoding integrated with mlx_lm's BatchGenerator.
+
+Subclasses BatchGenerator to add MTP drafting + S>1 verification with
+correct GDN state rollback via SpeculativeArraysCache.
+
+At BS=1: drafts γ tokens with MTP, verifies at S=γ+1, buffers accepted tokens.
+At BS>1: falls back to standard BatchGenerator (no speculative).
+
+Usage:
+    from mtp_batch_generator import MTPBatchGenerator
+    gen = MTPBatchGenerator(model, mtp_predictor, gamma=2, ...)
+    gen.insert([prompt_tokens])
+    while True:
+        responses = gen.next()
+"""
+
+import time
+
+import mlx.core as mx
+from mlx_lm.generate import BatchGenerator, generation_stream
+
+from .mtp_module import MTPPredictor, speculative_forward, draft_tokens
+
+
+class MTPBatchGenerator(BatchGenerator):
+    """BatchGenerator with MTP speculative decoding for BS=1."""
+
+    def __init__(
+        self,
+        model,
+        mtp_predictor: MTPPredictor,
+        gamma: int = 2,
+        temp: float = 0.0,
+        alpha: float = 1.0,
+        **kwargs,
+    ):
+        super().__init__(model, **kwargs)
+        self.mtp = mtp_predictor
+        self.gamma = gamma
+        self.temp = temp
+        self.alpha = alpha
+
+        self._token_buffer = {}       # uid → [(token, logprobs), ...]
+        self._captured = {}            # pre_norm / prompt_pre_norm from norm wrapper
+        self._mtp_pre_norm = {}        # uid → (B, 1, D) pre-norm hidden state
+        self._mtp_prefilled = set()    # uids with MTP cache prefilled
+        self._request_temp = {}        # uid → temperature from request
+
+        self._setup_hidden_capture()
+
+    def _setup_hidden_capture(self):
+        """Monkey-patch model's final norm to capture pre-norm hidden state.
+
+        Captures:
+        - pre_norm: hidden states before final RMSNorm (for MTP input)
+        - prompt_pre_norm: same but only when S>1 (prefill)
+        """
+        inner = getattr(self.model, 'model', None) or self.model.language_model.model
+        original_norm = inner.norm
+        captured = self._captured
+
+        class _CapturingNorm:
+            def __init__(self, orig):
+                self._orig = orig
+                self.weight = orig.weight
+
+            def __call__(self, x):
+                captured['pre_norm'] = x
+                if x.shape[1] > 1:
+                    captured['prompt_pre_norm'] = x
+                return self._orig(x)
+
+            def __getattr__(self, name):
+                return getattr(self._orig, name)
+
+        inner.norm = _CapturingNorm(original_norm)
+
+    def _next(self):
+        batch = self.active_batch
+
+        # Yield buffered tokens first
+        if batch is not None and len(batch) == 1:
+            uid = batch.uids[0]
+            if uid in self._token_buffer and self._token_buffer[uid]:
+                return self._yield_buffered(batch, uid)
+
+        # BS=1 speculative path
+        if (batch is not None
+                and len(batch) == 1
+                and self.gamma > 0
+                and len(self.unprocessed_prompts) == 0):
+            uid = batch.uids[0]
+            if uid not in self._mtp_prefilled:
+                return self._first_step_and_prefill(batch, uid)
+            return self._speculative_next()
+
+        # Standard path (BS>1 or no batch)
+        responses = super()._next()
+        if responses and batch is not None and len(batch) == 1:
+            if 'pre_norm' in self._captured:
+                uid = batch.uids[0]
+                self._mtp_pre_norm[uid] = self._captured['pre_norm'][:, -1:, :]
+        return responses
+
+    def _first_step_and_prefill(self, batch, uid):
+        """First decode step. MTP cache already prefilled by ExoBatchGenerator.submit()."""
+        responses = super()._next()
+        if not responses:
+            return responses
+
+        # Capture decode pre_norm from this standard step for first speculative cycle
+        decode_pre_norm = self._captured.get('pre_norm')
+        if decode_pre_norm is not None:
+            mx.eval(decode_pre_norm)
+            self._mtp_pre_norm[uid] = decode_pre_norm[:, -1:, :]
+
+        self._mtp_prefilled.add(uid)
+        return responses
+
+    def _speculative_next(self):
+        """Core speculative cycle with correct GDN rollback."""
+        tic = time.perf_counter()
+        batch = self.active_batch
+        uid = batch.uids[0]
+        y = batch.y          # (1,) — token from previous step, to be yielded
+        y_val = y[0].item()
+        y_logprobs = batch.logprobs[0]
+
+        # Append current y to token history
+        batch.tokens[0] = mx.concatenate((batch.tokens[0], y[0:1]))
+
+        pre_norm = self._mtp_pre_norm.get(uid)
+        if pre_norm is None:
+            return super()._next()
+
+        gamma = self.gamma
+        temp = self._request_temp.get(uid, self.temp)
+        alpha = self.alpha
+
+        # 1. Draft γ tokens (lazy chain, no eval)
+        next_token_arr = y.reshape(1, 1)
+        draft_ids, draft_probs = draft_tokens(
+            self.mtp, pre_norm, next_token_arr, gamma, temp)
+
+        # 2. Verify via speculative_forward (handles GDN cache wrapping + kernel swap)
+        draft_concat = mx.concatenate(
+            [d.reshape(1, 1) for d in draft_ids], axis=1)  # (1, γ)
+        verify_input = mx.concatenate(
+            [next_token_arr, draft_concat], axis=1)  # (1, γ+1)
+        verify_pre_norm, verify_logits = speculative_forward(
+            self.model, verify_input, batch.cache, speculative=True)
+
+        # 3. Build acceptance check lazily
+        target_tokens = mx.argmax(verify_logits[:, :gamma, :], axis=-1)
+
+        if temp == 0:
+            matches = mx.equal(target_tokens, draft_concat).squeeze(0)
+            all_next = mx.argmax(verify_logits[0], axis=-1)
+            logprobs_all = verify_logits[0] - mx.logsumexp(
+                verify_logits[0], axis=-1, keepdims=True)
+            mx.async_eval(matches, all_next, logprobs_all, verify_pre_norm)
+        else:
+            accept_ratios = []
+            for i in range(gamma):
+                p = mx.softmax(verify_logits[0, i] / temp, axis=-1)
+                q = draft_probs[i]
+                p_di = p[draft_ids[i].squeeze()]
+                q_di = q[0, draft_ids[i].squeeze()]
+                ratio = p_di / mx.maximum(q_di, 1e-10)
+                accept_ratios.append(mx.minimum(ratio ** alpha, 1.0))
+            uniforms = mx.random.uniform(shape=(gamma,))
+            corrections = []
+            for i in range(gamma):
+                p = mx.softmax(verify_logits[0, i] / temp, axis=-1)
+                q = draft_probs[i][0]
+                residual = mx.maximum(p - q, 0.0)
+                corrections.append(mx.random.categorical(mx.log(residual + 1e-10)))
+            bonus_token = mx.random.categorical(verify_logits[0, gamma] * (1.0 / temp))
+            logprobs_all = verify_logits[0] - mx.logsumexp(
+                verify_logits[0], axis=-1, keepdims=True)
+            mx.async_eval(accept_ratios, uniforms, corrections, bonus_token,
+                          logprobs_all, verify_pre_norm, draft_concat)
+
+        # 4. Determine acceptance
+        n_accepted = 0
+        for i in range(gamma):
+            if temp == 0:
+                if matches[i].item():
+                    n_accepted += 1
+                else:
+                    break
+            else:
+                if uniforms[i].item() < accept_ratios[i].item():
+                    n_accepted += 1
+                else:
+                    break
+
+        # 5. Rollback cache
+        rollback = gamma - n_accepted
+        if rollback > 0:
+            for c in batch.cache:
+                if hasattr(c, 'offset'):
+                    c.offset -= rollback
+                elif hasattr(c, 'rollback'):
+                    c.rollback(n_accepted)
+
+        # Unwrap SpeculativeArraysCache
+        for i, c in enumerate(batch.cache):
+            if hasattr(c, 'base'):
+                batch.cache[i] = c.base
+
+        # 6. Bonus/correction token + logprobs
+        if n_accepted == gamma:
+            if temp == 0:
+                bonus_val = all_next[gamma].item()
+            else:
+                bonus_val = bonus_token.item()
+            bonus_lp = logprobs_all[gamma]
+        else:
+            if temp == 0:
+                bonus_val = all_next[n_accepted].item()
+            else:
+                bonus_val = corrections[n_accepted].item()
+            bonus_lp = logprobs_all[n_accepted]
+
+        # 7. Update MTP pre_norm for next cycle
+        self._mtp_pre_norm[uid] = verify_pre_norm[
+            :, (gamma if n_accepted == gamma else n_accepted):
+               (gamma if n_accepted == gamma else n_accepted) + 1, :]
+
+        # 8. Build token list: current y + accepted drafts
+        draft_int_values = draft_concat[0].tolist()
+        all_tokens = [(y_val, y_logprobs)]
+        for i in range(n_accepted):
+            all_tokens.append((draft_int_values[i], logprobs_all[i]))
+
+        # 9. Set batch.y = bonus for next cycle
+        batch.y = mx.array([bonus_val])
+        batch.logprobs = [bonus_lp]
+
+        # Append accepted drafts to token history
+        if n_accepted > 0:
+            batch.tokens[0] = mx.concatenate(
+                (batch.tokens[0], mx.array([t for t, _ in all_tokens[1:]])))
+        batch.num_tokens[0] += len(all_tokens)
+
+        # 10. Check stop conditions — truncate at stop token
+        toc = time.perf_counter()
+        self._stats.generation_time += toc - tic
+        self._stats.generation_tokens += len(all_tokens)
+
+        # Find first stop token or length limit in all_tokens
+        stop_idx = None
+        for idx, (tok, _) in enumerate(all_tokens):
+            if tok in self.stop_tokens:
+                stop_idx = idx
+                break
+            if batch.num_tokens[0] >= batch.max_tokens[0]:
+                stop_idx = idx
+                break
+
+        first_tok, first_lp = all_tokens[0]
+
+        if stop_idx is not None:
+            # Tokens before the stop are valid output — buffer them
+            # The stop token itself triggers finish_reason
+            valid_tokens = all_tokens[:stop_idx]
+            if valid_tokens:
+                # Yield first, buffer rest + a final stop entry
+                if len(valid_tokens) > 1:
+                    self._token_buffer[uid] = valid_tokens[1:]
+                # Append stop marker as last buffered token
+                stop_tok, stop_lp = all_tokens[stop_idx]
+                if uid not in self._token_buffer:
+                    self._token_buffer[uid] = []
+                self._token_buffer[uid].append((stop_tok, stop_lp))
+                mx.async_eval(batch.y)
+                return [self.Response(uid, first_tok, first_lp, None, lambda: None)]
+            else:
+                # Stop token is the first token — finish immediately
+                cache = batch.extract_cache(0)
+                self.active_batch = None
+                self._cleanup_uid(uid)
+                return [self.Response(uid, first_tok, first_lp, "stop", cache)]
+
+        # Buffer remaining tokens
+        if len(all_tokens) > 1:
+            self._token_buffer[uid] = all_tokens[1:]
+
+        mx.async_eval(batch.y)
+        return [self.Response(uid, first_tok, first_lp, None, lambda: None)]
+
+    def _yield_buffered(self, batch, uid):
+        """Yield one buffered token from a previous speculative cycle."""
+        tic = time.perf_counter()
+        buf = self._token_buffer[uid]
+        tok, lp = buf.pop(0)
+
+        if not buf:
+            del self._token_buffer[uid]
+
+        finish_reason = None
+        if tok in self.stop_tokens:
+            finish_reason = "stop"
+        elif batch.num_tokens[0] >= batch.max_tokens[0]:
+            finish_reason = "length"
+
+        cache = None
+        if finish_reason:
+            cache = batch.extract_cache(0)
+            self.active_batch = None
+            self._cleanup_uid(uid)
+
+        toc = time.perf_counter()
+        self._stats.generation_time += toc - tic
+        return [self.Response(uid, tok, lp, finish_reason, cache or (lambda: None))]
+
+    def _cleanup_uid(self, uid):
+        """Clean up MTP state for a finished request."""
+        self._mtp_pre_norm.pop(uid, None)
+        self._mtp_prefilled.discard(uid)
+        self._token_buffer.pop(uid, None)
+        self._request_temp.pop(uid, None)
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
dmcc73	37ad1fb3ed	Call warmup_speculative at startup to pre-compile LpB kernels The warmup_speculative() function was defined but never called. Custom Metal kernels (LpB) require first-call compilation (~200ms). Without warmup, the first speculative cycle is slow, dragging down average TPS by 10-20% on short generations. In mlx_bench testing: cold 48 TPS → warm 60 TPS for DFlash, cold 39 TPS → warm 44 TPS for MTP. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-04-01 23:15:05 +01:00
dmcc73	b47a287f3e	Add EXO_DISABLE_LOGPROBS=1 to skip per-token logprobs extraction For profiling: extract_top_logprobs() does 11 .item() calls + argpartition on 248K vocab per token. Testing if this is the source of speculative overhead vs mlx_bench. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-30 23:41:47 +01:00
dmcc73	e1cf376e45	Add speculative warmup: compile MTP + verify kernels at startup The standard warmup only runs S=1 generation, leaving speculative kernels (S>1 verify, speculative GDN kernel, MTP draft) uncompiled. First real speculative cycle had compilation overhead. New warmup_speculative(): prefills a short prompt, runs 3 speculative cycles to compile all kernels before real requests arrive. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-30 22:39:34 +01:00
dmcc73	75932cbcca	Fix stop token dropping valid tokens before it When <\|im_end\|> appeared in accepted drafts, all preceding tokens in the cycle were returned with finish_reason="stop", causing exo to drop them (exo skips adding tokens with finish_reason="stop"). Symptom: γ=0 outputs "20", γ=1 outputs "2", γ=2 outputs nothing — losing exactly γ tokens at the end. Fix: yield tokens before the stop normally (no finish_reason), buffer the stop token, let _yield_buffered return it with finish_reason="stop". Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-30 18:44:01 +01:00
dmcc73	199a4ab7e0	EXO_SPECULATIVE_TEMP overrides model sampling temperature globally When set, overrides the request's temperature for both the model's sampler AND the speculative acceptance. This allows testing greedy baseline (γ=0) and greedy speculative (γ=2) with the same T=0. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-30 18:24:25 +01:00
dmcc73	4818b9a3db	EXO_SPECULATIVE_TEMP overrides request temp when set If EXO_SPECULATIVE_TEMP is explicitly set, use it (for testing greedy). If not set, use the request's temperature (production behavior). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-30 18:23:05 +01:00
dmcc73	f0433505a8	Fix speculative temp: use request temperature, not global env var The speculative cycle was using EXO_SPECULATIVE_TEMP (global) instead of the request's actual temperature. This caused greedy decoding in speculative while the model sampled at T=0.7, producing different (shorter) output and missing responses after </think>. Now passes task_params.temperature from submit() to MTPBatchGenerator per-request via _request_temp[uid]. Falls back to self.temp (env var) if not set. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-30 18:21:57 +01:00
dmcc73	88bc1656a2	Fix MTP prefill to use all captured positions Was using prompt_pre_norm[:, :-1, :] (missing last position). Now uses full prompt_pre_norm paired with all_prompt_tokens[1:S_pre+1], matching the mlx_bench MTPBatchGenerator's prefill behavior. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-30 18:06:53 +01:00
dmcc73	7bd1ba6605	Fix MTP prefill: do it in submit() with correct prompt tokens Bug: _CapturingEmbed was overwritten by BatchGenerator's 2-token insert, causing MTP prefill to silently skip (len check failed: 2 < N-1). MTP drafted without any prompt context → low acceptance → low TPS. Fix: Do MTP prefill in ExoBatchGenerator.submit() right after main model prefill, using all_prompt_tokens (available as local variable). Remove _CapturingEmbed entirely. Simplify _first_step_and_prefill to just capture decode pre_norm. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-30 17:42:20 +01:00
dmcc73	e7c5d56e83	Add LpB kernel patches for Qwen3.5 dense models (27B, 9B) Loop-over-B custom GEMV kernels for expanding projections (N > K): gate_proj, up_proj, down_proj, in_proj_qkv, in_proj_z, out_proj, q_proj. These reduce S>1 verification cost from ~7ms/token to ~3ms/token, critical for speculative decoding speedup. Auto-detected for model_type=qwen3_5 (dense models like 27B, 9B). MoE models (qwen3_5_moe) use the existing batched fused patches instead. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-30 17:00:02 +01:00
dmcc73	dd71182457	Fix MTP prefill for exo: capture prompt tokens via embed_tokens wrapper Exo does its own prefill outside BatchGenerator, so batch.tokens only has the last 2 tokens. Added _CapturingEmbed wrapper on embed_tokens to capture the full prompt token ids during prefill. MTP prefill now uses these captured tokens instead of batch.tokens. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-30 15:57:25 +01:00
dmcc73	09012d3799	Auto-extract MTP weights from HuggingFace model repo When EXO_SPECULATIVE=1, MTP weights are resolved in order: 1. EXO_MTP_WEIGHTS=/path/to/file (explicit path) 2. EXO_MTP_MODEL=Qwen/Qwen3.5-27B (explicit HF repo) 3. Auto-detect: if model has mtp_num_hidden_layers > 0 and is Qwen3.5, defaults to Qwen/Qwen3.5-27B Downloads safetensors from HF, extracts model.mtp.* tensors, caches to ~/.cache/exo/mtp_weights/ for future use. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-30 15:19:48 +01:00
dmcc73	ce19267d2d	Pass temperature and alpha to MTP speculative decoding Default temp=0.7 (matching exo's default) so probabilistic acceptance runs correctly. Configurable via EXO_SPECULATIVE_TEMP and EXO_SPECULATIVE_ALPHA env vars. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-30 15:12:59 +01:00
dmcc73	8a65a51569	Add MTP speculative decoding for Qwen3.5 models Integrates MTP-based speculative decoding into exo's BatchGenerator. When enabled via EXO_SPECULATIVE=1 and EXO_MTP_WEIGHTS=<path>, MTPBatchGenerator replaces the standard MlxBatchGenerator for BS=1 inference, drafting γ tokens with the model's built-in MTP head and verifying at S=γ+1. New files in speculative/: - mtp_module.py: MTPPredictor + speculative_forward (kernel swap for GDN rollback) + draft_tokens (lazy MTP chaining) - mtp_batch_generator.py: MTPBatchGenerator subclassing mlx_lm's BatchGenerator with token buffering and BS>1 fallback - speculative_cache.py: SpeculativeArraysCache for GDN state rollback - speculative_gdn_kernel.py: Metal kernel with per-step state output Environment variables: EXO_SPECULATIVE=1 Enable speculative decoding EXO_MTP_WEIGHTS=/path/to/file Path to MTP weights safetensors EXO_SPECULATIVE_GAMMA=2 Draft tokens per cycle (default: 2) MTP weights must be extracted from the original HF model (e.g. Qwen/Qwen3.5-27B) as they are stripped during MLX quantization. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-30 15:10:11 +01:00
dmcc73	a2de281c67	Replace GDN projections with register-sharing batched kernel Old kernel used grid z=B, loading weights B times independently. New kernel loads weights once into registers and computes B dot products. 11-14% faster at B=2-4 in full model benchmarks (194 vs 174 TPS at B=2). B=1 generates identical code, no regression. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-19 15:13:39 +00:00
dmcc73	9394d04f5f	Add LCB TPS benchmark script Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-18 17:37:36 +00:00
dmcc73	92c04b0aa5	Add batched fused Metal kernel patches for Qwen3.5 MoE decode Custom Metal kernels with register-level weight sharing for batch sizes 1-8. Fuses o_proj + RMSNorm + gate GEMV + softmax + topk + SwiGLU + down_proj + epilogue into 4 dispatches per MoE layer, plus fused GDN and GQA attention projections. Falls back to vanilla for B>8 or S>1 (prefill). Controlled by EXO_FUSED_KERNELS env var. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-18 17:36:28 +00:00