Call warmup_speculative at startup to pre-compile LpB kernels

The warmup_speculative() function was defined but never called. Custom Metal kernels (LpB) require first-call compilation (~200ms). Without warmup, the first speculative cycle is slow, dragging down average TPS by 10-20% on short generations. In mlx_bench testing: cold 48 TPS → warm 60 TPS for DFlash, cold 39 TPS → warm 44 TPS for MTP. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Add EXO_DISABLE_LOGPROBS=1 to skip per-token logprobs extraction
2026-04-01 23:15:05 +01:00 · 2026-03-30 23:41:47 +01:00 · 2026-03-30 22:39:34 +01:00 · 2026-03-30 18:44:01 +01:00 · 2026-03-30 18:24:25 +01:00 · 2026-03-30 18:23:05 +01:00
134 changed files with 6616 additions and 281 deletions
@@ -2396,7 +2396,7 @@ def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array:
        array: The angles in degrees.
    """

-def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array]):
+def depends[T](inputs: T, dependencies: array | Sequence[array]) -> T:
    """
    Insert dependencies between arrays in the graph. The outputs are
    identical to ``inputs`` but with dependencies on ``dependencies``.
@@ -1,9 +1,5 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from layers import *
-from utils import *
+from .layers import *
+from .utils import *

 from . import init as init
 from . import losses as losses
@@ -1,20 +1,16 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from activations import *
-from base import *
-from containers import *
-from convolution import *
-from convolution_transpose import *
-from distributed import *
-from dropout import *
-from embedding import *
-from linear import *
-from normalization import *
-from pooling import *
-from positional_encoding import *
-from quantized import *
-from recurrent import *
-from transformer import *
-from upsample import *
+from .activations import *
+from .base import *
+from .containers import *
+from .convolution import *
+from .convolution_transpose import *
+from .distributed import *
+from .dropout import *
+from .embedding import *
+from .linear import *
+from .normalization import *
+from .pooling import *
+from .positional_encoding import *
+from .quantized import *
+from .recurrent import *
+from .transformer import *
+from .upsample import *
@@ -53,7 +53,7 @@ class Module(dict):
        mx.eval(model.parameters())
    """

-    __call__: Callable
+    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
    def __init__(self) -> None:
        """Should be called by the subclasses of ``Module``."""

@@ -30,7 +30,7 @@ def str2bool(string):  # -> bool:
 def setup_arg_parser():  # -> ArgumentParser:
    """Set up and return the argument parser."""

-generation_stream = ...
+generation_stream: mx.Stream

@contextlib.contextmanager
 def wired_limit(
@@ -266,12 +266,12 @@ def _merge_caches(caches: Any) -> List[Any]: ...
 class Batch:
    uids: List[int]
    y: mx.array
-    logprobs: mx.array
+    logprobs: List[mx.array] | mx.array
    max_tokens: List[int]
    num_tokens: List[int]
    cache: List[Any]
-    samplers: List[Any]
-    logits_processors: List[Any]
+    samplers: List[Callable[[mx.array], mx.array] | None]
+    logits_processors: List[List[Callable[[mx.array, mx.array], mx.array]]]
    tokens: List[mx.array]
    def __len__(self) -> int: ...
    def filter(self, keep_idx: List[int]) -> None: ...
@@ -279,13 +279,18 @@ class Batch:
    def extract_cache(self, idx: int) -> List[Any]: ...

 class BatchGenerator:
-    model: Any
+    model: nn.Module
+    sampler: Callable[[mx.array], mx.array]
+    stop_tokens: set[int]
    max_kv_size: Optional[int]
    prefill_step_size: int
+    completion_batch_size: int
+    prefill_batch_size: int
    unprocessed_prompts: List[Any]
    active_batch: Optional[Batch]
    prompt_progress_callback: Callable[[List[Tuple[int, int, int]]], None]
    _stats: BatchStats
+    _next_count: int

    @dataclass
    class Response:
@@ -88,8 +88,8 @@ def create_attention_mask(
 ) -> array | Literal["causal"] | None: ...

 class _BaseCache(Cache):
-    keys: mx.array
-    values: mx.array
+    keys: mx.array | None
+    values: mx.array | None
    offset: int
    @property
    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
@@ -268,29 +268,14 @@ class CacheList(_BaseCache):
        """

 class BatchKVCache(_BaseCache):
-    step = ...
-    def __init__(self, left_padding: List[int]) -> None:
-        """
-        The BatchKV cache expects inputs to be left-padded.
-
-        E.g. the following prompts:
-
-            [1, 3, 5]
-            [7]
-            [2, 6, 8, 9]
-
-        Should be padded like so:
-
-            [0, 1, 3, 5]
-            [0, 0, 0, 7]
-            [2, 6, 8, 9]
-
-        And ``left_padding`` specifies the amount of padding for each.
-        In this case, ``left_padding = [1, 3, 0]``.
-        """
-
-    def update_and_fetch(self, keys, values):  # -> tuple[array | Any, array | Any]:
-        ...
+    step: int
+    keys: array | None
+    values: array | None
+    offset: array
+    left_padding: array
+    _idx: int
+    def __init__(self, left_padding: List[int]) -> None: ...
+    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
    @property
    def state(
        self,
@@ -316,12 +301,21 @@ class BatchKVCache(_BaseCache):
        """

 class BatchRotatingKVCache(_BaseCache):
-    step = ...
-    def __init__(self, max_size, left_padding: List[int]) -> None: ...
-    def update_and_fetch(
-        self, keys, values
-    ):  # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
-        ...
+    step: int
+    keys: array | None
+    values: array | None
+    offset: array
+    left_padding: array
+    max_size: int
+    _idx: int
+    _offset: int
+    rotated: bool
+    _lengths: array | None
+    def __init__(self, max_size: int, left_padding: List[int]) -> None: ...
+    def _trim(self, trim_size: int, v: array, append: array | None = ...) -> array: ...
+    def _update_in_place(self, keys: array, values: array) -> tuple[array, array]: ...
+    def _update_concat(self, keys: array, values: array) -> tuple[array, array]: ...
+    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
    @property
    def state(
        self,
@@ -0,0 +1,35 @@
+from typing import Optional
+
+import mlx.core as mx
+
+def compute_g(A_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array: ...
+def gated_delta_update(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    a: mx.array,
+    b: mx.array,
+    A_log: mx.array,
+    dt_bias: mx.array,
+    state: Optional[mx.array] = ...,
+    mask: Optional[mx.array] = ...,
+    use_kernel: bool = ...,
+) -> tuple[mx.array, mx.array]: ...
+def gated_delta_ops(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: Optional[mx.array] = ...,
+    mask: Optional[mx.array] = ...,
+) -> tuple[mx.array, mx.array]: ...
+def gated_delta_kernel(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: mx.array,
+    mask: Optional[mx.array] = ...,
+) -> tuple[mx.array, mx.array]: ...
@@ -0,0 +1,51 @@
+from typing import Any, Optional
+
+import mlx.nn as nn
+
+class YarnRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        scaling_factor: float = ...,
+        original_max_position_embeddings: int = ...,
+        beta_fast: float = ...,
+        beta_slow: float = ...,
+        mscale: float = ...,
+        mscale_all_dim: float = ...,
+    ) -> None: ...
+
+class Llama3RoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        scaling_factor: float = ...,
+        original_max_position_embeddings: int = ...,
+        low_freq_factor: float = ...,
+        high_freq_factor: float = ...,
+    ) -> None: ...
+
+class SuScaledRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        short_factor: Any = ...,
+        long_factor: Any = ...,
+        original_max_position_embeddings: int = ...,
+    ) -> None: ...
+
+def initialize_rope(
+    dims: int,
+    base: float = ...,
+    traditional: bool = ...,
+    scaling_config: Optional[dict[str, Any]] = ...,
+    max_position_embeddings: Optional[int] = ...,
+) -> nn.Module: ...
@@ -39,11 +39,11 @@ class StreamingDetokenizer:
    """

    __slots__ = ...
-    def reset(self): ...
-    def add_token(self, token): ...
-    def finalize(self): ...
+    def reset(self) -> None: ...
+    def add_token(self, token: int) -> None: ...
+    def finalize(self) -> None: ...
    @property
-    def last_segment(self):
+    def last_segment(self) -> str:
        """Return the last segment of readable text since last time this property was accessed."""

 class NaiveStreamingDetokenizer(StreamingDetokenizer):
@@ -496,20 +496,26 @@ def main() -> int:
                                all_rows.append(row)

                            if batch_results:
-                                agg_gen_tps = sum(
+                                valid_gen_tps = [
                                    x["stats"]["generation_tps"]
                                    for x, _ in batch_results
                                    if x["stats"]["generation_tps"] > 0
+                                ]
+                                per_req_tps = (
+                                    mean(valid_gen_tps) if valid_gen_tps else 0.0
                                )
+                                agg_gen_tps = per_req_tps * concurrency
                                logger.info(
                                    f"[concurrent {concurrency}x]  "
                                    f"agg_gen_tps={agg_gen_tps:.2f}  "
+                                    f"per_req_tps={per_req_tps:.2f}  "
                                    f"errors={batch_errors}"
                                )

                    if runs:
                        prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
-                        gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
+                        per_req_tps = mean(x["stats"]["generation_tps"] for x in runs)
+                        gen_tps = per_req_tps * concurrency
                        ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
                        gtok = mean(x["stats"]["generation_tokens"] for x in runs)
                        peak = mean(
@@ -0,0 +1,91 @@
+#!/usr/bin/env bash
+#
+# Run exo_bench.py for each model/mode from bench_params.json.
+#
+# For each entry, runs with:
+#   --pp 800 (fixed, representative LCB prompt length)
+#   --tg <mean completion tokens from vLLM>
+#   --sharding tensor --instance-meta jaccl
+#   --min-nodes 1 --max-nodes 4
+#   --repeat 1
+#   --danger-delete-downloads
+#   --settle-timeout 300
+#
+# Results go to bench/eval_results/<model_dir>/tps_<mode>.json
+#
+# Usage:
+#   bash bench/run_lcb_tps_bench.sh              # run all
+#   bash bench/run_lcb_tps_bench.sh --dry-run    # show what would run
+
+set -euo pipefail
+cd "$(dirname "$0")"
+
+PARAMS_FILE="eval_results/bench_params.json"
+PP=800
+HOST="${EXO_HOST:-s9}"
+DRY_RUN=false
+
+if [[ "${1:-}" == "--dry-run" ]]; then
+    DRY_RUN=true
+fi
+
+if [[ ! -f "$PARAMS_FILE" ]]; then
+    echo "ERROR: $PARAMS_FILE not found. Run compute_bench_params.py first."
+    exit 1
+fi
+
+# Parse bench_params.json and run each entry
+python3 -c "
+import json, sys
+data = json.load(open('$PARAMS_FILE'))
+for entry in data:
+    mlx_id = entry['mlx_model_id']
+    mode = entry['mode']
+    tg = entry['bench_params']['tg']
+    vllm_name = entry['vllm_name']
+    # Output dir: replace / with _
+    out_dir = 'eval_results/' + mlx_id.replace('/', '_')
+    out_file = out_dir + '/tps_' + mode + '.json'
+    print(f'{mlx_id}\t{mode}\t{tg}\t{out_file}\t{vllm_name}')
+" | while IFS=$'\t' read -r model mode tg out_file vllm_name; do
+    out_dir="$(dirname "$out_file")"
+    mkdir -p "$out_dir"
+
+    echo ""
+    echo "============================================================"
+    echo "Model:  $model"
+    echo "Mode:   $mode"
+    echo "vLLM:   $vllm_name"
+    echo "PP:     $PP"
+    echo "TG:     $tg"
+    echo "Output: $out_file"
+    echo "============================================================"
+
+    if [[ -f "$out_file" ]]; then
+        echo "SKIP: $out_file already exists"
+        continue
+    fi
+
+    if [[ "$DRY_RUN" == "true" ]]; then
+        echo "DRY-RUN: would run exo_bench.py"
+        continue
+    fi
+
+    uv run python exo_bench.py \
+        --host "$HOST" \
+        --model "$model" \
+        --pp "$PP" \
+        --tg "$tg" \
+        --repeat 1 \
+        --sharding tensor \
+        --instance-meta jaccl \
+        --min-nodes 1 \
+        --max-nodes 4 \
+        --settle-timeout 300 \
+        --force-download \
+        --danger-delete-downloads \
+        --json-out "$out_file" || echo "FAILED: $model ($mode)"
+done
+
+echo ""
+echo "All benchmarks complete."
@@ -1,9 +1,6 @@
 <script lang="ts">
  import {
    isLoading,
-    sendMessage,
-    generateImage,
-    editImage,
    editingImage,
    clearEditingImage,
    selectedChatModel,
@@ -28,7 +25,7 @@
    modelTasks?: Record<string, string[]>;
    modelCapabilities?: Record<string, string[]>;
    onSend?: () => void;
-    onAutoSend?: (
+    onAutoSend: (
      content: string,
      files?: {
        id: string;
@@ -216,49 +213,10 @@
    uploadedFiles = [];
    resetTextareaHeight();

-    // When onAutoSend is provided, the parent controls all send logic
-    // (including launching non-running models before sending)
-    if (onAutoSend) {
-      onAutoSend(content, files);
-      onSend?.();
-      setTimeout(() => textareaRef?.focus(), 10);
-      return;
-    }
-
-    // Use image editing if in edit mode
-    if (isEditMode && currentEditingImage && content) {
-      editImage(content, currentEditingImage.imageDataUrl);
-    }
-    // If user attached an image with an ImageToImage model, use edit endpoint
-    else if (
-      currentModel &&
-      modelSupportsImageEditing(currentModel) &&
-      files.length > 0 &&
-      content
-    ) {
-      // Use the first attached image for editing
-      const imageFile = files[0];
-      if (imageFile.preview) {
-        editImage(content, imageFile.preview);
-      }
-    } else if (
-      currentModel &&
-      modelSupportsTextToImage(currentModel) &&
-      content
-    ) {
-      // Use image generation for text-to-image models
-      generateImage(content);
-    } else {
-      sendMessage(
-        content,
-        files,
-        modelSupportsThinking() ? thinkingEnabled : null,
-      );
-    }
-
+    // Parent controls all send logic (including image routing,
+    // launching non-running models before sending, etc.)
+    onAutoSend(content, files);
    onSend?.();
-
-    // Refocus the textarea after sending
    setTimeout(() => textareaRef?.focus(), 10);
  }

@@ -82,6 +82,12 @@
      d="M12.025 1.13c-5.77 0-10.449 4.647-10.449 10.378 0 1.112.178 2.181.503 3.185.064-.222.203-.444.416-.577a.96.96 0 0 1 .524-.15c.293 0 .584.124.84.284.278.173.48.408.71.694.226.282.458.611.684.951v-.014c.017-.324.106-.622.264-.874s.403-.487.762-.543c.3-.047.596.06.787.203s.31.313.4.467c.15.257.212.468.233.542.01.026.653 1.552 1.657 2.54.616.605 1.01 1.223 1.082 1.912.055.537-.096 1.059-.38 1.572.637.121 1.294.187 1.967.187.657 0 1.298-.063 1.921-.178-.287-.517-.44-1.041-.384-1.581.07-.69.465-1.307 1.081-1.913 1.004-.987 1.647-2.513 1.657-2.539.021-.074.083-.285.233-.542.09-.154.208-.323.4-.467a1.08 1.08 0 0 1 .787-.203c.359.056.604.29.762.543s.247.55.265.874v.015c.225-.34.457-.67.683-.952.23-.286.432-.52.71-.694.257-.16.547-.284.84-.285a.97.97 0 0 1 .524.151c.228.143.373.388.43.625l.006.04a10.3 10.3 0 0 0 .534-3.273c0-5.731-4.678-10.378-10.449-10.378M8.327 6.583a1.5 1.5 0 0 1 .713.174 1.487 1.487 0 0 1 .617 2.013c-.183.343-.762-.214-1.102-.094-.38.134-.532.914-.917.71a1.487 1.487 0 0 1 .69-2.803m7.486 0a1.487 1.487 0 0 1 .689 2.803c-.385.204-.536-.576-.916-.71-.34-.12-.92.437-1.103.094a1.487 1.487 0 0 1 .617-2.013 1.5 1.5 0 0 1 .713-.174m-10.68 1.55a.96.96 0 1 1 0 1.921.96.96 0 0 1 0-1.92m13.838 0a.96.96 0 1 1 0 1.92.96.96 0 0 1 0-1.92M8.489 11.458c.588.01 1.965 1.157 3.572 1.164 1.607-.007 2.984-1.155 3.572-1.164.196-.003.305.12.305.454 0 .886-.424 2.328-1.563 3.202-.22-.756-1.396-1.366-1.63-1.32q-.011.001-.02.006l-.044.026-.01.008-.03.024q-.018.017-.035.036l-.032.04a1 1 0 0 0-.058.09l-.014.025q-.049.088-.11.19a1 1 0 0 1-.083.116 1.2 1.2 0 0 1-.173.18q-.035.029-.075.058a1.3 1.3 0 0 1-.251-.243 1 1 0 0 1-.076-.107c-.124-.193-.177-.363-.337-.444-.034-.016-.104-.008-.2.022q-.094.03-.216.087-.06.028-.125.063l-.13.074q-.067.04-.136.086a3 3 0 0 0-.135.096 3 3 0 0 0-.26.219 2 2 0 0 0-.12.121 2 2 0 0 0-.106.128l-.002.002a2 2 0 0 0-.09.132l-.001.001a1.2 1.2 0 0 0-.105.212q-.013.036-.024.073c-1.139-.875-1.563-2.317-1.563-3.203 0-.334.109-.457.305-.454m.836 10.354c.824-1.19.766-2.082-.365-3.194-1.13-1.112-1.789-2.738-1.789-2.738s-.246-.945-.806-.858-.97 1.499.202 2.362c1.173.864-.233 1.45-.685.64-.45-.812-1.683-2.896-2.322-3.295s-1.089-.175-.938.647 2.822 2.813 2.562 3.244-1.176-.506-1.176-.506-2.866-2.567-3.49-1.898.473 1.23 2.037 2.16c1.564.932 1.686 1.178 1.464 1.53s-3.675-2.511-4-1.297c-.323 1.214 3.524 1.567 3.287 2.405-.238.839-2.71-1.587-3.216-.642-.506.946 3.49 2.056 3.522 2.064 1.29.33 4.568 1.028 5.713-.624m5.349 0c-.824-1.19-.766-2.082.365-3.194 1.13-1.112 1.789-2.738 1.789-2.738s.246-.945.806-.858.97 1.499-.202 2.362c-1.173.864.233 1.45.685.64.451-.812 1.683-2.896 2.322-3.295s1.089-.175.938.647-2.822 2.813-2.562 3.244 1.176-.506 1.176-.506 2.866-2.567 3.49-1.898-.473 1.23-2.037 2.16c-1.564.932-1.686 1.178-1.464 1.53s3.675-2.511 4-1.297c.323 1.214-3.524 1.567-3.287 2.405.238.839 2.71-1.587 3.216-.642.506.946-3.49 2.056-3.522 2.064-1.29.33-4.568 1.028-5.713-.624"
    />
  </svg>
+{:else if family === "step"}
+  <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
+    <path
+      d="M22.012 0h1.032v.927H24v.968h-.956V3.78h-1.032V1.896h-1.878v-.97h1.878V0zM2.6 12.371V1.87h.969v10.502h-.97zm10.423.66h10.95v.918h-6.208v9.579h-4.742V13.03zM5.629 3.333v12.356H0v4.51h10.386V8L20.859 8l-.003-4.668-15.227.001z"
+    />
+  </svg>
 {:else}
  <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
    <path
@@ -1577,6 +1577,7 @@ class AppStore {
    // Remove messages after user message (including the user message for image requests
    // since generateImage/editImage will re-add it)
    this.messages = this.messages.slice(0, lastUserIndex);
+    this.updateActiveConversation();

    switch (requestType) {
      case "image-generation":
@@ -1792,6 +1793,14 @@ class AppStore {
            this.persistConversation(targetConversationId);
          }
        },
+        {
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+            }
+          },
+        },
      );

      // Final update
@@ -1989,6 +1998,14 @@ class AppStore {
            this.persistConversation(targetConversationId);
          }
        },
+        {
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+            }
+          },
+        },
      );

      // Final cleanup of the message (if conversation still exists)
@@ -2396,7 +2413,7 @@ class AppStore {

      let streamedContent = "";
      let streamedThinking = "";
-
+      let serverTpsReceived = false;
      interface ChatCompletionChunk {
        choices?: Array<{
          delta?: { content?: string; reasoning_content?: string };
@@ -2461,7 +2478,6 @@ class AppStore {
            tokenCount += 1;
            this.totalTokens = tokenCount;

-            // Update real-time TPS during streaming
            if (firstTokenTime !== null && tokenCount > 1) {
              const elapsed = performance.now() - firstTokenTime;
              this.tps = (tokenCount / elapsed) * 1000;
@@ -2512,16 +2528,24 @@ class AppStore {
              startedAt: this.prefillProgress?.startedAt ?? performance.now(),
            };
          },
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+              serverTpsReceived = true;
+            }
+          },
        },
      );

      // Clear prefill progress after stream ends
      this.prefillProgress = null;

-      // Calculate final TPS
-      if (firstTokenTime !== null && tokenCount > 1) {
+      // Use server-side TPS if available, otherwise fall back to client-side
+      if (!serverTpsReceived && firstTokenTime !== null && tokenCount > 1) {
        const totalGenerationTime = performance.now() - firstTokenTime;
-        this.tps = (tokenCount / totalGenerationTime) * 1000; // tokens per second
+        this.tps = (tokenCount / totalGenerationTime) * 1000;
      }

      // Final cleanup of the message (if conversation still exists)
@@ -42,6 +42,9 @@
    setSelectedChatModel,
    selectedChatModel,
    sendMessage,
+    generateImage,
+    editImage,
+    editingImage,
    messages,
    debugMode,
    toggleDebugMode,
@@ -834,6 +837,52 @@
    if (!model?.tasks) return false;
    return model.tasks.includes("ImageToImage");
  }
+
+  // Route a message to the correct endpoint based on model capabilities.
+  // Image models go to generateImage/editImage; text models go to sendMessage.
+  function routeMessage(
+    content: string,
+    files?: {
+      id: string;
+      name: string;
+      type: string;
+      textContent?: string;
+      preview?: string;
+    }[],
+  ) {
+    const model = selectedChatModel();
+    if (!model) {
+      sendMessage(content, files, null);
+      return;
+    }
+
+    const currentEditImage = editingImage();
+
+    // Image editing mode (explicit edit or attached image with ImageToImage model)
+    if (currentEditImage && content && modelSupportsImageEditing(model)) {
+      editImage(content, currentEditImage.imageDataUrl);
+      return;
+    }
+    if (
+      modelSupportsImageEditing(model) &&
+      files?.length &&
+      files[0].preview &&
+      content
+    ) {
+      editImage(content, files[0].preview);
+      return;
+    }
+
+    // Text-to-image generation
+    if (modelSupportsImageGeneration(model) && content) {
+      generateImage(content);
+      return;
+    }
+
+    // Default: text chat
+    sendMessage(content, files, null);
+  }
+
  let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
  type InstanceMeta = "MlxRing" | "MlxJaccl";

@@ -1527,7 +1576,11 @@
          downloadKind
        ] as Record<string, unknown>;

-        if (downloadKind !== "DownloadOngoing") continue;
+        if (
+          downloadKind !== "DownloadOngoing" &&
+          downloadKind !== "DownloadPending"
+        )
+          continue;
        if (!downloadPayload) continue;

        const downloadModelId = extractModelIdFromDownload(downloadPayload);
@@ -1542,9 +1595,38 @@
          if (downloadModelId !== modelId) continue;
        }

-        isDownloading = true;
+        // For DownloadPending with partial bytes (paused/resumed downloads),
+        // synthesize a progress object from the top-level downloaded/total fields
+        let progress: DownloadProgress | null;
+        if (downloadKind === "DownloadPending") {
+          const pendingDownloaded = getBytes(
+            downloadPayload.downloaded ??
+              downloadPayload.downloaded_bytes ??
+              downloadPayload.downloadedBytes,
+          );
+          const pendingTotal = getBytes(
+            downloadPayload.total ??
+              downloadPayload.total_bytes ??
+              downloadPayload.totalBytes,
+          );
+          if (pendingDownloaded <= 0 && pendingTotal <= 0) continue;
+          isDownloading = true;
+          progress = {
+            totalBytes: pendingTotal,
+            downloadedBytes: pendingDownloaded,
+            speed: 0,
+            etaMs: 0,
+            percentage:
+              pendingTotal > 0 ? (pendingDownloaded / pendingTotal) * 100 : 0,
+            completedFiles: 0,
+            totalFiles: 0,
+            files: [],
+          };
+        } else {
+          isDownloading = true;
+          progress = parseDownloadProgress(downloadPayload);
+        }

-        const progress = parseDownloadProgress(downloadPayload);
        if (progress) {
          // Sum all values across nodes - each node downloads independently
          totalBytes += progress.totalBytes;
@@ -1696,7 +1778,11 @@
          }
        }

-        if (downloadKind !== "DownloadOngoing") continue;
+        if (
+          downloadKind !== "DownloadOngoing" &&
+          downloadKind !== "DownloadPending"
+        )
+          continue;
        if (!downloadPayload) continue;

        // Check if this download is for this instance's model
@@ -1706,9 +1792,37 @@
          downloadModelId &&
          downloadModelId === instanceModelId
        ) {
-          isDownloading = true;
+          // For DownloadPending with partial bytes, synthesize progress
+          let progress: DownloadProgress | null;
+          if (downloadKind === "DownloadPending") {
+            const pendingDownloaded = getBytes(
+              downloadPayload.downloaded ??
+                downloadPayload.downloaded_bytes ??
+                downloadPayload.downloadedBytes,
+            );
+            const pendingTotal = getBytes(
+              downloadPayload.total ??
+                downloadPayload.total_bytes ??
+                downloadPayload.totalBytes,
+            );
+            if (pendingDownloaded <= 0 && pendingTotal <= 0) continue;
+            isDownloading = true;
+            progress = {
+              totalBytes: pendingTotal,
+              downloadedBytes: pendingDownloaded,
+              speed: 0,
+              etaMs: 0,
+              percentage:
+                pendingTotal > 0 ? (pendingDownloaded / pendingTotal) * 100 : 0,
+              completedFiles: 0,
+              totalFiles: 0,
+              files: [],
+            };
+          } else {
+            isDownloading = true;
+            progress = parseDownloadProgress(downloadPayload);
+          }

-          const progress = parseDownloadProgress(downloadPayload);
          if (progress) {
            // Sum all values across nodes - each node downloads independently
            totalBytes += progress.totalBytes;
@@ -2786,7 +2900,7 @@
        // Running model is same or better tier — use it directly
        setSelectedChatModel(bestRunning.id);
        if (!chatStarted) createConversation();
-        sendMessage(content, files);
+        routeMessage(content, files);
        return;
      }
    }
@@ -2803,7 +2917,7 @@
    if (hasRunningInstance(autoModel.id)) {
      setSelectedChatModel(autoModel.id);
      if (!chatStarted) createConversation();
-      sendMessage(content, files);
+      routeMessage(content, files);
      return;
    }

@@ -2956,7 +3070,7 @@
      if (pendingAutoMessage) {
        const msg = pendingAutoMessage;
        pendingAutoMessage = null;
-        sendMessage(msg.content, msg.files);
+        routeMessage(msg.content, msg.files);
      }
      return;
    }
@@ -3035,7 +3149,7 @@
    // Model is selected and running — send directly
    if (model && hasRunningInstance(model)) {
      chatLaunchState = "ready";
-      sendMessage(content, files, null);
+      routeMessage(content, files);
      return;
    }

@@ -117,12 +117,13 @@
              uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
              mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package);
              uvLockMlxVersion = mlxPackage.version;
+              uvLockMlxRev = builtins.elemAt (builtins.split "#" mlxPackage.source.git) 2;
            in
            {
              metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
              mlx = pkgs.callPackage ./nix/mlx.nix {
                inherit (self'.packages) metal-toolchain;
-                inherit uvLockMlxVersion;
+                inherit uvLockMlxVersion uvLockMlxRev;
              };
              default = self'.packages.exo;
            }
@@ -11,6 +11,7 @@
 , fmt
 , python313Packages
 , uvLockMlxVersion
+, uvLockMlxRev
 }:

 assert stdenv.isDarwin;
@@ -41,15 +42,13 @@ let

  mlx = stdenv.mkDerivation rec {
    pname = "mlx";
-    version = let v = "0.30.7.dev20260225+257d5692"; in
-      assert v == uvLockMlxVersion || throw "MLX version mismatch: nix/mlx.nix has ${v} but uv.lock has ${uvLockMlxVersion}. Update both the version and hash in nix/mlx.nix.";
-      v;
+    version = uvLockMlxVersion;
    pyproject = true;

    src = fetchFromGitHub {
      owner = "rltakashige";
      repo = "mlx-jaccl-fix-small-recv";
-      rev = "257d5692fc7af6bba3b8afaeb63c549b7d1e43d5";
+      rev = uvLockMlxRev;
      hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ=";
    };

@@ -25,8 +25,7 @@ dependencies = [
  "openai-harmony>=0.0.8",
  "httpx>=0.28.1",
  "tomlkit>=0.14.0",
-  "pillow>=11.0,<12.0",                        # compatibility with mflux
-  "mflux==0.15.5",
+  "mflux==0.16.9",
  "python-multipart>=0.0.21",
  "msgspec>=0.19.0",
  "zstandard>=0.23.0",
@@ -62,7 +61,7 @@ members = ["rust/exo_pyo3_bindings", "bench"]
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
-mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" }
+mlx-lm = { git = "https://github.com/ml-explore/mlx-lm", branch = "main" }
 # Uncomment to use local mlx/mlx-lm development versions:
 # mlx = { path = "/Users/Shared/mlx", editable=true }
 # mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
@@ -1,6 +1,7 @@
 model_id = "mlx-community/DeepSeek-V3.1-4bit"
 n_layers = 61
 hidden_size = 7168
+num_key_value_heads = 128
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "deepseek"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/DeepSeek-V3.1-8bit"
 n_layers = 61
 hidden_size = 7168
+num_key_value_heads = 128
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "deepseek"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.5-Air-8bit"
 n_layers = 46
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = false
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.5-Air-bf16"
 n_layers = 46
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.7-4bit"
 n_layers = 91
 hidden_size = 5120
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.7-6bit"
 n_layers = 91
 hidden_size = 5120
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.7-8bit-gs32"
 n_layers = 91
 hidden_size = 5120
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.7-Flash-4bit"
 n_layers = 47
 hidden_size = 2048
+num_key_value_heads = 20
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.7-Flash-5bit"
 n_layers = 47
 hidden_size = 2048
+num_key_value_heads = 20
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.7-Flash-6bit"
 n_layers = 47
 hidden_size = 2048
+num_key_value_heads = 20
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-4.7-Flash-8bit"
 n_layers = 47
 hidden_size = 2048
+num_key_value_heads = 20
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-5-8bit-MXFP8"
 n_layers = 78
 hidden_size = 6144
+num_key_value_heads = 64
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-5-MXFP4-Q8"
 n_layers = 78
 hidden_size = 6144
+num_key_value_heads = 64
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/GLM-5"
 n_layers = 78
 hidden_size = 6144
+num_key_value_heads = 64
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "glm"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Kimi-K2-Instruct-4bit"
 n_layers = 61
 hidden_size = 7168
+num_key_value_heads = 64
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "kimi"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Kimi-K2-Thinking"
 n_layers = 61
 hidden_size = 7168
+num_key_value_heads = 64
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "kimi"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Kimi-K2.5"
 n_layers = 61
 hidden_size = 7168
+num_key_value_heads = 64
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "kimi"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Llama-3.2-1B-Instruct-4bit"
 n_layers = 16
 hidden_size = 2048
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Llama-3.2-3B-Instruct-4bit"
 n_layers = 28
 hidden_size = 3072
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Llama-3.2-3B-Instruct-8bit"
 n_layers = 28
 hidden_size = 3072
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Llama-3.3-70B-Instruct-4bit"
 n_layers = 80
 hidden_size = 8192
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Llama-3.3-70B-Instruct-8bit"
 n_layers = 80
 hidden_size = 8192
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Meta-Llama-3.1-70B-Instruct-4bit"
 n_layers = 80
 hidden_size = 8192
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"
 n_layers = 32
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Meta-Llama-3.1-8B-Instruct-8bit"
 n_layers = 32
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Meta-Llama-3.1-8B-Instruct-bf16"
 n_layers = 32
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/MiniMax-M2.1-3bit"
 n_layers = 61
 hidden_size = 3072
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "minimax"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/MiniMax-M2.1-8bit"
 n_layers = 61
 hidden_size = 3072
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "minimax"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/MiniMax-M2.5-4bit"
 n_layers = 62
 hidden_size = 3072
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "minimax"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/MiniMax-M2.5-6bit"
 n_layers = 62
 hidden_size = 3072
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "minimax"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/MiniMax-M2.5-8bit"
 n_layers = 62
 hidden_size = 3072
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "minimax"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-0.6B-4bit"
 n_layers = 28
 hidden_size = 1024
+num_key_value_heads = 8
 supports_tensor = false
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-0.6B-8bit"
 n_layers = 28
 hidden_size = 1024
+num_key_value_heads = 8
 supports_tensor = false
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-235B-A22B-Instruct-2507-4bit"
 n_layers = 94
 hidden_size = 4096
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-235B-A22B-Instruct-2507-8bit"
 n_layers = 94
 hidden_size = 4096
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-30B-A3B-4bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-30B-A3B-8bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Coder-480B-A35B-Instruct-4bit"
 n_layers = 62
 hidden_size = 6144
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Coder-480B-A35B-Instruct-8bit"
 n_layers = 62
 hidden_size = 6144
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Coder-Next-4bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Coder-Next-5bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Coder-Next-6bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Coder-Next-8bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Coder-Next-bf16"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Next-80B-A3B-Instruct-4bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Next-80B-A3B-Instruct-8bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Next-80B-A3B-Thinking-4bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3-Next-80B-A3B-Thinking-8bit"
 n_layers = 48
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-122B-A10B-4bit"
 n_layers = 48
 hidden_size = 3072
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-122B-A10B-6bit"
 n_layers = 48
 hidden_size = 3072
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-122B-A10B-8bit"
 n_layers = 48
 hidden_size = 3072
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-122B-A10B-bf16"
 n_layers = 48
 hidden_size = 3072
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-27B-4bit"
 n_layers = 64
 hidden_size = 5120
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-27B-8bit"
 n_layers = 64
 hidden_size = 5120
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-2B-MLX-8bit"
 n_layers = 24
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-35B-A3B-4bit"
 n_layers = 40
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-35B-A3B-8bit"
 n_layers = 40
 hidden_size = 2048
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-397B-A17B-4bit"
 n_layers = 60
 hidden_size = 4096
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-397B-A17B-6bit"
 n_layers = 60
 hidden_size = 4096
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-397B-A17B-8bit"
 n_layers = 60
 hidden_size = 4096
+num_key_value_heads = 2
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-9B-4bit"
 n_layers = 32
 hidden_size = 4096
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Qwen3.5-9B-8bit"
 n_layers = 32
 hidden_size = 4096
+num_key_value_heads = 4
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "qwen"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Step-3.5-Flash-4bit"
 n_layers = 45
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "step"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Step-3.5-Flash-6bit"
 n_layers = 45
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "step"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/Step-3.5-Flash-8Bit"
 n_layers = 45
 hidden_size = 4096
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "step"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/gpt-oss-120b-MXFP4-Q8"
 n_layers = 36
 hidden_size = 2880
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "gpt-oss"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/gpt-oss-20b-MXFP4-Q8"
 n_layers = 24
 hidden_size = 2880
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "gpt-oss"
@@ -1,6 +1,7 @@
 model_id = "mlx-community/llama-3.3-70b-instruct-fp16"
 n_layers = 80
 hidden_size = 8192
+num_key_value_heads = 8
 supports_tensor = true
 tasks = ["TextGeneration"]
 family = "llama"
@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+"""Fetch num_key_value_heads from HuggingFace config.json and update TOML model cards.
+
+Usage:
+    # Update only cards missing num_key_value_heads
+    uv run python scripts/fetch_kv_heads.py --missing
+
+    # Update all cards (overwrite existing values)
+    uv run python scripts/fetch_kv_heads.py --all
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import urllib.request
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+
+import tomlkit
+
+CARDS_DIR = (
+    Path(__file__).resolve().parent.parent / "resources" / "inference_model_cards"
+)
+MAX_WORKERS = 5
+
+
+def fetch_kv_heads(model_id: str) -> int | None:
+    """Fetch num_key_value_heads from HuggingFace config.json."""
+    url = f"https://huggingface.co/{model_id}/raw/main/config.json"
+    try:
+        with urllib.request.urlopen(url, timeout=15) as resp:
+            config = json.loads(resp.read())
+    except Exception as e:
+        print(f"  ERROR fetching {url}: {e}", file=sys.stderr)
+        return None
+
+    for source in [config, config.get("text_config", {})]:
+        if "num_key_value_heads" in source:
+            return int(source["num_key_value_heads"])
+
+    return None
+
+
+def update_toml(path: Path, kv_heads: int) -> bool:
+    """Insert or update num_key_value_heads in a TOML file. Returns True if changed."""
+    content = path.read_text()
+    doc = tomlkit.parse(content)
+
+    if doc.get("num_key_value_heads") == kv_heads:
+        return False
+
+    # Insert after hidden_size if adding for the first time
+    if "num_key_value_heads" not in doc:
+        new_doc = tomlkit.document()
+        for key, value in doc.items():
+            new_doc[key] = value
+            if key == "hidden_size":
+                new_doc["num_key_value_heads"] = kv_heads
+        path.write_text(tomlkit.dumps(new_doc))
+    else:
+        doc["num_key_value_heads"] = kv_heads
+        path.write_text(tomlkit.dumps(doc))
+
+    return True
+
+
+def process_card(path: Path) -> tuple[str, str]:
+    """Fetch and update a single card. Returns (filename, status)."""
+    content = path.read_text()
+    doc = tomlkit.parse(content)
+    model_id = doc.get("model_id")
+    if not model_id:
+        return path.name, "SKIP (no model_id)"
+
+    kv_heads = fetch_kv_heads(str(model_id))
+    if kv_heads is None:
+        return path.name, "FAILED"
+
+    changed = update_toml(path, kv_heads)
+    return path.name, f"{kv_heads} ({'UPDATED' if changed else 'UNCHANGED'})"
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Fetch num_key_value_heads from HuggingFace and update TOML cards."
+    )
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        "--all",
+        action="store_true",
+        help="Update all model cards (overwrite existing values)",
+    )
+    group.add_argument(
+        "--missing",
+        action="store_true",
+        help="Only update cards missing num_key_value_heads",
+    )
+    args = parser.parse_args()
+
+    toml_files = sorted(CARDS_DIR.glob("*.toml"))
+    if not toml_files:
+        print(f"No TOML files found in {CARDS_DIR}", file=sys.stderr)
+        sys.exit(1)
+
+    to_process = []
+    skipped = 0
+
+    for path in toml_files:
+        if args.missing and "num_key_value_heads" in path.read_text():
+            skipped += 1
+            continue
+        to_process.append(path)
+
+    updated = 0
+    failed = 0
+
+    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
+        futures = {pool.submit(process_card, path): path for path in to_process}
+        for future in as_completed(futures):
+            name, status = future.result()
+            print(f"  {name}: {status}")
+            if "UPDATED" in status:
+                updated += 1
+            elif "FAILED" in status:
+                failed += 1
+
+    print(f"\nDone: {updated} updated, {skipped} skipped, {failed} failed")
+
+
+if __name__ == "__main__":
+    main()
@@ -1,4 +1,3 @@
-import asyncio
 from dataclasses import dataclass, field

 import anyio
@@ -47,7 +46,7 @@ class DownloadCoordinator:

    # Local state
    download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
-    active_downloads: dict[ModelId, asyncio.Task[None]] = field(default_factory=dict)
+    active_downloads: dict[ModelId, anyio.CancelScope] = field(default_factory=dict)

    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)

@@ -77,8 +76,6 @@ class DownloadCoordinator:
            await self.event_sender.send(
                NodeDownloadProgress(download_progress=completed)
            )
-            if model_id in self.active_downloads:
-                del self.active_downloads[model_id]
            self._last_progress_time.pop(model_id, None)
        elif (
            progress.status == "in_progress"
@@ -103,13 +100,9 @@ class DownloadCoordinator:
        logger.info(
            f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
        )
-        try:
-            async with self._tg as tg:
-                tg.start_soon(self._command_processor)
-                tg.start_soon(self._emit_existing_download_progress)
-        finally:
-            for task in self.active_downloads.values():
-                task.cancel()
+        async with self._tg as tg:
+            tg.start_soon(self._command_processor)
+            tg.start_soon(self._emit_existing_download_progress)

    def shutdown(self) -> None:
        self._tg.cancel_tasks()
@@ -132,7 +125,7 @@ class DownloadCoordinator:
    async def _cancel_download(self, model_id: ModelId) -> None:
        if model_id in self.active_downloads and model_id in self.download_status:
            logger.info(f"Cancelling download for {model_id}")
-            self.active_downloads.pop(model_id).cancel()
+            self.active_downloads[model_id].cancel()
            current_status = self.download_status[model_id]
            pending = DownloadPending(
                shard_metadata=current_status.shard_metadata,
@@ -236,9 +229,10 @@ class DownloadCoordinator:
        self.download_status[model_id] = status
        self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))

-        async def download_wrapper() -> None:
+        async def download_wrapper(cancel_scope: anyio.CancelScope) -> None:
            try:
-                await self.shard_downloader.ensure_shard(shard)
+                with cancel_scope:
+                    await self.shard_downloader.ensure_shard(shard)
            except Exception as e:
                logger.error(f"Download failed for {model_id}: {e}")
                failed = DownloadFailed(
@@ -251,12 +245,15 @@ class DownloadCoordinator:
                await self.event_sender.send(
                    NodeDownloadProgress(download_progress=failed)
                )
+            except anyio.get_cancelled_exc_class():
+                # ignore cancellation - let cleanup do its thing
+                pass
            finally:
-                if model_id in self.active_downloads:
-                    del self.active_downloads[model_id]
+                self.active_downloads.pop(model_id, None)

-        task = asyncio.create_task(download_wrapper())
-        self.active_downloads[model_id] = task
+        scope = anyio.CancelScope()
+        self._tg.start_soon(download_wrapper, scope)
+        self.active_downloads[model_id] = scope

    async def _delete_download(self, model_id: ModelId) -> None:
        # Protect read-only models (from EXO_MODELS_PATH) from deletion
@@ -272,7 +269,6 @@ class DownloadCoordinator:
        if model_id in self.active_downloads:
            logger.info(f"Cancelling active download for {model_id} before deletion")
            self.active_downloads[model_id].cancel()
-            del self.active_downloads[model_id]

        # Delete from disk
        logger.info(f"Deleting model files for {model_id}")
@@ -202,6 +202,8 @@ async def generate_chat_stream(
                    usage=last_usage,
                )
                yield f"data: {tool_response.model_dump_json()}\n\n"
+                if chunk.stats is not None:
+                    yield f": generation_stats {chunk.stats.model_dump_json()}\n\n"
                yield "data: [DONE]\n\n"
                return

@@ -216,6 +218,8 @@ async def generate_chat_stream(
                yield f"data: {chunk_response.model_dump_json()}\n\n"

                if chunk.finish_reason is not None:
+                    if chunk.stats is not None:
+                        yield f": generation_stats {chunk.stats.model_dump_json()}\n\n"
                    yield "data: [DONE]\n\n"


@@ -128,6 +128,28 @@ def responses_request_to_text_generation(
        effort_from_reasoning, request.enable_thinking
    )

+    # The responses API often does not provide tool args nested under a "function" field.
+    # Since we follow the chat completions format of tools in the backend (for MLX chat templates)
+    # we need to normalise to this format.
+    normalised_tools: list[dict[str, Any]] | None = None
+    if request.tools:
+        normalised_tools = []
+        for tool in request.tools:
+            if "function" in tool:
+                normalised_tools.append(tool)
+            else:
+                normalised_tools.append(
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": tool.get("name", ""),
+                            "description": tool.get("description", ""),
+                            "parameters": tool.get("parameters", {}),
+                            **({"strict": tool["strict"]} if "strict" in tool else {}),
+                        },
+                    }
+                )
+
    return TextGenerationTaskParams(
        model=request.model,
        input=input_value,
@@ -136,7 +158,7 @@ def responses_request_to_text_generation(
        temperature=request.temperature,
        top_p=request.top_p,
        stream=request.stream,
-        tools=request.tools,
+        tools=normalised_tools,
        top_k=request.top_k,
        stop=request.stop,
        seed=request.seed,
@@ -90,14 +90,19 @@ def place_instance(
                f"Requested Tensor sharding but this model does not support tensor parallelism: {command.model_card.model_id}"
            )
        # TODO: the condition here for tensor parallel is not correct, but it works good enough for now.
+        kv_heads = command.model_card.num_key_value_heads
        cycles_with_sufficient_memory = [
            cycle
            for cycle in cycles_with_sufficient_memory
            if command.model_card.hidden_size % len(cycle) == 0
+            and (kv_heads is None or kv_heads % len(cycle) == 0)
        ]
        if not cycles_with_sufficient_memory:
            raise ValueError(
-                f"No tensor sharding found for model with hidden_size {command.model_card.hidden_size} candidate cycles"
+                f"No tensor sharding found for model with "
+                f"hidden_size={command.model_card.hidden_size}"
+                f"{f', num_key_value_heads={kv_heads}' if kv_heads is not None else ''}"
+                f" across candidate cycles"
            )
    if command.sharding == Sharding.Pipeline and command.model_card.model_id == ModelId(
        "mlx-community/DeepSeek-V3.1-8bit"
@@ -83,6 +83,7 @@ class ModelCard(CamelCaseModel):
    n_layers: PositiveInt
    hidden_size: PositiveInt
    supports_tensor: bool
+    num_key_value_heads: PositiveInt | None = None
    tasks: list[ModelTask]
    components: list[ComponentInfo] | None = None
    family: str = ""
@@ -137,6 +138,7 @@ class ModelCard(CamelCaseModel):
            n_layers=num_layers,
            hidden_size=config_data.hidden_size or 0,
            supports_tensor=config_data.supports_tensor,
+            num_key_value_heads=config_data.num_key_value_heads,
            tasks=[ModelTask.TextGeneration],
            trust_remote_code=False,
        )
@@ -170,6 +172,7 @@ class ConfigData(BaseModel):

    architectures: list[str] | None = None
    hidden_size: Annotated[int, Field(ge=0)] | None = None
+    num_key_value_heads: PositiveInt | None = None
    layer_count: int = Field(
        validation_alias=AliasChoices(
            "num_hidden_layers",
@@ -209,6 +212,7 @@ class ConfigData(BaseModel):
        for field in [
            "architectures",
            "hidden_size",
+            "num_key_value_heads",
            "num_hidden_layers",
            "num_layers",
            "n_layer",
@@ -529,6 +529,9 @@ class InfoGatherer:
        if self.macmon_interval is None:
            return
        # macmon pipe --interval [interval in ms]
+        # Timeout: if macmon produces no output for this many seconds, restart it.
+        # macmon writes every macmon_interval seconds, so 10x that is generous.
+        read_timeout = max(self.macmon_interval * 10, 30)
        while True:
            try:
                async with await open_process(
@@ -542,10 +545,15 @@ class InfoGatherer:
                    if not p.stdout:
                        logger.critical("MacMon closed stdout")
                        return
-                    async for text in TextReceiveStream(
-                        BufferedByteReceiveStream(p.stdout)
-                    ):
+                    stream = TextReceiveStream(BufferedByteReceiveStream(p.stdout))
+                    while True:
+                        with fail_after(read_timeout):
+                            text = await stream.receive()
                        await self.info_sender.send(MacmonMetrics.from_raw_json(text))
+            except TimeoutError:
+                logger.warning(
+                    f"MacMon produced no output for {read_timeout}s, restarting"
+                )
            except CalledProcessError as e:
                stderr_msg = "no stderr"
                stderr_output = cast(bytes | str | None, e.stderr)
@@ -56,10 +56,10 @@ class QwenJointBlockWrapper(JointBlockWrapper[QwenTransformerBlock]):
        attn = self.block.attn

        img_mod_params = self.block.img_mod_linear(
-            self.block.img_mod_silu(text_embeddings)  # pyright: ignore[reportUnknownArgumentType]
+            self.block.img_mod_silu(text_embeddings)
        )
        txt_mod_params = self.block.txt_mod_linear(
-            self.block.txt_mod_silu(text_embeddings)  # pyright: ignore[reportUnknownArgumentType]
+            self.block.txt_mod_silu(text_embeddings)
        )

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

        return logits

@@ -1,3 +1,4 @@
+import os
 import time
 from dataclasses import dataclass, field
 from typing import Callable, cast
@@ -8,7 +9,7 @@ from mlx_lm.generate import (
 )
 from mlx_lm.models.cache import RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
-from mlx_lm.tokenizer_utils import TokenizerWrapper
+from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper

 from exo.shared.types.api import (
    CompletionTokensDetails,
@@ -57,6 +58,7 @@ class _EngineTask:
    prefix_hit_length: int
    matched_index: int | None
    cache_snapshots: list[CacheSnapshot] | None
+    detokenizer: StreamingDetokenizer
    on_generation_token: Callable[[], None] | None = None
    generated_text_parts: list[str] = field(default_factory=list)
    potential_stop_sequence_text: str = ""
@@ -64,6 +66,7 @@ class _EngineTask:
    generation_start_time: float = 0.0
    in_thinking: bool = False
    reasoning_tokens: int = 0
+    prefill_tps: float = 0.0


@dataclass(eq=False)
@@ -77,12 +80,173 @@ class ExoBatchGenerator:
    _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False)

    def __post_init__(self) -> None:
-        self._exo_gen = MlxBatchGenerator(
-            model=self.model,
-            stop_tokens=set(eos_ids_from_tokenizer(self.tokenizer)),
-            prefill_step_size=4096,
+        use_speculative = os.environ.get("EXO_SPECULATIVE", "0") == "1"
+        stop_tokens = set(eos_ids_from_tokenizer(self.tokenizer))
+
+        if use_speculative:
+            try:
+                from exo.worker.engines.mlx.speculative.mtp_module import MTPPredictor
+                from exo.worker.engines.mlx.speculative.mtp_batch_generator import MTPBatchGenerator
+
+                mtp_weights = self._resolve_mtp_weights()
+                gamma = int(os.environ.get("EXO_SPECULATIVE_GAMMA", "2"))
+
+                if mtp_weights:
+                    mtp = MTPPredictor(self.model, mtp_weights, quantize=False)
+                    temp = float(os.environ.get("EXO_SPECULATIVE_TEMP", "0.7"))
+                    alpha = float(os.environ.get("EXO_SPECULATIVE_ALPHA", "1.0"))
+                    self._exo_gen = MTPBatchGenerator(
+                        model=self.model,
+                        mtp_predictor=mtp,
+                        gamma=gamma,
+                        temp=temp,
+                        alpha=alpha,
+                        stop_tokens=stop_tokens,
+                        prefill_step_size=4096,
+                    )
+                    logger.info(f"MTP speculative decoding enabled (γ={gamma}, T={temp})")
+                    self.warmup_speculative(self.model, self.tokenizer)
+                else:
+                    logger.warning("EXO_SPECULATIVE=1 but could not find MTP weights. Falling back to standard generation.")
+                    self._exo_gen = MlxBatchGenerator(
+                        model=self.model,
+                        stop_tokens=stop_tokens,
+                        prefill_step_size=4096,
+                    )
+            except Exception as e:
+                logger.warning(f"Failed to initialize MTP speculative decoding: {e}. Falling back to standard generation.")
+                self._exo_gen = MlxBatchGenerator(
+                    model=self.model,
+                    stop_tokens=stop_tokens,
+                    prefill_step_size=4096,
+                )
+        else:
+            self._exo_gen = MlxBatchGenerator(
+                model=self.model,
+                stop_tokens=stop_tokens,
+                prefill_step_size=4096,
+            )
+
+    def _resolve_mtp_weights(self) -> str | None:
+        """Find MTP weights: explicit path, explicit HF model, or auto-extract."""
+        # 1. Explicit path
+        explicit_path = os.environ.get("EXO_MTP_WEIGHTS", "")
+        if explicit_path and os.path.exists(explicit_path):
+            return explicit_path
+
+        # 2. Explicit HF model repo containing MTP weights
+        mtp_model = os.environ.get("EXO_MTP_MODEL", "")
+
+        # 3. Auto-detect: if no EXO_MTP_MODEL set, try to infer from model config
+        if not mtp_model:
+            try:
+                inner = getattr(self.model, 'model', None) or self.model.language_model.model
+                args = getattr(inner, 'args', None)
+                if args and getattr(args, 'mtp_num_hidden_layers', 0) > 0:
+                    model_type = getattr(args, 'model_type', '')
+                    if 'qwen3_5' in model_type or 'qwen3.5' in str(type(self.model).__module__):
+                        # Default pairing for Qwen3.5-27B
+                        mtp_model = "Qwen/Qwen3.5-27B"
+                        logger.info(f"Auto-detected MTP model: {mtp_model}")
+            except Exception:
+                pass
+
+        if not mtp_model:
+            return None
+
+        # Download and extract MTP weights from HF repo
+        try:
+            return self._extract_mtp_from_hf(mtp_model)
+        except Exception as e:
+            logger.warning(f"Failed to extract MTP weights from {mtp_model}: {e}")
+            return None
+
+    def _extract_mtp_from_hf(self, repo_id: str) -> str:
+        """Download MTP tensors from HF repo and cache as a single safetensors file."""
+        import hashlib
+        from pathlib import Path
+        from huggingface_hub import snapshot_download
+        from safetensors.torch import load_file, save_file
+
+        cache_dir = Path.home() / ".cache" / "exo" / "mtp_weights"
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        cache_key = hashlib.md5(repo_id.encode()).hexdigest()[:12]
+        cached_path = cache_dir / f"mtp_{cache_key}.safetensors"
+
+        if cached_path.exists():
+            logger.info(f"Using cached MTP weights: {cached_path}")
+            return str(cached_path)
+
+        logger.info(f"Downloading MTP weights from {repo_id}...")
+        model_dir = snapshot_download(
+            repo_id,
+            allow_patterns=["*.safetensors", "*.json"],
        )

+        # Extract MTP tensors from all safetensors files
+        mtp_tensors = {}
+        model_path = Path(model_dir)
+        for sf_file in sorted(model_path.glob("*.safetensors")):
+            tensors = load_file(str(sf_file))
+            for k, v in tensors.items():
+                if k.startswith("model.mtp."):
+                    # Strip "model." prefix to match our MTPPredictor format
+                    clean_key = k[len("model."):]
+                    mtp_tensors[clean_key] = v
+
+        if not mtp_tensors:
+            raise ValueError(f"No MTP tensors found in {repo_id}")
+
+        save_file(mtp_tensors, str(cached_path))
+        logger.info(f"Extracted {len(mtp_tensors)} MTP tensors → {cached_path} ({cached_path.stat().st_size / 1e6:.0f}MB)")
+        return str(cached_path)
+
+    def warmup_speculative(self, model, tokenizer) -> None:
+        """Warm up the speculative decoding path (MTP draft + verify kernels)."""
+        if not hasattr(self._exo_gen, 'mtp'):
+            return
+
+        from mlx_lm.models import cache as cache_mod
+        from exo.worker.engines.mlx.speculative.mtp_module import speculative_forward, draft_tokens
+
+        logger.info("Warming up speculative decoding kernels...")
+        mtp = self._exo_gen.mtp
+        gamma = self._exo_gen.gamma
+
+        # Small warmup: prefill a short prompt, run a few speculative cycles
+        warmup_prompt = tokenizer.encode("Warm up speculative decoding.")
+        cache = cache_mod.make_prompt_cache(model)
+        mtp.reset_cache()
+
+        # Prefill
+        pre_norm, logits = speculative_forward(model, mx.array([warmup_prompt]), cache)
+        mx.eval(pre_norm, logits)
+        next_token = mx.argmax(logits[0, -1], axis=-1).item()
+
+        # MTP prefill
+        if pre_norm.shape[1] > 1:
+            _ = mtp.predict(pre_norm[:, :-1, :], mx.array([warmup_prompt[1:]]))
+            mx.eval(_)
+
+        # Run a few speculative cycles to compile kernels
+        last_pn = pre_norm[:, -1:, :]
+        next_arr = mx.array([[next_token]])
+        for _ in range(3):
+            draft_ids, _ = draft_tokens(mtp, last_pn, next_arr, gamma, 0.0)
+            draft_concat = mx.concatenate([d.reshape(1, 1) for d in draft_ids], axis=1)
+            verify_input = mx.concatenate([next_arr, draft_concat], axis=1)
+            vpn, vl = speculative_forward(model, verify_input, cache, speculative=True)
+            all_next = mx.argmax(vl[0], axis=-1)
+            mx.eval(vpn, all_next)
+            # Accept all for warmup (don't care about correctness)
+            next_arr = all_next[0].reshape(1, 1)
+            last_pn = vpn[:, 0:1, :]
+            for i, c in enumerate(cache):
+                if hasattr(c, 'base'):
+                    cache[i] = c.base
+
+        logger.info("Speculative warmup complete")
+
    @property
    def has_work(self) -> bool:
        return (
@@ -129,10 +293,16 @@ class ExoBatchGenerator:
        seed = task_params.seed if task_params.seed is not None else 42
        mx.random.seed(seed)

+        spec_temp_override = os.environ.get("EXO_SPECULATIVE_TEMP")
+        if spec_temp_override is not None:
+            sampling_temp = float(spec_temp_override)
+        elif task_params.temperature is not None:
+            sampling_temp = task_params.temperature
+        else:
+            sampling_temp = 0.7
+
        sampler = make_sampler(
-            temp=task_params.temperature
-            if task_params.temperature is not None
-            else 0.7,
+            temp=sampling_temp,
            top_p=task_params.top_p if task_params.top_p is not None else 1.0,
            min_p=task_params.min_p if task_params.min_p is not None else 0.05,
            top_k=task_params.top_k if task_params.top_k is not None else 0,
@@ -149,6 +319,23 @@ class ExoBatchGenerator:
            distributed_prompt_progress_callback,
        )

+        # MTP prefill: build MTP KV cache from prompt hidden states
+        # Pair position i with token i+1 (MTP predicts token t+2 from hidden[t] + embed[t+1])
+        if hasattr(self._exo_gen, 'mtp'):
+            prompt_pre_norm = self._exo_gen._captured.get('prompt_pre_norm')
+            if prompt_pre_norm is not None:
+                mx.eval(prompt_pre_norm)
+                self._exo_gen.mtp.reset_cache()
+                S_pre = prompt_pre_norm.shape[1]
+                if S_pre > 0 and len(all_prompt_tokens) > S_pre:
+                    mtp_toks = all_prompt_tokens[1:S_pre + 1].tolist()
+                    _ = self._exo_gen.mtp.predict(
+                        prompt_pre_norm,
+                        mx.array([mtp_toks])
+                    )
+                    mx.eval(_)
+                logger.info(f"MTP cache prefilled ({S_pre} positions)")
+
        # We need to clamp rotating kv caches to max size so that mlx lm's _merge_caches behaves
        for c in cache:
            if (
@@ -198,6 +385,16 @@ class ExoBatchGenerator:

        uid = uids[0]

+        # Pass request temperature to speculative cycle
+        # EXO_SPECULATIVE_TEMP overrides if set; otherwise use request temp
+        if hasattr(self._exo_gen, '_request_temp'):
+            env_temp = os.environ.get("EXO_SPECULATIVE_TEMP")
+            if env_temp is not None:
+                self._exo_gen._request_temp[uid] = float(env_temp)
+            else:
+                request_temp = task_params.temperature if task_params.temperature is not None else 0.7
+                self._exo_gen._request_temp[uid] = request_temp
+
        self._active_tasks[uid] = _EngineTask(
            uid=uid,
            task_params=task_params,
@@ -205,8 +402,10 @@ class ExoBatchGenerator:
            prefix_hit_length=prefix_hit_length,
            matched_index=matched_index,
            cache_snapshots=cache_snapshots or None,
+            detokenizer=self.tokenizer.detokenizer,
            on_generation_token=on_generation_token,
            generation_start_time=time.perf_counter(),
+            prefill_tps=_prefill_tps,
        )

        return uid
@@ -229,11 +428,11 @@ class ExoBatchGenerator:
            state = self._active_tasks[response.uid]
            if state.on_generation_token is not None:
                state.on_generation_token()
-            text = (
-                ""
-                if response.finish_reason == "stop"
-                else self.tokenizer.decode([response.token])
-            )
+            if response.finish_reason != "stop":
+                state.detokenizer.add_token(response.token)
+            if response.finish_reason is not None:
+                state.detokenizer.finalize()
+            text = state.detokenizer.last_segment
            state.completion_tokens += 1
            state.generated_text_parts.append(text)
            state.potential_stop_sequence_text += text
@@ -272,7 +471,7 @@ class ExoBatchGenerator:

            logprob: float | None = None
            top_logprobs: list[TopLogprobItem] | None = None
-            if task_params.logprobs:
+            if task_params.logprobs and os.environ.get("EXO_DISABLE_LOGPROBS") != "1":
                logprob, top_logprobs = extract_top_logprobs(
                    logprobs=response.logprobs,
                    tokenizer=self.tokenizer,
@@ -283,18 +482,22 @@ class ExoBatchGenerator:
            stats: GenerationStats | None = None
            usage: Usage | None = None
            if is_done:
-                generation_elapsed = time.perf_counter() - state.generation_start_time
-                generation_tps = (
-                    state.completion_tokens / generation_elapsed
-                    if generation_elapsed > 0
-                    else 0.0
-                )
-                mlx_stats = self._exo_gen.stats()
+                try:
+                    mlx_stats = self._exo_gen.stats()
+                    generation_tps = mlx_stats.generation_tps
+                except ZeroDivisionError:
+                    generation_elapsed = (
+                        time.perf_counter() - state.generation_start_time
+                    )
+                    generation_tps = (
+                        state.completion_tokens / generation_elapsed
+                        if generation_elapsed > 0
+                        else 0.0
+                    )
+
                stats = GenerationStats(
-                    prompt_tps=float(mlx_stats.prompt_tps)
-                    if mlx_stats.prompt_time > 0
-                    else 0.0,
-                    generation_tps=float(generation_tps),
+                    prompt_tps=state.prefill_tps,
+                    generation_tps=generation_tps,
                    prompt_tokens=len(state.all_prompt_tokens),
                    generation_tokens=state.completion_tokens,
                    peak_memory_usage=Memory.from_gb(mx.get_peak_memory() / 1e9),
@@ -179,7 +179,8 @@ def pipeline_parallel_prefill(
        flush_prefill_sends()

    assert _prompt_cache is not None
-    mx.eval([c.state for c in _prompt_cache])  # type: ignore
+    with mx.stream(generation_stream):
+        mx.eval([c.state for c in _prompt_cache])  # type: ignore

    # Final callback matching generate_step
    prompt_progress_callback(total, total)
@@ -398,52 +399,44 @@ def extract_top_logprobs(
    tokenizer: TokenizerWrapper,
    top_logprobs: int,
    selected_token: int,
+    precomputed_indices: list[int] | None = None,
+    precomputed_values: list[float] | None = None,
+    precomputed_selected: float | None = None,
 ) -> tuple[float, list[TopLogprobItem]]:
-    """Extract the selected token's logprob and top alternative tokens.
-
-    Args:
-        logprobs: Full vocabulary logprobs array from MLX
-        tokenizer: Tokenizer for decoding token IDs to strings
-        top_logprobs: Number of top alternatives to return
-        selected_token: The token ID that was actually sampled
-
-    Returns:
-        Tuple of (selected_token_logprob, list of TopLogprobItem for top alternatives)
-    """
-    # Get the logprob of the selected token
-    selected_logprob = float(logprobs[selected_token].item())
-
-    # Get top indices (most probable tokens)
-    # mx.argpartition gives indices that would partition the array
-    # We negate logprobs since argpartition finds smallest, and we want largest
-    top_logprobs = min(top_logprobs, logprobs.shape[0])  # Don't exceed vocab size
-    top_indices = mx.argpartition(-logprobs, top_logprobs)[:top_logprobs]
-
-    # Get the actual logprob values for these indices
-    top_values = logprobs[top_indices]
-
-    # Sort by logprob (descending) for consistent ordering
-    sort_order = mx.argsort(-top_values)
-    top_indices = top_indices[sort_order]
-    top_values = top_values[sort_order]
+    if (
+        precomputed_indices is not None
+        and precomputed_values is not None
+        and precomputed_selected is not None
+    ):
+        top_indices_list: list[int] = precomputed_indices[:top_logprobs]
+        top_values_list: list[float] = precomputed_values[:top_logprobs]
+        selected_logprob = precomputed_selected
+    else:
+        selected_logprob_arr = logprobs[selected_token]
+        top_logprobs = min(top_logprobs, logprobs.shape[0] - 1)
+        top_indices = mx.argpartition(-logprobs, top_logprobs)[:top_logprobs]
+        top_values = logprobs[top_indices]
+        sort_order = mx.argsort(-top_values)
+        top_indices = top_indices[sort_order]
+        top_values = top_values[sort_order]
+        mx.eval(selected_logprob_arr, top_indices, top_values)
+        selected_logprob = float(selected_logprob_arr.item())
+        top_indices_list = top_indices.tolist()  # type: ignore
+        top_values_list = top_values.tolist()  # type: ignore

    # Convert to list of TopLogprobItem
    top_logprob_items: list[TopLogprobItem] = []
-    for i in range(top_logprobs):
-        token_id = int(top_indices[i].item())
-        token_logprob = float(top_values[i].item())
+    for token_id, token_logprob in zip(top_indices_list, top_values_list, strict=True):
        if math.isnan(token_logprob):
            continue

        # Decode token ID to string
        token_str = tokenizer.decode([token_id])
-        # Get byte representation
-        token_bytes = list(token_str.encode("utf-8"))
        top_logprob_items.append(
            TopLogprobItem(
                token=token_str,
                logprob=token_logprob,
-                bytes=token_bytes,
+                bytes=list(token_str.encode("utf-8")),
            )
        )

@@ -624,12 +617,13 @@ def mlx_generate(
        logprob: float | None = None
        top_logprobs: list[TopLogprobItem] | None = None
        if task.logprobs:
-            logprob, top_logprobs = extract_top_logprobs(
-                logprobs=out.logprobs,
-                tokenizer=tokenizer,
-                top_logprobs=task.top_logprobs or DEFAULT_TOP_LOGPROBS,
-                selected_token=out.token,
-            )
+            with mx.stream(generation_stream):
+                logprob, top_logprobs = extract_top_logprobs(
+                    logprobs=out.logprobs,
+                    tokenizer=tokenizer,
+                    top_logprobs=task.top_logprobs or DEFAULT_TOP_LOGPROBS,
+                    selected_token=out.token,
+                )

        if is_done:
            # Log generation stats
@@ -0,0 +1,45 @@
+"""Model-specific kernel fusion patches for MLX inference.
+
+Detects model type after loading and applies optimized kernel patches.
+Currently supports:
+- Qwen3.5 MoE (model_type: qwen3_5_moe): batched fused oproj (GDN + GQA + MoE)
+
+Set EXO_FUSED_KERNELS=0 to disable all patches (vanilla mode).
+Default: EXO_FUSED_KERNELS=1 (enabled).
+"""
+
+import json
+import os
+from pathlib import Path
+
+import mlx.nn as nn
+from loguru import logger
+
+
+def maybe_apply_patches(model: nn.Module, model_path: Path) -> None:
+    """Detect model type and apply kernel fusion patches if available."""
+    fused_mode = os.environ.get("EXO_FUSED_KERNELS", "1")
+    if fused_mode == "0":
+        logger.info("Kernel fusion patches disabled (EXO_FUSED_KERNELS=0)")
+        return
+
+    config_path = model_path / "config.json"
+    if not config_path.exists():
+        return
+
+    with open(config_path) as f:
+        config = json.load(f)
+
+    model_type = config.get("model_type", "")
+
+    if model_type == "qwen3_5_moe":
+        from .qwen3_5_moe.apply import apply_qwen35_batched_fused_patches
+
+        logger.info("Detected Qwen3.5 MoE model, applying batched fused kernel patches")
+        apply_qwen35_batched_fused_patches(model)
+
+    elif model_type == "qwen3_5":
+        from .qwen3_5.lpb_patch import apply_lpb_patches
+
+        logger.info("Detected Qwen3.5 dense model, applying LpB kernel patches")
+        apply_lpb_patches(model, batch_size=4)
@@ -0,0 +1,27 @@
+import sys
+
+import mlx.core as mx
+from mlx_lm.models.gated_delta import compute_g
+
+
+def _compute_g_f32(a_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array:
+    return mx.exp(
+        -mx.exp(a_log.astype(mx.float32))
+        * mx.where(
+            (a + dt_bias).astype(mx.float32) > 20,
+            (a + dt_bias).astype(mx.float32),
+            mx.log1p(mx.exp((a + dt_bias).astype(mx.float32))),
+        )
+    ).astype(a.dtype)
+
+
+def patch_gdn_softplus() -> None:
+    from mlx_lm.models import gated_delta
+
+    gated_delta.compute_g = _compute_g_f32
+
+    for mod in list(sys.modules.values()):
+        if mod is gated_delta:
+            continue
+        if getattr(mod, "compute_g", None) is compute_g:
+            object.__setattr__(mod, "compute_g", _compute_g_f32)
@@ -0,0 +1,173 @@
+import time
+from typing import Any, cast
+
+import mlx.core as mx
+from mlx_lm.generate import BatchGenerator, generation_stream
+
+_PRECOMPUTE_TOP_K = 20
+
+_original_public_next = BatchGenerator.next
+
+_pending_topk_idx: mx.array | None = None
+_pending_topk_val: mx.array | None = None
+_pending_selected_lps: mx.array | None = None
+
+
+def _fast_next(self: BatchGenerator) -> list[BatchGenerator.Response]:
+    tic = time.perf_counter()
+    batch = self.active_batch
+    assert batch is not None
+    batch_size = len(batch)
+
+    prev_tokens = batch.y
+    prev_logprobs = batch.logprobs
+
+    has_processors = any(p for ps in batch.logits_processors for p in ps)
+    if has_processors:
+        for i, toks in enumerate(batch.tokens):
+            batch.tokens[i] = mx.concatenate([toks, prev_tokens[i : i + 1]])
+
+    logits = self.model(prev_tokens[:, None], cache=batch.cache)
+    logits = logits[:, -1, :]
+
+    if has_processors:
+        processed_logits: list[mx.array] = []
+        for e in range(batch_size):
+            sample_logits: mx.array = logits[e : e + 1]
+            for processor in batch.logits_processors[e]:
+                sample_logits = processor(batch.tokens[e], sample_logits)
+            processed_logits.append(sample_logits)
+        logits = mx.concatenate(processed_logits, axis=0)
+
+    logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+
+    if (
+        batch_size == 1
+        or any(batch.samplers)
+        and all(s is batch.samplers[0] for s in batch.samplers)
+    ):
+        sampler = batch.samplers[0] or self.sampler
+        batch.y = sampler(logprobs)
+    elif any(batch.samplers):
+        all_samples: list[mx.array] = []
+        for e in range(batch_size):
+            s = batch.samplers[e] or self.sampler
+            all_samples.append(s(logprobs[e : e + 1]))
+        batch.y = mx.concatenate(all_samples, axis=0)
+    else:
+        batch.y = self.sampler(logprobs)
+    batch.logprobs = list(logprobs)
+
+    global _pending_topk_idx, _pending_topk_val, _pending_selected_lps
+
+    emit_topk_indices: list[list[int]] = (
+        cast(list[list[int]], _pending_topk_idx.tolist())
+        if _pending_topk_idx is not None
+        else []
+    )
+    emit_topk_values: list[list[float]] = (
+        cast(list[list[float]], _pending_topk_val.tolist())
+        if _pending_topk_val is not None
+        else []
+    )
+    emit_selected_lps: list[float] = (
+        cast(list[float], _pending_selected_lps.tolist())
+        if _pending_selected_lps is not None
+        else []
+    )
+
+    needs_topk: bool = getattr(self, "_needs_topk", False)
+    if needs_topk:
+        k = min(_PRECOMPUTE_TOP_K, logprobs.shape[1])
+        _pending_topk_idx = mx.argpartition(-logprobs, k, axis=1)[:, :k]
+        _pending_topk_val = mx.take_along_axis(logprobs, _pending_topk_idx, axis=1)
+        sort_order = mx.argsort(-_pending_topk_val, axis=1)
+        _pending_topk_idx = mx.take_along_axis(_pending_topk_idx, sort_order, axis=1)
+        _pending_topk_val = mx.take_along_axis(_pending_topk_val, sort_order, axis=1)
+        _pending_selected_lps = logprobs[mx.arange(batch_size), batch.y]
+        mx.async_eval(
+            batch.y,
+            *batch.logprobs,
+            *batch.tokens,
+            _pending_topk_idx,
+            _pending_topk_val,
+            _pending_selected_lps,
+        )
+    else:
+        _pending_topk_idx = None
+        _pending_topk_val = None
+        _pending_selected_lps = None
+        mx.async_eval(batch.y, *batch.logprobs, *batch.tokens)
+
+    prev_token_list: list[int] = cast(list[int], prev_tokens.tolist())
+
+    toc = time.perf_counter()
+    self._stats.generation_time += toc - tic
+
+    keep_idx: list[int] = []
+    end_idx: list[int] = []
+    responses: list[Any] = []
+    stop_tokens = self.stop_tokens
+
+    for e in range(batch_size):
+        t = prev_token_list[e]
+        uid = batch.uids[e]
+        num_tok = batch.num_tokens[e] + 1
+        batch.num_tokens[e] = num_tok
+
+        if t in stop_tokens:
+            finish_reason = "stop"
+            end_idx.append(e)
+        elif num_tok >= batch.max_tokens[e]:
+            finish_reason = "length"
+            end_idx.append(e)
+        else:
+            finish_reason = None
+            keep_idx.append(e)
+
+        cache = None
+        if finish_reason is not None:
+            cache = batch.extract_cache(e)
+        response = self.Response(uid, t, prev_logprobs[e], finish_reason, cache)
+        if emit_topk_indices and e < len(emit_topk_indices):
+            response._topk_indices = emit_topk_indices[e]  # pyright: ignore[reportAttributeAccessIssue]
+            response._topk_values = emit_topk_values[e]  # pyright: ignore[reportAttributeAccessIssue]
+            response._selected_logprob = emit_selected_lps[e]  # pyright: ignore[reportAttributeAccessIssue]
+        responses.append(response)
+
+    if end_idx:
+        if keep_idx:
+            batch.filter(keep_idx)
+            if (
+                _pending_topk_idx is not None
+                and _pending_topk_val is not None
+                and _pending_selected_lps is not None
+            ):
+                ki = mx.array(keep_idx)
+                _pending_topk_idx = _pending_topk_idx[ki]
+                _pending_topk_val = _pending_topk_val[ki]
+                _pending_selected_lps = _pending_selected_lps[ki]
+        else:
+            self.active_batch = None
+            _pending_topk_idx = None
+            _pending_topk_val = None
+            _pending_selected_lps = None
+
+    self._next_count += 1
+    if self._next_count % 512 == 0:
+        mx.clear_cache()
+    self._stats.generation_tokens += len(responses)
+    return responses
+
+
+def _patched_public_next(self: BatchGenerator) -> list[BatchGenerator.Response]:
+    batch = self.active_batch
+    # Only do decode with fast_next
+    if batch is not None and not self.unprocessed_prompts:
+        with mx.stream(generation_stream):
+            return _fast_next(self)
+    return _original_public_next(self)
+
+
+def apply_batch_gen_patch() -> None:
+    BatchGenerator.next = _patched_public_next
--- a/Show More
+++ b/Show More