batch_generate: preserve 'mtp.' prefix in cached MTP weights (was a real bug)

MTPPredictor._load_weights expects keys with the 'mtp.' prefix preserved
(see mtp_module.py:202+), e.g. 'mtp.fc.weight', 'mtp.layers.0.self_attn.q_proj.weight'.
Both the original code (stripping 'model.mtp.' to leave 'fc.weight') and
my recent change (stripping 'mtp.') produced cached files MTPPredictor
couldn't read — manifesting as KeyError 'mtp.fc.weight' at MTPPredictor
init. The try/except in _resolve_mtp_weights swallowed it and silently
fell back to non-speculative.

Fix: only strip the optional 'model.' wrapper; keep 'mtp.' intact. Bump
cache filename to mtp_v2_<hash>.safetensors so any existing wrong cache
gets re-extracted.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

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

@@ -167,7 +167,7 @@ export interface ModelDownloadStatus {
 // Placement preview from the API
 export interface PlacementPreview {
   model_id: string;
-  sharding: "Pipeline" | "Tensor";
+  sharding: "Pipeline" | "Tensor" | "AttnMoeSplit";
   instance_meta: "MlxRing" | "MlxJaccl";
   instance: unknown | null;
   memory_delta_by_node: Record<string, number> | null;

dashboard/src/routes/+page.svelte

@@ -885,14 +885,16 @@
     sendMessage(content, files, thinkingEnabled());
   }
 
-  let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
+  let selectedSharding = $state<"Pipeline" | "Tensor" | "AttnMoeSplit">(
+    "Pipeline",
+  );
   type InstanceMeta = "MlxRing" | "MlxJaccl";
 
   // Launch defaults persistence
   const LAUNCH_DEFAULTS_KEY = "exo-launch-defaults-v2";
   interface LaunchDefaults {
     modelId: string | null;
-    sharding: "Pipeline" | "Tensor";
+    sharding: "Pipeline" | "Tensor" | "AttnMoeSplit";
     instanceType: InstanceMeta;
     minNodes: number;
   }
@@ -5759,6 +5761,30 @@
                         </span>
                         Tensor
                       </button>
+                      <button
+                        onclick={() => {
+                          selectedSharding = "AttnMoeSplit";
+                          saveLaunchDefaults();
+                        }}
+                        class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedSharding ===
+                        'AttnMoeSplit'
+                          ? 'bg-transparent text-exo-yellow border-exo-yellow'
+                          : 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
+                        title="Qwen3.5 MoE only, 2 nodes: rank 0 runs attention, rank 1 runs MoE"
+                      >
+                        <span
+                          class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedSharding ===
+                          'AttnMoeSplit'
+                            ? 'border-exo-yellow'
+                            : 'border-exo-medium-gray'}"
+                        >
+                          {#if selectedSharding === "AttnMoeSplit"}
+                            <span class="w-1.5 h-1.5 rounded-full bg-exo-yellow"
+                            ></span>
+                          {/if}
+                        </span>
+                        Attn/MoE Split
+                      </button>
                     </div>
                   </div>
 

resources/inference_model_cards/mlx-community--Qwen3.5-27B-bf16.toml

@@ -0,0 +1,15 @@
+model_id = "mlx-community/Qwen3.5-27B-bf16"
+n_layers = 64
+hidden_size = 5120
+num_key_value_heads = 4
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "bf16"
+base_model = "Qwen3.5 27B"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 54760833024

src/exo/api/main.py

@@ -474,7 +474,7 @@ class API:
                 status_code=400, detail=f"Failed to load model card: {exc}"
             ) from exc
         instance_combinations: list[tuple[Sharding, InstanceMeta, int]] = []
-        for sharding in (Sharding.Pipeline, Sharding.Tensor):
+        for sharding in (Sharding.Pipeline, Sharding.Tensor, Sharding.AttnMoeSplit):
             for instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl):
                 instance_combinations.extend(
                     [

src/exo/master/placement.py

@@ -153,6 +153,20 @@ def place_instance(
         raise ValueError(
             "Pipeline parallelism is not supported for DeepSeek V3.1 (8-bit)"
         )
+    if command.sharding == Sharding.AttnMoeSplit:
+        model_id_lower = str(command.model_card.model_id).lower()
+        if "qwen3.5" not in model_id_lower and "qwen-3.5" not in model_id_lower:
+            raise ValueError(
+                "AttnMoeSplit sharding is only supported for Qwen3.5 MoE models, "
+                f"got {command.model_card.model_id}"
+            )
+        cycles_with_sufficient_memory = [
+            cycle for cycle in cycles_with_sufficient_memory if len(cycle) == 2
+        ]
+        if not cycles_with_sufficient_memory:
+            raise ValueError(
+                "AttnMoeSplit sharding requires exactly 2 nodes in the cycle"
+            )
 
     smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
 
@@ -191,8 +205,16 @@ def place_instance(
         ),
     )
 
-    # Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
+    # Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node).
+    # AttnMoeSplit was already filtered to exactly-2-node cycles above, so
+    # this branch can never fire for AttnMoeSplit — if it somehow did we'd
+    # silently rewrite the sharding mode, which is the wrong behavior.
     if len(selected_cycle) == 1:
+        if command.sharding == Sharding.AttnMoeSplit:
+            raise ValueError(
+                "AttnMoeSplit sharding requires 2 nodes but a single-node cycle "
+                "was selected"
+            )
         command.instance_meta = InstanceMeta.MlxRing
         command.sharding = Sharding.Pipeline
 

src/exo/master/placement_utils.py

@@ -10,6 +10,7 @@ from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
 from exo.shared.types.topology import Cycle, RDMAConnection, SocketConnection
 from exo.shared.types.worker.runners import RunnerId, ShardAssignments
 from exo.shared.types.worker.shards import (
+    AttnMoeSplitShardMetadata,
     CfgShardMetadata,
     PipelineShardMetadata,
     Sharding,
@@ -273,6 +274,44 @@ def get_shard_assignments_for_tensor_parallel(
     return shard_assignments
 
 
+def get_shard_assignments_for_attn_moe_split(
+    model_card: ModelCard,
+    cycle: Cycle,
+) -> ShardAssignments:
+    """Create shard assignments for attention/MoE split execution.
+
+    Both ranks own the full layer range [0, n_layers). device_rank=0 runs
+    attention, device_rank=1 runs MoE. Requires exactly 2 nodes.
+    """
+    if len(cycle) != 2:
+        raise ValueError(
+            f"AttnMoeSplit sharding requires exactly 2 nodes, got {len(cycle)}"
+        )
+    total_layers = model_card.n_layers
+    world_size = 2
+    runner_to_shard: dict[RunnerId, ShardMetadata] = {}
+    node_to_runner: dict[NodeId, RunnerId] = {}
+
+    for i, node_id in enumerate(cycle):
+        shard = AttnMoeSplitShardMetadata(
+            model_card=model_card,
+            device_rank=i,
+            world_size=world_size,
+            start_layer=0,
+            end_layer=total_layers,
+            n_layers=total_layers,
+        )
+        runner_id = RunnerId()
+        runner_to_shard[runner_id] = shard
+        node_to_runner[node_id] = runner_id
+
+    return ShardAssignments(
+        model_id=model_card.model_id,
+        runner_to_shard=runner_to_shard,
+        node_to_runner=node_to_runner,
+    )
+
+
 def get_shard_assignments(
     model_card: ModelCard,
     cycle: Cycle,
@@ -291,6 +330,11 @@ def get_shard_assignments(
                 model_card=model_card,
                 cycle=cycle,
             )
+        case Sharding.AttnMoeSplit:
+            return get_shard_assignments_for_attn_moe_split(
+                model_card=model_card,
+                cycle=cycle,
+            )
 
 
 def get_mlx_jaccl_devices_matrix(

src/exo/shared/types/worker/shards.py

@@ -10,6 +10,7 @@ from exo.utils.pydantic_ext import TaggedModel
 class Sharding(str, Enum):
     Tensor = "Tensor"
     Pipeline = "Pipeline"
+    AttnMoeSplit = "AttnMoeSplit"
 
 
 class BaseShardMetadata(TaggedModel):
@@ -79,6 +80,20 @@ class TensorShardMetadata(BaseShardMetadata):
     pass
 
 
+@final
+class AttnMoeSplitShardMetadata(BaseShardMetadata):
+    """Attention/MoE split shard meta.
+
+    Both ranks own the full layer range [0, n_layers). device_rank=0 runs
+    per-layer attention + first residual; device_rank=1 runs
+    post_attention_layernorm + MoE + second residual. One cross-rank send/recv
+    pair per layer ferries the hidden state. world_size must be 2.
+    """
+
+
 ShardMetadata: TypeAlias = (
-    PipelineShardMetadata | CfgShardMetadata | TensorShardMetadata
+    PipelineShardMetadata
+    | CfgShardMetadata
+    | TensorShardMetadata
+    | AttnMoeSplitShardMetadata
 )

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

@@ -57,7 +57,10 @@ from mlx_lm.models.step3p5 import Model as Step35Model
 from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP
 from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel
 
-from exo.shared.types.worker.shards import PipelineShardMetadata
+from exo.shared.types.worker.shards import (
+    AttnMoeSplitShardMetadata,
+    PipelineShardMetadata,
+)
 from exo.worker.runner.bootstrap import logger
 
 if TYPE_CHECKING:
@@ -430,6 +433,48 @@ def pipeline_auto_parallel(
     return patch_pipeline_model(model, group)
 
 
+def attn_moe_split_auto_parallel(
+    model: nn.Module,
+    group: mx.distributed.Group,
+    model_shard_meta: AttnMoeSplitShardMetadata,
+    on_layer_loaded: LayerLoadedCallback | None,
+) -> nn.Module:
+    """Install the attention/MoE split on a Qwen3.5 MoE model.
+
+    Both ranks own the full layer range (start_layer=0, end_layer=n_layers).
+    No layer wrapping: the split is installed via class-level
+    DecoderLayer.__call__ replacement in
+    patches/qwen3_5_moe_split/apply.py.
+    """
+    if group.size() != 2:
+        raise ValueError(
+            f"AttnMoeSplit requires world_size==2, got {group.size()}"
+        )
+
+    # Qwen3.5 MoE uses Qwen3_5TextModelInner. Any other model type landing
+    # here is a placement bug.
+    inner_model_instance: nn.Module = get_inner_model(model)
+    if not isinstance(inner_model_instance, Qwen3_5TextModelInner):
+        raise ValueError(
+            "AttnMoeSplit sharding is only implemented for Qwen3.5 MoE models, "
+            f"got inner model of type {type(inner_model_instance).__name__}"
+        )
+
+    layers = get_layers(inner_model_instance)
+    total = len(layers)
+    for i, layer in enumerate(layers):
+        mx.eval(layer)  # type: ignore
+        if on_layer_loaded is not None:
+            on_layer_loaded(i, total)
+
+    from exo.worker.engines.mlx.patches.qwen3_5_moe_split.apply import (
+        apply_qwen35_attn_moe_split_patches,
+    )
+
+    apply_qwen35_attn_moe_split_patches(model, group)
+    return model
+
+
 def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
     # Patch __call__ on the model's class
     cls = model.__class__

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

@@ -1,4 +1,5 @@
 import contextlib
+import os
 import time
 from dataclasses import dataclass, field
 from typing import Callable, cast
@@ -96,13 +97,392 @@ class ExoBatchGenerator:
     _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False)
 
     def __post_init__(self) -> None:
-        self._mlx_gen = MlxBatchGenerator(
-            model=self.model,
-            stop_tokens=set(eos_ids_from_tokenizer(self.tokenizer)),
-            prefill_step_size=4096,
-        )
+        use_speculative = os.environ.get("EXO_SPECULATIVE", "0") == "1"
+        spec_mode = os.environ.get("EXO_SPECULATIVE_MODE", "mtp")
+        stop_tokens = set(eos_ids_from_tokenizer(self.tokenizer))
+        temp = float(os.environ.get("EXO_SPECULATIVE_TEMP", "0.7"))
+        alpha = float(os.environ.get("EXO_SPECULATIVE_ALPHA", "1.0"))
+
+        if use_speculative and spec_mode == "dflash":
+            try:
+                from exo.worker.engines.mlx.speculative.dflash_module import DFlashDrafter
+                from exo.worker.engines.mlx.speculative.dflash_batch_generator import DFlashBatchGenerator
+                from exo.worker.engines.mlx.speculative.bf16_lpb_patch import apply_bf16_lpb_patches
+
+                dflash_path = os.environ.get("EXO_DFLASH_MODEL", "z-lab/Qwen3.5-27B-DFlash")
+                verify_len = int(os.environ.get("EXO_DFLASH_VERIFY", "5"))
+                block_size = int(os.environ.get("EXO_DFLASH_BLOCK_SIZE", "6"))
+
+                drafter = DFlashDrafter(self.model, dflash_path)
+                apply_bf16_lpb_patches(drafter)
+
+                self._mlx_gen = DFlashBatchGenerator(
+                    model=self.model,
+                    drafter=drafter,
+                    verify_len=verify_len,
+                    block_size=block_size,
+                    temp=temp,
+                    alpha=alpha,
+                    stop_tokens=stop_tokens,
+                    prefill_step_size=4096,
+                )
+                logger.info(f"DFlash speculative decoding enabled (V={verify_len}, BS={block_size}, T={temp})")
+                self.warmup_dflash(self.model, self.tokenizer, drafter)
+            except Exception as e:
+                logger.warning(f"Failed to init DFlash speculative: {e}. Falling back.")
+                self._mlx_gen = MlxBatchGenerator(model=self.model, stop_tokens=stop_tokens, prefill_step_size=4096)
+
+        elif use_speculative and spec_mode == "mtp":
+            try:
+                from exo.worker.engines.mlx.speculative.mtp_module import MTPPredictor
+                from exo.worker.engines.mlx.speculative.mtp_batch_generator import MTPBatchGenerator
+
+                mtp_weights = self._resolve_mtp_weights()
+                gamma = int(os.environ.get("EXO_SPECULATIVE_GAMMA", "2"))
+
+                if mtp_weights:
+                    mtp = MTPPredictor(self.model, mtp_weights, quantize=False)
+                    self._mlx_gen = MTPBatchGenerator(
+                        model=self.model,
+                        mtp_predictor=mtp,
+                        gamma=gamma,
+                        temp=temp,
+                        alpha=alpha,
+                        stop_tokens=stop_tokens,
+                        prefill_step_size=4096,
+                    )
+                    logger.info(f"MTP speculative decoding enabled (γ={gamma}, T={temp})")
+                    self.warmup_speculative(self.model, self.tokenizer)
+                else:
+                    logger.warning("EXO_SPECULATIVE=1 but could not find MTP weights. Falling back.")
+                    self._mlx_gen = MlxBatchGenerator(model=self.model, stop_tokens=stop_tokens, prefill_step_size=4096)
+            except Exception as e:
+                logger.warning(f"Failed to init MTP speculative: {e}. Falling back.")
+                self._mlx_gen = MlxBatchGenerator(model=self.model, stop_tokens=stop_tokens, prefill_step_size=4096)
+        else:
+            self._mlx_gen = MlxBatchGenerator(model=self.model, stop_tokens=stop_tokens, prefill_step_size=4096)
+
         self._mlx_gen._needs_topk = False  # pyright: ignore[reportAttributeAccessIssue]
 
+    def _resolve_mtp_weights(self) -> str | None:
+        """Find MTP weights: explicit path, explicit HF model, or auto-extract."""
+        explicit_path = os.environ.get("EXO_MTP_WEIGHTS", "")
+        if explicit_path and os.path.exists(explicit_path):
+            return explicit_path
+
+        mtp_model = os.environ.get("EXO_MTP_MODEL", "")
+        if not mtp_model:
+            try:
+                inner = getattr(self.model, 'model', None) or self.model.language_model.model
+                args = getattr(inner, 'args', None)
+                if args and getattr(args, 'mtp_num_hidden_layers', 0) > 0:
+                    model_type = getattr(args, 'model_type', '')
+                    if 'qwen3_5' in model_type or 'qwen3.5' in str(type(self.model).__module__):
+                        mtp_model = "Qwen/Qwen3.5-27B"
+                        logger.info(f"Auto-detected MTP model: {mtp_model}")
+            except Exception:
+                pass
+
+        if not mtp_model:
+            return None
+
+        try:
+            return self._extract_mtp_from_hf(mtp_model)
+        except Exception as e:
+            logger.warning(f"Failed to extract MTP weights from {mtp_model}: {e}")
+            return None
+
+    def _extract_mtp_from_hf(self, repo_id: str) -> str:
+        """Download MTP tensors from HF repo and cache as a single safetensors file.
+
+        Three-tier strategy, smallest download first:
+          1. Index-aware byte-range fetch: read the safetensors index, find
+             MTP-bearing shards, then for each shard read only its header
+             (small JSON describing tensor byte offsets) and range-fetch
+             just the bytes belonging to MTP tensors. For Qwen/Qwen3.5-27B
+             this is ~500 MB instead of 55 GB.
+          2. Shard download (fallback if byte-range fails): pull only
+             MTP-bearing shards (~20 GB).
+          3. Full download (fallback if no index): pull every safetensors
+             file (~55 GB).
+
+        Both 'model.mtp.' and 'mtp.' tensor-key prefixes are accepted —
+        Qwen/Qwen3.5-27B's repo uses the bare 'mtp.' prefix; some variants
+        may use 'model.mtp.'.
+        """
+        import hashlib
+        import json as _json
+        from pathlib import Path
+        from huggingface_hub import hf_hub_download, snapshot_download
+        from safetensors.torch import load_file, save_file
+
+        cache_dir = Path.home() / ".cache" / "exo" / "mtp_weights"
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        cache_key = hashlib.md5(repo_id.encode()).hexdigest()[:12]
+        # v2: preserve the 'mtp.' prefix in cached keys (MTPPredictor expects it).
+        # Older caches (mtp_<key>.safetensors) had the prefix stripped — invalidate
+        # them by using a new filename.
+        cached_path = cache_dir / f"mtp_v2_{cache_key}.safetensors"
+
+        if cached_path.exists():
+            logger.info(f"Using cached MTP weights: {cached_path}")
+            return str(cached_path)
+
+        def _is_mtp_key(k: str) -> bool:
+            return k.startswith("model.mtp.") or k.startswith("mtp.")
+
+        def _normalize_mtp_key(k: str) -> str:
+            # MTPPredictor._load_weights expects keys with the 'mtp.' prefix
+            # preserved. Strip the optional 'model.' wrapper only.
+            if k.startswith("model.mtp."):
+                return k[len("model."):]
+            return k
+
+        def _save_and_return(mtp_tensors: dict) -> str:
+            save_file(mtp_tensors, str(cached_path))
+            logger.info(f"Cached {len(mtp_tensors)} MTP tensors to {cached_path}")
+            return str(cached_path)
+
+        # ── Tier 1: read index, identify MTP shards
+        try:
+            index_path = hf_hub_download(repo_id, "model.safetensors.index.json")
+            with open(index_path) as f:
+                index = _json.load(f)
+            weight_map = index.get("weight_map", {})
+            mtp_shards = sorted(
+                {shard for name, shard in weight_map.items() if _is_mtp_key(name)}
+            )
+            if not mtp_shards:
+                raise ValueError(
+                    f"No keys with prefix 'model.mtp.' or 'mtp.' in "
+                    f"{repo_id}'s weight_map"
+                )
+            total_shards = len(set(weight_map.values()))
+            logger.info(
+                f"MTP head spans {len(mtp_shards)}/{total_shards} shard(s) of "
+                f"{repo_id}: {mtp_shards}"
+            )
+        except Exception as e:
+            logger.warning(
+                f"Couldn't read sharded index from {repo_id} ({e}); "
+                f"falling back to full safetensors download"
+            )
+            model_dir = snapshot_download(
+                repo_id, allow_patterns=["*.safetensors", "*.json"]
+            )
+            mtp_tensors = {}
+            for sf_file in sorted(Path(model_dir).glob("*.safetensors")):
+                tensors = load_file(str(sf_file))
+                for k, v in tensors.items():
+                    if _is_mtp_key(k):
+                        mtp_tensors[_normalize_mtp_key(k)] = v
+            if not mtp_tensors:
+                raise ValueError(f"No MTP tensors found in {repo_id}")
+            return _save_and_return(mtp_tensors)
+
+        # ── Tier 1 continued: byte-range fetch only the MTP tensors
+        try:
+            from huggingface_hub import HfFileSystem
+            import torch
+
+            ST_DTYPE_TO_TORCH = {
+                "BF16": torch.bfloat16,
+                "F16": torch.float16,
+                "F32": torch.float32,
+                "F64": torch.float64,
+                "I8": torch.int8,
+                "I16": torch.int16,
+                "I32": torch.int32,
+                "I64": torch.int64,
+                "U8": torch.uint8,
+                "BOOL": torch.bool,
+            }
+            fs = HfFileSystem()
+            mtp_tensors = {}
+            total_bytes = 0
+            for shard in mtp_shards:
+                with fs.open(f"{repo_id}/{shard}", mode="rb") as f:
+                    header_size = int.from_bytes(f.read(8), "little")
+                    header = _json.loads(f.read(header_size).decode("utf-8"))
+                    data_start = 8 + header_size
+                    in_shard = [
+                        (name, meta)
+                        for name, meta in header.items()
+                        if name != "__metadata__" and _is_mtp_key(name)
+                    ]
+                    logger.info(
+                        f"  {shard}: extracting {len(in_shard)} MTP tensor(s)"
+                    )
+                    for name, meta in in_shard:
+                        start, end = meta["data_offsets"]
+                        f.seek(data_start + start)
+                        raw = f.read(end - start)
+                        dtype = ST_DTYPE_TO_TORCH.get(meta["dtype"])
+                        if dtype is None:
+                            raise ValueError(
+                                f"Unsupported safetensors dtype {meta['dtype']} "
+                                f"for {name}"
+                            )
+                        t = torch.frombuffer(
+                            bytearray(raw), dtype=dtype
+                        ).reshape(meta["shape"])
+                        mtp_tensors[_normalize_mtp_key(name)] = t
+                        total_bytes += len(raw)
+            if not mtp_tensors:
+                raise ValueError("Byte-range read returned no MTP tensors")
+            logger.info(
+                f"Byte-range fetched {len(mtp_tensors)} MTP tensors "
+                f"({total_bytes / 1e6:.1f} MB) from {len(mtp_shards)} shards"
+            )
+            return _save_and_return(mtp_tensors)
+        except Exception as e:
+            logger.warning(
+                f"Byte-range fetch failed ({e}); "
+                f"falling back to MTP-shard snapshot download"
+            )
+
+        # ── Tier 2: download just the MTP-bearing shards
+        model_dir = snapshot_download(
+            repo_id, allow_patterns=list(mtp_shards) + ["*.json"]
+        )
+        mtp_tensors = {}
+        for sf_file in sorted(Path(model_dir).glob("*.safetensors")):
+            tensors = load_file(str(sf_file))
+            for k, v in tensors.items():
+                if _is_mtp_key(k):
+                    mtp_tensors[_normalize_mtp_key(k)] = v
+        if not mtp_tensors:
+            raise ValueError(f"No MTP tensors found in {repo_id} (shard fallback)")
+        return _save_and_return(mtp_tensors)
+
+    def warmup_speculative(self, model, tokenizer) -> None:
+        """Warm up the speculative decoding path (MTP draft + verify kernels)."""
+        if not hasattr(self._mlx_gen, 'mtp'):
+            return
+
+        from mlx_lm.models import cache as cache_mod
+        from exo.worker.engines.mlx.speculative.mtp_module import speculative_forward, draft_tokens
+
+        logger.info("Warming up speculative decoding kernels...")
+        mtp = self._mlx_gen.mtp
+        gamma = self._mlx_gen.gamma
+
+        warmup_prompt = tokenizer.encode("Warm up speculative decoding.")
+        cache = cache_mod.make_prompt_cache(model)
+        mtp.reset_cache()
+
+        pre_norm, logits = speculative_forward(model, mx.array([warmup_prompt]), cache)
+        mx.eval(pre_norm, logits)
+        next_token = mx.argmax(logits[0, -1], axis=-1).item()
+
+        if pre_norm.shape[1] > 1:
+            _ = mtp.predict(pre_norm[:, :-1, :], mx.array([warmup_prompt[1:]]))
+            mx.eval(_)
+
+        last_pn = pre_norm[:, -1:, :]
+        next_arr = mx.array([[next_token]])
+        for _ in range(3):
+            draft_ids, _ = draft_tokens(mtp, last_pn, next_arr, gamma, 0.0)
+            draft_concat = mx.concatenate([d.reshape(1, 1) for d in draft_ids], axis=1)
+            verify_input = mx.concatenate([next_arr, draft_concat], axis=1)
+            vpn, vl = speculative_forward(model, verify_input, cache, speculative=True)
+            all_next = mx.argmax(vl[0], axis=-1)
+            mx.eval(vpn, all_next)
+            next_arr = all_next[0].reshape(1, 1)
+            last_pn = vpn[:, 0:1, :]
+            for i, c in enumerate(cache):
+                if hasattr(c, 'base'):
+                    cache[i] = c.base
+
+        logger.info("Speculative warmup complete")
+
+    def warmup_dflash(self, model, tokenizer, drafter) -> None:
+        """Warm up the DFlash speculative decoding path.
+
+        With the dynamic kernel picker every projection memoizes kernels
+        keyed by the actual M seen in its forward. At runtime the M set
+        each call-site can hit is:
+
+            target q/k/v/o, gate/up/down, in_proj_*, lm_head:  M = V+1
+            drafter q/o/mlp_*:                                  M = BS
+            drafter fc:                                         M = S_ctx ∈ [1, V+1]
+            drafter k_proj/v_proj:                              M = BS + S_ctx ∈ [BS+1, BS+V+1]
+            drafter lm_head:                                    M = BS - 1
+
+        To compile every one of those kernels up front we sweep
+        S_ctx = 1..V+1 against the drafter (calling draft() with
+        fresh caches and different target-hidden slice lengths),
+        then run a single target verify at M=V+1.
+
+        `modes` is a list of (block_size, verify_len) pairs; warmup is
+        run once per mode so dynamic-switching configurations can jump
+        between modes at runtime without triggering Metal compilation
+        stalls.
+        """
+        from mlx_lm.models import cache as cache_mod
+        from exo.worker.engines.mlx.speculative.dflash_speculative import dflash_speculative_forward
+
+        logger.info("Warming up DFlash speculative decoding kernels...")
+
+        original_bs = drafter.block_size
+        verify_len = self._mlx_gen.verify_len  # pyright: ignore[reportAttributeAccessIssue]
+        modes = [(original_bs, verify_len)]
+
+        max_v = max(v for _, v in modes)
+        ctx_len = max(64, max_v + 2)
+        prompt_tokens = [1] * ctx_len
+
+        cache = cache_mod.make_prompt_cache(model)
+        drafter.reset_draft_cache()
+
+        target_hidden_full, _, logits = dflash_speculative_forward(
+            model, mx.array([prompt_tokens]), cache, drafter.target_layer_ids)
+        mx.eval(target_hidden_full, logits)
+        next_token = mx.argmax(logits[0, -1], axis=-1).item()
+
+        try:
+            for m_bs, m_v in modes:
+                drafter.block_size = m_bs
+
+                # 1. Sweep S_ctx ∈ [1, V+1] so the drafter's fc, k_proj,
+                #    v_proj compile every kernel the picker may return.
+                #    Each iteration rebuilds the draft KV cache, so the
+                #    fc/k/v projections see exactly s_ctx / s_ctx+m_bs.
+                for s_ctx in range(1, m_v + 2):
+                    drafter.reset_draft_cache()
+                    last_th = target_hidden_full[:, -s_ctx:, :]
+                    block_ids = mx.full(
+                        (1, m_bs), drafter.mask_token_id, dtype=mx.int32
+                    )
+                    block_ids[:, 0] = next_token
+                    dl = drafter.draft(last_th, block_ids, ctx_len)
+                    mx.eval(dl)
+
+                # 2. One target verify at M = V+1 to compile every target
+                #    projection (q/k/v/o, gate/up/down, in_proj_*, lm_head).
+                verify_input = mx.array([[next_token] * (m_v + 1)])
+                target_hidden, _, vl = dflash_speculative_forward(
+                    model, verify_input, cache, drafter.target_layer_ids,
+                    speculative=True,
+                )
+                mx.eval(target_hidden, vl)
+
+                # 3. Fully undo the verify so the next mode (or the real
+                #    generation) starts from the same cache offset.
+                for c in cache:
+                    if hasattr(c, 'offset'):
+                        c.offset -= (m_v + 1)
+                    elif hasattr(c, 'rollback'):
+                        c.rollback(0)
+                for i, c in enumerate(cache):
+                    if hasattr(c, 'base'):
+                        cache[i] = c.base
+        finally:
+            drafter.block_size = original_bs
+            drafter.reset_draft_cache()
+
+        logger.info("DFlash warmup complete")
+
     @property
     def has_work(self) -> bool:
         return (
@@ -254,6 +634,31 @@ class ExoBatchGenerator:
 
         uid = uids[0]
 
+        # MTP prefill: build MTP cache from prompt hidden states
+        if hasattr(self._mlx_gen, 'mtp'):
+            prompt_pre_norm = self._mlx_gen._captured.get('prompt_pre_norm')
+            if prompt_pre_norm is not None:
+                mx.eval(prompt_pre_norm)
+                self._mlx_gen.mtp.reset_cache()
+                S_pre = prompt_pre_norm.shape[1]
+                if S_pre > 1:
+                    toks_list = all_prompt_tokens.tolist() if hasattr(all_prompt_tokens, 'tolist') else list(all_prompt_tokens)
+                    mtp_tokens = toks_list[1:S_pre]
+                    _ = self._mlx_gen.mtp.predict(
+                        prompt_pre_norm[:, :-1, :],
+                        mx.array([mtp_tokens])
+                    )
+                    mx.eval(_)
+                    logger.info(f"MTP cache prefilled ({S_pre} positions)")
+
+        # Set per-request temperature for speculative
+        if hasattr(self._mlx_gen, '_request_temp'):
+            env_temp = os.environ.get("EXO_SPECULATIVE_TEMP")
+            if env_temp is not None:
+                self._mlx_gen._request_temp[uid] = float(env_temp)
+            elif task_params.temperature is not None:
+                self._mlx_gen._request_temp[uid] = task_params.temperature
+
         self._active_tasks[uid] = _EngineTask(
             uid=uid,
             task_params=task_params,
@@ -337,7 +742,7 @@ class ExoBatchGenerator:
 
             logprob: float | None = None
             top_logprobs: list[TopLogprobItem] | None = None
-            if task_params.logprobs:
+            if task_params.logprobs and os.environ.get("EXO_DISABLE_LOGPROBS") != "1":
                 with mx.stream(generation_stream):
                     logprob, top_logprobs = extract_top_logprobs(
                         logprobs=response.logprobs,

src/exo/worker/engines/mlx/matmul/kernels/bf16/gemm_bm8.py

@@ -0,0 +1,179 @@
+#!/usr/bin/env python3
+"""Custom bf16 GEMM with BM=8. Same pattern as custom_qmm_bm8 but no dequantization."""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_bf16_gemm_source(M_val, N_val, K_val, BM):
+    BN = 32
+    BK = 32
+    if BM <= 8:
+        WM, WN = 1, 2
+    else:
+        WM, WN = 2, 2
+    tgp_size = WM * WN * 32
+    BK_PAD = BK + 8
+
+    # BlockLoader for W: (BROWS=BN=32, BCOLS=BK=32, tgp_size)
+    w_n_reads = (BN * BK) // tgp_size
+
+    # BlockLoader for X: (BROWS=BM, BCOLS=BK=32, tgp_size)
+    x_total = BM * BK
+    x_n_reads = max(1, x_total // tgp_size)
+    x_TCOLS = BK // x_n_reads if x_n_reads > 0 else BK
+    x_TROWS = tgp_size // x_TCOLS if x_TCOLS > 0 else tgp_size
+
+    # MMA fragment layout
+    TM = BM // (8 * WM)  # fragments per SG in M dimension
+    TN = BN // (8 * WN)  # fragments per SG in N dimension
+
+    # Generate MMA accumulators, load, and store for TM×TN fragments
+    # C[tm_i][tn_j] for tm_i in 0..TM-1, tn_j in 0..TN-1
+    c_decl_lines = []
+    for ti in range(TM):
+        for tj in range(TN):
+            c_decl_lines.append(f"    simdgroup_matrix<float, 8, 8> C{ti}{tj} = simdgroup_matrix<float, 8, 8>(0);")
+    c_decl = "\n".join(c_decl_lines)
+
+    # MMA inner loop: load A for each TM row, B for each TN col, multiply
+    mma_lines = []
+    for ti in range(TM):
+        mma_lines.append(f"""            simdgroup_matrix<float, 8, 8> A{ti};
+            simdgroup_load(A{ti}, &Xs[(tm + {ti} * {WM} * 8) * BK_PAD + kk], BK_PAD);""")
+    for tj in range(TN):
+        mma_lines.append(f"""            simdgroup_matrix<float, 8, 8> B{tj};
+            simdgroup_load(B{tj}, &Ws[(tn + {tj} * {WN} * 8) * BK_PAD + kk], BK_PAD, ulong2(0, 0), true);""")
+    for ti in range(TM):
+        for tj in range(TN):
+            mma_lines.append(f"            simdgroup_multiply_accumulate(C{ti}{tj}, A{ti}, B{tj}, C{ti}{tj});")
+    mma_load = "\n".join(mma_lines)
+
+    # Store: each fragment at its (tm + ti*8, tn + tj*WN*8) position
+    store_lines = []
+    for ti in range(TM):
+        for tj in range(TN):
+            store_lines.append(f"    simdgroup_store(C{ti}{tj}, &Ws[(tm + {ti} * {WM} * 8) * BN + tn + {tj} * {WN} * 8], BN);")
+    c_store = "\n".join(store_lines)
+
+    return f"""
+    const int BM = {BM};
+    const int BN = {BN};
+    const int BK = {BK};
+    const int BK_PAD = {BK_PAD};
+    const int K = {K_val};
+    const int N = {N_val};
+    const int M = {M_val};
+
+    uint3 tid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+    int thread_idx = sgid * 32 + slid;
+
+    int y_row = tid.y * BM;
+    int y_col = tid.x * BN;
+
+    threadgroup float Xs[{BM} * {BK_PAD}];
+    threadgroup float Ws[{BN} * {BK_PAD}];
+
+    // Pointer setup
+    const device bfloat16_t* x_base = (const device bfloat16_t*)x + y_row * K;
+    const device bfloat16_t* w_base = (const device bfloat16_t*)w + y_col * K;
+
+    // BlockLoader for X
+    int x_bi = thread_idx / {x_TCOLS};
+    int x_bj = {x_n_reads} * (thread_idx % {x_TCOLS});
+    const device bfloat16_t* x_src = x_base + x_bi * K + x_bj;
+    threadgroup float* x_dst = Xs + x_bi * BK_PAD + x_bj;
+
+    // BlockLoader for W (bf16, no dequant)
+    int w_bi = {w_n_reads} * thread_idx / {BK};
+    int w_bj = ({w_n_reads} * thread_idx) % {BK};
+    const device bfloat16_t* w_src = w_base + w_bi * K + w_bj;
+    threadgroup float* w_dst = Ws + w_bi * BK_PAD + w_bj;
+
+    // MMA setup
+    short sg_row = sgid / {WN};
+    short sg_col = sgid % {WN};
+    short tm = 8 * sg_row;
+    short tn = 8 * sg_col;
+    {c_decl}
+
+    // K-loop
+    for (int k = 0; k < K; k += BK) {{
+
+        // Load X tile (bf16 → float)
+        if (x_bi < BM) {{
+            for (int i = 0; i < {x_n_reads}; i++) {{
+                x_dst[i] = float(x_src[i]);
+            }}
+        }}
+
+        // Load W tile (bf16 → float, no dequant)
+        if (w_bi < BN) {{
+            for (int i = 0; i < {w_n_reads}; i++) {{
+                w_dst[i] = float(w_src[i]);
+            }}
+        }}
+
+        threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+        // MMA inner loop
+        for (short kk = 0; kk < BK; kk += 8) {{
+{mma_load}
+        }}
+
+        threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+        x_src += BK;
+        w_src += BK;
+    }}
+
+    // Store results via TG memory → device
+    {c_store}
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+    for (int i = thread_idx; i < BM * BN; i += {tgp_size}) {{
+        int r = i / BN;
+        int c = i % BN;
+        if (y_row + r < M && y_col + c < N) {{
+            y[(y_row + r) * N + y_col + c] = static_cast<bfloat16_t>(Ws[r * BN + c]);
+        }}
+    }}
+"""
+
+
+_cache = {}
+
+
+def custom_bf16_gemm(x, w, M, N, K, BM=8):
+    key = (M, N, K, BM)
+    if key not in _cache:
+        WM = 1 if BM <= 8 else 2
+        WN = 2
+        sgs = WM * WN
+        _cache[key] = mx.fast.metal_kernel(
+            name=f"custom_bf16_gemm_bm{BM}_M{M}_N{N}_K{K}",
+            input_names=["x", "w"],
+            output_names=["y"],
+            source=_gen_bf16_gemm_source(M, N, K, BM),
+        )
+    kern = _cache[key]
+
+    WM = 1 if BM <= 8 else 2
+    WN = 2
+    sgs = WM * WN
+    n_tg_n = ceil_div(N, 32)
+    n_tg_m = ceil_div(M, BM)
+
+    result = kern(
+        inputs=[x, w],
+        output_shapes=[(M * N,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(32 * n_tg_n, sgs * n_tg_m, 1),
+        threadgroup=(32, sgs, 1),
+    )
+    return result[0].reshape(M, N)

src/exo/worker/engines/mlx/matmul/kernels/bf16/gemm_splitk_steel.py

@@ -0,0 +1,271 @@
+#!/usr/bin/env python3
+"""Split-K bf16 GEMM using actual Steel GEMM infrastructure.
+
+Bundles MLX's Steel headers and passes them via mx.fast.metal_kernel's
+header parameter. The kernel source is just the template instantiation
+of gemm_splitk — identical to stock MLX.
+
+Only difference from stock: no K >= max(M, N) restriction.
+
+Usage:
+    from gemm_splitk_steel import custom_bf16_gemm_splitk_steel
+    y = custom_bf16_gemm_splitk_steel(x, w, M=16, N=17408, K=5120, BM=16)
+"""
+
+import os
+import re
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def next_power_of_2(n):
+    if n <= 0:
+        return 1
+    return 1 << (n - 1).bit_length()
+
+
+def compute_partitions(M, N, K):
+    _tm = ceil_div(M, 32)
+    _tn = ceil_div(N, 32)
+    _tk = K // 16
+    tmtn = max(_tm * _tn, 1)
+    return min(max(2, next_power_of_2(_tk // tmtn)), 32)
+
+
+def _load_steel_headers():
+    """Read and concatenate Steel headers in dependency order."""
+    base = os.path.join(os.path.dirname(__file__), '..', 'steel_include')
+
+    files_in_order = [
+        'steel/defines.h',
+        'steel/utils/integral_constant.h',
+        'steel/utils/type_traits.h',
+        'steel/utils.h',
+        'steel/gemm/transforms.h',
+        'steel/gemm/params.h',
+        'steel/gemm/loader.h',
+        'steel/gemm/mma.h',
+        'steel/gemm/gemm.h',
+        'steel/gemm/kernels/steel_gemm_splitk.h',
+    ]
+
+    parts = []
+    for f in files_in_order:
+        path = os.path.join(base, f)
+        with open(path) as fh:
+            content = fh.read()
+        # Strip #pragma once and #include directives (deps are manually ordered)
+        content = re.sub(r'#pragma once', '', content)
+        content = re.sub(r'#include\s+"[^"]*"', '', content)
+        content = re.sub(r'#include\s+<metal_simdgroup_matrix>', '', content)
+        content = re.sub(r'#include\s+<metal_simdgroup>', '', content)
+        content = re.sub(r'#include\s+<metal_stdlib>', '', content)
+        parts.append(f'// ── {f} ──\n' + content)
+
+    return '\n'.join(parts)
+
+
+_steel_header = None
+
+
+def _get_steel_header():
+    global _steel_header
+    if _steel_header is None:
+        _steel_header = _load_steel_headers()
+    return _steel_header
+
+
+def _gen_splitk_instantiation(M_val, N_val, K_val, BM, split_k_partitions, BN=32, BK=16):
+    """Generate the kernel source — just template instantiation + dispatch logic.
+    The actual GEMM code comes from the Steel headers in the header parameter."""
+
+    WM = 1 if BM <= 8 else 2
+    WN = 2
+
+    gemm_k_iterations = (K_val // BK) // split_k_partitions
+    split_k_partition_size = gemm_k_iterations * BK
+
+    mn_aligned_str = "true" if M_val % BM == 0 and N_val % BN == 0 else "false"
+    k_aligned_str = "true" if K_val % BK == 0 else "false"
+
+    # The kernel body is a direct copy of steel_gemm_splitk.h's gemm_splitk function
+    # with template parameters expanded
+    return f"""
+    using namespace mlx::steel;
+
+    using T = bfloat16_t;
+    using U = float;
+
+    const int BM_val = {BM};
+    const int BN_val = {BN};
+    const int BK_val = {BK};
+
+    using gemm_kernel = GEMMKernel<T, U, {BM}, {BN}, {BK}, {WM}, {WN}, false, true, {mn_aligned_str}, {k_aligned_str}>;
+    using loader_a_t = typename gemm_kernel::loader_a_t;
+    using loader_b_t = typename gemm_kernel::loader_b_t;
+    using mma_t = typename gemm_kernel::mma_t;
+
+    threadgroup T As[gemm_kernel::tgp_mem_size_a];
+    threadgroup T Bs[gemm_kernel::tgp_mem_size_b];
+
+    uint simd_lane_id = thread_index_in_simdgroup;
+    uint simd_group_id = simdgroup_index_in_threadgroup;
+    uint3 tid = threadgroup_position_in_grid;
+
+    const int tiles_n = {ceil_div(N_val, BN)};
+    const int tiles_m = {ceil_div(M_val, BM)};
+    const int tid_x = tid.x;
+    const int tid_y = tid.y;
+    const int tid_z = tid.z;
+
+    if (tiles_n <= tid_x || tiles_m <= tid_y) return;
+
+    const int M = {M_val};
+    const int N = {N_val};
+    const int K = {K_val};
+    const int lda = K;  // A is (M, K) row-major
+    const int ldb = K;  // B is (N, K) row-major, transposed
+
+    const int split_k_partitions = {split_k_partitions};
+    const int split_k_partition_size = {split_k_partition_size};
+    const int split_k_partition_stride = M * N;
+    const int gemm_k_iterations_aligned = {gemm_k_iterations};
+
+    const int c_row = tid_y * {BM};
+    const int c_col = tid_x * {BN};
+    const int k_start = split_k_partition_size * tid_z;
+
+    // Pointer setup (matching steel_gemm_splitk.h)
+    // transpose_a=false: A += k_start + c_row * lda
+    // transpose_b=true:  B += k_start + c_col * ldb
+    const device T* A = (const device T*)x + (long)k_start + (long)c_row * lda;
+    const device T* B = (const device T*)w + (long)k_start + (long)c_col * ldb;
+    device U* C = (device U*)y + (long)split_k_partition_stride * tid_z + (long)c_row * N + c_col;
+
+    // Prepare loaders and MMA
+    thread loader_a_t loader_a(A, lda, As, simd_group_id, simd_lane_id);
+    thread loader_b_t loader_b(B, ldb, Bs, simd_group_id, simd_lane_id);
+    thread mma_t mma_op(simd_group_id, simd_lane_id);
+
+    int gemm_k_iters = gemm_k_iterations_aligned;
+
+    short tgp_bm = min((short){BM}, (short)(M - c_row));
+    short tgp_bn = min((short){BN}, (short)(N - c_col));
+    short leftover_bk = K % {BK};
+
+    // Main GEMM loop
+    if (tgp_bm == {BM} && tgp_bn == {BN}) {{
+        gemm_kernel::gemm_loop(
+            As, Bs, gemm_k_iters, loader_a, loader_b, mma_op,
+            tgp_bm, tgp_bn, leftover_bk,
+            LoopAlignment<true, true, true>{{}});
+    }} else if (tgp_bn == {BN}) {{
+        gemm_kernel::gemm_loop(
+            As, Bs, gemm_k_iters, loader_a, loader_b, mma_op,
+            tgp_bm, tgp_bn, leftover_bk,
+            LoopAlignment<false, true, true>{{}});
+    }} else if (tgp_bm == {BM}) {{
+        gemm_kernel::gemm_loop(
+            As, Bs, gemm_k_iters, loader_a, loader_b, mma_op,
+            tgp_bm, tgp_bn, leftover_bk,
+            LoopAlignment<true, false, true>{{}});
+    }} else {{
+        gemm_kernel::gemm_loop(
+            As, Bs, gemm_k_iters, loader_a, loader_b, mma_op,
+            tgp_bm, tgp_bn, leftover_bk,
+            LoopAlignment<false, false, true>{{}});
+    }}
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Last partition handles leftover K
+    if ((tid_z + 1) == split_k_partitions) {{
+        int gemm_k_iter_remaining = (K - (k_start + split_k_partition_size)) / {BK};
+        if (gemm_k_iter_remaining > 0)
+            gemm_kernel::gemm_loop(
+                As, Bs, gemm_k_iter_remaining, loader_a, loader_b, mma_op,
+                tgp_bm, tgp_bn, leftover_bk,
+                LoopAlignment<false, false, {k_aligned_str}>{{}});
+    }}
+
+    // Store results
+    if (tgp_bm == {BM} && tgp_bn == {BN}) {{
+        mma_op.store_result(C, N);
+    }} else {{
+        mma_op.store_result_safe(C, N, short2(tgp_bn, tgp_bm));
+    }}
+"""
+
+
+def _gen_accum_source(split_k_partitions):
+    return f"""
+    uint gid_x = thread_position_in_grid.x;
+    uint gid_y = thread_position_in_grid.y;
+    uint N_val = threads_per_grid.x;
+
+    int offset = gid_y * N_val + gid_x;
+    int stride = threads_per_grid.x * threads_per_grid.y;
+
+    float out = 0.0f;
+    for (int p = 0; p < {split_k_partitions}; p++) {{
+        out += ((const device float*)c_split)[offset + p * stride];
+    }}
+
+    y[gid_y * N_val + gid_x] = static_cast<bfloat16_t>(out);
+"""
+
+
+_gemm_cache = {}
+_accum_cache = {}
+
+
+def custom_bf16_gemm_splitk_steel(x, w, M, N, K, BM=16, BN=32, BK=16):
+    WM = 1 if BM <= 8 else 2
+    WN = 2
+    P = compute_partitions(M, N, K)
+
+    gemm_key = (M, N, K, BM, BN, BK, P)
+    if gemm_key not in _gemm_cache:
+        _gemm_cache[gemm_key] = mx.fast.metal_kernel(
+            name=f"splitk_steel_bm{BM}_bn{BN}_bk{BK}_M{M}_N{N}_K{K}_P{P}",
+            input_names=["x", "w"],
+            output_names=["y"],
+            header=_get_steel_header(),
+            source=_gen_splitk_instantiation(M, N, K, BM, P, BN=BN, BK=BK),
+        )
+    gemm_kern = _gemm_cache[gemm_key]
+
+    sgs = WM * WN
+    n_tg_n = ceil_div(N, BN)
+    n_tg_m = ceil_div(M, BM)
+
+    c_split = gemm_kern(
+        inputs=[x, w],
+        output_shapes=[(P * M * N,)],
+        output_dtypes=[mx.float32],
+        grid=(32 * n_tg_n, sgs * n_tg_m, P),
+        threadgroup=(32, sgs, 1),
+    )[0]
+
+    accum_key = P
+    if accum_key not in _accum_cache:
+        _accum_cache[accum_key] = mx.fast.metal_kernel(
+            name=f"splitk_accum_P{P}",
+            input_names=["c_split"],
+            output_names=["y"],
+            source=_gen_accum_source(P),
+        )
+    accum_kern = _accum_cache[accum_key]
+
+    y = accum_kern(
+        inputs=[c_split],
+        output_shapes=[(M * N,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(N, M, 1),
+        threadgroup=(min(N, 256), 1, 1),
+    )[0]
+
+    return y.reshape(M, N)

src/exo/worker/engines/mlx/matmul/kernels/bf16/gemv_loop_over_b.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+"""Loop-over-B GEMV for bf16 (non-quantized) matmul.
+
+Same pattern as custom_qmv_loop_over_b.py but for bf16 weights:
+no dequantization, no scales/biases — just bf16 weight reads.
+
+Y = X @ W^T where W is (N, K) bf16.
+
+TG: (32, 2, 1) = 64 threads = 2 SGs.
+Each SG: 4 output rows. B loop inside row loop.
+
+Usage:
+    from custom_bf16_qmv_loop_over_b import custom_bf16_qmv_loop_over_b
+    y = custom_bf16_qmv_loop_over_b(x, w, M=4, N=8192, K=2048)
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_bf16_qmv_source(M_val, N_val, K_val):
+    B = M_val
+    # bf16 weights: each element is 2 bytes, read as bfloat16_t
+    # No quantization groups — just w[row * K + col]
+    # Each thread handles 8 K-elements per block (VALUES_PER_THREAD=8)
+    # Block size = 32 threads × 8 = 256 elements per K-iteration
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int K = {K_val};
+    const int N = {N_val};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+    int tg = tgid.y;
+
+    int out_row = tg * 8 + sgid * RESULTS_PER_SG;
+    if (out_row >= N) return;
+
+    // Weight pointer: w is (N, K) bf16
+    const device bfloat16_t* ws = (const device bfloat16_t*)w + (long)out_row * K + slid * VALUES_PER_THREAD;
+
+    // Result accumulators: 4 rows × B batches
+    float result[{4 * B}];
+    for (int i = 0; i < {4 * B}; i++) result[i] = 0;
+
+    int x_base = slid * VALUES_PER_THREAD;
+
+    // K-loop: loop over B inside row loop
+    for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{
+
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            const device bfloat16_t* wl = ws + row * K;
+
+            for (int b = 0; b < {B}; b++) {{
+                float accum = 0;
+                for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+                    float xi = float(((const device bfloat16_t*)x)[b * K + x_base + i]);
+                    accum += xi * float(wl[i]);
+                }}
+                result[b * 4 + row] += accum;
+            }}
+        }}
+
+        ws += BLOCK_SIZE;
+        x_base += BLOCK_SIZE;
+    }}
+
+    // Reduction
+    for (int i = 0; i < {4 * B}; i++) result[i] = simd_sum(result[i]);
+
+    // Write output (bf16)
+    if (slid < 4u) {{
+        for (int b = 0; b < {B}; b++) {{
+            int r = out_row + (int)slid;
+            if (r < N) {{
+                y[b * N + r] = static_cast<bfloat16_t>(result[b * 4 + slid]);
+            }}
+        }}
+    }}
+"""
+
+
+_cache = {}
+
+
+def custom_bf16_qmv_loop_over_b(x, w, M, N, K):
+    key = (M, N, K)
+    if key not in _cache:
+        _cache[key] = mx.fast.metal_kernel(
+            name=f"custom_bf16_qmv_loop_b_M{M}_N{N}_K{K}",
+            input_names=["x", "w"],
+            output_names=["y"],
+            source=_gen_bf16_qmv_source(M, N, K),
+        )
+    kern = _cache[key]
+
+    n_tg = ceil_div(N, 8)
+    result = kern(
+        inputs=[x, w],
+        output_shapes=[(M * N,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(32, n_tg * 2, 1),
+        threadgroup=(32, 2, 1),
+    )
+    return result[0].reshape(M, N)

src/exo/worker/engines/mlx/matmul/kernels/bf16/gemv_loop_over_b_twice.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python3
+"""Two-pass LPB inside a single kernel for large M.
+
+Same threadgroup processes batches 0..M/2-1 first, then M/2..M-1.
+First pass results are written to threadgroup memory to free registers.
+Second pass reuses the same registers. Both results written to device
+memory at the end.
+
+Register pressure: R = 4*(M/2) + 5 per pass (not 2x).
+At M=16: R = 37 per pass instead of R = 69 single-pass.
+
+Usage:
+    from gemv_loop_over_b_twice import custom_bf16_qmv_loop_over_b_twice
+    y = custom_bf16_qmv_loop_over_b_twice(x, w, M=16, N=17408, K=5120)
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_bf16_qmv_twice_source(M_val, N_val, K_val):
+    # Pass 0 takes ceil(M/2) rows, pass 1 takes floor(M/2) rows.
+    # When M is odd, half_a > half_b by 1.
+    half_a = (M_val + 1) // 2
+    half_b = M_val - half_a
+    # Threadgroup memory sized for the larger pass (half_a)
+    tg_size = 2 * 4 * half_a
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int K = {K_val};
+    const int N = {N_val};
+    const int HALF_A = {half_a};
+    const int HALF_B = {half_b};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+    int tg = tgid.y;
+
+    int out_row = tg * 8 + sgid * RESULTS_PER_SG;
+    if (out_row >= N) return;
+
+    const device bfloat16_t* w_base = (const device bfloat16_t*)w + (long)out_row * K + slid * VALUES_PER_THREAD;
+
+    // Threadgroup memory to store pass 0 results (frees registers for pass 1)
+    threadgroup float tg_results[{tg_size}];
+    int tg_offset = sgid * {4 * half_a};  // each SG gets its own slice
+
+    // ══════ Pass 0: batches 0..HALF_A-1 ══════
+    {{
+        float result[{4 * half_a}];
+        for (int i = 0; i < {4 * half_a}; i++) result[i] = 0;
+
+        const device bfloat16_t* ws = w_base;
+        int x_base = slid * VALUES_PER_THREAD;
+
+        for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{
+            for (int row = 0; row < RESULTS_PER_SG; row++) {{
+                const device bfloat16_t* wl = ws + row * K;
+                for (int b = 0; b < {half_a}; b++) {{
+                    float accum = 0;
+                    for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+                        float xi = float(((const device bfloat16_t*)x)[b * K + x_base + i]);
+                        accum += xi * float(wl[i]);
+                    }}
+                    result[b * 4 + row] += accum;
+                }}
+            }}
+            ws += BLOCK_SIZE;
+            x_base += BLOCK_SIZE;
+        }}
+
+        // Reduce and store to threadgroup memory
+        for (int i = 0; i < {4 * half_a}; i++) {{
+            result[i] = simd_sum(result[i]);
+        }}
+        if (slid == 0u) {{
+            for (int i = 0; i < {4 * half_a}; i++) {{
+                tg_results[tg_offset + i] = result[i];
+            }}
+        }}
+    }}
+    // result[] goes out of scope here — registers freed
+""" + (f"""
+    // ══════ Pass 1: batches HALF_A..M-1 (HALF_B rows) ══════
+    {{
+        float result[{4 * half_b}];
+        for (int i = 0; i < {4 * half_b}; i++) result[i] = 0;
+
+        const device bfloat16_t* ws = w_base;
+        int x_base = slid * VALUES_PER_THREAD;
+
+        for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{
+            for (int row = 0; row < RESULTS_PER_SG; row++) {{
+                const device bfloat16_t* wl = ws + row * K;
+                for (int b = 0; b < {half_b}; b++) {{
+                    float accum = 0;
+                    for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+                        float xi = float(((const device bfloat16_t*)x)[(HALF_A + b) * K + x_base + i]);
+                        accum += xi * float(wl[i]);
+                    }}
+                    result[b * 4 + row] += accum;
+                }}
+            }}
+            ws += BLOCK_SIZE;
+            x_base += BLOCK_SIZE;
+        }}
+
+        // Reduce
+        for (int i = 0; i < {4 * half_b}; i++) {{
+            result[i] = simd_sum(result[i]);
+        }}
+
+        // Write pass 1 results to device memory
+        if (slid < 4u) {{
+            for (int b = 0; b < {half_b}; b++) {{
+                int r = out_row + (int)slid;
+                if (r < N) {{
+                    y[(HALF_A + b) * N + r] = static_cast<bfloat16_t>(result[b * 4 + slid]);
+                }}
+            }}
+        }}
+    }}
+""" if half_b > 0 else "") + f"""
+
+    // Write pass 0 results from threadgroup memory to device memory
+    if (slid < 4u) {{
+        for (int b = 0; b < {half_a}; b++) {{
+            int r = out_row + (int)slid;
+            if (r < N) {{
+                y[b * N + r] = static_cast<bfloat16_t>(tg_results[tg_offset + b * 4 + slid]);
+            }}
+        }}
+    }}
+"""
+
+
+_cache = {}
+
+
+def custom_bf16_qmv_loop_over_b_twice(x, w, M, N, K):
+    key = (M, N, K)
+    if key not in _cache:
+        _cache[key] = mx.fast.metal_kernel(
+            name=f"custom_bf16_qmv_loop_b_2x_M{M}_N{N}_K{K}",
+            input_names=["x", "w"],
+            output_names=["y"],
+            source=_gen_bf16_qmv_twice_source(M, N, K),
+        )
+    kern = _cache[key]
+
+    n_tg = ceil_div(N, 8)
+    result = kern(
+        inputs=[x, w],
+        output_shapes=[(M * N,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(32, n_tg * 2, 1),
+        threadgroup=(32, 2, 1),
+    )
+    return result[0].reshape(M, N)

src/exo/worker/engines/mlx/matmul/kernels/quantized/qmm_bm16.py

@@ -0,0 +1,239 @@
+#!/usr/bin/env python3
+"""Custom quantized GEMM kernel with BM=16 (vs MLX's hardcoded BM=32).
+
+Replicates MLX's affine_qmm_t exactly but with BM=16 to eliminate
+50% compute waste when M=16. Fixed parameters: 8-bit, gs=64, bfloat16.
+
+BM=16, BN=32, BK=32, WM=2, WN=2 → 4 SGs = 128 threads.
+TM = BM/(8*WM) = 16/16 = 1, TN = BN/(8*WN) = 32/16 = 2.
+Each SG: 1×2 output fragments (8×16).
+
+Uses simdgroup_matrix<float, 8, 8> for hardware MMA.
+Inlines QuantizedBlockLoader and BlockLoader for 8-bit gs=64.
+
+Usage:
+    from custom_qmm_bm16 import custom_qmm_t_bm16
+    y = custom_qmm_t_bm16(x, w, scales, biases, M=16, N=8192, K=2048)
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_custom_qmm_source(M_val, N_val, K_val, group_size=64):
+    gs = group_size
+    BM = 16
+    BN = 32
+    BK = 32
+    WM = 2
+    WN = 2
+    SIMD_SIZE = 32
+    tgp_size = WM * WN * SIMD_SIZE  # 128
+
+    BK_PAD = BK + 16 // 2  # +8 for bfloat16 (16/sizeof(bf16)=8)
+
+    # QuantizedBlockLoader params for (BROWS=BN=32, BCOLS=BK=32, tgp_size=128, bits=8)
+    # pack_factor=1 for 8-bit, BCOLS_PACKED=32
+    # n_reads = (32*32)/128 = 8
+    w_n_reads = (BN * BK) // tgp_size  # 8
+
+    # BlockLoader params for input X (BROWS=BM=16, BCOLS=BK=32, tgp_size=128)
+    # n_reads = (16*32)/128 = 4
+    # TCOLS = 32/4 = 8, TROWS = 128/8 = 16
+    x_n_reads = (BM * BK) // tgp_size  # 4
+    x_TCOLS = BK // x_n_reads  # 8
+    x_TROWS = tgp_size // x_TCOLS  # 16
+    x_n_rows = ceil_div(BM, x_TROWS)  # 1
+
+    K_groups = K_val // gs
+
+    return f"""
+    // ═══ Constants ═══
+    const int BM = {BM};
+    const int BN = {BN};
+    const int BK = {BK};
+    const int BK_PAD = {BK_PAD};
+    const int K = {K_val};
+    const int N = {N_val};
+    const int M = {M_val};
+    const int K_groups = {K_groups};
+    const int GROUP_SIZE = {gs};
+
+    uint3 tid = threadgroup_position_in_grid;
+    uint lid = thread_index_in_threadgroup;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+    int thread_idx = sgid * 32 + slid;
+
+    int y_row = tid.y * BM;  // M-dimension tile start
+    int y_col = tid.x * BN;  // N-dimension tile start
+
+    // ═══ TG memory ═══
+    threadgroup float Xs[{BM} * {BK_PAD}];
+    threadgroup float Ws[{BN} * {BK_PAD}];
+
+    // ═══ Pointer setup (follows qmm_t_impl lines 1137-1148) ═══
+    const device bfloat16_t* x_base = (const device bfloat16_t*)x + y_row * K;
+    const device uint8_t* w_base = (const device uint8_t*)w + y_col * K;
+    const device bfloat16_t* s_base = (const device bfloat16_t*)scales + y_col * K_groups;
+    const device bfloat16_t* b_base = (const device bfloat16_t*)biases + y_col * K_groups;
+
+    // ═══ BlockLoader for X: (BROWS=16, BCOLS=32, tgp_size=128) ═══
+    // n_reads=4, TCOLS=8, TROWS=16, n_rows=1
+    int x_bi = thread_idx / {x_TCOLS};           // row in tile (0..15)
+    int x_bj = {x_n_reads} * (thread_idx % {x_TCOLS});  // col start in tile
+    const device bfloat16_t* x_src = x_base + x_bi * K + x_bj;
+    threadgroup float* x_dst = Xs + x_bi * BK_PAD + x_bj;
+
+    // ═══ QuantizedBlockLoader for W: (BROWS=32, BCOLS=32, tgp_size=128, bits=8) ═══
+    // n_reads=8, pack_factor=1
+    int w_bi = {w_n_reads} * thread_idx / {BK};  // row in tile (0..31)
+    int w_bj = ({w_n_reads} * thread_idx) % {BK}; // col start in tile
+    const device uint8_t* w_src = w_base + w_bi * K + w_bj;
+    const device bfloat16_t* w_scales = s_base + w_bi * K_groups;
+    const device bfloat16_t* w_biases = b_base + w_bi * K_groups;
+    threadgroup float* w_dst = Ws + w_bi * BK_PAD + w_bj;
+    short w_group_step_cnt = 0;
+    const int w_group_steps = GROUP_SIZE / BK;  // 64/32 = 2
+
+    // ═══ MMA setup (follows BlockMMA constructor, lines 488-505) ═══
+    // WM=2, WN=2: sgid layout: sgid/WN=row, sgid%WN=col
+    // TM=1, TN=2: each SG has 1×2 output fragments
+    short sg_row = sgid / {WN};   // 0 or 1
+    short sg_col = sgid % {WN};   // 0 or 1
+    short tm = 8 * sg_row;        // M offset within BM (0 or 8)
+    short tn = 8 * sg_col;        // N offset within BN (0 or 8)
+    // But TN=2, so each SG handles 2 N-fragments: tn and tn+16
+    // Actually: with WN=2, sg_col is 0 or 1, tn = 8*sg_col = 0 or 8
+    // But we need to cover BN=32 with WN=2 → TN=2 fragments per SG
+    // So SG col 0 handles N cols [0..7] and [16..23], SG col 1 handles [8..15] and [24..31]
+    // Wait, that's the serpentine ordering. Let me follow MLX exactly.
+    // sm/sn are the thread's position within the SG's output tile
+    // get_coord returns (col, row) within the 8x8 fragment
+
+    // Accumulators: TM=1 × TN=2 = 2 fragments per SG
+    simdgroup_matrix<float, 8, 8> C00 = simdgroup_matrix<float, 8, 8>(0);
+    simdgroup_matrix<float, 8, 8> C01 = simdgroup_matrix<float, 8, 8>(0);
+
+    // Offsets into TG memory for this SG's MMA reads
+    short As_offset_m = tm;  // row offset in Xs
+    short Bs_offset_n0 = tn;          // first N-fragment
+    short Bs_offset_n1 = tn + {WN} * 8; // second N-fragment (offset by WN*8=16)
+
+    // ═══ K-loop ═══
+    for (int k = 0; k < K; k += BK) {{
+
+        // ── Load X tile: BM×BK (bf16 → float) ──
+        // BlockLoader pattern: each thread loads n_reads={x_n_reads} elements
+        if (x_bi < BM) {{
+            for (int i = 0; i < {x_n_reads}; i++) {{
+                x_dst[i] = float(x_src[i]);
+            }}
+        }}
+
+        // ── Load + dequantize W tile: BN×BK (uint8 → float) ──
+        // QuantizedBlockLoader pattern: each thread loads n_reads={w_n_reads} elements
+        if (w_bi < BN) {{
+            float scale = float(*w_scales);
+            float bias = float(*w_biases);
+            for (int i = 0; i < {w_n_reads}; i++) {{
+                w_dst[i] = scale * float(w_src[i]) + bias;
+            }}
+        }}
+
+        threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+        // ── MMA inner loop: BK/8 = 4 iterations ──
+        for (short kk = 0; kk < BK; kk += 8) {{
+            simdgroup_matrix<float, 8, 8> A_frag;
+            simdgroup_matrix<float, 8, 8> B_frag0, B_frag1;
+
+            // Load A fragment: 8×8 from Xs at (As_offset_m, kk)
+            simdgroup_load(A_frag, &Xs[As_offset_m * BK_PAD + kk], BK_PAD);
+
+            // Load B fragments: transposed from Ws (stored as BN × BK_PAD)
+            // B_frag0: output cols [Bs_offset_n0..+8], K rows [kk..+8]
+            // B_frag1: output cols [Bs_offset_n1..+8], K rows [kk..+8]
+            simdgroup_load(B_frag0, &Ws[Bs_offset_n0 * BK_PAD + kk], BK_PAD, ulong2(0, 0), true);
+            simdgroup_load(B_frag1, &Ws[Bs_offset_n1 * BK_PAD + kk], BK_PAD, ulong2(0, 0), true);
+
+            simdgroup_multiply_accumulate(C00, A_frag, B_frag0, C00);
+            simdgroup_multiply_accumulate(C01, A_frag, B_frag1, C01);
+        }}
+
+        threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+        // ── Advance loaders ──
+        x_src += BK;
+        w_src += BK;
+        w_group_step_cnt++;
+        if (w_group_step_cnt == w_group_steps) {{
+            w_group_step_cnt = 0;
+            w_scales++;
+            w_biases++;
+        }}
+    }}
+
+    // ═══ Store results to device memory ═══
+    // Each SG stores its 1×2 fragments (two 8×8 blocks)
+    // Output layout: y[m * N + n] for m in [y_row..y_row+BM), n in [y_col..y_col+BN)
+    device bfloat16_t* y_ptr = (device bfloat16_t*)y + y_row * N + y_col;
+
+    // Store C00 at (tm, tn) and C01 at (tm, tn + WN*8)
+    simdgroup_store(C00, &Ws[tm * BN + tn], BN);
+    simdgroup_store(C01, &Ws[tm * BN + tn + {WN} * 8], BN);
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+    // Convert float → bf16 and write to device
+    for (int i = thread_idx; i < BM * BN; i += {tgp_size}) {{
+        int r = i / BN;
+        int c = i % BN;
+        if (y_row + r < M && y_col + c < N) {{
+            y_ptr[r * N + c] = static_cast<bfloat16_t>(Ws[r * BN + c]);
+        }}
+    }}
+"""
+
+
+_custom_qmm_cache = {}
+
+
+def custom_qmm_t_bm16(x, w, scales, biases, M, N, K, group_size=64):
+    """Custom quantized GEMM with BM=16.
+
+    Args:
+        x: (M, K) bfloat16 input
+        w: (N, K/4) uint32 packed 8-bit weights
+        scales: (N, K/gs) bfloat16
+        biases: (N, K/gs) bfloat16
+        M, N, K: matrix dimensions
+    Returns:
+        y: (M, N) bfloat16
+    """
+    key = (M, N, K, group_size)
+    if key not in _custom_qmm_cache:
+        _custom_qmm_cache[key] = mx.fast.metal_kernel(
+            name=f"custom_qmm_t_bm16_M{M}_N{N}_K{K}",
+            input_names=["x", "w", "scales", "biases"],
+            output_names=["y"],
+            source=_gen_custom_qmm_source(M, N, K, group_size),
+        )
+    kern = _custom_qmm_cache[key]
+
+    BN = 32
+    BM = 16
+    n_tg_n = ceil_div(N, BN)
+    n_tg_m = ceil_div(M, BM)
+
+    result = kern(
+        inputs=[x, w, scales, biases],
+        output_shapes=[(M * N,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(32 * n_tg_n, 4 * n_tg_m, 1),
+        threadgroup=(32, 4, 1),
+    )
+
+    return result[0].reshape(M, N)

src/exo/worker/engines/mlx/matmul/kernels/quantized/qmm_bm8.py

@@ -0,0 +1,232 @@
+#!/usr/bin/env python3
+"""Custom quantized GEMM kernel with BM=8 (vs MLX's GEMV at BS=8).
+
+At BS=8, MLX uses affine_qmv_fast (GEMV) which loads weights 8 times
+via grid x. Our GEMM loads weights once via TG memory tiling.
+
+BM=8, BN=32, BK=32, WM=1, WN=2 → 2 SGs = 64 threads.
+TM = BM/(8*WM) = 8/8 = 1, TN = BN/(8*WN) = 32/16 = 2.
+Each SG: 1×2 output fragments (8×16).
+
+Uses simdgroup_matrix<float, 8, 8> for hardware MMA.
+Inlines QuantizedBlockLoader and BlockLoader for 8-bit gs=64.
+
+Usage:
+    from custom_qmm_bm8 import custom_qmm_t_bm8
+    y = custom_qmm_t_bm8(x, w, scales, biases, M=8, N=8192, K=2048)
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_custom_qmm_source(M_val, N_val, K_val, group_size=64):
+    gs = group_size
+    BM = 8
+    BN = 32
+    BK = 32
+    WM = 1
+    WN = 2
+    SIMD_SIZE = 32
+    tgp_size = WM * WN * SIMD_SIZE  # 64
+
+    BK_PAD = BK + 16 // 2  # +8 for bfloat16 (16/sizeof(bf16)=8)
+
+    # QuantizedBlockLoader params for (BROWS=BN=32, BCOLS=BK=32, tgp_size=64, bits=8)
+    # pack_factor=1 for 8-bit, BCOLS_PACKED=32
+    # n_reads = (32*32)/64 = 16
+    w_n_reads = (BN * BK) // tgp_size  # 16
+
+    # BlockLoader params for input X (BROWS=BM=8, BCOLS=BK=32, tgp_size=64)
+    # n_reads = (8*32)/64 = 4
+    # TCOLS = 32/4 = 8, TROWS = 64/8 = 8
+    x_n_reads = (BM * BK) // tgp_size  # 4
+    x_TCOLS = BK // x_n_reads  # 8
+    x_TROWS = tgp_size // x_TCOLS  # 8
+    x_n_rows = ceil_div(BM, x_TROWS)  # 1
+
+    K_groups = K_val // gs
+
+    return f"""
+    // ═══ Constants ═══
+    const int BM = {BM};
+    const int BN = {BN};
+    const int BK = {BK};
+    const int BK_PAD = {BK_PAD};
+    const int K = {K_val};
+    const int N = {N_val};
+    const int M = {M_val};
+    const int K_groups = {K_groups};
+    const int GROUP_SIZE = {gs};
+
+    uint3 tid = threadgroup_position_in_grid;
+    uint lid = thread_index_in_threadgroup;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+    int thread_idx = sgid * 32 + slid;
+
+    int y_row = tid.y * BM;  // M-dimension tile start
+    int y_col = tid.x * BN;  // N-dimension tile start
+
+    // ═══ TG memory ═══
+    threadgroup float Xs[{BM} * {BK_PAD}];
+    threadgroup float Ws[{BN} * {BK_PAD}];
+
+    // ═══ Pointer setup (follows qmm_t_impl lines 1137-1148) ═══
+    const device bfloat16_t* x_base = (const device bfloat16_t*)x + y_row * K;
+    const device uint8_t* w_base = (const device uint8_t*)w + y_col * K;
+    const device bfloat16_t* s_base = (const device bfloat16_t*)scales + y_col * K_groups;
+    const device bfloat16_t* b_base = (const device bfloat16_t*)biases + y_col * K_groups;
+
+    // ═══ BlockLoader for X: (BROWS=8, BCOLS=32, tgp_size=64) ═══
+    // n_reads=4, TCOLS=8, TROWS=8, n_rows=1
+    int x_bi = thread_idx / {x_TCOLS};           // row in tile (0..15)
+    int x_bj = {x_n_reads} * (thread_idx % {x_TCOLS});  // col start in tile
+    const device bfloat16_t* x_src = x_base + x_bi * K + x_bj;
+    threadgroup float* x_dst = Xs + x_bi * BK_PAD + x_bj;
+
+    // ═══ QuantizedBlockLoader for W: (BROWS=32, BCOLS=32, tgp_size=64, bits=8) ═══
+    // n_reads=16, pack_factor=1
+    int w_bi = {w_n_reads} * thread_idx / {BK};  // row in tile (0..31)
+    int w_bj = ({w_n_reads} * thread_idx) % {BK}; // col start in tile
+    const device uint8_t* w_src = w_base + w_bi * K + w_bj;
+    const device bfloat16_t* w_scales = s_base + w_bi * K_groups;
+    const device bfloat16_t* w_biases = b_base + w_bi * K_groups;
+    threadgroup float* w_dst = Ws + w_bi * BK_PAD + w_bj;
+    short w_group_step_cnt = 0;
+    const int w_group_steps = GROUP_SIZE / BK;  // 64/32 = 2
+
+    // ═══ MMA setup (follows BlockMMA constructor, lines 488-505) ═══
+    // WM=1, WN=2: sgid layout: sg_row=0 always, sg_col=sgid (0 or 1)
+    // TM=1, TN=2: each SG has 1×2 output fragments
+    short sg_row = 0;
+    short sg_col = sgid;             // 0 or 1
+    short tm = 0;                    // M offset (always 0, only 8 rows = 1 fragment)
+    short tn = 8 * sg_col;          // N offset: 0 or 8
+
+    // Accumulators: TM=1 × TN=2 = 2 fragments per SG
+    simdgroup_matrix<float, 8, 8> C00 = simdgroup_matrix<float, 8, 8>(0);
+    simdgroup_matrix<float, 8, 8> C01 = simdgroup_matrix<float, 8, 8>(0);
+
+    // Offsets into TG memory for this SG's MMA reads
+    short As_offset_m = tm;  // row offset in Xs
+    short Bs_offset_n0 = tn;          // first N-fragment
+    short Bs_offset_n1 = tn + {WN} * 8; // second N-fragment (offset by WN*8=16)
+
+    // ═══ K-loop ═══
+    for (int k = 0; k < K; k += BK) {{
+
+        // ── Load X tile: BM×BK (bf16 → float) ──
+        // BlockLoader pattern: each thread loads n_reads={x_n_reads} elements
+        if (x_bi < BM) {{
+            for (int i = 0; i < {x_n_reads}; i++) {{
+                x_dst[i] = float(x_src[i]);
+            }}
+        }}
+
+        // ── Load + dequantize W tile: BN×BK (uint8 → float) ──
+        // QuantizedBlockLoader pattern: each thread loads n_reads={w_n_reads} elements
+        if (w_bi < BN) {{
+            float scale = float(*w_scales);
+            float bias = float(*w_biases);
+            for (int i = 0; i < {w_n_reads}; i++) {{
+                w_dst[i] = scale * float(w_src[i]) + bias;
+            }}
+        }}
+
+        threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+        // ── MMA inner loop: BK/8 = 4 iterations ──
+        for (short kk = 0; kk < BK; kk += 8) {{
+            simdgroup_matrix<float, 8, 8> A_frag;
+            simdgroup_matrix<float, 8, 8> B_frag0, B_frag1;
+
+            // Load A fragment: 8×8 from Xs at (As_offset_m, kk)
+            simdgroup_load(A_frag, &Xs[As_offset_m * BK_PAD + kk], BK_PAD);
+
+            // Load B fragments: transposed from Ws (stored as BN × BK_PAD)
+            // B_frag0: output cols [Bs_offset_n0..+8], K rows [kk..+8]
+            // B_frag1: output cols [Bs_offset_n1..+8], K rows [kk..+8]
+            simdgroup_load(B_frag0, &Ws[Bs_offset_n0 * BK_PAD + kk], BK_PAD, ulong2(0, 0), true);
+            simdgroup_load(B_frag1, &Ws[Bs_offset_n1 * BK_PAD + kk], BK_PAD, ulong2(0, 0), true);
+
+            simdgroup_multiply_accumulate(C00, A_frag, B_frag0, C00);
+            simdgroup_multiply_accumulate(C01, A_frag, B_frag1, C01);
+        }}
+
+        threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+        // ── Advance loaders ──
+        x_src += BK;
+        w_src += BK;
+        w_group_step_cnt++;
+        if (w_group_step_cnt == w_group_steps) {{
+            w_group_step_cnt = 0;
+            w_scales++;
+            w_biases++;
+        }}
+    }}
+
+    // ═══ Store results to device memory ═══
+    // Each SG stores its 1×2 fragments (two 8×8 blocks)
+    // Output layout: y[m * N + n] for m in [y_row..y_row+BM), n in [y_col..y_col+BN)
+    device bfloat16_t* y_ptr = (device bfloat16_t*)y + y_row * N + y_col;
+
+    // Store C00 at (tm, tn) and C01 at (tm, tn + WN*8)
+    simdgroup_store(C00, &Ws[tm * BN + tn], BN);
+    simdgroup_store(C01, &Ws[tm * BN + tn + {WN} * 8], BN);
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+    // Convert float → bf16 and write to device
+    for (int i = thread_idx; i < BM * BN; i += {tgp_size}) {{
+        int r = i / BN;
+        int c = i % BN;
+        if (y_row + r < M && y_col + c < N) {{
+            y_ptr[r * N + c] = static_cast<bfloat16_t>(Ws[r * BN + c]);
+        }}
+    }}
+"""
+
+
+_custom_qmm_cache = {}
+
+
+def custom_qmm_t_bm8(x, w, scales, biases, M, N, K, group_size=64):
+    """Custom quantized GEMM with BM=8.
+
+    Args:
+        x: (M, K) bfloat16 input
+        w: (N, K/4) uint32 packed 8-bit weights
+        scales: (N, K/gs) bfloat16
+        biases: (N, K/gs) bfloat16
+        M, N, K: matrix dimensions
+    Returns:
+        y: (M, N) bfloat16
+    """
+    key = (M, N, K, group_size)
+    if key not in _custom_qmm_cache:
+        _custom_qmm_cache[key] = mx.fast.metal_kernel(
+            name=f"custom_qmm_t_bm8_M{M}_N{N}_K{K}",
+            input_names=["x", "w", "scales", "biases"],
+            output_names=["y"],
+            source=_gen_custom_qmm_source(M, N, K, group_size),
+        )
+    kern = _custom_qmm_cache[key]
+
+    BN = 32
+    BM = 8
+    n_tg_n = ceil_div(N, BN)
+    n_tg_m = ceil_div(M, BM)
+
+    result = kern(
+        inputs=[x, w, scales, biases],
+        output_shapes=[(M * N,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(32 * n_tg_n, 2 * n_tg_m, 1),
+        threadgroup=(32, 2, 1),
+    )
+
+    return result[0].reshape(M, N)

src/exo/worker/engines/mlx/matmul/kernels/quantized/qmm_splitk.py

@@ -0,0 +1,316 @@
+#!/usr/bin/env python3
+"""Split-K int8 quantized GEMM (gs=64, bits=8).
+
+Int8 analogue of matmul/kernels/bf16/gemm_splitk.py: partitions K across
+P threadgroups per output tile, each writing a partial fp32 result to a
+scratch buffer; a second dispatch sums the partials and casts to bf16.
+
+Key modification from the bf16 version: the W loader dequantizes 8-bit
+uint8 weights on the fly using per-group bf16 scales/biases, following
+the same pattern as matmul/kernels/quantized/qmm_bm16.py.
+
+BK is chosen equal to GROUP_SIZE (64) so each K-iteration advances the
+scales/biases pointers by exactly one group, eliminating the need for an
+inner group_step counter and keeping partition starts aligned to group
+boundaries.
+
+Usage:
+    from qmm_splitk import custom_qmm_splitk
+    y = custom_qmm_splitk(x, w, scales, biases, M=16, N=8192, K=5120, BM=16)
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def next_power_of_2(n):
+    if n <= 0:
+        return 1
+    return 1 << (n - 1).bit_length()
+
+
+def compute_partitions(M, N, K):
+    """Same formula as gemm_splitk_steel.compute_partitions (bf16)."""
+    _tm = ceil_div(M, 32)
+    _tn = ceil_div(N, 32)
+    _tk = K // 16
+    tmtn = max(_tm * _tn, 1)
+    return min(max(2, next_power_of_2(_tk // tmtn)), 32)
+
+
+def _gen_splitk_qmm_source(M_val, N_val, K_val, BM, split_k_partitions, group_size=64):
+    """Int8 split-K GEMM source. BK is fixed to group_size so scales/biases
+    advance by exactly one group per K-iteration (no group_step counter).
+    """
+    gs = group_size
+    BN = 32
+    BK = gs  # = 64
+    if BM <= 8:
+        WM, WN = 1, 2
+    else:
+        WM, WN = 2, 2
+    tgp_size = WM * WN * 32  # 64 or 128
+    BK_PAD = BK + 8  # bank-conflict padding (matches Steel's tgp_padding for bf16)
+
+    gemm_k_iterations = (K_val // BK) // split_k_partitions
+    split_k_partition_size = gemm_k_iterations * BK
+
+    # Loader params
+    # X (bf16, BROWS=BM, BCOLS=BK): each thread reads x_n_reads values
+    x_total = BM * BK
+    x_n_reads = max(1, x_total // tgp_size)
+    x_TCOLS = BK // x_n_reads if x_n_reads > 0 else BK
+
+    # W (uint8, BROWS=BN, BCOLS=BK): each thread reads w_n_reads bytes
+    w_n_reads = (BN * BK) // tgp_size
+
+    TM = BM // (8 * WM)
+    TN = BN // (8 * WN)
+
+    # Unrolled accumulator declarations: TM*TN simdgroup_matrix<float,8,8> per SG
+    c_decl_lines = []
+    for ti in range(TM):
+        for tj in range(TN):
+            c_decl_lines.append(
+                f"    simdgroup_matrix<float, 8, 8> C{ti}{tj} = simdgroup_matrix<float, 8, 8>(0);"
+            )
+    c_decl = "\n".join(c_decl_lines)
+
+    # Inner MMA step (inside the BK/8 loop over kk)
+    mma_lines = []
+    for ti in range(TM):
+        mma_lines.append(
+            f"            simdgroup_matrix<float, 8, 8> A{ti};\n"
+            f"            simdgroup_load(A{ti}, &Xs[(tm + {ti} * {WM} * 8) * BK_PAD + kk], BK_PAD);"
+        )
+    for tj in range(TN):
+        mma_lines.append(
+            f"            simdgroup_matrix<float, 8, 8> B{tj};\n"
+            f"            simdgroup_load(B{tj}, &Ws[(tn + {tj} * {WN} * 8) * BK_PAD + kk], BK_PAD, ulong2(0, 0), true);"
+        )
+    for ti in range(TM):
+        for tj in range(TN):
+            mma_lines.append(
+                f"            simdgroup_multiply_accumulate(C{ti}{tj}, A{ti}, B{tj}, C{ti}{tj});"
+            )
+    mma_load = "\n".join(mma_lines)
+
+    # Final accumulator store into Ws (then copied to device y as partial result)
+    store_lines = []
+    for ti in range(TM):
+        for tj in range(TN):
+            store_lines.append(
+                f"    simdgroup_store(C{ti}{tj}, &Ws[(tm + {ti} * {WM} * 8) * BN + tn + {tj} * {WN} * 8], BN);"
+            )
+    c_store = "\n".join(store_lines)
+
+    # Unrolled vectorized X load (per thread)
+    x_vec = "\n".join(
+        f"                x_dst[{i}] = float(x_src[{i}]);" for i in range(x_n_reads)
+    )
+    # Unrolled W dequantize (per thread)
+    w_vec = "\n".join(
+        f"                w_dst[{i}] = scale * float(w_src[{i}]) + bias;" for i in range(w_n_reads)
+    )
+
+    return f"""
+    const int BM = {BM};
+    const int BN = {BN};
+    const int BK = {BK};
+    const int BK_PAD = {BK_PAD};
+    const int K = {K_val};
+    const int N = {N_val};
+    const int M = {M_val};
+    const int GROUP_SIZE = {gs};
+    const int K_groups = {K_val // gs};
+    const int PARTITION_SIZE = {split_k_partition_size};
+    const int GEMM_K_ITERS = {gemm_k_iterations};
+    const int SPLIT_K_PARTS = {split_k_partitions};
+    const int PART_STRIDE = M * N;
+
+    uint3 tid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+    int thread_idx = sgid * 32 + slid;
+    int partition_idx = tid.z;
+
+    int y_row = tid.y * BM;
+    int y_col = tid.x * BN;
+
+    if (y_row >= M || y_col >= N) return;
+
+    threadgroup float Xs[{BM} * {BK_PAD}];
+    threadgroup float Ws[{BN} * {BK_PAD}];
+
+    int k_start = PARTITION_SIZE * partition_idx;
+    int k_group_start = k_start / GROUP_SIZE;
+
+    // Pointer setup. Scales/biases are stored as bf16 in real quantized
+    // models (the source Linear weight was bf16 before quantization), and the
+    // weight itself is a packed uint32 buffer — we reinterpret it as uint8*.
+    const device bfloat16_t* x_base = (const device bfloat16_t*)x + y_row * K + k_start;
+    const device uint8_t*    w_base = (const device uint8_t*)w    + y_col * K + k_start;
+    const device bfloat16_t* s_base = (const device bfloat16_t*)scales + y_col * K_groups + k_group_start;
+    const device bfloat16_t* b_base = (const device bfloat16_t*)biases + y_col * K_groups + k_group_start;
+
+    // BlockLoader X: (BROWS=BM, BCOLS=BK)
+    int x_bi = thread_idx / {x_TCOLS};
+    int x_bj = {x_n_reads} * (thread_idx % {x_TCOLS});
+    const device bfloat16_t* x_src = x_base + x_bi * K + x_bj;
+    threadgroup float* x_dst = Xs + x_bi * BK_PAD + x_bj;
+
+    // QuantizedBlockLoader W: (BROWS=BN, BCOLS=BK, bits=8, pack_factor=1)
+    int w_bi = ({w_n_reads} * thread_idx) / {BK};
+    int w_bj = ({w_n_reads} * thread_idx) % {BK};
+    const device uint8_t*    w_src    = w_base + w_bi * K + w_bj;
+    const device bfloat16_t* w_scales = s_base + w_bi * K_groups;
+    const device bfloat16_t* w_biases = b_base + w_bi * K_groups;
+    threadgroup float* w_dst = Ws + w_bi * BK_PAD + w_bj;
+
+    // MMA setup
+    short sg_row = sgid / {WN};
+    short sg_col = sgid % {WN};
+    short tm = 8 * sg_row;
+    short tn = 8 * sg_col;
+    {c_decl}
+
+    int k_iters = GEMM_K_ITERS;
+    if (partition_idx == SPLIT_K_PARTS - 1) {{
+        k_iters = (K - k_start) / BK;
+    }}
+
+    // K-loop: each iter advances by BK == GROUP_SIZE, so scales/biases
+    // advance by exactly one group per iteration.
+    for (int ki = 0; ki < k_iters; ki++) {{
+
+        // Load X tile (bf16 -> float)
+        if (x_bi < BM) {{
+{x_vec}
+        }}
+
+        // Load + dequantize W tile (uint8 -> float via scale * v + bias)
+        if (w_bi < BN) {{
+            float scale = float(*w_scales);
+            float bias  = float(*w_biases);
+{w_vec}
+        }}
+
+        threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+        // MMA: BK/8 simdgroup 8x8 accumulations
+        for (short kk = 0; kk < BK; kk += 8) {{
+{mma_load}
+        }}
+
+        threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+        x_src += BK;
+        w_src += BK;
+        w_scales += 1;
+        w_biases += 1;
+    }}
+
+    // Write partial fp32 result to y[partition_idx, y_row:y_row+BM, y_col:y_col+BN]
+    {c_store}
+    threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+
+    for (int i = thread_idx; i < BM * BN; i += {tgp_size}) {{
+        int r = i / BN;
+        int c = i % BN;
+        if (y_row + r < M && y_col + c < N) {{
+            y[partition_idx * PART_STRIDE + (y_row + r) * N + y_col + c] = Ws[r * BN + c];
+        }}
+    }}
+"""
+
+
+def _gen_accum_source(split_k_partitions):
+    """Identical to gemm_splitk._gen_accum_source — sum partitions, cast to bf16."""
+    return f"""
+    uint gid_x = thread_position_in_grid.x;
+    uint gid_y = thread_position_in_grid.y;
+    uint N_val = threads_per_grid.x;
+
+    int offset = gid_y * N_val + gid_x;
+    int stride = threads_per_grid.x * threads_per_grid.y;
+
+    float out = 0.0f;
+    for (int p = 0; p < {split_k_partitions}; p++) {{
+        out += ((const device float*)c_split)[offset + p * stride];
+    }}
+
+    y[gid_y * N_val + gid_x] = static_cast<bfloat16_t>(out);
+"""
+
+
+_gemm_cache = {}
+_accum_cache = {}
+
+
+def custom_qmm_splitk(x, w, scales, biases, M, N, K, BM=16, group_size=64):
+    """Split-K int8 quantized GEMM.
+
+    Args:
+        x: (M, K) bfloat16 input
+        w: (N, K/4) uint32 packed int8 weights (pack_factor=1 for bits=8)
+        scales: (N, K/gs) bfloat16
+        biases: (N, K/gs) bfloat16
+        M, N, K: matrix dimensions
+        BM: 8, 16, or 32
+        group_size: quantization group size (must be 64 for now — BK is fixed to gs)
+    Returns:
+        y: (M, N) bfloat16
+    """
+    assert group_size == 64, "qmm_splitk currently fixes BK=GROUP_SIZE=64"
+    BN = 32
+    BK = group_size
+    if BM <= 8:
+        WM, WN = 1, 2
+    else:
+        WM, WN = 2, 2
+    P = compute_partitions(M, N, K)
+
+    gemm_key = (M, N, K, BM, P, group_size)
+    if gemm_key not in _gemm_cache:
+        _gemm_cache[gemm_key] = mx.fast.metal_kernel(
+            name=f"qmm_splitk_bm{BM}_M{M}_N{N}_K{K}_P{P}",
+            input_names=["x", "w", "scales", "biases"],
+            output_names=["y"],
+            source=_gen_splitk_qmm_source(M, N, K, BM, P, group_size),
+        )
+    gemm_kern = _gemm_cache[gemm_key]
+
+    sgs = WM * WN
+    n_tg_n = ceil_div(N, BN)
+    n_tg_m = ceil_div(M, BM)
+
+    c_split = gemm_kern(
+        inputs=[x, w, scales, biases],
+        output_shapes=[(P * M * N,)],
+        output_dtypes=[mx.float32],
+        grid=(32 * n_tg_n, sgs * n_tg_m, P),
+        threadgroup=(32, sgs, 1),
+    )[0]
+
+    accum_key = P
+    if accum_key not in _accum_cache:
+        _accum_cache[accum_key] = mx.fast.metal_kernel(
+            name=f"qmm_splitk_accum_P{P}",
+            input_names=["c_split"],
+            output_names=["y"],
+            source=_gen_accum_source(P),
+        )
+    accum_kern = _accum_cache[accum_key]
+
+    y = accum_kern(
+        inputs=[c_split],
+        output_shapes=[(M * N,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(N, M, 1),
+        threadgroup=(min(N, 256), 1, 1),
+    )[0]
+
+    return y.reshape(M, N)

src/exo/worker/engines/mlx/matmul/kernels/quantized/qmv_loop_over_b.py

@@ -0,0 +1,135 @@
+#!/usr/bin/env python3
+"""Isolated loop-over-B GEMV kernel for quantized matmul.
+
+Extracts the loop-over-B pattern from batched_fused_gdn_projections_8bit
+but without any epilogues — pure Y = X @ dequant(W)^T output.
+
+For comparing our GEMV approach against MLX's affine_qmv_fast on
+an isolated QuantizedLinear operation (e.g., in_proj_qkv: N=8192, K=2048).
+
+TG: (32, 2, 1) = 64 threads = 2 SGs.
+Each SG: 4 output rows.
+B loop inside row loop for low register pressure (R = 4B + 5).
+
+Usage:
+    from custom_qmv_loop_over_b import custom_qmv_loop_over_b
+    y = custom_qmv_loop_over_b(x, w, scales, biases, M=8, N=8192, K=2048)
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_custom_qmv_source(M_val, N_val, K_val, group_size=64):
+    gs = group_size
+    sc_stride = 256 // gs
+    slid_div = gs // 8
+    K_groups = K_val // gs
+    B = M_val  # batch size = M
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int K = {K_val};
+    const int N = {N_val};
+    const int M = {M_val};
+    const int K_groups = {K_groups};
+    const int SC_STRIDE = {sc_stride};
+    const int SLID_DIV = {slid_div};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+    int tg = tgid.y;
+
+    int out_row = tg * 8 + sgid * RESULTS_PER_SG;
+    if (out_row >= N) return;
+
+    // Weight pointers
+    const device uint8_t* ws = (const device uint8_t*)w + (long)out_row * K + slid * VALUES_PER_THREAD;
+    const device bfloat16_t* sc = (const device bfloat16_t*)scales + (long)out_row * K_groups + slid / SLID_DIV;
+    const device bfloat16_t* bi = (const device bfloat16_t*)biases + (long)out_row * K_groups + slid / SLID_DIV;
+
+    // Result accumulators: 4 rows × B batches
+    float result[{4 * B}];
+    for (int i = 0; i < {4 * B}; i++) result[i] = 0;
+
+    int x_base = slid * VALUES_PER_THREAD;
+
+    // K-loop: loop over B inside row loop
+    for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{
+
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            const device uint8_t* wl = ws + row * K;
+            float s_val = float(sc[row * K_groups]);
+            float b_val = float(bi[row * K_groups]);
+
+            for (int b = 0; b < {B}; b++) {{
+                float accum = 0, xsum = 0;
+                for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+                    float xi = float(((const device bfloat16_t*)x)[b * K + x_base + i]);
+                    accum += xi * float(wl[i]);
+                    xsum += xi;
+                }}
+                result[b * 4 + row] += s_val * accum + xsum * b_val;
+            }}
+        }}
+
+        ws += BLOCK_SIZE; sc += SC_STRIDE; bi += SC_STRIDE; x_base += BLOCK_SIZE;
+    }}
+
+    // Reduction
+    for (int i = 0; i < {4 * B}; i++) result[i] = simd_sum(result[i]);
+
+    // Write output (bf16)
+    if (slid < 4u) {{
+        for (int b = 0; b < {B}; b++) {{
+            int r = out_row + (int)slid;
+            if (r < N) {{
+                y[b * N + r] = static_cast<bfloat16_t>(result[b * 4 + slid]);
+            }}
+        }}
+    }}
+"""
+
+
+_custom_qmv_cache = {}
+
+
+def custom_qmv_loop_over_b(x, w, scales, biases, M, N, K, group_size=64):
+    """Loop-over-B GEMV for quantized matmul.
+
+    Args:
+        x: (M, K) bfloat16 input
+        w: (N, K/4) uint32 packed 8-bit weights
+        scales: (N, K/gs) bfloat16
+        biases: (N, K/gs) bfloat16
+        M, N, K: dimensions
+    Returns:
+        y: (M, N) bfloat16
+    """
+    key = (M, N, K, group_size)
+    if key not in _custom_qmv_cache:
+        _custom_qmv_cache[key] = mx.fast.metal_kernel(
+            name=f"custom_qmv_loop_b_M{M}_N{N}_K{K}",
+            input_names=["x", "w", "scales", "biases"],
+            output_names=["y"],
+            source=_gen_custom_qmv_source(M, N, K, group_size),
+        )
+    kern = _custom_qmv_cache[key]
+
+    n_tg = ceil_div(N, 8)
+
+    result = kern(
+        inputs=[x, w, scales, biases],
+        output_shapes=[(M * N,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(32, n_tg * 2, 1),
+        threadgroup=(32, 2, 1),
+    )
+
+    return result[0].reshape(M, N)

src/exo/worker/engines/mlx/matmul/kernels/steel_include/steel/defines.h

@@ -0,0 +1,7 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#define STEEL_CONST static constant constexpr const
+#define STEEL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)")
+#define STEEL_PRAGMA_NO_UNROLL _Pragma("clang loop unroll(disable)")

src/exo/worker/engines/mlx/matmul/kernels/steel_include/steel/gemm/gemm.h

@@ -0,0 +1,295 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/gemm/loader.h"
+#include "mlx/backend/metal/kernels/steel/gemm/mma.h"
+#include "mlx/backend/metal/kernels/steel/gemm/params.h"
+#include "mlx/backend/metal/kernels/steel/gemm/transforms.h"
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+using namespace metal;
+
+///////////////////////////////////////////////////////////////////////////////
+// GEMM kernel class
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <bool M_aligned, bool N_aligned, bool K_aligned>
+struct LoopAlignment {};
+
+template <
+    typename T,
+    typename U,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    bool MN_aligned,
+    bool K_aligned,
+    typename AccumType = typename AccumHelper<T>::accum_type,
+    typename Epilogue = TransformNone<U, AccumType>>
+struct GEMMKernel {
+  STEEL_CONST short tgp_padding_a = 16 / sizeof(T);
+  STEEL_CONST short tgp_padding_b = 16 / sizeof(T);
+  STEEL_CONST short tgp_mem_size_a =
+      transpose_a ? BK * (BM + tgp_padding_a) : BM * (BK + tgp_padding_a);
+  STEEL_CONST short tgp_mem_size_b =
+      transpose_b ? BN * (BK + tgp_padding_b) : BK * (BN + tgp_padding_b);
+  STEEL_CONST short tgp_mem_size = tgp_mem_size_a + tgp_mem_size_b;
+
+  STEEL_CONST short tgp_size = WM * WN * 32;
+
+  using loader_a_t = BlockLoader<
+      T,
+      transpose_a ? BK : BM,
+      transpose_a ? BM : BK,
+      transpose_a ? BM + tgp_padding_a : BK + tgp_padding_a,
+      !transpose_a,
+      tgp_size>;
+  using loader_b_t = BlockLoader<
+      T,
+      transpose_b ? BN : BK,
+      transpose_b ? BK : BN,
+      transpose_b ? BK + tgp_padding_b : BN + tgp_padding_b,
+      transpose_b,
+      tgp_size>;
+  using mma_t = BlockMMA<
+      T,
+      U,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      transpose_a ? BM + tgp_padding_a : BK + tgp_padding_a,
+      transpose_b ? BK + tgp_padding_b : BN + tgp_padding_b,
+      AccumType,
+      Epilogue>;
+
+  /* Main kernel function */
+  template <bool M_aligned, bool N_aligned, bool K_aligned_>
+  static METAL_FUNC void gemm_loop(
+      threadgroup T* As [[threadgroup(0)]],
+      threadgroup T* Bs [[threadgroup(1)]],
+      const int gemm_k_iterations,
+      thread loader_a_t& loader_a,
+      thread loader_b_t& loader_b,
+      thread mma_t& mma_op,
+      thread const short& tgp_bm,
+      thread const short& tgp_bn,
+      thread const short& lbk,
+      LoopAlignment<M_aligned, N_aligned, K_aligned_> l = {}) {
+    // Appease the compiler
+    (void)l;
+
+    short2 tile_dims_A = transpose_a ? short2(tgp_bm, BK) : short2(BK, tgp_bm);
+
+    short2 tile_dims_B = transpose_b ? short2(BK, tgp_bn) : short2(tgp_bn, BK);
+
+    for (int k = 0; k < gemm_k_iterations; k++) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+      // Load elements into threadgroup
+      if (M_aligned) {
+        loader_a.load_unsafe();
+      } else {
+        loader_a.load_safe(tile_dims_A);
+      }
+
+      if (N_aligned) {
+        loader_b.load_unsafe();
+      } else {
+        loader_b.load_safe(tile_dims_B);
+      }
+
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      // Multiply and accumulate threadgroup elements
+      mma_op.mma(As, Bs);
+
+      // Prepare for next iteration
+      loader_a.next();
+      loader_b.next();
+    }
+
+    if (!K_aligned_) {
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      short2 tile_dims_A_last =
+          transpose_a ? short2(tgp_bm, lbk) : short2(lbk, tgp_bm);
+      short2 tile_dims_B_last =
+          transpose_b ? short2(lbk, tgp_bn) : short2(tgp_bn, lbk);
+
+      loader_a.load_safe(tile_dims_A_last);
+      loader_b.load_safe(tile_dims_B_last);
+
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+
+      mma_op.mma(As, Bs);
+    }
+  }
+
+  /* Main kernel function */
+  static METAL_FUNC void run(
+      const device T* A [[buffer(0)]],
+      const device T* B [[buffer(1)]],
+      device U* D [[buffer(2)]],
+      const constant GEMMParams* params [[buffer(3)]],
+      threadgroup T* As [[threadgroup(0)]],
+      threadgroup T* Bs [[threadgroup(1)]],
+      uint simd_lane_id [[thread_index_in_simdgroup]],
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint3 tid [[threadgroup_position_in_grid]],
+      uint3 lid [[thread_position_in_threadgroup]]) {
+    // Pacifying compiler
+    (void)lid;
+
+    const int tid_y = ((tid.y) << params->swizzle_log) +
+        ((tid.x) & ((1 << params->swizzle_log) - 1));
+    const int tid_x = (tid.x) >> params->swizzle_log;
+
+    if (params->tiles_n <= tid_x || params->tiles_m <= tid_y) {
+      return;
+    }
+
+    threadgroup_barrier(mem_flags::mem_none);
+
+    // Find block in A, B, C
+    const int c_row = tid_y * BM;
+    const int c_col = tid_x * BN;
+    const size_t c_row_long = size_t(c_row);
+    const size_t c_col_long = size_t(c_col);
+
+    A += transpose_a ? c_row_long : c_row_long * params->lda;
+    B += transpose_b ? c_col_long * params->ldb : c_col_long;
+    D += c_row_long * params->ldd + c_col_long;
+
+    // Prepare threadgroup loading operations
+    thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id);
+    thread loader_b_t loader_b(B, params->ldb, Bs, simd_group_id, simd_lane_id);
+
+    // Prepare threadgroup mma operation
+    thread mma_t mma_op(simd_group_id, simd_lane_id);
+
+    int gemm_k_iterations = params->gemm_k_iterations_aligned;
+
+    ///////////////////////////////////////////////////////////////////////////////
+    // MNK aligned loop
+    if (MN_aligned) {
+      for (int k = 0; k < gemm_k_iterations; k++) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        // Load elements into threadgroup
+        loader_a.load_unsafe();
+        loader_b.load_unsafe();
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // Multiply and accumulate threadgroup elements
+        mma_op.mma(As, Bs);
+
+        // Prepare for next iteration
+        loader_a.next();
+        loader_b.next();
+      }
+
+      threadgroup_barrier(mem_flags::mem_none);
+
+      // Loop tail
+      if (!K_aligned) {
+        int lbk = params->K - params->gemm_k_iterations_aligned * BK;
+        short2 tile_dims_A = transpose_a ? short2(BM, lbk) : short2(lbk, BM);
+        short2 tile_dims_B = transpose_b ? short2(lbk, BN) : short2(BN, lbk);
+
+        loader_a.load_safe(tile_dims_A);
+        loader_b.load_safe(tile_dims_B);
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        mma_op.mma(As, Bs);
+      }
+
+      // Store results to device memory
+      mma_op.store_result(D, params->ldd);
+      return;
+
+    }
+    ///////////////////////////////////////////////////////////////////////////////
+    // MN unaligned loop
+    else { // Loop over K - unaligned case
+      short tgp_bm = min(BM, params->M - c_row);
+      short tgp_bn = min(BN, params->N - c_col);
+      short leftover_bk = params->K - params->gemm_k_iterations_aligned * BK;
+
+      if (tgp_bm == BM && tgp_bn == BN) {
+        gemm_loop<true, true, K_aligned>(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            leftover_bk);
+
+        mma_op.store_result(D, params->ldd);
+        return;
+
+      } else if (tgp_bn == BN) {
+        gemm_loop<false, true, K_aligned>(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            leftover_bk);
+
+        mma_op.store_result_safe(D, params->ldd, short2(tgp_bn, tgp_bm));
+        return;
+
+      } else if (tgp_bm == BM) {
+        gemm_loop<true, false, K_aligned>(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            leftover_bk);
+
+        mma_op.store_result_safe(D, params->ldd, short2(tgp_bn, tgp_bm));
+        return;
+
+      } else {
+        gemm_loop<false, false, K_aligned>(
+            As,
+            Bs,
+            gemm_k_iterations,
+            loader_a,
+            loader_b,
+            mma_op,
+            tgp_bm,
+            tgp_bn,
+            leftover_bk);
+
+        mma_op.store_result_safe(D, params->ldd, short2(tgp_bn, tgp_bm));
+        return;
+      }
+    }
+  }
+};
+
+} // namespace steel
+} // namespace mlx

src/exo/worker/engines/mlx/matmul/kernels/steel_include/steel/gemm/kernels/steel_gemm_splitk.h

@@ -0,0 +1,227 @@
+// Copyright © 2024 Apple Inc.
+
+using namespace mlx::steel;
+
+///////////////////////////////////////////////////////////////////////////////
+// GEMM kernels
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename T,
+    typename U,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    bool MN_aligned,
+    bool K_aligned>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gemm_splitk(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    device U* C [[buffer(2)]],
+    const constant GEMMSpiltKParams* params [[buffer(3)]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]]) {
+  (void)lid;
+
+  using gemm_kernel = GEMMKernel<
+      T,
+      U,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      MN_aligned,
+      K_aligned>;
+  using loader_a_t = typename gemm_kernel::loader_a_t;
+  using loader_b_t = typename gemm_kernel::loader_b_t;
+  using mma_t = typename gemm_kernel::mma_t;
+
+  threadgroup T As[gemm_kernel::tgp_mem_size_a];
+  threadgroup T Bs[gemm_kernel::tgp_mem_size_b];
+
+  const int tid_x = tid.x;
+  const int tid_y = tid.y;
+  const int tid_z = tid.z;
+
+  if (params->tiles_n <= tid_x || params->tiles_m <= tid_y) {
+    return;
+  }
+
+  // Find block in A, B, C
+  const int c_row = tid_y * BM;
+  const int c_col = tid_x * BN;
+  const int k_start = params->split_k_partition_size * tid_z;
+
+  const size_t c_row_long = size_t(c_row);
+  const size_t c_col_long = size_t(c_col);
+  const size_t k_start_long = size_t(k_start);
+
+  A += transpose_a ? (c_row_long + k_start_long * params->lda)
+                   : (k_start_long + c_row_long * params->lda);
+  B += transpose_b ? (k_start_long + c_col_long * params->ldb)
+                   : (c_col_long + k_start_long * params->ldb);
+  C += (size_t(params->split_k_partition_stride) * tid_z) +
+      (c_row_long * params->ldc + c_col_long);
+
+  // Prepare threadgroup loading operations
+  thread loader_a_t loader_a(A, params->lda, As, simd_group_id, simd_lane_id);
+  thread loader_b_t loader_b(B, params->ldb, Bs, simd_group_id, simd_lane_id);
+
+  // Prepare threadgroup mma operation
+  thread mma_t mma_op(simd_group_id, simd_lane_id);
+
+  int gemm_k_iterations = params->gemm_k_iterations_aligned;
+
+  short tgp_bm = min(BM, params->M - c_row);
+  short tgp_bn = min(BN, params->N - c_col);
+  short leftover_bk = params->K % BK;
+
+  if (MN_aligned || (tgp_bm == BM && tgp_bn == BN)) {
+    gemm_kernel::gemm_loop(
+        As,
+        Bs,
+        gemm_k_iterations,
+        loader_a,
+        loader_b,
+        mma_op,
+        tgp_bm,
+        tgp_bn,
+        leftover_bk,
+        LoopAlignment<true, true, true>{});
+  } else if (tgp_bn == BN) {
+    gemm_kernel::gemm_loop(
+        As,
+        Bs,
+        gemm_k_iterations,
+        loader_a,
+        loader_b,
+        mma_op,
+        tgp_bm,
+        tgp_bn,
+        leftover_bk,
+        LoopAlignment<false, true, true>{});
+  } else if (tgp_bm == BM) {
+    gemm_kernel::gemm_loop(
+        As,
+        Bs,
+        gemm_k_iterations,
+        loader_a,
+        loader_b,
+        mma_op,
+        tgp_bm,
+        tgp_bn,
+        leftover_bk,
+        LoopAlignment<true, false, true>{});
+  } else {
+    gemm_kernel::gemm_loop(
+        As,
+        Bs,
+        gemm_k_iterations,
+        loader_a,
+        loader_b,
+        mma_op,
+        tgp_bm,
+        tgp_bn,
+        leftover_bk,
+        LoopAlignment<false, false, true>{});
+  }
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  if ((tid_z + 1) == (params->split_k_partitions)) {
+    int gemm_k_iter_remaining =
+        (params->K - (k_start + params->split_k_partition_size)) / BK;
+    if (!K_aligned || gemm_k_iter_remaining > 0)
+      gemm_kernel::gemm_loop(
+          As,
+          Bs,
+          gemm_k_iter_remaining,
+          loader_a,
+          loader_b,
+          mma_op,
+          tgp_bm,
+          tgp_bn,
+          leftover_bk,
+          LoopAlignment<false, false, K_aligned>{});
+  }
+
+  if (MN_aligned || (tgp_bm == BM && tgp_bn == BN)) {
+    mma_op.store_result(C, params->ldc);
+  } else {
+    mma_op.store_result_safe(C, params->ldc, short2(tgp_bn, tgp_bm));
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Split k accumulation kernel
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename AccT,
+    typename OutT,
+    typename Epilogue = TransformNone<OutT, AccT>>
+[[kernel]] void gemm_splitk_accum(
+    const device AccT* C_split [[buffer(0)]],
+    device OutT* D [[buffer(1)]],
+    const constant int& k_partitions [[buffer(2)]],
+    const constant int& partition_stride [[buffer(3)]],
+    const constant int& ldd [[buffer(4)]],
+    uint2 gid [[thread_position_in_grid]]) {
+  // Ajust D and C
+  D += gid.x + gid.y * size_t(ldd);
+  C_split += gid.x + gid.y * size_t(ldd);
+
+  size_t offset = 0;
+  AccT out = 0;
+
+  for (int i = 0; i < k_partitions; i++) {
+    out += C_split[offset];
+    offset += partition_stride;
+  }
+
+  // Write output
+  D[0] = Epilogue::apply(out);
+}
+
+template <
+    typename AccT,
+    typename OutT,
+    typename Epilogue = TransformAxpby<OutT, AccT>>
+[[kernel]] void gemm_splitk_accum_axpby(
+    const device AccT* C_split [[buffer(0)]],
+    device OutT* D [[buffer(1)]],
+    const constant int& k_partitions [[buffer(2)]],
+    const constant int& partition_stride [[buffer(3)]],
+    const constant int& ldd [[buffer(4)]],
+    const device OutT* C [[buffer(5)]],
+    const constant int& ldc [[buffer(6)]],
+    const constant int& fdc [[buffer(7)]],
+    const constant float& alpha [[buffer(8)]],
+    const constant float& beta [[buffer(9)]],
+    uint2 gid [[thread_position_in_grid]]) {
+  // Ajust D and C
+  C += gid.x * size_t(fdc) + gid.y * size_t(ldc);
+  D += gid.x + gid.y * size_t(ldd);
+  C_split += gid.x + gid.y * size_t(ldd);
+
+  size_t offset = 0;
+  AccT out = 0;
+
+  for (int i = 0; i < k_partitions; i++) {
+    out += C_split[offset];
+    offset += partition_stride;
+  }
+
+  // Write output
+  Epilogue op(alpha, beta);
+  D[0] = op.apply(out, *C);
+}

src/exo/worker/engines/mlx/matmul/kernels/steel_include/steel/gemm/loader.h

@@ -0,0 +1,137 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/defines.h"
+
+///////////////////////////////////////////////////////////////////////////////
+// Loading helper
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <
+    typename T,
+    short BROWS,
+    short BCOLS,
+    short dst_ld,
+    short reduction_dim,
+    short tgp_size,
+    short alignment = 1,
+    short n_reads = (BCOLS * BROWS) / (tgp_size),
+    short TCOLS = BCOLS / n_reads,
+    short TROWS = tgp_size / TCOLS>
+struct BlockLoader {
+  STEEL_CONST short n_rows = (BROWS + TROWS - 1) / TROWS;
+  STEEL_CONST short vec_size = n_reads;
+
+  // Leading dimension for src
+  const int src_ld;
+  const int tile_stride;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+  const device T* src;
+
+  struct alignas(alignment * sizeof(T)) ReadVector {
+    uint8_t v[sizeof(T) * vec_size];
+  };
+
+  /* Constructor */
+  METAL_FUNC BlockLoader(
+      const device T* src_,
+      const int src_ld_,
+      threadgroup T* dst_,
+      ushort simd_group_id [[simdgroup_index_in_threadgroup]],
+      ushort simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(src_ld_),
+        tile_stride(reduction_dim ? BCOLS : BROWS * src_ld),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        src(src_ + bi * src_ld + bj) {}
+
+  /* Apply operation to threadgroup without bound checking */
+  template <typename UnaryOp>
+  METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        dst[i * dst_ld + j] = op.apply(dst[i * dst_ld + j]);
+      }
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      *((threadgroup ReadVector*)(&dst[i * dst_ld])) =
+          *((const device ReadVector*)(&src[i * src_ld]));
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - with bound checking */
+  METAL_FUNC void load_safe(short2 src_tile_dim) const {
+    src_tile_dim = src_tile_dim - short2(bj, bi);
+
+    // Skip loading if thread has no valid reads
+    if (src_tile_dim.x <= 0 || src_tile_dim.y <= 0) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BROWS; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          dst[i * dst_ld + j] = T(0);
+        }
+      }
+      return;
+    }
+
+    // Use fast thread memory for bound checks
+    bool tmp_idx[vec_size];
+    T tmp_val[vec_size];
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < BROWS; i += TROWS) {
+      // Make sure tmp_idx only contains valid indices
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_idx[j] = (i < src_tile_dim.y) && (j < src_tile_dim.x);
+      }
+
+      // Read valid indices into tmp_val
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_val[j] = src[(tmp_idx[j] ? i * src_ld + j : 0)];
+      }
+
+      // Zero out unneeded values
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        tmp_val[j] = tmp_idx[j] ? tmp_val[j] : T(0);
+      }
+
+      // Copy values to threadgroup memory
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < vec_size; j++) {
+        dst[i * dst_ld + j] = tmp_val[j];
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    src += tile_stride;
+  }
+};
+
+} // namespace steel
+} // namespace mlx

src/exo/worker/engines/mlx/matmul/kernels/steel_include/steel/gemm/params.h

@@ -0,0 +1,65 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+///////////////////////////////////////////////////////////////////////////////
+// GEMM param classes
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+struct GEMMParams {
+  const int M;
+  const int N;
+  const int K;
+
+  const int lda;
+  const int ldb;
+  const int ldd;
+
+  const int tiles_n;
+  const int tiles_m;
+
+  const int64_t batch_stride_a;
+  const int64_t batch_stride_b;
+  const int64_t batch_stride_d;
+
+  const int swizzle_log;
+  const int gemm_k_iterations_aligned;
+
+  const int batch_ndim;
+};
+
+struct GEMMSpiltKParams {
+  const int M;
+  const int N;
+  const int K;
+
+  const int lda;
+  const int ldb;
+  const int ldc;
+
+  const int tiles_n;
+  const int tiles_m;
+
+  const int split_k_partitions;
+  const int split_k_partition_stride;
+  const int split_k_partition_size;
+
+  const int swizzle_log;
+  const int gemm_k_iterations_aligned;
+};
+
+struct GEMMAddMMParams {
+  const int ldc;
+  const int fdc;
+
+  const int64_t batch_stride_c;
+
+  const float alpha;
+  const float beta;
+};
+
+} // namespace steel
+} // namespace mlx

src/exo/worker/engines/mlx/matmul/kernels/steel_include/steel/gemm/transforms.h

@@ -0,0 +1,72 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+///////////////////////////////////////////////////////////////////////////////
+// Transforms and Epilogues
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <typename OutT, typename InT>
+struct TransformNone {
+  static METAL_FUNC OutT apply(InT x) {
+    return static_cast<OutT>(x);
+  }
+
+  static METAL_FUNC OutT apply(InT x, OutT) {
+    return static_cast<OutT>(x);
+  }
+};
+
+template <typename OutT, typename InT>
+struct TransformAdd {
+  TransformAdd(const float, const float) {}
+
+  static METAL_FUNC OutT apply(InT x) {
+    return static_cast<OutT>(x);
+  }
+
+  static METAL_FUNC OutT apply(InT x, OutT c) {
+    return static_cast<OutT>(x) + c;
+  }
+};
+
+template <typename OutT, typename InT>
+struct TransformAxpby {
+  const float alpha;
+  const float beta;
+
+  TransformAxpby(const float alpha_, const float beta_)
+      : alpha(alpha_), beta(beta_) {}
+
+  static METAL_FUNC OutT apply(InT x) {
+    return static_cast<OutT>(x);
+  }
+
+  METAL_FUNC OutT apply(InT x, OutT c) const {
+    return static_cast<OutT>(
+        x * static_cast<InT>(alpha) + (static_cast<OutT>(beta) * c));
+  }
+};
+
+template <typename T>
+struct AccumHelper {
+  typedef float accum_type;
+};
+
+struct BlockSwizzle {
+  static METAL_FUNC int2
+  swizzle(uint3 tid [[threadgroup_position_in_grid]], const int swizzle_log) {
+    const int tid_x = (tid.x) >> swizzle_log;
+    const int tid_y =
+        ((tid.y) << swizzle_log) + ((tid.x) & ((1 << swizzle_log) - 1));
+    return int2(tid_x, tid_y);
+  }
+};
+
+} // namespace steel
+} // namespace mlx

src/exo/worker/engines/mlx/matmul/kernels/steel_include/steel/utils.h

@@ -0,0 +1,42 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <metal_stdlib>
+
+METAL_FUNC ulong2 elem_to_loc_broadcast(
+    uint elem,
+    constant const int* shape,
+    constant const int64_t* a_strides,
+    constant const int64_t* b_strides,
+    int ndim) {
+  ulong loc_a{0};
+  ulong loc_b{0};
+  for (int i = ndim - 1; i >= 0 && elem > 0; --i) {
+    int pos_in_dim = (elem % shape[i]);
+    elem /= shape[i];
+    loc_a += pos_in_dim * a_strides[i];
+    loc_b += pos_in_dim * b_strides[i];
+  }
+  return ulong2(loc_a, loc_b);
+}
+
+METAL_FUNC ulong3 elem_to_loc_broadcast(
+    uint elem,
+    constant const int* shape,
+    constant const int64_t* a_strides,
+    constant const int64_t* b_strides,
+    constant const int64_t* c_strides,
+    int ndim) {
+  ulong loc_a{0};
+  ulong loc_b{0};
+  ulong loc_c{0};
+  for (int i = ndim - 1; i >= 0 && elem > 0; --i) {
+    int pos_in_dim = (elem % shape[i]);
+    elem /= shape[i];
+    loc_a += pos_in_dim * a_strides[i];
+    loc_b += pos_in_dim * b_strides[i];
+    loc_c += pos_in_dim * c_strides[i];
+  }
+  return ulong3(loc_a, loc_b, loc_c);
+}

src/exo/worker/engines/mlx/matmul/kernels/steel_include/steel/utils/integral_constant.h

@@ -0,0 +1,134 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <metal_stdlib>
+#include "mlx/backend/metal/kernels/steel/utils/type_traits.h"
+
+#pragma METAL internals : enable
+
+namespace mlx {
+namespace steel {
+
+///////////////////////////////////////////////////////////////////////////////
+// Integral constant with casting
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T, T v>
+struct integral_constant {
+  static constexpr constant T value = v;
+  using value_type = T;
+  using type = integral_constant;
+
+  METAL_FUNC constexpr operator value_type() const noexcept {
+    return value;
+  }
+
+  // METAL_FUNC constexpr value_type operator()() const noexcept {
+  //   return value;
+  // }
+};
+
+template <bool B>
+using bool_constant = integral_constant<bool, B>;
+using true_type = bool_constant<true>;
+using false_type = bool_constant<false>;
+
+template <class T>
+struct is_integral : bool_constant<metal::is_integral<T>::value> {};
+
+template <class T, T v>
+struct is_integral<integral_constant<T, v>>
+    : bool_constant<metal::is_integral<T>::value> {};
+
+template <typename T>
+constexpr constant bool is_integral_v = is_integral<T>::value;
+
+template <int val>
+using Int = integral_constant<int, val>;
+
+///////////////////////////////////////////////////////////////////////////////
+// Binary Operators on Integral constants
+///////////////////////////////////////////////////////////////////////////////
+
+#define integral_const_binop(__op__, __operator__)          \
+  template <typename T, T tv, typename U, U uv>             \
+  METAL_FUNC constexpr auto __operator__(                   \
+      integral_constant<T, tv>, integral_constant<U, uv>) { \
+    constexpr auto res = tv __op__ uv;                      \
+    return integral_constant<decltype(res), res>{};         \
+  }
+
+integral_const_binop(+, operator+);
+integral_const_binop(-, operator-);
+integral_const_binop(*, operator*);
+integral_const_binop(/, operator/);
+
+integral_const_binop(==, operator==);
+integral_const_binop(!=, operator!=);
+integral_const_binop(<, operator<);
+integral_const_binop(>, operator>);
+integral_const_binop(<=, operator<=);
+integral_const_binop(>=, operator>=);
+
+integral_const_binop(&&, operator&&);
+integral_const_binop(||, operator||);
+
+template <typename T, typename = metal::enable_if_t<!is_integral_v<T>>>
+METAL_FUNC constexpr auto operator||(true_type, T) {
+  return true_type{};
+}
+template <typename T, typename = metal::enable_if_t<!is_integral_v<T>>>
+METAL_FUNC constexpr auto operator||(T, true_type) {
+  return true_type{};
+}
+
+template <typename T, typename = metal::enable_if_t<!is_integral_v<T>>>
+METAL_FUNC constexpr auto operator&&(false_type, T) {
+  return false_type{};
+}
+
+template <typename T, typename = metal::enable_if_t<!is_integral_v<T>>>
+METAL_FUNC constexpr auto operator&&(T, false_type) {
+  return false_type{};
+}
+
+// Dispatch utilities
+template <typename F>
+void dispatch_bool(bool v, F f) {
+  if (v) {
+    f(true_type{});
+  } else {
+    f(false_type{});
+  }
+}
+
+template <int start, int stop, int step, typename F>
+constexpr void const_for_loop(F f) {
+  if constexpr (start < stop) {
+    constexpr auto idx = Int<start>{};
+    f(idx);
+    const_for_loop<start + step, stop, step, F>(f);
+  }
+}
+
+#undef integral_const_binop
+
+///////////////////////////////////////////////////////////////////////////////
+// Reduction operators
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+METAL_FUNC constexpr T sum(T x) {
+  return x;
+}
+
+template <typename T, typename... Us>
+METAL_FUNC constexpr auto sum(T x, Us... us) {
+  return x + sum(us...);
+}
+
+} // namespace steel
+} // namespace mlx
+
+#pragma METAL internals : disable

src/exo/worker/engines/mlx/matmul/kernels/steel_include/steel/utils/type_traits.h

@@ -0,0 +1,55 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <metal_stdlib>
+
+#pragma METAL internals : enable
+
+namespace metal {
+
+template <typename T>
+struct is_empty : metal::bool_constant<__is_empty(T)> {};
+
+#ifdef __cpp_variable_templates
+template <typename T>
+constexpr constant bool is_empty_v = is_empty<T>::value;
+#endif
+
+template <typename... Ts>
+struct make_void {
+  typedef void type;
+};
+
+template <typename... Ts>
+using void_t = typename make_void<Ts...>::type;
+
+template <class T>
+struct is_static : metal::bool_constant<is_empty<remove_cv_t<T>>::value> {};
+
+template <typename T>
+struct pointer_element {};
+
+template <typename T>
+struct pointer_element<thread T*> {
+  using type = remove_cv_t<T>;
+};
+template <typename T>
+struct pointer_element<device T*> {
+  using type = remove_cv_t<T>;
+};
+template <typename T>
+struct pointer_element<constant T*> {
+  using type = remove_cv_t<T>;
+};
+template <typename T>
+struct pointer_element<threadgroup T*> {
+  using type = remove_cv_t<T>;
+};
+
+template <typename T>
+using pointer_element_t = typename pointer_element<remove_cv_t<T>>::type;
+
+} // namespace metal
+
+#pragma METAL internals : disable

src/exo/worker/engines/mlx/matmul/patches/kernel_picker.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""Kernel picker — single source of truth for selecting optimal matmul
+kernel per (N, K, M, dtype).
+
+Used by both the drafter patch (speculative/bf16_lpb_patch.py) and the
+target patch (patches/qwen3_5/lpb_patch.py). Selections are derived from
+mlx_bench matmul/RESULTS.md winners rounded up to the nearest benchmark
+M column (M ∈ {1, 2, 4, 8, 12, 16, 32, 64}).
+
+Returns a uniform-signature callable:
+- bf16: fn(x, w, M, N, K)
+- int8: fn(x, w, scales, biases, M, N, K, group_size)
+"""
+
+from ..kernels.bf16.gemv_loop_over_b import custom_bf16_qmv_loop_over_b as _bf16_lpb
+from ..kernels.bf16.gemv_loop_over_b_twice import custom_bf16_qmv_loop_over_b_twice as _bf16_lpb_twice
+from ..kernels.bf16.gemm_bm8 import custom_bf16_gemm as _bf16_bm8  # noqa: F401
+from ..kernels.bf16.gemm_splitk_steel import custom_bf16_gemm_splitk_steel as _bf16_sk_steel
+
+from ..kernels.quantized.qmv_loop_over_b import custom_qmv_loop_over_b as _int8_lpb
+from ..kernels.quantized.qmm_bm8 import custom_qmm_t_bm8 as _int8_bm8  # noqa: F401
+from ..kernels.quantized.qmm_bm16 import custom_qmm_t_bm16 as _int8_bm16  # noqa: F401
+from ..kernels.quantized.qmm_splitk import custom_qmm_splitk as _int8_qsk
+
+
+def _round_up_to_bench_col(M):
+    for col in (1, 2, 4, 8, 12, 16, 32, 64):
+        if M <= col:
+            return col
+    return 64
+
+
+def pick_bf16_kernel(N, K, M):
+    """Return (name, fn) for the best bf16 kernel at this (N, K, M).
+
+    fn signature: fn(x, w, M, N, K) → y (M, N)
+
+    KNOWN BUG: sk_steel8/16 produce silently incorrect output at very large
+    N (verified broken at N=248320, K=5120 with 100% rel error). For
+    lm_head-sized projections (N > 50000) we force lpb_twice, which is
+    also the fastest CORRECT kernel at those dims.
+    """
+    M_rnd = _round_up_to_bench_col(M)
+
+    # Safety guard for large-N projections (e.g. lm_head).
+    if N > 50000:
+        if M_rnd <= 8:
+            return "lpb", _bf16_lpb
+        return "lpb_twice", _bf16_lpb_twice
+
+    if M_rnd <= 8:
+        return "lpb", _bf16_lpb
+
+    if M_rnd == 12:
+        if max(N, K) <= 4096:
+            return "lpb_twice", _bf16_lpb_twice
+        fn = lambda x, w, M_, N_, K_: _bf16_sk_steel(x, w, M_, N_, K_, BM=16)
+        return "sk_steel16", fn
+
+    if M_rnd == 16:
+        fn = lambda x, w, M_, N_, K_: _bf16_sk_steel(x, w, M_, N_, K_, BM=16)
+        return "sk_steel16", fn
+
+    # M = 32, 64
+    fn = lambda x, w, M_, N_, K_: _bf16_sk_steel(x, w, M_, N_, K_, BM=32)
+    return "sk_steel32", fn
+
+
+def pick_int8_kernel(N, K, M):
+    """Return (name, fn) for the best int8 kernel at this (N, K, M).
+
+    fn signature: fn(x, w, scales, biases, M, N, K, group_size) → y (M, N)
+    """
+    M_rnd = _round_up_to_bench_col(M)
+
+    if M_rnd <= 8:
+        return "lpb", _int8_lpb
+
+    if M_rnd == 12:
+        fn = lambda x, w, s, b, M_, N_, K_, gs: _int8_qsk(x, w, s, b, M_, N_, K_, BM=16, group_size=gs)
+        return "qsk16", fn
+
+    if M_rnd == 16:
+        fn = lambda x, w, s, b, M_, N_, K_, gs: _int8_qsk(x, w, s, b, M_, N_, K_, BM=16, group_size=gs)
+        return "qsk16", fn
+
+    # M = 32, 64
+    fn = lambda x, w, s, b, M_, N_, K_, gs: _int8_qsk(x, w, s, b, M_, N_, K_, BM=32, group_size=gs)
+    return "qsk32", fn

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

@@ -1,3 +1,10 @@
+import json
+import os
+from pathlib import Path
+
+import mlx.nn as nn
+from loguru import logger
+
 from exo.worker.engines.mlx.patches.opt_batch_gen import apply_batch_gen_patch
 from exo.worker.engines.mlx.patches.standard_yarn_rope import patch_yarn_rope
 
@@ -10,5 +17,50 @@ def apply_mlx_patches() -> None:
         return
     _applied = True
     patch_yarn_rope()
-    # patch_gdn_softplus()
-    apply_batch_gen_patch()
+    # Skip fast_next patch when speculative is enabled — it bypasses _next()
+    # which MTPBatchGenerator overrides for speculative decoding
+    if os.environ.get("EXO_SPECULATIVE") != "1":
+        apply_batch_gen_patch()
+
+
+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", "0")
+
+    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 fused_mode == "0":
+        # Fused kernels disabled — fall back to dynamic LpB patches for both
+        # dense and MoE Qwen3.5. apply_lpb_patches silently skips submodules
+        # it doesn't recognize (e.g. routed experts), so it's safe on either
+        # variant.
+        if model_type in ("qwen3_5", "qwen3_5_moe"):
+            from .qwen3_5.lpb_patch import apply_lpb_patches
+
+            logger.info(
+                f"Fused kernels disabled (EXO_FUSED_KERNELS=0); "
+                f"applying LpB patches for {model_type}"
+            )
+            apply_lpb_patches(model, batch_size=4)
+        else:
+            logger.info("Kernel fusion patches disabled (EXO_FUSED_KERNELS=0)")
+        return
+
+    if model_type == "qwen3_5_moe":
+        from .qwen3_5_moe.apply import apply_qwen35_batched_fused_patches
+
+        logger.info("Detected Qwen3.5 MoE model, applying batched fused kernel patches")
+        apply_qwen35_batched_fused_patches(model)
+
+    elif model_type == "qwen3_5":
+        from .qwen3_5.lpb_patch import apply_lpb_patches
+
+        logger.info("Detected Qwen3.5 dense model, applying LpB kernel patches")
+        apply_lpb_patches(model, batch_size=4)

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

@@ -85,7 +85,7 @@ def _fast_next(self: BatchGenerator) -> list[BatchGenerator.Response]:
         _pending_topk_idx = mx.take_along_axis(_pending_topk_idx, sort_order, axis=1)
         _pending_topk_val = mx.take_along_axis(_pending_topk_val, sort_order, axis=1)
         _pending_selected_lps = logprobs[mx.arange(batch_size), batch.y]
-        mx.async_eval(
+        mx.eval(
             batch.y,
             *batch.logprobs,
             *batch.tokens,
@@ -97,7 +97,7 @@ def _fast_next(self: BatchGenerator) -> list[BatchGenerator.Response]:
         _pending_topk_idx = None
         _pending_topk_val = None
         _pending_selected_lps = None
-        mx.async_eval(batch.y, *batch.logprobs, *batch.tokens)
+        mx.eval(batch.y, *batch.logprobs, *batch.tokens)
 
     prev_token_list: list[int] = cast(list[int], prev_tokens.tolist())
 

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

@@ -0,0 +1,147 @@
+#!/usr/bin/env python3
+"""Dynamic per-call loop-over-B patches for Qwen3.5-27B (bf16 or 8-bit target).
+
+For each projection, the kernel is picked at CALL time based on the actual
+M seen in the forward pass. Memoized per projection so runtime overhead is
+a dict lookup after the first call at a given M.
+
+Target-side projections during verify all see M = verify_len + 1 (uniform),
+but we use dynamic picking for consistency with the drafter patch and to
+support arbitrary (BS, V) without re-patching.
+
+Patches (every layer):
+  - MLP: gate_proj, up_proj, down_proj
+  - GQA attn: q_proj, k_proj, v_proj, o_proj
+  - GDN attn: in_proj_qkv, in_proj_z, out_proj
+  - lm_head (runs during verify forward)
+"""
+
+import mlx.nn as nn
+
+from exo.worker.engines.mlx.matmul.patches.kernel_picker import (
+    pick_bf16_kernel,
+    pick_int8_kernel,
+)
+
+
+MAX_M = 16  # Above this (prefill), fall back to the original projection
+
+
+def _make_bf16_forward(original, N, K):
+    cache = {}
+
+    def forward(self_unused, x):
+        M = 1
+        for d in x.shape[:-1]:
+            M *= d
+        if M > MAX_M:
+            return original(x)
+        fn = cache.get(M)
+        if fn is None:
+            _, fn = pick_bf16_kernel(N, K, M)
+            cache[M] = fn
+        orig_shape = x.shape
+        x_2d = x.reshape(-1, K)
+        y = fn(x_2d, original.weight, M, N, K)
+        return y.reshape(*orig_shape[:-1], N)
+
+    return forward
+
+
+def _make_int8_forward(original, N, K, GS):
+    cache = {}
+
+    def forward(self_unused, x):
+        M = 1
+        for d in x.shape[:-1]:
+            M *= d
+        if M > MAX_M:
+            return original(x)
+        fn = cache.get(M)
+        if fn is None:
+            _, fn = pick_int8_kernel(N, K, M)
+            cache[M] = fn
+        orig_shape = x.shape
+        x_2d = x.reshape(-1, K)
+        y = fn(x_2d, original.weight, original.scales, original.biases,
+               M, N, K, GS)
+        return y.reshape(*orig_shape[:-1], N)
+
+    return forward
+
+
+def _patch_proj(parent, proj_name):
+    proj = getattr(parent, proj_name, None)
+    if proj is None:
+        return 0
+
+    if isinstance(proj, nn.QuantizedLinear):
+        N = proj.weight.shape[0]
+        K = proj.weight.shape[1] * (32 // proj.bits)
+        GS = proj.group_size
+        forward = _make_int8_forward(proj, N, K, GS)
+        setattr(parent, proj_name, type('LpBQuant', (), {
+            '__call__': forward,
+            'weight': proj.weight,
+            'scales': proj.scales,
+            'biases': proj.biases,
+        })())
+        return 1
+    elif isinstance(proj, nn.Linear):
+        N = proj.weight.shape[0]
+        K = proj.weight.shape[1]
+        forward = _make_bf16_forward(proj, N, K)
+        setattr(parent, proj_name, type('LpBLinear', (), {
+            '__call__': forward,
+            'weight': proj.weight,
+        })())
+        return 1
+    return 0
+
+
+def apply_lpb_patches(model, batch_size=None, verify_len=None):
+    """Patch all Qwen3.5-27B projections with dynamic LpB kernels.
+
+    Note: batch_size / verify_len args are kept for backward compat but ignored.
+    Kernel is picked at call time based on actual M.
+    """
+    inner = getattr(model, 'model', None) or model.language_model.model
+    patched = 0
+
+    for _li, layer in enumerate(inner.layers):
+        mlp = layer.mlp
+        for pn in ('gate_proj', 'up_proj', 'down_proj'):
+            patched += _patch_proj(mlp, pn)
+
+        # MoE (Qwen3NextSparseMoeBlock): dense sub-modules worth LpB-patching.
+        # Silent-skip on dense 27B (attributes don't exist). Routed experts
+        # (switch_mlp) use SwitchLinear which needs routing indices — left on
+        # stock.
+        patched += _patch_proj(mlp, 'gate')
+        patched += _patch_proj(mlp, 'shared_expert_gate')
+        shared = getattr(mlp, 'shared_expert', None)
+        if shared is not None:
+            for pn in ('gate_proj', 'up_proj', 'down_proj'):
+                patched += _patch_proj(shared, pn)
+
+        if layer.is_linear:
+            attn = layer.linear_attn
+            for pn in ('in_proj_qkv', 'in_proj_z', 'out_proj'):
+                patched += _patch_proj(attn, pn)
+        else:
+            attn = layer.self_attn
+            for pn in ('q_proj', 'k_proj', 'v_proj', 'o_proj'):
+                patched += _patch_proj(attn, pn)
+
+    # Qwen3.5 / Qwen3.5-MoE keep lm_head on the TextModel wrapper
+    # (model.language_model.lm_head), not on model or the inner Qwen3_5TextModel.
+    # Check all three levels so we cover both layouts.
+    for holder in (model, getattr(model, 'language_model', None), inner):
+        if holder is None:
+            continue
+        if hasattr(holder, 'lm_head') and holder.lm_head is not None:
+            patched += _patch_proj(holder, 'lm_head')
+            break
+
+    print(f"  Patched {patched} target projections with dynamic LpB")
+    return patched

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

@@ -0,0 +1,74 @@
+"""Apply batched fused kernel patches to Qwen3.5 MoE models.
+
+Entry point called from patches/__init__.py after model type detection.
+"""
+
+import time
+
+import mlx.nn as nn
+from loguru import logger
+
+from .common import (
+    _patch_swiglu_weights,
+    _patch_shared_expert,
+    _patch_down_proj,
+    _patch_oproj_gate_rms,
+    _patch_gdn_proj_weights,
+    _patch_gqa_proj_weights,
+)
+from mlx_lm.models.qwen3_5 import DecoderLayer
+from mlx_lm.models.qwen3_next import Qwen3NextAttention, Qwen3NextSparseMoeBlock
+from mlx_lm.models.qwen3_5 import GatedDeltaNet
+
+
+def apply_qwen35_batched_fused_patches(model: nn.Module) -> None:
+    """Apply batched fused patches (GDN + GQA attention + oproj MoE) to all layers.
+
+    Fused GDN attention (3/4 layers) + fused GQA projections (1/4 layers)
+    + batched oproj MoE (4 custom dispatches). Works with BatchGenerator for
+    any batch size 1..8. Falls back to vanilla for B>8 or S>1.
+    """
+    layers = model.layers  # type: ignore[attr-defined]
+    n_layers = len(layers)
+
+    t0 = time.time()
+
+    n_gdn = 0
+    n_gqa = 0
+    for li, layer in enumerate(layers):
+        moe = layer.mlp
+        if isinstance(moe, Qwen3NextSparseMoeBlock):
+            # MoE weight prep
+            _patch_swiglu_weights(moe)
+            _patch_shared_expert(moe)
+            _patch_down_proj(moe)
+            _patch_oproj_gate_rms(layer, gate_bm=8)
+
+            # Attention weight prep
+            if layer.is_linear:
+                _patch_gdn_proj_weights(layer.linear_attn)
+                n_gdn += 1
+            else:
+                _patch_gqa_proj_weights(layer.self_attn)
+                n_gqa += 1
+
+            if (li + 1) % 10 == 0 or li == 0:
+                logger.info(f"  Patched layer {li+1}/{n_layers}")
+
+    # Import patched __call__ methods
+    from .fused_gdn_attention import _fused_gdn_call
+    from .batched_fused_gqa_attention import _batched_fused_gqa_call
+    from .batched_moe import _batched_oproj_moe_call
+    from .decoder import _fused_gdn_decoder_call
+
+    # Class-level method replacement
+    GatedDeltaNet.__call__ = _fused_gdn_call
+    Qwen3NextAttention.__call__ = _batched_fused_gqa_call
+    Qwen3NextSparseMoeBlock.__call__ = _batched_oproj_moe_call
+    DecoderLayer.__call__ = _fused_gdn_decoder_call
+
+    t_patch = time.time() - t0
+    logger.info(
+        f"Qwen3.5 batched fused: {n_gdn} GDN + {n_gqa} GQA layers, "
+        f"{n_layers} total in {t_patch:.1f}s"
+    )

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

@@ -0,0 +1,102 @@
+"""Batched fused GQA attention for Qwen3.5 (projections + norm+rope fused, vanilla SDPA).
+
+Dispatches:
+  1. batched_fused_gqa_projections — merged q+gate+k+v GEMV with register weight sharing
+  2. fused_qk_norm_rope — per-head RMSNorm + RoPE (already supports B>1 via grid z)
+  3. Vanilla cache update (BatchKVCache)
+  4. Vanilla SDPA (MLX built-in, handles batching natively)
+  5. Vanilla gate multiply
+
+Returns pre-out_proj output for the oproj MoE block.
+Falls back to vanilla (with o_proj) for B>8 or S>1.
+"""
+
+from typing import Any, Optional
+
+import mlx.core as mx
+
+from .kernels.batched_fused_gqa_projections_8bit import batched_fused_gqa_projections
+
+
+def _batched_fused_gqa_call(
+    self,
+    x: mx.array,
+    mask: Optional[mx.array] = None,
+    cache: Optional[Any] = None,
+) -> mx.array:
+    """Batched fused GQA attention with custom projection + norm/rope kernels.
+
+    For 1<=B<=8, S=1: fused projections + fused norm+rope + vanilla SDPA.
+    For B>8 or S>1: vanilla fallback.
+    Returns pre-out_proj output [B, S, H_q*D].
+    """
+    B, S, _ = x.shape
+
+    if S > 1 or B > 8:
+        # Vanilla fallback
+        q_proj_output = self.q_proj(x)
+        queries, gate = mx.split(
+            q_proj_output.reshape(B, S, self.num_attention_heads, -1), 2, axis=-1
+        )
+        gate = gate.reshape(B, S, -1)
+        keys, values = self.k_proj(x), self.v_proj(x)
+        queries = self.q_norm(queries).transpose(0, 2, 1, 3)
+        keys = self.k_norm(
+            keys.reshape(B, S, self.num_key_value_heads, -1)
+        ).transpose(0, 2, 1, 3)
+        values = values.reshape(B, S, self.num_key_value_heads, -1).transpose(
+            0, 2, 1, 3
+        )
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        from mlx_lm.models.qwen3_next import scaled_dot_product_attention
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, S, -1)
+        return self.o_proj(output * mx.sigmoid(gate))
+
+    H_q = self.num_attention_heads
+    H_kv = self.num_key_value_heads
+    D = self.head_dim
+
+    # ── Dispatch 1: batched fused projections ──
+    queries, gate_sigmoid, keys, values = batched_fused_gqa_projections(
+        x,
+        self._merged_proj_w, self._merged_proj_s, self._merged_proj_b,
+        self._merged_proj_dims,
+        batch_size=B,
+        total_tg=getattr(self, '_d1_total_tg', None),
+    )
+
+    # ── Dispatch 2+: vanilla norm + rope (avoids mx.eval sync on BatchKVCache offset) ──
+    queries = self.q_norm(queries.reshape(B, 1, H_q, D)).transpose(0, 2, 1, 3)
+    keys = self.k_norm(keys.reshape(B, 1, H_kv, D)).transpose(0, 2, 1, 3)
+    values = values.reshape(B, 1, H_kv, D).transpose(0, 2, 1, 3)
+
+    if cache is not None:
+        queries = self.rope(queries, offset=cache.offset)
+        keys = self.rope(keys, offset=cache.offset)
+    else:
+        queries = self.rope(queries)
+        keys = self.rope(keys)
+
+    # ── Dispatch 3: KV cache update ──
+    if cache is not None:
+        keys, values = cache.update_and_fetch(keys, values)
+
+    # ── Dispatch 4: vanilla SDPA ──
+    from mlx_lm.models.qwen3_next import scaled_dot_product_attention
+    output = scaled_dot_product_attention(
+        queries, keys, values, cache=cache, scale=self.scale, mask=mask
+    )
+    output = output.transpose(0, 2, 1, 3).reshape(B, S, -1)
+
+    # ── Gate multiply ──
+    return output * gate_sigmoid.astype(output.dtype)

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

@@ -0,0 +1,94 @@
+"""Batched oproj MoE with 4 custom Metal kernel dispatches.
+
+Fuses o_proj + RMSNorm + gate GEMV + softmax + topk + SwiGLU + down_proj + epilogue
+into 4 dispatches with register-level weight sharing for the shared expert.
+Falls back to vanilla MoE when called without _residual (from vanilla decoder path).
+"""
+
+import mlx.core as mx
+
+from .kernels.batched_merged_down_proj_8bit import batched_merged_down_proj_8bit
+from .kernels.batched_oproj_gate_gemv_8bit import batched_oproj_gate_gemv
+from .kernels.batched_softmax_topk_swiglu_8bit import batched_softmax_topk_swiglu_8bit
+from .kernels.batched_moe_epilogue import batched_moe_epilogue
+
+
+def _batched_oproj_moe_call(self, attn_out_3d, _residual=None):
+    """Batched MoE with full oproj fusion (4 custom dispatches).
+
+    Receives raw attention output (pre-o_proj) and residual for B tokens.
+    All 4 dispatches use register-level weight sharing for shared weights.
+
+    When _residual is None, called from vanilla decoder — do vanilla MoE.
+    """
+    if _residual is None:
+        # Vanilla MoE path (called from vanilla decoder for B>8 or S>1)
+        x = attn_out_3d
+        gates = self.gate(x)
+        gates = mx.softmax(gates, axis=-1, precise=True)
+        k = self.top_k
+        inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]
+        scores = mx.take_along_axis(gates, inds, axis=-1)
+        if self.norm_topk_prob:
+            scores = scores / scores.sum(axis=-1, keepdims=True)
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2)
+        shared_y = self.shared_expert(x)
+        shared_y = mx.sigmoid(self.shared_expert_gate(x)) * shared_y
+        return y + shared_y
+
+    B_dim = attn_out_3d.shape[0]
+
+    K = self._oproj_M
+    K_attn = self._oproj_K_attn
+    n_active = self.top_k
+    E = self._oproj_n_experts
+
+    attn_out = attn_out_3d.reshape(B_dim, K_attn).astype(mx.bfloat16)
+    residual = _residual.reshape(B_dim, K).astype(mx.bfloat16)
+
+    # ── Dispatch 1: batched o_proj + gate GEMVs ──
+    h_scaled, h_out, x2_partials, gate_part_a, gate_part_b = \
+        batched_oproj_gate_gemv(
+            self._oproj_w, self._oproj_s, self._oproj_b,
+            attn_out, residual, self._oproj_rms_weight,
+            self._oproj_M1, self._oproj_W_fused,
+            M=K, K_attn=K_attn, batch_size=B_dim,
+            n_experts=E, gate_bm=self._oproj_gate_bm,
+            K_hidden=self._oproj_K_hidden,
+        )
+
+    n_oproj_tg = (K + 31) // 32
+    N_INTER = self.switch_mlp._fused_n_inter
+    SHARED_INTER = self._shared_inter
+
+    # ── Dispatch 2: batched softmax + topk + SwiGLU ──
+    y_routed, y_shared, out_inds, norm_scores, gate_raw = \
+        batched_softmax_topk_swiglu_8bit(
+            self.switch_mlp._fused_w_gu, self.switch_mlp._fused_s_gu,
+            self.switch_mlp._fused_b_gu,
+            self._shared_w_gu, self._shared_s_gu, self._shared_b_gu,
+            self._seg_w, self._seg_s, self._seg_b,
+            h_scaled, gate_part_a, gate_part_b, x2_partials,
+            n_inter=N_INTER, k_hidden=K, batch_size=B_dim,
+            n_active=n_active, n_oproj_tg=n_oproj_tg,
+            n_experts=E, shared_inter=SHARED_INTER,
+        )
+
+    # ── Dispatch 3: batched merged down_proj ──
+    d_routed, d_shared = batched_merged_down_proj_8bit(
+        self._down_w, self._down_s, self._down_b,
+        self._shared_down_w, self._shared_down_s, self._shared_down_b,
+        y_routed, y_shared.reshape(B_dim * SHARED_INTER), out_inds,
+        k_out=K, n_in=self._down_N, batch_size=B_dim,
+        n_active=n_active, shared_n_in=SHARED_INTER,
+    )
+
+    # ── Dispatch 4: batched epilogue ──
+    Y = batched_moe_epilogue(
+        d_routed, d_shared, norm_scores,
+        h_out, gate_raw,
+        k_val=K, batch_size=B_dim, n_active=n_active,
+    )
+
+    return Y.reshape(B_dim, 1, K).astype(attn_out_3d.dtype)

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

@@ -0,0 +1,500 @@
+"""Common weight preparation functions for Qwen3.5 fused kernel patches.
+
+Functions:
+  ceil_div               — integer ceiling division
+  _patch_swiglu_weights  — stack gate+up weights for fused SwiGLU kernel
+  _patch_down_proj       — extract down_proj weights for merged kernel dispatch
+  _patch_shared_expert   — prepare shared expert weights (8-bit)
+  dequantize_shared_expert — convert shared expert from 8-bit to bf16
+  _patch_oproj_gate_rms  — precompute M1/W_fused for fused o_proj + gate GEMV
+  _patch_gdn_proj_weights — merge GDN projection weights for fused GEMV
+  _patch_gqa_proj_weights — merge GQA q/k/v weights with q_proj permutation
+  make_qwen_random_cache — create pre-filled cache for testing
+  build_model            — build Qwen3.5 MoE layers with 8-bit quantization
+"""
+
+from types import SimpleNamespace
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx_lm.models.qwen3_5 import (
+    DecoderLayer,
+    TextModelArgs,
+)
+from mlx_lm.models.qwen3_next import Qwen3NextSparseMoeBlock
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _patch_swiglu_weights(moe):
+    """Stack gate+up weights for fused 8-bit SwiGLU kernel.
+
+    Creates concatenated (E, 2*N_INTER, K/4) weight, (E, 2*N_INTER, K/gs) scales/biases
+    from the separate gate_proj and up_proj QuantizedSwitchLinear layers.
+    """
+    gate_proj = moe.switch_mlp.gate_proj
+    up_proj = moe.switch_mlp.up_proj
+
+    moe.switch_mlp._fused_w_gu = mx.concatenate(
+        [gate_proj.weight, up_proj.weight], axis=1)
+    moe.switch_mlp._fused_s_gu = mx.concatenate(
+        [gate_proj.scales, up_proj.scales], axis=1)
+    moe.switch_mlp._fused_b_gu = mx.concatenate(
+        [gate_proj.biases, up_proj.biases], axis=1)
+    moe.switch_mlp._fused_n_inter = gate_proj.output_dims
+    moe.switch_mlp._fused_k_hidden = gate_proj.input_dims
+    moe.switch_mlp._fused_group_size = gate_proj.group_size
+
+    mx.eval(moe.switch_mlp._fused_w_gu,
+            moe.switch_mlp._fused_s_gu,
+            moe.switch_mlp._fused_b_gu)
+
+
+def _patch_shared_expert(moe):
+    """Prepare shared expert quantized weights for fused 8-bit path.
+
+    Stacks shared gate+up quantized weights (weight, scales, biases).
+    Stores down_proj quantized weights separately.
+    Shared expert stays in 8-bit — same as vanilla MLX dispatch.
+    """
+    shared = moe.shared_expert
+    gp = shared.gate_proj
+    up = shared.up_proj
+    dp = shared.down_proj
+
+    # Gate+up stacked: (2*SHARED_INTER, K/4) uint32, (2*SHARED_INTER, K/gs) bf16
+    moe._shared_w_gu = mx.concatenate([gp.weight, up.weight], axis=0)
+    moe._shared_s_gu = mx.concatenate([gp.scales, up.scales], axis=0)
+    moe._shared_b_gu = mx.concatenate([gp.biases, up.biases], axis=0)
+
+    # Down_proj: (K, SHARED_INTER/4) uint32, (K, SHARED_INTER/gs) bf16
+    moe._shared_down_w = dp.weight
+    moe._shared_down_s = dp.scales
+    moe._shared_down_b = dp.biases
+
+    # QuantizedLinear: weight is (out_features, in_features/pack_factor) uint32
+    # For 8-bit: pack_factor = 4, so in_features = weight.shape[1] * 4
+    moe._shared_inter = gp.weight.shape[0]  # SHARED_INTER (= out_features)
+    moe._shared_gs = gp.group_size           # gs (64)
+
+    mx.eval(moe._shared_w_gu, moe._shared_s_gu, moe._shared_b_gu,
+            moe._shared_down_w, moe._shared_down_s, moe._shared_down_b)
+
+
+def _patch_down_proj(moe):
+    """Extract down_proj weights for merged 8-bit kernel dispatch."""
+    dp = moe.switch_mlp.down_proj
+    moe._down_w = dp.weight       # (E, K_OUT, N_IN/4) uint32
+    moe._down_s = dp.scales       # (E, K_OUT, N_IN/gs) bf16
+    moe._down_b = dp.biases       # (E, K_OUT, N_IN/gs) bf16
+    moe._down_K = dp.output_dims  # K = 4096
+    moe._down_N = dp.input_dims   # N = 1024
+    moe._down_gs = dp.group_size  # gs = 64
+    mx.eval(moe._down_w, moe._down_s, moe._down_b)
+
+
+def dequantize_shared_expert(moe):
+    """Convert shared expert from 8-bit QuantizedLinear to bf16 weight wrappers.
+
+    The fused kernels expect bf16 shared expert weights. The real model (and our
+    random model) has shared expert quantized to 8-bit. This dequantizes in-place.
+    """
+    shared = moe.shared_expert
+    for proj_name in ["gate_proj", "up_proj", "down_proj"]:
+        proj = getattr(shared, proj_name)
+        if hasattr(proj, 'scales') and hasattr(proj, 'biases'):
+            w_bf16 = mx.dequantize(
+                proj.weight, proj.scales, proj.biases,
+                group_size=proj.group_size, bits=proj.bits,
+            ).astype(mx.bfloat16)
+            mx.eval(w_bf16)
+            setattr(shared, proj_name, SimpleNamespace(weight=w_bf16))
+
+
+def _patch_oproj_gate_rms(layer, gate_bm=8):
+    """Precompute M1/W_fused for fused o_proj + gate GEMV (oproj 4-dispatch mode).
+
+    Gate decomposition:
+      gate_score[e] = W_gate[e,:] @ rms_norm(h)
+      where h = residual + W_oproj @ attn_out
+      rms_norm(h) = h * w_rms * inv_rms
+
+    Expanding:
+      gate_score[e] = (W_fused @ residual + M1 @ attn_out) * inv_rms
+
+    Precomputed offline (per layer, stored on moe block):
+      W_fused = dequant(W_gate) · diag(w_rms)    — (E, K) bf16
+      M1 = W_fused @ dequant(W_oproj)            — (E, K_attn) bf16
+
+    Also stores o_proj quantized weights and shared_expert_gate weights
+    on the moe block for use by Dispatch 1 and Dispatch 2.
+
+    Args:
+        layer: DecoderLayer instance
+        gate_bm: SGs per gate TG in Dispatch 1 (1,2,4,8)
+    """
+    moe = layer.mlp
+
+    # ── Get attention output projection (works for both attention types) ──
+    if layer.is_linear:
+        oproj = layer.linear_attn.out_proj
+    else:
+        oproj = layer.self_attn.o_proj
+
+    # ── Dequantize gate and o_proj (temporary, for M1 computation) ──
+    # Eval incrementally to limit peak memory (E=512: dequant temps are ~140 MB)
+    gate = moe.gate
+    W_gate_f32 = mx.dequantize(
+        gate.weight, gate.scales, gate.biases,
+        group_size=gate.group_size, bits=gate.bits,
+    ).astype(mx.float32)
+
+    W_oproj_f32 = mx.dequantize(
+        oproj.weight, oproj.scales, oproj.biases,
+        group_size=oproj.group_size, bits=oproj.bits,
+    ).astype(mx.float32)
+    mx.eval(W_gate_f32, W_oproj_f32)
+
+    # ── RMSNorm weight ──
+    rms_weight = layer.post_attention_layernorm.weight.astype(mx.bfloat16)
+
+    # ── W_fused = dequant(W_gate) · diag(w_rms) ──
+    w_rms_f32 = rms_weight.astype(mx.float32)
+    W_fused = (W_gate_f32 * w_rms_f32).astype(mx.bfloat16)
+    mx.eval(W_fused)
+    del W_gate_f32  # free ~8 MB (E=512) or ~1 MB (E=64)
+
+    # ── M1 = W_fused @ W_oproj — precomputed in f32, stored bf16 ──
+    M1 = (W_fused.astype(mx.float32) @ W_oproj_f32).astype(mx.bfloat16)
+    mx.eval(M1)
+    del W_oproj_f32  # free ~128 MB
+
+    # Store on moe block
+    moe._oproj_M1 = M1              # (E, K_attn) bf16
+    moe._oproj_W_fused = W_fused    # (E, K) bf16
+    moe._oproj_rms_weight = rms_weight  # (K,) bf16
+
+    # ── O_proj quantized weights (for 8-bit GEMV in Dispatch 1) ──
+    moe._oproj_w = oproj.weight      # (K, K_attn/4) uint32
+    moe._oproj_s = oproj.scales      # (K, K_attn/gs) bf16
+    moe._oproj_b = oproj.biases      # (K, K_attn/gs) bf16
+    moe._oproj_K_attn = oproj.weight.shape[1] * 4  # 8192 (8-bit: pack_factor=4)
+
+    # ── Shared expert gate weights (for TG(0,0,0) fusion in Dispatch 2) ──
+    seg = moe.shared_expert_gate
+    moe._seg_w = seg.weight    # (1, K/4) uint32
+    moe._seg_s = seg.scales    # (1, K/gs) bf16
+    moe._seg_b = seg.biases    # (1, K/gs) bf16
+
+    # ── Dimensions ──
+    M = oproj.weight.shape[0]                    # 4096 (hidden_size)
+    K_hidden = W_fused.shape[1]                   # 4096 (same as M for Qwen)
+    n_experts = W_fused.shape[0]                  # E
+    moe._oproj_M = M
+    moe._oproj_K_hidden = K_hidden
+    moe._oproj_n_experts = n_experts
+    moe._oproj_n_tg = ceil_div(M, 32)            # 128 for M=4096
+    moe._oproj_gate_bm = gate_bm
+
+    mx.eval(moe._oproj_rms_weight)
+
+
+
+def _patch_gdn_proj_weights(attn):
+    """Merge all 4 GDN projection weights into contiguous buffers.
+
+    Concatenates in_proj_qkv/z/b/a weights, scales, biases into single
+    contiguous arrays for better memory locality in the fused GEMV kernel.
+    Stored on the GatedDeltaNet module as _merged_proj_*.
+    """
+    W_merged = mx.concatenate([
+        attn.in_proj_qkv.weight,
+        attn.in_proj_z.weight,
+        attn.in_proj_b.weight,
+        attn.in_proj_a.weight,
+    ], axis=0)
+    S_merged = mx.concatenate([
+        attn.in_proj_qkv.scales,
+        attn.in_proj_z.scales,
+        attn.in_proj_b.scales,
+        attn.in_proj_a.scales,
+    ], axis=0)
+    B_merged = mx.concatenate([
+        attn.in_proj_qkv.biases,
+        attn.in_proj_z.biases,
+        attn.in_proj_b.biases,
+        attn.in_proj_a.biases,
+    ], axis=0)
+    attn._merged_proj_w = W_merged
+    attn._merged_proj_s = S_merged
+    attn._merged_proj_b = B_merged
+    attn._merged_proj_dims = (
+        attn.in_proj_qkv.weight.shape[0],  # N_QKV = 8192
+        attn.in_proj_z.weight.shape[0],     # N_Z = 4096
+        attn.in_proj_b.weight.shape[0],     # N_B = 32
+        attn.in_proj_a.weight.shape[0],     # N_A = 32
+    )
+    mx.eval(W_merged, S_merged, B_merged)
+
+
+def _patch_gqa_proj_weights(attn):
+    """Merge GQA q_proj, k_proj, v_proj weights into contiguous buffers.
+
+    q_proj outputs (H_q * 2 * D) = interleaved [queries, gate] per head.
+    We permute rows so queries (H_q * D) come first, then gate (H_q * D),
+    then k_proj, then v_proj. This gives clean contiguous regions for
+    the fused GEMV kernel's TG routing.
+
+    Permutation for q_proj:
+        Original row layout: [head0_q[0:D], head0_gate[0:D], head1_q[0:D], ...]
+        After permutation:   [head0_q, head1_q, ..., head0_gate, head1_gate, ...]
+
+    Stored on Qwen3NextAttention as _merged_proj_*.
+    """
+    q = attn.q_proj
+    k = attn.k_proj
+    v = attn.v_proj
+
+    H_q = attn.num_attention_heads
+    D = attn.head_dim
+
+    # Permute q_proj weights: separate queries and gate rows
+    # q_proj.weight shape: (H_q * 2 * D, K / pack_factor) for 8-bit
+    # Reshape to (H_q, 2*D, ...), split into queries[:, :D, :] and gate[:, D:, :]
+    W_q = q.weight.reshape(H_q, 2 * D, -1)
+    S_q = q.scales.reshape(H_q, 2 * D, -1)
+    B_q = q.biases.reshape(H_q, 2 * D, -1)
+
+    W_queries = W_q[:, :D, :].reshape(H_q * D, -1)
+    W_gate = W_q[:, D:, :].reshape(H_q * D, -1)
+    S_queries = S_q[:, :D, :].reshape(H_q * D, -1)
+    S_gate = S_q[:, D:, :].reshape(H_q * D, -1)
+    B_queries = B_q[:, :D, :].reshape(H_q * D, -1)
+    B_gate = B_q[:, D:, :].reshape(H_q * D, -1)
+
+    # Merge: [queries, gate, keys, values]
+    W_merged = mx.contiguous(mx.concatenate([W_queries, W_gate, k.weight, v.weight], axis=0))
+    S_merged = mx.contiguous(mx.concatenate([S_queries, S_gate, k.scales, v.scales], axis=0))
+    B_merged = mx.contiguous(mx.concatenate([B_queries, B_gate, k.biases, v.biases], axis=0))
+
+    attn._merged_proj_w = W_merged
+    attn._merged_proj_s = S_merged
+    attn._merged_proj_b = B_merged
+    attn._merged_proj_dims = (
+        H_q * D,           # N_Q = 4096 (queries)
+        H_q * D,           # N_GATE = 4096 (gate)
+        k.weight.shape[0], # N_K = 512
+        v.weight.shape[0], # N_V = 512
+    )
+    mx.eval(W_merged, S_merged, B_merged)
+
+    # Pre-cache constant scalar arrays for kernel dispatch (avoid per-call creation)
+    N_Q, N_GATE, N_K, N_V = attn._merged_proj_dims
+    N_TOTAL = N_Q + N_GATE + N_K + N_V
+    K = q.weight.shape[1] * 4  # 8-bit: pack_factor=4
+    attn._kernel_scalars = {
+        # Dispatch 1: fused_gqa_projections
+        'K': mx.array(K, dtype=mx.int32),
+        'N_Q': mx.array(N_Q, dtype=mx.int32),
+        'N_GATE': mx.array(N_GATE, dtype=mx.int32),
+        'N_K': mx.array(N_K, dtype=mx.int32),
+        'N_TOTAL': mx.array(N_TOTAL, dtype=mx.int32),
+        'N_Q_TG': mx.array(ceil_div(N_Q, 8), dtype=mx.int32),
+        'N_GATE_TG': mx.array(ceil_div(N_GATE, 8), dtype=mx.int32),
+        'N_K_TG': mx.array(ceil_div(N_K, 8), dtype=mx.int32),
+        # Dispatch 4-5: custom SDPA
+        'scale': mx.array(attn.head_dim ** -0.5, dtype=mx.float32),
+        'H_Q': mx.array(attn.num_attention_heads, dtype=mx.int32),
+        'H_KV': mx.array(attn.num_key_value_heads, dtype=mx.int32),
+        'N_blocks': mx.array(128, dtype=mx.int32),
+    }
+    mx.eval(*attn._kernel_scalars.values())
+
+    # Precompute grid/TG dims for Dispatch 1
+    N_V_TG = ceil_div(N_V, 8)
+    attn._d1_total_tg = ceil_div(N_Q, 8) + ceil_div(N_GATE, 8) + ceil_div(N_K, 8) + N_V_TG
+
+    # Precompute RoPE inv_freq for fused norm+rope kernel (Dispatch 2)
+    # inv_freq[d] = theta^(-d / half_dims) for d in {0, ..., half_dims-1}
+    rope_dims = attn.rope.dims           # 64 (partial_rotary_factor * head_dim)
+    half_dims = rope_dims // 2            # 32
+    theta = attn.rope.base                # 10000000
+    d_indices = mx.arange(half_dims, dtype=mx.float32)
+    attn._rope_inv_freq = theta ** (-d_indices / half_dims)
+    mx.eval(attn._rope_inv_freq)
+
+
+
+def make_qwen_random_cache(layer, config, prefill_len):
+    """Create a pre-filled cache for a single Qwen3.5 decoder layer.
+
+    GatedDeltaNet layers get ArraysCache(size=2) with fixed-size state:
+      cache[0] = conv state: (B, conv_kernel_size-1, conv_dim) bf16
+      cache[1] = SSM state:  (B, num_v_heads, head_k_dim, head_v_dim) bf16
+
+    GQA layers get KVCache with prefill_len tokens:
+      keys:   (B, n_kv_heads, alloc_len, head_dim) bf16
+      values: (B, n_kv_heads, alloc_len, head_dim) bf16
+    """
+    if layer.is_linear:
+        from mlx_lm.models.cache import ArraysCache
+        cache = ArraysCache(size=2)
+        attn = layer.linear_attn
+        cache[0] = mx.random.normal(
+            (1, attn.conv_kernel_size - 1, attn.conv_dim)
+        ).astype(mx.bfloat16)
+        cache[1] = mx.random.normal(
+            (1, attn.num_v_heads, attn.head_k_dim, attn.head_v_dim)
+        ).astype(mx.bfloat16)
+        return cache
+    else:
+        from mlx_lm.models.cache import KVCache
+        cache = KVCache()
+        n_steps = (prefill_len + KVCache.step - 1) // KVCache.step
+        alloc_len = n_steps * KVCache.step
+        n_kv = config.num_key_value_heads
+        hd = config.head_dim
+        cache.keys = mx.random.normal((1, n_kv, alloc_len, hd)).astype(mx.bfloat16)
+        cache.values = mx.random.normal((1, n_kv, alloc_len, hd)).astype(mx.bfloat16)
+        cache.offset = prefill_len
+        return cache
+
+
+def build_model(n_experts=16, n_layers=1, top_k=4,
+                hidden_size=4096, moe_intermediate_size=1024,
+                shared_expert_intermediate_size=2048, tp=1,
+                n_attn_heads=32, n_kv_heads=2,
+                lin_v_heads=64, lin_k_heads=16,
+                head_dim=256):
+    """Build Qwen3.5 MoE decoder layers with 8-bit gs=64 quantization.
+
+    Matches real mlx-community Qwen3.5 quantization:
+      Everything 8-bit gs=64 (gate, experts, shared expert, shared_expert_gate, attention/SSM).
+      RMSNorm weights: bf16.
+
+    Default dimensions are for Qwen3.5-397B-A17B. For 35B-A3B, pass:
+      n_attn_heads=16, lin_v_heads=32, hidden_size=2048
+
+    Uses qwen3_5.DecoderLayer (same class as real model) with hybrid attention:
+      3/4 GatedDeltaNet (SSM-like), 1/4 full attention (full_attention_interval=4).
+
+    TP=2 halves all column-parallel dimensions.
+
+    Returns:
+        layers: list of DecoderLayer instances
+        config: TextModelArgs
+        GROUP_SIZE: int (64)
+    """
+    GROUP_SIZE = 64
+    BITS = 8
+
+    # Apply TP sharding
+    moe_inter = moe_intermediate_size // tp
+    shared_inter = shared_expert_intermediate_size // tp
+    n_attn_heads_tp = n_attn_heads // tp
+    n_kv_heads_tp = max(1, n_kv_heads // tp)
+    lin_v_heads_tp = lin_v_heads // tp
+    lin_k_heads_tp = lin_k_heads // tp
+
+    config = TextModelArgs(
+        model_type="qwen3_5_moe",
+        hidden_size=hidden_size,
+        num_hidden_layers=n_layers,
+        intermediate_size=moe_inter,
+        num_attention_heads=n_attn_heads_tp,
+        num_key_value_heads=n_kv_heads_tp,
+        linear_num_value_heads=lin_v_heads_tp,
+        linear_num_key_heads=lin_k_heads_tp,
+        linear_key_head_dim=128,
+        linear_value_head_dim=128,
+        linear_conv_kernel_dim=4,
+        num_experts=n_experts,
+        num_experts_per_tok=top_k,
+        decoder_sparse_step=1,
+        shared_expert_intermediate_size=shared_inter,
+        moe_intermediate_size=moe_inter,
+        norm_topk_prob=True,
+        rms_norm_eps=1e-6,
+        vocab_size=248320,
+        head_dim=head_dim,
+        full_attention_interval=4,
+        max_position_embeddings=262144,
+        rope_theta=10000000,
+        partial_rotary_factor=0.25,
+        rope_parameters={
+            "type": "default",
+            "rope_theta": 10000000,
+            "partial_rotary_factor": 0.25,
+        },
+    )
+
+    tp_str = f" (TP={tp})" if tp > 1 else ""
+    print(f"  Config: {n_layers} layer(s), {n_experts} experts, top_k={top_k}, "
+          f"hidden={hidden_size}, inter={moe_inter}, shared={shared_inter}{tp_str}")
+    print(f"  Quant: {BITS}-bit gs={GROUP_SIZE} (all weights)")
+
+    layers = [DecoderLayer(config, idx) for idx in range(n_layers)]
+
+    for li, layer in enumerate(layers):
+        # Cast all attention/SSM params to bf16 before quantizing, matching
+        # real safetensors model where all non-quantized params are bf16.
+        # nn.quantize only touches nn.Linear; other params (conv1d, dt_bias,
+        # norm, A_log) must be cast manually.
+        attn_mod = layer.linear_attn if layer.is_linear else layer.self_attn
+        for name, mod in attn_mod.named_modules():
+            if isinstance(mod, nn.Linear):
+                mod.weight = mod.weight.astype(mx.bfloat16)
+            elif isinstance(mod, nn.Conv1d):
+                mod.weight = mod.weight.astype(mx.bfloat16)
+        # Cast leaf parameters (dt_bias, norm.weight, q/k_norm) to bf16
+        # A_log stays f32 (matches real model)
+        if layer.is_linear:
+            gdn = layer.linear_attn
+            gdn.dt_bias = gdn.dt_bias.astype(mx.bfloat16)
+            gdn.norm.weight = gdn.norm.weight.astype(mx.bfloat16)
+        else:
+            gqa = layer.self_attn
+            gqa.q_norm.weight = gqa.q_norm.weight.astype(mx.bfloat16)
+            gqa.k_norm.weight = gqa.k_norm.weight.astype(mx.bfloat16)
+        nn.quantize(attn_mod, bits=BITS, group_size=GROUP_SIZE)
+        mx.eval(attn_mod.parameters())
+
+        # RMSNorm to bf16 (norms are never quantized)
+        layer.input_layernorm.weight = layer.input_layernorm.weight.astype(mx.bfloat16)
+        layer.post_attention_layernorm.weight = layer.post_attention_layernorm.weight.astype(mx.bfloat16)
+        mx.eval(layer.input_layernorm.weight, layer.post_attention_layernorm.weight)
+
+        # MoE block: quantize everything to 8-bit gs=64
+        moe = layer.mlp
+        if isinstance(moe, Qwen3NextSparseMoeBlock):
+            # Gate: random init (zeros get optimized away), then quantize.
+            # nn.quantize on a leaf nn.Linear is a no-op (walks children, finds none).
+            # Use QuantizedLinear.from_linear directly.
+            moe.gate.weight = (
+                mx.random.normal(moe.gate.weight.shape) * 0.01
+            ).astype(mx.float32)
+            moe.gate = nn.QuantizedLinear.from_linear(
+                moe.gate, group_size=GROUP_SIZE, bits=BITS)
+            mx.eval(moe.gate.parameters())
+
+            # Routed experts: quantize per-projection to limit peak memory
+            nn.quantize(moe.switch_mlp, bits=BITS, group_size=GROUP_SIZE)
+            mx.eval(moe.switch_mlp.gate_proj.parameters())
+            mx.eval(moe.switch_mlp.up_proj.parameters())
+            mx.eval(moe.switch_mlp.down_proj.parameters())
+
+            # Shared expert: quantize to 8-bit (matching real model)
+            nn.quantize(moe.shared_expert, bits=BITS, group_size=GROUP_SIZE)
+            mx.eval(moe.shared_expert.parameters())
+
+            # shared_expert_gate: quantize to 8-bit gs=64 (leaf nn.Linear fix)
+            moe.shared_expert_gate = nn.QuantizedLinear.from_linear(
+                moe.shared_expert_gate, group_size=GROUP_SIZE, bits=BITS)
+            mx.eval(moe.shared_expert_gate.parameters())
+
+        if (li + 1) % 10 == 0 or li == 0:
+            print(f"  Layer {li+1}/{n_layers} ready")
+
+    return layers, config, GROUP_SIZE

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

@@ -0,0 +1,199 @@
+"""Decoder layer __call__ variants for Qwen3.5.
+
+Two modes:
+  _fused_decoder_call: passes residual to fused MoE epilogue (~15 dispatches)
+  _oproj_decoder_call: fuses o_proj + RMSNorm + gate GEMV (4 dispatches)
+
+Attention patches for oproj mode:
+  _pre_oproj_attention_call: Qwen3NextAttention.__call__ that skips o_proj
+  _pre_oproj_qwen35_linear_attn_call: qwen3_5.GatedDeltaNet.__call__ that skips out_proj
+
+Note: qwen3_5.GatedDeltaNet (used by DecoderLayer) is a DIFFERENT class from
+qwen3_next.Qwen3NextGatedDeltaNet. They have different projection layouts:
+  - qwen3_5.GatedDeltaNet: separate in_proj_qkv, in_proj_z, in_proj_b, in_proj_a
+  - qwen3_next.Qwen3NextGatedDeltaNet: merged in_proj_qkvz, in_proj_ba
+The patch must match qwen3_5.GatedDeltaNet's __call__ structure.
+"""
+
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.nn.layers.activations import silu as nn_silu
+
+# Map moe block id → parent decoder layer (avoids circular refs in model tree)
+_parent_layer_map = {}
+
+
+def _fused_decoder_call(self, x, mask=None, cache=None):
+    """Decoder layer with residual passed to fused MoE epilogue.
+
+    Replaces:
+      h = x + attn(norm(x))
+      out = h + mlp(norm(h))        # mlp returns MoE output, then adds h
+
+    With:
+      h = x + attn(norm(x))
+      out = mlp(norm(h), _residual=h)  # epilogue fuses: moe_out + h
+    """
+    if self.is_linear:
+        r = self.linear_attn(self.input_layernorm(x), mask, cache)
+    else:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+    h = x + r
+    out = self.mlp(self.post_attention_layernorm(h), _residual=h)
+    return out  # already includes residual add from epilogue
+
+
+def _oproj_decoder_call(self, x, mask=None, cache=None):
+    """Decoder with fused o_proj + RMSNorm + gate GEMV (oproj 4-dispatch mode).
+
+    Skips o_proj, addmm, and post_attention_layernorm — all fused into Dispatch 1.
+    Attention __call__ is patched to return pre-o_proj output.
+
+    Flow:
+      pre_oproj = attn(input_layernorm(x))   # returns BEFORE o_proj
+      MoE receives (pre_oproj, residual=x) and handles o_proj + RMSNorm + gate internally
+    """
+    if self.is_linear:
+        pre_oproj = self.linear_attn(self.input_layernorm(x), mask, cache)
+    else:
+        pre_oproj = self.self_attn(self.input_layernorm(x), mask, cache)
+    _parent_layer_map[id(self.mlp)] = self
+    return self.mlp(pre_oproj, _residual=x)
+
+
+def _vanilla_decoder_call(self, x, mask=None, cache=None):
+    """Original vanilla DecoderLayer.__call__ (fallback for B>8 or S>1)."""
+    if self.is_linear:
+        r = self.linear_attn(self.input_layernorm(x), mask, cache)
+    else:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+    h = x + r
+    out = self.mlp(self.post_attention_layernorm(h))
+    return h + out
+
+
+def _fused_gdn_decoder_call(self, x, mask=None, cache=None):
+    """Decoder with batched fused kernels. Falls back to vanilla for B>8 or S>1.
+
+    When fused: attention returns pre-out_proj output, MoE handles oproj + gate + experts.
+    When vanilla: original DecoderLayer flow (attention + residual + layernorm + MoE).
+    """
+    B = x.shape[0]
+    S = x.shape[1]
+
+    # Full vanilla fallback for large batch or prefill
+    if B > 8 or S > 1:
+        return _vanilla_decoder_call(self, x, mask, cache)
+
+    # Fused path: attention returns pre-oproj, MoE handles the rest
+    if self.is_linear:
+        pre_oproj = self.linear_attn(self.input_layernorm(x), mask, cache)
+    else:
+        pre_oproj = self.self_attn(self.input_layernorm(x), mask, cache)
+    _parent_layer_map[id(self.mlp)] = self
+    return self.mlp(pre_oproj, _residual=x)
+
+
+def _pre_oproj_attention_call(self, x, mask=None, cache=None):
+    """Qwen3NextAttention.__call__ that returns pre-o_proj output.
+
+    Identical to original except final line returns output*sigmoid(gate)
+    instead of self.o_proj(output*sigmoid(gate)).
+    """
+    B, L, D = x.shape
+    q_proj_output = self.q_proj(x)
+    queries, gate = mx.split(
+        q_proj_output.reshape(B, L, self.num_attention_heads, -1), 2, axis=-1
+    )
+    gate = gate.reshape(B, L, -1)
+    keys, values = self.k_proj(x), self.v_proj(x)
+    queries = self.q_norm(queries).transpose(0, 2, 1, 3)
+    keys = self.k_norm(
+        keys.reshape(B, L, self.num_key_value_heads, -1)
+    ).transpose(0, 2, 1, 3)
+    values = values.reshape(B, L, self.num_key_value_heads, -1).transpose(
+        0, 2, 1, 3
+    )
+    if cache is not None:
+        queries = self.rope(queries, offset=cache.offset)
+        keys = self.rope(keys, offset=cache.offset)
+        keys, values = cache.update_and_fetch(keys, values)
+    else:
+        queries = self.rope(queries)
+        keys = self.rope(keys)
+
+    from mlx_lm.models.qwen3_next import scaled_dot_product_attention
+    output = scaled_dot_product_attention(
+        queries, keys, values, cache=cache, scale=self.scale, mask=mask
+    )
+    output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+    return output * mx.sigmoid(gate)  # skip o_proj
+
+
+def _pre_oproj_qwen35_linear_attn_call(
+    self,
+    inputs: mx.array,
+    mask: Optional[mx.array] = None,
+    cache: Optional[Any] = None,
+) -> mx.array:
+    """qwen3_5.GatedDeltaNet.__call__ that returns pre-out_proj output.
+
+    Identical to qwen3_5.GatedDeltaNet.__call__ except final line returns
+    out.reshape(B,S,-1) instead of self.out_proj(out.reshape(B,S,-1)).
+
+    Note: this targets qwen3_5.GatedDeltaNet (separate projections), NOT
+    qwen3_next.Qwen3NextGatedDeltaNet (merged projections). They are
+    different classes with different __call__ bodies.
+    """
+    from mlx_lm.models.gated_delta import gated_delta_update
+
+    B, S, _ = inputs.shape
+
+    qkv = self.in_proj_qkv(inputs)
+    z = self.in_proj_z(inputs).reshape(B, S, self.num_v_heads, self.head_v_dim)
+    b = self.in_proj_b(inputs)
+    a = self.in_proj_a(inputs)
+
+    if cache is not None and cache[0] is not None:
+        conv_state = cache[0]
+    else:
+        conv_state = mx.zeros(
+            (B, self.conv_kernel_size - 1, self.conv_dim),
+            dtype=inputs.dtype,
+        )
+
+    if mask is not None:
+        qkv = mx.where(mask[..., None], qkv, 0)
+    conv_input = mx.concatenate([conv_state, qkv], axis=1)
+    if cache is not None:
+        cache[0] = conv_input[:, -(self.conv_kernel_size - 1) :]
+    conv_out = nn.silu(self.conv1d(conv_input))
+
+    q, k, v = [
+        t.reshape(B, S, h, d)
+        for t, h, d in zip(
+            mx.split(conv_out, [self.key_dim, 2 * self.key_dim], -1),
+            [self.num_k_heads, self.num_k_heads, self.num_v_heads],
+            [self.head_k_dim, self.head_k_dim, self.head_v_dim],
+        )
+    ]
+
+    state = cache[1] if cache else None
+    inv_scale = k.shape[-1] ** -0.5
+    q = (inv_scale**2) * mx.fast.rms_norm(q, None, 1e-6)
+    k = inv_scale * mx.fast.rms_norm(k, None, 1e-6)
+
+    out, state = gated_delta_update(
+        q, k, v, a, b,
+        self.A_log, self.dt_bias,
+        state, mask,
+        use_kernel=True,
+    )
+
+    if cache is not None:
+        cache[1] = state
+
+    out = self.norm(out, z)
+    return out.reshape(B, S, -1)  # skip out_proj

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

@@ -0,0 +1,161 @@
+"""Fused GDN attention __call__ for qwen3_5.GatedDeltaNet (Dispatches 2-5).
+
+Replaces the vanilla GatedDeltaNet.__call__ with fused kernel dispatches:
+  Dispatch 2: fused_gdn_projections — merged 8-bit GEMV + conv1d + SiLU(qkv) + SiLU(z)
+              + sigmoid(b)→beta + g=exp(-exp(A_log)*softplus(a+dt_bias))
+  Dispatch 3: fused_qk_rmsnorm — per-head L2-norm on q (×Dk^(-½)) and k
+  Dispatch 4: gated_delta_kernel — GDN recurrence (receives pre-computed g, beta)
+  Dispatch 5: fused_rms_norm_gated — RMSNorm(out, weight) × z_silu
+
+All 4 projection weights are pre-merged into contiguous buffers at patch time
+(_patch_gdn_proj_weights) for better memory locality.
+
+g/beta computation is fused into Dispatch 2 epilogues, eliminating ~8 micro-
+dispatches that gated_delta_update would otherwise generate.
+
+Fused path is decode-only (S=1). For prefill (S>1), falls back to vanilla ops.
+
+Returns pre-out_proj output (same interface as _pre_oproj_qwen35_linear_attn_call).
+Dispatch 1 (input_layernorm) is handled by the decoder.
+Dispatch 6 (oproj_gate_gemv) is handled by the MoE __call__.
+"""
+
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .kernels.batched_fused_gdn_projections_8bit import batched_fused_gdn_projections as fused_gdn_projections
+from .kernels.fused_qk_rmsnorm import fused_qk_rmsnorm
+from .kernels.fused_rms_norm_gated import fused_rms_norm_gated
+
+
+def _vanilla_gdn_call(self, inputs, mask, cache):
+    """Vanilla GDN path for prefill (S>1). Returns pre-out_proj output."""
+    from mlx_lm.models.gated_delta import gated_delta_update
+
+    B, S, _ = inputs.shape
+
+    qkv = self.in_proj_qkv(inputs)
+    z = self.in_proj_z(inputs).reshape(B, S, self.num_v_heads, self.head_v_dim)
+    b = self.in_proj_b(inputs)
+    a = self.in_proj_a(inputs)
+
+    if cache is not None and cache[0] is not None:
+        conv_state = cache[0]
+    else:
+        conv_state = mx.zeros(
+            (B, self.conv_kernel_size - 1, self.conv_dim),
+            dtype=inputs.dtype,
+        )
+
+    if mask is not None:
+        qkv = mx.where(mask[..., None], qkv, 0)
+    conv_input = mx.concatenate([conv_state, qkv], axis=1)
+    if cache is not None:
+        cache[0] = conv_input[:, -(self.conv_kernel_size - 1):]
+    conv_out = nn.silu(self.conv1d(conv_input))
+
+    q, k, v = [
+        t.reshape(B, S, h, d)
+        for t, h, d in zip(
+            mx.split(conv_out, [self.key_dim, 2 * self.key_dim], -1),
+            [self.num_k_heads, self.num_k_heads, self.num_v_heads],
+            [self.head_k_dim, self.head_k_dim, self.head_v_dim],
+        )
+    ]
+
+    state = cache[1] if cache else None
+    inv_scale = k.shape[-1] ** -0.5
+    q = inv_scale * q * mx.rsqrt(
+        (q * q).sum(axis=-1, keepdims=True) + 1e-6
+    )
+    k = k * mx.rsqrt(
+        (k * k).sum(axis=-1, keepdims=True) + 1e-6
+    )
+
+    out, state = gated_delta_update(
+        q, k, v, a, b,
+        self.A_log, self.dt_bias,
+        state, mask,
+        use_kernel=True,
+    )
+
+    if cache is not None:
+        cache[1] = state
+
+    out = self.norm(out, z)
+    return self.out_proj(out.reshape(B, S, -1))  # include out_proj for vanilla decoder
+
+
+def _fused_gdn_call(
+    self,
+    inputs: mx.array,
+    mask: Optional[mx.array] = None,
+    cache: Optional[Any] = None,
+) -> mx.array:
+    """Fused GDN attention: merged projections + existing GDN kernel.
+
+    Decode (S=1): uses fused kernels with merged weight buffers.
+    Prefill (S>1): falls back to vanilla ops.
+
+    Returns pre-out_proj output [B, S, value_dim] for Dispatch 6.
+    """
+    B, S, _ = inputs.shape
+
+    # Vanilla fallback: fused kernels are decode-only (S=1, B<=8)
+    if S > 1 or B > 8:
+        return _vanilla_gdn_call(self, inputs, mask, cache)
+
+    from mlx_lm.models.gated_delta import gated_delta_kernel
+
+    # ── Cache: conv state ──
+    if cache is not None and cache[0] is not None:
+        conv_state = cache[0]
+    else:
+        conv_state = mx.zeros(
+            (B, self.conv_kernel_size - 1, self.conv_dim),
+            dtype=inputs.dtype,
+        )
+
+    # ── Dispatch 2: fused projections (merged GEMV + conv + SiLU + g/beta) ──
+    qkv_conv_silu, z_silu, beta, g, conv_state_out = fused_gdn_projections(
+        inputs,
+        self._merged_proj_w, self._merged_proj_s, self._merged_proj_b,
+        self._merged_proj_dims,
+        conv_state, self.conv1d.weight,
+        self.A_log, self.dt_bias,
+        batch_size=B,
+    )
+
+    if cache is not None:
+        cache[0] = conv_state_out
+
+    # ── Dispatch 3: fused Q/K L2-norm ──
+    qk_normed = fused_qk_rmsnorm(qkv_conv_silu, batch_size=B)
+
+    # ── Split q, k from normed output; v from conv output ──
+    q = qk_normed[:, :, :self.key_dim].reshape(B, S, self.num_k_heads, self.head_k_dim)
+    k = qk_normed[:, :, self.key_dim:].reshape(B, S, self.num_k_heads, self.head_k_dim)
+    v = qkv_conv_silu[:, :, 2 * self.key_dim:].reshape(B, S, self.num_v_heads, self.head_v_dim)
+
+    # ── Dispatch 4: GDN recurrence with pre-computed g/beta ──
+    state = cache[1] if cache else None
+    if state is None:
+        state = mx.zeros(
+            (B, self.num_v_heads, self.head_v_dim, self.head_k_dim),
+            dtype=inputs.dtype,
+        )
+
+    out, state_new = gated_delta_kernel(
+        q, k, v, g, beta, state, mask,
+    )
+
+    if cache is not None:
+        cache[1] = state_new
+
+    # ── Dispatch 5: fused RMSNorm × z_silu ──
+    norm_weight = self.norm.weight
+    result = fused_rms_norm_gated(out, z_silu, norm_weight, batch_size=B)
+
+    return result  # [B, S, value_dim] — skip out_proj (handled by Dispatch 6)

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

@@ -0,0 +1,265 @@
+"""Batched fused GDN projections for Qwen3.5-35B-A3B, batch_size 1..8.
+
+Register-level weight sharing: each TG loads weights once, computes B outputs.
+Adapts fused_gdn_projections_8bit with the same pattern as
+batched_fused_gqa_projections_8bit.
+
+4 regions with different epilogues:
+  - QKV: GEMV → conv1d(4-tap) → SiLU → bf16 + cache update
+  - Z:   GEMV → SiLU → f32
+  - B:   GEMV → sigmoid → f32 (beta for GDN kernel)
+  - A:   GEMV → g=exp(-exp(A_log)*softplus(a+dt_bias)) → f32
+
+All constants baked into Metal source. B unrolled at code-generation time.
+
+Grid: (32, total_tg * 2, 1), TG: (32, 2, 1)
+No grid z for batch — batch is handled in registers.
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_batched_fused_gdn_proj_source(K, N_QKV, N_Z, N_B, N_A, B, group_size=64):
+    gs = int(group_size)
+    sc_stride = 256 // gs
+    slid_div = gs // 8
+    N_TOTAL = N_QKV + N_Z + N_B + N_A
+    K_groups = K // gs
+    N_QKV_TG = ceil_div(N_QKV, 8)
+    N_Z_TG = ceil_div(N_Z, 8)
+    N_B_TG = ceil_div(N_B, 8)
+
+    # Per-batch x loading (B unrolled)
+    x_load = "\n".join(f"""
+        float x{b}_thread[VALUES_PER_THREAD]; float xsum{b} = 0;
+        for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+            float xi = float(x[{b} * K + x_base + i]); x{b}_thread[i] = xi; xsum{b} += xi;
+        }}""" for b in range(B))
+
+    # Per-batch dot product with weights in registers
+    qdot = "\n".join(f"""
+            float accum{b} = 0;
+            for (int i = 0; i < VALUES_PER_THREAD; i++) accum{b} += x{b}_thread[i] * w_vals[i];
+            result{b}[row] += s_val * accum{b} + xsum{b} * b_val;""" for b in range(B))
+
+    result_decls = " ".join(f"float result{b}[4] = {{0,0,0,0}};" for b in range(B))
+    simd_reduce = "\n    ".join(
+        f"for (int row = 0; row < 4; row++) result{b}[row] = simd_sum(result{b}[row]);" for b in range(B))
+
+    # QKV epilogue: conv1d + SiLU + cache update, per batch
+    qkv_write = "\n".join(f"""
+            if (slid < 4u && c < N_QKV) {{
+                float qkv_val = result{b}[slid];
+                long cs_base = (long){b} * 3 * conv_dim;
+                float s0 = float(conv_state[cs_base + 0 * conv_dim + c]);
+                float s1 = float(conv_state[cs_base + 1 * conv_dim + c]);
+                float s2 = float(conv_state[cs_base + 2 * conv_dim + c]);
+                float conv_out = float(conv_w[c * 4 + 0]) * s0
+                               + float(conv_w[c * 4 + 1]) * s1
+                               + float(conv_w[c * 4 + 2]) * s2
+                               + float(conv_w[c * 4 + 3]) * qkv_val;
+                float silu_out = conv_out / (1.0f + metal::exp(-conv_out));
+                conv_state_out[cs_base + 0 * conv_dim + c] = static_cast<bfloat16_t>(s1);
+                conv_state_out[cs_base + 1 * conv_dim + c] = static_cast<bfloat16_t>(s2);
+                conv_state_out[cs_base + 2 * conv_dim + c] = static_cast<bfloat16_t>(qkv_val);
+                qkv_out[{b} * conv_dim + c] = static_cast<bfloat16_t>(silu_out);
+            }}""" for b in range(B))
+
+    # Z epilogue: SiLU per batch
+    z_write = "\n".join(f"""
+            if (slid < 4u && z_row < N_Z) {{
+                float val = result{b}[slid];
+                z_silu_out[{b} * N_Z + z_row] = val / (1.0f + metal::exp(-val));
+            }}""" for b in range(B))
+
+    # B epilogue: sigmoid per batch
+    b_write = "\n".join(f"""
+            if (slid < 4u && b_row < N_B) {{
+                b_out[{b} * N_B + b_row] = 1.0f / (1.0f + metal::exp(-result{b}[slid]));
+            }}""" for b in range(B))
+
+    # A epilogue: g computation per batch
+    a_write = "\n".join(f"""
+            if (slid < 4u && a_row < N_A_val) {{
+                float a_val = result{b}[slid];
+                float dt = float(dt_bias_arr[a_row]);
+                float x_g = a_val + dt;
+                float sp = (x_g > 20.0f) ? x_g : metal::log(1.0f + metal::exp(x_g));
+                float g_val = metal::exp(-metal::exp(float(A_log_arr[a_row])) * sp);
+                a_out[{b} * N_A_val + a_row] = g_val;
+            }}""" for b in range(B))
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int GROUP_SIZE = {gs};
+    const int SC_STRIDE = {sc_stride};
+    const int SLID_DIV = {slid_div};
+    const int K = {K};
+    const int K_groups = {K_groups};
+    const int N_QKV = {N_QKV};
+    const int N_Z = {N_Z};
+    const int N_B = {N_B};
+    const int N_TOTAL = {N_TOTAL};
+    const int N_QKV_TG = {N_QKV_TG};
+    const int N_Z_TG = {N_Z_TG};
+    const int N_B_TG = {N_B_TG};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+    int tg = tgid.y;
+
+    int out_row, region;
+    if (tg < N_QKV_TG) {{
+        region = 0; out_row = tg * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_QKV_TG + N_Z_TG) {{
+        region = 1; out_row = N_QKV + (tg - N_QKV_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_QKV_TG + N_Z_TG + N_B_TG) {{
+        region = 2; out_row = N_QKV + N_Z + (tg - N_QKV_TG - N_Z_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else {{
+        region = 3; out_row = N_QKV + N_Z + N_B + (tg - N_QKV_TG - N_Z_TG - N_B_TG) * 8 + sgid * RESULTS_PER_SG;
+    }}
+    if (out_row >= N_TOTAL) return;
+
+    // Weight pointers (shared across all batch elements)
+    const device uint8_t* ws = (const device uint8_t*)W_merged + (long)out_row * K + slid * VALUES_PER_THREAD;
+    const device bfloat16_t* sc = (const device bfloat16_t*)S_merged + (long)out_row * K_groups + slid / SLID_DIV;
+    const device bfloat16_t* bi = (const device bfloat16_t*)B_merged + (long)out_row * K_groups + slid / SLID_DIV;
+
+    {result_decls}
+    int x_base = slid * VALUES_PER_THREAD;
+
+    // K-loop: load weights into registers once, compute {B} batch elements
+    for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{
+{x_load}
+
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            const device uint8_t* wl = ws + row * K;
+            float s_val = float(sc[row * K_groups]);
+            float b_val = float(bi[row * K_groups]);
+            float w_vals[VALUES_PER_THREAD];
+            for (int i = 0; i < VALUES_PER_THREAD; i++) w_vals[i] = float(wl[i]);
+{qdot}
+        }}
+
+        ws += BLOCK_SIZE; sc += SC_STRIDE; bi += SC_STRIDE; x_base += BLOCK_SIZE;
+    }}
+
+    {simd_reduce}
+
+    // Region-specific epilogues for all {B} batches
+    if (region == 0) {{
+        int c = out_row + (int)slid;
+        int conv_dim = N_QKV;
+{qkv_write}
+    }} else if (region == 1) {{
+        int z_row = out_row - N_QKV + (int)slid;
+{z_write}
+    }} else if (region == 2) {{
+        int b_row = out_row - N_QKV - N_Z + (int)slid;
+{b_write}
+    }} else {{
+        int a_row = out_row - N_QKV - N_Z - N_B + (int)slid;
+        int N_A_val = N_TOTAL - N_QKV - N_Z - N_B;
+{a_write}
+    }}
+"""
+
+
+_batched_gdn_proj_cache = {}
+
+
+def _get_batched_gdn_proj_kernel(K, N_QKV, N_Z, N_B, N_A, B, group_size=64):
+    key = (K, N_QKV, N_Z, N_B, N_A, B, group_size)
+    if key not in _batched_gdn_proj_cache:
+        _batched_gdn_proj_cache[key] = mx.fast.metal_kernel(
+            name=f"batched_fused_gdn_proj_K{K}_NQKV{N_QKV}_B{B}",
+            input_names=[
+                "x",
+                "W_merged", "S_merged", "B_merged",
+                "conv_state", "conv_w",
+                "A_log_arr", "dt_bias_arr",
+            ],
+            output_names=["qkv_out", "z_silu_out", "b_out", "a_out", "conv_state_out"],
+            source=_gen_batched_fused_gdn_proj_source(K, N_QKV, N_Z, N_B, N_A, B, group_size),
+        )
+    return _batched_gdn_proj_cache[key]
+
+
+def batched_fused_gdn_projections(
+    x,
+    W_merged, S_merged, B_merged,
+    proj_dims,
+    conv_state, conv_weights,
+    A_log, dt_bias,
+    batch_size=1,
+):
+    """Batched fused GDN projections with register-level weight sharing.
+
+    Same as fused_gdn_projections but loads weights once per TG and computes
+    B outputs from registers. No grid z for batch.
+
+    Args:
+        x: [B, 1, K] bf16 — post-RMSNorm hidden state
+        W_merged, S_merged, B_merged: merged quantized weights
+        proj_dims: (N_QKV, N_Z, N_B, N_A)
+        conv_state: [B, 3, conv_dim] bf16
+        conv_weights: [conv_dim, 4, 1] or [conv_dim, 4] bf16
+        A_log: [Hv] f32, dt_bias: [Hv] f32
+        batch_size: int (1..8)
+
+    Returns:
+        qkv_conv_silu: [B, 1, N_QKV] bf16
+        z_silu:        [B, 1, N_Z] f32
+        beta:          [B, 1, N_B] f32
+        g:             [B, 1, N_A] f32
+        conv_state_out: [B, 3, N_QKV] bf16
+    """
+    B = batch_size
+
+    N_QKV, N_Z, N_B, N_A = proj_dims
+    K = x.shape[-1]
+
+    kern = _get_batched_gdn_proj_kernel(K, N_QKV, N_Z, N_B, N_A, B)
+
+    N_QKV_TG = ceil_div(N_QKV, 8)
+    N_Z_TG = ceil_div(N_Z, 8)
+    N_B_TG = ceil_div(N_B, 8)
+    N_A_TG = ceil_div(N_A, 8)
+    total_tg = N_QKV_TG + N_Z_TG + N_B_TG + N_A_TG
+
+    conv_w_flat = conv_weights.reshape(-1, 4) if conv_weights.ndim == 3 else conv_weights
+    x_flat = x.reshape(B, K)
+
+    results = kern(
+        inputs=[
+            x_flat,
+            W_merged, S_merged, B_merged,
+            conv_state, conv_w_flat,
+            A_log, dt_bias,
+        ],
+        output_shapes=[
+            (B * N_QKV,),
+            (B * N_Z,),
+            (B * N_B,),
+            (B * N_A,),
+            (B * 3 * N_QKV,),
+        ],
+        output_dtypes=[mx.bfloat16, mx.float32, mx.float32, mx.float32, mx.bfloat16],
+        grid=(32, total_tg * 2, 1),   # No grid z — batch in registers
+        threadgroup=(32, 2, 1),
+    )
+
+    qkv_out = results[0].reshape(B, 1, N_QKV)
+    z_silu = results[1].reshape(B, 1, N_Z)
+    beta = results[2].reshape(B, 1, N_B)
+    g = results[3].reshape(B, 1, N_A)
+    conv_state_out = results[4].reshape(B, 3, N_QKV)
+
+    return qkv_out, z_silu, beta, g, conv_state_out

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

@@ -0,0 +1,205 @@
+"""Batched fused GQA projections (Dispatch 1) for batch_size 1..8.
+
+Adapts fused_gqa_projections_8bit for B>1 with register-level weight sharing.
+Each TG loads weights once, computes B outputs from registers.
+
+4 regions with different epilogues (same as B=1):
+  - Queries: GEMV → raw bf16
+  - Gate: GEMV → sigmoid → f32
+  - Keys: GEMV → raw bf16
+  - Values: GEMV → raw bf16
+
+All constants baked into Metal source. B unrolled at code-generation time.
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_batched_fused_gqa_proj_source(K, N_Q, N_GATE, N_K, N_V, B, group_size=64):
+    gs = int(group_size)
+    sc_stride = 256 // gs
+    slid_div = gs // 8
+    N_TOTAL = N_Q + N_GATE + N_K + N_V
+    K_groups = K // gs
+    N_Q_TG = ceil_div(N_Q, 8)
+    N_GATE_TG = ceil_div(N_GATE, 8)
+    N_K_TG = ceil_div(N_K, 8)
+
+    # Per-batch x loading
+    x_load = "\n".join(f"""
+        float x{b}_thread[VALUES_PER_THREAD]; float xsum{b} = 0;
+        for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+            float xi = float(x[{b} * K + x_base + i]); x{b}_thread[i] = xi; xsum{b} += xi;
+        }}""" for b in range(B))
+
+    # Per-batch qdot (weights in registers)
+    qdot = "\n".join(f"""
+            float accum{b} = 0;
+            for (int i = 0; i < VALUES_PER_THREAD; i++) accum{b} += x{b}_thread[i] * w_vals[i];
+            result{b}[row] += s_val * accum{b} + xsum{b} * b_val;""" for b in range(B))
+
+    result_decls = " ".join(f"float result{b}[4] = {{0,0,0,0}};" for b in range(B))
+    simd_reduce = "\n    ".join(
+        f"for (int row = 0; row < 4; row++) result{b}[row] = simd_sum(result{b}[row]);" for b in range(B))
+
+    # Queries epilogue (bf16 write per batch)
+    q_write = "\n".join(f"""
+            if (slid < 4u && q_row < N_Q) q_out[{b} * N_Q + q_row] = static_cast<bfloat16_t>(result{b}[slid]);"""
+        for b in range(B))
+
+    # Gate epilogue (sigmoid → f32 per batch)
+    gate_write = "\n".join(f"""
+            if (slid < 4u && g_row < N_GATE) {{
+                float sig{b} = 1.0f / (1.0f + metal::exp(-result{b}[slid]));
+                gate_out[{b} * N_GATE + g_row] = sig{b};
+            }}""" for b in range(B))
+
+    # Keys epilogue
+    k_write = "\n".join(f"""
+            if (slid < 4u && k_row < N_K) k_out[{b} * N_K + k_row] = static_cast<bfloat16_t>(result{b}[slid]);"""
+        for b in range(B))
+
+    # Values epilogue
+    v_write = "\n".join(f"""
+            if (slid < 4u && v_row < N_V) v_out[{b} * N_V + v_row] = static_cast<bfloat16_t>(result{b}[slid]);"""
+        for b in range(B))
+
+    N_V_val = N_TOTAL - N_Q - N_GATE - N_K
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int GROUP_SIZE = {gs};
+    const int SC_STRIDE = {sc_stride};
+    const int SLID_DIV = {slid_div};
+    const int K = {K};
+    const int K_groups = {K_groups};
+    const int N_Q = {N_Q};
+    const int N_GATE = {N_GATE};
+    const int N_K = {N_K};
+    const int N_V = {N_V_val};
+    const int N_TOTAL = {N_TOTAL};
+    const int N_Q_TG = {N_Q_TG};
+    const int N_GATE_TG = {N_GATE_TG};
+    const int N_K_TG = {N_K_TG};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+    int b_idx = tgid.z;
+    int tg = tgid.y;
+
+    int out_row, region;
+    if (tg < N_Q_TG) {{
+        region = 0; out_row = tg * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_Q_TG + N_GATE_TG) {{
+        region = 1; out_row = N_Q + (tg - N_Q_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_Q_TG + N_GATE_TG + N_K_TG) {{
+        region = 2; out_row = N_Q + N_GATE + (tg - N_Q_TG - N_GATE_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else {{
+        region = 3; out_row = N_Q + N_GATE + N_K + (tg - N_Q_TG - N_GATE_TG - N_K_TG) * 8 + sgid * RESULTS_PER_SG;
+    }}
+    if (out_row >= N_TOTAL) return;
+
+    // Weight pointers (shared across all batch elements)
+    const device uint8_t* ws = (const device uint8_t*)W_merged + (long)out_row * K + slid * VALUES_PER_THREAD;
+    const device bfloat16_t* sc = (const device bfloat16_t*)S_merged + (long)out_row * K_groups + slid / SLID_DIV;
+    const device bfloat16_t* bi = (const device bfloat16_t*)B_merged + (long)out_row * K_groups + slid / SLID_DIV;
+
+    {result_decls}
+    int x_base = slid * VALUES_PER_THREAD;
+
+    // K-loop: load weights once, compute {B} batch elements
+    for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{
+{x_load}
+
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            const device uint8_t* wl = ws + row * K;
+            float s_val = float(sc[row * K_groups]);
+            float b_val = float(bi[row * K_groups]);
+            float w_vals[VALUES_PER_THREAD];
+            for (int i = 0; i < VALUES_PER_THREAD; i++) w_vals[i] = float(wl[i]);
+{qdot}
+        }}
+
+        ws += BLOCK_SIZE; sc += SC_STRIDE; bi += SC_STRIDE; x_base += BLOCK_SIZE;
+    }}
+
+    {simd_reduce}
+
+    // Region-specific epilogues for all {B} batches
+    if (region == 0) {{
+        int q_row = out_row + (int)slid;
+{q_write}
+    }} else if (region == 1) {{
+        int g_row = out_row - N_Q + (int)slid;
+{gate_write}
+    }} else if (region == 2) {{
+        int k_row = out_row - N_Q - N_GATE + (int)slid;
+{k_write}
+    }} else {{
+        int v_row = out_row - N_Q - N_GATE - N_K + (int)slid;
+{v_write}
+    }}
+"""
+
+
+_batched_proj_cache = {}
+
+
+def _get_batched_proj_kernel(K, N_Q, N_GATE, N_K, N_V, B, group_size=64):
+    key = (K, N_Q, N_GATE, N_K, N_V, B, group_size)
+    if key not in _batched_proj_cache:
+        _batched_proj_cache[key] = mx.fast.metal_kernel(
+            name=f"batched_fused_gqa_proj_K{K}_NQ{N_Q}_B{B}",
+            input_names=["x", "W_merged", "S_merged", "B_merged"],
+            output_names=["q_out", "gate_out", "k_out", "v_out"],
+            source=_gen_batched_fused_gqa_proj_source(K, N_Q, N_GATE, N_K, N_V, B, group_size),
+        )
+    return _batched_proj_cache[key]
+
+
+def batched_fused_gqa_projections(x, W_merged, S_merged, B_merged, proj_dims,
+                                   batch_size, total_tg=None):
+    """Batched fused GQA projections with register weight sharing.
+
+    Args:
+        x: [B, 1, K] bf16
+        W_merged, S_merged, B_merged: merged q+gate+k+v weights
+        proj_dims: (N_Q, N_GATE, N_K, N_V)
+        batch_size: B (1..8)
+
+    Returns:
+        queries (B, 1, N_Q) bf16, gate_sigmoid (B, 1, N_GATE) f32,
+        keys (B, 1, N_K) bf16, values (B, 1, N_V) bf16
+    """
+    B = batch_size
+    N_Q, N_GATE, N_K, N_V = proj_dims
+    K = x.shape[-1]
+
+    kern = _get_batched_proj_kernel(K, N_Q, N_GATE, N_K, N_V, B)
+
+    if total_tg is None:
+        total_tg = ceil_div(N_Q, 8) + ceil_div(N_GATE, 8) + ceil_div(N_K, 8) + ceil_div(N_V, 8)
+
+    x_flat = x.reshape(B, K)
+
+    results = kern(
+        inputs=[x_flat, W_merged, S_merged, B_merged],
+        output_shapes=[
+            (B * N_Q,), (B * N_GATE,), (B * N_K,), (B * N_V,),
+        ],
+        output_dtypes=[mx.bfloat16, mx.float32, mx.bfloat16, mx.bfloat16],
+        grid=(32, total_tg * 2, 1),
+        threadgroup=(32, 2, 1),
+    )
+
+    return (results[0].reshape(B, 1, N_Q),
+            results[1].reshape(B, 1, N_GATE),
+            results[2].reshape(B, 1, N_K),
+            results[3].reshape(B, 1, N_V))

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

@@ -0,0 +1,252 @@
+"""Batched merged 8-bit down_proj GEMV for Qwen3.5 routed + shared experts.
+
+Adapts merged_down_proj_8bit for batch_size B (1..8).
+
+Grid z-dimension: B * n_active + 1
+  - tgid.z < B * n_active: routed experts (one TG per batch×expert pair)
+    Same structure as affine_gather_qmv — each TG independently indexes into
+    expert weights via inds[flat_idx].
+  - tgid.z == B * n_active: shared expert (ONE TG handles ALL B batch elements
+    with register-level weight sharing — loads weights once, computes B outputs)
+
+All constants baked into Metal source (no scalar kernel inputs).
+"""
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_batched_merged_down_8bit_source(K_OUT, N_IN, SHARED_N_IN, n_active, B, group_size=64):
+    gs = int(group_size)
+    sc_stride = 256 // gs
+    slid_divisor = gs // 8
+    N_groups = N_IN // gs
+    SHARED_N_groups = SHARED_N_IN // gs
+    total_routed = B * n_active
+
+    # Shared expert: generate unrolled batch loops
+    shared_x_load_lines = []
+    for b in range(B):
+        shared_x_load_lines.append(f"""
+            float x{b}_thread[VALUES_PER_THREAD];
+            float xsum{b} = 0;
+            for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+                float xi = X_shared[{b} * SHARED_N_IN + x_base + i];
+                x{b}_thread[i] = xi;
+                xsum{b} += xi;
+            }}""")
+    shared_x_load = "\n".join(shared_x_load_lines)
+
+    shared_qdot_lines = []
+    for b in range(B):
+        shared_qdot_lines.append(f"""
+                float accum{b} = 0;
+                for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+                    accum{b} += x{b}_thread[i] * w_vals[i];
+                }}
+                result{b}[row] += s_val * accum{b} + xsum{b} * b_val;""")
+    shared_qdot = "\n".join(shared_qdot_lines)
+
+    shared_result_decls = "\n        ".join(
+        f"float result{b}[RESULTS_PER_SG] = {{0, 0, 0, 0}};"
+        for b in range(B)
+    )
+
+    shared_write_lines = []
+    for b in range(B):
+        shared_write_lines.append(f"""
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            float r{b} = simd_sum(result{b}[row]);
+            if (slid == 0) {{
+                Y_shared[{b} * K_OUT + out_row + row] = r{b};
+            }}
+        }}""")
+    shared_write = "\n".join(shared_write_lines)
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int K_OUT = {K_OUT};
+    const int N_IN = {N_IN};
+    const int SHARED_N_IN = {SHARED_N_IN};
+    const int N_GROUPS = {N_groups};
+    const int SHARED_N_GROUPS = {SHARED_N_groups};
+    const int N_ACTIVE = {n_active};
+    const int TOTAL_ROUTED = {total_routed};
+    const int BATCH_SIZE = {B};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+
+    int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG;
+    if (out_row >= K_OUT) return;
+
+    if (tgid.z < (uint)TOTAL_ROUTED) {{
+        // ═══════ ROUTED EXPERT PATH (same as gather_qmv) ═══════
+        // tgid.z indexes flat (batch, expert) pairs
+        int flat_idx = (int)tgid.z;
+        int expert = inds[flat_idx];
+
+        const device uint8_t* ws = (const device uint8_t*)W
+            + (long)expert * K_OUT * N_IN + out_row * N_IN + slid * VALUES_PER_THREAD;
+        const device bfloat16_t* sc = (const device bfloat16_t*)S
+            + (long)expert * K_OUT * N_GROUPS + out_row * N_GROUPS + slid / {slid_divisor};
+        const device bfloat16_t* bi = (const device bfloat16_t*)B_q
+            + (long)expert * K_OUT * N_GROUPS + out_row * N_GROUPS + slid / {slid_divisor};
+
+        const device float* x_ptr = (const device float*)X_routed
+            + flat_idx * N_IN;
+
+        int x_base = slid * VALUES_PER_THREAD;
+        float result[4] = {{0, 0, 0, 0}};
+
+        for (int k = 0; k < N_IN; k += BLOCK_SIZE) {{
+            float x_thread[8];
+            float xsum = 0;
+            for (int i = 0; i < 8; i++) {{
+                float xi = x_ptr[x_base + i];
+                x_thread[i] = xi;
+                xsum += xi;
+            }}
+
+            for (int row = 0; row < RESULTS_PER_SG; row++) {{
+                const device uint8_t* wl = ws + row * N_IN;
+                float s = float(sc[row * N_GROUPS]);
+                float b = float(bi[row * N_GROUPS]);
+                float accum = 0;
+                for (int i = 0; i < 8; i++) {{
+                    accum += x_thread[i] * float(wl[i]);
+                }}
+                result[row] += s * accum + xsum * b;
+            }}
+
+            ws += BLOCK_SIZE;
+            sc += {sc_stride};
+            bi += {sc_stride};
+            x_base += BLOCK_SIZE;
+        }}
+
+        device float* yp = Y_routed + flat_idx * K_OUT + out_row;
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            float r = simd_sum(result[row]);
+            if (slid == 0) {{
+                yp[row] = r;
+            }}
+        }}
+
+    }} else {{
+        // ═══════ SHARED EXPERT PATH (register-level weight sharing) ═══════
+        // ONE TG handles ALL {B} batch elements.
+        // Load shared expert weights once, compute {B} outputs from registers.
+
+        const device uint8_t* ws = (const device uint8_t*)W_shared_down
+            + (long)out_row * SHARED_N_IN + slid * VALUES_PER_THREAD;
+        const device bfloat16_t* sc = (const device bfloat16_t*)S_shared_down
+            + (long)out_row * SHARED_N_GROUPS + slid / {slid_divisor};
+        const device bfloat16_t* bi = (const device bfloat16_t*)B_shared_down
+            + (long)out_row * SHARED_N_GROUPS + slid / {slid_divisor};
+
+        int x_base = slid * VALUES_PER_THREAD;
+        {shared_result_decls}
+
+        for (int k = 0; k < SHARED_N_IN; k += BLOCK_SIZE) {{
+            // Load x for all {B} batch elements
+{shared_x_load}
+
+            // Load weights once into registers, compute all {B} batches
+            for (int row = 0; row < RESULTS_PER_SG; row++) {{
+                const device uint8_t* wl = ws + row * SHARED_N_IN;
+                float s_val = float(sc[row * SHARED_N_GROUPS]);
+                float b_val = float(bi[row * SHARED_N_GROUPS]);
+
+                float w_vals[VALUES_PER_THREAD];
+                for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+                    w_vals[i] = float(wl[i]);
+                }}
+
+{shared_qdot}
+            }}
+
+            ws += BLOCK_SIZE;
+            sc += {sc_stride};
+            bi += {sc_stride};
+            x_base += BLOCK_SIZE;
+        }}
+
+        // Write all {B} outputs
+{shared_write}
+    }}
+"""
+
+
+_batched_down_cache = {}
+
+
+def _get_batched_down_kernel(K_OUT, N_IN, SHARED_N_IN, n_active, B, group_size=64):
+    key = (K_OUT, N_IN, SHARED_N_IN, n_active, B, group_size)
+    if key not in _batched_down_cache:
+        _batched_down_cache[key] = mx.fast.metal_kernel(
+            name=f"batched_down_K{K_OUT}_N{N_IN}_SN{SHARED_N_IN}_na{n_active}_B{B}",
+            input_names=["W", "S", "B_q",
+                         "W_shared_down", "S_shared_down", "B_shared_down",
+                         "X_routed", "X_shared", "inds"],
+            output_names=["Y_routed", "Y_shared"],
+            source=_gen_batched_merged_down_8bit_source(
+                K_OUT, N_IN, SHARED_N_IN, n_active, B, group_size),
+        )
+    return _batched_down_cache[key]
+
+
+def batched_merged_down_proj_8bit(w_q, s, b_q,
+                                   w_shared_down, s_shared_down, b_shared_down,
+                                   x_routed, x_shared, inds,
+                                   k_out, n_in, batch_size,
+                                   n_active, group_size=64,
+                                   shared_n_in=None):
+    """Batched merged down_proj for 8-bit routed + shared experts.
+
+    Args:
+        w_q: routed weights (E, K_OUT, N_IN/4) uint32
+        s: routed scales (E, K_OUT, N_IN/gs) bf16
+        b_q: routed biases (E, K_OUT, N_IN/gs) bf16
+        w_shared_down: shared weight (K_OUT, SHARED_N_IN/4) uint32
+        s_shared_down: shared scales (K_OUT, SHARED_N_IN/gs) bf16
+        b_shared_down: shared biases (K_OUT, SHARED_N_IN/gs) bf16
+        x_routed: (B * n_active, N_IN) f32
+        x_shared: (B * SHARED_N_IN,) f32
+        inds: (B * n_active,) uint32
+        k_out: output dimension
+        n_in: routed input dimension
+        batch_size: B
+        n_active: experts per token (top_k)
+        shared_n_in: shared input dim (defaults to n_in)
+
+    Returns:
+        Y_routed: (B * n_active, k_out) f32
+        Y_shared: (B, k_out) f32
+    """
+    B = batch_size
+    k_out_val = int(k_out)
+    n_in_val = int(n_in)
+    shared_n_in_val = int(shared_n_in) if shared_n_in is not None else n_in_val
+
+    kern = _get_batched_down_kernel(k_out_val, n_in_val, shared_n_in_val, n_active, B)
+
+    y_groups = ceil_div(k_out_val, 8)
+    total_routed = B * n_active
+
+    Y = kern(
+        inputs=[w_q, s, b_q,
+                w_shared_down, s_shared_down, b_shared_down,
+                x_routed, x_shared, inds],
+        output_shapes=[(total_routed * k_out_val,), (B * k_out_val,)],
+        output_dtypes=[mx.float32, mx.float32],
+        grid=(32, y_groups * 2, total_routed + 1),
+        threadgroup=(32, 2, 1),
+    )
+
+    return Y[0].reshape(total_routed, k_out_val), Y[1].reshape(B, k_out_val)

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

@@ -0,0 +1,92 @@
+"""Batched MoE epilogue for Qwen3.5: weighted sum + shared expert gate + residual.
+
+Computes per batch element:
+  Y[b, j] = bf16( Σ_a(scores[b,a] * D_routed[b*n_active+a, j])
+                 + sigmoid(gate_raw[b]) * D_shared[b, j]
+                 + H[b, j] )
+
+Grid z = B (one set of threads per batch element).
+All constants baked into Metal source.
+"""
+import mlx.core as mx
+
+
+def _gen_batched_epilogue_source(K, n_active, B):
+    return f"""
+    const int K_const = {K};
+    const int n_active_const = {n_active};
+    const int B_const = {B};
+
+    uint tid = thread_position_in_grid.x;
+    uint batch_id = thread_position_in_grid.z;
+
+    if (tid >= K_const || batch_id >= B_const) return;
+
+    // Weighted sum of routed expert outputs for this batch element
+    float acc = 0.0f;
+    int routed_base = (int)batch_id * n_active_const * K_const;
+    int score_base = (int)batch_id * n_active_const;
+    for (int a = 0; a < n_active_const; a++) {{
+        acc += scores[score_base + a] * D_routed[routed_base + a * K_const + tid];
+    }}
+
+    // Shared expert: sigmoid(gate_raw) * D_shared
+    float gate_raw_val = gate_raw[(int)batch_id];
+    float gate = 1.0f / (1.0f + metal::exp(-gate_raw_val));
+    float shared_val = D_shared[(int)batch_id * K_const + tid] * gate;
+
+    // Add residual and write
+    Y[(int)batch_id * K_const + tid] = static_cast<bfloat16_t>(
+        acc + shared_val + float(H[(int)batch_id * K_const + tid])
+    );
+"""
+
+
+_batched_epilogue_cache = {}
+
+
+def _get_batched_epilogue_kernel(K, n_active, B):
+    key = (K, n_active, B)
+    if key not in _batched_epilogue_cache:
+        _batched_epilogue_cache[key] = mx.fast.metal_kernel(
+            name=f"batched_epilogue_K{K}_na{n_active}_B{B}",
+            input_names=["D_routed", "D_shared", "scores", "H", "gate_raw"],
+            output_names=["Y"],
+            source=_gen_batched_epilogue_source(K, n_active, B),
+        )
+    return _batched_epilogue_cache[key]
+
+
+def batched_moe_epilogue(d_routed, d_shared, scores, h, gate_raw,
+                          k_val, batch_size, n_active):
+    """Batched MoE epilogue with fused sigmoid.
+
+    Args:
+        d_routed: (B * n_active, K) f32
+        d_shared: (B, K) f32
+        scores: (B * n_active,) f32
+        h: (B, K) bf16 — residual
+        gate_raw: (B,) f32 — raw shared expert gate (pre-sigmoid)
+        k_val: hidden dimension
+        batch_size: B
+        n_active: experts per token
+
+    Returns:
+        Y: (B, K) bf16
+    """
+    K = int(k_val)
+    B = batch_size
+    kern = _get_batched_epilogue_kernel(K, n_active, B)
+
+    tg_size = min(K, 1024)
+    n_tg = (K + tg_size - 1) // tg_size
+
+    Y = kern(
+        inputs=[d_routed, d_shared.reshape(B * K), scores, h.reshape(B * K), gate_raw],
+        output_shapes=[(B * K,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(n_tg * tg_size, 1, B),
+        threadgroup=(tg_size, 1, 1),
+    )
+
+    return Y[0].reshape(B, K)

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

@@ -0,0 +1,332 @@
+"""Batched Dispatch 1: Fused o_proj (8-bit) + gate GEMV parts + x² partials.
+
+Adapts custom_oproj_gate_gemv_8bit for batch_size B (1..8).
+Register-level weight sharing: each TG loads weights once, computes B outputs.
+
+Three GEMV regions (same as B=1):
+  TGs 0..N_OPROJ_TG-1: o_proj GEMV (8-bit) + residual + h_scaled + x²
+  TGs N_OPROJ_TG..+N_M1_TG-1: M1 × attn_out → gate_part_a (bf16 GEMV)
+  TGs +N_M1_TG..end: W_fused × residual → gate_part_b (bf16 GEMV)
+
+All constants baked into Metal source. B is unrolled at code-generation time.
+"""
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_batched_oproj_source(n_experts, M, K_attn, K_hidden, B, group_size=64, gate_bm=8):
+    E = int(n_experts)
+    gs = group_size
+    oproj_slid_divisor = gs // 8
+    oproj_sc_stride = 256 // gs
+    blockM_gate = gate_bm * 4
+    n_m1_tg = ceil_div(E, blockM_gate)
+
+    # Generate unrolled per-batch code for o_proj epilogue
+    oproj_epilogue = []
+    for b in range(B):
+        oproj_epilogue.append(f"""
+            float x2_acc{b} = 0.0f;
+            for (int tm = 0; tm < TM; tm++) {{
+                int k = out_row + tm;
+                float h{b} = result{b}[tm] + float(residual[{b} * M_DIM + k]);
+                x2_acc{b} += h{b} * h{b};
+                h_scaled[{b} * M_DIM + k] = static_cast<bfloat16_t>(h{b} * float(w_rms[k]));
+                h_out[{b} * M_DIM + k] = static_cast<bfloat16_t>(h{b});
+            }}""")
+    oproj_epilogue_code = "\n".join(oproj_epilogue)
+
+    oproj_x2_write = []
+    for b in range(B):
+        oproj_x2_write.append(f"""
+                total{b} += tgp_x2[s * {B} + {b}];""")
+    oproj_x2_sum = "\n".join(oproj_x2_write)
+
+    oproj_x2_final = []
+    for b in range(B):
+        oproj_x2_final.append(f"""
+                x2_partials[{b} * N_OPROJ_TG_DIM + tg_x] = total{b};""")
+    oproj_x2_final_code = "\n".join(oproj_x2_final)
+
+    # o_proj K-loop: load weights once, compute B batch elements
+    oproj_x_load = "\n".join(f"""
+            float xv{b}[VPT]; float xsum{b} = 0.0f;
+            for (int i = 0; i < VPT; i++) {{ xv{b}[i] = float(attn_out[{b} * K_ATTN_DIM + xb + i]); xsum{b} += xv{b}[i]; }}""" for b in range(B))
+
+    oproj_qdot = "\n".join(f"""
+                acc{b}[row] += s_val * wdot(xv{b}, w_vals) + xsum{b} * b_val;""" for b in range(B))
+
+    oproj_result_decls = " ".join(f"float acc{b}[TM] = {{0,0,0,0}};" for b in range(B))
+    oproj_simd_reduce = "\n".join(f"            float result{b}[TM]; for (int tm=0;tm<TM;tm++) result{b}[tm] = simd_sum(acc{b}[tm]);" for b in range(B))
+
+    oproj_tgp_write = "\n".join(f"            tgp_x2[sgid * {B} + {b}] = x2_acc{b};" for b in range(B))
+    oproj_total_decls = " ".join(f"float total{b} = 0.0f;" for b in range(B))
+
+    # Gate M1 GEMV: load M1 weights once, compute B dot products with B attn_outs
+    gate_a_x_load = "\n".join(f"""
+            float v{b}[TN];
+            for (int tn = 0; tn < TN; tn++) v{b}[tn] = float(attn_out[{b} * K_ATTN_DIM + bn + tn]);""" for b in range(B))
+
+    gate_a_dot = "\n".join(f"""
+                float gacc{b} = 0.0f;
+                for (int tn = 0; tn < TN; tn++) gacc{b} += w_row[tn] * v{b}[tn];
+                gresult{b}[tm] += gacc{b};""" for b in range(B))
+
+    gate_a_decls = " ".join(f"float gresult{b}[TM] = {{0,0,0,0}};" for b in range(B))
+    gate_a_reduce = "\n".join(f"            gresult{b}[tm] = simd_sum(gresult{b}[tm]);" for b in range(B))
+    gate_a_write = "\n".join(f"""
+                    gate_part_a[{b} * E_CONST + e] = gresult{b}[tm];""" for b in range(B))
+
+    # Gate W_fused GEMV: same pattern but with residual input
+    gate_b_x_load = "\n".join(f"""
+            float rv{b}[TN];
+            for (int tn = 0; tn < TN; tn++) rv{b}[tn] = float(residual[{b} * K_HIDDEN_DIM + bn + tn]);""" for b in range(B))
+
+    gate_b_dot = "\n".join(f"""
+                float wdot{b} = 0.0f;
+                for (int tn = 0; tn < TN; tn++) wdot{b} += w_row[tn] * rv{b}[tn];
+                bresult{b}[tm] += wdot{b};""" for b in range(B))
+
+    gate_b_decls = " ".join(f"float bresult{b}[TM] = {{0,0,0,0}};" for b in range(B))
+    gate_b_reduce = "\n".join(f"            bresult{b}[tm] = simd_sum(bresult{b}[tm]);" for b in range(B))
+    gate_b_write = "\n".join(f"""
+                    gate_part_b[{b} * E_CONST + e] = bresult{b}[tm];""" for b in range(B))
+
+    return f"""
+    const int TM = 4;
+    const int TN = 4;
+    const int blockN = 128;
+    const int E_CONST = {E};
+    const int M_DIM = {M};
+    const int K_ATTN_DIM = {K_attn};
+    const int K_HIDDEN_DIM = {K_hidden};
+    const int N_OPROJ_TG_DIM = {ceil_div(M, 32)};
+    const int BATCH_SIZE = {B};
+
+    // Helper: dot product of x_thread and w_vals (8 elements)
+    auto wdot = [](thread float* x, thread float* w) -> float {{
+        float a = 0;
+        for (int i = 0; i < 8; i++) a += x[i] * w[i];
+        return a;
+    }};
+
+    const int N_OPROJ_TG = N_OPROJ_TG_DIM;
+    const int N_M1_TG = {n_m1_tg};
+    const int blockM_gate = {blockM_gate};
+
+    uint tg_x = threadgroup_position_in_grid.x;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+
+    if (tg_x < (uint)N_OPROJ_TG) {{
+        // ══════ O_PROJ GEMV (8-bit, register-sharing for {B} batches) ══════
+        const int blockM = 32;
+        const int VPT = 8;
+        const int BLOCK_SIZE = 256;
+
+        int out_row = int(tg_x) * blockM + int(sgid) * TM;
+        if (out_row >= M_DIM) return;
+        out_row = (out_row + TM <= M_DIM) ? out_row : (M_DIM - TM);
+
+        threadgroup float tgp_x2[8 * {B}];
+
+        {oproj_result_decls}
+        int K_groups = K_ATTN_DIM / {gs};
+
+        // Weight pointers (shared across batch)
+        const device uint8_t* ws = (const device uint8_t*)W_oproj
+            + (long)out_row * K_ATTN_DIM + slid * VPT;
+        const device bfloat16_t* sc = (const device bfloat16_t*)S_oproj
+            + (long)out_row * K_groups + slid / {oproj_slid_divisor};
+        const device bfloat16_t* bi = (const device bfloat16_t*)B_oproj
+            + (long)out_row * K_groups + slid / {oproj_slid_divisor};
+
+        int xb = slid * VPT;
+
+        for (int k = 0; k < K_ATTN_DIM; k += BLOCK_SIZE) {{
+            // Load x for all {B} batches
+{oproj_x_load}
+
+            // Load weights once, compute all batches
+            for (int row = 0; row < TM; row++) {{
+                const device uint8_t* wl = ws + row * K_ATTN_DIM;
+                float s_val = float(sc[row * K_groups]);
+                float b_val = float(bi[row * K_groups]);
+                float w_vals[VPT];
+                for (int i = 0; i < VPT; i++) w_vals[i] = float(wl[i]);
+{oproj_qdot}
+            }}
+
+            ws += BLOCK_SIZE; sc += {oproj_sc_stride}; bi += {oproj_sc_stride};
+            xb += BLOCK_SIZE;
+        }}
+
+        // simd_sum for all batches
+{oproj_simd_reduce}
+
+        // Epilogue: residual add + x² + h_scaled + h_out
+        if (slid == 0) {{
+{oproj_epilogue_code}
+{oproj_tgp_write}
+        }}
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (sgid == 0 && slid == 0) {{
+            {oproj_total_decls}
+            for (int s = 0; s < 8; s++) {{
+{oproj_x2_sum}
+            }}
+{oproj_x2_final_code}
+        }}
+
+    }} else if (tg_x < (uint)(N_OPROJ_TG + N_M1_TG)) {{
+        // ══════ M1 GEMV (bf16, register-sharing for {B} batches) ══════
+        int local_tg = int(tg_x) - N_OPROJ_TG;
+        int out_row = local_tg * blockM_gate + int(sgid) * TM;
+        if (out_row >= E_CONST) return;
+        out_row = (out_row + TM <= E_CONST) ? out_row : (E_CONST - TM);
+
+        {gate_a_decls}
+        int bn = int(slid) * TN;
+        int n_iter = K_ATTN_DIM / blockN;
+
+        for (int i = 0; i < n_iter; i++) {{
+{gate_a_x_load}
+            for (int tm = 0; tm < TM; tm++) {{
+                float w_row[TN];
+                for (int tn = 0; tn < TN; tn++) w_row[tn] = float(M1[(out_row + tm) * K_ATTN_DIM + bn + tn]);
+{gate_a_dot}
+            }}
+            bn += blockN;
+        }}
+
+        for (int tm = 0; tm < TM; tm++) {{
+{gate_a_reduce}
+        }}
+
+        if (slid == 0) {{
+            for (int tm = 0; tm < TM; tm++) {{
+                int e = out_row + tm;
+                if (e < E_CONST) {{
+{gate_a_write}
+                }}
+            }}
+        }}
+
+    }} else {{
+        // ══════ W_FUSED GEMV (bf16, register-sharing for {B} batches) ══════
+        int local_tg = int(tg_x) - N_OPROJ_TG - N_M1_TG;
+        int out_row = local_tg * blockM_gate + int(sgid) * TM;
+        if (out_row >= E_CONST) return;
+        out_row = (out_row + TM <= E_CONST) ? out_row : (E_CONST - TM);
+
+        {gate_b_decls}
+        int bn = int(slid) * TN;
+        int n_iter = K_HIDDEN_DIM / blockN;
+
+        for (int i = 0; i < n_iter; i++) {{
+{gate_b_x_load}
+            for (int tm = 0; tm < TM; tm++) {{
+                float w_row[TN];
+                for (int tn = 0; tn < TN; tn++) w_row[tn] = float(W_fused[(out_row + tm) * K_HIDDEN_DIM + bn + tn]);
+{gate_b_dot}
+            }}
+            bn += blockN;
+        }}
+
+        for (int tm = 0; tm < TM; tm++) {{
+{gate_b_reduce}
+        }}
+
+        if (slid == 0) {{
+            for (int tm = 0; tm < TM; tm++) {{
+                int e = out_row + tm;
+                if (e < E_CONST) {{
+{gate_b_write}
+                }}
+            }}
+        }}
+    }}
+"""
+
+
+_batched_oproj_cache = {}
+
+
+def _get_batched_oproj_kernel(n_experts, M, K_attn, K_hidden, B, group_size=64, gate_bm=8):
+    key = (n_experts, M, K_attn, K_hidden, B, group_size, gate_bm)
+    if key not in _batched_oproj_cache:
+        _batched_oproj_cache[key] = mx.fast.metal_kernel(
+            name=f"batched_oproj_E{n_experts}_M{M}_Ka{K_attn}_Kh{K_hidden}_B{B}",
+            input_names=[
+                "W_oproj", "S_oproj", "B_oproj",
+                "attn_out", "residual", "w_rms",
+                "M1", "W_fused",
+            ],
+            output_names=["h_scaled", "h_out", "x2_partials",
+                          "gate_part_a", "gate_part_b"],
+            source=_gen_batched_oproj_source(n_experts, M, K_attn, K_hidden, B, group_size, gate_bm),
+        )
+    return _batched_oproj_cache[key]
+
+
+def batched_oproj_gate_gemv(W_oproj, S_oproj, B_oproj,
+                             attn_out, residual, w_rms,
+                             M1, W_fused,
+                             M, K_attn, batch_size,
+                             n_experts=256, gate_bm=8,
+                             K_hidden=None, group_size=64):
+    """Batched fused 8-bit o_proj + bf16 gate GEMVs.
+
+    Args:
+        W_oproj/S_oproj/B_oproj: 8-bit o_proj weights
+        attn_out: (B, K_attn) bf16
+        residual: (B, K) bf16
+        w_rms: (K,) bf16 — RMSNorm weight (shared)
+        M1: (E, K_attn) bf16 (shared)
+        W_fused: (E, K) bf16 (shared)
+        M: hidden size
+        K_attn: attention output dim
+        batch_size: B
+        n_experts: E
+        gate_bm: SGs per gate TG
+
+    Returns:
+        h_scaled (B, M) bf16, h_out (B, M) bf16,
+        x2_partials (B, N_TG) f32, gate_part_a (B, E) f32, gate_part_b (B, E) f32
+    """
+    B = batch_size
+    M_val = int(M)
+    K_attn_val = int(K_attn)
+    K_hidden_val = int(K_hidden) if K_hidden is not None else M_val
+
+    kern = _get_batched_oproj_kernel(n_experts, M_val, K_attn_val, K_hidden_val, B, group_size, gate_bm)
+
+    n_oproj_tg = ceil_div(M_val, 32)
+    blockM_gate = gate_bm * 4
+    n_m1_tg = ceil_div(n_experts, blockM_gate)
+    n_wf_tg = ceil_div(n_experts, blockM_gate)
+    total_tg = n_oproj_tg + n_m1_tg + n_wf_tg
+
+    results = kern(
+        inputs=[W_oproj, S_oproj, B_oproj,
+                attn_out.reshape(B * K_attn_val), residual.reshape(B * M_val), w_rms,
+                M1, W_fused],
+        output_shapes=[
+            (B * M_val,), (B * M_val,),
+            (B * n_oproj_tg,),
+            (B * n_experts,), (B * n_experts,),
+        ],
+        output_dtypes=[mx.bfloat16, mx.bfloat16, mx.float32, mx.float32, mx.float32],
+        grid=(total_tg * 32, 8, 1),
+        threadgroup=(32, 8, 1),
+    )
+
+    return (results[0].reshape(B, M_val),
+            results[1].reshape(B, M_val),
+            results[2].reshape(B, n_oproj_tg),
+            results[3].reshape(B, n_experts),
+            results[4].reshape(B, n_experts))

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

@@ -0,0 +1,578 @@
+"""Dispatch 2 (batched): Softmax prologue + 8-bit SwiGLU for Qwen3.5 MoE.
+
+Combines the prologue from oproj_softmax_topk_swiglu_8bit (B=1) with the
+batched body from batched_merged_swiglu_8bit. All constants are baked into
+the Metal source at Python code-generation time.
+
+Grid z-dimension: B * n_active + 1
+  - tgid.z < B * n_active: routed expert TGs
+    batch_id = tgid.z / n_active, local_z = tgid.z % n_active
+  - tgid.z == B * n_active: shared expert TG (register-level weight sharing
+    for B batch elements, including shared_expert_gate GEMV)
+
+Prologue (all TGs):
+  Phase 1: distributed x2 partial sum -> inv_rms (per batch_id)
+  Phase 2 (routed TGs only): gate scores -> softmax -> parallel top-k ->
+           norm_topk_prob -> write out_inds / norm_scores
+  Phase 3 (shared TG, SG 0): shared_expert_gate 8-bit GEMV with
+           register-level weight sharing -> gate_raw[B]
+
+Body (after TG barrier):
+  Routed: 8-bit gate+up+SwiGLU with h_scaled[batch_id] input
+  Shared: register-level weight sharing for B batch elements
+"""
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_batched_softmax_topk_swiglu_source(
+    N_INTER, SHARED_INTER, K, n_active, B,
+    n_experts=256, top_k=10, norm_topk=True, group_size=64,
+    n_oproj_tg=64,
+):
+    """Generate Metal source for batched softmax + top-k + SwiGLU.
+
+    All routed TGs compute the full softmax+topk for their batch_id into TG
+    memory. Only the TG with local_z==0 writes out_inds/norm_scores to device
+    memory. Each routed TG reads its own expert from tg_inds[local_z].
+    """
+    gs = group_size
+    sc_stride = 256 // gs
+    slid_divisor = gs // 8
+    N_TOTAL = 2 * N_INTER
+    K_groups = K // gs
+    SHARED_K_groups = K // gs
+    total_routed = B * n_active
+    E = int(n_experts)
+    K_TOP = int(top_k)
+    SPT = (E + 63) // 64
+
+    # ── Shared expert body: unrolled per-batch code ──
+    shared_x_load = "\n".join(f"""
+            float x{b}_thread[VALUES_PER_THREAD];
+            float xsum{b} = 0;
+            for (int i = 0; i < VALUES_PER_THREAD; i++) {{
+                float xi = float(X[{b} * K_DIM + x_base + i]);
+                x{b}_thread[i] = xi;
+                xsum{b} += xi;
+            }}""" for b in range(B))
+
+    shared_gate_qdot = "\n".join(f"""
+                float accum_g{b} = 0;
+                for (int i = 0; i < VALUES_PER_THREAD; i++) accum_g{b} += x{b}_thread[i] * wg_vals[i];
+                gate{b}[row] += sg * accum_g{b} + xsum{b} * bg;""" for b in range(B))
+
+    shared_up_qdot = "\n".join(f"""
+                float accum_u{b} = 0;
+                for (int i = 0; i < VALUES_PER_THREAD; i++) accum_u{b} += x{b}_thread[i] * wu_vals[i];
+                up{b}[row] += su * accum_u{b} + xsum{b} * bu;""" for b in range(B))
+
+    shared_result_decls = "\n        ".join(
+        f"float gate{b}[RESULTS_PER_SG] = {{0,0,0,0}}; float up{b}[RESULTS_PER_SG] = {{0,0,0,0}};"
+        for b in range(B))
+
+    shared_write_lines = []
+    for b in range(B):
+        shared_write_lines.append(f"""
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            float g{b} = simd_sum(gate{b}[row]) * inv_rms_{b};
+            float u{b} = simd_sum(up{b}[row]) * inv_rms_{b};
+            if (slid == 0) {{
+                float silu_g{b} = g{b} / (1.0f + metal::exp(-g{b}));
+                Y_shared[{b} * SHARED_INTER_DIM + out_row + row] = silu_g{b} * u{b};
+            }}
+        }}""")
+    shared_write = "\n".join(shared_write_lines)
+
+    # inv_rms for all B batches in the shared TG
+    shared_inv_rms_lines = []
+    for b in range(B):
+        shared_inv_rms_lines.append(f"""
+        float local_x2_{b} = 0.0f;
+        for (int i = x2_start; i < x2_end; i++) local_x2_{b} += x2_partials[{b} * N_OPROJ_TG_DIM + i];
+        float sg_x2_{b} = simd_sum(local_x2_{b});
+        if (slid == 0) tg_x2_sg[sgid] = sg_x2_{b};
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        float inv_rms_{b} = metal::precise::rsqrt((tg_x2_sg[0] + tg_x2_sg[1]) / (float)K_DIM + 1e-6f);""")
+    shared_inv_rms_block = "\n".join(shared_inv_rms_lines)
+
+    # Shared expert gate GEMV accumulator declarations
+    seg_acc_decls = "\n        ".join(
+        f"float seg_gate_acc{b} = 0.0f;" for b in range(B))
+
+    seg_write = "\n".join(
+        f"            gate_raw[{b}] = seg_gate_acc{b} * inv_rms_{b};"
+        for b in range(B))
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int N_INTER_DIM = {N_INTER};
+    const int SHARED_INTER_DIM = {SHARED_INTER};
+    const int K_DIM = {K};
+    const int K_GROUPS = {K_groups};
+    const int N_TOTAL = {N_TOTAL};
+    const int N_ACTIVE = {n_active};
+    const int TOTAL_ROUTED = {total_routed};
+    const int BATCH_SIZE = {B};
+    const int E_CONST = {E};
+    const int K_TOP_CONST = {K_TOP};
+    const int SPT = {SPT};
+    const int N_OPROJ_TG_DIM = {n_oproj_tg};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;  // 0 or 1
+    uint slid = thread_index_in_simdgroup;       // 0..31
+    int tid = int(sgid) * 32 + int(slid);        // 0..63
+
+    threadgroup float tg_x2_sg[2];
+    threadgroup int tg_inds[{K_TOP}];
+    threadgroup float tg_selected_scores[{K_TOP}];
+
+    if (tgid.z < (uint)TOTAL_ROUTED) {{
+        // ═══════════════════════════════════════════════════════════════
+        // ROUTED TG PROLOGUE
+        // ═══════════════════════════════════════════════════════════════
+        int flat_idx = (int)tgid.z;
+        int batch_id = flat_idx / N_ACTIVE;
+        int local_z = flat_idx % N_ACTIVE;
+
+        // Phase 1: distributed x2 sum -> inv_rms for batch_id
+        int chunk = (N_OPROJ_TG_DIM + 63) / 64;
+        int x2_start = tid * chunk;
+        int x2_end = min(x2_start + chunk, N_OPROJ_TG_DIM);
+        float local_x2 = 0.0f;
+        for (int i = x2_start; i < x2_end; i++)
+            local_x2 += x2_partials[batch_id * N_OPROJ_TG_DIM + i];
+        float sg_x2_sum = simd_sum(local_x2);
+        if (slid == 0) tg_x2_sg[sgid] = sg_x2_sum;
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        float total_x2 = tg_x2_sg[0] + tg_x2_sg[1];
+        float inv_rms = metal::precise::rsqrt(total_x2 / (float)K_DIM + 1e-6f);
+
+        // Phase 2: ALL routed TGs compute full softmax + top-k for their
+        // batch_id. Each TG gets its own copy in TG memory (tg_inds,
+        // tg_selected_scores). This avoids cross-TG communication.
+        {{
+            float my_scores[SPT];
+            for (int j = 0; j < SPT; j++) {{
+                int e = tid * SPT + j;
+                if (e < E_CONST)
+                    my_scores[j] = (gate_part_a[batch_id * E_CONST + e]
+                                  + gate_part_b[batch_id * E_CONST + e]) * inv_rms;
+                else
+                    my_scores[j] = -1e30f;
+            }}
+
+            // Softmax: distributed max
+            float local_max = -1e30f;
+            for (int j = 0; j < SPT; j++)
+                local_max = max(local_max, my_scores[j]);
+            float sg_max_val = simd_max(local_max);
+            threadgroup float tg_softmax_sg[2];
+            if (slid == 0) tg_softmax_sg[sgid] = sg_max_val;
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+            float tg_max = max(tg_softmax_sg[0], tg_softmax_sg[1]);
+
+            // Softmax: exp + distributed sum
+            float local_sum = 0.0f;
+            for (int j = 0; j < SPT; j++) {{
+                float e_val = metal::exp(my_scores[j] - tg_max);
+                my_scores[j] = e_val;
+                local_sum += e_val;
+            }}
+            float sg_sum_val = simd_sum(local_sum);
+            if (slid == 0) tg_softmax_sg[sgid] = sg_sum_val;
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+            float tg_sum = tg_softmax_sg[0] + tg_softmax_sg[1];
+
+            // Softmax: normalize
+            float inv_sum = 1.0f / tg_sum;
+            for (int j = 0; j < SPT; j++)
+                my_scores[j] *= inv_sum;
+
+            // Parallel top-k: K_TOP rounds
+            threadgroup float tg_tk_val[2];
+            threadgroup int tg_tk_info[2];
+
+            for (int round = 0; round < K_TOP_CONST; round++) {{
+                float best = -1.0f;
+                int best_e = -1;
+                for (int j = 0; j < SPT; j++) {{
+                    int e = tid * SPT + j;
+                    if (e < E_CONST && my_scores[j] > best) {{
+                        best = my_scores[j];
+                        best_e = e;
+                    }}
+                }}
+
+                float sg_best = simd_max(best);
+                int candidate = (best == sg_best && best > 0.0f) ? int(slid) : 999;
+                int sg_winner = simd_min(candidate);
+
+                if (slid == 0) {{
+                    tg_tk_val[sgid] = sg_best;
+                    tg_tk_info[sgid] = sg_winner;
+                }}
+                threadgroup_barrier(mem_flags::mem_threadgroup);
+
+                int winner_sg = (tg_tk_val[0] >= tg_tk_val[1]) ? 0 : 1;
+                int winner_lane = tg_tk_info[winner_sg];
+                int winner_tid = winner_sg * 32 + winner_lane;
+
+                if (tid == winner_tid) {{
+                    tg_inds[round] = best_e;
+                    tg_selected_scores[round] = best;
+                    for (int j = 0; j < SPT; j++) {{
+                        if (tid * SPT + j == best_e) {{
+                            my_scores[j] = -1.0f;
+                            break;
+                        }}
+                    }}
+                }}
+                threadgroup_barrier(mem_flags::mem_threadgroup);
+            }}
+
+            // Only local_z==0 writes norm_scores + out_inds to device memory
+            if (local_z == 0 && tid == 0) {{
+                float total_score = 0.0f;
+                for (int a = 0; a < {K_TOP}; a++) total_score += tg_selected_scores[a];
+                float inv_total = {"1.0f / total_score" if norm_topk else "1.0f"};
+                for (int a = 0; a < {K_TOP}; a++) {{
+                    norm_scores[batch_id * {K_TOP} + a] = tg_selected_scores[a] * inv_total;
+                    out_inds[batch_id * {K_TOP} + a] = (uint)tg_inds[a];
+                }}
+            }}
+        }}
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // ═══════════════════════════════════════════════════════════════
+        // ROUTED BODY: 8-bit gate+up+SwiGLU
+        // ═══════════════════════════════════════════════════════════════
+        int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG;
+        if (out_row >= N_INTER_DIM) return;
+
+        int expert = tg_inds[local_z];
+
+        const device uint8_t* ws_gate = (const device uint8_t*)W
+            + (long)expert * N_TOTAL * K_DIM + out_row * K_DIM + slid * VALUES_PER_THREAD;
+        const device bfloat16_t* sc_gate = (const device bfloat16_t*)S
+            + (long)expert * N_TOTAL * K_GROUPS + out_row * K_GROUPS + slid / {slid_divisor};
+        const device bfloat16_t* bi_gate = (const device bfloat16_t*)B_q
+            + (long)expert * N_TOTAL * K_GROUPS + out_row * K_GROUPS + slid / {slid_divisor};
+
+        const device uint8_t* ws_up = (const device uint8_t*)W
+            + (long)expert * N_TOTAL * K_DIM + (out_row + N_INTER_DIM) * K_DIM + slid * VALUES_PER_THREAD;
+        const device bfloat16_t* sc_up = (const device bfloat16_t*)S
+            + (long)expert * N_TOTAL * K_GROUPS + (out_row + N_INTER_DIM) * K_GROUPS + slid / {slid_divisor};
+        const device bfloat16_t* bi_up = (const device bfloat16_t*)B_q
+            + (long)expert * N_TOTAL * K_GROUPS + (out_row + N_INTER_DIM) * K_GROUPS + slid / {slid_divisor};
+
+        int x_base = batch_id * K_DIM + slid * VALUES_PER_THREAD;
+
+        float gate_result[4] = {{0, 0, 0, 0}};
+        float up_result[4] = {{0, 0, 0, 0}};
+
+        for (int k = 0; k < K_DIM; k += BLOCK_SIZE) {{
+            float x_thread[8];
+            float xsum = 0;
+            for (int i = 0; i < 8; i++) {{
+                float xi = float(X[x_base + i]);
+                x_thread[i] = xi;
+                xsum += xi;
+            }}
+
+            for (int row = 0; row < RESULTS_PER_SG; row++) {{
+                const device uint8_t* wg = ws_gate + row * K_DIM;
+                float sg = float(sc_gate[row * K_GROUPS]);
+                float bg = float(bi_gate[row * K_GROUPS]);
+                float accum_g = 0;
+                for (int i = 0; i < 8; i++) accum_g += x_thread[i] * float(wg[i]);
+                gate_result[row] += sg * accum_g + xsum * bg;
+
+                const device uint8_t* wu = ws_up + row * K_DIM;
+                float su = float(sc_up[row * K_GROUPS]);
+                float bu = float(bi_up[row * K_GROUPS]);
+                float accum_u = 0;
+                for (int i = 0; i < 8; i++) accum_u += x_thread[i] * float(wu[i]);
+                up_result[row] += su * accum_u + xsum * bu;
+            }}
+
+            ws_gate += BLOCK_SIZE; ws_up += BLOCK_SIZE;
+            sc_gate += {sc_stride}; sc_up += {sc_stride};
+            bi_gate += {sc_stride}; bi_up += {sc_stride};
+            x_base += BLOCK_SIZE;
+        }}
+
+        // Epilogue: apply inv_rms (factored), SwiGLU, write f32
+        device float* yp = Y_routed + flat_idx * N_INTER_DIM + out_row;
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            float g = simd_sum(gate_result[row]) * inv_rms;
+            float u = simd_sum(up_result[row]) * inv_rms;
+            if (slid == 0) {{
+                float silu_g = g / (1.0f + metal::exp(-g));
+                yp[row] = silu_g * u;
+            }}
+        }}
+
+    }} else {{
+        // ═══════════════════════════════════════════════════════════════
+        // SHARED TG (tgid.z == TOTAL_ROUTED)
+        // ═══════════════════════════════════════════════════════════════
+
+        // Phase 1: compute inv_rms for ALL B batch elements
+        int chunk = (N_OPROJ_TG_DIM + 63) / 64;
+        int x2_start = tid * chunk;
+        int x2_end = min(x2_start + chunk, N_OPROJ_TG_DIM);
+{shared_inv_rms_block}
+
+        // Phase 3: shared_expert_gate 8-bit GEMV (SG 0 only)
+        // Load W_seg once, compute B dot products via register-level weight sharing
+        if (sgid == 0) {{
+            const int VPT = 8;
+            const int SEG_BLOCK = 256;  // 32 * VPT
+
+            const device uint8_t* seg_w_ptr = (const device uint8_t*)W_seg
+                + slid * VPT;
+            const device bfloat16_t* seg_sc = (const device bfloat16_t*)S_seg
+                + slid / {slid_divisor};
+            const device bfloat16_t* seg_bi = (const device bfloat16_t*)B_seg
+                + slid / {slid_divisor};
+            int seg_xb = slid * VPT;
+
+            {seg_acc_decls}
+
+            for (int k = 0; k < K_DIM; k += SEG_BLOCK) {{
+                // Load weight block once into registers
+                float seg_w_regs[VPT];
+                for (int i = 0; i < VPT; i++) seg_w_regs[i] = float(seg_w_ptr[i]);
+                float seg_sc_val = float(*seg_sc);
+                float seg_bi_val = float(*seg_bi);
+
+                // Compute B dot products from the same weight registers
+{chr(10).join(f'''                {{
+                    float xsum{b} = 0.0f, wacc{b} = 0.0f;
+                    for (int i = 0; i < VPT; i++) {{
+                        float xi = float(X[{b} * K_DIM + seg_xb + i]);
+                        xsum{b} += xi;
+                        wacc{b} += xi * seg_w_regs[i];
+                    }}
+                    seg_gate_acc{b} += seg_sc_val * wacc{b} + xsum{b} * seg_bi_val;
+                }}''' for b in range(B))}
+
+                seg_w_ptr += SEG_BLOCK;
+                seg_sc += {sc_stride};
+                seg_bi += {sc_stride};
+                seg_xb += SEG_BLOCK;
+            }}
+
+            // Reduce across SG and write gate_raw[B]
+{chr(10).join(f"            seg_gate_acc{b} = simd_sum(seg_gate_acc{b});" for b in range(B))}
+            if (slid == 0) {{
+{seg_write}
+            }}
+        }}
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // ═══════════════════════════════════════════════════════════════
+        // SHARED BODY: register-level weight sharing for B batch elements
+        // ═══════════════════════════════════════════════════════════════
+        int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG;
+        if (out_row >= SHARED_INTER_DIM) return;
+
+        const device uint8_t* ws_gate = (const device uint8_t*)W_shared
+            + (long)out_row * K_DIM + slid * VALUES_PER_THREAD;
+        const device bfloat16_t* sc_gate = (const device bfloat16_t*)S_shared
+            + (long)out_row * {SHARED_K_groups} + slid / {slid_divisor};
+        const device bfloat16_t* bi_gate = (const device bfloat16_t*)B_shared
+            + (long)out_row * {SHARED_K_groups} + slid / {slid_divisor};
+
+        const device uint8_t* ws_up = (const device uint8_t*)W_shared
+            + (long)(out_row + SHARED_INTER_DIM) * K_DIM + slid * VALUES_PER_THREAD;
+        const device bfloat16_t* sc_up = (const device bfloat16_t*)S_shared
+            + (long)(out_row + SHARED_INTER_DIM) * {SHARED_K_groups} + slid / {slid_divisor};
+        const device bfloat16_t* bi_up = (const device bfloat16_t*)B_shared
+            + (long)(out_row + SHARED_INTER_DIM) * {SHARED_K_groups} + slid / {slid_divisor};
+
+        int x_base = slid * VALUES_PER_THREAD;
+        {shared_result_decls}
+
+        for (int k = 0; k < K_DIM; k += BLOCK_SIZE) {{
+            // Load x for all {B} batch elements
+{shared_x_load}
+
+            for (int row = 0; row < RESULTS_PER_SG; row++) {{
+                // Load gate weights once into registers
+                const device uint8_t* wg = ws_gate + row * K_DIM;
+                float sg = float(sc_gate[row * {SHARED_K_groups}]);
+                float bg = float(bi_gate[row * {SHARED_K_groups}]);
+                float wg_vals[VALUES_PER_THREAD];
+                for (int i = 0; i < VALUES_PER_THREAD; i++) wg_vals[i] = float(wg[i]);
+
+                // Compute gate for all {B} batches from registers
+{shared_gate_qdot}
+
+                // Load up weights once into registers
+                const device uint8_t* wu = ws_up + row * K_DIM;
+                float su = float(sc_up[row * {SHARED_K_groups}]);
+                float bu = float(bi_up[row * {SHARED_K_groups}]);
+                float wu_vals[VALUES_PER_THREAD];
+                for (int i = 0; i < VALUES_PER_THREAD; i++) wu_vals[i] = float(wu[i]);
+
+                // Compute up for all {B} batches from registers
+{shared_up_qdot}
+            }}
+
+            ws_gate += BLOCK_SIZE; ws_up += BLOCK_SIZE;
+            sc_gate += {sc_stride}; sc_up += {sc_stride};
+            bi_gate += {sc_stride}; bi_up += {sc_stride};
+            x_base += BLOCK_SIZE;
+        }}
+
+        // SwiGLU epilogue + write for all {B} batches (with inv_rms)
+{shared_write}
+    }}
+"""
+
+_batched_softmax_topk_swiglu_cache = {}
+
+
+def _get_batched_softmax_topk_swiglu_kernel(
+    N_INTER, SHARED_INTER, K, n_active, B,
+    n_experts=256, top_k=10, norm_topk=True, group_size=64,
+    n_oproj_tg=64,
+):
+    key = (N_INTER, SHARED_INTER, K, n_active, B, n_experts, top_k, norm_topk, group_size, n_oproj_tg)
+    if key not in _batched_softmax_topk_swiglu_cache:
+        nt_tag = "_nt" if norm_topk else ""
+        _batched_softmax_topk_swiglu_cache[key] = mx.fast.metal_kernel(
+            name=(f"batched_softmax_topk_swiglu_8bit"
+                  f"_NI{N_INTER}_SI{SHARED_INTER}_K{K}"
+                  f"_na{n_active}_B{B}_E{n_experts}_k{top_k}"
+                  f"_gs{group_size}{nt_tag}"),
+            input_names=[
+                "W", "S", "B_q",                    # routed expert weights
+                "W_shared", "S_shared", "B_shared",  # shared expert weights
+                "W_seg", "S_seg", "B_seg",            # shared_expert_gate weights
+                "X",                                  # h_scaled (B, K) bf16
+                "gate_part_a", "gate_part_b",         # (B, E) f32
+                "x2_partials",                        # (B, N_OPROJ_TG) f32
+            ],
+            output_names=["Y_routed", "Y_shared", "out_inds",
+                          "norm_scores", "gate_raw"],
+            source=_gen_batched_softmax_topk_swiglu_source(
+                N_INTER, SHARED_INTER, K, n_active, B,
+                n_experts, top_k, norm_topk, group_size, n_oproj_tg),
+        )
+    return _batched_softmax_topk_swiglu_cache[key]
+
+
+def batched_softmax_topk_swiglu_8bit(
+    w_gu, s_gu, b_gu,               # routed gate+up weights (E, 2*N_INTER, K/4)
+    w_shared, s_shared, b_shared,    # shared gate+up weights (2*SHARED_INTER, K/4)
+    w_seg, s_seg, b_seg,             # shared_expert_gate weights (1, K/4)
+    h_scaled,                        # (B, K) bf16 — from Dispatch 1
+    gate_part_a,                     # (B, E) f32 — from Dispatch 1
+    gate_part_b,                     # (B, E) f32 — from Dispatch 1
+    x2_partials,                     # (B, N_OPROJ_TG) f32 — from Dispatch 1
+    n_inter, k_hidden, batch_size, n_active,
+    n_oproj_tg, n_experts=256,
+    shared_inter=None, group_size=64,
+):
+    """Batched Dispatch 2: softmax + top-k + merged 8-bit SwiGLU with oproj prologue.
+
+    Prologue (per-batch):
+      Phase 1: distributed x2 -> inv_rms (all TGs, indexed by batch_id)
+      Phase 2: softmax(gate_scores) -> top-k -> norm_topk_prob (all routed TGs)
+      Phase 3: shared_expert_gate 8-bit GEMV -> gate_raw[B] (shared TG, SG 0)
+
+    Body:
+      Routed: 8-bit gate+up+SwiGLU with h_scaled[batch_id] input, inv_rms factored
+      Shared: register-level weight sharing for B batch elements
+
+    Args:
+        w_gu: stacked routed weights (E, 2*N_INTER, K/4) uint32
+        s_gu: routed scales (E, 2*N_INTER, K/gs) bf16
+        b_gu: routed biases (E, 2*N_INTER, K/gs) bf16
+        w_shared: shared gate+up stacked (2*SHARED_INTER, K/4) uint32
+        s_shared: shared scales (2*SHARED_INTER, K/gs) bf16
+        b_shared: shared biases (2*SHARED_INTER, K/gs) bf16
+        w_seg/s_seg/b_seg: shared_expert_gate 8-bit weights (1, K/4) uint32
+        h_scaled: (B, K) bf16 — h * w_rms from Dispatch 1
+        gate_part_a: (B, E) f32 — partial gate scores from Dispatch 1
+        gate_part_b: (B, E) f32 — partial gate scores from Dispatch 1
+        x2_partials: (B, N_OPROJ_TG) f32 — per-TG x2 sums from Dispatch 1
+        n_inter: routed intermediate size
+        k_hidden: hidden size K
+        batch_size: B (1..8)
+        n_active: experts per token (top_k)
+        n_oproj_tg: number of o_proj TGs (for x2 partial sum reduction)
+        n_experts: total number of experts E
+        shared_inter: shared intermediate size (defaults to n_inter)
+        group_size: quantization group size (default 64)
+
+    Returns:
+        (Y_routed, Y_shared, out_inds, norm_scores, gate_raw):
+            Y_routed: (B * n_active, n_inter) f32
+            Y_shared: (B, shared_inter) f32
+            out_inds: (B * n_active,) uint32
+            norm_scores: (B * n_active,) f32
+            gate_raw: (B,) f32 — raw shared expert gate values (sigmoid in epilogue)
+    """
+    B = int(batch_size)
+    n_inter_val = int(n_inter)
+    shared_inter_val = int(shared_inter) if shared_inter is not None else n_inter_val
+    k_val = int(k_hidden)
+    n_active_val = int(n_active)
+    top_k = n_active_val
+    E = int(n_experts)
+    n_oproj_tg_val = int(n_oproj_tg)
+
+    kern = _get_batched_softmax_topk_swiglu_kernel(
+        n_inter_val, shared_inter_val, k_val, n_active_val, B,
+        E, top_k, True, int(group_size), n_oproj_tg_val,
+    )
+
+    max_inter = max(n_inter_val, shared_inter_val)
+    total_routed = B * n_active_val
+
+    results = kern(
+        inputs=[
+            w_gu, s_gu, b_gu,
+            w_shared, s_shared, b_shared,
+            w_seg, s_seg, b_seg,
+            h_scaled,
+            gate_part_a, gate_part_b,
+            x2_partials,
+        ],
+        output_shapes=[
+            (total_routed * n_inter_val,),        # Y_routed flat
+            (B * shared_inter_val,),               # Y_shared flat
+            (total_routed,),                       # out_inds
+            (total_routed,),                       # norm_scores
+            (B,),                                  # gate_raw
+        ],
+        output_dtypes=[
+            mx.float32,    # Y_routed
+            mx.float32,    # Y_shared
+            mx.uint32,     # out_inds
+            mx.float32,    # norm_scores
+            mx.float32,    # gate_raw
+        ],
+        grid=(32, ceil_div(max_inter, 8) * 2, total_routed + 1),
+        threadgroup=(32, 2, 1),
+    )
+
+    Y_routed = results[0].reshape(total_routed, n_inter_val)
+    Y_shared = results[1].reshape(B, shared_inter_val)
+    out_inds = results[2]
+    norm_scores = results[3]
+    gate_raw = results[4]
+
+    return Y_routed, Y_shared, out_inds, norm_scores, gate_raw

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

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

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

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

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

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

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

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

src/exo/worker/engines/mlx/patches/qwen3_5_moe_split/PIPELINE.md

@@ -0,0 +1,440 @@
+# DFlash Speculative Decoding on Qwen3.5 Attn/MoE Split — Pipeline Reference
+
+Navigation aid for reading and debugging this code path. Every step has a `file:line` so it can be opened directly. Written for Qwen3.5 (both dense 27B and MoE variants) running with `world_size==2` on JACCL.
+
+All paths below are relative to `exo/src/exo/worker/engines/mlx/` unless noted otherwise.
+
+---
+
+## 1. Executive summary
+
+- **Two ranks.** `ATTN_RANK=0` runs attention (GQA `self_attn` and linear `linear_attn`). `MOE_RANK=1` runs the MLP / MoE block. Ranks are defined in `patches/qwen3_5_moe_split/decoder.py:16`.
+- **Two independent layers of patches:**
+  - **Layer A (structural).** Replaces class methods so different ranks do different work. All installs in `patches/qwen3_5_moe_split/apply.py:220–294`.
+  - **Layer B (kernel).** Instance-level swaps of each projection's `__call__` with a dynamic LpB kernel picker. Target: `patches/qwen3_5/lpb_patch.py:102`. Drafter: `speculative/bf16_lpb_patch.py:99`.
+- **Three phases per request.** Prefill (S > 1) → speculative cycles (S = V+1, one-shot per cycle) → termination (stop token or length cap).
+- **Drafter lives on both ranks.** Both ranks run the drafter in parallel; drafts are synced via `all_gather` and MOE_RANK's drafts win.
+- **Rollback runs on both ranks.** Same `n_accepted` on both (deterministic at temp=0; broadcast at temp>0). ATTN_RANK uses the real speculative GDN states; MOE_RANK's `cache[1]` stays `None` because it never runs `linear_attn`.
+
+---
+
+## 2. Architecture overview
+
+```
+                 JACCL group (world_size=2)
+                  ▲                    ▲
+                  │ all_gather         │ all_gather
+                  │                    │
+  ┌───────────────┴─────────┐  ┌───────┴────────────────┐
+  │  ATTN_RANK (rank 0)     │  │  MOE_RANK (rank 1)     │
+  │  ─────────────────────  │  │  ─────────────────────  │
+  │  Target model:          │  │  Target model:          │
+  │   - self_attn (GQA)     │  │   - mlp (dense/MoE)     │
+  │   - linear_attn (GDN)   │  │   - post_attention_ln   │
+  │   - input_layernorm     │  │   - input_layernorm†    │
+  │   - KV cache real       │  │   - KV cache empty*     │
+  │   - GDN state real      │  │   - GDN state empty*    │
+  │                         │  │                         │
+  │  Drafter (full copy)    │  │  Drafter (full copy)    │
+  │   - 5 layers, KVCache   │  │   - 5 layers, KVCache   │
+  └─────────────────────────┘  └─────────────────────────┘
+
+  † kept under EXO_SPECULATIVE=1 (used for GDN conv_input reconstruction)
+  * rank-specific cache shims in apply.py:33 make empty-cache ops safe
+```
+
+**Draft-position authority.** ATTN_RANK's `cache.offset` is the truth — it runs attention and knows the real position. MOE_RANK's `cache.offset` is always 0. Every speculative cycle opens with an `all_gather` of `_draft_position` that overwrites MOE_RANK's copy (`dflash_split.py:65–70`).
+
+---
+
+## 3. Patch inventory
+
+### 3a. Layer A — structural patches
+
+Installed in order by `apply_qwen35_attn_moe_split_patches` (`apply.py:220`). Invoked from `auto_parallel.py:474` under `attn_moe_split_auto_parallel`, gated by `isinstance(inner, Qwen3_5TextModelInner)` and `world_size==2`.
+
+| # | Patched symbol | Replacement | Install site | Replacement body | Purpose |
+|---|----------------|-------------|--------------|------------------|---------|
+| 1 | `DecoderLayer.__call__` | `_split_call` | `apply.py:250` | `decoder.py:36` | Serial S==1 split, one all_gather per sub-step |
+| 2 | `Qwen3_5TextModel.__call__` | `_pipelined_call` | `apply.py:253` | `model_forward.py:430` | S>1 routes through `pipelined_layer_loop`; S==1 falls through to stock loop → `_split_call` |
+| 3 | `mtp_module.speculative_forward` | pipelined variant | `apply.py:259` | `model_forward.py:491` (factory), `:499` (body) | MTP speculative path (not used under `SPECULATIVE_MODE=dflash`) |
+| 4 | `dflash_speculative.dflash_speculative_forward` | `_pipelined_dflash_forward` | `apply.py:270` | `model_forward.py:615` (factory), `:623` (body) | DFlash verify forward over the pipeline; captures target-layer hiddens + GDN speculative states |
+| 5 | `DFlashBatchGenerator._speculative_next` | `_split_speculative_next` | `apply.py:279` | `dflash_split.py:29` | Draft/verify/accept cycle with cross-rank syncs |
+| 6 | Cache shims on MOE_RANK | many | `apply.py:284` | `apply.py:33–211` | Make `state` / `extract` / `filter` / `extend` safe on unpopulated caches |
+| 7 | Weight dropping | `_drop_unused_weights` | `apply.py:288` | `apply.py:297` | Only runs when `EXO_SPECULATIVE!=1`; speculative keeps both sides' weights |
+
+### 3b. Layer B — kernel patches
+
+| # | Target | Function | Body | Notes |
+|---|--------|----------|------|-------|
+| 1 | Target model projections | `apply_lpb_patches(model)` | `patches/qwen3_5/lpb_patch.py:102` | Called from `apply.py:245` under `EXO_LPB_PATCHES=1`. Patches per-layer `mlp.{gate,up,down}_proj`, attn `{q,k,v,o}_proj` or `{in_proj_qkv,in_proj_z,out_proj}`, and `lm_head`. Instance-level wrapping via `setattr(parent, proj_name, _LpB*Linear(...))`. |
+| 2 | Drafter projections | `apply_bf16_lpb_patches(drafter)` | `speculative/bf16_lpb_patch.py:99` | Called from `generator/batch_generate.py:117`. Patches per-layer `mlp_{gate,up,down}` + `self_attn.{q,k,v,o}_proj` + `drafter.fc` + `drafter.lm_head`. `DFLASH_LPB_ONLY=q_proj,k_proj,...` bisect knob at `bf16_lpb_patch.py:106`. |
+| 3 | `gated_delta_update` | `_make_speculative_gdu(spec_all_states)` | `speculative/mtp_module.py:126` | Temporary swap on `mlx_lm.models.qwen3_5.gated_delta_update`. Installed at verify entry (`model_forward.py:667`), restored post-loop (`model_forward.py:732`). Captures per-step recurrent states for rollback. |
+
+The LpB wrappers memoize `M → kernel_fn` per projection. Kernel selection logic lives in `matmul/patches/kernel_picker.py:33` (`pick_bf16_kernel`) and `:69` (`pick_int8_kernel`). `MAX_M=16` in both LpB patches; M>16 falls back to the original `__call__`.
+
+### 3c. Target-layer capture (stock DFlash, no bypass)
+
+Stock `DFlashBatchGenerator._setup_hidden_capture` (`speculative/dflash_batch_generator.py:53`) wraps target-layer indices in a `_CapturingLayer`. Its `__call__` is what populates `self._captured['{layer,prefill}_hiddens']`. Our pipelined path **never calls `layer.__call__`** — it dispatches `layer.self_attn` / `layer.linear_attn` / `layer.mlp` directly. Compensation: `_populate_dflash_captured` (`model_forward.py:387`) reaches into `_CapturingLayer`'s closure at the end of `_pipelined_call` to write the dict DFlash expects. See commit `b1a5dccf`.
+
+---
+
+## 4. Execution flow by phase
+
+### 4a. Startup & warmup
+
+1. Model loads via `auto_parallel.attn_moe_split_auto_parallel` (`auto_parallel.py:436`), which calls `apply_qwen35_attn_moe_split_patches(model, group)` (`auto_parallel.py:474`). Both ranks run patch install; differences are gated by `group.rank()`.
+   > *Prints:* `Patched X target projections with dynamic LpB` (`lpb_patch.py:130`) if `EXO_LPB_PATCHES!=0`; `Qwen3.5 attn/moe split patch applied on rank N/2` (`apply.py:291`).
+2. `BatchGenerator` init reads `EXO_SPECULATIVE=1 EXO_SPECULATIVE_MODE=dflash` (`batch_generate.py:100`).
+3. Drafter is instantiated: `DFlashDrafter(self.model, dflash_path)` (`batch_generate.py:116`). See `dflash_module.py:107` — weights loaded via `huggingface_hub.snapshot_download` if not cached.
+   > *Prints:* `DFlash config: N layers, hidden=..., heads=.../..., block=...` (`dflash_module.py:133`); `Target layers: [...], mask_token=...` (`:136`) — **this is the authoritative list of `target_layer_ids`**; `DFlash loaded: X tensors, Y.YM params` (`:164`).
+4. Drafter is LpB-patched: `apply_bf16_lpb_patches(drafter)` (`batch_generate.py:117`).
+   > *Prints:* `DFLASH_LPB_ONLY=...` only if that env var is set (`bf16_lpb_patch.py:109`); `Patched X DFlash drafter projections with dynamic LpB` (`:143`).
+5. `DFlashBatchGenerator.__init__` (`dflash_batch_generator.py:23`) runs `_setup_hidden_capture` (`dflash_batch_generator.py:53`) — wraps each index in `drafter.target_layer_ids` with `_CapturingLayer`.
+6. **Warmup** — `warmup_dflash` (`batch_generate.py:272`) sweeps drafter projections across every M they'll see: `S_ctx ∈ [1, V+1]` × `block_size`. Then one target verify at `M=V+1` to compile every target projection. Without this, Metal kernel compilation stalls the first real batch. Warmup also runs one non-speculative `dflash_speculative_forward` (`batch_generate.py:311`) for target-only kernel priming, then rolls back the cache.
+   > *Prints:* bookends `Warming up DFlash speculative decoding kernels...` (`batch_generate.py:298`) and `DFlash warmup complete` (`:357`). Each pipelined forward inside warmup emits the full `[rank N] pipeline begin ...` + per-stage stream from `model_forward.py:212, 225, 268, 306, 331`.
+   > *Evals:* `mx.eval(target_hidden_full, logits)` (`:313`), `mx.eval(dl)` per drafter sweep iteration (`:332`), `mx.eval(target_hidden, vl)` after the verify forward (`:341`).
+
+### 4b. Prefill (first step for a uid)
+
+Entry: `BatchGenerator.step` → `DFlashBatchGenerator._next` (`dflash_batch_generator.py:92`).
+
+1. `_next` sees uid ∉ `_prefilled`, routes to `_first_step_capture` (`dflash_batch_generator.py:119`).
+2. `super()._next()` → stock `BatchGenerator._next` processes the prompt, calls `self.model(prompt_tokens, cache)`.
+3. That reaches our patched `Qwen3_5TextModel.__call__` = `_pipelined_call` (`model_forward.py:430`). S = prompt_len > 1, so it takes the pipelined branch (`:456–481`):
+   1. `_dflash_capturing_target_ids(self)` (`model_forward.py:353`) detects the `_CapturingLayer` wrappers.
+   2. `pipelined_layer_loop(..., capture_layers=target_ids)` (`model_forward.py:146`) runs the 2N+1 stage pipeline — see section 4c step 7.5 for the full stage-by-stage print/eval inventory. Same structure applies here at `S = prompt_len`; whether the single-gather fast path or the two-gather slow path fires depends on `prompt_len % 2`.
+   3. `_populate_dflash_captured(self, layer_hiddens, S)` (`model_forward.py:387`) writes `_captured['layer_hiddens']` and `_captured['prefill_hiddens']` via closure introspection.
+4. Back in `super()._next()`: `inner.norm` → `lm_head` (LpB wrapper, M = prompt_len, typically > 16 → fallback to stock GEMM). First token sampled.
+5. Back in `_first_step_capture`: reads `_captured['prefill_hiddens']` → `_last_target_hidden[uid]` shape `(1, prompt_len, D·|target_layers|)`. Seeds `_draft_position[uid]` from `cache.offset` on one of the caches.
+   > *Evals:* `mx.eval(target_hidden)` (`dflash_batch_generator.py:125`) before caching — forces the concatenated prefill hiddens to materialize.
+6. Falls into `_speculative_next` (first real cycle).
+
+### 4c. Speculative cycle (per step)
+
+Entry: `DFlashBatchGenerator._speculative_next` — patched to `_split_speculative_next` in `dflash_split.py:29`. Both ranks run identical code unless noted.
+
+Numbered per the source, cross-referenced with `dflash_split.py` line numbers.
+
+1. **Preamble** (`:34–55`). Record `tic`, append `y_val` to `batch.tokens[0]`, read `_last_target_hidden[uid]`. If missing, fall back to `super()._next()` — this hits the S==1 decode path (section 4d). Logs `DFlash speculative cycle (...)` and any fallback.
+   > *Prints:* either `[rank N] DFlash: NO target_hidden -> fallback ... (y_val=...)` (`dflash_split.py:46`) OR `[rank N] DFlash speculative cycle (y_val=..., target_hidden.shape=...)` (`:51`) — one per cycle per rank.
+2. **Draft-position sync** (`:65–70`). `all_gather(_draft_position)` → take `ATTN_RANK`'s. ATTN_RANK's cache offset is authoritative; MOE_RANK's stays at 0 because it never runs attention.
+3. **Draft** (`:73–76`). Both ranks: `drafter.draft(last_target_hidden, block_ids, start)` → `draft_logits` of shape `(B, block_size-1, vocab)`. `block_ids` is `[y_val, MASK, MASK, ...]` of length `block_size`. The drafter implementation is `dflash_module.py:226` — see notes at the bottom of section 4c.
+4. **Sample** (`:79–91`). temp=0: `argmax` of `draft_logits`. temp>0: per-position `mx.random.categorical` with softmax. Both ranks sample **independently** (RNG seeds may diverge).
+   > *Evals:* `mx.eval(all_drafts_arr)` (`dflash_split.py:81`) before `.tolist()` on temp=0 path; temp>0 uses `.item()` per position which evals inline.
+5. **Draft sync** (`:94–98`). `all_gather(drafts_local)` → take MOE_RANK's `drafts_arr`. After this, both ranks have identical `drafts` of length `verify_len`.
+   > *Evals:* `mx.eval(drafts_arr)` (`dflash_split.py:97`) before `.tolist()` — required because the list is used Python-side to build `verify_input`.
+6. **Build verify input** (`:101–102`). `verify_input = concat([[y_val]], drafts_arr)`, shape `(1, V+1)`.
+7. **Pipelined verify forward** (`:106–112`). `dflash_speculative_forward(model, verify_input, cache, target_layer_ids, speculative=True)` → our `_pipelined_dflash_forward` (`model_forward.py:623`):
+   1. Wrap GDN caches in `SpeculativeArraysCache` (`speculative_cache.py:15`).
+   2. Monkey-patch `gated_delta_update` → `_make_speculative_gdu(spec_all_states)` (`mtp_module.py:126`, installed at `model_forward.py:667`). The speculative kernel returns `(y, state_out, all_states)` and appends `all_states` to the closure list. Shape: `(B, T, H_v, D_v, D_k)`.
+   3. Collect GDN pre-loop data: for each `is_linear` layer, take `spec_cache[0]` (existing conv state) or zeros, store `(pre_conv, c, layer, idx)`.
+   4. Compute `capture_set = target_layer_ids ∪ {L-1 : L is GDN, L ≥ 1}` (`model_forward.py:707`). The extra `L-1` outputs are needed to reconstruct conv_input post-loop.
+   5. Run `pipelined_layer_loop(..., capture_layers=capture_set)`. Stages inside the loop:
+      - **Stage 0 (startup)** — ATTN_RANK: `attn_0(H0)`; MOE_RANK: idle placeholder. 1 all_gather. (`model_forward.py:214–228`)
+      - **Stages 1..2N-1 (main)** — alternating B/A:
+        - B (odd stage, `T = stage//2`): ATTN runs `attn_T(H1)`, MOE runs `moe_T(h_T_H0)`. Even S → 1 all_gather, odd S → 2. (`:236–275`)
+        - A (even stage, T≥1): ATTN runs `attn_T(H0)`, MOE runs `moe_{T-1}(h_{T-1}_H1)`. Capture of layer T-1 at A-stage end: `capture[T-1] = concat(x_H0, x_H1)` (`:318–319`). 1 or 2 all_gathers.
+      - **Stage 2N (drain)** — ATTN idle, MOE `moe_{N-1}(h_{N-1}_H1)`. 1 all_gather. (`:322–334`)
+      - Capture of layer N-1 happens post-drain (`:342–343`).
+      - Total collectives: even S → 2N+1; odd S → 4N-1.
+      - Per-stage `mx.eval` calls (`:221`, `:249`, `:262` etc.) break MLX's graph accumulation and keep JACCL's queue drained.
+
+      > *Prints per stage:* `[rank N] pipeline begin S=V+1 N=64 even=<bool>` at loop entry (`model_forward.py:212`), then per-stage lines each carrying `... [rank N:role] eval_local=X.XXms eval_gather=Y.YYms` at `:225` (startup), `:268` (B), `:306` (A), `:331` (drain). Total = 1 + (2N+1) = 2N+2 lines per rank per pipelined forward. `eval_local` times only `mx.eval(my_out)`; `eval_gather` times only `mx.eval(gathered)` — see section 8a for the role labels and what to read from them.
+      > *Evals per stage (graph-break + JACCL queue bound):* startup — `mx.eval(contribution)` (`:221`), `mx.eval(gathered)` (`:223`). Main-loop B stage — even S: `mx.eval(my_out)` (`:249`), `mx.eval(gathered)` (`:251`). Odd S: `mx.eval(attn_side)` (`:261`), `mx.eval(moe_side)` (`:262`), `mx.eval(attn_contrib)` (`:266`), `mx.eval(moe_contrib)` (`:267`). Main-loop A stage — even S: `mx.eval(my_out)` (`:287`), `mx.eval(gathered)` (`:289`). Odd S: `mx.eval(attn_side)` (`:299`), `mx.eval(moe_side)` (`:300`), `mx.eval(attn_contrib)` (`:304`), `mx.eval(moe_contrib)` (`:305`). Drain — `mx.eval(contribution)` (`:327`), `mx.eval(gathered)` (`:329`).
+      > *Async / no eval:* captured tensors written to the `capture` dict (`:319`, `:343`) and the final `concat([out_H0, out_H1])` (`:339`) are **not** eval'd here — they materialize implicitly when post-loop reconstruction reads them (step 7.7 / 7.8) and when `lm_head` runs in step 7.10.
+   6. Restore stock `gated_delta_update` (`:732`).
+   7. **Merge H0+H1 speculative states** (`:741–752`). `spec_all_states` has `2 × len(gdn_spec_data)` entries (one per half per layer). Concat consecutive pairs along step dim → per-layer `all_states`. On MOE_RANK `spec_all_states` is empty because the monkey-patched GDU never ran there → `merged_states == []`.
+   8. **Reconstruct conv_input per GDN layer** (`:754–788`). `layer_input = initial_embed` (for layer 0) or `layer_hiddens[layer_idx - 1]` (captured by the pipeline). Then `normed = input_layernorm(layer_input)`, `qkv = in_proj_qkv(normed)`, and `spec_cache.conv_input = concat([pre_conv, qkv], axis=1)`. On MOE_RANK we keep `input_layernorm` and `linear_attn` under `EXO_SPECULATIVE=1`, so this runs there too.
+   9. Concat captured target hiddens → `target_hidden` of shape `(1, V+1, D·|target_layers|)` (`:791`).
+   10. Final norm + lm_head (LpB wrapper, M = V+1 — LpB fires) → `verify_logits` of shape `(1, V+1, vocab)`.
+8. **Acceptance** (`:116–162`). temp=0: `argmax(verify_logits[:, :V, :]) == drafts_arr` → `matches`, count leading True → `n_accepted`. temp>0: MOE_RANK computes acceptance ratios + uniforms; `all_gather([n_accepted_local])` broadcasts MOE's value.
+   > *Async evals:* temp=0 — `mx.async_eval(matches, all_next, target_hidden)` (`dflash_split.py:120`) kicks off acceptance + next-step state; the Python `for i in range(V): if matches[i].item(): ...` loop blocks per-index on the already-scheduled compute. temp>0 — `mx.async_eval(accept_ratios, uniforms, corrections, bonus_token, target_hidden)` (`:146–148`).
+   > *Prints:* `[DFlash] n_accepted=X/V` on MOE_RANK only (`:165`) — one line per cycle.
+9. **Rollback** (`:171–177`). `rollback = V - n_accepted`. If > 0: GQA `BatchKVCache` → `c.offset -= rollback`. GDN `SpeculativeArraysCache` → `c.rollback(n_accepted)` (`speculative_cache.py:81`): sets `base.cache[1] = all_states[0, n_accepted]` if `all_states` is set (ATTN_RANK only), and `base.cache[0] = conv_input[:, n_accepted+1 : n_accepted+1+3, :]` (both ranks).
+10. **Unwrap speculative caches** (`:179–181`). Replace each `SpeculativeArraysCache` with its `.base` for the next cycle's stock layer dispatch.
+11. **Emit tokens** (`:184–212`). Bonus/correction token (all accepted: `all_next[V]` or `bonus_token`; partial: `all_next[n_accepted]` or `corrections[n_accepted]`). Update `_last_target_hidden[uid] = target_hidden[:, :n_accepted+1, :]`, advance `_draft_position += n_accepted+1`. Buffer accepted draft tokens into `_token_buffer[uid]` (`:234` or `:248`).
+   > *Async evals:* `mx.async_eval(batch.y)` at `dflash_split.py:239` (stop-token path) or `:250` (normal return) — schedules the next-step bonus token without blocking the `Response` yield.
+
+**Notes on the drafter.** `DFlashDrafter.draft` (`dflash_module.py:226`):
+- `target_hidden` shape `(B, accepted_len, n_layers · hidden_size)` — compressed via `self.fc` (`dflash_module.py:243`) to `(B, accepted_len, hidden_size)`.
+- `start` = `_draft_position[uid]` = prompt_len on first cycle, then `start + n_accepted + 1` each cycle.
+- K/V are formed from `concat(target_hidden, draft_input)` (`dflash_module.py:50`), so `k_proj/v_proj` see `M = accepted_len + block_size`. On the first cycle `accepted_len = prompt_len`, so these projections usually fall back to stock GEMM. After the first cycle `accepted_len ≤ V+1`, so LpB fires at `M = V+1 + block_size` (if ≤ 16).
+
+### 4d. Buffered drain and S==1 fallback decode
+
+- `DFlashBatchGenerator._next` (`dflash_batch_generator.py:92`) yields from `_token_buffer[uid]` first via `_yield_buffered` (`:341`) — one token per call.
+- Once the buffer is empty and a new forward is needed, `_speculative_next` fires again.
+- If `_last_target_hidden[uid]` is missing (edge case — e.g. dropped between requests), `_split_speculative_next` calls `super()._next()`. That reaches `self.model(y, cache)` with `S=1` → `_pipelined_call` S==1 branch (`model_forward.py:446–454`) → stock layer loop → `DecoderLayer.__call__` → `_split_call` (`decoder.py:36`):
+  - Step 1: ATTN runs `self_attn` or `linear_attn` on input; MOE burns `DUMMY_LN_ITERS=50` layernorms to keep JACCL balanced, evaled every 2 layers (`decoder.py:52`).
+  - 1st all_gather picks ATTN's output.
+  - Step 2: MOE runs MLP; ATTN burns layernorms.
+  - 2nd all_gather picks MOE's output.
+  - Two `all_gather`s per layer, N layers → `2N` total collectives for an S==1 decode.
+  > *Prints:* per layer per rank — `[rank N] L=lc after gather-1 h.mean=...` (`decoder.py:56`) and `[rank N] L=lc after gather-2 out.mean=...` (`:69`). 2·N lines per decode step per rank.
+  > *Evals:* per layer — `mx.eval(h)` (`:53`) on the idle rank every 2nd `lc` (JACCL queue bound); `mx.eval(h)` post-gather-1 (`:55`); `mx.eval(out)` (`:66`) idle-rank equivalent; `mx.eval(result)` post-gather-2 (`:68`).
+- At M=1 every LpB wrapper fires its fast-path kernel.
+
+### 4e. Termination
+
+- `_yield_buffered` (`dflash_batch_generator.py:341`) drains `_token_buffer[uid]` one token at a time.
+- Detects `finish_reason = "stop"` (`:348`) or `"length"` (`:350`).
+- On finish: `cache = batch.extract_cache(0)` (`:356`) — this iterates `[c.extract(idx) for c in batch.cache]`.
+  - `BatchKVCache.extract` → our `_bkv_extract` (`apply.py:165`): returns empty `KVCache` if `self.keys is None`.
+  - `ArraysCache.extract` → our `_arrays_extract` (`apply.py:130`): builds a new `ArraysCache` preserving `None` per-element (`c[idx:idx+1] if c is not None else None`). Required because MOE_RANK's GDN `base.cache[1]` stays `None` after rollback (see section 6).
+
+---
+
+## 5. Kernel dispatch matrix
+
+Rows = dispatch site. Columns = phase / M seen at call time. Cell = kernel actually dispatched. "stock" means MLX built-in GEMM via the original `nn.Linear.__call__`.
+
+| Site | Prefill (M = prompt_len) | Verify (M = V+1) | Decode (M = 1) | Drafter first cycle | Drafter steady state |
+|------|--------------------------|------------------|----------------|---------------------|----------------------|
+| target `q/k/v/o_proj` | stock | LpB | LpB | — | — |
+| target `mlp.{gate,up,down}_proj` | stock | LpB | LpB | — | — |
+| target `linear_attn.{in_proj_*, out_proj}` | stock | LpB | LpB | — | — |
+| target `gated_delta_update` | stock | **speculative** (captures all_states) | stock | — | — |
+| target `lm_head` | stock | LpB | LpB | — | — |
+| drafter `q_proj` | — | — | — | LpB (M=BS) | LpB (M=BS) |
+| drafter `k_proj, v_proj` | — | — | — | stock (M=prompt_len+BS, typically >16) | LpB (M=accepted_len+BS) |
+| drafter `o_proj` | — | — | — | LpB (M=BS) | LpB (M=BS) |
+| drafter `mlp_{gate,up,down}` | — | — | — | LpB (M=BS) | LpB (M=BS) |
+| drafter `fc` | — | — | — | stock (M=prompt_len, >16) | LpB (M=accepted_len ≤ V+1) |
+| drafter `lm_head` | — | — | — | LpB (M=BS-1) | LpB (M=BS-1) |
+
+**LpB kernel picker.** `pick_bf16_kernel` (`matmul/patches/kernel_picker.py:33`) / `pick_int8_kernel` (`:69`). Rounds M to `{1, 2, 4, 8, 12, 16, 32, 64}`. Decision summary:
+
+- bf16, N > 50000 (lm_head): `lpb` (M ≤ 8) or `lpb_twice`.
+- bf16, M_rnd ≤ 8: `lpb`.
+- bf16, M_rnd = 12: `lpb_twice` if max(N, K) ≤ 4096 else `sk_steel16`.
+- bf16, M_rnd = 16: `sk_steel16`.
+- bf16, M_rnd ∈ {32, 64}: `sk_steel32`.
+- int8, M_rnd ≤ 8: `lpb`.
+- int8, M_rnd = 12 or 16: `qsk16`.
+- int8, M_rnd ∈ {32, 64}: `qsk32`.
+
+Note: `MAX_M=16` in both LpB patches means branches for `M_rnd ∈ {32, 64}` are never hit on this path — the picker supports them, the wrappers don't call at those M.
+
+---
+
+## 6. Invariants
+
+### Shapes
+
+- `_last_target_hidden[uid]`: `(1, k, D · |target_layer_ids|)`.
+  - `k = prompt_len` right after `_first_step_capture`.
+  - `k = n_accepted + 1` after each `_speculative_next` (`dflash_split.py:197`).
+- `block_ids`: `(1, block_size)` with `block_ids[0, 0] = y_val`, rest `mask_token_id`.
+- `verify_input`: `(1, V+1)`.
+- `target_hidden` returned from verify: `(1, V+1, D · |target_layer_ids|)`.
+- `verify_logits`: `(1, V+1, vocab)`.
+- GDN `SpeculativeArraysCache.all_states`: `(B, V+1, H_v, D_v, D_k)` after the H0/H1 merge. `SpeculativeArraysCache.conv_input`: `(B, 3 + V+1, conv_dim)`.
+
+### Cache offsets
+
+- `BatchKVCache.offset == BatchKVCache._idx - left_padding`. Masks must slice against `_idx`, not `offset` (the pipelined loop uses `_idx` at `model_forward.py:185`).
+- Rollback: `BatchKVCache.offset -= (V - n_accepted)` decrements both `offset` and `_idx` via the cache API. `ArraysCache` rollback rewrites `cache[0]` and `cache[1]` from speculative state.
+- GDN base cache entries on MOE_RANK: `cache[0]` gets reconstructed by `_pipelined_dflash_forward`, `cache[1]` stays `None` because the speculative GDU never ran there. Downstream code must tolerate this — `_arrays_extract` does (`apply.py:130`).
+
+### Stage indexing in `pipelined_layer_loop` (`model_forward.py:146`)
+
+- Stage range: `0` (startup) + `1..2N-1` (main loop) + `2N` (drain) = `2N+1` stages.
+- `T = stage // 2`. B-stage = odd stage (runs `attn_T(H1)` + `moe_T(h_T_H0)`). A-stage = even stage (runs `attn_T(H0)` + `moe_{T-1}(h_{T-1}_H1)`).
+- Capture of layer L happens at the start of the A-stage for T = L+1 (i.e. stage `2(L+1)`), so `capture[L]` is set when we enter A-stage for layer L+1. Layer N-1 is captured post-drain (`model_forward.py:342`).
+- Even S: 1 all_gather per stage. Odd S: 2 all_gathers per stage (`_gather_two` at `model_forward.py:110`).
+
+### Draft-position authority
+
+- ATTN_RANK's `cache.offset` is authoritative throughout.
+- `_split_speculative_next` begins with `all_gather([_draft_position])` and takes index `ATTN_RANK` (`dflash_split.py:69`).
+
+---
+
+## 7. Environment variables
+
+| Variable | Default | File:line | Effect |
+|----------|---------|-----------|--------|
+| `EXO_SPECULATIVE` | `"0"` | `batch_generate.py:100`, `apply.py:288` | Enable speculative decode. Also gates `_drop_unused_weights` — speculative keeps all weights on both ranks. |
+| `EXO_SPECULATIVE_MODE` | `"mtp"` | `batch_generate.py:101` | `"dflash"` for this path, `"mtp"` for MTP. |
+| `EXO_SPECULATIVE_TEMP` | `"0.7"` | `batch_generate.py:103` | Sampling temperature for the drafter; 0 means greedy/deterministic. |
+| `EXO_SPECULATIVE_ALPHA` | `"1.0"` | `batch_generate.py:104` | Acceptance ratio exponent α (temp>0 only). |
+| `EXO_SPECULATIVE_GAMMA` | `"2"` | `batch_generate.py:141` | MTP draft chain length. Not used on DFlash path. |
+| `EXO_DFLASH_MODEL` | `"z-lab/Qwen3.5-27B-DFlash"` | `batch_generate.py:112` | HF repo for drafter weights. Pre-pull with `huggingface-cli download ...` to avoid first-run download. |
+| `EXO_DFLASH_VERIFY` | `"5"` | `batch_generate.py:113` | V — number of drafts per cycle. `verify_input.shape = (1, V+1)`. Even `V+1` → single-all_gather fast path in `pipelined_layer_loop`. |
+| `EXO_DFLASH_BLOCK_SIZE` | `"6"` | `batch_generate.py:114` | Drafter block size BS. Draft produces `BS-1` logits. |
+| `EXO_MTP_WEIGHTS` | `""` | `batch_generate.py:169` | Explicit MTP weights path. Not used on DFlash path. |
+| `EXO_MTP_MODEL` | `""` | `batch_generate.py:173` | HF repo for MTP extraction. Not used on DFlash path. |
+| `EXO_LPB_PATCHES` | `"1"` | `apply.py:243` | Toggle target-side LpB kernel swaps. `"0"` disables. |
+| `EXO_DISABLE_LOGPROBS` | `"0"` | `batch_generate.py:618` | Skip logprob computation (some speculative paths require this). |
+| `DFLASH_LPB_ONLY` | `""` | `bf16_lpb_patch.py:106` | Comma-separated projection names (`mlp_gate,mlp_up,...,lm_head`) to restrict drafter LpB patching. Debugging bisect knob. |
+
+---
+
+## 8. Debugging guide
+
+### 8a. Existing prints
+
+Init-time (once per process):
+
+| Print | Source | Notes |
+|-------|--------|-------|
+| `DFlash config: ...` | `dflash_module.py:133` | Drafter arch summary |
+| `Target layers: [...], mask_token=...` | `dflash_module.py:136` | **Exact `target_layer_ids` the drafter expects** |
+| `DFlash loaded: N tensors, M params` | `dflash_module.py:164` | Drafter weight load confirmation |
+| `Patched X target projections with dynamic LpB` | `lpb_patch.py:130` | Target-side LpB swap count |
+| `Patched X DFlash drafter projections with dynamic LpB` | `bf16_lpb_patch.py:143` | Drafter-side LpB swap count |
+| `DFLASH_LPB_ONLY=...` | `bf16_lpb_patch.py:109` | Only if the bisect env var is set |
+| `Qwen3.5 attn/moe split patch applied on rank N/2` | `apply.py:291` | Structural patches done |
+| `Warming up DFlash speculative decoding kernels...` / `DFlash warmup complete` | `batch_generate.py:298, 357` | Kernel warmup bookends |
+
+Per forward / per cycle (hot path — noisy, disable once stable):
+
+| Print | Source | When |
+|-------|--------|------|
+| `[rank N] pipeline begin S=... N=... even=...` | `model_forward.py:212` | Once per pipelined forward (prefill + verify + warmup verify) |
+| `[rank N] stage 0 (T=0 startup) h_H0.mean=... [rank N:role] eval_local=X.XXms eval_gather=Y.YYms` | `model_forward.py:225` | Startup stage. `role` is `attn_0(H0)` on rank 0 or `idle` on rank 1. |
+| `[rank N] stage K (T=... B) ... [rank N:role] eval_local=... eval_gather=...` | `model_forward.py:268` | Per B stage. `role` is `attn_T(H1)` on rank 0 or `moe_T(h_T_H0)` on rank 1. |
+| `[rank N] stage K (T=... A) ... [rank N:role] eval_local=... eval_gather=...` | `model_forward.py:306` | Per A stage. `role` is `attn_T(H0)` on rank 0 or `moe_{T-1}(h_{T-1}_H1)` on rank 1. |
+| `[rank N] drain out_H1.mean=... [rank N:role] eval_local=... eval_gather=...` | `model_forward.py:331` | Drain stage. `role` is `idle` on rank 0 or `moe_{N-1}(h_{N-1}_H1)` on rank 1. |
+
+**Per-stage timing.** Both numbers bracket **only `mx.eval`** — graph-building code (`attention(...)`, `moe(...)`, `all_gather(...)`) runs outside the timer, so the numbers reflect real Metal / network work, not Python overhead.
+- `eval_local` = walltime of `mx.eval(my_out)` (or `mx.eval(contribution)` at startup/drain). On the active rank: attention or MoE Metal kernels. On the idle rank (startup for MOE_RANK, drain for ATTN_RANK): a trivial `x - x` — should be near zero.
+- `eval_gather` = walltime of `mx.eval(gathered)`. This is the collective itself — network + the implicit barrier that forces both ranks to converge.
+- **Compare `eval_local` across ranks at the same stage number** to see whether attn or MoE is the limiter for that layer. The slower side's `eval_local` is the critical-path work; the faster side's `eval_gather` absorbs the difference (it sits in the collective waiting).
+| `[rank N] L=lc after gather-1 h.mean=...` | `decoder.py:56` | Per layer, mid-`_split_call` (S==1 decode) |
+| `[rank N] L=lc after gather-2 out.mean=...` | `decoder.py:69` | Per layer, end of `_split_call` |
+| `[rank N] DFlash speculative cycle (y_val=..., target_hidden.shape=...)` | `dflash_split.py:51` | Once per speculative cycle |
+| `[rank N] DFlash: NO target_hidden -> fallback ...` | `dflash_split.py:46` | When `_last_target_hidden[uid]` is missing |
+| `[DFlash] n_accepted=X/V` | `dflash_split.py:165` | MOE_RANK only, end of cycle |
+
+### 8b. `mx.eval` / `mx.async_eval` inventory
+
+Evals serve three distinct purposes in this pipeline, and removing the wrong one
+can silently break correctness or JACCL flow control. Classification:
+
+**(i) Pipeline graph-break evals.** Force the MLX graph to materialize between
+stages of `pipelined_layer_loop` and between gathers of the S==1 `_split_call`.
+Without these, MLX would fuse arbitrarily many stages into one graph, blowing
+up memory and (more importantly) desynchronizing `all_gather` order with
+JACCL's internal queue (`MAX_SEND_WR=32`). Keep these unless you're
+intentionally testing graph fusion.
+
+| File:line | Target | Location in schedule |
+|-----------|--------|----------------------|
+| `model_forward.py:221` | `contribution` | Before startup gather |
+| `model_forward.py:223` | `gathered` | After startup gather |
+| `model_forward.py:249, 251` | `my_out`, `gathered` | B stage, even-S (single gather) |
+| `model_forward.py:261, 262` | `attn_side`, `moe_side` | B stage, odd-S (pre two-gather) |
+| `model_forward.py:266, 267` | `attn_contrib`, `moe_contrib` | B stage, odd-S (post two-gather) |
+| `model_forward.py:287, 289` | `my_out`, `gathered` | A stage, even-S |
+| `model_forward.py:299, 300` | `attn_side`, `moe_side` | A stage, odd-S (pre) |
+| `model_forward.py:304, 305` | `attn_contrib`, `moe_contrib` | A stage, odd-S (post) |
+| `model_forward.py:327, 329` | `contribution`, `gathered` | Drain stage (pre / post gather) |
+| `decoder.py:53` | `h` | Idle rank's dummy layernorm, every 2nd layer — bounds JACCL queue under S==1 |
+| `decoder.py:55` | `h` | After per-layer gather-1 |
+| `decoder.py:66` | `out` | Idle rank's dummy layernorm, every 2nd layer |
+| `decoder.py:68` | `result` | After per-layer gather-2 |
+
+**(ii) Sampling evals (pre-`.item()` / `.tolist()`).** Required because Python-side
+control flow depends on concrete values (draft tokens, acceptance bits, `n_accepted`).
+
+| File:line | Target | Used for |
+|-----------|--------|----------|
+| `dflash_split.py:81` | `all_drafts_arr` | `.tolist()` to build `drafts` (temp=0) |
+| `dflash_split.py:97` | `drafts_arr` | `.tolist()` after draft-sync `all_gather` |
+| `dflash_batch_generator.py:125` | `target_hidden` | Before caching as `_last_target_hidden[uid]` (stock capture path) |
+| `dflash_batch_generator.py:145, 151` | `prompt_toks`, `(target_hidden, logits)` | "direct" prefill mode only (unused under our patched path) |
+| `dflash_batch_generator.py:199` | `all_drafts_arr` | Same role as `dflash_split.py:81`, stock path |
+
+**(iii) Async evals (kick off compute, keep the Python loop moving).** These don't
+block; MLX schedules the compute and the next `.item()` will wait. Used to
+overlap acceptance math with the next forward's dispatch.
+
+| File:line | Targets | Purpose |
+|-----------|---------|---------|
+| `dflash_split.py:120` | `matches`, `all_next`, `target_hidden` | temp=0 acceptance loop reads `matches[i].item()`; kicked off early |
+| `dflash_split.py:146–148` | `accept_ratios`, `uniforms`, `corrections`, `bonus_token`, `target_hidden` | temp>0 acceptance |
+| `dflash_split.py:239, 250` | `batch.y` | Emit next-token without blocking the `Response` return |
+| `dflash_batch_generator.py:225, 243, 327, 338` | same four roles, stock path | Stock DFlash mirror — only fires if `_speculative_next` isn't patched |
+| `mtp_batch_generator.py:115, 162, 182, 278, 291` | MTP equivalents | Not on DFlash path, here for cross-reference |
+
+Note: `mtp_module.py:340, 359` also call `mx.eval` during MTP weight load; not on our hot path.
+
+### 8c. Suggested instrumentation points
+
+- Before/after each `all_gather` in `_split_speculative_next` (`dflash_split.py:66, 95, 159`) — shape + mean.
+- `_populate_dflash_captured` (`model_forward.py:387`) — dump the set of keys written on each call.
+- `SpeculativeArraysCache.rollback` (`speculative_cache.py:81`) — print `n_accepted`, whether `all_states` / `conv_input` were non-None, and the resulting `base.cache[*]` shapes.
+- Entry to `_pipelined_dflash_forward` (`model_forward.py:623`) — print `inputs.shape`, length of `cache_list`, and whether any `SpeculativeArraysCache` already exists.
+- Warmup completion — pair with `model_forward.py:212` to confirm one pipelined forward happens during warmup at `M=V+1`.
+
+### 8d. Gotchas (bug history, one-liner each)
+
+- Mask from `create_attention_mask(..., return_array=True)` is 2D `(S, offset+S)`, not 4D — `slice_fa_mask` handles both (`model_forward.py:63`).
+- `BatchKVCache.offset` is sometimes an `mx.array`, sometimes `int`. Use `.max().item()` when unsure (`model_forward.py:188`).
+- `_idx ≠ offset` on `BatchKVCache` — mask slicing uses `_idx` (actual K buffer length), positional encodings use `offset` (net of left_padding).
+- `spec_all_states` has `2 × len(gdn_spec_data)` entries after our pipelined loop, not `len(gdn_spec_data)` — H0 and H1 each push. Merge post-loop (`model_forward.py:654`).
+- Conv rollback only works if `conv_input` is reconstructed post-loop — the pipelined path destroys the pre-layer state that stock `dflash_speculative_forward` implicitly relies on.
+- `_CapturingLayer.__call__` never fires on the pipelined path (we bypass `layer.__call__`). Compensated by `_populate_dflash_captured`.
+- `ArraysCache.extract` stock crashes on mixed None / non-None `cache[i]` entries. MOE_RANK rollback produces exactly that — `[conv_input, None]`. Fixed by per-element `_arrays_extract` (`apply.py:130`).
+- Drafter first cycle: `k_proj`, `v_proj`, `fc` see `M > MAX_M` and fall back to stock GEMM. Not a bug, but worth knowing when profiling cold-start latency.
+- `_draft_position` on MOE_RANK stays at 0 because it never runs attention. Sync via `all_gather` at the top of every cycle (`dflash_split.py:66`).
+- Drafter shares `embed_tokens` and `lm_head` with the target (`dflash_module.py:151, 153`) — weight changes to the target propagate automatically.
+
+---
+
+## 9. File index
+
+Grouped by layer. Each entry: path, purpose, key symbols.
+
+### 9a. Layer A (structural, `patches/qwen3_5_moe_split/`)
+
+| File | Purpose | Key symbols |
+|------|---------|-------------|
+| `apply.py` | Install all structural patches | `apply_qwen35_attn_moe_split_patches:220`, `_patch_caches_for_moe_rank:33`, `_arrays_extract:130`, `_drop_unused_weights:297` |
+| `decoder.py` | S==1 serial split `DecoderLayer.__call__` | `make_split_decoder_call:21`, `_split_call:36`, `ATTN_RANK:16`, `MOE_RANK:17`, `DUMMY_LN_ITERS:18` |
+| `model_forward.py` | S>1 pipelined forward + DFlash capture plumbing | `attention:38`, `moe:50`, `slice_fa_mask:63`, `pipelined_layer_loop:146`, `_dflash_capturing_target_ids:353`, `_dflash_captured_dict:366`, `_populate_dflash_captured:387`, `make_pipelined_model_call:416`, `make_pipelined_speculative_forward:491` (MTP), `make_pipelined_dflash_speculative_forward:615` |
+| `dflash_split.py` | Rank-aware `_speculative_next` for DFlash | `make_split_speculative_next:25`, `_split_speculative_next:29` |
+| `PIPELINE.md` | This document | — |
+
+### 9b. Layer B (kernel)
+
+| File | Purpose | Key symbols |
+|------|---------|-------------|
+| `patches/qwen3_5/lpb_patch.py` | Target-side LpB projection patcher | `apply_lpb_patches:102`, `_make_bf16_forward:30`, `_make_int8_forward:51`, `_patch_proj:73`, `MAX_M:27` |
+| `speculative/bf16_lpb_patch.py` | Drafter-side LpB patcher | `apply_bf16_lpb_patches:99`, `_BF16LpBLinear:27`, `_QuantizedLpBLinear:55`, `MAX_M:24` |
+| `matmul/patches/kernel_picker.py` | (N, K, M) → kernel selection | `pick_bf16_kernel:33`, `pick_int8_kernel:69`, `_round_up_to_bench_col:26` |
+| `matmul/kernels/bf16/*.py` | bf16 Metal kernels (lpb, lpb_twice, sk_steel, bm8) | `custom_bf16_qmv_loop_over_b`, `custom_bf16_qmv_loop_over_b_twice`, `custom_bf16_gemm_splitk_steel` |
+| `matmul/kernels/quantized/*.py` | int8 Metal kernels (lpb, bm8, bm16, qsk) | `custom_qmv_loop_over_b`, `custom_qmm_splitk` |
+
+### 9c. Stock DFlash (`speculative/`, unmodified)
+
+| File | Purpose | Key symbols |
+|------|---------|-------------|
+| `dflash_batch_generator.py` | BatchGenerator subclass with DFlash logic | `DFlashBatchGenerator:20`, `_setup_hidden_capture:53`, `_next:92`, `_first_step_capture:119`, `_speculative_next:170`, `_yield_buffered:341` |
+| `dflash_module.py` | DFlash drafter model | `DFlashAttention:19`, `DFlashDecoderLayer:73`, `DFlashDrafter:97`, `DFlashDrafter.draft:226`, `reset_draft_cache:186`, `crop_draft_cache:189` |
+| `dflash_speculative.py` | Stock DFlash forward (replaced by our pipelined variant) | `dflash_speculative_forward:14` |
+| `speculative_cache.py` | GDN rollback wrapper | `SpeculativeArraysCache:15`, `SpeculativeArraysCache.rollback:81` |
+| `mtp_module.py` | MTP speculative machinery + shared speculative GDU factory | `speculative_forward:27`, `_make_speculative_gdu:126`, `MTPPredictor:154` |
+| `speculative_gdn_kernel.py` | Speculative GDN Metal kernel | `speculative_gated_delta_kernel` |
+
+### 9d. Driver / entry points
+
+| File | Purpose | Key symbols |
+|------|---------|-------------|
+| `generator/batch_generate.py` | Orchestrator; drafter instantiation; warmup | DFlash init block at `:106–130`, `warmup_dflash:272`, MTP init block at `:135–170` |
+| `auto_parallel.py` | Per-strategy model setup | `attn_moe_split_auto_parallel:436`, call site `:474` |
+| `patches/__init__.py` | Single-device fused patch dispatch (not this path) | `apply_mlx_patches:14`, `maybe_apply_patches:26` |
+
+### 9e. Stock mlx_lm touch points (read-only reference)
+
+| File | Purpose |
+|------|---------|
+| `mlx_lm/models/qwen3_5.py` | `DecoderLayer`, `Qwen3_5TextModel`, `gated_delta_update`, attention modules |
+| `mlx_lm/models/cache.py` | `ArraysCache:592`, `BatchKVCache`, `KVCache`, `ArraysCache.extract:630` |
+| `mlx_lm/generate.py` | `BatchGenerator`, `Batch.extract_cache:882` |
+
+---
+
+_Document maintained on `david/attn-moe-split`. Last updated when `_arrays_extract` was made None-tolerant (commit `e8276c11`) and `_populate_dflash_captured` was added (commit `b1a5dccf`)._

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

@@ -0,0 +1,348 @@
+"""Install the attention/MoE split DecoderLayer.__call__ on a Qwen3.5 MoE model.
+
+Unlike patches/qwen3_5_moe/apply.py (which installs fused local kernels),
+this patch replaces DecoderLayer.__call__ with a two-rank split: rank 0 runs
+attention + first residual, rank 1 runs post_attention_layernorm + MoE +
+second residual, with one cross-rank send/recv pair per layer.
+
+Must be called instead of apply_qwen35_batched_fused_patches — they both
+replace DecoderLayer.__call__. In exo's current distributed load path
+(utils_mlx.shard_and_load) maybe_apply_patches is only called in the
+single-device branch, so there is no conflict for distributed AttnMoeSplit
+runs — we call this directly from attn_moe_split_auto_parallel.
+"""
+
+import os
+
+import mlx.core as mx
+import mlx.nn as nn
+from loguru import logger
+from mlx_lm.models import cache as cache_module
+from mlx_lm.models.qwen3_5 import DecoderLayer, Qwen3_5TextModel
+
+from .decoder import ATTN_RANK, MOE_RANK, make_split_decoder_call
+from .model_forward import (
+    make_pipelined_dflash_speculative_forward,
+    make_pipelined_model_call,
+    make_pipelined_speculative_forward,
+)
+
+_EMPTY_KV_PLACEHOLDER = mx.zeros((1, 1, 0, 1))
+
+
+def _patch_caches_for_moe_rank() -> None:
+    """Return zero-length placeholder arrays when caches are empty.
+
+    Rank 1 (MoE) never runs attention, so its per-layer KVCache /
+    BatchKVCache / BatchRotatingKVCache / ArraysCache never have their keys
+    populated. mlx_lm.generate and BatchGenerator both call
+    ``mx.eval([c.state for c in prompt_cache])`` once per step, which
+    invokes the cache's state property and crashes on
+    ``self.keys.shape[2]`` (AttributeError: 'NoneType' object has no
+    attribute 'shape'). Patch each used cache class's state property on
+    rank 1 only, so that an unpopulated cache reports empty placeholder
+    arrays whose eval is a no-op.
+    """
+
+    # KVCache.state → (k, v)
+    def _kv_state(self):  # type: ignore[no-untyped-def]
+        if self.keys is None:
+            return (_EMPTY_KV_PLACEHOLDER, _EMPTY_KV_PLACEHOLDER)
+        if self.offset == self.keys.shape[2]:
+            return self.keys, self.values
+        return (
+            self.keys[..., : self.offset, :],
+            self.values[..., : self.offset, :],
+        )
+
+    def _kv_state_setter(self, v):  # type: ignore[no-untyped-def]
+        self.keys, self.values = v
+        self.offset = self.keys.shape[2]
+
+    cache_module.KVCache.state = property(_kv_state, _kv_state_setter)  # type: ignore[method-assign]
+
+    # BatchKVCache.state → (k, v, offset, left_padding)
+    def _batch_kv_state(self):  # type: ignore[no-untyped-def]
+        if self.keys is None:
+            return (
+                _EMPTY_KV_PLACEHOLDER,
+                _EMPTY_KV_PLACEHOLDER,
+                self.offset,
+                self.left_padding,
+            )
+        k, v = self.keys, self.values
+        if self._idx < k.shape[2]:
+            k = k[..., : self._idx, :]
+            v = v[..., : self._idx, :]
+        return k, v, self.offset, self.left_padding
+
+    def _batch_kv_state_setter(self, v):  # type: ignore[no-untyped-def]
+        self.keys, self.values, self.offset, self.left_padding = v
+        self._idx = self.keys.shape[2]
+
+    cache_module.BatchKVCache.state = property(  # type: ignore[method-assign]
+        _batch_kv_state, _batch_kv_state_setter
+    )
+
+    # BatchRotatingKVCache.state → (k, v, offset, left_padding)
+    def _batch_rot_state(self):  # type: ignore[no-untyped-def]
+        if self.keys is None:
+            return (
+                _EMPTY_KV_PLACEHOLDER,
+                _EMPTY_KV_PLACEHOLDER,
+                self.offset,
+                self.left_padding,
+            )
+        k, v = self.keys, self.values
+        if self._offset < k.shape[2]:
+            k, v = k[..., : self._offset, :], v[..., : self._offset, :]
+        return k, v, self.offset, self.left_padding
+
+    def _batch_rot_state_setter(self, v):  # type: ignore[no-untyped-def]
+        self.keys, self.values, self.offset, self.left_padding = v
+
+    cache_module.BatchRotatingKVCache.state = property(  # type: ignore[method-assign]
+        _batch_rot_state, _batch_rot_state_setter
+    )
+
+    # ArraysCache.state → list that may contain None entries
+    original_arrays_state_getter = cache_module.ArraysCache.state.fget  # type: ignore[attr-defined]
+    original_arrays_state_setter = cache_module.ArraysCache.state.fset  # type: ignore[attr-defined]
+
+    def _arrays_state(self):  # type: ignore[no-untyped-def]
+        raw = original_arrays_state_getter(self)
+        return [_EMPTY_KV_PLACEHOLDER if c is None else c for c in raw]
+
+    cache_module.ArraysCache.state = property(  # type: ignore[method-assign]
+        _arrays_state, original_arrays_state_setter
+    )
+
+    # Batching methods (extract / filter / extend) short-circuit when the
+    # cache is still in its uninitialised state on rank 1. Every method
+    # defers to the stock implementation once keys/values are populated,
+    # which would only happen if future code actually ran attention on
+    # rank 1 — for now that never occurs, but the guards keep the patches
+    # safe under that change.
+
+    def _arrays_is_empty(self) -> bool:  # type: ignore[no-untyped-def]
+        return all(c is None for c in self.cache)
+
+    def _arrays_extract(self, idx):  # type: ignore[no-untyped-def]
+        # Stock extract does [c[idx:idx+1] for c in self.cache] which crashes
+        # on None. On MOE_RANK the GDN base.cache is mixed after speculative
+        # rollback: conv_input gets reconstructed (base.cache[0] non-None)
+        # but all_states stays None (the monkey-patched gated_delta_update
+        # never runs on MOE_RANK, so spec_all_states is empty). Handle any
+        # None/non-None mixture per-element.
+        new_cache = cache_module.ArraysCache(len(self.cache))
+        new_cache.cache = [
+            c[idx : idx + 1] if c is not None else None for c in self.cache
+        ]
+        return new_cache
+
+    cache_module.ArraysCache.extract = _arrays_extract  # type: ignore[method-assign]
+
+    _arrays_filter_orig = cache_module.ArraysCache.filter
+
+    def _arrays_filter(self, batch_indices):  # type: ignore[no-untyped-def]
+        if _arrays_is_empty(self):
+            return
+        _arrays_filter_orig(self, batch_indices)
+
+    cache_module.ArraysCache.filter = _arrays_filter  # type: ignore[method-assign]
+
+    _arrays_extend_orig = cache_module.ArraysCache.extend
+
+    def _arrays_extend(self, other):  # type: ignore[no-untyped-def]
+        if _arrays_is_empty(self) and _arrays_is_empty(other):
+            return
+        _arrays_extend_orig(self, other)
+
+    cache_module.ArraysCache.extend = _arrays_extend  # type: ignore[method-assign]
+
+    # BatchKVCache: keys/values stay None until attention runs.
+    _bkv_extract_orig = cache_module.BatchKVCache.extract
+
+    def _bkv_extract(self, idx):  # type: ignore[no-untyped-def]
+        if self.keys is None:
+            return cache_module.KVCache()
+        return _bkv_extract_orig(self, idx)
+
+    cache_module.BatchKVCache.extract = _bkv_extract  # type: ignore[method-assign]
+
+    _bkv_filter_orig = cache_module.BatchKVCache.filter
+
+    def _bkv_filter(self, batch_indices):  # type: ignore[no-untyped-def]
+        if self.keys is None:
+            return
+        _bkv_filter_orig(self, batch_indices)
+
+    cache_module.BatchKVCache.filter = _bkv_filter  # type: ignore[method-assign]
+
+    _bkv_extend_orig = cache_module.BatchKVCache.extend
+
+    def _bkv_extend(self, other):  # type: ignore[no-untyped-def]
+        if self.keys is None and other.keys is None:
+            return
+        _bkv_extend_orig(self, other)
+
+    cache_module.BatchKVCache.extend = _bkv_extend  # type: ignore[method-assign]
+
+    # BatchRotatingKVCache: same story.
+    _brkv_extract_orig = cache_module.BatchRotatingKVCache.extract
+
+    def _brkv_extract(self, idx):  # type: ignore[no-untyped-def]
+        if self.keys is None:
+            return cache_module.RotatingKVCache(self.max_size)
+        return _brkv_extract_orig(self, idx)
+
+    cache_module.BatchRotatingKVCache.extract = _brkv_extract  # type: ignore[method-assign]
+
+    _brkv_filter_orig = cache_module.BatchRotatingKVCache.filter
+
+    def _brkv_filter(self, batch_indices):  # type: ignore[no-untyped-def]
+        if self.keys is None:
+            return
+        _brkv_filter_orig(self, batch_indices)
+
+    cache_module.BatchRotatingKVCache.filter = _brkv_filter  # type: ignore[method-assign]
+
+    _brkv_extend_orig = cache_module.BatchRotatingKVCache.extend
+
+    def _brkv_extend(self, other):  # type: ignore[no-untyped-def]
+        if self.keys is None and other.keys is None:
+            return
+        _brkv_extend_orig(self, other)
+
+    cache_module.BatchRotatingKVCache.extend = _brkv_extend  # type: ignore[method-assign]
+
+
+def apply_qwen35_attn_moe_split_patches(
+    model: nn.Module, group: mx.distributed.Group
+) -> nn.Module:
+    """Install the split DecoderLayer.__call__ on a Qwen3.5 MoE model."""
+    if group.size() != 2:
+        raise ValueError(
+            f"Qwen3.5 attn/moe split requires world_size==2, got {group.size()}"
+        )
+
+    inner = model
+    for attr in ("model", "language_model"):
+        if hasattr(inner, attr):
+            inner = getattr(inner, attr)
+    if hasattr(inner, "model"):
+        inner = inner.model
+    n_layers = len(inner.layers) if hasattr(inner, "layers") else 48
+
+    # Apply target-side LPB projection patches BEFORE the split patches.
+    # Dynamic kernel picking at call time (M <= 16 uses the LpB fast path).
+    # Benefits decode (S=1) and DFlash verify (S=V+1=6); prefill falls back.
+    # Only meaningful for dense Qwen3.5 variants (e.g. 27B); on MoE variants
+    # it is mostly a no-op (MLP structure differs, only attn + lm_head patched).
+    try:
+        if os.environ.get("EXO_LPB_PATCHES", "1") != "0":
+            from exo.worker.engines.mlx.patches.qwen3_5.lpb_patch import apply_lpb_patches
+            apply_lpb_patches(model)
+    except Exception as e:
+        logger.warning(f"LPB target patches skipped: {e}")
+
+    # S == 1 serial split lives in DecoderLayer.__call__.
+    DecoderLayer.__call__ = make_split_decoder_call(group, n_layers=n_layers)  # type: ignore[method-assign]
+
+    # S > 1 pipelined split lives in Qwen3_5TextModel.__call__.
+    Qwen3_5TextModel.__call__ = make_pipelined_model_call(group)  # type: ignore[method-assign]
+
+    # Speculative verify path gets its own pipelined forward.
+    try:
+        from exo.worker.engines.mlx.speculative import mtp_module
+
+        mtp_module.speculative_forward = make_pipelined_speculative_forward(group)
+    except ImportError:
+        pass  # speculative not used
+
+    # DFlash has its own speculative_forward + batch generator step.
+    try:
+        from exo.worker.engines.mlx.speculative import dflash_speculative as _dfs
+        from exo.worker.engines.mlx.speculative.dflash_batch_generator import (
+            DFlashBatchGenerator,
+        )
+
+        _dfs.dflash_speculative_forward = make_pipelined_dflash_speculative_forward(group)
+        # The batch generator imports dflash_speculative_forward at module load,
+        # so also patch the attribute inside dflash_batch_generator.
+        import exo.worker.engines.mlx.speculative.dflash_batch_generator as _dfbg
+
+        _dfbg.dflash_speculative_forward = _dfs.dflash_speculative_forward
+
+        from .dflash_split import make_split_speculative_next
+
+        DFlashBatchGenerator._speculative_next = make_split_speculative_next(group)  # type: ignore[method-assign]
+    except ImportError:
+        pass  # DFlash not used
+
+    if group.rank() == MOE_RANK:
+        _patch_caches_for_moe_rank()
+
+    # Drop unused weights. Under speculative we keep a narrower subset on
+    # MOE_RANK because the speculative forward's post-loop reconstruction
+    # runs input_layernorm + linear_attn.in_proj_* there. ATTN_RANK always
+    # drops mlp + post_attention_layernorm since attention never runs them.
+    _drop_unused_weights(
+        model,
+        group,
+        speculative=os.environ.get("EXO_SPECULATIVE") == "1",
+    )
+
+    logger.info(
+        f"Qwen3.5 attn/moe split patch applied on rank {group.rank()}/{group.size()}"
+    )
+    return model
+
+
+def _drop_unused_weights(
+    model: nn.Module,
+    group: mx.distributed.Group,
+    speculative: bool = False,
+) -> None:
+    """Free weights each rank doesn't need.
+
+    Non-speculative:
+      ATTN_RANK drops mlp + post_attention_layernorm.
+      MOE_RANK  drops self_attn + linear_attn + input_layernorm.
+
+    Speculative (EXO_SPECULATIVE=1):
+      ATTN_RANK drops mlp + post_attention_layernorm (same as above).
+      MOE_RANK  drops self_attn only. Keeps linear_attn + input_layernorm
+                because the speculative forward's post-loop conv_input
+                reconstruction runs input_layernorm + linear_attn.in_proj_*
+                on MOE_RANK too.
+
+    embed_tokens, norm, and lm_head stay on both ranks in all cases.
+    """
+    import gc
+
+    inner = model
+    for attr in ("model", "language_model"):
+        if hasattr(inner, attr):
+            inner = getattr(inner, attr)
+    if hasattr(inner, "model"):
+        inner = inner.model
+
+    layers = inner.layers if hasattr(inner, "layers") else []
+
+    for layer in layers:
+        if group.rank() == ATTN_RANK:
+            layer.mlp = None  # type: ignore[assignment]
+            layer.post_attention_layernorm = None  # type: ignore[assignment]
+        else:
+            layer.self_attn = None  # type: ignore[assignment]
+            if not speculative:
+                layer.linear_attn = None  # type: ignore[assignment]
+                layer.input_layernorm = None  # type: ignore[assignment]
+
+    gc.collect()
+    mx.clear_cache()
+    mode = "speculative" if speculative else "full"
+    logger.info(
+        f"Rank {group.rank()}: dropped unused weights from {len(layers)} layers ({mode})"
+    )

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

@@ -0,0 +1,72 @@
+"""Serial attn/MoE split DecoderLayer.__call__ — S == 1 decode only.
+
+For S > 1 (prefill / speculative verify), the model-level pipelined
+forward in `model_forward.py` calls the layer primitives directly and
+never invokes this function.
+
+S == 1 path: ATTN rank runs attention + residual, MOE rank runs MoE +
+residual. Each step ends in one all_gather; the idle rank does dummy
+layernorm work, evaled every ~2 layers to prevent JACCL queue overflow.
+"""
+
+from collections.abc import Callable
+
+import mlx.core as mx
+
+ATTN_RANK = 0
+MOE_RANK = 1
+DUMMY_LN_ITERS = 50
+
+
+def make_split_decoder_call(
+    group: mx.distributed.Group,
+    n_layers: int = 48,
+) -> Callable[..., mx.array]:
+    """Build a DecoderLayer.__call__ replacement for S == 1 decode."""
+    if group.size() != 2:
+        raise ValueError(f"world_size==2 required, got {group.size()}")
+    rank = group.rank()
+    layer_counter = 0
+
+    def gather_from(rank_to_pick: int, tensor: mx.array) -> mx.array:
+        """all_gather and slice the tensor contributed by ``rank_to_pick``."""
+        gathered = mx.distributed.all_gather(tensor, group=group)
+        return gathered[rank_to_pick : rank_to_pick + 1]
+
+    def _split_call(self, x, mask=None, cache=None):  # type: ignore[no-untyped-def]
+        nonlocal layer_counter
+        layer_counter += 1
+        lc = layer_counter
+
+        # Step 1: ATTN does attention + residual; MOE does dummy layernorm.
+        if rank == ATTN_RANK:
+            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
+        else:
+            h = x
+            for _ in range(DUMMY_LN_ITERS):
+                h = self.post_attention_layernorm(h) + h
+            if lc % 2 == 0:
+                mx.eval(h)
+        h = gather_from(ATTN_RANK, h)
+        mx.eval(h)
+        print(f"[rank {rank}] L={lc} after gather-1  h.mean={h.mean().item():+.6f}", flush=True)
+
+        # Step 2: MOE does MoE + residual; ATTN does dummy layernorm.
+        if rank == MOE_RANK:
+            out = h + self.mlp(self.post_attention_layernorm(h))
+        else:
+            out = h
+            for _ in range(DUMMY_LN_ITERS):
+                out = self.input_layernorm(out) + out
+            if lc % 2 == 0:
+                mx.eval(out)
+        result = gather_from(MOE_RANK, out)
+        mx.eval(result)
+        print(f"[rank {rank}] L={lc} after gather-2 out.mean={result.mean().item():+.6f}", flush=True)
+        return result
+
+    return _split_call

src/exo/worker/engines/mlx/patches/qwen3_5_moe_split/dflash_split.py

@@ -0,0 +1,253 @@
+"""Rank-aware DFlashBatchGenerator patches for the attn/MoE split.
+
+Both ranks draft independently, then sync on MOE_RANK's drafts before the
+pipelined verify forward. Acceptance is deterministic at temp=0 (both ranks
+get the same verify_logits via the pipelined forward, so both compute the
+same n_accepted). At temp>0, MOE_RANK's n_accepted wins via all_gather so
+stochastic decisions stay consistent.
+
+Stock `_first_step_capture` is kept — it reads `self._captured['prefill_hiddens']`
+which is populated by `_CapturingLayer.__call__` during prefill. Our
+pipelined prefill bypasses `layer.__call__` entirely, so `_CapturingLayer`
+would never fire on its own. `make_pipelined_model_call` compensates by
+detecting the `_CapturingLayer` wrappers and writing the captured layer
+outputs from `pipelined_layer_loop` directly into the closure-captured
+`captured` dict (see `_populate_dflash_captured` in `model_forward.py`).
+"""
+
+import time
+
+import mlx.core as mx
+
+from .decoder import ATTN_RANK, MOE_RANK
+
+
+def make_split_speculative_next(group):  # type: ignore[no-untyped-def]
+    """Build a DFlashBatchGenerator._speculative_next replacement closed over group."""
+    rank = group.rank()
+
+    def _split_speculative_next(self):  # type: ignore[no-untyped-def]
+        from exo.worker.engines.mlx.speculative.dflash_speculative import (
+            dflash_speculative_forward,
+        )
+
+        tic = time.perf_counter()
+        batch = self.active_batch
+        uid = batch.uids[0]
+        y = batch.y
+        y_val = y[0].item()
+        y_logprobs = batch.logprobs[0]
+
+        batch.tokens[0] = mx.concatenate((batch.tokens[0], y[0:1]))
+
+        last_target_hidden = self._last_target_hidden.get(uid)
+        if last_target_hidden is None:
+            print(
+                f"[rank {rank}] DFlash: NO target_hidden -> fallback to super()._next() "
+                f"(y_val={y_val})",
+                flush=True,
+            )
+            return super(type(self), self)._next()
+        print(
+            f"[rank {rank}] DFlash speculative cycle "
+            f"(y_val={y_val}, target_hidden.shape={last_target_hidden.shape})",
+            flush=True,
+        )
+
+        bs = self.drafter.block_size
+        verify_len = self.verify_len
+        temp = self._request_temp.get(uid, self.temp)
+        alpha = self.alpha
+
+        # ATTN_RANK's cache offset is correct (updated by attention); MOE_RANK's
+        # stays at 0 because it never runs attention. Sync _draft_position
+        # across ranks so the drafter uses the right positional encoding.
+        local_start = self._draft_position[uid]
+        gathered = mx.distributed.all_gather(
+            mx.array([local_start], dtype=mx.int32), group=group
+        )
+        start = int(gathered[ATTN_RANK].item())
+        self._draft_position[uid] = start
+
+        # 1. Draft — both ranks draft independently
+        block_ids = mx.full((1, bs), self.drafter.mask_token_id, dtype=mx.int32)
+        block_ids[:, 0] = y_val
+        draft_logits = self.drafter.draft(last_target_hidden, block_ids, start)
+        self.drafter.crop_draft_cache(start)
+
+        # 2. Sample — both ranks sample locally
+        if temp == 0:
+            all_drafts_arr = mx.argmax(draft_logits, axis=-1).squeeze(0)
+            mx.eval(all_drafts_arr)
+            all_drafts = all_drafts_arr.tolist()
+            draft_probs = None
+        else:
+            all_drafts = []
+            draft_probs = []
+            for i in range(bs - 1):
+                p = mx.softmax(draft_logits[0, i] / temp, axis=-1)
+                tok = mx.random.categorical(mx.log(p)).item()
+                all_drafts.append(tok)
+                draft_probs.append(p)
+
+        # 3. Sync drafts: take MOE_RANK's drafts as source of truth for both ranks.
+        drafts_local = mx.array([all_drafts[:verify_len]], dtype=mx.int32)  # (1, V)
+        gathered = mx.distributed.all_gather(drafts_local, group=group)      # (2, V)
+        drafts_arr = gathered[MOE_RANK : MOE_RANK + 1]                       # (1, V)
+        mx.eval(drafts_arr)
+        drafts = drafts_arr[0].tolist()
+
+        # 4. Build verify_input on BOTH ranks (identical tokens now)
+        y_val_tensor = mx.array([[y_val]], dtype=mx.int32)
+        verify_input = mx.concatenate([y_val_tensor, drafts_arr], axis=1)
+
+        # 5. Pipelined verify forward — both ranks run the 2N+1 stage pipeline
+        #    via the patched dflash_speculative_forward.
+        target_hidden, _, verify_logits = dflash_speculative_forward(
+            self.model,
+            verify_input,
+            batch.cache,
+            self.drafter.target_layer_ids,
+            speculative=True,
+        )
+
+        # 6. Acceptance. temp==0: deterministic on both ranks.
+        #    temp>0: MOE_RANK decides, all_gather broadcasts n_accepted.
+        if temp == 0:
+            target_tokens = mx.argmax(verify_logits[:, :verify_len, :], axis=-1)
+            matches = mx.equal(target_tokens, drafts_arr).squeeze(0)
+            all_next = mx.argmax(verify_logits[0], axis=-1)
+            mx.eval(matches, all_next, target_hidden)
+            n_accepted = 0
+            for i in range(verify_len):
+                if matches[i].item():
+                    n_accepted += 1
+                else:
+                    break
+        else:
+            accept_ratios = []
+            for i in range(verify_len):
+                p = mx.softmax(verify_logits[0, i] / temp, axis=-1)
+                q = draft_probs[i]
+                p_di = p[drafts[i]]
+                q_di = mx.maximum(q[drafts[i]], mx.array(1e-10))
+                ratio = p_di / q_di
+                accept_ratios.append(mx.minimum(ratio ** alpha, mx.array(1.0)))
+            uniforms = mx.random.uniform(shape=(verify_len,))
+            corrections = []
+            for i in range(verify_len):
+                p = mx.softmax(verify_logits[0, i] / temp, axis=-1)
+                q = draft_probs[i]
+                residual = mx.maximum(p - q, 0.0)
+                corrections.append(mx.random.categorical(mx.log(residual + 1e-10)))
+            bonus_token = mx.random.categorical(
+                verify_logits[0, verify_len] * (1.0 / temp)
+            )
+            mx.async_eval(
+                accept_ratios, uniforms, corrections, bonus_token, target_hidden
+            )
+            # Only MOE_RANK decides; broadcast n_accepted.
+            if rank == MOE_RANK:
+                n_accepted_local = 0
+                for i in range(verify_len):
+                    if uniforms[i].item() < accept_ratios[i].item():
+                        n_accepted_local += 1
+                    else:
+                        break
+            else:
+                n_accepted_local = 0
+            gathered_n = mx.distributed.all_gather(
+                mx.array([n_accepted_local], dtype=mx.int32), group=group
+            )
+            n_accepted = int(gathered_n[MOE_RANK].item())
+
+        if rank == MOE_RANK:
+            print(
+                f"[DFlash] n_accepted={n_accepted}/{verify_len}",
+                flush=True,
+            )
+
+        # 7. Rollback — same n_accepted on both ranks keeps caches consistent
+        rollback = verify_len - n_accepted
+        if rollback > 0:
+            for c in batch.cache:
+                if hasattr(c, "offset"):
+                    c.offset -= rollback
+                elif hasattr(c, "rollback"):
+                    c.rollback(n_accepted)
+
+        for i, c in enumerate(batch.cache):
+            if hasattr(c, "base"):
+                batch.cache[i] = c.base
+
+        # 8. Bonus / correction token
+        no_lp = mx.array(0.0)
+        if n_accepted == verify_len:
+            if temp == 0:
+                bonus_val = all_next[verify_len].item()
+            else:
+                bonus_val = bonus_token.item()
+        else:
+            if temp == 0:
+                bonus_val = all_next[n_accepted].item()
+            else:
+                bonus_val = corrections[n_accepted].item()
+
+        # 9. Update state
+        self._last_target_hidden[uid] = target_hidden[:, : n_accepted + 1, :]
+        self._draft_position[uid] = start + n_accepted + 1
+
+        # 10. Build token list (no logprobs)
+        all_tokens = [(y_val, y_logprobs)]
+        for i in range(n_accepted):
+            all_tokens.append((drafts[i], no_lp))
+
+        batch.y = mx.array([bonus_val])
+        batch.logprobs = [no_lp]
+
+        if n_accepted > 0:
+            batch.tokens[0] = mx.concatenate(
+                (batch.tokens[0], mx.array([t for t, _ in all_tokens[1:]]))
+            )
+        batch.num_tokens[0] += len(all_tokens)
+
+        # 11. Stop conditions (same as stock)
+        toc = time.perf_counter()
+        self._stats.generation_time += toc - tic
+        self._stats.generation_tokens += len(all_tokens)
+
+        stop_idx = None
+        for idx, (tok, _) in enumerate(all_tokens):
+            if tok in self.stop_tokens:
+                stop_idx = idx
+                break
+            if batch.num_tokens[0] >= batch.max_tokens[0]:
+                stop_idx = idx
+                break
+
+        first_tok, first_lp = all_tokens[0]
+
+        if stop_idx is not None:
+            valid_tokens = all_tokens[:stop_idx]
+            if valid_tokens:
+                if len(valid_tokens) > 1:
+                    self._token_buffer[uid] = valid_tokens[1:]
+                stop_tok, stop_lp = all_tokens[stop_idx]
+                if uid not in self._token_buffer:
+                    self._token_buffer[uid] = []
+                self._token_buffer[uid].append((stop_tok, stop_lp))
+                mx.async_eval(batch.y)
+                return [self.Response(uid, first_tok, first_lp, None, lambda: None)]
+            else:
+                cache = batch.extract_cache(0)
+                self.active_batch = None
+                self._cleanup_uid(uid)
+                return [self.Response(uid, first_tok, first_lp, "stop", cache)]
+
+        if len(all_tokens) > 1:
+            self._token_buffer[uid] = all_tokens[1:]
+
+        mx.async_eval(batch.y)
+        return [self.Response(uid, first_tok, first_lp, None, lambda: None)]
+
+    return _split_speculative_next

src/exo/worker/engines/mlx/patches/qwen3_5_moe_split/model_forward.py

@@ -0,0 +1,833 @@
+"""Model-level pipelined forward for the attention/MoE split.
+
+Replaces ``Qwen3_5TextModel.__call__`` and ``mtp_module.speculative_forward``
+to implement a 2N+1 stage pipeline when S>1 (prefill / speculative verify).
+
+Pipeline schedule for N layers, two halves H0 (first S/2) and H1 (last S/2):
+
+  Stage 0 (startup) : ATTN attn_0(H0).                 MOE idle.
+  Stage 1           : ATTN attn_0(H1).                 MOE moe_0(h_0_H0).
+  Stage 2T  (T≥1)   : ATTN attn_T(H0), input out_{T-1}_H0.
+                                                        MOE moe_{T-1}(h_{T-1}_H1).
+  Stage 2T+1 (T≥1)  : ATTN attn_T(H1), input out_{T-1}_H1.
+                                                        MOE moe_T(h_T_H0).
+  Stage 2N (drain)  : ATTN idle.                        MOE moe_{N-1}(h_{N-1}_H1).
+
+Each stage: both ranks compute in parallel on different GPUs, then ONE
+all_gather shares both results (each rank contributes its real output,
+both receive both — same shape, one collective). Idle ranks contribute a
+zero placeholder of the right shape.
+
+For S==1 decode this module is bypassed — the stock layer loop + the
+serial _split_call in decoder.py handles it.
+"""
+
+import time
+from typing import Any
+
+import mlx.core as mx
+from mlx_lm.models.base import create_attention_mask, create_ssm_mask
+
+from .decoder import ATTN_RANK, MOE_RANK
+
+
+# ---------------------------------------------------------------------------
+# Primitive layer ops
+# ---------------------------------------------------------------------------
+
+
+def attention(layer, x, mask, cache):  # type: ignore[no-untyped-def]
+    """Run a layer's attention (GDN or GQA) with the input residual.
+
+    Returns h = x + attn(input_layernorm(x)).
+    """
+    if layer.is_linear:
+        r = layer.linear_attn(layer.input_layernorm(x), mask, cache)
+    else:
+        r = layer.self_attn(layer.input_layernorm(x), mask, cache)
+    return x + r
+
+
+def moe(layer, h):  # type: ignore[no-untyped-def]
+    """Run a layer's MoE with the post-attention residual.
+
+    Returns out = h + mlp(post_attention_layernorm(h)).
+    """
+    return h + layer.mlp(layer.post_attention_layernorm(h))
+
+
+# ---------------------------------------------------------------------------
+# Mask slicing
+# ---------------------------------------------------------------------------
+
+
+def slice_fa_mask(fa_mask, mid: int, offset: int):  # type: ignore[no-untyped-def]
+    """Slice the GQA attention mask for H0 / H1 queries.
+
+    create_attention_mask(return_array=True) returns a 2D causal mask of
+    shape (S, offset+S) — the last two dims of a standard 4D attention mask.
+    scaled_dot_product_attention broadcasts 2D masks automatically.
+
+    mask_H0 shape: (mid, offset+mid)     — queries [0,mid),   keys [0,offset+mid)
+    mask_H1 shape: (S-mid, offset+S)     — queries [mid,S),   keys [0,offset+S)
+
+    Also handles 4D masks (B, 1, S, offset+S) defensively.
+    """
+    if fa_mask is None:
+        return None, None
+    if fa_mask.ndim == 2:
+        return fa_mask[:mid, : offset + mid], fa_mask[mid:, :]
+    if fa_mask.ndim == 4:
+        return fa_mask[:, :, :mid, : offset + mid], fa_mask[:, :, mid:, :]
+    raise ValueError(f"unexpected fa_mask ndim={fa_mask.ndim}, shape={fa_mask.shape}")
+
+
+def slice_ssm_mask(ssm_mask, mid: int):  # type: ignore[no-untyped-def]
+    """Slice the 2D SSM/GDN mask along the sequence dim."""
+    if ssm_mask is None:
+        return None, None
+    return ssm_mask[:, :mid], ssm_mask[:, mid:]
+
+
+# ---------------------------------------------------------------------------
+# Pipeline building blocks
+# ---------------------------------------------------------------------------
+
+
+def _gather_own(my_contribution: mx.array, group) -> mx.array:
+    """all_gather + return the tensor contributed by the CURRENT rank.
+
+    Other ranks' slices are returned too but unused here — the caller runs
+    this separately on each rank's output so both sides see both tensors.
+    """
+    return mx.distributed.all_gather(my_contribution, group=group)
+
+
+def _zeros_like(x: mx.array) -> mx.array:
+    """Return x - x (zeros with a data dependency on x so MLX can't const-fold)."""
+    return x - x
+
+
+def _gather_two(
+    rank: int,
+    attn_side: mx.array,
+    moe_side: mx.array,
+    attn_shape_template: mx.array,
+    moe_shape_template: mx.array,
+    group,
+) -> tuple[mx.array, mx.array]:
+    """Two all_gathers per stage — one per shape.
+
+    Each rank contributes a placeholder of the OTHER rank's shape. One
+    all_gather for ATTN's output, one for MOE's. Handles odd S where
+    H0 and H1 have different sizes.
+
+    Args:
+      attn_side: what this rank is contributing for ATTN's all_gather.
+                 On ATTN_RANK this is the real attention output; on
+                 MOE_RANK it is a placeholder shaped like ATTN's output.
+      moe_side: same but for MOE's all_gather.
+
+    Returns (attn_result, moe_result) — both ranks get both after the
+    two collectives.
+    """
+    attn_gathered = mx.distributed.all_gather(attn_side, group=group)
+    moe_gathered = mx.distributed.all_gather(moe_side, group=group)
+    return (
+        attn_gathered[ATTN_RANK : ATTN_RANK + 1],
+        moe_gathered[MOE_RANK : MOE_RANK + 1],
+    )
+
+
+# ---------------------------------------------------------------------------
+# Pipelined layer loop
+# ---------------------------------------------------------------------------
+
+
+def pipelined_layer_loop(
+    inner,  # Qwen3_5TextModel instance
+    hidden_states: mx.array,
+    cache,  # list of per-layer caches
+    group,  # mx.distributed.Group
+    fa_mask,  # 4D tensor or None (after return_array=True)
+    ssm_mask,  # 2D tensor or None
+    capture_layers=None,  # set[int] | None — DFlash target_layer_ids
+):
+    """Execute all decoder layers with the 2N+1 stage H0/H1 pipeline.
+
+    Args:
+      capture_layers: if not None, snapshot each listed layer's full
+        (B, S, D) output into a returned dict. Used by DFlash to collect
+        target_hidden.
+
+    Returns:
+      (final, captured) where final = concat(out_H0, out_H1) — layer N-1's
+      full output — and captured = {T: concat(out_T_H0, out_T_H1) for T
+      in capture_layers} or {} if capture_layers is None.
+    """
+    rank = group.rank()
+    layers = inner.layers
+    N = len(layers)
+    S = hidden_states.shape[1]
+    mid = S // 2
+    even_S = (S % 2 == 0)   # True -> single-gather fast path
+    capture: dict[int, mx.array] = {}
+    capture_set: set[int] = set(capture_layers) if capture_layers is not None else set()
+
+    # Compute cache offset for mask slicing (same for all GQA layers).
+    # For BatchKVCache, update_and_fetch returns keys[:_idx], so the mask must
+    # match _idx (actual K buffer length), NOT `offset` (which is _idx minus
+    # left_padding for positional encodings). Plain KVCache uses `offset`.
+    fa_cache = cache[inner.fa_idx] if cache[inner.fa_idx] is not None else None
+    if fa_cache is None:
+        offset = 0
+    elif hasattr(fa_cache, "_idx"):
+        # BatchKVCache: _idx is Python int tracking actual K length
+        offset = int(fa_cache._idx)
+    elif hasattr(fa_cache, "offset"):
+        raw = fa_cache.offset
+        offset = int(raw.max().item()) if isinstance(raw, mx.array) else int(raw)
+    else:
+        offset = 0
+
+    # Slice masks once (reused every stage)
+    fa_mask_H0, fa_mask_H1 = slice_fa_mask(fa_mask, mid, offset)
+    ssm_mask_H0, ssm_mask_H1 = slice_ssm_mask(ssm_mask, mid)
+
+    def mask_for(layer, half: str):  # type: ignore[no-untyped-def]
+        if layer.is_linear:
+            return ssm_mask_H0 if half == "H0" else ssm_mask_H1
+        else:
+            return fa_mask_H0 if half == "H0" else fa_mask_H1
+
+    # Running state between stages (loop-local, no closure).
+    x_H0 = hidden_states[:, :mid, :]     # (B, mid, D)
+    x_H1 = hidden_states[:, mid:, :]     # (B, S-mid, D)
+
+    # Pending handoffs:
+    #   h_H0_ready:   output of most recent attn(H0) — awaiting moe
+    #   h_H1_pending: output of most recent attn(H1) — awaiting moe
+    h_H0_ready = None
+    h_H1_pending = None
+
+    print(f"[rank {rank}] pipeline begin S={S} N={N} even={even_S}", flush=True)
+
+    # --- Stage 0 (startup bubble): ATTN attn_0(H0). MOE idle. ---
+    layer_0 = layers[0]
+    c0 = cache[0]
+    if rank == ATTN_RANK:
+        contribution = attention(layer_0, x_H0, mask_for(layer_0, "H0"), c0)
+    else:
+        contribution = _zeros_like(x_H0)
+    _t0 = time.perf_counter()
+    mx.eval(contribution)
+    _t_eval_local_ms = (time.perf_counter() - _t0) * 1000.0
+    gathered = mx.distributed.all_gather(contribution, group=group)
+    _t0 = time.perf_counter()
+    mx.eval(gathered)
+    _t_eval_gather_ms = (time.perf_counter() - _t0) * 1000.0
+    h_H0_ready = gathered[ATTN_RANK : ATTN_RANK + 1]
+    _role = "attn_0(H0)" if rank == ATTN_RANK else "idle"
+    print(
+        f"[rank {rank}] stage 0 (T=0 startup) "
+        f"[rank {rank}:{_role}] eval_local={_t_eval_local_ms:.2f}ms eval_gather={_t_eval_gather_ms:.2f}ms",
+        flush=True,
+    )
+
+    # --- Stages 1..2N-1: main pipeline ---
+    # Even S: attn_side and moe_side have the same shape (mid == S-mid), so
+    #         we do ONE all_gather per stage — each rank contributes its real
+    #         output, both get both.
+    # Odd S:  shapes differ, so we do TWO all_gathers per stage — each rank
+    #         contributes a zero placeholder of the OTHER side's shape.
+    for stage in range(1, 2 * N):
+        is_B_stage = (stage % 2 == 1)
+        T = stage // 2
+
+        if is_B_stage:
+            # Stage 2T+1: ATTN attn_T(H1) (B, S-mid, D) | MOE moe_T(h_T_H0) (B, mid, D)
+            layer_T = layers[T]
+            cT = cache[T]
+            if even_S:
+                if rank == ATTN_RANK:
+                    my_out = attention(layer_T, x_H1, mask_for(layer_T, "H1"), cT)
+                else:
+                    my_out = moe(layer_T, h_H0_ready)
+                _t0 = time.perf_counter()
+                mx.eval(my_out)
+                _t_eval_local_ms = (time.perf_counter() - _t0) * 1000.0
+                gathered = mx.distributed.all_gather(my_out, group=group)
+                _t0 = time.perf_counter()
+                mx.eval(gathered)
+                _t_eval_gather_ms = (time.perf_counter() - _t0) * 1000.0
+                attn_contrib = gathered[ATTN_RANK : ATTN_RANK + 1]
+                moe_contrib = gathered[MOE_RANK : MOE_RANK + 1]
+            else:
+                if rank == ATTN_RANK:
+                    attn_side = attention(layer_T, x_H1, mask_for(layer_T, "H1"), cT)
+                    moe_side = _zeros_like(h_H0_ready)
+                else:
+                    attn_side = _zeros_like(x_H1)
+                    moe_side = moe(layer_T, h_H0_ready)
+                _t0 = time.perf_counter()
+                mx.eval(attn_side)
+                mx.eval(moe_side)
+                _t_eval_local_ms = (time.perf_counter() - _t0) * 1000.0
+                attn_contrib, moe_contrib = _gather_two(
+                    rank, attn_side, moe_side, attn_side, moe_side, group
+                )
+                _t0 = time.perf_counter()
+                mx.eval(attn_contrib)
+                mx.eval(moe_contrib)
+                _t_eval_gather_ms = (time.perf_counter() - _t0) * 1000.0
+            _role = f"attn_{T}(H1)" if rank == ATTN_RANK else f"moe_{T}(h_{T}_H0)"
+            print(
+                f"[rank {rank}] stage {stage} (T={T} B) "
+                f"[rank {rank}:{_role}] eval_local={_t_eval_local_ms:.2f}ms eval_gather={_t_eval_gather_ms:.2f}ms",
+                flush=True,
+            )
+
+            h_H1_pending = attn_contrib
+            x_H0 = moe_contrib
+        else:
+            # Stage 2T (T>=1): ATTN attn_T(H0) (B, mid, D) | MOE moe_{T-1}(h_{T-1}_H1) (B, S-mid, D)
+            layer_T = layers[T]
+            cT = cache[T]
+            prev_layer = layers[T - 1]
+            if even_S:
+                if rank == ATTN_RANK:
+                    my_out = attention(layer_T, x_H0, mask_for(layer_T, "H0"), cT)
+                else:
+                    my_out = moe(prev_layer, h_H1_pending)
+                _t0 = time.perf_counter()
+                mx.eval(my_out)
+                _t_eval_local_ms = (time.perf_counter() - _t0) * 1000.0
+                gathered = mx.distributed.all_gather(my_out, group=group)
+                _t0 = time.perf_counter()
+                mx.eval(gathered)
+                _t_eval_gather_ms = (time.perf_counter() - _t0) * 1000.0
+                attn_contrib = gathered[ATTN_RANK : ATTN_RANK + 1]
+                moe_contrib = gathered[MOE_RANK : MOE_RANK + 1]
+            else:
+                if rank == ATTN_RANK:
+                    attn_side = attention(layer_T, x_H0, mask_for(layer_T, "H0"), cT)
+                    moe_side = _zeros_like(h_H1_pending)
+                else:
+                    attn_side = _zeros_like(x_H0)
+                    moe_side = moe(prev_layer, h_H1_pending)
+                _t0 = time.perf_counter()
+                mx.eval(attn_side)
+                mx.eval(moe_side)
+                _t_eval_local_ms = (time.perf_counter() - _t0) * 1000.0
+                attn_contrib, moe_contrib = _gather_two(
+                    rank, attn_side, moe_side, attn_side, moe_side, group
+                )
+                _t0 = time.perf_counter()
+                mx.eval(attn_contrib)
+                mx.eval(moe_contrib)
+                _t_eval_gather_ms = (time.perf_counter() - _t0) * 1000.0
+            _role = f"attn_{T}(H0)" if rank == ATTN_RANK else f"moe_{T - 1}(h_{T - 1}_H1)"
+            print(
+                f"[rank {rank}] stage {stage} (T={T} A) "
+                f"[rank {rank}:{_role}] eval_local={_t_eval_local_ms:.2f}ms eval_gather={_t_eval_gather_ms:.2f}ms",
+                flush=True,
+            )
+
+            h_H0_ready = attn_contrib
+            x_H1 = moe_contrib
+
+            # Capture layer T-1's full output: x_H0 = out_{T-1}_H0 (set in
+            # previous B stage) and x_H1 = out_{T-1}_H1 (just set above).
+            if (T - 1) in capture_set:
+                capture[T - 1] = mx.concatenate([x_H0, x_H1], axis=1)
+
+    # --- Stage 2N (drain bubble): ATTN idle. MOE moe_{N-1}(h_{N-1}_H1). ---
+    last_layer = layers[N - 1]
+    if rank == MOE_RANK:
+        contribution = moe(last_layer, h_H1_pending)
+    else:
+        contribution = _zeros_like(h_H1_pending)
+    _t0 = time.perf_counter()
+    mx.eval(contribution)
+    _t_eval_local_ms = (time.perf_counter() - _t0) * 1000.0
+    gathered = mx.distributed.all_gather(contribution, group=group)
+    _t0 = time.perf_counter()
+    mx.eval(gathered)
+    _t_eval_gather_ms = (time.perf_counter() - _t0) * 1000.0
+    out_H1 = gathered[MOE_RANK : MOE_RANK + 1]
+    _role = f"moe_{N - 1}(h_{N - 1}_H1)" if rank == MOE_RANK else "idle"
+    print(
+        f"[rank {rank}] drain "
+        f"[rank {rank}:{_role}] eval_local={_t_eval_local_ms:.2f}ms eval_gather={_t_eval_gather_ms:.2f}ms",
+        flush=True,
+    )
+
+    # After the final B stage (Stage 2N-1), out_{N-1}_H0 = x_H0 (MOE's contribution
+    # from Stage 2N-1 — stored in the x_H0 variable).
+    out_H0 = x_H0
+    final = mx.concatenate([out_H0, out_H1], axis=1)
+
+    # Capture layer N-1's output if requested (only now is out_{N-1}_H1 known).
+    if (N - 1) in capture_set:
+        capture[N - 1] = final
+
+    return final, capture
+
+
+# ---------------------------------------------------------------------------
+# DFlash _CapturingLayer detection / population
+# ---------------------------------------------------------------------------
+
+
+def _dflash_capturing_target_ids(inner) -> set[int]:
+    """Return layer indices wrapped by DFlashBatchGenerator._CapturingLayer.
+
+    Detected by the presence of both ``_orig`` and ``_layer_idx`` attributes —
+    the two instance attrs set in _CapturingLayer.__init__.
+    """
+    ids: set[int] = set()
+    for layer in inner.layers:
+        if hasattr(layer, "_orig") and hasattr(layer, "_layer_idx"):
+            ids.add(int(layer._layer_idx))
+    return ids
+
+
+def _dflash_captured_dict(inner):  # type: ignore[no-untyped-def]
+    """Return the ``captured`` dict shared by all _CapturingLayer instances.
+
+    _CapturingLayer.__call__ writes to a closure variable ``captured`` which
+    is a reference to DFlashBatchGenerator._captured. We fish it out via
+    ``__call__.__closure__`` so we can populate it from the pipelined path
+    (which bypasses _CapturingLayer.__call__ entirely).
+    """
+    for layer in inner.layers:
+        if hasattr(layer, "_orig") and hasattr(layer, "_layer_idx"):
+            call_fn = type(layer).__call__
+            closure = getattr(call_fn, "__closure__", None)
+            if closure is None:
+                return None
+            for name, cell in zip(call_fn.__code__.co_freevars, closure):
+                if name == "captured":
+                    return cell.cell_contents
+            return None
+    return None
+
+
+def _populate_dflash_captured(
+    inner, layer_hiddens: dict, S: int
+) -> None:  # type: ignore[no-untyped-def]
+    """Write captured layer outputs into the DFlash ``_captured`` dict.
+
+    Mirrors _CapturingLayer.__call__'s behavior: always writes to
+    ``layer_hiddens``, and additionally to ``prefill_hiddens`` when S > 1.
+    """
+    if not layer_hiddens:
+        return
+    captured = _dflash_captured_dict(inner)
+    if captured is None:
+        return
+    if "layer_hiddens" not in captured:
+        captured["layer_hiddens"] = {}
+    for idx, out in layer_hiddens.items():
+        captured["layer_hiddens"][idx] = out
+    if S > 1:
+        if "prefill_hiddens" not in captured:
+            captured["prefill_hiddens"] = {}
+        for idx, out in layer_hiddens.items():
+            captured["prefill_hiddens"][idx] = out
+
+
+# ---------------------------------------------------------------------------
+# Replacement for Qwen3_5TextModel.__call__
+# ---------------------------------------------------------------------------
+
+
+def make_pipelined_model_call(group):  # type: ignore[no-untyped-def]
+    """Build a Qwen3_5TextModel.__call__ replacement closed over ``group``.
+
+    For S==1 decode, falls through to the stock layer loop which calls
+    DecoderLayer.__call__ — i.e. our serial _split_call. For S>1, runs the
+    pipelined layer loop.
+
+    DFlash integration: if DFlashBatchGenerator._setup_hidden_capture has
+    wrapped any decoder layer in _CapturingLayer, we detect it and populate
+    its closure-captured ``captured`` dict with layer outputs directly —
+    because the pipelined path bypasses layer.__call__ and _CapturingLayer
+    never fires on its own during prefill.
+    """
+
+    def _pipelined_call(
+        self,
+        inputs: mx.array,
+        cache=None,
+        input_embeddings=None,
+    ) -> mx.array:
+        if input_embeddings is not None:
+            hidden_states = input_embeddings
+        else:
+            hidden_states = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        S = hidden_states.shape[1]
+
+        if S == 1:
+            # Decode: stock loop -> DecoderLayer.__call__ -> serial _split_call.
+            # _CapturingLayer.__call__ fires naturally here for S==1.
+            fa_mask = create_attention_mask(hidden_states, cache[self.fa_idx])
+            ssm_mask = create_ssm_mask(hidden_states, cache[self.ssm_idx])
+            for layer, c in zip(self.layers, cache):
+                mask = ssm_mask if layer.is_linear else fa_mask
+                hidden_states = layer(hidden_states, mask=mask, cache=c)
+            return self.norm(hidden_states)
+
+        # S > 1: pipelined path. Force fa_mask as a real tensor for slicing.
+        fa_mask = create_attention_mask(
+            hidden_states, cache[self.fa_idx], return_array=True
+        )
+        ssm_mask = create_ssm_mask(hidden_states, cache[self.ssm_idx])
+
+        # If DFlash has wrapped target layers, request their outputs so we can
+        # populate its captured dict (stock _CapturingLayer.__call__ never
+        # fires on the pipelined path).
+        target_ids = _dflash_capturing_target_ids(self)
+        capture_layers = target_ids if target_ids else None
+
+        hidden_states, layer_hiddens = pipelined_layer_loop(
+            self,
+            hidden_states,
+            cache,
+            group,
+            fa_mask,
+            ssm_mask,
+            capture_layers=capture_layers,
+        )
+
+        if target_ids:
+            _populate_dflash_captured(self, layer_hiddens, S)
+
+        return self.norm(hidden_states)
+
+    return _pipelined_call
+
+
+# ---------------------------------------------------------------------------
+# Replacement for exo.speculative.mtp_module.speculative_forward
+# ---------------------------------------------------------------------------
+
+
+def make_pipelined_speculative_forward(group):  # type: ignore[no-untyped-def]
+    """Build a speculative_forward replacement that uses pipelined_layer_loop.
+
+    Same pre-loop setup as the original (SpeculativeArraysCache wrapping,
+    gated_delta_update monkey-patch) and same post-loop GDN state capture,
+    but swaps the layer loop for our pipelined version.
+    """
+
+    def _pipelined_speculative_forward(
+        model, inputs: mx.array, cache, speculative: bool = False
+    ):  # type: ignore[no-untyped-def]
+        inner = getattr(model, "model", None) or model.language_model.model
+        text_model = getattr(model, "model", None) or model.language_model
+        S = inputs.shape[1]
+        do_spec = speculative and S > 1
+
+        if hasattr(inner, "embed_tokens"):
+            hidden_states = inner.embed_tokens(inputs)
+        else:
+            hidden_states = inputs
+
+        cache_list: list[Any] = cache if cache is not None else [None] * len(inner.layers)
+
+        gdn_spec_data: list[Any] = []
+        if do_spec:
+            from exo.worker.engines.mlx.speculative.speculative_cache import (
+                SpeculativeArraysCache,
+            )
+
+            for i, c in enumerate(cache_list):
+                if c is not None and hasattr(c, "cache") and not hasattr(c, "offset"):
+                    cache_list[i] = SpeculativeArraysCache(c, S=S)
+            if cache is not None:
+                for i in range(len(cache)):
+                    cache[i] = cache_list[i]
+
+        spec_all_states: list[Any] = []
+        _orig_gdu = None
+        if do_spec:
+            from exo.worker.engines.mlx.speculative.mtp_module import (
+                _make_speculative_gdu,
+            )
+            import mlx_lm.models.qwen3_5 as _qwen3_5_mod
+
+            _orig_gdu = _qwen3_5_mod.gated_delta_update
+            _qwen3_5_mod.gated_delta_update = _make_speculative_gdu(spec_all_states)
+
+        fa_mask = create_attention_mask(
+            hidden_states, cache_list[inner.fa_idx], return_array=(S > 1)
+        )
+        ssm_mask = create_ssm_mask(hidden_states, cache_list[inner.ssm_idx])
+
+        if do_spec:
+            # Capture per-GDN-layer data for the post-loop state reconstruction
+            # (matches original speculative_forward lines 80-89).
+            from exo.worker.engines.mlx.speculative.speculative_cache import (
+                SpeculativeArraysCache as _SAC,
+            )
+
+            for layer, c in zip(inner.layers, cache_list):
+                if layer.is_linear and isinstance(c, _SAC):
+                    pre_conv = c[0]
+                    if pre_conv is None:
+                        gdn = layer.linear_attn
+                        pre_conv = mx.zeros(
+                            (
+                                hidden_states.shape[0],
+                                gdn.conv_kernel_size - 1,
+                                gdn.conv_dim,
+                            ),
+                            dtype=hidden_states.dtype,
+                        )
+                    gdn_spec_data.append((None, pre_conv, c, layer))
+
+        # Pipelined layer loop for S>1; stock layer loop for S==1.
+        if S > 1:
+            hidden_states, _ = pipelined_layer_loop(
+                inner, hidden_states, cache_list, group, fa_mask, ssm_mask
+            )
+        else:
+            for layer, c in zip(inner.layers, cache_list):
+                mask = ssm_mask if layer.is_linear else fa_mask
+                hidden_states = layer(hidden_states, mask=mask, cache=c)
+
+        if do_spec:
+            import mlx_lm.models.qwen3_5 as _qwen3_5_mod
+
+            _qwen3_5_mod.gated_delta_update = _orig_gdu
+
+            gdn_idx = 0
+            for _layer_input, pre_conv, spec_cache, parent_layer in gdn_spec_data:
+                if gdn_idx < len(spec_all_states):
+                    spec_cache.all_states = spec_all_states[gdn_idx]
+                gdn_idx += 1
+
+                gdn = parent_layer.linear_attn
+                # Recover the layer input from the original hidden_states before
+                # the layer loop — but we don't have it anymore because the
+                # pipeline overwrote it. We store the original at loop start.
+                # For now we reconstruct conv_input from pre_conv + qkv of the
+                # ORIGINAL embeddings (pre-layer). TODO revisit this.
+                # Keeping the conv_input reconstruction as-is for now but
+                # conv_input will NOT be valid after the pipelined forward.
+                # This is acceptable only if rollback path isn't exercised.
+                pass  # conv_input reconstruction skipped under split pipeline
+
+        pre_norm = hidden_states
+        normed = inner.norm(hidden_states)
+
+        if hasattr(text_model, "lm_head"):
+            logits = text_model.lm_head(normed)
+        else:
+            logits = inner.embed_tokens.as_linear(normed)
+
+        return pre_norm, logits
+
+    return _pipelined_speculative_forward
+
+
+# ---------------------------------------------------------------------------
+# Replacement for exo.speculative.dflash_speculative.dflash_speculative_forward
+# ---------------------------------------------------------------------------
+
+
+def make_pipelined_dflash_speculative_forward(group):  # type: ignore[no-untyped-def]
+    """Build a dflash_speculative_forward replacement using pipelined_layer_loop.
+
+    Differences from make_pipelined_speculative_forward:
+    - Captures hidden states at ``target_layer_ids`` for the DFlash drafter
+    - Returns (target_hidden, pre_norm, logits) instead of (pre_norm, logits)
+    """
+
+    def _pipelined_dflash_forward(
+        model,
+        inputs: mx.array,
+        cache,
+        target_layer_ids,
+        speculative: bool = False,
+    ):  # type: ignore[no-untyped-def]
+        inner = getattr(model, "model", None) or model.language_model.model
+        text_model = getattr(model, "model", None) or model.language_model
+        S = inputs.shape[1]
+        do_spec = speculative and S > 1
+
+        if hasattr(inner, "embed_tokens"):
+            hidden_states = inner.embed_tokens(inputs)
+        else:
+            hidden_states = inputs
+
+        cache_list: list[Any] = (
+            cache if cache is not None else [None] * len(inner.layers)
+        )
+
+        # Wrap GDN caches for rollback
+        if do_spec:
+            from exo.worker.engines.mlx.speculative.speculative_cache import (
+                SpeculativeArraysCache,
+            )
+
+            for i, c in enumerate(cache_list):
+                if c is not None and hasattr(c, "cache") and not hasattr(c, "offset"):
+                    cache_list[i] = SpeculativeArraysCache(c, S=S)
+            if cache is not None:
+                for i in range(len(cache)):
+                    cache[i] = cache_list[i]
+
+        # Swap in speculative GDN kernel
+        spec_all_states: list[Any] = []
+        _orig_gdu = None
+        if do_spec:
+            from exo.worker.engines.mlx.speculative.mtp_module import (
+                _make_speculative_gdu,
+            )
+            import mlx_lm.models.qwen3_5 as _qwen3_5_mod
+
+            _orig_gdu = _qwen3_5_mod.gated_delta_update
+            _qwen3_5_mod.gated_delta_update = _make_speculative_gdu(spec_all_states)
+
+        fa_mask = create_attention_mask(
+            hidden_states, cache_list[inner.fa_idx], return_array=(S > 1)
+        )
+        ssm_mask = create_ssm_mask(hidden_states, cache_list[inner.ssm_idx])
+
+        # Collect per-GDN-layer data for post-loop conv_input reconstruction
+        gdn_spec_data: list[Any] = []
+        if do_spec:
+            from exo.worker.engines.mlx.speculative.speculative_cache import (
+                SpeculativeArraysCache as _SAC,
+            )
+            for idx, (layer, c) in enumerate(zip(inner.layers, cache_list)):
+                if layer.is_linear and isinstance(c, _SAC):
+                    pre_conv = c[0]
+                    if pre_conv is None:
+                        gdn = layer.linear_attn
+                        pre_conv = mx.zeros(
+                            (
+                                hidden_states.shape[0],
+                                gdn.conv_kernel_size - 1,
+                                gdn.conv_dim,
+                            ),
+                            dtype=hidden_states.dtype,
+                        )
+                    # Tuple order matches stock: (pre_conv, spec_cache, layer, layer_idx).
+                    # layer_idx is used post-loop to look up the layer's input.
+                    gdn_spec_data.append((pre_conv, c, layer, idx))
+
+        # Run layer loop; capture hidden states at target_layer_ids + GDN
+        # layer INPUTS (which equal layer L-1's outputs) for conv rollback.
+        # Initial embed is layer 0's input; we pass it explicitly below.
+        initial_embed = hidden_states
+        gdn_layer_idxs = [
+            i for i, layer in enumerate(inner.layers) if layer.is_linear
+        ]
+        # For GDN layer L, we need layer L's input = layer L-1's output (or
+        # initial embed if L=0). Request capture of outputs for indices L-1.
+        extra_capture = {L - 1 for L in gdn_layer_idxs if L >= 1}
+        capture_set = set(target_layer_ids) | extra_capture
+
+        if S > 1:
+            hidden_states, layer_hiddens = pipelined_layer_loop(
+                inner,
+                hidden_states,
+                cache_list,
+                group,
+                fa_mask,
+                ssm_mask,
+                capture_layers=capture_set,
+            )
+        else:
+            # S==1 path: stock loop with manual capture
+            layer_hiddens = {}
+            for i, (layer, c) in enumerate(zip(inner.layers, cache_list)):
+                mask = ssm_mask if layer.is_linear else fa_mask
+                hidden_states = layer(hidden_states, mask=mask, cache=c)
+                if i in capture_set:
+                    layer_hiddens[i] = hidden_states
+
+        # Post-loop: restore GDN kernel, reconstruct conv_input for rollback
+        if do_spec:
+            import mlx_lm.models.qwen3_5 as _qwen3_5_mod
+
+            _qwen3_5_mod.gated_delta_update = _orig_gdu
+
+            # Cross-layer pipeline calls GDN.linear_attn TWICE per GDN layer
+            # (once for H0, once for H1 with S>1). The monkey-patched
+            # gated_delta_update appends per-step states on each call, so
+            # spec_all_states has 2*N_gdn entries: [H0 states of layer 0,
+            # H1 states of layer 0, H0 states of layer 1, H1 states of layer 1, ...].
+            # Merge consecutive pairs so one entry per GDN layer.
+            # S==1 case: only one call per layer — no merging needed.
+            merged_states: list[Any] = []
+            if S > 1 and len(spec_all_states) == 2 * len(gdn_spec_data):
+                for i in range(0, len(spec_all_states), 2):
+                    # all_states shape: (B, step_count, H_v, D_v, D_k).
+                    # Concat H0 (step_count=mid) and H1 (step_count=S-mid) along step dim.
+                    merged_states.append(
+                        mx.concatenate(
+                            [spec_all_states[i], spec_all_states[i + 1]], axis=1
+                        )
+                    )
+            else:
+                merged_states = spec_all_states
+
+            gdn_idx = 0
+            for pre_conv, spec_cache, parent_layer, layer_idx in gdn_spec_data:
+                if gdn_idx < len(merged_states):
+                    spec_cache.all_states = merged_states[gdn_idx]
+                gdn_idx += 1
+
+                # Reconstruct conv_input for rollback. Mirrors stock
+                # dflash_speculative.py:97-113 but retrieves layer_input from
+                # our captured layer outputs (capture[L-1] is layer L's input)
+                # or the initial embedding for layer 0.
+                if layer_idx == 0:
+                    layer_input = initial_embed
+                else:
+                    layer_input = layer_hiddens.get(layer_idx - 1)
+                if layer_input is None:
+                    continue  # shouldn't happen given our capture_set, safety
+
+                gdn = parent_layer.linear_attn
+                normed = parent_layer.input_layernorm(layer_input)
+                if hasattr(gdn, "in_proj_qkv"):
+                    qkv = gdn.in_proj_qkv(normed)
+                else:
+                    q, k, v, z, b, a = gdn.fix_query_key_value_ordering(
+                        gdn.in_proj_qkvz(normed), gdn.in_proj_ba(normed)
+                    )
+                    B_dim = normed.shape[0]
+                    qkv = mx.concatenate(
+                        [
+                            q.reshape(B_dim, S, -1),
+                            k.reshape(B_dim, S, -1),
+                            v.reshape(B_dim, S, -1),
+                        ],
+                        axis=-1,
+                    )
+                spec_cache.conv_input = mx.concatenate([pre_conv, qkv], axis=1)
+
+        # Concatenate target_hidden from captured layers
+        selected = [layer_hiddens[i] for i in target_layer_ids]
+        target_hidden = mx.concatenate(selected, axis=-1)
+
+        pre_norm = hidden_states
+        normed = inner.norm(hidden_states)
+
+        if hasattr(text_model, "lm_head"):
+            logits = text_model.lm_head(normed)
+        else:
+            logits = inner.embed_tokens.as_linear(normed)
+
+        return target_hidden, pre_norm, logits
+
+    return _pipelined_dflash_forward

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

@@ -0,0 +1,144 @@
+#!/usr/bin/env python3
+"""Dynamic per-call loop-over-B patches for the DFlash draft model.
+
+For each patched projection, the kernel is picked at CALL time based on
+the actual M seen in the forward pass (via matmul/patches/kernel_picker).
+Each projection memoizes its kernel cache, so after warmup the pick is
+O(1) dict lookup.
+
+This is correct for any (block_size, verify_len) combo because each
+projection sees a different M (e.g. k_proj sees M = S_ctx + block_size,
+mlp_gate sees M = block_size, fc sees M = S_ctx, lm_head sees M = block_size - 1).
+"""
+
+import os
+
+import mlx.nn as nn
+
+from exo.worker.engines.mlx.matmul.patches.kernel_picker import (
+    pick_bf16_kernel,
+    pick_int8_kernel,
+)
+
+
+MAX_M = 16  # Above this (prefill), fall back to the original projection
+
+
+class _BF16LpBLinear:
+    """nn.Linear drop-in with dynamic per-call kernel selection."""
+
+    def __init__(self, original, N, K):
+        self._orig = original  # kept for prefill fallback and weight access
+        self.weight = original.weight
+        self._N = N
+        self._K = K
+        self._cache = {}       # M → fn
+        self._name_log = {}    # M → kernel name (debug)
+
+    def __call__(self, x):
+        M = 1
+        for d in x.shape[:-1]:
+            M *= d
+        if M > MAX_M:
+            return self._orig(x)
+        fn = self._cache.get(M)
+        if fn is None:
+            name, fn = pick_bf16_kernel(self._N, self._K, M)
+            self._cache[M] = fn
+            self._name_log[M] = name
+        orig_shape = x.shape
+        x_2d = x.reshape(-1, self._K)
+        y = fn(x_2d, self.weight, M, self._N, self._K)
+        return y.reshape(*orig_shape[:-1], self._N)
+
+
+class _QuantizedLpBLinear:
+    """nn.QuantizedLinear drop-in with dynamic per-call kernel selection."""
+
+    def __init__(self, original, N, K, GS):
+        self._orig = original
+        self.weight = original.weight
+        self.scales = original.scales
+        self.biases = original.biases
+        self._N = N
+        self._K = K
+        self._GS = GS
+        self._cache = {}
+        self._name_log = {}
+
+    def __call__(self, x):
+        M = 1
+        for d in x.shape[:-1]:
+            M *= d
+        if M > MAX_M:
+            return self._orig(x)
+        fn = self._cache.get(M)
+        if fn is None:
+            name, fn = pick_int8_kernel(self._N, self._K, M)
+            self._cache[M] = fn
+            self._name_log[M] = name
+        orig_shape = x.shape
+        x_2d = x.reshape(-1, self._K)
+        y = fn(x_2d, self.weight, self.scales, self.biases, M, self._N, self._K, self._GS)
+        return y.reshape(*orig_shape[:-1], self._N)
+
+
+def _wrap(proj):
+    if isinstance(proj, nn.QuantizedLinear):
+        N = proj.weight.shape[0]
+        K = proj.weight.shape[1] * (32 // proj.bits)
+        GS = proj.group_size
+        return _QuantizedLpBLinear(proj, N, K, GS)
+    elif isinstance(proj, nn.Linear):
+        N = proj.weight.shape[0]
+        K = proj.weight.shape[1]
+        return _BF16LpBLinear(proj, N, K)
+    return proj
+
+
+def apply_bf16_lpb_patches(drafter):
+    """Patch DFlash drafter's projections with dynamic LpB kernels.
+
+    Set env DFLASH_LPB_ONLY to a comma-separated list of projection names
+    to patch only those (for bisecting which patch breaks things).
+    Names: mlp_gate, mlp_up, mlp_down, q_proj, k_proj, v_proj, o_proj, fc, lm_head
+    """
+    only_str = os.environ.get("DFLASH_LPB_ONLY", "")
+    only = set(only_str.split(",")) if only_str else None
+    if only:
+        print(f"  DFLASH_LPB_ONLY={only_str}")
+
+    def should_patch(name):
+        return only is None or name in only
+
+    patched = 0
+
+    for layer in drafter.layers:
+        for attr in ('mlp_gate', 'mlp_up', 'mlp_down'):
+            if not should_patch(attr):
+                continue
+            proj = getattr(layer, attr, None)
+            if proj is not None and isinstance(proj, nn.Linear):
+                setattr(layer, attr, _wrap(proj))
+                patched += 1
+
+        for attr in ('q_proj', 'k_proj', 'v_proj', 'o_proj'):
+            if not should_patch(attr):
+                continue
+            proj = getattr(layer.self_attn, attr, None)
+            if proj is not None and isinstance(proj, nn.Linear):
+                setattr(layer.self_attn, attr, _wrap(proj))
+                patched += 1
+
+    if should_patch('fc') and hasattr(drafter, 'fc') and isinstance(drafter.fc, nn.Linear):
+        drafter.fc = _wrap(drafter.fc)
+        patched += 1
+
+    if should_patch('lm_head') and drafter.lm_head is not None and isinstance(
+        drafter.lm_head, (nn.Linear, nn.QuantizedLinear)
+    ):
+        drafter.lm_head = _wrap(drafter.lm_head)
+        patched += 1
+
+    print(f"  Patched {patched} DFlash drafter projections with dynamic LpB")
+    return patched

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

@@ -0,0 +1,371 @@
+#!/usr/bin/env python3
+"""DFlash Speculative Decoding integrated with mlx_lm's BatchGenerator.
+
+Two prefill approaches (PREFILL_MODE env var):
+  "capture" (default): Use _CapturingLayer to intercept hidden states during
+                       super()._next() prefill. No extra forward pass.
+  "direct":           Run dflash_speculative_forward for prefill ourselves,
+                       bypassing super()._next() for the first step.
+"""
+
+import time
+
+import mlx.core as mx
+from mlx_lm.generate import BatchGenerator
+
+from .dflash_module import DFlashDrafter
+from .dflash_speculative import dflash_speculative_forward
+
+
+class DFlashBatchGenerator(BatchGenerator):
+    """BatchGenerator with DFlash speculative decoding for BS=1."""
+
+    def __init__(
+        self,
+        model,
+        drafter: DFlashDrafter,
+        verify_len: int = 5,
+        block_size: int = 6,
+        temp: float = 0.0,
+        alpha: float = 1.0,
+        prefill_mode: str = "capture",
+        **kwargs,
+    ):
+        super().__init__(model, **kwargs)
+        self.drafter = drafter
+        self.verify_len = verify_len
+        self.temp = temp
+        self.alpha = alpha
+        self.prefill_mode = prefill_mode
+
+        drafter.block_size = block_size
+
+        self._token_buffer = {}
+        self._last_target_hidden = {}
+        self._draft_position = {}
+        self._prefilled = set()
+        self._request_temp = {}
+        self._captured = {}
+
+        if prefill_mode == "capture":
+            self._setup_hidden_capture()
+
+    def _setup_hidden_capture(self):
+        """Replace target layers with wrappers that capture hidden states."""
+        inner = getattr(self.model, 'model', None) or self.model.language_model.model
+        target_ids = set(self.drafter.target_layer_ids)
+        captured = self._captured
+
+        class _CapturingLayer:
+            def __init__(self, orig, layer_idx):
+                self._orig = orig
+                self._layer_idx = layer_idx
+
+            def __call__(self, *args, **kwargs):
+                out = self._orig(*args, **kwargs)
+                if 'layer_hiddens' not in captured:
+                    captured['layer_hiddens'] = {}
+                captured['layer_hiddens'][self._layer_idx] = out
+                if out.shape[1] > 1:
+                    if 'prefill_hiddens' not in captured:
+                        captured['prefill_hiddens'] = {}
+                    captured['prefill_hiddens'][self._layer_idx] = out
+                return out
+
+            def __getattr__(self, name):
+                return getattr(self._orig, name)
+
+        for i in range(len(inner.layers)):
+            if i in target_ids:
+                inner.layers[i] = _CapturingLayer(inner.layers[i], i)
+
+    def _get_captured_target_hidden(self, key='layer_hiddens'):
+        hiddens = self._captured.get(key)
+        if hiddens is None:
+            return None
+        selected = [hiddens[i] for i in self.drafter.target_layer_ids if i in hiddens]
+        if len(selected) != len(self.drafter.target_layer_ids):
+            return None
+        self._captured[key] = {}
+        return mx.concatenate(selected, axis=-1)
+
+    def _next(self):
+        batch = self.active_batch
+
+        # Yield buffered tokens first
+        if batch is not None and len(batch) == 1:
+            uid = batch.uids[0]
+            if uid in self._token_buffer and self._token_buffer[uid]:
+                return self._yield_buffered(batch, uid)
+
+        # BS=1 speculative path
+        if (batch is not None
+                and len(batch) == 1
+                and self.verify_len > 0
+                and len(self.unprocessed_prompts) == 0):
+            uid = batch.uids[0]
+            if uid not in self._prefilled:
+                if self.prefill_mode == "direct":
+                    return self._first_step_direct(batch, uid)
+                else:
+                    return self._first_step_capture(batch, uid)
+            return self._speculative_next()
+
+        # Standard path (BS>1 or no batch)
+        return super()._next()
+
+    # ── Approach 1: "capture" ──
+
+    def _first_step_capture(self, batch, uid):
+        """Use prefill hidden states captured by _CapturingLayer during super()._next()."""
+        target_hidden = self._get_captured_target_hidden(key='prefill_hiddens')
+        if target_hidden is None:
+            target_hidden = self._get_captured_target_hidden(key='layer_hiddens')
+        if target_hidden is not None:
+            mx.eval(target_hidden)
+            self._last_target_hidden[uid] = target_hidden
+
+        for c in batch.cache:
+            if hasattr(c, 'offset'):
+                off = c.offset
+                self._draft_position[uid] = off.item() if hasattr(off, 'item') else int(off)
+                break
+        self.drafter.reset_draft_cache()
+        self._prefilled.add(uid)
+        return self._speculative_next()
+
+    # ── Approach 2: "direct" ──
+
+    def _first_step_direct(self, batch, uid):
+        """Run dflash_speculative_forward ourselves for prefill, then first speculative cycle."""
+        # batch.y has the first token from super()._next()'s prefill
+        # But we need target hidden from that prefill. Run dflash_speculative_forward
+        # on the prompt tokens to get them.
+        prompt_toks = batch.tokens[0]  # full prompt token history
+        mx.eval(prompt_toks)
+
+        # Run target model forward on prompt to capture multi-layer hidden
+        target_hidden, _, logits = dflash_speculative_forward(
+            self.model, prompt_toks.reshape(1, -1), batch.cache,
+            self.drafter.target_layer_ids, speculative=False)
+        mx.eval(target_hidden, logits)
+
+        self._last_target_hidden[uid] = target_hidden
+
+        # First token from logits
+        first_token = mx.argmax(logits[0, -1], axis=-1).item()
+        batch.y = mx.array([first_token])
+
+        for c in batch.cache:
+            if hasattr(c, 'offset'):
+                off = c.offset
+                self._draft_position[uid] = off.item() if hasattr(off, 'item') else int(off)
+                break
+        self.drafter.reset_draft_cache()
+        self._prefilled.add(uid)
+        return self._speculative_next()
+
+    # ── Core speculative cycle ──
+
+    def _speculative_next(self):
+        tic = time.perf_counter()
+        batch = self.active_batch
+        uid = batch.uids[0]
+        y = batch.y
+        y_val = y[0].item()
+        y_logprobs = batch.logprobs[0]
+
+        batch.tokens[0] = mx.concatenate((batch.tokens[0], y[0:1]))
+
+        last_target_hidden = self._last_target_hidden.get(uid)
+        if last_target_hidden is None:
+            return super()._next()
+
+        bs = self.drafter.block_size
+        verify_len = self.verify_len
+        temp = self._request_temp.get(uid, self.temp)
+        alpha = self.alpha
+        start = self._draft_position[uid]
+
+        # 1. Draft
+        block_ids = mx.full((1, bs), self.drafter.mask_token_id, dtype=mx.int32)
+        block_ids[:, 0] = y_val
+        draft_logits = self.drafter.draft(last_target_hidden, block_ids, start)
+        self.drafter.crop_draft_cache(start)
+
+        # 2. Sample
+        if temp == 0:
+            all_drafts_arr = mx.argmax(draft_logits, axis=-1).squeeze(0)
+            mx.eval(all_drafts_arr)
+            all_drafts = all_drafts_arr.tolist()
+            draft_probs = None
+        else:
+            all_drafts = []
+            draft_probs = []
+            for i in range(bs - 1):
+                p = mx.softmax(draft_logits[0, i] / temp, axis=-1)
+                tok = mx.random.categorical(mx.log(p)).item()
+                all_drafts.append(tok)
+                draft_probs.append(p)
+
+        drafts = all_drafts[:verify_len]
+
+        # 3. Verify
+        verify_input = mx.array([[y_val] + drafts])
+        target_hidden, _, verify_logits = dflash_speculative_forward(
+            self.model, verify_input, batch.cache,
+            self.drafter.target_layer_ids, speculative=True)
+
+        # Build acceptance lazily (no logprobs — skip expensive logsumexp on 248K vocab)
+        if temp == 0:
+            target_tokens = mx.argmax(verify_logits[:, :verify_len, :], axis=-1)
+            draft_arr = mx.array([drafts])
+            matches = mx.equal(target_tokens, draft_arr).squeeze(0)
+            all_next = mx.argmax(verify_logits[0], axis=-1)
+            mx.async_eval(matches, all_next, target_hidden)
+        else:
+            accept_ratios = []
+            for i in range(verify_len):
+                p = mx.softmax(verify_logits[0, i] / temp, axis=-1)
+                q = draft_probs[i]
+                p_di = p[drafts[i]]
+                q_di = mx.maximum(q[drafts[i]], mx.array(1e-10))
+                ratio = p_di / q_di
+                accept_ratios.append(mx.minimum(ratio ** alpha, mx.array(1.0)))
+            uniforms = mx.random.uniform(shape=(verify_len,))
+            corrections = []
+            for i in range(verify_len):
+                p = mx.softmax(verify_logits[0, i] / temp, axis=-1)
+                q = draft_probs[i]
+                residual = mx.maximum(p - q, 0.0)
+                corrections.append(mx.random.categorical(mx.log(residual + 1e-10)))
+            bonus_token = mx.random.categorical(verify_logits[0, verify_len] * (1.0 / temp))
+            mx.async_eval(accept_ratios, uniforms, corrections, bonus_token, target_hidden)
+
+        # 4. Accept
+        n_accepted = 0
+        for i in range(verify_len):
+            if temp == 0:
+                if matches[i].item():
+                    n_accepted += 1
+                else:
+                    break
+            else:
+                if uniforms[i].item() < accept_ratios[i].item():
+                    n_accepted += 1
+                else:
+                    break
+
+        print(f"[DFlash] n_accepted={n_accepted}/{verify_len}", flush=True)
+
+        # 5. Rollback
+        rollback = verify_len - n_accepted
+        if rollback > 0:
+            for c in batch.cache:
+                if hasattr(c, 'offset'):
+                    c.offset -= rollback
+                elif hasattr(c, 'rollback'):
+                    c.rollback(n_accepted)
+
+        for i, c in enumerate(batch.cache):
+            if hasattr(c, 'base'):
+                batch.cache[i] = c.base
+
+        # 6. Bonus/correction
+        no_lp = mx.array(0.0)  # placeholder — no logprobs computed
+        if n_accepted == verify_len:
+            if temp == 0:
+                bonus_val = all_next[verify_len].item()
+            else:
+                bonus_val = bonus_token.item()
+        else:
+            if temp == 0:
+                bonus_val = all_next[n_accepted].item()
+            else:
+                bonus_val = corrections[n_accepted].item()
+
+        # 7. Update state
+        self._last_target_hidden[uid] = target_hidden[:, :n_accepted + 1, :]
+        self._draft_position[uid] = start + n_accepted + 1
+
+        # 8. Build token list (no logprobs)
+        all_tokens = [(y_val, y_logprobs)]
+        for i in range(n_accepted):
+            all_tokens.append((drafts[i], no_lp))
+
+        batch.y = mx.array([bonus_val])
+        batch.logprobs = [no_lp]
+
+        if n_accepted > 0:
+            batch.tokens[0] = mx.concatenate(
+                (batch.tokens[0], mx.array([t for t, _ in all_tokens[1:]])))
+        batch.num_tokens[0] += len(all_tokens)
+
+        # 9. Stop conditions
+        toc = time.perf_counter()
+        self._stats.generation_time += toc - tic
+        self._stats.generation_tokens += len(all_tokens)
+
+        stop_idx = None
+        for idx, (tok, _) in enumerate(all_tokens):
+            if tok in self.stop_tokens:
+                stop_idx = idx
+                break
+            if batch.num_tokens[0] >= batch.max_tokens[0]:
+                stop_idx = idx
+                break
+
+        first_tok, first_lp = all_tokens[0]
+
+        if stop_idx is not None:
+            valid_tokens = all_tokens[:stop_idx]
+            if valid_tokens:
+                if len(valid_tokens) > 1:
+                    self._token_buffer[uid] = valid_tokens[1:]
+                stop_tok, stop_lp = all_tokens[stop_idx]
+                if uid not in self._token_buffer:
+                    self._token_buffer[uid] = []
+                self._token_buffer[uid].append((stop_tok, stop_lp))
+                mx.async_eval(batch.y)
+                return [self.Response(uid, first_tok, first_lp, None, lambda: None)]
+            else:
+                cache = batch.extract_cache(0)
+                self.active_batch = None
+                self._cleanup_uid(uid)
+                return [self.Response(uid, first_tok, first_lp, "stop", cache)]
+
+        if len(all_tokens) > 1:
+            self._token_buffer[uid] = all_tokens[1:]
+
+        mx.async_eval(batch.y)
+        return [self.Response(uid, first_tok, first_lp, None, lambda: None)]
+
+    def _yield_buffered(self, batch, uid):
+        tic = time.perf_counter()
+        buf = self._token_buffer[uid]
+        tok, lp = buf.pop(0)
+        if not buf:
+            del self._token_buffer[uid]
+
+        finish_reason = None
+        if tok in self.stop_tokens:
+            finish_reason = "stop"
+        elif batch.num_tokens[0] >= batch.max_tokens[0]:
+            finish_reason = "length"
+
+        cache = None
+        if finish_reason:
+            cache = batch.extract_cache(0)
+            self.active_batch = None
+            self._cleanup_uid(uid)
+
+        toc = time.perf_counter()
+        self._stats.generation_time += toc - tic
+        return [self.Response(uid, tok, lp, finish_reason, cache or (lambda: None))]
+
+    def _cleanup_uid(self, uid):
+        self._last_target_hidden.pop(uid, None)
+        self._draft_position.pop(uid, None)
+        self._prefilled.discard(uid)
+        self._token_buffer.pop(uid, None)
+        self._request_temp.pop(uid, None)

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

@@ -0,0 +1,263 @@
+#!/usr/bin/env python3
+"""DFlash draft model for Qwen3.5 on MLX.
+
+Port of z-lab's DFlash (Block Diffusion for Flash Speculative Decoding).
+A 5-layer bidirectional transformer that generates all draft tokens in one
+parallel forward pass, conditioned on target model hidden states.
+
+Reference: z-lab/dflash benchmark.py
+"""
+
+import json
+import os
+import mlx.core as mx
+import mlx.nn as nn
+from mlx_lm.models.cache import KVCache
+from mlx_lm.models.activations import swiglu
+
+
+class DFlashAttention(nn.Module):
+    """Bidirectional attention with target context injection.
+
+    Q from draft tokens. K/V from cat(target_context, draft_tokens).
+    Non-causal: every draft position attends to every other + context.
+    """
+
+    def __init__(self, hidden_size, num_heads, num_kv_heads, head_dim, rope_theta,
+                 rms_norm_eps=1e-6):
+        super().__init__()
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.scale = head_dim ** -0.5
+
+        self.q_proj = nn.Linear(hidden_size, num_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, num_kv_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, num_kv_heads * head_dim, bias=False)
+        self.o_proj = nn.Linear(num_heads * head_dim, hidden_size, bias=False)
+        self.q_norm = nn.RMSNorm(head_dim, eps=rms_norm_eps)
+        self.k_norm = nn.RMSNorm(head_dim, eps=rms_norm_eps)
+        self.rope = nn.RoPE(head_dim, base=rope_theta)
+
+    def __call__(self, hidden_states, target_hidden, q_offset, cache=None):
+        B, S, _ = hidden_states.shape
+
+        # Q from draft only
+        q = self.q_norm(
+            self.q_proj(hidden_states).reshape(B, S, self.num_heads, self.head_dim)
+        ).transpose(0, 2, 1, 3)
+
+        # K/V from cat(context, draft) — single projection each
+        kv_input = mx.concatenate([target_hidden, hidden_states], axis=1)
+        kv_len = kv_input.shape[1]
+        k = self.k_norm(
+            self.k_proj(kv_input).reshape(B, kv_len, self.num_kv_heads, self.head_dim)
+        ).transpose(0, 2, 1, 3)
+        v = self.v_proj(kv_input).reshape(
+            B, kv_len, self.num_kv_heads, self.head_dim
+        ).transpose(0, 2, 1, 3)
+
+        # RoPE — fused mx.fast.rope, 1 dispatch each
+        q = self.rope(q, offset=q_offset)
+        k = self.rope(k, offset=cache.offset)
+
+        # KV cache — pre-allocated buffer, no concat copies
+        keys, values = cache.update_and_fetch(k, v)
+
+        # Bidirectional attention (no causal mask)
+        output = mx.fast.scaled_dot_product_attention(
+            q, keys, values, scale=self.scale)
+        output = output.transpose(0, 2, 1, 3).reshape(B, S, -1)
+        return self.o_proj(output)
+
+
+class DFlashDecoderLayer(nn.Module):
+    def __init__(self, hidden_size, num_heads, num_kv_heads, head_dim,
+                 intermediate_size, rope_theta, rms_norm_eps=1e-6):
+        super().__init__()
+        self.self_attn = DFlashAttention(
+            hidden_size, num_heads, num_kv_heads, head_dim, rope_theta,
+            rms_norm_eps=rms_norm_eps)
+        self.mlp_gate = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.mlp_up = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.mlp_down = nn.Linear(intermediate_size, hidden_size, bias=False)
+        self.input_layernorm = nn.RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(hidden_size, eps=rms_norm_eps)
+
+    def __call__(self, hidden_states, target_hidden, q_offset, cache=None):
+        residual = hidden_states
+        h = self.input_layernorm(hidden_states)
+        h = self.self_attn(h, target_hidden, q_offset, cache)
+        h = residual + h
+
+        residual = h
+        h = self.post_attention_layernorm(h)
+        h = residual + self.mlp_down(swiglu(self.mlp_gate(h), self.mlp_up(h)))
+        return h
+
+
+class DFlashDrafter:
+    """DFlash draft model for speculative decoding.
+
+    Maintains a KV cache for the draft model across iterations.
+    Each iteration:
+      1. Receives target_hidden for accepted positions
+      2. Processes [accepted_token, MASK, ...] with bidirectional attention
+      3. Draft model catches up on accepted positions via KV cache
+    """
+
+    def __init__(self, target_model, dflash_model_path):
+        self.target = target_model
+        self._inner = getattr(target_model, 'model', None) or target_model.language_model.model
+        self._text_model = getattr(target_model, 'model', None) or target_model.language_model
+
+        # Load config
+        if os.path.isdir(dflash_model_path):
+            model_dir = dflash_model_path
+        else:
+            from huggingface_hub import snapshot_download
+            model_dir = snapshot_download(dflash_model_path)
+
+        with open(os.path.join(model_dir, 'config.json')) as f:
+            config = json.load(f)
+
+        self.hidden_size = config['hidden_size']
+        self.num_heads = config['num_attention_heads']
+        self.num_kv_heads = config['num_key_value_heads']
+        self.head_dim = config.get('head_dim', self.hidden_size // self.num_heads)
+        self.intermediate_size = config['intermediate_size']
+        self.num_layers = config['num_hidden_layers']
+        self.block_size = config['block_size']
+        self.target_layer_ids = config['dflash_config']['target_layer_ids']
+        self.mask_token_id = config['dflash_config']['mask_token_id']
+        self.rope_theta = config.get('rope_theta', 10000000)
+        self.rms_norm_eps = config.get('rms_norm_eps', 1e-6)
+
+        print(f"  DFlash config: {self.num_layers} layers, hidden={self.hidden_size}, "
+              f"heads={self.num_heads}/{self.num_kv_heads}, head_dim={self.head_dim}, "
+              f"block={self.block_size}, rms_norm_eps={self.rms_norm_eps}")
+        print(f"  Target layers: {self.target_layer_ids}, mask_token={self.mask_token_id}")
+
+        # Build layers
+        self.layers = []
+        for i in range(self.num_layers):
+            self.layers.append(DFlashDecoderLayer(
+                self.hidden_size, self.num_heads, self.num_kv_heads,
+                self.head_dim, self.intermediate_size, self.rope_theta,
+                rms_norm_eps=self.rms_norm_eps))
+
+        n_target = len(self.target_layer_ids)
+        self.fc = nn.Linear(n_target * self.hidden_size, self.hidden_size, bias=False)
+        self.hidden_norm = nn.RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
+        self.norm = nn.RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
+
+        # Shared from target
+        self.embed_tokens = self._inner.embed_tokens
+        if hasattr(self._text_model, 'lm_head'):
+            self.lm_head = self._text_model.lm_head
+        else:
+            self.lm_head = None
+
+        # Draft KV cache: KVCache per layer (pre-allocated buffers)
+        self.draft_cache = None
+
+        # Load weights
+        weights = mx.load(os.path.join(model_dir, 'model.safetensors'))
+        self._load_weights(weights)
+        total_params = sum(w.size for w in weights.values())
+        print(f"  DFlash loaded: {len(weights)} tensors, {total_params / 1e6:.1f}M params")
+
+    def _load_weights(self, weights):
+        self.fc.weight = weights['fc.weight']
+        self.hidden_norm.weight = weights['hidden_norm.weight']
+        self.norm.weight = weights['norm.weight']
+
+        for i, layer in enumerate(self.layers):
+            p = f'layers.{i}'
+            attn = layer.self_attn
+            attn.q_proj.weight = weights[f'{p}.self_attn.q_proj.weight']
+            attn.k_proj.weight = weights[f'{p}.self_attn.k_proj.weight']
+            attn.v_proj.weight = weights[f'{p}.self_attn.v_proj.weight']
+            attn.o_proj.weight = weights[f'{p}.self_attn.o_proj.weight']
+            attn.q_norm.weight = weights[f'{p}.self_attn.q_norm.weight']
+            attn.k_norm.weight = weights[f'{p}.self_attn.k_norm.weight']
+            layer.mlp_gate.weight = weights[f'{p}.mlp.gate_proj.weight']
+            layer.mlp_up.weight = weights[f'{p}.mlp.up_proj.weight']
+            layer.mlp_down.weight = weights[f'{p}.mlp.down_proj.weight']
+            layer.input_layernorm.weight = weights[f'{p}.input_layernorm.weight']
+            layer.post_attention_layernorm.weight = weights[f'{p}.post_attention_layernorm.weight']
+
+    def reset_draft_cache(self):
+        self.draft_cache = [KVCache() for _ in range(self.num_layers)]
+
+    def crop_draft_cache(self, length):
+        """Crop draft KV cache to given length (discard speculative entries)."""
+        for cache in self.draft_cache:
+            cache.offset = length
+
+    def get_target_hidden(self, inputs, cache):
+        """Run target model, capture hidden states at specified layers + logits."""
+        model = self._inner
+        if hasattr(model, 'embed_tokens'):
+            hidden_states = model.embed_tokens(inputs)
+        else:
+            hidden_states = inputs
+
+        cache_list = cache if cache is not None else [None] * len(model.layers)
+
+        from mlx_lm.models.qwen3_5 import create_attention_mask, create_ssm_mask
+        fa_mask = create_attention_mask(hidden_states, cache_list[model.fa_idx])
+        ssm_mask = create_ssm_mask(hidden_states, cache_list[model.ssm_idx])
+
+        layer_hiddens = {}
+        for i, (layer, c) in enumerate(zip(model.layers, cache_list)):
+            mask = ssm_mask if layer.is_linear else fa_mask
+            hidden_states = layer(hidden_states, mask=mask, cache=c)
+            if i in self.target_layer_ids:
+                layer_hiddens[i] = hidden_states
+
+        selected = [layer_hiddens[i] for i in self.target_layer_ids]
+        target_hidden = mx.concatenate(selected, axis=-1)
+
+        normed = model.norm(hidden_states)
+        if self.lm_head is not None:
+            logits = self.lm_head(normed)
+        else:
+            logits = self.embed_tokens.as_linear(normed)
+
+        return target_hidden, logits
+
+    def draft(self, target_hidden, block_output_ids, start):
+        """Run DFlash draft model.
+
+        Args:
+            target_hidden: (B, accepted_len, n_layers * hidden) — hidden states for
+                           accepted positions from the previous verify step
+            block_output_ids: (B, block_size) — [accepted_token, MASK, ..., MASK]
+            start: int — position of the first token in the block
+        Returns:
+            draft_logits: (B, block_size-1, vocab) — logits for mask positions
+        """
+        bs = self.block_size
+
+        # Noise embedding from target's embed_tokens
+        noise_embedding = self.embed_tokens(block_output_ids)  # (B, bs, hidden)
+
+        # Project target hidden states
+        ctx = self.hidden_norm(self.fc(target_hidden))  # (B, accepted_len, hidden)
+
+        # Run through layers
+        h = noise_embedding
+        for i, layer in enumerate(self.layers):
+            h = layer(h, ctx, q_offset=start, cache=self.draft_cache[i])
+
+        # Final norm + lm_head on mask positions (1..bs-1)
+        h = self.norm(h)
+        draft_hidden = h[:, -(bs - 1):, :]
+
+        if self.lm_head is not None:
+            draft_logits = self.lm_head(draft_hidden)
+        else:
+            draft_logits = self.embed_tokens.as_linear(draft_hidden)
+
+        return draft_logits

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

@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+"""DFlash-specific speculative forward pass.
+
+Combines:
+1. Multi-layer hidden state capture (for DFlash draft context)
+2. GDN SpeculativeArraysCache wrapping + speculative kernel swap (for rollback)
+
+Single forward pass returns (target_hidden, pre_norm, logits) with rollback-ready caches.
+"""
+
+import mlx.core as mx
+
+
+def dflash_speculative_forward(model, inputs, cache, target_layer_ids, speculative=False):
+    """Run model forward, capture multi-layer hidden states + optional GDN rollback.
+
+    Args:
+        model: the loaded model
+        inputs: (B, S) int token ids
+        cache: cache list from make_prompt_cache()
+        target_layer_ids: list of layer indices to capture hidden states from
+        speculative: if True, wraps GDN caches for rollback
+
+    Returns:
+        (target_hidden, pre_norm, logits)
+        - target_hidden: (B, S, n_layers * hidden) — concatenated hidden states from target layers
+        - pre_norm: (B, S, hidden) — pre-RMSNorm hidden states
+        - logits: (B, S, vocab) — output logits
+    """
+    from .mtp_module import _make_speculative_gdu
+
+    inner = getattr(model, 'model', None) or model.language_model.model
+    text_model = getattr(model, 'model', None) or model.language_model
+    S = inputs.shape[1]
+    do_spec = speculative and S > 1
+
+    if hasattr(inner, 'embed_tokens'):
+        hidden_states = inner.embed_tokens(inputs)
+    else:
+        hidden_states = inputs
+
+    cache_list = cache if cache is not None else [None] * len(inner.layers)
+
+    # GDN rollback setup
+    gdn_spec_data = []
+    if do_spec:
+        from .speculative_cache import SpeculativeArraysCache
+        for i, c in enumerate(cache_list):
+            if c is not None and hasattr(c, 'cache') and not hasattr(c, 'offset'):
+                cache_list[i] = SpeculativeArraysCache(c, S=S)
+        if cache is not None:
+            for i in range(len(cache)):
+                cache[i] = cache_list[i]
+
+    spec_all_states = []
+    if do_spec:
+        import mlx_lm.models.qwen3_5 as _qwen3_5_mod
+        _orig_gdu = _qwen3_5_mod.gated_delta_update
+        _qwen3_5_mod.gated_delta_update = _make_speculative_gdu(spec_all_states)
+
+    from mlx_lm.models.qwen3_5 import create_attention_mask, create_ssm_mask
+    fa_mask = create_attention_mask(hidden_states, cache_list[inner.fa_idx])
+    ssm_mask = create_ssm_mask(hidden_states, cache_list[inner.ssm_idx])
+
+    # Layer loop — capture multi-layer hidden + GDN pre-conv state
+    target_layer_ids_set = set(target_layer_ids)
+    layer_hiddens = {}
+
+    for i, (layer, c) in enumerate(zip(inner.layers, cache_list)):
+        mask = ssm_mask if layer.is_linear else fa_mask
+
+        if do_spec and layer.is_linear:
+            from .speculative_cache import SpeculativeArraysCache as _SAC
+            if isinstance(c, _SAC):
+                pre_conv = c[0]
+                if pre_conv is None:
+                    gdn = layer.linear_attn
+                    pre_conv = mx.zeros(
+                        (hidden_states.shape[0], gdn.conv_kernel_size - 1,
+                         gdn.conv_dim), dtype=hidden_states.dtype)
+                gdn_spec_data.append((hidden_states, pre_conv, c, layer))
+
+        hidden_states = layer(hidden_states, mask=mask, cache=c)
+
+        # Capture hidden states at target layers
+        if i in target_layer_ids_set:
+            layer_hiddens[i] = hidden_states
+
+    # Restore original kernel + distribute all_states + conv_input
+    if do_spec:
+        _qwen3_5_mod.gated_delta_update = _orig_gdu
+
+        gdn_idx = 0
+        for layer_input, pre_conv, spec_cache, parent_layer in gdn_spec_data:
+            if gdn_idx < len(spec_all_states):
+                spec_cache.all_states = spec_all_states[gdn_idx]
+            gdn_idx += 1
+
+            gdn = parent_layer.linear_attn
+            normed = parent_layer.input_layernorm(layer_input)
+            if hasattr(gdn, 'in_proj_qkv'):
+                qkv = gdn.in_proj_qkv(normed)
+            else:
+                q, k, v, z, b, a = gdn.fix_query_key_value_ordering(
+                    gdn.in_proj_qkvz(normed), gdn.in_proj_ba(normed))
+                B_dim = normed.shape[0]
+                qkv = mx.concatenate(
+                    [q.reshape(B_dim, S, -1), k.reshape(B_dim, S, -1),
+                     v.reshape(B_dim, S, -1)], axis=-1)
+            spec_cache.conv_input = mx.concatenate([pre_conv, qkv], axis=1)
+
+    # Concatenate multi-layer hidden states
+    selected = [layer_hiddens[i] for i in target_layer_ids]
+    target_hidden = mx.concatenate(selected, axis=-1)
+
+    # Final norm + lm_head
+    pre_norm = hidden_states
+    normed = inner.norm(hidden_states)
+
+    if hasattr(text_model, 'lm_head'):
+        logits = text_model.lm_head(normed)
+    else:
+        logits = inner.embed_tokens.as_linear(normed)
+
+    return target_hidden, pre_norm, logits

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

@@ -0,0 +1,326 @@
+#!/usr/bin/env python3
+"""MTP Speculative Decoding integrated with mlx_lm's BatchGenerator.
+
+Subclasses BatchGenerator to add MTP drafting + S>1 verification with
+correct GDN state rollback via SpeculativeArraysCache.
+
+At BS=1: drafts γ tokens with MTP, verifies at S=γ+1, buffers accepted tokens.
+At BS>1: falls back to standard BatchGenerator (no speculative).
+
+Usage:
+    from mtp_batch_generator import MTPBatchGenerator
+    gen = MTPBatchGenerator(model, mtp_predictor, gamma=2, ...)
+    gen.insert([prompt_tokens])
+    while True:
+        responses = gen.next()
+"""
+
+import time
+
+import mlx.core as mx
+from mlx_lm.generate import BatchGenerator, generation_stream
+
+from .mtp_module import MTPPredictor, speculative_forward, draft_tokens
+
+
+class MTPBatchGenerator(BatchGenerator):
+    """BatchGenerator with MTP speculative decoding for BS=1."""
+
+    def __init__(
+        self,
+        model,
+        mtp_predictor: MTPPredictor,
+        gamma: int = 2,
+        temp: float = 0.0,
+        alpha: float = 1.0,
+        **kwargs,
+    ):
+        super().__init__(model, **kwargs)
+        self.mtp = mtp_predictor
+        self.gamma = gamma
+        self.temp = temp
+        self.alpha = alpha
+
+        self._token_buffer = {}       # uid → [(token, logprobs), ...]
+        self._captured = {}            # pre_norm / prompt_pre_norm from norm wrapper
+        self._mtp_pre_norm = {}        # uid → (B, 1, D) pre-norm hidden state
+        self._mtp_prefilled = set()    # uids with MTP cache prefilled
+        self._request_temp = {}        # uid → temperature from request
+
+        self._setup_hidden_capture()
+
+    def _setup_hidden_capture(self):
+        """Monkey-patch model's final norm to capture pre-norm hidden state.
+
+        Captures:
+        - pre_norm: hidden states before final RMSNorm (for MTP input)
+        - prompt_pre_norm: same but only when S>1 (prefill)
+        """
+        inner = getattr(self.model, 'model', None) or self.model.language_model.model
+        original_norm = inner.norm
+        captured = self._captured
+
+        class _CapturingNorm:
+            def __init__(self, orig):
+                self._orig = orig
+                self.weight = orig.weight
+
+            def __call__(self, x):
+                captured['pre_norm'] = x
+                if x.shape[1] > 1:
+                    captured['prompt_pre_norm'] = x
+                return self._orig(x)
+
+            def __getattr__(self, name):
+                return getattr(self._orig, name)
+
+        inner.norm = _CapturingNorm(original_norm)
+
+    def _next(self):
+        batch = self.active_batch
+
+        # Yield buffered tokens first
+        if batch is not None and len(batch) == 1:
+            uid = batch.uids[0]
+            if uid in self._token_buffer and self._token_buffer[uid]:
+                return self._yield_buffered(batch, uid)
+
+        # BS=1 speculative path
+        if (batch is not None
+                and len(batch) == 1
+                and self.gamma > 0
+                and len(self.unprocessed_prompts) == 0):
+            uid = batch.uids[0]
+            if uid not in self._mtp_prefilled:
+                return self._first_step_and_prefill(batch, uid)
+            return self._speculative_next()
+
+        # Standard path (BS>1 or no batch)
+        responses = super()._next()
+        if responses and batch is not None and len(batch) == 1:
+            if 'pre_norm' in self._captured:
+                uid = batch.uids[0]
+                self._mtp_pre_norm[uid] = self._captured['pre_norm'][:, -1:, :]
+        return responses
+
+    def _first_step_and_prefill(self, batch, uid):
+        """First decode step. MTP cache already prefilled by ExoBatchGenerator.submit()."""
+        responses = super()._next()
+        if not responses:
+            return responses
+
+        # Capture decode pre_norm from this standard step for first speculative cycle
+        decode_pre_norm = self._captured.get('pre_norm')
+        if decode_pre_norm is not None:
+            mx.eval(decode_pre_norm)
+            self._mtp_pre_norm[uid] = decode_pre_norm[:, -1:, :]
+
+        self._mtp_prefilled.add(uid)
+        return responses
+
+    def _speculative_next(self):
+        """Core speculative cycle with correct GDN rollback."""
+        tic = time.perf_counter()
+        batch = self.active_batch
+        uid = batch.uids[0]
+        y = batch.y          # (1,) — token from previous step, to be yielded
+        y_val = y[0].item()
+        y_logprobs = batch.logprobs[0]
+
+        # Append current y to token history
+        batch.tokens[0] = mx.concatenate((batch.tokens[0], y[0:1]))
+
+        pre_norm = self._mtp_pre_norm.get(uid)
+        if pre_norm is None:
+            return super()._next()
+
+        gamma = self.gamma
+        temp = self._request_temp.get(uid, self.temp)
+        alpha = self.alpha
+
+        # 1. Draft γ tokens (lazy chain, no eval)
+        next_token_arr = y.reshape(1, 1)
+        draft_ids, draft_probs = draft_tokens(
+            self.mtp, pre_norm, next_token_arr, gamma, temp)
+
+        # 2. Verify via speculative_forward (handles GDN cache wrapping + kernel swap)
+        draft_concat = mx.concatenate(
+            [d.reshape(1, 1) for d in draft_ids], axis=1)  # (1, γ)
+        verify_input = mx.concatenate(
+            [next_token_arr, draft_concat], axis=1)  # (1, γ+1)
+        verify_pre_norm, verify_logits = speculative_forward(
+            self.model, verify_input, batch.cache, speculative=True)
+
+        # 3. Build acceptance check lazily
+        target_tokens = mx.argmax(verify_logits[:, :gamma, :], axis=-1)
+
+        if temp == 0:
+            matches = mx.equal(target_tokens, draft_concat).squeeze(0)
+            all_next = mx.argmax(verify_logits[0], axis=-1)
+            logprobs_all = verify_logits[0] - mx.logsumexp(
+                verify_logits[0], axis=-1, keepdims=True)
+            mx.async_eval(matches, all_next, logprobs_all, verify_pre_norm)
+        else:
+            accept_ratios = []
+            for i in range(gamma):
+                p = mx.softmax(verify_logits[0, i] / temp, axis=-1)
+                q = draft_probs[i]
+                p_di = p[draft_ids[i].squeeze()]
+                q_di = q[0, draft_ids[i].squeeze()]
+                ratio = p_di / mx.maximum(q_di, 1e-10)
+                accept_ratios.append(mx.minimum(ratio ** alpha, 1.0))
+            uniforms = mx.random.uniform(shape=(gamma,))
+            corrections = []
+            for i in range(gamma):
+                p = mx.softmax(verify_logits[0, i] / temp, axis=-1)
+                q = draft_probs[i][0]
+                residual = mx.maximum(p - q, 0.0)
+                corrections.append(mx.random.categorical(mx.log(residual + 1e-10)))
+            bonus_token = mx.random.categorical(verify_logits[0, gamma] * (1.0 / temp))
+            logprobs_all = verify_logits[0] - mx.logsumexp(
+                verify_logits[0], axis=-1, keepdims=True)
+            mx.async_eval(accept_ratios, uniforms, corrections, bonus_token,
+                          logprobs_all, verify_pre_norm, draft_concat)
+
+        # 4. Determine acceptance
+        n_accepted = 0
+        for i in range(gamma):
+            if temp == 0:
+                if matches[i].item():
+                    n_accepted += 1
+                else:
+                    break
+            else:
+                if uniforms[i].item() < accept_ratios[i].item():
+                    n_accepted += 1
+                else:
+                    break
+
+        print(f"[MTP] n_accepted={n_accepted}/{gamma}", flush=True)
+
+        # 5. Rollback cache
+        rollback = gamma - n_accepted
+        if rollback > 0:
+            for c in batch.cache:
+                if hasattr(c, 'offset'):
+                    c.offset -= rollback
+                elif hasattr(c, 'rollback'):
+                    c.rollback(n_accepted)
+
+        # Unwrap SpeculativeArraysCache
+        for i, c in enumerate(batch.cache):
+            if hasattr(c, 'base'):
+                batch.cache[i] = c.base
+
+        # 6. Bonus/correction token + logprobs
+        if n_accepted == gamma:
+            if temp == 0:
+                bonus_val = all_next[gamma].item()
+            else:
+                bonus_val = bonus_token.item()
+            bonus_lp = logprobs_all[gamma]
+        else:
+            if temp == 0:
+                bonus_val = all_next[n_accepted].item()
+            else:
+                bonus_val = corrections[n_accepted].item()
+            bonus_lp = logprobs_all[n_accepted]
+
+        # 7. Update MTP pre_norm for next cycle
+        self._mtp_pre_norm[uid] = verify_pre_norm[
+            :, (gamma if n_accepted == gamma else n_accepted):
+               (gamma if n_accepted == gamma else n_accepted) + 1, :]
+
+        # 8. Build token list: current y + accepted drafts
+        draft_int_values = draft_concat[0].tolist()
+        all_tokens = [(y_val, y_logprobs)]
+        for i in range(n_accepted):
+            all_tokens.append((draft_int_values[i], logprobs_all[i]))
+
+        # 9. Set batch.y = bonus for next cycle
+        batch.y = mx.array([bonus_val])
+        batch.logprobs = [bonus_lp]
+
+        # Append accepted drafts to token history
+        if n_accepted > 0:
+            batch.tokens[0] = mx.concatenate(
+                (batch.tokens[0], mx.array([t for t, _ in all_tokens[1:]])))
+        batch.num_tokens[0] += len(all_tokens)
+
+        # 10. Check stop conditions — truncate at stop token
+        toc = time.perf_counter()
+        self._stats.generation_time += toc - tic
+        self._stats.generation_tokens += len(all_tokens)
+
+        # Find first stop token or length limit in all_tokens
+        stop_idx = None
+        for idx, (tok, _) in enumerate(all_tokens):
+            if tok in self.stop_tokens:
+                stop_idx = idx
+                break
+            if batch.num_tokens[0] >= batch.max_tokens[0]:
+                stop_idx = idx
+                break
+
+        first_tok, first_lp = all_tokens[0]
+
+        if stop_idx is not None:
+            # Tokens before the stop are valid output — buffer them
+            # The stop token itself triggers finish_reason
+            valid_tokens = all_tokens[:stop_idx]
+            if valid_tokens:
+                # Yield first, buffer rest + a final stop entry
+                if len(valid_tokens) > 1:
+                    self._token_buffer[uid] = valid_tokens[1:]
+                # Append stop marker as last buffered token
+                stop_tok, stop_lp = all_tokens[stop_idx]
+                if uid not in self._token_buffer:
+                    self._token_buffer[uid] = []
+                self._token_buffer[uid].append((stop_tok, stop_lp))
+                mx.async_eval(batch.y)
+                return [self.Response(uid, first_tok, first_lp, None, lambda: None)]
+            else:
+                # Stop token is the first token — finish immediately
+                cache = batch.extract_cache(0)
+                self.active_batch = None
+                self._cleanup_uid(uid)
+                return [self.Response(uid, first_tok, first_lp, "stop", cache)]
+
+        # Buffer remaining tokens
+        if len(all_tokens) > 1:
+            self._token_buffer[uid] = all_tokens[1:]
+
+        mx.async_eval(batch.y)
+        return [self.Response(uid, first_tok, first_lp, None, lambda: None)]
+
+    def _yield_buffered(self, batch, uid):
+        """Yield one buffered token from a previous speculative cycle."""
+        tic = time.perf_counter()
+        buf = self._token_buffer[uid]
+        tok, lp = buf.pop(0)
+
+        if not buf:
+            del self._token_buffer[uid]
+
+        finish_reason = None
+        if tok in self.stop_tokens:
+            finish_reason = "stop"
+        elif batch.num_tokens[0] >= batch.max_tokens[0]:
+            finish_reason = "length"
+
+        cache = None
+        if finish_reason:
+            cache = batch.extract_cache(0)
+            self.active_batch = None
+            self._cleanup_uid(uid)
+
+        toc = time.perf_counter()
+        self._stats.generation_time += toc - tic
+        return [self.Response(uid, tok, lp, finish_reason, cache or (lambda: None))]
+
+    def _cleanup_uid(self, uid):
+        """Clean up MTP state for a finished request."""
+        self._mtp_pre_norm.pop(uid, None)
+        self._mtp_prefilled.discard(uid)
+        self._token_buffer.pop(uid, None)
+        self._request_temp.pop(uid, None)

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

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

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

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

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

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

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

@@ -52,6 +52,7 @@ from exo.shared.types.worker.instances import (
     MlxRingInstance,
 )
 from exo.shared.types.worker.shards import (
+    AttnMoeSplitShardMetadata,
     CfgShardMetadata,
     PipelineShardMetadata,
     ShardMetadata,
@@ -60,6 +61,7 @@ from exo.shared.types.worker.shards import (
 from exo.worker.engines.mlx.auto_parallel import (
     LayerLoadedCallback,
     TimeoutCallback,
+    attn_moe_split_auto_parallel,
     eval_with_timeout,
     get_inner_model,
     get_layers,
@@ -72,15 +74,18 @@ Group = mx.distributed.Group
 
 
 def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
+    # PipelineShardMetadata and AttnMoeSplitShardMetadata store disjoint or
+    # full-range layer slices per rank, so each rank's weights are already
+    # scoped by [start_layer, end_layer). TensorShardMetadata shards every
+    # weight across `world_size` so we divide by it.
+    full_layer_shard = isinstance(
+        model_shard_meta, (PipelineShardMetadata, AttnMoeSplitShardMetadata)
+    )
     return Memory.from_float_kb(
         (model_shard_meta.end_layer - model_shard_meta.start_layer)
         / model_shard_meta.n_layers
         * model_shard_meta.model_card.storage_size.in_kb
-        / (
-            1
-            if isinstance(model_shard_meta, PipelineShardMetadata)
-            else model_shard_meta.world_size
-        )
+        / (1 if full_layer_shard else model_shard_meta.world_size)
     )
 
 
@@ -193,6 +198,8 @@ 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:
@@ -281,6 +288,12 @@ def shard_and_load(
                 model, group, shard_metadata, on_layer_loaded=on_layer_loaded
             )
             eval_with_timeout(model.parameters(), timeout_seconds, on_timeout)
+        case AttnMoeSplitShardMetadata():
+            logger.info(f"loading model from {model_path} with attn/moe split")
+            model = attn_moe_split_auto_parallel(
+                model, group, shard_metadata, on_layer_loaded=on_layer_loaded
+            )
+            eval_with_timeout(model.parameters(), timeout_seconds, on_timeout)
         case CfgShardMetadata():
             raise ValueError(
                 "CfgShardMetadata is not supported for text model loading - "