implement engine interface for mlx and mflux

## 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>

.github/workflows/build-app.yml

@@ -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

.mlx_typings/mlx/core/__init__.pyi

@@ -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:
     """

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

@@ -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."""

.mlx_typings/mlx_lm/models/base.pyi

@@ -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] = ...,

.mlx_typings/mlx_lm/models/gpt_oss.pyi

@@ -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]: ...

.mlx_typings/mlx_lm/models/minimax.pyi

@@ -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]: ...

.mlx_typings/mlx_lm/models/nemotron_h.pyi

@@ -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: ...

.mlx_typings/mlx_lm/models/qwen3_next.pyi

@@ -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

app/EXO/EXO/ContentView.swift

@@ -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)

app/EXO/EXO/Info.plist

@@ -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>

bench/eval_tool_calls.py

@@ -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}")

bench/harness.py

@@ -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(

dashboard/src/routes/integrations/+page.svelte

@@ -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"

package-lock.json

@@ -0,0 +1,6 @@
+{
+  "name": "exo",
+  "lockfileVersion": 3,
+  "requires": true,
+  "packages": {}
+}

pyproject.toml

@@ -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",
@@ -49,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'",
 ]
 
@@ -75,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" }
 
@@ -149,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"]

python/parts.nix

@@ -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
+        '';
       };
     };
 }

resources/inference_model_cards/mlx-community--GLM-5.1-DQ4plus-q8.toml

@@ -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

resources/inference_model_cards/mlx-community--GLM-5.1-MXFP4-Q8.toml

@@ -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

resources/inference_model_cards/mlx-community--GLM-5.1.toml

@@ -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

resources/inference_model_cards/mlx-community--Kimi-K2.6-mlx-DQ3_K_M-q8.toml

@@ -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

resources/inference_model_cards/mlx-community--Qwen3.6-27B-4bit.toml

@@ -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

resources/inference_model_cards/mlx-community--Qwen3.6-27B-8bit.toml

@@ -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

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

@@ -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

resources/inference_model_cards/moonshotai--Kimi-K2.6.toml

@@ -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

src/exo/api/adapters/responses.py

@@ -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(

src/exo/api/main.py

@@ -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(

src/exo/download/coordinator.py

@@ -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,

src/exo/download/download_utils.py

@@ -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] = {}

src/exo/download/impl_shard_downloader.py

@@ -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

src/exo/master/main.py

@@ -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)
 

src/exo/shared/types/chunks.py

@@ -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

src/exo/shared/types/events.py

@@ -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):

src/exo/shared/types/tasks.py

@@ -101,3 +101,6 @@ Task = (
     | ImageEdits
     | Shutdown
 )
+TextTask = TextGeneration
+ImageTask = ImageGeneration | ImageEdits
+GenerationTask = TextTask | ImageTask

src/exo/shared/types/text_generation.py

@@ -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

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

@@ -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

src/exo/worker/engines/base.py

@@ -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: ...

src/exo/worker/engines/image/__init__.py

@@ -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",
 ]

src/exo/worker/engines/image/builder.py

@@ -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

src/exo/worker/engines/image/distributed_model.py

@@ -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)

src/exo/worker/engines/image/generate.py

@@ -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))

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

@@ -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

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

@@ -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,
+            )

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

@@ -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)

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

@@ -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(

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

@@ -457,6 +457,7 @@ class ExoBatchGenerator:
 
     def close(self) -> None:
         self._mlx_gen.close()
+        mx.clear_cache()
 
     def _save_prefix_cache(
         self,

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

@@ -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(

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

@@ -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

src/exo/worker/main.py

@@ -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={

src/exo/worker/plan.py

@@ -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

src/exo/worker/runner/bootstrap.py

@@ -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")

src/exo/worker/runner/image_models/runner.py

@@ -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=}"
-                )

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

@@ -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

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

@@ -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

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

@@ -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,
-        )

src/exo/worker/runner/runner.py

@@ -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))

src/exo/worker/tests/unittests/test_mlx/conftest.py

@@ -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. "

src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py

@@ -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."""

src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py

@@ -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)

src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py

@@ -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:

src/exo/worker/tests/unittests/test_mlx/test_tp_bit_exact.py

@@ -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)

src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py

@@ -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(),
     )

src/exo/worker/tests/unittests/test_plan/test_runner_lifecycle.py

@@ -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(),
     )

src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py

@@ -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(),
     )

src/exo/worker/tests/unittests/test_plan/test_warmup.py

@@ -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(),
     )

src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.py

@@ -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)]

src/exo/worker/tests/unittests/test_runner/test_event_ordering.py

@@ -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),

src/exo/worker/tests/unittests/test_runner/test_finish_reason_sse.py

@@ -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

src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py

@@ -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