Add tests

dashboard/src/lib/components/HeaderNav.svelte

@@ -297,5 +297,26 @@
       </svg>
       <span class="hidden sm:inline">Integrations</span>
     </a>
+    <a
+      href="/#/advanced"
+      class="text-xs md:text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-1.5 md:gap-2 cursor-pointer"
+      title="Advanced cluster settings"
+    >
+      <svg
+        class="w-4 h-4"
+        viewBox="0 0 24 24"
+        fill="none"
+        stroke="currentColor"
+        stroke-width="2"
+        stroke-linecap="round"
+        stroke-linejoin="round"
+      >
+        <circle cx="12" cy="12" r="3" />
+        <path
+          d="M19.4 15a1.65 1.65 0 0 0 .33 1.82l.06.06a2 2 0 0 1 0 2.83 2 2 0 0 1-2.83 0l-.06-.06a1.65 1.65 0 0 0-1.82-.33 1.65 1.65 0 0 0-1 1.51V21a2 2 0 0 1-4 0v-.09A1.65 1.65 0 0 0 9 19.4a1.65 1.65 0 0 0-1.82.33l-.06.06a2 2 0 0 1-2.83 0 2 2 0 0 1 0-2.83l.06-.06a1.65 1.65 0 0 0 .33-1.82 1.65 1.65 0 0 0-1.51-1H3a2 2 0 0 1 0-4h.09A1.65 1.65 0 0 0 4.6 9a1.65 1.65 0 0 0-.33-1.82l-.06-.06a2 2 0 0 1 0-2.83 2 2 0 0 1 2.83 0l.06.06a1.65 1.65 0 0 0 1.82.33H9a1.65 1.65 0 0 0 1-1.51V3a2 2 0 0 1 4 0v.09a1.65 1.65 0 0 0 1 1.51 1.65 1.65 0 0 0 1.82-.33l.06-.06a2 2 0 0 1 2.83 0 2 2 0 0 1 0 2.83l-.06.06a1.65 1.65 0 0 0-.33 1.82V9a1.65 1.65 0 0 0 1.51 1H21a2 2 0 0 1 0 4h-.09a1.65 1.65 0 0 0-1.51 1z"
+        />
+      </svg>
+      <span class="hidden sm:inline">Advanced</span>
+    </a>
   </nav>
 </header>

dashboard/src/lib/components/PrefillDecodeDisaggregation.svelte

@@ -0,0 +1,580 @@
+<script lang="ts">
+  import { onMount, onDestroy } from "svelte";
+  import FamilyLogos from "$lib/components/FamilyLogos.svelte";
+  import {
+    instances,
+    instanceLinks,
+    nodeIdentities,
+    refreshState,
+    createInstanceLink,
+    updateInstanceLink,
+    deleteInstanceLink,
+    type Instance,
+  } from "$lib/stores/app.svelte";
+  import { deriveBaseModel, deriveFamily } from "$lib/utils/model_family";
+
+  type InstanceWrapper = {
+    MlxRingInstance?: Instance;
+    MlxJacclInstance?: Instance;
+    VllmInstance?: Instance;
+  };
+
+  let interval: ReturnType<typeof setInterval> | null = null;
+
+  onMount(() => {
+    refreshState();
+    interval = setInterval(refreshState, 3000);
+  });
+  onDestroy(() => {
+    if (interval) clearInterval(interval);
+  });
+
+  type InstanceRow = {
+    id: string;
+    modelId: string;
+    family: string;
+    baseModel: string;
+    nodeNames: string[];
+  };
+
+  const instanceRows = $derived.by<InstanceRow[]>(() => {
+    const rows: InstanceRow[] = [];
+    const ids = nodeIdentities();
+    for (const [id, raw] of Object.entries(instances())) {
+      const wrapper = raw as InstanceWrapper;
+      const inst =
+        wrapper.MlxRingInstance ??
+        wrapper.MlxJacclInstance ??
+        wrapper.VllmInstance;
+      const modelId = inst?.shardAssignments?.modelId ?? "";
+      const nodeToRunner = inst?.shardAssignments?.nodeToRunner ?? {};
+      const nodeNames = Object.keys(nodeToRunner)
+        .map((nodeId) => ids[nodeId]?.friendlyName ?? nodeId.slice(0, 6))
+        .filter((name) => !!name);
+      rows.push({
+        id,
+        modelId,
+        family: deriveFamily(modelId),
+        baseModel: deriveBaseModel(modelId),
+        nodeNames,
+      });
+    }
+    rows.sort((a, b) => a.modelId.localeCompare(b.modelId));
+    return rows;
+  });
+
+  const instanceById = $derived(
+    Object.fromEntries(instanceRows.map((r) => [r.id, r])),
+  );
+
+  type LinkRow = {
+    linkId: string;
+    prefill: string[];
+    decode: string[];
+    families: string[];
+  };
+
+  const linkRows = $derived.by<LinkRow[]>(() => {
+    const rows: LinkRow[] = [];
+    for (const [, link] of Object.entries(instanceLinks())) {
+      const fams = new Set<string>();
+      for (const id of [...link.prefillInstances, ...link.decodeInstances]) {
+        const r = instanceById[id];
+        if (r && r.baseModel) fams.add(r.baseModel.toLowerCase());
+      }
+      rows.push({
+        linkId: link.linkId,
+        prefill: link.prefillInstances,
+        decode: link.decodeInstances,
+        families: Array.from(fams),
+      });
+    }
+    return rows;
+  });
+
+  let editingLinkId = $state<string | null>(null);
+  let editingPrefill = $state<Set<string>>(new Set());
+  let editingDecode = $state<Set<string>>(new Set());
+  let saving = $state(false);
+  let errorMessage = $state<string | null>(null);
+
+  function startCreate() {
+    editingLinkId = "new";
+    editingPrefill = new Set();
+    editingDecode = new Set();
+    errorMessage = null;
+  }
+
+  function startEdit(row: LinkRow) {
+    editingLinkId = row.linkId;
+    editingPrefill = new Set(row.prefill);
+    editingDecode = new Set(row.decode);
+    errorMessage = null;
+  }
+
+  function cancelEdit() {
+    editingLinkId = null;
+    editingPrefill = new Set();
+    editingDecode = new Set();
+    errorMessage = null;
+  }
+
+  type Role = "prefill" | "decode" | "none";
+
+  function roleOf(id: string): Role {
+    if (editingPrefill.has(id)) return "prefill";
+    if (editingDecode.has(id)) return "decode";
+    return "none";
+  }
+
+  function setRole(id: string, role: Role) {
+    const p = new Set(editingPrefill);
+    const d = new Set(editingDecode);
+    p.delete(id);
+    d.delete(id);
+    if (role === "prefill") p.add(id);
+    if (role === "decode") d.add(id);
+    editingPrefill = p;
+    editingDecode = d;
+  }
+
+  const editingFamilies = $derived.by<string[]>(() => {
+    const fams = new Set<string>();
+    for (const id of [...editingPrefill, ...editingDecode]) {
+      const r = instanceById[id];
+      if (r && r.baseModel) fams.add(r.baseModel.toLowerCase());
+    }
+    return Array.from(fams);
+  });
+
+  const editingMismatch = $derived(editingFamilies.length > 1);
+  const canSave = $derived(
+    editingLinkId !== null &&
+      editingPrefill.size > 0 &&
+      editingDecode.size > 0 &&
+      !saving,
+  );
+
+  async function save() {
+    if (editingLinkId === null) return;
+    saving = true;
+    errorMessage = null;
+    try {
+      const prefill = Array.from(editingPrefill);
+      const decode = Array.from(editingDecode);
+      if (editingLinkId === "new") {
+        await createInstanceLink(prefill, decode);
+      } else {
+        await updateInstanceLink(editingLinkId, prefill, decode);
+      }
+      cancelEdit();
+      await refreshState();
+    } catch (err) {
+      errorMessage = err instanceof Error ? err.message : String(err);
+    } finally {
+      saving = false;
+    }
+  }
+
+  async function remove(linkId: string) {
+    if (!confirm("Remove this routing?")) return;
+    try {
+      await deleteInstanceLink(linkId);
+      if (editingLinkId === linkId) cancelEdit();
+      await refreshState();
+    } catch (err) {
+      errorMessage = err instanceof Error ? err.message : String(err);
+    }
+  }
+</script>
+
+<div class="tab-content">
+  <p class="tab-intro">
+    Route prefill from one set of instances to another. The decode worker
+    decides per-request whether to use a remote prefill (gated on prefix-cache
+    miss size). Linked instances should be running the same model family.
+  </p>
+
+  {#if errorMessage}
+    <div class="error">{errorMessage}</div>
+  {/if}
+
+  <section>
+    <div class="section-head">
+      <h2>Existing routes</h2>
+      <button
+        class="primary"
+        onclick={startCreate}
+        disabled={editingLinkId !== null}
+      >
+        + New route
+      </button>
+    </div>
+
+    {#if linkRows.length === 0}
+      <p class="empty">No routes yet. Create one to enable remote prefill.</p>
+    {:else}
+      <div class="link-grid">
+        {#each linkRows as row (row.linkId)}
+          <article class="link-card">
+            {#if row.families.length > 1}
+              <div class="warn-banner">
+                ⚠ Mixed model families: {row.families.join(", ")}
+              </div>
+            {/if}
+            <div class="route">
+              <div class="route-side">
+                <span class="role-tag prefill">PREFILL</span>
+                <ul class="instance-list">
+                  {#each row.prefill as id (id)}
+                    {@const r = instanceById[id]}
+                    {#if r}
+                      <li class="chip">
+                        <FamilyLogos family={r.family} />
+                        <div>
+                          <div class="model-name">
+                            {r.baseModel || r.modelId}
+                          </div>
+                          <div class="node-name">
+                            {r.nodeNames.join(", ") || "?"}
+                          </div>
+                        </div>
+                      </li>
+                    {/if}
+                  {/each}
+                </ul>
+              </div>
+              <div class="arrow" aria-hidden="true">→</div>
+              <div class="route-side">
+                <span class="role-tag decode">DECODE</span>
+                <ul class="instance-list">
+                  {#each row.decode as id (id)}
+                    {@const r = instanceById[id]}
+                    {#if r}
+                      <li class="chip">
+                        <FamilyLogos family={r.family} />
+                        <div>
+                          <div class="model-name">
+                            {r.baseModel || r.modelId}
+                          </div>
+                          <div class="node-name">
+                            {r.nodeNames.join(", ") || "?"}
+                          </div>
+                        </div>
+                      </li>
+                    {/if}
+                  {/each}
+                </ul>
+              </div>
+            </div>
+            <div class="link-actions">
+              <button
+                onclick={() => startEdit(row)}
+                disabled={editingLinkId !== null}
+              >
+                Edit
+              </button>
+              <button class="danger" onclick={() => remove(row.linkId)}
+                >Remove</button
+              >
+            </div>
+          </article>
+        {/each}
+      </div>
+    {/if}
+  </section>
+
+  {#if editingLinkId !== null}
+    <section class="editor">
+      <h2>{editingLinkId === "new" ? "New route" : "Edit route"}</h2>
+
+      {#if editingMismatch}
+        <div class="warn-banner">
+          ⚠ Selected instances span multiple model families:
+          <strong>{editingFamilies.join(", ")}</strong>. Linking across families
+          produces a corrupt KV cache.
+        </div>
+      {/if}
+
+      <p class="hint">
+        Pick a role for each instance: <strong>Prefill</strong> serves KV cache,
+        <strong>Decode</strong> consumes it.
+      </p>
+
+      {#if instanceRows.length === 0}
+        <p class="empty">No instances available.</p>
+      {:else}
+        <div class="instance-picker">
+          {#each instanceRows as row (row.id)}
+            {@const role = roleOf(row.id)}
+            <div
+              class="picker-card"
+              class:picker-prefill={role === "prefill"}
+              class:picker-decode={role === "decode"}
+            >
+              <div class="picker-meta">
+                <FamilyLogos family={row.family} />
+                <div>
+                  <div class="model-name">{row.baseModel || row.modelId}</div>
+                  <div class="node-name">{row.nodeNames.join(", ") || "?"}</div>
+                </div>
+              </div>
+              <div class="role-toggle">
+                <button
+                  class="role-btn"
+                  class:active={role === "none"}
+                  onclick={() => setRole(row.id, "none")}
+                  title="Don't include in this route">–</button
+                >
+                <button
+                  class="role-btn prefill"
+                  class:active={role === "prefill"}
+                  onclick={() => setRole(row.id, "prefill")}>Prefill</button
+                >
+                <button
+                  class="role-btn decode"
+                  class:active={role === "decode"}
+                  onclick={() => setRole(row.id, "decode")}>Decode</button
+                >
+              </div>
+            </div>
+          {/each}
+        </div>
+      {/if}
+
+      <div class="editor-actions">
+        <button class="primary" onclick={save} disabled={!canSave}>
+          {saving ? "Saving..." : "Save route"}
+        </button>
+        <button onclick={cancelEdit} disabled={saving}>Cancel</button>
+      </div>
+    </section>
+  {/if}
+</div>
+
+<style>
+  .tab-content {
+    color: var(--text-primary, #eee);
+  }
+  .tab-intro {
+    margin: 0 0 1.5rem 0;
+    color: var(--text-secondary, #aaa);
+    max-width: 70ch;
+  }
+  .page-header h1 {
+    margin: 0 0 0.5rem 0;
+  }
+  .page-header p {
+    margin: 0;
+    color: var(--text-secondary, #aaa);
+    max-width: 70ch;
+  }
+  .error {
+    margin: 1rem 0;
+    padding: 0.75rem 1rem;
+    background: rgba(255, 80, 80, 0.15);
+    border: 1px solid rgba(255, 80, 80, 0.5);
+    border-radius: 6px;
+  }
+  .section-head {
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+    margin: 2rem 0 1rem 0;
+  }
+  .section-head h2,
+  .editor h2 {
+    margin: 0;
+  }
+  button {
+    background: rgba(255, 255, 255, 0.08);
+    color: inherit;
+    border: 1px solid rgba(255, 255, 255, 0.15);
+    border-radius: 4px;
+    padding: 0.45rem 0.9rem;
+    cursor: pointer;
+    font: inherit;
+  }
+  button:hover:not(:disabled) {
+    background: rgba(255, 255, 255, 0.14);
+  }
+  button:disabled {
+    opacity: 0.5;
+    cursor: not-allowed;
+  }
+  button.primary {
+    background: rgba(80, 180, 255, 0.2);
+    border-color: rgba(80, 180, 255, 0.5);
+  }
+  button.danger {
+    background: rgba(255, 80, 80, 0.18);
+    border-color: rgba(255, 80, 80, 0.4);
+  }
+  .empty {
+    color: var(--text-secondary, #888);
+    font-style: italic;
+  }
+  .link-grid {
+    display: flex;
+    flex-direction: column;
+    gap: 1rem;
+  }
+  .link-card {
+    border: 1px solid rgba(255, 255, 255, 0.1);
+    border-radius: 8px;
+    padding: 1rem 1.25rem;
+    background: rgba(255, 255, 255, 0.02);
+  }
+  .route {
+    display: grid;
+    grid-template-columns: 1fr auto 1fr;
+    gap: 1rem;
+    align-items: center;
+  }
+  .route-side {
+    min-width: 0;
+  }
+  .arrow {
+    font-size: 1.5rem;
+    color: var(--text-secondary, #888);
+  }
+  .role-tag {
+    display: inline-block;
+    font-size: 0.7rem;
+    letter-spacing: 0.08em;
+    padding: 0.15rem 0.5rem;
+    border-radius: 3px;
+    margin-bottom: 0.5rem;
+  }
+  .role-tag.prefill {
+    background: rgba(80, 180, 255, 0.2);
+    color: #6cb6ff;
+  }
+  .role-tag.decode {
+    background: rgba(150, 220, 100, 0.2);
+    color: #b3e07d;
+  }
+  .instance-list {
+    list-style: none;
+    padding: 0;
+    margin: 0;
+    display: flex;
+    flex-direction: column;
+    gap: 0.5rem;
+  }
+  .chip {
+    display: flex;
+    gap: 0.6rem;
+    align-items: center;
+    background: rgba(255, 255, 255, 0.04);
+    border: 1px solid rgba(255, 255, 255, 0.08);
+    border-radius: 6px;
+    padding: 0.45rem 0.6rem;
+  }
+  .chip :global(svg) {
+    flex-shrink: 0;
+    color: var(--text-secondary, #ccc);
+  }
+  .model-name {
+    font-weight: 500;
+    overflow: hidden;
+    text-overflow: ellipsis;
+    white-space: nowrap;
+  }
+  .node-name {
+    font-size: 0.8em;
+    color: var(--text-secondary, #999);
+  }
+  .link-actions {
+    margin-top: 0.85rem;
+    display: flex;
+    gap: 0.5rem;
+    justify-content: flex-end;
+  }
+  .editor {
+    margin-top: 2rem;
+    padding: 1.25rem 1.5rem;
+    border: 1px solid rgba(255, 255, 255, 0.15);
+    border-radius: 8px;
+    background: rgba(255, 255, 255, 0.03);
+  }
+  .editor h2 {
+    margin: 0 0 0.5rem 0;
+  }
+  .editor .hint {
+    color: var(--text-secondary, #aaa);
+    margin: 0 0 1rem 0;
+    font-size: 0.95em;
+  }
+  .editor-actions {
+    margin-top: 1.25rem;
+    display: flex;
+    gap: 0.5rem;
+  }
+  .warn-banner {
+    margin-bottom: 0.85rem;
+    padding: 0.6rem 0.85rem;
+    background: rgba(255, 183, 77, 0.1);
+    border: 1px solid rgba(255, 183, 77, 0.4);
+    border-radius: 6px;
+    color: #ffb74d;
+    font-size: 0.9em;
+  }
+  .instance-picker {
+    display: grid;
+    grid-template-columns: repeat(auto-fill, minmax(360px, 1fr));
+    gap: 0.75rem;
+  }
+  .picker-card {
+    border: 1px solid rgba(255, 255, 255, 0.12);
+    border-radius: 8px;
+    padding: 0.75rem 0.9rem;
+    display: flex;
+    flex-direction: column;
+    gap: 0.65rem;
+    background: rgba(255, 255, 255, 0.02);
+    transition:
+      border-color 0.15s,
+      background 0.15s;
+  }
+  .picker-card.picker-prefill {
+    border-color: rgba(80, 180, 255, 0.5);
+    background: rgba(80, 180, 255, 0.06);
+  }
+  .picker-card.picker-decode {
+    border-color: rgba(150, 220, 100, 0.5);
+    background: rgba(150, 220, 100, 0.06);
+  }
+  .picker-meta {
+    display: flex;
+    gap: 0.7rem;
+    align-items: center;
+  }
+  .picker-meta :global(svg) {
+    flex-shrink: 0;
+    color: var(--text-secondary, #ddd);
+  }
+  .role-toggle {
+    display: flex;
+    gap: 0.35rem;
+  }
+  .role-btn {
+    flex: 1;
+    padding: 0.35rem 0.5rem;
+    font-size: 0.85em;
+  }
+  .role-btn.active {
+    background: rgba(255, 255, 255, 0.18);
+    border-color: rgba(255, 255, 255, 0.35);
+  }
+  .role-btn.prefill.active {
+    background: rgba(80, 180, 255, 0.3);
+    border-color: rgba(80, 180, 255, 0.7);
+    color: #fff;
+  }
+  .role-btn.decode.active {
+    background: rgba(150, 220, 100, 0.3);
+    border-color: rgba(150, 220, 100, 0.7);
+    color: #fff;
+  }
+</style>

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

@@ -74,6 +74,12 @@ export interface Instance {
   };
 }
 
+export interface RawInstanceLink {
+  linkId: string;
+  prefillInstances: string[];
+  decodeInstances: string[];
+}
+
 // Granular node state types from the new state structure
 interface RawNodeIdentity {
   modelId?: string;
@@ -223,6 +229,7 @@ interface RawStateResponse {
     }
   >;
   runners?: Record<string, unknown>;
+  instanceLinks?: Record<string, RawInstanceLink>;
   downloads?: Record<string, unknown[]>;
   // New granular node state fields
   nodeIdentities?: Record<string, RawNodeIdentity>;
@@ -541,6 +548,7 @@ class AppStore {
   topologyData = $state<TopologyData | null>(null);
   instances = $state<Record<string, unknown>>({});
   runners = $state<Record<string, unknown>>({});
+  instanceLinks = $state<Record<string, RawInstanceLink>>({});
   downloads = $state<Record<string, unknown[]>>({});
   nodeDisk = $state<
     Record<
@@ -1310,6 +1318,11 @@ class AppStore {
       if (data.runners) {
         this.runners = data.runners;
       }
+      if (data.instanceLinks) {
+        this.instanceLinks = data.instanceLinks;
+      } else {
+        this.instanceLinks = {};
+      }
       if (data.downloads) {
         this.downloads = data.downloads;
       }
@@ -3281,6 +3294,60 @@ class AppStore {
     }
   }
 
+  async createInstanceLink(
+    prefillInstances: string[],
+    decodeInstances: string[],
+  ): Promise<void> {
+    const response = await fetch("/v1/instance-links", {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify({
+        prefill_instances: prefillInstances,
+        decode_instances: decodeInstances,
+      }),
+    });
+    if (!response.ok) {
+      throw new Error(
+        `Failed to create instance link: ${response.status} ${await response.text()}`,
+      );
+    }
+  }
+
+  async updateInstanceLink(
+    linkId: string,
+    prefillInstances: string[],
+    decodeInstances: string[],
+  ): Promise<void> {
+    const response = await fetch(
+      `/v1/instance-links/${encodeURIComponent(linkId)}`,
+      {
+        method: "PUT",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({
+          prefill_instances: prefillInstances,
+          decode_instances: decodeInstances,
+        }),
+      },
+    );
+    if (!response.ok) {
+      throw new Error(
+        `Failed to update instance link: ${response.status} ${await response.text()}`,
+      );
+    }
+  }
+
+  async deleteInstanceLink(linkId: string): Promise<void> {
+    const response = await fetch(
+      `/v1/instance-links/${encodeURIComponent(linkId)}`,
+      { method: "DELETE" },
+    );
+    if (!response.ok) {
+      throw new Error(
+        `Failed to delete instance link: ${response.status} ${await response.text()}`,
+      );
+    }
+  }
+
   /**
    * Delete a downloaded model from a specific node
    */
@@ -3379,6 +3446,18 @@ export const prefillProgress = () => appStore.prefillProgress;
 export const topologyData = () => appStore.topologyData;
 export const instances = () => appStore.instances;
 export const runners = () => appStore.runners;
+export const instanceLinks = () => appStore.instanceLinks;
+export const createInstanceLink = (
+  prefillInstances: string[],
+  decodeInstances: string[],
+) => appStore.createInstanceLink(prefillInstances, decodeInstances);
+export const updateInstanceLink = (
+  linkId: string,
+  prefillInstances: string[],
+  decodeInstances: string[],
+) => appStore.updateInstanceLink(linkId, prefillInstances, decodeInstances);
+export const deleteInstanceLink = (linkId: string) =>
+  appStore.deleteInstanceLink(linkId);
 export const downloads = () => appStore.downloads;
 export const nodeDisk = () => appStore.nodeDisk;
 export const placementPreviews = () => appStore.placementPreviews;

dashboard/src/lib/utils/model_family.ts

@@ -0,0 +1,44 @@
+// Mirrors src/exo/shared/models/model_cards.py:derive_base_model
+const QUANT_SUFFIXES = new RegExp(
+  "[-_ ](?:MLX|MXFP[0-9]+|NVFP[0-9]+|GPTQ|AWQ|GGUF|fp16|bf16|fp8|int[0-9]+|[0-9]+(?:\\.[0-9]+)?bit|Q[0-9]+(?:_[A-Z0-9]+)?|gs[0-9]+)" +
+    "(?:[-_ ](?:MLX|Q[0-9]+|Int[0-9]+|[A-Z0-9]+|gs[0-9]+))*$",
+  "i",
+);
+
+function normalize(s: string): string {
+  return s
+    .replaceAll("-", " ")
+    .replaceAll("_", " ")
+    .replaceAll("  ", " ")
+    .trim();
+}
+
+export function deriveBaseModel(modelId: string): string {
+  const short = modelId.includes("/")
+    ? (modelId.split("/").pop() ?? modelId)
+    : modelId;
+  const stripped = short.replace(QUANT_SUFFIXES, "");
+  return normalize(stripped);
+}
+
+export function baseModelsCompatible(a: string, b: string): boolean {
+  return deriveBaseModel(a).toLowerCase() === deriveBaseModel(b).toLowerCase();
+}
+
+// Mirrors src/exo/shared/models/model_cards.py:derive_family
+export function deriveFamily(modelId: string): string {
+  const short = modelId.includes("/")
+    ? (modelId.split("/").pop() ?? modelId)
+    : modelId;
+  const stripped = short
+    .replace(QUANT_SUFFIXES, "")
+    .toLowerCase()
+    .replaceAll("_", "-");
+  const parts = stripped.split(/[-.]/);
+  const familyParts: string[] = [];
+  for (const p of parts) {
+    if (/^\d+$/.test(p) || /^\d+[bm]?$/i.test(p)) break;
+    familyParts.push(p);
+  }
+  return familyParts.length > 0 ? familyParts.join("-") : stripped;
+}

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

@@ -0,0 +1,110 @@
+<script lang="ts">
+  import HeaderNav from "$lib/components/HeaderNav.svelte";
+  import PrefillDecodeDisaggregation from "$lib/components/PrefillDecodeDisaggregation.svelte";
+
+  type TabId = "prefill-decode";
+
+  type Tab = {
+    id: TabId;
+    label: string;
+  };
+
+  const tabs: Tab[] = [
+    { id: "prefill-decode", label: "Prefill / Decode disaggregation" },
+  ];
+
+  let activeTab = $state<TabId>(tabs[0].id);
+</script>
+
+<HeaderNav />
+
+<main class="page">
+  <header class="page-header">
+    <h1>Advanced</h1>
+    <p>Cluster-level configuration. Most users don't need anything here.</p>
+  </header>
+
+  <div class="layout">
+    <nav class="tab-list" aria-label="Advanced settings sections">
+      {#each tabs as tab (tab.id)}
+        <button
+          class="tab"
+          class:active={activeTab === tab.id}
+          onclick={() => (activeTab = tab.id)}
+        >
+          {tab.label}
+        </button>
+      {/each}
+    </nav>
+
+    <section class="tab-panel">
+      {#if activeTab === "prefill-decode"}
+        <PrefillDecodeDisaggregation />
+      {/if}
+    </section>
+  </div>
+</main>
+
+<style>
+  .page {
+    max-width: 1100px;
+    margin: 0 auto;
+    padding: 1.5rem;
+    color: var(--text-primary, #eee);
+  }
+  .page-header h1 {
+    margin: 0 0 0.5rem 0;
+  }
+  .page-header p {
+    margin: 0 0 1.5rem 0;
+    color: var(--text-secondary, #aaa);
+  }
+  .layout {
+    display: grid;
+    grid-template-columns: 240px 1fr;
+    gap: 1.5rem;
+    align-items: start;
+  }
+  @media (max-width: 720px) {
+    .layout {
+      grid-template-columns: 1fr;
+    }
+  }
+  .tab-list {
+    display: flex;
+    flex-direction: column;
+    gap: 0.25rem;
+    border-right: 1px solid rgba(255, 255, 255, 0.08);
+    padding-right: 1rem;
+  }
+  @media (max-width: 720px) {
+    .tab-list {
+      flex-direction: row;
+      flex-wrap: wrap;
+      border-right: none;
+      border-bottom: 1px solid rgba(255, 255, 255, 0.08);
+      padding-right: 0;
+      padding-bottom: 0.5rem;
+    }
+  }
+  .tab {
+    background: transparent;
+    color: inherit;
+    border: 1px solid transparent;
+    border-radius: 4px;
+    padding: 0.5rem 0.75rem;
+    text-align: left;
+    cursor: pointer;
+    font: inherit;
+  }
+  .tab:hover:not(.active) {
+    background: rgba(255, 255, 255, 0.05);
+  }
+  .tab.active {
+    background: rgba(80, 180, 255, 0.15);
+    border-color: rgba(80, 180, 255, 0.5);
+  }
+  .tab-panel {
+    min-width: 0;
+  }
+</style>

src/exo/api/main.py

@@ -79,6 +79,8 @@ from exo.api.types import (
     ImageListItem,
     ImageListResponse,
     ImageSize,
+    InstanceLinkBody,
+    InstanceLinkResponse,
     ModelList,
     ModelListModel,
     PlaceInstanceParams,
@@ -154,6 +156,7 @@ from exo.shared.types.commands import (
     DeleteCustomModelCard,
     DeleteDownload,
     DeleteInstance,
+    DeleteInstanceLink,
     DownloadCommand,
     ForwarderCommand,
     ForwarderDownloadCommand,
@@ -161,6 +164,7 @@ from exo.shared.types.commands import (
     ImageGeneration,
     PlaceInstance,
     SendInputChunk,
+    SetInstanceLink,
     StartDownload,
     TaskCancelled,
     TaskFinished,
@@ -174,6 +178,7 @@ from exo.shared.types.events import (
     InstanceDeleted,
     TracesMerged,
 )
+from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.memory import Memory
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
@@ -328,6 +333,10 @@ class API:
         self.app.get("/instance/previews")(self.get_placement_previews)
         self.app.get("/instance/{instance_id}")(self.get_instance)
         self.app.delete("/instance/{instance_id}")(self.delete_instance)
+        self.app.get("/v1/instance-links")(self.list_instance_links)
+        self.app.post("/v1/instance-links")(self.create_instance_link)
+        self.app.put("/v1/instance-links/{link_id}")(self.update_instance_link)
+        self.app.delete("/v1/instance-links/{link_id}")(self.delete_instance_link)
         self.app.get("/models")(self.get_models)
         self.app.get("/v1/models")(self.get_models)
         self.app.post("/models/add")(self.add_custom_model)
@@ -615,6 +624,41 @@ class API:
             instance_id=instance_id,
         )
 
+    async def list_instance_links(self) -> list[InstanceLink]:
+        return list(self.state.instance_links.values())
+
+    async def create_instance_link(
+        self, body: InstanceLinkBody
+    ) -> InstanceLinkResponse:
+        return await self._set_instance_link(None, body)
+
+    async def update_instance_link(
+        self, link_id: InstanceLinkId, body: InstanceLinkBody
+    ) -> InstanceLinkResponse:
+        return await self._set_instance_link(link_id, body)
+
+    async def _set_instance_link(
+        self, link_id: InstanceLinkId | None, body: InstanceLinkBody
+    ) -> InstanceLinkResponse:
+        command = SetInstanceLink(
+            link_id=link_id,
+            prefill_instances=list(body.prefill_instances),
+            decode_instances=list(body.decode_instances),
+        )
+        await self._send(command)
+        return InstanceLinkResponse(
+            message="Command received.", command_id=command.command_id
+        )
+
+    async def delete_instance_link(
+        self, link_id: InstanceLinkId
+    ) -> InstanceLinkResponse:
+        command = DeleteInstanceLink(link_id=link_id)
+        await self._send(command)
+        return InstanceLinkResponse(
+            message="Command received.", command_id=command.command_id
+        )
+
     async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
         """Cancel an active command by closing its stream and notifying workers."""
         sender = self._text_generation_queues.get(

src/exo/api/types/__init__.py

@@ -34,6 +34,8 @@ from .api import ImageGenerationTaskParams as ImageGenerationTaskParams
 from .api import ImageListItem as ImageListItem
 from .api import ImageListResponse as ImageListResponse
 from .api import ImageSize as ImageSize
+from .api import InstanceLinkBody as InstanceLinkBody
+from .api import InstanceLinkResponse as InstanceLinkResponse
 from .api import Logprobs as Logprobs
 from .api import LogprobsContentItem as LogprobsContentItem
 from .api import ModelList as ModelList

src/exo/api/types/api.py

@@ -295,6 +295,16 @@ class CancelCommandResponse(BaseModel):
     command_id: CommandId
 
 
+class InstanceLinkBody(BaseModel):
+    prefill_instances: list[InstanceId]
+    decode_instances: list[InstanceId]
+
+
+class InstanceLinkResponse(BaseModel):
+    message: str
+    command_id: CommandId
+
+
 ImageSize = Literal[
     "auto",
     "512x512",

src/exo/master/main.py

@@ -17,6 +17,7 @@ from exo.shared.types.commands import (
     CreateInstance,
     DeleteCustomModelCard,
     DeleteInstance,
+    DeleteInstanceLink,
     ForwarderCommand,
     ForwarderDownloadCommand,
     ImageEdits,
@@ -24,6 +25,7 @@ from exo.shared.types.commands import (
     PlaceInstance,
     RequestEventLog,
     SendInputChunk,
+    SetInstanceLink,
     TaskCancelled,
     TaskFinished,
     TestCommand,
@@ -38,6 +40,8 @@ from exo.shared.types.events import (
     IndexedEvent,
     InputChunkReceived,
     InstanceDeleted,
+    InstanceLinkCreated,
+    InstanceLinkDeleted,
     LocalForwarderEvent,
     NodeGatheredInfo,
     NodeTimedOut,
@@ -48,6 +52,7 @@ from exo.shared.types.events import (
     TracesCollected,
     TracesMerged,
 )
+from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
     ImageEdits as ImageEditsTask,
@@ -63,12 +68,57 @@ from exo.shared.types.tasks import (
     TextGeneration as TextGenerationTask,
 )
 from exo.shared.types.worker.instances import InstanceId
+from exo.shared.types.worker.runners import RunnerReady
 from exo.utils.channels import Receiver, Sender
 from exo.utils.disk_event_log import DiskEventLog
 from exo.utils.event_buffer import MultiSourceBuffer
 from exo.utils.task_group import TaskGroup
 
 
+def _prefill_endpoints_for(state: State, decode_instance_id: InstanceId) -> list[str]:
+    from exo.master.placement_utils import (
+        _find_ip_prioritised as find_ip_prioritised,  # pyright: ignore[reportPrivateUsage]
+    )
+
+    decode = state.instances.get(decode_instance_id)
+    if decode is None:
+        return []
+    decode_node = next(iter(decode.shard_assignments.node_to_runner.keys()), None)
+    if decode_node is None:
+        return []
+
+    sources: set[InstanceId] = set()
+    for link in state.instance_links.values():
+        if decode_instance_id in link.decode_instances:
+            sources.update(link.prefill_instances)
+    sources.discard(decode_instance_id)
+
+    endpoints: list[str] = []
+    seen: set[str] = set()
+    for src_id in sources:
+        instance = state.instances.get(src_id)
+        if instance is None:
+            continue
+        for node_id, runner_id in instance.shard_assignments.node_to_runner.items():
+            status = state.runners.get(runner_id)
+            if not isinstance(status, RunnerReady):
+                continue
+            if status.prefill_server_port is None:
+                continue
+            ip = find_ip_prioritised(
+                decode_node, node_id, state.topology, state.node_network, ring=True
+            )
+            if ip is None:
+                continue
+            endpoint = f"{ip}:{status.prefill_server_port}"
+            if endpoint in seen:
+                continue
+            seen.add(endpoint)
+            endpoints.append(endpoint)
+            break
+    return endpoints
+
+
 class Master:
     def __init__(
         self,
@@ -154,20 +204,27 @@ class Master:
                                 ],
                             )
 
+                            decode_instance_id = available_instance_ids[0]
                             task_id = TaskId()
+                            params = command.task_params.model_copy(
+                                update={
+                                    "prefill_endpoints": _prefill_endpoints_for(
+                                        self.state, decode_instance_id
+                                    ),
+                                }
+                            )
                             generated_events.append(
                                 TaskCreated(
                                     task_id=task_id,
                                     task=TextGenerationTask(
                                         task_id=task_id,
                                         command_id=command.command_id,
-                                        instance_id=available_instance_ids[0],
+                                        instance_id=decode_instance_id,
                                         task_status=TaskStatus.Pending,
-                                        task_params=command.task_params,
+                                        task_params=params,
                                     ),
                                 )
                             )
-
                             self.command_task_mapping[command.command_id] = task_id
                         case ImageGeneration():
                             for instance in self.state.instances.values():
@@ -357,6 +414,21 @@ class Master:
                             generated_events.append(
                                 CustomModelCardDeleted(model_id=command.model_id)
                             )
+                        case SetInstanceLink():
+                            link = InstanceLink(
+                                link_id=command.link_id or InstanceLinkId(),
+                                prefill_instances=list(
+                                    dict.fromkeys(command.prefill_instances)
+                                ),
+                                decode_instances=list(
+                                    dict.fromkeys(command.decode_instances)
+                                ),
+                            )
+                            generated_events.append(InstanceLinkCreated(link=link))
+                        case DeleteInstanceLink():
+                            generated_events.append(
+                                InstanceLinkDeleted(link_id=command.link_id)
+                            )
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             # rate limit to 1000 at a time

src/exo/master/placement.py

@@ -1,4 +1,3 @@
-import random
 from collections.abc import Mapping
 from copy import deepcopy
 from typing import Sequence
@@ -46,11 +45,7 @@ from exo.shared.types.worker.instances import (
     MlxRingInstance,
 )
 from exo.shared.types.worker.shards import Sharding
-
-
-def random_ephemeral_port() -> int:
-    port = random.randint(49153, 65535)
-    return port - 1 if port <= 52415 else port
+from exo.utils.ports import random_ephemeral_port
 
 
 def add_instance_to_placements(

src/exo/shared/apply.py

@@ -14,6 +14,8 @@ from exo.shared.types.events import (
     InputChunkReceived,
     InstanceCreated,
     InstanceDeleted,
+    InstanceLinkCreated,
+    InstanceLinkDeleted,
     NodeDownloadProgress,
     NodeGatheredInfo,
     NodeTimedOut,
@@ -29,6 +31,7 @@ from exo.shared.types.events import (
     TracesCollected,
     TracesMerged,
 )
+from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
     NodeIdentity,
     NodeNetworkInfo,
@@ -95,6 +98,10 @@ def event_apply(event: Event, state: State) -> State:
             return apply_topology_edge_created(event, state)
         case TopologyEdgeDeleted():
             return apply_topology_edge_deleted(event, state)
+        case InstanceLinkCreated():
+            return apply_instance_link_created(event, state)
+        case InstanceLinkDeleted():
+            return apply_instance_link_deleted(event, state)
 
 
 def apply(state: State, event: IndexedEvent) -> State:
@@ -194,7 +201,38 @@ def apply_instance_deleted(event: InstanceDeleted, state: State) -> State:
     new_instances: Mapping[InstanceId, Instance] = {
         iid: inst for iid, inst in state.instances.items() if iid != event.instance_id
     }
-    return state.model_copy(update={"instances": new_instances})
+    new_links: dict[InstanceLinkId, InstanceLink] = {}
+    for link_id, link in state.instance_links.items():
+        prefill = [i for i in link.prefill_instances if i != event.instance_id]
+        decode = [i for i in link.decode_instances if i != event.instance_id]
+        if not prefill or not decode:
+            continue
+        if prefill == list(link.prefill_instances) and decode == list(
+            link.decode_instances
+        ):
+            new_links[link_id] = link
+        else:
+            new_links[link_id] = link.model_copy(
+                update={"prefill_instances": prefill, "decode_instances": decode}
+            )
+    return state.model_copy(
+        update={"instances": new_instances, "instance_links": new_links}
+    )
+
+
+def apply_instance_link_created(event: InstanceLinkCreated, state: State) -> State:
+    new_links: Mapping[InstanceLinkId, InstanceLink] = {
+        **state.instance_links,
+        event.link.link_id: event.link,
+    }
+    return state.model_copy(update={"instance_links": new_links})
+
+
+def apply_instance_link_deleted(event: InstanceLinkDeleted, state: State) -> State:
+    new_links: Mapping[InstanceLinkId, InstanceLink] = {
+        lid: link for lid, link in state.instance_links.items() if lid != event.link_id
+    }
+    return state.model_copy(update={"instance_links": new_links})
 
 
 def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State:

src/exo/shared/tests/test_apply/test_apply_instance_link.py

@@ -0,0 +1,72 @@
+from exo.shared.apply import (
+    apply_instance_deleted,
+    apply_instance_link_created,
+    apply_instance_link_deleted,
+)
+from exo.shared.types.events import (
+    InstanceDeleted,
+    InstanceLinkCreated,
+    InstanceLinkDeleted,
+)
+from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
+from exo.shared.types.state import State
+from exo.shared.types.worker.instances import InstanceId
+
+
+def _link(
+    prefill: list[InstanceId],
+    decode: list[InstanceId],
+    link_id: InstanceLinkId | None = None,
+) -> InstanceLink:
+    return InstanceLink(
+        link_id=link_id or InstanceLinkId(),
+        prefill_instances=prefill,
+        decode_instances=decode,
+    )
+
+
+def test_create_link() -> None:
+    state = State()
+    link = _link([InstanceId("a")], [InstanceId("b")])
+    new_state = apply_instance_link_created(InstanceLinkCreated(link=link), state)
+    assert new_state.instance_links == {link.link_id: link}
+
+
+def test_update_replaces_existing_link() -> None:
+    a, b, c = InstanceId("a"), InstanceId("b"), InstanceId("c")
+    link = _link([a], [b])
+    state = State(instance_links={link.link_id: link})
+
+    updated = link.model_copy(update={"decode_instances": [b, c]})
+    new_state = apply_instance_link_created(InstanceLinkCreated(link=updated), state)
+    assert set(new_state.instance_links[link.link_id].decode_instances) == {b, c}
+
+
+def test_delete_link() -> None:
+    link = _link([InstanceId("a")], [InstanceId("b")])
+    state = State(instance_links={link.link_id: link})
+
+    new_state = apply_instance_link_deleted(
+        InstanceLinkDeleted(link_id=link.link_id), state
+    )
+    assert new_state.instance_links == {}
+
+
+def test_instance_deleted_strips_from_links() -> None:
+    a, b, c = InstanceId("a"), InstanceId("b"), InstanceId("c")
+    link = _link([a, c], [b])
+    state = State(instance_links={link.link_id: link})
+
+    new_state = apply_instance_deleted(InstanceDeleted(instance_id=a), state)
+    remaining = new_state.instance_links[link.link_id]
+    assert remaining.prefill_instances == [c]
+    assert remaining.decode_instances == [b]
+
+
+def test_instance_deleted_drops_link_when_role_empties() -> None:
+    a, b = InstanceId("a"), InstanceId("b")
+    link = _link([a], [b])
+    state = State(instance_links={link.link_id: link})
+
+    new_state = apply_instance_deleted(InstanceDeleted(instance_id=a), state)
+    assert link.link_id not in new_state.instance_links

src/exo/shared/types/commands.py

@@ -7,6 +7,7 @@ from exo.api.types import (
 from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.common import CommandId, NodeId, SystemId
+from exo.shared.types.instance_link import InstanceLinkId
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding, ShardMetadata
@@ -89,6 +90,16 @@ class DeleteCustomModelCard(BaseCommand):
     model_id: ModelId
 
 
+class SetInstanceLink(BaseCommand):
+    link_id: InstanceLinkId | None = None
+    prefill_instances: list[InstanceId]
+    decode_instances: list[InstanceId]
+
+
+class DeleteInstanceLink(BaseCommand):
+    link_id: InstanceLinkId
+
+
 DownloadCommand = StartDownload | DeleteDownload | CancelDownload
 
 
@@ -106,6 +117,8 @@ Command = (
     | SendInputChunk
     | AddCustomModelCard
     | DeleteCustomModelCard
+    | SetInstanceLink
+    | DeleteInstanceLink
 )
 
 

src/exo/shared/types/events.py

@@ -7,6 +7,7 @@ from exo.shared.models.model_cards import ModelCard
 from exo.shared.topology import Connection
 from exo.shared.types.chunks import GenerationChunk, InputImageChunk
 from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId, SystemId
+from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -137,6 +138,14 @@ class TracesMerged(BaseEvent):
     traces: list[TraceEventData]
 
 
+class InstanceLinkCreated(BaseEvent):
+    link: InstanceLink
+
+
+class InstanceLinkDeleted(BaseEvent):
+    link_id: InstanceLinkId
+
+
 Event = (
     TestEvent
     | TaskCreated
@@ -158,6 +167,8 @@ Event = (
     | TracesMerged
     | CustomModelCardAdded
     | CustomModelCardDeleted
+    | InstanceLinkCreated
+    | InstanceLinkDeleted
 )
 
 

src/exo/shared/types/instance_link.py

@@ -0,0 +1,13 @@
+from exo.shared.types.common import Id
+from exo.shared.types.worker.instances import InstanceId
+from exo.utils.pydantic_ext import FrozenModel
+
+
+class InstanceLinkId(Id):
+    pass
+
+
+class InstanceLink(FrozenModel):
+    link_id: InstanceLinkId
+    prefill_instances: list[InstanceId]
+    decode_instances: list[InstanceId]

src/exo/shared/types/state.py

@@ -7,6 +7,7 @@ from pydantic.alias_generators import to_camel
 
 from exo.shared.topology import Topology, TopologySnapshot
 from exo.shared.types.common import NodeId
+from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
     DiskUsage,
     MemoryUsage,
@@ -61,6 +62,8 @@ class State(FrozenModel):
     # Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes)
     thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = []
 
+    instance_links: Mapping[InstanceLinkId, InstanceLink] = {}
+
     @field_serializer("topology", mode="plain")
     def _encode_topology(self, value: Topology) -> TopologySnapshot:
         return value.to_snapshot()

src/exo/shared/types/text_generation.py

@@ -115,6 +115,8 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     images: list[Base64Image] = Field(default_factory=list)
     image_hashes: dict[int, Base64ImageHash] = Field(default_factory=dict)
 
+    prefill_endpoints: list[str] = Field(default_factory=list)
+
     def with_card_sampling_defaults(self) -> "TextGenerationTaskParams":
         from exo.shared.models.model_cards import get_card
 

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

@@ -47,7 +47,7 @@ class RunnerWarmingUp(BaseRunnerStatus):
 
 
 class RunnerReady(BaseRunnerStatus):
-    pass
+    prefill_server_port: int | None = None
 
 
 class RunnerRunning(BaseRunnerStatus):

src/exo/utils/ports.py

@@ -0,0 +1,6 @@
+import random
+
+
+def random_ephemeral_port() -> int:
+    port = random.randint(49153, 65535)
+    return port - 1 if port <= 52415 else port

src/exo/worker/disaggregated/protocol.py

@@ -0,0 +1,154 @@
+from typing import BinaryIO, Literal
+
+import msgspec
+
+DType = Literal["bfloat16", "float16", "float32"]
+Layout = Literal["NHD"]
+
+
+class ProtocolError(Exception):
+    pass
+
+
+class Header(msgspec.Struct):
+    request_id: str = ""
+    model_id: str = ""
+    num_layers: int = 0
+    dtype: DType = "bfloat16"
+    layout: Layout = "NHD"
+    start_pos: int = 0
+
+
+class TensorBlob(msgspec.Struct):
+    dtype: DType
+    shape: tuple[int, ...]
+    data: bytes
+
+
+class KVChunk(msgspec.Struct, tag="kv_chunk"):
+    layer_idx: int
+    num_tokens: int
+    n_heads: int
+    head_dim: int
+    dtype: DType
+    keys: bytes
+    values: bytes
+
+    @property
+    def shape(self) -> tuple[int, int, int]:
+        return (self.num_tokens, self.n_heads, self.head_dim)
+
+
+class ArraysState(msgspec.Struct, tag="arrays_state"):
+    layer_idx: int
+    arrays: list[TensorBlob] = []
+
+
+class Done(msgspec.Struct, tag="done"):
+    total_tokens: int
+
+
+class ErrorMessage(msgspec.Struct, tag="error"):
+    code: int
+    message: str
+
+
+Message = KVChunk | ArraysState | Done | ErrorMessage
+
+_msg_encoder = msgspec.msgpack.Encoder()
+_msg_decoder: msgspec.msgpack.Decoder[Message] = msgspec.msgpack.Decoder(Message)
+_header_encoder = msgspec.msgpack.Encoder()
+_header_decoder: msgspec.msgpack.Decoder[Header] = msgspec.msgpack.Decoder(Header)
+
+
+def _read_exactly(stream: BinaryIO, n: int) -> bytes:
+    buf = bytearray()
+    while len(buf) < n:
+        chunk = stream.read(n - len(buf))
+        if not chunk:
+            if len(buf) == 0:
+                return b""
+            raise ConnectionError(f"Connection closed after {len(buf)}/{n} bytes")
+        buf.extend(chunk)
+    return bytes(buf)
+
+
+def _write_frame(stream: BinaryIO, payload: bytes) -> None:
+    stream.write(len(payload).to_bytes(4, "big"))
+    stream.write(payload)
+    stream.flush()
+
+
+def _read_frame(stream: BinaryIO) -> bytes:
+    raw = _read_exactly(stream, 4)
+    if not raw:
+        return b""
+    length = int.from_bytes(raw, "big")
+    return _read_exactly(stream, length)
+
+
+def write_header(stream: BinaryIO, header: Header) -> None:
+    _write_frame(stream, _header_encoder.encode(header))
+
+
+def read_header(stream: BinaryIO) -> Header:
+    payload = _read_frame(stream)
+    if not payload:
+        raise ConnectionError("No header received")
+    try:
+        return _header_decoder.decode(payload)
+    except msgspec.DecodeError as exc:
+        raise ProtocolError(f"Bad header: {exc}") from exc
+
+
+def write_message(stream: BinaryIO, msg: Message) -> None:
+    _write_frame(stream, _msg_encoder.encode(msg))
+
+
+def read_message(stream: BinaryIO) -> Message | None:
+    payload = _read_frame(stream)
+    if not payload:
+        return None
+    try:
+        return _msg_decoder.decode(payload)
+    except msgspec.DecodeError as exc:
+        raise ProtocolError(f"Bad message: {exc}") from exc
+
+
+def write_kv_chunk(
+    stream: BinaryIO,
+    *,
+    layer_idx: int,
+    num_tokens: int,
+    n_heads: int,
+    head_dim: int,
+    dtype: DType,
+    keys: bytes,
+    values: bytes,
+) -> None:
+    write_message(
+        stream,
+        KVChunk(
+            layer_idx=layer_idx,
+            num_tokens=num_tokens,
+            n_heads=n_heads,
+            head_dim=head_dim,
+            dtype=dtype,
+            keys=keys,
+            values=values,
+        ),
+    )
+
+
+def write_arrays_state(
+    stream: BinaryIO, layer_idx: int, arrays: list[TensorBlob]
+) -> None:
+    write_message(stream, ArraysState(layer_idx=layer_idx, arrays=arrays))
+
+
+def write_done(stream: BinaryIO, total_tokens: int) -> None:
+    write_message(stream, Done(total_tokens=total_tokens))
+
+
+def write_error(stream: BinaryIO, code: int, message: str) -> None:
+    write_message(stream, ErrorMessage(code=code, message=message))

src/exo/worker/disaggregated/server.py

@@ -0,0 +1,151 @@
+import json
+import socket
+import socketserver
+import threading
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any, BinaryIO, cast
+
+from loguru import logger
+
+from exo.worker.disaggregated.protocol import (
+    Header,
+    write_error,
+    write_header,
+)
+
+
+@dataclass
+class PrefillJob:
+    request_id: str
+    model_id: str
+    token_ids: list[int]
+    start_pos: int
+
+
+ResolveHandler = Callable[[PrefillJob], bytes | None]
+
+
+def _parse_request_line(line: bytes) -> PrefillJob:
+    parsed = cast(object, json.loads(line.decode("utf-8")))
+    if not isinstance(parsed, dict):
+        raise ValueError(f"Expected JSON object, got {type(parsed).__name__}")
+    d = cast(dict[str, object], parsed)
+
+    def _str(key: str, default: str = "") -> str:
+        v = d.get(key, default)
+        return v if isinstance(v, str) else default
+
+    def _int(key: str, default: int = 0) -> int:
+        v = d.get(key, default)
+        return v if isinstance(v, int) else default
+
+    raw_tokens = d.get("token_ids")
+    if not isinstance(raw_tokens, list):
+        raise ValueError("Missing token_ids in request")
+    token_ids: list[int] = []
+    for t in cast(list[object], raw_tokens):
+        if not isinstance(t, int):
+            raise ValueError("token_ids must be list[int]")
+        token_ids.append(t)
+
+    return PrefillJob(
+        request_id=_str("request_id"),
+        model_id=_str("model"),
+        token_ids=token_ids,
+        start_pos=_int("start_pos"),
+    )
+
+
+def _send_error(wfile: BinaryIO, code: int, message: str) -> None:
+    try:
+        write_header(wfile, Header(num_layers=0, dtype="float32"))
+        write_error(wfile, code=code, message=message)
+    except Exception:
+        pass
+
+
+class _PrefillTCPServer(socketserver.ThreadingTCPServer):
+    allow_reuse_address = True
+    daemon_threads = True
+    resolve: ResolveHandler
+
+
+class _PrefillHandler(socketserver.StreamRequestHandler):
+    def setup(self) -> None:
+        super().setup()
+        sock = cast(socket.socket, self.request)
+        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 4 * 1024 * 1024)
+
+    def handle(self) -> None:
+        server = cast(_PrefillTCPServer, self.server)
+        wfile: BinaryIO = cast(BinaryIO, cast(object, self.wfile))
+        line: bytes = self.rfile.readline()
+        if not line:
+            return
+        try:
+            job = _parse_request_line(line)
+        except (ValueError, json.JSONDecodeError) as exc:
+            _send_error(wfile, 400, f"Bad request: {exc}")
+            return
+        try:
+            payload = server.resolve(job)
+        except Exception as exc:  # noqa: BLE001
+            logger.opt(exception=True).warning(
+                f"Prefill resolve error for request_id={job.request_id}"
+            )
+            _send_error(wfile, 500, str(exc))
+            return
+        if payload is None:
+            _send_error(
+                wfile, 503, f"No payload ready for request_id={job.request_id!r}"
+            )
+            return
+        try:
+            wfile.write(payload)
+            wfile.flush()
+        except Exception:
+            logger.opt(exception=True).warning(
+                f"Failed to write payload for request_id={job.request_id}"
+            )
+
+
+@dataclass
+class PrefillServer:
+    resolve: ResolveHandler
+    host: str = "0.0.0.0"
+    port: int = 0
+    _server: _PrefillTCPServer | None = field(default=None, init=False, repr=False)
+    _thread: threading.Thread | None = field(default=None, init=False, repr=False)
+    _bound_port: int = field(default=0, init=False, repr=False)
+
+    @property
+    def bound_port(self) -> int:
+        return self._bound_port
+
+    def start(self) -> int:
+        if self._server is not None:
+            return self._bound_port
+
+        self._server = _PrefillTCPServer((self.host, self.port), _PrefillHandler)
+        self._server.resolve = self.resolve
+        sock = cast(socket.socket, cast(Any, self._server).socket)
+        addr = cast(tuple[str, int], sock.getsockname())
+        self._bound_port = int(addr[1])
+        self._thread = threading.Thread(
+            target=self._server.serve_forever, name="prefill-server", daemon=True
+        )
+        self._thread.start()
+        logger.info(f"Prefill server listening on {self.host}:{self._bound_port}")
+        return self._bound_port
+
+    def stop(self) -> None:
+        if self._server is not None:
+            self._server.shutdown()
+            self._server.server_close()
+            self._server = None
+        if self._thread is not None:
+            self._thread.join(timeout=5)
+            self._thread = None
+        self._bound_port = 0

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

@@ -16,8 +16,8 @@ from mlx_lm.models.cache import (
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType, Model
 from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
+from exo.worker.engines.mlx.types import KVCacheType, Model
 from exo.worker.runner.bootstrap import logger
 
 if TYPE_CHECKING:

src/exo/worker/engines/mlx/disaggregated/adapter.py

@@ -0,0 +1,227 @@
+import io
+from typing import BinaryIO
+
+import mlx.core as mx
+import numpy as np
+from mlx_lm.models.cache import (
+    ArraysCache,
+    CacheList,
+    KVCache,
+    QuantizedKVCache,
+    RotatingKVCache,
+)
+
+from exo.worker.disaggregated.protocol import (
+    DType,
+    Header,
+    KVChunk,
+    TensorBlob,
+    write_arrays_state,
+    write_done,
+    write_header,
+    write_kv_chunk,
+)
+from exo.worker.engines.mlx.types import KVCacheType
+
+_STR_TO_MX: dict[DType, mx.Dtype] = {
+    "bfloat16": mx.bfloat16,
+    "float16": mx.float16,
+    "float32": mx.float32,
+}
+
+_MX_TO_STR: dict[mx.Dtype, DType] = {v: k for k, v in _STR_TO_MX.items()}
+
+
+def mx_dtype_to_str(dtype: mx.Dtype) -> DType:
+    if dtype not in _MX_TO_STR:
+        raise ValueError(f"Unsupported mlx dtype on wire: {dtype}")
+    return _MX_TO_STR[dtype]
+
+
+def wire_dtype_from_cache(caches: KVCacheType) -> DType:
+    for c in caches:
+        keys: mx.array | None = getattr(c, "keys", None)
+        if keys is None:
+            continue
+        if keys.dtype in _MX_TO_STR:
+            return _MX_TO_STR[keys.dtype]
+        break
+    return "bfloat16"
+
+
+def str_to_mx_dtype(dtype: DType) -> mx.Dtype:
+    if dtype not in _STR_TO_MX:
+        raise ValueError(f"Unsupported wire dtype: {dtype!r}")
+    return _STR_TO_MX[dtype]
+
+
+def array_to_bytes(t: mx.array) -> bytes:
+    # bf16 has no native numpy dtype; bitcast through uint16.
+    if t.dtype == mx.bfloat16:
+        return np.asarray(t.view(mx.uint16)).tobytes()
+    if t.dtype in (mx.float16, mx.float32):
+        return np.asarray(t).tobytes()
+    raise ValueError(f"Unsupported mlx dtype for wire: {t.dtype}")
+
+
+def bytes_to_array(data: bytes, shape: tuple[int, ...], dtype: DType) -> mx.array:
+    if dtype == "bfloat16":
+        arr = np.frombuffer(data, dtype=np.uint16).reshape(shape).copy()
+        return mx.array(arr).view(mx.bfloat16)
+    if dtype == "float16":
+        arr = np.frombuffer(data, dtype=np.float16).reshape(shape).copy()
+        return mx.array(arr)
+    if dtype == "float32":
+        arr = np.frombuffer(data, dtype=np.float32).reshape(shape).copy()
+        return mx.array(arr)
+    raise ValueError(f"Unsupported wire dtype for mlx: {dtype!r}")
+
+
+def bhsd_to_nhd(t: mx.array) -> mx.array:
+    if t.ndim != 4 or int(t.shape[0]) != 1:
+        raise ValueError(f"Expected BHSD with B=1, got shape={tuple(t.shape)}")
+    return mx.transpose(t[0], (1, 0, 2))
+
+
+def nhd_to_bhsd(t: mx.array) -> mx.array:
+    if t.ndim != 3:
+        raise ValueError(f"Expected NHD (3D), got shape={tuple(t.shape)}")
+    return mx.expand_dims(mx.transpose(t, (1, 0, 2)), 0)
+
+
+def send_mlx_kv_cache(
+    stream: BinaryIO,
+    caches: KVCacheType,
+    *,
+    dtype: DType,
+    start_pos: int = 0,
+    max_tokens: int | None = None,
+) -> int:
+    tokens_sent = 0
+    for layer_idx, c in enumerate(caches):
+        if isinstance(c, (QuantizedKVCache, CacheList)):
+            continue
+        if isinstance(c, (KVCache, RotatingKVCache)):
+            keys = c.keys
+            values = c.values
+            if keys is None or values is None:
+                continue
+            offset = int(c.offset)
+            if max_tokens is not None:
+                offset = min(offset, max_tokens)
+            if offset <= start_pos:
+                continue
+            with mx.stream(mx.Device(mx.cpu)):
+                k = mx.array(keys[:, :, start_pos:offset, :])
+                v = mx.array(values[:, :, start_pos:offset, :])
+                k_nhd = bhsd_to_nhd(k)
+                v_nhd = bhsd_to_nhd(v)
+                mx.eval(k_nhd, v_nhd)
+            num_tokens = int(k_nhd.shape[0])
+            n_heads = int(k_nhd.shape[1])
+            head_dim = int(k_nhd.shape[2])
+            write_kv_chunk(
+                stream,
+                layer_idx=layer_idx,
+                num_tokens=num_tokens,
+                n_heads=n_heads,
+                head_dim=head_dim,
+                dtype=dtype,
+                keys=array_to_bytes(k_nhd),
+                values=array_to_bytes(v_nhd),
+            )
+            tokens_sent = max(tokens_sent, num_tokens)
+        else:
+            blobs: list[TensorBlob] = []
+            for a in c.state:
+                if a is None:
+                    continue
+                with mx.stream(mx.Device(mx.cpu)):
+                    a_cpu = mx.array(a)
+                    mx.eval(a_cpu)
+                blobs.append(
+                    TensorBlob(
+                        dtype=mx_dtype_to_str(a_cpu.dtype),
+                        shape=tuple(int(d) for d in a_cpu.shape),
+                        data=array_to_bytes(a_cpu),
+                    )
+                )
+            if blobs:
+                write_arrays_state(stream, layer_idx, blobs)
+    return tokens_sent
+
+
+def chunk_to_mlx_nhd(chunk: KVChunk) -> tuple[mx.array, mx.array]:
+    shape = chunk.shape
+    return (
+        bytes_to_array(chunk.keys, shape, chunk.dtype),
+        bytes_to_array(chunk.values, shape, chunk.dtype),
+    )
+
+
+def blob_to_mlx(blob: TensorBlob) -> mx.array:
+    return bytes_to_array(blob.data, blob.shape, blob.dtype)
+
+
+def inject_kv_chunk(
+    cache: KVCache,
+    keys_nhd: mx.array,
+    values_nhd: mx.array,
+    offset: int,
+    *,
+    start_pos: int = 0,
+    existing_k: mx.array | None = None,
+    existing_v: mx.array | None = None,
+) -> None:
+    k_bhsd = nhd_to_bhsd(keys_nhd)
+    v_bhsd = nhd_to_bhsd(values_nhd)
+    if start_pos > 0 and existing_k is not None and existing_v is not None:
+        cache.keys = mx.concatenate([existing_k[:, :, :start_pos, :], k_bhsd], axis=2)
+        cache.values = mx.concatenate([existing_v[:, :, :start_pos, :], v_bhsd], axis=2)
+    else:
+        cache.keys = k_bhsd
+        cache.values = v_bhsd
+    cache.offset = offset
+
+
+def inject_rotating_kv_chunk(
+    cache: RotatingKVCache,
+    keys_nhd: mx.array,
+    values_nhd: mx.array,
+    offset: int,
+) -> None:
+    k_bhsd = nhd_to_bhsd(keys_nhd)
+    v_bhsd = nhd_to_bhsd(values_nhd)
+    cache.keys = k_bhsd
+    cache.values = v_bhsd
+    cache.offset = offset
+    cache._idx = int(k_bhsd.shape[2])
+
+
+def inject_arrays_cache(cache: ArraysCache, blobs: list[TensorBlob]) -> None:
+    cache.state = [blob_to_mlx(b) for b in blobs]
+
+
+def serialize_mlx_cache_to_payload(
+    caches: KVCacheType,
+    *,
+    dtype: DType,
+    model_id: str = "",
+    request_id: str = "",
+    start_pos: int = 0,
+    max_tokens: int | None = None,
+) -> bytes:
+    buf = io.BytesIO()
+    header = Header(
+        request_id=request_id,
+        model_id=model_id,
+        num_layers=len(caches),
+        dtype=dtype,
+        start_pos=start_pos,
+    )
+    write_header(buf, header)
+    tokens_sent = send_mlx_kv_cache(
+        buf, caches, dtype=dtype, start_pos=start_pos, max_tokens=max_tokens
+    )
+    write_done(buf, tokens_sent)
+    return buf.getvalue()

src/exo/worker/engines/mlx/disaggregated/client.py

@@ -0,0 +1,165 @@
+import json
+import socket
+from collections import defaultdict
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import BinaryIO, cast
+
+import mlx.core as mx
+from loguru import logger
+from mlx_lm.models.cache import ArraysCache, KVCache, RotatingKVCache
+
+from exo.worker.engines.mlx.disaggregated.adapter import (
+    chunk_to_mlx_nhd,
+    inject_arrays_cache,
+    inject_kv_chunk,
+    inject_rotating_kv_chunk,
+)
+from exo.worker.disaggregated.protocol import (
+    ArraysState,
+    Done,
+    Header,
+    KVChunk,
+    TensorBlob,
+    read_header,
+    read_message,
+)
+
+DEFAULT_PORT = 8900
+_SOCKET_TIMEOUT_SECS = 60
+_RECV_BUFFER_BYTES = 4 * 1024 * 1024
+
+
+@dataclass
+class PrefillRequest:
+    model_id: str
+    token_ids: list[int]
+    start_pos: int = 0
+    request_id: str = ""
+
+
+@dataclass
+class PrefillResult:
+    header: Header
+    kv_chunks: dict[int, list[KVChunk]] = field(
+        default_factory=dict[int, list[KVChunk]]
+    )
+    arrays: dict[int, list[TensorBlob]] = field(
+        default_factory=dict[int, list[TensorBlob]]
+    )
+    total_tokens: int = 0
+
+
+def _parse_endpoint(endpoint: str) -> tuple[str, int]:
+    if ":" in endpoint:
+        host, port_str = endpoint.rsplit(":", 1)
+        return host, int(port_str)
+    return endpoint, DEFAULT_PORT
+
+
+def remote_prefill_fetch(
+    endpoint: str,
+    request: PrefillRequest,
+    on_header: Callable[[Header], None] | None = None,
+    on_kv_chunk: Callable[[KVChunk, int], None] | None = None,
+    timeout_secs: float = _SOCKET_TIMEOUT_SECS,
+) -> PrefillResult:
+    host, port = _parse_endpoint(endpoint)
+    logger.info(
+        f"Connecting to prefill server at {host}:{port} "
+        f"({len(request.token_ids)} tokens, start_pos={request.start_pos})"
+    )
+
+    sock = socket.create_connection((host, port), timeout=timeout_secs)
+    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, _RECV_BUFFER_BYTES)
+    try:
+        req_bytes = (
+            json.dumps(
+                {
+                    "request_id": request.request_id,
+                    "model": request.model_id,
+                    "token_ids": request.token_ids,
+                    "start_pos": request.start_pos,
+                }
+            ).encode("utf-8")
+            + b"\n"
+        )
+        sock.sendall(req_bytes)
+
+        raw_stream = sock.makefile("rb", buffering=256 * 1024)
+        stream: BinaryIO = cast(BinaryIO, cast(object, raw_stream))
+
+        header = read_header(stream)
+        if on_header is not None:
+            on_header(header)
+
+        result = PrefillResult(header=header)
+        kv_by_layer: dict[int, list[KVChunk]] = defaultdict(list)
+        chunks_received = 0
+
+        while True:
+            msg = read_message(stream)
+            if msg is None:
+                break
+            if isinstance(msg, KVChunk):
+                kv_by_layer[msg.layer_idx].append(msg)
+                chunks_received += 1
+                if on_kv_chunk is not None:
+                    on_kv_chunk(msg, chunks_received)
+            elif isinstance(msg, ArraysState):
+                result.arrays[msg.layer_idx] = msg.arrays
+            elif isinstance(msg, Done):
+                result.total_tokens = msg.total_tokens
+                break
+            else:
+                raise RuntimeError(f"Prefill server error [{msg.code}]: {msg.message}")
+
+        result.kv_chunks = dict(kv_by_layer)
+        return result
+    finally:
+        sock.close()
+
+
+def ingest_into_mlx_cache(
+    result: PrefillResult,
+    caches: list[KVCache | RotatingKVCache | ArraysCache],
+    *,
+    start_pos: int = 0,
+) -> int:
+    max_received = max(
+        (sum(c.num_tokens for c in chunks) for chunks in result.kv_chunks.values()),
+        default=0,
+    )
+    final_offset = start_pos + max_received
+
+    for i, cache in enumerate(caches):
+        if i in result.kv_chunks:
+            chunks = result.kv_chunks[i]
+            if len(chunks) == 1:
+                k_nhd, v_nhd = chunk_to_mlx_nhd(chunks[0])
+            else:
+                decoded = [chunk_to_mlx_nhd(c) for c in chunks]
+                k_nhd = mx.concatenate([k for k, _ in decoded], axis=0)
+                v_nhd = mx.concatenate([v for _, v in decoded], axis=0)
+
+            if isinstance(cache, RotatingKVCache):
+                inject_rotating_kv_chunk(cache, k_nhd, v_nhd, final_offset)
+            elif isinstance(cache, KVCache):
+                if start_pos > 0:
+                    inject_kv_chunk(
+                        cache,
+                        k_nhd,
+                        v_nhd,
+                        final_offset,
+                        start_pos=start_pos,
+                        existing_k=cache.keys,
+                        existing_v=cache.values,
+                    )
+                else:
+                    inject_kv_chunk(cache, k_nhd, v_nhd, final_offset)
+
+        if i in result.arrays and isinstance(cache, ArraysCache):
+            inject_arrays_cache(cache, result.arrays[i])
+
+    return final_offset

src/exo/worker/engines/mlx/disaggregated/tests/test_end_to_end.py

@@ -0,0 +1,108 @@
+import mlx.core as mx
+import numpy as np
+import pytest
+from mlx_lm.models.cache import KVCache
+
+from exo.worker.engines.mlx.disaggregated.adapter import serialize_mlx_cache_to_payload
+from exo.worker.engines.mlx.disaggregated.client import (
+    PrefillRequest,
+    ingest_into_mlx_cache,
+    remote_prefill_fetch,
+)
+from exo.worker.engines.mlx.disaggregated.server import (
+    PrefillJob,
+    PrefillPayloadLookup,
+    PrefillServer,
+)
+
+
+def _equal(a: mx.array, b: mx.array) -> bool:
+    if a.dtype != b.dtype or tuple(a.shape) != tuple(b.shape):
+        return False
+    if a.dtype == mx.bfloat16:
+        return bool(
+            np.array_equal(np.asarray(a.view(mx.uint16)), np.asarray(b.view(mx.uint16)))
+        )
+    return bool(np.array_equal(np.asarray(a), np.asarray(b)))
+
+
+def _make_cache(seq_len: int, n_heads: int, head_dim: int) -> KVCache:
+    mx.random.seed(0)
+    cache = KVCache()
+    cache.keys = (mx.random.uniform(shape=(1, n_heads, seq_len, head_dim)) * 10).astype(
+        mx.bfloat16
+    )
+    cache.values = (
+        mx.random.uniform(shape=(1, n_heads, seq_len, head_dim)) * 10
+    ).astype(mx.bfloat16)
+    cache.offset = seq_len
+    return cache
+
+
+@pytest.mark.slow
+def test_server_client_roundtrip() -> None:
+    seq_len = 5
+    n_heads = 2
+    head_dim = 4
+    gold = _make_cache(seq_len, n_heads, head_dim)
+
+    lookup = PrefillPayloadLookup()
+    payload = serialize_mlx_cache_to_payload(
+        [gold], dtype="bfloat16", model_id="test-model", request_id="req-1"
+    )
+    lookup.register("req-1", payload)
+
+    def resolve(job: PrefillJob) -> bytes | None:
+        return lookup.pop(job.request_id)
+
+    server = PrefillServer(resolve=resolve, host="127.0.0.1", port=0)
+    port = server.start()
+    try:
+        result = remote_prefill_fetch(
+            endpoint=f"127.0.0.1:{port}",
+            request=PrefillRequest(
+                model_id="test-model",
+                token_ids=list(range(seq_len)),
+                request_id="req-1",
+            ),
+        )
+        assert result.total_tokens == seq_len
+        assert 0 in result.kv_chunks
+
+        dst = KVCache()
+        final_offset = ingest_into_mlx_cache(result, [dst])
+        assert final_offset == seq_len
+        assert dst.offset == seq_len
+        dst_k = dst.keys
+        dst_v = dst.values
+        gold_k = gold.keys
+        gold_v = gold.values
+        assert dst_k is not None and dst_v is not None
+        assert gold_k is not None and gold_v is not None
+        assert _equal(dst_k, gold_k)
+        assert _equal(dst_v, gold_v)
+    finally:
+        server.stop()
+
+
+@pytest.mark.slow
+def test_server_reports_unknown_request() -> None:
+    lookup = PrefillPayloadLookup()
+
+    def resolve(job: PrefillJob) -> bytes | None:
+        return lookup.pop(job.request_id)
+
+    server = PrefillServer(resolve=resolve, host="127.0.0.1", port=0)
+    port = server.start()
+    try:
+        with pytest.raises(RuntimeError, match="No payload ready"):
+            _ = remote_prefill_fetch(
+                endpoint=f"127.0.0.1:{port}",
+                request=PrefillRequest(
+                    model_id="test-model",
+                    token_ids=[1, 2, 3],
+                    request_id="never-registered",
+                ),
+            )
+    finally:
+        server.stop()

src/exo/worker/engines/mlx/disaggregated/tests/test_mlx_adapter.py

@@ -0,0 +1,138 @@
+import io
+
+import mlx.core as mx
+import numpy as np
+from mlx_lm.models.cache import ArraysCache, KVCache
+
+from exo.worker.engines.mlx.disaggregated.adapter import (
+    array_to_bytes,
+    bhsd_to_nhd,
+    bytes_to_array,
+    chunk_to_mlx_nhd,
+    inject_arrays_cache,
+    inject_kv_chunk,
+    nhd_to_bhsd,
+    send_mlx_kv_cache,
+)
+from exo.worker.engines.mlx.disaggregated.protocol import (
+    ArraysState,
+    Done,
+    KVChunk,
+    TensorBlob,
+    make_header,
+    read_header,
+    read_message,
+    write_done,
+    write_header,
+)
+
+
+def _equal(a: mx.array, b: mx.array) -> bool:
+    if a.dtype != b.dtype or tuple(a.shape) != tuple(b.shape):
+        return False
+    if a.dtype == mx.bfloat16:
+        return bool(
+            np.array_equal(np.asarray(a.view(mx.uint16)), np.asarray(b.view(mx.uint16)))
+        )
+    return bool(np.array_equal(np.asarray(a), np.asarray(b)))
+
+
+def _rand(shape: tuple[int, ...], dtype: mx.Dtype) -> mx.array:
+    mx.random.seed(0)
+    return (mx.random.uniform(shape=shape) * 10).astype(dtype)
+
+
+def test_bytes_roundtrip_bf16() -> None:
+    x = _rand((2, 3, 4), mx.bfloat16)
+    y = bytes_to_array(array_to_bytes(x), (2, 3, 4), "bfloat16")
+    assert _equal(x, y)
+
+
+def test_bytes_roundtrip_f16() -> None:
+    x = _rand((5,), mx.float16)
+    y = bytes_to_array(array_to_bytes(x), (5,), "float16")
+    assert _equal(x, y)
+
+
+def test_bytes_roundtrip_f32() -> None:
+    x = _rand((2, 2), mx.float32)
+    y = bytes_to_array(array_to_bytes(x), (2, 2), "float32")
+    assert _equal(x, y)
+
+
+def test_bhsd_nhd_roundtrip() -> None:
+    bhsd = _rand((1, 4, 7, 8), mx.float32)
+    nhd = bhsd_to_nhd(bhsd)
+    assert tuple(nhd.shape) == (7, 4, 8)
+    back = nhd_to_bhsd(nhd)
+    assert _equal(bhsd, back)
+
+
+def test_kv_cache_inject_roundtrip() -> None:
+    n_heads, seq_len, head_dim = 3, 5, 4
+    k_bhsd = _rand((1, n_heads, seq_len, head_dim), mx.float32)
+    v_bhsd = _rand((1, n_heads, seq_len, head_dim), mx.float32)
+    k_nhd = bhsd_to_nhd(k_bhsd)
+    v_nhd = bhsd_to_nhd(v_bhsd)
+
+    cache = KVCache()
+    inject_kv_chunk(cache, k_nhd, v_nhd, offset=seq_len)
+    assert cache.offset == seq_len
+    assert cache.keys is not None and cache.values is not None
+    assert _equal(cache.keys, k_bhsd)
+    assert _equal(cache.values, v_bhsd)
+
+
+def test_arrays_cache_inject() -> None:
+    a = _rand((3,), mx.float32)
+    b = _rand((2, 2), mx.bfloat16)
+    blobs = [
+        TensorBlob(dtype="float32", shape=(3,), data=array_to_bytes(a)),
+        TensorBlob(dtype="bfloat16", shape=(2, 2), data=array_to_bytes(b)),
+    ]
+    cache = ArraysCache(size=2)
+    inject_arrays_cache(cache, blobs)
+    s0 = cache.state[0]
+    s1 = cache.state[1]
+    assert s0 is not None and s1 is not None
+    assert _equal(s0, a)
+    assert _equal(s1, b)
+
+
+def test_send_mlx_cache_end_to_end() -> None:
+    n_heads, head_dim = 2, 4
+    seq_len = 3
+
+    k_bhsd = _rand((1, n_heads, seq_len, head_dim), mx.bfloat16)
+    v_bhsd = _rand((1, n_heads, seq_len, head_dim), mx.bfloat16)
+
+    src = KVCache()
+    src.keys = k_bhsd
+    src.values = v_bhsd
+    src.offset = seq_len
+
+    buf = io.BytesIO()
+    hdr = make_header(num_layers=1, dtype="bfloat16")
+    write_header(buf, hdr)
+    tokens = send_mlx_kv_cache(buf, [src])
+    write_done(buf, tokens)
+    buf.seek(0)
+
+    got_hdr = read_header(buf)
+    assert got_hdr["num_layers"] == 1
+
+    msg = read_message(buf, got_hdr)
+    assert isinstance(msg, KVChunk)
+    k_nhd, v_nhd = chunk_to_mlx_nhd(msg)
+    dst = KVCache()
+    inject_kv_chunk(dst, k_nhd, v_nhd, offset=msg.num_tokens)
+
+    done = read_message(buf, got_hdr)
+    assert isinstance(done, Done)
+    assert done.total_tokens == seq_len
+
+    assert dst.offset == seq_len
+    assert dst.keys is not None and dst.values is not None
+    assert _equal(dst.keys, k_bhsd)
+    assert _equal(dst.values, v_bhsd)
+    _ = ArraysState

src/exo/worker/engines/mlx/disaggregated/tests/test_protocol_roundtrip.py

@@ -0,0 +1,157 @@
+import io
+
+import pytest
+
+from exo.worker.engines.mlx.disaggregated.protocol import (
+    ArraysState,
+    Done,
+    ErrorMessage,
+    KVChunk,
+    ProtocolError,
+    TensorBlob,
+    dtype_size,
+    header_dtype,
+    header_int,
+    make_header,
+    read_header,
+    read_message,
+    write_arrays_state,
+    write_done,
+    write_error,
+    write_header,
+    write_kv_chunk,
+)
+
+
+def _mk_bytes(n: int) -> bytes:
+    return bytes(i & 0xFF for i in range(n))
+
+
+def test_header_roundtrip() -> None:
+    hdr = make_header(
+        num_layers=32, dtype="bfloat16", model_id="m", request_id="r", start_pos=42
+    )
+    buf = io.BytesIO()
+    write_header(buf, hdr)
+    buf.seek(0)
+    got = read_header(buf)
+    assert got == hdr
+    assert header_dtype(got) == "bfloat16"
+    assert header_int(got, "num_layers") == 32
+    assert header_int(got, "start_pos") == 42
+
+
+def test_header_dtype_rejects_unknown() -> None:
+    with pytest.raises(ProtocolError):
+        header_dtype({"dtype": "int4"})
+
+
+def test_kv_chunk_roundtrip() -> None:
+    dtype = "bfloat16"
+    num_tokens, n_heads, head_dim = 7, 4, 8
+    n_bytes = num_tokens * n_heads * head_dim * dtype_size(dtype)
+    keys = _mk_bytes(n_bytes)
+    values = _mk_bytes(n_bytes)[::-1]
+
+    buf = io.BytesIO()
+    write_kv_chunk(
+        buf,
+        layer_idx=3,
+        num_tokens=num_tokens,
+        n_heads=n_heads,
+        head_dim=head_dim,
+        keys=keys,
+        values=values,
+    )
+    buf.seek(0)
+    msg = read_message(buf, make_header(num_layers=1, dtype=dtype))
+    assert isinstance(msg, KVChunk)
+    assert msg.layer_idx == 3
+    assert msg.shape == (num_tokens, n_heads, head_dim)
+    assert msg.dtype == dtype
+    assert msg.keys == keys
+    assert msg.values == values
+
+
+def test_arrays_state_roundtrip() -> None:
+    arrs = [
+        TensorBlob(dtype="float32", shape=(2, 3), data=_mk_bytes(2 * 3 * 4)),
+        TensorBlob(dtype="bfloat16", shape=(5,), data=_mk_bytes(5 * 2)),
+    ]
+    buf = io.BytesIO()
+    write_arrays_state(buf, layer_idx=9, arrays=arrs)
+    buf.seek(0)
+    msg = read_message(buf, make_header(num_layers=1, dtype="float32"))
+    assert isinstance(msg, ArraysState)
+    assert msg.layer_idx == 9
+    assert len(msg.arrays) == 2
+    assert msg.arrays[0].dtype == "float32"
+    assert msg.arrays[0].shape == (2, 3)
+    assert msg.arrays[0].data == arrs[0].data
+    assert msg.arrays[1].dtype == "bfloat16"
+    assert msg.arrays[1].shape == (5,)
+    assert msg.arrays[1].data == arrs[1].data
+
+
+def test_done_roundtrip() -> None:
+    buf = io.BytesIO()
+    write_done(buf, 1234)
+    buf.seek(0)
+    msg = read_message(buf, make_header(num_layers=1, dtype="float32"))
+    assert isinstance(msg, Done)
+    assert msg.total_tokens == 1234
+
+
+def test_error_roundtrip() -> None:
+    buf = io.BytesIO()
+    write_error(buf, code=42, message="boom")
+    buf.seek(0)
+    msg = read_message(buf, make_header(num_layers=1, dtype="float32"))
+    assert isinstance(msg, ErrorMessage)
+    assert msg.code == 42
+    assert msg.message == "boom"
+
+
+def test_stream_of_messages() -> None:
+    hdr = make_header(num_layers=2, dtype="float32")
+    buf = io.BytesIO()
+    write_header(buf, hdr)
+    write_kv_chunk(
+        buf,
+        layer_idx=0,
+        num_tokens=1,
+        n_heads=1,
+        head_dim=2,
+        keys=_mk_bytes(1 * 1 * 2 * 4),
+        values=_mk_bytes(1 * 1 * 2 * 4),
+    )
+    write_arrays_state(
+        buf,
+        layer_idx=1,
+        arrays=[TensorBlob(dtype="float32", shape=(1,), data=_mk_bytes(4))],
+    )
+    write_done(buf, total_tokens=1)
+    buf.seek(0)
+
+    got_hdr = read_header(buf)
+    assert got_hdr == hdr
+
+    m1 = read_message(buf, got_hdr)
+    m2 = read_message(buf, got_hdr)
+    m3 = read_message(buf, got_hdr)
+    m4 = read_message(buf, got_hdr)
+    assert isinstance(m1, KVChunk)
+    assert isinstance(m2, ArraysState)
+    assert isinstance(m3, Done)
+    assert m4 is None
+
+
+def test_unknown_message_type_raises() -> None:
+    # Construct a stream with a valid header and then an unknown tag.
+    buf = io.BytesIO()
+    write_header(buf, make_header(num_layers=1, dtype="float32"))
+    buf.write(bytes([0xFF]))
+    buf.seek(0)
+    hdr = read_header(buf)
+    with pytest.raises(ProtocolError):
+        _ = read_message(buf, hdr)

src/exo/worker/engines/mlx/disaggregated/tests/test_server_drain.py

@@ -0,0 +1,148 @@
+"""End-to-end producer test: server thread receives request, main thread
+drains queue, runs the prefill callable, returns bytes."""
+
+import io
+import queue
+import threading
+from collections.abc import Callable
+
+import mlx.core as mx
+import numpy as np
+import pytest
+from mlx_lm.models.cache import KVCache
+
+from exo.utils.ports import random_ephemeral_port
+from exo.worker.engines.mlx.disaggregated.adapter import (
+    chunk_to_mlx_nhd,
+    serialize_mlx_cache_to_payload,
+)
+from exo.worker.engines.mlx.disaggregated.client import (
+    PrefillRequest,
+    PrefillResult,
+    ingest_into_mlx_cache,
+    remote_prefill_fetch,
+)
+from exo.worker.engines.mlx.disaggregated.protocol import (
+    Done,
+    KVChunk,
+    make_header,
+    read_header,
+    read_message,
+    write_done,
+    write_header,
+)
+from exo.worker.engines.mlx.disaggregated.server import (
+    PrefillJob,
+    PrefillServer,
+)
+
+
+def _equal(a: mx.array, b: mx.array) -> bool:
+    if a.dtype != b.dtype or tuple(a.shape) != tuple(b.shape):
+        return False
+    if a.dtype == mx.bfloat16:
+        return bool(
+            np.array_equal(np.asarray(a.view(mx.uint16)), np.asarray(b.view(mx.uint16)))
+        )
+    return bool(np.array_equal(np.asarray(a), np.asarray(b)))
+
+
+def _make_cache(seq_len: int, n_heads: int, head_dim: int) -> KVCache:
+    mx.random.seed(0)
+    cache = KVCache()
+    cache.keys = (mx.random.uniform(shape=(1, n_heads, seq_len, head_dim)) * 10).astype(
+        mx.bfloat16
+    )
+    cache.values = (
+        mx.random.uniform(shape=(1, n_heads, seq_len, head_dim)) * 10
+    ).astype(mx.bfloat16)
+    cache.offset = seq_len
+    return cache
+
+
+@pytest.mark.slow
+def test_server_drains_via_main_thread() -> None:
+    seq_len = 4
+    n_heads = 2
+    head_dim = 4
+    gold = _make_cache(seq_len, n_heads, head_dim)
+
+    # Mimic the runner's queue-and-drain pattern.
+    request_queue: queue.Queue[
+        tuple[PrefillJob, threading.Event, list[bytes | None]]
+    ] = queue.Queue()
+
+    def resolve(job: PrefillJob) -> bytes | None:
+        event = threading.Event()
+        holder: list[bytes | None] = [None]
+        request_queue.put((job, event, holder))
+        if not event.wait(timeout=5):
+            return None
+        return holder[0]
+
+    server = PrefillServer(
+        resolve=resolve, host="127.0.0.1", port=random_ephemeral_port()
+    )
+    port = server.start()
+
+    def serve_one() -> bytes:
+        return serialize_mlx_cache_to_payload(
+            [gold], dtype="bfloat16", model_id="m", request_id="req-1"
+        )
+
+    drained_job: list[PrefillJob] = []
+    fetch_result: list[PrefillResult] = []
+
+    def fetcher() -> None:
+        fetch_result.append(
+            remote_prefill_fetch(
+                endpoint=f"127.0.0.1:{port}",
+                request=PrefillRequest(
+                    model_id="m", token_ids=list(range(seq_len)), request_id="req-1"
+                ),
+            )
+        )
+
+    fetch = threading.Thread(target=fetcher, daemon=True)
+    fetch.start()
+    try:
+        # Test thread acts as the runner's main thread: drain the queue.
+        job, event, holder = request_queue.get(timeout=5)
+        drained_job.append(job)
+        try:
+            holder[0] = serve_one()
+        finally:
+            event.set()
+        fetch.join(timeout=5)
+        assert fetch_result, "fetcher did not return"
+        result = fetch_result[0]
+        assert drained_job[0].request_id == "req-1"
+        assert result.total_tokens == seq_len
+
+        dst = KVCache()
+        ingest_into_mlx_cache(result, [dst])
+        assert dst.offset == seq_len
+        dst_k = dst.keys
+        dst_v = dst.values
+        gold_k = gold.keys
+        gold_v = gold.values
+        assert dst_k is not None and dst_v is not None
+        assert gold_k is not None and gold_v is not None
+        assert _equal(dst_k, gold_k)
+        assert _equal(dst_v, gold_v)
+    finally:
+        server.stop()
+
+    # Silence unused imports referenced only for completeness.
+    _ = (
+        Callable,
+        KVChunk,
+        Done,
+        chunk_to_mlx_nhd,
+        io.BytesIO,
+        make_header,
+        read_header,
+        read_message,
+        write_done,
+        write_header,
+    )

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

@@ -1,5 +1,6 @@
 import contextlib
 import time
+import uuid
 from dataclasses import dataclass, field
 from typing import Callable, Literal, cast
 
@@ -23,7 +24,6 @@ from exo.api.types import (
     Usage,
 )
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType, Model
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.runner_response import GenerationResponse
 from exo.worker.engines.mlx.cache import (
@@ -40,10 +40,12 @@ from exo.worker.engines.mlx.generator.generate import (
     patch_embed_tokens,
     prefill,
 )
+from exo.worker.engines.mlx.generator.remote_prefill import remote_prefill
 from exo.worker.engines.mlx.patches.opt_batch_gen import (
     set_needs_topk,
     take_ready_topk,
 )
+from exo.worker.engines.mlx.types import KVCacheType, Model
 from exo.worker.engines.mlx.utils_mlx import (
     fix_unmatched_think_end_tokens,
     system_prompt_token_count,
@@ -57,6 +59,7 @@ from exo.worker.engines.mlx.vision import (
 from exo.worker.runner.bootstrap import logger
 
 _MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5
+REMOTE_PREFILL_MIN_TOKENS = 1000
 
 
 def _stop_sequences(task_params: TextGenerationTaskParams) -> list[str]:
@@ -199,17 +202,51 @@ class ExoBatchGenerator:
             if vision is not None
             else contextlib.nullcontext()
         )
+        uncached_count = len(prompt_tokens)
+        use_remote = (
+            uncached_count > REMOTE_PREFILL_MIN_TOKENS
+            and bool(task_params.prefill_endpoints)
+            and not is_bench
+        )
+
+        _prefill_tps: float = 0.0
+        _prefill_tokens: int = 0
+        cache_snapshots: list[CacheSnapshot] = []
+        remote_prefilled = False
         with vision_ctx:
-            _prefill_tps, _prefill_tokens, cache_snapshots = prefill(
-                self.model,
-                self.tokenizer,
-                sampler,
-                prompt_tokens[:-1],
-                cache,
-                self.group,
-                on_prefill_progress,
-                distributed_prompt_progress_callback,
-            )
+            if use_remote:
+                try:
+                    _prefill_tps, _prefill_tokens, cache_snapshots = remote_prefill(
+                        self.model,
+                        self.tokenizer,
+                        sampler,
+                        prompt_tokens[:-1],
+                        cache,
+                        self.group,
+                        on_prefill_progress,
+                        distributed_prompt_progress_callback,
+                        endpoint=task_params.prefill_endpoints[0],
+                        request_id=str(uuid.uuid4()),
+                        model_id=str(task_params.model),
+                        start_pos=prefix_hit_length,
+                    )
+                    remote_prefilled = True
+                except Exception:
+                    logger.opt(exception=True).warning(
+                        "Remote prefill failed, falling back to local prefill"
+                    )
+
+            if not remote_prefilled:
+                _prefill_tps, _prefill_tokens, cache_snapshots = prefill(
+                    self.model,
+                    self.tokenizer,
+                    sampler,
+                    prompt_tokens[:-1],
+                    cache,
+                    self.group,
+                    on_prefill_progress,
+                    distributed_prompt_progress_callback,
+                )
 
         prefix_cache_hit: Literal["none", "partial", "exact"] = "none"
         if matched_index is not None and prefix_hit_length > 0:

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

@@ -24,7 +24,6 @@ from exo.api.types import (
 )
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType, Model
 from exo.shared.types.text_generation import (
     InputMessage,
     InputMessageContent,
@@ -55,6 +54,7 @@ from exo.worker.engines.mlx.constants import (
     KV_GROUP_SIZE,
     MAX_TOKENS,
 )
+from exo.worker.engines.mlx.types import KVCacheType, Model
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
     fix_unmatched_think_end_tokens,

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

@@ -0,0 +1,79 @@
+import time
+from collections.abc import Callable
+from typing import cast
+
+import mlx.core as mx
+from mlx_lm.models.cache import ArraysCache, KVCache, RotatingKVCache
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+from exo.worker.engines.mlx.cache import CacheSnapshot, snapshot_ssm_states
+from exo.worker.engines.mlx.disaggregated.client import (
+    PrefillRequest,
+    ingest_into_mlx_cache,
+    remote_prefill_fetch,
+)
+from exo.worker.disaggregated.protocol import Header, KVChunk
+from exo.worker.engines.mlx.types import KVCacheType, Model
+from exo.worker.runner.bootstrap import logger
+
+
+def remote_prefill(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    sampler: Callable[[mx.array], mx.array],
+    prompt_tokens: mx.array,
+    cache: KVCacheType,
+    group: mx.distributed.Group | None,
+    on_prefill_progress: Callable[[int, int], None] | None,
+    distributed_prompt_progress_callback: Callable[[], None] | None,
+    *,
+    endpoint: str,
+    request_id: str,
+    model_id: str,
+    start_pos: int = 0,
+) -> tuple[float, int, list[CacheSnapshot]]:
+    del model, tokenizer, sampler, group, distributed_prompt_progress_callback
+
+    t0 = time.perf_counter()
+    total_prompt_tokens = int(prompt_tokens.shape[0])
+    num_layers_box: list[int] = [0]
+
+    def _on_header(header: Header) -> None:
+        num_layers_box[0] = header.num_layers
+
+    def _on_chunk(_chunk: KVChunk, chunks_received: int) -> None:
+        if on_prefill_progress is None:
+            return
+        num_layers = num_layers_box[0]
+        if num_layers > 0 and chunks_received % num_layers == 0:
+            tokens_so_far = chunks_received // num_layers
+            on_prefill_progress(
+                min(tokens_so_far, total_prompt_tokens),
+                total_prompt_tokens,
+            )
+
+    request = PrefillRequest(
+        model_id=model_id,
+        token_ids=cast(list[int], prompt_tokens.tolist()),
+        start_pos=start_pos,
+        request_id=request_id,
+    )
+    result = remote_prefill_fetch(
+        endpoint, request, on_header=_on_header, on_kv_chunk=_on_chunk
+    )
+    t_received = time.perf_counter()
+
+    caches = cast(list[KVCache | RotatingKVCache | ArraysCache], list(cache))
+    final_offset = ingest_into_mlx_cache(result, caches, start_pos=start_pos)
+    t_done = time.perf_counter()
+
+    num_tokens = final_offset - start_pos
+    tps = num_tokens / max(t_done - t0, 0.001)
+
+    logger.info(
+        f"Remote prefill: {num_tokens} tokens (start_pos={start_pos}, "
+        f"final_offset={final_offset}) at {tps:.0f} tok/s, "
+        f"transfer={(t_received - t0) * 1000:.0f}ms, "
+        f"inject={(t_done - t_received) * 1000:.0f}ms"
+    )
+    return tps, num_tokens, [snapshot_ssm_states(cache)]

src/exo/worker/engines/mlx/tests/test_batch_generate.py

@@ -24,9 +24,9 @@ from transformers import AutoTokenizer
 
 # Import batch_generate to activate the right-padding BatchKVCache patch
 import exo.worker.engines.mlx.generator.batch_generate  # noqa: F401
-from exo.shared.types.mlx import Model
 from exo.worker.engines.mlx.cache import encode_prompt, make_kv_cache
 from exo.worker.engines.mlx.generator.generate import prefill
+from exo.worker.engines.mlx.types import Model
 
 NUM_STEPS = 20
 

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

@@ -1,5 +1,3 @@
-"""Shared types for MLX-related functionality."""
-
 from collections.abc import Sequence
 
 from mlx import core as mx
@@ -12,14 +10,11 @@ from mlx_lm.models.cache import (
     RotatingKVCache,
 )
 
-# This list contains one cache entry per transformer layer
 KVCacheType = Sequence[
     KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList
 ]
 
 
-# Model is a wrapper function to fix the fact that mlx is not strongly typed in the same way that EXO is.
-# For example - MLX has no guarantee of the interface that nn.Module will expose. But we need a guarantee that it has a __call__() function
 class Model(nn.Module):
     layers: list[nn.Module]
 

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

@@ -44,7 +44,6 @@ from pydantic import RootModel
 from exo.download.download_utils import build_model_path
 from exo.shared.types.common import Host
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import Model
 from exo.shared.types.tasks import TaskId, TextGeneration
 from exo.shared.types.text_generation import ChatTemplateValue, TextGenerationTaskParams
 from exo.shared.types.worker.instances import (
@@ -65,6 +64,7 @@ from exo.worker.engines.mlx.auto_parallel import (
     pipeline_auto_parallel,
     tensor_auto_parallel,
 )
+from exo.worker.engines.mlx.types import Model
 from exo.worker.runner.bootstrap import logger
 
 
@@ -610,7 +610,7 @@ def apply_chat_template(
             messages.append({"role": msg.role, "content": msg.content})
 
     prompt = render_chat_template(tokenizer, messages, task_params)
-    logger.info(prompt)
+    logger.debug(prompt)
 
     return prompt
 

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

@@ -25,9 +25,9 @@ from transformers import AutoImageProcessor
 from exo.download.download_utils import build_model_path
 from exo.shared.models.model_cards import VisionCardConfig
 from exo.shared.types.common import ModelId
-from exo.shared.types.mlx import Model
 from exo.shared.types.text_generation import Base64Image, TextGenerationTaskParams
 from exo.worker.engines.mlx.cache import encode_prompt
+from exo.worker.engines.mlx.types import Model
 from exo.worker.engines.mlx.utils_mlx import (
     fix_unmatched_think_end_tokens,
     render_chat_template,

src/exo/worker/plan.py

@@ -68,7 +68,13 @@ def plan(
         or _init_distributed_backend(runners, all_runners)
         or _load_model(runners, all_runners, global_download_status)
         or _ready_to_warmup(runners, all_runners)
-        or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer, image_cache)
+        or _pending_tasks(
+            runners,
+            tasks,
+            all_runners,
+            input_chunk_buffer,
+            image_cache,
+        )
     )
 
 

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

@@ -12,7 +12,6 @@ from exo.shared.constants import EXO_MAX_CONCURRENT_REQUESTS
 from exo.shared.types.chunks import ErrorChunk, GenerationChunk, PrefillProgressChunk
 from exo.shared.types.common import ModelId
 from exo.shared.types.events import ChunkGenerated, Event
-from exo.shared.types.mlx import Model
 from exo.shared.types.tasks import CANCEL_ALL_TASKS, TaskId, TextGeneration
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.runner_response import GenerationResponse
@@ -24,6 +23,7 @@ from exo.worker.engines.mlx.generator.generate import (
     mlx_generate,
     warmup_inference,
 )
+from exo.worker.engines.mlx.types import Model
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
     mx_all_gather_tasks,
@@ -62,6 +62,11 @@ class GeneratorQueue[T]:
 class InferenceGenerator(ABC):
     _cancelled_tasks: set[TaskId]
 
+    model: Model
+    tokenizer: TokenizerWrapper
+    group: mx.distributed.Group | None
+    kv_prefix_cache: KVPrefixCache | None
+
     def should_cancel(self, task_id: TaskId) -> bool:
         return (
             task_id in self._cancelled_tasks

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

@@ -21,8 +21,8 @@ from exo.shared.types.chunks import (
     ToolCallChunk,
 )
 from exo.shared.types.common import ModelId
-from exo.shared.types.mlx import Model
 from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
+from exo.worker.engines.mlx.types import Model
 from exo.worker.engines.mlx.utils_mlx import (
     detect_thinking_prompt_suffix,
 )

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

@@ -1,11 +1,14 @@
 import os
+import queue
+import threading
 import time
 from collections.abc import Generator
 from dataclasses import dataclass
 from enum import Enum
 
 import mlx.core as mx
-from anyio import WouldBlock
+from anyio import ClosedResourceError, EndOfStream
+from mlx_lm.sample_utils import make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.models.model_cards import ModelTask
@@ -18,7 +21,6 @@ from exo.shared.types.events import (
     TaskAcknowledged,
     TaskStatusUpdated,
 )
-from exo.shared.types.mlx import Model
 from exo.shared.types.tasks import (
     ConnectToGroup,
     LoadModel,
@@ -47,8 +49,25 @@ from exo.shared.types.worker.runners import (
     RunnerWarmingUp,
 )
 from exo.utils.channels import MpReceiver, MpSender
-from exo.worker.engines.mlx.cache import KVPrefixCache
+from exo.utils.ports import random_ephemeral_port
+from exo.worker.engines.mlx.cache import (
+    KVPrefixCache,
+    cache_length,
+    make_kv_cache,
+    snapshot_ssm_states,
+)
+from exo.worker.engines.mlx.disaggregated.adapter import (
+    serialize_mlx_cache_to_payload,
+    wire_dtype_from_cache,
+)
+from exo.worker.disaggregated.server import (
+    PrefillJob,
+    PrefillServer,
+)
+from exo.worker.engines.mlx.generator.generate import prefill as mlx_prefill
+from exo.worker.engines.mlx.types import Model
 from exo.worker.engines.mlx.utils_mlx import (
+    fix_unmatched_think_end_tokens,
     initialize_mlx,
     load_mlx_items,
 )
@@ -63,6 +82,21 @@ from exo.worker.runner.llm_inference.batch_generator import (
 from .batch_generator import Cancelled, Finished
 from .tool_parsers import make_mlx_parser
 
+PREFILL_PICKUP_TIMEOUT_SECONDS = 3
+PREFILL_FINISH_TIMEOUT_SECONDS = 300
+
+
+@dataclass
+class _PrefillRequest:
+    job: PrefillJob
+    started: threading.Event
+    done: threading.Event
+    holder: list[bytes | None]
+
+
+_TaskStreamClosed = object()
+WorkItem = Task | _PrefillRequest | object
+
 
 class ExitCode(str, Enum):
     AllTasksComplete = "AllTasksComplete"
@@ -110,9 +144,160 @@ class Runner:
             TextGeneration,
         ] = {}
 
+        self._prefill_server: PrefillServer | None = None
+        self._prefill_server_port: int | None = None
+        self._work_queue: queue.Queue[WorkItem] = queue.Queue()
+        self._task_reader_thread: threading.Thread | None = None
+
         logger.info("runner created")
         self.update_status(RunnerIdle())
 
+    def _start_prefill_server(self) -> int | None:
+        if self._prefill_server_port is not None:
+            return self._prefill_server_port
+
+        def resolve(job: PrefillJob) -> bytes | None:
+            req = _PrefillRequest(
+                job=job,
+                started=threading.Event(),
+                done=threading.Event(),
+                holder=[None],
+            )
+            self._work_queue.put(req)
+            if not req.started.wait(timeout=PREFILL_PICKUP_TIMEOUT_SECONDS):
+                logger.warning(
+                    f"Prefill request {job.request_id} not picked up within "
+                    f"{PREFILL_PICKUP_TIMEOUT_SECONDS}s — runner busy"
+                )
+                return None
+            if not req.done.wait(timeout=PREFILL_FINISH_TIMEOUT_SECONDS):
+                logger.warning(
+                    f"Prefill request {job.request_id} did not finish within "
+                    f"{PREFILL_FINISH_TIMEOUT_SECONDS}s"
+                )
+                return None
+            return req.holder[0]
+
+        self._prefill_server = PrefillServer(
+            resolve=resolve, host="0.0.0.0", port=random_ephemeral_port()
+        )
+        self._prefill_server.start()
+        self._prefill_server_port = self._prefill_server.bound_port
+        return self._prefill_server_port
+
+    def _start_task_reader(self) -> None:
+        if self._task_reader_thread is not None:
+            return
+
+        def loop() -> None:
+            try:
+                with self.task_receiver:
+                    for task in self.task_receiver:
+                        self._work_queue.put(task)
+            except (EndOfStream, ClosedResourceError):
+                pass
+            finally:
+                self._work_queue.put(_TaskStreamClosed)
+
+        self._task_reader_thread = threading.Thread(
+            target=loop, name="task-reader", daemon=True
+        )
+        self._task_reader_thread.start()
+
+    def _serve_prefill(self, req: _PrefillRequest) -> None:
+        req.started.set()
+        was_ready = isinstance(self.current_status, RunnerReady)
+        if was_ready:
+            self.update_status(RunnerRunning())
+        try:
+            req.holder[0] = self._serve_prefill_request(req.job)
+        except Exception:
+            logger.opt(exception=True).warning(
+                f"Failed to serve prefill request {req.job.request_id}"
+            )
+            req.holder[0] = None
+        finally:
+            req.done.set()
+            if was_ready:
+                self.update_status(
+                    RunnerReady(prefill_server_port=self._prefill_server_port)
+                )
+
+    def _serve_prefill_request(self, job: PrefillJob) -> bytes:
+        assert isinstance(self.generator, InferenceGenerator)
+        model = self.generator.model
+        tokenizer = self.generator.tokenizer
+        group = self.generator.group
+        kv_prefix_cache = self.generator.kv_prefix_cache
+
+        prompt_tokens = mx.array(job.token_ids)
+        prompt_tokens = fix_unmatched_think_end_tokens(prompt_tokens, tokenizer)
+        n_tokens = int(prompt_tokens.shape[0])
+        t0 = time.perf_counter()
+
+        matched_index: int | None = None
+        prefix_hit_length = 0
+        if kv_prefix_cache is not None:
+            cache, remaining, matched_index, _ = kv_prefix_cache.get_kv_cache(
+                model, prompt_tokens
+            )
+            prefix_hit_length = n_tokens - int(remaining.shape[0])
+        else:
+            cache = make_kv_cache(model)
+            remaining = prompt_tokens
+
+        remaining_n = int(remaining.shape[0])
+        prefill_input = remaining[:-2] if remaining_n > 2 else remaining
+        if int(prefill_input.shape[0]) > 0:
+            sampler = make_sampler(temp=1.0)
+            _ = mlx_prefill(
+                model=model,
+                tokenizer=tokenizer,
+                sampler=sampler,
+                prompt_tokens=prefill_input,
+                cache=cache,
+                group=group,
+                on_prefill_progress=None,
+                distributed_prompt_progress_callback=None,
+            )
+
+        if kv_prefix_cache is not None:
+            try:
+                cache_snapshots = [snapshot_ssm_states(cache)]
+                hit_ratio = prefix_hit_length / n_tokens if n_tokens > 0 else 0.0
+                if matched_index is not None and hit_ratio >= 0.5:
+                    kv_prefix_cache.update_kv_cache(
+                        matched_index,
+                        prompt_tokens,
+                        cache,
+                        cache_snapshots,
+                        restore_pos=prefix_hit_length,
+                    )
+                else:
+                    kv_prefix_cache.add_kv_cache(prompt_tokens, cache, cache_snapshots)
+            except Exception:
+                logger.opt(exception=True).warning(
+                    "Failed to save prefix cache on prefill server"
+                )
+
+        final_offset = cache_length(cache)
+        payload = serialize_mlx_cache_to_payload(
+            cache,
+            dtype=wire_dtype_from_cache(cache),
+            model_id=job.model_id,
+            request_id=job.request_id,
+            start_pos=job.start_pos,
+        )
+        elapsed = time.perf_counter() - t0
+        sent = max(0, final_offset - job.start_pos)
+        logger.info(
+            f"Served prefill: request_id={job.request_id} "
+            f"{n_tokens} tokens (prefix_hit={prefix_hit_length}, "
+            f"client_start_pos={job.start_pos}, sent={sent}) "
+            f"in {elapsed * 1000:.0f}ms"
+        )
+        return payload
+
     def update_status(self, status: RunnerStatus):
         self.current_status = status
         self.event_sender.send(
@@ -130,15 +315,22 @@ class Runner:
         self.event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
     def main(self):
-        with self.task_receiver:
-            for task in self.task_receiver:
-                if task.task_id in self.seen:
-                    logger.warning("repeat task - potential error")
-                    continue
-                self.seen.add(task.task_id)
-                self.handle_first_task(task)
-                if isinstance(self.current_status, RunnerShutdown):
-                    break
+        self._start_task_reader()
+        while True:
+            item = self._work_queue.get()
+            if item is _TaskStreamClosed:
+                break
+            if isinstance(item, _PrefillRequest):
+                self._serve_prefill(item)
+                continue
+            task: Task = item  # type: ignore[assignment]
+            if task.task_id in self.seen:
+                logger.warning("repeat task - potential error")
+                continue
+            self.seen.add(task.task_id)
+            self.handle_first_task(task)
+            if isinstance(self.current_status, RunnerShutdown):
+                break
 
     def handle_first_task(self, task: Task):
         self.send_task_status(task.task_id, TaskStatus.Running)
@@ -219,8 +411,9 @@ class Runner:
                     f"runner initialized in {time.time() - self.setup_start_time} seconds"
                 )
 
+                prefill_port = self._start_prefill_server()
                 self.send_task_status(task.task_id, TaskStatus.Complete)
-                self.update_status(RunnerReady())
+                self.update_status(RunnerReady(prefill_server_port=prefill_port))
                 logger.info("runner ready")
 
             case TextGeneration() if isinstance(self.current_status, RunnerReady):
@@ -285,29 +478,37 @@ class Runner:
                 self.active_tasks.pop(task_id, None)
 
             try:
-                task = self.task_receiver.receive_nowait()
+                item = self._work_queue.get_nowait()
+            except queue.Empty:
+                continue
+            if item is _TaskStreamClosed:
+                # Task stream closed mid-generation. Bail out.
+                return ExitCode.Shutdown
+            if isinstance(item, _PrefillRequest):
+                # Refuse — runner is mid-generation. Client picks up the 3s
+                # pickup timeout and 503s.
+                item.started.set()
+                item.holder[0] = None
+                item.done.set()
+                continue
+            task: Task = item  # type: ignore[assignment]
+            if task.task_id in self.seen:
+                logger.warning("repeat task - potential error")
+                continue
+            self.seen.add(task.task_id)
+            match task:
+                case TextGeneration():
+                    self.acknowledge_task(task)
+                    self.submit_text_generation(task)
+                case Shutdown():
+                    self.shutdown(task)
+                    return ExitCode.Shutdown
+                case _:
+                    raise ValueError(
+                        f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}"
+                    )
 
-                if task.task_id in self.seen:
-                    logger.warning("repeat task - potential error")
-                    continue
-                self.seen.add(task.task_id)
-
-                match task:
-                    case TextGeneration():
-                        self.acknowledge_task(task)
-                        self.submit_text_generation(task)
-                    case Shutdown():
-                        self.shutdown(task)
-                        return ExitCode.Shutdown
-                    case _:
-                        raise ValueError(
-                            f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}"
-                        )
-
-            except WouldBlock:
-                pass
-
-        self.update_status(RunnerReady())
+        self.update_status(RunnerReady(prefill_server_port=self._prefill_server_port))
         logger.info("runner ready")
 
         return ExitCode.AllTasksComplete

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

@@ -14,10 +14,10 @@ from exo.shared.constants import EXO_DEFAULT_MODELS_DIR
 from exo.shared.models.model_cards import ModelCard, ModelTask
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import Model
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata
 from exo.worker.engines.mlx.generator.generate import mlx_generate
+from exo.worker.engines.mlx.types import Model
 from exo.worker.engines.mlx.utils_mlx import apply_chat_template, shard_and_load
 
 

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

@@ -9,7 +9,6 @@ from mlx_lm.models.cache import KVCache
 from mlx_lm.sample_utils import make_sampler
 
 from exo.shared.types.common import ModelId
-from exo.shared.types.mlx import Model
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.worker.engines.mlx.cache import (
     KVPrefixCache,
@@ -19,6 +18,7 @@ from exo.worker.engines.mlx.cache import (
     make_kv_cache,
 )
 from exo.worker.engines.mlx.generator.generate import mlx_generate, prefill
+from exo.worker.engines.mlx.types import Model
 from exo.worker.engines.mlx.utils_mlx import apply_chat_template
 from exo.worker.tests.unittests.test_mlx.conftest import (
     DEFAULT_GPT_OSS_CONFIG,

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

@@ -17,7 +17,6 @@ from mlx_lm.tokenizer_utils import TokenizerWrapper
 from exo.download.download_utils import resolve_existing_model
 from exo.shared.constants import EXO_MODELS_DIRS, EXO_MODELS_READ_ONLY_DIRS
 from exo.shared.types.common import ModelId
-from exo.shared.types.mlx import Model
 from exo.shared.types.text_generation import (
     InputMessage,
     InputMessageContent,
@@ -25,6 +24,7 @@ from exo.shared.types.text_generation import (
 )
 from exo.worker.engines.mlx.cache import KVPrefixCache
 from exo.worker.engines.mlx.generator.generate import mlx_generate
+from exo.worker.engines.mlx.types import Model
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
     load_tokenizer_for_model_id,

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

@@ -113,6 +113,16 @@ CHAT_TASK = TextGeneration(
 def assert_events_equal(test_events: Iterable[Event], true_events: Iterable[Event]):
     for test_event, true_event in zip(test_events, true_events, strict=True):
         test_event = test_event.model_copy(update={"event_id": true_event.event_id})
+        if isinstance(test_event, RunnerStatusUpdated) and isinstance(
+            test_event.runner_status, (RunnerReady, RunnerRunning)
+        ):
+            test_event = test_event.model_copy(
+                update={
+                    "runner_status": test_event.runner_status.model_copy(
+                        update={"prefill_server_port": None}
+                    )
+                }
+            )
         assert test_event == true_event, f"{test_event} != {true_event}"
 
 
@@ -152,6 +162,12 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
     )
     monkeypatch.setattr(mlx_batch_generator, "ExoBatchGenerator", FakeExoBatchGenerator)
 
+    # Don't bind a real TCP port for the prefill server in event-ordering tests.
+    def _no_prefill_server(_self: mlx_runner.Runner) -> int | None:
+        return None
+
+    monkeypatch.setattr(mlx_runner.Runner, "_start_prefill_server", _no_prefill_server)
+
 
 class FakeExoBatchGenerator:
     def __init__(self, *_args: object, **_kwargs: object) -> None: