implement engine interface for mlx and mflux

Add more model cards (#1970 )
chore(app): hardcode bug report presigned-URL endpoint (#1971 )
2026-09-08 19:41:32 -04:00 · 2026-04-25 01:52:29 +01:00 · 2026-04-23 15:28:40 +01:00 · 2026-04-23 14:08:16 +00:00 · 2026-04-23 11:39:36 +01:00 · 2026-04-23 01:50:39 +00:00
74 changed files with 2853 additions and 2373 deletions
@@ -32,7 +32,6 @@ jobs:
      SPARKLE_ED25519_PRIVATE: ${{ secrets.SPARKLE_ED25519_PRIVATE }}
      SPARKLE_S3_BUCKET: ${{ secrets.SPARKLE_S3_BUCKET }}
      SPARKLE_S3_PREFIX: ${{ secrets.SPARKLE_S3_PREFIX }}
-      EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT: ${{ secrets.EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT }}
      AWS_REGION: ${{ secrets.AWS_REGION }}
      EXO_BUILD_NUMBER: ${{ github.run_number }}
      EXO_LIBP2P_NAMESPACE: ${{ github.ref_name }}
@@ -347,7 +346,6 @@ jobs:
            EXO_BUILD_COMMIT="$GITHUB_SHA" \
            SPARKLE_FEED_URL="$SPARKLE_FEED_URL" \
            SPARKLE_ED25519_PUBLIC="$SPARKLE_ED25519_PUBLIC" \
-            EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT="$EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT" \
            CODE_SIGNING_IDENTITY="$SIGNING_IDENTITY" \
            CODE_SIGN_INJECT_BASE_ENTITLEMENTS=YES
          mkdir -p ../../output
@@ -1767,12 +1767,12 @@ def clip(
        array: The clipped array.
    """

-def compile(
-    fun: Callable,
+def compile[F: Callable[..., object]](
+    fun: F,
    inputs: object | None = ...,
    outputs: object | None = ...,
    shapeless: bool = ...,
-) -> Callable:
+) -> F:
    """
    Returns a compiled function which produces the same output as ``fun``.

@@ -2915,8 +2915,8 @@ def gather_mm(
    a: array,
    b: array,
    /,
-    lhs_indices: array,
-    rhs_indices: array,
+    lhs_indices: array | None = ...,
+    rhs_indices: array | None = ...,
    *,
    sorted_indices: bool = ...,
    stream: Stream | Device | None = ...,
@@ -4707,6 +4707,7 @@ def softmax(
    /,
    axis: int | Sequence[int] | None = ...,
    *,
+    precise: bool = ...,
    stream: Stream | Device | None = ...,
 ) -> array:
    """
@@ -57,6 +57,10 @@ class Module(dict):
    def __init__(self) -> None:
        """Should be called by the subclasses of ``Module``."""

+    def __getitem__(self, key: str) -> mx.array | Module: ...
+    def get(
+        self, key: str, default: mx.array | Module | None = ...
+    ) -> mx.array | Module | None: ...
    @property
    def training(self):  # -> bool:
        """Boolean indicating if the model is in training mode."""
@@ -3,7 +3,7 @@ This type stub file was generated by pyright.
 """

 from dataclasses import dataclass
-from typing import Optional
+from typing import Any, Optional

 import mlx.core as mx

@@ -37,10 +37,10 @@ def quantized_scaled_dot_product_attention(
    bits: int = ...,
 ) -> mx.array: ...
 def scaled_dot_product_attention(
-    queries,
-    keys,
-    values,
-    cache,
+    queries: mx.array,
+    keys: mx.array,
+    values: mx.array,
+    cache: Optional[Any],
    scale: float,
    mask: Optional[mx.array],
    sinks: Optional[mx.array] = ...,
@@ -0,0 +1,103 @@
+"""Type stubs for mlx_lm.models.gpt_oss"""
+
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import KVCache
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    vocab_size: int
+    rms_norm_eps: float
+    sliding_window: int
+    layer_types: Optional[List[str]]
+
+def mlx_topk(a: mx.array, k: int, axis: int = -1) -> tuple[mx.array, mx.array]: ...
+
+class AttentionBlock(nn.Module):
+    head_dim: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_key_value_groups: int
+    sinks: mx.array
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    sm_scale: float
+    rope: nn.Module
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class TransformerBlock(nn.Module):
+    self_attn: AttentionBlock
+    mlp: MLPBlock
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MLPBlock(nn.Module):
+    hidden_size: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    experts: SwitchGLU
+    router: nn.Linear
+    sharding_group: Optional[mx.distributed.Group]
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class GptOssMoeModel(nn.Module):
+    embed_tokens: nn.Embedding
+    norm: nn.RMSNorm
+    layer_types: List[str]
+    layers: list[TransformerBlock]
+    window_size: int
+    swa_idx: int
+    ga_idx: int
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    model_type: str
+    model: GptOssMoeModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[nn.Module]: ...
+    def make_cache(self) -> list[KVCache]: ...
@@ -0,0 +1,94 @@
+"""Type stubs for mlx_lm.models.minimax"""
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_local_experts: int
+    num_experts_per_tok: int
+    max_position_embeddings: int
+
+class MiniMaxAttention(nn.Module):
+    num_heads: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    scale: float
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    q_norm: nn.Module
+    k_norm: nn.Module
+    rope: nn.Module
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MiniMaxSparseMoeBlock(nn.Module):
+    num_experts_per_tok: int
+    gate: nn.Linear
+    switch_mlp: SwitchGLU
+    e_score_correction_bias: mx.array
+    sharding_group: Optional[mx.distributed.Group]
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class MiniMaxDecoderLayer(nn.Module):
+    self_attn: MiniMaxAttention
+    block_sparse_moe: MiniMaxSparseMoeBlock
+    input_layernorm: nn.RMSNorm
+    post_attention_layernorm: nn.RMSNorm
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class MiniMaxModel(nn.Module):
+    embed_tokens: nn.Embedding
+    layers: list[MiniMaxDecoderLayer]
+    norm: nn.RMSNorm
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    model_type: str
+    model: MiniMaxModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[MiniMaxDecoderLayer]: ...
@@ -92,6 +92,15 @@ class NemotronHAttention(nn.Module):
        cache: Optional[KVCache] = None,
    ) -> mx.array: ...

+class MoEGate(nn.Module):
+    config: ModelArgs
+    top_k: int
+    norm_topk_prob: bool
+    weight: mx.array
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
+
 class NemotronHMLP(nn.Module):
    up_proj: nn.Linear
    down_proj: nn.Linear
@@ -102,9 +111,14 @@ class NemotronHMLP(nn.Module):
    def __call__(self, x: mx.array) -> mx.array: ...

 class NemotronHMoE(nn.Module):
+    config: ModelArgs
    num_experts_per_tok: int
+    moe_latent_size: Optional[int]
    switch_mlp: SwitchMLP
+    gate: MoEGate
    shared_experts: NemotronHMLP
+    fc1_latent_proj: nn.Linear
+    fc2_latent_proj: nn.Linear

    def __init__(self, config: ModelArgs) -> None: ...
    def __call__(self, x: mx.array) -> mx.array: ...
@@ -71,6 +71,7 @@ class Qwen3NextAttention(nn.Module):
 class Qwen3NextSparseMoeBlock(nn.Module):
    norm_topk_prob: bool
    num_experts: int
+    num_experts_per_tok: int
    top_k: int
    gate: nn.Linear
    switch_mlp: SwitchGLU
@@ -584,9 +584,18 @@ struct ContentView: View {

            case .prompting:
                VStack(alignment: .leading, spacing: 6) {
-                    Text("What's the issue? (optional)")
+                    VStack(alignment: .leading, spacing: 2) {
+                        Text("Tell us what went wrong (optional)")
+                            .font(.caption2)
+                            .foregroundColor(.secondary)
+                        Text(
+                            "A quick description of what you were doing and what happened helps us track down the bug for you."
+                        )
                        .font(.caption2)
                        .foregroundColor(.secondary)
+                        .opacity(0.8)
+                        .fixedSize(horizontal: false, vertical: true)
+                    }
                    TextEditor(text: $bugReportUserDescription)
                        .font(.caption2)
                        .frame(height: 60)
@@ -9,7 +9,7 @@
 	<key>EXOBuildCommit</key>
 	<string>$(EXO_BUILD_COMMIT)</string>
 	<key>EXOBugReportPresignedUrlEndpoint</key>
-	<string>$(EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT)</string>
+	<string>https://reports.exolabs.net/presigned-urls</string>
 	<key>NSLocalNetworkUsageDescription</key>
 	<string>EXO needs local network access to discover and connect to other devices in your cluster for distributed AI inference.</string>
 	<key>NSBonjourServices</key>
@@ -3,11 +3,13 @@ from __future__ import annotations

 import argparse
 import contextlib
+import io
 import json
 import os
 import sys
 import time
 import tomllib
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Literal
@@ -209,7 +211,7 @@ def _openai_build_request(
        "model": model,
        "messages": messages,
        "tools": tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
        "temperature": 0.0,
    }
    return "/v1/chat/completions", body
@@ -276,7 +278,7 @@ def _openai_build_followup(
        "model": model,
        "messages": followup_messages,
        "tools": tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
        "temperature": 0.0,
    }
    return "/v1/chat/completions", body
@@ -379,7 +381,7 @@ def _claude_build_request(
        "model": model,
        "messages": claude_messages,
        "tools": claude_tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
        "temperature": 0.0,
    }
    if system_content is not None:
@@ -489,7 +491,7 @@ def _claude_build_followup(
        "model": model,
        "messages": claude_messages,
        "tools": claude_tools,
-        "max_tokens": 16384,
+        "max_tokens": 4096,
        "temperature": 0.0,
    }
    if system_content is not None:
@@ -913,6 +915,12 @@ Examples:
        default=1,
        help="Repeat each scenario N times (default: 1)",
    )
+    parser.add_argument(
+        "--concurrency",
+        type=int,
+        default=1,
+        help="Run up to N scenarios in parallel against the same instance (default: 1)",
+    )
    parser.add_argument(
        "--scenarios",
        nargs="*",
@@ -935,6 +943,13 @@ Examples:
    )
    args = parser.parse_args()

+    if args.concurrency < 1:
+        print(
+            f"--concurrency must be >= 1 (got {args.concurrency})",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+
    all_scenarios = load_scenarios(SCENARIOS_PATH)
    if args.scenarios:
        scenarios = [s for s in all_scenarios if s.name in args.scenarios]
@@ -1010,42 +1025,72 @@ Examples:
    cluster_snapshot = capture_cluster_snapshot(exo)
    all_results: list[ScenarioResult] = []

+    tasks: list[tuple[int, Scenario, ApiName]] = [
+        (run_idx, scenario, api_name)
+        for run_idx in range(args.repeat)
+        for scenario in scenarios
+        for api_name in api_names
+    ]
+
+    def _run_one(
+        http_client: httpx.Client,
+        task: tuple[int, Scenario, ApiName],
+    ) -> tuple[tuple[int, Scenario, ApiName], list[ScenarioResult], str]:
+        run_idx, scenario, api_name = task
+        buf = io.StringIO()
+        run_tag = f"[run {run_idx + 1}/{args.repeat}]" if args.repeat > 1 else ""
+        print(
+            f"\n  {run_tag}[{api_name:>9}] {scenario.name}: {scenario.description}",
+            file=buf,
+        )
+        scenario_results = run_scenario(
+            http_client,
+            args.host,
+            args.port,
+            full_model_id,
+            scenario,
+            api_name,
+            args.timeout,
+            args.verbose,
+        )
+        for r in scenario_results:
+            status = "PASS" if r.passed else "FAIL"
+            print(
+                f"    [{r.phase:>10}] {status}  ({r.latency_ms:.0f}ms)",
+                file=buf,
+            )
+            for check_name, check_ok in r.checks.items():
+                mark = "+" if check_ok else "-"
+                print(f"      {mark} {check_name}", file=buf)
+            if r.error:
+                print(f"      ! {r.error}", file=buf)
+        return task, scenario_results, buf.getvalue()
+
    try:
        with httpx.Client() as http_client:
-            for run_idx in range(args.repeat):
-                if args.repeat > 1:
-                    print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
-
-                for scenario in scenarios:
-                    for api_name in api_names:
-                        print(
-                            f"\n  [{api_name:>9}] {scenario.name}: {scenario.description}",
-                            file=log,
-                        )
-
-                        scenario_results = run_scenario(
-                            http_client,
-                            args.host,
-                            args.port,
-                            full_model_id,
-                            scenario,
-                            api_name,
-                            args.timeout,
-                            args.verbose,
-                        )
+            if args.concurrency == 1:
+                current_run = -1
+                for task in tasks:
+                    run_idx = task[0]
+                    if args.repeat > 1 and run_idx != current_run:
+                        print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
+                        current_run = run_idx
+                    _, scenario_results, buffered = _run_one(http_client, task)
+                    all_results.extend(scenario_results)
+                    log.write(buffered)
+                    log.flush()
+            else:
+                print(
+                    f"Running {len(tasks)} tasks with concurrency={args.concurrency}",
+                    file=log,
+                )
+                with ThreadPoolExecutor(max_workers=args.concurrency) as pool:
+                    futures = [pool.submit(_run_one, http_client, t) for t in tasks]
+                    for fut in as_completed(futures):
+                        _, scenario_results, buffered = fut.result()
                        all_results.extend(scenario_results)
-
-                        for r in scenario_results:
-                            status = "PASS" if r.passed else "FAIL"
-                            print(
-                                f"    [{r.phase:>10}] {status}  ({r.latency_ms:.0f}ms)",
-                                file=log,
-                            )
-                            for check_name, check_ok in r.checks.items():
-                                mark = "+" if check_ok else "-"
-                                print(f"      {mark} {check_name}", file=log)
-                            if r.error:
-                                print(f"      ! {r.error}", file=log)
+                        log.write(buffered)
+                        log.flush()
    finally:
        try:
            exo.request_json("DELETE", f"/instance/{instance_id}")
@@ -564,7 +564,7 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
    ap.add_argument(
        "--settle-timeout",
        type=float,
-        default=0,
+        default=60.0,
        help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
    )
    ap.add_argument(
@@ -88,10 +88,12 @@
  let codexModel = $state("");
  let codexMcpPath = $state("/Users/username");
  let openClawModel = $state("");
+  let piModel = $state("");
  $effect(() => {
    const def = modelsBySize.length > 0 ? modelsBySize[0] : "your-model-id";
    codexModel = def;
    openClawModel = def;
+    piModel = def;
  });

  const claudeShellCommand = $derived(
@@ -218,6 +220,55 @@
    ),
  );

+  const piModelsJson = $derived.by(() => {
+    const models: Record<string, unknown>[] = [];
+    for (const modelId of runningModels) {
+      const caps = modelCapabilities[modelId] || [];
+      const ctxLen = modelContextLengths[modelId] || 0;
+      const entry: Record<string, unknown> = { id: modelId };
+      if (caps.includes("vision")) {
+        entry.input = ["text", "image"];
+      }
+      // Mark thinking-capable models so pi surfaces its thinking-level selector
+      // for them. exo capability strings: "thinking" (model emits reasoning
+      // content) and "thinking_toggle" (user can turn it on/off).
+      if (caps.includes("thinking") || caps.includes("thinking_toggle")) {
+        entry.reasoning = true;
+      }
+      if (ctxLen > 0) {
+        entry.contextWindow = ctxLen;
+      }
+      models.push(entry);
+    }
+    if (models.length === 0) {
+      models.push({ id: "your-model-id" });
+    }
+    return JSON.stringify(
+      {
+        providers: {
+          exo: {
+            baseUrl: `${apiUrl}/v1`,
+            api: "openai-completions",
+            apiKey: "exo",
+            compat: {
+              supportsDeveloperRole: false,
+              // exo's OpenAI surface takes a boolean `enable_thinking` toggle,
+              // not graded effort levels, so disable pi's `reasoning_effort`
+              // parameter and use the matching top-level-boolean format.
+              supportsReasoningEffort: false,
+              thinkingFormat: "qwen",
+            },
+            models,
+          },
+        },
+      },
+      null,
+      2,
+    );
+  });
+
+  const piShellCommand = $derived(`pi --provider exo --model ${piModel}`);
+
  const ollamaCommand = $derived(
    `OLLAMA_HOST=${apiUrl}/ollama ollama run ${modelsBySize.length > 0 ? modelsBySize[0] : "your-model-id"}`,
  );
@@ -277,6 +328,7 @@
    "OpenCode",
    "Codex",
    "OpenClaw",
+    "Pi",
    "Open WebUI",
    "n8n",
    "Firefox",
@@ -515,6 +567,33 @@
          config={`openclaw doctor --fix${(modelCapabilities[openClawModel] || []).includes("vision") ? `\nopenclaw models set-image exo/${openClawModel}` : ""}\nopenclaw gateway &\nopenclaw dashboard`}
          language="bash"
        />
+      {:else if activeTab === "Pi"}
+        {#if runningModels.length > 1}
+          <div class="text-xs">
+            <span
+              class="text-exo-light-gray/50 text-[10px] uppercase tracking-wider block mb-1"
+              >Model</span
+            >
+            <select bind:value={piModel} class={selectClass}>
+              {#each runningModels as model}
+                <option value={model}>{model.split("/").pop()}</option>
+              {/each}
+            </select>
+          </div>
+        {/if}
+        <IntegrationCard
+          title="Models Config"
+          subtitle="~/.pi/agent/models.json"
+          description="Register exo as a custom provider in pi. Create or edit this file, then run pi and pick an exo model via /model. Install pi with: npm install -g @mariozechner/pi-coding-agent"
+          config={piModelsJson}
+        />
+        <IntegrationCard
+          title="Shell Command"
+          subtitle="Run in terminal"
+          description="Launch pi directly with the exo provider and model selected."
+          config={piShellCommand}
+          language="bash"
+        />
      {:else if activeTab === "Open WebUI"}
        <IntegrationCard
          title="1. Start Open WebUI"
@@ -0,0 +1,6 @@
+{
+  "name": "exo",
+  "lockfileVersion": 3,
+  "requires": true,
+  "packages": {}
+}
@@ -3,7 +3,7 @@ name = "exo"
 version = "0.3.70"
 description = "Exo"
 readme = "README.md"
-requires-python = ">=3.13"
+requires-python = "==3.13.*"
 dependencies = [
  "aiofiles>=24.1.0",
  "aiohttp>=3.12.14",
@@ -17,8 +17,8 @@ dependencies = [
  "loguru>=0.7.3",
  "exo-pyo3-bindings",                       # rust bindings
  "anyio==4.11.0",
-  "mlx==0.31.1; sys_platform == 'darwin'",
-  "mlx-lm",
+  "mlx==0.31.2; sys_platform == 'darwin'",
+  "mlx-lm; sys_platform=='darwin'",
  "tiktoken>=0.12.0",                        # required for kimi k2 tokenizer
  "hypercorn>=0.18.0",
  "openai-harmony>=0.0.8",
@@ -30,7 +30,6 @@ dependencies = [
  "zstandard>=0.23.0",
  "mlx-vlm>=0.3.11",
  "transformers>=5.0.0,<5.4.0",
-  "pydantic-settings>=2.13.1",
 ]

 [project.scripts]
@@ -50,15 +49,21 @@ dev = [
 [project.optional-dependencies]
 build = ["nanobind"]
 cpu = [
+  "mlx==0.31.1; sys_platform == 'linux'",
  "mlx-cpu==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
  "torch>=2.10.0; sys_platform == 'linux'",
 ]
 cuda12 = [
+  "mlx==0.31.1; sys_platform == 'linux'",
  "mlx-cuda-12==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
  "torch>=2.10.0; sys_platform == 'linux'",
 ]
 cuda13 = [
+  "mlx==0.31.1; sys_platform == 'linux'",
  "mlx-cuda-13==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
  "torch>=2.10.0; sys_platform == 'linux'",
 ]

@@ -76,7 +81,7 @@ mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/fix-arra
 torch = [
  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'cuda13' and extra != 'cpu' and extra != 'cuda12'" },
  { index = "pytorch-cu120", marker = "sys_platform == 'linux' and extra == 'cuda12' and extra != 'cpu' and extra != 'cuda13'" },
-  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'cpu' and extra != 'cuda12' and extra != 'cuda13'" },
+  { index = "pytorch-cpu", marker = "(extra != 'cuda12' and extra != 'cuda13' and sys_platform == 'linux') or sys_platform == 'darwin'" },
 ]
 vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }

@@ -150,6 +155,10 @@ prerelease = "allow"
 environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
 conflicts = [[{ extra = "cuda12" }, { extra = "cuda13" }, { extra = "cpu" }]]
 constraint-dependencies = ["transformers>=5.0.0,<5.4.0"]
+override-dependencies = [
+  "mlx==0.31.1; sys_platform=='linux'",
+  "mlx; sys_platform=='darwin'",
+]

 [tool.uv.extra-build-dependencies]
 miniaudio = ["setuptools", "cffi", "pycparser"]
@@ -5,14 +5,13 @@ let
    workspaceRoot = ../.;
  };

-  mkPythonSet = { pkgs, lib, self' }:
+  mkPythonSet = { pkgs, lib, self', members }:
    let
      inherit (pkgs.stdenv.hostPlatform) isLinux isDarwin isx86_64;
      inherit (pkgs.config) cudaSupport;
      inherit (pkgs) cudaPackages;
      cuda13Support = cudaSupport && cudaPackages.cudaMajorVersion == "13";
      libmlx_source = if cuda13Support then "mlx-cuda-13" else if cudaSupport then "mlx-cuda-12" else "mlx-cpu";
-      uv_extra = if cuda13Support then "cuda13" else if cudaSupport then "cuda12" else "cpu";
      python = pkgs.python313;
      cudaLibs = with cudaPackages; [
        cuda_cudart
@@ -51,7 +50,7 @@ let
          '';
        };
      };
-      buildSystemsOverlay = final: prev: { } //
+      buildSystemsOverlay = final: prev:
        lib.optionalAttrs isDarwin
          {
            mlx = prev.mlx.overrideAttrs (old:
@@ -81,7 +80,7 @@ let
                nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake self'.packages.metal-toolchain ];
                # TODO: non-sdk_26 support
                buildInputs = (old.buildInputs or [ ])
-                  ++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.apple-sdk_26 ];
+                ++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.apple-sdk_26 ];
                patches = [
                  (pkgs.replaceVars ../nix/darwin-build-fixes.patch {
                    sdkVersion = pkgs.apple-sdk_26.version;
@@ -113,42 +112,42 @@ let
                MACOSX_DEPLOYMENT_TARGET = pkgs.apple-sdk_26.version;
              });
          } // lib.optionalAttrs isLinux {
-        mlx = prev.mlx.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
-          autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
-          postInstall = (old.postInstall or "") + ''
-            cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
-          '';
-        });
-      } // lib.optionalAttrs cudaSupport {
-        "${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ cudaLibs;
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ cudaLibs;
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ [ cudaLibs ];
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-        torch = prev.torch.overrideAttrs (old: {
-          buildInputs = old.buildInputs ++ cudaLibs;
-          autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-        });
-      };
+          mlx = prev.mlx.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
+            autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
+            postInstall = ''
+              cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
+            '';
+          });
+        } // lib.optionalAttrs cudaSupport {
+          "${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ [ cudaLibs ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+          torch = prev.torch.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+          });
+        };
      pyprojectOverlay = workspace.mkPyprojectOverlay {
        sourcePreference = "wheel";
-        dependencies = { exo = [ uv_extra ]; exo-bench = [ ]; };
+        dependencies = members;
      };
      editableOverlay = workspace.mkEditablePyprojectOverlay {
        # Use environment variable pointing to editable root directory
@@ -165,8 +164,8 @@ let
          buildSystemsOverlay
        ]
      );
-
-      mkApp = cmd: name: members: pkgs.writeShellApplication {
+      venv = name: (pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; });
+      mkApp = cmd: name: pkgs.writeShellApplication {
        inherit name;
        runtimeEnv = {
          EXO_DASHBOARD_DIR = self'.packages.dashboard;
@@ -174,17 +173,17 @@ let
        };
        runtimeInputs = [
          # mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
-          ((pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; }))
+          (venv name)
        ]
        ++ lib.optionals isDarwin [ pkgs.macmon ];
        text = "exec " + lib.optionalString cudaSupport "${lib.getExe pkgs.nix-gl-host} " + cmd;
      };
    in
    {
-      inherit pythonSet;
+      inherit venv;
      editablePythonSet = pythonSet.overrideScope editableOverlay;
-      mkPythonScript = members: name: path: mkApp ''python ${path} "$@"'' name members;
-      mkExo = name: members: mkApp ''exo "$@"'' name members;
+      mkPythonScript = path: mkApp ''python ${path} "$@"'';
+      mkExo = mkApp ''exo "$@"'';
    };
 in
 {
@@ -192,16 +191,21 @@ in
    { self', pkgs, unfreePkgs, lib, ... }:
    let
      inherit (pkgs.stdenv.hostPlatform) isLinux;
-      inherit (mkPythonSet { inherit self' pkgs lib; }) pythonSet editablePythonSet mkPythonScript mkExo;
-
-      exoVenv = pythonSet.mkVirtualEnv "exo-env" { exo = lib.optionals isLinux [ "cpu" ]; };
+      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "cpu" ]; }; }) editablePythonSet mkExo;

      # Virtual environment with dev dependencies for testing
-      testVenv = pythonSet.mkVirtualEnv "exo-test-env" {
-        exo = [ "dev" ] ++ lib.optionals isLinux [ "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
+      testVenv = (mkPythonSet {
+        inherit self' pkgs lib; members = {
+        exo = [ "dev" "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
      };
+      }).venv "exo-test";

-      mkBenchScript = mkPythonScript { exo-bench = [ ]; };
+      mkBenchScript = (mkPythonSet {
+        inherit self' pkgs lib; members = {
+        exo = [ "cpu" ];
+        exo-bench = [ ]; # Include pytest, pytest-asyncio, pytest-env
+      };
+      }).mkPythonScript;

      mkSimplePythonScript = name: path: pkgs.writeShellApplication {
        inherit name;
@@ -212,9 +216,7 @@ in
    in
    {
      packages = {
-        exo = mkExo "exo" { exo = lib.optionals isLinux [ "cpu" ]; };
-        # for devShell
-        exo-venv = exoVenv;
+        exo = mkExo "exo";
        editableVenv = editablePythonSet.mkVirtualEnv "exo-dev-env" { exo = [ "dev" ]; };
        # for running tests in ci
        exo-test-env = testVenv;
@@ -224,8 +226,8 @@ in
        # used by ./tests/run_exo_on.sh
        exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
      } // lib.optionalAttrs isLinux {
-        exo-cuda-12 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; }).mkExo "exo-cuda-12" { exo = [ "cuda12" ]; };
-        exo-cuda-13 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; }).mkExo "exo-cuda-13" { exo = [ "cuda13" ]; };
+        exo-cuda-12 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "cuda12" ]; }; }).mkExo "exo-cuda-12";
+        exo-cuda-13 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "cuda13" ]; }; }).mkExo "exo-cuda-13";
      };

      checks = {
@@ -235,19 +237,11 @@ in
          touch $out
        '';

-        typecheck = pkgs.runCommand "typecheck"
-          {
-            nativeBuildInputs = [
-              testVenv
-              pkgs.basedpyright
-            ];
-          }
-          ''
-            cd ${inputs.self}
-            export HOME=$TMPDIR
-            basedpyright --pythonpath ${testVenv}/bin/python --project ${inputs.self}/pyproject.toml
-            touch $out
-          '';
+        typecheck = pkgs.runCommand "typecheck" { nativeBuildInputs = [ testVenv ]; } ''
+          cd ${inputs.self}
+          basedpyright
+          touch $out
+        '';
      };
    };
 }
@@ -0,0 +1,21 @@
+model_id = "mlx-community/GLM-5.1-DQ4plus-q8"
+n_layers = 78
+hidden_size = 6144
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "glm"
+quantization = "8bit"
+base_model = "GLM-5.1"
+capabilities = ["text", "thinking"]
+
+context_length = 202752
+
+[storage_size]
+in_bytes = 465173655552
+
+# Source: https://huggingface.co/zai-org/GLM-5.1
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
@@ -0,0 +1,21 @@
+model_id = "mlx-community/GLM-5.1-MXFP4-Q8"
+n_layers = 78
+hidden_size = 6144
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "glm"
+quantization = "MXFP4-Q8"
+base_model = "GLM-5.1"
+capabilities = ["text", "thinking"]
+
+context_length = 202752
+
+[storage_size]
+in_bytes = 405480321024
+
+# Source: https://huggingface.co/zai-org/GLM-5.1
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
@@ -0,0 +1,21 @@
+model_id = "mlx-community/GLM-5.1"
+n_layers = 78
+hidden_size = 6144
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "glm"
+quantization = "bf16"
+base_model = "GLM-5.1"
+capabilities = ["text", "thinking"]
+
+context_length = 202752
+
+[storage_size]
+in_bytes = 1487822475264
+
+# Source: https://huggingface.co/zai-org/GLM-5.1
+# Source: https://docs.z.ai/api-reference/llm/chat-completion
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
@@ -0,0 +1,33 @@
+model_id = "mlx-community/Kimi-K2.6-mlx-DQ3_K_M-q8"
+n_layers = 61
+hidden_size = 7168
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "kimi"
+quantization = "3bit"
+base_model = "Kimi K2.6"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 470628683776
+
+[vision]
+image_token_id = 163605
+model_type = "kimi_vl"
+weights_repo = "exolabs/Kimi-K2.6-vision"
+processor_repo = "moonshotai/Kimi-K2.6"
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+min_p = 0.01
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults.non_thinking]
+temperature = 0.6
+top_p = 0.95
+min_p = 0.01
@@ -0,0 +1,35 @@
+model_id = "mlx-community/Qwen3.6-27B-4bit"
+n_layers = 64
+hidden_size = 5120
+num_key_value_heads = 4
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.6 27B"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 16054262240
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults.non_thinking]
+temperature = 0.7
+top_p = 0.8
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
@@ -0,0 +1,35 @@
+model_id = "mlx-community/Qwen3.6-27B-8bit"
+n_layers = 64
+hidden_size = 5120
+num_key_value_heads = 4
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.6 27B"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 29500938720
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults.non_thinking]
+temperature = 0.7
+top_p = 0.8
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
@@ -0,0 +1,35 @@
+model_id = "mlx-community/Qwen3.6-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.6 27B"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 54713457120
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
+
+# Source: https://huggingface.co/Qwen/Qwen3.6-27B#best-practices
+# Source: https://unsloth.ai/docs/models/qwen3.5
+[sampling_defaults.non_thinking]
+temperature = 0.7
+top_p = 0.8
+top_k = 20
+min_p = 0.0
+repetition_penalty = 1.0
+presence_penalty = 1.5
@@ -0,0 +1,33 @@
+model_id = "moonshotai/Kimi-K2.6"
+n_layers = 61
+hidden_size = 7168
+num_key_value_heads = 64
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "kimi"
+quantization = ""
+base_model = "Kimi K2.6"
+capabilities = ["text", "thinking", "thinking_toggle", "vision"]
+
+context_length = 262144
+
+[storage_size]
+in_bytes = 595148192736
+
+[vision]
+image_token_id = 163605
+model_type = "kimi_vl"
+weights_repo = "exolabs/Kimi-K2.6-vision"
+processor_repo = "moonshotai/Kimi-K2.6"
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults]
+temperature = 1.0
+top_p = 0.95
+min_p = 0.01
+
+# Source: https://huggingface.co/moonshotai/Kimi-K2.6
+[sampling_defaults.non_thinking]
+temperature = 0.6
+top_p = 0.95
+min_p = 0.01
@@ -113,6 +113,23 @@ def _extract_content(content: str | list[ResponseContentPart]) -> str:
    )


+def _append_tool_call(
+    chat_template_messages: list[dict[str, Any]], tool_call: dict[str, Any]
+) -> None:
+    if chat_template_messages:
+        prev = chat_template_messages[-1]
+        if prev.get("role") == "assistant" and isinstance(prev.get("content"), str):
+            existing: list[dict[str, Any]] | None = prev.get("tool_calls")
+            if existing is None:
+                prev["tool_calls"] = [tool_call]
+            else:
+                existing.append(tool_call)
+            return
+    chat_template_messages.append(
+        {"role": "assistant", "content": "", "tool_calls": [tool_call]}
+    )
+
+
 async def responses_request_to_text_generation(
    request: ResponsesRequest,
 ) -> TextGenerationTaskParams:
@@ -182,59 +199,44 @@ async def responses_request_to_text_generation(
                    | McpCallInputItem()
                    | CustomToolCallInputItem()
                ):
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": item.name,
-                                        "arguments": item.arguments,
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": item.name,
+                                "arguments": item.arguments,
+                            },
+                        },
                    )
                case (
                    LocalShellCallInputItem()
                    | ShellCallInputItem()
                    | ComputerCallInputItem()
                ):
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": item.type,
-                                        "arguments": json.dumps(item.action),
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": item.type,
+                                "arguments": json.dumps(item.action),
+                            },
+                        },
                    )
                case ApplyPatchCallInputItem():
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": "apply_patch",
-                                        "arguments": json.dumps({"patch": item.patch}),
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": "apply_patch",
+                                "arguments": json.dumps({"patch": item.patch}),
+                            },
+                        },
                    )
                case (
                    WebSearchCallInputItem()
@@ -254,21 +256,16 @@ async def responses_request_to_text_generation(
                        args = {"prompt": item.prompt}
                    else:
                        args = {"query": item.query}
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": item.type,
-                                        "arguments": json.dumps(args),
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": item.type,
+                                "arguments": json.dumps(args),
+                            },
+                        },
                    )
                case (
                    FunctionCallOutputInputItem()
@@ -320,21 +317,16 @@ async def responses_request_to_text_generation(
                            }
                        )
                case McpApprovalRequestInputItem():
-                    chat_template_messages.append(
+                    _append_tool_call(
+                        chat_template_messages,
                        {
-                            "role": "assistant",
-                            "content": "",
-                            "tool_calls": [
-                                {
-                                    "id": item.call_id,
-                                    "type": "function",
-                                    "function": {
-                                        "name": item.name,
-                                        "arguments": item.arguments,
-                                    },
-                                }
-                            ],
-                        }
+                            "id": item.call_id,
+                            "type": "function",
+                            "function": {
+                                "name": item.name,
+                                "arguments": item.arguments,
+                            },
+                        },
                    )
                case McpApprovalResponseInputItem():
                    chat_template_messages.append(
@@ -185,7 +185,10 @@ from exo.shared.types.tasks import (
 from exo.shared.types.tasks import (
    TextGeneration as TextGenerationTask,
 )
-from exo.shared.types.text_generation import Base64Image, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+    Base64ImageHash,
+    TextGenerationTaskParams,
+)
 from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
@@ -234,6 +237,7 @@ class API:
        self.node_id: NodeId = node_id
        self.last_completed_election: int = 0
        self.port = port
+        self._sent_image_hashes: set[str] = set()

        self.paused: bool = False
        self.paused_ev: anyio.Event = anyio.Event()
@@ -283,6 +287,7 @@ class API:
        self.event_receiver.close()
        self.event_receiver = event_receiver
        self._tg.start_soon(self._apply_state)
+        self._sent_image_hashes = set()

    def unpause(self, result_clock: int):
        logger.info("Unpausing API")
@@ -737,8 +742,6 @@ class API:
            "TODO: we should send a notification to the user to download the model"
        )

-    _sent_image_hashes: set[str] = set()
-
    async def _send_text_generation_with_images(
        self, task_params: TextGenerationTaskParams
    ) -> TextGeneration:
@@ -750,23 +753,19 @@ class API:
            return command

        hashes = [hashlib.sha256(img.encode("ascii")).hexdigest() for img in images]
+        all_hashes = {idx: Base64ImageHash(h) for idx, h in enumerate(hashes)}
+        task_params = task_params.model_copy(
+            update={"images": [], "image_hashes": all_hashes}
+        )
+        command = TextGeneration(task_params=task_params)

-        cached_hashes: dict[int, str] = {}
        new_images: list[tuple[int, str]] = []
        for idx, (img, h) in enumerate(zip(images, hashes, strict=True)):
-            if h in self._sent_image_hashes:
-                cached_hashes[idx] = h
-            else:
+            if h not in self._sent_image_hashes:
                self._sent_image_hashes.add(h)
                new_images.append((idx, img))

-        wrapped_hashes = {idx: Base64Image(h) for idx, h in cached_hashes.items()}
-
        if not new_images:
-            task_params = task_params.model_copy(
-                update={"images": [], "image_hashes": wrapped_hashes}
-            )
-            command = TextGeneration(task_params=task_params)
            await self._send(command)
            return command

@@ -775,16 +774,6 @@ class API:
            for i in range(0, len(img_data), EXO_MAX_CHUNK_SIZE):
                all_chunks.append((img_idx, img_data[i : i + EXO_MAX_CHUNK_SIZE]))

-        task_params = task_params.model_copy(
-            update={
-                "images": [],
-                "image_hashes": wrapped_hashes,
-                "total_input_chunks": len(all_chunks),
-                "image_count": len(new_images),
-            }
-        )
-        command = TextGeneration(task_params=task_params)
-
        for global_idx, (img_idx, chunk_data) in enumerate(all_chunks):
            await self._send(
                SendInputChunk(
@@ -88,7 +88,9 @@ class DownloadCoordinator:

        try:
            if progress.status == "complete":
-                found = await to_thread.run_sync(resolve_existing_model, model_id)
+                found = await to_thread.run_sync(
+                    resolve_existing_model, model_id, callback_shard.model_card
+                )
                if found is not None:
                    completed = self._completed_from_path(
                        callback_shard, found, progress.total
@@ -193,7 +195,9 @@ class DownloadCoordinator:
                return

        # Check all model directories for pre-existing complete models
-        found_path = await to_thread.run_sync(resolve_existing_model, model_id)
+        found_path = await to_thread.run_sync(
+            resolve_existing_model, model_id, shard.model_card
+        )
        if found_path is not None:
            logger.info(f"DownloadCoordinator: Model {model_id} found at {found_path}")
            completed = self._completed_from_path(
@@ -220,7 +224,9 @@ class DownloadCoordinator:
        )

        if initial_progress.status == "complete":
-            found = await to_thread.run_sync(resolve_existing_model, model_id)
+            found = await to_thread.run_sync(
+                resolve_existing_model, model_id, shard.model_card
+            )
            if found is not None:
                completed = self._completed_from_path(
                    shard, found, initial_progress.total
@@ -351,7 +357,9 @@ class DownloadCoordinator:

                    if progress.status == "complete":
                        found = await to_thread.run_sync(
-                            resolve_existing_model, model_id
+                            resolve_existing_model,
+                            model_id,
+                            progress.shard.model_card,
                        )
                        if found is not None:
                            status: DownloadProgress = self._completed_from_path(
@@ -380,7 +388,9 @@ class DownloadCoordinator:
                        # (is_model_directory_complete) which validates that all
                        # safetensors weight files are present.
                        found = await to_thread.run_sync(
-                            resolve_existing_model, model_id
+                            resolve_existing_model,
+                            model_id,
+                            progress.shard.model_card,
                        )
                        if found is not None:
                            status = self._completed_from_path(
@@ -421,7 +431,9 @@ class DownloadCoordinator:
                            (DownloadCompleted, DownloadOngoing, DownloadFailed),
                        ):
                            continue
-                        found = await to_thread.run_sync(resolve_existing_model, mid)
+                        found = await to_thread.run_sync(
+                            resolve_existing_model, mid, card
+                        )
                        if found is not None and is_read_only_model_dir(found):
                            path_shard = PipelineShardMetadata(
                                model_card=card,
@@ -35,7 +35,7 @@ from exo.shared.constants import (
    EXO_MODELS_DIRS,
    EXO_MODELS_READ_ONLY_DIRS,
 )
-from exo.shared.models.model_cards import ModelTask
+from exo.shared.models.model_cards import ModelCard, ModelTask
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.downloads import (
@@ -118,7 +118,9 @@ class InsufficientDiskSpaceError(Exception):
    """Raised when no writable model directory has enough free space."""


-def resolve_existing_model(model_id: ModelId) -> Path | None:
+def resolve_existing_model(
+    model_id: ModelId, card: ModelCard | None = None
+) -> Path | None:
    """Search all model directories for a complete, pre-existing model.

    Checks read-only directories first, then writable directories.
@@ -128,7 +130,7 @@ def resolve_existing_model(model_id: ModelId) -> Path | None:
    normalized = model_id.normalize()
    for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
        candidate = search_dir / normalized
-        if candidate.is_dir() and is_model_directory_complete(candidate):
+        if candidate.is_dir() and is_model_directory_complete(candidate, card):
            return candidate
    return None

@@ -165,6 +167,29 @@ def select_download_dir(required_bytes: int) -> Path:
    )


+async def select_download_dir_for_shard(
+    model_id: ModelId,
+    filtered_file_list: list[FileListEntry],
+    total_size: int,
+) -> Path:
+    for candidate_dir in EXO_MODELS_DIRS:
+        if not candidate_dir.exists():
+            continue
+        sub = candidate_dir / model_id.normalize()
+        if not await aios.path.isdir(sub):
+            continue
+        existing_bytes = 0
+        for file_entry in filtered_file_list:
+            existing_bytes += await get_downloaded_size(sub / file_entry.path)
+        remaining = max(total_size - existing_bytes, 0)
+        try:
+            if shutil.disk_usage(candidate_dir).free >= remaining:
+                return candidate_dir
+        except OSError:
+            continue
+    return select_download_dir(total_size)
+
+
 async def resolve_model_dir(model_id: ModelId) -> Path:
    """Return the directory for a model's files, creating it if needed.

@@ -279,10 +304,26 @@ def _scan_model_directory(
    return list(entries_by_path.values())


-def is_model_directory_complete(model_dir: Path) -> bool:
-    """Check if a model directory contains all required weight files."""
+def is_model_directory_complete(model_dir: Path, card: ModelCard | None = None) -> bool:
+    """Check if a model directory contains all required weight files.
+    Also checks for sibling weights repo.
+    """
    file_list = _scan_model_directory(model_dir, recursive=True)
-    return file_list is not None and all(f.size is not None for f in file_list)
+    if file_list is None or not all(f.size is not None for f in file_list):
+        return False
+    if (
+        card is not None
+        and card.vision is not None
+        and card.vision.weights_repo != str(card.model_id)
+    ):
+        vision_id = ModelId(card.vision.weights_repo)
+        normalized = vision_id.normalize()
+        for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
+            candidate = search_dir / normalized
+            if candidate.is_dir() and is_model_directory_complete(candidate):
+                return True
+        return False
+    return True


 async def _build_file_list_from_local_directory(
@@ -834,7 +875,9 @@ async def download_shard(
            else EXO_DEFAULT_MODELS_DIR / model_id.normalize()
        )
    else:
-        models_dir = select_download_dir(total_size)
+        models_dir = await select_download_dir_for_shard(
+            model_id, filtered_file_list, total_size
+        )
        target_dir = models_dir / model_id.normalize()
        await aios.makedirs(target_dir, exist_ok=True)
    file_progress: dict[str, RepoFileDownloadProgress] = {}
@@ -117,40 +117,39 @@ class ResumableShardDownloader(ShardDownloader):
    ) -> Path:
        allow_patterns = ["config.json"] if config_only else None

+        has_vision_sibling = (
+            not config_only
+            and not self.offline
+            and shard.model_card.vision is not None
+            and shard.model_card.vision.weights_repo != str(shard.model_card.model_id)
+        )
+
+        async def main_progress(
+            cb_shard: ShardMetadata, progress: RepoDownloadProgress
+        ) -> None:
+            if has_vision_sibling and progress.status == "complete":
+                return
+            await self.on_progress_wrapper(cb_shard, progress)
+
        target_dir, _ = await download_shard(
            shard,
-            self.on_progress_wrapper,
+            main_progress,
            max_parallel_downloads=self.max_parallel_downloads,
            allow_patterns=allow_patterns,
            skip_internet=self.offline,
        )

-        if (
-            not config_only
-            and not self.offline
-            and shard.model_card.vision
-            and shard.model_card.vision.weights_repo != str(shard.model_card.model_id)
-        ):
-            vision_repo = shard.model_card.vision.weights_repo
-            vision_card = ModelCard(
-                model_id=ModelId(vision_repo),
-                storage_size=Memory.from_bytes(0),
-                n_layers=1,
-                hidden_size=1,
-                supports_tensor=False,
-                tasks=[ModelTask.TextGeneration],
-            )
-            vision_shard = PipelineShardMetadata(
-                model_card=vision_card,
-                device_rank=0,
-                world_size=1,
-                start_layer=0,
-                end_layer=1,
-                n_layers=1,
-            )
+        if has_vision_sibling:
+            vision_shard = self._build_vision_shard(shard)
+
+            async def vision_progress(
+                _cb_shard: ShardMetadata, progress: RepoDownloadProgress
+            ) -> None:
+                await self.on_progress_wrapper(shard, progress)
+
            await download_shard(
                vision_shard,
-                self.on_progress_wrapper,
+                vision_progress,
                max_parallel_downloads=self.max_parallel_downloads,
                allow_patterns=["*.safetensors", "config.json"],
                skip_internet=self.offline,
@@ -158,6 +157,87 @@ class ResumableShardDownloader(ShardDownloader):

        return target_dir

+    async def _status_for_shard(
+        self, shard: ShardMetadata
+    ) -> tuple[Path, RepoDownloadProgress]:
+        async def _noop(
+            _cb_shard: ShardMetadata, _progress: RepoDownloadProgress
+        ) -> None:
+            return
+
+        path, main_progress = await download_shard(
+            shard,
+            _noop,
+            skip_download=True,
+            skip_internet=self.offline,
+        )
+
+        has_vision_sibling = (
+            shard.model_card.vision is not None
+            and shard.model_card.vision.weights_repo != str(shard.model_card.model_id)
+        )
+        if not has_vision_sibling:
+            return path, main_progress
+
+        vision_shard = self._build_vision_shard(shard)
+        _, vision_progress = await download_shard(
+            vision_shard,
+            _noop,
+            skip_download=True,
+            skip_internet=self.offline,
+        )
+        combined = self._combine_progress(shard, main_progress, vision_progress)
+        return path, combined
+
+    @staticmethod
+    def _build_vision_shard(shard: ShardMetadata) -> PipelineShardMetadata:
+        assert shard.model_card.vision is not None
+        vision_card = ModelCard(
+            model_id=ModelId(shard.model_card.vision.weights_repo),
+            storage_size=Memory.from_bytes(0),
+            n_layers=1,
+            hidden_size=1,
+            supports_tensor=False,
+            tasks=[ModelTask.TextGeneration],
+        )
+        return PipelineShardMetadata(
+            model_card=vision_card,
+            device_rank=0,
+            world_size=1,
+            start_layer=0,
+            end_layer=1,
+            n_layers=1,
+        )
+
+    @staticmethod
+    def _combine_progress(
+        shard: ShardMetadata,
+        main: RepoDownloadProgress,
+        vision: RepoDownloadProgress,
+    ) -> RepoDownloadProgress:
+        status_rank = {"not_started": 0, "in_progress": 1, "complete": 2}
+        combined_status = min(
+            (main.status, vision.status), key=lambda s: status_rank[s]
+        )
+        file_progress = dict(main.file_progress)
+        for file_path, fp in vision.file_progress.items():
+            file_progress[f"{vision.repo_id}/{file_path}"] = fp
+        return RepoDownloadProgress(
+            repo_id=main.repo_id,
+            repo_revision=main.repo_revision,
+            shard=shard,
+            completed_files=main.completed_files + vision.completed_files,
+            total_files=main.total_files + vision.total_files,
+            downloaded=main.downloaded + vision.downloaded,
+            downloaded_this_session=main.downloaded_this_session
+            + vision.downloaded_this_session,
+            total=main.total + vision.total,
+            overall_speed=main.overall_speed + vision.overall_speed,
+            overall_eta=max(main.overall_eta, vision.overall_eta),
+            status=combined_status,
+            file_progress=file_progress,
+        )
+
    async def get_shard_download_status(
        self,
    ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
@@ -166,12 +246,7 @@ class ResumableShardDownloader(ShardDownloader):
        ) -> tuple[Path, RepoDownloadProgress]:
            """Helper coroutine that builds the shard for a model and gets its download status."""
            shard = await build_full_shard(model_id)
-            return await download_shard(
-                shard,
-                self.on_progress_wrapper,
-                skip_download=True,
-                skip_internet=self.offline,
-            )
+            return await self._status_for_shard(shard)

        semaphore = asyncio.Semaphore(self.max_parallel_downloads)

@@ -195,10 +270,5 @@ class ResumableShardDownloader(ShardDownloader):
    async def get_shard_download_status_for_shard(
        self, shard: ShardMetadata
    ) -> RepoDownloadProgress:
-        _, progress = await download_shard(
-            shard,
-            self.on_progress_wrapper,
-            skip_download=True,
-            skip_internet=self.offline,
-        )
+        _, progress = await self._status_for_shard(shard)
        return progress
@@ -410,8 +410,6 @@ class Master:
                        continue

                    logger.debug(f"Master indexing event: {str(event)[:100]}")
-                    indexed = IndexedEvent(event=event, idx=len(self._event_log))
-                    self.state = apply(self.state, indexed)

                    event = event.model_copy(
                        update={"_master_time_stamp": datetime.now(tz=timezone.utc)}
@@ -421,6 +419,9 @@ class Master:
                            update={"when": str(datetime.now(tz=timezone.utc))}
                        )

+                    indexed = IndexedEvent(event=event, idx=len(self._event_log))
+                    self.state = apply(self.state, indexed)
+
                    self._event_log.append(event)
                    await self._send_event(indexed)

@@ -68,6 +68,7 @@ DASHBOARD_DIR = (
 # Log files (data/logs or cache)
 EXO_LOG_DIR = EXO_CACHE_HOME / "exo_log"
 EXO_LOG = EXO_LOG_DIR / "exo.log"
+EXO_TEST_LOG = EXO_CACHE_HOME / "exo_test.log"

 # Identity (config)
 EXO_NODE_ID_KEYPAIR = EXO_CONFIG_HOME / "node_id.keypair"
@@ -1,158 +0,0 @@
-from pathlib import Path
-from collections.abc import Sequence
-import tomlkit
-from exo.utils.pydantic_ext import FrozenModel
-from typing import Self, Any
-from pydantic import Field, BaseModel, model_validator
-from pydantic_settings import BaseSettings, SettingsConfigDict, PydanticBaseSettingsSource, TomlConfigSettingsSource
-from exo.shared.types.common import NodeId, ModelId
-from exo.shared.types.worker.instances import InstanceId
-from exo.shared.constants import EXO_CONFIG_HOME, EXO_DATA_HOME, EXO_CACHE_HOME
-from exo.utils.dashboard_path import find_dashboard, find_resources
-
-
-def default_merge[T: BaseModel](left: T, right: T) -> T:
-    if left == right:
-        return left
-    merged_dict = {}
-    for key in type(left).model_fields:
-        try:
-            merged_dict[key] = getattr(left, key).merge(  # pyright: ignore[reportAny]
-                getattr(right, key, None)
-            )
-        except AttributeError:
-            raise NotImplementedError("Cluster Option using default implementation incorrectly")
-
-    return type(left).model_validate(merged_dict)
-
-
-def _parse_colon_separated_dirs(obj: Any) -> set[Path]:  # pyright: ignore[reportAny]
-        if isinstance(obj, (list, set)):
-            return set(Path(d).expanduser() for d in obj)  # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType]
-        else:
-            return set(Path(d).expanduser() for d in str(obj).split(":"))  # pyright: ignore[reportAny]
-
-class ModelDirsSettings(BaseModel, frozen=True):
-    # env: EXO_MODEL_DIRS_DEFAULT prepends to WRITEABLE, defaults to EXO_DATA_HOME/models
-    # env: EXO_MODEL_DIRS_WRITEABLE, defaults to []
-    writeable: list[Path] = []
-    # env: EXO_MODEL_DIRS_READONLY, defaults to []
-    readonly: list[Path] = []
-
-    @model_validator(mode="before")
-    @classmethod
-    def build_defaults(cls, data: Any) -> Any:  # pyright: ignore[reportAny]
-        if not isinstance(data, dict):
-            return data  # pyright: ignore[reportAny]
-        default = Path(data.get("default", EXO_DATA_HOME / "models")).expanduser()  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
-        readonly = _parse_colon_separated_dirs(data.get("readonly", []))  # pyright: ignore[reportUnknownMemberType]
-        writeable = _parse_colon_separated_dirs(data.get("writeable", [])).difference(readonly)  # pyright: ignore[reportUnknownMemberType]
-        if default not in readonly:
-            writeable = [default, *writeable]
-        return {**data, "writeable": writeable, "readonly": readonly}  # pyright: ignore[reportUnknownVariableType]
-
-class RuntimeDirsSettings(BaseModel, frozen=True):
-    dashboard: Path = Field(default_factory=find_dashboard)
-    resources: Path = Field(default_factory=find_resources)
-    logs: Path = EXO_CACHE_HOME / "log"
-    log_file: str = "latest.log"
-
-    def log_file_path(self):
-        return self.logs / self.log_file
-
-# doesnt require merge
-class LocalSettings(FrozenModel):
-    runtime_dirs: RuntimeDirsSettings
-    model_dirs: ModelDirsSettings
-
-class InstanceSettings(FrozenModel):
-    # env: EXO_INSTANCE_DEFAULTS_BATCH_CONCURRENCY
-    batch_concurrency: int
-
-    def merge(self, other: Self) -> Self:
-        return type(self)(batch_concurrency=min(self.batch_concurrency, other.batch_concurrency))
-
-class ClusterSettings(FrozenModel):
-    instance_defaults: InstanceSettings = InstanceSettings(batch_concurrency=8)
-    model_settings_overrides: dict[ModelId, InstanceSettings] = {}
-
-    def merge(self, other: Self) -> Self:
-        return default_merge(self, other)
-
-class SettingsFile(BaseSettings):
-    model_config = SettingsConfigDict(
-        extra='ignore',
-        frozen=True,
-        toml_file=EXO_CONFIG_HOME / "config.toml",
-        env_prefix="EXO_",
-        env_nested_delimiter="_",
-        env_ignore_empty=True,
-    )
-
-    model_dirs: ModelDirsSettings
-    runtime_dirs: RuntimeDirsSettings
-    model_settings_overrides: dict[ModelId, InstanceSettings] = {}
-    instance_defaults: InstanceSettings
-
-    def get_local(self) -> LocalSettings:
-        ...
-    def get_cluster(self) -> ClusterSettings:
-        ...
-
-    @classmethod
-    def settings_customise_sources(
-        cls,
-        settings_cls: type[BaseSettings],
-        init_settings: PydanticBaseSettingsSource,
-        env_settings: PydanticBaseSettingsSource,
-        dotenv_settings: PydanticBaseSettingsSource,
-        file_secret_settings: PydanticBaseSettingsSource,
-    ) -> tuple[PydanticBaseSettingsSource, ...]:
-        return (init_settings, env_settings, TomlConfigSettingsSource(settings_cls),)
-
-    def sync(self):
-        """nb: only call this once per save"""
-        cfg_path = type(self).model_config.get("toml_file", None)
-        if isinstance(cfg_path, Sequence):
-            cfg_path=cfg_path[0]
-        if cfg_path:
-            with open(cfg_path, "w") as fp:
-                tomlkit.dump(self.model_dump(exclude_defaults=True), fp) # pyright: ignore[reportUnknownMemberType]
-
-
-
-class StateSettings(FrozenModel):
-    per_node: dict[NodeId, LocalSettings]
-    per_instance: dict[InstanceId, InstanceSettings]
-    cluster: ClusterSettings
-    
-    def model_merge_local(self, node_id: NodeId, settings: LocalSettings) -> Self:
-        return self.model_copy(update={
-            "per_node": {
-                **self.per_node,
-                node_id: settings
-            }
-        })
-
-    def model_merge_cluster(self, settings: ClusterSettings) -> Self:
-        return self.model_copy(update={
-            "cluster": self.cluster.merge(settings)
-        })
-
-    def settings_for(self, node_id: NodeId) -> StoredSettings:
-        merged = {}
-        for key, val in self.cluster.model_dump(exclude_defaults=True).items():  # pyright: ignore[reportAny]
-            merged[key] = val
-
-        if (local := self.per_node.get(node_id, None)) is not None:
-            for key, val in local.model_dump(exclude_defaults=True).items():  # pyright: ignore[reportAny]
-                merged[key] = val
-
-        return StoredSettings.model_validate(merged)
-
-    def sync(self, node_id: NodeId):
-        """nb: only call this once per save"""
-        toml_file=EXO_CONFIG_HOME / "config.toml"
-        with open(toml_file, "w") as fp:
-            tomlkit.dump(self.settings_for(node_id).model_dump(exclude_defaults=True), fp) # pyright: ignore[reportUnknownMemberType]
-
@@ -85,6 +85,6 @@ class PrefillProgressChunk(BaseChunk):
    total_tokens: int


-GenerationChunk = (
-    TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk | PrefillProgressChunk
-)
+StatusChunk = PrefillProgressChunk
+GenerationChunk = TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk
+Chunk = StatusChunk | GenerationChunk
@@ -5,7 +5,7 @@ from pydantic import Field

 from exo.shared.models.model_cards import ModelCard
 from exo.shared.topology import Connection
-from exo.shared.types.chunks import GenerationChunk, InputImageChunk
+from exo.shared.types.chunks import Chunk, InputImageChunk
 from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId, SystemId
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
@@ -91,7 +91,7 @@ class NodeDownloadProgress(BaseEvent):

 class ChunkGenerated(BaseEvent):
    command_id: CommandId
-    chunk: GenerationChunk
+    chunk: Chunk


 class InputChunkReceived(BaseEvent):
@@ -101,3 +101,6 @@ Task = (
    | ImageEdits
    | Shutdown
 )
+TextTask = TextGeneration
+ImageTask = ImageGeneration | ImageEdits
+GenerationTask = TextTask | ImageTask
@@ -114,8 +114,6 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
    frequency_penalty: float | None = None
    images: list[Base64Image] = Field(default_factory=list)
    image_hashes: dict[int, Base64ImageHash] = Field(default_factory=dict)
-    total_input_chunks: int = 0
-    image_count: int = 0

    def with_card_sampling_defaults(self) -> "TextGenerationTaskParams":
        from exo.shared.models.model_cards import get_card
@@ -16,10 +16,6 @@ class BaseRunnerResponse(TaggedModel):
    pass


-class TokenizedResponse(BaseRunnerResponse):
-    prompt_tokens: int
-
-
 class GenerationResponse(BaseRunnerResponse):
    text: str
    token: int
@@ -70,6 +66,15 @@ class FinishedResponse(BaseRunnerResponse):
    pass


+class ModelLoadingResponse(BaseRunnerResponse):
+    layers_loaded: int
+    total: int
+
+
+class CancelledResponse(BaseRunnerResponse):
+    pass
+
+
 class PrefillProgressResponse(BaseRunnerResponse):
    processed_tokens: int
    total_tokens: int
@@ -1,10 +1,8 @@
 import sys
 from pathlib import Path
 from typing import cast
-from functools import cache


-@cache
 def find_resources() -> Path:
    resources = _find_resources_in_repo() or _find_resources_in_bundle()
    if resources is None:
@@ -33,7 +31,6 @@ def _find_resources_in_bundle() -> Path | None:
    return None


-@cache
 def find_dashboard() -> Path:
    dashboard = _find_dashboard_in_repo() or _find_dashboard_in_bundle()
    if not dashboard:
@@ -0,0 +1,55 @@
+from abc import ABC, abstractmethod
+from collections.abc import Generator, Iterable
+
+from exo.shared.types.chunks import Chunk
+from exo.shared.types.tasks import CANCEL_ALL_TASKS, GenerationTask, TaskId
+from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.runner_response import (
+    CancelledResponse,
+    FinishedResponse,
+    ModelLoadingResponse,
+)
+
+
+class Engine(ABC):
+    _cancelled_tasks: set[TaskId]
+
+    def should_cancel(self, task_id: TaskId) -> bool:
+        return (
+            task_id in self._cancelled_tasks
+            or CANCEL_ALL_TASKS in self._cancelled_tasks
+        )
+
+    @abstractmethod
+    def warmup(self) -> None: ...
+
+    @abstractmethod
+    def submit(
+        self,
+        task: GenerationTask,
+    ) -> None: ...
+
+    @abstractmethod
+    def step(
+        self,
+    ) -> Iterable[tuple[TaskId, Chunk | CancelledResponse | FinishedResponse]]: ...
+
+    @abstractmethod
+    def close(self) -> None: ...
+
+
+class Builder(ABC):
+    @abstractmethod
+    def connect(self, bound_instance: BoundInstance) -> None: ...
+
+    @abstractmethod
+    def load(
+        self,
+        bound_instance: BoundInstance,
+    ) -> Generator[ModelLoadingResponse]: ...
+
+    @abstractmethod
+    def build(self) -> Engine: ...
+
+    @abstractmethod
+    def close(self) -> None: ...
@@ -1,12 +1,16 @@
+from exo.worker.engines.image.builder import (
+    ImageEngine,
+    MfluxBuilder,
+)
 from exo.worker.engines.image.distributed_model import (
    DistributedImageModel,
-    initialize_image_model,
 )
 from exo.worker.engines.image.generate import generate_image, warmup_image_generator

 __all__ = [
+    "MfluxBuilder",
+    "ImageEngine",
    "DistributedImageModel",
    "generate_image",
-    "initialize_image_model",
    "warmup_image_generator",
 ]
@@ -0,0 +1,212 @@
+import contextlib
+from collections import deque
+from collections.abc import Generator, Iterable
+from dataclasses import dataclass, field
+
+import mlx.core as mx
+from loguru import logger
+
+from exo.api.types import ImageEditsTaskParams, ImageGenerationTaskParams
+from exo.shared.constants import EXO_TRACING_ENABLED
+from exo.shared.tracing import clear_trace_buffer, get_trace_buffer
+from exo.shared.types.chunks import Chunk, ErrorChunk
+from exo.shared.types.events import (
+    Event,
+    TraceEventData,
+    TracesCollected,
+)
+from exo.shared.types.tasks import (
+    GenerationTask,
+    ImageEdits,
+    ImageGeneration,
+    ImageTask,
+    TaskId,
+)
+from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.runner_response import (
+    CancelledResponse,
+    FinishedResponse,
+    ModelLoadingResponse,
+)
+from exo.shared.types.worker.shards import (
+    CfgShardMetadata,
+    PipelineShardMetadata,
+    ShardMetadata,
+)
+from exo.utils.channels import MpReceiver, MpSender
+from exo.worker.engines.base import Builder, Engine
+from exo.worker.engines.image.distributed_model import (
+    DistributedImageModel,
+)
+from exo.worker.engines.image.generate import (
+    generate_image,
+    warmup_image_generator,
+)
+from exo.worker.engines.mlx.utils_mlx import (
+    initialize_mlx,
+)
+
+
+def _is_primary_output_node(shard_metadata: ShardMetadata) -> bool:
+    """Check if this node is the primary output node for image generation.
+
+    For CFG models: the last pipeline stage in CFG group 0 (positive prompt).
+    For non-CFG models: the last pipeline stage.
+    """
+    if isinstance(shard_metadata, CfgShardMetadata):
+        is_pipeline_last = (
+            shard_metadata.pipeline_rank == shard_metadata.pipeline_world_size - 1
+        )
+        return is_pipeline_last and shard_metadata.cfg_rank == 0
+    elif isinstance(shard_metadata, PipelineShardMetadata):
+        return shard_metadata.device_rank == shard_metadata.world_size - 1
+    return False
+
+
+def _send_traces_if_enabled(
+    event_sender: MpSender[Event],
+    task_id: TaskId,
+    rank: int,
+) -> None:
+    if not EXO_TRACING_ENABLED:
+        return
+
+    traces = get_trace_buffer()
+    if traces:
+        trace_data = [
+            TraceEventData(
+                name=t.name,
+                start_us=t.start_us,
+                duration_us=t.duration_us,
+                rank=t.rank,
+                category=t.category,
+            )
+            for t in traces
+        ]
+        event_sender.send(
+            TracesCollected(
+                task_id=task_id,
+                rank=rank,
+                traces=trace_data,
+            )
+        )
+    clear_trace_buffer()
+
+
+@dataclass
+class MfluxBuilder(Builder):
+    event_sender: MpSender[Event]
+    cancel_receiver: MpReceiver[TaskId]
+    shard_metadata: ShardMetadata | None = None
+    image_model: DistributedImageModel | None = None
+    group: mx.distributed.Group | None = None
+
+    def connect(self, bound_instance: BoundInstance) -> None:
+        self.group = initialize_mlx(bound_instance)
+
+    def load(self, bound_instance: BoundInstance) -> Generator[ModelLoadingResponse]:
+        self.shard_metadata = bound_instance.bound_shard
+        self.image_model = DistributedImageModel.from_shard_metadata(
+            bound_instance.bound_shard, self.group
+        )
+        return
+        # very important!
+        yield
+
+    def close(self) -> None:
+        with contextlib.suppress(NameError, AttributeError):
+            del self.image_model, self.group
+
+    def build(
+        self,
+    ) -> Engine:
+        assert self.image_model
+        assert self.shard_metadata
+
+        return ImageEngine(
+            self.image_model,
+            self.shard_metadata,
+            self.event_sender,
+            self.cancel_receiver,
+        )
+
+
+@dataclass
+class ImageEngine(Engine):
+    image_model: DistributedImageModel
+    shard_metadata: ShardMetadata
+    event_sender: MpSender[Event]
+    cancel_receiver: MpReceiver[TaskId]
+    current_gen: Generator[tuple[TaskId, Chunk]] | None = field(
+        init=False, default=None
+    )
+    queue: deque[ImageTask] = field(init=False, default_factory=deque)
+
+    def warmup(self) -> None:
+        image = warmup_image_generator(model=self.image_model)
+        if image is not None:
+            logger.info(f"warmed up by generating {image.size} image")
+        else:
+            logger.info("warmup completed (non-primary node)")
+
+    def submit(
+        self,
+        task: GenerationTask,
+    ) -> None:
+        assert isinstance(task, (ImageGeneration, ImageEdits))
+        self.queue.append(task)
+
+    def step(
+        self,
+    ) -> Iterable[tuple[TaskId, Chunk | CancelledResponse | FinishedResponse]]:
+        resp = None
+        if self.current_gen is not None:
+            resp = next(self.current_gen, None)
+        if resp is None and len(self.queue) > 0:
+            task = self.queue.popleft()
+            self.current_gen = self._run_image_task(task.task_id, task.task_params)
+            resp = next(self.current_gen, None)
+        return (resp,) if resp is not None else ()
+
+    def close(self) -> None:
+        with contextlib.suppress(NameError, AttributeError):
+            del self.image_model
+
+    def _run_image_task(
+        self,
+        task_id: TaskId,
+        task_params: ImageGenerationTaskParams | ImageEditsTaskParams,
+    ) -> Generator[tuple[TaskId, Chunk]]:
+        assert self.image_model
+        logger.info(f"received image task: {str(task_params)[:500]}")
+
+        def cancel_checker() -> bool:
+            for cancel_id in self.cancel_receiver.collect():
+                self._cancelled_tasks.add(cancel_id)
+            return self.should_cancel(task_id)
+
+        try:
+            for response in generate_image(
+                model=self.image_model,
+                task=task_params,
+                cancel_checker=cancel_checker,
+            ):
+                if _is_primary_output_node(self.shard_metadata):
+                    yield (task_id, response)
+        except Exception as e:
+            if _is_primary_output_node(self.shard_metadata):
+                yield (
+                    task_id,
+                    ErrorChunk(
+                        model=self.shard_metadata.model_card.model_id,
+                        finish_reason="error",
+                        error_message=str(e),
+                    ),
+                )
+            raise
+        finally:
+            _send_traces_if_enabled(
+                self.event_sender, task_id, self.shard_metadata.device_rank
+            )
+
+        return
@@ -1,6 +1,6 @@
 from collections.abc import Callable, Generator
 from pathlib import Path
-from typing import Any, Literal, Optional
+from typing import Any, Literal

 import mlx.core as mx
 from mflux.models.common.config.config import Config
@@ -8,8 +8,12 @@ from PIL import Image

 from exo.api.types import AdvancedImageParams
 from exo.download.download_utils import build_model_path
-from exo.shared.types.worker.instances import BoundInstance
-from exo.shared.types.worker.shards import CfgShardMetadata, PipelineShardMetadata
+from exo.shared.types.common import ModelId
+from exo.shared.types.worker.shards import (
+    CfgShardMetadata,
+    PipelineShardMetadata,
+    ShardMetadata,
+)
 from exo.worker.engines.image.config import ImageModelConfig
 from exo.worker.engines.image.models import (
    create_adapter_for_model,
@@ -17,21 +21,22 @@ from exo.worker.engines.image.models import (
 )
 from exo.worker.engines.image.models.base import ModelAdapter
 from exo.worker.engines.image.pipeline import DiffusionRunner
-from exo.worker.engines.mlx.utils_mlx import mlx_distributed_init, mx_barrier
+from exo.worker.engines.mlx.utils_mlx import mx_barrier
 from exo.worker.runner.bootstrap import logger


 class DistributedImageModel:
+    model_id: ModelId
    _config: ImageModelConfig
    _adapter: ModelAdapter[Any, Any]
    _runner: DiffusionRunner

    def __init__(
        self,
-        model_id: str,
+        model_id: ModelId,
        local_path: Path,
        shard_metadata: PipelineShardMetadata | CfgShardMetadata,
-        group: Optional[mx.distributed.Group] = None,
+        group: mx.distributed.Group | None,
        quantize: int | None = None,
    ):
        config = get_config_for_model(model_id)
@@ -68,37 +73,27 @@ class DistributedImageModel:
        else:
            logger.info("Single-node initialization")

+        self.model_id = model_id
        self._config = config
        self._adapter = adapter
        self._runner = runner

    @classmethod
-    def from_bound_instance(
-        cls, bound_instance: BoundInstance
+    def from_shard_metadata(
+        cls, shard: ShardMetadata, group: mx.distributed.Group | None
    ) -> "DistributedImageModel":
-        model_id = bound_instance.bound_shard.model_card.model_id
+        model_id = shard.model_card.model_id
        model_path = build_model_path(model_id)

-        shard_metadata = bound_instance.bound_shard
-        if not isinstance(shard_metadata, (PipelineShardMetadata, CfgShardMetadata)):
+        if not isinstance(shard, (PipelineShardMetadata, CfgShardMetadata)):
            raise ValueError(
                "Expected PipelineShardMetadata or CfgShardMetadata for image generation"
            )

-        is_distributed = (
-            len(bound_instance.instance.shard_assignments.node_to_runner) > 1
-        )
-
-        if is_distributed:
-            logger.info("Starting distributed init for image model")
-            group = mlx_distributed_init(bound_instance)
-        else:
-            group = None
-
        return cls(
            model_id=model_id,
            local_path=model_path,
-            shard_metadata=shard_metadata,
+            shard_metadata=shard,
            group=group,
        )

@@ -173,7 +168,3 @@ class DistributedImageModel:
            else:
                logger.info("generated image")
                yield result
-
-
-def initialize_image_model(bound_instance: BoundInstance) -> DistributedImageModel:
-    return DistributedImageModel.from_bound_instance(bound_instance)
@@ -3,7 +3,7 @@ import io
 import random
 import tempfile
 import time
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from pathlib import Path
 from typing import Generator, Literal

@@ -17,11 +17,10 @@ from exo.api.types import (
    ImageGenerationTaskParams,
    ImageSize,
 )
+from exo.shared.constants import EXO_MAX_CHUNK_SIZE
+from exo.shared.types.chunks import ImageChunk
+from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.worker.runner_response import (
-    ImageGenerationResponse,
-    PartialImageResponse,
-)
 from exo.worker.engines.image.distributed_model import DistributedImageModel


@@ -71,16 +70,8 @@ def generate_image(
    model: DistributedImageModel,
    task: ImageGenerationTaskParams | ImageEditsTaskParams,
    cancel_checker: Callable[[], bool] | None = None,
-) -> Generator[ImageGenerationResponse | PartialImageResponse, None, None]:
-    """Generate image(s), optionally yielding partial results.
-
-    When partial_images > 0 or stream=True, yields PartialImageResponse for
-    intermediate images, then ImageGenerationResponse for the final image.
-
-    Yields:
-        PartialImageResponse for intermediate images (if partial_images > 0, first image only)
-        ImageGenerationResponse for final complete images
-    """
+) -> Generator[ImageChunk, None, None]:
+    """Generate image(s), optionally yielding partial results."""
    width, height = parse_size(task.size)
    quality: Literal["low", "medium", "high"] = task.quality or "medium"

@@ -142,12 +133,14 @@ def generate_image(
                        image = image.convert("RGB")
                    image.save(buffer, format=image_format)

-                    yield PartialImageResponse(
+                    yield from _process_image_response(
                        image_data=buffer.getvalue(),
-                        format=task.output_format,
+                        image_format=task.output_format,
                        partial_index=partial_idx,
                        total_partials=total_partials,
                        image_index=image_num,
+                        model_id=model.model_id,
+                        stats=None,
                    )
                else:
                    image = result
@@ -189,9 +182,54 @@ def generate_image(
                        image = image.convert("RGB")
                    image.save(buffer, format=image_format)

-                    yield ImageGenerationResponse(
+                    yield from _process_image_response(
                        image_data=buffer.getvalue(),
-                        format=task.output_format,
+                        image_format=task.output_format,
                        stats=stats,
                        image_index=image_num,
+                        model_id=model.model_id,
+                        partial_index=None,
+                        total_partials=None,
                    )
+
+
+def _process_image_response(
+    image_data: bytes,
+    image_index: int,
+    image_format: Literal["png", "jpeg", "webp"],
+    partial_index: int | None,
+    total_partials: int | None,
+    stats: ImageGenerationStats | None,
+    model_id: ModelId,
+) -> Iterator[ImageChunk]:
+    """Process a single image response and send chunks."""
+    is_partial = partial_index is not None
+    encoded_data = base64.b64encode(image_data).decode("utf-8")
+    # Extract stats from final ImageGenerationResponse if available
+    data_chunks = [
+        encoded_data[i : i + EXO_MAX_CHUNK_SIZE]
+        for i in range(0, len(encoded_data), EXO_MAX_CHUNK_SIZE)
+    ]
+    total_chunks = len(data_chunks)
+
+    def _data_to_chunk(item: tuple[int, str]) -> ImageChunk:
+        chunk_index, chunk_data = item
+        # Only include stats on the last chunk of the final image
+        chunk_stats = (
+            stats if chunk_index == total_chunks - 1 and not is_partial else None
+        )
+
+        return ImageChunk(
+            model=model_id,
+            data=chunk_data,
+            chunk_index=chunk_index,
+            total_chunks=total_chunks,
+            image_index=image_index,
+            is_partial=is_partial,
+            partial_index=partial_index,
+            total_partials=total_partials,
+            stats=chunk_stats,
+            format=image_format,
+        )
+
+    return map(_data_to_chunk, enumerate(data_chunks))
@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from collections.abc import Callable
+from collections.abc import Callable, Generator
 from functools import partial
 from inspect import signature
 from typing import TYPE_CHECKING, Literal, Protocol, cast
@@ -12,7 +12,7 @@ from mlx.nn.layers.distributed import (
    sum_gradients,
 )
 from mlx_lm.models.base import (
-    scaled_dot_product_attention,  # pyright: ignore[reportUnknownVariableType]
+    scaled_dot_product_attention,
 )
 from mlx_lm.models.cache import ArraysCache, KVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3MLP
@@ -59,14 +59,13 @@ 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.runner_response import ModelLoadingResponse
 from exo.shared.types.worker.shards import PipelineShardMetadata
 from exo.worker.runner.bootstrap import logger

 if TYPE_CHECKING:
    from mlx_lm.models.cache import Cache

-LayerLoadedCallback = Callable[[int, int], None]  # (layers_loaded, total_layers)
-

 _pending_prefill_sends: list[tuple[mx.array, int, mx.distributed.Group]] = []

@@ -276,8 +275,7 @@ def pipeline_auto_parallel(
    model: nn.Module,
    group: mx.distributed.Group,
    model_shard_meta: PipelineShardMetadata,
-    on_layer_loaded: LayerLoadedCallback | None,
-) -> nn.Module:
+) -> Generator[ModelLoadingResponse, None, nn.Module]:
    """
    Automatically parallelize a model across multiple devices.
    Args:
@@ -297,8 +295,7 @@ def pipeline_auto_parallel(
    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)
+        yield ModelLoadingResponse(layers_loaded=i, total=total)

    layers[0] = PipelineFirstLayer(layers[0], device_rank, group=group)
    layers[-1] = PipelineLastLayer(
@@ -309,24 +306,20 @@ def pipeline_auto_parallel(
    )

    if isinstance(inner_model_instance, GptOssMoeModel):
-        inner_model_instance.layer_types = inner_model_instance.layer_types[  # type: ignore
+        inner_model_instance.layer_types = inner_model_instance.layer_types[
            start_layer:end_layer
        ]
        # We can assume the model has at least one layer thanks to placement.
        # If a layer type doesn't exist, we can set it to 0.
        inner_model_instance.swa_idx = (
            0
-            if "sliding_attention" not in inner_model_instance.layer_types  # type: ignore
-            else inner_model_instance.layer_types.index(  # type: ignore
-                "sliding_attention"
-            )
+            if "sliding_attention" not in inner_model_instance.layer_types
+            else inner_model_instance.layer_types.index("sliding_attention")
        )
        inner_model_instance.ga_idx = (
            0
-            if "full_attention" not in inner_model_instance.layer_types  # type: ignore
-            else inner_model_instance.layer_types.index(  # type: ignore
-                "full_attention"
-            )
+            if "full_attention" not in inner_model_instance.layer_types
+            else inner_model_instance.layer_types.index("full_attention")
        )

    if isinstance(inner_model_instance, Step35InnerModel):
@@ -460,8 +453,7 @@ def patch_tensor_model[T](model: T) -> T:
 def tensor_auto_parallel(
    model: nn.Module,
    group: mx.distributed.Group,
-    on_layer_loaded: LayerLoadedCallback | None,
-) -> nn.Module:
+) -> Generator[ModelLoadingResponse, None, nn.Module]:
    all_to_sharded_linear = partial(
        shard_linear,
        sharding="all-to-sharded",
@@ -595,7 +587,7 @@ def tensor_auto_parallel(
    else:
        raise ValueError(f"Unsupported model type: {type(model)}")

-    model = tensor_parallel_sharding_strategy.shard_model(model, on_layer_loaded)
+    model = yield from tensor_parallel_sharding_strategy.shard_model(model)
    return patch_tensor_model(model)


@@ -619,16 +611,14 @@ class TensorParallelShardingStrategy(ABC):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module: ...
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]: ...


 class LlamaShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(LlamaModel, model)
        total = len(model.layers)
        for i, layer in enumerate(model.layers):
@@ -646,8 +636,8 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
            layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
            layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model


@@ -681,8 +671,7 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(DeepseekV3Model, model)
        total = len(model.layers)

@@ -738,8 +727,8 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
                layer.mlp.sharding_group = self.group

            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)

        return model

@@ -764,8 +753,7 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(GLM4MoeLiteModel, model)
        total = len(model.layers)  # type: ignore
        for i, layer in enumerate(model.layers):  # type: ignore
@@ -816,8 +804,8 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
                layer.mlp = ShardedMoE(layer.mlp)  # type: ignore
                layer.mlp.sharding_group = self.group  # type: ignore
            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)

        return model

@@ -891,7 +879,7 @@ class WrappedMiniMaxAttention(CustomMlxLayer):
            keys,
            values,
            cache=cache,
-            scale=self._original_layer.scale,  # type: ignore
+            scale=self._original_layer.scale,
            mask=mask,
        )

@@ -904,8 +892,7 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(MiniMaxModel, model)
        total = len(model.layers)
        for i, layer in enumerate(model.layers):
@@ -931,11 +918,11 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
            self.all_to_sharded_linear_in_place(
                layer.block_sparse_moe.switch_mlp.up_proj
            )
-            layer.block_sparse_moe = ShardedMoE(layer.block_sparse_moe)  # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
-            layer.block_sparse_moe.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
+            layer.block_sparse_moe = ShardedMoE(layer.block_sparse_moe)  # type: ignore
+            layer.block_sparse_moe.sharding_group = self.group
            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model


@@ -943,8 +930,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(
            Qwen3Model
            | Qwen3MoeModel
@@ -1099,8 +1085,8 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)

            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model


@@ -1108,8 +1094,7 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(Glm4MoeModel, model)
        total = len(model.layers)
        for i, layer in enumerate(model.layers):
@@ -1145,8 +1130,8 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
                layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)

            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model


@@ -1154,8 +1139,7 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(GptOssMoeModel, model)
        total = len(model.layers)

@@ -1184,10 +1168,10 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
            self.all_to_sharded_linear_in_place(layer.mlp.experts.up_proj)

            layer.mlp = ShardedMoE(layer.mlp)  # type: ignore
-            layer.mlp.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
+            layer.mlp.sharding_group = self.group
            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model


@@ -1195,8 +1179,7 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(Step35Model, model)
        total = len(model.layers)

@@ -1229,8 +1212,8 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
                self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)

            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model


@@ -1238,8 +1221,7 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(NemotronHModel, model)
        rank = self.group.rank()
        total = len(model.layers)
@@ -1272,8 +1254,7 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
                layer.mixer = mixer  # pyright: ignore[reportAttributeAccessIssue]

            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model

    def _shard_mamba2_mixer(self, mixer: NemotronHMamba2Mixer, rank: int) -> None:
@@ -1380,8 +1361,7 @@ class Gemma4ShardingStrategy(TensorParallelShardingStrategy):
    def shard_model(
        self,
        model: nn.Module,
-        on_layer_loaded: LayerLoadedCallback | None,
-    ) -> nn.Module:
+    ) -> Generator[ModelLoadingResponse, None, nn.Module]:
        model = cast(Gemma4Model, model)
        layers = model.language_model.model.layers
        total = len(layers)
@@ -1390,9 +1370,11 @@ class Gemma4ShardingStrategy(TensorParallelShardingStrategy):

            attn = layer.self_attn
            attn.q_proj = self.all_to_sharded_linear(attn.q_proj)
-            attn.k_proj = self.all_to_sharded_linear(attn.k_proj)
-            if not attn.use_k_eq_v:
-                attn.v_proj = self.all_to_sharded_linear(attn.v_proj)
+            has_kv: bool = cast(bool, attn.has_kv)
+            if has_kv:
+                attn.k_proj = self.all_to_sharded_linear(attn.k_proj)
+                if not attn.use_k_eq_v:
+                    attn.v_proj = self.all_to_sharded_linear(attn.v_proj)
            attn.o_proj = self.sharded_to_all_linear(attn.o_proj)
            attn.n_heads //= self.N
            attn.n_kv_heads //= self.N
@@ -1409,6 +1391,5 @@ class Gemma4ShardingStrategy(TensorParallelShardingStrategy):
                layer.experts.sharding_group = self.group

            mx.eval(layer)
-            if on_layer_loaded is not None:
-                on_layer_loaded(i, total)
+            yield ModelLoadingResponse(layers_loaded=i, total=total)
        return model
@@ -0,0 +1,108 @@
+import contextlib
+import os
+from collections.abc import Generator
+from dataclasses import dataclass
+
+import mlx.core as mx
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+from exo.shared.types.common import ModelId
+from exo.shared.types.events import Event
+from exo.shared.types.mlx import Model
+from exo.shared.types.tasks import TaskId
+from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.runner_response import ModelLoadingResponse
+from exo.utils.channels import MpReceiver, MpSender
+from exo.worker.engines.base import Builder, Engine
+from exo.worker.engines.mlx.cache import KVPrefixCache
+from exo.worker.engines.mlx.utils_mlx import (
+    initialize_mlx,
+    load_mlx_items,
+)
+from exo.worker.engines.mlx.vision import VisionProcessor
+from exo.worker.runner.bootstrap import logger
+from exo.worker.runner.llm_inference.batch_generator import (
+    BatchGenerator,
+    SequentialGenerator,
+)
+from exo.worker.runner.llm_inference.tool_parsers import make_mlx_parser
+
+
+@dataclass
+class MlxBuilder(Builder):
+    model_id: ModelId
+    event_sender: MpSender[Event]
+    cancel_receiver: MpReceiver[TaskId]
+    inference_model: Model | None = None
+    tokenizer: TokenizerWrapper | None = None
+    group: mx.distributed.Group | None = None
+    vision_processor: VisionProcessor | None = None
+
+    def connect(self, bound_instance: BoundInstance) -> None:
+        self.group = initialize_mlx(bound_instance)
+
+    def load(self, bound_instance: BoundInstance) -> Generator[ModelLoadingResponse]:
+        (
+            self.inference_model,
+            self.tokenizer,
+            self.vision_processor,
+        ) = yield from load_mlx_items(bound_instance, self.group)
+
+    def close(self) -> None:
+        with contextlib.suppress(NameError, AttributeError):
+            del self.inference_model, self.tokenizer, self.group
+
+    def build(
+        self,
+    ) -> Engine:
+        assert self.inference_model
+        assert self.tokenizer
+
+        vision_processor = self.vision_processor
+
+        tool_parser = None
+        logger.info(
+            f"model has_tool_calling={self.tokenizer.has_tool_calling} using tokens {self.tokenizer.tool_call_start}, {self.tokenizer.tool_call_end}"
+        )
+        if (
+            self.tokenizer.tool_call_start
+            and self.tokenizer.tool_call_end
+            and self.tokenizer.tool_parser  # type: ignore
+        ):
+            tool_parser = make_mlx_parser(
+                self.tokenizer.tool_call_start,
+                self.tokenizer.tool_call_end,
+                self.tokenizer.tool_parser,  # type: ignore
+            )
+
+        kv_prefix_cache = KVPrefixCache(self.group)
+
+        device_rank = 0 if self.group is None else self.group.rank()
+        if os.environ.get("EXO_NO_BATCH"):
+            logger.info("using SequentialGenerator (batching disabled)")
+            return SequentialGenerator(
+                model=self.inference_model,
+                tokenizer=self.tokenizer,
+                group=self.group,
+                tool_parser=tool_parser,
+                kv_prefix_cache=kv_prefix_cache,
+                model_id=self.model_id,
+                device_rank=device_rank,
+                cancel_receiver=self.cancel_receiver,
+                event_sender=self.event_sender,
+                vision_processor=vision_processor,
+            )
+        else:
+            logger.info("using BatchGenerator")
+            return BatchGenerator(
+                model=self.inference_model,
+                tokenizer=self.tokenizer,
+                group=self.group,
+                tool_parser=tool_parser,
+                kv_prefix_cache=kv_prefix_cache,
+                model_id=self.model_id,
+                device_rank=device_rank,
+                cancel_receiver=self.cancel_receiver,
+                event_sender=self.event_sender,
+                vision_processor=vision_processor,
+            )
@@ -249,7 +249,12 @@ class KVPrefixCache:
        # For partial match: trim to best_length, remaining has suffix to prefill
        # This ensures stream_generate always has at least one token to start with
        has_ssm = has_non_kv_caches(self.caches[best_index])
-        target = (max_length - 1) if is_exact and not has_ssm else best_length
+        cached_length = cache_length(self.caches[best_index])
+        if has_ssm:
+            target = best_length
+        else:
+            desired = (max_length - 1) if is_exact else best_length
+            target = min(cached_length, desired)
        restore_pos, restore_snap = self._get_snapshot(best_index, target)

        # No usable snapshot — need fresh cache
@@ -257,7 +262,6 @@ class KVPrefixCache:
            return make_kv_cache(model), prompt_tokens, None, False

        prompt_cache = deepcopy(self.caches[best_index])
-        cached_length = cache_length(self.caches[best_index])
        tokens_to_trim = cached_length - restore_pos
        if tokens_to_trim > 0:
            trim_cache(prompt_cache, tokens_to_trim, restore_snap)
@@ -17,6 +17,13 @@ TOOL_CALLS_START = f"<{DSML_TOKEN}function_calls>"
 TOOL_CALLS_END = f"</{DSML_TOKEN}function_calls>"
 _ORPHAN_THINK_END = ASSISTANT_TOKEN + THINKING_END
 _FIXED_THINK_BLOCK = ASSISTANT_TOKEN + THINKING_START + "\n" + THINKING_END
+_FUNCTION_RESULTS_CLOSE = "</function_results>"
+_ORPHAN_TOOL_RESULT_SUFFIX = _FUNCTION_RESULTS_CLOSE + "\n\n" + THINKING_END
+_EMPTY_THINK_BLOCKS = (
+    THINKING_START + "\n\n" + THINKING_END,
+    THINKING_START + "\n" + THINKING_END,
+    THINKING_START + THINKING_END,
+)


 def encode_messages(
@@ -35,7 +42,11 @@ def encode_messages(
        add_default_bos_token=add_default_bos_token,
        tools=tools,
    )
-    return prompt.replace(_ORPHAN_THINK_END, _FIXED_THINK_BLOCK)
+    prompt = prompt.replace(_ORPHAN_TOOL_RESULT_SUFFIX, _FUNCTION_RESULTS_CLOSE)
+    prompt = prompt.replace(_ORPHAN_THINK_END, _FIXED_THINK_BLOCK)
+    for empty in _EMPTY_THINK_BLOCKS:
+        prompt = prompt.replace(empty, "")
+    return prompt


 _INVOKE_PATTERN = re.compile(
@@ -457,6 +457,7 @@ class ExoBatchGenerator:

    def close(self) -> None:
        self._mlx_gen.close()
+        mx.clear_cache()

    def _save_prefix_cache(
        self,
@@ -4,6 +4,7 @@ import re
 import sys
 import tempfile
 import time
+from collections.abc import Generator
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast

@@ -51,6 +52,7 @@ from exo.shared.types.worker.instances import (
    MlxJacclInstance,
    MlxRingInstance,
 )
+from exo.shared.types.worker.runner_response import ModelLoadingResponse
 from exo.shared.types.worker.shards import (
    CfgShardMetadata,
    PipelineShardMetadata,
@@ -58,7 +60,6 @@ from exo.shared.types.worker.shards import (
    TensorShardMetadata,
 )
 from exo.worker.engines.mlx.auto_parallel import (
-    LayerLoadedCallback,
    get_inner_model,
    get_layers,
    pipeline_auto_parallel,
@@ -66,8 +67,6 @@ from exo.worker.engines.mlx.auto_parallel import (
 )
 from exo.worker.runner.bootstrap import logger

-Group = mx.distributed.Group
-

 def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
    return Memory.from_float_kb(
@@ -90,7 +89,7 @@ class HostList(RootModel[list[str]]):

 def mlx_distributed_init(
    bound_instance: BoundInstance,
-) -> Group:
+) -> mx.distributed.Group:
    """
    Initialize MLX distributed.
    """
@@ -149,7 +148,7 @@ def mlx_distributed_init(

 def initialize_mlx(
    bound_instance: BoundInstance,
-) -> Group:
+) -> mx.distributed.Group:
    # should we unseed it?
    # TODO: pass in seed from params
    mx.random.seed(42)
@@ -162,9 +161,10 @@ def initialize_mlx(

 def load_mlx_items(
    bound_instance: BoundInstance,
-    group: Group | None,
-    on_layer_loaded: LayerLoadedCallback | None,
-) -> "tuple[Model, TokenizerWrapper, VisionProcessor | None]":
+    group: mx.distributed.Group | None,
+) -> Generator[
+    ModelLoadingResponse, None, tuple[Model, TokenizerWrapper, "VisionProcessor | None"]
+]:
    if group is None:
        logger.info(f"Single device used for {bound_instance.instance}")
        model_path = build_model_path(bound_instance.bound_shard.model_card.model_id)
@@ -177,8 +177,7 @@ def load_mlx_items(
            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)
+                yield ModelLoadingResponse(layers_loaded=i, total=total)
        except ValueError as e:
            logger.opt(exception=e).debug(
                "Model architecture doesn't support layer-by-layer progress tracking",
@@ -191,10 +190,9 @@ def load_mlx_items(
    else:
        logger.info("Starting distributed init")
        start_time = time.perf_counter()
-        model, tokenizer = shard_and_load(
+        model, tokenizer = yield from shard_and_load(
            bound_instance.bound_shard,
            group=group,
-            on_layer_loaded=on_layer_loaded,
        )
        end_time = time.perf_counter()
        logger.info(
@@ -210,9 +208,20 @@ def load_mlx_items(
    if vision_config is not None:
        from exo.worker.engines.mlx.vision import VisionProcessor

-        vision_processor: VisionProcessor | None = VisionProcessor(
-            vision_config, bound_instance.bound_shard.model_card.model_id
-        )
+        vision_start_time = time.perf_counter()
+        try:
+            vision_processor: VisionProcessor | None = VisionProcessor(
+                vision_config, bound_instance.bound_shard.model_card.model_id
+            )
+            vision_processor.load()
+            logger.info(
+                f"Time taken to load vision weights: {(time.perf_counter() - vision_start_time):.2f}s"
+            )
+        except Exception as e:
+            logger.opt(exception=e).error(
+                "Failed to load vision weights — disabling vision for this runner"
+            )
+            vision_processor = None
    else:
        vision_processor = None

@@ -221,9 +230,8 @@ def load_mlx_items(

 def shard_and_load(
    shard_metadata: ShardMetadata,
-    group: Group,
-    on_layer_loaded: LayerLoadedCallback | None,
-) -> tuple[nn.Module, TokenizerWrapper]:
+    group: mx.distributed.Group,
+) -> Generator[ModelLoadingResponse, None, tuple[nn.Module, TokenizerWrapper]]:
    model_path = build_model_path(shard_metadata.model_card.model_id)

    model, _ = load_model(model_path, lazy=True, strict=False)
@@ -254,12 +262,10 @@ def shard_and_load(
    match shard_metadata:
        case TensorShardMetadata():
            logger.info(f"loading model from {model_path} with tensor parallelism")
-            model = tensor_auto_parallel(model, group, on_layer_loaded)
+            model = yield from tensor_auto_parallel(model, group)
        case PipelineShardMetadata():
            logger.info(f"loading model from {model_path} with pipeline parallelism")
-            model = pipeline_auto_parallel(
-                model, group, shard_metadata, on_layer_loaded=on_layer_loaded
-            )
+            model = yield from pipeline_auto_parallel(model, group, shard_metadata)
            mx.eval(model.parameters())
        case CfgShardMetadata():
            raise ValueError(
@@ -541,7 +547,6 @@ def render_chat_template(
        )
        if partial_assistant_content:
            prompt += partial_assistant_content
-        logger.info(prompt)
        return prompt

    for msg in formatted_messages:
@@ -748,7 +753,9 @@ def set_wired_limit_for_model(model_size: Memory):


 def mlx_cleanup(
-    model: Model | None, tokenizer: TokenizerWrapper | None, group: Group | None
+    model: Model | None,
+    tokenizer: TokenizerWrapper | None,
+    group: mx.distributed.Group | None,
 ) -> None:
    del model, tokenizer, group
    mx.clear_cache()
@@ -757,7 +764,7 @@ def mlx_cleanup(
    gc.collect()


-def mx_any(bool_: bool, group: Group | None) -> bool:
+def mx_any(bool_: bool, group: mx.distributed.Group | None) -> bool:
    if group is None:
        return bool_
    num_true = mx.distributed.all_sum(
@@ -767,7 +774,7 @@ def mx_any(bool_: bool, group: Group | None) -> bool:
    return num_true.item() > 0


-def mx_barrier(group: Group | None):
+def mx_barrier(group: mx.distributed.Group | None):
    if group is None:
        return
    mx.eval(
@@ -36,6 +36,19 @@ from exo.worker.runner.bootstrap import logger

 _video_processor_patched = False

+_MLX_VLM_MODEL_TYPE_ALIASES: dict[str, str] = {
+    "kimi_k25": "kimi_vl",
+    "kimi_k26": "kimi_vl",
+}
+
+
+def _torch_tensor_to_mx(
+    tensor: Any,  # pyright: ignore[reportAny]
+) -> mx.array:
+    if str(tensor.dtype) == "torch.bfloat16":  # type: ignore
+        return mx.array(tensor.float().numpy(), dtype=mx.bfloat16)  # type: ignore
+    return mx.array(tensor.numpy())  # type: ignore
+

 def _filter_config(cls: type, d: dict[str, Any]) -> dict[str, Any]:
    valid = set(inspect.signature(cls.__init__).parameters.keys()) - {"self"}
@@ -85,6 +98,8 @@ def _instantiate_projector(
    params = {n: p for n, p in init_sig.parameters.items() if n != "self"}
    kwargs: dict[str, Any] = {}

+    if "config" in params:
+        kwargs["config"] = model_config
    if "embedding_dim" in params:
        kwargs["embedding_dim"] = vision_config.hidden_size  # pyright: ignore[reportAny]
    if "text_hidden_size" in params:
@@ -205,7 +220,9 @@ class VisionEncoder:
        return {}

    def _import_mlx_vlm(self, *submodules: str) -> Any:  # type: ignore
-        mt = self._config.model_type
+        mt = _MLX_VLM_MODEL_TYPE_ALIASES.get(
+            self._config.model_type, self._config.model_type
+        )
        results: list[Any] = []
        for sub in submodules:
            name = f"mlx_vlm.models.{mt}.{sub}"
@@ -238,7 +255,7 @@ class VisionEncoder:
    def _load_image_processor_from_module(self, repo: str) -> "ImageProcessor | None":
        # mlx_vlm.utils.load_image_processor only works for models that set
        # `Model.ImageProcessor = <cls>`, but Gemma4 just uses
-        # `Gemma4ImageProcessor` from the package `__init__.py`
+        # `Gemma4ImageProcessor` from the package `__init__.py`.
        try:
            pkg: Any = importlib.import_module(
                f"mlx_vlm.models.{self._config.model_type}"
@@ -319,10 +336,16 @@ class VisionEncoder:
        else:
            self._load_weights_from_model_repo()

-        repo = processor_repo or str(self._model_path)
-        image_proc = load_image_processor(
-            repo
-        ) or self._load_image_processor_from_module(repo)
+        if processor_repo:
+            repo = str(build_model_path(ModelId(processor_repo)))
+        else:
+            repo = str(self._model_path)
+        try:
+            image_proc = load_image_processor(repo)
+        except ValueError:
+            image_proc = None
+        if image_proc is None:
+            image_proc = self._load_image_processor_from_module(repo)
        if image_proc is not None:
            self._processor = image_proc
        else:
@@ -339,39 +362,42 @@ class VisionEncoder:
        if not safetensors_files:
            raise FileNotFoundError(f"No safetensors files found in {self._model_path}")

-        weights: dict[str, mx.array] = {}
-        for sf_path in safetensors_files:
-            with safe_open(str(sf_path), framework="pt") as f:
-                keys = f.keys()
-                for key in keys:
-                    tensor = f.get_tensor(key)  # type: ignore
-                    np_tensor = tensor.float().numpy()  # type: ignore
-                    weights[key] = mx.array(np_tensor, dtype=mx.bfloat16)  # type: ignore
-
        vision_weights: dict[str, mx.array] = {}
        projector_weights: dict[str, mx.array] = {}
-        for key, val in weights.items():
-            if key.startswith("vision_tower."):
-                short_key = key[len("vision_tower.") :]
-                if short_key.startswith("encoder."):
-                    short_key = short_key[len("encoder.") :]
-                m = re.match(r"^(blocks\.\d+)\.(wqkv|wo)\.(weight|bias)$", short_key)
-                if m:
-                    short_key = f"{m.group(1)}.attn.{m.group(2)}.{m.group(3)}"
-                if short_key == "patch_embed.proj.weight" and val.ndim == 4:
-                    val = val.transpose(0, 2, 3, 1)
-                vision_weights[short_key] = val
-            elif key.startswith(("mm_projector.", "multi_modal_projector.")):
-                if key.startswith("multi_modal_projector."):
-                    short_key = key[len("multi_modal_projector.") :]
-                    if short_key.startswith("mm_projector."):
-                        short_key = short_key[len("mm_projector.") :]
-                else:
-                    short_key = key[len("mm_projector.") :]
-                short_key = short_key.replace("proj.0.", "linear_1.").replace(
-                    "proj.2.", "linear_2."
-                )
-                projector_weights[short_key] = val
+
+        for sf_path in safetensors_files:
+            with safe_open(str(sf_path), framework="pt") as f:
+                keys = cast(list[str], list(f.keys()))  # type: ignore
+                for key in keys:
+                    if key.startswith("vision_tower."):
+                        short_key = key[len("vision_tower.") :]
+                        if short_key.startswith("encoder."):
+                            short_key = short_key[len("encoder.") :]
+                        m = re.match(
+                            r"^(blocks\.\d+)\.(wqkv|wo)\.(weight|bias)$", short_key
+                        )
+                        if m:
+                            short_key = f"{m.group(1)}.attn.{m.group(2)}.{m.group(3)}"
+                        tensor = f.get_tensor(key)  # type: ignore
+                        val = mx.array(tensor.float().numpy(), dtype=mx.bfloat16)  # type: ignore
+                        if short_key == "patch_embed.proj.weight" and val.ndim == 4:
+                            val = val.transpose(0, 2, 3, 1)
+                        vision_weights[short_key] = val
+                    elif key.startswith(("mm_projector.", "multi_modal_projector.")):
+                        if key.startswith("multi_modal_projector."):
+                            short_key = key[len("multi_modal_projector.") :]
+                            if short_key.startswith("mm_projector."):
+                                short_key = short_key[len("mm_projector.") :]
+                        else:
+                            short_key = key[len("mm_projector.") :]
+                        short_key = short_key.replace("proj.0.", "linear_1.").replace(
+                            "proj.2.", "linear_2."
+                        )
+                        tensor = f.get_tensor(key)  # type: ignore
+                        projector_weights[short_key] = mx.array(
+                            tensor.float().numpy(),  # type: ignore
+                            dtype=mx.bfloat16,
+                        )

        assert self._vision_tower is not None
        self._vision_tower.load_weights(list(vision_weights.items()))
@@ -407,18 +433,26 @@ class VisionEncoder:
        needs_sanitize = False

        for sf_path in safetensors_files:
-            file_weights: dict[str, mx.array] = mx.load(str(sf_path))  # type: ignore
-            for key, val in file_weights.items():
-                for prefix in vision_prefixes:
-                    if key.startswith(prefix):
-                        vision_weights[key[len(prefix) :]] = val
-                        if prefix == "model.visual.":
-                            needs_sanitize = True
-                        break
-                else:
+            with safe_open(str(sf_path), framework="pt") as f:
+                keys = cast(list[str], list(f.keys()))  # type: ignore
+                for key in keys:
+                    matched = False
+                    for prefix in vision_prefixes:
+                        if key.startswith(prefix):
+                            vision_weights[key[len(prefix) :]] = _torch_tensor_to_mx(
+                                f.get_tensor(key)
+                            )
+                            if prefix == "model.visual.":
+                                needs_sanitize = True
+                            matched = True
+                            break
+                    if matched:
+                        continue
                    for prefix in projector_prefixes:
                        if key.startswith(prefix):
-                            projector_weights[key[len(prefix) :]] = val
+                            projector_weights[key[len(prefix) :]] = _torch_tensor_to_mx(
+                                f.get_tensor(key)
+                            )
                            break

        if not vision_weights:
@@ -463,7 +497,12 @@ class VisionEncoder:
        grid_thw: mx.array | None
        n_tokens_per_image: list[int]

-        if self._config.processor_repo:
+        is_kimi_vl_processor = any(
+            "mlx_vlm.models.kimi_vl" in cls.__module__
+            for cls in type(self._processor).__mro__
+        )
+
+        if self._config.processor_repo and not is_kimi_vl_processor:
            processed = self._processor.preprocess(
                [{"type": "image", "image": img} for img in pil_images],
                return_tensors="np",
@@ -481,6 +520,24 @@ class VisionEncoder:
                int(mx.prod(grid_thw[i]).item()) // merge_length
                for i in range(grid_thw.shape[0])
            ]
+        elif is_kimi_vl_processor:
+            proc: Any = self._processor
+            raw_processed = proc.preprocess(pil_images, return_tensors="np")  # type: ignore
+            stacked_pixels = mx.array(raw_processed["pixel_values"])  # type: ignore
+            if stacked_pixels.ndim == 3:
+                stacked_pixels = stacked_pixels[None]
+            per_image_pixels = [
+                stacked_pixels[i : i + 1] for i in range(stacked_pixels.shape[0])
+            ]
+            grid_raw = raw_processed.get("image_grid_hws")  # type: ignore
+            if grid_raw is None:
+                grid_raw = raw_processed["grid_thws"]  # type: ignore
+            grid_thw = mx.array(grid_raw)  # type: ignore
+            merge_length = int(np.prod(self._merge_kernel_size or [2, 2]))
+            n_tokens_per_image = [
+                int(mx.prod(grid_thw[i]).item()) // merge_length
+                for i in range(grid_thw.shape[0])
+            ]
        else:
            batch, tokens_override = _run_processor(self._processor, pil_images)
            # `Gemma4ImageProcessor` returns pixel_values as a plain ndarray
@@ -57,7 +57,7 @@ from exo.utils.info_gatherer.net_profile import check_reachable
 from exo.utils.keyed_backoff import KeyedBackoff
 from exo.utils.task_group import TaskGroup
 from exo.worker.plan import plan
-from exo.worker.runner.runner_supervisor import RunnerSupervisor
+from exo.worker.runner.supervisor import RunnerSupervisor


 class Worker:
@@ -152,6 +152,26 @@ class Worker:
                        event.chunk
                    )

+                    if (
+                        len(self.input_chunk_buffer[cmd_id])
+                        == self.input_chunk_counts[cmd_id]
+                    ):
+                        per_image: defaultdict[int, list[InputImageChunk]] = (
+                            defaultdict(list)
+                        )
+                        for chunk in self.input_chunk_buffer[cmd_id].values():
+                            per_image[chunk.image_index].append(chunk)
+                        for chunks_for_image in per_image.values():
+                            sorted_chunks = sorted(
+                                chunks_for_image, key=lambda c: c.chunk_index
+                            )
+                            img = Base64Image("".join(c.data for c in sorted_chunks))
+                            self.image_cache[
+                                Base64ImageHash(
+                                    hashlib.sha256(img.encode("ascii")).hexdigest()
+                                )
+                            ] = img
+
                if isinstance(event, CustomModelCardAdded):
                    await event.model_card.save_to_custom_dir()
                    add_to_card_cache(event.model_card)
@@ -170,6 +190,7 @@ class Worker:
                self.state.runners,
                self.state.tasks,
                self.input_chunk_buffer,
+                self.image_cache,
                self._instance_backoff,
                self._download_backoff,
            )
@@ -209,7 +230,7 @@ class Worker:
                    self._download_backoff.record_attempt(model_id)

                    found_path = await to_thread.run_sync(
-                        resolve_existing_model, model_id
+                        resolve_existing_model, model_id, shard.model_card
                    )
                    if found_path is not None:
                        logger.info(f"Model {model_id} found at {found_path}")
@@ -307,42 +328,11 @@ class Worker:
                        del self.input_chunk_counts[cmd_id]
                    await self._start_runner_task(modified_task)

-                case TextGeneration() if (
-                    task.task_params.image_hashes
-                    or task.task_params.total_input_chunks > 0
-                ):
+                case TextGeneration() if task.task_params.image_hashes:
                    cmd_id = task.command_id
-                    by_index: dict[int, Base64Image] = {}
-
-                    for idx, h in task.task_params.image_hashes.items():
-                        assert h in self.image_cache
-                        by_index[idx] = self.image_cache[h]
-
-                    if task.task_params.total_input_chunks > 0:
-                        chunk_buffer = self.input_chunk_buffer.get(cmd_id, {})
-                        per_image: defaultdict[int, list[InputImageChunk]] = (
-                            defaultdict(list)
-                        )
-                        for chunk in chunk_buffer.values():
-                            per_image[chunk.image_index].append(chunk)
-                        for img_idx in sorted(per_image):
-                            sorted_chunks = sorted(
-                                per_image[img_idx], key=lambda c: c.chunk_index
-                            )
-                            img = Base64Image("".join(c.data for c in sorted_chunks))
-                            self.image_cache[
-                                Base64ImageHash(
-                                    hashlib.sha256(img.encode("ascii")).hexdigest()
-                                )
-                            ] = img
-                            by_index[img_idx] = img
-                        logger.info(
-                            f"Assembled {len(per_image)} VLM image(s) "
-                            f"from {len(chunk_buffer)} chunks"
-                        )
-
                    resolved_images = [
-                        Base64Image(by_index[i]) for i in sorted(by_index)
+                        self.image_cache[h]
+                        for _, h in sorted(task.task_params.image_hashes.items())
                    ]
                    modified_task = task.model_copy(
                        update={
@@ -19,6 +19,7 @@ from exo.shared.types.tasks import (
    TaskStatus,
    TextGeneration,
 )
+from exo.shared.types.text_generation import Base64Image, Base64ImageHash
 from exo.shared.types.worker.downloads import (
    DownloadCompleted,
    DownloadFailed,
@@ -40,7 +41,7 @@ from exo.shared.types.worker.runners import (
    RunnerWarmingUp,
 )
 from exo.utils.keyed_backoff import KeyedBackoff
-from exo.worker.runner.runner_supervisor import RunnerSupervisor
+from exo.worker.runner.supervisor import RunnerSupervisor


 def plan(
@@ -52,6 +53,7 @@ def plan(
    all_runners: Mapping[RunnerId, RunnerStatus],  # all global
    tasks: Mapping[TaskId, Task],
    input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
+    image_cache: Mapping[Base64ImageHash, Base64Image],
    instance_backoff: KeyedBackoff[InstanceId],
    download_backoff: KeyedBackoff[ModelId],
 ) -> Task | None:
@@ -66,7 +68,7 @@ def plan(
        or _init_distributed_backend(runners, all_runners)
        or _load_model(runners, all_runners, global_download_status)
        or _ready_to_warmup(runners, all_runners)
-        or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer)
+        or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer, image_cache)
    )


@@ -300,6 +302,7 @@ def _pending_tasks(
    tasks: Mapping[TaskId, Task],
    all_runners: Mapping[RunnerId, RunnerStatus],
    input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
+    image_cache: Mapping[Base64ImageHash, Base64Image],
 ) -> Task | None:
    for task in tasks.values():
        # for now, just forward chat completions
@@ -309,16 +312,20 @@ def _pending_tasks(
        if task.task_status not in (TaskStatus.Pending, TaskStatus.Running):
            continue

-        # For tasks with images, verify all input chunks have been received
-        expected_image_chunks = 0
-        if isinstance(task, (ImageEdits, TextGeneration)):
-            expected_image_chunks = task.task_params.total_input_chunks
-        if expected_image_chunks > 0:
-            cmd_id = task.command_id
-            received = len(input_chunk_buffer.get(cmd_id, {}))
-            if received < expected_image_chunks:
+        if isinstance(task, ImageEdits) and task.task_params.total_input_chunks > 0:
+            received = len(input_chunk_buffer.get(task.command_id, {}))
+            if received < task.task_params.total_input_chunks:
                continue  # Wait for all chunks to arrive

+        if (
+            isinstance(task, TextGeneration)
+            and task.task_params.image_hashes
+            and not all(
+                h in image_cache for h in task.task_params.image_hashes.values()
+            )
+        ):
+            continue  # Wait for all images to be assembled into the cache
+
        for runner in runners.values():
            if task.instance_id != runner.bound_instance.instance.instance_id:
                continue
@@ -8,6 +8,7 @@ from exo.shared.types.tasks import Task, TaskId
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
+from exo.worker.engines.base import Builder

 logger: "loguru.Logger" = loguru.logger

@@ -35,23 +36,32 @@ def entrypoint(

    # Import main after setting global logger - this lets us just import logger from this module
    try:
-        if bound_instance.is_image_model:
-            from exo.worker.runner.image_models.runner import Runner as ImageRunner
+        from exo.worker.runner.runner import Runner

-            runner = ImageRunner(
-                bound_instance, event_sender, task_receiver, cancel_receiver
+        builder: Builder
+
+        if bound_instance.is_image_model:
+            from exo.worker.engines.image.builder import MfluxBuilder
+
+            builder = MfluxBuilder(
+                event_sender, cancel_receiver, bound_instance.bound_shard
            )
-            runner.main()
        else:
            from exo.worker.engines.mlx.patches import apply_mlx_patches
-            from exo.worker.runner.llm_inference.runner import Runner

            apply_mlx_patches()

-            runner = Runner(
-                bound_instance, event_sender, task_receiver, cancel_receiver
+            from exo.worker.engines.mlx.builder import MlxBuilder
+
+            # evil sharing of the event sender
+            builder = MlxBuilder(
+                model_id=bound_instance.bound_shard.model_card.model_id,
+                event_sender=event_sender,
+                cancel_receiver=cancel_receiver,
            )
-            runner.main()
+
+        runner = Runner(bound_instance, builder, event_sender, task_receiver)
+        runner.main()

    except ClosedResourceError:
        logger.warning("Runner communication closed unexpectedly")
@@ -1,403 +0,0 @@
-import base64
-import time
-from typing import TYPE_CHECKING, Literal
-
-import mlx.core as mx
-
-from exo.api.types import (
-    ImageEditsTaskParams,
-    ImageGenerationStats,
-    ImageGenerationTaskParams,
-)
-from exo.shared.constants import EXO_MAX_CHUNK_SIZE, EXO_TRACING_ENABLED
-from exo.shared.models.model_cards import ModelTask
-from exo.shared.tracing import clear_trace_buffer, get_trace_buffer
-from exo.shared.types.chunks import ErrorChunk, ImageChunk
-from exo.shared.types.common import CommandId, ModelId
-from exo.shared.types.events import (
-    ChunkGenerated,
-    Event,
-    RunnerStatusUpdated,
-    TaskAcknowledged,
-    TaskStatusUpdated,
-    TraceEventData,
-    TracesCollected,
-)
-from exo.shared.types.tasks import (
-    CANCEL_ALL_TASKS,
-    ConnectToGroup,
-    ImageEdits,
-    ImageGeneration,
-    LoadModel,
-    Shutdown,
-    StartWarmup,
-    Task,
-    TaskId,
-    TaskStatus,
-)
-from exo.shared.types.worker.instances import BoundInstance
-from exo.shared.types.worker.runner_response import (
-    ImageGenerationResponse,
-    PartialImageResponse,
-)
-from exo.shared.types.worker.runners import (
-    RunnerConnected,
-    RunnerConnecting,
-    RunnerIdle,
-    RunnerLoaded,
-    RunnerLoading,
-    RunnerReady,
-    RunnerRunning,
-    RunnerShutdown,
-    RunnerShuttingDown,
-    RunnerStatus,
-    RunnerWarmingUp,
-)
-from exo.shared.types.worker.shards import (
-    CfgShardMetadata,
-    PipelineShardMetadata,
-    ShardMetadata,
-)
-from exo.utils.channels import MpReceiver, MpSender
-from exo.worker.engines.image import (
-    DistributedImageModel,
-    generate_image,
-    initialize_image_model,
-    warmup_image_generator,
-)
-from exo.worker.engines.mlx.utils_mlx import (
-    initialize_mlx,
-)
-from exo.worker.runner.bootstrap import logger
-
-
-def _is_primary_output_node(shard_metadata: ShardMetadata) -> bool:
-    """Check if this node is the primary output node for image generation.
-
-    For CFG models: the last pipeline stage in CFG group 0 (positive prompt).
-    For non-CFG models: the last pipeline stage.
-    """
-    if isinstance(shard_metadata, CfgShardMetadata):
-        is_pipeline_last = (
-            shard_metadata.pipeline_rank == shard_metadata.pipeline_world_size - 1
-        )
-        return is_pipeline_last and shard_metadata.cfg_rank == 0
-    elif isinstance(shard_metadata, PipelineShardMetadata):
-        return shard_metadata.device_rank == shard_metadata.world_size - 1
-    return False
-
-
-def _process_image_response(
-    response: ImageGenerationResponse | PartialImageResponse,
-    command_id: CommandId,
-    shard_metadata: ShardMetadata,
-    event_sender: MpSender[Event],
-    image_index: int,
-) -> None:
-    """Process a single image response and send chunks."""
-    encoded_data = base64.b64encode(response.image_data).decode("utf-8")
-    is_partial = isinstance(response, PartialImageResponse)
-    # Extract stats from final ImageGenerationResponse if available
-    stats = response.stats if isinstance(response, ImageGenerationResponse) else None
-    _send_image_chunk(
-        encoded_data=encoded_data,
-        command_id=command_id,
-        model_id=shard_metadata.model_card.model_id,
-        event_sender=event_sender,
-        image_index=response.image_index,
-        is_partial=is_partial,
-        partial_index=response.partial_index if is_partial else None,
-        total_partials=response.total_partials if is_partial else None,
-        stats=stats,
-        image_format=response.format,
-    )
-
-
-def _send_traces_if_enabled(
-    event_sender: MpSender[Event],
-    task_id: TaskId,
-    rank: int,
-) -> None:
-    if not EXO_TRACING_ENABLED:
-        return
-
-    traces = get_trace_buffer()
-    if traces:
-        trace_data = [
-            TraceEventData(
-                name=t.name,
-                start_us=t.start_us,
-                duration_us=t.duration_us,
-                rank=t.rank,
-                category=t.category,
-            )
-            for t in traces
-        ]
-        event_sender.send(
-            TracesCollected(
-                task_id=task_id,
-                rank=rank,
-                traces=trace_data,
-            )
-        )
-    clear_trace_buffer()
-
-
-def _send_image_chunk(
-    encoded_data: str,
-    command_id: CommandId,
-    model_id: ModelId,
-    event_sender: MpSender[Event],
-    image_index: int,
-    is_partial: bool,
-    partial_index: int | None = None,
-    total_partials: int | None = None,
-    stats: ImageGenerationStats | None = None,
-    image_format: Literal["png", "jpeg", "webp"] | None = None,
-) -> None:
-    """Send base64-encoded image data as chunks via events."""
-    data_chunks = [
-        encoded_data[i : i + EXO_MAX_CHUNK_SIZE]
-        for i in range(0, len(encoded_data), EXO_MAX_CHUNK_SIZE)
-    ]
-    total_chunks = len(data_chunks)
-    for chunk_index, chunk_data in enumerate(data_chunks):
-        # Only include stats on the last chunk of the final image
-        chunk_stats = (
-            stats if chunk_index == total_chunks - 1 and not is_partial else None
-        )
-        event_sender.send(
-            ChunkGenerated(
-                command_id=command_id,
-                chunk=ImageChunk(
-                    model=model_id,
-                    data=chunk_data,
-                    chunk_index=chunk_index,
-                    total_chunks=total_chunks,
-                    image_index=image_index,
-                    is_partial=is_partial,
-                    partial_index=partial_index,
-                    total_partials=total_partials,
-                    stats=chunk_stats,
-                    format=image_format,
-                ),
-            )
-        )
-
-
-class Runner:
-    def __init__(
-        self,
-        bound_instance: BoundInstance,
-        event_sender: MpSender[Event],
-        task_receiver: MpReceiver[Task],
-        cancel_receiver: MpReceiver[TaskId],
-    ):
-        self.event_sender = event_sender
-        self.task_receiver = task_receiver
-        self.cancel_receiver = cancel_receiver
-        self.bound_instance = bound_instance
-
-        self.instance, self.runner_id, self.shard_metadata = (
-            bound_instance.instance,
-            bound_instance.bound_runner_id,
-            bound_instance.bound_shard,
-        )
-        self.device_rank = self.shard_metadata.device_rank
-
-        logger.info("hello from the runner")
-        if getattr(self.shard_metadata, "immediate_exception", False):
-            raise Exception("Fake exception - runner failed to spin up.")
-        if timeout := getattr(self.shard_metadata, "should_timeout", 0):
-            time.sleep(timeout)
-
-        self.setup_start_time = time.time()
-        self.cancelled_tasks = set[TaskId]()
-
-        self.image_model: DistributedImageModel | None = None
-        self.group = None
-
-        self.current_status: RunnerStatus = RunnerIdle()
-        logger.info("runner created")
-        self.update_status(RunnerIdle())
-        self.seen = set[TaskId]()
-
-    def update_status(self, status: RunnerStatus):
-        self.current_status = status
-        self.event_sender.send(
-            RunnerStatusUpdated(
-                runner_id=self.runner_id, runner_status=self.current_status
-            )
-        )
-
-    def send_task_status(self, task: Task, status: TaskStatus):
-        self.event_sender.send(
-            TaskStatusUpdated(task_id=task.task_id, task_status=status)
-        )
-
-    def acknowledge_task(self, task: Task):
-        self.event_sender.send(TaskAcknowledged(task_id=task.task_id))
-
-    def _check_cancelled(self, task_id: TaskId) -> bool:
-        for cancel_id in self.cancel_receiver.collect():
-            self.cancelled_tasks.add(cancel_id)
-        return (
-            task_id in self.cancelled_tasks or CANCEL_ALL_TASKS in self.cancelled_tasks
-        )
-
-    def _run_image_task(
-        self,
-        task: Task,
-        task_params: ImageGenerationTaskParams | ImageEditsTaskParams,
-        command_id: CommandId,
-    ) -> None:
-        assert self.image_model
-        logger.info(f"received image task: {str(task)[:500]}")
-        logger.info("runner running")
-        self.update_status(RunnerRunning())
-        self.acknowledge_task(task)
-
-        def cancel_checker() -> bool:
-            return self._check_cancelled(task.task_id)
-
-        try:
-            image_index = 0
-            for response in generate_image(
-                model=self.image_model,
-                task=task_params,
-                cancel_checker=cancel_checker,
-            ):
-                if _is_primary_output_node(self.shard_metadata):
-                    match response:
-                        case PartialImageResponse():
-                            logger.info(
-                                f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
-                            )
-                            _process_image_response(
-                                response,
-                                command_id,
-                                self.shard_metadata,
-                                self.event_sender,
-                                image_index,
-                            )
-                        case ImageGenerationResponse():
-                            logger.info("sending final ImageChunk")
-                            _process_image_response(
-                                response,
-                                command_id,
-                                self.shard_metadata,
-                                self.event_sender,
-                                image_index,
-                            )
-                            image_index += 1
-        except Exception as e:
-            if _is_primary_output_node(self.shard_metadata):
-                self.event_sender.send(
-                    ChunkGenerated(
-                        command_id=command_id,
-                        chunk=ErrorChunk(
-                            model=self.shard_metadata.model_card.model_id,
-                            finish_reason="error",
-                            error_message=str(e),
-                        ),
-                    )
-                )
-            raise
-        finally:
-            _send_traces_if_enabled(self.event_sender, task.task_id, self.device_rank)
-
-        self.current_status = RunnerReady()
-        logger.info("runner ready")
-
-    def main(self):
-        with self.task_receiver as tasks:
-            for task in tasks:
-                if task.task_id in self.seen:
-                    logger.warning("repeat task - potential error")
-                self.seen.add(task.task_id)
-                self.cancelled_tasks.discard(CANCEL_ALL_TASKS)
-                self.send_task_status(task, TaskStatus.Running)
-                self.handle_task(task)
-                was_cancelled = (task.task_id in self.cancelled_tasks) or (
-                    CANCEL_ALL_TASKS in self.cancelled_tasks
-                )
-                if not was_cancelled:
-                    self.send_task_status(task, TaskStatus.Complete)
-                self.update_status(self.current_status)
-
-                if isinstance(self.current_status, RunnerShutdown):
-                    break
-
-    def handle_task(self, task: Task):
-        match task:
-            case ConnectToGroup() if isinstance(self.current_status, RunnerIdle):
-                logger.info("runner connecting")
-                self.update_status(RunnerConnecting())
-                self.acknowledge_task(task)
-                self.group = initialize_mlx(self.bound_instance)
-
-                logger.info("runner connected")
-                self.current_status = RunnerConnected()
-
-            # we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
-            case LoadModel() if (
-                isinstance(self.current_status, RunnerConnected)
-                and self.group is not None
-            ) or (isinstance(self.current_status, RunnerIdle) and self.group is None):
-                logger.info("runner loading")
-                self.update_status(RunnerLoading())
-                self.acknowledge_task(task)
-
-                assert (
-                    ModelTask.TextToImage in self.shard_metadata.model_card.tasks
-                    or ModelTask.ImageToImage in self.shard_metadata.model_card.tasks
-                ), f"Incorrect model task(s): {self.shard_metadata.model_card.tasks}"
-
-                self.image_model = initialize_image_model(self.bound_instance)
-                self.current_status = RunnerLoaded()
-                logger.info("runner loaded")
-
-            case StartWarmup() if isinstance(self.current_status, RunnerLoaded):
-                logger.info("runner warming up")
-                self.update_status(RunnerWarmingUp())
-                self.acknowledge_task(task)
-
-                logger.info(f"warming up inference for instance: {self.instance}")
-
-                assert self.image_model
-                image = warmup_image_generator(model=self.image_model)
-                if image is not None:
-                    logger.info(f"warmed up by generating {image.size} image")
-                else:
-                    logger.info("warmup completed (non-primary node)")
-
-                logger.info(
-                    f"runner initialized in {time.time() - self.setup_start_time} seconds"
-                )
-
-                self.current_status = RunnerReady()
-                logger.info("runner ready")
-
-            case (
-                ImageGeneration(task_params=task_params, command_id=command_id)
-                | ImageEdits(task_params=task_params, command_id=command_id)
-            ) if isinstance(self.current_status, RunnerReady):
-                self._run_image_task(task, task_params, command_id)
-
-            case Shutdown():
-                logger.info("runner shutting down")
-                if not TYPE_CHECKING:
-                    del self.image_model, self.group
-                    mx.clear_cache()
-                    import gc
-
-                    gc.collect()
-
-                self.update_status(RunnerShuttingDown())
-                self.acknowledge_task(task)
-
-                self.current_status = RunnerShutdown()
-            case _:
-                raise ValueError(
-                    f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}"
-                )
@@ -1,22 +1,31 @@
 import itertools
 import time
-from abc import ABC, abstractmethod
 from collections import deque
-from collections.abc import Generator, Iterable
+from collections.abc import Generator, Iterator
 from dataclasses import dataclass, field

 import mlx.core as mx
 from mlx_lm.tokenizer_utils import TokenizerWrapper

 from exo.shared.constants import EXO_MAX_CONCURRENT_REQUESTS
-from exo.shared.types.chunks import ErrorChunk, PrefillProgressChunk
+from exo.shared.types.chunks import ErrorChunk, GenerationChunk, PrefillProgressChunk
 from exo.shared.types.common import ModelId
 from exo.shared.types.events import ChunkGenerated, Event
 from exo.shared.types.mlx import Model
-from exo.shared.types.tasks import CANCEL_ALL_TASKS, TaskId, TextGeneration
+from exo.shared.types.tasks import (
+    CANCEL_ALL_TASKS,
+    GenerationTask,
+    TaskId,
+    TextGeneration,
+)
 from exo.shared.types.text_generation import TextGenerationTaskParams
-from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
+from exo.shared.types.worker.runner_response import (
+    CancelledResponse,
+    FinishedResponse,
+    GenerationResponse,
+)
 from exo.utils.channels import MpReceiver, MpSender
+from exo.worker.engines.base import Engine
 from exo.worker.engines.mlx.cache import KVPrefixCache
 from exo.worker.engines.mlx.generator.batch_generate import ExoBatchGenerator
 from exo.worker.engines.mlx.generator.generate import (
@@ -32,18 +41,10 @@ from exo.worker.engines.mlx.utils_mlx import (
 from exo.worker.engines.mlx.vision import VisionProcessor
 from exo.worker.runner.bootstrap import logger

-from .model_output_parsers import apply_all_parsers
+from .model_output_parsers import apply_all_parsers, map_responses_to_chunks
 from .tool_parsers import ToolParser


-class Cancelled:
-    pass
-
-
-class Finished:
-    pass
-
-
 class GeneratorQueue[T]:
    def __init__(self):
        self._q = deque[T]()
@@ -59,35 +60,6 @@ class GeneratorQueue[T]:
                yield self._q.popleft()


-class InferenceGenerator(ABC):
-    _cancelled_tasks: set[TaskId]
-
-    def should_cancel(self, task_id: TaskId) -> bool:
-        return (
-            task_id in self._cancelled_tasks
-            or CANCEL_ALL_TASKS in self._cancelled_tasks
-        )
-
-    @abstractmethod
-    def warmup(self) -> None: ...
-
-    @abstractmethod
-    def submit(
-        self,
-        task: TextGeneration,
-    ) -> None: ...
-
-    @abstractmethod
-    def step(
-        self,
-    ) -> Iterable[
-        tuple[TaskId, ToolCallResponse | GenerationResponse | Cancelled | Finished]
-    ]: ...
-
-    @abstractmethod
-    def close(self) -> None: ...
-
-
 EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
 EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM"
 EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"
@@ -111,7 +83,7 @@ def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None:


@dataclass(eq=False)
-class SequentialGenerator(InferenceGenerator):
+class SequentialGenerator(Engine):
    model: Model
    tokenizer: TokenizerWrapper
    group: mx.distributed.Group | None
@@ -137,7 +109,7 @@ class SequentialGenerator(InferenceGenerator):
            # queue that the 1st generator should push to and 3rd generator should pull from
            GeneratorQueue[GenerationResponse],
            # generator to get parsed outputs
-            Generator[GenerationResponse | ToolCallResponse | None],
+            Iterator[GenerationChunk | None],
        ]
        | None
    ) = field(default=None, init=False)
@@ -152,8 +124,9 @@ class SequentialGenerator(InferenceGenerator):

    def submit(
        self,
-        task: TextGeneration,
+        task: GenerationTask,
    ) -> None:
+        assert isinstance(task, TextGeneration)
        self._cancelled_tasks.discard(CANCEL_ALL_TASKS)
        self._all_tasks[task.task_id] = task
        self._maybe_queue.append(task)
@@ -183,8 +156,8 @@ class SequentialGenerator(InferenceGenerator):

    def step(
        self,
-    ) -> Iterable[
-        tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
+    ) -> Iterator[
+        tuple[TaskId, GenerationChunk | FinishedResponse | CancelledResponse]
    ]:
        if self._active is None:
            self.agree_on_tasks()
@@ -192,23 +165,25 @@ class SequentialGenerator(InferenceGenerator):
            if self._queue:
                self._start_next()
            else:
-                return map(lambda task: (task, Cancelled()), self._cancelled_tasks)
+                return map(
+                    lambda task: (task, CancelledResponse()), self._cancelled_tasks
+                )

        assert self._active is not None

-        task, mlx_gen, queue, output_generator = self._active
+        task, gen, queue, output_generator = self._active
        output: list[
-            tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
+            tuple[TaskId, GenerationChunk | CancelledResponse | FinishedResponse]
        ] = []
        try:
-            response = next(mlx_gen)
+            response = next(gen)
            queue.push(response)
            # drain potentially many responses every time
            while (parsed := next(output_generator, None)) is not None:
                output.append((task.task_id, parsed))

        except (StopIteration, PrefillCancelled):
-            output.append((task.task_id, Finished()))
+            output.append((task.task_id, FinishedResponse()))
            self._active = None
            if self._queue:
                self._start_next()
@@ -220,20 +195,22 @@ class SequentialGenerator(InferenceGenerator):

        return itertools.chain(
            output,
-            map(lambda task: (task, Cancelled()), self._cancelled_tasks),
+            map(lambda task: (task, CancelledResponse()), self._cancelled_tasks),
        )

    def _start_next(self) -> None:
        task = self._queue.popleft()
        try:
-            mlx_gen = self._build_generator(task)
+            gen = self._build_generator(task)
        except Exception as e:
            self._send_error(task, e)
            raise
        queue = GeneratorQueue[GenerationResponse]()

        if task.task_params.bench:
-            output_generator = queue.gen()
+            output_generator: Iterator[GenerationChunk | None] = map(
+                lambda r: map_responses_to_chunks(r, self.model_id), queue.gen()
+            )
        else:
            output_generator = apply_all_parsers(
                queue.gen(),
@@ -244,7 +221,7 @@ class SequentialGenerator(InferenceGenerator):
                self.model_id,
                task.task_params.tools,
            )
-        self._active = (task, mlx_gen, queue, output_generator)
+        self._active = (task, gen, queue, output_generator)

    def _send_error(self, task: TextGeneration, e: Exception) -> None:
        if self.device_rank == 0:
@@ -314,7 +291,7 @@ class SequentialGenerator(InferenceGenerator):


@dataclass(eq=False)
-class BatchGenerator(InferenceGenerator):
+class BatchGenerator(Engine):
    model: Model
    tokenizer: TokenizerWrapper
    group: mx.distributed.Group | None
@@ -332,18 +309,18 @@ class BatchGenerator(InferenceGenerator):
    _maybe_cancel: list[TextGeneration] = field(default_factory=list, init=False)
    _all_tasks: dict[TaskId, TextGeneration] = field(default_factory=dict, init=False)
    _queue: deque[TextGeneration] = field(default_factory=deque, init=False)
-    _mlx_gen: ExoBatchGenerator = field(init=False)
+    _gen: ExoBatchGenerator = field(init=False)
    _active_tasks: dict[
        int,
        tuple[
            TextGeneration,
            GeneratorQueue[GenerationResponse],
-            Generator[GenerationResponse | ToolCallResponse | None],
+            Iterator[GenerationChunk | None],
        ],
    ] = field(default_factory=dict, init=False)

    def __post_init__(self) -> None:
-        self._mlx_gen = ExoBatchGenerator(
+        self._gen = ExoBatchGenerator(
            model=self.model,
            tokenizer=self.tokenizer,
            group=self.group,
@@ -361,8 +338,9 @@ class BatchGenerator(InferenceGenerator):

    def submit(
        self,
-        task: TextGeneration,
+        task: GenerationTask,
    ) -> None:
+        assert isinstance(task, TextGeneration)
        self._cancelled_tasks.discard(CANCEL_ALL_TASKS)
        self._all_tasks[task.task_id] = task
        self._maybe_queue.append(task)
@@ -392,8 +370,8 @@ class BatchGenerator(InferenceGenerator):

    def step(
        self,
-    ) -> Iterable[
-        tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
+    ) -> Iterator[
+        tuple[TaskId, GenerationChunk | CancelledResponse | FinishedResponse]
    ]:
        if not self._queue:
            self.agree_on_tasks()
@@ -411,7 +389,9 @@ class BatchGenerator(InferenceGenerator):

            queue = GeneratorQueue[GenerationResponse]()
            if task.task_params.bench:
-                output_generator = queue.gen()
+                output_generator: Iterator[GenerationChunk | None] = map(
+                    lambda r: map_responses_to_chunks(r, self.model_id), queue.gen()
+                )
            else:
                output_generator = apply_all_parsers(
                    queue.gen(),
@@ -424,13 +404,13 @@ class BatchGenerator(InferenceGenerator):
                )
            self._active_tasks[uid] = (task, queue, output_generator)

-        if not self._mlx_gen.has_work:
+        if not self._gen.has_work:
            return self._apply_cancellations()

-        results = self._mlx_gen.step()
+        results = self._gen.step()

        output: list[
-            tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
+            tuple[TaskId, GenerationChunk | CancelledResponse | FinishedResponse]
        ] = []
        for uid, response in results:
            if uid not in self._active_tasks:
@@ -446,38 +426,38 @@ class BatchGenerator(InferenceGenerator):

            # check if original response was terminal and append a Finished()
            if response.finish_reason is not None:
-                output.append((task.task_id, Finished()))
+                output.append((task.task_id, FinishedResponse()))
                del self._active_tasks[uid]

        return itertools.chain(output, self._apply_cancellations())

    def _apply_cancellations(
        self,
-    ) -> list[tuple[TaskId, Cancelled]]:
+    ) -> Iterator[tuple[TaskId, CancelledResponse]]:
        if not self._cancelled_tasks:
-            return []
+            return iter([])

        cancel_all = CANCEL_ALL_TASKS in self._cancelled_tasks

        uids_to_cancel: list[int] = []
-        results: list[tuple[TaskId, Cancelled]] = []
+        results: list[tuple[TaskId, CancelledResponse]] = []

        for uid, (task, _, _) in list(self._active_tasks.items()):
            if task.task_id in self._cancelled_tasks or cancel_all:
                uids_to_cancel.append(uid)
-                results.append((task.task_id, Cancelled()))
+                results.append((task.task_id, CancelledResponse()))
                del self._active_tasks[uid]

        if uids_to_cancel:
-            self._mlx_gen.cancel(uids_to_cancel)
+            self._gen.cancel(uids_to_cancel)

        already_cancelled = {tid for tid, _ in results}
        for tid in self._cancelled_tasks:
            if tid != CANCEL_ALL_TASKS and tid not in already_cancelled:
-                results.append((tid, Cancelled()))
+                results.append((tid, CancelledResponse()))

        self._cancelled_tasks.clear()
-        return results
+        return iter(results)

    def _send_error(self, task: TextGeneration, e: Exception) -> None:
        if self.device_rank == 0:
@@ -529,7 +509,7 @@ class BatchGenerator(InferenceGenerator):

                self.agree_on_tasks()

-        return self._mlx_gen.submit(
+        return self._gen.submit(
            task_params=task.task_params,
            prompt=prompt,
            on_prefill_progress=on_prefill_progress,
@@ -538,5 +518,5 @@ class BatchGenerator(InferenceGenerator):
        )

    def close(self) -> None:
-        self._mlx_gen.close()
+        self._gen.close()
        del self.model, self.tokenizer, self.group
@@ -1,4 +1,4 @@
-from collections.abc import Generator
+from collections.abc import Generator, Iterator
 from functools import cache
 from typing import Any

@@ -14,6 +14,12 @@ from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
 )

 from exo.api.types import ToolCallItem
+from exo.shared.types.chunks import (
+    ErrorChunk,
+    GenerationChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
 from exo.shared.types.common import ModelId
 from exo.shared.types.mlx import Model
 from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
@@ -64,29 +70,77 @@ def apply_all_parsers(
    model_type: type[Model],
    model_id: ModelId,
    tools: list[dict[str, Any]] | None,
-) -> Generator[GenerationResponse | ToolCallResponse | None]:
-    mlx_generator = receiver
+) -> Iterator[GenerationChunk | None]:
+    generator = receiver

    if issubclass(model_type, GptOssModel):
-        mlx_generator = parse_gpt_oss(mlx_generator)
+        generator = parse_gpt_oss(generator)
    elif (
        issubclass(model_type, DeepseekV32Model)
        and "deepseek" in model_id.normalize().lower()
    ):
-        mlx_generator = parse_deepseek_v32(mlx_generator)
+        if tokenizer.has_thinking:
+            generator = parse_thinking_models(
+                generator,
+                tokenizer.think_start,
+                tokenizer.think_end,
+                starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer),
+            )
+        generator = parse_deepseek_v32(generator)
    else:
        if tokenizer.has_thinking:
-            mlx_generator = parse_thinking_models(
-                mlx_generator,
+            generator = parse_thinking_models(
+                generator,
                tokenizer.think_start,
                tokenizer.think_end,
                starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer),
            )

        if tool_parser:
-            mlx_generator = parse_tool_calls(mlx_generator, tool_parser, tools)
+            generator = parse_tool_calls(generator, tool_parser, tools)

-    return count_reasoning_tokens(mlx_generator)
+    generator = count_reasoning_tokens(generator)
+
+    return map(lambda r: map_responses_to_chunks(r, model_id), generator)
+
+
+def map_responses_to_chunks(
+    response: GenerationResponse | ToolCallResponse | None, model_id: ModelId
+) -> GenerationChunk | None:
+    match response:
+        case None:
+            return None
+        case GenerationResponse():
+            if response.finish_reason == "error":
+                return ErrorChunk(
+                    error_message=response.text,
+                    model=model_id,
+                )
+            else:
+                finish_reason = response.finish_reason
+                assert finish_reason not in (
+                    "error",
+                    "tool_calls",
+                    "function_call",
+                )
+                return TokenChunk(
+                    model=model_id,
+                    text=response.text,
+                    token_id=response.token,
+                    usage=response.usage,
+                    finish_reason=finish_reason,
+                    stats=response.stats,
+                    logprob=response.logprob,
+                    top_logprobs=response.top_logprobs,
+                    is_thinking=response.is_thinking,
+                )
+        case ToolCallResponse():
+            return ToolCallChunk(
+                tool_calls=response.tool_calls,
+                model=model_id,
+                usage=response.usage,
+                stats=response.stats,
+            )


 def parse_gpt_oss(
@@ -163,11 +217,10 @@ def parse_deepseek_v32(

    Uses accumulated-text matching (not per-token marker checks) because
    DSML markers like <｜DSML｜function_calls> may span multiple tokens.
-    Also handles <think>...</think> blocks for thinking mode.
+    Thinking tag handling is delegated to parse_thinking_models, which
+    wraps this parser in apply_all_parsers.
    """
    from exo.worker.engines.mlx.dsml_encoding import (
-        THINKING_END,
-        THINKING_START,
        TOOL_CALLS_END,
        TOOL_CALLS_START,
        parse_dsml_output,
@@ -175,7 +228,6 @@ def parse_deepseek_v32(

    accumulated = ""
    in_tool_call = False
-    thinking = False
    # Tokens buffered while we detect the start of a DSML block
    pending_buffer: list[GenerationResponse] = []
    # Text accumulated during a tool call block
@@ -217,29 +269,6 @@ def parse_deepseek_v32(
                yield response
            break

-        # ── Handle thinking tags ──
-        if not thinking and THINKING_START in response.text:
-            thinking = True
-            # Yield any text before the <think> tag
-            before = response.text[: response.text.index(THINKING_START)]
-            if before:
-                yield response.model_copy(update={"text": before})
-            continue
-
-        if thinking and THINKING_END in response.text:
-            thinking = False
-            # Yield any text after the </think> tag
-            after = response.text[
-                response.text.index(THINKING_END) + len(THINKING_END) :
-            ]
-            if after:
-                yield response.model_copy(update={"text": after, "is_thinking": False})
-            continue
-
-        if thinking:
-            yield response.model_copy(update={"is_thinking": True})
-            continue
-
        # ── Handle tool call accumulation ──
        if in_tool_call:
            tool_call_text += response.text
@@ -1,434 +0,0 @@
-import os
-import time
-from dataclasses import dataclass
-from enum import Enum
-
-import mlx.core as mx
-from anyio import WouldBlock
-from mlx_lm.tokenizer_utils import TokenizerWrapper
-
-from exo.shared.models.model_cards import ModelTask
-from exo.shared.types.chunks import (
-    ErrorChunk,
-    TokenChunk,
-    ToolCallChunk,
-)
-from exo.shared.types.common import CommandId, ModelId
-from exo.shared.types.events import (
-    ChunkGenerated,
-    Event,
-    RunnerStatusUpdated,
-    TaskAcknowledged,
-    TaskStatusUpdated,
-)
-from exo.shared.types.mlx import Model
-from exo.shared.types.tasks import (
-    ConnectToGroup,
-    LoadModel,
-    Shutdown,
-    StartWarmup,
-    Task,
-    TaskId,
-    TaskStatus,
-    TextGeneration,
-)
-from exo.shared.types.worker.instances import BoundInstance
-from exo.shared.types.worker.runner_response import (
-    GenerationResponse,
-    ToolCallResponse,
-)
-from exo.shared.types.worker.runners import (
-    RunnerConnected,
-    RunnerConnecting,
-    RunnerIdle,
-    RunnerLoaded,
-    RunnerLoading,
-    RunnerReady,
-    RunnerRunning,
-    RunnerShutdown,
-    RunnerShuttingDown,
-    RunnerStatus,
-    RunnerWarmingUp,
-)
-from exo.utils.channels import MpReceiver, MpSender
-from exo.worker.engines.mlx.cache import KVPrefixCache
-from exo.worker.engines.mlx.utils_mlx import (
-    initialize_mlx,
-    load_mlx_items,
-)
-from exo.worker.engines.mlx.vision import VisionProcessor
-from exo.worker.runner.bootstrap import logger
-from exo.worker.runner.llm_inference.batch_generator import (
-    BatchGenerator,
-    InferenceGenerator,
-    SequentialGenerator,
-)
-
-from .batch_generator import Cancelled, Finished
-from .tool_parsers import make_mlx_parser
-
-
-class ExitCode(str, Enum):
-    AllTasksComplete = "AllTasksComplete"
-    Shutdown = "Shutdown"
-
-
-class Runner:
-    def __init__(
-        self,
-        bound_instance: BoundInstance,
-        event_sender: MpSender[Event],
-        task_receiver: MpReceiver[Task],
-        cancel_receiver: MpReceiver[TaskId],
-    ):
-        self.event_sender = event_sender
-        self.task_receiver = task_receiver
-        self.cancel_receiver = cancel_receiver
-        self.bound_instance = bound_instance
-
-        self.instance, self.runner_id, self.shard_metadata = (
-            self.bound_instance.instance,
-            self.bound_instance.bound_runner_id,
-            self.bound_instance.bound_shard,
-        )
-        self.model_id = self.shard_metadata.model_card.model_id
-        self.device_rank = self.shard_metadata.device_rank
-
-        logger.info("hello from the runner")
-        if getattr(self.shard_metadata, "immediate_exception", False):
-            raise Exception("Fake exception - runner failed to spin up.")
-        if timeout := getattr(self.shard_metadata, "should_timeout", 0):
-            time.sleep(timeout)
-
-        self.setup_start_time = time.time()
-
-        self.generator: Builder | InferenceGenerator = Builder(
-            self.model_id,
-            self.event_sender,
-            self.cancel_receiver,
-        )
-
-        self.seen: set[TaskId] = set()
-        self.active_tasks: dict[
-            TaskId,
-            TextGeneration,
-        ] = {}
-
-        logger.info("runner created")
-        self.update_status(RunnerIdle())
-
-    def update_status(self, status: RunnerStatus):
-        self.current_status = status
-        self.event_sender.send(
-            RunnerStatusUpdated(
-                runner_id=self.runner_id, runner_status=self.current_status
-            )
-        )
-
-    def send_task_status(self, task_id: TaskId, task_status: TaskStatus):
-        self.event_sender.send(
-            TaskStatusUpdated(task_id=task_id, task_status=task_status)
-        )
-
-    def acknowledge_task(self, task: Task):
-        self.event_sender.send(TaskAcknowledged(task_id=task.task_id))
-
-    def main(self):
-        with self.task_receiver:
-            for task in self.task_receiver:
-                if task.task_id in self.seen:
-                    logger.warning("repeat task - potential error")
-                    continue
-                self.seen.add(task.task_id)
-                self.handle_first_task(task)
-                if isinstance(self.current_status, RunnerShutdown):
-                    break
-
-    def handle_first_task(self, task: Task):
-        self.send_task_status(task.task_id, TaskStatus.Running)
-
-        match task:
-            case ConnectToGroup() if isinstance(self.current_status, RunnerIdle):
-                assert isinstance(self.generator, Builder)
-                logger.info("runner connecting")
-                self.update_status(RunnerConnecting())
-                self.acknowledge_task(task)
-
-                self.generator.group = initialize_mlx(self.bound_instance)
-
-                self.send_task_status(task.task_id, TaskStatus.Complete)
-                self.update_status(RunnerConnected())
-                logger.info("runner connected")
-
-            # we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
-            case LoadModel() if isinstance(self.generator, Builder) and (
-                (
-                    isinstance(self.current_status, RunnerConnected)
-                    and self.generator.group is not None
-                )
-                or (
-                    isinstance(self.current_status, RunnerIdle)
-                    and self.generator.group is None
-                )
-            ):
-                total_layers = (
-                    self.shard_metadata.end_layer - self.shard_metadata.start_layer
-                )
-                logger.info("runner loading")
-
-                self.update_status(
-                    RunnerLoading(layers_loaded=0, total_layers=total_layers)
-                )
-                self.acknowledge_task(task)
-
-                def on_layer_loaded(layers_loaded: int, total: int) -> None:
-                    self.update_status(
-                        RunnerLoading(layers_loaded=layers_loaded, total_layers=total)
-                    )
-
-                assert (
-                    ModelTask.TextGeneration in self.shard_metadata.model_card.tasks
-                ), f"Incorrect model task(s): {self.shard_metadata.model_card.tasks}"
-                (
-                    self.generator.inference_model,
-                    self.generator.tokenizer,
-                    self.generator.vision_processor,
-                ) = load_mlx_items(
-                    self.bound_instance,
-                    self.generator.group,
-                    on_layer_loaded=on_layer_loaded,
-                )
-
-                self.generator = self.generator.build()
-
-                self.send_task_status(task.task_id, TaskStatus.Complete)
-                self.update_status(RunnerLoaded())
-                logger.info("runner loaded")
-
-            case StartWarmup() if isinstance(self.current_status, RunnerLoaded):
-                assert isinstance(self.generator, InferenceGenerator)
-                logger.info("runner warming up")
-
-                self.update_status(RunnerWarmingUp())
-                self.acknowledge_task(task)
-
-                self.generator.warmup()
-
-                logger.info(
-                    f"runner initialized in {time.time() - self.setup_start_time} seconds"
-                )
-
-                self.send_task_status(task.task_id, TaskStatus.Complete)
-                self.update_status(RunnerReady())
-                logger.info("runner ready")
-
-            case TextGeneration() if isinstance(self.current_status, RunnerReady):
-                return_code = self.handle_generation_tasks(starting_task=task)
-                if return_code == ExitCode.Shutdown:
-                    return
-
-            case Shutdown():
-                self.shutdown(task)
-                return
-
-            case _:
-                raise ValueError(
-                    f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}"
-                )
-
-    def shutdown(self, task: Task):
-        logger.info("runner shutting down")
-        self.update_status(RunnerShuttingDown())
-        self.acknowledge_task(task)
-        if isinstance(self.generator, InferenceGenerator):
-            self.generator.close()
-        mx.clear_cache()
-        import gc
-
-        gc.collect()
-        self.send_task_status(task.task_id, TaskStatus.Complete)
-        self.update_status(RunnerShutdown())
-
-    def submit_text_generation(self, task: TextGeneration):
-        assert isinstance(self.generator, InferenceGenerator)
-        self.active_tasks[task.task_id] = task
-        self.generator.submit(task)
-
-    def handle_generation_tasks(self, starting_task: TextGeneration):
-        assert isinstance(self.current_status, RunnerReady)
-        assert isinstance(self.generator, InferenceGenerator)
-
-        logger.info(f"received chat request: {starting_task}")
-        self.update_status(RunnerRunning())
-        logger.info("runner running")
-        self.acknowledge_task(starting_task)
-        self.seen.add(starting_task.task_id)
-
-        self.submit_text_generation(starting_task)
-
-        while self.active_tasks:
-            results = self.generator.step()
-
-            finished: list[TaskId] = []
-            for task_id, result in results:
-                match result:
-                    case Cancelled():
-                        finished.append(task_id)
-                    case Finished():
-                        self.send_task_status(task_id, TaskStatus.Complete)
-                        finished.append(task_id)
-                    case _:
-                        self.send_response(
-                            result, self.active_tasks[task_id].command_id
-                        )
-
-            for task_id in finished:
-                self.active_tasks.pop(task_id, None)
-
-            try:
-                task = self.task_receiver.receive_nowait()
-
-                if task.task_id in self.seen:
-                    logger.warning("repeat task - potential error")
-                    continue
-                self.seen.add(task.task_id)
-
-                match task:
-                    case TextGeneration():
-                        self.acknowledge_task(task)
-                        self.submit_text_generation(task)
-                    case Shutdown():
-                        self.shutdown(task)
-                        return ExitCode.Shutdown
-                    case _:
-                        raise ValueError(
-                            f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}"
-                        )
-
-            except WouldBlock:
-                pass
-
-        self.update_status(RunnerReady())
-        logger.info("runner ready")
-
-        return ExitCode.AllTasksComplete
-
-    def send_response(
-        self,
-        response: GenerationResponse | ToolCallResponse,
-        command_id: CommandId,
-    ):
-        match response:
-            case GenerationResponse():
-                if self.device_rank == 0 and response.finish_reason == "error":
-                    self.event_sender.send(
-                        ChunkGenerated(
-                            command_id=command_id,
-                            chunk=ErrorChunk(
-                                error_message=response.text,
-                                model=self.model_id,
-                            ),
-                        )
-                    )
-
-                elif self.device_rank == 0:
-                    assert response.finish_reason not in (
-                        "error",
-                        "tool_calls",
-                        "function_call",
-                    )
-                    self.event_sender.send(
-                        ChunkGenerated(
-                            command_id=command_id,
-                            chunk=TokenChunk(
-                                model=self.model_id,
-                                text=response.text,
-                                token_id=response.token,
-                                usage=response.usage,
-                                finish_reason=response.finish_reason,
-                                stats=response.stats,
-                                logprob=response.logprob,
-                                top_logprobs=response.top_logprobs,
-                                is_thinking=response.is_thinking,
-                            ),
-                        )
-                    )
-            case ToolCallResponse():
-                if self.device_rank == 0:
-                    self.event_sender.send(
-                        ChunkGenerated(
-                            command_id=command_id,
-                            chunk=ToolCallChunk(
-                                tool_calls=response.tool_calls,
-                                model=self.model_id,
-                                usage=response.usage,
-                                stats=response.stats,
-                            ),
-                        )
-                    )
-
-
-@dataclass
-class Builder:
-    model_id: ModelId
-    event_sender: MpSender[Event]
-    cancel_receiver: MpReceiver[TaskId]
-    inference_model: Model | None = None
-    tokenizer: TokenizerWrapper | None = None
-    group: mx.distributed.Group | None = None
-    vision_processor: VisionProcessor | None = None
-
-    def build(
-        self,
-    ) -> InferenceGenerator:
-        assert self.model_id
-        assert self.inference_model
-        assert self.tokenizer
-
-        vision_processor = self.vision_processor
-
-        tool_parser = None
-        logger.info(
-            f"model has_tool_calling={self.tokenizer.has_tool_calling} using tokens {self.tokenizer.tool_call_start}, {self.tokenizer.tool_call_end}"
-        )
-        if (
-            self.tokenizer.tool_call_start
-            and self.tokenizer.tool_call_end
-            and self.tokenizer.tool_parser  # type: ignore
-        ):
-            tool_parser = make_mlx_parser(
-                self.tokenizer.tool_call_start,
-                self.tokenizer.tool_call_end,
-                self.tokenizer.tool_parser,  # type: ignore
-            )
-
-        kv_prefix_cache = KVPrefixCache(self.group)
-
-        device_rank = 0 if self.group is None else self.group.rank()
-        if os.environ.get("EXO_NO_BATCH"):
-            logger.info("using SequentialGenerator (batching disabled)")
-            return SequentialGenerator(
-                model=self.inference_model,
-                tokenizer=self.tokenizer,
-                group=self.group,
-                tool_parser=tool_parser,
-                kv_prefix_cache=kv_prefix_cache,
-                model_id=self.model_id,
-                device_rank=device_rank,
-                cancel_receiver=self.cancel_receiver,
-                event_sender=self.event_sender,
-                vision_processor=vision_processor,
-            )
-        logger.info("using BatchGenerator")
-        return BatchGenerator(
-            model=self.inference_model,
-            tokenizer=self.tokenizer,
-            group=self.group,
-            tool_parser=tool_parser,
-            kv_prefix_cache=kv_prefix_cache,
-            model_id=self.model_id,
-            device_rank=device_rank,
-            cancel_receiver=self.cancel_receiver,
-            event_sender=self.event_sender,
-            vision_processor=vision_processor,
-        )
@@ -0,0 +1,279 @@
+import time
+from enum import Enum
+
+from anyio import WouldBlock
+
+from exo.shared.types.chunks import Chunk
+from exo.shared.types.common import CommandId
+from exo.shared.types.events import (
+    ChunkGenerated,
+    Event,
+    RunnerStatusUpdated,
+    TaskAcknowledged,
+    TaskStatusUpdated,
+)
+from exo.shared.types.tasks import (
+    ConnectToGroup,
+    GenerationTask,
+    ImageEdits,
+    ImageGeneration,
+    LoadModel,
+    Shutdown,
+    StartWarmup,
+    Task,
+    TaskId,
+    TaskStatus,
+    TextGeneration,
+)
+from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.runner_response import (
+    CancelledResponse,
+    FinishedResponse,
+)
+from exo.shared.types.worker.runners import (
+    RunnerConnected,
+    RunnerConnecting,
+    RunnerIdle,
+    RunnerLoaded,
+    RunnerLoading,
+    RunnerReady,
+    RunnerRunning,
+    RunnerShutdown,
+    RunnerShuttingDown,
+    RunnerStatus,
+    RunnerWarmingUp,
+)
+from exo.utils.channels import MpReceiver, MpSender
+from exo.worker.engines.base import Builder, Engine
+
+from .bootstrap import logger
+
+
+class ExitCode(str, Enum):
+    AllTasksComplete = "AllTasksComplete"
+    Shutdown = "Shutdown"
+
+
+class Runner:
+    def __init__(
+        self,
+        bound_instance: BoundInstance,
+        builder: Builder,
+        event_sender: MpSender[Event],
+        task_receiver: MpReceiver[Task],
+    ):
+        self.event_sender = event_sender
+        self.task_receiver = task_receiver
+        self.bound_instance = bound_instance
+
+        self.instance, self.runner_id, self.shard_metadata = (
+            self.bound_instance.instance,
+            self.bound_instance.bound_runner_id,
+            self.bound_instance.bound_shard,
+        )
+        self.model_id = self.shard_metadata.model_card.model_id
+        self.device_rank = self.shard_metadata.device_rank
+
+        logger.info("hello from the runner")
+        if getattr(self.shard_metadata, "immediate_exception", False):
+            raise Exception("Fake exception - runner failed to spin up.")
+        if timeout := getattr(self.shard_metadata, "should_timeout", 0):
+            time.sleep(timeout)
+
+        self.setup_start_time = time.time()
+
+        self.generator: Builder | Engine = builder
+
+        self.seen: set[TaskId] = set()
+        self.active_tasks: dict[
+            TaskId,
+            GenerationTask,
+        ] = {}
+
+        logger.info("runner created")
+        self.update_status(RunnerIdle())
+
+    def update_status(self, status: RunnerStatus):
+        self.current_status = status
+        self.event_sender.send(
+            RunnerStatusUpdated(
+                runner_id=self.runner_id, runner_status=self.current_status
+            )
+        )
+
+    def send_task_status(self, task_id: TaskId, task_status: TaskStatus):
+        self.event_sender.send(
+            TaskStatusUpdated(task_id=task_id, task_status=task_status)
+        )
+
+    def acknowledge_task(self, task: Task):
+        self.event_sender.send(TaskAcknowledged(task_id=task.task_id))
+
+    def main(self):
+        with self.task_receiver:
+            for task in self.task_receiver:
+                if task.task_id in self.seen:
+                    logger.warning("repeat task - potential error")
+                    continue
+                self.seen.add(task.task_id)
+                self.handle_first_task(task)
+                if isinstance(self.current_status, RunnerShutdown):
+                    break
+
+    def handle_first_task(self, task: Task):
+        self.send_task_status(task.task_id, TaskStatus.Running)
+
+        match task:
+            case ConnectToGroup() if isinstance(self.current_status, RunnerIdle):
+                assert isinstance(self.generator, Builder)
+                logger.info("runner connecting")
+                self.update_status(RunnerConnecting())
+                self.acknowledge_task(task)
+
+                self.generator.connect(self.bound_instance)
+
+                self.send_task_status(task.task_id, TaskStatus.Complete)
+                self.update_status(RunnerConnected())
+                logger.info("runner connected")
+
+            # we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
+            case LoadModel() if isinstance(self.generator, Builder) and (
+                isinstance(self.current_status, (RunnerConnected, RunnerIdle))
+            ):
+                total_layers = (
+                    self.shard_metadata.end_layer - self.shard_metadata.start_layer
+                )
+                logger.info("runner loading")
+
+                self.update_status(
+                    RunnerLoading(layers_loaded=0, total_layers=total_layers)
+                )
+                self.acknowledge_task(task)
+
+                for load_progress in self.generator.load(self.bound_instance):
+                    self.update_status(
+                        RunnerLoading(
+                            layers_loaded=load_progress.layers_loaded,
+                            total_layers=load_progress.total,
+                        )
+                    )
+
+                self.generator = self.generator.build()
+
+                self.send_task_status(task.task_id, TaskStatus.Complete)
+                self.update_status(RunnerLoaded())
+                logger.info("runner loaded")
+
+            case StartWarmup() if isinstance(self.current_status, RunnerLoaded):
+                assert isinstance(self.generator, Engine)
+                logger.info("runner warming up")
+
+                self.update_status(RunnerWarmingUp())
+                self.acknowledge_task(task)
+
+                self.generator.warmup()
+
+                logger.info(
+                    f"runner initialized in {time.time() - self.setup_start_time} seconds"
+                )
+
+                self.send_task_status(task.task_id, TaskStatus.Complete)
+                self.update_status(RunnerReady())
+                logger.info("runner ready")
+
+            case TextGeneration() | ImageEdits() | ImageGeneration() if isinstance(
+                self.current_status, RunnerReady
+            ):
+                return_code = self.handle_generation_tasks(starting_task=task)
+                if return_code == ExitCode.Shutdown:
+                    return
+
+            case Shutdown():
+                self.shutdown(task)
+                return
+
+            case _:
+                raise ValueError(
+                    f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}"
+                )
+
+    def shutdown(self, task: Task):
+        logger.info("runner shutting down")
+        self.update_status(RunnerShuttingDown())
+        self.acknowledge_task(task)
+        self.generator.close()
+        import gc
+
+        gc.collect()
+        self.send_task_status(task.task_id, TaskStatus.Complete)
+        self.update_status(RunnerShutdown())
+
+    def submit_generation(self, task: GenerationTask):
+        assert isinstance(self.generator, Engine)
+        self.active_tasks[task.task_id] = task
+        self.generator.submit(task)
+
+    def handle_generation_tasks(self, starting_task: GenerationTask):
+        assert isinstance(self.current_status, RunnerReady)
+        assert isinstance(self.generator, Engine)
+
+        logger.info(f"received chat request: {starting_task}")
+        self.update_status(RunnerRunning())
+        logger.info("runner running")
+        self.acknowledge_task(starting_task)
+        self.seen.add(starting_task.task_id)
+
+        self.submit_generation(starting_task)
+
+        while self.active_tasks:
+            results = self.generator.step()
+
+            finished: list[TaskId] = []
+            for task_id, result in results:
+                match result:
+                    case CancelledResponse():
+                        finished.append(task_id)
+                    case FinishedResponse():
+                        self.send_task_status(task_id, TaskStatus.Complete)
+                        finished.append(task_id)
+                    case other:
+                        self.send_chunk(other, self.active_tasks[task_id].command_id)
+
+            for task_id in finished:
+                self.active_tasks.pop(task_id, None)
+
+            try:
+                task = self.task_receiver.receive_nowait()
+
+                if task.task_id in self.seen:
+                    logger.warning("repeat task - potential error")
+                    continue
+                self.seen.add(task.task_id)
+
+                match task:
+                    case TextGeneration() | ImageEdits() | ImageGeneration():
+                        self.acknowledge_task(task)
+                        self.submit_generation(task)
+                    case Shutdown():
+                        self.shutdown(task)
+                        return ExitCode.Shutdown
+                    case _:
+                        raise ValueError(
+                            f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}"
+                        )
+
+            except WouldBlock:
+                pass
+
+        self.update_status(RunnerReady())
+        logger.info("runner ready")
+
+        return ExitCode.AllTasksComplete
+
+    def send_chunk(
+        self,
+        chunk: Chunk,
+        command_id: CommandId,
+    ):
+        if self.device_rank == 0:
+            self.event_sender.send(ChunkGenerated(command_id=command_id, chunk=chunk))
@@ -96,7 +96,12 @@ def run_gpt_oss_pipeline_device(
            n_layers=24,
        )

-        model, tokenizer = shard_and_load(shard_meta, group, on_layer_loaded=None)
+        gen = shard_and_load(shard_meta, group)
+        try:
+            while True:
+                next(gen)
+        except StopIteration as stop:
+            model, tokenizer = stop.value
        model = cast(Model, model)

        # Generate a prompt of exact token length
@@ -172,7 +177,12 @@ def run_gpt_oss_tensor_parallel_device(
            n_layers=24,
        )

-        model, tokenizer = shard_and_load(shard_meta, group, on_layer_loaded=None)
+        gen = shard_and_load(shard_meta, group)
+        try:
+            while True:
+                next(gen)
+        except StopIteration as stop:
+            model, tokenizer = stop.value
        model = cast(Model, model)

        base_text = "The quick brown fox jumps over the lazy dog. "
@@ -343,7 +343,7 @@ class TestKVPrefixCacheWithModel:
            )

    def test_mlx_generate_populates_cache(self, model_and_tokenizer):
-        """mlx_generate should save the cache after generation completes."""
+        """mlx_generate should save the post-prefill cache (before the decode loop)."""
        model, tokenizer = model_and_tokenizer

        kv_prefix_cache = KVPrefixCache(None)
@@ -356,7 +356,6 @@ class TestKVPrefixCacheWithModel:
        prompt_tokens = encode_prompt(tokenizer, prompt)

        # Consume the entire generator so the cache-saving code after yield runs
-        generated_tokens = 0
        for _response in mlx_generate(
            model=model,
            tokenizer=tokenizer,
@@ -365,13 +364,14 @@ class TestKVPrefixCacheWithModel:
            kv_prefix_cache=kv_prefix_cache,
            group=None,
        ):
-            generated_tokens += 1
+            pass

        assert len(kv_prefix_cache.prompts) == 1
        assert len(kv_prefix_cache.caches) == 1
-        # Cache should contain prompt + generated tokens
-        expected_length = len(prompt_tokens) + generated_tokens
-        assert cache_length(kv_prefix_cache.caches[0]) == expected_length
+        # add_kv_cache is called before the decode loop and stores a deepcopy of
+        # the cache as it is just after prefill + trim(2). Generation tokens are
+        # never written into the stored entry.
+        assert cache_length(kv_prefix_cache.caches[0]) == len(prompt_tokens) - 2

    def test_mlx_generate_second_call_gets_prefix_hit(self, model_and_tokenizer):
        """Second mlx_generate call with same prompt should get a prefix hit from stored cache."""
@@ -174,7 +174,12 @@ def _run_pipeline_device(
            n_layers=TOTAL_LAYERS,
        )

-        model, tokenizer = shard_and_load(shard_meta, group, on_layer_loaded=None)
+        gen = shard_and_load(shard_meta, group)
+        try:
+            while True:
+                next(gen)
+        except StopIteration as stop:
+            model, tokenizer = stop.value
        model = cast(Any, model)

        prompt, task = _build_prompt(tokenizer, prompt_tokens)
@@ -14,6 +14,8 @@ import pytest
 from mlx.utils import tree_flatten, tree_unflatten
 from mlx_lm.tokenizer_utils import TokenizerWrapper

+from exo.download.download_utils import resolve_existing_model
+from exo.shared.constants import EXO_MODELS_DIRS, EXO_MODELS_READ_ONLY_DIRS
 from exo.shared.types.common import ModelId
 from exo.shared.types.mlx import Model
 from exo.shared.types.text_generation import (
@@ -28,8 +30,6 @@ from exo.worker.engines.mlx.utils_mlx import (
    load_tokenizer_for_model_id,
 )

-HF_CACHE = Path.home() / ".cache" / "huggingface" / "hub"
-
 # ── Config reduction ──────────────────────────────────────────────────────── #

 _REDUCE = {
@@ -100,12 +100,21 @@ def _reduce_config(cfg: dict[str, Any]) -> dict[str, Any]:


 def _find_snapshot(hub_name: str) -> Path | None:
-    model_dir = HF_CACHE / f"models--mlx-community--{hub_name}"
-    snaps = model_dir / "snapshots"
-    if not snaps.exists():
-        return None
-    children = sorted(snaps.iterdir())
-    return children[0] if children else None
+    """Locate a model directory under exo's models dirs.
+
+    Uses resolve_existing_model for fully-downloaded models; falls back to any
+    existing directory (even partial) so that tokenizer-only copies still work.
+    """
+    model_id = ModelId(f"mlx-community/{hub_name}")
+    found = resolve_existing_model(model_id)
+    if found is not None:
+        return found
+    normalized = model_id.normalize()
+    for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
+        candidate = search_dir / normalized
+        if candidate.is_dir():
+            return candidate
+    return None


 def _copy_tokenizer(src: Path, dst: Path) -> None:
@@ -192,13 +201,31 @@ ARCHITECTURES: list[ArchSpec] = [
 ]


+def _has_chat_template(model_dir: Path) -> bool:
+    """Check if a model dir has a usable chat template (inline or separate)."""
+    if (model_dir / "chat_template.jinja").exists():
+        return True
+    cfg = model_dir / "tokenizer_config.json"
+    if not cfg.exists():
+        return False
+    try:
+        data = cast(dict[str, Any], json.loads(cfg.read_text()))
+    except (OSError, json.JSONDecodeError):
+        return False
+    return bool(data.get("chat_template"))
+
+
 def _arch_available(spec: ArchSpec) -> bool:
    snap = _find_snapshot(spec.hub_name)
    if snap is None or not (snap / "config.json").exists():
        return False
+    tokenizer_snap = snap
    if spec.tokenizer_hub is not None:
-        return _find_snapshot(spec.tokenizer_hub) is not None
-    return True
+        alt = _find_snapshot(spec.tokenizer_hub)
+        if alt is None:
+            return False
+        tokenizer_snap = alt
+    return _has_chat_template(tokenizer_snap)


 def _make_task() -> TextGenerationTaskParams:
@@ -0,0 +1,395 @@
+# type: ignore
+"""uv run pytest -v -m "" src/exo/worker/tests/unittests/test_mlx/test_tp_bit_exact.py"""
+
+import importlib
+import json
+import multiprocessing as mp
+import os
+import sys
+import tempfile
+import traceback
+
+import numpy as np
+import pytest
+
+MODEL_CONFIGS = {
+    "llama": dict(
+        module="mlx_lm.models.llama",
+        args=dict(
+            model_type="llama",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=4,
+            rms_norm_eps=1e-6,
+            vocab_size=512,
+            max_position_embeddings=128,
+            head_dim=32,
+            rope_theta=10000.0,
+        ),
+    ),
+    "qwen3_5_moe": dict(
+        module="mlx_lm.models.qwen3_5_moe",
+        args=dict(
+            model_type="qwen3_5_moe",
+            text_config=dict(
+                model_type="qwen3_5_moe",
+                vocab_size=512,
+                hidden_size=512,
+                intermediate_size=1024,
+                num_hidden_layers=4,
+                num_attention_heads=16,
+                num_key_value_heads=4,
+                head_dim=32,
+                max_position_embeddings=128,
+                rms_norm_eps=1e-6,
+                tie_word_embeddings=False,
+                attention_bias=False,
+                full_attention_interval=2,
+                linear_num_value_heads=32,
+                linear_num_key_heads=16,
+                linear_key_head_dim=32,
+                linear_value_head_dim=32,
+                linear_conv_kernel_dim=4,
+                num_experts=16,
+                num_experts_per_tok=2,
+                decoder_sparse_step=1,
+                shared_expert_intermediate_size=256,
+                moe_intermediate_size=256,
+                norm_topk_prob=True,
+                rope_parameters={
+                    "type": "default",
+                    "rope_theta": 10000.0,
+                    "partial_rotary_factor": 0.25,
+                    "mrope_section": [11, 11, 10],
+                },
+            ),
+        ),
+    ),
+    "qwen3_next": dict(
+        module="mlx_lm.models.qwen3_next",
+        args=dict(
+            model_type="qwen3_next",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=4,
+            num_attention_heads=16,
+            num_key_value_heads=4,
+            head_dim=32,
+            max_position_embeddings=128,
+            rms_norm_eps=1e-6,
+            vocab_size=512,
+            attention_bias=False,
+            full_attention_interval=2,
+            linear_num_value_heads=32,
+            linear_num_key_heads=16,
+            linear_key_head_dim=32,
+            linear_value_head_dim=32,
+            linear_conv_kernel_dim=4,
+            num_experts=16,
+            num_experts_per_tok=2,
+            decoder_sparse_step=1,
+            shared_expert_intermediate_size=256,
+            moe_intermediate_size=256,
+            norm_topk_prob=True,
+            mlp_only_layers=[],
+            rope_theta=10000.0,
+            partial_rotary_factor=0.25,
+        ),
+    ),
+    "deepseek_v3": dict(
+        module="mlx_lm.models.deepseek_v3",
+        args=dict(
+            model_type="deepseek_v3",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=16,
+            vocab_size=512,
+            max_position_embeddings=128,
+            rms_norm_eps=1e-6,
+            n_routed_experts=8,
+            n_shared_experts=1,
+            num_experts_per_tok=2,
+            moe_intermediate_size=256,
+            moe_layer_freq=1,
+            first_k_dense_replace=0,
+            n_group=1,
+            topk_group=1,
+            routed_scaling_factor=1.0,
+            q_lora_rank=None,
+            kv_lora_rank=16,
+            qk_nope_head_dim=16,
+            qk_rope_head_dim=16,
+            v_head_dim=32,
+            rope_theta=10000.0,
+            rope_scaling={},
+            attention_bias=False,
+            norm_topk_prob=True,
+            scoring_func="sigmoid",
+            topk_method="noaux_tc",
+        ),
+    ),
+    "deepseek_v3_q4": dict(
+        module="mlx_lm.models.deepseek_v3",
+        quantize=dict(group_size=32, bits=4, mode="affine"),
+        args=dict(
+            model_type="deepseek_v3",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=16,
+            vocab_size=512,
+            max_position_embeddings=128,
+            rms_norm_eps=1e-6,
+            n_routed_experts=8,
+            n_shared_experts=1,
+            num_experts_per_tok=2,
+            moe_intermediate_size=256,
+            moe_layer_freq=1,
+            first_k_dense_replace=0,
+            n_group=1,
+            topk_group=1,
+            routed_scaling_factor=1.0,
+            q_lora_rank=None,
+            kv_lora_rank=64,
+            qk_nope_head_dim=32,
+            qk_rope_head_dim=32,
+            v_head_dim=32,
+            rope_theta=10000.0,
+            rope_scaling={},
+            attention_bias=False,
+            norm_topk_prob=True,
+            scoring_func="sigmoid",
+            topk_method="noaux_tc",
+        ),
+    ),
+    "glm4_moe_lite": dict(
+        module="mlx_lm.models.glm4_moe_lite",
+        args=dict(
+            model_type="glm4_moe_lite",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=16,
+            vocab_size=512,
+            max_position_embeddings=128,
+            rms_norm_eps=1e-6,
+            n_routed_experts=8,
+            n_shared_experts=1,
+            num_experts_per_tok=2,
+            moe_intermediate_size=256,
+            first_k_dense_replace=1,
+            n_group=1,
+            topk_group=1,
+            routed_scaling_factor=1.0,
+            rope_theta=10000.0,
+            attention_bias=False,
+            q_lora_rank=None,
+            kv_lora_rank=16,
+            qk_rope_head_dim=16,
+            qk_nope_head_dim=16,
+            v_head_dim=32,
+        ),
+    ),
+    "minimax": dict(
+        module="mlx_lm.models.minimax",
+        args=dict(
+            model_type="minimax",
+            hidden_size=512,
+            intermediate_size=1024,
+            num_attention_heads=16,
+            num_key_value_heads=4,
+            max_position_embeddings=128,
+            num_experts_per_tok=2,
+            num_local_experts=8,
+            shared_intermediate_size=256,
+            num_hidden_layers=2,
+            rms_norm_eps=1e-6,
+            rope_theta=10000.0,
+            rotary_dim=32,
+            vocab_size=512,
+        ),
+    ),
+    "gpt_oss": dict(
+        module="mlx_lm.models.gpt_oss",
+        args=dict(
+            model_type="gpt_oss",
+            hidden_size=512,
+            intermediate_size=256,
+            num_hidden_layers=2,
+            num_attention_heads=16,
+            num_key_value_heads=4,
+            vocab_size=512,
+            head_dim=32,
+            rms_norm_eps=1e-6,
+            num_local_experts=8,
+            num_experts_per_tok=2,
+            layer_types=["sliding_attention", "full_attention"],
+            sliding_window=64,
+            rope_theta=10000.0,
+        ),
+    ),
+    "gemma4": dict(
+        module="mlx_lm.models.gemma4",
+        args=dict(
+            model_type="gemma4",
+            vocab_size=512,
+            text_config=dict(
+                vocab_size=512,
+                hidden_size=512,
+                intermediate_size=1024,
+                num_hidden_layers=4,
+                num_attention_heads=16,
+                num_key_value_heads=4,
+                head_dim=32,
+                global_head_dim=32,
+                num_kv_shared_layers=0,
+                vocab_size_per_layer_input=512,
+                hidden_size_per_layer_input=512,
+                rms_norm_eps=1e-6,
+                max_position_embeddings=128,
+                sliding_window=64,
+                sliding_window_pattern=2,
+                layer_types=[
+                    "sliding_attention",
+                    "full_attention",
+                    "sliding_attention",
+                    "full_attention",
+                ],
+                enable_moe_block=True,
+                num_experts=8,
+                top_k_experts=2,
+                moe_intermediate_size=256,
+            ),
+        ),
+    ),
+}
+
+_PROMPT = [[1, 23, 45, 67, 89, 12, 34, 56]]
+
+
+def _build(name):
+    import mlx.core as mx
+    import mlx.nn as nn
+    from mlx.utils import tree_map_with_path
+
+    import exo.worker.engines.mlx.auto_parallel  # noqa: F401
+
+    cfg = MODEL_CONFIGS[name]
+    module = importlib.import_module(cfg["module"])
+    model_cls = module.Model
+    model_args_cls = module.ModelArgs
+
+    mx.random.seed(0)
+    args = model_args_cls(**cfg["args"])
+    m = model_cls(args)
+
+    def _to_bf16(_p, v):
+        if hasattr(v, "dtype") and v.dtype in (mx.float16, mx.float32, mx.bfloat16):
+            return v.astype(mx.bfloat16)
+        return v
+
+    m.update(tree_map_with_path(_to_bf16, m.parameters()))
+    if "quantize" in cfg:
+        nn.quantize(m, **cfg["quantize"])
+    mx.eval(m.parameters())
+    return mx, m
+
+
+def _run(name, out_path, shard):
+    import mlx.core as mx
+
+    if shard:
+        g = mx.distributed.init(backend="ring", strict=True)
+    mx_, m = _build(name)
+    if shard:
+        from exo.worker.engines.mlx.auto_parallel import tensor_auto_parallel
+
+        m = tensor_auto_parallel(m, g, on_layer_loaded=None)
+        mx_.eval(m.parameters())
+    inputs = mx_.array(_PROMPT, dtype=mx_.int32)
+    logits = m(inputs)
+    mx_.eval(logits)
+    np.savez(out_path, logits=np.asarray(logits.astype(mx_.float32)))
+
+
+def _ref_worker(name, out_path, q):
+    try:
+        _run(name, out_path, shard=False)
+        q.put(True)
+    except BaseException as e:
+        q.put(f"{e}\n{traceback.format_exc()}")
+
+
+def _tp_worker(name, rank, hf, out_path, q):
+    os.environ["MLX_HOSTFILE"] = hf
+    os.environ["MLX_RANK"] = str(rank)
+    try:
+        path = out_path if rank == 0 else out_path + f".r{rank}"
+        _run(name, path, shard=True)
+        q.put((rank, True, None))
+    except BaseException as e:
+        q.put((rank, False, f"{e}\n{traceback.format_exc()}"))
+
+
+def _run_compare(name, world_size, port_base):
+    d = tempfile.mkdtemp()
+    ref_path = f"{d}/ref.npz"
+    tp_path = f"{d}/tp.npz"
+    ctx = mp.get_context("spawn")
+    q = ctx.Queue()
+
+    p = ctx.Process(target=_ref_worker, args=(name, ref_path, q))
+    p.start()
+    p.join(300)
+    r = q.get(timeout=10)
+    if r is not True:
+        pytest.fail(f"[{name}] ref FAIL: {str(r)[:500]}")
+
+    hosts = [f"127.0.0.1:{port_base + i}" for i in range(world_size)]
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(hosts, f)
+        hf = f.name
+    ps = [
+        ctx.Process(target=_tp_worker, args=(name, rank, hf, tp_path, q))
+        for rank in range(world_size)
+    ]
+    for pp in ps:
+        pp.start()
+    results = [q.get(timeout=300) for _ in range(world_size)]
+    for pp in ps:
+        pp.join(60)
+    for rank, ok, payload in results:
+        if not ok:
+            pytest.fail(f"[{name}] rank {rank} FAIL: {payload[:500]}")
+
+    ref = np.load(ref_path)["logits"]
+    tp = np.load(tp_path)["logits"]
+    diff = np.abs(ref - tp)
+    max_diff = float(diff.max())
+    mean_diff = float(diff.mean())
+    assert max_diff == 0.0, (
+        f"[{name} TP={world_size}] not bit-exact: max={max_diff} mean={mean_diff}"
+    )
+
+
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.skipif(
+        sys.platform != "darwin", reason="MLX distributed requires Metal"
+    ),
+]
+
+
+@pytest.mark.skip("TP=2 is currently very different to TP=1. This test will not pass")
+@pytest.mark.parametrize("world_size", [2, 4])
+@pytest.mark.parametrize("name", list(MODEL_CONFIGS))
+def test_tp_bit_exact(name, world_size):
+    name_idx = list(MODEL_CONFIGS).index(name)
+    port = 32000 + name_idx * 20 + world_size
+    _run_compare(name, world_size, port)
@@ -54,6 +54,7 @@ def test_plan_requests_download_when_waiting_and_shard_not_downloaded():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -109,6 +110,7 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -154,6 +156,7 @@ def test_plan_does_not_request_download_when_shard_already_downloaded():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -204,6 +207,7 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -227,6 +231,7 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -54,6 +54,7 @@ def test_plan_kills_runner_when_instance_missing():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -96,6 +97,7 @@ def test_plan_kills_runner_when_sibling_failed():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -130,6 +132,7 @@ def test_plan_creates_runner_when_missing_for_node():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -171,6 +174,7 @@ def test_plan_does_not_create_runner_when_supervisor_already_present():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -203,6 +207,7 @@ def test_plan_does_not_create_runner_for_unassigned_node():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -78,6 +78,7 @@ def test_plan_forwards_pending_chat_completion_when_runner_ready():
        all_runners=all_runners,
        tasks={TASK_1_ID: task},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -131,6 +132,7 @@ def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
        all_runners=all_runners,
        tasks={TASK_1_ID: task},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -181,6 +183,7 @@ def test_plan_does_not_forward_tasks_for_other_instances():
        all_runners=all_runners,
        tasks={foreign_task.task_id: foreign_task},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -249,6 +252,7 @@ def test_plan_ignores_non_pending_or_non_chat_tasks():
        all_runners=all_runners,
        tasks={TASK_1_ID: completed_task, other_task_id: other_task},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -291,6 +295,7 @@ def test_plan_returns_none_when_nothing_to_do():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -63,6 +63,7 @@ def test_plan_starts_warmup_for_accepting_rank_when_all_loaded_or_warming():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -107,6 +108,7 @@ def test_plan_starts_warmup_for_rank_zero_after_others_warming():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -150,6 +152,7 @@ def test_plan_does_not_start_warmup_for_non_zero_rank_until_all_loaded_or_warmin
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -197,6 +200,7 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -216,6 +220,7 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -262,6 +267,7 @@ def test_plan_starts_warmup_for_connecting_rank_after_others_warming():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -307,6 +313,7 @@ def test_plan_does_not_start_warmup_for_accepting_rank_until_all_loaded_or_warmi
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -351,6 +358,7 @@ def test_plan_does_not_start_warmup_for_connecting_rank_until_others_warming():
        all_runners=all_runners,
        tasks={},
        input_chunk_buffer={},
+        image_cache={},
        instance_backoff=KeyedBackoff(),
        download_backoff=KeyedBackoff(),
    )
@@ -20,7 +20,25 @@ from exo.worker.engines.mlx.dsml_encoding import (
    encode_messages,
    parse_dsml_output,
 )
-from exo.worker.runner.llm_inference.model_output_parsers import parse_deepseek_v32
+from exo.worker.runner.llm_inference.model_output_parsers import (
+    parse_deepseek_v32,
+    parse_thinking_models,
+)
+
+
+def _parse_deepseek_with_thinking(
+    source: Generator[GenerationResponse | None],
+    starts_in_thinking: bool = False,
+) -> Generator[GenerationResponse | ToolCallResponse | None]:
+    return parse_deepseek_v32(
+        parse_thinking_models(
+            source,
+            think_start=THINKING_START,
+            think_end=THINKING_END,
+            starts_in_thinking=starts_in_thinking,
+        )
+    )
+

 # ── Shared fixtures ──────────────────────────────────────────────

@@ -333,9 +351,7 @@ class TestE2EThinkingAndToolCall:
        assert prompt.endswith(THINKING_START)

        # Simulate: model outputs <think>, thinks, closes thinking, then tool call.
-        # In the full pipeline, parse_thinking_models handles the case where
-        # <think> is in the prompt. Here we test parse_deepseek_v32 directly,
-        # which detects <think>/<think> markers in the stream.
+        # Use the full production chain (parse_thinking_models → parse_deepseek_v32).
        model_tokens = [
            THINKING_START,
            "The user wants weather",
@@ -353,7 +369,7 @@ class TestE2EThinkingAndToolCall:
            TOOL_CALLS_END,
        ]

-        results = list(parse_deepseek_v32(_simulate_tokens(model_tokens)))
+        results = list(_parse_deepseek_with_thinking(_simulate_tokens(model_tokens)))

        gen_results = [r for r in results if isinstance(r, GenerationResponse)]
        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
@@ -387,7 +403,7 @@ class TestE2EThinkingAndToolCall:
        prompt_no_think = encode_messages(
            messages, tools=_WEATHER_TOOLS, thinking_mode="chat"
        )
-        assert prompt_no_think.endswith(THINKING_END)
+        assert not prompt_no_think.endswith(THINKING_START)

        # Both should have the same tool definitions
        assert "get_weather" in prompt_think
@@ -597,7 +613,9 @@ class TestE2EFullRoundTrip:
            f"</{DSML_TOKEN}invoke>\n",
            TOOL_CALLS_END,
        ]
-        results_1 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_1)))
+        results_1 = list(
+            _parse_deepseek_with_thinking(_simulate_tokens(model_tokens_1))
+        )

        # Verify: thinking tokens + tool call
        gen_1 = [r for r in results_1 if isinstance(r, GenerationResponse)]
@@ -660,7 +678,9 @@ class TestE2EFullRoundTrip:
            THINKING_END,
            "The weather in Hangzhou is currently cloudy with temperatures between 7°C and 13°C.",
        ]
-        results_2 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_2)))
+        results_2 = list(
+            _parse_deepseek_with_thinking(_simulate_tokens(model_tokens_2))
+        )

        gen_2 = [r for r in results_2 if isinstance(r, GenerationResponse)]
        tool_2 = [r for r in results_2 if isinstance(r, ToolCallResponse)]
@@ -1,14 +1,13 @@
 # Check tasks are complete before runner is ever ready.
-import unittest.mock
 from collections.abc import Iterable
+from dataclasses import dataclass
 from typing import Callable

-import mlx.core as mx
 import pytest

+import exo.worker.engines.mlx.builder as mlx_builder
 import exo.worker.runner.llm_inference.batch_generator as mlx_batch_generator
 import exo.worker.runner.llm_inference.model_output_parsers as mlx_model_output_parsers
-import exo.worker.runner.llm_inference.runner as mlx_runner
 from exo.shared.types.chunks import TokenChunk
 from exo.shared.types.events import (
    ChunkGenerated,
@@ -46,6 +45,8 @@ from exo.shared.types.worker.runners import (
    RunnerWarmingUp,
 )
 from exo.utils.channels import mp_channel
+from exo.worker.engines.mlx.builder import MlxBuilder
+from exo.worker.runner.runner import Runner

 from ...constants import (
    CHAT_COMPLETION_TASK_ID,
@@ -115,13 +116,22 @@ def assert_events_equal(test_events: Iterable[Event], true_events: Iterable[Even
        assert test_event == true_event, f"{test_event} != {true_event}"


+@dataclass
+class MockLoadOutput:
+    layers_loaded: int
+    total: int
+
+
@pytest.fixture
 def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
    # initialize_mlx returns a mock group
-    monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MockGroup()))
-    monkeypatch.setattr(
-        mlx_runner, "load_mlx_items", make_nothin((1, MockTokenizer, None))
-    )
+    monkeypatch.setattr(mlx_builder, "initialize_mlx", make_nothin(MockGroup()))
+
+    def lmi_gen():
+        yield MockLoadOutput(1, 1)
+        return (1, MockTokenizer, None)
+
+    monkeypatch.setattr(mlx_builder, "load_mlx_items", make_nothin(lmi_gen()))
    monkeypatch.setattr(mlx_batch_generator, "warmup_inference", make_nothin(1))
    monkeypatch.setattr(mlx_batch_generator, "_check_for_debug_prompts", nothin)
    monkeypatch.setattr(mlx_batch_generator, "mx_any", make_nothin(False))
@@ -264,17 +274,18 @@ def _run(tasks: Iterable[Task], send_after_ready: list[Task] | None = None):
        # this is some c++ nonsense
        task_receiver.close = nothin
        task_receiver.join = nothin
-        with unittest.mock.patch(
-            "exo.worker.runner.llm_inference.runner.mx.distributed.all_gather",
-            make_nothin(mx.array([1])),
-        ):
-            runner = mlx_runner.Runner(
-                bound_instance,
-                event_sender,  # pyright: ignore[reportArgumentType]
-                task_receiver,
-                cancel_receiver,
-            )
-            runner.main()
+        builder = MlxBuilder(
+            bound_instance.bound_shard.model_card.model_id,
+            event_sender,  # pyright: ignore[reportArgumentType]
+            cancel_receiver,
+        )
+        runner = Runner(
+            bound_instance,
+            builder,
+            event_sender,  # pyright: ignore[reportArgumentType]
+            task_receiver,
+        )
+        runner.main()

        return event_sender.events

@@ -318,6 +329,10 @@ def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
                runner_status=RunnerLoading(layers_loaded=0, total_layers=32),
            ),
            TaskAcknowledged(task_id=LOAD_TASK_ID),
+            RunnerStatusUpdated(
+                runner_id=RUNNER_1_ID,
+                runner_status=RunnerLoading(layers_loaded=1, total_layers=1),
+            ),
            TaskStatusUpdated(task_id=LOAD_TASK_ID, task_status=TaskStatus.Complete),
            RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerLoaded()),
            TaskStatusUpdated(task_id=WARMUP_TASK_ID, task_status=TaskStatus.Running),
@@ -380,6 +380,110 @@ class TestGenericToolCallsFinishReason:
 # ── Double parser chain (parse_thinking_models → parse_deepseek_v32) ──


+class TestDeepSeekV32StartsInThinking:
+    """Regression tests for deepseek v3.2 where the chat template appends
+    <think> to the prompt so the model starts already inside a thinking block.
+    """
+
+    def test_reasoning_tagged_when_starts_in_thinking(self):
+        tokens = [
+            _make_response("let me", 0),
+            _make_response(" think", 1),
+            _make_response(THINKING_END, 2),
+            _make_response("\n", 3),
+            _make_response("42", 4, finish_reason="stop"),
+        ]
+        thinking = parse_thinking_models(
+            _queue_source(tokens),
+            think_start=THINKING_START,
+            think_end=THINKING_END,
+            starts_in_thinking=True,
+        )
+        results = _step_until_finish(parse_deepseek_v32(thinking))
+        gens = [
+            r
+            for r in results
+            if isinstance(r, GenerationResponse) and r.finish_reason is None
+        ]
+        texts = [(r.text, r.is_thinking) for r in gens]
+        assert texts == [("let me", True), (" think", True), ("\n", False)]
+        final = [
+            r
+            for r in results
+            if isinstance(r, GenerationResponse) and r.finish_reason is not None
+        ]
+        assert len(final) == 1
+        assert final[0].text == "42"
+        assert final[0].is_thinking is False
+
+    def test_starts_in_thinking_then_tool_call(self):
+        tokens = [
+            _make_response("need weather", 0),
+            _make_response(THINKING_END, 1),
+            _make_response("\n\n", 2),
+            _make_response(TOOL_CALLS_START, 3),
+            _make_response("\n", 4),
+            _make_response(f'<{DSML_TOKEN}invoke name="get_weather">\n', 5),
+            _make_response(
+                f'<{DSML_TOKEN}parameter name="city" string="true">NYC</{DSML_TOKEN}parameter>\n',
+                6,
+            ),
+            _make_response(f"</{DSML_TOKEN}invoke>\n", 7),
+            _make_response(TOOL_CALLS_END, 8, finish_reason="stop"),
+        ]
+        thinking = parse_thinking_models(
+            _queue_source(tokens),
+            think_start=THINKING_START,
+            think_end=THINKING_END,
+            starts_in_thinking=True,
+        )
+        results = _step_until_finish(parse_deepseek_v32(thinking))
+        reasoning_gens = [
+            r
+            for r in results
+            if isinstance(r, GenerationResponse)
+            and r.finish_reason is None
+            and r.is_thinking
+        ]
+        assert [r.text for r in reasoning_gens] == ["need weather"]
+        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+        assert len(tool_results) == 1
+        assert tool_results[0].tool_calls[0].name == "get_weather"
+
+    def test_reasoning_tokens_counted_starts_in_thinking(self):
+        usage = Usage(
+            prompt_tokens=10,
+            completion_tokens=5,
+            total_tokens=15,
+            prompt_tokens_details=PromptTokensDetails(cached_tokens=0),
+            completion_tokens_details=CompletionTokensDetails(reasoning_tokens=0),
+        )
+        tokens = [
+            _make_response("reasoning", 0),
+            _make_response(" more", 1),
+            _make_response(THINKING_END, 2),
+            _make_response("\n", 3),
+            GenerationResponse(text="42", token=4, finish_reason="stop", usage=usage),
+        ]
+        thinking = parse_thinking_models(
+            _queue_source(tokens),
+            think_start=THINKING_START,
+            think_end=THINKING_END,
+            starts_in_thinking=True,
+        )
+        results = _step_until_finish(
+            count_reasoning_tokens(parse_deepseek_v32(thinking))
+        )
+        final = [
+            r
+            for r in results
+            if isinstance(r, GenerationResponse) and r.finish_reason is not None
+        ]
+        assert len(final) == 1
+        assert final[0].usage is not None
+        assert final[0].usage.completion_tokens_details.reasoning_tokens == 2
+
+
 class TestBatchGeneratorSingleNext:
    def test_finish_reason_with_buffered_tokens_drain_loop(self):
        from exo.worker.runner.llm_inference.batch_generator import GeneratorQueue
@@ -17,7 +17,7 @@ from exo.shared.types.text_generation import (
 from exo.shared.types.worker.instances import BoundInstance, InstanceId
 from exo.shared.types.worker.runners import RunnerFailed, RunnerId
 from exo.utils.channels import channel, mp_channel
-from exo.worker.runner.runner_supervisor import RunnerSupervisor
+from exo.worker.runner.supervisor import RunnerSupervisor
 from exo.worker.tests.unittests.conftest import get_bound_mlx_ring_instance
Author	SHA1	Message	Date
Evan	4216ca541a	implement engine interface for mlx and mflux	2026-04-25 01:52:29 +01:00
rltakashige	fd707de30b	Add more model cards (#1970 )	2026-04-23 15:28:40 +01:00
Alex CheemaandClaude Opus 4.7	45248c5c85	chore(app): hardcode bug report presigned-URL endpoint (#1971 ) ## Motivation The bug-report presigned-URL endpoint (`https://reports.exolabs.net/presigned-urls`) was injected at build time from the `EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT` GitHub Actions secret into `Info.plist`, then read at runtime by `BugReportService`. It isn't actually a secret — the POST body is just `{"keys":[...]}` with no credential (see `app/EXO/EXO/Services/BugReportService.swift:136-142`), abuse prevention lives server-side on the lambda, and the URL is already visible in every publicly-distributed DMG's `Info.plist`. Treating it as a repo secret added plumbing with no security benefit and broke local dev builds — hitting Send Bug Report on an uncustomised `just build-app` raised "Bug report endpoint is invalid". ## Changes - `app/EXO/EXO/Info.plist`: replace `$(EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT)` with the literal URL. - `.github/workflows/build-app.yml`: drop the `EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT` job-level env var and the xcodebuild build-setting passthrough. No other workflow changes. Swift code is unchanged — `BugReportService` still reads from `Info.plist`, which leaves an escape hatch if anyone ever needs to override via `xcodebuild EXOBugReportPresignedUrlEndpoint=...` without recompiling. Follow-up: the `EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT` repo secret can now be deleted in the GitHub Actions settings UI. ## Why It Works `Info.plist` variable substitution turns `$(FOO)` into whatever build setting `FOO` resolves to. CI was setting `FOO` via xcodebuild; local dev wasn't, so the key resolved to an empty string, which `BugReportService.fetchPresignedUploadUrls` rejects via the `!trimmedEndpointString.isEmpty` guard at `BugReportService.swift:131`. Hardcoding the literal string removes the substitution entirely, so every build — local or CI — gets the right value. ## Test Plan ### Manual Testing <!-- Hardware: MacBook Pro (macOS app build via Xcode) --> - `just build-app` with no extra env vars (reproduces the failure path on `main`). - `/usr/libexec/PlistBuddy -c "Print :EXOBugReportPresignedUrlEndpoint" app/EXO/build/Build/Products/Debug/EXO.app/Contents/Info.plist` → returns `https://reports.exolabs.net/presigned-urls` (was empty before this change). - `open app/EXO/build/Build/Products/Debug/EXO.app` → menubar → Debug Info → Send Bug Report → type a description → Send → upload succeeds and the Create GitHub Issue button appears (was failing with "Bug report endpoint is invalid" before). - Cross-check on the Slack side that the uploaded `report.json` lands under `reports/YYYY/MM/DD/<ts>/` as before. ### Automated Testing <!-- Describe changes to automated tests, or how existing tests cover this change --> - No new tests. This is a single-string change to `Info.plist` plus a workflow cleanup. `nix flake check` in CI verifies formatting/lint for the rest of the tree. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>	2026-04-23 14:08:16 +00:00
rltakashige	290e3fd927	Keep image cache fresh (temporary fix) (#1961 ) ## Motivation When a new node joins, it might not have the cache. Caveat: This is potentially fallible if a new node joins and updates real topology, but the API topology hasn't caught up with this fact and the user queues up a new text generation. In practice, there is only a split second where this is the case, and this is only for users of the dashboard interface. We should fix this properly after the release.	2026-04-23 11:39:36 +01:00
rltakashige	3894cf134e	Fix Gemma 4 E2B TP + DeepSeek V32 thinking parsing (#1967 )	2026-04-23 01:50:39 +00:00
Alex CheemaandClaude Opus 4.7	8993ccaf09	feat(app): add friendly context message to bug report prompt (#1959 ) ## Motivation When a user clicks Send Bug Report in the macOS app, we already give them the option to add more context via an optional text field. But the current prompt is just a terse label — `"What's the issue? (optional)"` — which doesn't tell the user why bothering to fill it in matters. A friendly one-line explanation increases the chance they'll describe what went wrong, which is the single most useful signal when we triage the resulting diagnostic bundle. ## Changes - `app/EXO/EXO/ContentView.swift`: In the `.prompting` phase of `sendBugReportButton`, replace the single label with a two-line hierarchy: - Primary: `Tell us what went wrong (optional)` - Helper: `A quick description of what you were doing and what happened helps us track down the bug for you.` - The helper uses `.caption2` + `.secondary` + `.opacity(0.8)` + `.fixedSize(horizontal: false, vertical: true)` so it stays visually subordinate and wraps cleanly inside the 340pt popover. No changes to `BugReportService`, the `user_description` payload, or any other flow. ## Why It Works The optional description is already plumbed end-to-end (text editor → `bugReportUserDescription` state → `BugReportService.sendReport(..., userDescription:)` → `report.json`'s `user_description` field → GitHub issue pre-fill). The only gap was user-facing motivation, so this is purely a copy/layout tweak inside the existing `.prompting` case — no new state, bindings, or service changes. ## Test Plan ### Manual Testing <!-- Hardware: MacBook Pro (macOS app build via Xcode) --> - Build the macOS app in Xcode (`app/EXO/EXO.xcodeproj`) and launch it. - Open the menubar popover → expand Debug Info → click Send Bug Report. - Verify the new primary label and helper sentence both appear above the text editor and wrap cleanly within the popover width. - Leave the field empty → click Send → upload should succeed (no `user_description` in payload, same as before). - Fill in a description → click Send → upload succeeds and the success card with Create GitHub Issue appears; clicking it opens GitHub with the description pre-filled. - Click Cancel from the prompting state → returns to idle. ### Automated Testing <!-- Describe changes to automated tests, or how existing tests cover this change --> - No new automated tests. This is a SwiftUI copy/layout change; existing `EXOTests` are smoke-level and don't cover `ContentView` view bodies, and UI snapshot tests aren't worth adding for a two-line copy tweak. - `nix fmt` reports 0 files changed after the edit; `nix flake check` in CI will verify formatting/lint for the rest of the tree. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>	2026-04-22 17:39:36 +00:00
Nadeem Hilal Wani	4939fbe995	feat(dashboard): add Pi integration tab (#1925 ) ## Summary Adds a new Pi tab to the Integrations page (`/#/integrations`) alongside the existing Claude Code, OpenCode, Codex, OpenClaw, Open WebUI, n8n, and Firefox tabs. [pi](https://pi.dev) (`@mariozechner/pi-coding-agent`) is a terminal coding agent that supports custom OpenAI-compatible providers via `~/.pi/agent/models.json`. This tab gives users a copy-pasteable config to wire pi up to their exo cluster. ## What's in the tab - Model selector (shown when multiple models are running) — picks the default model for the generated shell command. - Models Config card — generates `~/.pi/agent/models.json` registering `exo` as a custom provider: - `baseUrl` → `<apiUrl>/v1` - `api` → `openai-completions` - `apiKey` → `"exo"` (placeholder; exo ignores it) - `compat.supportsDeveloperRole: false` and `compat.supportsReasoningEffort: false`, per pi docs recommendation for local OpenAI-compatible servers - Auto-populates every running model with `id`, `contextWindow` (from `/v1/models`), and `input: ["text", "image"]` for vision-capable models - Shell Command card — `pi --provider exo --model <model>` for quick launch. The tab gracefully falls back to `your-model-id` when no models are running, matching the behavior of the other tabs. ## Usage 1. `npm install -g @mariozechner/pi-coding-agent` 2. Paste the generated config into `~/.pi/agent/models.json` 3. Run `pi` and pick an exo model via `/model` — or run the shell command directly ## Changes - `dashboard/src/routes/integrations/+page.svelte` — adds `"Pi"` to the `tabs` tuple, `piModel` state, `piModelsJson` + `piShellCommand` derivations, and the tab content block. Single-file, scoped change — no backend or type changes. ## Testing - `cd dashboard && npm run build` — ✅ builds cleanly - `svelte-check` on the edited file — no new errors - Manually verified the tab renders, the model selector updates the generated JSON, and the config reflects `/v1/models` capabilities (vision → `input: ["text","image"]`, `context_length` → `contextWindow`). ## Screenshots <img width="1545" height="1236" alt="pi-tab" src="https://github.com/user-attachments/assets/38aa179f-4ed9-4a1e-9783-d3baa7738263" />	2026-04-22 17:29:47 +00:00
rltakashige	73782ecc65	Fix event mutation causing indexed vs event mismatch (#1964 ) Fixes small issue with #1957	2026-04-22 16:12:24 +00:00
rltakashige	f6e418ed23	Cleanup on #1952 (#1960 )	2026-04-22 17:05:49 +01:00
rltakashigeandEvan	7a312a177b	Misc fixes: upstream JACCL all_sum, API, etc. + Add Kimi K2.6 (#1952 ) ## Motivation This fixes a bunch of observed model quality issues introduced upstream in JACCL, as well as API issues and prefix cache calculation. ## Test Plan ### Manual Testing Tested a bunch ### Automated Testing Added a test, automated eval tool calls on Kimi K2.6, Minimax M2.7, GPT OSS and Qwen3.6 models. --------- Co-authored-by: Evan <evanev7@gmail.com>	2026-04-22 15:43:27 +00:00
Evan Quiney	0a549f8846	remove layer loading callback (#1890 ) first part of modularising the backend is simplifying some of the control flow. more tbd.	2026-04-22 14:03:31 +01:00