Eval noise before CFG exchange

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

@@ -178,36 +178,6 @@ interface ImageApiResponse {
   data: Array<{ b64_json?: string; url?: string }>;
 }
 
-// Trace API response types
-export interface TraceCategoryStats {
-  totalUs: number;
-  count: number;
-  minUs: number;
-  maxUs: number;
-  avgUs: number;
-}
-
-export interface TraceRankStats {
-  byCategory: Record<string, TraceCategoryStats>;
-}
-
-export interface TraceStatsResponse {
-  taskId: string;
-  totalWallTimeUs: number;
-  byCategory: Record<string, TraceCategoryStats>;
-  byRank: Record<number, TraceRankStats>;
-}
-
-export interface TraceListItem {
-  taskId: string;
-  createdAt: string;
-  fileSize: number;
-}
-
-export interface TraceListResponse {
-  traces: TraceListItem[];
-}
-
 interface RawStateResponse {
   topology?: RawTopology;
   instances?: Record<
@@ -2585,49 +2555,6 @@ class AppStore {
       throw error;
     }
   }
-
-  /**
-   * List all available traces
-   */
-  async listTraces(): Promise<TraceListResponse> {
-    const response = await fetch("/v1/traces");
-    if (!response.ok) {
-      throw new Error(`Failed to list traces: ${response.status}`);
-    }
-    return (await response.json()) as TraceListResponse;
-  }
-
-  /**
-   * Check if a trace exists for a given task ID
-   */
-  async checkTraceExists(taskId: string): Promise<boolean> {
-    try {
-      const response = await fetch(`/v1/traces/${encodeURIComponent(taskId)}`);
-      return response.ok;
-    } catch {
-      return false;
-    }
-  }
-
-  /**
-   * Get computed statistics for a task's trace
-   */
-  async fetchTraceStats(taskId: string): Promise<TraceStatsResponse> {
-    const response = await fetch(
-      `/v1/traces/${encodeURIComponent(taskId)}/stats`,
-    );
-    if (!response.ok) {
-      throw new Error(`Failed to fetch trace stats: ${response.status}`);
-    }
-    return (await response.json()) as TraceStatsResponse;
-  }
-
-  /**
-   * Get the URL for the raw trace file (for Perfetto)
-   */
-  getTraceRawUrl(taskId: string): string {
-    return `/v1/traces/${encodeURIComponent(taskId)}/raw`;
-  }
 }
 
 export const appStore = new AppStore();
@@ -2739,12 +2666,3 @@ export const startDownload = (nodeId: string, shardMetadata: object) =>
   appStore.startDownload(nodeId, shardMetadata);
 export const deleteDownload = (nodeId: string, modelId: string) =>
   appStore.deleteDownload(nodeId, modelId);
-
-// Trace actions
-export const listTraces = () => appStore.listTraces();
-export const checkTraceExists = (taskId: string) =>
-  appStore.checkTraceExists(taskId);
-export const fetchTraceStats = (taskId: string) =>
-  appStore.fetchTraceStats(taskId);
-export const getTraceRawUrl = (taskId: string) =>
-  appStore.getTraceRawUrl(taskId);

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

@@ -1,172 +0,0 @@
-<script lang="ts">
-  import { onMount } from "svelte";
-  import {
-    listTraces,
-    getTraceRawUrl,
-    type TraceListItem,
-  } from "$lib/stores/app.svelte";
-  import HeaderNav from "$lib/components/HeaderNav.svelte";
-
-  let traces = $state<TraceListItem[]>([]);
-  let loading = $state(true);
-  let error = $state<string | null>(null);
-
-  function formatBytes(bytes: number): string {
-    if (!bytes || bytes <= 0) return "0B";
-    const units = ["B", "KB", "MB", "GB"];
-    const i = Math.min(
-      Math.floor(Math.log(bytes) / Math.log(1024)),
-      units.length - 1,
-    );
-    const val = bytes / Math.pow(1024, i);
-    return `${val.toFixed(val >= 10 ? 0 : 1)}${units[i]}`;
-  }
-
-  function formatDate(isoString: string): string {
-    const date = new Date(isoString);
-    return date.toLocaleString();
-  }
-
-  async function openInPerfetto(taskId: string) {
-    // Fetch trace data from our local API
-    const response = await fetch(getTraceRawUrl(taskId));
-    const traceData = await response.arrayBuffer();
-
-    // Open Perfetto UI
-    const perfettoWindow = window.open("https://ui.perfetto.dev");
-    if (!perfettoWindow) {
-      alert("Failed to open Perfetto. Please allow popups.");
-      return;
-    }
-
-    // Wait for Perfetto to be ready, then send trace via postMessage
-    const onMessage = (e: MessageEvent) => {
-      if (e.data === "PONG") {
-        window.removeEventListener("message", onMessage);
-        perfettoWindow.postMessage(
-          {
-            perfetto: {
-              buffer: traceData,
-              title: `Trace ${taskId}`,
-            },
-          },
-          "https://ui.perfetto.dev",
-        );
-      }
-    };
-    window.addEventListener("message", onMessage);
-
-    // Ping Perfetto until it responds
-    const pingInterval = setInterval(() => {
-      perfettoWindow.postMessage("PING", "https://ui.perfetto.dev");
-    }, 50);
-
-    // Clean up after 10 seconds
-    setTimeout(() => {
-      clearInterval(pingInterval);
-      window.removeEventListener("message", onMessage);
-    }, 10000);
-  }
-
-  async function refresh() {
-    loading = true;
-    error = null;
-    try {
-      const response = await listTraces();
-      traces = response.traces;
-    } catch (e) {
-      error = e instanceof Error ? e.message : "Failed to load traces";
-    } finally {
-      loading = false;
-    }
-  }
-
-  onMount(() => {
-    refresh();
-  });
-</script>
-
-<div class="min-h-screen bg-exo-dark-gray text-white">
-  <HeaderNav showHome={true} />
-  <div class="max-w-7xl mx-auto px-4 lg:px-8 py-6 space-y-6">
-    <div class="flex items-center justify-between gap-4 flex-wrap">
-      <div>
-        <h1
-          class="text-2xl font-mono tracking-[0.2em] uppercase text-exo-yellow"
-        >
-          Traces
-        </h1>
-      </div>
-      <div class="flex items-center gap-3">
-        <button
-          type="button"
-          class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-2 py-1 rounded"
-          onclick={refresh}
-          disabled={loading}
-        >
-          Refresh
-        </button>
-      </div>
-    </div>
-
-    {#if loading}
-      <div
-        class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-6 text-center text-exo-light-gray"
-      >
-        <div class="text-sm">Loading traces...</div>
-      </div>
-    {:else if error}
-      <div
-        class="rounded border border-red-500/30 bg-red-500/10 p-6 text-center text-red-400"
-      >
-        <div class="text-sm">{error}</div>
-      </div>
-    {:else if traces.length === 0}
-      <div
-        class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-6 text-center text-exo-light-gray space-y-2"
-      >
-        <div class="text-sm">No traces found.</div>
-        <div class="text-xs text-exo-light-gray/70">
-          Run exo with EXO_TRACING_ENABLED=1 to collect traces.
-        </div>
-      </div>
-    {:else}
-      <div class="space-y-3">
-        {#each traces as trace}
-          <div
-            class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-4 flex items-center justify-between gap-4"
-          >
-            <div class="min-w-0 flex-1">
-              <a
-                href="#/traces/{trace.taskId}"
-                class="text-sm font-mono text-white hover:text-exo-yellow transition-colors truncate block"
-              >
-                {trace.taskId}
-              </a>
-              <div class="text-xs text-exo-light-gray font-mono mt-1">
-                {formatDate(trace.createdAt)} &bull; {formatBytes(
-                  trace.fileSize,
-                )}
-              </div>
-            </div>
-            <div class="flex items-center gap-2 shrink-0">
-              <a
-                href="#/traces/{trace.taskId}"
-                class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-2 py-1 rounded"
-              >
-                View Stats
-              </a>
-              <button
-                type="button"
-                class="text-xs font-mono text-exo-dark-gray bg-exo-yellow hover:bg-exo-yellow/90 transition-colors uppercase px-2 py-1 rounded font-semibold"
-                onclick={() => openInPerfetto(trace.taskId)}
-              >
-                View Trace
-              </button>
-            </div>
-          </div>
-        {/each}
-      </div>
-    {/if}
-  </div>
-</div>

dashboard/src/routes/traces/[taskId]/+page.svelte

@@ -1,347 +0,0 @@
-<script lang="ts">
-  import { page } from "$app/stores";
-  import { onMount } from "svelte";
-  import {
-    fetchTraceStats,
-    getTraceRawUrl,
-    type TraceStatsResponse,
-    type TraceCategoryStats,
-  } from "$lib/stores/app.svelte";
-  import HeaderNav from "$lib/components/HeaderNav.svelte";
-
-  const taskId = $derived($page.params.taskId);
-
-  let stats = $state<TraceStatsResponse | null>(null);
-  let loading = $state(true);
-  let error = $state<string | null>(null);
-
-  function formatDuration(us: number): string {
-    if (us < 1000) return `${us.toFixed(0)}us`;
-    if (us < 1_000_000) return `${(us / 1000).toFixed(2)}ms`;
-    return `${(us / 1_000_000).toFixed(2)}s`;
-  }
-
-  function formatPercentage(part: number, total: number): string {
-    if (total === 0) return "0.0%";
-    return `${((part / total) * 100).toFixed(1)}%`;
-  }
-
-  // Parse hierarchical categories like "sync/compute" into phases
-  type PhaseData = {
-    name: string;
-    subcategories: { name: string; stats: TraceCategoryStats }[];
-    totalUs: number; // From outer span (e.g., "sync" category)
-    stepCount: number; // Count of outer span events
-  };
-
-  function parsePhases(
-    byCategory: Record<string, TraceCategoryStats>,
-  ): PhaseData[] {
-    const phases = new Map<
-      string,
-      {
-        subcats: Map<string, TraceCategoryStats>;
-        outerStats: TraceCategoryStats | null;
-      }
-    >();
-
-    for (const [category, catStats] of Object.entries(byCategory)) {
-      if (category.includes("/")) {
-        const [phase, subcat] = category.split("/", 2);
-        if (!phases.has(phase)) {
-          phases.set(phase, { subcats: new Map(), outerStats: null });
-        }
-        phases.get(phase)!.subcats.set(subcat, catStats);
-      } else {
-        // Outer span - this IS the phase total
-        if (!phases.has(category)) {
-          phases.set(category, { subcats: new Map(), outerStats: null });
-        }
-        phases.get(category)!.outerStats = catStats;
-      }
-    }
-
-    return Array.from(phases.entries())
-      .filter(([_, data]) => data.outerStats !== null) // Only phases with outer spans
-      .map(([name, data]) => ({
-        name,
-        subcategories: Array.from(data.subcats.entries())
-          .map(([subName, subStats]) => ({ name: subName, stats: subStats }))
-          .sort((a, b) => b.stats.totalUs - a.stats.totalUs),
-        totalUs: data.outerStats!.totalUs, // Outer span total
-        stepCount: data.outerStats!.count, // Number of steps
-      }))
-      .sort((a, b) => b.totalUs - a.totalUs);
-  }
-
-  async function openInPerfetto() {
-    if (!taskId) return;
-
-    // Fetch trace data from our local API
-    const response = await fetch(getTraceRawUrl(taskId));
-    const traceData = await response.arrayBuffer();
-
-    // Open Perfetto UI
-    const perfettoWindow = window.open("https://ui.perfetto.dev");
-    if (!perfettoWindow) {
-      alert("Failed to open Perfetto. Please allow popups.");
-      return;
-    }
-
-    // Wait for Perfetto to be ready, then send trace via postMessage
-    const onMessage = (e: MessageEvent) => {
-      if (e.data === "PONG") {
-        window.removeEventListener("message", onMessage);
-        perfettoWindow.postMessage(
-          {
-            perfetto: {
-              buffer: traceData,
-              title: `Trace ${taskId}`,
-            },
-          },
-          "https://ui.perfetto.dev",
-        );
-      }
-    };
-    window.addEventListener("message", onMessage);
-
-    // Ping Perfetto until it responds
-    const pingInterval = setInterval(() => {
-      perfettoWindow.postMessage("PING", "https://ui.perfetto.dev");
-    }, 50);
-
-    // Clean up after 10 seconds
-    setTimeout(() => {
-      clearInterval(pingInterval);
-      window.removeEventListener("message", onMessage);
-    }, 10000);
-  }
-
-  onMount(async () => {
-    if (!taskId) {
-      error = "No task ID provided";
-      loading = false;
-      return;
-    }
-
-    try {
-      stats = await fetchTraceStats(taskId);
-    } catch (e) {
-      error = e instanceof Error ? e.message : "Failed to load trace";
-    } finally {
-      loading = false;
-    }
-  });
-
-  const phases = $derived(stats ? parsePhases(stats.byCategory) : []);
-  const sortedRanks = $derived(
-    stats
-      ? Object.keys(stats.byRank)
-          .map(Number)
-          .sort((a, b) => a - b)
-      : [],
-  );
-  const nodeCount = $derived(sortedRanks.length || 1);
-</script>
-
-<div class="min-h-screen bg-exo-dark-gray text-white">
-  <HeaderNav showHome={true} />
-  <div class="max-w-7xl mx-auto px-4 lg:px-8 py-6 space-y-6">
-    <div class="flex items-center justify-between gap-4 flex-wrap">
-      <div>
-        <h1
-          class="text-2xl font-mono tracking-[0.2em] uppercase text-exo-yellow"
-        >
-          Trace
-        </h1>
-        <p class="text-sm text-exo-light-gray font-mono truncate max-w-lg">
-          {taskId}
-        </p>
-      </div>
-      <div class="flex items-center gap-3">
-        <a
-          href="#/traces"
-          class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-3 py-1.5 rounded"
-        >
-          All Traces
-        </a>
-        <button
-          type="button"
-          class="text-xs font-mono text-exo-dark-gray bg-exo-yellow hover:bg-exo-yellow/90 transition-colors uppercase px-3 py-1.5 rounded font-semibold"
-          onclick={openInPerfetto}
-          disabled={loading || !!error}
-        >
-          View Trace
-        </button>
-      </div>
-    </div>
-
-    {#if loading}
-      <div
-        class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-6 text-center text-exo-light-gray"
-      >
-        <div class="text-sm">Loading trace data...</div>
-      </div>
-    {:else if error}
-      <div
-        class="rounded border border-red-500/30 bg-red-500/10 p-6 text-center text-red-400"
-      >
-        <div class="text-sm">{error}</div>
-      </div>
-    {:else if stats}
-      <!-- Wall Time Summary -->
-      <div
-        class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-4 space-y-2"
-      >
-        <h2
-          class="text-sm font-mono uppercase tracking-wider text-exo-light-gray"
-        >
-          Summary
-        </h2>
-        <div class="text-3xl font-mono text-exo-yellow">
-          {formatDuration(stats.totalWallTimeUs)}
-        </div>
-        <div class="text-xs text-exo-light-gray">Total wall time</div>
-      </div>
-
-      <!-- By Phase -->
-      {#if phases.length > 0}
-        <div
-          class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-4 space-y-4"
-        >
-          <h2
-            class="text-sm font-mono uppercase tracking-wider text-exo-light-gray"
-          >
-            By Phase <span class="text-exo-light-gray/50">(avg per node)</span>
-          </h2>
-          <div class="space-y-4">
-            {#each phases as phase}
-              {@const normalizedTotal = phase.totalUs / nodeCount}
-              {@const normalizedStepCount = phase.stepCount / nodeCount}
-              <div class="space-y-2">
-                <div class="flex items-center justify-between">
-                  <span class="text-sm font-mono text-white">{phase.name}</span>
-                  <span class="text-sm font-mono">
-                    <span class="text-exo-yellow"
-                      >{formatDuration(normalizedTotal)}</span
-                    >
-                    <span class="text-exo-light-gray ml-2">
-                      ({normalizedStepCount} steps, {formatDuration(
-                        normalizedTotal / normalizedStepCount,
-                      )}/step)
-                    </span>
-                  </span>
-                </div>
-                {#if phase.subcategories.length > 0}
-                  <div class="pl-4 space-y-1.5">
-                    {#each phase.subcategories as subcat}
-                      {@const normalizedSubcat =
-                        subcat.stats.totalUs / nodeCount}
-                      {@const pct = formatPercentage(
-                        normalizedSubcat,
-                        normalizedTotal,
-                      )}
-                      {@const perStep = normalizedSubcat / normalizedStepCount}
-                      <div
-                        class="flex items-center justify-between text-xs font-mono"
-                      >
-                        <span class="text-exo-light-gray">{subcat.name}</span>
-                        <span class="text-white">
-                          {formatDuration(normalizedSubcat)}
-                          <span class="text-exo-light-gray ml-2">({pct})</span>
-                          <span class="text-exo-light-gray/60 ml-2"
-                            >{formatDuration(perStep)}/step</span
-                          >
-                        </span>
-                      </div>
-                      <!-- Progress bar -->
-                      <div
-                        class="relative h-1.5 bg-exo-black/60 rounded-sm overflow-hidden"
-                      >
-                        <div
-                          class="absolute inset-y-0 left-0 bg-gradient-to-r from-exo-yellow to-exo-yellow/70 transition-all duration-300"
-                          style="width: {pct}"
-                        ></div>
-                      </div>
-                    {/each}
-                  </div>
-                {/if}
-              </div>
-            {/each}
-          </div>
-        </div>
-      {/if}
-
-      <!-- By Rank -->
-      {#if sortedRanks.length > 0}
-        <div
-          class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-4 space-y-4"
-        >
-          <h2
-            class="text-sm font-mono uppercase tracking-wider text-exo-light-gray"
-          >
-            By Rank
-          </h2>
-          <div class="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 gap-4">
-            {#each sortedRanks as rank}
-              {@const rankStats = stats.byRank[rank]}
-              {@const rankPhases = parsePhases(rankStats.byCategory)}
-              <div
-                class="rounded border border-exo-medium-gray/20 bg-exo-dark-gray/60 p-3 space-y-3"
-              >
-                <div class="text-sm font-mono text-exo-yellow">
-                  Rank {rank}
-                </div>
-                <div class="space-y-2">
-                  {#each rankPhases as phase}
-                    <div class="space-y-1">
-                      <div class="flex items-center justify-between text-xs">
-                        <span class="font-mono text-exo-light-gray"
-                          >{phase.name}</span
-                        >
-                        <span class="font-mono text-white">
-                          {formatDuration(phase.totalUs)}
-                          <span class="text-exo-light-gray/50 ml-1">
-                            ({phase.stepCount}x)
-                          </span>
-                        </span>
-                      </div>
-                      {#if phase.subcategories.length > 0}
-                        <div class="pl-2 space-y-0.5">
-                          {#each phase.subcategories as subcat}
-                            {@const pct = formatPercentage(
-                              subcat.stats.totalUs,
-                              phase.totalUs,
-                            )}
-                            {@const perStep =
-                              subcat.stats.totalUs / phase.stepCount}
-                            <div
-                              class="flex items-center justify-between text-[10px] font-mono"
-                            >
-                              <span class="text-exo-light-gray/70"
-                                >{subcat.name}</span
-                              >
-                              <span class="text-exo-light-gray">
-                                {formatDuration(subcat.stats.totalUs)}
-                                <span class="text-exo-light-gray/50"
-                                  >({pct})</span
-                                >
-                                <span class="text-exo-light-gray/30 ml-1"
-                                  >{formatDuration(perStep)}/step</span
-                                >
-                              </span>
-                            </div>
-                          {/each}
-                        </div>
-                      {/if}
-                    </div>
-                  {/each}
-                </div>
-              </div>
-            {/each}
-          </div>
-        </div>
-      {/if}
-    {/if}
-  </div>
-</div>

src/exo/download/impl_shard_downloader.py

@@ -21,7 +21,7 @@ def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
 
 
 async def build_base_shard(model_id: ModelId) -> ShardMetadata:
-    model_card = await ModelCard.load(model_id)
+    model_card = await ModelCard.from_hf(model_id)
     return PipelineShardMetadata(
         model_card=model_card,
         device_rank=0,
@@ -166,8 +166,9 @@ class ResumableShardDownloader(ShardDownloader):
         for task in asyncio.as_completed(tasks):
             try:
                 yield await task
+            # TODO: except Exception
             except Exception as e:
-                logger.warning(f"Error downloading shard: {type(e).__name__}")
+                logger.error("Error downloading shard:", e)
 
     async def get_shard_download_status_for_shard(
         self, shard: ShardMetadata

src/exo/master/api.py

@@ -1,11 +1,10 @@
 import base64
 import contextlib
 import json
+import random
 import time
 from collections.abc import AsyncGenerator
-from datetime import datetime, timezone
 from http import HTTPStatus
-from pathlib import Path
 from typing import Annotated, Literal, cast
 from uuid import uuid4
 
@@ -35,7 +34,6 @@ from exo.shared.models.model_cards import (
     ModelCard,
     ModelId,
 )
-from exo.shared.tracing import compute_stats, load_trace_file
 from exo.shared.types.api import (
     AdvancedImageParams,
     BenchChatCompletionResponse,
@@ -68,16 +66,7 @@ from exo.shared.types.api import (
     StartDownloadParams,
     StartDownloadResponse,
     StreamingChoiceResponse,
-    StreamOptions,
     ToolCall,
-    TraceCategoryStats,
-    TraceEventResponse,
-    TraceListItem,
-    TraceListResponse,
-    TraceRankStats,
-    TraceResponse,
-    TraceStatsResponse,
-    Usage,
 )
 from exo.shared.types.chunks import (
     ErrorChunk,
@@ -124,10 +113,17 @@ def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None)
     return f"image/{image_format or 'png'}"
 
 
+def _ensure_seed(params: AdvancedImageParams | None) -> AdvancedImageParams:
+    """Ensure advanced params has a seed set for distributed consistency."""
+    if params is None:
+        return AdvancedImageParams(seed=random.randint(0, 2**32 - 1))
+    if params.seed is None:
+        return params.model_copy(update={"seed": random.randint(0, 2**32 - 1)})
+    return params
+
+
 def chunk_to_response(
-    chunk: TokenChunk | ToolCallChunk,
-    command_id: CommandId,
-    usage: Usage | None,
+    chunk: TokenChunk | ToolCallChunk, command_id: CommandId
 ) -> ChatCompletionResponse:
     return ChatCompletionResponse(
         id=command_id,
@@ -152,10 +148,21 @@ def chunk_to_response(
                 finish_reason=chunk.finish_reason,
             )
         ],
-        usage=usage,
     )
 
 
+async def resolve_model_card(model_id: ModelId) -> ModelCard:
+    if model_id in MODEL_CARDS:
+        model_card = MODEL_CARDS[model_id]
+        return model_card
+
+    for card in MODEL_CARDS.values():
+        if card.model_id == ModelId(model_id):
+            return card
+
+    return await ModelCard.from_hf(model_id)
+
+
 class API:
     def __init__(
         self,
@@ -274,14 +281,10 @@ class API:
         self.app.get("/events")(lambda: self._event_log)
         self.app.post("/download/start")(self.start_download)
         self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download)
-        self.app.get("/v1/traces")(self.list_traces)
-        self.app.get("/v1/traces/{task_id}")(self.get_trace)
-        self.app.get("/v1/traces/{task_id}/stats")(self.get_trace_stats)
-        self.app.get("/v1/traces/{task_id}/raw")(self.get_trace_raw)
 
     async def place_instance(self, payload: PlaceInstanceParams):
         command = PlaceInstance(
-            model_card=await ModelCard.load(payload.model_id),
+            model_card=await resolve_model_card(payload.model_id),
             sharding=payload.sharding,
             instance_meta=payload.instance_meta,
             min_nodes=payload.min_nodes,
@@ -298,7 +301,7 @@ class API:
         self, payload: CreateInstanceParams
     ) -> CreateInstanceResponse:
         instance = payload.instance
-        model_card = await ModelCard.load(instance.shard_assignments.model_id)
+        model_card = await resolve_model_card(instance.shard_assignments.model_id)
         required_memory = model_card.storage_size
         available_memory = self._calculate_total_available_memory()
 
@@ -326,7 +329,7 @@ class API:
         instance_meta: InstanceMeta = InstanceMeta.MlxRing,
         min_nodes: int = 1,
     ) -> Instance:
-        model_card = await ModelCard.load(model_id)
+        model_card = await resolve_model_card(model_id)
 
         try:
             placements = get_instance_placements(
@@ -529,10 +532,9 @@ class API:
                 del self._chat_completion_queues[command_id]
 
     async def _generate_chat_stream(
-        self, command_id: CommandId, stream_options: StreamOptions | None = None
+        self, command_id: CommandId
     ) -> AsyncGenerator[str, None]:
         """Generate chat completion stream as JSON strings."""
-        include_usage = stream_options.include_usage if stream_options else False
 
         async for chunk in self._chat_chunk_stream(command_id):
             assert not isinstance(chunk, ImageChunk)
@@ -548,10 +550,8 @@ class API:
                 yield "data: [DONE]\n\n"
                 return
 
-            usage = chunk.usage if include_usage else None
-
             chunk_response: ChatCompletionResponse = chunk_to_response(
-                chunk, command_id, usage=usage
+                chunk, command_id
             )
             logger.debug(f"chunk_response: {chunk_response}")
 
@@ -567,9 +567,8 @@ class API:
 
         text_parts: list[str] = []
         tool_calls: list[ToolCall] = []
-        model: ModelId | None = None
+        model: str | None = None
         finish_reason: FinishReason | None = None
-        usage: Usage | None = None
 
         async for chunk in self._chat_chunk_stream(command_id):
             if isinstance(chunk, ErrorChunk):
@@ -594,9 +593,6 @@ class API:
                     for i, tool in enumerate(chunk.tool_calls)
                 )
 
-            if chunk.usage is not None:
-                usage = chunk.usage
-
             if chunk.finish_reason is not None:
                 finish_reason = chunk.finish_reason
 
@@ -618,7 +614,6 @@ class API:
                     finish_reason=finish_reason,
                 )
             ],
-            usage=usage,
         )
 
     async def _collect_chat_completion_with_stats(
@@ -626,7 +621,7 @@ class API:
     ) -> BenchChatCompletionResponse:
         text_parts: list[str] = []
         tool_calls: list[ToolCall] = []
-        model: ModelId | None = None
+        model: str | None = None
         finish_reason: FinishReason | None = None
 
         stats: GenerationStats | None = None
@@ -679,7 +674,7 @@ class API:
         )
         return resp
 
-    async def _trigger_notify_user_to_download_model(self, model_id: ModelId) -> None:
+    async def _trigger_notify_user_to_download_model(self, model_id: str) -> None:
         logger.warning(
             "TODO: we should send a notification to the user to download the model"
         )
@@ -688,7 +683,7 @@ class API:
         self, payload: ChatCompletionTaskParams
     ) -> ChatCompletionResponse | StreamingResponse:
         """Handle chat completions, supporting both streaming and non-streaming responses."""
-        model_card = await ModelCard.load(ModelId(payload.model))
+        model_card = await resolve_model_card(ModelId(payload.model))
         payload.model = model_card.model_id
 
         if not any(
@@ -706,7 +701,7 @@ class API:
         await self._send(command)
         if payload.stream:
             return StreamingResponse(
-                self._generate_chat_stream(command.command_id, payload.stream_options),
+                self._generate_chat_stream(command.command_id),
                 media_type="text/event-stream",
             )
 
@@ -715,7 +710,7 @@ class API:
     async def bench_chat_completions(
         self, payload: BenchChatCompletionTaskParams
     ) -> BenchChatCompletionResponse:
-        model_card = await ModelCard.load(ModelId(payload.model))
+        model_card = await resolve_model_card(ModelId(payload.model))
         payload.model = model_card.model_id
 
         if not any(
@@ -735,12 +730,12 @@ class API:
         response = await self._collect_chat_completion_with_stats(command.command_id)
         return response
 
-    async def _validate_image_model(self, model: ModelId) -> ModelId:
+    async def _validate_image_model(self, model: str) -> ModelId:
         """Validate model exists and return resolved model ID.
 
         Raises HTTPException 404 if no instance is found for the model.
         """
-        model_card = await ModelCard.load(model)
+        model_card = await resolve_model_card(ModelId(model))
         resolved_model = model_card.model_id
         if not any(
             instance.shard_assignments.model_id == resolved_model
@@ -786,7 +781,10 @@ class API:
         When stream=True and partial_images > 0, returns a StreamingResponse
         with SSE-formatted events for partial and final images.
         """
-        payload.model = await self._validate_image_model(ModelId(payload.model))
+        payload.model = await self._validate_image_model(payload.model)
+        payload = payload.model_copy(
+            update={"advanced_params": _ensure_seed(payload.advanced_params)}
+        )
 
         command = ImageGeneration(
             request_params=payload,
@@ -1031,10 +1029,13 @@ class API:
     async def bench_image_generations(
         self, request: Request, payload: BenchImageGenerationTaskParams
     ) -> BenchImageGenerationResponse:
-        payload.model = await self._validate_image_model(ModelId(payload.model))
+        payload.model = await self._validate_image_model(payload.model)
 
         payload.stream = False
         payload.partial_images = 0
+        payload = payload.model_copy(
+            update={"advanced_params": _ensure_seed(payload.advanced_params)}
+        )
 
         command = ImageGeneration(
             request_params=payload,
@@ -1052,7 +1053,7 @@ class API:
         self,
         image: UploadFile,
         prompt: str,
-        model: ModelId,
+        model: str,
         n: int,
         size: str,
         response_format: Literal["url", "b64_json"],
@@ -1066,6 +1067,7 @@ class API:
     ) -> ImageEdits:
         """Prepare and send an image edits command with chunked image upload."""
         resolved_model = await self._validate_image_model(model)
+        advanced_params = _ensure_seed(advanced_params)
 
         image_content = await image.read()
         image_data = base64.b64encode(image_content).decode("utf-8")
@@ -1147,7 +1149,7 @@ class API:
         command = await self._send_image_edits_command(
             image=image,
             prompt=prompt,
-            model=ModelId(model),
+            model=model,
             n=n,
             size=size,
             response_format=response_format,
@@ -1203,7 +1205,7 @@ class API:
         command = await self._send_image_edits_command(
             image=image,
             prompt=prompt,
-            model=ModelId(model),
+            model=model,
             n=n,
             size=size,
             response_format=response_format,
@@ -1350,110 +1352,3 @@ class API:
         )
         await self._send_download(command)
         return DeleteDownloadResponse(command_id=command.command_id)
-
-    def _get_traces_dir(self) -> Path:
-        return Path.home() / ".exo" / "traces"
-
-    def _get_trace_path(self, task_id: str) -> Path:
-        return self._get_traces_dir() / f"trace_{task_id}.json"
-
-    async def list_traces(self) -> TraceListResponse:
-        traces_dir = self._get_traces_dir()
-        traces: list[TraceListItem] = []
-
-        if not traces_dir.exists():
-            return TraceListResponse(traces=[])
-
-        for trace_file in sorted(
-            traces_dir.glob("trace_*.json"),
-            key=lambda p: p.stat().st_mtime,
-            reverse=True,
-        ):
-            # Extract task_id from filename (trace_{task_id}.json)
-            task_id = trace_file.stem.removeprefix("trace_")
-            stat = trace_file.stat()
-            created_at = datetime.fromtimestamp(
-                stat.st_mtime, tz=timezone.utc
-            ).isoformat()
-            traces.append(
-                TraceListItem(
-                    task_id=task_id,
-                    created_at=created_at,
-                    file_size=stat.st_size,
-                )
-            )
-
-        return TraceListResponse(traces=traces)
-
-    async def get_trace(self, task_id: str) -> TraceResponse:
-        trace_path = self._get_trace_path(task_id)
-
-        if not trace_path.exists():
-            raise HTTPException(status_code=404, detail=f"Trace not found: {task_id}")
-
-        trace_events = load_trace_file(trace_path)
-
-        return TraceResponse(
-            task_id=task_id,
-            traces=[
-                TraceEventResponse(
-                    name=event.name,
-                    start_us=event.start_us,
-                    duration_us=event.duration_us,
-                    rank=event.rank,
-                    category=event.category,
-                )
-                for event in trace_events
-            ],
-        )
-
-    async def get_trace_stats(self, task_id: str) -> TraceStatsResponse:
-        trace_path = self._get_trace_path(task_id)
-
-        if not trace_path.exists():
-            raise HTTPException(status_code=404, detail=f"Trace not found: {task_id}")
-
-        trace_events = load_trace_file(trace_path)
-        stats = compute_stats(trace_events)
-
-        return TraceStatsResponse(
-            task_id=task_id,
-            total_wall_time_us=stats.total_wall_time_us,
-            by_category={
-                category: TraceCategoryStats(
-                    total_us=cat_stats.total_us,
-                    count=cat_stats.count,
-                    min_us=cat_stats.min_us,
-                    max_us=cat_stats.max_us,
-                    avg_us=cat_stats.avg_us,
-                )
-                for category, cat_stats in stats.by_category.items()
-            },
-            by_rank={
-                rank: TraceRankStats(
-                    by_category={
-                        category: TraceCategoryStats(
-                            total_us=cat_stats.total_us,
-                            count=cat_stats.count,
-                            min_us=cat_stats.min_us,
-                            max_us=cat_stats.max_us,
-                            avg_us=cat_stats.avg_us,
-                        )
-                        for category, cat_stats in rank_stats.items()
-                    }
-                )
-                for rank, rank_stats in stats.by_rank.items()
-            },
-        )
-
-    async def get_trace_raw(self, task_id: str) -> FileResponse:
-        trace_path = self._get_trace_path(task_id)
-
-        if not trace_path.exists():
-            raise HTTPException(status_code=404, detail=f"Trace not found: {task_id}")
-
-        return FileResponse(
-            path=trace_path,
-            media_type="application/json",
-            filename=f"trace_{task_id}.json",
-        )

src/exo/master/main.py

@@ -1,5 +1,4 @@
 from datetime import datetime, timedelta, timezone
-from pathlib import Path
 
 import anyio
 from anyio.abc import TaskGroup
@@ -12,7 +11,6 @@ from exo.master.placement import (
     place_instance,
 )
 from exo.shared.apply import apply
-from exo.shared.tracing import TraceEvent, export_trace, is_tracing_enabled
 from exo.shared.types.commands import (
     ChatCompletion,
     CreateInstance,
@@ -37,8 +35,6 @@ from exo.shared.types.events import (
     NodeTimedOut,
     TaskCreated,
     TaskDeleted,
-    TraceEventData,
-    TracesCollected,
 )
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
@@ -90,8 +86,6 @@ class Master:
         self._multi_buffer = MultiSourceBuffer[NodeId, Event]()
         # TODO: not have this
         self._event_log: list[Event] = []
-        self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {}
-        self._expected_ranks: dict[TaskId, set[int]] = {}
 
     async def run(self):
         logger.info("Starting Master")
@@ -193,14 +187,13 @@ class Master:
                             )
 
                             task_id = TaskId()
-                            selected_instance_id = available_instance_ids[0]
                             generated_events.append(
                                 TaskCreated(
                                     task_id=task_id,
                                     task=ImageGenerationTask(
                                         task_id=task_id,
                                         command_id=command.command_id,
-                                        instance_id=selected_instance_id,
+                                        instance_id=available_instance_ids[0],
                                         task_status=TaskStatus.Pending,
                                         task_params=command.request_params,
                                     ),
@@ -208,17 +201,6 @@ class Master:
                             )
 
                             self.command_task_mapping[command.command_id] = task_id
-
-                            if is_tracing_enabled():
-                                selected_instance = self.state.instances.get(
-                                    selected_instance_id
-                                )
-                                if selected_instance:
-                                    ranks = set(
-                                        shard.device_rank
-                                        for shard in selected_instance.shard_assignments.runner_to_shard.values()
-                                    )
-                                    self._expected_ranks[task_id] = ranks
                         case ImageEdits():
                             for instance in self.state.instances.values():
                                 if (
@@ -247,14 +229,13 @@ class Master:
                             )
 
                             task_id = TaskId()
-                            selected_instance_id = available_instance_ids[0]
                             generated_events.append(
                                 TaskCreated(
                                     task_id=task_id,
                                     task=ImageEditsTask(
                                         task_id=task_id,
                                         command_id=command.command_id,
-                                        instance_id=selected_instance_id,
+                                        instance_id=available_instance_ids[0],
                                         task_status=TaskStatus.Pending,
                                         task_params=command.request_params,
                                     ),
@@ -262,17 +243,6 @@ class Master:
                             )
 
                             self.command_task_mapping[command.command_id] = task_id
-
-                            if is_tracing_enabled():
-                                selected_instance = self.state.instances.get(
-                                    selected_instance_id
-                                )
-                                if selected_instance:
-                                    ranks = set(
-                                        shard.device_rank
-                                        for shard in selected_instance.shard_assignments.runner_to_shard.values()
-                                    )
-                                    self._expected_ranks[task_id] = ranks
                         case DeleteInstance():
                             placement = delete_instance(command, self.state.instances)
                             transition_events = get_transition_events(
@@ -365,10 +335,6 @@ class Master:
                     local_event.origin,
                 )
                 for event in self._multi_buffer.drain():
-                    if isinstance(event, TracesCollected):
-                        self._handle_traces_collected(event)
-                        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)
@@ -407,38 +373,3 @@ class Master:
                 event=event.event,
             )
         )
-
-    def _handle_traces_collected(self, event: TracesCollected) -> None:
-        task_id = event.task_id
-        if task_id not in self._pending_traces:
-            self._pending_traces[task_id] = {}
-        self._pending_traces[task_id][event.rank] = event.traces
-
-        if (
-            task_id in self._expected_ranks
-            and set(self._pending_traces[task_id].keys())
-            >= self._expected_ranks[task_id]
-        ):
-            self._merge_and_save_traces(task_id)
-
-    def _merge_and_save_traces(self, task_id: TaskId) -> None:
-        all_traces: list[TraceEvent] = []
-        for trace_data in self._pending_traces[task_id].values():
-            for t in trace_data:
-                all_traces.append(
-                    TraceEvent(
-                        name=t.name,
-                        start_us=t.start_us,
-                        duration_us=t.duration_us,
-                        rank=t.rank,
-                        category=t.category,
-                    )
-                )
-
-        output_path = Path.home() / ".exo" / "traces" / f"trace_{task_id}.json"
-        export_trace(all_traces, output_path)
-        logger.info(f"Merged traces saved to {output_path}")
-
-        del self._pending_traces[task_id]
-        if task_id in self._expected_ranks:
-            del self._expected_ranks[task_id]

src/exo/master/placement_utils.py

@@ -94,20 +94,35 @@ def get_shard_assignments_for_pipeline_parallel(
     runner_to_shard: dict[RunnerId, ShardMetadata] = {}
     node_to_runner: dict[NodeId, RunnerId] = {}
 
+    # Determine CFG parallelism topology
+    # CFG parallel only for even node counts with CFG models (2+ nodes)
+    use_cfg_parallel = model_card.uses_cfg and world_size >= 2 and world_size % 2 == 0
+    cfg_world_size = 2 if use_cfg_parallel else 1
+    pipeline_world_size = world_size // cfg_world_size
+
+    # For CFG parallel, we only need to allocate layers for one pipeline group
+    # (both CFG groups run the same layers). Use the first pipeline group's nodes.
+    pipeline_node_ids = cycle.node_ids[:pipeline_world_size]
+    pipeline_memory = sum(
+        (node_memory[node_id].ram_available for node_id in pipeline_node_ids),
+        start=Memory(),
+    )
+
     layer_allocations = allocate_layers_proportionally(
         total_layers=total_layers,
         memory_fractions=[
-            node_memory[node_id].ram_available.in_bytes / cycle_memory.in_bytes
-            for node_id in cycle.node_ids
+            node_memory[node_id].ram_available.in_bytes / pipeline_memory.in_bytes
+            for node_id in pipeline_node_ids
         ],
     )
 
-    # Validate each node has sufficient memory for its assigned layers
-    memory_per_layer = model_card.storage_size.in_bytes / total_layers
-    for i, (node_id, node_layers) in enumerate(
-        zip(cycle.node_ids, layer_allocations, strict=True)
-    ):
-        required_memory = node_layers * memory_per_layer
+    # Validate each pipeline node has sufficient memory for its assigned layers
+    # Use integer arithmetic to avoid floating point precision issues
+    total_storage_bytes = model_card.storage_size.in_bytes
+    for i, node_id in enumerate(pipeline_node_ids):
+        node_layers = layer_allocations[i]
+        # Integer division then multiply to get conservative estimate
+        required_memory = (total_storage_bytes * node_layers) // total_layers
         available_memory = node_memory[node_id].ram_available.in_bytes
         if required_memory > available_memory:
             raise ValueError(
@@ -116,24 +131,69 @@ def get_shard_assignments_for_pipeline_parallel(
                 f"but only has {available_memory / (1024**3):.2f} GB available"
             )
 
-    layers_assigned = 0
-    for i, (node_id, node_layers) in enumerate(
-        zip(cycle.node_ids, layer_allocations, strict=True)
-    ):
+    # CFG group 0: pipeline ranks in ascending order (0, 1, 2, ...)
+    # CFG group 1: pipeline ranks in descending order (reversed)
+    # This places both "last stages" as ring neighbors for CFG exchange.
+    position_to_cfg_pipeline = [(0, r) for r in range(pipeline_world_size)] + [
+        (1, r) for r in reversed(range(pipeline_world_size))
+    ]
+
+    cfg_pipeline_to_device: dict[tuple[int, int], int] = {
+        (cfg_rank, pipeline_rank): i
+        for i, (cfg_rank, pipeline_rank) in enumerate(position_to_cfg_pipeline)
+    }
+
+    for i, node_id in enumerate(cycle.node_ids):
+        cfg_rank, pipeline_rank = position_to_cfg_pipeline[i]
+
+        layers_before = sum(layer_allocations[:pipeline_rank])
+        node_layers = layer_allocations[pipeline_rank]
+
+        is_first_stage = pipeline_rank == 0
+        is_last_stage = pipeline_rank == pipeline_world_size - 1
+
+        if is_last_stage:
+            next_pipeline_device = None
+        else:
+            next_pipeline_device = cfg_pipeline_to_device[(cfg_rank, pipeline_rank + 1)]
+
+        if is_first_stage:
+            prev_pipeline_device = None
+        else:
+            prev_pipeline_device = cfg_pipeline_to_device[(cfg_rank, pipeline_rank - 1)]
+
+        if is_last_stage and use_cfg_parallel:
+            other_cfg_rank = 1 - cfg_rank
+            cfg_peer_device = cfg_pipeline_to_device[(other_cfg_rank, pipeline_rank)]
+        else:
+            cfg_peer_device = None
+
+        first_pipeline_device = cfg_pipeline_to_device[(cfg_rank, 0)]
+        last_pipeline_device = cfg_pipeline_to_device[
+            (cfg_rank, pipeline_world_size - 1)
+        ]
+
         runner_id = RunnerId()
 
         shard = PipelineShardMetadata(
             model_card=model_card,
             device_rank=i,
             world_size=world_size,
-            start_layer=layers_assigned,
-            end_layer=layers_assigned + node_layers,
+            start_layer=layers_before,
+            end_layer=layers_before + node_layers,
             n_layers=total_layers,
+            cfg_rank=cfg_rank,
+            cfg_world_size=cfg_world_size,
+            explicit_pipeline_rank=pipeline_rank,
+            next_pipeline_device=next_pipeline_device,
+            prev_pipeline_device=prev_pipeline_device,
+            cfg_peer_device=cfg_peer_device,
+            first_pipeline_device=first_pipeline_device,
+            last_pipeline_device=last_pipeline_device,
         )
 
         runner_to_shard[runner_id] = shard
         node_to_runner[node_id] = runner_id
-        layers_assigned += node_layers
 
     shard_assignments = ShardAssignments(
         model_id=model_card.model_id,

src/exo/master/tests/test_placement_utils.py

@@ -5,6 +5,7 @@ from exo.master.placement_utils import (
     filter_cycles_by_memory,
     get_mlx_jaccl_coordinators,
     get_shard_assignments,
+    get_shard_assignments_for_pipeline_parallel,
     get_smallest_cycles,
 )
 from exo.master.tests.conftest import (
@@ -20,7 +21,7 @@ from exo.shared.types.profiling import (
     NodeNetworkInfo,
 )
 from exo.shared.types.topology import Connection, SocketConnection
-from exo.shared.types.worker.shards import Sharding
+from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
 
 
 def test_filter_cycles_by_memory():
@@ -487,3 +488,195 @@ def test_get_shard_assignments_insufficient_memory_raises():
         get_shard_assignments(
             model_card, selected_cycle, Sharding.Pipeline, node_memory
         )
+
+
+class TestCfgParallelPlacement:
+    def _create_ring_topology(self, node_ids: list[NodeId]) -> Topology:
+        topology = Topology()
+        for node_id in node_ids:
+            topology.add_node(node_id)
+
+        for i, node_id in enumerate(node_ids):
+            next_node = node_ids[(i + 1) % len(node_ids)]
+            conn = Connection(
+                source=node_id,
+                sink=next_node,
+                edge=create_socket_connection(i + 1),
+            )
+            topology.add_connection(conn)
+
+        return topology
+
+    def test_two_nodes_cfg_model_uses_cfg_parallel(self):
+        """Two nodes with CFG model should use CFG parallel (no pipeline)."""
+        node_a = NodeId()
+        node_b = NodeId()
+
+        topology = self._create_ring_topology([node_a, node_b])
+        cycles = [c for c in topology.get_cycles() if len(c) == 2]
+        cycle = cycles[0]
+
+        node_memory = {
+            node_a: create_node_memory(1000 * 1024),
+            node_b: create_node_memory(1000 * 1024),
+        }
+
+        model_card = ModelCard(
+            model_id=ModelId("qwen-image-test"),
+            n_layers=60,
+            storage_size=Memory.from_kb(1000),
+            hidden_size=1,
+            supports_tensor=False,
+            uses_cfg=True,
+            tasks=[ModelTask.TextToImage],
+        )
+
+        assignments = get_shard_assignments_for_pipeline_parallel(
+            model_card, cycle, node_memory
+        )
+
+        shards = list(assignments.runner_to_shard.values())
+        assert len(shards) == 2
+
+        # Both nodes should have all layers (no pipeline split)
+        for shard in shards:
+            assert isinstance(shard, PipelineShardMetadata)
+            assert shard.start_layer == 0
+            assert shard.end_layer == 60
+            assert shard.cfg_world_size == 2
+
+        cfg_ranks = sorted(
+            s.cfg_rank for s in shards if isinstance(s, PipelineShardMetadata)
+        )
+        assert cfg_ranks == [0, 1]
+
+    def test_four_nodes_cfg_model_uses_hybrid(self):
+        """Four nodes with CFG model should use 2 CFG groups x 2 pipeline stages."""
+        nodes = [NodeId() for _ in range(4)]
+
+        topology = self._create_ring_topology(nodes)
+        cycles = [c for c in topology.get_cycles() if len(c) == 4]
+        cycle = cycles[0]
+
+        node_memory = {n: create_node_memory(1000 * 1024) for n in nodes}
+
+        model_card = ModelCard(
+            model_id=ModelId("qwen-image-test"),
+            n_layers=60,
+            storage_size=Memory.from_kb(1000),
+            hidden_size=1,
+            supports_tensor=False,
+            uses_cfg=True,
+            tasks=[ModelTask.TextToImage],
+        )
+
+        assignments = get_shard_assignments_for_pipeline_parallel(
+            model_card, cycle, node_memory
+        )
+
+        shards = list(assignments.runner_to_shard.values())
+        assert len(shards) == 4
+
+        for shard in shards:
+            assert isinstance(shard, PipelineShardMetadata)
+            assert shard.cfg_world_size == 2
+            assert shard.pipeline_world_size == 2
+
+        # Check we have 2 nodes in each CFG group
+        cfg_0_shards = [
+            s
+            for s in shards
+            if isinstance(s, PipelineShardMetadata) and s.cfg_rank == 0
+        ]
+        cfg_1_shards = [
+            s
+            for s in shards
+            if isinstance(s, PipelineShardMetadata) and s.cfg_rank == 1
+        ]
+        assert len(cfg_0_shards) == 2
+        assert len(cfg_1_shards) == 2
+
+        # Both CFG groups should have the same layer assignments
+        cfg_0_layers = [(s.start_layer, s.end_layer) for s in cfg_0_shards]
+        cfg_1_layers = [(s.start_layer, s.end_layer) for s in cfg_1_shards]
+        assert sorted(cfg_0_layers) == sorted(cfg_1_layers)
+
+    def test_three_nodes_cfg_model_uses_sequential_cfg(self):
+        """Three nodes (odd) with CFG model should use sequential CFG."""
+        nodes = [NodeId() for _ in range(3)]
+
+        topology = self._create_ring_topology(nodes)
+        cycles = [c for c in topology.get_cycles() if len(c) == 3]
+        cycle = cycles[0]
+
+        node_memory = {n: create_node_memory(1000 * 1024) for n in nodes}
+
+        model_card = ModelCard(
+            model_id=ModelId("qwen-image-test"),
+            n_layers=60,
+            storage_size=Memory.from_kb(1000),
+            hidden_size=1,
+            supports_tensor=False,
+            uses_cfg=True,
+            tasks=[ModelTask.TextToImage],
+        )
+
+        assignments = get_shard_assignments_for_pipeline_parallel(
+            model_card, cycle, node_memory
+        )
+
+        shards = list(assignments.runner_to_shard.values())
+        assert len(shards) == 3
+
+        for shard in shards:
+            assert isinstance(shard, PipelineShardMetadata)
+            # cfg_world_size = 1 means sequential CFG
+            assert shard.cfg_world_size == 1
+            assert shard.cfg_rank == 0
+
+    def test_two_nodes_non_cfg_model_uses_pipeline(self):
+        """Two nodes with non-CFG model should use pure pipeline."""
+        node_a = NodeId()
+        node_b = NodeId()
+
+        topology = self._create_ring_topology([node_a, node_b])
+        cycles = [c for c in topology.get_cycles() if len(c) == 2]
+        cycle = cycles[0]
+
+        node_memory = {
+            node_a: create_node_memory(1000 * 1024),
+            node_b: create_node_memory(1000 * 1024),
+        }
+
+        model_card = ModelCard(
+            model_id=ModelId("flux-test"),
+            n_layers=57,
+            storage_size=Memory.from_kb(1000),
+            hidden_size=1,
+            supports_tensor=False,
+            uses_cfg=False,  # Non-CFG model
+            tasks=[ModelTask.TextToImage],
+        )
+
+        assignments = get_shard_assignments_for_pipeline_parallel(
+            model_card, cycle, node_memory
+        )
+
+        shards = list(assignments.runner_to_shard.values())
+        assert len(shards) == 2
+
+        for shard in shards:
+            assert isinstance(shard, PipelineShardMetadata)
+            # cfg_world_size = 1 means no CFG parallel
+            assert shard.cfg_world_size == 1
+            assert shard.cfg_rank == 0
+
+        # Should have actual layer sharding (pipeline)
+        layer_ranges = sorted(
+            (s.start_layer, s.end_layer)
+            for s in shards
+            if isinstance(s, PipelineShardMetadata)
+        )
+        # First shard starts at 0, last shard ends at 57
+        assert layer_ranges[0][0] == 0
+        assert layer_ranges[-1][1] == 57

src/exo/routing/router.py

@@ -216,8 +216,6 @@ def get_node_id_keypair(
     Obtains the :class:`Keypair` associated with this node-ID.
     Obtain the :class:`PeerId` by from it.
     """
-    # TODO(evan): bring back node id persistence once we figure out how to deal with duplicates
-    return Keypair.generate_ed25519()
 
     def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path:
         return Path(str(path) + ".lock")

src/exo/shared/apply.py

@@ -25,7 +25,6 @@ from exo.shared.types.events import (
     TestEvent,
     TopologyEdgeCreated,
     TopologyEdgeDeleted,
-    TracesCollected,
 )
 from exo.shared.types.profiling import (
     NodeIdentity,
@@ -56,11 +55,7 @@ def event_apply(event: Event, state: State) -> State:
     """Apply an event to state."""
     match event:
         case (
-            TestEvent()
-            | ChunkGenerated()
-            | TaskAcknowledged()
-            | InputChunkReceived()
-            | TracesCollected()
+            TestEvent() | ChunkGenerated() | TaskAcknowledged() | InputChunkReceived()
         ):  # Pass-through events that don't modify state
             return state
         case InstanceCreated():

src/exo/shared/constants.py

@@ -53,9 +53,3 @@ EXO_IMAGE_CACHE_DIR = EXO_CACHE_HOME / "images"
 EXO_ENABLE_IMAGE_MODELS = (
     os.getenv("EXO_ENABLE_IMAGE_MODELS", "false").lower() == "true"
 )
-
-EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "").lower() in (
-    "1",
-    "true",
-    "yes",
-)

src/exo/shared/models/model_cards.py

@@ -47,6 +47,7 @@ class ModelCard(CamelCaseModel):
     supports_tensor: bool
     tasks: list[ModelTask]
     components: list[ComponentInfo] | None = None
+    uses_cfg: bool = False
 
     @field_validator("tasks", mode="before")
     @classmethod
@@ -562,6 +563,7 @@ _IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
         hidden_size=1,
         supports_tensor=False,
         tasks=[ModelTask.TextToImage],
+        uses_cfg=True,
         components=[
             ComponentInfo(
                 component_name="text_encoder",
@@ -596,6 +598,7 @@ _IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
         hidden_size=1,
         supports_tensor=False,
         tasks=[ModelTask.ImageToImage],
+        uses_cfg=True,
         components=[
             ComponentInfo(
                 component_name="text_encoder",
@@ -681,6 +684,7 @@ def _generate_image_model_quant_variants(
             hidden_size=base_card.hidden_size,
             supports_tensor=base_card.supports_tensor,
             tasks=base_card.tasks,
+            uses_cfg=base_card.uses_cfg,
             components=with_transformer_size(transformer_bytes),
         )
     }
@@ -700,6 +704,7 @@ def _generate_image_model_quant_variants(
             hidden_size=base_card.hidden_size,
             supports_tensor=base_card.supports_tensor,
             tasks=base_card.tasks,
+            uses_cfg=base_card.uses_cfg,
             components=with_transformer_size(quant_transformer_bytes),
         )
 

src/exo/shared/tests/test_node_id_persistence.py

@@ -8,7 +8,7 @@ from multiprocessing.synchronize import Event as EventT
 from multiprocessing.synchronize import Semaphore as SemaphoreT
 
 from loguru import logger
-from pytest import LogCaptureFixture, mark
+from pytest import LogCaptureFixture
 
 from exo.routing.router import get_node_id_keypair
 from exo.shared.constants import EXO_NODE_ID_KEYPAIR
@@ -74,7 +74,6 @@ def _delete_if_exists(p: str | bytes | os.PathLike[str] | os.PathLike[bytes]):
         os.remove(p)
 
 
-@mark.skip(reason="this functionality is currently disabled but may return in future")
 def test_node_id_fetching(caplog: LogCaptureFixture):
     reps = 10
 

src/exo/shared/tracing.py

@@ -1,450 +0,0 @@
-from __future__ import annotations
-
-import json
-import logging
-import time
-from collections import defaultdict
-from collections.abc import Generator
-from contextlib import contextmanager
-from contextvars import ContextVar
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import cast, final
-
-from exo.shared.constants import EXO_TRACING_ENABLED
-
-logger = logging.getLogger(__name__)
-
-# Context variable to track the current trace category for hierarchical nesting
-_current_category: ContextVar[str | None] = ContextVar("current_category", default=None)
-
-
-@final
-@dataclass(frozen=True)
-class TraceEvent:
-    name: str
-    start_us: int
-    duration_us: int
-    rank: int
-    category: str
-
-
-@final
-@dataclass
-class CategoryStats:
-    total_us: int = 0
-    count: int = 0
-    min_us: int = 0
-    max_us: int = 0
-
-    def add(self, duration_us: int) -> None:
-        if self.count == 0:
-            self.min_us = duration_us
-            self.max_us = duration_us
-        else:
-            self.min_us = min(self.min_us, duration_us)
-            self.max_us = max(self.max_us, duration_us)
-        self.total_us += duration_us
-        self.count += 1
-
-    @property
-    def avg_us(self) -> float:
-        return self.total_us / self.count if self.count > 0 else 0.0
-
-
-@final
-@dataclass
-class TraceStats:
-    total_wall_time_us: int = 0
-    by_category: dict[str, CategoryStats] = field(default_factory=dict)
-    by_rank: dict[int, dict[str, CategoryStats]] = field(default_factory=dict)
-
-
-# Global trace buffer - each rank accumulates traces here
-_trace_buffer: list[TraceEvent] = []
-
-
-def is_tracing_enabled() -> bool:
-    """Check if tracing is enabled via environment variable."""
-    return EXO_TRACING_ENABLED
-
-
-def _record_span(
-    name: str, start_us: int, duration_us: int, rank: int, category: str
-) -> None:
-    _trace_buffer.append(
-        TraceEvent(
-            name=name,
-            start_us=start_us,
-            duration_us=duration_us,
-            rank=rank,
-            category=category,
-        )
-    )
-
-
-@contextmanager
-def trace(
-    name: str,
-    rank: int,
-    category: str = "compute",
-) -> Generator[None, None, None]:
-    """Context manager to trace any operation.
-
-    Nested traces automatically inherit the parent category, creating hierarchical
-    categories like "sync/compute" or "async/comms".
-
-    Args:
-        name: Name of the operation (e.g., "recv 0", "send 1", "joint_blocks")
-        rank: This rank's ID
-        category: Category for grouping in trace viewer ("comm", "compute", "step")
-
-    Example:
-        with trace(f"sync {t}", rank, "sync"):
-            with trace("joint_blocks", rank, "compute"):
-                # Recorded with category "sync/compute"
-                hidden_states = some_computation(...)
-    """
-    if not is_tracing_enabled():
-        yield
-        return
-
-    # Combine with parent category if nested
-    parent = _current_category.get()
-    full_category = f"{parent}/{category}" if parent else category
-
-    # Set as current for nested traces
-    token = _current_category.set(full_category)
-
-    try:
-        start_us = int(time.time() * 1_000_000)
-        start_perf = time.perf_counter()
-        yield
-        duration_us = int((time.perf_counter() - start_perf) * 1_000_000)
-        _record_span(name, start_us, duration_us, rank, full_category)
-    finally:
-        _current_category.reset(token)
-
-
-def get_trace_buffer() -> list[TraceEvent]:
-    return list(_trace_buffer)
-
-
-def clear_trace_buffer() -> None:
-    _trace_buffer.clear()
-
-
-def export_trace(traces: list[TraceEvent], output_path: Path) -> None:
-    trace_events: list[dict[str, object]] = []
-
-    for event in traces:
-        # Chrome trace format uses "X" for complete events (with duration)
-        chrome_event: dict[str, object] = {
-            "name": event.name,
-            "cat": event.category,
-            "ph": "X",
-            "ts": event.start_us,
-            "dur": event.duration_us,
-            "pid": 0,
-            "tid": event.rank,
-            "args": {"rank": event.rank},
-        }
-        trace_events.append(chrome_event)
-
-    ranks_seen = set(t.rank for t in traces)
-    for rank in ranks_seen:
-        trace_events.append(
-            {
-                "name": "thread_name",
-                "ph": "M",  # Metadata event
-                "pid": 0,
-                "tid": rank,
-                "args": {"name": f"Rank {rank}"},
-            }
-        )
-
-    chrome_trace = {"traceEvents": trace_events}
-
-    try:
-        output_path.parent.mkdir(parents=True, exist_ok=True)
-        with open(output_path, "w") as f:
-            json.dump(chrome_trace, f, indent=2)
-    except OSError as e:
-        logger.warning("Failed to export trace to %s: %s", output_path, e)
-
-
-def export_local_traces(rank: int) -> None:
-    if not is_tracing_enabled():
-        return
-
-    local_traces = get_trace_buffer()
-    if local_traces:
-        output_path = Path.home() / ".exo" / "traces" / f"trace_{rank}.json"
-        try:
-            export_trace(local_traces, output_path)
-        except Exception as e:
-            logger.warning("Failed to export local traces for rank %d: %s", rank, e)
-
-    clear_trace_buffer()
-
-
-def merge_trace_files(trace_dir: Path | None = None) -> Path | None:
-    if trace_dir is None:
-        trace_dir = Path.home() / ".exo" / "traces"
-
-    if not trace_dir.exists():
-        return None
-
-    trace_files = sorted(trace_dir.glob("trace_*.json"))
-
-    if not trace_files:
-        return None
-
-    merged_events: list[dict[str, object]] = []
-    for trace_file in trace_files:
-        file_rank = int(trace_file.stem.split("_")[1])
-
-        with open(trace_file) as f:
-            raw = f.read()
-            data = cast(dict[str, list[dict[str, object]]], json.loads(raw))
-            events: list[dict[str, object]] = data.get("traceEvents", [])
-            for event in events:
-                event["tid"] = file_rank
-                if "args" in event and isinstance(event["args"], dict):
-                    event["args"]["rank"] = file_rank
-            merged_events.extend(events)
-
-    output_path = Path.home() / ".exo" / "traces" / "merged_trace.json"
-    try:
-        output_path.parent.mkdir(parents=True, exist_ok=True)
-        with open(output_path, "w") as f:
-            json.dump({"traceEvents": merged_events}, f, indent=2)
-    except OSError as e:
-        logger.warning("Failed to write merged trace to %s: %s", output_path, e)
-        return None
-
-    return output_path
-
-
-def _format_duration(us: int | float) -> str:
-    if us < 1000:
-        return f"{us:.0f}µs"
-    elif us < 1_000_000:
-        return f"{us / 1000:.2f}ms"
-    else:
-        return f"{us / 1_000_000:.2f}s"
-
-
-def load_trace_file(path: Path) -> list[TraceEvent]:
-    """Load a Chrome Trace Format JSON file into TraceEvent objects."""
-    with open(path) as f:
-        data = cast(dict[str, list[dict[str, object]]], json.load(f))
-
-    events = data.get("traceEvents", [])
-    traces: list[TraceEvent] = []
-
-    for event in events:
-        # Skip metadata events
-        if event.get("ph") == "M":
-            continue
-
-        name = str(event.get("name", ""))
-        category = str(event.get("cat", ""))
-        ts_value = event.get("ts", 0)
-        dur_value = event.get("dur", 0)
-        tid_value = event.get("tid", 0)
-        start_us = int(ts_value) if isinstance(ts_value, (int, float, str)) else 0
-        duration_us = int(dur_value) if isinstance(dur_value, (int, float, str)) else 0
-
-        # Get rank from tid or args
-        rank = int(tid_value) if isinstance(tid_value, (int, float, str)) else 0
-        args = event.get("args")
-        if isinstance(args, dict):
-            args_dict = cast(dict[str, object], args)
-            rank_from_args = args_dict.get("rank")
-            if isinstance(rank_from_args, (int, float, str)):
-                rank = int(rank_from_args)
-
-        traces.append(
-            TraceEvent(
-                name=name,
-                start_us=start_us,
-                duration_us=duration_us,
-                rank=rank,
-                category=category,
-            )
-        )
-
-    return traces
-
-
-def compute_stats(traces: list[TraceEvent]) -> TraceStats:
-    """Compute comprehensive statistics from trace events."""
-    stats = TraceStats()
-
-    if not traces:
-        return stats
-
-    # Calculate wall time from earliest start to latest end
-    min_start = min(t.start_us for t in traces)
-    max_end = max(t.start_us + t.duration_us for t in traces)
-    stats.total_wall_time_us = max_end - min_start
-
-    # Initialize nested dicts
-    by_category: dict[str, CategoryStats] = defaultdict(CategoryStats)
-    by_rank: dict[int, dict[str, CategoryStats]] = defaultdict(
-        lambda: defaultdict(CategoryStats)
-    )
-
-    for event in traces:
-        # By category
-        by_category[event.category].add(event.duration_us)
-
-        # By rank and category
-        by_rank[event.rank][event.category].add(event.duration_us)
-
-    stats.by_category = dict(by_category)
-    stats.by_rank = {k: dict(v) for k, v in by_rank.items()}
-
-    return stats
-
-
-def print_stats(stats: TraceStats) -> None:
-    """Print formatted trace statistics."""
-    print("=== Trace Statistics ===")
-    print()
-    print(f"Wall Time: {_format_duration(stats.total_wall_time_us)}")
-    print()
-
-    # Parse hierarchical categories (e.g., "sync/compute" -> phase="sync", subcat="compute")
-    if stats.by_category:
-        phases: dict[str, dict[str, CategoryStats]] = defaultdict(dict)
-        has_hierarchical = False
-
-        for cat, cat_stats in stats.by_category.items():
-            if "/" in cat:
-                phase, subcat = cat.split("/", 1)
-                phases[phase][subcat] = cat_stats
-                has_hierarchical = True
-            else:
-                phases[cat]["_total"] = cat_stats
-
-        if has_hierarchical:
-            print("By Phase:")
-            for phase in sorted(phases.keys()):
-                subcats = phases[phase]
-                # Skip phases that only have _total (non-hierarchical top-level categories)
-                non_total_subcats = {k: v for k, v in subcats.items() if k != "_total"}
-                if not non_total_subcats:
-                    continue
-
-                phase_total = sum(s.total_us for s in non_total_subcats.values())
-                print(f"  {phase}:")
-                for subcat, subcat_stats in sorted(
-                    non_total_subcats.items(),
-                    key=lambda x: x[1].total_us,
-                    reverse=True,
-                ):
-                    pct = (
-                        subcat_stats.total_us / phase_total * 100 if phase_total else 0
-                    )
-                    # Use parent phase's step count for per-step average
-                    phase_step_count = subcats.get("_total", CategoryStats()).count
-                    if phase_step_count > 0:
-                        avg_per_step = subcat_stats.total_us / phase_step_count
-                    else:
-                        avg_per_step = subcat_stats.avg_us  # fallback
-                    print(
-                        f"    {subcat:12s} {_format_duration(subcat_stats.total_us):>10s} "
-                        f"({pct:5.1f}%)  avg: {_format_duration(avg_per_step)}"
-                    )
-            print()
-        else:
-            # Fall back to flat category display if no hierarchical categories
-            print("By Category:")
-            total_time = sum(c.total_us for c in stats.by_category.values())
-            for category, cat_stats in sorted(
-                stats.by_category.items(), key=lambda x: x[1].total_us, reverse=True
-            ):
-                pct = (cat_stats.total_us / total_time * 100) if total_time > 0 else 0
-                print(
-                    f"  {category:12s} {_format_duration(cat_stats.total_us):>10s} "
-                    f"({pct:5.1f}%)  avg: {_format_duration(cat_stats.avg_us):>8s}  "
-                    f"count: {cat_stats.count}"
-                )
-            print()
-
-    # By Rank
-    if stats.by_rank:
-        print("By Rank:")
-        for rank in sorted(stats.by_rank.keys()):
-            rank_stats = stats.by_rank[rank]
-            print(f"  Rank {rank}:")
-
-            # Parse hierarchical categories for this rank
-            rank_phases: dict[str, dict[str, CategoryStats]] = defaultdict(dict)
-            has_hierarchical = False
-            for cat, cat_stats in rank_stats.items():
-                if "/" in cat:
-                    phase, subcat = cat.split("/", 1)
-                    rank_phases[phase][subcat] = cat_stats
-                    has_hierarchical = True
-                else:
-                    rank_phases[cat]["_total"] = cat_stats
-
-            if has_hierarchical:
-                for phase in sorted(rank_phases.keys()):
-                    subcats = rank_phases[phase]
-                    non_total_subcats = {
-                        k: v for k, v in subcats.items() if k != "_total"
-                    }
-                    if not non_total_subcats:
-                        continue
-
-                    phase_total = sum(s.total_us for s in non_total_subcats.values())
-                    print(f"    {phase}:")
-                    for subcat, subcat_stats in sorted(
-                        non_total_subcats.items(),
-                        key=lambda x: x[1].total_us,
-                        reverse=True,
-                    ):
-                        pct = (
-                            subcat_stats.total_us / phase_total * 100
-                            if phase_total
-                            else 0
-                        )
-                        # Use parent phase's step count for per-step average
-                        phase_step_count = subcats.get("_total", CategoryStats()).count
-                        if phase_step_count > 0:
-                            avg_per_step = subcat_stats.total_us / phase_step_count
-                        else:
-                            avg_per_step = subcat_stats.avg_us  # fallback
-                        print(
-                            f"      {subcat:12s} {_format_duration(subcat_stats.total_us):>10s} "
-                            f"({pct:5.1f}%)  avg: {_format_duration(avg_per_step)}"
-                        )
-            else:
-                # Flat display fallback
-                for category, cat_stats in sorted(
-                    rank_stats.items(), key=lambda x: x[1].total_us, reverse=True
-                ):
-                    print(f"    {category}: {_format_duration(cat_stats.total_us)}")
-        print()
-
-
-if __name__ == "__main__":
-    import sys
-
-    path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("trace.json")
-
-    if not path.exists():
-        print(f"Error: File not found: {path}")
-        sys.exit(1)
-
-    traces = load_trace_file(path)
-    if not traces:
-        print("No trace events found in file.")
-        sys.exit(0)
-
-    computed_stats = compute_stats(traces)
-    print_stats(computed_stats)

src/exo/shared/types/api.py

@@ -11,7 +11,7 @@ from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding, ShardMetadata
-from exo.utils.pydantic_ext import CamelCaseModel, ConfigDict, TaggedModel
+from exo.utils.pydantic_ext import CamelCaseModel
 
 FinishReason = Literal[
     "stop", "length", "tool_calls", "content_filter", "function_call", "error"
@@ -116,8 +116,8 @@ class Usage(BaseModel):
     prompt_tokens: int
     completion_tokens: int
     total_tokens: int
-    prompt_tokens_details: PromptTokensDetails
-    completion_tokens_details: CompletionTokensDetails
+    prompt_tokens_details: PromptTokensDetails | None = None
+    completion_tokens_details: CompletionTokensDetails | None = None
 
 
 class StreamingChoiceResponse(BaseModel):
@@ -170,13 +170,7 @@ class BenchChatCompletionResponse(ChatCompletionResponse):
     generation_stats: GenerationStats | None = None
 
 
-class StreamOptions(BaseModel):
-    include_usage: bool = False
-
-
-class ChatCompletionTaskParams(TaggedModel):
-    model_config = ConfigDict(extra="ignore")
-
+class ChatCompletionTaskParams(BaseModel):
     model: str
     frequency_penalty: float | None = None
     messages: list[ChatCompletionMessage]
@@ -190,7 +184,6 @@ class ChatCompletionTaskParams(TaggedModel):
     seed: int | None = None
     stop: str | list[str] | None = None
     stream: bool = False
-    stream_options: StreamOptions | None = None
     temperature: float | None = None
     top_p: float | None = None
     tools: list[dict[str, Any]] | None = None
@@ -373,45 +366,3 @@ class StartDownloadResponse(CamelCaseModel):
 
 class DeleteDownloadResponse(CamelCaseModel):
     command_id: CommandId
-
-
-class TraceEventResponse(CamelCaseModel):
-    name: str
-    start_us: int
-    duration_us: int
-    rank: int
-    category: str
-
-
-class TraceResponse(CamelCaseModel):
-    task_id: str
-    traces: list[TraceEventResponse]
-
-
-class TraceCategoryStats(CamelCaseModel):
-    total_us: int
-    count: int
-    min_us: int
-    max_us: int
-    avg_us: float
-
-
-class TraceRankStats(CamelCaseModel):
-    by_category: dict[str, TraceCategoryStats]
-
-
-class TraceStatsResponse(CamelCaseModel):
-    task_id: str
-    total_wall_time_us: int
-    by_category: dict[str, TraceCategoryStats]
-    by_rank: dict[int, TraceRankStats]
-
-
-class TraceListItem(CamelCaseModel):
-    task_id: str
-    created_at: str
-    file_size: int
-
-
-class TraceListResponse(CamelCaseModel):
-    traces: list[TraceListItem]

src/exo/shared/types/chunks.py

@@ -2,7 +2,7 @@ from collections.abc import Generator
 from typing import Any, Literal
 
 from exo.shared.models.model_cards import ModelId
-from exo.shared.types.api import GenerationStats, ImageGenerationStats, Usage
+from exo.shared.types.api import GenerationStats, ImageGenerationStats
 from exo.utils.pydantic_ext import TaggedModel
 
 from .api import FinishReason
@@ -17,7 +17,6 @@ class BaseChunk(TaggedModel):
 class TokenChunk(BaseChunk):
     text: str
     token_id: int
-    usage: Usage | None
     finish_reason: Literal["stop", "length", "content_filter"] | None = None
     stats: GenerationStats | None = None
 
@@ -29,7 +28,6 @@ class ErrorChunk(BaseChunk):
 
 class ToolCallChunk(BaseChunk):
     tool_calls: list[ToolCallItem]
-    usage: Usage | None
     finish_reason: Literal["tool_calls"] = "tool_calls"
     stats: GenerationStats | None = None
 

src/exo/shared/types/commands.py

@@ -2,7 +2,6 @@ from pydantic import Field
 
 from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.api import (
-    BenchChatCompletionTaskParams,
     ChatCompletionTaskParams,
     ImageEditsInternalParams,
     ImageGenerationTaskParams,
@@ -23,7 +22,7 @@ class TestCommand(BaseCommand):
 
 
 class ChatCompletion(BaseCommand):
-    request_params: ChatCompletionTaskParams | BenchChatCompletionTaskParams
+    request_params: ChatCompletionTaskParams
 
 
 class ImageGeneration(BaseCommand):

src/exo/shared/types/events.py

@@ -1,5 +1,4 @@
 from datetime import datetime
-from typing import final
 
 from pydantic import Field
 
@@ -11,7 +10,7 @@ from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo
-from exo.utils.pydantic_ext import CamelCaseModel, FrozenModel, TaggedModel
+from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
 
 class EventId(Id):
@@ -110,22 +109,6 @@ class TopologyEdgeDeleted(BaseEvent):
     conn: Connection
 
 
-@final
-class TraceEventData(FrozenModel):
-    name: str
-    start_us: int
-    duration_us: int
-    rank: int
-    category: str
-
-
-@final
-class TracesCollected(BaseEvent):
-    task_id: TaskId
-    rank: int
-    traces: list[TraceEventData]
-
-
 Event = (
     TestEvent
     | TaskCreated
@@ -144,7 +127,6 @@ Event = (
     | InputChunkReceived
     | TopologyEdgeCreated
     | TopologyEdgeDeleted
-    | TracesCollected
 )
 
 

src/exo/shared/types/tasks.py

@@ -3,7 +3,6 @@ from enum import Enum
 from pydantic import Field
 
 from exo.shared.types.api import (
-    BenchChatCompletionTaskParams,
     ChatCompletionTaskParams,
     ImageEditsInternalParams,
     ImageGenerationTaskParams,
@@ -55,7 +54,7 @@ class StartWarmup(BaseTask):  # emitted by Worker
 
 class ChatCompletion(BaseTask):  # emitted by Master
     command_id: CommandId
-    task_params: ChatCompletionTaskParams | BenchChatCompletionTaskParams
+    task_params: ChatCompletionTaskParams
 
     error_type: str | None = Field(default=None)
     error_message: str | None = Field(default=None)

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

@@ -6,7 +6,6 @@ from exo.shared.types.api import (
     GenerationStats,
     ImageGenerationStats,
     ToolCallItem,
-    Usage,
 )
 from exo.utils.pydantic_ext import TaggedModel
 
@@ -25,7 +24,6 @@ class GenerationResponse(BaseRunnerResponse):
     # logprobs: list[float] | None = None # too big. we can change to be top-k
     finish_reason: FinishReason | None = None
     stats: GenerationStats | None = None
-    usage: Usage | None
 
 
 class ImageGenerationResponse(BaseRunnerResponse):
@@ -59,7 +57,6 @@ class PartialImageResponse(BaseRunnerResponse):
 
 class ToolCallResponse(BaseRunnerResponse):
     tool_calls: list[ToolCallItem]
-    usage: Usage | None
 
 
 class FinishedResponse(BaseRunnerResponse):

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

@@ -57,8 +57,62 @@ class PipelineShardMetadata(BaseShardMetadata):
 
     Layers are represented as a half-open interval [start_layer, end_layer),
     where start_layer is inclusive and end_layer is exclusive.
+
+    CFG parallelism fields:
+    - cfg_rank: 0 = positive branch, 1 = negative branch (or 0 if no CFG parallel)
+    - cfg_world_size: 1 = sequential CFG, 2 = parallel CFG
+
+    Communication rank fields (explicit to support ring topology):
+    - next_pipeline_device: device to send to in pipeline forward pass
+    - prev_pipeline_device: device to receive from in pipeline forward pass
+    - cfg_peer_device: device for CFG exchange (last stage only)
+    - first_pipeline_device: device of first stage in same CFG group (for latent return)
     """
 
+    cfg_rank: int = 0
+    cfg_world_size: int = 1
+
+    # Explicit pipeline position (CFG group 1 uses reversed pipeline order)
+    explicit_pipeline_rank: int | None = None
+
+    next_pipeline_device: int | None = None
+    prev_pipeline_device: int | None = None
+    cfg_peer_device: int | None = None
+    first_pipeline_device: int | None = None
+    last_pipeline_device: int | None = None
+
+    @property
+    def pipeline_world_size(self) -> int:
+        return self.world_size // self.cfg_world_size
+
+    @property
+    def pipeline_rank(self) -> int:
+        if self.explicit_pipeline_rank is not None:
+            return self.explicit_pipeline_rank
+        return self.device_rank % self.pipeline_world_size
+
+    @property
+    def is_pipeline_first(self) -> bool:
+        return self.pipeline_rank == 0
+
+    @property
+    def is_pipeline_last(self) -> bool:
+        return self.pipeline_rank == self.pipeline_world_size - 1
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                self.model_card.model_id,
+                self.start_layer,
+                self.end_layer,
+                self.n_layers,
+                self.device_rank,
+                self.world_size,
+                self.cfg_rank,
+                self.cfg_world_size,
+            )
+        )
+
 
 class TensorShardMetadata(BaseShardMetadata):
     pass

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

@@ -37,7 +37,12 @@ class DistributedImageModel:
         config = get_config_for_model(model_id)
         adapter = create_adapter_for_model(config, model_id, local_path, quantize)
 
-        if group is not None:
+        has_layer_sharding = (
+            shard_metadata.start_layer != 0
+            or shard_metadata.end_layer != shard_metadata.n_layers
+        )
+
+        if group is not None and has_layer_sharding:
             adapter.slice_transformer_blocks(
                 start_layer=shard_metadata.start_layer,
                 end_layer=shard_metadata.end_layer,

src/exo/worker/engines/image/models/base.py

@@ -86,6 +86,27 @@ class PromptData(ABC):
         """
         ...
 
+    @abstractmethod
+    def get_cfg_branch_data(
+        self, positive: bool
+    ) -> tuple[mx.array, mx.array | None, mx.array | None, mx.array | None]:
+        """Get embeddings for a single CFG branch (positive or negative).
+
+        Used for sequential CFG and CFG parallel modes where we process
+        one branch at a time instead of batching.
+
+        Args:
+            positive: True for positive prompt, False for negative prompt
+
+        Returns:
+            Tuple of:
+            - embeds: [1, seq, hidden] prompt embeddings
+            - mask: [1, seq] attention mask or None
+            - pooled: [1, hidden] pooled embeddings or None
+            - conditioning_latents: [1, latent_seq, latent_dim] or None
+        """
+        ...
+
 
 class ModelAdapter(ABC, Generic[ModelT, TransformerT]):
     _config: ImageModelConfig

src/exo/worker/engines/image/models/flux/adapter.py

@@ -64,6 +64,12 @@ class FluxPromptData(PromptData):
     ) -> tuple[mx.array, mx.array, mx.array | None, mx.array | None] | None:
         return None
 
+    def get_cfg_branch_data(
+        self, positive: bool
+    ) -> tuple[mx.array, mx.array | None, mx.array | None, mx.array | None]:
+        """Flux doesn't use CFG, but we return positive data for compatibility."""
+        return (self._prompt_embeds, None, self._pooled_prompt_embeds, None)
+
 
 class FluxModelAdapter(ModelAdapter[Flux1, Transformer]):
     def __init__(

src/exo/worker/engines/image/models/qwen/adapter.py

@@ -133,6 +133,24 @@ class QwenPromptData(PromptData):
 
         return batched_embeds, batched_mask, None, cond_latents
 
+    def get_cfg_branch_data(
+        self, positive: bool
+    ) -> tuple[mx.array, mx.array | None, mx.array | None, mx.array | None]:
+        if positive:
+            return (
+                self._prompt_embeds,
+                self._prompt_mask,
+                None,
+                self.conditioning_latents,
+            )
+        else:
+            return (
+                self._negative_prompt_embeds,
+                self._negative_prompt_mask,
+                None,
+                self.conditioning_latents,
+            )
+
 
 class QwenModelAdapter(ModelAdapter[QwenImage, QwenTransformer]):
     """Adapter for Qwen-Image model.

src/exo/worker/engines/image/models/qwen/edit_adapter.py

@@ -153,6 +153,24 @@ class QwenEditPromptData(PromptData):
 
         return batched_embeds, batched_mask, None, batched_cond_latents
 
+    def get_cfg_branch_data(
+        self, positive: bool
+    ) -> tuple[mx.array, mx.array | None, mx.array | None, mx.array | None]:
+        if positive:
+            return (
+                self._prompt_embeds,
+                self._prompt_mask,
+                None,
+                self._conditioning_latents,
+            )
+        else:
+            return (
+                self._negative_prompt_embeds,
+                self._negative_prompt_mask,
+                None,
+                self._conditioning_latents,
+            )
+
 
 class QwenEditModelAdapter(ModelAdapter[QwenImageEdit, QwenTransformer]):
     """Adapter for Qwen-Image-Edit model.

src/exo/worker/engines/image/pipeline/runner.py

@@ -1,16 +1,13 @@
+from collections.abc import Iterator
+from dataclasses import dataclass
 from math import ceil
-from typing import Any, Optional
+from typing import Any, Optional, final
 
 import mlx.core as mx
 from mflux.models.common.config.config import Config
 from mflux.utils.exceptions import StopImageGenerationException
 from tqdm import tqdm
 
-from exo.shared.tracing import (
-    clear_trace_buffer,
-    is_tracing_enabled,
-    trace,
-)
 from exo.shared.types.worker.shards import PipelineShardMetadata
 from exo.worker.engines.image.config import ImageModelConfig
 from exo.worker.engines.image.models.base import (
@@ -25,6 +22,16 @@ from exo.worker.engines.image.pipeline.block_wrapper import (
 )
 
 
+@final
+@dataclass
+class CfgBranch:
+    positive: bool
+    embeds: mx.array
+    mask: mx.array | None
+    pooled: mx.array | None
+    cond_latents: mx.array | None
+
+
 def calculate_patch_heights(
     latent_height: int, num_patches: int
 ) -> tuple[list[int], int]:
@@ -77,22 +84,11 @@ class DiffusionRunner:
         self.adapter = adapter
         self.group = group
 
-        if group is None:
-            self.rank = 0
-            self.world_size = 1
-            self.next_rank = 0
-            self.prev_rank = 0
-            self.start_layer = 0
-            self.end_layer = config.total_blocks
-        else:
-            self.rank = shard_metadata.device_rank
-            self.world_size = shard_metadata.world_size
-            self.next_rank = (self.rank + 1) % self.world_size
-            self.prev_rank = (self.rank - 1 + self.world_size) % self.world_size
-            self.start_layer = shard_metadata.start_layer
-            self.end_layer = shard_metadata.end_layer
+        self._init_cfg_topology(shard_metadata)
 
-        self.num_patches = num_patches if num_patches else max(1, self.world_size)
+        self.num_patches = (
+            num_patches if num_patches else max(1, self.pipeline_world_size)
+        )
 
         self.total_joint = config.joint_block_count
         self.total_single = config.single_block_count
@@ -102,6 +98,48 @@ class DiffusionRunner:
 
         self._compute_assigned_blocks()
 
+    def _init_cfg_topology(self, shard_metadata: PipelineShardMetadata) -> None:
+        """Initialize CFG and pipeline topology from shard metadata."""
+        if self.group is None:
+            self.rank = 0
+            self.world_size = 1
+            self.start_layer = 0
+            self.end_layer = self.config.total_blocks
+
+            self.cfg_rank = 0
+            self.cfg_world_size = 1
+            self.cfg_parallel = False
+
+            self.pipeline_world_size = 1
+            self.pipeline_rank = 0
+
+            self.next_pipeline_rank: int | None = None
+            self.prev_pipeline_rank: int | None = None
+            self.cfg_peer_rank: int | None = None
+            self.first_pipeline_rank: int = 0
+            self.last_pipeline_rank: int = 0
+        else:
+            self.rank = shard_metadata.device_rank
+            self.world_size = shard_metadata.world_size
+            self.start_layer = shard_metadata.start_layer
+            self.end_layer = shard_metadata.end_layer
+
+            self.cfg_rank = shard_metadata.cfg_rank
+            self.cfg_world_size = shard_metadata.cfg_world_size
+            self.cfg_parallel = self.cfg_world_size > 1
+
+            self.pipeline_world_size = shard_metadata.pipeline_world_size
+            self.pipeline_rank = shard_metadata.pipeline_rank
+
+            self.next_pipeline_rank = shard_metadata.next_pipeline_device
+            self.prev_pipeline_rank = shard_metadata.prev_pipeline_device
+            self.cfg_peer_rank = shard_metadata.cfg_peer_device
+
+            assert shard_metadata.first_pipeline_device is not None
+            assert shard_metadata.last_pipeline_device is not None
+            self.first_pipeline_rank = shard_metadata.first_pipeline_device
+            self.last_pipeline_rank = shard_metadata.last_pipeline_device
+
     def _compute_assigned_blocks(self) -> None:
         """Determine which joint/single blocks this stage owns."""
         start = self.start_layer
@@ -138,11 +176,11 @@ class DiffusionRunner:
 
     @property
     def is_first_stage(self) -> bool:
-        return self.rank == 0
+        return self.pipeline_rank == 0
 
     @property
     def is_last_stage(self) -> bool:
-        return self.rank == self.world_size - 1
+        return self.pipeline_rank == self.pipeline_world_size - 1
 
     @property
     def is_distributed(self) -> bool:
@@ -153,6 +191,97 @@ class DiffusionRunner:
             return self._guidance_override
         return self.config.guidance_scale
 
+    def _get_cfg_branches(self, prompt_data: PromptData) -> Iterator[CfgBranch]:
+        """Yield the CFG branches this node should process.
+
+        - No CFG: yields one branch (positive)
+        - CFG parallel: yields one branch (our assigned branch)
+        - Sequential CFG: yields two branches (positive, then negative)
+        """
+        if not self.adapter.needs_cfg:
+            embeds, mask, pooled, cond = prompt_data.get_cfg_branch_data(positive=True)
+            yield CfgBranch(
+                positive=True,
+                embeds=embeds,
+                mask=mask,
+                pooled=pooled,
+                cond_latents=cond,
+            )
+        elif self.cfg_parallel:
+            positive = self.cfg_rank == 0
+            embeds, mask, pooled, cond = prompt_data.get_cfg_branch_data(positive)
+            yield CfgBranch(
+                positive=positive,
+                embeds=embeds,
+                mask=mask,
+                pooled=pooled,
+                cond_latents=cond,
+            )
+        else:
+            pos_embeds, pos_mask, pos_pooled, pos_cond = (
+                prompt_data.get_cfg_branch_data(positive=True)
+            )
+            yield CfgBranch(
+                positive=True,
+                embeds=pos_embeds,
+                mask=pos_mask,
+                pooled=pos_pooled,
+                cond_latents=pos_cond,
+            )
+            neg_embeds, neg_mask, neg_pooled, neg_cond = (
+                prompt_data.get_cfg_branch_data(positive=False)
+            )
+            yield CfgBranch(
+                positive=False,
+                embeds=neg_embeds,
+                mask=neg_mask,
+                pooled=neg_pooled,
+                cond_latents=neg_cond,
+            )
+
+    def _combine_cfg_results(self, results: list[tuple[bool, mx.array]]) -> mx.array:
+        if len(results) == 1:
+            positive, noise = results[0]
+            if self.cfg_parallel and self.is_last_stage:
+                # TODO(ciaran): try to remove
+                mx.eval(noise)
+                return self._exchange_and_apply_guidance(noise, positive)
+            return noise
+
+        noise_neg = next(n for p, n in results if not p)
+        noise_pos = next(n for p, n in results if p)
+        return self._apply_guidance(noise_pos, noise_neg)
+
+    def _exchange_and_apply_guidance(
+        self, noise: mx.array, is_positive: bool
+    ) -> mx.array:
+        assert self.group is not None
+        assert self.cfg_peer_rank is not None
+
+        if is_positive:
+            noise = mx.distributed.send(noise, self.cfg_peer_rank, group=self.group)
+            mx.async_eval(noise)
+            noise_neg = mx.distributed.recv_like(
+                noise, self.cfg_peer_rank, group=self.group
+            )
+            mx.eval(noise_neg)
+            noise_pos = noise
+        else:
+            noise_pos = mx.distributed.recv_like(
+                noise, self.cfg_peer_rank, group=self.group
+            )
+            mx.eval(noise_pos)
+            noise = mx.distributed.send(noise, self.cfg_peer_rank, group=self.group)
+            mx.async_eval(noise)
+            noise_neg = noise
+
+        return self._apply_guidance(noise_pos, noise_neg)
+
+    def _apply_guidance(self, noise_pos: mx.array, noise_neg: mx.array) -> mx.array:
+        scale = self._get_effective_guidance_scale()
+        assert scale is not None
+        return self.adapter.apply_guidance(noise_pos, noise_neg, scale)
+
     def _ensure_wrappers(
         self,
         text_seq_len: int,
@@ -329,7 +458,6 @@ class DiffusionRunner:
             capture_steps = set()
 
         self._reset_all_caches()
-        clear_trace_buffer()
 
         time_steps = tqdm(range(runtime_config.num_inference_steps))
 
@@ -470,23 +598,23 @@ class DiffusionRunner:
     ) -> mx.array:
         if self.group is None:
             return self._single_node_step(t, config, latents, prompt_data)
-        elif t < config.init_time_step + num_sync_steps:
-            with trace(name=f"sync {t}", rank=self.rank, category="sync"):
-                return self._sync_pipeline_step(
-                    t,
-                    config,
-                    latents,
-                    prompt_data,
-                )
+        elif (
+            self.pipeline_world_size == 1 or t < config.init_time_step + num_sync_steps
+        ):
+            return self._sync_pipeline_step(
+                t,
+                config,
+                latents,
+                prompt_data,
+            )
         else:
-            with trace(name=f"async {t}", rank=self.rank, category="async"):
-                return self._async_pipeline_step(
-                    t,
-                    config,
-                    latents,
-                    prompt_data,
-                    is_first_async_step=t == config.init_time_step + num_sync_steps,
-                )
+            return self._async_pipeline_step(
+                t,
+                config,
+                latents,
+                prompt_data,
+                is_first_async_step=t == config.init_time_step + num_sync_steps,
+            )
 
     def _single_node_step(
         self,
@@ -496,42 +624,29 @@ class DiffusionRunner:
         prompt_data: PromptData,
     ) -> mx.array:
         cond_image_grid = prompt_data.cond_image_grid
-        needs_cfg = self.adapter.needs_cfg
+        results: list[tuple[bool, mx.array]] = []
+
+        for branch in self._get_cfg_branches(prompt_data):
+            # Reset caches before each branch to ensure no state contamination
+            self._reset_all_caches()
 
-        if needs_cfg:
-            batched_data = prompt_data.get_batched_cfg_data()
-            assert batched_data is not None, "CFG model must provide batched data"
-            prompt_embeds, encoder_mask, batched_pooled, cond_latents = batched_data
             pooled_embeds = (
-                batched_pooled if batched_pooled is not None else prompt_embeds
-            )
-            step_latents = mx.concatenate([latents, latents], axis=0)
-        else:
-            prompt_embeds = prompt_data.prompt_embeds
-            pooled_embeds = prompt_data.pooled_prompt_embeds
-            encoder_mask = prompt_data.get_encoder_hidden_states_mask(positive=True)
-            cond_latents = prompt_data.conditioning_latents
-            step_latents = latents
-
-        noise = self._forward_pass(
-            step_latents,
-            prompt_embeds,
-            pooled_embeds,
-            t=t,
-            config=config,
-            encoder_hidden_states_mask=encoder_mask,
-            cond_image_grid=cond_image_grid,
-            conditioning_latents=cond_latents,
-        )
-
-        if needs_cfg:
-            noise_pos, noise_neg = mx.split(noise, 2, axis=0)
-            guidance_scale = self._get_effective_guidance_scale()
-            assert guidance_scale is not None
-            noise = self.adapter.apply_guidance(
-                noise_pos, noise_neg, guidance_scale=guidance_scale
+                branch.pooled if branch.pooled is not None else branch.embeds
             )
 
+            noise = self._forward_pass(
+                latents,
+                branch.embeds,
+                pooled_embeds,
+                t=t,
+                config=config,
+                encoder_hidden_states_mask=branch.mask,
+                cond_image_grid=cond_image_grid,
+                conditioning_latents=branch.cond_latents,
+            )
+            results.append((branch.positive, noise))
+
+        noise = self._combine_cfg_results(results)
         return config.scheduler.step(noise=noise, timestep=t, latents=latents)  # pyright: ignore[reportAny]
 
     def _create_patches(
@@ -582,7 +697,7 @@ class DiffusionRunner:
             )
 
         text_embeddings = self.adapter.compute_text_embeddings(
-            t, config, pooled_prompt_embeds
+            t, config, pooled_prompt_embeds, hidden_states=hidden_states
         )
         image_rotary_embeddings = self.adapter.compute_rotary_embeddings(
             prompt_embeds,
@@ -594,41 +709,31 @@ class DiffusionRunner:
 
         if self.has_joint_blocks:
             if not self.is_first_stage:
-                with trace(
-                    name=f"recv {self.prev_rank}", rank=self.rank, category="comms"
-                ):
-                    hidden_states = mx.distributed.recv(
-                        (batch_size, num_img_tokens, hidden_dim),
-                        dtype,
-                        self.prev_rank,
-                        group=self.group,
-                    )
-                    encoder_hidden_states = mx.distributed.recv(
-                        (batch_size, text_seq_len, hidden_dim),
-                        dtype,
-                        self.prev_rank,
-                        group=self.group,
-                    )
-                    mx.eval(hidden_states, encoder_hidden_states)
+                assert self.prev_pipeline_rank is not None
+                hidden_states = mx.distributed.recv(
+                    (batch_size, num_img_tokens, hidden_dim),
+                    dtype,
+                    self.prev_pipeline_rank,
+                    group=self.group,
+                )
+                encoder_hidden_states = mx.distributed.recv(
+                    (batch_size, text_seq_len, hidden_dim),
+                    dtype,
+                    self.prev_pipeline_rank,
+                    group=self.group,
+                )
+                mx.eval(hidden_states, encoder_hidden_states)
 
             assert self.joint_block_wrappers is not None
             assert encoder_hidden_states is not None
-            with trace(
-                name="joint_blocks",
-                rank=self.rank,
-                category="compute",
-            ):
-                for wrapper in self.joint_block_wrappers:
-                    wrapper.set_patch(BlockWrapperMode.CACHING)
-                    encoder_hidden_states, hidden_states = wrapper(
-                        hidden_states=hidden_states,
-                        encoder_hidden_states=encoder_hidden_states,
-                        text_embeddings=text_embeddings,
-                        rotary_embeddings=image_rotary_embeddings,
-                    )
-
-                if is_tracing_enabled():
-                    mx.eval(encoder_hidden_states, hidden_states)
+            for wrapper in self.joint_block_wrappers:
+                wrapper.set_patch(BlockWrapperMode.CACHING)
+                encoder_hidden_states, hidden_states = wrapper(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    text_embeddings=text_embeddings,
+                    rotary_embeddings=image_rotary_embeddings,
+                )
 
         if self.owns_concat_stage:
             assert encoder_hidden_states is not None
@@ -639,63 +744,49 @@ class DiffusionRunner:
             if self.has_single_blocks or self.is_last_stage:
                 hidden_states = concatenated
             else:
-                with trace(
-                    name=f"send {self.next_rank}", rank=self.rank, category="comms"
-                ):
-                    concatenated = mx.distributed.send(
-                        concatenated, self.next_rank, group=self.group
-                    )
-                    mx.async_eval(concatenated)
+                assert self.next_pipeline_rank is not None
+                concatenated = mx.distributed.send(
+                    concatenated, self.next_pipeline_rank, group=self.group
+                )
+                mx.async_eval(concatenated)
 
         elif self.has_joint_blocks and not self.is_last_stage:
             assert encoder_hidden_states is not None
-            with trace(name=f"send {self.next_rank}", rank=self.rank, category="comms"):
-                hidden_states = mx.distributed.send(
-                    hidden_states, self.next_rank, group=self.group
-                )
-                encoder_hidden_states = mx.distributed.send(
-                    encoder_hidden_states, self.next_rank, group=self.group
-                )
-                mx.async_eval(hidden_states, encoder_hidden_states)
+            assert self.next_pipeline_rank is not None
+            hidden_states = mx.distributed.send(
+                hidden_states, self.next_pipeline_rank, group=self.group
+            )
+            encoder_hidden_states = mx.distributed.send(
+                encoder_hidden_states, self.next_pipeline_rank, group=self.group
+            )
+            mx.async_eval(hidden_states, encoder_hidden_states)
 
         if self.has_single_blocks:
             if not self.owns_concat_stage and not self.is_first_stage:
-                with trace(
-                    name=f"recv {self.prev_rank}", rank=self.rank, category="comms"
-                ):
-                    hidden_states = mx.distributed.recv(
-                        (batch_size, text_seq_len + num_img_tokens, hidden_dim),
-                        dtype,
-                        self.prev_rank,
-                        group=self.group,
-                    )
-                    mx.eval(hidden_states)
+                assert self.prev_pipeline_rank is not None
+                hidden_states = mx.distributed.recv(
+                    (batch_size, text_seq_len + num_img_tokens, hidden_dim),
+                    dtype,
+                    self.prev_pipeline_rank,
+                    group=self.group,
+                )
+                mx.eval(hidden_states)
 
             assert self.single_block_wrappers is not None
-            with trace(
-                name="single blocks",
-                rank=self.rank,
-                category="compute",
-            ):
-                for wrapper in self.single_block_wrappers:
-                    wrapper.set_patch(BlockWrapperMode.CACHING)
-                    hidden_states = wrapper(
-                        hidden_states=hidden_states,
-                        text_embeddings=text_embeddings,
-                        rotary_embeddings=image_rotary_embeddings,
-                    )
-
-                if is_tracing_enabled():
-                    mx.eval(hidden_states)
+            for wrapper in self.single_block_wrappers:
+                wrapper.set_patch(BlockWrapperMode.CACHING)
+                hidden_states = wrapper(
+                    hidden_states=hidden_states,
+                    text_embeddings=text_embeddings,
+                    rotary_embeddings=image_rotary_embeddings,
+                )
 
             if not self.is_last_stage:
-                with trace(
-                    name=f"send {self.next_rank}", rank=self.rank, category="comms"
-                ):
-                    hidden_states = mx.distributed.send(
-                        hidden_states, self.next_rank, group=self.group
-                    )
-                    mx.async_eval(hidden_states)
+                assert self.next_pipeline_rank is not None
+                hidden_states = mx.distributed.send(
+                    hidden_states, self.next_pipeline_rank, group=self.group
+                )
+                mx.async_eval(hidden_states)
 
         hidden_states = hidden_states[:, text_seq_len:, ...]
 
@@ -716,83 +807,67 @@ class DiffusionRunner:
         kontext_image_ids: mx.array | None = None,
     ) -> mx.array:
         prev_latents = hidden_states
-        needs_cfg = self.adapter.needs_cfg
         cond_image_grid = prompt_data.cond_image_grid
 
         scaled_hidden_states = config.scheduler.scale_model_input(hidden_states, t)  # pyright: ignore[reportAny]
         original_latent_tokens: int = scaled_hidden_states.shape[1]  # pyright: ignore[reportAny]
 
-        if needs_cfg:
-            batched_data = prompt_data.get_batched_cfg_data()
-            assert batched_data is not None, "CFG model must provide batched data"
-            prompt_embeds, encoder_mask, batched_pooled, cond_latents = batched_data
+        results: list[tuple[bool, mx.array]] = []
+
+        for branch in self._get_cfg_branches(prompt_data):
             pooled_embeds = (
-                batched_pooled if batched_pooled is not None else prompt_embeds
+                branch.pooled if branch.pooled is not None else branch.embeds
             )
-            step_latents = mx.concatenate(
-                [scaled_hidden_states, scaled_hidden_states], axis=0
+
+            cond_latents = branch.cond_latents
+            if cond_latents is not None:
+                num_img_tokens: int = original_latent_tokens + cond_latents.shape[1]
+            else:
+                num_img_tokens = original_latent_tokens
+
+            step_latents: mx.array = scaled_hidden_states  # pyright: ignore[reportAny]
+            if self.is_first_stage and cond_latents is not None:
+                step_latents = mx.concatenate([step_latents, cond_latents], axis=1)
+
+            text_seq_len = branch.embeds.shape[1]
+            self._ensure_wrappers(text_seq_len, branch.mask)
+
+            noise = self._run_sync_pass(
+                t,
+                config,
+                step_latents,
+                branch.embeds,
+                pooled_embeds,
+                branch.mask,
+                cond_image_grid,
+                kontext_image_ids,
+                num_img_tokens,
+                original_latent_tokens,
+                cond_latents,
             )
-        else:
-            prompt_embeds = prompt_data.prompt_embeds
-            pooled_embeds = prompt_data.pooled_prompt_embeds
-            encoder_mask = prompt_data.get_encoder_hidden_states_mask(positive=True)
-            cond_latents = prompt_data.conditioning_latents
-            step_latents = scaled_hidden_states  # pyright: ignore[reportAny]
 
-        if cond_latents is not None:
-            num_img_tokens: int = original_latent_tokens + cond_latents.shape[1]
-        else:
-            num_img_tokens = original_latent_tokens
-
-        if self.is_first_stage and cond_latents is not None:
-            step_latents = mx.concatenate([step_latents, cond_latents], axis=1)
-
-        text_seq_len = prompt_embeds.shape[1]
-        self._ensure_wrappers(text_seq_len, encoder_mask)
-
-        noise = self._run_sync_pass(
-            t,
-            config,
-            step_latents,
-            prompt_embeds,
-            pooled_embeds,
-            encoder_mask,
-            cond_image_grid,
-            kontext_image_ids,
-            num_img_tokens,
-            original_latent_tokens,
-            cond_latents,
-        )
+            if self.is_last_stage:
+                assert noise is not None
+                results.append((branch.positive, noise))
 
         if self.is_last_stage:
-            assert noise is not None
-            if needs_cfg:
-                noise_pos, noise_neg = mx.split(noise, 2, axis=0)
-                guidance_scale = self._get_effective_guidance_scale()
-                assert guidance_scale is not None
-                noise = self.adapter.apply_guidance(
-                    noise_pos, noise_neg, guidance_scale
-                )
+            noise = self._combine_cfg_results(results)
 
             hidden_states = config.scheduler.step(  # pyright: ignore[reportAny]
                 noise=noise, timestep=t, latents=prev_latents
             )
 
             if not self.is_first_stage:
-                with trace(name="send 0", rank=self.rank, category="comms"):
-                    hidden_states = mx.distributed.send(
-                        hidden_states, 0, group=self.group
-                    )
-                    mx.async_eval(hidden_states)
+                hidden_states = mx.distributed.send(
+                    hidden_states, self.first_pipeline_rank, group=self.group
+                )
+                mx.async_eval(hidden_states)
 
         elif self.is_first_stage:
-            with trace(
-                name=f"recv {self.world_size - 1}", rank=self.rank, category="comms"
-            ):
-                hidden_states = mx.distributed.recv_like(
-                    prev_latents, src=self.world_size - 1, group=self.group
-                )
-                mx.eval(hidden_states)
+            hidden_states = mx.distributed.recv_like(
+                prev_latents, src=self.last_pipeline_rank, group=self.group
+            )
+            mx.eval(hidden_states)
 
         else:
             hidden_states = prev_latents
@@ -809,39 +884,10 @@ class DiffusionRunner:
         kontext_image_ids: mx.array | None = None,
     ) -> mx.array:
         patch_latents, token_indices = self._create_patches(latents, config)
-        needs_cfg = self.adapter.needs_cfg
         cond_image_grid = prompt_data.cond_image_grid
 
-        if needs_cfg:
-            batched_data = prompt_data.get_batched_cfg_data()
-            assert batched_data is not None, "CFG model must provide batched data"
-            prompt_embeds, encoder_mask, batched_pooled, _ = batched_data
-            pooled_embeds = (
-                batched_pooled if batched_pooled is not None else prompt_embeds
-            )
-        else:
-            prompt_embeds = prompt_data.prompt_embeds
-            pooled_embeds = prompt_data.pooled_prompt_embeds
-            encoder_mask = prompt_data.get_encoder_hidden_states_mask(positive=True)
-
-        text_seq_len = prompt_embeds.shape[1]
-        self._ensure_wrappers(text_seq_len, encoder_mask)
-        self._set_text_seq_len(text_seq_len)
-
-        if self.joint_block_wrappers:
-            for wrapper in self.joint_block_wrappers:
-                wrapper.set_encoder_mask(encoder_mask)
-
-        text_embeddings = self.adapter.compute_text_embeddings(t, config, pooled_embeds)
-        image_rotary_embeddings = self.adapter.compute_rotary_embeddings(
-            prompt_embeds,
-            config,
-            encoder_hidden_states_mask=encoder_mask,
-            cond_image_grid=cond_image_grid,
-            kontext_image_ids=kontext_image_ids,
-        )
-
         prev_patch_latents = [p for p in patch_latents]
+
         encoder_hidden_states: mx.array | None = None
 
         for patch_idx in range(len(patch_latents)):
@@ -852,35 +898,53 @@ class DiffusionRunner:
                 and not self.is_last_stage
                 and not is_first_async_step
             ):
-                with trace(
-                    name=f"recv {self.prev_rank}", rank=self.rank, category="comms"
-                ):
-                    patch = mx.distributed.recv_like(
-                        patch, src=self.prev_rank, group=self.group
-                    )
-                    mx.eval(patch)
+                patch = mx.distributed.recv_like(
+                    patch, src=self.last_pipeline_rank, group=self.group
+                )
+                mx.eval(patch)
 
-            step_patch = mx.concatenate([patch, patch], axis=0) if needs_cfg else patch
+            results: list[tuple[bool, mx.array]] = []
 
-            noise, encoder_hidden_states = self._run_single_patch_pass(
-                patch=step_patch,
-                patch_idx=patch_idx,
-                token_indices=token_indices[patch_idx],
-                prompt_embeds=prompt_embeds,
-                text_embeddings=text_embeddings,
-                image_rotary_embeddings=image_rotary_embeddings,
-                encoder_hidden_states=encoder_hidden_states,
-            )
+            for branch in self._get_cfg_branches(prompt_data):
+                pooled_embeds = (
+                    branch.pooled if branch.pooled is not None else branch.embeds
+                )
+
+                text_seq_len = branch.embeds.shape[1]
+                self._ensure_wrappers(text_seq_len, branch.mask)
+                self._set_text_seq_len(text_seq_len)
+
+                if self.joint_block_wrappers:
+                    for wrapper in self.joint_block_wrappers:
+                        wrapper.set_encoder_mask(branch.mask)
+
+                text_embeddings = self.adapter.compute_text_embeddings(
+                    t, config, pooled_embeds
+                )
+                image_rotary_embeddings = self.adapter.compute_rotary_embeddings(
+                    branch.embeds,
+                    config,
+                    encoder_hidden_states_mask=branch.mask,
+                    cond_image_grid=cond_image_grid,
+                    kontext_image_ids=kontext_image_ids,
+                )
+
+                noise, encoder_hidden_states = self._run_single_patch_pass(
+                    patch=patch,
+                    patch_idx=patch_idx,
+                    token_indices=token_indices[patch_idx],
+                    prompt_embeds=branch.embeds,
+                    text_embeddings=text_embeddings,
+                    image_rotary_embeddings=image_rotary_embeddings,
+                    encoder_hidden_states=encoder_hidden_states,
+                )
+
+                if self.is_last_stage:
+                    assert noise is not None
+                    results.append((branch.positive, noise))
 
             if self.is_last_stage:
-                assert noise is not None
-                if needs_cfg:
-                    noise_pos, noise_neg = mx.split(noise, 2, axis=0)
-                    guidance_scale = self._get_effective_guidance_scale()
-                    assert guidance_scale is not None
-                    noise = self.adapter.apply_guidance(
-                        noise_pos, noise_neg, guidance_scale
-                    )
+                noise = self._combine_cfg_results(results)
 
                 patch_latents[patch_idx] = config.scheduler.step(  # pyright: ignore[reportAny]
                     noise=noise,
@@ -889,13 +953,12 @@ class DiffusionRunner:
                 )
 
                 if not self.is_first_stage and t != config.num_inference_steps - 1:
-                    with trace(
-                        name=f"send {self.next_rank}", rank=self.rank, category="comms"
-                    ):
-                        patch_latents[patch_idx] = mx.distributed.send(
-                            patch_latents[patch_idx], self.next_rank, group=self.group
-                        )
-                        mx.async_eval(patch_latents[patch_idx])
+                    patch_latents[patch_idx] = mx.distributed.send(
+                        patch_latents[patch_idx],
+                        self.first_pipeline_rank,
+                        group=self.group,
+                    )
+                    mx.async_eval(patch_latents[patch_idx])
 
         return mx.concatenate(patch_latents, axis=1)
 
@@ -933,29 +996,24 @@ class DiffusionRunner:
 
         if self.has_joint_blocks:
             if not self.is_first_stage:
+                assert self.prev_pipeline_rank is not None
                 patch_len = patch.shape[1]
-                with trace(
-                    name=f"recv {self.prev_rank}", rank=self.rank, category="comms"
-                ):
-                    patch = mx.distributed.recv(
-                        (batch_size, patch_len, hidden_dim),
-                        patch.dtype,
-                        self.prev_rank,
-                        group=self.group,
-                    )
-                    mx.eval(patch)
+                patch = mx.distributed.recv(
+                    (batch_size, patch_len, hidden_dim),
+                    patch.dtype,
+                    self.prev_pipeline_rank,
+                    group=self.group,
+                )
+                mx.eval(patch)
 
                 if patch_idx == 0:
-                    with trace(
-                        name=f"recv {self.prev_rank}", rank=self.rank, category="comms"
-                    ):
-                        encoder_hidden_states = mx.distributed.recv(
-                            (batch_size, text_seq_len, hidden_dim),
-                            patch.dtype,
-                            self.prev_rank,
-                            group=self.group,
-                        )
-                        mx.eval(encoder_hidden_states)
+                    encoder_hidden_states = mx.distributed.recv(
+                        (batch_size, text_seq_len, hidden_dim),
+                        patch.dtype,
+                        self.prev_pipeline_rank,
+                        group=self.group,
+                    )
+                    mx.eval(encoder_hidden_states)
 
             if self.is_first_stage:
                 patch, encoder_hidden_states = self.adapter.compute_embeddings(
@@ -964,22 +1022,14 @@ class DiffusionRunner:
 
             assert self.joint_block_wrappers is not None
             assert encoder_hidden_states is not None
-            with trace(
-                name=f"joint patch {patch_idx}",
-                rank=self.rank,
-                category="compute",
-            ):
-                for wrapper in self.joint_block_wrappers:
-                    wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token)
-                    encoder_hidden_states, patch = wrapper(
-                        hidden_states=patch,
-                        encoder_hidden_states=encoder_hidden_states,
-                        text_embeddings=text_embeddings,
-                        rotary_embeddings=image_rotary_embeddings,
-                    )
-
-                if is_tracing_enabled():
-                    mx.eval(encoder_hidden_states, patch)
+            for wrapper in self.joint_block_wrappers:
+                wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token)
+                encoder_hidden_states, patch = wrapper(
+                    hidden_states=patch,
+                    encoder_hidden_states=encoder_hidden_states,
+                    text_embeddings=text_embeddings,
+                    rotary_embeddings=image_rotary_embeddings,
+                )
 
         if self.owns_concat_stage:
             assert encoder_hidden_states is not None
@@ -988,70 +1038,57 @@ class DiffusionRunner:
             if self.has_single_blocks or self.is_last_stage:
                 patch = patch_concat
             else:
-                with trace(
-                    name=f"send {self.next_rank}", rank=self.rank, category="comms"
-                ):
-                    patch_concat = mx.distributed.send(
-                        patch_concat, self.next_rank, group=self.group
-                    )
-                    mx.async_eval(patch_concat)
+                assert self.next_pipeline_rank is not None
+                patch_concat = mx.distributed.send(
+                    patch_concat, self.next_pipeline_rank, group=self.group
+                )
+                mx.async_eval(patch_concat)
 
         elif self.has_joint_blocks and not self.is_last_stage:
-            with trace(name=f"send {self.next_rank}", rank=self.rank, category="comms"):
-                patch = mx.distributed.send(patch, self.next_rank, group=self.group)
-                mx.async_eval(patch)
+            assert self.next_pipeline_rank is not None
+            patch = mx.distributed.send(
+                patch, self.next_pipeline_rank, group=self.group
+            )
+            mx.async_eval(patch)
 
             if patch_idx == 0:
                 assert encoder_hidden_states is not None
-                with trace(
-                    name=f"send {self.next_rank}", rank=self.rank, category="comms"
-                ):
-                    encoder_hidden_states = mx.distributed.send(
-                        encoder_hidden_states, self.next_rank, group=self.group
-                    )
-                    mx.async_eval(encoder_hidden_states)
+                encoder_hidden_states = mx.distributed.send(
+                    encoder_hidden_states, self.next_pipeline_rank, group=self.group
+                )
+                mx.async_eval(encoder_hidden_states)
 
         if self.has_single_blocks:
             if not self.owns_concat_stage and not self.is_first_stage:
+                assert self.prev_pipeline_rank is not None
                 patch_len = patch.shape[1]
-                with trace(
-                    name=f"recv {self.prev_rank}", rank=self.rank, category="comms"
-                ):
-                    patch = mx.distributed.recv(
-                        (batch_size, text_seq_len + patch_len, hidden_dim),
-                        patch.dtype,
-                        self.prev_rank,
-                        group=self.group,
-                    )
-                    mx.eval(patch)
+                patch = mx.distributed.recv(
+                    (batch_size, text_seq_len + patch_len, hidden_dim),
+                    patch.dtype,
+                    self.prev_pipeline_rank,
+                    group=self.group,
+                )
+                mx.eval(patch)
 
             assert self.single_block_wrappers is not None
-            with trace(
-                name=f"single patch {patch_idx}",
-                rank=self.rank,
-                category="compute",
-            ):
-                for wrapper in self.single_block_wrappers:
-                    wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token)
-                    patch = wrapper(
-                        hidden_states=patch,
-                        text_embeddings=text_embeddings,
-                        rotary_embeddings=image_rotary_embeddings,
-                    )
-
-                if is_tracing_enabled():
-                    mx.eval(patch)
+            for wrapper in self.single_block_wrappers:
+                wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token)
+                patch = wrapper(
+                    hidden_states=patch,
+                    text_embeddings=text_embeddings,
+                    rotary_embeddings=image_rotary_embeddings,
+                )
 
             if not self.is_last_stage:
-                with trace(
-                    name=f"send {self.next_rank}", rank=self.rank, category="comms"
-                ):
-                    patch = mx.distributed.send(patch, self.next_rank, group=self.group)
-                    mx.async_eval(patch)
+                assert self.next_pipeline_rank is not None
+                patch = mx.distributed.send(
+                    patch, self.next_pipeline_rank, group=self.group
+                )
+                mx.async_eval(patch)
 
         noise: mx.array | None = None
         if self.is_last_stage:
-            patch = patch[:, text_seq_len:, :]
-            noise = self.adapter.final_projection(patch, text_embeddings)
+            patch_img_only = patch[:, text_seq_len:, :]
+            noise = self.adapter.final_projection(patch_img_only, text_embeddings)
 
         return noise, encoder_hidden_states

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

@@ -3,7 +3,6 @@ from copy import deepcopy
 from typing import Any, cast
 
 import mlx.core as mx
-import psutil
 from mlx_lm.models.cache import (
     KVCache,
     QuantizedKVCache,
@@ -13,29 +12,25 @@ from mlx_lm.models.cache import (
 from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
-from exo.shared.types.memory import Memory
 from exo.shared.types.mlx import KVCacheType
 from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
 from exo.worker.runner.bootstrap import logger
 
 # Fraction of device memory above which LRU eviction kicks in
-_DEFAULT_MEMORY_THRESHOLD = 0.9
+_DEFAULT_MEMORY_THRESHOLD = 0.85
 _MEMORY_THRESHOLD = float(
     os.environ.get("EXO_MEMORY_THRESHOLD", _DEFAULT_MEMORY_THRESHOLD)
 )
 
 
 class KVPrefixCache:
-    def __init__(
-        self, tokenizer: TokenizerWrapper, group: mx.distributed.Group | None = None
-    ):
+    def __init__(self, tokenizer: TokenizerWrapper):
         self.prompts: list[mx.array] = []  # mx array of tokens (ints)
         self.caches: list[KVCacheType] = []
         self._last_used: list[int] = []  # monotonic counter of last access per entry
         self._access_counter: int = 0
         self._tokenizer: TokenizerWrapper = tokenizer
-        self._group = group
 
     def clear(self):
         """Clear all cached prompts and caches."""
@@ -86,13 +81,13 @@ class KVPrefixCache:
         best_snapshot_index, best_snapshot_length = None, 0
 
         for i, cached_prompt in enumerate(self.prompts):
-            length = get_prefix_length(tokenized_prompt, cached_prompt)
+            length = _get_prefix_length(tokenized_prompt, cached_prompt)
 
             if length == max_length:
                 # Exact match - cached prompt starts with our entire prompt
                 # Trim cache to prompt length - 1, return last token for stream_generate
                 prompt_cache = deepcopy(self.caches[i])
-                cached_length = cache_length(self.caches[i])
+                cached_length = _cache_length(self.caches[i])
                 tokens_to_trim = cached_length - (max_length - 1)
                 if tokens_to_trim > 0:
                     trim_prompt_cache(cast(list[Any], prompt_cache), tokens_to_trim)
@@ -114,7 +109,7 @@ class KVPrefixCache:
             prompt_cache = deepcopy(self.caches[best_snapshot_index])
 
             # Trim removes tokens from the end, so we trim (cached_length - prefix_length) to keep the prefix
-            cached_length = cache_length(self.caches[best_snapshot_index])
+            cached_length = _cache_length(self.caches[best_snapshot_index])
             tokens_to_trim = cached_length - best_snapshot_length
             if tokens_to_trim > 0:
                 trim_prompt_cache(cast(list[Any], prompt_cache), tokens_to_trim)
@@ -136,37 +131,29 @@ class KVPrefixCache:
             return prompt_cache, tokenized_prompt, None
 
     def _evict_if_needed(self):
-        """Evict least recently used entries while memory usage is high."""
+        """Evict least recently used entries while memory pressure is high."""
         if len(self.caches) == 0:
             return
 
+        active: int = mx.metal.get_active_memory()
+        limit = int(mx.metal.device_info()["max_recommended_working_set_size"])
+        if active < limit * _MEMORY_THRESHOLD:
+            return
+
         # Evict LRU entries until below threshold or only one entry left
-        while (
-            len(self.caches) > 1
-            and self.get_memory_used_percentage() > _MEMORY_THRESHOLD
-        ):
+        while len(self.caches) > 0:
             lru_index = self._last_used.index(min(self._last_used))
             evicted_tokens = len(self.prompts[lru_index])
             self.prompts.pop(lru_index)
             self.caches.pop(lru_index)
             self._last_used.pop(lru_index)
             logger.info(
-                f"KV cache evicted LRU entry ({evicted_tokens} tokens) due to memory usage"
+                f"KV cache evicted LRU entry ({evicted_tokens} tokens) due to memory pressure"
             )
 
-    def get_memory_used_percentage(self) -> float:
-        local_pressure: float = get_memory_used_percentage()
-
-        if self._group is None:
-            return local_pressure
-
-        all_pressure = mx.distributed.all_gather(
-            mx.array([local_pressure], dtype=mx.float32),
-            group=self._group,
-        )
-        # .item() evals.
-        max_pressure = float(mx.max(all_pressure).item())
-        return max_pressure
+            active = mx.metal.get_active_memory()
+            if active < limit * _MEMORY_THRESHOLD:
+                break
 
 
 def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
@@ -181,13 +168,13 @@ def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
     return mx.array(tokenized_prompt)
 
 
-def cache_length(cache: KVCacheType) -> int:
+def _cache_length(cache: KVCacheType) -> int:
     """Get the number of tokens in a KV cache."""
     # Use .offset attribute which all cache types have (len() not implemented in older QuantizedKVCache)
     return max(c.offset for c in cache)  # type: ignore
 
 
-def get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
+def _get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
     """Find the length of the common prefix between two token arrays."""
     n = min(int(prompt.shape[0]), int(cached_prompt.shape[0]))
     if n == 0:
@@ -198,17 +185,6 @@ def get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
     return int(mx.sum(prefix_mask).item())
 
 
-def get_available_memory() -> Memory:
-    mem: int = psutil.virtual_memory().available
-    return Memory.from_bytes(mem)
-
-
-def get_memory_used_percentage() -> float:
-    mem = psutil.virtual_memory()
-    # percent is 0-100
-    return float(mem.percent / 100)
-
-
 def make_kv_cache(
     model: Model, max_kv_size: int | None = None, keep: int = 0
 ) -> KVCacheType:

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

@@ -10,11 +10,8 @@ from mlx_lm.tokenizer_utils import TokenizerWrapper
 from exo.shared.types.api import (
     BenchChatCompletionTaskParams,
     ChatCompletionMessage,
-    CompletionTokensDetails,
     FinishReason,
     GenerationStats,
-    PromptTokensDetails,
-    Usage,
 )
 from exo.shared.types.memory import Memory
 from exo.shared.types.mlx import KVCacheType
@@ -42,7 +39,7 @@ def prefill(
     sampler: Callable[[mx.array], mx.array],
     prompt_tokens: mx.array,
     cache: KVCacheType,
-) -> tuple[float, int]:
+) -> float:
     """Prefill the KV cache with prompt tokens.
 
     This runs the model over the prompt tokens to populate the cache,
@@ -53,7 +50,7 @@ def prefill(
     """
     num_tokens = len(prompt_tokens)
     if num_tokens == 0:
-        return 0.0, 0
+        return 0.0
 
     logger.debug(f"Prefilling {num_tokens} tokens...")
     start_time = time.perf_counter()
@@ -88,7 +85,7 @@ def prefill(
         f"Prefill complete: {num_tokens} tokens in {elapsed:.2f}s "
         f"({tokens_per_sec:.1f} tok/s)"
     )
-    return tokens_per_sec, num_tokens
+    return tokens_per_sec
 
 
 def warmup_inference(
@@ -172,8 +169,6 @@ def mlx_generate(
     mx.reset_peak_memory()
     is_bench: bool = isinstance(task, BenchChatCompletionTaskParams)
 
-    logger.info(f"{is_bench=}")
-
     # Currently we support chat-completion tasks only.
     logger.debug(f"task_params: {task}")
 
@@ -209,9 +204,7 @@ def mlx_generate(
     )
 
     # Prefill cache with all tokens except the last one
-    prefill_tps, prefill_tokens = prefill(
-        model, tokenizer, sampler, prompt_tokens[:-1], caches
-    )
+    prefill_tps = prefill(model, tokenizer, sampler, prompt_tokens[:-1], caches)
 
     # stream_generate starts from the last token
     last_token = prompt_tokens[-1:]
@@ -219,43 +212,28 @@ def mlx_generate(
     max_tokens = task.max_tokens or MAX_TOKENS
     generated_text_parts: list[str] = []
     generation_start_time = time.perf_counter()
-    usage: Usage | None = None
-    in_thinking = False
-    reasoning_tokens = 0
-    think_start = tokenizer.think_start
-    think_end = tokenizer.think_end
-    for completion_tokens, out in enumerate(
-        stream_generate(
-            model=model,
-            tokenizer=tokenizer,
-            prompt=last_token,
-            max_tokens=max_tokens,
-            sampler=sampler,
-            logits_processors=logits_processors,
-            prompt_cache=caches,
-            # TODO: Dynamically change prefill step size to be the maximum possible without timing out.
-            prefill_step_size=2048,
-            kv_group_size=KV_GROUP_SIZE,
-            kv_bits=KV_BITS,
-        ),
-        start=1,
+    for out in stream_generate(
+        model=model,
+        tokenizer=tokenizer,
+        prompt=last_token,
+        max_tokens=max_tokens,
+        sampler=sampler,
+        logits_processors=logits_processors,
+        prompt_cache=caches,
+        # TODO: Dynamically change prefill step size to be the maximum possible without timing out.
+        prefill_step_size=2048,
+        kv_group_size=KV_GROUP_SIZE,
+        kv_bits=KV_BITS,
     ):
         generated_text_parts.append(out.text)
         logger.info(out.text)
 
-        if think_start is not None and out.text == think_start:
-            in_thinking = True
-        elif think_end is not None and out.text == think_end:
-            in_thinking = False
-        if in_thinking:
-            reasoning_tokens += 1
-
         stats: GenerationStats | None = None
         if out.finish_reason is not None:
             stats = GenerationStats(
                 prompt_tps=float(prefill_tps or out.prompt_tps),
                 generation_tps=float(out.generation_tps),
-                prompt_tokens=int(prefill_tokens + out.prompt_tokens),
+                prompt_tokens=int(out.prompt_tokens),
                 generation_tokens=int(out.generation_tokens),
                 peak_memory_usage=Memory.from_gb(out.peak_memory),
             )
@@ -267,24 +245,11 @@ def mlx_generate(
                     f"Model generated unexpected finish_reason: {out.finish_reason}"
                 )
 
-            usage = Usage(
-                prompt_tokens=int(out.prompt_tokens),
-                completion_tokens=completion_tokens,
-                total_tokens=int(out.prompt_tokens) + completion_tokens,
-                prompt_tokens_details=PromptTokensDetails(
-                    cached_tokens=prefix_hit_length
-                ),
-                completion_tokens_details=CompletionTokensDetails(
-                    reasoning_tokens=reasoning_tokens
-                ),
-            )
-
         yield GenerationResponse(
             text=out.text,
             token=out.token,
             finish_reason=cast(FinishReason | None, out.finish_reason),
             stats=stats,
-            usage=usage,
         )
 
         if out.finish_reason is not None:

src/exo/worker/runner/runner.py

@@ -18,7 +18,6 @@ from pydantic import ValidationError
 
 from exo.shared.constants import EXO_MAX_CHUNK_SIZE
 from exo.shared.models.model_cards import ModelId, ModelTask
-from exo.shared.tracing import clear_trace_buffer, get_trace_buffer, is_tracing_enabled
 from exo.shared.types.api import ChatCompletionMessageText, ImageGenerationStats
 from exo.shared.types.chunks import ErrorChunk, ImageChunk, TokenChunk, ToolCallChunk
 from exo.shared.types.common import CommandId
@@ -28,8 +27,6 @@ from exo.shared.types.events import (
     RunnerStatusUpdated,
     TaskAcknowledged,
     TaskStatusUpdated,
-    TraceEventData,
-    TracesCollected,
 )
 from exo.shared.types.tasks import (
     ChatCompletion,
@@ -40,7 +37,6 @@ from exo.shared.types.tasks import (
     Shutdown,
     StartWarmup,
     Task,
-    TaskId,
     TaskStatus,
 )
 from exo.shared.types.worker.instances import BoundInstance
@@ -65,7 +61,7 @@ from exo.shared.types.worker.runners import (
     RunnerStatus,
     RunnerWarmingUp,
 )
-from exo.shared.types.worker.shards import ShardMetadata
+from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
 from exo.utils.channels import MpReceiver, MpSender
 from exo.worker.engines.image import (
     DistributedImageModel,
@@ -115,12 +111,8 @@ def main(
     event_sender.send(
         RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
     )
-    seen = set[TaskId]()
     with task_receiver as tasks:
         for task in tasks:
-            if task.task_id in seen:
-                logger.warning("repeat task - potential error")
-            seen.add(task.task_id)
             event_sender.send(
                 TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
             )
@@ -171,7 +163,7 @@ def main(
                         logger.info(
                             f"model has_tool_calling={tokenizer.has_tool_calling}"
                         )
-                        kv_prefix_cache = KVPrefixCache(tokenizer, group)
+                        kv_prefix_cache = KVPrefixCache(tokenizer)
 
                     elif (
                         ModelTask.TextToImage in shard_metadata.model_card.tasks
@@ -285,11 +277,9 @@ def main(
                                 tokenizer.tool_parser,  # pyright: ignore[reportAny]
                             )
 
-                        completion_tokens = 0
                         for response in mlx_generator:
                             match response:
                                 case GenerationResponse():
-                                    completion_tokens += 1
                                     if (
                                         device_rank == 0
                                         and response.finish_reason == "error"
@@ -317,7 +307,6 @@ def main(
                                                     model=shard_metadata.model_card.model_id,
                                                     text=response.text,
                                                     token_id=response.token,
-                                                    usage=response.usage,
                                                     finish_reason=response.finish_reason,
                                                     stats=response.stats,
                                                 ),
@@ -331,7 +320,6 @@ def main(
                                                 chunk=ToolCallChunk(
                                                     tool_calls=response.tool_calls,
                                                     model=shard_metadata.model_card.model_id,
-                                                    usage=response.usage,
                                                 ),
                                             )
                                         )
@@ -372,8 +360,9 @@ def main(
                         image_index = 0
                         for response in generate_image(model=model, task=task_params):
                             if (
-                                shard_metadata.device_rank
-                                == shard_metadata.world_size - 1
+                                isinstance(shard_metadata, PipelineShardMetadata)
+                                and shard_metadata.is_pipeline_last
+                                and shard_metadata.cfg_rank == 0
                             ):
                                 match response:
                                     case PartialImageResponse():
@@ -399,7 +388,11 @@ def main(
                                         image_index += 1
                     # can we make this more explicit?
                     except Exception as e:
-                        if shard_metadata.device_rank == shard_metadata.world_size - 1:
+                        if (
+                            isinstance(shard_metadata, PipelineShardMetadata)
+                            and shard_metadata.is_pipeline_last
+                            and shard_metadata.cfg_rank == 0
+                        ):
                             event_sender.send(
                                 ChunkGenerated(
                                     command_id=command_id,
@@ -411,10 +404,6 @@ def main(
                                 )
                             )
                         raise
-                    finally:
-                        _send_traces_if_enabled(
-                            event_sender, task.task_id, shard_metadata.device_rank
-                        )
 
                     current_status = RunnerReady()
                     logger.info("runner ready")
@@ -435,8 +424,9 @@ def main(
                         image_index = 0
                         for response in generate_image(model=model, task=task_params):
                             if (
-                                shard_metadata.device_rank
-                                == shard_metadata.world_size - 1
+                                isinstance(shard_metadata, PipelineShardMetadata)
+                                and shard_metadata.is_pipeline_last
+                                and shard_metadata.cfg_rank == 0
                             ):
                                 match response:
                                     case PartialImageResponse():
@@ -461,7 +451,11 @@ def main(
                                         )
                                         image_index += 1
                     except Exception as e:
-                        if shard_metadata.device_rank == shard_metadata.world_size - 1:
+                        if (
+                            isinstance(shard_metadata, PipelineShardMetadata)
+                            and shard_metadata.is_pipeline_last
+                            and shard_metadata.cfg_rank == 0
+                        ):
                             event_sender.send(
                                 ChunkGenerated(
                                     command_id=command_id,
@@ -473,10 +467,6 @@ def main(
                                 )
                             )
                         raise
-                    finally:
-                        _send_traces_if_enabled(
-                            event_sender, task.task_id, shard_metadata.device_rank
-                        )
 
                     current_status = RunnerReady()
                     logger.info("runner ready")
@@ -555,10 +545,10 @@ def parse_gpt_oss(
                             name=current_tool_name,
                             arguments="".join(tool_arg_parts).strip(),
                         )
-                    ],
-                    usage=response.usage,
+                    ]
                 )
                 tool_arg_parts = []
+                break
             current_tool_name = recipient
 
         # If inside a tool call, accumulate arguments
@@ -651,36 +641,6 @@ def _send_image_chunk(
         )
 
 
-def _send_traces_if_enabled(
-    event_sender: MpSender[Event],
-    task_id: TaskId,
-    rank: int,
-) -> None:
-    if not is_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 _process_image_response(
     response: ImageGenerationResponse | PartialImageResponse,
     command_id: CommandId,
@@ -734,7 +694,7 @@ def parse_tool_calls(
                     tools = [_validate_single_tool(tool) for tool in parsed]
                 else:
                     tools = [_validate_single_tool(parsed)]
-                yield ToolCallResponse(tool_calls=tools, usage=response.usage)
+                yield ToolCallResponse(tool_calls=tools)
 
             except (
                 json.JSONDecodeError,

src/exo/worker/runner/runner_supervisor.py

@@ -127,25 +127,20 @@ class RunnerSupervisor:
         self._tg.cancel_scope.cancel()
 
     async def start_task(self, task: Task):
-        if task.task_id in self.pending:
-            logger.warning(
-                f"Skipping invalid task {task} as it has already been submitted"
-            )
-            return
         if task.task_id in self.completed:
-            logger.warning(
+            logger.info(
                 f"Skipping invalid task {task} as it has already been completed"
             )
-            return
         logger.info(f"Starting task {task}")
         event = anyio.Event()
         self.pending[task.task_id] = event
         try:
-            await self._task_sender.send_async(task)
+            self._task_sender.send(task)
         except ClosedResourceError:
             logger.warning(f"Task {task} dropped, runner closed communication.")
             return
         await event.wait()
+        logger.info(f"Finished task {task}")
 
     async def _forward_events(self):
         with self._ev_recv as events:

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

@@ -14,9 +14,9 @@ from exo.shared.types.tasks import ChatCompletionTaskParams
 from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.cache import (
     KVPrefixCache,
-    cache_length,
+    _cache_length,
+    _get_prefix_length,
     encode_prompt,
-    get_prefix_length,
     make_kv_cache,
 )
 from exo.worker.engines.mlx.generator.generate import mlx_generate, prefill
@@ -35,47 +35,47 @@ class TestGetPrefixLength:
     def test_identical_arrays(self):
         a = mx.array([1, 2, 3, 4, 5])
         b = mx.array([1, 2, 3, 4, 5])
-        assert get_prefix_length(a, b) == 5
+        assert _get_prefix_length(a, b) == 5
 
     def test_no_common_prefix(self):
         a = mx.array([1, 2, 3])
         b = mx.array([4, 5, 6])
-        assert get_prefix_length(a, b) == 0
+        assert _get_prefix_length(a, b) == 0
 
     def test_partial_prefix(self):
         a = mx.array([1, 2, 3, 4, 5])
         b = mx.array([1, 2, 3, 7, 8])
-        assert get_prefix_length(a, b) == 3
+        assert _get_prefix_length(a, b) == 3
 
     def test_prompt_longer_than_cached(self):
         a = mx.array([1, 2, 3, 4, 5])
         b = mx.array([1, 2, 3])
-        assert get_prefix_length(a, b) == 3
+        assert _get_prefix_length(a, b) == 3
 
     def test_cached_longer_than_prompt(self):
         a = mx.array([1, 2, 3])
         b = mx.array([1, 2, 3, 4, 5])
-        assert get_prefix_length(a, b) == 3
+        assert _get_prefix_length(a, b) == 3
 
     def test_single_token_match(self):
         a = mx.array([1, 2, 3])
         b = mx.array([1, 5, 6])
-        assert get_prefix_length(a, b) == 1
+        assert _get_prefix_length(a, b) == 1
 
     def test_empty_prompt(self):
         a = mx.array([]).astype(mx.int32)
         b = mx.array([1, 2, 3])
-        assert get_prefix_length(a, b) == 0
+        assert _get_prefix_length(a, b) == 0
 
     def test_empty_cached(self):
         a = mx.array([1, 2, 3])
         b = mx.array([]).astype(mx.int32)
-        assert get_prefix_length(a, b) == 0
+        assert _get_prefix_length(a, b) == 0
 
     def test_both_empty(self):
         a = mx.array([]).astype(mx.int32)
         b = mx.array([]).astype(mx.int32)
-        assert get_prefix_length(a, b) == 0
+        assert _get_prefix_length(a, b) == 0
 
 
 class TestKVPrefix:
@@ -146,7 +146,7 @@ class TestKVPrefixCacheWithModel:
         prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
 
         # Cache should now hold the prompt tokens
-        assert cache_length(cache) == len(tokens)
+        assert _cache_length(cache) == len(tokens)
 
     def test_add_and_get_exact_match(self, model_and_tokenizer):
         model, tokenizer = model_and_tokenizer
@@ -166,7 +166,7 @@ class TestKVPrefixCacheWithModel:
         kv_prefix_cache.add_kv_cache(prompt, cache)
 
         assert len(kv_prefix_cache.prompts) == 1
-        stored_length = cache_length(kv_prefix_cache.caches[0])
+        stored_length = _cache_length(kv_prefix_cache.caches[0])
         assert stored_length > 0
 
         # Retrieve with same prompt: exact match
@@ -209,7 +209,7 @@ class TestKVPrefixCacheWithModel:
         long_tokens = encode_prompt(tokenizer, long_prompt)
 
         # The prompts share a prefix (chat template preamble + "Hi")
-        expected_prefix = get_prefix_length(long_tokens, short_tokens)
+        expected_prefix = _get_prefix_length(long_tokens, short_tokens)
         assert expected_prefix > 0, (
             "Prompts should share a prefix from the chat template"
         )
@@ -243,7 +243,7 @@ class TestKVPrefixCacheWithModel:
         kv_prefix_cache = KVPrefixCache(tokenizer)
         kv_prefix_cache.add_kv_cache(prompt, cache)
 
-        stored_length = cache_length(kv_prefix_cache.caches[0])
+        stored_length = _cache_length(kv_prefix_cache.caches[0])
 
         # Get cache and mutate it (simulating what generation does)
         result_cache, _, matched_index = kv_prefix_cache.get_kv_cache(model, prompt)
@@ -259,7 +259,7 @@ class TestKVPrefixCacheWithModel:
         mx.eval([c.keys for c in result_cache])
 
         # Stored cache must be unchanged
-        assert cache_length(kv_prefix_cache.caches[0]) == stored_length
+        assert _cache_length(kv_prefix_cache.caches[0]) == stored_length
 
     def test_stored_cache_survives_repeated_get_mutate_cycles(
         self, model_and_tokenizer
@@ -281,7 +281,7 @@ class TestKVPrefixCacheWithModel:
         kv_prefix_cache = KVPrefixCache(tokenizer)
         kv_prefix_cache.add_kv_cache(prompt, cache)
 
-        stored_length = cache_length(kv_prefix_cache.caches[0])
+        stored_length = _cache_length(kv_prefix_cache.caches[0])
 
         for i in range(3):
             result_cache, _, _ = kv_prefix_cache.get_kv_cache(model, prompt)
@@ -293,7 +293,7 @@ class TestKVPrefixCacheWithModel:
                 layer_cache.update_and_fetch(extra, extra)
             mx.eval([c.keys for c in result_cache])
 
-            assert cache_length(kv_prefix_cache.caches[0]) == stored_length, (
+            assert _cache_length(kv_prefix_cache.caches[0]) == stored_length, (
                 f"Failed on loop {i}"
             )
 
@@ -325,7 +325,7 @@ class TestKVPrefixCacheWithModel:
         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
+        assert _cache_length(kv_prefix_cache.caches[0]) == expected_length
 
     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."""
@@ -400,7 +400,7 @@ class TestKVPrefixCacheWithModel:
         first_gen_time = time.perf_counter() - t0
 
         assert len(kv_prefix_cache.prompts) == 1
-        first_cache_length = cache_length(kv_prefix_cache.caches[0])
+        first_cache_length = _cache_length(kv_prefix_cache.caches[0])
 
         # Second generation: same long prompt + extra content (simulating multi-turn)
         task2 = ChatCompletionTaskParams(
@@ -416,7 +416,7 @@ class TestKVPrefixCacheWithModel:
         prompt2_tokens = encode_prompt(tokenizer, prompt2)
 
         # Verify the prompts share a long prefix
-        prefix_len = get_prefix_length(prompt2_tokens, prompt1_tokens)
+        prefix_len = _get_prefix_length(prompt2_tokens, prompt1_tokens)
         assert prefix_len > 1000, "Prompts must share > 1000 token prefix"
 
         # Second generation should reuse the cached prefix (only prefill new tokens)
@@ -440,7 +440,7 @@ class TestKVPrefixCacheWithModel:
         # With prefix_hit > 1000, should update in-place (not add a second entry)
         assert len(kv_prefix_cache.prompts) == 1
         # Updated cache should be longer (prompt2 + generated > prompt1 + generated)
-        updated_cache_length = cache_length(kv_prefix_cache.caches[0])
+        updated_cache_length = _cache_length(kv_prefix_cache.caches[0])
         assert updated_cache_length > first_cache_length
 
     def test_mlx_generate_stored_cache_not_mutated(self, model_and_tokenizer):
@@ -465,7 +465,7 @@ class TestKVPrefixCacheWithModel:
         ):
             pass
 
-        firstcache_length = cache_length(kv_prefix_cache.caches[0])
+        first_cache_length = _cache_length(kv_prefix_cache.caches[0])
 
         # Second generation gets the cache and mutates it during generation
         for _response in mlx_generate(
@@ -478,7 +478,7 @@ class TestKVPrefixCacheWithModel:
             pass
 
         # The first stored cache must not have been mutated by the second generation
-        assert cache_length(kv_prefix_cache.caches[0]) == firstcache_length
+        assert _cache_length(kv_prefix_cache.caches[0]) == first_cache_length
 
     def test_evicts_lru_entry_under_memory_pressure(self, model_and_tokenizer):
         """Under memory pressure, adding a new cache entry evicts the least recently used one."""
@@ -540,6 +540,6 @@ class TestKVPrefixCacheWithModel:
         assert len(kv_prefix_cache.prompts) == 1
         # The surviving entry should be the newly added one
         new_tokens = encode_prompt(tokenizer, prompt)
-        assert get_prefix_length(kv_prefix_cache.prompts[0], new_tokens) == len(
+        assert _get_prefix_length(kv_prefix_cache.prompts[0], new_tokens) == len(
             new_tokens
         )

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

@@ -109,8 +109,8 @@ def assert_events_equal(test_events: Iterable[Event], true_events: Iterable[Even
 
 @pytest.fixture
 def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
-    # initialize_mlx returns a mock group
-    monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MockGroup()))
+    # initialize_mlx returns a "group" equal to 1
+    monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(1))
     monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, MockTokenizer)))
     monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1))
     monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", nothin)
@@ -120,7 +120,7 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(mlx_runner, "detect_thinking_prompt_suffix", make_nothin(False))
 
     def fake_generate(*_1: object, **_2: object):
-        yield GenerationResponse(token=0, text="hi", finish_reason="stop", usage=None)
+        yield GenerationResponse(token=0, text="hi", finish_reason="stop")
 
     monkeypatch.setattr(mlx_runner, "mlx_generate", fake_generate)
 
@@ -147,14 +147,6 @@ class MockTokenizer:
     has_tool_calling = False
 
 
-class MockGroup:
-    def rank(self) -> int:
-        return 0
-
-    def size(self) -> int:
-        return 1
-
-
 def _run(tasks: Iterable[Task]):
     bound_instance = get_bound_mlx_ring_instance(
         instance_id=INSTANCE_1_ID,
@@ -190,8 +182,6 @@ def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
             text="hi",
             token_id=0,
             finish_reason="stop",
-            usage=None,
-            stats=None,
         ),
     )
 

tmp/config_examples/opencode.json

@@ -1,18 +0,0 @@
-{
-  "$schema": "https://opencode.ai/config.json",
-  "model": "exo/mlx-community/gpt-oss-120b-MXFP4-Q8",
-  "provider": {
-    "exo": {
-      "api": "http://localhost:52415/v1",
-      "models": {
-        "mlx-community/gpt-oss-120b-MXFP4-Q8": {
-          "name": "GPT OSS 120B",
-          "limit": {
-            "context": 32768,
-            "output": 8192
-          }
-        }
-      }
-    }
-  }
-}