flush old addrs

.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[T](inputs: T, dependencies: array | Sequence[array]) -> T:
+def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array]):
     """
     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,5 +1,9 @@
-from .layers import *
-from .utils import *
+"""
+This type stub file was generated by pyright.
+"""
+
+from layers import *
+from utils import *
 
 from . import init as init
 from . import losses as losses

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

@@ -1,16 +1,20 @@
-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 *
+"""
+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 *

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

@@ -53,7 +53,7 @@ class Module(dict):
         mx.eval(model.parameters())
     """
 
-    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
+    __call__: Callable
     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: mx.Stream
+generation_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: List[mx.array] | mx.array
+    logprobs: mx.array
     max_tokens: List[int]
     num_tokens: List[int]
     cache: List[Any]
-    samplers: List[Callable[[mx.array], mx.array] | None]
-    logits_processors: List[List[Callable[[mx.array, mx.array], mx.array]]]
+    samplers: List[Any]
+    logits_processors: List[Any]
     tokens: List[mx.array]
     def __len__(self) -> int: ...
     def filter(self, keep_idx: List[int]) -> None: ...
@@ -279,18 +279,13 @@ class Batch:
     def extract_cache(self, idx: int) -> List[Any]: ...
 
 class BatchGenerator:
-    model: nn.Module
-    sampler: Callable[[mx.array], mx.array]
-    stop_tokens: set[int]
+    model: Any
     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 | None
-    values: mx.array | None
+    keys: mx.array
+    values: mx.array
     offset: int
     @property
     def state(self) -> tuple[mx.array | None, mx.array | None]: ...
@@ -268,14 +268,29 @@ class CacheList(_BaseCache):
         """
 
 class BatchKVCache(_BaseCache):
-    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]: ...
+    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]:
+        ...
     @property
     def state(
         self,
@@ -301,21 +316,12 @@ class BatchKVCache(_BaseCache):
         """
 
 class BatchRotatingKVCache(_BaseCache):
-    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]: ...
+    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]:
+        ...
     @property
     def state(
         self,

.mlx_typings/mlx_lm/models/gated_delta.pyi

@@ -1,35 +0,0 @@
-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

@@ -1,51 +0,0 @@
-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,6 +1,11 @@
 [workspace]
 resolver = "3"
-members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
+members = [
+  "rust/networking",
+  "rust/exo_pyo3_bindings",
+  "rust/util",
+  "rust/babblerd",
+]
 
 [workspace.package]
 version = "0.0.1"
@@ -43,7 +48,6 @@ 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,21 +501,23 @@ def main() -> int:
                                     for x, _ in batch_results
                                     if x["stats"]["generation_tps"] > 0
                                 ]
-                                per_req_tps = (
+                                agg_gen_tps = (
                                     mean(valid_gen_tps) if valid_gen_tps else 0.0
                                 )
-                                agg_gen_tps = per_req_tps * concurrency
+                                gen_tps = agg_gen_tps / concurrency
                                 logger.info(
                                     f"[concurrent {concurrency}x]  "
                                     f"agg_gen_tps={agg_gen_tps:.2f}  "
-                                    f"per_req_tps={per_req_tps:.2f}  "
+                                    f"gen_tps={gen_tps:.2f}  "
                                     f"errors={batch_errors}"
                                 )
 
                     if runs:
                         prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
-                        per_req_tps = mean(x["stats"]["generation_tps"] for x in runs)
-                        gen_tps = per_req_tps * concurrency
+                        gen_tps = mean(
+                            x["stats"]["generation_tps"] / x["concurrency"]
+                            for x in runs
+                        )
                         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

@@ -1,91 +0,0 @@
-#!/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

@@ -0,0 +1,377 @@
+# 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,14 +1793,6 @@ 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
@@ -1998,14 +1990,6 @@ 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)
@@ -2413,7 +2397,7 @@ class AppStore {
 
       let streamedContent = "";
       let streamedThinking = "";
-      let serverTpsReceived = false;
+
       interface ChatCompletionChunk {
         choices?: Array<{
           delta?: { content?: string; reasoning_content?: string };
@@ -2478,6 +2462,7 @@ 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;
@@ -2528,24 +2513,16 @@ 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;
 
-      // Use server-side TPS if available, otherwise fall back to client-side
-      if (!serverTpsReceived && firstTokenTime !== null && tokenCount > 1) {
+      // Calculate final TPS
+      if (firstTokenTime !== null && tokenCount > 1) {
         const totalGenerationTime = performance.now() - firstTokenTime;
-        this.tps = (tokenCount / totalGenerationTime) * 1000;
+        this.tps = (tokenCount / totalGenerationTime) * 1000; // tokens per second
       }
 
       // Final cleanup of the message (if conversation still exists)

flake.nix

@@ -112,7 +112,9 @@
             };
           };
 
-          packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
+          packages = {
+            babeld = pkgs.callPackage ./nix/babeld.nix { };
+          } // 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);
@@ -157,9 +159,6 @@
                 just
                 jq
               ]
-              ++ lib.optionals stdenv.isLinux [
-                unixtools.ifconfig
-              ]
               ++ lib.optionals stdenv.isDarwin [
                 macmon
               ];

nix/babeld.nix

@@ -0,0 +1,25 @@
+{ 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,7 +25,8 @@ dependencies = [
   "openai-harmony>=0.0.8",
   "httpx>=0.28.1",
   "tomlkit>=0.14.0",
-  "mflux==0.16.9",
+  "pillow>=11.0,<12.0",                        # compatibility with mflux
+  "mflux==0.15.5",
   "python-multipart>=0.0.21",
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
@@ -61,7 +62,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/ml-explore/mlx-lm", branch = "main" }
+mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" }
 # 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,6 +185,7 @@
         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

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

rust/babblerd/Cargo.toml

@@ -0,0 +1,19 @@
+[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

@@ -0,0 +1,548 @@
+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

@@ -0,0 +1,167 @@
+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', pkgs, lib, ... }:
+    { inputs', self', pkgs, lib, ... }:
     let
       # Fenix nightly toolchain with all components
       rustToolchain = inputs'.fenix.packages.stable.withComponents [
@@ -79,7 +79,7 @@
       };
 
       config = {
-        packages = {
+        packages = rec {
           # Python bindings wheel via maturin
           exo_pyo3_bindings = craneLib.buildPackage (
             commonArgs
@@ -110,6 +110,19 @@
               '';
             }
           );
+          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,8 +202,6 @@ 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
 
@@ -218,8 +216,6 @@ 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)
+            self.block.img_mod_silu(text_embeddings)  # pyright: ignore[reportUnknownArgumentType]
         )
         txt_mod_params = self.block.txt_mod_linear(
-            self.block.txt_mod_silu(text_embeddings)
+            self.block.txt_mod_silu(text_embeddings)  # pyright: ignore[reportUnknownArgumentType]
         )
 
         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]
+                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # pyright: ignore[reportAny,reportUnknownMemberType]
 
         return logits
 

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

@@ -1,4 +1,3 @@
-import os
 import time
 from dataclasses import dataclass, field
 from typing import Callable, cast
@@ -80,173 +79,12 @@ class ExoBatchGenerator:
     _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False)
 
     def __post_init__(self) -> None:
-        use_speculative = os.environ.get("EXO_SPECULATIVE", "0") == "1"
-        stop_tokens = set(eos_ids_from_tokenizer(self.tokenizer))
-
-        if use_speculative:
-            try:
-                from exo.worker.engines.mlx.speculative.mtp_module import MTPPredictor
-                from exo.worker.engines.mlx.speculative.mtp_batch_generator import MTPBatchGenerator
-
-                mtp_weights = self._resolve_mtp_weights()
-                gamma = int(os.environ.get("EXO_SPECULATIVE_GAMMA", "2"))
-
-                if mtp_weights:
-                    mtp = MTPPredictor(self.model, mtp_weights, quantize=False)
-                    temp = float(os.environ.get("EXO_SPECULATIVE_TEMP", "0.7"))
-                    alpha = float(os.environ.get("EXO_SPECULATIVE_ALPHA", "1.0"))
-                    self._exo_gen = MTPBatchGenerator(
-                        model=self.model,
-                        mtp_predictor=mtp,
-                        gamma=gamma,
-                        temp=temp,
-                        alpha=alpha,
-                        stop_tokens=stop_tokens,
-                        prefill_step_size=4096,
-                    )
-                    logger.info(f"MTP speculative decoding enabled (γ={gamma}, T={temp})")
-                    self.warmup_speculative(self.model, self.tokenizer)
-                else:
-                    logger.warning("EXO_SPECULATIVE=1 but could not find MTP weights. Falling back to standard generation.")
-                    self._exo_gen = MlxBatchGenerator(
-                        model=self.model,
-                        stop_tokens=stop_tokens,
-                        prefill_step_size=4096,
-                    )
-            except Exception as e:
-                logger.warning(f"Failed to initialize MTP speculative decoding: {e}. Falling back to standard generation.")
-                self._exo_gen = MlxBatchGenerator(
-                    model=self.model,
-                    stop_tokens=stop_tokens,
-                    prefill_step_size=4096,
-                )
-        else:
-            self._exo_gen = MlxBatchGenerator(
-                model=self.model,
-                stop_tokens=stop_tokens,
-                prefill_step_size=4096,
-            )
-
-    def _resolve_mtp_weights(self) -> str | None:
-        """Find MTP weights: explicit path, explicit HF model, or auto-extract."""
-        # 1. Explicit path
-        explicit_path = os.environ.get("EXO_MTP_WEIGHTS", "")
-        if explicit_path and os.path.exists(explicit_path):
-            return explicit_path
-
-        # 2. Explicit HF model repo containing MTP weights
-        mtp_model = os.environ.get("EXO_MTP_MODEL", "")
-
-        # 3. Auto-detect: if no EXO_MTP_MODEL set, try to infer from model config
-        if not mtp_model:
-            try:
-                inner = getattr(self.model, 'model', None) or self.model.language_model.model
-                args = getattr(inner, 'args', None)
-                if args and getattr(args, 'mtp_num_hidden_layers', 0) > 0:
-                    model_type = getattr(args, 'model_type', '')
-                    if 'qwen3_5' in model_type or 'qwen3.5' in str(type(self.model).__module__):
-                        # Default pairing for Qwen3.5-27B
-                        mtp_model = "Qwen/Qwen3.5-27B"
-                        logger.info(f"Auto-detected MTP model: {mtp_model}")
-            except Exception:
-                pass
-
-        if not mtp_model:
-            return None
-
-        # Download and extract MTP weights from HF repo
-        try:
-            return self._extract_mtp_from_hf(mtp_model)
-        except Exception as e:
-            logger.warning(f"Failed to extract MTP weights from {mtp_model}: {e}")
-            return None
-
-    def _extract_mtp_from_hf(self, repo_id: str) -> str:
-        """Download MTP tensors from HF repo and cache as a single safetensors file."""
-        import hashlib
-        from pathlib import Path
-        from huggingface_hub import snapshot_download
-        from safetensors.torch import load_file, save_file
-
-        cache_dir = Path.home() / ".cache" / "exo" / "mtp_weights"
-        cache_dir.mkdir(parents=True, exist_ok=True)
-        cache_key = hashlib.md5(repo_id.encode()).hexdigest()[:12]
-        cached_path = cache_dir / f"mtp_{cache_key}.safetensors"
-
-        if cached_path.exists():
-            logger.info(f"Using cached MTP weights: {cached_path}")
-            return str(cached_path)
-
-        logger.info(f"Downloading MTP weights from {repo_id}...")
-        model_dir = snapshot_download(
-            repo_id,
-            allow_patterns=["*.safetensors", "*.json"],
+        self._exo_gen = MlxBatchGenerator(
+            model=self.model,
+            stop_tokens=set(eos_ids_from_tokenizer(self.tokenizer)),
+            prefill_step_size=4096,
         )
 
-        # 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 (
@@ -293,16 +131,10 @@ 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=sampling_temp,
+            temp=task_params.temperature
+            if task_params.temperature is not None
+            else 0.7,
             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,
@@ -319,23 +151,6 @@ 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 (
@@ -385,16 +200,6 @@ 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,
@@ -471,7 +276,7 @@ class ExoBatchGenerator:
 
             logprob: float | None = None
             top_logprobs: list[TopLogprobItem] | None = None
-            if task_params.logprobs and os.environ.get("EXO_DISABLE_LOGPROBS") != "1":
+            if task_params.logprobs:
                 logprob, top_logprobs = extract_top_logprobs(
                     logprobs=response.logprobs,
                     tokenizer=self.tokenizer,

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

@@ -179,8 +179,7 @@ def pipeline_parallel_prefill(
         flush_prefill_sends()
 
     assert _prompt_cache is not None
-    with mx.stream(generation_stream):
-        mx.eval([c.state for c in _prompt_cache])  # type: ignore
+    mx.eval([c.state for c in _prompt_cache])  # type: ignore
 
     # Final callback matching generate_step
     prompt_progress_callback(total, total)
@@ -399,44 +398,52 @@ 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]]:
-    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
+    """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]
 
     # Convert to list of TopLogprobItem
     top_logprob_items: list[TopLogprobItem] = []
-    for token_id, token_logprob in zip(top_indices_list, top_values_list, strict=True):
+    for i in range(top_logprobs):
+        token_id = int(top_indices[i].item())
+        token_logprob = float(top_values[i].item())
         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=list(token_str.encode("utf-8")),
+                bytes=token_bytes,
             )
         )
 
@@ -617,13 +624,12 @@ def mlx_generate(
         logprob: float | None = None
         top_logprobs: list[TopLogprobItem] | None = None
         if task.logprobs:
-            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,
-                )
+            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

@@ -1,45 +0,0 @@
-"""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

@@ -1,27 +0,0 @@
-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

@@ -1,173 +0,0 @@
-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

@@ -1,135 +0,0 @@
-#!/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

@@ -1,108 +0,0 @@
-#!/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

@@ -1,74 +0,0 @@
-"""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

@@ -1,102 +0,0 @@
-"""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

@@ -1,94 +0,0 @@
-"""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

@@ -1,500 +0,0 @@
-"""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

@@ -1,199 +0,0 @@
-"""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

@@ -1,161 +0,0 @@
-"""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

@@ -1,265 +0,0 @@
-"""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

@@ -1,205 +0,0 @@
-"""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

@@ -1,252 +0,0 @@
-"""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

@@ -1,92 +0,0 @@
-"""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

@@ -1,332 +0,0 @@
-"""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

@@ -1,578 +0,0 @@
-"""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

@@ -1,288 +0,0 @@
-"""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

@@ -1,128 +0,0 @@
-"""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

@@ -1,117 +0,0 @@
-"""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

@@ -1,177 +0,0 @@
-"""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

@@ -1,118 +0,0 @@
-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

@@ -1,324 +0,0 @@
-#!/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

@@ -1,515 +0,0 @@
-#!/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

@@ -1,98 +0,0 @@
-#!/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

@@ -1,180 +0,0 @@
-#!/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

@@ -1,290 +0,0 @@
-# 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,11 +190,6 @@ 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:
@@ -643,7 +638,6 @@ 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,7 +8,6 @@ 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
 
@@ -46,8 +45,6 @@ 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,8 +345,6 @@ 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,
@@ -429,8 +427,7 @@ class BatchGenerator(InferenceGenerator):
 
             task, queue, output_generator = self._active_tasks[uid]
             queue.push(response)
-            # If a generator fails to parse for some reason and returns early, we should not crash
-            parsed = next(output_generator, None)
+            parsed = next(output_generator)
 
             if parsed is not None:
                 output.append((task.task_id, parsed))

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

@@ -319,9 +319,7 @@ 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)
-            return x
+                x = layer(x, *args, **kwargs)  # pyright: ignore[reportUnknownVariableType]
+            return x  # pyright: ignore[reportUnknownVariableType]
 
     try:
         group = mx.distributed.init(backend="ring", strict=True)

uv.lock

@@ -489,6 +489,7 @@ 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'" },
@@ -521,12 +522,13 @@ 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.16.9" },
+    { name = "mflux", specifier = "==0.15.5" },
     { 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/ml-explore/mlx-lm?branch=main" },
+    { name = "mlx-lm", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Feval-left-padding-in-batched-rotation" },
     { 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" },
@@ -763,36 +765,6 @@ 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"
@@ -1354,12 +1326,11 @@ wheels = [
 
 [[package]]
 name = "mflux"
-version = "0.16.9"
+version = "0.15.5"
 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'" },
@@ -1381,9 +1352,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/e3/dc/73dfccae395df22623fbb6f748657700fd9372be22904a0dafb825dc42d2/mflux-0.16.9.tar.gz", hash = "sha256:ff35e9386b5f026a6b97c786b3426e5341105de3356e55ec1b0c5951ff0ac8c8", size = 764296, upload-time = "2026-03-07T11:54:45.031Z" }
+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" }
 wheels = [
-    { 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" },
+    { 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" },
 ]
 
 [[package]]
@@ -1444,8 +1415,8 @@ wheels = [
 
 [[package]]
 name = "mlx-lm"
-version = "0.31.2"
-source = { git = "https://github.com/ml-explore/mlx-lm?branch=main#ed7884cb80968e0e77fce6cde5d1597952bbd524" }
+version = "0.31.0"
+source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Feval-left-padding-in-batched-rotation#5e0c484cfb5c68d281a71409927ee1bb75adaae2" }
 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'" },
@@ -1964,48 +1935,45 @@ wheels = [
 
 [[package]]
 name = "pillow"
-version = "12.1.1"
+version = "11.3.0"
 source = { registry = "https://pypi.org/simple" }
-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" }
+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" }
 wheels = [
-    { 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" },
+    { 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" },
 ]
 
 [[package]]