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>

.mlx_typings/mlx/core/__init__.pyi

@@ -2396,7 +2396,7 @@ def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array:
         array: The angles in degrees.
     """
 
-def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array]):
+def depends[T](inputs: T, dependencies: array | Sequence[array]) -> T:
     """
     Insert dependencies between arrays in the graph. The outputs are
     identical to ``inputs`` but with dependencies on ``dependencies``.

.mlx_typings/mlx/nn/__init__.pyi

@@ -1,9 +1,5 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from layers import *
-from utils import *
+from .layers import *
+from .utils import *
 
 from . import init as init
 from . import losses as losses

.mlx_typings/mlx/nn/layers/__init__.pyi

@@ -1,20 +1,16 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from activations import *
-from base import *
-from containers import *
-from convolution import *
-from convolution_transpose import *
-from distributed import *
-from dropout import *
-from embedding import *
-from linear import *
-from normalization import *
-from pooling import *
-from positional_encoding import *
-from quantized import *
-from recurrent import *
-from transformer import *
-from upsample import *
+from .activations import *
+from .base import *
+from .containers import *
+from .convolution import *
+from .convolution_transpose import *
+from .distributed import *
+from .dropout import *
+from .embedding import *
+from .linear import *
+from .normalization import *
+from .pooling import *
+from .positional_encoding import *
+from .quantized import *
+from .recurrent import *
+from .transformer import *
+from .upsample import *

.mlx_typings/mlx/nn/layers/base.pyi

@@ -53,7 +53,7 @@ class Module(dict):
         mx.eval(model.parameters())
     """
 
-    __call__: Callable
+    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
     def __init__(self) -> None:
         """Should be called by the subclasses of ``Module``."""
 

.mlx_typings/mlx_lm/generate.pyi

@@ -30,7 +30,7 @@ def str2bool(string):  # -> bool:
 def setup_arg_parser():  # -> ArgumentParser:
     """Set up and return the argument parser."""
 
-generation_stream = ...
+generation_stream: mx.Stream
 
 @contextlib.contextmanager
 def wired_limit(
@@ -266,12 +266,12 @@ def _merge_caches(caches: Any) -> List[Any]: ...
 class Batch:
     uids: List[int]
     y: mx.array
-    logprobs: mx.array
+    logprobs: List[mx.array] | mx.array
     max_tokens: List[int]
     num_tokens: List[int]
     cache: List[Any]
-    samplers: List[Any]
-    logits_processors: List[Any]
+    samplers: List[Callable[[mx.array], mx.array] | None]
+    logits_processors: List[List[Callable[[mx.array, mx.array], mx.array]]]
     tokens: List[mx.array]
     def __len__(self) -> int: ...
     def filter(self, keep_idx: List[int]) -> None: ...
@@ -279,13 +279,18 @@ class Batch:
     def extract_cache(self, idx: int) -> List[Any]: ...
 
 class BatchGenerator:
-    model: Any
+    model: nn.Module
+    sampler: Callable[[mx.array], mx.array]
+    stop_tokens: set[int]
     max_kv_size: Optional[int]
     prefill_step_size: int
+    completion_batch_size: int
+    prefill_batch_size: int
     unprocessed_prompts: List[Any]
     active_batch: Optional[Batch]
     prompt_progress_callback: Callable[[List[Tuple[int, int, int]]], None]
     _stats: BatchStats
+    _next_count: int
 
     @dataclass
     class Response:

.mlx_typings/mlx_lm/models/cache.pyi

@@ -88,8 +88,8 @@ def create_attention_mask(
 ) -> array | Literal["causal"] | None: ...
 
 class _BaseCache(Cache):
-    keys: mx.array
-    values: mx.array
+    keys: mx.array | None
+    values: mx.array | None
     offset: int
     @property
     def state(self) -> tuple[mx.array | None, mx.array | None]: ...
@@ -268,29 +268,14 @@ class CacheList(_BaseCache):
         """
 
 class BatchKVCache(_BaseCache):
-    step = ...
-    def __init__(self, left_padding: List[int]) -> None:
-        """
-        The BatchKV cache expects inputs to be left-padded.
-
-        E.g. the following prompts:
-
-            [1, 3, 5]
-            [7]
-            [2, 6, 8, 9]
-
-        Should be padded like so:
-
-            [0, 1, 3, 5]
-            [0, 0, 0, 7]
-            [2, 6, 8, 9]
-
-        And ``left_padding`` specifies the amount of padding for each.
-        In this case, ``left_padding = [1, 3, 0]``.
-        """
-
-    def update_and_fetch(self, keys, values):  # -> tuple[array | Any, array | Any]:
-        ...
+    step: int
+    keys: array | None
+    values: array | None
+    offset: array
+    left_padding: array
+    _idx: int
+    def __init__(self, left_padding: List[int]) -> None: ...
+    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
     @property
     def state(
         self,
@@ -316,12 +301,21 @@ class BatchKVCache(_BaseCache):
         """
 
 class BatchRotatingKVCache(_BaseCache):
-    step = ...
-    def __init__(self, max_size, left_padding: List[int]) -> None: ...
-    def update_and_fetch(
-        self, keys, values
-    ):  # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
-        ...
+    step: int
+    keys: array | None
+    values: array | None
+    offset: array
+    left_padding: array
+    max_size: int
+    _idx: int
+    _offset: int
+    rotated: bool
+    _lengths: array | None
+    def __init__(self, max_size: int, left_padding: List[int]) -> None: ...
+    def _trim(self, trim_size: int, v: array, append: array | None = ...) -> array: ...
+    def _update_in_place(self, keys: array, values: array) -> tuple[array, array]: ...
+    def _update_concat(self, keys: array, values: array) -> tuple[array, array]: ...
+    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
     @property
     def state(
         self,

.mlx_typings/mlx_lm/models/gated_delta.pyi

@@ -0,0 +1,35 @@
+from typing import Optional
+
+import mlx.core as mx
+
+def compute_g(A_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array: ...
+def gated_delta_update(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    a: mx.array,
+    b: mx.array,
+    A_log: mx.array,
+    dt_bias: mx.array,
+    state: Optional[mx.array] = ...,
+    mask: Optional[mx.array] = ...,
+    use_kernel: bool = ...,
+) -> tuple[mx.array, mx.array]: ...
+def gated_delta_ops(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: Optional[mx.array] = ...,
+    mask: Optional[mx.array] = ...,
+) -> tuple[mx.array, mx.array]: ...
+def gated_delta_kernel(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: mx.array,
+    mask: Optional[mx.array] = ...,
+) -> tuple[mx.array, mx.array]: ...

.mlx_typings/mlx_lm/models/rope_utils.pyi

@@ -0,0 +1,51 @@
+from typing import Any, Optional
+
+import mlx.nn as nn
+
+class YarnRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        scaling_factor: float = ...,
+        original_max_position_embeddings: int = ...,
+        beta_fast: float = ...,
+        beta_slow: float = ...,
+        mscale: float = ...,
+        mscale_all_dim: float = ...,
+    ) -> None: ...
+
+class Llama3RoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        scaling_factor: float = ...,
+        original_max_position_embeddings: int = ...,
+        low_freq_factor: float = ...,
+        high_freq_factor: float = ...,
+    ) -> None: ...
+
+class SuScaledRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        short_factor: Any = ...,
+        long_factor: Any = ...,
+        original_max_position_embeddings: int = ...,
+    ) -> None: ...
+
+def initialize_rope(
+    dims: int,
+    base: float = ...,
+    traditional: bool = ...,
+    scaling_config: Optional[dict[str, Any]] = ...,
+    max_position_embeddings: Optional[int] = ...,
+) -> nn.Module: ...

bench/exo_bench.py

@@ -501,23 +501,21 @@ def main() -> int:
                                     for x, _ in batch_results
                                     if x["stats"]["generation_tps"] > 0
                                 ]
-                                agg_gen_tps = (
+                                per_req_tps = (
                                     mean(valid_gen_tps) if valid_gen_tps else 0.0
                                 )
-                                gen_tps = agg_gen_tps / concurrency
+                                agg_gen_tps = per_req_tps * concurrency
                                 logger.info(
                                     f"[concurrent {concurrency}x]  "
                                     f"agg_gen_tps={agg_gen_tps:.2f}  "
-                                    f"gen_tps={gen_tps:.2f}  "
+                                    f"per_req_tps={per_req_tps:.2f}  "
                                     f"errors={batch_errors}"
                                 )
 
                     if runs:
                         prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
-                        gen_tps = mean(
-                            x["stats"]["generation_tps"] / x["concurrency"]
-                            for x in runs
-                        )
+                        per_req_tps = mean(x["stats"]["generation_tps"] for x in runs)
+                        gen_tps = per_req_tps * concurrency
                         ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
                         gtok = mean(x["stats"]["generation_tokens"] for x in runs)
                         peak = mean(

bench/run_lcb_tps_bench.sh

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

dashboard/src/lib/stores/app.svelte.ts

@@ -1793,6 +1793,14 @@ class AppStore {
             this.persistConversation(targetConversationId);
           }
         },
+        {
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+            }
+          },
+        },
       );
 
       // Final update
@@ -1990,6 +1998,14 @@ class AppStore {
             this.persistConversation(targetConversationId);
           }
         },
+        {
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+            }
+          },
+        },
       );
 
       // Final cleanup of the message (if conversation still exists)
@@ -2397,7 +2413,7 @@ class AppStore {
 
       let streamedContent = "";
       let streamedThinking = "";
-
+      let serverTpsReceived = false;
       interface ChatCompletionChunk {
         choices?: Array<{
           delta?: { content?: string; reasoning_content?: string };
@@ -2462,7 +2478,6 @@ class AppStore {
             tokenCount += 1;
             this.totalTokens = tokenCount;
 
-            // Update real-time TPS during streaming
             if (firstTokenTime !== null && tokenCount > 1) {
               const elapsed = performance.now() - firstTokenTime;
               this.tps = (tokenCount / elapsed) * 1000;
@@ -2513,16 +2528,24 @@ class AppStore {
               startedAt: this.prefillProgress?.startedAt ?? performance.now(),
             };
           },
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+              serverTpsReceived = true;
+            }
+          },
         },
       );
 
       // Clear prefill progress after stream ends
       this.prefillProgress = null;
 
-      // Calculate final TPS
-      if (firstTokenTime !== null && tokenCount > 1) {
+      // Use server-side TPS if available, otherwise fall back to client-side
+      if (!serverTpsReceived && firstTokenTime !== null && tokenCount > 1) {
         const totalGenerationTime = performance.now() - firstTokenTime;
-        this.tps = (tokenCount / totalGenerationTime) * 1000; // tokens per second
+        this.tps = (tokenCount / totalGenerationTime) * 1000;
       }
 
       // Final cleanup of the message (if conversation still exists)

pyproject.toml

@@ -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/eval-left-padding-in-batched-rotation" }
+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 }

src/exo/master/adapters/chat_completions.py

@@ -202,6 +202,8 @@ async def generate_chat_stream(
                     usage=last_usage,
                 )
                 yield f"data: {tool_response.model_dump_json()}\n\n"
+                if chunk.stats is not None:
+                    yield f": generation_stats {chunk.stats.model_dump_json()}\n\n"
                 yield "data: [DONE]\n\n"
                 return
 
@@ -216,6 +218,8 @@ async def generate_chat_stream(
                 yield f"data: {chunk_response.model_dump_json()}\n\n"
 
                 if chunk.finish_reason is not None:
+                    if chunk.stats is not None:
+                        yield f": generation_stats {chunk.stats.model_dump_json()}\n\n"
                     yield "data: [DONE]\n\n"
 
 

src/exo/worker/engines/image/models/qwen/wrappers.py

@@ -56,10 +56,10 @@ class QwenJointBlockWrapper(JointBlockWrapper[QwenTransformerBlock]):
         attn = self.block.attn
 
         img_mod_params = self.block.img_mod_linear(
-            self.block.img_mod_silu(text_embeddings)  # pyright: ignore[reportUnknownArgumentType]
+            self.block.img_mod_silu(text_embeddings)
         )
         txt_mod_params = self.block.txt_mod_linear(
-            self.block.txt_mod_silu(text_embeddings)  # pyright: ignore[reportUnknownArgumentType]
+            self.block.txt_mod_silu(text_embeddings)
         )
 
         img_mod1, img_mod2 = mx.split(img_mod_params, 2, axis=-1)

src/exo/worker/engines/mlx/auto_parallel.py

@@ -480,7 +480,7 @@ def patch_tensor_model[T](model: T) -> T:
             last = cache[-1]  # pyright: ignore[reportAny]
             dep_cache = last[0] if hasattr(last, "caches") else last  # pyright: ignore[reportAny]
             if hasattr(dep_cache, "keys"):  # type: ignore
-                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # pyright: ignore[reportAny,reportUnknownMemberType]
+                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # pyright: ignore[reportAny]
 
         return logits
 

src/exo/worker/engines/mlx/generator/batch_generate.py

@@ -1,3 +1,4 @@
+import os
 import time
 from dataclasses import dataclass, field
 from typing import Callable, cast
@@ -79,12 +80,173 @@ class ExoBatchGenerator:
     _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False)
 
     def __post_init__(self) -> None:
-        self._exo_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"],
         )
 
+        # 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 (
@@ -131,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,
@@ -151,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 (
@@ -200,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,
@@ -276,7 +471,7 @@ class ExoBatchGenerator:
 
             logprob: float | None = None
             top_logprobs: list[TopLogprobItem] | None = None
-            if task_params.logprobs:
+            if task_params.logprobs and os.environ.get("EXO_DISABLE_LOGPROBS") != "1":
                 logprob, top_logprobs = extract_top_logprobs(
                     logprobs=response.logprobs,
                     tokenizer=self.tokenizer,

src/exo/worker/engines/mlx/generator/generate.py

@@ -179,7 +179,8 @@ def pipeline_parallel_prefill(
         flush_prefill_sends()
 
     assert _prompt_cache is not None
-    mx.eval([c.state for c in _prompt_cache])  # type: ignore
+    with mx.stream(generation_stream):
+        mx.eval([c.state for c in _prompt_cache])  # type: ignore
 
     # Final callback matching generate_step
     prompt_progress_callback(total, total)
@@ -398,52 +399,44 @@ def extract_top_logprobs(
     tokenizer: TokenizerWrapper,
     top_logprobs: int,
     selected_token: int,
+    precomputed_indices: list[int] | None = None,
+    precomputed_values: list[float] | None = None,
+    precomputed_selected: float | None = None,
 ) -> tuple[float, list[TopLogprobItem]]:
-    """Extract the selected token's logprob and top alternative tokens.
-
-    Args:
-        logprobs: Full vocabulary logprobs array from MLX
-        tokenizer: Tokenizer for decoding token IDs to strings
-        top_logprobs: Number of top alternatives to return
-        selected_token: The token ID that was actually sampled
-
-    Returns:
-        Tuple of (selected_token_logprob, list of TopLogprobItem for top alternatives)
-    """
-    # Get the logprob of the selected token
-    selected_logprob = float(logprobs[selected_token].item())
-
-    # Get top indices (most probable tokens)
-    # mx.argpartition gives indices that would partition the array
-    # We negate logprobs since argpartition finds smallest, and we want largest
-    top_logprobs = min(top_logprobs, logprobs.shape[0])  # Don't exceed vocab size
-    top_indices = mx.argpartition(-logprobs, top_logprobs)[:top_logprobs]
-
-    # Get the actual logprob values for these indices
-    top_values = logprobs[top_indices]
-
-    # Sort by logprob (descending) for consistent ordering
-    sort_order = mx.argsort(-top_values)
-    top_indices = top_indices[sort_order]
-    top_values = top_values[sort_order]
+    if (
+        precomputed_indices is not None
+        and precomputed_values is not None
+        and precomputed_selected is not None
+    ):
+        top_indices_list: list[int] = precomputed_indices[:top_logprobs]
+        top_values_list: list[float] = precomputed_values[:top_logprobs]
+        selected_logprob = precomputed_selected
+    else:
+        selected_logprob_arr = logprobs[selected_token]
+        top_logprobs = min(top_logprobs, logprobs.shape[0] - 1)
+        top_indices = mx.argpartition(-logprobs, top_logprobs)[:top_logprobs]
+        top_values = logprobs[top_indices]
+        sort_order = mx.argsort(-top_values)
+        top_indices = top_indices[sort_order]
+        top_values = top_values[sort_order]
+        mx.eval(selected_logprob_arr, top_indices, top_values)
+        selected_logprob = float(selected_logprob_arr.item())
+        top_indices_list = top_indices.tolist()  # type: ignore
+        top_values_list = top_values.tolist()  # type: ignore
 
     # Convert to list of TopLogprobItem
     top_logprob_items: list[TopLogprobItem] = []
-    for i in range(top_logprobs):
-        token_id = int(top_indices[i].item())
-        token_logprob = float(top_values[i].item())
+    for token_id, token_logprob in zip(top_indices_list, top_values_list, strict=True):
         if math.isnan(token_logprob):
             continue
 
         # Decode token ID to string
         token_str = tokenizer.decode([token_id])
-        # Get byte representation
-        token_bytes = list(token_str.encode("utf-8"))
         top_logprob_items.append(
             TopLogprobItem(
                 token=token_str,
                 logprob=token_logprob,
-                bytes=token_bytes,
+                bytes=list(token_str.encode("utf-8")),
             )
         )
 
@@ -624,12 +617,13 @@ def mlx_generate(
         logprob: float | None = None
         top_logprobs: list[TopLogprobItem] | None = None
         if task.logprobs:
-            logprob, top_logprobs = extract_top_logprobs(
-                logprobs=out.logprobs,
-                tokenizer=tokenizer,
-                top_logprobs=task.top_logprobs or DEFAULT_TOP_LOGPROBS,
-                selected_token=out.token,
-            )
+            with mx.stream(generation_stream):
+                logprob, top_logprobs = extract_top_logprobs(
+                    logprobs=out.logprobs,
+                    tokenizer=tokenizer,
+                    top_logprobs=task.top_logprobs or DEFAULT_TOP_LOGPROBS,
+                    selected_token=out.token,
+                )
 
         if is_done:
             # Log generation stats

src/exo/worker/engines/mlx/patches/__init__.py

@@ -0,0 +1,45 @@
+"""Model-specific kernel fusion patches for MLX inference.
+
+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 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
+
+    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)

src/exo/worker/engines/mlx/patches/high_precision_gdn_softplus.py

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

src/exo/worker/engines/mlx/patches/opt_batch_gen.py

@@ -0,0 +1,173 @@
+import time
+from typing import Any, cast
+
+import mlx.core as mx
+from mlx_lm.generate import BatchGenerator, generation_stream
+
+_PRECOMPUTE_TOP_K = 20
+
+_original_public_next = BatchGenerator.next
+
+_pending_topk_idx: mx.array | None = None
+_pending_topk_val: mx.array | None = None
+_pending_selected_lps: mx.array | None = None
+
+
+def _fast_next(self: BatchGenerator) -> list[BatchGenerator.Response]:
+    tic = time.perf_counter()
+    batch = self.active_batch
+    assert batch is not None
+    batch_size = len(batch)
+
+    prev_tokens = batch.y
+    prev_logprobs = batch.logprobs
+
+    has_processors = any(p for ps in batch.logits_processors for p in ps)
+    if has_processors:
+        for i, toks in enumerate(batch.tokens):
+            batch.tokens[i] = mx.concatenate([toks, prev_tokens[i : i + 1]])
+
+    logits = self.model(prev_tokens[:, None], cache=batch.cache)
+    logits = logits[:, -1, :]
+
+    if has_processors:
+        processed_logits: list[mx.array] = []
+        for e in range(batch_size):
+            sample_logits: mx.array = logits[e : e + 1]
+            for processor in batch.logits_processors[e]:
+                sample_logits = processor(batch.tokens[e], sample_logits)
+            processed_logits.append(sample_logits)
+        logits = mx.concatenate(processed_logits, axis=0)
+
+    logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+
+    if (
+        batch_size == 1
+        or any(batch.samplers)
+        and all(s is batch.samplers[0] for s in batch.samplers)
+    ):
+        sampler = batch.samplers[0] or self.sampler
+        batch.y = sampler(logprobs)
+    elif any(batch.samplers):
+        all_samples: list[mx.array] = []
+        for e in range(batch_size):
+            s = batch.samplers[e] or self.sampler
+            all_samples.append(s(logprobs[e : e + 1]))
+        batch.y = mx.concatenate(all_samples, axis=0)
+    else:
+        batch.y = self.sampler(logprobs)
+    batch.logprobs = list(logprobs)
+
+    global _pending_topk_idx, _pending_topk_val, _pending_selected_lps
+
+    emit_topk_indices: list[list[int]] = (
+        cast(list[list[int]], _pending_topk_idx.tolist())
+        if _pending_topk_idx is not None
+        else []
+    )
+    emit_topk_values: list[list[float]] = (
+        cast(list[list[float]], _pending_topk_val.tolist())
+        if _pending_topk_val is not None
+        else []
+    )
+    emit_selected_lps: list[float] = (
+        cast(list[float], _pending_selected_lps.tolist())
+        if _pending_selected_lps is not None
+        else []
+    )
+
+    needs_topk: bool = getattr(self, "_needs_topk", False)
+    if needs_topk:
+        k = min(_PRECOMPUTE_TOP_K, logprobs.shape[1])
+        _pending_topk_idx = mx.argpartition(-logprobs, k, axis=1)[:, :k]
+        _pending_topk_val = mx.take_along_axis(logprobs, _pending_topk_idx, axis=1)
+        sort_order = mx.argsort(-_pending_topk_val, axis=1)
+        _pending_topk_idx = mx.take_along_axis(_pending_topk_idx, sort_order, axis=1)
+        _pending_topk_val = mx.take_along_axis(_pending_topk_val, sort_order, axis=1)
+        _pending_selected_lps = logprobs[mx.arange(batch_size), batch.y]
+        mx.async_eval(
+            batch.y,
+            *batch.logprobs,
+            *batch.tokens,
+            _pending_topk_idx,
+            _pending_topk_val,
+            _pending_selected_lps,
+        )
+    else:
+        _pending_topk_idx = None
+        _pending_topk_val = None
+        _pending_selected_lps = None
+        mx.async_eval(batch.y, *batch.logprobs, *batch.tokens)
+
+    prev_token_list: list[int] = cast(list[int], prev_tokens.tolist())
+
+    toc = time.perf_counter()
+    self._stats.generation_time += toc - tic
+
+    keep_idx: list[int] = []
+    end_idx: list[int] = []
+    responses: list[Any] = []
+    stop_tokens = self.stop_tokens
+
+    for e in range(batch_size):
+        t = prev_token_list[e]
+        uid = batch.uids[e]
+        num_tok = batch.num_tokens[e] + 1
+        batch.num_tokens[e] = num_tok
+
+        if t in stop_tokens:
+            finish_reason = "stop"
+            end_idx.append(e)
+        elif num_tok >= batch.max_tokens[e]:
+            finish_reason = "length"
+            end_idx.append(e)
+        else:
+            finish_reason = None
+            keep_idx.append(e)
+
+        cache = None
+        if finish_reason is not None:
+            cache = batch.extract_cache(e)
+        response = self.Response(uid, t, prev_logprobs[e], finish_reason, cache)
+        if emit_topk_indices and e < len(emit_topk_indices):
+            response._topk_indices = emit_topk_indices[e]  # pyright: ignore[reportAttributeAccessIssue]
+            response._topk_values = emit_topk_values[e]  # pyright: ignore[reportAttributeAccessIssue]
+            response._selected_logprob = emit_selected_lps[e]  # pyright: ignore[reportAttributeAccessIssue]
+        responses.append(response)
+
+    if end_idx:
+        if keep_idx:
+            batch.filter(keep_idx)
+            if (
+                _pending_topk_idx is not None
+                and _pending_topk_val is not None
+                and _pending_selected_lps is not None
+            ):
+                ki = mx.array(keep_idx)
+                _pending_topk_idx = _pending_topk_idx[ki]
+                _pending_topk_val = _pending_topk_val[ki]
+                _pending_selected_lps = _pending_selected_lps[ki]
+        else:
+            self.active_batch = None
+            _pending_topk_idx = None
+            _pending_topk_val = None
+            _pending_selected_lps = None
+
+    self._next_count += 1
+    if self._next_count % 512 == 0:
+        mx.clear_cache()
+    self._stats.generation_tokens += len(responses)
+    return responses
+
+
+def _patched_public_next(self: BatchGenerator) -> list[BatchGenerator.Response]:
+    batch = self.active_batch
+    # Only do decode with fast_next
+    if batch is not None and not self.unprocessed_prompts:
+        with mx.stream(generation_stream):
+            return _fast_next(self)
+    return _original_public_next(self)
+
+
+def apply_batch_gen_patch() -> None:
+    BatchGenerator.next = _patched_public_next

src/exo/worker/engines/mlx/patches/qwen3_5/custom_qmv_loop_over_b.py

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

src/exo/worker/engines/mlx/patches/qwen3_5/lpb_patch.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/apply.py

@@ -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"
+    )

src/exo/worker/engines/mlx/patches/qwen3_5_moe/batched_fused_gqa_attention.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/batched_moe.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/common.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/decoder.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/fused_gdn_attention.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_fused_gdn_projections_8bit.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_fused_gqa_projections_8bit.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_merged_down_proj_8bit.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_moe_epilogue.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_oproj_gate_gemv_8bit.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_softmax_topk_swiglu_8bit.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_gdn_projections_8bit.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_qk_rmsnorm.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_rms_norm_gated.py

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/gdn_step_precomputed.py

@@ -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],
+    )

src/exo/worker/engines/mlx/patches/standard_yarn_rope.py

@@ -0,0 +1,118 @@
+import math
+
+import mlx.core as mx
+from mlx_lm.models import rope_utils
+
+_original_YarnRoPE_init = rope_utils.YarnRoPE.__init__  # noqa: N816
+_original_initialize_rope = rope_utils.initialize_rope
+
+
+def _patched_yarn_init(
+    self: rope_utils.YarnRoPE,
+    dims: int,
+    traditional: bool = False,
+    max_position_embeddings: int = 2048,
+    base: float = 10000,
+    scaling_factor: float = 1.0,
+    original_max_position_embeddings: int = 4096,
+    beta_fast: float = 32,
+    beta_slow: float = 1,
+    mscale: float = 1,
+    mscale_all_dim: float = 0,
+    truncate: bool = True,
+) -> None:
+    """Patch mlx_lm's YarnRoPE to match vLLM's inverse-frequency blending formula for compatability."""
+
+    super(rope_utils.YarnRoPE, self).__init__()
+
+    def yarn_find_correction_dim(num_rotations: float) -> float:
+        return (
+            dims
+            * math.log(original_max_position_embeddings / (num_rotations * 2 * math.pi))
+        ) / (2 * math.log(base))
+
+    def yarn_find_correction_range() -> tuple[float, float]:
+        low: float = yarn_find_correction_dim(beta_fast)
+        high: float = yarn_find_correction_dim(beta_slow)
+        if truncate:
+            low = math.floor(low)
+            high = math.ceil(high)
+        return max(low, 0), min(high, dims - 1)
+
+    def yarn_get_mscale(scale: float = 1, ms: float = 1) -> float:
+        if scale <= 1:
+            return 1.0
+        return 0.1 * ms * math.log(scale) + 1.0
+
+    def yarn_linear_ramp_mask(min_val: float, max_val: float, dim: int) -> mx.array:
+        if min_val == max_val:
+            max_val += 0.001
+        linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (max_val - min_val)
+        return mx.clip(linear_func, 0, 1)
+
+    self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
+        scaling_factor, mscale_all_dim
+    )
+    pos_freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
+    inv_freq_extrapolation = 1.0 / pos_freqs
+    inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
+    low, high = yarn_find_correction_range()
+    inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dims // 2)
+    inv_freq = (
+        inv_freq_interpolation * (1 - inv_freq_mask)
+        + inv_freq_extrapolation * inv_freq_mask
+    )
+    self._freqs = 1.0 / inv_freq
+    self.dims = dims
+    self.traditional = traditional
+
+
+def _patched_initialize_rope(
+    dims: int,
+    base: float,
+    traditional: bool,
+    scaling_config: dict[str, str | int | float | bool] | None = None,
+    max_position_embeddings: int | None = None,
+) -> object:
+    rope_type = "default"
+    if scaling_config is not None:
+        rope_type = str(
+            scaling_config.get("type") or scaling_config.get("rope_type", "default")
+        )
+
+    # All the yarn rope types supported in mlx lm
+    if rope_type in ("yarn", "deepseek_yarn"):
+        assert scaling_config is not None
+        cfg = scaling_config
+
+        def _float(key: str, default: float) -> float:
+            v = cfg.get(key)
+            return float(v) if v is not None else default
+
+        def _int(key: str, default: int) -> int:
+            v = cfg.get(key)
+            return int(v) if v is not None else default
+
+        return rope_utils.YarnRoPE(
+            dims=dims,
+            max_position_embeddings=max_position_embeddings or 2048,
+            traditional=traditional,
+            scaling_factor=_float("factor", 1.0),
+            base=base,
+            original_max_position_embeddings=_int(
+                "original_max_position_embeddings", 4096
+            ),
+            beta_fast=_float("beta_fast", 32),
+            beta_slow=_float("beta_slow", 1),
+            mscale=_float("mscale", 1),
+            mscale_all_dim=_float("mscale_all_dim", 0),
+        )
+
+    return _original_initialize_rope(
+        dims, base, traditional, scaling_config, max_position_embeddings
+    )
+
+
+def patch_yarn_rope() -> None:
+    rope_utils.YarnRoPE.__init__ = _patched_yarn_init
+    rope_utils.initialize_rope = _patched_initialize_rope

src/exo/worker/engines/mlx/speculative/mtp_batch_generator.py

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

src/exo/worker/engines/mlx/speculative/mtp_module.py

@@ -0,0 +1,515 @@
+#!/usr/bin/env python3
+"""MTP (Multi-Token Prediction) module for Qwen3.5-27B.
+
+Architecture (from llama.cpp build_mtp_head + HuggingFace config):
+  1. Normalize: pre_fc_norm_hidden(hidden_state) || pre_fc_norm_embedding(embed(token))
+  2. Combine: fc(concat([e_norm, h_norm])) → 5120
+  3. 1 GQA decoder layer (same config as main model's full-attention layers)
+     - Attention with Q/K RMSNorm + partial RoPE + output gate
+  4. Final norm → shared lm_head → vocab logits
+
+Predicts token t+2 given the main model's hidden state at position t
+and the token sampled at position t+1.
+
+Usage:
+    from .mtp_module import MTPPredictor
+    mtp = MTPPredictor(model, "mtp_weights.safetensors")
+    # During decode:
+    pre_norm, normed = mtp.get_hidden_state(input_tokens, cache)
+    logits_t1 = mtp.apply_lm_head(normed)            # token t+1
+    logits_t2 = mtp.predict(pre_norm, token_t1)       # token t+2
+"""
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+def speculative_forward(model, inputs, cache, speculative=False):
+    """Run model forward pass, optionally capturing GDN per-step states for rollback.
+
+    This is the shared core for both MTP and draft-model speculative decoding.
+    It manually iterates model layers to:
+    1. Wrap GDN caches in SpeculativeArraysCache when speculative=True
+    2. Patch gated_delta_update to use the speculative kernel
+    3. Capture per-step recurrent states and reconstruct conv_input
+
+    Args:
+        model: the loaded model (e.g. from mlx_lm.load)
+        inputs: (B, S) int token ids
+        cache: cache list from make_prompt_cache()
+        speculative: if True, saves per-step GDN states for rollback
+
+    Returns:
+        (pre_norm, logits) — pre-RMSNorm hidden states and vocab logits
+    """
+    inner = getattr(model, 'model', None) or model.language_model.model
+    text_model = getattr(model, 'model', None) or model.language_model
+    S = inputs.shape[1]
+    do_spec = speculative and S > 1
+
+    if hasattr(inner, 'embed_tokens'):
+        hidden_states = inner.embed_tokens(inputs)
+    else:
+        hidden_states = inputs
+
+    cache_list = cache if cache is not None else [None] * len(inner.layers)
+
+    gdn_spec_data = []
+    if do_spec:
+        from .speculative_cache import SpeculativeArraysCache
+        for i, c in enumerate(cache_list):
+            if c is not None and hasattr(c, 'cache') and not hasattr(c, 'offset'):
+                cache_list[i] = SpeculativeArraysCache(c, S=S)
+        if cache is not None:
+            for i in range(len(cache)):
+                cache[i] = cache_list[i]
+
+    spec_all_states = []
+    if do_spec:
+        import mlx_lm.models.qwen3_5 as _qwen3_5_mod
+        _orig_gdu = _qwen3_5_mod.gated_delta_update
+        _qwen3_5_mod.gated_delta_update = _make_speculative_gdu(spec_all_states)
+
+    from mlx_lm.models.qwen3_5 import create_attention_mask, create_ssm_mask
+    fa_mask = create_attention_mask(hidden_states, cache_list[inner.fa_idx])
+    ssm_mask = create_ssm_mask(hidden_states, cache_list[inner.ssm_idx])
+
+    for layer, c in zip(inner.layers, cache_list):
+        mask = ssm_mask if layer.is_linear else fa_mask
+
+        if do_spec and layer.is_linear:
+            from .speculative_cache import SpeculativeArraysCache as _SAC
+            if isinstance(c, _SAC):
+                pre_conv = c[0]
+                if pre_conv is None:
+                    gdn = layer.linear_attn
+                    pre_conv = mx.zeros(
+                        (hidden_states.shape[0], gdn.conv_kernel_size - 1,
+                         gdn.conv_dim), dtype=hidden_states.dtype)
+                gdn_spec_data.append((hidden_states, pre_conv, c, layer))
+
+        hidden_states = layer(hidden_states, mask=mask, cache=c)
+
+    if do_spec:
+        _qwen3_5_mod.gated_delta_update = _orig_gdu
+
+        gdn_idx = 0
+        for layer_input, pre_conv, spec_cache, parent_layer in gdn_spec_data:
+            if gdn_idx < len(spec_all_states):
+                spec_cache.all_states = spec_all_states[gdn_idx]
+            gdn_idx += 1
+
+            gdn = parent_layer.linear_attn
+            normed = parent_layer.input_layernorm(layer_input)
+            if hasattr(gdn, 'in_proj_qkv'):
+                qkv = gdn.in_proj_qkv(normed)
+            else:
+                q, k, v, z, b, a = gdn.fix_query_key_value_ordering(
+                    gdn.in_proj_qkvz(normed), gdn.in_proj_ba(normed))
+                B_dim = normed.shape[0]
+                qkv = mx.concatenate(
+                    [q.reshape(B_dim, S, -1), k.reshape(B_dim, S, -1),
+                     v.reshape(B_dim, S, -1)], axis=-1)
+            spec_cache.conv_input = mx.concatenate([pre_conv, qkv], axis=1)
+
+    pre_norm = hidden_states
+    normed = inner.norm(hidden_states)
+
+    if hasattr(text_model, 'lm_head'):
+        logits = text_model.lm_head(normed)
+    else:
+        logits = inner.embed_tokens.as_linear(normed)
+
+    return pre_norm, logits
+
+
+def _make_speculative_gdu(all_states_list):
+    """Create a gated_delta_update replacement that uses the speculative kernel.
+
+    The speculative kernel is identical to the original but also outputs
+    per-step recurrent states (all_states). These are appended to
+    all_states_list for later assignment to SpeculativeArraysCache wrappers.
+
+    Returns (y, state_out) — same interface as original gated_delta_update.
+    """
+    from .speculative_gdn_kernel import speculative_gated_delta_kernel
+    from mlx_lm.models.gated_delta import compute_g
+
+    def speculative_gated_delta_update(q, k, v, a, b, A_log, dt_bias,
+                                        state=None, mask=None, use_kernel=True):
+        beta = mx.sigmoid(b)
+        g = compute_g(A_log, a, dt_bias)
+        if state is None:
+            B, _, Hk, Dk = q.shape
+            Hv, Dv = v.shape[-2:]
+            state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
+        y, state_out, all_states = speculative_gated_delta_kernel(
+            q, k, v, g, beta, state, mask)
+        all_states_list.append(all_states)
+        return y, state_out
+
+    return speculative_gated_delta_update
+
+
+class MTPPredictor:
+    """MTP draft predictor for speculative decoding.
+
+    Wraps the MTP module weights and provides:
+      - get_hidden_state(): extract pre-lm_head hidden state from main model
+      - predict(): run MTP to get next-next-token logits
+    """
+
+    def __init__(self, model, mtp_weights_path, quantize=True):
+        """Load MTP weights and attach to the main model.
+
+        Args:
+            model: loaded Qwen3.5-27B model
+            mtp_weights_path: path to mtp_weights.safetensors
+            quantize: quantize MTP linears to 8-bit gs=64
+        """
+        self.model = model
+        self._inner = getattr(model, 'model', None) or model.language_model.model
+        self._text_model = getattr(model, 'model', None) or model.language_model
+
+        # Shared components
+        self.embed_tokens = self._inner.embed_tokens
+        if hasattr(self._text_model, 'lm_head'):
+            self.lm_head = self._text_model.lm_head
+        else:
+            # tie_word_embeddings case
+            self.lm_head = None
+
+        # Load MTP weights
+        weights = mx.load(mtp_weights_path)
+
+        # ---- Sanitize norm weights ----
+        # CRITICAL: Qwen3.5 HuggingFace format stores ALL norm weights as (actual - 1.0).
+        # mlx-lm's TextModel.sanitize() adds +1.0 back for the main model norms, but
+        # MTP weights are stripped before sanitize runs. We must apply the same shift
+        # to ALL 1-D norm weights in the MTP.
+        #
+        # Evidence: pre_fc_norm_hidden has mean=-0.17 raw → 0.83 after shift (plausible).
+        # Linear projection weights (2-D) are NOT shifted.
+        shifted = []
+        for k in list(weights.keys()):
+            if weights[k].ndim == 1:
+                weights[k] = weights[k] + 1.0
+                shifted.append(k)
+        if shifted:
+            print(f"  Sanitized {len(shifted)} norm weights (+1.0 shift)")
+
+        # Infer all dimensions from weight shapes (works for any Qwen3.5 size)
+        fc_w = weights['mtp.fc.weight']
+        hidden_size = fc_w.shape[0]                    # 4096 (9B) or 5120 (27B)
+        fc_in = fc_w.shape[1]                          # 2 * hidden_size
+
+        q_w = weights['mtp.layers.0.self_attn.q_proj.weight']
+        q_out = q_w.shape[0]                           # num_heads * head_dim * 2 (gate)
+        k_w = weights['mtp.layers.0.self_attn.k_proj.weight']
+        kv_out = k_w.shape[0]                          # num_kv_heads * head_dim
+        o_w = weights['mtp.layers.0.self_attn.o_proj.weight']
+        o_in = o_w.shape[1]                            # num_heads * head_dim
+
+        # Detect MoE vs dense MLP
+        self.is_moe = 'mtp.layers.0.mlp.gate.weight' in weights
+
+        if not self.is_moe:
+            gate_w = weights['mtp.layers.0.mlp.gate_proj.weight']
+            intermediate = gate_w.shape[0]
+        else:
+            intermediate = 0  # MoE experts handle this
+
+        # head_dim from q_norm weight (always per-head)
+        head_dim = weights.get('mtp.layers.0.self_attn.q_norm.weight',
+                               mx.ones(256)).shape[0]
+        num_heads = o_in // head_dim
+        num_kv_heads = kv_out // head_dim
+
+        print(f"  Dims: hidden={hidden_size}, heads={num_heads}, kv_heads={num_kv_heads}, "
+              f"head_dim={head_dim}, MLP={'MoE' if self.is_moe else f'dense({intermediate})'}")
+
+        # Build layers from weights — all dimension-agnostic
+        def make_linear(w):
+            out_dim, in_dim = w.shape
+            l = nn.Linear(in_dim, out_dim, bias=False)
+            l.weight = w
+            return l
+
+        self.pre_fc_norm_hidden = nn.RMSNorm(hidden_size)
+        self.pre_fc_norm_hidden.weight = weights['mtp.pre_fc_norm_hidden.weight']
+
+        self.pre_fc_norm_embedding = nn.RMSNorm(hidden_size)
+        self.pre_fc_norm_embedding.weight = weights['mtp.pre_fc_norm_embedding.weight']
+
+        self.fc = make_linear(fc_w)
+        self.q_proj = make_linear(q_w)
+        self.k_proj = make_linear(k_w)
+        self.v_proj = make_linear(weights['mtp.layers.0.self_attn.v_proj.weight'])
+        self.o_proj = make_linear(o_w)
+
+        self.q_norm = nn.RMSNorm(head_dim)
+        self.k_norm = nn.RMSNorm(head_dim)
+        q_norm_key = 'mtp.layers.0.self_attn.q_norm.weight'
+        k_norm_key = 'mtp.layers.0.self_attn.k_norm.weight'
+        if q_norm_key in weights:
+            self.q_norm.weight = weights[q_norm_key]
+            self.k_norm.weight = weights[k_norm_key]
+
+        self.input_layernorm = nn.RMSNorm(hidden_size)
+        self.input_layernorm.weight = weights['mtp.layers.0.input_layernorm.weight']
+
+        self.post_attention_layernorm = nn.RMSNorm(hidden_size)
+        self.post_attention_layernorm.weight = weights['mtp.layers.0.post_attention_layernorm.weight']
+
+        if self.is_moe:
+            # Reuse mlx-lm's SparseMoeBlock from the target model
+            moe_layer = None
+            for layer in self._inner.layers:
+                if hasattr(layer, 'mlp') and hasattr(layer.mlp, 'gate'):
+                    moe_layer = layer.mlp
+                    break
+            if moe_layer is None:
+                raise RuntimeError("MTP has MoE weights but target model has no MoE layer")
+
+            # Create a new MoE block with same class/config as target
+            moe_class = type(moe_layer)
+            args = getattr(self._text_model, 'args', None)
+            if args is None and hasattr(self._text_model, 'model'):
+                args = getattr(self._text_model.model, 'args', None)
+            self.mlp = moe_class(args)
+
+            # Load MTP MoE weights — remap HF expert names to mlx-lm SwitchLinear
+            prefix = 'mtp.layers.0.mlp.'
+            # Direct weights: gate, shared_expert, shared_expert_gate
+            direct_keys = {}
+            expert_weights = {}  # {proj_name: {expert_idx: weight}}
+            for k, v in weights.items():
+                if not k.startswith(prefix):
+                    continue
+                name = k[len(prefix):]
+                # Check if it's an individual expert weight
+                if name.startswith('experts.'):
+                    # experts.N.{gate,up,down}_proj.weight → stack into switch_mlp
+                    parts = name.split('.')
+                    idx = int(parts[1])
+                    proj = parts[2]  # gate_proj, up_proj, down_proj
+                    key = f'{proj}.{parts[3]}'  # gate_proj.weight
+                    if key not in expert_weights:
+                        expert_weights[key] = {}
+                    expert_weights[key][idx] = v
+                else:
+                    direct_keys[name] = v
+
+            # Stack individual expert weights into SwitchLinear format
+            moe_weights = []
+            for proj_key, idx_map in expert_weights.items():
+                n_experts = max(idx_map.keys()) + 1
+                stacked = mx.stack([idx_map[i] for i in range(n_experts)])
+                moe_weights.append((f'switch_mlp.{proj_key}', stacked))
+
+            # Add direct weights
+            for name, v in direct_keys.items():
+                moe_weights.append((name, v))
+
+            self.mlp.load_weights(moe_weights)
+            print(f"  MoE MLP: {len(moe_weights)} weight groups loaded "
+                  f"({len(expert_weights)} stacked expert projections)")
+        else:
+            self.gate_proj = make_linear(gate_w)
+            self.up_proj = make_linear(weights['mtp.layers.0.mlp.up_proj.weight'])
+            self.down_proj = make_linear(weights['mtp.layers.0.mlp.down_proj.weight'])
+
+        self.norm = nn.RMSNorm(hidden_size)
+        self.norm.weight = weights['mtp.norm.weight']
+
+        # RoPE from main model's GQA layers
+        for layer in self._inner.layers:
+            if not layer.is_linear:
+                self.rope = layer.self_attn.rope
+                break
+
+        # GQA config
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.scale = head_dim ** -0.5
+
+        # MTP KV cache (separate from main model)
+        self.kv_cache = None
+
+        mx.eval(self.pre_fc_norm_hidden.weight, self.pre_fc_norm_embedding.weight,
+                self.fc.weight, self.input_layernorm.weight,
+                self.post_attention_layernorm.weight, self.norm.weight,
+                self.q_norm.weight, self.k_norm.weight)
+
+        if quantize:
+            self._quantize_linears()
+
+        total_params = sum(w.size for w in weights.values())
+        print(f"  MTP loaded: {len(weights)} tensors, {total_params / 1e6:.1f}M params"
+              f"{' (quantized 8-bit gs=64)' if quantize else ' (bf16)'}")
+
+    def _quantize_linears(self):
+        """Quantize all MTP linear layers to 8-bit gs=64."""
+        for name in ['fc', 'q_proj', 'k_proj', 'v_proj', 'o_proj',
+                      'gate_proj', 'up_proj', 'down_proj']:
+            linear = getattr(self, name)
+            linear.weight = linear.weight.astype(mx.bfloat16)
+            q = nn.QuantizedLinear.from_linear(linear, group_size=64, bits=8)
+            mx.eval(q.parameters())
+            setattr(self, name, q)
+
+    def reset_cache(self):
+        """Reset the MTP KV cache (call at start of generation)."""
+        from mlx_lm.models.cache import KVCache
+        self.kv_cache = KVCache()
+
+    def get_hidden_state(self, inputs, cache, speculative=False):
+        """Run main model and return pre-norm hidden states + logits.
+
+        Delegates to the shared speculative_forward() function.
+        """
+        return speculative_forward(self.model, inputs, cache, speculative)
+
+    def _attn_mlp(self, h):
+        """Run GQA attention + MLP. Shared by predict, predict_hidden, predict_from_hidden."""
+        B, S = h.shape[0], h.shape[1]
+
+        residual = h
+        h = self.input_layernorm(h)
+
+        q_out = self.q_proj(h)
+        q_out, gate = mx.split(
+            q_out.reshape(B, S, self.num_heads, -1), 2, axis=-1
+        )
+        gate = gate.reshape(B, S, -1)
+
+        queries = self.q_norm(q_out).transpose(0, 2, 1, 3)
+        keys = self.k_norm(
+            self.k_proj(h).reshape(B, S, self.num_kv_heads, self.head_dim)
+        ).transpose(0, 2, 1, 3)
+        values = self.v_proj(h).reshape(
+            B, S, self.num_kv_heads, self.head_dim
+        ).transpose(0, 2, 1, 3)
+
+        if self.kv_cache is not None:
+            offset = self.kv_cache.offset
+            queries = self.rope(queries, offset=offset)
+            keys = self.rope(keys, offset=offset)
+            keys, values = self.kv_cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        mask = None
+        if S > 1:
+            total_kv = keys.shape[2]
+            q_pos = mx.arange(S) + (total_kv - S)
+            k_pos = mx.arange(total_kv)
+            mask = mx.where(k_pos[None, :] <= q_pos[:, None],
+                           mx.array(0, dtype=queries.dtype),
+                           mx.array(-1e9, dtype=queries.dtype))
+
+        output = mx.fast.scaled_dot_product_attention(
+            queries, keys, values, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, S, -1)
+
+        h = residual + self.o_proj(output * mx.sigmoid(gate))
+
+        residual = h
+        h = self.post_attention_layernorm(h)
+        if self.is_moe:
+            h = residual + self.mlp(h)
+        else:
+            h = residual + self.down_proj(nn.silu(self.gate_proj(h)) * self.up_proj(h))
+
+        return h  # post-FFN, pre-norm
+
+    def _combine(self, hidden_state, token_ids):
+        """Combine hidden state + token embedding → fc input."""
+        B, S = hidden_state.shape[0], hidden_state.shape[1]
+        embed = self.embed_tokens(token_ids.reshape(B, S))
+        h_norm = self.pre_fc_norm_hidden(hidden_state)
+        e_norm = self.pre_fc_norm_embedding(embed)
+        return self.fc(mx.concatenate([e_norm, h_norm], axis=-1))
+
+    def predict(self, hidden_state, token_ids, return_hidden=False, draft_mode=False):
+        """Predict next-next-token logits using MTP.
+
+        Args:
+            hidden_state: (B, S, D) bf16 — PRE-NORM hidden states
+            token_ids: (B, S) or (S,) int — tokens at each position
+            return_hidden: if True, also return pre-norm hidden for chaining
+            draft_mode: if True, use truncated lm_head (32K vocab) for speed
+        Returns:
+            logits: (B, S, vocab_size) if S>1, (B, vocab_size) if S=1
+            If return_hidden: (logits, hidden)
+        """
+        S = hidden_state.shape[1]
+        h = self._combine(hidden_state, token_ids)
+        pre_norm_out = self._attn_mlp(h)
+
+        normed = self.norm(pre_norm_out)
+        if draft_mode:
+            logits = normed @ self.draft_lm_head_weight.T
+        elif self.lm_head is not None:
+            logits = self.lm_head(normed)
+        else:
+            logits = self.embed_tokens.as_linear(normed)
+
+        if S == 1:
+            logits = logits.squeeze(1)
+
+        if return_hidden:
+            return logits, pre_norm_out
+        return logits
+
+    def predict_hidden(self, hidden_state, token_ids):
+        """Like predict() but returns only post-FFN hidden state (no lm_head)."""
+        h = self._combine(hidden_state, token_ids)
+        return self._attn_mlp(h)
+
+    def predict_from_hidden(self, prev_hidden):
+        """MTP step using post_norm of prev_hidden instead of token embedding.
+
+        Replaces embed_tokens + pre_fc_norm_embedding with just norm(prev_hidden).
+        This skips the lm_head → argmax → embed_tokens roundtrip.
+        """
+        post_norm = self.norm(prev_hidden)
+        h_norm = self.pre_fc_norm_hidden(prev_hidden)
+        h = self.fc(mx.concatenate([post_norm, h_norm], axis=-1))
+        return self._attn_mlp(h)
+
+
+def draft_tokens(mtp_pred, hidden, first_token_arr, gamma, temp, fast_lm_head=False):
+    """Draft γ tokens by chaining MTP predictions — fully lazy, no mx.eval.
+
+    The entire chain stays in the MLX computation graph. Draft token ids
+    are lazy mx.arrays (argmax/categorical results), not Python ints.
+
+    Args:
+        first_token_arr: mx.array of shape (1,1) — the token to start from
+    Returns: (draft_ids, draft_probs) where draft_ids[i] is a lazy mx.array
+             scalar, draft_probs[i] is the full draft distribution (or None if greedy)
+    """
+    draft_ids = []
+    draft_probs = []
+    h = hidden
+    tok_arr = first_token_arr
+
+    for i in range(gamma):
+        logits, h = mtp_pred.predict(h, tok_arr, return_hidden=True,
+                                      draft_mode=fast_lm_head)
+
+        if temp == 0:
+            tok_arr = mx.argmax(logits, axis=-1).reshape(1, 1)
+            draft_ids.append(tok_arr.reshape(-1))
+            draft_probs.append(None)
+        else:
+            q = mx.softmax(logits / temp, axis=-1)
+            tok_arr = mx.random.categorical(logits * (1.0 / temp)).reshape(1, 1)
+            draft_ids.append(tok_arr.reshape(-1))
+            draft_probs.append(q)
+
+    return draft_ids, draft_probs

src/exo/worker/engines/mlx/speculative/speculative_cache.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+"""SpeculativeArraysCache — wraps ArraysCache for correct GDN rollback.
+
+During speculative verification (S>1), captures:
+- all_states: per-step recurrent states from the speculative kernel
+- conv_input: full conv_input tensor for conv state rollback
+
+On rejection, rollback(n_accepted) restores both recurrent and conv
+state to the correct intermediate position.
+"""
+
+import mlx.core as mx
+
+
+class SpeculativeArraysCache:
+    """Wrapper around ArraysCache that supports rollback for speculative decode.
+
+    Delegates all normal cache operations to the underlying ArraysCache.
+    Adds all_states/conv_input storage and a rollback() method.
+    """
+
+    def __init__(self, base_cache, S, conv_kernel_size=4):
+        self.base = base_cache
+        self._S = S
+        self.n_keep = conv_kernel_size - 1  # typically 3
+        self.all_states = None    # [B, T, Hv, Dv, Dk] from speculative kernel
+        self.conv_input = None    # [B, n_keep+S, conv_dim] for conv rollback
+
+    # Delegate cache operations
+    def __getitem__(self, idx):
+        return self.base[idx]
+
+    def __setitem__(self, idx, val):
+        self.base[idx] = val
+
+    @property
+    def cache(self):
+        return self.base.cache
+
+    @cache.setter
+    def cache(self, v):
+        self.base.cache = v
+
+    @property
+    def state(self):
+        return self.base.state
+
+    @state.setter
+    def state(self, v):
+        self.base.state = v
+
+    @property
+    def lengths(self):
+        return self.base.lengths
+
+    @lengths.setter
+    def lengths(self, v):
+        self.base.lengths = v
+
+    @property
+    def left_padding(self):
+        return self.base.left_padding
+
+    @left_padding.setter
+    def left_padding(self, v):
+        self.base.left_padding = v
+
+    def advance(self, N):
+        self.base.advance(N)
+
+    def make_mask(self, N):
+        return self.base.make_mask(N)
+
+    def empty(self):
+        return self.base.empty()
+
+    @property
+    def nbytes(self):
+        return self.base.nbytes
+
+    def rollback(self, n_accepted):
+        """Roll back to state after processing n_accepted+1 tokens.
+
+        Args:
+            n_accepted: number of accepted draft tokens (0 = all rejected,
+                        only the first token 'y' was processed correctly)
+        """
+        # Recurrent state: restore intermediate state at accepted position
+        if self.all_states is not None:
+            self.base.cache[1] = self.all_states[0, n_accepted]
+
+        # Conv state: slice conv_input to the correct window
+        if self.conv_input is not None:
+            self.base.cache[0] = self.conv_input[:, n_accepted + 1: n_accepted + 1 + self.n_keep, :]
+
+        # Clear stored states
+        self.all_states = None
+        self.conv_input = None

src/exo/worker/engines/mlx/speculative/speculative_gdn_kernel.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+"""Speculative variant of the GatedDeltaNet kernel.
+
+Identical to the original gated_delta_kernel but also outputs per-step
+recurrent states for rollback during speculative decoding.
+
+The original kernel only writes the FINAL state. This variant writes
+the state at EVERY timestep to an extra output buffer `all_states`.
+"""
+
+from typing import Optional, Tuple
+import mlx.core as mx
+
+
+def _make_speculative_gated_delta_kernel(has_mask=False, vectorized=False):
+    mask_source = "mask[b_idx * T + t]" if has_mask else "true"
+
+    if vectorized:
+        g_comment = "// g: [B, T, Hv, Dk]"
+        g_setup = "auto g_ = g + (b_idx * T * Hv + hv_idx) * Dk;"
+        g_access = "g_[s_idx]"
+        g_advance = "g_ += Hv * Dk;"
+    else:
+        g_comment = "// g: [B, T, Hv]"
+        g_setup = "auto g_ = g + b_idx * T * Hv;"
+        g_access = "g_[hv_idx]"
+        g_advance = "g_ += Hv;"
+
+    source = f"""
+        auto n = thread_position_in_grid.z;
+        auto b_idx = n / Hv;
+        auto hv_idx = n % Hv;
+        auto hk_idx = hv_idx / (Hv / Hk);
+        constexpr int n_per_t = Dk / 32;
+
+        // q, k: [B, T, Hk, Dk]
+        auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
+        auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
+
+        // v, y: [B, T, Hv, Dv]
+        auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
+        y += b_idx * T * Hv * Dv + hv_idx * Dv;
+
+        auto dk_idx = thread_position_in_threadgroup.x;
+        auto dv_idx = thread_position_in_grid.y;
+
+        // state_in, state_out: [B, Hv, Dv, Dk]
+        auto i_state = state_in + (n * Dv + dv_idx) * Dk;
+        auto o_state = state_out + (n * Dv + dv_idx) * Dk;
+
+        // all_states: [B, T, Hv, Dv, Dk] — per-step state output
+        auto a_state = all_states + (b_idx * T * Hv * Dv + hv_idx * Dv + dv_idx) * Dk;
+        auto a_stride = Hv * Dv * Dk;
+
+        float state[n_per_t];
+        for (int i = 0; i < n_per_t; ++i) {{
+          auto s_idx = n_per_t * dk_idx + i;
+          state[i] = static_cast<float>(i_state[s_idx]);
+        }}
+
+        {g_comment}
+        {g_setup}
+        auto beta_ = beta + b_idx * T * Hv;
+
+        for (int t = 0; t < T; ++t) {{
+          if ({mask_source}) {{
+            float kv_mem = 0.0f;
+            for (int i = 0; i < n_per_t; ++i) {{
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = state[i] * {g_access};
+              kv_mem += state[i] * k_[s_idx];
+            }}
+            kv_mem = simd_sum(kv_mem);
+
+            auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
+
+            float out = 0.0f;
+            for (int i = 0; i < n_per_t; ++i) {{
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = state[i] + k_[s_idx] * delta;
+              out += state[i] * q_[s_idx];
+            }}
+            out = simd_sum(out);
+            if (thread_index_in_simdgroup == 0) {{
+              y[dv_idx] = static_cast<InT>(out);
+            }}
+          }}
+
+          // Save per-step state for speculative rollback
+          for (int i = 0; i < n_per_t; ++i) {{
+            auto s_idx = n_per_t * dk_idx + i;
+            a_state[s_idx] = static_cast<InT>(state[i]);
+          }}
+          a_state += a_stride;
+
+          q_ += Hk * Dk;
+          k_ += Hk * Dk;
+          v_ += Hv * Dv;
+          y += Hv * Dv;
+          {g_advance}
+          beta_ += Hv;
+        }}
+        // Write final state (same as original kernel)
+        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 = "_spec"
+    if vectorized:
+        suffix += "_vec"
+    if has_mask:
+        suffix += "_mask"
+
+    return mx.fast.metal_kernel(
+        name=f"gated_delta_step{suffix}",
+        input_names=inputs,
+        output_names=["y", "state_out", "all_states"],
+        source=source,
+    )
+
+
+# Pre-build kernel variants
+_spec_kernel = _make_speculative_gated_delta_kernel(has_mask=False, vectorized=False)
+_spec_kernel_masked = _make_speculative_gated_delta_kernel(has_mask=True, vectorized=False)
+_spec_kernel_vec = _make_speculative_gated_delta_kernel(has_mask=False, vectorized=True)
+_spec_kernel_vec_masked = _make_speculative_gated_delta_kernel(has_mask=True, vectorized=True)
+
+
+def speculative_gated_delta_kernel(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: mx.array,
+    mask: Optional[mx.array] = None,
+) -> Tuple[mx.array, mx.array, mx.array]:
+    """Like gated_delta_kernel but also returns per-step states.
+
+    Returns:
+        y: [B, T, Hv, Dv] — output (same as original)
+        state_out: [B, Hv, Dv, Dk] — final state (same as original)
+        all_states: [B, T, Hv, Dv, Dk] — state after each timestep
+    """
+    B, T, Hk, Dk = k.shape
+    Hv, Dv = v.shape[2:]
+    input_type = q.dtype
+
+    if g.ndim == 4:
+        kernel = _spec_kernel_vec
+        inputs = [q, k, v, g, beta, state, T]
+        if mask is not None:
+            kernel = _spec_kernel_vec_masked
+            inputs.append(mask)
+    else:
+        kernel = _spec_kernel
+        inputs = [q, k, v, g, beta, state, T]
+        if mask is not None:
+            kernel = _spec_kernel_masked
+            inputs.append(mask)
+
+    return kernel(
+        inputs=inputs,
+        template=[
+            ("InT", input_type),
+            ("Dk", Dk),
+            ("Dv", Dv),
+            ("Hk", Hk),
+            ("Hv", Hv),
+        ],
+        grid=(32, Dv, B * Hv),
+        threadgroup=(32, 4, 1),
+        output_shapes=[(B, T, Hv, Dv), state.shape, (B, T, Hv, Dv, Dk)],
+        output_dtypes=[input_type, input_type, input_type],
+    )

src/exo/worker/engines/mlx/tests/test_extract_top_logprobs.py

@@ -0,0 +1,290 @@
+# type: ignore
+import math
+from unittest.mock import MagicMock
+
+import mlx.core as mx
+import mlx.nn as nn
+import pytest
+from mlx_lm.generate import BatchGenerator
+
+from exo.worker.engines.mlx.generator.generate import extract_top_logprobs
+from exo.worker.engines.mlx.patches.opt_batch_gen import (
+    _PRECOMPUTE_TOP_K,
+    apply_batch_gen_patch,
+)
+
+
+def _mock_tokenizer() -> MagicMock:
+    tok = MagicMock()
+    tok.decode = lambda ids: f"tok_{ids[0]}"
+    return tok
+
+
+def _make_logprobs(values: list[float]) -> mx.array:
+    arr = mx.array(values, dtype=mx.float32)
+    mx.eval(arr)
+    return arr
+
+
+class TestExtractTopLogprobsFallback:
+    def test_returns_correct_selected_logprob(self) -> None:
+        lp = _make_logprobs([-1.0, -2.0, -0.5, -3.0, -4.0])
+        selected, _ = extract_top_logprobs(
+            lp, _mock_tokenizer(), top_logprobs=3, selected_token=2
+        )
+        assert selected == pytest.approx(-0.5)
+
+    def test_returns_top_k_sorted_descending(self) -> None:
+        lp = _make_logprobs([-1.0, -2.0, -0.5, -3.0, -4.0])
+        _, items = extract_top_logprobs(
+            lp, _mock_tokenizer(), top_logprobs=3, selected_token=0
+        )
+        logprob_values = [item.logprob for item in items]
+        assert logprob_values == sorted(logprob_values, reverse=True)
+        assert len(items) == 3
+
+    def test_top_tokens_are_most_probable(self) -> None:
+        lp = _make_logprobs([-5.0, -1.0, -3.0, -0.1, -2.0])
+        _, items = extract_top_logprobs(
+            lp, _mock_tokenizer(), top_logprobs=2, selected_token=0
+        )
+        token_ids = [int(item.token.split("_")[1]) for item in items]
+        assert 3 in token_ids
+        assert 1 in token_ids
+
+    def test_top_logprobs_clamped_to_vocab_size(self) -> None:
+        lp = _make_logprobs([-1.0, -2.0, -3.0, -4.0, -5.0])
+        _, items = extract_top_logprobs(
+            lp, _mock_tokenizer(), top_logprobs=10, selected_token=0
+        )
+        assert len(items) == 4
+
+    def test_nan_logprobs_filtered(self) -> None:
+        lp = _make_logprobs([-1.0, float("nan"), -0.5])
+        _, items = extract_top_logprobs(
+            lp, _mock_tokenizer(), top_logprobs=3, selected_token=0
+        )
+        for item in items:
+            assert not math.isnan(item.logprob)
+
+    def test_token_bytes_correct(self) -> None:
+        tok = MagicMock()
+        tok.decode = lambda ids: "hello"
+        lp = _make_logprobs([-1.0, -2.0])
+        _, items = extract_top_logprobs(lp, tok, top_logprobs=2, selected_token=0)
+        assert items[0].bytes == list("hello".encode("utf-8"))
+
+
+class TestExtractTopLogprobsPrecomputed:
+    def test_uses_precomputed_data(self) -> None:
+        lp = _make_logprobs([-99.0])
+        selected, items = extract_top_logprobs(
+            lp,
+            _mock_tokenizer(),
+            top_logprobs=2,
+            selected_token=0,
+            precomputed_indices=[3, 1, 0],
+            precomputed_values=[-0.1, -1.0, -5.0],
+            precomputed_selected=-0.1,
+        )
+        assert selected == pytest.approx(-0.1)
+        assert len(items) == 2
+        assert items[0].token == "tok_3"
+        assert items[0].logprob == pytest.approx(-0.1)
+        assert items[1].token == "tok_1"
+        assert items[1].logprob == pytest.approx(-1.0)
+
+    def test_slices_precomputed_to_requested_k(self) -> None:
+        lp = _make_logprobs([-99.0])
+        _, items = extract_top_logprobs(
+            lp,
+            _mock_tokenizer(),
+            top_logprobs=1,
+            selected_token=0,
+            precomputed_indices=[3, 1, 0, 2, 4],
+            precomputed_values=[-0.1, -1.0, -2.0, -3.0, -4.0],
+            precomputed_selected=-0.1,
+        )
+        assert len(items) == 1
+        assert items[0].token == "tok_3"
+
+    def test_falls_back_when_precomputed_partial(self) -> None:
+        lp = _make_logprobs([-1.0, -2.0, -0.5])
+        selected, items = extract_top_logprobs(
+            lp,
+            _mock_tokenizer(),
+            top_logprobs=2,
+            selected_token=2,
+            precomputed_indices=[0, 2],
+            precomputed_values=None,
+            precomputed_selected=None,
+        )
+        assert selected == pytest.approx(-0.5)
+        assert len(items) == 2
+
+    def test_precomputed_matches_fallback(self) -> None:
+        lp = _make_logprobs([-1.0, -0.3, -2.5, -0.1, -4.0, -0.8, -3.0, -1.5])
+        tok = _mock_tokenizer()
+
+        selected_fb, items_fb = extract_top_logprobs(
+            lp, tok, top_logprobs=5, selected_token=1
+        )
+
+        pre_indices = [item.token.split("_")[1] for item in items_fb]
+        pre_indices_int = [int(x) for x in pre_indices]
+        pre_values = [item.logprob for item in items_fb]
+
+        selected_pc, items_pc = extract_top_logprobs(
+            lp,
+            tok,
+            top_logprobs=5,
+            selected_token=1,
+            precomputed_indices=pre_indices_int,
+            precomputed_values=pre_values,
+            precomputed_selected=selected_fb,
+        )
+
+        assert selected_pc == pytest.approx(selected_fb)
+        assert len(items_pc) == len(items_fb)
+        for a, b in zip(items_pc, items_fb, strict=True):
+            assert a.token == b.token
+            assert a.logprob == pytest.approx(b.logprob)
+
+
+def _tiny_model() -> nn.Module:
+    from mlx_lm.models.llama import Model, ModelArgs
+
+    mx.random.seed(42)
+    args = ModelArgs(
+        model_type="llama",
+        hidden_size=64,
+        num_hidden_layers=2,
+        intermediate_size=128,
+        num_attention_heads=2,
+        num_key_value_heads=1,
+        rms_norm_eps=1e-6,
+        vocab_size=256,
+        rope_theta=10000.0,
+        tie_word_embeddings=True,
+    )
+    model = Model(args)
+    mx.eval(model.parameters())
+    return model
+
+
+@pytest.mark.slow
+class TestBatchedTopKPrecompute:
+    @pytest.fixture(autouse=True)
+    def _reset_globals(self) -> None:
+        import exo.worker.engines.mlx.patches.opt_batch_gen as _mod
+
+        _mod._pending_topk_idx = None
+        _mod._pending_topk_val = None
+        _mod._pending_selected_lps = None
+
+    def _run_generator(
+        self, model: nn.Module, prompts: list[list[int]], steps: int, needs_topk: bool
+    ) -> list[list[BatchGenerator.Response]]:
+        apply_batch_gen_patch()
+        gen = BatchGenerator(model=model, stop_tokens=set(), prefill_step_size=512)
+        gen._needs_topk = needs_topk
+        gen.insert(prompts)
+        all_responses: list[list[BatchGenerator.Response]] = []
+        for _ in range(steps + len(prompts)):
+            responses = gen.next()
+            if responses:
+                all_responses.append(responses)
+            if gen.active_batch is None and not gen.unprocessed_prompts:
+                break
+        gen.close()
+        return all_responses
+
+    def test_precomputed_topk_attached_to_responses(self) -> None:
+        model = _tiny_model()
+        steps = self._run_generator(model, [[1, 2, 3]], 5, needs_topk=True)
+        found_precomputed = False
+        for step_responses in steps:
+            for resp in step_responses:
+                if hasattr(resp, "_topk_indices"):
+                    found_precomputed = True
+                    assert hasattr(resp, "_topk_values"), (
+                        "Response missing _topk_values"
+                    )
+                    assert hasattr(resp, "_selected_logprob"), (
+                        "Response missing _selected_logprob"
+                    )
+                    assert len(resp._topk_indices) == _PRECOMPUTE_TOP_K
+                    assert len(resp._topk_values) == _PRECOMPUTE_TOP_K
+        assert found_precomputed, "No responses had precomputed topk"
+
+    def test_no_topk_when_not_needed(self) -> None:
+        model = _tiny_model()
+        steps = self._run_generator(model, [[1, 2, 3]], 5, needs_topk=False)
+        for step_responses in steps:
+            for resp in step_responses:
+                assert not hasattr(resp, "_topk_indices")
+
+    def test_precomputed_matches_fallback_in_batch(self) -> None:
+        model = _tiny_model()
+        tok = _mock_tokenizer()
+        steps = self._run_generator(model, [[1, 2, 3]], 10, needs_topk=True)
+        for step_responses in steps[1:]:
+            for resp in step_responses:
+                if not hasattr(resp, "_topk_indices"):
+                    continue
+                selected_fb, items_fb = extract_top_logprobs(
+                    resp.logprobs, tok, top_logprobs=5, selected_token=resp.token
+                )
+                selected_pc, items_pc = extract_top_logprobs(
+                    resp.logprobs,
+                    tok,
+                    top_logprobs=5,
+                    selected_token=resp.token,
+                    precomputed_indices=resp._topk_indices,
+                    precomputed_values=resp._topk_values,
+                    precomputed_selected=resp._selected_logprob,
+                )
+                assert selected_pc == pytest.approx(selected_fb, abs=1e-5)
+                for a, b in zip(items_pc, items_fb, strict=True):
+                    assert a.token == b.token
+                    assert a.logprob == pytest.approx(b.logprob, abs=1e-5)
+
+    def test_topk_correct_after_batch_shrink(self) -> None:
+        model = _tiny_model()
+        tok = _mock_tokenizer()
+        apply_batch_gen_patch()
+        gen = BatchGenerator(
+            model=model, stop_tokens={0}, prefill_step_size=512, max_tokens=3
+        )
+        gen._needs_topk = True
+        gen.insert([[1, 2, 3], [4, 5, 6]], max_tokens=[3, 20])
+
+        seen_shrink = False
+        for _ in range(30):
+            responses = gen.next()
+            for resp in responses:
+                if resp.finish_reason is not None:
+                    seen_shrink = True
+                    continue
+                if not hasattr(resp, "_topk_indices"):
+                    continue
+                selected_fb, items_fb = extract_top_logprobs(
+                    resp.logprobs, tok, top_logprobs=5, selected_token=resp.token
+                )
+                selected_pc, _ = extract_top_logprobs(
+                    resp.logprobs,
+                    tok,
+                    top_logprobs=5,
+                    selected_token=resp.token,
+                    precomputed_indices=resp._topk_indices,
+                    precomputed_values=resp._topk_values,
+                    precomputed_selected=resp._selected_logprob,
+                )
+                assert selected_pc == pytest.approx(selected_fb, abs=1e-5), (
+                    f"Mismatch after batch shrink: precomputed={selected_pc}, fallback={selected_fb}"
+                )
+            if gen.active_batch is None and not gen.unprocessed_prompts:
+                break
+
+        gen.close()
+        assert seen_shrink, "Expected at least one request to finish (batch shrink)"

src/exo/worker/engines/mlx/utils_mlx.py

@@ -190,6 +190,11 @@ def load_mlx_items(
         mx.eval(model)
         end_time = time.perf_counter()
         logger.info(f"Time taken to load model: {(end_time - start_time):.2f}s")
+
+        # Apply kernel fusion patches if available for this model
+        from exo.worker.engines.mlx.patches import maybe_apply_patches
+        maybe_apply_patches(model, model_path)
+
         tokenizer = get_tokenizer(model_path, bound_instance.bound_shard)
 
     else:
@@ -638,6 +643,7 @@ class NullKVCache(KVCache):
     @property
     def state(self) -> tuple[mx.array, mx.array]:
         # matches what mx.save_safetensors / mx.eval expect
+        assert self.keys is not None and self.values is not None
         return self.keys, self.values
 
     @state.setter

src/exo/worker/runner/bootstrap.py

@@ -8,6 +8,7 @@ from exo.shared.types.tasks import Task, TaskId
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
+from exo.worker.engines.mlx.patches import apply_mlx_patches
 
 logger: "loguru.Logger" = loguru.logger
 
@@ -45,6 +46,8 @@ def entrypoint(
         else:
             from exo.worker.runner.llm_inference.runner import Runner
 
+            apply_mlx_patches()
+
             runner = Runner(
                 bound_instance, event_sender, task_receiver, cancel_receiver
             )

src/exo/worker/runner/llm_inference/batch_generator.py

@@ -345,6 +345,8 @@ class BatchGenerator(InferenceGenerator):
             group=self.group,
             model_id=self.model_id,
         )
+        # Warm up speculative cycle if MTP is enabled
+        self._mlx_gen.warmup_speculative(self.model, self.tokenizer)
 
     def submit(
         self,
@@ -427,7 +429,8 @@ class BatchGenerator(InferenceGenerator):
 
             task, queue, output_generator = self._active_tasks[uid]
             queue.push(response)
-            parsed = next(output_generator)
+            # If a generator fails to parse for some reason and returns early, we should not crash
+            parsed = next(output_generator, None)
 
             if parsed is not None:
                 output.append((task.task_id, parsed))

src/exo/worker/runner/llm_inference/runner.py

@@ -319,7 +319,9 @@ class Runner:
         return ExitCode.AllTasksComplete
 
     def send_response(
-        self, response: GenerationResponse | ToolCallResponse, command_id: CommandId
+        self,
+        response: GenerationResponse | ToolCallResponse,
+        command_id: CommandId,
     ):
         match response:
             case GenerationResponse():

src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py

@@ -43,8 +43,8 @@ def run_pipeline_device(
 
         def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
             for layer in self.layers:
-                x = layer(x, *args, **kwargs)  # pyright: ignore[reportUnknownVariableType]
-            return x  # pyright: ignore[reportUnknownVariableType]
+                x = layer(x, *args, **kwargs)
+            return x
 
     try:
         group = mx.distributed.init(backend="ring", strict=True)

uv.lock

@@ -213,14 +213,20 @@ sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8
 wheels = [
     { url = "https://files.pythonhosted.org/packages/b0/1e/d22cc63332bd59b06481ceaac49d6c507598642e2230f201649058a7e704/cffi-2.0.0-cp313-cp313-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:07b271772c100085dd28b74fa0cd81c8fb1a3ba18b21e03d7c27f3436a10606b", size = 212446, upload-time = "2025-09-08T23:23:03.472Z" },
     { url = "https://files.pythonhosted.org/packages/a9/f5/a2c23eb03b61a0b8747f211eb716446c826ad66818ddc7810cc2cc19b3f2/cffi-2.0.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d48a880098c96020b02d5a1f7d9251308510ce8858940e6fa99ece33f610838b", size = 220101, upload-time = "2025-09-08T23:23:04.792Z" },
+    { url = "https://files.pythonhosted.org/packages/f2/7f/e6647792fc5850d634695bc0e6ab4111ae88e89981d35ac269956605feba/cffi-2.0.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:f93fd8e5c8c0a4aa1f424d6173f14a892044054871c771f8566e4008eaa359d2", size = 207948, upload-time = "2025-09-08T23:23:06.127Z" },
+    { url = "https://files.pythonhosted.org/packages/cb/1e/a5a1bd6f1fb30f22573f76533de12a00bf274abcdc55c8edab639078abb6/cffi-2.0.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:dd4f05f54a52fb558f1ba9f528228066954fee3ebe629fc1660d874d040ae5a3", size = 206422, upload-time = "2025-09-08T23:23:07.753Z" },
     { url = "https://files.pythonhosted.org/packages/98/df/0a1755e750013a2081e863e7cd37e0cdd02664372c754e5560099eb7aa44/cffi-2.0.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c8d3b5532fc71b7a77c09192b4a5a200ea992702734a2e9279a37f2478236f26", size = 219499, upload-time = "2025-09-08T23:23:09.648Z" },
     { url = "https://files.pythonhosted.org/packages/50/e1/a969e687fcf9ea58e6e2a928ad5e2dd88cc12f6f0ab477e9971f2309b57c/cffi-2.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d9b29c1f0ae438d5ee9acb31cadee00a58c46cc9c0b2f9038c6b0b3470877a8c", size = 222928, upload-time = "2025-09-08T23:23:10.928Z" },
     { url = "https://files.pythonhosted.org/packages/36/54/0362578dd2c9e557a28ac77698ed67323ed5b9775ca9d3fe73fe191bb5d8/cffi-2.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6d50360be4546678fc1b79ffe7a66265e28667840010348dd69a314145807a1b", size = 221302, upload-time = "2025-09-08T23:23:12.42Z" },
     { url = "https://files.pythonhosted.org/packages/d6/43/0e822876f87ea8a4ef95442c3d766a06a51fc5298823f884ef87aaad168c/cffi-2.0.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:24b6f81f1983e6df8db3adc38562c83f7d4a0c36162885ec7f7b77c7dcbec97b", size = 220049, upload-time = "2025-09-08T23:23:20.853Z" },
+    { url = "https://files.pythonhosted.org/packages/b4/89/76799151d9c2d2d1ead63c2429da9ea9d7aac304603de0c6e8764e6e8e70/cffi-2.0.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:12873ca6cb9b0f0d3a0da705d6086fe911591737a59f28b7936bdfed27c0d47c", size = 207793, upload-time = "2025-09-08T23:23:22.08Z" },
+    { url = "https://files.pythonhosted.org/packages/bb/dd/3465b14bb9e24ee24cb88c9e3730f6de63111fffe513492bf8c808a3547e/cffi-2.0.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:d9b97165e8aed9272a6bb17c01e3cc5871a594a446ebedc996e2397a1c1ea8ef", size = 206300, upload-time = "2025-09-08T23:23:23.314Z" },
     { url = "https://files.pythonhosted.org/packages/47/d9/d83e293854571c877a92da46fdec39158f8d7e68da75bf73581225d28e90/cffi-2.0.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:afb8db5439b81cf9c9d0c80404b60c3cc9c3add93e114dcae767f1477cb53775", size = 219244, upload-time = "2025-09-08T23:23:24.541Z" },
     { url = "https://files.pythonhosted.org/packages/2b/0f/1f177e3683aead2bb00f7679a16451d302c436b5cbf2505f0ea8146ef59e/cffi-2.0.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:737fe7d37e1a1bffe70bd5754ea763a62a066dc5913ca57e957824b72a85e205", size = 222828, upload-time = "2025-09-08T23:23:26.143Z" },
     { url = "https://files.pythonhosted.org/packages/c6/0f/cafacebd4b040e3119dcb32fed8bdef8dfe94da653155f9d0b9dc660166e/cffi-2.0.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:38100abb9d1b1435bc4cc340bb4489635dc2f0da7456590877030c9b3d40b0c1", size = 220926, upload-time = "2025-09-08T23:23:27.873Z" },
     { url = "https://files.pythonhosted.org/packages/be/b4/c56878d0d1755cf9caa54ba71e5d049479c52f9e4afc230f06822162ab2f/cffi-2.0.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7cc09976e8b56f8cebd752f7113ad07752461f48a58cbba644139015ac24954c", size = 221593, upload-time = "2025-09-08T23:23:31.91Z" },
+    { url = "https://files.pythonhosted.org/packages/e0/0d/eb704606dfe8033e7128df5e90fee946bbcb64a04fcdaa97321309004000/cffi-2.0.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:92b68146a71df78564e4ef48af17551a5ddd142e5190cdf2c5624d0c3ff5b2e8", size = 209354, upload-time = "2025-09-08T23:23:33.214Z" },
+    { url = "https://files.pythonhosted.org/packages/d8/19/3c435d727b368ca475fb8742ab97c9cb13a0de600ce86f62eab7fa3eea60/cffi-2.0.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:b1e74d11748e7e98e2f426ab176d4ed720a64412b6a15054378afdb71e0f37dc", size = 208480, upload-time = "2025-09-08T23:23:34.495Z" },
     { url = "https://files.pythonhosted.org/packages/d0/44/681604464ed9541673e486521497406fadcc15b5217c3e326b061696899a/cffi-2.0.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:28a3a209b96630bca57cce802da70c266eb08c6e97e5afd61a75611ee6c64592", size = 221584, upload-time = "2025-09-08T23:23:36.096Z" },
     { url = "https://files.pythonhosted.org/packages/25/8e/342a504ff018a2825d395d44d63a767dd8ebc927ebda557fecdaca3ac33a/cffi-2.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:7553fb2090d71822f02c629afe6042c299edf91ba1bf94951165613553984512", size = 224443, upload-time = "2025-09-08T23:23:37.328Z" },
     { url = "https://files.pythonhosted.org/packages/e1/5e/b666bacbbc60fbf415ba9988324a132c9a7a0448a9a8f125074671c0f2c3/cffi-2.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6c6c373cfc5c83a975506110d17457138c8c63016b563cc9ed6e056a82f13ce4", size = 223437, upload-time = "2025-09-08T23:23:38.945Z" },
@@ -344,8 +350,10 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/5c/49/498c86566a1d80e978b42f0d702795f69887005548c041636df6ae1ca64c/cryptography-46.0.3-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:01ca9ff2885f3acc98c29f1860552e37f6d7c7d013d7334ff2a9de43a449315d", size = 4450807, upload-time = "2025-10-15T23:16:56.414Z" },
     { url = "https://files.pythonhosted.org/packages/4b/0a/863a3604112174c8624a2ac3c038662d9e59970c7f926acdcfaed8d61142/cryptography-46.0.3-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:6eae65d4c3d33da080cff9c4ab1f711b15c1d9760809dad6ea763f3812d254cb", size = 4299615, upload-time = "2025-10-15T23:16:58.442Z" },
     { url = "https://files.pythonhosted.org/packages/64/02/b73a533f6b64a69f3cd3872acb6ebc12aef924d8d103133bb3ea750dc703/cryptography-46.0.3-cp311-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:e5bf0ed4490068a2e72ac03d786693adeb909981cc596425d09032d372bcc849", size = 4016800, upload-time = "2025-10-15T23:17:00.378Z" },
+    { url = "https://files.pythonhosted.org/packages/25/d5/16e41afbfa450cde85a3b7ec599bebefaef16b5c6ba4ec49a3532336ed72/cryptography-46.0.3-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:5ecfccd2329e37e9b7112a888e76d9feca2347f12f37918facbb893d7bb88ee8", size = 4984707, upload-time = "2025-10-15T23:17:01.98Z" },
     { url = "https://files.pythonhosted.org/packages/c9/56/e7e69b427c3878352c2fb9b450bd0e19ed552753491d39d7d0a2f5226d41/cryptography-46.0.3-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:a2c0cd47381a3229c403062f764160d57d4d175e022c1df84e168c6251a22eec", size = 4482541, upload-time = "2025-10-15T23:17:04.078Z" },
     { url = "https://files.pythonhosted.org/packages/78/f6/50736d40d97e8483172f1bb6e698895b92a223dba513b0ca6f06b2365339/cryptography-46.0.3-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:549e234ff32571b1f4076ac269fcce7a808d3bf98b76c8dd560e42dbc66d7d91", size = 4299464, upload-time = "2025-10-15T23:17:05.483Z" },
+    { url = "https://files.pythonhosted.org/packages/00/de/d8e26b1a855f19d9994a19c702fa2e93b0456beccbcfe437eda00e0701f2/cryptography-46.0.3-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:c0a7bb1a68a5d3471880e264621346c48665b3bf1c3759d682fc0864c540bd9e", size = 4950838, upload-time = "2025-10-15T23:17:07.425Z" },
     { url = "https://files.pythonhosted.org/packages/8f/29/798fc4ec461a1c9e9f735f2fc58741b0daae30688f41b2497dcbc9ed1355/cryptography-46.0.3-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:10b01676fc208c3e6feeb25a8b83d81767e8059e1fe86e1dc62d10a3018fa926", size = 4481596, upload-time = "2025-10-15T23:17:09.343Z" },
     { url = "https://files.pythonhosted.org/packages/15/8d/03cd48b20a573adfff7652b76271078e3045b9f49387920e7f1f631d125e/cryptography-46.0.3-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0abf1ffd6e57c67e92af68330d05760b7b7efb243aab8377e583284dbab72c71", size = 4426782, upload-time = "2025-10-15T23:17:11.22Z" },
     { url = "https://files.pythonhosted.org/packages/fa/b1/ebacbfe53317d55cf33165bda24c86523497a6881f339f9aae5c2e13e57b/cryptography-46.0.3-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a04bee9ab6a4da801eb9b51f1b708a1b5b5c9eb48c03f74198464c66f0d344ac", size = 4698381, upload-time = "2025-10-15T23:17:12.829Z" },
@@ -353,8 +361,10 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/c5/fd/bc1daf8230eaa075184cbbf5f8cd00ba9db4fd32d63fb83da4671b72ed8a/cryptography-46.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:39b6755623145ad5eff1dab323f4eae2a32a77a7abef2c5089a04a3d04366715", size = 4435078, upload-time = "2025-10-15T23:17:23.042Z" },
     { url = "https://files.pythonhosted.org/packages/82/98/d3bd5407ce4c60017f8ff9e63ffee4200ab3e23fe05b765cab805a7db008/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:db391fa7c66df6762ee3f00c95a89e6d428f4d60e7abc8328f4fe155b5ac6e54", size = 4293460, upload-time = "2025-10-15T23:17:24.885Z" },
     { url = "https://files.pythonhosted.org/packages/26/e9/e23e7900983c2b8af7a08098db406cf989d7f09caea7897e347598d4cd5b/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:78a97cf6a8839a48c49271cdcbd5cf37ca2c1d6b7fdd86cc864f302b5e9bf459", size = 3995237, upload-time = "2025-10-15T23:17:26.449Z" },
+    { url = "https://files.pythonhosted.org/packages/91/15/af68c509d4a138cfe299d0d7ddb14afba15233223ebd933b4bbdbc7155d3/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:dfb781ff7eaa91a6f7fd41776ec37c5853c795d3b358d4896fdbb5df168af422", size = 4967344, upload-time = "2025-10-15T23:17:28.06Z" },
     { url = "https://files.pythonhosted.org/packages/ca/e3/8643d077c53868b681af077edf6b3cb58288b5423610f21c62aadcbe99f4/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:6f61efb26e76c45c4a227835ddeae96d83624fb0d29eb5df5b96e14ed1a0afb7", size = 4466564, upload-time = "2025-10-15T23:17:29.665Z" },
     { url = "https://files.pythonhosted.org/packages/0e/43/c1e8726fa59c236ff477ff2b5dc071e54b21e5a1e51aa2cee1676f1c986f/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:23b1a8f26e43f47ceb6d6a43115f33a5a37d57df4ea0ca295b780ae8546e8044", size = 4292415, upload-time = "2025-10-15T23:17:31.686Z" },
+    { url = "https://files.pythonhosted.org/packages/42/f9/2f8fefdb1aee8a8e3256a0568cffc4e6d517b256a2fe97a029b3f1b9fe7e/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:b419ae593c86b87014b9be7396b385491ad7f320bde96826d0dd174459e54665", size = 4931457, upload-time = "2025-10-15T23:17:33.478Z" },
     { url = "https://files.pythonhosted.org/packages/79/30/9b54127a9a778ccd6d27c3da7563e9f2d341826075ceab89ae3b41bf5be2/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:50fc3343ac490c6b08c0cf0d704e881d0d660be923fd3076db3e932007e726e3", size = 4466074, upload-time = "2025-10-15T23:17:35.158Z" },
     { url = "https://files.pythonhosted.org/packages/ac/68/b4f4a10928e26c941b1b6a179143af9f4d27d88fe84a6a3c53592d2e76bf/cryptography-46.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:22d7e97932f511d6b0b04f2bfd818d73dcd5928db509460aaf48384778eb6d20", size = 4420569, upload-time = "2025-10-15T23:17:37.188Z" },
     { url = "https://files.pythonhosted.org/packages/a3/49/3746dab4c0d1979888f125226357d3262a6dd40e114ac29e3d2abdf1ec55/cryptography-46.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:d55f3dffadd674514ad19451161118fd010988540cee43d8bc20675e775925de", size = 4681941, upload-time = "2025-10-15T23:17:39.236Z" },
@@ -362,8 +372,10 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/26/42/fa8389d4478368743e24e61eea78846a0006caffaf72ea24a15159215a14/cryptography-46.0.3-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:15ab9b093e8f09daab0f2159bb7e47532596075139dd74365da52ecc9cb46c5d", size = 4440029, upload-time = "2025-10-15T23:17:49.837Z" },
     { url = "https://files.pythonhosted.org/packages/5f/eb/f483db0ec5ac040824f269e93dd2bd8a21ecd1027e77ad7bdf6914f2fd80/cryptography-46.0.3-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:46acf53b40ea38f9c6c229599a4a13f0d46a6c3fa9ef19fc1a124d62e338dfa0", size = 4297222, upload-time = "2025-10-15T23:17:51.357Z" },
     { url = "https://files.pythonhosted.org/packages/fd/cf/da9502c4e1912cb1da3807ea3618a6829bee8207456fbbeebc361ec38ba3/cryptography-46.0.3-cp38-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:10ca84c4668d066a9878890047f03546f3ae0a6b8b39b697457b7757aaf18dbc", size = 4012280, upload-time = "2025-10-15T23:17:52.964Z" },
+    { url = "https://files.pythonhosted.org/packages/6b/8f/9adb86b93330e0df8b3dcf03eae67c33ba89958fc2e03862ef1ac2b42465/cryptography-46.0.3-cp38-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:36e627112085bb3b81b19fed209c05ce2a52ee8b15d161b7c643a7d5a88491f3", size = 4978958, upload-time = "2025-10-15T23:17:54.965Z" },
     { url = "https://files.pythonhosted.org/packages/d1/a0/5fa77988289c34bdb9f913f5606ecc9ada1adb5ae870bd0d1054a7021cc4/cryptography-46.0.3-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:1000713389b75c449a6e979ffc7dcc8ac90b437048766cef052d4d30b8220971", size = 4473714, upload-time = "2025-10-15T23:17:56.754Z" },
     { url = "https://files.pythonhosted.org/packages/14/e5/fc82d72a58d41c393697aa18c9abe5ae1214ff6f2a5c18ac470f92777895/cryptography-46.0.3-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:b02cf04496f6576afffef5ddd04a0cb7d49cf6be16a9059d793a30b035f6b6ac", size = 4296970, upload-time = "2025-10-15T23:17:58.588Z" },
+    { url = "https://files.pythonhosted.org/packages/78/06/5663ed35438d0b09056973994f1aec467492b33bd31da36e468b01ec1097/cryptography-46.0.3-cp38-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:71e842ec9bc7abf543b47cf86b9a743baa95f4677d22baa4c7d5c69e49e9bc04", size = 4940236, upload-time = "2025-10-15T23:18:00.897Z" },
     { url = "https://files.pythonhosted.org/packages/fc/59/873633f3f2dcd8a053b8dd1d38f783043b5fce589c0f6988bf55ef57e43e/cryptography-46.0.3-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:402b58fc32614f00980b66d6e56a5b4118e6cb362ae8f3fda141ba4689bd4506", size = 4472642, upload-time = "2025-10-15T23:18:02.749Z" },
     { url = "https://files.pythonhosted.org/packages/3d/39/8e71f3930e40f6877737d6f69248cf74d4e34b886a3967d32f919cc50d3b/cryptography-46.0.3-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ef639cb3372f69ec44915fafcd6698b6cc78fbe0c2ea41be867f6ed612811963", size = 4423126, upload-time = "2025-10-15T23:18:04.85Z" },
     { url = "https://files.pythonhosted.org/packages/cd/c7/f65027c2810e14c3e7268353b1681932b87e5a48e65505d8cc17c99e36ae/cryptography-46.0.3-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:3b51b8ca4f1c6453d8829e1eb7299499ca7f313900dd4d89a24b8b87c0a780d4", size = 4686573, upload-time = "2025-10-15T23:18:06.908Z" },
@@ -473,7 +485,7 @@ dependencies = [
     { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260303+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
     { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "openai-harmony", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -512,7 +524,7 @@ requires-dist = [
     { name = "mflux", specifier = "==0.16.9" },
     { name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" },
     { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.6" },
-    { name = "mlx-lm", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Feval-left-padding-in-batched-rotation" },
+    { name = "mlx-lm", git = "https://github.com/ml-explore/mlx-lm?branch=main" },
     { name = "msgspec", specifier = ">=0.19.0" },
     { name = "openai-harmony", specifier = ">=0.0.8" },
     { name = "psutil", specifier = ">=7.0.0" },
@@ -1351,7 +1363,7 @@ dependencies = [
     { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260303+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1399,7 +1411,7 @@ cuda13 = [
 
 [[package]]
 name = "mlx"
-version = "0.30.7.dev20260303+257d5692"
+version = "0.30.7.dev20260225+257d5692"
 source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }
 resolution-markers = [
     "python_full_version >= '3.14' and sys_platform == 'darwin'",
@@ -1432,11 +1444,11 @@ wheels = [
 
 [[package]]
 name = "mlx-lm"
-version = "0.31.0"
-source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Feval-left-padding-in-batched-rotation#5e0c484cfb5c68d281a71409927ee1bb75adaae2" }
+version = "0.31.2"
+source = { git = "https://github.com/ml-explore/mlx-lm?branch=main#ed7884cb80968e0e77fce6cde5d1597952bbd524" }
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260303+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },