Bake scale into SDPA pass1 source, remove last scalar input

This reverts commit 97d7f03032.

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

@@ -0,0 +1,47 @@
+"""Model-specific kernel fusion patches for MLX inference.
+
+Detects model type after loading and applies optimized kernel patches
+where available. Currently supports:
+- Qwen3.5 MoE (model_type: qwen3_5_moe):
+    EXO_FUSED_KERNELS=1: oproj mode (4-dispatch MoE fusion)
+    EXO_FUSED_KERNELS=2: fused_gqa_gdn mode (full GDN + GQA + MoE fusion, default)
+
+Set EXO_FUSED_KERNELS=0 to disable all patches (vanilla mode).
+Default: EXO_FUSED_KERNELS=2 (fused_gqa_gdn, best performance).
+"""
+
+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", "2")
+    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":
+        if fused_mode == "1":
+            from .qwen3_5_moe.apply import apply_qwen35_oproj_patches
+
+            logger.info("Detected Qwen3.5 MoE model, applying oproj fusion patches")
+            apply_qwen35_oproj_patches(model)
+        else:
+            from .qwen3_5_moe.apply import apply_qwen35_fused_gqa_gdn_patches
+
+            logger.info("Detected Qwen3.5 MoE model, applying fused GQA+GDN patches")
+            apply_qwen35_fused_gqa_gdn_patches(model)

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

@@ -0,0 +1,44 @@
+"""Apply kernel fusion patches to Qwen3.5 MoE models.
+
+Entry point called from patches/__init__.py after model type detection.
+Supports oproj mode (MoE only) and fused_gqa_gdn mode (full attention + MoE).
+"""
+
+import time
+
+import mlx.nn as nn
+from loguru import logger
+
+from .common import apply_oproj_fused_patches, apply_fused_gqa_gdn_patches
+
+
+def apply_qwen35_oproj_patches(model: nn.Module) -> None:
+    """Apply oproj 4-dispatch fusion to all layers of a Qwen3.5 MoE model.
+
+    Patches decoder, attention, and MoE __call__ methods to use custom Metal
+    kernels for decode (seq_len=1). Prefill (seq_len>1) falls back to vanilla.
+    """
+    layers = model.layers  # type: ignore[attr-defined]
+    n_layers = len(layers)
+
+    t0 = time.time()
+    apply_oproj_fused_patches(layers, gate_bm=8, free_originals=False)
+    t_patch = time.time() - t0
+
+    logger.info(f"Qwen3.5 oproj fusion: patched {n_layers} layers in {t_patch:.1f}s")
+
+
+def apply_qwen35_fused_gqa_gdn_patches(model: nn.Module) -> None:
+    """Apply full fusion (GDN + GQA attention + oproj MoE) to all layers.
+
+    Fused GDN attention (3/4 layers) + fused GQA attention (1/4 layers)
+    + oproj MoE (all layers). 44% faster than vanilla on Qwen3.5-35B-A3B.
+    """
+    layers = model.layers  # type: ignore[attr-defined]
+    n_layers = len(layers)
+
+    t0 = time.time()
+    apply_fused_gqa_gdn_patches(layers, gate_bm=8, free_originals=False)
+    t_patch = time.time() - t0
+
+    logger.info(f"Qwen3.5 fused GQA+GDN: patched {n_layers} layers in {t_patch:.1f}s")

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

@@ -0,0 +1,346 @@
+"""Weight preparation and patch orchestration for Qwen3.5 kernel fusion.
+
+Supports:
+  apply_oproj_fused_patches    — oproj mode (4-dispatch MoE only)
+  apply_fused_gqa_gdn_patches  — full fusion (fused GDN + GQA attention + oproj MoE)
+
+Adapted from mlx_bench/model_patches/qwen/common.py.
+"""
+
+import mlx.core as mx
+import mlx.nn as nn
+from loguru import logger
+from mlx_lm.models.qwen3_5 import DecoderLayer
+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."""
+    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."""
+    shared = moe.shared_expert
+    gp = shared.gate_proj
+    up = shared.up_proj
+    dp = shared.down_proj
+
+    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)
+
+    moe._shared_down_w = dp.weight
+    moe._shared_down_s = dp.scales
+    moe._shared_down_b = dp.biases
+
+    moe._shared_inter = gp.weight.shape[0]
+    moe._shared_gs = gp.group_size
+
+    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
+    moe._down_s = dp.scales
+    moe._down_b = dp.biases
+    moe._down_K = dp.output_dims
+    moe._down_N = dp.input_dims
+    moe._down_gs = dp.group_size
+    mx.eval(moe._down_w, moe._down_s, moe._down_b)
+
+
+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
+    """
+    moe = layer.mlp
+
+    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)
+    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)
+
+    rms_weight = layer.post_attention_layernorm.weight.astype(mx.bfloat16)
+
+    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
+
+    M1 = (W_fused.astype(mx.float32) @ W_oproj_f32).astype(mx.bfloat16)
+    mx.eval(M1)
+    del W_oproj_f32
+
+    moe._oproj_M1 = M1
+    moe._oproj_W_fused = W_fused
+    moe._oproj_rms_weight = rms_weight
+
+    moe._oproj_w = oproj.weight
+    moe._oproj_s = oproj.scales
+    moe._oproj_b = oproj.biases
+    moe._oproj_K_attn = oproj.weight.shape[1] * 4
+
+    seg = moe.shared_expert_gate
+    moe._seg_w = seg.weight
+    moe._seg_s = seg.scales
+    moe._seg_b = seg.biases
+
+    M = oproj.weight.shape[0]
+    K_hidden = W_fused.shape[1]
+    n_experts = W_fused.shape[0]
+    moe._oproj_M = M
+    moe._oproj_K_hidden = K_hidden
+    moe._oproj_n_experts = n_experts
+    moe._oproj_n_tg = ceil_div(M, 32)
+    moe._oproj_gate_bm = gate_bm
+
+    mx.eval(moe._oproj_rms_weight)
+
+
+def apply_oproj_fused_patches(layers, gate_bm=8, free_originals=False):
+    """Apply oproj-mode fused MoE patches (4-dispatch) to all layers.
+
+    1. Prepare SwiGLU weights (_patch_swiglu_weights)
+    2. Prepare shared expert 8-bit weights (_patch_shared_expert)
+    3. Prepare down_proj weights (_patch_down_proj)
+    4. Precompute M1/W_fused + store o_proj/gate weights (_patch_oproj_gate_rms)
+    5. Replace __call__ methods for decoder + MoE + attention
+    """
+    from .moe import _oproj_moe_call
+    from .decoder import (
+        _oproj_decoder_call,
+        _pre_oproj_attention_call,
+        _pre_oproj_qwen35_linear_attn_call,
+    )
+    from mlx_lm.models.qwen3_next import Qwen3NextAttention
+    from mlx_lm.models.qwen3_5 import GatedDeltaNet
+
+    n_patched = 0
+    for li, layer in enumerate(layers):
+        moe = layer.mlp
+        if isinstance(moe, Qwen3NextSparseMoeBlock):
+            _patch_swiglu_weights(moe)
+            _patch_shared_expert(moe)
+            _patch_down_proj(moe)
+            _patch_oproj_gate_rms(layer, gate_bm=gate_bm)
+
+            if free_originals:
+                for attr in ('weight', 'scales', 'biases'):
+                    for proj in (moe.switch_mlp.gate_proj,
+                                 moe.switch_mlp.up_proj):
+                        try:
+                            delattr(proj, attr)
+                        except AttributeError:
+                            pass
+
+            n_patched += 1
+            if (li + 1) % 10 == 0 or li == 0:
+                logger.info(f"  Patched layer {li+1}/{len(layers)} (oproj mode)")
+
+    Qwen3NextAttention.__call__ = _pre_oproj_attention_call
+    GatedDeltaNet.__call__ = _pre_oproj_qwen35_linear_attn_call
+    Qwen3NextSparseMoeBlock.__call__ = _oproj_moe_call
+    DecoderLayer.__call__ = _oproj_decoder_call
+    logger.info(f"  Patched {n_patched} MoE blocks (oproj mode, 4 dispatches)")
+
+
+def _patch_gdn_proj_weights(attn):
+    """Merge all 4 GDN projection weights into contiguous buffers."""
+    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],
+        attn.in_proj_z.weight.shape[0],
+        attn.in_proj_b.weight.shape[0],
+        attn.in_proj_a.weight.shape[0],
+    )
+    mx.eval(W_merged, S_merged, B_merged)
+
+
+def _patch_gqa_proj_weights(attn):
+    """Merge GQA q_proj, k_proj, v_proj weights with q_proj row permutation.
+
+    q_proj rows are interleaved [head0_q, head0_gate, head1_q, ...].
+    Permute so queries come first, then gate, then k, then v.
+    Pre-cache constant scalar arrays for kernel dispatch.
+    """
+    q = attn.q_proj
+    k = attn.k_proj
+    v = attn.v_proj
+
+    H_q = attn.num_attention_heads
+    D = attn.head_dim
+
+    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)
+
+    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
+    N_Q = H_q * D
+    N_GATE = H_q * D
+    N_K = k.weight.shape[0]
+    N_V = v.weight.shape[0]
+    attn._merged_proj_dims = (N_Q, N_GATE, N_K, N_V)
+    mx.eval(W_merged, S_merged, B_merged)
+
+    # Pre-cache constant scalar arrays for kernel dispatch
+    N_TOTAL = N_Q + N_GATE + N_K + N_V
+    K_dim = q.weight.shape[1] * 4  # 8-bit: pack_factor=4
+    attn._kernel_scalars = {
+        'K': mx.array(K_dim, 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),
+        '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())
+
+    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
+    rope_dims = attn.rope.dims
+    half_dims = rope_dims // 2
+    theta = attn.rope.base
+    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 apply_fused_gqa_gdn_patches(layers, gate_bm=8, free_originals=False):
+    """Apply all fusions: fused GDN + fused GQA + oproj MoE.
+
+    Combines:
+      - Fused GDN attention (3/4 layers: GatedDeltaNet)
+      - Fused GQA attention (1/4 layers: Qwen3NextAttention)
+      - Oproj MoE (all layers: oproj_gate_gemv + fused MoE dispatches)
+    """
+    from .moe import _oproj_moe_call
+    from .decoder import _fused_gdn_decoder_call
+    from .fused_gdn_attention import _fused_gdn_call
+    from .fused_gqa_attention import _fused_gqa_call
+    from mlx_lm.models.qwen3_next import Qwen3NextAttention
+    from mlx_lm.models.qwen3_5 import GatedDeltaNet
+
+    n_patched = 0
+    n_gdn = 0
+    n_gqa = 0
+    for li, layer in enumerate(layers):
+        moe = layer.mlp
+        if isinstance(moe, Qwen3NextSparseMoeBlock):
+            _patch_swiglu_weights(moe)
+            _patch_shared_expert(moe)
+            _patch_down_proj(moe)
+            _patch_oproj_gate_rms(layer, gate_bm=gate_bm)
+
+            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 free_originals:
+                for attr in ('weight', 'scales', 'biases'):
+                    for proj in (moe.switch_mlp.gate_proj,
+                                 moe.switch_mlp.up_proj):
+                        try:
+                            delattr(proj, attr)
+                        except AttributeError:
+                            pass
+
+            n_patched += 1
+            if (li + 1) % 10 == 0 or li == 0:
+                logger.info(f"  Patched layer {li+1}/{len(layers)} (fused GQA+GDN mode)")
+
+    GatedDeltaNet.__call__ = _fused_gdn_call
+    Qwen3NextAttention.__call__ = _fused_gqa_call
+    Qwen3NextSparseMoeBlock.__call__ = _oproj_moe_call
+    DecoderLayer.__call__ = _fused_gdn_decoder_call
+    logger.info(f"  Patched {n_patched} MoE blocks ({n_gdn} GDN + {n_gqa} GQA, fused GQA+GDN mode)")

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

@@ -0,0 +1,186 @@
+"""Decoder layer __call__ variants for Qwen3.5.
+
+Three modes:
+  _fused_decoder_call: passes residual to fused MoE epilogue (~15 dispatches)
+  _oproj_decoder_call: fuses o_proj + RMSNorm + gate GEMV (4 dispatches)
+  _fused_gdn_decoder_call: fused GDN/GQA attention + oproj MoE (6 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.
+
+Adapted from mlx_bench/model_patches/qwen/decoder.py.
+"""
+
+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 _fused_gdn_decoder_call(self, x, mask=None, cache=None):
+    """Decoder with fused GDN/GQA attention + oproj MoE (6-dispatch mode).
+
+    GDN layers use fused kernels, GQA layers use fused or vanilla attention.
+    Both return pre-out_proj output. MoE handles oproj_gate_gemv + MoE dispatches.
+
+    Flow is identical to oproj mode — the difference is that GatedDeltaNet.__call__
+    and/or Qwen3NextAttention.__call__ are patched with fused kernel implementations.
+    """
+    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=not self.training,
+    )
+
+    if cache is not None:
+        cache[1] = state
+
+    out = self.norm(out, z)
+    return out.reshape(B, S, -1)  # skip out_proj

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

@@ -0,0 +1,161 @@
+"""Fused GDN attention __call__ for qwen3_5.GatedDeltaNet (Dispatches 2-5).
+
+Replaces the vanilla GatedDeltaNet.__call__ with fused kernel dispatches:
+  Dispatch 2: fused_gdn_projections — merged 8-bit GEMV + conv1d + SiLU(qkv) + SiLU(z)
+              + sigmoid(b)→beta + g=exp(-exp(A_log)*softplus(a+dt_bias))
+  Dispatch 3: fused_qk_rmsnorm — per-head L2-norm on q (×Dk^(-½)) and k
+  Dispatch 4: gated_delta_kernel — GDN recurrence (receives pre-computed g, beta)
+  Dispatch 5: fused_rms_norm_gated — RMSNorm(out, weight) × z_silu
+
+All 4 projection weights are pre-merged into contiguous buffers at patch time
+(_patch_gdn_proj_weights) for better memory locality.
+
+g/beta computation is fused into Dispatch 2 epilogues, eliminating ~8 micro-
+dispatches that gated_delta_update would otherwise generate.
+
+Fused path is decode-only (S=1). For prefill (S>1), falls back to vanilla ops.
+
+Returns pre-out_proj output (same interface as _pre_oproj_qwen35_linear_attn_call).
+Dispatch 1 (input_layernorm) is handled by the decoder.
+Dispatch 6 (oproj_gate_gemv) is handled by the MoE __call__.
+"""
+
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .kernels.fused_gdn_projections_8bit import 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 out.reshape(B, S, -1)  # skip out_proj
+
+
+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
+
+    # Prefill fallback: fused kernels are decode-only (S=1)
+    if S > 1:
+        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/fused_gqa_attention.py

@@ -0,0 +1,130 @@
+"""Fused GQA attention __call__ for qwen3_next.Qwen3NextAttention.
+
+Replaces vanilla Qwen3NextAttention.__call__ with fused kernel dispatches:
+  Dispatch 1: fused_gqa_projections — merged 8-bit GEMV (q+gate+k+v) + sigmoid(gate)
+  Dispatch 2: fused_qk_norm_rope — RMSNorm + RoPE (TODO: custom kernel)
+  Dispatch 3: KV cache update (MLX built-in)
+  Dispatch 4: custom_sdpa_pass1 (TODO: custom kernel)
+  Dispatch 5: custom_sdpa_pass2_gate (TODO: custom kernel, includes gate multiply)
+  Dispatch 6: oproj_gate_gemv (existing, handled by MoE __call__)
+
+Dispatches 1-2, 4-5 are custom kernels, Dispatch 3 is MLX built-in.
+Returns pre-out_proj output (output * sigmoid(gate)) for Dispatch 6.
+
+Fused path is decode-only (S=1). For prefill (S>1), falls back to vanilla ops.
+"""
+
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .kernels.fused_gqa_projections_8bit import fused_gqa_projections
+from .kernels.fused_qk_norm_rope import fused_qk_norm_rope
+from .kernels.custom_sdpa_pass1 import custom_sdpa_pass1
+from .kernels.custom_sdpa_pass2_gate import custom_sdpa_pass2_gate
+
+
+def _vanilla_gqa_call(self, x, mask, cache):
+    """Vanilla GQA path for prefill (S>1). Returns pre-out_proj output."""
+    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 _fused_gqa_call(
+    self,
+    x: mx.array,
+    mask: Optional[mx.array] = None,
+    cache: Optional[Any] = None,
+) -> mx.array:
+    """Fused GQA attention: custom projection + norm/rope kernels.
+
+    Decode (S=1): Dispatch 1 (fused GEMV) + Dispatch 2 (fused norm+rope)
+                  + vanilla cache + SDPA (Dispatches 3-5).
+    Prefill (S>1): falls back to fully vanilla ops.
+
+    Returns pre-out_proj output [B, S, H_q*D] for Dispatch 6.
+    """
+    B, S, _ = x.shape
+
+    # Prefill fallback: fused kernels are decode-only (S=1)
+    if S > 1:
+        return _vanilla_gqa_call(self, x, mask, cache)
+
+    H_q = self.num_attention_heads
+    H_kv = self.num_key_value_heads
+    D = self.head_dim
+
+    # ── Dispatch 1: fused projections (merged GEMV + sigmoid(gate)) ──
+    queries, gate_sigmoid, keys, values = 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: fused Q/K RMSNorm + RoPE ──
+    queries, keys = fused_qk_norm_rope(
+        queries, keys,
+        self.q_norm.weight, self.k_norm.weight,
+        self._rope_inv_freq, cache.offset,
+        H_q, H_kv, D, batch_size=B,
+    )
+    # queries: [B, H_q, 1, D], keys: [B, H_kv, 1, D]
+    values = values.reshape(B, H_kv, 1, D)
+
+    # ── Dispatch 3: KV cache update ──
+    cache.update_and_fetch(keys, values)
+    N = cache.offset
+    alloc_len = cache.keys.shape[2]
+
+    # ── Dispatch 4: SDPA Pass 1 ──
+    blocks = 128
+    if N < 1024:
+        k_sliced = cache.keys[:, :, :N, :]
+        v_sliced = cache.values[:, :, :N, :]
+        from mlx_lm.models.qwen3_next import scaled_dot_product_attention
+        output = scaled_dot_product_attention(
+            queries, k_sliced, v_sliced, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, S, -1)
+        return output * gate_sigmoid.astype(output.dtype)
+
+    o_partials, sums, maxs = custom_sdpa_pass1(
+        queries, cache.keys, cache.values, self.scale,
+        H_q, H_kv, D, blocks=blocks, batch_size=B,
+        N=N, alloc_len=alloc_len,
+    )
+
+    # ── Dispatch 5: SDPA Pass 2 + gate multiply ──
+    return custom_sdpa_pass2_gate(
+        o_partials, sums, maxs, gate_sigmoid,
+        H_q, D, blocks=blocks, V_SPLIT=4, batch_size=B,
+    )

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

@@ -0,0 +1,324 @@
+"""Dispatch 1: Fused o_proj (8-bit) + gate GEMV parts + x² partials for Qwen3.5.
+
+Port of Kimi's custom_oproj_gate_gemv.py adapted for 8-bit quantized o_proj.
+
+Single dispatch with 3 GEMV types sharing 256-thread TGs (8 SGs of 32):
+
+TGs 0 to N_OPROJ_TG-1: o_proj GEMV (8-bit, M=4096, K=8192)
+  8-bit affine dequant: result = scale * Σ(x*w_uint8) + bias * Σ(x)
+  Epilogue: h = oproj_result + residual, h_scaled = h*w_rms, h_out = h, x²_acc
+  TG reduction: 8 SG x² → 1 float per TG.
+
+TGs N_OPROJ_TG to +N_M1_TG-1: M1 GEMV (bf16, E × K_attn)
+  M1 = W_fused @ W_oproj (pre-computed). Input = attn_out.
+  Output: gate_part_a (E,) f32.
+
+TGs +N_M1_TG to end: W_fused GEMV (bf16, E × K)
+  W_fused × residual → gate_part_b (E,) f32.
+"""
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_custom_oproj_8bit_source(n_experts=64, group_size=64, scale_bf16=True):
+    """Generate Metal source for fused 8-bit o_proj + bf16 gate GEMVs."""
+    E = int(n_experts)
+    gs = int(group_size)
+    sc_t = "bfloat16_t" if scale_bf16 else "float"
+    # 8-bit dequant params for o_proj
+    oproj_slid_divisor = gs // 8     # 64/8 = 8
+    oproj_sc_stride = 256 // gs      # 256/64 = 4
+
+    return f"""
+    // ── Constants ──
+    const int TM = 4;       // rows per SG for bf16 GEMVs
+    const int TN = 4;       // elements per thread per iter for bf16 GEMVs
+    const int blockN = 128; // 32 threads × TN=4
+    const int E_CONST = {E};
+
+    int M = M_val;                   // 4096 (hidden_size)
+    int K_attn = K_attn_val;         // 8192 (o_proj input dim)
+    int K_hidden = K_hidden_val;     // 4096 (hidden_size, same as M for Qwen)
+    int N_OPROJ_TG = N_OPROJ_TG_val;
+    int N_M1_TG = N_M1_TG_val;
+    int blockM_gate = BM_GATE_val * TM;  // rows per gate GEMV TG
+
+    uint tg_x = threadgroup_position_in_grid.x;
+    uint sgid = simdgroup_index_in_threadgroup;  // 0..7
+    uint slid = thread_index_in_simdgroup;       // 0..31
+
+    if (tg_x < (uint)N_OPROJ_TG) {{
+        // ════════════════════════════════════════════════════════════════
+        // O_PROJ GEMV: 8-bit quantized, M=4096, K_attn=8192
+        // Each TG handles 32 rows (8 SGs × 4 rows each)
+        // 8-bit affine dequant: result = scale * Σ(x*w) + bias * Σ(x)
+        // ════════════════════════════════════════════════════════════════
+        const int blockM = 32;  // 8 SGs × TM=4
+        const int VPT = 8;     // values per thread per iteration
+        const int BLOCK_SIZE = 256;  // 32 threads × 8 values
+
+        int out_row = int(tg_x) * blockM + int(sgid) * TM;
+        if (out_row >= M) return;
+        out_row = (out_row + TM <= M) ? out_row : (M - TM);
+
+        threadgroup float tgp_x2[8];
+
+        // 8-bit GEMV K-loop: K-outer, tm-inner (input loaded once, reused across TM rows)
+        float acc[TM] = {{0.0f, 0.0f, 0.0f, 0.0f}};
+        float result[TM];
+        int K_groups = K_attn / {gs};
+
+        // Per-row weight/scale/bias pointers
+        const device uint8_t* ws0 = (const device uint8_t*)W_oproj + (long)(out_row + 0) * K_attn + slid * VPT;
+        const device uint8_t* ws1 = (const device uint8_t*)W_oproj + (long)(out_row + 1) * K_attn + slid * VPT;
+        const device uint8_t* ws2 = (const device uint8_t*)W_oproj + (long)(out_row + 2) * K_attn + slid * VPT;
+        const device uint8_t* ws3 = (const device uint8_t*)W_oproj + (long)(out_row + 3) * K_attn + slid * VPT;
+
+        const device {sc_t}* sc0 = (const device {sc_t}*)S_oproj + (long)(out_row + 0) * K_groups + slid / {oproj_slid_divisor};
+        const device {sc_t}* sc1 = (const device {sc_t}*)S_oproj + (long)(out_row + 1) * K_groups + slid / {oproj_slid_divisor};
+        const device {sc_t}* sc2 = (const device {sc_t}*)S_oproj + (long)(out_row + 2) * K_groups + slid / {oproj_slid_divisor};
+        const device {sc_t}* sc3 = (const device {sc_t}*)S_oproj + (long)(out_row + 3) * K_groups + slid / {oproj_slid_divisor};
+
+        const device {sc_t}* bi0 = (const device {sc_t}*)B_oproj + (long)(out_row + 0) * K_groups + slid / {oproj_slid_divisor};
+        const device {sc_t}* bi1 = (const device {sc_t}*)B_oproj + (long)(out_row + 1) * K_groups + slid / {oproj_slid_divisor};
+        const device {sc_t}* bi2 = (const device {sc_t}*)B_oproj + (long)(out_row + 2) * K_groups + slid / {oproj_slid_divisor};
+        const device {sc_t}* bi3 = (const device {sc_t}*)B_oproj + (long)(out_row + 3) * K_groups + slid / {oproj_slid_divisor};
+
+        int xb = slid * VPT;
+
+        for (int k = 0; k < K_attn; k += BLOCK_SIZE) {{
+            // Load input ONCE per K-block
+            float xv[VPT];
+            float xsum = 0.0f;
+            for (int i = 0; i < VPT; i++) {{
+                xv[i] = float(attn_out[xb + i]);
+                xsum += xv[i];
+            }}
+
+            // Row 0
+            {{ float wacc = 0.0f;
+              for (int i = 0; i < VPT; i++) wacc += xv[i] * float(ws0[i]);
+              acc[0] += float(*sc0) * wacc + xsum * float(*bi0);
+              ws0 += BLOCK_SIZE; sc0 += {oproj_sc_stride}; bi0 += {oproj_sc_stride}; }}
+            // Row 1
+            {{ float wacc = 0.0f;
+              for (int i = 0; i < VPT; i++) wacc += xv[i] * float(ws1[i]);
+              acc[1] += float(*sc1) * wacc + xsum * float(*bi1);
+              ws1 += BLOCK_SIZE; sc1 += {oproj_sc_stride}; bi1 += {oproj_sc_stride}; }}
+            // Row 2
+            {{ float wacc = 0.0f;
+              for (int i = 0; i < VPT; i++) wacc += xv[i] * float(ws2[i]);
+              acc[2] += float(*sc2) * wacc + xsum * float(*bi2);
+              ws2 += BLOCK_SIZE; sc2 += {oproj_sc_stride}; bi2 += {oproj_sc_stride}; }}
+            // Row 3
+            {{ float wacc = 0.0f;
+              for (int i = 0; i < VPT; i++) wacc += xv[i] * float(ws3[i]);
+              acc[3] += float(*sc3) * wacc + xsum * float(*bi3);
+              ws3 += BLOCK_SIZE; sc3 += {oproj_sc_stride}; bi3 += {oproj_sc_stride}; }}
+
+            xb += BLOCK_SIZE;
+        }}
+
+        for (int tm = 0; tm < TM; tm++)
+            result[tm] = simd_sum(acc[tm]);
+
+        // Epilogue: addmm + x² + h_scaled + h_out
+        float x2_acc = 0.0f;
+        if (slid == 0) {{
+            for (int tm = 0; tm < TM; tm++) {{
+                int k = out_row + tm;
+                float h = result[tm] + float(residual[k]);
+                x2_acc += h * h;
+                h_scaled[k] = static_cast<bfloat16_t>(h * float(w_rms[k]));
+                h_out[k] = static_cast<bfloat16_t>(h);
+            }}
+        }}
+
+        // TG x² reduction: 8 SGs → 1 value
+        if (slid == 0) tgp_x2[sgid] = x2_acc;
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (sgid == 0 && slid == 0) {{
+            float total = 0.0f;
+            for (int s = 0; s < 8; s++) total += tgp_x2[s];
+            x2_partials[tg_x] = total;
+        }}
+
+    }} else if (tg_x < (uint)(N_OPROJ_TG + N_M1_TG)) {{
+        // ════════════════════════════════════════════════════════════════
+        // M1 GEMV: bf16, M=E, K=K_attn — M1 × attn_out → gate_part_a
+        // ════════════════════════════════════════════════════════════════
+        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);
+
+        float result[TM] = {{0.0f, 0.0f, 0.0f, 0.0f}};
+        int bn = int(slid) * TN;
+        int n_iter = K_attn / blockN;
+
+        for (int i = 0; i < n_iter; i++) {{
+            float v[TN];
+            for (int tn = 0; tn < TN; tn++)
+                v[tn] = float(attn_out[bn + tn]);
+            for (int tm = 0; tm < TM; tm++) {{
+                float acc = 0.0f;
+                for (int tn = 0; tn < TN; tn++)
+                    acc += float(M1[(out_row + tm) * K_attn + bn + tn]) * v[tn];
+                result[tm] += acc;
+            }}
+            bn += blockN;
+        }}
+
+        // Handle remainder if K_attn not divisible by blockN
+        if (K_attn > n_iter * blockN) {{
+            for (int tm = 0; tm < TM; tm++) {{
+                for (int tn = 0; tn < TN; tn++) {{
+                    if (bn + tn < K_attn)
+                        result[tm] += float(M1[(out_row + tm) * K_attn + bn + tn])
+                                    * float(attn_out[bn + tn]);
+                }}
+            }}
+        }}
+
+        for (int tm = 0; tm < TM; tm++)
+            result[tm] = simd_sum(result[tm]);
+
+        if (slid == 0) {{
+            for (int tm = 0; tm < TM; tm++) {{
+                int e = out_row + tm;
+                if (e < E_CONST)
+                    gate_part_a[e] = result[tm];
+            }}
+        }}
+
+    }} else {{
+        // ════════════════════════════════════════════════════════════════
+        // W_FUSED GEMV: bf16, M=E, K=K_hidden — W_fused × residual → gate_part_b
+        // ════════════════════════════════════════════════════════════════
+        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);
+
+        float result[TM] = {{0.0f, 0.0f, 0.0f, 0.0f}};
+        int bn = int(slid) * TN;
+        int n_iter = K_hidden / blockN;
+
+        for (int i = 0; i < n_iter; i++) {{
+            float v[TN];
+            for (int tn = 0; tn < TN; tn++)
+                v[tn] = float(residual[bn + tn]);
+            for (int tm = 0; tm < TM; tm++) {{
+                float acc = 0.0f;
+                for (int tn = 0; tn < TN; tn++)
+                    acc += float(W_fused[(out_row + tm) * K_hidden + bn + tn]) * v[tn];
+                result[tm] += acc;
+            }}
+            bn += blockN;
+        }}
+
+        if (K_hidden > n_iter * blockN) {{
+            for (int tm = 0; tm < TM; tm++) {{
+                for (int tn = 0; tn < TN; tn++) {{
+                    if (bn + tn < K_hidden)
+                        result[tm] += float(W_fused[(out_row + tm) * K_hidden + bn + tn])
+                                    * float(residual[bn + tn]);
+                }}
+            }}
+        }}
+
+        for (int tm = 0; tm < TM; tm++)
+            result[tm] = simd_sum(result[tm]);
+
+        if (slid == 0) {{
+            for (int tm = 0; tm < TM; tm++) {{
+                int e = out_row + tm;
+                if (e < E_CONST)
+                    gate_part_b[e] = result[tm];
+            }}
+        }}
+    }}
+"""
+
+
+_custom_oproj_8bit_kernels = {}
+
+
+def _get_custom_oproj_8bit_kernel(n_experts=64, group_size=64, scale_bf16=True):
+    key = (n_experts, group_size, scale_bf16)
+    if key not in _custom_oproj_8bit_kernels:
+        sc_tag = "_bf16sc" if scale_bf16 else ""
+        _custom_oproj_8bit_kernels[key] = mx.fast.metal_kernel(
+            name=f"custom_oproj_gate_gemv_8bit_e{n_experts}_gs{group_size}{sc_tag}",
+            input_names=[
+                "W_oproj", "S_oproj", "B_oproj",   # o_proj 8-bit weights
+                "attn_out",                          # (K_attn,) bf16
+                "residual",                          # (K,) bf16
+                "w_rms",                             # (K,) bf16 — RMSNorm weight
+                "M1",                                # (E, K_attn) bf16
+                "W_fused",                           # (E, K) bf16
+                "M_val", "K_attn_val", "K_hidden_val",
+                "N_OPROJ_TG_val", "N_M1_TG_val", "BM_GATE_val",
+            ],
+            output_names=["h_scaled", "h_out", "x2_partials",
+                          "gate_part_a", "gate_part_b"],
+            source=_gen_custom_oproj_8bit_source(n_experts, group_size, scale_bf16),
+        )
+    return _custom_oproj_8bit_kernels[key]
+
+
+def fused_custom_oproj_8bit(W_oproj, S_oproj, B_oproj,
+                             attn_out, residual, w_rms,
+                             M1, W_fused,
+                             M, K_attn, K_hidden=None,
+                             n_experts=64, gate_bm=8,
+                             group_size=64):
+    """Dispatch the fused 8-bit o_proj + bf16 gate GEMVs kernel.
+
+    Three GEMV types in one dispatch:
+    - TGs 0..N_OPROJ_TG-1: 8-bit o_proj GEMV → h_scaled, h_out, x2_partials
+    - TGs N_OPROJ_TG..+N_M1_TG: bf16 M1 GEMV → gate_part_a
+    - TGs +N_M1_TG..end: bf16 W_fused GEMV → gate_part_b
+
+    Args:
+        W_oproj/S_oproj/B_oproj: 8-bit quantized o_proj weights
+        attn_out: (K_attn,) bf16 — pre-o_proj attention output
+        residual: (K,) bf16 — input residual
+        w_rms: (K,) bf16 — post_attention_layernorm weight
+        M1: (E, K_attn) bf16 — precomputed W_fused @ W_oproj
+        W_fused: (E, K) bf16 — precomputed dequant(W_gate) * w_rms
+        M: hidden size (4096)
+        K_attn: attention output dim (8192)
+        K_hidden: hidden size for W_fused (defaults to M)
+        gate_bm: SGs per gate TG (1,2,4,8). Controls TG count.
+
+    Returns:
+        (h_scaled, h_out, x2_partials, gate_part_a, gate_part_b)
+    """
+    M = int(M)
+    K_attn = int(K_attn)
+    K_hidden = int(K_hidden) if K_hidden is not None else M
+    scale_bf16 = (S_oproj.dtype == mx.bfloat16)
+
+    kern = _get_custom_oproj_8bit_kernel(n_experts, group_size, scale_bf16)
+
+    n_oproj_tg = ceil_div(M, 32)  # 128 for M=4096
+    blockM_gate = gate_bm * 4     # rows per gate GEMV TG
+    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, residual, w_rms, M1, W_fused,
+                M, K_attn, K_hidden, n_oproj_tg, n_m1_tg, gate_bm],
+        output_shapes=[(M,), (M,), (n_oproj_tg,), (n_experts,), (n_experts,)],
+        output_dtypes=[mx.bfloat16, mx.bfloat16, mx.float32,
+                       mx.float32, mx.float32],
+        grid=(total_tg * 32, 8, 1),    # total threads: total_tg * 256
+        threadgroup=(32, 8, 1),          # 256 threads per TG (8 SGs of 32)
+    )
+    return results[0], results[1], results[2], results[3], results[4]

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

@@ -0,0 +1,136 @@
+"""Custom SDPA Pass 1 for GQA attention (Dispatch 4).
+
+Online softmax over N key positions, strided by blocks.
+Constants baked into Metal source. Per-call varying values (N_keys, alloc_len)
+passed via a params array.
+"""
+
+import mlx.core as mx
+
+
+def _gen_sdpa_pass1_source(D=256, H_q=16, H_kv=2, blocks=128):
+    EPT = D // 32
+    gqa_factor = H_q // H_kv
+    scale = D ** -0.5
+
+    return f"""
+    const int D_DIM = {D};
+    const int EPT = {EPT};
+    const int H_Q = {H_q};
+    const int H_KV = {H_kv};
+    const int N_BLOCKS = {blocks};
+    const int GQA_FACTOR = {gqa_factor};
+    const float SCALE = {scale}f;
+
+    uint simd_lid = thread_index_in_simdgroup;
+    uint kv_head_idx = threadgroup_position_in_grid.x;
+    uint gqa_lane = thread_index_in_threadgroup / 32;
+
+    uint block_idx = threadgroup_position_in_grid.z % (uint)N_BLOCKS;
+    uint batch_idx = threadgroup_position_in_grid.z / (uint)N_BLOCKS;
+
+    int q_head_idx = (int)kv_head_idx * GQA_FACTOR + (int)gqa_lane;
+    int N_val = params[0];
+    int alloc = params[1];
+
+    int q_offset = ((int)batch_idx * H_Q + q_head_idx) * D_DIM;
+
+    int kv_head_offset = ((int)batch_idx * H_KV + (int)kv_head_idx) * alloc * D_DIM;
+    int k_offset = kv_head_offset + (int)block_idx * D_DIM;
+    int v_offset = kv_head_offset + (int)block_idx * D_DIM;
+
+    int out_head_offset = ((int)batch_idx * H_Q + q_head_idx);
+    int o_offset = (out_head_offset * N_BLOCKS + (int)block_idx) * D_DIM;
+    int s_offset = out_head_offset * N_BLOCKS + (int)block_idx;
+
+    float q[{EPT}];
+    for (int i = 0; i < EPT; i++) {{
+        q[i] = SCALE * (float)queries[q_offset + (int)simd_lid * EPT + i];
+    }}
+
+    float max_score = -__FLT_MAX__;
+    float sum_exp = 0.0f;
+    float o[{EPT}] = {{0}};
+
+    int k_stride = N_BLOCKS * D_DIM;
+
+    for (int pos = (int)block_idx; pos < N_val; pos += N_BLOCKS) {{
+        float score = 0.0f;
+        for (int i = 0; i < EPT; i++) {{
+            score += q[i] * (float)keys[k_offset + (int)simd_lid * EPT + i];
+        }}
+        score = simd_sum(score);
+
+        float new_max = metal::max(max_score, score);
+        float factor = metal::fast::exp(max_score - new_max);
+        float exp_score = metal::fast::exp(score - new_max);
+
+        max_score = new_max;
+        sum_exp = sum_exp * factor + exp_score;
+
+        for (int i = 0; i < EPT; i++) {{
+            o[i] = o[i] * factor + exp_score * (float)values[v_offset + (int)simd_lid * EPT + i];
+        }}
+
+        k_offset += k_stride;
+        v_offset += k_stride;
+    }}
+
+    if (simd_lid == 0) {{
+        sums[s_offset] = sum_exp;
+        maxs[s_offset] = max_score;
+    }}
+
+    for (int i = 0; i < EPT; i++) {{
+        o_partials[o_offset + (int)simd_lid * EPT + i] = static_cast<bfloat16_t>(o[i]);
+    }}
+"""
+
+
+_pass1_cache = {}
+
+
+def _get_pass1_kernel(D, H_q, H_kv, blocks, gqa_factor):
+    key = (D, H_q, H_kv, blocks)
+    if key not in _pass1_cache:
+        _pass1_cache[key] = mx.fast.metal_kernel(
+            name=f"sdpa_pass1_D{D}_Hq{H_q}_Hkv{H_kv}_b{blocks}",
+            input_names=["queries", "keys", "values", "params"],
+            output_names=["o_partials", "sums", "maxs"],
+            source=_gen_sdpa_pass1_source(D, H_q, H_kv, blocks),
+        )
+    return _pass1_cache[key]
+
+
+def custom_sdpa_pass1(queries, keys, values, scale, H_q, H_kv, D,
+                      blocks=128, batch_size=1, N=None, alloc_len=None,
+                      scalars=None):
+    B = batch_size
+    gqa_factor = H_q // H_kv
+
+    kern = _get_pass1_kernel(D, H_q, H_kv, blocks, gqa_factor)
+
+    if N is None:
+        N = keys.shape[2]
+    if alloc_len is None:
+        alloc_len = N
+
+    # Pack per-call varying values into a params array (always a pointer)
+    params = mx.array([int(N), int(alloc_len)], dtype=mx.int32)
+
+    results = kern(
+        inputs=[queries, keys, values, params],
+        output_shapes=[
+            (B * H_q * blocks * D,),
+            (B * H_q * blocks,),
+            (B * H_q * blocks,),
+        ],
+        output_dtypes=[mx.bfloat16, mx.float32, mx.float32],
+        grid=(H_kv * 32, gqa_factor, blocks * B),
+        threadgroup=(32, gqa_factor, 1),
+    )
+
+    o_partials = results[0].reshape(B, H_q, blocks, D)
+    sums = results[1].reshape(B, H_q, blocks)
+    maxs = results[2].reshape(B, H_q, blocks)
+    return o_partials, sums, maxs

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

@@ -0,0 +1,127 @@
+"""Custom SDPA Pass 2 + gate multiply for GQA attention (Dispatch 5).
+
+32-SG block-parallel reduction + V_SPLIT for high GPU utilization.
+All constants baked into Metal source (no scalar kernel inputs).
+"""
+
+import mlx.core as mx
+
+
+def _gen_sdpa_pass2_gate_source(D=256, V_SPLIT=4, H_q=16, blocks=128):
+    BN = 32
+    V_PER_SPLIT = D // V_SPLIT
+    EPT = V_PER_SPLIT // BN
+
+    return f"""
+    const int D_DIM = {D};
+    const int V_SPLIT = {V_SPLIT};
+    const int V_PER_SPLIT = {V_PER_SPLIT};
+    const int EPT = {EPT};
+    const int BN = {BN};
+    const int H_Q = {H_q};
+    const int N_BLOCKS = {blocks};
+
+    uint tg_idx = threadgroup_position_in_grid.x;
+    uint simd_gid = simdgroup_index_in_threadgroup;
+    uint simd_lid = thread_index_in_simdgroup;
+    uint b_idx = threadgroup_position_in_grid.z;
+
+    int head_idx = (int)tg_idx / V_SPLIT;
+    int split_idx = (int)tg_idx % V_SPLIT;
+    int v_offset = split_idx * V_PER_SPLIT;
+
+    int head_base = ((int)b_idx * H_Q + head_idx);
+
+    int p_base = head_base * N_BLOCKS * D_DIM
+               + (int)simd_gid * D_DIM
+               + v_offset + (int)simd_lid * EPT;
+    int ms_base = head_base * N_BLOCKS;
+
+    // ── Phase 1: Find global max across all blocks ──
+    float local_max = -__FLT_MAX__;
+    for (int b = 0; b < N_BLOCKS / BN; ++b) {{
+        local_max = metal::max(local_max, maxs[ms_base + (int)simd_lid + BN * b]);
+    }}
+    float global_max = simd_max(local_max);
+
+    // ── Phase 2: Compute global sum_exp ──
+    float local_sum = 0.0f;
+    for (int b = 0; b < N_BLOCKS / BN; ++b) {{
+        float factor = metal::fast::exp(maxs[ms_base + (int)simd_lid + BN * b] - global_max);
+        local_sum += factor * sums[ms_base + (int)simd_lid + BN * b];
+    }}
+    float global_sum = simd_sum(local_sum);
+
+    // ── Phase 3: Accumulate V-partials (block-parallel via SGs) ──
+    float o[EPT] = {{0}};
+    for (int b = 0; b < N_BLOCKS / BN; ++b) {{
+        float factor = metal::fast::exp(maxs[ms_base + (int)simd_gid] - global_max);
+        for (int i = 0; i < EPT; i++) {{
+            o[i] += factor * (float)o_partials[p_base + i];
+        }}
+        ms_base += BN;
+        p_base += BN * D_DIM;
+    }}
+
+    // ── Phase 4: Shared memory transpose + reduce across SGs ──
+    threadgroup float tg_mem[BN * BN];
+
+    for (int i = 0; i < EPT; i++) {{
+        tg_mem[(int)simd_lid * BN + (int)simd_gid] = o[i];
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        o[i] = simd_sum(tg_mem[(int)simd_gid * BN + (int)simd_lid]);
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }}
+
+    // ── Phase 5: Normalize + gate multiply + write (simd_lid == 0 only) ──
+    if (simd_lid == 0) {{
+        float inv_sum = (global_sum > 0.0f) ? (1.0f / global_sum) : 0.0f;
+        int out_base = (int)b_idx * H_Q * D_DIM + head_idx * D_DIM
+                     + v_offset + (int)simd_gid * EPT;
+        for (int i = 0; i < EPT; i++) {{
+            float val = o[i] * inv_sum;
+            float gate = gate_sigmoid[out_base + i];
+            attn_output[out_base + i] = static_cast<bfloat16_t>(val * gate);
+        }}
+    }}
+"""
+
+
+_pass2_cache = {}
+
+
+def _get_pass2_kernel(D, V_SPLIT, H_q, blocks):
+    key = (D, V_SPLIT, H_q, blocks)
+    if key not in _pass2_cache:
+        _pass2_cache[key] = mx.fast.metal_kernel(
+            name=f"sdpa_pass2_gate_D{D}_vs{V_SPLIT}_Hq{H_q}_b{blocks}",
+            input_names=["o_partials", "sums", "maxs", "gate_sigmoid"],
+            output_names=["attn_output"],
+            source=_gen_sdpa_pass2_gate_source(D, V_SPLIT, H_q, blocks),
+        )
+    return _pass2_cache[key]
+
+
+def custom_sdpa_pass2_gate(o_partials, sums, maxs, gate_sigmoid,
+                           H_q, D, blocks=128, V_SPLIT=4, batch_size=1,
+                           scalars=None):
+    B = batch_size
+
+    kern = _get_pass2_kernel(D, V_SPLIT, H_q, blocks)
+
+    o_flat = o_partials.reshape(B * H_q * blocks * D)
+    s_flat = sums.reshape(B * H_q * blocks)
+    m_flat = maxs.reshape(B * H_q * blocks)
+    g_flat = gate_sigmoid.reshape(B * H_q * D)
+
+    n_tgs = H_q * V_SPLIT
+    BN = 32
+    results = kern(
+        inputs=[o_flat, s_flat, m_flat, g_flat],
+        output_shapes=[(B * H_q * D,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(n_tgs * BN, BN, B),
+        threadgroup=(BN, BN, 1),
+    )
+
+    return results[0].reshape(B, 1, H_q * D)

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

@@ -0,0 +1,288 @@
+"""Fused GDN projections for Qwen3.5-35B-A3B (Dispatch 2).
+
+Single dispatch fuses 4 quantized 8-bit GEMVs + depthwise conv1d + activations:
+  - in_proj_qkv (8192×2048): GEMV → conv1d(4-tap) → SiLU → write bf16 + cache update
+  - in_proj_z   (4096×2048): GEMV → SiLU → write f32
+  - in_proj_b   (32×2048):   GEMV → sigmoid → write f32 (beta for GDN kernel)
+  - in_proj_a   (32×2048):   GEMV → g=exp(-exp(A_log)*softplus(a+dt_bias)) → write f32
+
+All 4 projection weight matrices are pre-merged into one contiguous buffer
+(W_merged, S_merged, B_merged) for better memory locality and cache behavior.
+Merging is done offline at patch time by _patch_gdn_proj_weights().
+
+B/A epilogues compute g and beta in-kernel, eliminating ~8 micro-dispatches
+that gated_delta_update would otherwise generate (sigmoid, exp, log, etc.).
+The caller can pass g/beta directly to gated_delta_kernel.
+
+TG-level multiplexing: tgid.y routes to different epilogues.
+Each TG: 64 threads = 2 SGs of 32, produces 8 output rows (4 per SG).
+Standard 8-bit affine GEMV: result = scale * Σ(x[i]*w[i]) + bias * Σ(x[i])
+
+Grid: (32, total_tg * 2, B), TG: (32, 2, 1)
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_fused_gdn_projections_source(K, N_QKV, N_Z, N_B, N_A, group_size=64):
+    """Generate Metal source for fused GDN projections with merged weights.
+
+    All constants baked into Metal source (no scalar kernel inputs).
+    """
+    gs = int(group_size)
+    sc_stride = 256 // gs
+    slid_div = gs // 8
+    N_TOTAL = N_QKV + N_Z + N_B + N_A
+    K_groups = K // gs
+    N_QKV_TG = ceil_div(N_QKV, 8)
+    N_Z_TG = ceil_div(N_Z, 8)
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int GROUP_SIZE = {gs};
+    const int SC_STRIDE = {sc_stride};
+    const int SLID_DIV = {slid_div};
+    const int K = {K};
+    const int K_groups = {K_groups};
+    const int N_QKV = {N_QKV};
+    const int N_Z = {N_Z};
+    const int N_B = {N_B};
+    const int N_TOTAL = {N_TOTAL};
+    const int N_QKV_TG = {N_QKV_TG};
+    const int N_Z_TG = {N_Z_TG};
+    const int N_B_TG = {ceil_div(N_B, 8)};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;  // 0 or 1
+    uint slid = thread_index_in_simdgroup;       // 0..31
+    int b_idx = tgid.z;
+
+    int tg = tgid.y;
+
+    // ─── Determine region and absolute out_row in merged matrix ───
+    int out_row;
+    int region;  // 0=QKV, 1=Z, 2=B, 3=A
+
+    if (tg < N_QKV_TG) {{
+        region = 0;
+        out_row = tg * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_QKV_TG + N_Z_TG) {{
+        region = 1;
+        out_row = N_QKV + (tg - N_QKV_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_QKV_TG + N_Z_TG + N_B_TG) {{
+        region = 2;
+        out_row = N_QKV + N_Z + (tg - N_QKV_TG - N_Z_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else {{
+        region = 3;
+        out_row = N_QKV + N_Z + N_B + (tg - N_QKV_TG - N_Z_TG - N_B_TG) * 8 + sgid * RESULTS_PER_SG;
+    }}
+
+    if (out_row >= N_TOTAL) return;
+
+    // ─── Single pointer into merged weight buffer ───
+    const device uint8_t* ws = (const device uint8_t*)W_merged + (long)out_row * K + slid * VALUES_PER_THREAD;
+    const device bfloat16_t* sc = (const device bfloat16_t*)S_merged + (long)out_row * K_groups + slid / SLID_DIV;
+    const device bfloat16_t* bi = (const device bfloat16_t*)B_merged + (long)out_row * K_groups + slid / SLID_DIV;
+
+    // ─── 8-bit GEMV K-loop (unified for all regions) ───
+    float result[4] = {{0, 0, 0, 0}};
+    int x_base = b_idx * K + slid * VALUES_PER_THREAD;
+
+    for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{
+        float x_thread[8];
+        float xsum = 0;
+        for (int i = 0; i < 8; i++) {{
+            float xi = float(x[x_base + i]);
+            x_thread[i] = xi;
+            xsum += xi;
+        }}
+
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            const device uint8_t* w = ws + row * K;
+            float s_val = float(sc[row * K_groups]);
+            float b_val = float(bi[row * K_groups]);
+            float accum = 0;
+            for (int i = 0; i < 8; i++) {{
+                accum += x_thread[i] * float(w[i]);
+            }}
+            result[row] += s_val * accum + xsum * b_val;
+        }}
+
+        ws += BLOCK_SIZE;
+        sc += SC_STRIDE;
+        bi += SC_STRIDE;
+        x_base += BLOCK_SIZE;
+    }}
+
+    // ─── Reduction ───
+    for (int row = 0; row < RESULTS_PER_SG; row++) {{
+        result[row] = simd_sum(result[row]);
+    }}
+
+    // ─── Region-specific epilogues ───
+    // After simd_sum, all 32 threads have result[0..3].
+    // Threads 0-3 each handle one output row.
+
+    if (region == 0) {{
+        // ═══ QKV: conv1d(4-tap) + SiLU + cache update ═══
+        int c = out_row + (int)slid;  // channel index (= absolute row for QKV)
+        if (slid < (uint)RESULTS_PER_SG && c < N_QKV) {{
+            float qkv_val = result[slid];
+
+            int conv_dim = N_QKV;
+            long cs_base = (long)b_idx * 3 * conv_dim;
+            float s0 = float(conv_state[cs_base + 0 * conv_dim + c]);
+            float s1 = float(conv_state[cs_base + 1 * conv_dim + c]);
+            float s2 = float(conv_state[cs_base + 2 * conv_dim + c]);
+
+            float conv_out = float(conv_w[c * 4 + 0]) * s0
+                           + float(conv_w[c * 4 + 1]) * s1
+                           + float(conv_w[c * 4 + 2]) * s2
+                           + float(conv_w[c * 4 + 3]) * qkv_val;
+
+            float silu_out = conv_out / (1.0f + metal::exp(-conv_out));
+
+            conv_state_out[cs_base + 0 * conv_dim + c] = static_cast<bfloat16_t>(s1);
+            conv_state_out[cs_base + 1 * conv_dim + c] = static_cast<bfloat16_t>(s2);
+            conv_state_out[cs_base + 2 * conv_dim + c] = static_cast<bfloat16_t>(qkv_val);
+
+            qkv_out[b_idx * conv_dim + c] = static_cast<bfloat16_t>(silu_out);
+        }}
+
+    }} else if (region == 1) {{
+        // ═══ Z: SiLU → write f32 ═══
+        int z_row = out_row - N_QKV + (int)slid;
+        if (slid < (uint)RESULTS_PER_SG && z_row < N_Z) {{
+            float val = result[slid];
+            float silu_val = val / (1.0f + metal::exp(-val));
+            z_silu_out[b_idx * N_Z + z_row] = silu_val;
+        }}
+
+    }} else if (region == 2) {{
+        // ═══ B: sigmoid(result) → beta (f32) ═══
+        int b_row = out_row - N_QKV - N_Z + (int)slid;
+        if (slid < (uint)RESULTS_PER_SG && b_row < N_B) {{
+            float val = result[slid];
+            float beta = 1.0f / (1.0f + metal::exp(-val));
+            b_out[b_idx * N_B + b_row] = beta;
+        }}
+
+    }} else {{
+        // ═══ A: g = exp(-exp(A_log) * softplus(a + dt_bias)) → f32 ═══
+        int a_row = out_row - N_QKV - N_Z - N_B + (int)slid;
+        int N_A = N_TOTAL - N_QKV - N_Z - N_B;
+        if (slid < (uint)RESULTS_PER_SG && a_row < N_A) {{
+            float a_val = result[slid];
+            float dt = float(dt_bias_arr[a_row]);
+            float x_g = a_val + dt;
+            // softplus(x) = log(1 + exp(x)), with x>20 shortcut for numerical stability
+            float sp = (x_g > 20.0f) ? x_g : metal::log(1.0f + metal::exp(x_g));
+            float g_val = metal::exp(-metal::exp(float(A_log_arr[a_row])) * sp);
+            a_out[b_idx * N_A + a_row] = g_val;
+        }}
+    }}
+"""
+
+
+_fused_gdn_proj_cache = {}
+
+
+def _get_fused_gdn_proj_kernel(K, N_QKV, N_Z, N_B, N_A, group_size=64):
+    key = (K, N_QKV, N_Z, N_B, N_A, group_size)
+    if key not in _fused_gdn_proj_cache:
+        _fused_gdn_proj_cache[key] = mx.fast.metal_kernel(
+            name=f"fused_gdn_proj_K{K}_NQKV{N_QKV}_NZ{N_Z}_NB{N_B}_NA{N_A}",
+            input_names=[
+                "x",
+                "W_merged", "S_merged", "B_merged",
+                "conv_state", "conv_w",
+                "A_log_arr", "dt_bias_arr",
+            ],
+            output_names=["qkv_out", "z_silu_out", "b_out", "a_out", "conv_state_out"],
+            source=_gen_fused_gdn_projections_source(K, N_QKV, N_Z, N_B, N_A, group_size),
+        )
+    return _fused_gdn_proj_cache[key]
+
+
+def fused_gdn_projections(
+    x,
+    W_merged, S_merged, B_merged,
+    proj_dims,
+    conv_state, conv_weights,
+    A_log, dt_bias,
+    batch_size=1,
+):
+    """Fused GDN projections: 4 GEMVs + conv1d + activations + g/beta.
+
+    Uses pre-merged contiguous weight buffers for all 4 projections.
+    B epilogue computes beta = sigmoid(b) in f32.
+    A epilogue computes g = exp(-exp(A_log) * softplus(a + dt_bias)) in f32.
+    Caller passes g/beta directly to gated_delta_kernel (no micro-dispatches).
+
+    Args:
+        x: [B, 1, K] bf16 — post-RMSNorm hidden state
+        W_merged: [N_TOTAL, K/4] uint32 — merged quantized weights
+        S_merged: [N_TOTAL, K/gs] bf16 — merged scales
+        B_merged: [N_TOTAL, K/gs] bf16 — merged biases
+        proj_dims: (N_QKV, N_Z, N_B, N_A) — per-projection output dims
+        conv_state:  [B, 3, conv_dim] bf16 — previous 3 timesteps
+        conv_weights: [conv_dim, 4, 1] or [conv_dim, 4] bf16 — depthwise conv filters
+        A_log: [Hv] f32 — GDN decay log-parameter
+        dt_bias: [Hv] f32 — GDN time constant bias
+        batch_size: int
+
+    Returns:
+        qkv_conv_silu: [B, 1, N_QKV] bf16 — post-conv, post-SiLU
+        z_silu:        [B, 1, N_Z] f32  — post-SiLU
+        beta:          [B, 1, N_B] f32  — sigmoid(b), ready for GDN kernel
+        g:             [B, 1, N_A] f32  — gating, ready for GDN kernel
+        conv_state_out: [B, 3, N_QKV] bf16
+    """
+    B = batch_size
+
+    N_QKV, N_Z, N_B, N_A = proj_dims
+    K = x.shape[-1]
+
+    kern = _get_fused_gdn_proj_kernel(K, N_QKV, N_Z, N_B, N_A)
+
+    N_QKV_TG = ceil_div(N_QKV, 8)
+    N_Z_TG = ceil_div(N_Z, 8)
+    N_B_TG = ceil_div(N_B, 8)
+    N_A_TG = ceil_div(N_A, 8)
+    total_tg = N_QKV_TG + N_Z_TG + N_B_TG + N_A_TG
+
+    conv_w_flat = conv_weights.reshape(-1, 4) if conv_weights.ndim == 3 else conv_weights
+    x_flat = x.reshape(B, K)
+
+    results = kern(
+        inputs=[
+            x_flat,
+            W_merged, S_merged, B_merged,
+            conv_state, conv_w_flat,
+            A_log, dt_bias,
+        ],
+        output_shapes=[
+            (B * N_QKV,),              # qkv_out
+            (B * N_Z,),                # z_silu_out
+            (B * N_B,),                # beta_out (f32)
+            (B * N_A,),                # g_out (f32)
+            (B * 3 * N_QKV,),          # conv_state_out
+        ],
+        output_dtypes=[mx.bfloat16, mx.float32, mx.float32, mx.float32, mx.bfloat16],
+        grid=(32, total_tg * 2, B),
+        threadgroup=(32, 2, 1),
+    )
+
+    qkv_out = results[0].reshape(B, 1, N_QKV)
+    z_silu = results[1].reshape(B, 1, N_Z)
+    beta = results[2].reshape(B, 1, N_B)
+    g = results[3].reshape(B, 1, N_A)
+    conv_state_out = results[4].reshape(B, 3, N_QKV)
+
+    return qkv_out, z_silu, beta, g, conv_state_out

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

@@ -0,0 +1,178 @@
+"""Fused GQA projections for Qwen3.5 (Dispatch 1).
+
+Single dispatch fuses 4 quantized 8-bit GEMVs with region-specific epilogues:
+  - q_proj queries: GEMV → raw bf16 write
+  - q_proj gate:    GEMV → sigmoid → f32 write
+  - k_proj:         GEMV → raw bf16 write
+  - v_proj:         GEMV → raw bf16 write
+
+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_fused_gqa_projections_source(K, N_Q, N_GATE, N_K, N_V, 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)
+
+    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_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;
+    int 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;
+
+    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;
+
+    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;
+    }}
+
+    for (int row = 0; row < RESULTS_PER_SG; row++) {{
+        result[row] = simd_sum(result[row]);
+    }}
+
+    if (region == 0) {{
+        int q_row = out_row + (int)slid;
+        if (slid < (uint)RESULTS_PER_SG && q_row < N_Q) {{
+            q_out[b_idx * N_Q + q_row] = static_cast<bfloat16_t>(result[slid]);
+        }}
+    }} else if (region == 1) {{
+        int g_row = out_row - N_Q + (int)slid;
+        if (slid < (uint)RESULTS_PER_SG && g_row < N_GATE) {{
+            float val = result[slid];
+            float sig = 1.0f / (1.0f + metal::exp(-val));
+            gate_out[b_idx * N_GATE + g_row] = sig;
+        }}
+    }} else if (region == 2) {{
+        int k_row = out_row - N_Q - N_GATE + (int)slid;
+        if (slid < (uint)RESULTS_PER_SG && k_row < N_K) {{
+            k_out[b_idx * N_K + k_row] = static_cast<bfloat16_t>(result[slid]);
+        }}
+    }} else {{
+        int v_row = out_row - N_Q - N_GATE - N_K + (int)slid;
+        int N_V = N_TOTAL - N_Q - N_GATE - N_K;
+        if (slid < (uint)RESULTS_PER_SG && v_row < N_V) {{
+            v_out[b_idx * N_V + v_row] = static_cast<bfloat16_t>(result[slid]);
+        }}
+    }}
+"""
+
+
+_fused_gqa_proj_cache = {}
+
+
+def _get_fused_gqa_proj_kernel(K, N_Q, N_GATE, N_K, N_V, group_size=64):
+    key = (K, N_Q, N_GATE, N_K, N_V, group_size)
+    if key not in _fused_gqa_proj_cache:
+        _fused_gqa_proj_cache[key] = mx.fast.metal_kernel(
+            name=f"fused_gqa_proj_K{K}_NQ{N_Q}_NG{N_GATE}_NK{N_K}_NV{N_V}",
+            input_names=["x", "W_merged", "S_merged", "B_merged"],
+            output_names=["q_out", "gate_out", "k_out", "v_out"],
+            source=_gen_fused_gqa_projections_source(K, N_Q, N_GATE, N_K, N_V, group_size),
+        )
+    return _fused_gqa_proj_cache[key]
+
+
+def fused_gqa_projections(
+    x, W_merged, S_merged, B_merged, proj_dims, batch_size=1,
+    scalars=None, total_tg=None,
+):
+    B = batch_size
+    N_Q, N_GATE, N_K, N_V = proj_dims
+    K = x.shape[-1]
+
+    kern = _get_fused_gqa_proj_kernel(K, N_Q, N_GATE, N_K, N_V)
+
+    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, B),
+        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/fused_moe_epilogue_qwen.py

@@ -0,0 +1,97 @@
+"""Fused MoE epilogue for Qwen3.5: weighted sum + shared expert gate + residual.
+
+Computes:
+  Y[j] = bf16( Σ_a(scores[a] * D_routed[a,j]) + gate_shared * D_shared[j] + H[j] )
+
+Two modes:
+  fuse_sigmoid=False: gate_shared is pre-computed sigmoid output (scalar f32)
+  fuse_sigmoid=True:  gate_raw is raw dot product; sigmoid computed internally
+"""
+import mlx.core as mx
+
+
+def _gen_epilogue_qwen_source(K, n_active, fuse_sigmoid=False):
+    """Metal source for fused MoE epilogue with shared expert gate."""
+    if fuse_sigmoid:
+        gate_line = """
+    // Compute sigmoid from raw gate value
+    float gate = 1.0f / (1.0f + metal::exp(-shared_gate_val));"""
+    else:
+        gate_line = """
+    float gate = shared_gate_val;"""
+
+    return f"""
+    const int K_const = {K};
+    const int n_active_const = {n_active};
+
+    uint tid = thread_position_in_grid.x;
+    if (tid >= K_const) return;
+
+    // Weighted sum of routed expert outputs
+    float acc = 0.0f;
+    for (int a = 0; a < n_active_const; a++) {{
+        acc += scores[a] * D_routed[a * K_const + tid];
+    }}
+{gate_line}
+
+    // Shared expert: multiply by gate, add to accumulator
+    float shared_val = float(D_shared[tid]) * gate;
+
+    // Add residual and write
+    Y[tid] = static_cast<bfloat16_t>(acc + shared_val + float(H[tid]));
+"""
+
+
+_epilogue_qwen_kernels = {}
+
+
+def _get_epilogue_qwen_kernel(K, n_active, fuse_sigmoid=False):
+    key = (K, n_active, fuse_sigmoid)
+    if key not in _epilogue_qwen_kernels:
+        sig_tag = "_fsig" if fuse_sigmoid else ""
+        _epilogue_qwen_kernels[key] = mx.fast.metal_kernel(
+            name=f"fused_moe_epilogue_qwen_K{K}_n{n_active}{sig_tag}",
+            input_names=["D_routed", "D_shared", "scores", "H",
+                         "shared_gate_val"],
+            output_names=["Y"],
+            source=_gen_epilogue_qwen_source(K, n_active, fuse_sigmoid),
+        )
+    return _epilogue_qwen_kernels[key]
+
+
+def fused_moe_epilogue_qwen(d_routed, d_shared, scores, h,
+                             shared_gate, k_val, fuse_sigmoid=False):
+    """Fused MoE epilogue with shared expert gate.
+
+    Args:
+        d_routed: routed expert outputs (n_active, K) float32
+        d_shared: shared expert output (K,) float32
+        scores: normalized routing scores (n_active,) float32
+        h: residual hidden state (K,) bfloat16
+        shared_gate: sigmoid output (fuse_sigmoid=False) or raw dot product
+                     (fuse_sigmoid=True), scalar float32
+        k_val: hidden dimension
+        fuse_sigmoid: if True, compute sigmoid(shared_gate) internally
+
+    Returns:
+        Y: (K,) bfloat16 — final layer output
+    """
+    K = int(k_val)
+    n_active = scores.shape[0]
+    kern = _get_epilogue_qwen_kernel(K, n_active, fuse_sigmoid)
+
+    # shared_gate is a scalar — pass as 0-d array
+    if shared_gate.ndim > 0:
+        shared_gate = shared_gate.reshape(())
+
+    tg_size = min(K, 1024)
+    n_tg = (K + tg_size - 1) // tg_size
+
+    Y = kern(
+        inputs=[d_routed, d_shared, scores, h, shared_gate],
+        output_shapes=[(K,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(n_tg * tg_size, 1, 1),
+        threadgroup=(tg_size, 1, 1),
+    )
+    return Y[0]

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

@@ -0,0 +1,152 @@
+"""Fused Q/K RMSNorm + RoPE for GQA attention (Dispatch 2).
+
+Performs per-head RMSNorm with learned weight (head_dim=256) then applies
+RoPE on the first 64 dims (partial_rotary_factor=0.25) using non-traditional
+pairing: element p pairs with p+32 (not p+1).
+
+cos/sin values precomputed in Python from inv_freq * position,
+passed as array inputs (no scalar kernel inputs).
+
+All constants baked into Metal source (no scalar kernel inputs).
+"""
+
+import mlx.core as mx
+
+
+def _gen_fused_qk_norm_rope_source(H_q=16, H_kv=2, D=256, rope_dims=64):
+    N_READS = D // 32
+    ROPE_HALF = rope_dims // 2
+    ROPE_THREADS = rope_dims // N_READS
+    FIRST_HALF = ROPE_THREADS // 2
+    PARTNER_XOR = FIRST_HALF
+
+    return f"""
+    const int H_Q = {H_q};
+    const int H_KV = {H_kv};
+    const int D_DIM = {D};
+    const int N_READS = {N_READS};
+    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;
+
+    bool is_q = (head_idx < (uint)H_Q);
+    int head_local = is_q ? (int)head_idx : ((int)head_idx - H_Q);
+
+    // ── Phase 1: Load N_READS elements + partial sum of squares ──
+    int in_base = is_q
+        ? (int)(b_idx * H_Q * D_DIM + (int)head_idx * D_DIM)
+        : (int)(b_idx * H_KV * D_DIM + head_local * D_DIM);
+
+    int elem_base = (int)slid * N_READS;
+    float vals[{N_READS}];
+    float partial_sq = 0.0f;
+
+    for (int i = 0; i < N_READS; i++) {{
+        float xi;
+        if (is_q)
+            xi = (float)queries[in_base + elem_base + i];
+        else
+            xi = (float)keys[in_base + elem_base + i];
+        vals[i] = xi;
+        partial_sq += xi * xi;
+    }}
+
+    // ── Phase 2: RMSNorm reduction ──
+    float sum_sq = simd_sum(partial_sq);
+    float inv_rms = metal::precise::rsqrt(sum_sq / (float)D_DIM + EPS);
+
+    // ── Phase 3: Normalize with learned weight ──
+    for (int i = 0; i < N_READS; i++) {{
+        float w;
+        if (is_q)
+            w = (float)q_norm_w[elem_base + i];
+        else
+            w = (float)k_norm_w[elem_base + i];
+        vals[i] = vals[i] * inv_rms * w;
+    }}
+
+    // ── Phase 4: RoPE on first {rope_dims} dims ──
+    if (slid < {ROPE_THREADS}u) {{
+        ushort partner = (ushort)(slid ^ {PARTNER_XOR}u);
+        int cos_base = (int)(slid & {FIRST_HALF - 1}u) * N_READS;
+
+        // Load precomputed cos/sin (computed in Python from position * inv_freq)
+        float cos_arr[{N_READS}], sin_arr[{N_READS}];
+        for (int i = 0; i < N_READS; i++) {{
+            cos_arr[i] = rope_cos[cos_base + i];
+            sin_arr[i] = rope_sin[cos_base + i];
+        }}
+
+        float partner_vals[{N_READS}];
+        for (int i = 0; i < N_READS; i++)
+            partner_vals[i] = simd_shuffle(vals[i], partner);
+
+        if (slid < {FIRST_HALF}u) {{
+            for (int i = 0; i < N_READS; i++)
+                vals[i] = vals[i] * cos_arr[i] - partner_vals[i] * sin_arr[i];
+        }} else {{
+            for (int i = 0; i < N_READS; i++)
+                vals[i] = partner_vals[i] * sin_arr[i] + vals[i] * cos_arr[i];
+        }}
+    }}
+
+    // ── Phase 5: Write output ──
+    int out_base = is_q
+        ? (int)(b_idx * H_Q * D_DIM + (int)head_idx * D_DIM)
+        : (int)(b_idx * H_KV * D_DIM + head_local * D_DIM);
+
+    for (int i = 0; i < N_READS; i++) {{
+        if (is_q)
+            q_out[out_base + elem_base + i] = static_cast<bfloat16_t>(vals[i]);
+        else
+            k_out[out_base + elem_base + i] = static_cast<bfloat16_t>(vals[i]);
+    }}
+"""
+
+
+_kernel_cache = {}
+
+
+def _get_kernel(H_q, H_kv, D, rope_dims):
+    key = (H_q, H_kv, D, rope_dims)
+    if key not in _kernel_cache:
+        _kernel_cache[key] = mx.fast.metal_kernel(
+            name=f"fused_qk_norm_rope_Hq{H_q}_Hkv{H_kv}_D{D}_rd{rope_dims}",
+            input_names=["queries", "keys", "q_norm_w", "k_norm_w",
+                         "rope_cos", "rope_sin"],
+            output_names=["q_out", "k_out"],
+            source=_gen_fused_qk_norm_rope_source(H_q, H_kv, D, rope_dims),
+        )
+    return _kernel_cache[key]
+
+
+def fused_qk_norm_rope(queries, keys, q_norm_weight, k_norm_weight,
+                        inv_freq, cache_offset, H_q, H_kv, D,
+                        batch_size=1):
+    B = batch_size
+    rope_dims = inv_freq.shape[0] * 2
+
+    kern = _get_kernel(H_q, H_kv, D, rope_dims)
+
+    q_flat = queries.reshape(B, H_q * D)
+    k_flat = keys.reshape(B, H_kv * D)
+
+    # Precompute cos/sin in Python (avoids scalar kernel input)
+    angles = float(cache_offset) * inv_freq
+    rope_cos = mx.cos(angles).astype(mx.float32)
+    rope_sin = mx.sin(angles).astype(mx.float32)
+
+    n_heads = H_q + H_kv
+    results = kern(
+        inputs=[q_flat, k_flat, q_norm_weight, k_norm_weight, rope_cos, rope_sin],
+        output_shapes=[(B * H_q * D,), (B * H_kv * D,)],
+        output_dtypes=[mx.bfloat16, mx.bfloat16],
+        grid=(n_heads * 32, 1, B),
+        threadgroup=(32, 1, 1),
+    )
+
+    q_out = results[0].reshape(B, H_q, 1, D)
+    k_out = results[1].reshape(B, H_kv, 1, D)
+    return q_out, k_out

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

@@ -0,0 +1,128 @@
+"""Fused Q/K per-head L2-norm for GDN attention (Dispatch 3).
+
+Performs per-head L2 normalization on q and k vectors with different scaling.
+Matches vLLM and latest mlx-lm (qwen3_5.py) which use rsqrt(sum(x²) + eps),
+NOT rms_norm which uses rsqrt(mean(x²) + eps).
+
+From qwen3_5.py (updated to match vLLM):
+    inv_scale = Dk^(-0.5) = 128^(-0.5)
+    q = inv_scale * q * rsqrt(sum(q²) + 1e-6)   → L2-normalize then scale by 1/√Dk
+    k = k * rsqrt(sum(k²) + 1e-6)                → L2-normalize only (no extra scale)
+
+Grid: (32 heads × 32 threads, 1, B).
+Each TG = 32 threads = 1 SG, handles one 128-dim head.
+Dk=128 = 32 threads × 4 elements → exactly 1 SG, no cross-SG reduction.
+"""
+
+import mlx.core as mx
+
+
+def _gen_fused_qk_rmsnorm_source():
+    """Generate Metal source for fused Q/K per-head L2-norm.
+
+    Input: qkv [B, 8192] bf16 (flattened from [B, 1, 8192])
+      - [0, 2048): q = 16 heads × 128
+      - [2048, 4096): k = 16 heads × 128
+      - [4096, 8192): v (untouched)
+
+    Output: qk_out [B, 4096] bf16
+      - [0, 2048): q L2-normalized then scaled by 1/√Dk
+      - [2048, 4096): k L2-normalized (no extra scale)
+
+    Grid: (32 * 32, 1, B), TG: (32, 1, 1)
+      tgid.x 0..15:  q heads → scale = 1/√128
+      tgid.x 16..31: k heads → scale = 1.0
+      tgid.z: batch index
+    """
+    return """
+    const int N_READS = 4;
+    const int DK = 128;
+    const int HK = 16;
+    const float EPS = 1e-6f;
+    const float Q_SCALE = rsqrt(128.0f);        // inv_scale = Dk^(-0.5)
+    const float K_SCALE = 1.0f;                  // no extra scale for k
+
+    uint head_idx = threadgroup_position_in_grid.x;
+    uint slid = thread_index_in_simdgroup;
+    uint b_idx = thread_position_in_grid.z;
+
+    bool is_q = (head_idx < (uint)HK);
+
+    // Input offset: q heads at [0, 2048), k heads at [2048, 4096)
+    int in_base = is_q
+        ? (b_idx * 8192 + head_idx * DK)
+        : (b_idx * 8192 + 2048 + (head_idx - HK) * DK);
+
+    // Output offset: q at [0, 2048), k at [2048, 4096)
+    int out_base = b_idx * 4096 + head_idx * DK;
+
+    // ── Phase 1: Load 4 elements + sum of squares ──
+    float vals[4];
+    float partial_sq = 0.0f;
+    int elem_base = slid * N_READS;
+
+    for (int i = 0; i < N_READS; i++) {
+        float xi = float(qkv[in_base + elem_base + i]);
+        vals[i] = xi;
+        partial_sq += xi * xi;
+    }
+
+    // ── Phase 2: simd reduction (32 threads → full sum of 128 elements) ──
+    float sum_sq = simd_sum(partial_sq);
+
+    // ── Phase 3: compute L2 inv-norm (NOT rms_norm — no /Dk) ──
+    float inv_rms = metal::precise::rsqrt(sum_sq + EPS);
+
+    // ── Phase 4: scale and write ──
+    float scale = is_q ? Q_SCALE : K_SCALE;
+    float combined = inv_rms * scale;
+
+    for (int i = 0; i < N_READS; i++) {
+        qk_out[out_base + elem_base + i] = static_cast<bfloat16_t>(vals[i] * combined);
+    }
+"""
+
+
+_fused_qk_rmsnorm_kernel = None
+
+
+def _get_fused_qk_rmsnorm_kernel():
+    """Get or compile the fused Q/K RMSNorm kernel."""
+    global _fused_qk_rmsnorm_kernel
+    if _fused_qk_rmsnorm_kernel is None:
+        _fused_qk_rmsnorm_kernel = mx.fast.metal_kernel(
+            name="fused_qk_rmsnorm",
+            input_names=["qkv"],
+            output_names=["qk_out"],
+            source=_gen_fused_qk_rmsnorm_source(),
+        )
+    return _fused_qk_rmsnorm_kernel
+
+
+def fused_qk_rmsnorm(qkv_conv_silu, batch_size=1):
+    """Fused Q/K per-head RMSNorm for GDN attention.
+
+    Args:
+        qkv_conv_silu: [B, 1, 8192] bf16 — post-conv, post-SiLU output from Dispatch 2.
+            First 2048 = q (16 heads × 128), next 2048 = k, last 4096 = v.
+        batch_size: int — batch dimension.
+
+    Returns:
+        qk_normed: [B, 1, 4096] bf16 — normalized q (first 2048) and k (next 2048).
+            v is NOT copied; Dispatch 4 reads v directly from qkv_conv_silu[:, :, 4096:].
+    """
+    B = batch_size
+    kern = _get_fused_qk_rmsnorm_kernel()
+
+    # Flatten to [B, 8192] for kernel
+    qkv_flat = qkv_conv_silu.reshape(B, 8192)
+
+    n_heads = 32  # 16 q + 16 k
+    results = kern(
+        inputs=[qkv_flat],
+        output_shapes=[(B * 4096,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(n_heads * 32, 1, B),
+        threadgroup=(32, 1, 1),
+    )
+    return results[0].reshape(B, 1, 4096)

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

@@ -0,0 +1,117 @@
+"""Fused RMSNormGated for GDN attention (Dispatch 5).
+
+Fuses RMSNorm(out, weight) × z_silu into one kernel.
+SiLU on z was already applied in Dispatch 2, so z_silu arrives as f32.
+
+From qwen3_next.py (Qwen3NextRMSNormGated):
+    x = rms_norm(hidden_states, weight, eps)     # weight: [Dv=128]
+    gate = silu(z.float())                        # already done in Dispatch 2
+    return (gate * x).to(hidden_states.dtype)
+
+Grid: (32 heads × 32 threads, 1, B).
+Each TG = 32 threads = 1 SG, handles one 128-dim head.
+Dv=128 = 32 threads × 4 elements → exactly 1 SG.
+"""
+
+import mlx.core as mx
+
+
+def _gen_fused_rms_norm_gated_source():
+    """Generate Metal source for fused RMSNormGated.
+
+    Inputs:
+      gdn_out: [B, Hv*Dv] bf16  — GDN output, flattened (Hv=32, Dv=128)
+      z_silu:  [B, Hv*Dv] f32   — post-SiLU z from Dispatch 2
+      weight:  [Dv] f32         — RMSNormGated learned weight (128 elements)
+
+    Output:
+      out: [B, Hv*Dv] bf16      — result = z_silu * rms_norm(gdn_out, weight)
+
+    Grid: (32 * 32, 1, B), TG: (32, 1, 1)
+      tgid.x: head index (0..31)
+      tgid.z: batch index
+    """
+    return """
+    const int N_READS = 4;
+    const int DV = 128;
+    const int HV = 32;
+    const float EPS = 1e-6f;
+
+    uint head_idx = threadgroup_position_in_grid.x;
+    uint slid = thread_index_in_simdgroup;
+    uint b_idx = thread_position_in_grid.z;
+
+    int base = b_idx * HV * DV + head_idx * DV;
+    int elem_base = slid * N_READS;
+
+    // ── Phase 1: Load gdn_out elements + sum of squares ──
+    float gdn_vals[4];
+    float partial_sq = 0.0f;
+
+    for (int i = 0; i < N_READS; i++) {
+        float xi = float(gdn_out[base + elem_base + i]);
+        gdn_vals[i] = xi;
+        partial_sq += xi * xi;
+    }
+
+    // ── Phase 2: simd reduction (32 threads → full sum of 128 elements) ──
+    float sum_sq = simd_sum(partial_sq);
+
+    // ── Phase 3: compute inv_rms ──
+    float inv_rms = metal::precise::rsqrt(sum_sq / float(DV) + EPS);
+
+    // ── Phase 4: RMSNorm × z_silu, write bf16 ──
+    for (int i = 0; i < N_READS; i++) {
+        int idx = elem_base + i;
+        float w = float(weight[idx]);                           // learned weight[Dv]
+        float normed = gdn_vals[i] * inv_rms * w;              // RMSNorm
+        float z_val = z_silu[base + idx];                       // already f32, post-SiLU
+        out[base + idx] = static_cast<bfloat16_t>(z_val * normed);
+    }
+"""
+
+
+_fused_rms_norm_gated_kernel = None
+
+
+def _get_fused_rms_norm_gated_kernel():
+    """Get or compile the fused RMSNormGated kernel."""
+    global _fused_rms_norm_gated_kernel
+    if _fused_rms_norm_gated_kernel is None:
+        _fused_rms_norm_gated_kernel = mx.fast.metal_kernel(
+            name="fused_rms_norm_gated",
+            input_names=["gdn_out", "z_silu", "weight"],
+            output_names=["out"],
+            source=_gen_fused_rms_norm_gated_source(),
+        )
+    return _fused_rms_norm_gated_kernel
+
+
+def fused_rms_norm_gated(gdn_out, z_silu, weight, batch_size=1):
+    """Fused RMSNormGated: RMSNorm(out, weight) × z_silu.
+
+    Args:
+        gdn_out: [B, 1, Hv, Dv] bf16 — GDN recurrence output (Hv=32, Dv=128).
+        z_silu:  [B, 1, 4096] f32 — post-SiLU z from Dispatch 2.
+        weight:  [128] f32 — RMSNormGated learned weight (Dv elements).
+        batch_size: int.
+
+    Returns:
+        out: [B, 1, 4096] bf16 — ready for out_proj in Dispatch 6.
+    """
+    B = batch_size
+    kern = _get_fused_rms_norm_gated_kernel()
+
+    # Flatten to [B, 4096]
+    gdn_flat = gdn_out.reshape(B, 4096)
+    z_flat = z_silu.reshape(B, 4096)
+
+    n_heads = 32  # Hv
+    results = kern(
+        inputs=[gdn_flat, z_flat, weight],
+        output_shapes=[(B * 4096,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(n_heads * 32, 1, B),
+        threadgroup=(32, 1, 1),
+    )
+    return results[0].reshape(B, 1, 4096)

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

@@ -0,0 +1,177 @@
+"""GDN recurrence with pre-computed g and beta (Dispatch 4).
+
+Modified version of gated_delta_step from mlx-lm-fork/mlx_lm/models/gated_delta.py.
+Instead of computing g = exp(-exp(A_log) * softplus(a + dt_bias)) and beta = sigmoid(b)
+inside the kernel, accepts them as pre-computed f32 inputs from Dispatch 2.
+
+Non-vectorized only (Qwen3.5-35B-A3B uses scalar gating per head).
+
+Grid: (32, Dv, B*Hv) = (32, 128, B*32), TG: (32, 4, 1)
+"""
+
+from typing import Optional, Tuple
+
+import mlx.core as mx
+
+
+def _make_gdn_precomputed_kernel(has_mask=False):
+    """Build the GDN kernel with pre-computed g and beta."""
+    if not mx.metal.is_available():
+        return None
+
+    mask_source = "mask[b_idx * T + t]" if has_mask else "true"
+
+    source = f"""
+        auto n = thread_position_in_grid.z;
+        auto b_idx = n / Hv;
+        auto hv_idx = n % Hv;
+        auto hk_idx = hv_idx / (Hv / Hk);
+        constexpr int n_per_t = Dk / 32;
+
+        // q, k: [B, T, Hk, Dk]
+        auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
+        auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
+
+        // v, y: [B, T, Hv, Dv]
+        auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
+        y += b_idx * T * Hv * Dv + hv_idx * Dv;
+
+        auto dk_idx = thread_position_in_threadgroup.x;
+        auto dv_idx = thread_position_in_grid.y;
+
+        // state_in, state_out: [B, Hv, Dv, Dk]
+        auto i_state = state_in + (n * Dv + dv_idx) * Dk;
+        auto o_state = state_out + (n * Dv + dv_idx) * Dk;
+
+        float state[n_per_t];
+        for (int i = 0; i < n_per_t; ++i) {{
+          auto s_idx = n_per_t * dk_idx + i;
+          state[i] = static_cast<float>(i_state[s_idx]);
+        }}
+
+        // g: [B, T, Hv] f32 — pre-computed decay gate
+        auto g_ = g + b_idx * T * Hv;
+        // beta: [B, T, Hv] f32 — pre-computed sigmoid(b)
+        auto beta_ = beta + b_idx * T * Hv;
+
+        for (int t = 0; t < T; ++t) {{
+          if ({mask_source}) {{
+            // Pre-computed g and beta (no softplus/exp/sigmoid needed)
+            float g_val = g_[hv_idx];
+            float beta_val = beta_[hv_idx];
+
+            float kv_mem = 0.0f;
+            for (int i = 0; i < n_per_t; ++i) {{
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = state[i] * g_val;
+              kv_mem += state[i] * k_[s_idx];
+            }}
+            kv_mem = simd_sum(kv_mem);
+
+            auto delta = (v_[dv_idx] - kv_mem) * beta_val;
+
+            float out = 0.0f;
+            for (int i = 0; i < n_per_t; ++i) {{
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = state[i] + k_[s_idx] * delta;
+              out += state[i] * q_[s_idx];
+            }}
+            out = simd_sum(out);
+            if (thread_index_in_simdgroup == 0) {{
+              y[dv_idx] = static_cast<InT>(out);
+            }}
+          }}
+          // Increment data pointers to next time step
+          q_ += Hk * Dk;
+          k_ += Hk * Dk;
+          v_ += Hv * Dv;
+          y += Hv * Dv;
+          g_ += Hv;
+          beta_ += Hv;
+        }}
+        for (int i = 0; i < n_per_t; ++i) {{
+          auto s_idx = n_per_t * dk_idx + i;
+          o_state[s_idx] = static_cast<InT>(state[i]);
+        }}
+    """
+
+    inputs = ["q", "k", "v", "g", "beta", "state_in", "T"]
+    if has_mask:
+        inputs.append("mask")
+
+    suffix = "_precomputed"
+    if has_mask:
+        suffix += "_mask"
+
+    return mx.fast.metal_kernel(
+        name=f"gated_delta_step{suffix}",
+        input_names=inputs,
+        output_names=["y", "state_out"],
+        source=source,
+    )
+
+
+_gdn_precomputed_kernel = None
+_gdn_precomputed_kernel_masked = None
+
+
+def _get_gdn_precomputed_kernel(has_mask=False):
+    """Get or compile the pre-computed GDN kernel."""
+    global _gdn_precomputed_kernel, _gdn_precomputed_kernel_masked
+    if has_mask:
+        if _gdn_precomputed_kernel_masked is None:
+            _gdn_precomputed_kernel_masked = _make_gdn_precomputed_kernel(has_mask=True)
+        return _gdn_precomputed_kernel_masked
+    else:
+        if _gdn_precomputed_kernel is None:
+            _gdn_precomputed_kernel = _make_gdn_precomputed_kernel(has_mask=False)
+        return _gdn_precomputed_kernel
+
+
+def gated_delta_update_precomputed(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: mx.array,
+    mask: Optional[mx.array] = None,
+) -> Tuple[mx.array, mx.array]:
+    """GDN recurrence with pre-computed g and beta.
+
+    Args:
+        q: [B, T, Hk, Dk] bf16 — normalized q from Dispatch 3
+        k: [B, T, Hk, Dk] bf16 — normalized k from Dispatch 3
+        v: [B, T, Hv, Dv] bf16 — v from Dispatch 2 (qkv_conv_silu[:, :, 4096:])
+        g: [B, T, Hv] f32 — pre-computed decay gate from Dispatch 2
+        beta: [B, T, Hv] f32 — pre-computed sigmoid(b) from Dispatch 2
+        state: [B, Hv, Dv, Dk] bf16 — recurrent state from cache
+        mask: [B, T] optional
+
+    Returns:
+        y: [B, T, Hv, Dv] bf16
+        new_state: [B, Hv, Dv, Dk] bf16
+    """
+    B, T, Hk, Dk = k.shape
+    Hv, Dv = v.shape[2:]
+    input_type = q.dtype
+
+    kernel = _get_gdn_precomputed_kernel(has_mask=mask is not None)
+    inputs = [q, k, v, g, beta, state, T]
+    if mask is not None:
+        inputs.append(mask)
+
+    return kernel(
+        inputs=inputs,
+        template=[
+            ("InT", input_type),
+            ("Dk", Dk),
+            ("Dv", Dv),
+            ("Hk", Hk),
+            ("Hv", Hv),
+        ],
+        grid=(32, Dv, B * Hv),
+        threadgroup=(32, 4, 1),
+        output_shapes=[(B, T, Hv, Dv), state.shape],
+        output_dtypes=[input_type, input_type],
+    )

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

@@ -0,0 +1,218 @@
+"""Merged 8-bit down_proj GEMV for Qwen3.5 routed + shared experts.
+
+Port of the 4-bit down_proj kernel, adapted for 8-bit quantization (gs=64).
+Maps intermediate → hidden: (n_active, N_IN) → (n_active, K_OUT).
+
+Two paths via tgid.z:
+- 8-bit routed experts (z < n_active): uint8 dequant with expert index lookup, f32 output
+- 8-bit shared expert (z == n_active): uint8 dequant (no expert index), f32 output
+"""
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_merged_down_8bit_source(group_size=64, scale_bf16=True):
+    """Metal source for merged 8-bit down_proj GEMV.
+
+    Same structure as gate GEMV: 2 SGs × 4 rows/SG = 8 output rows per TG.
+    8-bit dequant: result = scale * Σ(x[i]*w[i]) + bias * Σ(x[i])
+    Both routed and shared paths use 8-bit dequantization.
+    """
+    gs = int(group_size)
+    sc_stride = 256 // gs  # groups per K-block (256/64 = 4)
+    slid_divisor = gs // 8  # threads per group (64/8 = 8)
+    sc_t = "bfloat16_t" if scale_bf16 else "float"
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+
+    int K_OUT = K_OUT_val;
+    int N_IN = N_IN_val;
+    int SHARED_N_IN = SHARED_N_IN_val;
+    int n_active = n_active_val;
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;  // 0 or 1
+    uint slid = thread_index_in_simdgroup;       // 0..31
+
+    int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG;
+    if (out_row >= K_OUT) return;
+
+    if (tgid.z < (uint)n_active) {{
+        // ═══════ 8-BIT ROUTED EXPERT PATH ═══════
+        int N_groups = N_IN / {gs};
+
+        int expert = inds[tgid.z];
+
+        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 {sc_t}* sc = (const device {sc_t}*)S
+            + (long)expert * K_OUT * N_groups + out_row * N_groups + slid / {slid_divisor};
+        const device {sc_t}* bi = (const device {sc_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
+            + tgid.z * 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 + tgid.z * 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 {{
+        // ═══════ 8-BIT SHARED EXPERT PATH ═══════
+        // Same dequant as routed path, but no expert index lookup.
+        // W_shared_down is (K_OUT, SHARED_N_IN/4) uint32 → (K_OUT, SHARED_N_IN) uint8
+        int N_groups = SHARED_N_IN / {gs};
+
+        const device uint8_t* ws = (const device uint8_t*)W_shared_down
+            + (long)out_row * SHARED_N_IN + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc = (const device {sc_t}*)S_shared_down
+            + (long)out_row * N_groups + slid / {slid_divisor};
+        const device {sc_t}* bi = (const device {sc_t}*)B_shared_down
+            + (long)out_row * N_groups + slid / {slid_divisor};
+
+        // X_shared is float32 (output of Kernel 1 shared path)
+        const device float* x_ptr = (const device float*)X_shared;
+
+        int x_base = slid * VALUES_PER_THREAD;
+        float result[4] = {{0, 0, 0, 0}};
+
+        for (int k = 0; k < SHARED_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 * SHARED_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_shared + out_row;
+        for (int tm = 0; tm < RESULTS_PER_SG; tm++) {{
+            float r = simd_sum(result[tm]);
+            if (slid == 0) {{
+                yp[tm] = r;
+            }}
+        }}
+    }}
+"""
+
+
+_merged_down_8bit_kernels = {}
+
+
+def _get_merged_down_8bit_kernel(group_size=64, scale_bf16=True):
+    key = (group_size, scale_bf16)
+    if key not in _merged_down_8bit_kernels:
+        sc_tag = "_bf16sc" if scale_bf16 else ""
+        _merged_down_8bit_kernels[key] = mx.fast.metal_kernel(
+            name=f"merged_down_proj_8bit_gs{group_size}{sc_tag}",
+            input_names=["W", "S", "B_q",
+                         "W_shared_down", "S_shared_down", "B_shared_down",
+                         "X_routed", "X_shared", "inds",
+                         "K_OUT_val", "N_IN_val", "SHARED_N_IN_val", "n_active_val"],
+            output_names=["Y_routed", "Y_shared"],
+            source=_gen_merged_down_8bit_source(group_size, scale_bf16),
+        )
+    return _merged_down_8bit_kernels[key]
+
+
+def fused_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, group_size=64,
+                                 shared_n_in=None):
+    """Single-dispatch merged down_proj for 8-bit routed + 8-bit shared experts.
+
+    Args:
+        w_q: routed quantized down weights (E, K_OUT, N_IN/4) uint32
+        s: routed scales (E, K_OUT, N_IN/gs) bfloat16
+        b_q: routed biases (E, K_OUT, N_IN/gs) bfloat16
+        w_shared_down: shared expert down weight (K_OUT, SHARED_N_IN/4) uint32
+        s_shared_down: shared expert down scales (K_OUT, SHARED_N_IN/gs) bfloat16
+        b_shared_down: shared expert down biases (K_OUT, SHARED_N_IN/gs) bfloat16
+        x_routed: routed SwiGLU output (n_active, N_IN) float32
+        x_shared: shared SwiGLU output (SHARED_N_IN,) float32
+        inds: selected expert indices (n_active,) uint32
+        k_out: output dimension (4096 for Qwen3.5)
+        n_in: routed expert input dimension (1024 for Qwen3.5)
+        group_size: quantization group size (64)
+        shared_n_in: shared expert input dimension (defaults to n_in)
+
+    Returns:
+        (Y_routed, Y_shared):
+            Y_routed: (n_active, k_out) float32
+            Y_shared: (k_out,) float32
+    """
+    scale_bf16 = (s.dtype == mx.bfloat16)
+    kern = _get_merged_down_8bit_kernel(group_size, scale_bf16)
+    n_active = inds.shape[0]
+    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
+    y_groups = ceil_div(k_out_val, 8)
+    Y = kern(
+        inputs=[w_q, s, b_q,
+                w_shared_down, s_shared_down, b_shared_down,
+                x_routed, x_shared, inds,
+                k_out_val, n_in_val, shared_n_in_val, n_active],
+        output_shapes=[(n_active, k_out_val), (k_out_val,)],
+        output_dtypes=[mx.float32, mx.float32],
+        grid=(32, y_groups * 2, n_active + 1),
+        threadgroup=(32, 2, 1),
+    )
+    return Y[0], Y[1]

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

@@ -0,0 +1,274 @@
+"""Merged 8-bit fused gate+up+SwiGLU for Qwen3.5 routed + shared experts.
+
+Port of the 4-bit kernel for Kimi K2.5, adapted for:
+- 8-bit quantization (direct uint8 byte reads, no nibble extraction)
+- group_size=64 (one scale+bias per 64 elements)
+- Qwen3.5 dimensions (K=4096, N_INTER=1024, E=512, top_k=10)
+- No score normalization in kernel (handled by gate dispatch)
+
+Two paths via tgid.z:
+- 8-bit routed experts (z < n_active): uint8 dequant with expert index lookup, f32 output
+- 8-bit shared expert (z == n_active): uint8 dequant (no expert index), f32 output
+"""
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_merged_8bit_source(group_size=64, scale_bf16=True):
+    """Generate Metal source for merged 8-bit fused gate+up+SwiGLU.
+
+    Both routed and shared expert paths use 8-bit dequant:
+        result = scale * Σ(x[i]*w[i]) + bias * Σ(x[i])
+
+    Each thread processes VALUES_PER_THREAD=8 elements of K per iteration.
+    BLOCK_SIZE = 32 * 8 = 256 elements per K-block.
+    """
+    gs = int(group_size)
+    sc_stride = 256 // gs  # groups consumed per K-block (256/64 = 4)
+    slid_divisor = gs // 8  # threads per group (64/8 = 8)
+    sc_t = "bfloat16_t" if scale_bf16 else "float"
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+
+    int N_INTER = N_INTER_val;
+    int SHARED_INTER = SHARED_INTER_val;
+    int K = K_val;
+    int n_active = n_active_val;
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;  // 0 or 1
+    uint slid = thread_index_in_simdgroup;       // 0..31
+
+    int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG;
+
+    // Routed path uses N_INTER rows, shared path uses SHARED_INTER rows
+    int row_limit = (tgid.z < (uint)n_active) ? N_INTER : SHARED_INTER;
+    if (out_row >= row_limit) return;
+
+    float gate_result[4] = {{0, 0, 0, 0}};
+    float up_result[4] = {{0, 0, 0, 0}};
+
+    if (tgid.z < (uint)n_active) {{
+        // ═══════ 8-BIT ROUTED EXPERT PATH ═══════
+        const int VALUES_PER_THREAD = 8;
+        const int BLOCK_SIZE = 256;  // 32 * 8
+
+        int N_TOTAL = 2 * N_INTER;  // gate + up stacked
+        int K_groups = K / {gs};
+
+        int expert = inds[tgid.z];
+
+        // W is stored as uint32 (4 bytes per uint32), cast to uint8_t
+        // Layout: (E, N_TOTAL, K/4) uint32 → (E, N_TOTAL, K) uint8
+        const device uint8_t* ws_gate = (const device uint8_t*)W
+            + (long)expert * N_TOTAL * K + out_row * K + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_gate = (const device {sc_t}*)S
+            + (long)expert * N_TOTAL * K_groups + out_row * K_groups + slid / {slid_divisor};
+        const device {sc_t}* bi_gate = (const device {sc_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 + (out_row + N_INTER) * K + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_up = (const device {sc_t}*)S
+            + (long)expert * N_TOTAL * K_groups + (out_row + N_INTER) * K_groups + slid / {slid_divisor};
+        const device {sc_t}* bi_up = (const device {sc_t}*)B_q
+            + (long)expert * N_TOTAL * K_groups + (out_row + N_INTER) * K_groups + slid / {slid_divisor};
+
+        int x_base = slid * VALUES_PER_THREAD;
+
+        for (int k = 0; k < K; k += BLOCK_SIZE) {{
+            // Load 8 x values and compute sum
+            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++) {{
+                // Gate projection
+                const device uint8_t* wg = ws_gate + row * K;
+                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;
+
+                // Up projection
+                const device uint8_t* wu = ws_up + row * K;
+                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: SwiGLU + write f32 to Y_routed
+        device float* yp = Y_routed + tgid.z * N_INTER + out_row;
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            float g = simd_sum(gate_result[row]);
+            float u = simd_sum(up_result[row]);
+            if (slid == 0) {{
+                float silu_g = g / (1.0f + metal::exp(-g));
+                yp[row] = silu_g * u;
+            }}
+        }}
+
+    }} else {{
+        // ═══════ 8-BIT SHARED EXPERT PATH ═══════
+        // Same dequant as routed path, but no expert index lookup.
+        // W_shared is (2*SHARED_INTER, K/4) uint32 → (2*SHARED_INTER, K) uint8
+        const int VALUES_PER_THREAD = 8;
+        const int BLOCK_SIZE = 256;  // 32 * 8
+
+        int K_groups = K / {gs};
+
+        const device uint8_t* ws_gate = (const device uint8_t*)W_shared
+            + (long)out_row * K + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_gate = (const device {sc_t}*)S_shared
+            + (long)out_row * K_groups + slid / {slid_divisor};
+        const device {sc_t}* bi_gate = (const device {sc_t}*)B_shared
+            + (long)out_row * K_groups + slid / {slid_divisor};
+
+        const device uint8_t* ws_up = (const device uint8_t*)W_shared
+            + (long)(out_row + SHARED_INTER) * K + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_up = (const device {sc_t}*)S_shared
+            + (long)(out_row + SHARED_INTER) * K_groups + slid / {slid_divisor};
+        const device {sc_t}* bi_up = (const device {sc_t}*)B_shared
+            + (long)(out_row + SHARED_INTER) * K_groups + slid / {slid_divisor};
+
+        int x_base = slid * VALUES_PER_THREAD;
+
+        for (int k = 0; k < K; 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++) {{
+                // Gate projection
+                const device uint8_t* wg = ws_gate + row * K;
+                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;
+
+                // Up projection
+                const device uint8_t* wu = ws_up + row * K;
+                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: SwiGLU + write f32 to Y_shared
+        device float* yp = Y_shared + out_row;
+        for (int tm = 0; tm < RESULTS_PER_SG; tm++) {{
+            float g = simd_sum(gate_result[tm]);
+            float u = simd_sum(up_result[tm]);
+            if (slid == 0) {{
+                float silu_g = g / (1.0f + metal::exp(-g));
+                yp[tm] = silu_g * u;
+            }}
+        }}
+    }}
+"""
+
+
+_merged_8bit_kernels = {}
+
+
+def _get_merged_8bit_kernel(group_size=64, scale_bf16=True):
+    key = (group_size, scale_bf16)
+    if key not in _merged_8bit_kernels:
+        sc_tag = "_bf16sc" if scale_bf16 else ""
+        _merged_8bit_kernels[key] = mx.fast.metal_kernel(
+            name=f"merged_routed_shared_swiglu_8bit_gs{group_size}{sc_tag}",
+            input_names=["W", "S", "B_q", "W_shared", "S_shared", "B_shared",
+                         "X", "inds",
+                         "N_INTER_val", "SHARED_INTER_val", "K_val", "n_active_val"],
+            output_names=["Y_routed", "Y_shared"],
+            source=_gen_merged_8bit_source(group_size, scale_bf16),
+        )
+    return _merged_8bit_kernels[key]
+
+
+def fused_merged_gate_up_swiglu_8bit(w_q, s, b_q,
+                                      w_shared, s_shared, b_shared,
+                                      x, inds,
+                                      n_inter, k_hidden, group_size=64,
+                                      shared_inter=None):
+    """Single-dispatch merged gate+up+SwiGLU for 8-bit routed + 8-bit shared experts.
+
+    Args:
+        w_q: stacked routed quantized weights (E, 2*N_INTER, K/4) uint32
+        s: routed scales (E, 2*N_INTER, K/gs) bfloat16
+        b_q: routed biases (E, 2*N_INTER, K/gs) bfloat16
+        w_shared: shared expert gate+up stacked (2*SHARED_INTER, K/4) uint32
+        s_shared: shared expert scales (2*SHARED_INTER, K/gs) bfloat16
+        b_shared: shared expert biases (2*SHARED_INTER, K/gs) bfloat16
+        x: input vector (K,) bfloat16
+        inds: selected expert indices (n_active,) uint32
+        n_inter: routed expert intermediate size (1024 for Qwen3.5)
+        k_hidden: hidden size (4096 for Qwen3.5)
+        group_size: quantization group size (64)
+        shared_inter: shared expert intermediate size (defaults to n_inter)
+
+    Returns:
+        (Y_routed, Y_shared):
+            Y_routed: (n_active, n_inter) float32
+            Y_shared: (shared_inter,) float32
+    """
+    scale_bf16 = (s.dtype == mx.bfloat16)
+    kern = _get_merged_8bit_kernel(group_size, scale_bf16)
+    n_active = inds.shape[0]
+    n_inter_val = int(n_inter)
+    shared_inter_val = int(shared_inter) if shared_inter is not None else n_inter_val
+    # Grid y must cover max(n_inter, shared_inter)
+    max_inter = max(n_inter_val, shared_inter_val)
+    Y = kern(
+        inputs=[w_q, s, b_q, w_shared, s_shared, b_shared,
+                x, inds,
+                n_inter_val, shared_inter_val, int(k_hidden), n_active],
+        output_shapes=[(n_active, n_inter_val), (shared_inter_val,)],
+        output_dtypes=[mx.float32, mx.float32],
+        grid=(32, ceil_div(max_inter, 8) * 2, n_active + 1),
+        threadgroup=(32, 2, 1),
+    )
+    return Y[0], Y[1]

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

@@ -0,0 +1,493 @@
+"""Dispatch 2: SwiGLU with softmax prologue for Qwen3.5 oproj fusion.
+
+Port of Kimi's oproj_topk_v2_swiglu.py adapted for:
+- 8-bit quantized weights with gs=64 (Kimi uses 4-bit gs=32)
+- Softmax routing (Kimi uses sigmoid)
+- E up to 512 with multiple scores per thread (SPT = ceil(E/64))
+- Shared expert gate GEMV hidden in TG(0,0,0) SG 0
+- Both routed and shared paths use 8-bit quantized weights
+
+Prologue (all TGs):
+  Phase 1: distributed x² sum → inv_rms
+  Phase 2 (routed TGs only): gate scores → softmax → parallel top-k → norm_topk_prob
+  Phase 3 (TG(0,0,0) SG 0): shared_expert_gate 8-bit GEMV → gate_raw
+
+Main SwiGLU (after prologue):
+  Routed (z < n_active): 8-bit gate+up+SwiGLU with h_scaled input (inv_rms factored out)
+  Shared (z == n_active): 8-bit SwiGLU for shared expert
+"""
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_oproj_softmax_topk_swiglu_8bit_source(group_size=64, scale_bf16=True,
+                                                 n_experts=64, top_k=10,
+                                                 norm_topk=True):
+    """Generate Metal source for softmax + top-k + SwiGLU with oproj prologue."""
+    gs = int(group_size)
+    sc_stride = 256 // gs       # groups per K-block (256/64 = 4)
+    slid_divisor = gs // 8      # threads per group (64/8 = 8)
+    sc_t = "bfloat16_t" if scale_bf16 else "float"
+    E = int(n_experts)
+    K_TOP = int(top_k)
+    SPT = (E + 63) // 64       # scores per thread
+
+    # Score normalization block (TG(0,0,0) thread 0 only)
+    if norm_topk:
+        score_norm_block = f"""
+    // ── norm_topk_prob + write indices (TG(0,0,0) thread 0) ──
+    if (tgid.y == 0 && tgid.z == 0 && tid == 0) {{
+        float total = 0.0f;
+        for (int a = 0; a < {K_TOP}; a++) total += tg_selected_scores[a];
+        float inv_total = 1.0f / total;
+        for (int a = 0; a < {K_TOP}; a++) {{
+            norm_scores[a] = tg_selected_scores[a] * inv_total;
+            out_inds[a] = (uint)tg_inds[a];
+        }}
+    }}"""
+    else:
+        score_norm_block = f"""
+    // ── Write indices and raw scores (TG(0,0,0) thread 0) ──
+    if (tgid.y == 0 && tgid.z == 0 && tid == 0) {{
+        for (int a = 0; a < {K_TOP}; a++) {{
+            norm_scores[a] = tg_selected_scores[a];
+            out_inds[a] = (uint)tg_inds[a];
+        }}
+    }}"""
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int E_CONST = {E};
+    const int K_TOP_CONST = {K_TOP};
+    const int SPT = {SPT};
+
+    int N_INTER = N_INTER_val;
+    int SHARED_INTER = SHARED_INTER_val;
+    int K = K_val;
+    int n_active = n_active_val;
+    int N_OPROJ_TG = N_OPROJ_TG_val;
+
+    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
+
+    // ═══════════════════════════════════════════════════════════════
+    // PROLOGUE PHASE 1: distributed x² sum → inv_rms (ALL TGs)
+    // ═══════════════════════════════════════════════════════════════
+    int chunk = (N_OPROJ_TG + 63) / 64;
+    int x2_start = tid * chunk;
+    int x2_end = min(x2_start + chunk, N_OPROJ_TG);
+    float local_x2 = 0.0f;
+    for (int i = x2_start; i < x2_end; i++) local_x2 += x2_partials[i];
+    float sg_x2_sum = simd_sum(local_x2);
+
+    threadgroup float tg_x2_sg[2];
+    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 + 1e-6f);
+
+    // ═══════════════════════════════════════════════════════════════
+    // PROLOGUE PHASE 2: Softmax + Top-k (routed TGs only)
+    // ═══════════════════════════════════════════════════════════════
+    threadgroup int tg_inds[{K_TOP}];
+    threadgroup float tg_selected_scores[{K_TOP}];
+
+    if (tgid.z < (uint)n_active) {{
+        // Load gate scores and apply inv_rms
+        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[e] + gate_part_b[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++) {{
+            // Find local best among SPT scores
+            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;
+                }}
+            }}
+
+            // SG-level max + winner identification
+            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);
+
+            // Determine global winner (SG with higher max; tie-break: SG 0)
+            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;
+
+            // Winner writes expert index and score to TG memory
+            if (tid == winner_tid) {{
+                tg_inds[round] = best_e;
+                tg_selected_scores[round] = best;
+                // Disable the winning score in register
+                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);
+        }}
+    }}
+{score_norm_block}
+
+    // ═══════════════════════════════════════════════════════════════
+    // PROLOGUE PHASE 3: Shared expert gate (TG(0,0,0) SG 0 only)
+    // 8-bit GEMV: W_seg (1, K) × X (K,) → gate_raw scalar
+    // Input is h_scaled; multiply result by inv_rms for true gate value
+    // ═══════════════════════════════════════════════════════════════
+    if (tgid.y == 0 && tgid.z == 0 && sgid == 0) {{
+        const int VPT = 8;
+        const int BLOCK = 256;  // 32 * VPT
+        int K_groups_seg = K / {gs};
+
+        const device uint8_t* wg_seg = (const device uint8_t*)W_seg
+            + slid * VPT;
+        const device {sc_t}* sc_seg = (const device {sc_t}*)S_seg
+            + slid / {slid_divisor};
+        const device {sc_t}* bi_seg = (const device {sc_t}*)B_seg
+            + slid / {slid_divisor};
+        int xb = slid * VPT;
+
+        float gate_acc = 0.0f;
+        for (int k = 0; k < K; k += BLOCK) {{
+            float xsum = 0.0f, wacc = 0.0f;
+            for (int i = 0; i < VPT; i++) {{
+                float xi = float(X[xb + i]);
+                xsum += xi;
+                wacc += xi * float(wg_seg[i]);
+            }}
+            gate_acc += float(*sc_seg) * wacc + xsum * float(*bi_seg);
+            wg_seg += BLOCK;
+            sc_seg += {sc_stride};
+            bi_seg += {sc_stride};
+            xb += BLOCK;
+        }}
+        gate_acc = simd_sum(gate_acc);
+        if (slid == 0) gate_raw[0] = gate_acc * inv_rms;
+    }}
+
+    // ═══════════════════════════════════════════════════════════════
+    // MAIN SwiGLU BODY
+    // ═══════════════════════════════════════════════════════════════
+    int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG;
+    int row_limit = (tgid.z < (uint)n_active) ? N_INTER : SHARED_INTER;
+    if (out_row >= row_limit) return;
+
+    float gate_result[4] = {{0, 0, 0, 0}};
+    float up_result[4] = {{0, 0, 0, 0}};
+
+    if (tgid.z < (uint)n_active) {{
+        // ═══════ 8-BIT ROUTED EXPERT PATH ═══════
+        const int VALUES_PER_THREAD = 8;
+        const int BLOCK_SIZE = 256;
+
+        int N_TOTAL = 2 * N_INTER;
+        int K_groups = K / {gs};
+
+        int expert = tg_inds[tgid.z];
+
+        const device uint8_t* ws_gate = (const device uint8_t*)W
+            + (long)expert * N_TOTAL * K + out_row * K + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_gate = (const device {sc_t}*)S
+            + (long)expert * N_TOTAL * K_groups + out_row * K_groups
+            + slid / {slid_divisor};
+        const device {sc_t}* bi_gate = (const device {sc_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 + (out_row + N_INTER) * K
+            + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_up = (const device {sc_t}*)S
+            + (long)expert * N_TOTAL * K_groups + (out_row + N_INTER) * K_groups
+            + slid / {slid_divisor};
+        const device {sc_t}* bi_up = (const device {sc_t}*)B_q
+            + (long)expert * N_TOTAL * K_groups + (out_row + N_INTER) * K_groups
+            + slid / {slid_divisor};
+
+        int x_base = slid * VALUES_PER_THREAD;
+
+        for (int k = 0; k < K; 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++) {{
+                // Gate projection
+                const device uint8_t* wg = ws_gate + row * K;
+                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;
+
+                // Up projection
+                const device uint8_t* wu = ws_up + row * K;
+                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 + tgid.z * N_INTER + 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 {{
+        // ═══════ 8-BIT SHARED EXPERT PATH ═══════
+        const int VALUES_PER_THREAD = 8;
+        const int BLOCK_SIZE = 256;
+
+        int K_groups = K / {gs};
+
+        const device uint8_t* ws_gate = (const device uint8_t*)W_shared
+            + (long)out_row * K + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_gate = (const device {sc_t}*)S_shared
+            + (long)out_row * K_groups + slid / {slid_divisor};
+        const device {sc_t}* bi_gate = (const device {sc_t}*)B_shared
+            + (long)out_row * K_groups + slid / {slid_divisor};
+
+        const device uint8_t* ws_up = (const device uint8_t*)W_shared
+            + (long)(out_row + SHARED_INTER) * K + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_up = (const device {sc_t}*)S_shared
+            + (long)(out_row + SHARED_INTER) * K_groups + slid / {slid_divisor};
+        const device {sc_t}* bi_up = (const device {sc_t}*)B_shared
+            + (long)(out_row + SHARED_INTER) * K_groups + slid / {slid_divisor};
+
+        int x_base = slid * VALUES_PER_THREAD;
+
+        for (int k = 0; k < K; 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_s = ws_gate + row * K;
+                float sg_s = float(sc_gate[row * K_groups]);
+                float bg_s = float(bi_gate[row * K_groups]);
+                float accum_g = 0;
+                for (int i = 0; i < 8; i++)
+                    accum_g += x_thread[i] * float(wg_s[i]);
+                gate_result[row] += sg_s * accum_g + xsum * bg_s;
+
+                const device uint8_t* wu_s = ws_up + row * K;
+                float su_s = float(sc_up[row * K_groups]);
+                float bu_s = float(bi_up[row * K_groups]);
+                float accum_u = 0;
+                for (int i = 0; i < 8; i++)
+                    accum_u += x_thread[i] * float(wu_s[i]);
+                up_result[row] += su_s * accum_u + xsum * bu_s;
+            }}
+
+            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_shared + out_row;
+        for (int tm = 0; tm < RESULTS_PER_SG; tm++) {{
+            float g = simd_sum(gate_result[tm]) * inv_rms;
+            float u = simd_sum(up_result[tm]) * inv_rms;
+            if (slid == 0) {{
+                float silu_g = g / (1.0f + metal::exp(-g));
+                yp[tm] = silu_g * u;
+            }}
+        }}
+    }}
+"""
+
+
+_oproj_softmax_swiglu_8bit_kernels = {}
+
+
+def _get_oproj_softmax_swiglu_8bit_kernel(group_size=64, scale_bf16=True,
+                                            n_experts=64, top_k=10,
+                                            norm_topk=True):
+    key = (group_size, scale_bf16, n_experts, top_k, norm_topk)
+    if key not in _oproj_softmax_swiglu_8bit_kernels:
+        sc_tag = "_bf16sc" if scale_bf16 else ""
+        nt_tag = "_nt" if norm_topk else ""
+        _oproj_softmax_swiglu_8bit_kernels[key] = mx.fast.metal_kernel(
+            name=(f"oproj_softmax_topk_swiglu_8bit_gs{group_size}"
+                  f"_e{n_experts}_k{top_k}{sc_tag}{nt_tag}"),
+            input_names=[
+                "W", "S", "B_q",                    # routed expert weights
+                "W_shared", "S_shared", "B_shared",  # shared expert weights
+                "X",                                  # h_scaled (K,) bf16
+                "gate_part_a", "gate_part_b",         # (E,) f32
+                "x2_partials",                        # (N_OPROJ_TG,) f32
+                "W_seg", "S_seg", "B_seg",            # shared_expert_gate weights
+                "N_INTER_val", "SHARED_INTER_val",
+                "K_val", "n_active_val", "N_OPROJ_TG_val",
+            ],
+            output_names=["Y_routed", "Y_shared", "out_inds",
+                          "norm_scores", "gate_raw"],
+            source=_gen_oproj_softmax_topk_swiglu_8bit_source(
+                group_size, scale_bf16, n_experts, top_k, norm_topk),
+        )
+    return _oproj_softmax_swiglu_8bit_kernels[key]
+
+
+def fused_oproj_softmax_topk_swiglu_8bit(
+    w_q, s, b_q,                         # routed expert weights
+    w_shared, s_shared, b_shared,         # shared expert weights
+    h_scaled,                             # (K,) bf16 — h * w_rms
+    gate_part_a, gate_part_b,             # (E,) f32 — gate decomposition
+    x2_partials,                          # (N_OPROJ_TG,) f32 — per-TG x²
+    w_seg, s_seg, b_seg,                  # shared_expert_gate 8-bit weights
+    n_inter, k_hidden,                    # MoE dimensions
+    n_experts, top_k,                     # routing params
+    n_oproj_tg,                           # number of o_proj TGs (for x²)
+    group_size=64,
+    shared_inter=None,
+    norm_topk=True,
+):
+    """Single-dispatch softmax + top-k + merged 8-bit SwiGLU with oproj prologue.
+
+    Prologue:
+      Phase 1: distributed x² → inv_rms (all TGs)
+      Phase 2: softmax(gate_part_a + gate_part_b) → top-k → norm (routed TGs)
+      Phase 3: shared_expert_gate 8-bit GEMV → gate_raw (TG(0,0,0) SG 0)
+
+    Main: 8-bit gate+up+SwiGLU for routed + shared experts.
+    inv_rms is factored out of the inner loop and applied once in epilogue.
+
+    Args:
+        h_scaled: (K,) bf16 — h * w_rms from Dispatch 1
+        gate_part_a: (E,) f32 — M1 @ attn_out from Dispatch 1
+        gate_part_b: (E,) f32 — W_fused @ residual from Dispatch 1
+        x2_partials: (N_OPROJ_TG,) f32 — per-TG Σh² from Dispatch 1
+        w_seg/s_seg/b_seg: shared_expert_gate 8-bit weights (1, K/4) uint32
+
+    Returns:
+        (Y_routed, Y_shared, out_inds, norm_scores, gate_raw):
+            Y_routed: (top_k, n_inter) f32
+            Y_shared: (shared_inter,) f32
+            out_inds: (top_k,) uint32
+            norm_scores: (top_k,) f32
+            gate_raw: (1,) f32 — raw shared expert gate value (sigmoid in epilogue)
+    """
+    scale_bf16 = (s.dtype == mx.bfloat16)
+    kern = _get_oproj_softmax_swiglu_8bit_kernel(
+        group_size, scale_bf16, n_experts, top_k, norm_topk)
+
+    n_inter_val = int(n_inter)
+    shared_inter_val = int(shared_inter) if shared_inter is not None else n_inter_val
+    n_active = int(top_k)
+    max_inter = max(n_inter_val, shared_inter_val)
+
+    results = kern(
+        inputs=[
+            w_q, s, b_q,
+            w_shared, s_shared, b_shared,
+            h_scaled,
+            gate_part_a, gate_part_b,
+            x2_partials,
+            w_seg, s_seg, b_seg,
+            n_inter_val, shared_inter_val,
+            int(k_hidden), n_active, int(n_oproj_tg),
+        ],
+        output_shapes=[
+            (n_active, n_inter_val),     # Y_routed
+            (shared_inter_val,),         # Y_shared
+            (n_active,),                 # out_inds
+            (n_active,),                 # norm_scores
+            (1,),                        # 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, n_active + 1),
+        threadgroup=(32, 2, 1),
+    )
+    return results[0], results[1], results[2], results[3], results[4]

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

@@ -0,0 +1,245 @@
+"""MoE __call__ variants for Qwen3.5.
+
+Two modes:
+  _fused_moe_call: gate→SwiGLU→down_proj→epilogue (~15 dispatches, fused to ~4+MLX)
+  _oproj_moe_call: o_proj+gate→SwiGLU(w/softmax prologue)→down_proj→epilogue (4 dispatches)
+
+Kernels used:
+  merged_routed_shared_swiglu_8bit.py    — 8-bit gate+up+SwiGLU (_fused_moe_call)
+  merged_down_proj_8bit.py               — 8-bit down_proj (both modes)
+  fused_moe_epilogue_qwen.py             — weighted sum + sigmoid gate + residual (both modes)
+  custom_oproj_gate_gemv_8bit.py         — 8-bit o_proj + bf16 gate GEMVs (_oproj_moe_call)
+  oproj_softmax_topk_swiglu_8bit.py      — SwiGLU with softmax prologue (_oproj_moe_call)
+
+Adapted from mlx_bench/model_patches/qwen/moe.py.
+"""
+
+import mlx.core as mx
+
+from .kernels.merged_routed_shared_swiglu_8bit import fused_merged_gate_up_swiglu_8bit
+from .kernels.merged_down_proj_8bit import fused_merged_down_proj_8bit
+from .kernels.fused_moe_epilogue_qwen import fused_moe_epilogue_qwen
+from .kernels.custom_oproj_gate_gemv_8bit import fused_custom_oproj_8bit
+from .kernels.oproj_softmax_topk_swiglu_8bit import fused_oproj_softmax_topk_swiglu_8bit
+
+
+def _vanilla_moe_call(self, x):
+    """Original Qwen3NextSparseMoeBlock.__call__ (for prefill fallback)."""
+    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
+
+
+def _fused_moe_call(self, x, _residual=None):
+    """Qwen3.5 MoE with fused kernels (4 custom dispatches).
+
+    Falls back to vanilla for prefill (seq_len > 1).
+
+    Args:
+        x: (B, S, K) bf16 — post-layernorm hidden state
+        _residual: (B, S, K) bf16 — pre-layernorm hidden state for residual add.
+                   If None, epilogue skips residual (returns MoE output only).
+    """
+    # Fused kernels are decode-only (seq_len=1). Fall back for prefill.
+    if x.shape[-2] > 1:
+        # Try vanilla switch_mlp path if weights still exist
+        has_switch_weights = hasattr(self.switch_mlp, 'gate_proj') and \
+            hasattr(self.switch_mlp.gate_proj, 'weight')
+        if has_switch_weights:
+            out = _vanilla_moe_call(self, x)
+        else:
+            # Weights were freed — process tokens one by one through fused path
+            outs = []
+            for t in range(x.shape[-2]):
+                xt = x[:, t:t+1, :]
+                res_t = _residual[:, t:t+1, :] if _residual is not None else None
+                outs.append(_fused_moe_call(self, xt, _residual=res_t))
+            return mx.concatenate(outs, axis=1)
+        if _residual is not None:
+            out = out + _residual
+        return out
+
+    # ── Gate routing (vanilla MLX ops) ──
+    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)
+
+    x_flat = x.reshape(-1).astype(mx.bfloat16)  # (K,)
+    inds_flat = inds.reshape(-1).astype(mx.uint32)
+    scores_flat = scores.reshape(-1).astype(mx.float32)
+
+    # ── Dispatch 1: Merged gate+up+SwiGLU (8-bit routed + 8-bit shared) ──
+    y_routed, y_shared = fused_merged_gate_up_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,
+        x_flat,
+        inds_flat,
+        self.switch_mlp._fused_n_inter,
+        self.switch_mlp._fused_k_hidden,
+        group_size=self.switch_mlp._fused_group_size,
+        shared_inter=self._shared_inter,
+    )
+
+    # ── Dispatch 2: Merged down_proj (8-bit routed + 8-bit shared) ──
+    y_down_routed, y_down_shared = fused_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,
+        inds_flat,
+        self._down_K,
+        self._down_N,
+        group_size=self._down_gs,
+        shared_n_in=self._shared_inter,
+    )
+
+    # ── Dispatch 3: Shared expert gate (small GEMV) ──
+    shared_gate_out = self.shared_expert_gate(x.reshape(1, 1, -1))
+    shared_gate_val = mx.sigmoid(shared_gate_out.reshape(()))
+
+    # ── Dispatch 4: Fused epilogue ──
+    if _residual is not None:
+        h_flat = _residual.reshape(-1).astype(mx.bfloat16)
+    else:
+        h_flat = mx.zeros((self.switch_mlp._fused_k_hidden,), dtype=mx.bfloat16)
+
+    y = fused_moe_epilogue_qwen(
+        y_down_routed,    # (n_active, K) f32
+        y_down_shared,    # (K,) f32
+        scores_flat,      # (n_active,) f32
+        h_flat,           # (K,) bf16
+        shared_gate_val,  # scalar f32
+        self._down_K,     # K
+    )
+    return y.reshape(1, 1, -1)
+
+
+def _oproj_moe_call(self, attn_out_3d, _residual=None):
+    """Qwen3.5 MoE with fused o_proj + gate GEMVs (4 custom dispatches).
+
+    Receives raw attention output (pre-o_proj) and residual.
+    Fuses o_proj, RMSNorm, gate softmax, top-k, score norm, shared_expert_gate,
+    SwiGLU, down_proj, and epilogue into 4 dispatches.
+
+    Dispatch 1: 8-bit o_proj + bf16 M1/W_fused GEMVs → h_scaled, h_out, x2_partials,
+                gate_part_a, gate_part_b
+    Dispatch 2: SwiGLU with softmax prologue (inv_rms + softmax + top-k + score norm
+                + shared_expert_gate in TG(0,0,0)) → y_routed, y_shared, inds, scores,
+                gate_raw
+    Dispatch 3: Merged down_proj (unchanged)
+    Dispatch 4: Fused epilogue with sigmoid(gate_raw)
+
+    Args:
+        attn_out_3d: (1, 1, K_attn) bf16 — raw attention output (pre-o_proj)
+        _residual: (1, 1, K) bf16 — input residual (before attention)
+    """
+    # Prefill fallback (S > 1): restore o_proj + vanilla MoE
+    if attn_out_3d.shape[-2] > 1:
+        from .decoder import _parent_layer_map
+        parent = _parent_layer_map[id(self)]
+        if parent.is_linear:
+            oproj_mod = parent.linear_attn.out_proj
+        else:
+            oproj_mod = parent.self_attn.o_proj
+        projected = oproj_mod(attn_out_3d)
+        h = _residual + projected
+        out = _vanilla_moe_call(self, parent.post_attention_layernorm(h))
+        return out + h
+
+    attn_out = attn_out_3d.reshape(-1).astype(mx.bfloat16)  # (K_attn,)
+    residual = _residual.reshape(-1).astype(mx.bfloat16)     # (K,)
+
+    K = self._oproj_M
+    K_attn = self._oproj_K_attn
+
+    # ── Dispatch 1: fused o_proj + M1 GEMV + W_fused GEMV ──
+    h_scaled, h_out, x2_partials, gate_part_a, gate_part_b = \
+        fused_custom_oproj_8bit(
+            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, K_hidden=self._oproj_K_hidden,
+            n_experts=self._oproj_n_experts,
+            gate_bm=self._oproj_gate_bm,
+        )
+
+    # ── Dispatch 2: SwiGLU with softmax prologue ──
+    y_routed, y_shared, inds, scores, gate_raw = \
+        fused_oproj_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,
+            h_scaled,
+            gate_part_a,
+            gate_part_b,
+            x2_partials,
+            self._seg_w,
+            self._seg_s,
+            self._seg_b,
+            n_inter=self.switch_mlp._fused_n_inter,
+            k_hidden=K,
+            n_experts=self._oproj_n_experts,
+            top_k=self.top_k,
+            n_oproj_tg=self._oproj_n_tg,
+            group_size=self.switch_mlp._fused_group_size,
+            shared_inter=self._shared_inter,
+            norm_topk=self.norm_topk_prob,
+        )
+
+    # ── Dispatch 3: Merged down_proj (8-bit routed + 8-bit shared) ──
+    y_down_routed, y_down_shared = fused_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,
+        inds,
+        self._down_K,
+        self._down_N,
+        group_size=self._down_gs,
+        shared_n_in=self._shared_inter,
+    )
+
+    # ── Dispatch 4: Fused epilogue with sigmoid(gate_raw) ──
+    y = fused_moe_epilogue_qwen(
+        y_down_routed,    # (n_active, K) f32
+        y_down_shared,    # (K,) f32
+        scores,           # (n_active,) f32 — norm_topk_prob from prologue
+        h_out,            # (K,) bf16 — post-o_proj hidden for residual add
+        gate_raw,         # scalar f32 — raw dot product, sigmoid fused
+        K,
+        fuse_sigmoid=True,
+    )
+    return y.reshape(1, 1, -1)

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

@@ -190,6 +190,10 @@ 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")
+
+        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: