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: ...

Cargo.toml

@@ -1,11 +1,6 @@
 [workspace]
 resolver = "3"
-members = [
-  "rust/networking",
-  "rust/exo_pyo3_bindings",
-  "rust/util",
-  "rust/babblerd",
-]
+members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
 
 [workspace.package]
 version = "0.0.1"
@@ -48,6 +43,7 @@ log = "0.4"
 
 # networking
 libp2p = "0.56"
+libp2p-tcp = "0.44"
 
 [workspace.lints.rust]
 static_mut_refs = "warn"      # Or use "warn" instead of deny

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

bench/test_mlx_bandwidth.py

@@ -1,377 +0,0 @@
-# type: ignore
-import argparse
-import json
-import os
-import statistics
-import sys
-import tempfile
-import time
-
-import mlx.core as mx
-
-DTYPE_MAP = {
-    "float32": (mx.float32, 4),
-    "float16": (mx.float16, 2),
-    "bfloat16": (mx.bfloat16, 2),
-}
-
-SIZES = [
-    1 * 1024,
-    4 * 1024,
-    16 * 1024,
-    64 * 1024,
-    256 * 1024,
-    1 * 1024 * 1024,
-    4 * 1024 * 1024,
-    16 * 1024 * 1024,
-    64 * 1024 * 1024,
-    256 * 1024 * 1024,
-    1 * 1024 * 1024 * 1024,
-    2 * 1024 * 1024 * 1024,
-    4 * 1024 * 1024 * 1024,
-    8 * 1024 * 1024 * 1024,
-]
-
-
-def format_bytes(n: int) -> str:
-    if n >= 1024 * 1024 * 1024:
-        return f"{n / (1024 * 1024 * 1024):.0f} GB"
-    if n >= 1024 * 1024:
-        return f"{n / (1024 * 1024):.0f} MB"
-    if n >= 1024:
-        return f"{n / 1024:.0f} KB"
-    return f"{n} B"
-
-
-def format_time(seconds: float) -> str:
-    if seconds >= 1.0:
-        return f"{seconds:.3f} s"
-    if seconds >= 0.001:
-        return f"{seconds * 1000:.2f} ms"
-    return f"{seconds * 1_000_000:.1f} us"
-
-
-def format_bandwidth(bytes_per_sec: float) -> str:
-    if bytes_per_sec >= 1024 * 1024 * 1024:
-        return f"{bytes_per_sec / (1024 * 1024 * 1024):.2f} GB/s"
-    if bytes_per_sec >= 1024 * 1024:
-        return f"{bytes_per_sec / (1024 * 1024):.1f} MB/s"
-    return f"{bytes_per_sec / 1024:.1f} KB/s"
-
-
-def barrier(group: mx.distributed.Group) -> None:
-    mx.eval(mx.distributed.all_sum(mx.array(1.0), group=group))
-
-
-def init_ring(
-    rank: int, self_ip: str, peer_ip: str, port: int, tmpdir: str
-) -> mx.distributed.Group:
-    if rank == 0:
-        hosts = [f"{self_ip}:{port}", f"{peer_ip}:{port}"]
-    else:
-        hosts = [f"{peer_ip}:{port}", f"{self_ip}:{port}"]
-
-    hostfile = os.path.join(tmpdir, "hosts.json")
-    with open(hostfile, "w") as f:
-        json.dump(hosts, f)
-
-    for var in ("MLX_HOSTFILE", "MLX_RANK", "MLX_IBV_DEVICES", "MLX_JACCL_COORDINATOR"):
-        os.environ.pop(var, None)
-
-    os.environ["MLX_HOSTFILE"] = hostfile
-    os.environ["MLX_RANK"] = str(rank)
-    return mx.distributed.init(backend="ring", strict=True)
-
-
-def init_jaccl(
-    rank: int, interface: str, coordinator: str, port: int, tmpdir: str
-) -> mx.distributed.Group:
-    devices = [[None, interface], [interface, None]]
-    devfile = os.path.join(tmpdir, "devices.json")
-    with open(devfile, "w") as f:
-        json.dump(devices, f)
-
-    for var in ("MLX_HOSTFILE", "MLX_RANK", "MLX_IBV_DEVICES", "MLX_JACCL_COORDINATOR"):
-        os.environ.pop(var, None)
-
-    os.environ["MLX_IBV_DEVICES"] = devfile
-    os.environ["MLX_RANK"] = str(rank)
-    if rank == 0:
-        os.environ["MLX_JACCL_COORDINATOR"] = f"0.0.0.0:{port}"
-    else:
-        os.environ["MLX_JACCL_COORDINATOR"] = coordinator
-
-    return mx.distributed.init(backend="jaccl", strict=True)
-
-
-def bench_unidirectional(
-    group: mx.distributed.Group,
-    rank: int,
-    size_bytes: int,
-    dtype: mx.Dtype,
-    element_size: int,
-    warmup: int,
-    iterations: int,
-) -> list[float]:
-    n_elements = size_bytes // element_size
-    tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
-    mx.eval(tensor)
-
-    for _ in range(warmup):
-        if rank == 0:
-            sent = mx.distributed.send(tensor, dst=1, group=group)
-            mx.eval(sent)
-        else:
-            received = mx.distributed.recv_like(tensor, src=0, group=group)
-            mx.eval(received)
-        barrier(group)
-
-    times: list[float] = []
-    for _ in range(iterations):
-        barrier(group)
-        t0 = time.perf_counter()
-        if rank == 0:
-            sent = mx.distributed.send(tensor, dst=1, group=group)
-            mx.eval(sent)
-        else:
-            received = mx.distributed.recv_like(tensor, src=0, group=group)
-            mx.eval(received)
-        barrier(group)
-        t1 = time.perf_counter()
-        times.append(t1 - t0)
-
-    return times
-
-
-def bench_rtt(
-    group: mx.distributed.Group,
-    rank: int,
-    size_bytes: int,
-    dtype: mx.Dtype,
-    element_size: int,
-    warmup: int,
-    iterations: int,
-) -> list[float]:
-    n_elements = size_bytes // element_size
-    tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
-    mx.eval(tensor)
-
-    for _ in range(warmup):
-        if rank == 0:
-            sent = mx.distributed.send(tensor, dst=1, group=group)
-            mx.eval(sent)
-            received = mx.distributed.recv_like(tensor, src=1, group=group)
-            mx.eval(received)
-        else:
-            received = mx.distributed.recv_like(tensor, src=0, group=group)
-            mx.eval(received)
-            sent = mx.distributed.send(received, dst=0, group=group)
-            mx.eval(sent)
-        barrier(group)
-
-    times: list[float] = []
-    for _ in range(iterations):
-        barrier(group)
-        t0 = time.perf_counter()
-        if rank == 0:
-            sent = mx.distributed.send(tensor, dst=1, group=group)
-            mx.eval(sent)
-            received = mx.distributed.recv_like(tensor, src=1, group=group)
-            mx.eval(received)
-        else:
-            received = mx.distributed.recv_like(tensor, src=0, group=group)
-            mx.eval(received)
-            sent = mx.distributed.send(received, dst=0, group=group)
-            mx.eval(sent)
-        barrier(group)
-        t1 = time.perf_counter()
-        times.append(t1 - t0)
-
-    return times
-
-
-def bench_all_gather(
-    group: mx.distributed.Group,
-    rank: int,
-    size_bytes: int,
-    dtype: mx.Dtype,
-    element_size: int,
-    warmup: int,
-    iterations: int,
-) -> list[float]:
-    n_elements = (size_bytes // 2) // element_size
-    tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
-    mx.eval(tensor)
-
-    for _ in range(warmup):
-        gathered = mx.distributed.all_gather(tensor, group=group)
-        mx.eval(gathered)
-        barrier(group)
-
-    times: list[float] = []
-    for _ in range(iterations):
-        barrier(group)
-        t0 = time.perf_counter()
-        gathered = mx.distributed.all_gather(tensor, group=group)
-        mx.eval(gathered)
-        t1 = time.perf_counter()
-        times.append(t1 - t0)
-
-    return times
-
-
-def print_table(title: str, rows: list[dict[str, str]]) -> None:
-    print(f"\n=== {title} ===")
-    headers = ["Size", "Median", "Min", "Max", "Bandwidth"]
-    widths = [
-        max(len(h), max((len(r[h]) for r in rows), default=0)) + 2 for h in headers
-    ]
-    header_line = "".join(h.ljust(w) for h, w in zip(headers, widths, strict=True))
-    print(header_line)
-    print("-" * len(header_line))
-    for row in rows:
-        print("".join(row[h].ljust(w) for h, w in zip(headers, widths, strict=True)))
-
-
-def run_bench(
-    name: str,
-    bench_fn,
-    group: mx.distributed.Group,
-    rank: int,
-    dtype: mx.Dtype,
-    element_size: int,
-    warmup: int,
-    iterations: int,
-    bw_multiplier: int = 1,
-) -> None:
-    rows: list[dict[str, str]] = []
-    for size in SIZES:
-        if rank == 0:
-            print(f"  {name}: {format_bytes(size)}...", end="", flush=True)
-        times = bench_fn(group, rank, size, dtype, element_size, warmup, iterations)
-        if rank == 0:
-            med = statistics.median(times)
-            mn = min(times)
-            mx_ = max(times)
-            bw = (size * bw_multiplier) / med
-            rows.append(
-                {
-                    "Size": format_bytes(size),
-                    "Median": format_time(med),
-                    "Min": format_time(mn),
-                    "Max": format_time(mx_),
-                    "Bandwidth": format_bandwidth(bw),
-                }
-            )
-            print(f" {format_bandwidth(bw)}")
-    if rank == 0:
-        print_table(name, rows)
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(
-        description="MLX Distributed Communication Benchmark"
-    )
-    subparsers = parser.add_subparsers(dest="backend", required=True)
-
-    ring_parser = subparsers.add_parser("ring")
-    ring_parser.add_argument("--rank", type=int, required=True, choices=[0, 1])
-    ring_parser.add_argument("--self-ip", required=True)
-    ring_parser.add_argument("--peer-ip", required=True)
-    ring_parser.add_argument("--port", type=int, default=5555)
-
-    jaccl_parser = subparsers.add_parser("jaccl")
-    jaccl_parser.add_argument("--rank", type=int, required=True, choices=[0, 1])
-    jaccl_parser.add_argument("--interface", required=True)
-    jaccl_parser.add_argument(
-        "--coordinator",
-        type=str,
-        default=None,
-        help="IP:PORT of rank 0 (required for rank 1)",
-    )
-    jaccl_parser.add_argument(
-        "--port", type=int, default=9999, help="Coordinator port (rank 0 only)"
-    )
-
-    for p in [ring_parser, jaccl_parser]:
-        p.add_argument("--warmup", type=int, default=3)
-        p.add_argument("--iterations", type=int, default=10)
-        p.add_argument("--dtype", choices=list(DTYPE_MAP.keys()), default="float32")
-
-    args = parser.parse_args()
-
-    if args.backend == "jaccl" and args.rank == 1 and args.coordinator is None:
-        jaccl_parser.error("--coordinator is required for rank 1")
-
-    return args
-
-
-def main() -> int:
-    args = parse_args()
-    dtype, element_size = DTYPE_MAP[args.dtype]
-
-    with tempfile.TemporaryDirectory() as tmpdir:
-        if args.backend == "ring":
-            print(f"Initializing ring backend (rank {args.rank})...")
-            group = init_ring(args.rank, args.self_ip, args.peer_ip, args.port, tmpdir)
-        else:
-            print(f"Initializing jaccl backend (rank {args.rank})...")
-            group = init_jaccl(
-                args.rank, args.interface, args.coordinator or "", args.port, tmpdir
-            )
-
-        print(f"Rank {group.rank()} of {group.size()} initialized")
-        barrier(group)
-
-        if args.rank == 0:
-            print("\nMLX Distributed Communication Benchmark")
-            print(
-                f"Backend: {args.backend} | Dtype: {args.dtype} | Warmup: {args.warmup} | Iterations: {args.iterations}"
-            )
-
-        run_bench(
-            "Unidirectional (rank 0 -> rank 1)",
-            bench_unidirectional,
-            group,
-            args.rank,
-            dtype,
-            element_size,
-            args.warmup,
-            args.iterations,
-        )
-        run_bench(
-            "Round-Trip (ping-pong)",
-            bench_rtt,
-            group,
-            args.rank,
-            dtype,
-            element_size,
-            args.warmup,
-            args.iterations,
-            bw_multiplier=2,
-        )
-        run_bench(
-            "All-Gather",
-            bench_all_gather,
-            group,
-            args.rank,
-            dtype,
-            element_size,
-            args.warmup,
-            args.iterations,
-        )
-
-        if args.rank == 0:
-            print("\nDone.")
-        else:
-            print("Rank 1 complete.")
-
-    return 0
-
-
-if __name__ == "__main__":
-    try:
-        sys.exit(main())
-    except KeyboardInterrupt:
-        print("\nInterrupted.")
-        sys.exit(1)

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)

flake.nix

@@ -112,9 +112,7 @@
             };
           };
 
-          packages = {
-            babeld = pkgs.callPackage ./nix/babeld.nix { };
-          } // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
+          packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
             let
               uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
               mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package);
@@ -159,6 +157,9 @@
                 just
                 jq
               ]
+              ++ lib.optionals stdenv.isLinux [
+                unixtools.ifconfig
+              ]
               ++ lib.optionals stdenv.isDarwin [
                 macmon
               ];

nix/babeld.nix

@@ -1,25 +0,0 @@
-{ stdenv
-, lib
-, fetchurl
-}:
-stdenv.mkDerivation rec {
-  pname = "babeld";
-  version = "1.13.1";
-
-  src = fetchurl {
-    url = "https://www.irif.fr/~jch/software/files/${pname}-${version}.tar.gz";
-    hash = "sha256-FfJNJtoMz8Bzq83vAwnygeRoTyqnESb4JlcsTIRejdk=";
-  };
-
-  outputs = [
-    "out"
-    "man"
-  ];
-
-  makeFlags = [
-    "PREFIX=${placeholder "out"}"
-    "ETCDIR=${placeholder "out"}/etc"
-  ]
-  ++ lib.optional stdenv.isDarwin "LDLIBS=''";
-}
-

pyproject.toml

@@ -25,8 +25,7 @@ dependencies = [
   "openai-harmony>=0.0.8",
   "httpx>=0.28.1",
   "tomlkit>=0.14.0",
-  "pillow>=11.0,<12.0",                        # compatibility with mflux
-  "mflux==0.15.5",
+  "mflux==0.16.9",
   "python-multipart>=0.0.21",
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
@@ -62,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 }

python/parts.nix

@@ -185,7 +185,6 @@
         exo-eval = mkBenchScript "exo-eval" (inputs.self + /bench/exo_eval.py);
         exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
         exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
-        exo-mlx-bandwidth-test = mkPythonScript "exo-mlx-bandwidth-test" (inputs.self + /bench/test_mlx_bandwidth.py);
       };
 
       checks = {

result

@@ -1 +0,0 @@
-/nix/store/41rd4g2p69a2106qscl3xvkzmh6j7nry-babblerd

rust/babblerd/Cargo.toml

@@ -1,19 +0,0 @@
-[package]
-name = "babblerd"
-version.workspace = true
-edition.workspace = true
-
-[dependencies]
-color-eyre = "0.6.5"
-futures-lite.workspace = true
-ipnet = "2.12.0"
-libc = "0.2.183"
-n0-watcher = "0.6.1"
-netwatch = "0.14.0"
-rand = "0.10.0"
-tokio = { workspace = true, features = ["full"] }
-tracing = "0.1.44"
-tracing-subscriber = { version = "0.3.22", features = ["env-filter"] }
-
-[lints]
-workspace = true

rust/babblerd/src/lib.rs

@@ -1,548 +0,0 @@
-pub use babel::{babel, handle_listener};
-pub use error::{BabbleError, Result};
-pub use if_watcher::{PREFIX, watch};
-pub mod error {
-    use std::io;
-
-    pub type Result<T> = core::result::Result<T, BabbleError>;
-    #[derive(Debug)]
-    pub enum BabbleError {
-        Io(io::Error),
-        Unspecified,
-        BabeldCrashed(Option<i32>),
-        FailedToSetIp,
-        Other(String),
-    }
-    impl std::fmt::Display for BabbleError {
-        // use the debug display for now
-        #[inline]
-        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-            write!(f, "{self:?}")
-        }
-    }
-    impl std::error::Error for BabbleError {}
-    impl From<io::Error> for BabbleError {
-        #[inline]
-        fn from(value: io::Error) -> Self {
-            Self::Io(value)
-        }
-    }
-}
-pub mod if_watcher {
-    #[cfg(target_os = "linux")]
-    use std::path::PathBuf;
-    use std::{
-        collections::HashSet,
-        net::{IpAddr, Ipv6Addr},
-    };
-
-    use futures_lite::StreamExt;
-    use ipnet::Ipv6Net;
-    use n0_watcher::Watcher;
-    use netwatch::interfaces::{Interface, IpNet};
-    use tokio::sync::mpsc;
-
-    use crate::{
-        BabbleError, Result,
-        babel::Babble,
-        ip_manager::{add_ip, remove_ip},
-    };
-
-    pub const PREFIX: Ipv6Net = Ipv6Net::new_assert(
-        Ipv6Addr::new(0xfde0, 0x20c6, 0x1fa7, 0xffff, 0, 0, 0, 0),
-        48,
-    );
-
-    trait IfaceExt {
-        fn has_link_local_v6(&self) -> bool;
-        fn is_real_interface(&self) -> bool;
-        fn will_babel(&self) -> bool;
-        fn get_v6_in(&self, prefix: Ipv6Net) -> Option<Ipv6Net>;
-    }
-    impl IfaceExt for Interface {
-        fn will_babel(&self) -> bool {
-            self.has_link_local_v6() && self.is_real_interface() && self.is_up()
-        }
-
-        fn has_link_local_v6(&self) -> bool {
-            let mut has = false;
-            for addr in self.addrs() {
-                let IpAddr::V6(a) = addr.addr() else {
-                    continue;
-                };
-                if a.is_unicast_link_local() {
-                    has = true;
-                    break;
-                }
-            }
-            has
-        }
-        fn is_real_interface(&self) -> bool {
-            // macos is weird. en0 & en1 are ethernet & wifi (varies which is which by device). en3+ is thunderbolt, but at some point becomes usb ethernet.
-            if self.name().strip_prefix("en").is_none()
-            //.and_then(|s| s.parse::<u8>().ok())
-            //.is_none_or(|_n| false)
-            {
-                return false;
-            }
-            #[cfg(target_os = "linux")]
-            {
-                if !PathBuf::from(format!("/sys/class/net/{}/device", self.name())).exists() {
-                    tracing::debug!(
-                        "skipping interface {} as it doesn't correspond to a physical link",
-                        self.name()
-                    );
-                    return false;
-                }
-                let dev_type_path = PathBuf::from(format!("/sys/class/net/{}/type", self.name()));
-                if !dev_type_path.exists() {
-                    tracing::debug!(
-                        "skipping interface {} with no type file at {:?}",
-                        self.name(),
-                        dev_type_path.to_str()
-                    );
-                    return false;
-                }
-                let Ok(dev_type) = std::fs::read_to_string(dev_type_path) else {
-                    return false;
-                };
-                if dev_type.trim() != "1" {
-                    tracing::debug!(
-                        "skipping interface {} with type {:?}",
-                        self.name(),
-                        dev_type
-                    );
-                    return false;
-                }
-            }
-            true
-        }
-        fn get_v6_in(&self, prefix: Ipv6Net) -> Option<Ipv6Net> {
-            for addr in self.addrs() {
-                if let IpNet::V6(v6) = addr
-                    && prefix.contains(&v6.addr())
-                {
-                    return Some(v6);
-                }
-            }
-            None
-        }
-    }
-
-    #[tracing::instrument(skip(send))]
-    pub async fn watch(my_range: Ipv6Net, send: mpsc::Sender<Babble>) -> Result<()> {
-        assert!(PREFIX.contains(&my_range));
-        let mut ready_ifaces = HashSet::new();
-        // 0 is reserved for the first loopback address
-        let mut iface_num: u16 = 1;
-
-        tracing::info!("starting interface monitor");
-        let mon = netwatch::netmon::Monitor::new()
-            .await
-            .map_err(|_| BabbleError::Unspecified)?;
-        let mut mon_stream = mon.interface_state().stream();
-
-        while let Some(s) = mon_stream.next().await {
-            for iface in s.interfaces.values() {
-                if !iface.is_real_interface() {
-                    continue;
-                }
-                for addr in iface.addrs() {
-                    if let IpNet::V6(v6) = addr
-                        && PREFIX.contains(&v6.addr())
-                        && !my_range.contains(&v6.addr())
-                    {
-                        tracing::info!("removing stale ip {v6} from {}", iface.name());
-                        // don't really care if this fails
-                        if let Err(e) = remove_ip(v6, iface).await {
-                            tracing::warn!(%e, "failed to remove ip");
-                        }
-                    }
-                }
-                if !iface.will_babel() {
-                    continue;
-                }
-                if iface.get_v6_in(my_range).is_none() {
-                    let addr = Ipv6Net::new_assert(
-                        Ipv6Addr::from_bits(my_range.addr().to_bits() | u128::from(iface_num)),
-                        128,
-                    );
-                    iface_num += 1;
-                    assert!(iface_num < u16::MAX, "Really? u16::MAX interfaces?");
-                    tracing::info!("adding new ip {addr} to {}", iface.name());
-                    add_ip(addr, iface).await?;
-                }
-
-                if ready_ifaces.insert(iface.name().to_owned()) {
-                    tracing::info!("telling babeld to watch {}", iface.name());
-                    let Ok(()) = send.send(Babble::AddIface(iface.name().to_owned())).await else {
-                        return Ok(());
-                    };
-                }
-            }
-        }
-        tracing::info!("stopping interface monitor");
-
-        Ok(())
-    }
-}
-pub mod babel {
-    #[tracing::instrument(skip_all)]
-    pub async fn handle_listener(sock: UnixStream, mut receiver: broadcast::Receiver<String>) {
-        tracing::info!("new socket conn");
-        let (reader, mut write) = sock.into_split();
-        let mut reader = BufReader::new(reader).lines();
-        loop {
-            tokio::select! {
-                read = reader.next_line() => {
-                    let Ok(Some(_)) = read else { break; };
-                }
-                recv = receiver.recv() => {
-                    let res = match recv {
-                        Err(broadcast::error::RecvError::Lagged(_)) => { tracing::warn!("receiver lagged"); continue; },
-                        Err(broadcast::error::RecvError::Closed) => { tracing::debug!("receiver closed, dropping connection"); break; },
-                        Ok(s) => write.write_all(format!("{s}\n").as_bytes()).await,
-                    };
-                    if let Err(e) = res { tracing::warn!(error=%e, "failed to write to socket"); break; }
-                }
-            };
-        }
-        tracing::info!("closing socket conn");
-        _ = write.shutdown().await;
-    }
-
-    #[cfg(target_os = "macos")]
-    const PRIVATE_SOCK_PATH: &str = "/var/run/babbler/private/babeld.sock";
-    #[cfg(target_os = "linux")]
-    const PRIVATE_SOCK_PATH: &str = "/run/babbler/private/babeld.sock";
-    #[cfg(target_os = "macos")]
-    const PRIVATE_DIR: &str = "/var/run/babbler/private";
-    #[cfg(target_os = "linux")]
-    const PRIVATE_DIR: &str = "/run/babbler/private";
-
-    use ipnet::Ipv6Net;
-    use std::fs::Permissions;
-    use std::io;
-    use std::os::unix::fs::PermissionsExt;
-    use tokio::time::Duration;
-
-    use futures_lite::FutureExt;
-    use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, Lines};
-    use tokio::net::UnixStream;
-    use tokio::net::unix::{OwnedReadHalf, OwnedWriteHalf};
-    use tokio::process::Command;
-    use tokio::sync::{broadcast, mpsc};
-
-    use crate::{BabbleError, Result};
-
-    #[derive(Debug)]
-    pub enum Babble {
-        AddIface(String),
-    }
-
-    pub struct BabeldProcess {
-        proc: tokio::process::Child,
-    }
-    impl Drop for BabeldProcess {
-        #[inline]
-        fn drop(&mut self) {
-            // emergency sigkill babeld process to prevent leakage
-            match self.proc.try_wait() {
-                Ok(None) => {}
-                Ok(Some(sc)) => {
-                    if !sc.success() {
-                        _ = self.proc.start_kill();
-                    }
-                }
-                _ => {
-                    _ = self.proc.start_kill();
-                }
-            }
-        }
-    }
-    impl BabeldProcess {
-        #[tracing::instrument]
-        async fn spawn(my_range: Ipv6Net, iface: String) -> Result<Self> {
-            tokio::fs::create_dir_all(PRIVATE_DIR).await?;
-            tokio::fs::set_permissions(PRIVATE_DIR, Permissions::from_mode(0o0700)).await?;
-            tracing::info!("spawning babeld socket in {PRIVATE_SOCK_PATH}");
-            let res = match Command::new("babeld")
-                .arg("-G")
-                .arg(PRIVATE_SOCK_PATH)
-                .arg("-I")
-                .arg(format!("{PRIVATE_DIR}/babeld.pid"))
-                .arg("-C")
-                .arg(format!("redistribute local ip {my_range}"))
-                .arg("-C")
-                .arg("redistribute local deny")
-                .arg(iface)
-                .spawn()
-            {
-                Ok(proc) => Ok(Self { proc }),
-                Err(e) => {
-                    tracing::warn!(error=%e, "failed to spawn babeld");
-                    Err(e.into())
-                }
-            };
-            tokio::time::sleep(Duration::from_millis(10)).await;
-            while !matches!(tokio::fs::try_exists(PRIVATE_SOCK_PATH).await, Ok(true)) {
-                tracing::info!("where is the sock");
-                tokio::time::sleep(Duration::from_secs(1)).await;
-            }
-            res
-        }
-
-        #[tracing::instrument(skip(read, write, send))]
-        async fn query(
-            read: &mut Lines<BufReader<OwnedReadHalf>>,
-            write: &mut OwnedWriteHalf,
-            send: &broadcast::Sender<String>,
-            cmd: &str,
-        ) -> io::Result<Option<bool>> {
-            write.write_all(cmd.as_bytes()).await?;
-            loop {
-                let Some(line) = read.next_line().await? else {
-                    tracing::warn!("babeld closed unexpectedly");
-                    return Ok(None);
-                };
-                tracing::info!("[babel] {:?}", line);
-                let ret = match line.as_str() {
-                    "ok" => Ok(Some(true)),
-                    "bad" => {
-                        tracing::warn!("malformed message sent to babeld");
-                        Ok(Some(false))
-                    }
-                    _ if line.starts_with("no") => {
-                        tracing::warn!("message rejected");
-                        Ok(Some(false))
-                    }
-                    _ => Ok(None),
-                };
-                let Ok(_) = send.send(line) else {
-                    return Ok(None);
-                };
-                if !matches!(ret, Ok(None)) {
-                    return ret;
-                }
-            }
-        }
-
-        #[tracing::instrument(skip_all)]
-        async fn supervise(
-            &self,
-            mut recv: mpsc::Receiver<Babble>,
-            send: broadcast::Sender<String>,
-        ) -> Result<()> {
-            std::fs::set_permissions(PRIVATE_SOCK_PATH, Permissions::from_mode(0o0600))?;
-            tracing::debug!("connecting to babeld.sock");
-            let (reader, mut writer) = UnixStream::connect(PRIVATE_SOCK_PATH).await?.into_split();
-            let mut babel_lines = BufReader::new(reader).lines();
-            while let Some(s) = babel_lines.next_line().await? {
-                tracing::debug!("[babeld] {}", s);
-                if s == "ok" {
-                    break;
-                }
-            }
-            tracing::info!("babeld ok");
-            /* TODO(evan): push rather than pull
-            if Self::query(&mut babel_lines, &mut writer, "monitor\n")
-                .await?
-                .is_none()
-            {
-                return Ok(());
-            };
-            */
-
-            let mut interval = tokio::time::interval(Duration::from_secs(5));
-            interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
-            loop {
-                tokio::select! {
-                    _ = interval.tick() => {
-                        if Self::query(&mut babel_lines, &mut writer, &send, "dump\n")
-                            .await?
-                            .is_none()
-                        {
-                            return Ok(());
-                        }
-                    },
-                    babble = recv.recv() => {
-                        tracing::debug!("[babble] {:?}", babble);
-                        let Some(babble) = babble else {
-                            break;
-                        };
-                        match babble {
-                            Babble::AddIface(iface) => {
-                                Self::query(&mut babel_lines, &mut writer, &send, format!("interface {iface}\n").as_ref()).await?;
-                            }
-                        }
-                    },
-                    line = babel_lines.next_line() => {
-                        let Ok(Some(line)) = line else {
-                            break;
-                        };
-                        tracing::debug!("[babeld] {}", line);
-                        let Ok(_) = send.send(line) else {
-                            break;
-                        };
-                    },
-                }
-            }
-            Ok(())
-        }
-
-        async fn shutdown(mut self) -> Result<()> {
-            let kill_res = if let Some(pid) = self.proc.id() {
-                let pid: libc::pid_t = pid.try_into().expect("pid overflow");
-                // SAFETY: pid >= 0, freshly checked.
-                let rc = unsafe { libc::kill(pid, libc::SIGINT) };
-                let rc_err = if rc != 0 && rc != libc::ESRCH {
-                    Err(io::Error::last_os_error().into())
-                } else {
-                    Ok(())
-                };
-                let exit_code = async { Some(self.proc.wait().await) }
-                    .or(async {
-                        tokio::time::sleep(Duration::from_secs(5)).await;
-                        None
-                    })
-                    .await;
-                match exit_code {
-                    Some(Ok(code)) => {
-                        if code.success() {
-                            rc_err
-                        } else {
-                            rc_err.and_then(|()| Err(BabbleError::BabeldCrashed(code.code())))
-                        }
-                    }
-                    Some(Err(e)) => Err(e.into()),
-                    None => {
-                        self.proc.kill().await?;
-                        rc_err.and(Err(BabbleError::BabeldCrashed(None)))
-                    }
-                }
-            } else {
-                Ok(())
-            };
-            let rem_res = match std::fs::remove_file(PRIVATE_SOCK_PATH) {
-                Ok(()) => Ok(()),
-                Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
-                Err(e) => Err(e.into()),
-            };
-            kill_res.and(rem_res)
-        }
-    }
-
-    #[tracing::instrument(skip(send, recv))]
-    pub async fn babel(
-        my_range: Ipv6Net,
-        mut recv: mpsc::Receiver<Babble>,
-        send: broadcast::Sender<String>,
-    ) -> Result<()> {
-        let iface = loop {
-            match recv.recv().await {
-                Some(Babble::AddIface(iface)) => {
-                    break iface;
-                }
-                None => return Ok(()),
-            }
-        };
-
-        let babel = BabeldProcess::spawn(my_range, iface).await?;
-        let res1 = babel.supervise(recv, send).await;
-        let res2 = babel.shutdown().await;
-        res1.and(res2)
-    }
-}
-pub(crate) mod ip_manager {
-    pub use sys::add_ip;
-    pub use sys::remove_ip;
-
-    #[cfg(target_os = "linux")]
-    mod sys {
-        use ipnet::Ipv6Net;
-        use netwatch::interfaces::Interface;
-
-        use crate::{BabbleError, Result};
-        use tokio::process::Command;
-
-        #[tracing::instrument]
-        pub async fn add_ip(subnet: Ipv6Net, iface: &Interface) -> Result<()> {
-            let out = Command::new("ip")
-                .arg("addr")
-                .arg("add")
-                .arg(format!("{subnet}"))
-                .arg("dev")
-                .arg(iface.name())
-                .output()
-                .await?;
-            if out.status.success() {
-                Ok(())
-            } else {
-                Err(BabbleError::FailedToSetIp)
-            }
-        }
-
-        #[tracing::instrument]
-        pub async fn remove_ip(v6: Ipv6Net, iface: &Interface) -> Result<()> {
-            let out = Command::new("ip")
-                .arg("addr")
-                .arg("del")
-                .arg(format!("{v6}"))
-                .arg("dev")
-                .arg(iface.name())
-                .output()
-                .await?;
-            if out.status.success() {
-                Ok(())
-            } else {
-                let std_err = String::from_utf8_lossy(&out.stdout);
-                tracing::debug!(%std_err);
-                Err(BabbleError::FailedToSetIp)
-            }
-        }
-    }
-
-    #[cfg(target_os = "macos")]
-    mod sys {
-        use ipnet::Ipv6Net;
-        use netwatch::interfaces::Interface;
-
-        use crate::BabbleError;
-        use crate::Result;
-        use tokio::process::Command;
-
-        #[tracing::instrument]
-        pub async fn add_ip(subnet: Ipv6Net, iface: &Interface) -> Result<()> {
-            let out = Command::new("ifconfig")
-                .arg(iface.name())
-                .arg("inet6")
-                .arg(format!("{subnet}"))
-                .arg("add")
-                .output()
-                .await?;
-            if out.status.success() {
-                Ok(())
-            } else {
-                Err(BabbleError::FailedToSetIp)
-            }
-        }
-
-        #[tracing::instrument]
-        pub async fn remove_ip(v6: Ipv6Net, iface: &Interface) -> Result<()> {
-            let out = Command::new("ifconfig")
-                .arg(iface.name())
-                .arg("inet6")
-                .arg(format!("{v6}"))
-                .arg("delete")
-                .output()
-                .await?;
-            if out.status.success() {
-                Ok(())
-            } else {
-                let std_err = String::from_utf8_lossy(&out.stdout);
-                tracing::debug!(%std_err);
-                Err(BabbleError::FailedToSetIp)
-            }
-        }
-    }
-}

rust/babblerd/src/main.rs

@@ -1,167 +0,0 @@
-use std::{fs::Permissions, io, net::Ipv6Addr, os::unix::fs::PermissionsExt};
-
-use babblerd::babel::handle_listener;
-use color_eyre::eyre::{WrapErr, eyre};
-use ipnet::Ipv6Net;
-use tokio::{
-    net::UnixListener,
-    signal,
-    sync::{broadcast, mpsc},
-    task::{JoinHandle, JoinSet},
-};
-
-#[cfg(target_os = "macos")]
-const PUBLIC_DIR: &str = "/var/run/babbler";
-#[cfg(target_os = "linux")]
-const PUBLIC_DIR: &str = "/run/babbler";
-#[cfg(target_os = "macos")]
-const PUBLIC_SOCK_PATH: &str = "/var/run/babbler/babblerd.sock";
-#[cfg(target_os = "linux")]
-const PUBLIC_SOCK_PATH: &str = "/run/babbler/babblerd.sock";
-
-enum State {
-    Idle,
-    Active {
-        recv: broadcast::Receiver<String>,
-        babel: JoinHandle<babblerd::Result<()>>,
-        watcher: JoinHandle<babblerd::Result<()>>,
-        listeners: JoinSet<()>,
-    },
-}
-
-#[tokio::main]
-async fn main() -> color_eyre::Result<()> {
-    color_eyre::install()?;
-    tracing_subscriber::fmt()
-        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
-        .init();
-    // cleanup old data
-    match std::fs::remove_file(PUBLIC_SOCK_PATH) {
-        Err(e) if e.kind() != io::ErrorKind::NotFound => {
-            return Err(e.into());
-        }
-        Ok(()) => {
-            tracing::info!("cleaned up old file at {PUBLIC_SOCK_PATH}");
-        }
-        _ => {}
-    }
-    std::fs::create_dir_all(PUBLIC_DIR)?;
-    // make our directory world readable
-    if let Err(e) = std::fs::set_permissions(PUBLIC_DIR, Permissions::from_mode(0o0755)) {
-        if e.kind() == io::ErrorKind::PermissionDenied {
-            return Err(eyre!(
-                "Insufficient permissions to run daemon -- did you forget sudo?"
-            ));
-        }
-        return Err(e.into());
-    }
-
-    let res = inner_main().await;
-    _ = std::fs::remove_file(PUBLIC_SOCK_PATH);
-    res
-}
-
-async fn inner_main() -> color_eyre::Result<()> {
-    tracing::info!("creating socket at {PUBLIC_SOCK_PATH}");
-    let public_socket = UnixListener::bind(PUBLIC_SOCK_PATH)?;
-    // make our socket world accessible
-    std::fs::set_permissions(PUBLIC_SOCK_PATH, Permissions::from_mode(0o0666))?;
-    let mut babbler: State = State::Idle;
-
-    loop {
-        babbler = match babbler {
-            State::Idle => {
-                tracing::info!("waiting for connection");
-                tokio::select! {
-                    sig = signal::ctrl_c() => {
-                        sig?;
-                        break;
-                    }
-                    sock = public_socket.accept() => {
-                        let sock = sock?.0;
-                        tracing::info!("starting babeld");
-                        let (br_send, br_recv) = broadcast::channel(1024);
-                        let (mp_send, mp_recv) = mpsc::channel(32);
-                        // node id is a PREFIX/64, NODE_ID/48 and an INTERFACE_ID/16
-                        let ip_node_id = u128::from(rand::random::<u64>() & (u64::MAX>>16)) << 16;
-                        let my_range = Ipv6Net::new_assert(
-                            Ipv6Addr::from_bits(babblerd::PREFIX.addr().to_bits() | ip_node_id),
-                            112,
-                        );
-                        let babel = tokio::spawn(babblerd::babel(my_range, mp_recv, br_send));
-                        let watcher = tokio::spawn(babblerd::watch(my_range, mp_send));
-                        let mut listeners = JoinSet::new();
-                        listeners.spawn(handle_listener(sock, br_recv.resubscribe()));
-                        State::Active { recv: br_recv, babel, watcher, listeners }
-                    }
-                }
-            }
-            State::Active {
-                recv,
-                mut babel,
-                mut watcher,
-                mut listeners,
-            } => {
-                tokio::select! {
-                    sig = signal::ctrl_c() => {
-                        sig?;
-                        drop(recv);
-                        watcher.abort();
-                        if let Ok(e) = watcher.await {
-                            e.wrap_err("while ctrl-c")?;
-                        }
-                        babel.await?.wrap_err("while ctrl-c")?;
-                        while let Some(res) = listeners.join_next().await {
-                            res.wrap_err("while ctrl-c")?;
-                        }
-                        break;
-                    }
-                    next_join_result = listeners.join_next(), if !listeners.is_empty() => {
-                        next_join_result.expect("checked")?;
-                        tracing::info!("dropped a listener");
-                        if listeners.is_empty() {
-                            tracing::info!("stopping babeld");
-                            drop(recv);
-
-                            watcher.abort();
-                            if let Ok(e) = watcher.await {
-                                e.wrap_err("while closing listeners")?;
-                            }
-                            babel.await?.wrap_err("while closing listeners")?;
-
-                            State::Idle
-                        } else {
-                            State::Active { recv, babel, watcher, listeners }
-                        }
-                    }
-                    sock = public_socket.accept() => {
-                        listeners.spawn(handle_listener(sock?.0, recv.resubscribe()));
-                        State::Active { recv, babel, watcher, listeners }
-                    }
-                    res = &mut watcher => {
-                        res??;
-                        drop(recv);
-                        babel.await?.wrap_err("while closing watcher")?;
-                        while let Some(res2) = listeners.join_next().await {
-                            res2.wrap_err("while closing watcher")?;
-                        }
-                        State::Idle
-                    }
-                    res = &mut babel => {
-                        res??;
-                        drop(recv);
-                        watcher.abort();
-                        if let Ok(e) = watcher.await {
-                            e.wrap_err("while closing babeld")?;
-                        }
-                        while let Some(res2) = listeners.join_next().await {
-                            res2.wrap_err("while closing babeld")?;
-                        }
-                        State::Idle
-                    }
-                }
-            }
-        };
-    }
-    Ok(())
-}

rust/parts.nix

@@ -1,7 +1,7 @@
 { inputs, ... }:
 {
   perSystem =
-    { inputs', self', pkgs, lib, ... }:
+    { inputs', pkgs, lib, ... }:
     let
       # Fenix nightly toolchain with all components
       rustToolchain = inputs'.fenix.packages.stable.withComponents [
@@ -79,7 +79,7 @@
       };
 
       config = {
-        packages = rec {
+        packages = {
           # Python bindings wheel via maturin
           exo_pyo3_bindings = craneLib.buildPackage (
             commonArgs
@@ -110,19 +110,6 @@
               '';
             }
           );
-          babblerd-unwrapped = craneLib.buildPackage (
-            commonArgs // {
-              inherit cargoArtifacts;
-              pname = "babblerd-unwrapped";
-            }
-          );
-          babblerd = pkgs.writeShellApplication {
-            name = "babblerd";
-            runtimeInputs = [ self'.packages.babeld ];
-            text = ''
-              exec ${babblerd-unwrapped}/bin/babblerd "$@"
-            '';
-          };
         };
 
         checks = {

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

@@ -489,7 +489,6 @@ dependencies = [
     { 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'" },
-    { name = "pillow", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "psutil", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "pydantic", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "python-multipart", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -522,13 +521,12 @@ requires-dist = [
     { name = "huggingface-hub", specifier = ">=0.33.4" },
     { name = "hypercorn", specifier = ">=0.18.0" },
     { name = "loguru", specifier = ">=0.7.3" },
-    { name = "mflux", specifier = "==0.15.5" },
+    { 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 = "pillow", specifier = ">=11.0,<12.0" },
     { name = "psutil", specifier = ">=7.0.0" },
     { name = "pydantic", specifier = ">=2.11.7" },
     { name = "python-multipart", specifier = ">=0.0.21" },
@@ -765,6 +763,36 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/69/b2/119f6e6dcbd96f9069ce9a2665e0146588dc9f88f29549711853645e736a/h2-4.3.0-py3-none-any.whl", hash = "sha256:c438f029a25f7945c69e0ccf0fb951dc3f73a5f6412981daee861431b70e2bdd", size = 61779, upload-time = "2025-08-23T18:12:17.779Z" },
 ]
 
+[[package]]
+name = "hf-transfer"
+version = "0.1.9"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/1a/eb/8fc64f40388c29ce8ce3b2b180a089d4d6b25b1d0d232d016704cb852104/hf_transfer-0.1.9.tar.gz", hash = "sha256:035572865dab29d17e783fbf1e84cf1cb24f3fcf8f1b17db1cfc7fdf139f02bf", size = 25201, upload-time = "2025-01-07T10:05:12.947Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/a4/78/0dce00208f585fae675f40033ef9a30dedfa83665d5ac79f16beb4a0a6c2/hf_transfer-0.1.9-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:6e94e8822da79573c9b6ae4d6b2f847c59a7a06c5327d7db20751b68538dc4f6", size = 1386084, upload-time = "2025-01-07T10:04:47.874Z" },
+    { url = "https://files.pythonhosted.org/packages/ea/2e/3d60b1a9e9f29a2152aa66c823bf5e399ae7be3fef310ff0de86779c5d2d/hf_transfer-0.1.9-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3ebc4ab9023414880c8b1d3c38174d1c9989eb5022d37e814fa91a3060123eb0", size = 1343558, upload-time = "2025-01-07T10:04:42.313Z" },
+    { url = "https://files.pythonhosted.org/packages/fb/38/130a5ac3747f104033591bcac1c961cb1faadfdc91704f59b09c0b465ff2/hf_transfer-0.1.9-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8674026f21ed369aa2a0a4b46000aca850fc44cd2b54af33a172ce5325b4fc82", size = 3726676, upload-time = "2025-01-07T10:04:11.539Z" },
+    { url = "https://files.pythonhosted.org/packages/15/a1/f4e27c5ad17aac616ae0849e2aede5aae31db8267a948c6b3eeb9fd96446/hf_transfer-0.1.9-cp313-cp313t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3a736dfbb2c84f5a2c975478ad200c0c8bfcb58a25a35db402678fb87ce17fa4", size = 3062920, upload-time = "2025-01-07T10:04:16.297Z" },
+    { url = "https://files.pythonhosted.org/packages/8d/0d/727abdfba39bc3f1132cfa4c970588c2c0bb0d82fe2d645cc10f4e2f8e0b/hf_transfer-0.1.9-cp313-cp313t-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:504b8427fd785dd8546d53b9fafe6e436bd7a3adf76b9dce556507650a7b4567", size = 3578681, upload-time = "2025-01-07T10:04:29.702Z" },
+    { url = "https://files.pythonhosted.org/packages/50/d0/2b213eb1ea8b1252ccaf1a6c804d0aba03fea38aae4124df6a3acb70511a/hf_transfer-0.1.9-cp313-cp313t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2c7fc1b85f4d0f76e452765d7648c9f4bfd0aedb9ced2ae1ebfece2d8cfaf8e2", size = 3398837, upload-time = "2025-01-07T10:04:22.778Z" },
+    { url = "https://files.pythonhosted.org/packages/8c/8a/79dbce9006e0bd6b74516f97451a7b7c64dbbb426df15d901dd438cfeee3/hf_transfer-0.1.9-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0d991376f0eac70a60f0cbc95602aa708a6f7c8617f28b4945c1431d67b8e3c8", size = 3546986, upload-time = "2025-01-07T10:04:36.415Z" },
+    { url = "https://files.pythonhosted.org/packages/a9/f7/9ac239b6ee6fe0bad130325d987a93ea58c4118e50479f0786f1733b37e8/hf_transfer-0.1.9-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:e6ac4eddcd99575ed3735ed911ddf9d1697e2bd13aa3f0ad7e3904dd4863842e", size = 4071715, upload-time = "2025-01-07T10:04:53.224Z" },
+    { url = "https://files.pythonhosted.org/packages/d8/a3/0ed697279f5eeb7a40f279bd783cf50e6d0b91f24120dcf66ef2cf8822b4/hf_transfer-0.1.9-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:57fd9880da1ee0f47250f735f791fab788f0aa1ee36afc49f761349869c8b4d9", size = 3388081, upload-time = "2025-01-07T10:04:57.818Z" },
+    { url = "https://files.pythonhosted.org/packages/dc/eb/47e477bdf1d784f31c7540db6cc8c354b777e51a186897a7abda34517f36/hf_transfer-0.1.9-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:5d561f0520f493c66b016d99ceabe69c23289aa90be38dd802d2aef279f15751", size = 3658654, upload-time = "2025-01-07T10:05:03.168Z" },
+    { url = "https://files.pythonhosted.org/packages/45/07/6661e43fbee09594a8a5e9bb778107d95fe38dac4c653982afe03d32bd4d/hf_transfer-0.1.9-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:a5b366d34cd449fe9b20ef25941e6eef0460a2f74e7389f02e673e1f88ebd538", size = 3690551, upload-time = "2025-01-07T10:05:09.238Z" },
+    { url = "https://files.pythonhosted.org/packages/81/f5/461d2e5f307e5048289b1168d5c642ae3bb2504e88dff1a38b92ed990a21/hf_transfer-0.1.9-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:e66acf91df4a8b72f60223059df3003062a5ae111757187ed1a06750a30e911b", size = 1393046, upload-time = "2025-01-07T10:04:51.003Z" },
+    { url = "https://files.pythonhosted.org/packages/41/ba/8d9fd9f1083525edfcb389c93738c802f3559cb749324090d7109c8bf4c2/hf_transfer-0.1.9-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:8669dbcc7a3e2e8d61d42cd24da9c50d57770bd74b445c65123291ca842a7e7a", size = 1348126, upload-time = "2025-01-07T10:04:45.712Z" },
+    { url = "https://files.pythonhosted.org/packages/8e/a2/cd7885bc9959421065a6fae0fe67b6c55becdeda4e69b873e52976f9a9f0/hf_transfer-0.1.9-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8fd0167c4407a3bc4cdd0307e65ada2294ec04f1813d8a69a5243e379b22e9d8", size = 3728604, upload-time = "2025-01-07T10:04:14.173Z" },
+    { url = "https://files.pythonhosted.org/packages/f6/2e/a072cf196edfeda3310c9a5ade0a0fdd785e6154b3ce24fc738c818da2a7/hf_transfer-0.1.9-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:ee8b10afedcb75f71091bcc197c526a6ebf5c58bbbadb34fdeee6160f55f619f", size = 3064995, upload-time = "2025-01-07T10:04:18.663Z" },
+    { url = "https://files.pythonhosted.org/packages/c2/84/aec9ef4c0fab93c1ea2b1badff38c78b4b2f86f0555b26d2051dbc920cde/hf_transfer-0.1.9-cp38-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:5828057e313de59300dd1abb489444bc452efe3f479d3c55b31a8f680936ba42", size = 3580908, upload-time = "2025-01-07T10:04:32.834Z" },
+    { url = "https://files.pythonhosted.org/packages/29/63/b560d39651a56603d64f1a0212d0472a44cbd965db2fa62b99d99cb981bf/hf_transfer-0.1.9-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:fc6bd19e1cc177c66bdef15ef8636ad3bde79d5a4f608c158021153b4573509d", size = 3400839, upload-time = "2025-01-07T10:04:26.122Z" },
+    { url = "https://files.pythonhosted.org/packages/d6/d8/f87ea6f42456254b48915970ed98e993110521e9263472840174d32c880d/hf_transfer-0.1.9-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cdca9bfb89e6f8f281890cc61a8aff2d3cecaff7e1a4d275574d96ca70098557", size = 3552664, upload-time = "2025-01-07T10:04:40.123Z" },
+    { url = "https://files.pythonhosted.org/packages/d6/56/1267c39b65fc8f4e2113b36297320f102718bf5799b544a6cbe22013aa1d/hf_transfer-0.1.9-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:89a23f58b7b7effbc047b8ca286f131b17728c99a9f972723323003ffd1bb916", size = 4073732, upload-time = "2025-01-07T10:04:55.624Z" },
+    { url = "https://files.pythonhosted.org/packages/82/1a/9c748befbe3decf7cb415e34f8a0c3789a0a9c55910dea73d581e48c0ce5/hf_transfer-0.1.9-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:dc7fff1345980d6c0ebb92c811d24afa4b98b3e07ed070c8e38cc91fd80478c5", size = 3390096, upload-time = "2025-01-07T10:04:59.98Z" },
+    { url = "https://files.pythonhosted.org/packages/72/85/4c03da147b6b4b7cb12e074d3d44eee28604a387ed0eaf7eaaead5069c57/hf_transfer-0.1.9-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:1a6bd16c667ebe89a069ca163060127a794fa3a3525292c900b8c8cc47985b0d", size = 3664743, upload-time = "2025-01-07T10:05:05.416Z" },
+    { url = "https://files.pythonhosted.org/packages/e7/6e/e597b04f753f1b09e6893075d53a82a30c13855cbaa791402695b01e369f/hf_transfer-0.1.9-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:d2fde99d502093ade3ab1b53f80da18480e9902aa960dab7f74fb1b9e5bc5746", size = 3695243, upload-time = "2025-01-07T10:05:11.411Z" },
+]
+
 [[package]]
 name = "hf-xet"
 version = "1.2.1rc0"
@@ -1326,11 +1354,12 @@ wheels = [
 
 [[package]]
 name = "mflux"
-version = "0.15.5"
+version = "0.16.9"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "fonttools", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+    { name = "hf-transfer", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { 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'" },
@@ -1352,9 +1381,9 @@ dependencies = [
     { name = "twine", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "urllib3", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/35/8e/f20de51bf9dc0a986535d9a825db4ae314163421b3d3ddaa90a2b959b9fd/mflux-0.15.5.tar.gz", hash = "sha256:9a3372bd64d51c4caff4ff9e7d7d698bea5833242fd849c59cbb0c92f7d7aa3b", size = 743700, upload-time = "2026-01-26T12:41:45.272Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/e3/dc/73dfccae395df22623fbb6f748657700fd9372be22904a0dafb825dc42d2/mflux-0.16.9.tar.gz", hash = "sha256:ff35e9386b5f026a6b97c786b3426e5341105de3356e55ec1b0c5951ff0ac8c8", size = 764296, upload-time = "2026-03-07T11:54:45.031Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/ac/bb/ef936eae2ae78a47cd92ddffc18fc06ad3fd5f438a0915fb62d8bb9508ec/mflux-0.15.5-py3-none-any.whl", hash = "sha256:c94891d4a518047a818863bb099c755e93af90c524ced358baf5b31502c09e82", size = 990939, upload-time = "2026-01-26T12:41:42.898Z" },
+    { url = "https://files.pythonhosted.org/packages/47/69/091df2d0f4a2677d5e2a8b326d568bd884d7ca37f15d3ff78e574aca2271/mflux-0.16.9-py3-none-any.whl", hash = "sha256:988e925653a024e81b46401fd5a2e423dd283428ba64949af4d9257e7df8ddd9", size = 1018804, upload-time = "2026-03-07T11:54:43.47Z" },
 ]
 
 [[package]]
@@ -1415,8 +1444,8 @@ 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.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
@@ -1935,45 +1964,48 @@ wheels = [
 
 [[package]]
 name = "pillow"
-version = "11.3.0"
+version = "12.1.1"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/f3/0d/d0d6dea55cd152ce3d6767bb38a8fc10e33796ba4ba210cbab9354b6d238/pillow-11.3.0.tar.gz", hash = "sha256:3828ee7586cd0b2091b6209e5ad53e20d0649bbe87164a459d0676e035e8f523", size = 47113069, upload-time = "2025-07-01T09:16:30.666Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/1f/42/5c74462b4fd957fcd7b13b04fb3205ff8349236ea74c7c375766d6c82288/pillow-12.1.1.tar.gz", hash = "sha256:9ad8fa5937ab05218e2b6a4cff30295ad35afd2f83ac592e68c0d871bb0fdbc4", size = 46980264, upload-time = "2026-02-11T04:23:07.146Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/1e/93/0952f2ed8db3a5a4c7a11f91965d6184ebc8cd7cbb7941a260d5f018cd2d/pillow-11.3.0-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:1c627742b539bba4309df89171356fcb3cc5a9178355b2727d1b74a6cf155fbd", size = 2128328, upload-time = "2025-07-01T09:14:35.276Z" },
-    { url = "https://files.pythonhosted.org/packages/4b/e8/100c3d114b1a0bf4042f27e0f87d2f25e857e838034e98ca98fe7b8c0a9c/pillow-11.3.0-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:30b7c02f3899d10f13d7a48163c8969e4e653f8b43416d23d13d1bbfdc93b9f8", size = 2170652, upload-time = "2025-07-01T09:14:37.203Z" },
-    { url = "https://files.pythonhosted.org/packages/aa/86/3f758a28a6e381758545f7cdb4942e1cb79abd271bea932998fc0db93cb6/pillow-11.3.0-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:7859a4cc7c9295f5838015d8cc0a9c215b77e43d07a25e460f35cf516df8626f", size = 2227443, upload-time = "2025-07-01T09:14:39.344Z" },
-    { url = "https://files.pythonhosted.org/packages/01/f4/91d5b3ffa718df2f53b0dc109877993e511f4fd055d7e9508682e8aba092/pillow-11.3.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:ec1ee50470b0d050984394423d96325b744d55c701a439d2bd66089bff963d3c", size = 5278474, upload-time = "2025-07-01T09:14:41.843Z" },
-    { url = "https://files.pythonhosted.org/packages/f9/0e/37d7d3eca6c879fbd9dba21268427dffda1ab00d4eb05b32923d4fbe3b12/pillow-11.3.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7db51d222548ccfd274e4572fdbf3e810a5e66b00608862f947b163e613b67dd", size = 4686038, upload-time = "2025-07-01T09:14:44.008Z" },
-    { url = "https://files.pythonhosted.org/packages/ff/b0/3426e5c7f6565e752d81221af9d3676fdbb4f352317ceafd42899aaf5d8a/pillow-11.3.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2d6fcc902a24ac74495df63faad1884282239265c6839a0a6416d33faedfae7e", size = 5864407, upload-time = "2025-07-03T13:10:15.628Z" },
-    { url = "https://files.pythonhosted.org/packages/fc/c1/c6c423134229f2a221ee53f838d4be9d82bab86f7e2f8e75e47b6bf6cd77/pillow-11.3.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f0f5d8f4a08090c6d6d578351a2b91acf519a54986c055af27e7a93feae6d3f1", size = 7639094, upload-time = "2025-07-03T13:10:21.857Z" },
-    { url = "https://files.pythonhosted.org/packages/ba/c9/09e6746630fe6372c67c648ff9deae52a2bc20897d51fa293571977ceb5d/pillow-11.3.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c37d8ba9411d6003bba9e518db0db0c58a680ab9fe5179f040b0463644bc9805", size = 5973503, upload-time = "2025-07-01T09:14:45.698Z" },
-    { url = "https://files.pythonhosted.org/packages/d5/1c/a2a29649c0b1983d3ef57ee87a66487fdeb45132df66ab30dd37f7dbe162/pillow-11.3.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:13f87d581e71d9189ab21fe0efb5a23e9f28552d5be6979e84001d3b8505abe8", size = 6642574, upload-time = "2025-07-01T09:14:47.415Z" },
-    { url = "https://files.pythonhosted.org/packages/36/de/d5cc31cc4b055b6c6fd990e3e7f0f8aaf36229a2698501bcb0cdf67c7146/pillow-11.3.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:023f6d2d11784a465f09fd09a34b150ea4672e85fb3d05931d89f373ab14abb2", size = 6084060, upload-time = "2025-07-01T09:14:49.636Z" },
-    { url = "https://files.pythonhosted.org/packages/d5/ea/502d938cbaeec836ac28a9b730193716f0114c41325db428e6b280513f09/pillow-11.3.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:45dfc51ac5975b938e9809451c51734124e73b04d0f0ac621649821a63852e7b", size = 6721407, upload-time = "2025-07-01T09:14:51.962Z" },
-    { url = "https://files.pythonhosted.org/packages/dd/80/a8a2ac21dda2e82480852978416cfacd439a4b490a501a288ecf4fe2532d/pillow-11.3.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:4c834a3921375c48ee6b9624061076bc0a32a60b5532b322cc0ea64e639dd50e", size = 5281110, upload-time = "2025-07-01T09:14:59.79Z" },
-    { url = "https://files.pythonhosted.org/packages/44/d6/b79754ca790f315918732e18f82a8146d33bcd7f4494380457ea89eb883d/pillow-11.3.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:5e05688ccef30ea69b9317a9ead994b93975104a677a36a8ed8106be9260aa6d", size = 4689547, upload-time = "2025-07-01T09:15:01.648Z" },
-    { url = "https://files.pythonhosted.org/packages/49/20/716b8717d331150cb00f7fdd78169c01e8e0c219732a78b0e59b6bdb2fd6/pillow-11.3.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1019b04af07fc0163e2810167918cb5add8d74674b6267616021ab558dc98ced", size = 5901554, upload-time = "2025-07-03T13:10:27.018Z" },
-    { url = "https://files.pythonhosted.org/packages/74/cf/a9f3a2514a65bb071075063a96f0a5cf949c2f2fce683c15ccc83b1c1cab/pillow-11.3.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f944255db153ebb2b19c51fe85dd99ef0ce494123f21b9db4877ffdfc5590c7c", size = 7669132, upload-time = "2025-07-03T13:10:33.01Z" },
-    { url = "https://files.pythonhosted.org/packages/98/3c/da78805cbdbee9cb43efe8261dd7cc0b4b93f2ac79b676c03159e9db2187/pillow-11.3.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1f85acb69adf2aaee8b7da124efebbdb959a104db34d3a2cb0f3793dbae422a8", size = 6005001, upload-time = "2025-07-01T09:15:03.365Z" },
-    { url = "https://files.pythonhosted.org/packages/6c/fa/ce044b91faecf30e635321351bba32bab5a7e034c60187fe9698191aef4f/pillow-11.3.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:05f6ecbeff5005399bb48d198f098a9b4b6bdf27b8487c7f38ca16eeb070cd59", size = 6668814, upload-time = "2025-07-01T09:15:05.655Z" },
-    { url = "https://files.pythonhosted.org/packages/7b/51/90f9291406d09bf93686434f9183aba27b831c10c87746ff49f127ee80cb/pillow-11.3.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:a7bc6e6fd0395bc052f16b1a8670859964dbd7003bd0af2ff08342eb6e442cfe", size = 6113124, upload-time = "2025-07-01T09:15:07.358Z" },
-    { url = "https://files.pythonhosted.org/packages/cd/5a/6fec59b1dfb619234f7636d4157d11fb4e196caeee220232a8d2ec48488d/pillow-11.3.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:83e1b0161c9d148125083a35c1c5a89db5b7054834fd4387499e06552035236c", size = 6747186, upload-time = "2025-07-01T09:15:09.317Z" },
-    { url = "https://files.pythonhosted.org/packages/73/f4/04905af42837292ed86cb1b1dabe03dce1edc008ef14c473c5c7e1443c5d/pillow-11.3.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:d9da3df5f9ea2a89b81bb6087177fb1f4d1c7146d583a3fe5c672c0d94e55e12", size = 5278520, upload-time = "2025-07-01T09:15:17.429Z" },
-    { url = "https://files.pythonhosted.org/packages/41/b0/33d79e377a336247df6348a54e6d2a2b85d644ca202555e3faa0cf811ecc/pillow-11.3.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:0b275ff9b04df7b640c59ec5a3cb113eefd3795a8df80bac69646ef699c6981a", size = 4686116, upload-time = "2025-07-01T09:15:19.423Z" },
-    { url = "https://files.pythonhosted.org/packages/49/2d/ed8bc0ab219ae8768f529597d9509d184fe8a6c4741a6864fea334d25f3f/pillow-11.3.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0743841cabd3dba6a83f38a92672cccbd69af56e3e91777b0ee7f4dba4385632", size = 5864597, upload-time = "2025-07-03T13:10:38.404Z" },
-    { url = "https://files.pythonhosted.org/packages/b5/3d/b932bb4225c80b58dfadaca9d42d08d0b7064d2d1791b6a237f87f661834/pillow-11.3.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2465a69cf967b8b49ee1b96d76718cd98c4e925414ead59fdf75cf0fd07df673", size = 7638246, upload-time = "2025-07-03T13:10:44.987Z" },
-    { url = "https://files.pythonhosted.org/packages/09/b5/0487044b7c096f1b48f0d7ad416472c02e0e4bf6919541b111efd3cae690/pillow-11.3.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:41742638139424703b4d01665b807c6468e23e699e8e90cffefe291c5832b027", size = 5973336, upload-time = "2025-07-01T09:15:21.237Z" },
-    { url = "https://files.pythonhosted.org/packages/a8/2d/524f9318f6cbfcc79fbc004801ea6b607ec3f843977652fdee4857a7568b/pillow-11.3.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:93efb0b4de7e340d99057415c749175e24c8864302369e05914682ba642e5d77", size = 6642699, upload-time = "2025-07-01T09:15:23.186Z" },
-    { url = "https://files.pythonhosted.org/packages/6f/d2/a9a4f280c6aefedce1e8f615baaa5474e0701d86dd6f1dede66726462bbd/pillow-11.3.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7966e38dcd0fa11ca390aed7c6f20454443581d758242023cf36fcb319b1a874", size = 6083789, upload-time = "2025-07-01T09:15:25.1Z" },
-    { url = "https://files.pythonhosted.org/packages/fe/54/86b0cd9dbb683a9d5e960b66c7379e821a19be4ac5810e2e5a715c09a0c0/pillow-11.3.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:98a9afa7b9007c67ed84c57c9e0ad86a6000da96eaa638e4f8abe5b65ff83f0a", size = 6720386, upload-time = "2025-07-01T09:15:27.378Z" },
-    { url = "https://files.pythonhosted.org/packages/1a/33/c88376898aff369658b225262cd4f2659b13e8178e7534df9e6e1fa289f6/pillow-11.3.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:59a03cdf019efbfeeed910bf79c7c93255c3d54bc45898ac2a4140071b02b4ae", size = 5281129, upload-time = "2025-07-01T09:15:35.194Z" },
-    { url = "https://files.pythonhosted.org/packages/1f/70/d376247fb36f1844b42910911c83a02d5544ebd2a8bad9efcc0f707ea774/pillow-11.3.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:f8a5827f84d973d8636e9dc5764af4f0cf2318d26744b3d902931701b0d46653", size = 4689580, upload-time = "2025-07-01T09:15:37.114Z" },
-    { url = "https://files.pythonhosted.org/packages/eb/1c/537e930496149fbac69efd2fc4329035bbe2e5475b4165439e3be9cb183b/pillow-11.3.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ee92f2fd10f4adc4b43d07ec5e779932b4eb3dbfbc34790ada5a6669bc095aa6", size = 5902860, upload-time = "2025-07-03T13:10:50.248Z" },
-    { url = "https://files.pythonhosted.org/packages/bd/57/80f53264954dcefeebcf9dae6e3eb1daea1b488f0be8b8fef12f79a3eb10/pillow-11.3.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c96d333dcf42d01f47b37e0979b6bd73ec91eae18614864622d9b87bbd5bbf36", size = 7670694, upload-time = "2025-07-03T13:10:56.432Z" },
-    { url = "https://files.pythonhosted.org/packages/70/ff/4727d3b71a8578b4587d9c276e90efad2d6fe0335fd76742a6da08132e8c/pillow-11.3.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4c96f993ab8c98460cd0c001447bff6194403e8b1d7e149ade5f00594918128b", size = 6005888, upload-time = "2025-07-01T09:15:39.436Z" },
-    { url = "https://files.pythonhosted.org/packages/05/ae/716592277934f85d3be51d7256f3636672d7b1abfafdc42cf3f8cbd4b4c8/pillow-11.3.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:41342b64afeba938edb034d122b2dda5db2139b9a4af999729ba8818e0056477", size = 6670330, upload-time = "2025-07-01T09:15:41.269Z" },
-    { url = "https://files.pythonhosted.org/packages/e7/bb/7fe6cddcc8827b01b1a9766f5fdeb7418680744f9082035bdbabecf1d57f/pillow-11.3.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:068d9c39a2d1b358eb9f245ce7ab1b5c3246c7c8c7d9ba58cfa5b43146c06e50", size = 6114089, upload-time = "2025-07-01T09:15:43.13Z" },
-    { url = "https://files.pythonhosted.org/packages/8b/f5/06bfaa444c8e80f1a8e4bff98da9c83b37b5be3b1deaa43d27a0db37ef84/pillow-11.3.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:a1bc6ba083b145187f648b667e05a2534ecc4b9f2784c2cbe3089e44868f2b9b", size = 6748206, upload-time = "2025-07-01T09:15:44.937Z" },
+    { url = "https://files.pythonhosted.org/packages/d5/11/6db24d4bd7685583caeae54b7009584e38da3c3d4488ed4cd25b439de486/pillow-12.1.1-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:d242e8ac078781f1de88bf823d70c1a9b3c7950a44cdf4b7c012e22ccbcd8e4e", size = 4062689, upload-time = "2026-02-11T04:21:06.804Z" },
+    { url = "https://files.pythonhosted.org/packages/33/c0/ce6d3b1fe190f0021203e0d9b5b99e57843e345f15f9ef22fcd43842fd21/pillow-12.1.1-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:02f84dfad02693676692746df05b89cf25597560db2857363a208e393429f5e9", size = 4138535, upload-time = "2026-02-11T04:21:08.452Z" },
+    { url = "https://files.pythonhosted.org/packages/a0/c6/d5eb6a4fb32a3f9c21a8c7613ec706534ea1cf9f4b3663e99f0d83f6fca8/pillow-12.1.1-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:e65498daf4b583091ccbb2556c7000abf0f3349fcd57ef7adc9a84a394ed29f6", size = 3601364, upload-time = "2026-02-11T04:21:10.194Z" },
+    { url = "https://files.pythonhosted.org/packages/14/a1/16c4b823838ba4c9c52c0e6bbda903a3fe5a1bdbf1b8eb4fff7156f3e318/pillow-12.1.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:6c6db3b84c87d48d0088943bf33440e0c42370b99b1c2a7989216f7b42eede60", size = 5262561, upload-time = "2026-02-11T04:21:11.742Z" },
+    { url = "https://files.pythonhosted.org/packages/bb/ad/ad9dc98ff24f485008aa5cdedaf1a219876f6f6c42a4626c08bc4e80b120/pillow-12.1.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8b7e5304e34942bf62e15184219a7b5ad4ff7f3bb5cca4d984f37df1a0e1aee2", size = 4657460, upload-time = "2026-02-11T04:21:13.786Z" },
+    { url = "https://files.pythonhosted.org/packages/9e/1b/f1a4ea9a895b5732152789326202a82464d5254759fbacae4deea3069334/pillow-12.1.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:18e5bddd742a44b7e6b1e773ab5db102bd7a94c32555ba656e76d319d19c3850", size = 6232698, upload-time = "2026-02-11T04:21:15.949Z" },
+    { url = "https://files.pythonhosted.org/packages/95/f4/86f51b8745070daf21fd2e5b1fe0eb35d4db9ca26e6d58366562fb56a743/pillow-12.1.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fc44ef1f3de4f45b50ccf9136999d71abb99dca7706bc75d222ed350b9fd2289", size = 8041706, upload-time = "2026-02-11T04:21:17.723Z" },
+    { url = "https://files.pythonhosted.org/packages/29/9b/d6ecd956bb1266dd1045e995cce9b8d77759e740953a1c9aad9502a0461e/pillow-12.1.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5a8eb7ed8d4198bccbd07058416eeec51686b498e784eda166395a23eb99138e", size = 6346621, upload-time = "2026-02-11T04:21:19.547Z" },
+    { url = "https://files.pythonhosted.org/packages/71/24/538bff45bde96535d7d998c6fed1a751c75ac7c53c37c90dc2601b243893/pillow-12.1.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:47b94983da0c642de92ced1702c5b6c292a84bd3a8e1d1702ff923f183594717", size = 7038069, upload-time = "2026-02-11T04:21:21.378Z" },
+    { url = "https://files.pythonhosted.org/packages/94/0e/58cb1a6bc48f746bc4cb3adb8cabff73e2742c92b3bf7a220b7cf69b9177/pillow-12.1.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:518a48c2aab7ce596d3bf79d0e275661b846e86e4d0e7dec34712c30fe07f02a", size = 6460040, upload-time = "2026-02-11T04:21:23.148Z" },
+    { url = "https://files.pythonhosted.org/packages/6c/57/9045cb3ff11eeb6c1adce3b2d60d7d299d7b273a2e6c8381a524abfdc474/pillow-12.1.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:a550ae29b95c6dc13cf69e2c9dc5747f814c54eeb2e32d683e5e93af56caa029", size = 7164523, upload-time = "2026-02-11T04:21:25.01Z" },
+    { url = "https://files.pythonhosted.org/packages/19/2a/b9d62794fc8a0dd14c1943df68347badbd5511103e0d04c035ffe5cf2255/pillow-12.1.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0330d233c1a0ead844fc097a7d16c0abff4c12e856c0b325f231820fee1f39da", size = 5264880, upload-time = "2026-02-11T04:21:32.865Z" },
+    { url = "https://files.pythonhosted.org/packages/26/9d/e03d857d1347fa5ed9247e123fcd2a97b6220e15e9cb73ca0a8d91702c6e/pillow-12.1.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:5dae5f21afb91322f2ff791895ddd8889e5e947ff59f71b46041c8ce6db790bc", size = 4660616, upload-time = "2026-02-11T04:21:34.97Z" },
+    { url = "https://files.pythonhosted.org/packages/f7/ec/8a6d22afd02570d30954e043f09c32772bfe143ba9285e2fdb11284952cd/pillow-12.1.1-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2e0c664be47252947d870ac0d327fea7e63985a08794758aa8af5b6cb6ec0c9c", size = 6269008, upload-time = "2026-02-11T04:21:36.623Z" },
+    { url = "https://files.pythonhosted.org/packages/3d/1d/6d875422c9f28a4a361f495a5f68d9de4a66941dc2c619103ca335fa6446/pillow-12.1.1-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:691ab2ac363b8217f7d31b3497108fb1f50faab2f75dfb03284ec2f217e87bf8", size = 8073226, upload-time = "2026-02-11T04:21:38.585Z" },
+    { url = "https://files.pythonhosted.org/packages/a1/cd/134b0b6ee5eda6dc09e25e24b40fdafe11a520bc725c1d0bbaa5e00bf95b/pillow-12.1.1-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e9e8064fb1cc019296958595f6db671fba95209e3ceb0c4734c9baf97de04b20", size = 6380136, upload-time = "2026-02-11T04:21:40.562Z" },
+    { url = "https://files.pythonhosted.org/packages/7a/a9/7628f013f18f001c1b98d8fffe3452f306a70dc6aba7d931019e0492f45e/pillow-12.1.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:472a8d7ded663e6162dafdf20015c486a7009483ca671cece7a9279b512fcb13", size = 7067129, upload-time = "2026-02-11T04:21:42.521Z" },
+    { url = "https://files.pythonhosted.org/packages/1e/f8/66ab30a2193b277785601e82ee2d49f68ea575d9637e5e234faaa98efa4c/pillow-12.1.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:89b54027a766529136a06cfebeecb3a04900397a3590fd252160b888479517bf", size = 6491807, upload-time = "2026-02-11T04:21:44.22Z" },
+    { url = "https://files.pythonhosted.org/packages/da/0b/a877a6627dc8318fdb84e357c5e1a758c0941ab1ddffdafd231983788579/pillow-12.1.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:86172b0831b82ce4f7877f280055892b31179e1576aa00d0df3bb1bbf8c3e524", size = 7190954, upload-time = "2026-02-11T04:21:46.114Z" },
+    { url = "https://files.pythonhosted.org/packages/03/d0/bebb3ffbf31c5a8e97241476c4cf8b9828954693ce6744b4a2326af3e16b/pillow-12.1.1-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:417423db963cb4be8bac3fc1204fe61610f6abeed1580a7a2cbb2fbda20f12af", size = 4062652, upload-time = "2026-02-11T04:21:53.19Z" },
+    { url = "https://files.pythonhosted.org/packages/2d/c0/0e16fb0addda4851445c28f8350d8c512f09de27bbb0d6d0bbf8b6709605/pillow-12.1.1-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:b957b71c6b2387610f556a7eb0828afbe40b4a98036fc0d2acfa5a44a0c2036f", size = 4138823, upload-time = "2026-02-11T04:22:03.088Z" },
+    { url = "https://files.pythonhosted.org/packages/6b/fb/6170ec655d6f6bb6630a013dd7cf7bc218423d7b5fa9071bf63dc32175ae/pillow-12.1.1-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:097690ba1f2efdeb165a20469d59d8bb03c55fb6621eb2041a060ae8ea3e9642", size = 3601143, upload-time = "2026-02-11T04:22:04.909Z" },
+    { url = "https://files.pythonhosted.org/packages/59/04/dc5c3f297510ba9a6837cbb318b87dd2b8f73eb41a43cc63767f65cb599c/pillow-12.1.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:2815a87ab27848db0321fb78c7f0b2c8649dee134b7f2b80c6a45c6831d75ccd", size = 5266254, upload-time = "2026-02-11T04:22:07.656Z" },
+    { url = "https://files.pythonhosted.org/packages/05/30/5db1236b0d6313f03ebf97f5e17cda9ca060f524b2fcc875149a8360b21c/pillow-12.1.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:f7ed2c6543bad5a7d5530eb9e78c53132f93dfa44a28492db88b41cdab885202", size = 4657499, upload-time = "2026-02-11T04:22:09.613Z" },
+    { url = "https://files.pythonhosted.org/packages/6f/18/008d2ca0eb612e81968e8be0bbae5051efba24d52debf930126d7eaacbba/pillow-12.1.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:652a2c9ccfb556235b2b501a3a7cf3742148cd22e04b5625c5fe057ea3e3191f", size = 6232137, upload-time = "2026-02-11T04:22:11.434Z" },
+    { url = "https://files.pythonhosted.org/packages/70/f1/f14d5b8eeb4b2cd62b9f9f847eb6605f103df89ef619ac68f92f748614ea/pillow-12.1.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d6e4571eedf43af33d0fc233a382a76e849badbccdf1ac438841308652a08e1f", size = 8042721, upload-time = "2026-02-11T04:22:13.321Z" },
+    { url = "https://files.pythonhosted.org/packages/5a/d6/17824509146e4babbdabf04d8171491fa9d776f7061ff6e727522df9bd03/pillow-12.1.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b574c51cf7d5d62e9be37ba446224b59a2da26dc4c1bb2ecbe936a4fb1a7cb7f", size = 6347798, upload-time = "2026-02-11T04:22:15.449Z" },
+    { url = "https://files.pythonhosted.org/packages/d1/ee/c85a38a9ab92037a75615aba572c85ea51e605265036e00c5b67dfafbfe2/pillow-12.1.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a37691702ed687799de29a518d63d4682d9016932db66d4e90c345831b02fb4e", size = 7039315, upload-time = "2026-02-11T04:22:17.24Z" },
+    { url = "https://files.pythonhosted.org/packages/ec/f3/bc8ccc6e08a148290d7523bde4d9a0d6c981db34631390dc6e6ec34cacf6/pillow-12.1.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:f95c00d5d6700b2b890479664a06e754974848afaae5e21beb4d83c106923fd0", size = 6462360, upload-time = "2026-02-11T04:22:19.111Z" },
+    { url = "https://files.pythonhosted.org/packages/f6/ab/69a42656adb1d0665ab051eec58a41f169ad295cf81ad45406963105408f/pillow-12.1.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:559b38da23606e68681337ad74622c4dbba02254fc9cb4488a305dd5975c7eeb", size = 7165438, upload-time = "2026-02-11T04:22:21.041Z" },
+    { url = "https://files.pythonhosted.org/packages/6c/9d/efd18493f9de13b87ede7c47e69184b9e859e4427225ea962e32e56a49bc/pillow-12.1.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:1f90cff8aa76835cba5769f0b3121a22bd4eb9e6884cfe338216e557a9a548b8", size = 5268612, upload-time = "2026-02-11T04:22:29.884Z" },
+    { url = "https://files.pythonhosted.org/packages/f8/f1/4f42eb2b388eb2ffc660dcb7f7b556c1015c53ebd5f7f754965ef997585b/pillow-12.1.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1f1be78ce9466a7ee64bfda57bdba0f7cc499d9794d518b854816c41bf0aa4e9", size = 4660567, upload-time = "2026-02-11T04:22:31.799Z" },
+    { url = "https://files.pythonhosted.org/packages/01/54/df6ef130fa43e4b82e32624a7b821a2be1c5653a5fdad8469687a7db4e00/pillow-12.1.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:42fc1f4677106188ad9a55562bbade416f8b55456f522430fadab3cef7cd4e60", size = 6269951, upload-time = "2026-02-11T04:22:33.921Z" },
+    { url = "https://files.pythonhosted.org/packages/a9/48/618752d06cc44bb4aae8ce0cd4e6426871929ed7b46215638088270d9b34/pillow-12.1.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:98edb152429ab62a1818039744d8fbb3ccab98a7c29fc3d5fcef158f3f1f68b7", size = 8074769, upload-time = "2026-02-11T04:22:35.877Z" },
+    { url = "https://files.pythonhosted.org/packages/c3/bd/f1d71eb39a72fa088d938655afba3e00b38018d052752f435838961127d8/pillow-12.1.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d470ab1178551dd17fdba0fef463359c41aaa613cdcd7ff8373f54be629f9f8f", size = 6381358, upload-time = "2026-02-11T04:22:37.698Z" },
+    { url = "https://files.pythonhosted.org/packages/64/ef/c784e20b96674ed36a5af839305f55616f8b4f8aa8eeccf8531a6e312243/pillow-12.1.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6408a7b064595afcab0a49393a413732a35788f2a5092fdc6266952ed67de586", size = 7068558, upload-time = "2026-02-11T04:22:39.597Z" },
+    { url = "https://files.pythonhosted.org/packages/73/cb/8059688b74422ae61278202c4e1ad992e8a2e7375227be0a21c6b87ca8d5/pillow-12.1.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:5d8c41325b382c07799a3682c1c258469ea2ff97103c53717b7893862d0c98ce", size = 6493028, upload-time = "2026-02-11T04:22:42.73Z" },
+    { url = "https://files.pythonhosted.org/packages/c6/da/e3c008ed7d2dd1f905b15949325934510b9d1931e5df999bb15972756818/pillow-12.1.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:c7697918b5be27424e9ce568193efd13d925c4481dd364e43f5dff72d33e10f8", size = 7191940, upload-time = "2026-02-11T04:22:44.543Z" },
 ]
 
 [[package]]