[object Object]

← back to Exo

MLX P/D (#1993)

f0d1371d89a7f899e96977014cce7307a53682f2 · 2026-04-28 01:12:42 +0100 · rltakashige

## Motivation

MLX only prefill server for Apple Silicon

Files touched

Diff

commit f0d1371d89a7f899e96977014cce7307a53682f2
Author: rltakashige <rl.takashige@gmail.com>
Date:   Tue Apr 28 01:12:42 2026 +0100

    MLX P/D (#1993)
    
    ## Motivation
    
    MLX only prefill server for Apple Silicon
---
 bench/harness.py                                   |  19 +-
 bench/prefill-decode.toml                          |  36 +
 bench/prefill_decode_bench.py                      | 785 +++++++++++++++++++++
 dashboard/src/lib/components/HeaderNav.svelte      |  26 +
 .../components/PrefillDecodeDisaggregation.svelte  | 565 +++++++++++++++
 dashboard/src/lib/stores/app.svelte.ts             |  92 +++
 dashboard/src/lib/utils/model_family.ts            |  44 ++
 dashboard/src/routes/advanced/+page.svelte         |  81 +++
 src/exo/api/main.py                                |  65 ++
 src/exo/api/types/__init__.py                      |   2 +
 src/exo/api/types/api.py                           |  10 +
 src/exo/master/main.py                             |  81 ++-
 src/exo/master/placement.py                        |   7 +-
 src/exo/master/placement_utils.py                  |   6 +-
 src/exo/shared/apply.py                            |  67 +-
 src/exo/shared/constants.py                        |   2 +
 .../tests/test_apply/test_apply_instance_link.py   |  72 ++
 src/exo/shared/types/commands.py                   |  13 +
 src/exo/shared/types/events.py                     |  11 +
 src/exo/shared/types/instance_link.py              |  13 +
 src/exo/shared/types/state.py                      |   4 +
 src/exo/shared/types/text_generation.py            |   2 +
 src/exo/shared/types/worker/runners.py             |   2 +-
 src/exo/utils/ports.py                             |   6 +
 src/exo/worker/disaggregated/__init__.py           |   0
 src/exo/worker/disaggregated/protocol.py           | 152 ++++
 src/exo/worker/disaggregated/server.py             | 105 +++
 src/exo/worker/engines/mlx/auto_parallel.py        |   4 +-
 src/exo/worker/engines/mlx/cache.py                |   2 +-
 .../worker/engines/mlx/disaggregated/__init__.py   |   0
 .../worker/engines/mlx/disaggregated/adapter.py    | 233 ++++++
 src/exo/worker/engines/mlx/disaggregated/client.py | 147 ++++
 src/exo/worker/engines/mlx/disaggregated/serve.py  |  86 +++
 .../engines/mlx/disaggregated/tests/__init__.py    |   0
 .../mlx/disaggregated/tests/test_end_to_end.py     | 167 +++++
 .../mlx/disaggregated/tests/test_mlx_adapter.py    | 270 +++++++
 .../disaggregated/tests/test_protocol_roundtrip.py | 154 ++++
 .../mlx/disaggregated/tests/test_server_drain.py   | 120 ++++
 .../worker/engines/mlx/generator/batch_generate.py |  53 +-
 src/exo/worker/engines/mlx/generator/generate.py   |  51 +-
 .../worker/engines/mlx/generator/remote_prefill.py |  72 ++
 .../engines/mlx/tests/test_batch_generate.py       |   2 +-
 .../types/mlx.py => worker/engines/mlx/types.py}   |   0
 src/exo/worker/engines/mlx/utils_mlx.py            |   4 +-
 src/exo/worker/engines/mlx/vision.py               |   2 +-
 .../worker/runner/llm_inference/batch_generator.py |  41 +-
 .../runner/llm_inference/model_output_parsers.py   |   2 +-
 src/exo/worker/runner/llm_inference/runner.py      | 168 ++++-
 .../worker/tests/unittests/test_mlx/conftest.py    |   2 +-
 .../unittests/test_mlx/test_kv_prefix_cache.py     |   2 +-
 .../test_mlx/test_prefix_cache_architectures.py    |   2 +-
 .../unittests/test_runner/test_event_ordering.py   |   5 +
 .../unittests/test_runner/test_serve_prefill.py    | 227 ++++++
 53 files changed, 4000 insertions(+), 84 deletions(-)

diff --git a/bench/harness.py b/bench/harness.py
index 3285becf..ba6d0a77 100644
--- a/bench/harness.py
+++ b/bench/harness.py
@@ -293,11 +293,15 @@ def sharding_filter(sharding: str, wanted: str) -> bool:
 
 
 def fetch_and_filter_placements(
-    client: ExoClient, full_model_id: str, args: argparse.Namespace
+    client: ExoClient,
+    full_model_id: str,
+    args: argparse.Namespace,
+    node_id: str | None = None,
 ) -> list[dict[str, Any]]:
-    previews_resp = client.request_json(
-        "GET", "/instance/previews", params={"model_id": full_model_id}
-    )
+    params: dict[str, str] = {"model_id": full_model_id}
+    if node_id is not None:
+        params["node_ids"] = node_id
+    previews_resp = client.request_json("GET", "/instance/previews", params=params)
     previews = previews_resp.get("previews") or []
 
     selected: list[dict[str, Any]] = []
@@ -357,8 +361,9 @@ def settle_and_fetch_placements(
     full_model_id: str,
     args: argparse.Namespace,
     settle_timeout: float = 0,
+    node_id: str | None = None,
 ) -> list[dict[str, Any]]:
-    selected = fetch_and_filter_placements(client, full_model_id, args)
+    selected = fetch_and_filter_placements(client, full_model_id, args, node_id=node_id)
 
     if not selected and settle_timeout > 0:
         backoff = _SETTLE_INITIAL_BACKOFF_S
@@ -371,7 +376,9 @@ def settle_and_fetch_placements(
             )
             time.sleep(min(backoff, remaining))
             backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
-            selected = fetch_and_filter_placements(client, full_model_id, args)
+            selected = fetch_and_filter_placements(
+                client, full_model_id, args, node_id=node_id
+            )
 
     return selected
 
diff --git a/bench/prefill-decode.toml b/bench/prefill-decode.toml
new file mode 100644
index 00000000..a26d0200
--- /dev/null
+++ b/bench/prefill-decode.toml
@@ -0,0 +1,36 @@
+# Prefill/Decode disaggregation benchmark config.
+#
+# Top-level keys are bench-wide. [prefill] and [decode] sections set per-side
+# placement filters and (optionally) per-side model.
+#
+# Example:
+#   uv run python bench/prefill_decode_bench.py --config bench/prefill-decode.toml
+
+host = "james"
+port = 52415
+timeout = 7200.0
+settle_timeout = 60.0
+
+# Workload
+pp = [4096]
+tg = [512]
+repeat = 1
+warmup = 0
+
+json_out = "bench/prefill_decode_results.json"
+
+[prefill]
+model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
+node = "mike"
+instance_meta = "ring"
+sharding = "pipeline"
+min_nodes = 1
+max_nodes = 1
+
+[decode]
+model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
+node = "james"
+instance_meta = "ring"
+sharding = "pipeline"
+min_nodes = 1
+max_nodes = 1
diff --git a/bench/prefill_decode_bench.py b/bench/prefill_decode_bench.py
new file mode 100644
index 00000000..375b66e8
--- /dev/null
+++ b/bench/prefill_decode_bench.py
@@ -0,0 +1,785 @@
+# type: ignore
+#!/usr/bin/env python3
+"""Disaggregated prefill-decode benchmark for exo (MLX → MLX).
+
+Spins up two MLX instances on the cluster, marks one as Prefill source and
+the other as Decode target via /v1/instance-links, then sends chat
+completions to the API. The master routes the request to the decode
+instance and stamps `prefill_endpoint` pointing at the prefill instance —
+the worker decides per-request whether to ship prefill remotely
+(uncached_count > REMOTE_PREFILL_MIN_TOKENS).
+
+Usage:
+    uv run python bench/prefill_decode_bench.py --model <id> --pp 2048,8192 --tg 128
+    uv run python bench/prefill_decode_bench.py --model <id> --pp 4096 --tg 128 --repeat 3
+    uv run python bench/prefill_decode_bench.py --model <id> --pp 2048 --tg 128 --dry-run
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import copy
+import itertools
+import json
+import sys
+import time
+import tomllib
+from pathlib import Path
+from statistics import mean
+from typing import Any
+
+from exo_bench import (
+    PromptSizer,
+    format_peak_memory,
+    load_tokenizer_for_bench,
+    parse_int_list,
+)
+from harness import (
+    ExoClient,
+    ExoHttpError,
+    add_common_instance_args,
+    instance_id_from_instance,
+    node_ids_from_instance,
+    nodes_used_in_instance,
+    resolve_model_short_id,
+    run_planning_phase,
+    settle_and_fetch_placements,
+    unwrap_instance,
+    wait_for_instance_gone,
+    wait_for_instance_ready,
+)
+from loguru import logger
+
+
+def _node_id_to_friendly(client: ExoClient) -> dict[str, str]:
+    identities = client.get_node_identities() or {}
+    out: dict[str, str] = {}
+    for node_id, identity in identities.items():
+        if isinstance(identity, dict):
+            name = identity.get("friendlyName") or identity.get("friendly_name")
+            if isinstance(name, str):
+                out[str(node_id)] = name
+    return out
+
+
+def _placement_node_friendly_names(
+    placement: dict[str, Any], id_to_friendly: dict[str, str]
+) -> list[str]:
+    instance = placement["instance"]
+    return [id_to_friendly.get(nid, nid) for nid in node_ids_from_instance(instance)]
+
+
+def _filter_by_node(
+    placements: list[dict[str, Any]],
+    friendly_name: str,
+    id_to_friendly: dict[str, str],
+) -> list[dict[str, Any]]:
+    target = friendly_name.lower()
+    matched: list[dict[str, Any]] = []
+    for p in placements:
+        names = [n.lower() for n in _placement_node_friendly_names(p, id_to_friendly)]
+        if any(target == n or target in n for n in names):
+            matched.append(p)
+    return matched
+
+
+def _node_id_by_friendly(id_to_friendly: dict[str, str], target: str) -> str | None:
+    target_lc = target.lower()
+    for nid, name in id_to_friendly.items():
+        if target_lc == name.lower() or target_lc in name.lower():
+            return nid
+    return None
+
+
+def _load_toml(path: str) -> dict[str, Any]:
+    with Path(path).open("rb") as f:
+        return tomllib.load(f)
+
+
+_TOP_LEVEL_TOML_KEYS = {
+    "host",
+    "port",
+    "timeout",
+    "settle_timeout",
+    "model",
+    "pp",
+    "tg",
+    "repeat",
+    "warmup",
+    "json_out",
+    "instance_meta",
+    "sharding",
+    "min_nodes",
+    "max_nodes",
+    "force_download",
+    "danger_delete_downloads",
+    "all_combinations",
+}
+
+
+def _inject_toml_into_argv() -> None:
+    """If --config X is in sys.argv, pre-load it and inject required CLI args
+    (--model, --pp, --tg) so argparse's required=True checks pass."""
+    argv = sys.argv
+    if "--config" not in argv:
+        return
+    idx = argv.index("--config")
+    if idx + 1 >= len(argv):
+        return
+    cfg_path = argv[idx + 1]
+    cfg = _load_toml(cfg_path)
+    decode = cfg.get("decode", {})
+
+    def _has(flag: str) -> bool:
+        return any(a == flag or a.startswith(flag + "=") for a in argv)
+
+    # --model: prefer top-level, then [decode].model
+    if not _has("--model"):
+        model = cfg.get("model") or decode.get("model")
+        if model:
+            argv += ["--model", str(model)]
+    if not _has("--pp"):
+        pp = cfg.get("pp")
+        if pp:
+            argv += (
+                ["--pp", *(str(x) for x in pp)]
+                if isinstance(pp, list)
+                else [
+                    "--pp",
+                    str(pp),
+                ]
+            )
+    if not _has("--tg"):
+        tg = cfg.get("tg")
+        if tg:
+            argv += (
+                ["--tg", *(str(x) for x in tg)]
+                if isinstance(tg, list)
+                else [
+                    "--tg",
+                    str(tg),
+                ]
+            )
+
+
+def _merge_toml_into_args(args: argparse.Namespace, cfg: dict[str, Any]) -> None:
+    """Apply top-level toml keys onto args namespace where args has a default."""
+    for key, value in cfg.items():
+        if key in {"prefill", "decode"}:
+            continue
+        if key not in _TOP_LEVEL_TOML_KEYS:
+            continue
+        attr = key
+        current = getattr(args, attr, None)
+        if current in (None, [], False):
+            setattr(args, attr, value)
+
+
+def _side_args(
+    base: argparse.Namespace, overrides: dict[str, Any]
+) -> argparse.Namespace:
+    out = copy.copy(base)
+    for k in (
+        "instance_meta",
+        "sharding",
+        "min_nodes",
+        "max_nodes",
+        "skip_pipeline_jaccl",
+        "skip_tensor_ring",
+    ):
+        if k in overrides:
+            setattr(out, k, overrides[k])
+    return out
+
+
+def _pick_two_distinct_placements(
+    placements: list[dict[str, Any]],
+) -> tuple[dict[str, Any], dict[str, Any]] | None:
+    if len(placements) < 2:
+        return None
+    seen_nodes: set[tuple[str, ...]] = set()
+    chosen: list[dict[str, Any]] = []
+    for p in placements:
+        nodes = tuple(sorted(str(n) for n in p.get("nodes", [])))
+        if nodes in seen_nodes:
+            continue
+        seen_nodes.add(nodes)
+        chosen.append(p)
+        if len(chosen) == 2:
+            return chosen[0], chosen[1]
+    return None
+
+
+def _create_instance_link(
+    client: ExoClient,
+    prefill_instance_id: str,
+    decode_instance_id: str,
+) -> str:
+    out = client.request_json(
+        "POST",
+        "/v1/instance-links",
+        body={
+            "prefill_instances": [prefill_instance_id],
+            "decode_instances": [decode_instance_id],
+        },
+    )
+    return str(out.get("commandId", ""))
+
+
+def _list_instance_links(client: ExoClient) -> list[dict[str, Any]]:
+    out = client.request_json("GET", "/v1/instance-links")
+    return out if isinstance(out, list) else []
+
+
+def _delete_instance_link(client: ExoClient, link_id: str) -> None:
+    client.request_json("DELETE", f"/v1/instance-links/{link_id}")
+
+
+def run_one(
+    client: ExoClient,
+    model_id: str,
+    pp_hint: int,
+    tg: int,
+    prompt_sizer: PromptSizer,
+) -> tuple[dict[str, Any], int]:
+    content, pp_tokens = prompt_sizer.build(pp_hint)
+    payload: dict[str, Any] = {
+        "model": model_id,
+        "messages": [{"role": "user", "content": content}],
+        "stream": False,
+        "max_tokens": tg,
+    }
+
+    t0 = time.perf_counter()
+    out = client.post_bench_chat_completions(payload)
+    elapsed = time.perf_counter() - t0
+
+    stats = out.get("generation_stats")
+    choices = out.get("choices") or [{}]
+    message = choices[0].get("message", {}) if choices else {}
+    text = message.get("content") or ""
+    preview = text[:200] if text else ""
+
+    return {
+        "elapsed_s": elapsed,
+        "output_text_preview": preview,
+        "stats": stats,
+    }, pp_tokens
+
+
+def _run_phase(
+    *,
+    client: ExoClient,
+    label: str,
+    pp_tg_pairs: list[tuple[int, int]],
+    model_id: str,
+    prompt_sizer: PromptSizer,
+    warmup: int,
+    repeat: int,
+    common_meta: dict[str, Any],
+) -> list[dict[str, Any]]:
+    logger.info(f"=== phase: {label} (model={model_id}) ===")
+    rows: list[dict[str, Any]] = []
+    for i in range(warmup):
+        run_one(client, model_id, pp_tg_pairs[0][0], pp_tg_pairs[0][1], prompt_sizer)
+        logger.debug(f"  warmup {i + 1}/{warmup} done")
+
+    for pp, tg in pp_tg_pairs:
+        logger.info(f"--- {label}: pp={pp} tg={tg} ---")
+        runs: list[dict[str, Any]] = []
+        for r in range(repeat):
+            time.sleep(2)
+            try:
+                row, actual_pp_tokens = run_one(client, model_id, pp, tg, prompt_sizer)
+            except Exception as e:
+                logger.error(e)
+                continue
+            row.update(common_meta)
+            row.update(
+                {
+                    "phase": label,
+                    "phase_model_id": model_id,
+                    "pp_tokens": actual_pp_tokens,
+                    "tg": tg,
+                    "repeat_index": r,
+                }
+            )
+            runs.append(row)
+            rows.append(row)
+
+        if runs:
+            prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
+            gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
+            ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
+            gtok = mean(x["stats"]["generation_tokens"] for x in runs)
+            peak = mean(x["stats"]["peak_memory_usage"]["inBytes"] for x in runs)
+            avg_elapsed = mean(x["elapsed_s"] for x in runs)
+            logger.info(
+                f"[{label}] prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
+                f"prompt_tokens={ptok} gen_tokens={gtok}    "
+                f"peak_memory={format_peak_memory(peak)}    "
+                f"avg_elapsed={avg_elapsed:.2f}s"
+            )
+        time.sleep(2)
+    return rows
+
+
+def _summarise(rows: list[dict[str, Any]]) -> dict[tuple[int, int], dict[str, float]]:
+    grouped: dict[tuple[int, int], list[dict[str, Any]]] = {}
+    for r in rows:
+        key = (int(r["pp_tokens"]), int(r["tg"]))
+        grouped.setdefault(key, []).append(r)
+    out: dict[tuple[int, int], dict[str, float]] = {}
+    for key, runs in grouped.items():
+        out[key] = {
+            "prompt_tps": mean(x["stats"]["prompt_tps"] for x in runs),
+            "gen_tps": mean(x["stats"]["generation_tps"] for x in runs),
+            "elapsed_s": mean(x["elapsed_s"] for x in runs),
+        }
+    return out
+
+
+def _print_diff(
+    disagg_rows: list[dict[str, Any]],
+    decode_alone_rows: list[dict[str, Any]],
+    prefill_alone_rows: list[dict[str, Any]],
+) -> None:
+    disagg = _summarise(disagg_rows)
+    decode_alone = _summarise(decode_alone_rows)
+    prefill_alone = _summarise(prefill_alone_rows)
+    keys = set(disagg.keys()) | set(decode_alone.keys()) | set(prefill_alone.keys())
+
+    width = 64
+    for key in sorted(keys):
+        pp, tg = key
+        logger.info("─" * width)
+        logger.info(f"  pp={pp}  tg={tg}")
+        logger.info("─" * width)
+        logger.info(
+            f"  {'phase':<16} {'elapsed':>10}  {'prompt_tps':>11}  {'gen_tps':>9}"
+        )
+        for label, summary in (
+            ("disaggregated", disagg.get(key)),
+            ("decode_alone", decode_alone.get(key)),
+            ("prefill_alone", prefill_alone.get(key)),
+        ):
+            if summary is None:
+                logger.info(f"  {label:<16} {'—':>10}  {'—':>11}  {'—':>9}")
+                continue
+            logger.info(
+                f"  {label:<16} "
+                f"{summary['elapsed_s']:>9.2f}s  "
+                f"{summary['prompt_tps']:>11.1f}  "
+                f"{summary['gen_tps']:>9.2f}"
+            )
+
+        d = disagg.get(key)
+        da = decode_alone.get(key)
+        pa = prefill_alone.get(key)
+        if d and da and d["elapsed_s"] > 0:
+            logger.info(
+                f"  speedup vs decode_alone:  {da['elapsed_s'] / d['elapsed_s']:.2f}x"
+            )
+        if d and pa and d["elapsed_s"] > 0:
+            logger.info(
+                f"  speedup vs prefill_alone: {pa['elapsed_s'] / d['elapsed_s']:.2f}x"
+            )
+    logger.info("─" * width)
+
+
+def main() -> int:
+    _inject_toml_into_argv()
+    ap = argparse.ArgumentParser(
+        prog="prefill-decode-bench",
+        description="Benchmark MLX-MLX disaggregated prefill/decode via instance links.",
+    )
+    add_common_instance_args(ap)
+    ap.add_argument(
+        "--pp",
+        nargs="+",
+        required=True,
+        help="Prompt-size hints (ints, must be >1000). Accepts commas.",
+    )
+    ap.add_argument(
+        "--tg",
+        nargs="+",
+        required=True,
+        help="Generation lengths (ints). Accepts commas.",
+    )
+    ap.add_argument(
+        "--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair."
+    )
+    ap.add_argument(
+        "--warmup",
+        type=int,
+        default=0,
+        help="Warmup runs (uses first pp/tg).",
+    )
+    ap.add_argument(
+        "--json-out",
+        default="bench/prefill_decode_results.json",
+        help="Write raw per-run results JSON to this path.",
+    )
+    ap.add_argument("--stdout", action="store_true", help="Write results to stdout")
+    ap.add_argument(
+        "--dry-run", action="store_true", help="List selected placements and exit."
+    )
+    ap.add_argument(
+        "--all-combinations",
+        action="store_true",
+        help="Force all pp×tg combinations even when lists have equal length.",
+    )
+    ap.add_argument(
+        "--prefill-model",
+        default=None,
+        help="Model id for the prefill instance. Defaults to --model.",
+    )
+    ap.add_argument(
+        "--prefill-node",
+        default=None,
+        help="friendly_name of the node hosting the prefill instance.",
+    )
+    ap.add_argument(
+        "--decode-node",
+        default=None,
+        help="friendly_name of the node hosting the decode instance.",
+    )
+    ap.add_argument(
+        "--config",
+        default=None,
+        help="TOML config file. CLI flags override toml values.",
+    )
+    ap.add_argument(
+        "--compare-baseline",
+        action="store_true",
+        help="Also run each (pp,tg) pair without the prefill/decode link "
+        "(decode instance does its own prefill) and report the diff.",
+    )
+    args = ap.parse_args()
+    cfg = _load_toml(args.config) if args.config else {}
+    _merge_toml_into_args(args, cfg)
+    prefill_overrides = cfg.get("prefill", {}) if cfg else {}
+    decode_overrides = cfg.get("decode", {}) if cfg else {}
+    if args.prefill_model is None and "model" in prefill_overrides:
+        args.prefill_model = prefill_overrides["model"]
+    if args.prefill_node is None and "node" in prefill_overrides:
+        args.prefill_node = prefill_overrides["node"]
+    if args.decode_node is None and "node" in decode_overrides:
+        args.decode_node = decode_overrides["node"]
+    if "model" in decode_overrides and not args.model:
+        args.model = decode_overrides["model"]
+
+    pp_list = parse_int_list(args.pp)
+    tg_list = parse_int_list(args.tg)
+    if not pp_list or not tg_list:
+        logger.error("pp and tg lists must be non-empty")
+        return 2
+    for pp in pp_list:
+        if pp <= 1000:
+            logger.error(
+                f"pp={pp} must be >1000 (remote prefill triggers when uncached >1000)"
+            )
+            return 2
+    if args.repeat <= 0:
+        logger.error("--repeat must be >= 1")
+        return 2
+
+    use_combinations = args.all_combinations or len(pp_list) != len(tg_list)
+    if use_combinations:
+        logger.info(
+            f"pp/tg mode: combinations (product) — {len(pp_list) * len(tg_list)} pairs"
+        )
+    else:
+        logger.info(f"pp/tg mode: tandem (zip) — {len(pp_list)} pairs")
+
+    client = ExoClient(args.host, args.port, timeout_s=args.timeout)
+
+    decode_short_id, decode_full_id = resolve_model_short_id(
+        client, args.model, force_download=args.force_download
+    )
+    if args.prefill_model:
+        prefill_short_id, prefill_full_id = resolve_model_short_id(
+            client, args.prefill_model, force_download=args.force_download
+        )
+    else:
+        prefill_short_id, prefill_full_id = decode_short_id, decode_full_id
+
+    tokenizer = load_tokenizer_for_bench(decode_full_id)
+    if tokenizer is None:
+        raise RuntimeError("[prefill-decode-bench] decode tokenizer load failed")
+    try:
+        decode_prompt_sizer = PromptSizer(tokenizer)
+    except Exception:
+        logger.error("[prefill-decode-bench] decode prompt sizing failed")
+        raise
+
+    if prefill_full_id == decode_full_id:
+        prefill_prompt_sizer = decode_prompt_sizer
+    else:
+        prefill_tokenizer = load_tokenizer_for_bench(prefill_full_id)
+        if prefill_tokenizer is None:
+            raise RuntimeError("[prefill-decode-bench] prefill tokenizer load failed")
+        prefill_prompt_sizer = PromptSizer(prefill_tokenizer)
+
+    id_to_friendly = _node_id_to_friendly(client)
+
+    prefill_args = _side_args(args, prefill_overrides)
+    decode_args = _side_args(args, decode_overrides)
+
+    if prefill_full_id == decode_full_id and prefill_overrides == decode_overrides:
+        placements = settle_and_fetch_placements(
+            client, decode_full_id, args, settle_timeout=args.settle_timeout
+        )
+        prefill_candidates = (
+            _filter_by_node(placements, args.prefill_node, id_to_friendly)
+            if args.prefill_node
+            else placements
+        )
+        decode_candidates = (
+            _filter_by_node(placements, args.decode_node, id_to_friendly)
+            if args.decode_node
+            else placements
+        )
+        if args.prefill_node and not prefill_candidates:
+            logger.error(f"No placement on prefill node {args.prefill_node!r}.")
+            return 1
+        if args.decode_node and not decode_candidates:
+            logger.error(f"No placement on decode node {args.decode_node!r}.")
+            return 1
+        if args.prefill_node and args.decode_node:
+            prefill_p = prefill_candidates[0]
+            decode_p = decode_candidates[0]
+        else:
+            pair = _pick_two_distinct_placements(placements)
+            if pair is None:
+                logger.error(
+                    "Need at least two distinct-node MLX placements for the same model."
+                )
+                return 1
+            prefill_p, decode_p = pair
+            if args.prefill_node:
+                prefill_p = prefill_candidates[0]
+            if args.decode_node:
+                decode_p = decode_candidates[0]
+    else:
+        prefill_node_id = (
+            _node_id_by_friendly(id_to_friendly, args.prefill_node)
+            if args.prefill_node
+            else None
+        )
+        decode_node_id = (
+            _node_id_by_friendly(id_to_friendly, args.decode_node)
+            if args.decode_node
+            else None
+        )
+        if args.prefill_node and prefill_node_id is None:
+            logger.error(f"Unknown node {args.prefill_node!r}.")
+            return 1
+        if args.decode_node and decode_node_id is None:
+            logger.error(f"Unknown node {args.decode_node!r}.")
+            return 1
+        prefill_placements = settle_and_fetch_placements(
+            client,
+            prefill_full_id,
+            prefill_args,
+            settle_timeout=args.settle_timeout,
+            node_id=prefill_node_id,
+        )
+        decode_placements = settle_and_fetch_placements(
+            client,
+            decode_full_id,
+            decode_args,
+            settle_timeout=args.settle_timeout,
+            node_id=decode_node_id,
+        )
+        if not prefill_placements:
+            logger.error(
+                f"No placement found for prefill model {prefill_full_id}"
+                f"{f' on node {args.prefill_node!r}' if args.prefill_node else ''}."
+            )
+            return 1
+        if not decode_placements:
+            logger.error(
+                f"No placement found for decode model {decode_full_id}"
+                f"{f' on node {args.decode_node!r}' if args.decode_node else ''}."
+            )
+            return 1
+        prefill_p = prefill_placements[0]
+        decode_p = decode_placements[0]
+
+    prefill_node_names = _placement_node_friendly_names(prefill_p, id_to_friendly)
+    decode_node_names = _placement_node_friendly_names(decode_p, id_to_friendly)
+    _ = unwrap_instance
+
+    prefill_instance = prefill_p["instance"]
+    decode_instance = decode_p["instance"]
+    prefill_id = instance_id_from_instance(prefill_instance)
+    decode_id = instance_id_from_instance(decode_instance)
+    prefill_meta = str(prefill_p.get("instance_meta", ""))
+    decode_meta = str(decode_p.get("instance_meta", ""))
+    prefill_nodes = nodes_used_in_instance(prefill_instance)
+    decode_nodes = nodes_used_in_instance(decode_instance)
+
+    logger.info("=" * 80)
+    logger.info(
+        f"PREFILL: {prefill_meta} / nodes={prefill_nodes} ({','.join(prefill_node_names)}) "
+        f"/ {prefill_short_id} ({prefill_full_id}) / instance_id={prefill_id}"
+    )
+    logger.info(
+        f"DECODE:  {decode_meta} / nodes={decode_nodes} ({','.join(decode_node_names)}) "
+        f"/ {decode_short_id} ({decode_full_id}) / instance_id={decode_id}"
+    )
+
+    if args.dry_run:
+        return 0
+
+    settle_deadline = (
+        time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
+    )
+
+    logger.info("Planning phase: prefill...")
+    run_planning_phase(
+        client,
+        prefill_full_id,
+        prefill_p,
+        args.danger_delete_downloads,
+        args.timeout,
+        settle_deadline,
+    )
+    logger.info("Planning phase: decode...")
+    run_planning_phase(
+        client,
+        decode_full_id,
+        decode_p,
+        args.danger_delete_downloads,
+        args.timeout,
+        settle_deadline,
+    )
+
+    if use_combinations:
+        pp_tg_pairs = list(itertools.product(pp_list, tg_list))
+    else:
+        pp_tg_pairs = list(zip(pp_list, tg_list, strict=True))
+
+    common_meta = {
+        "decode_model_short_id": decode_short_id,
+        "decode_model_id": decode_full_id,
+        "prefill_model_short_id": prefill_short_id,
+        "prefill_model_id": prefill_full_id,
+        "prefill_instance_id": prefill_id,
+        "prefill_instance_meta": prefill_meta,
+        "prefill_nodes": prefill_nodes,
+        "decode_instance_id": decode_id,
+        "decode_instance_meta": decode_meta,
+        "decode_nodes": decode_nodes,
+    }
+
+    all_rows: list[dict[str, Any]] = []
+    disagg_rows: list[dict[str, Any]] = []
+    decode_alone_rows: list[dict[str, Any]] = []
+    prefill_alone_rows: list[dict[str, Any]] = []
+    link_id = ""
+    prefill_alive = False
+    decode_alive = False
+    try:
+        logger.info("Creating prefill instance...")
+        client.request_json("POST", "/instance", body={"instance": prefill_instance})
+        wait_for_instance_ready(client, prefill_id)
+        prefill_alive = True
+        logger.info("Prefill instance ready")
+
+        if args.compare_baseline:
+            time.sleep(2)
+            prefill_alone_rows = _run_phase(
+                client=client,
+                label="prefill_alone",
+                pp_tg_pairs=pp_tg_pairs,
+                model_id=prefill_full_id,
+                prompt_sizer=prefill_prompt_sizer,
+                warmup=args.warmup,
+                repeat=args.repeat,
+                common_meta=common_meta,
+            )
+            all_rows.extend(prefill_alone_rows)
+
+        logger.info("Creating decode instance...")
+        client.request_json("POST", "/instance", body={"instance": decode_instance})
+        wait_for_instance_ready(client, decode_id)
+        decode_alive = True
+        logger.info("Decode instance ready")
+
+        logger.info("Linking instances (prefill → decode)...")
+        _create_instance_link(client, prefill_id, decode_id)
+        time.sleep(1)
+        links = _list_instance_links(client)
+        if not links:
+            logger.error("Link did not appear in state.")
+            return 1
+        link_id = str(links[-1].get("linkId") or links[-1].get("link_id") or "")
+        logger.info(f"Link created: {link_id}")
+        time.sleep(2)
+
+        disagg_rows = _run_phase(
+            client=client,
+            label="disaggregated",
+            pp_tg_pairs=pp_tg_pairs,
+            model_id=decode_full_id,
+            prompt_sizer=decode_prompt_sizer,
+            warmup=args.warmup,
+            repeat=args.repeat,
+            common_meta=common_meta,
+        )
+        all_rows.extend(disagg_rows)
+
+        if args.compare_baseline:
+            logger.info("Removing link and prefill instance to isolate decode_alone.")
+            with contextlib.suppress(ExoHttpError):
+                if link_id:
+                    _delete_instance_link(client, link_id)
+                    link_id = ""
+            with contextlib.suppress(ExoHttpError):
+                client.request_json("DELETE", f"/instance/{prefill_id}")
+            wait_for_instance_gone(client, prefill_id)
+            prefill_alive = False
+            time.sleep(2)
+
+            decode_alone_rows = _run_phase(
+                client=client,
+                label="decode_alone",
+                pp_tg_pairs=pp_tg_pairs,
+                model_id=decode_full_id,
+                prompt_sizer=decode_prompt_sizer,
+                warmup=args.warmup,
+                repeat=args.repeat,
+                common_meta=common_meta,
+            )
+            all_rows.extend(decode_alone_rows)
+
+            _print_diff(disagg_rows, decode_alone_rows, prefill_alone_rows)
+    finally:
+        with contextlib.suppress(ExoHttpError):
+            if link_id:
+                _delete_instance_link(client, link_id)
+        if decode_alive:
+            with contextlib.suppress(ExoHttpError):
+                client.request_json("DELETE", f"/instance/{decode_id}")
+            wait_for_instance_gone(client, decode_id)
+        if prefill_alive:
+            with contextlib.suppress(ExoHttpError):
+                client.request_json("DELETE", f"/instance/{prefill_id}")
+            wait_for_instance_gone(client, prefill_id)
+        logger.debug("Deleted both instances")
+
+    if args.stdout:
+        json.dump(all_rows, sys.stdout, indent=2, ensure_ascii=False)
+    elif args.json_out:
+        with open(args.json_out, "w", encoding="utf-8") as f:
+            json.dump(all_rows, f, indent=2, ensure_ascii=False)
+        logger.debug(f"\nWrote results JSON: {args.json_out}")
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/dashboard/src/lib/components/HeaderNav.svelte b/dashboard/src/lib/components/HeaderNav.svelte
index d7dab78d..36e5b659 100644
--- a/dashboard/src/lib/components/HeaderNav.svelte
+++ b/dashboard/src/lib/components/HeaderNav.svelte
@@ -1,5 +1,8 @@
 <script lang="ts">
   import { browser } from "$app/environment";
+  import { featureFlags } from "$lib/stores/app.svelte";
+
+  const showAdvanced = $derived(featureFlags()["disaggregation"] === true);
 
   interface Props {
     showHome?: boolean;
@@ -297,5 +300,28 @@
       </svg>
       <span class="hidden sm:inline">Integrations</span>
     </a>
+    {#if showAdvanced}
+      <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>
+    {/if}
   </nav>
 </header>
diff --git a/dashboard/src/lib/components/PrefillDecodeDisaggregation.svelte b/dashboard/src/lib/components/PrefillDecodeDisaggregation.svelte
new file mode 100644
index 00000000..7b9c1f65
--- /dev/null
+++ b/dashboard/src/lib/components/PrefillDecodeDisaggregation.svelte
@@ -0,0 +1,565 @@
+<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[];
+    nodeCount: number;
+  };
+
+  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 nodeIds = Object.keys(nodeToRunner);
+      const nodeNames = nodeIds
+        .map((nodeId) => ids[nodeId]?.friendlyName ?? nodeId.slice(0, 6))
+        .filter((name) => !!name);
+      rows.push({
+        id,
+        modelId,
+        family: deriveFamily(modelId),
+        baseModel: deriveBaseModel(modelId),
+        nodeNames,
+        nodeCount: nodeIds.length,
+      });
+    }
+    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[];
+    multiNode: boolean;
+  };
+
+  const linkRows = $derived.by<LinkRow[]>(() => {
+    const rows: LinkRow[] = [];
+    for (const [, link] of Object.entries(instanceLinks())) {
+      const fams = new Set<string>();
+      let multiNode = false;
+      for (const id of [...link.prefillInstances, ...link.decodeInstances]) {
+        const r = instanceById[id];
+        if (r && r.baseModel) fams.add(r.baseModel.toLowerCase());
+        if (r && r.nodeCount > 1) multiNode = true;
+      }
+      rows.push({
+        linkId: link.linkId,
+        prefill: link.prefillInstances,
+        decode: link.decodeInstances,
+        families: Array.from(fams),
+        multiNode,
+      });
+    }
+    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 editingMultiNode = $derived.by<string[]>(() => {
+    const names: string[] = [];
+    for (const id of [...editingPrefill, ...editingDecode]) {
+      const r = instanceById[id];
+      if (r && r.nodeCount > 1) {
+        names.push(r.baseModel || r.modelId);
+      }
+    }
+    return names;
+  });
+
+  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="font-mono text-foreground">
+  <div class="mb-6 space-y-4">
+    <details open class="group [&_summary::-webkit-details-marker]:hidden">
+      <summary
+        class="cursor-pointer list-none text-exo-yellow text-xs font-mono tracking-widest uppercase flex items-center gap-2 hover:opacity-80 transition-opacity"
+      >
+        <span
+          class="inline-block transition-transform group-open:rotate-90 text-exo-light-gray"
+          >▶</span
+        >
+        Prefill vs Decode
+      </summary>
+      <div class="mt-2 text-white/80 text-sm leading-relaxed">
+        Prefill is the compute-heavy pass that consumes the entire prompt and
+        builds a KV cache. Decode is the memory-bandwidth-bound loop that emits
+        tokens sequentially from that cache. The two phases have very different
+        bottlenecks, so running them on different hardware can be substantially
+        faster than doing both on one node.
+      </div>
+    </details>
+    <details class="group [&_summary::-webkit-details-marker]:hidden">
+      <summary
+        class="cursor-pointer list-none text-exo-yellow text-xs font-mono tracking-widest uppercase flex items-center gap-2 hover:opacity-80 transition-opacity"
+      >
+        <span
+          class="inline-block transition-transform group-open:rotate-90 text-exo-light-gray"
+          >▶</span
+        >
+        Linking Instances
+      </summary>
+      <div class="mt-2 text-white/80 text-sm leading-relaxed space-y-2">
+        <p>
+          A linked route here tells the cluster: when a request is sent to a
+          model in that cluster, the decode node (or the least active one if
+          there are multiple) will handle it. If it decides it must do a lot of
+          prefill not already cached in the prefix cache, it routes the request
+          to the prefill node over TCP IP. The prefill node streams the KV cache
+          back to the decode node which picks up from there.
+        </p>
+        <p>
+          Linked instances must be running the same model family — KV layouts
+          differ across architectures. More on the <a
+            class="text-exo-yellow underline underline-offset-2 hover:text-exo-yellow-darker transition-colors"
+            href="https://blog.exolabs.net/nvidia-dgx-spark/"
+            target="_blank"
+            rel="noreferrer noopener">blog</a
+          >.
+        </p>
+      </div>
+    </details>
+  </div>
+
+  {#if errorMessage}
+    <div
+      class="mb-4 px-4 py-3 bg-red-500/10 border border-red-500/40 text-red-300 text-sm"
+    >
+      {errorMessage}
+    </div>
+  {/if}
+
+  <section class="mt-12">
+    <h2
+      class="text-exo-yellow text-xs font-mono tracking-widest uppercase m-0 mb-3"
+    >
+      Existing routes
+    </h2>
+
+    {#if linkRows.length === 0}
+      {#if editingLinkId === null}
+        <div class="flex items-center justify-between">
+          <p class="text-exo-light-gray italic text-sm m-0">
+            No routes yet. Create one to enable remote prefill.
+          </p>
+          <button
+            class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-yellow/15 border border-exo-yellow/50 text-exo-yellow hover:bg-exo-yellow/25 hover:border-exo-yellow/80 transition-colors"
+            onclick={startCreate}
+          >
+            + New route
+          </button>
+        </div>
+      {/if}
+    {:else}
+      {#if editingLinkId === null}
+        <div class="flex justify-end mb-3">
+          <button
+            class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-yellow/15 border border-exo-yellow/50 text-exo-yellow hover:bg-exo-yellow/25 hover:border-exo-yellow/80 transition-colors"
+            onclick={startCreate}
+          >
+            + New route
+          </button>
+        </div>
+      {/if}
+      <div
+        class="bg-exo-dark-gray/60 border border-exo-medium-gray/40 flex flex-col"
+      >
+        {#each linkRows as row (row.linkId)}
+          {#if editingLinkId !== row.linkId}
+            <article
+              class="p-4 border-b border-exo-light-gray/25 last:border-b-0"
+            >
+              {#if row.multiNode}
+                <div
+                  class="mb-3 px-3 py-2 bg-red-500/10 border border-red-500/40 text-red-300 text-xs tracking-wide"
+                >
+                  ⚠ Multi-node instance detected. Remote prefill currently only
+                  works on single-node (rank-0) instances. This route will not
+                  function until that's supported.
+                </div>
+              {/if}
+              {#if row.families.length > 1}
+                <div
+                  class="mb-3 px-3 py-2 bg-amber-500/10 border border-amber-500/40 text-amber-300 text-xs tracking-wide"
+                >
+                  ⚠ Mixed model families: {row.families.join(", ")}
+                </div>
+              {/if}
+              <div
+                class="grid grid-cols-[1fr_auto_1fr_auto] items-center gap-x-3 gap-y-2"
+              >
+                <span
+                  class="inline-block justify-self-start text-[10px] font-mono tracking-widest uppercase px-2 py-0.5 bg-exo-yellow/15 border border-exo-yellow/40 text-exo-yellow"
+                  >Prefill</span
+                >
+                <span></span>
+                <span
+                  class="inline-block justify-self-start text-[10px] font-mono tracking-widest uppercase px-2 py-0.5 bg-exo-medium-gray/40 border border-exo-medium-gray/60 text-foreground"
+                  >Decode</span
+                >
+                <span></span>
+                <div class="min-w-0">
+                  <ul class="list-none p-0 m-0 flex flex-col gap-2">
+                    {#each row.prefill as id (id)}
+                      {@const r = instanceById[id]}
+                      {#if r}
+                        <li
+                          class="flex items-center gap-2 px-2.5 py-2 bg-exo-medium-gray/20 border border-exo-medium-gray/40"
+                        >
+                          <FamilyLogos family={r.family} />
+                          <div class="min-w-0 flex-1">
+                            <div
+                              class="text-exo-yellow text-xs font-mono truncate"
+                            >
+                              {r.baseModel || r.modelId}
+                            </div>
+                            <div
+                              class="text-exo-light-gray text-[11px] truncate"
+                            >
+                              {r.nodeNames.join(", ") || "?"}{r.nodeCount > 1
+                                ? ` (${r.nodeCount} nodes)`
+                                : ""}
+                            </div>
+                            <div
+                              class="text-exo-light-gray/40 text-[10px] font-mono truncate"
+                              title={r.id}
+                            >
+                              {r.id.slice(0, 8)}
+                            </div>
+                          </div>
+                        </li>
+                      {/if}
+                    {/each}
+                  </ul>
+                </div>
+                <div class="text-exo-yellow/60 text-xl px-2" aria-hidden="true">
+                  →
+                </div>
+                <div class="min-w-0">
+                  <ul class="list-none p-0 m-0 flex flex-col gap-2">
+                    {#each row.decode as id (id)}
+                      {@const r = instanceById[id]}
+                      {#if r}
+                        <li
+                          class="flex items-center gap-2 px-2.5 py-2 bg-exo-medium-gray/20 border border-exo-medium-gray/40"
+                        >
+                          <FamilyLogos family={r.family} />
+                          <div class="min-w-0 flex-1">
+                            <div
+                              class="text-exo-yellow text-xs font-mono truncate"
+                            >
+                              {r.baseModel || r.modelId}
+                            </div>
+                            <div
+                              class="text-exo-light-gray text-[11px] truncate"
+                            >
+                              {r.nodeNames.join(", ") || "?"}{r.nodeCount > 1
+                                ? ` (${r.nodeCount} nodes)`
+                                : ""}
+                            </div>
+                            <div
+                              class="text-exo-light-gray/40 text-[10px] font-mono truncate"
+                              title={r.id}
+                            >
+                              {r.id.slice(0, 8)}
+                            </div>
+                          </div>
+                        </li>
+                      {/if}
+                    {/each}
+                  </ul>
+                </div>
+                <div class="flex gap-2 pl-3">
+                  <button
+                    class="px-2 py-0.5 text-[11px] font-mono tracking-wider uppercase bg-exo-medium-gray/30 border border-exo-medium-gray/60 rounded text-foreground hover:border-exo-yellow/60 hover:text-exo-yellow disabled:opacity-40 disabled:cursor-not-allowed transition-colors"
+                    onclick={() => startEdit(row)}
+                    disabled={editingLinkId !== null}
+                  >
+                    Edit
+                  </button>
+                  <button
+                    class="px-2 py-0.5 text-[11px] font-mono tracking-wider uppercase bg-red-500/15 border border-red-500/40 rounded text-red-300 hover:bg-red-500/25 transition-colors"
+                    onclick={() => remove(row.linkId)}
+                  >
+                    Remove
+                  </button>
+                </div>
+              </div>
+            </article>
+          {/if}
+        {/each}
+      </div>
+    {/if}
+  </section>
+
+  {#if editingLinkId !== null && instanceRows.length === 0}
+    <section
+      class="mt-6 bg-exo-dark-gray/60 border border-exo-yellow/30 px-4 py-2.5 flex items-center justify-between gap-3"
+    >
+      <span class="text-exo-light-gray italic text-sm font-mono"
+        >No instances available.</span
+      >
+      <button
+        class="px-3 py-1 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 border border-exo-medium-gray/60 rounded text-foreground hover:border-exo-yellow/60 transition-colors"
+        onclick={cancelEdit}
+      >
+        Cancel
+      </button>
+    </section>
+  {:else if editingLinkId !== null}
+    <section class="mt-6 bg-exo-dark-gray/60 border border-exo-yellow/30 p-5">
+      <h2
+        class="text-exo-yellow text-xs font-mono tracking-widest uppercase m-0 mb-3"
+      >
+        {editingLinkId === "new" ? "New route" : "Edit route"}
+      </h2>
+
+      {#if editingMismatch}
+        <div
+          class="mb-3 px-3 py-2 bg-amber-500/10 border border-amber-500/40 text-amber-300 text-xs tracking-wide"
+        >
+          ⚠ Selected instances span multiple model families: <strong
+            >{editingFamilies.join(", ")}</strong
+          >. Linking across families produces a corrupt KV cache.
+        </div>
+      {/if}
+
+      {#if editingMultiNode.length > 0}
+        <div
+          class="mb-3 px-3 py-2 bg-red-500/10 border border-red-500/40 text-red-300 text-xs tracking-wide"
+        >
+          ⚠ Multi-node instance(s) selected: <strong
+            >{editingMultiNode.join(", ")}</strong
+          >. Remote prefill currently only works on single-node instances. This
+          route will not function until multi-node support lands.
+        </div>
+      {/if}
+
+      <p class="text-exo-light-gray text-xs mb-4">
+        Pick a role for each instance:
+        <span class="text-exo-yellow">Prefill</span>
+        serves KV cache,
+        <span class="text-foreground">Decode</span> consumes it.
+      </p>
+      <div
+        class="grid gap-2.5"
+        style="grid-template-columns: repeat(auto-fill, minmax(360px, 1fr));"
+      >
+        {#each instanceRows as row (row.id)}
+          {@const role = roleOf(row.id)}
+          <div
+            class="border p-3 flex flex-col gap-2.5 transition-colors {role ===
+            'prefill'
+              ? 'border-exo-yellow/60 bg-exo-dark-gray/60'
+              : role === 'decode'
+                ? 'border-exo-light-gray/60 bg-exo-dark-gray/60'
+                : 'border-exo-medium-gray/40 bg-exo-dark-gray/40'}"
+          >
+            <div class="flex items-center gap-2">
+              <FamilyLogos family={row.family} />
+              <div class="min-w-0 flex-1">
+                <div class="text-exo-yellow text-xs font-mono truncate">
+                  {row.baseModel || row.modelId}
+                </div>
+                <div class="text-exo-light-gray text-[11px] truncate">
+                  {row.nodeNames.join(", ") || "?"}{row.nodeCount > 1
+                    ? ` (${row.nodeCount} nodes)`
+                    : ""}
+                </div>
+                <div
+                  class="text-exo-light-gray/40 text-[10px] font-mono truncate"
+                  title={row.id}
+                >
+                  {row.id.slice(0, 8)}
+                </div>
+              </div>
+              {#if row.nodeCount > 1}
+                <span
+                  class="text-[9px] font-mono tracking-widest uppercase px-1.5 py-0.5 bg-red-500/15 border border-red-500/40 text-red-300"
+                  title="Multi-node instances are not supported by remote prefill yet."
+                  >Unsupported</span
+                >
+              {/if}
+            </div>
+            <div
+              class="flex rounded-md overflow-hidden border border-exo-light-gray/40 divide-x divide-exo-light-gray/40"
+            >
+              <button
+                class="flex-1 px-2 py-1 text-[11px] font-mono tracking-wider uppercase transition-colors {role ===
+                'prefill'
+                  ? 'bg-exo-yellow/20 text-exo-yellow'
+                  : 'bg-transparent text-white/80 hover:text-exo-yellow'}"
+                onclick={() =>
+                  setRole(row.id, role === "prefill" ? "none" : "prefill")}
+                >Prefill</button
+              >
+              <button
+                class="flex-1 px-2 py-1 text-[11px] font-mono tracking-wider uppercase transition-colors {role ===
+                'decode'
+                  ? 'bg-exo-medium-gray/50 text-foreground'
+                  : 'bg-transparent text-white/80 hover:text-foreground'}"
+                onclick={() =>
+                  setRole(row.id, role === "decode" ? "none" : "decode")}
+                >Decode</button
+              >
+            </div>
+          </div>
+        {/each}
+      </div>
+
+      <div class="flex gap-2 mt-5 justify-end">
+        <button
+          class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-yellow/15 border border-exo-yellow/50 text-exo-yellow hover:bg-exo-yellow/25 hover:border-exo-yellow/80 disabled:opacity-40 disabled:cursor-not-allowed transition-colors"
+          onclick={save}
+          disabled={!canSave}
+        >
+          {saving ? "Saving..." : "Save route"}
+        </button>
+        <button
+          class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 border border-exo-medium-gray/60 text-foreground hover:border-exo-yellow/60 disabled:opacity-40 disabled:cursor-not-allowed transition-colors"
+          onclick={cancelEdit}
+          disabled={saving}
+        >
+          Cancel
+        </button>
+      </div>
+    </section>
+  {/if}
+</div>
diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts
index dd6bd854..5b28b3f0 100644
--- a/dashboard/src/lib/stores/app.svelte.ts
+++ b/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,8 @@ class AppStore {
   topologyData = $state<TopologyData | null>(null);
   instances = $state<Record<string, unknown>>({});
   runners = $state<Record<string, unknown>>({});
+  instanceLinks = $state<Record<string, RawInstanceLink>>({});
+  featureFlags = $state<Record<string, boolean>>({});
   downloads = $state<Record<string, unknown[]>>({});
   nodeDisk = $state<
     Record<
@@ -1274,6 +1283,7 @@ class AppStore {
 
   startPolling() {
     this.fetchState();
+    this.fetchFeatureFlags();
     this.fetchInterval = setInterval(() => this.fetchState(), 1000);
   }
 
@@ -1285,6 +1295,16 @@ class AppStore {
     this.stopPreviewsPolling();
   }
 
+  async fetchFeatureFlags() {
+    try {
+      const response = await fetch("/v1/feature-flags");
+      if (!response.ok) return;
+      this.featureFlags = await response.json();
+    } catch {
+      // Silently ignore — defaults to all-disabled.
+    }
+  }
+
   async fetchState() {
     try {
       const response = await fetch("/state");
@@ -1310,6 +1330,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;
       }
@@ -3307,6 +3332,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
    */
@@ -3405,6 +3484,19 @@ 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 featureFlags = () => appStore.featureFlags;
+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;
diff --git a/dashboard/src/lib/utils/model_family.ts b/dashboard/src/lib/utils/model_family.ts
new file mode 100644
index 00000000..d5b1fd37
--- /dev/null
+++ b/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;
+}
diff --git a/dashboard/src/routes/advanced/+page.svelte b/dashboard/src/routes/advanced/+page.svelte
new file mode 100644
index 00000000..3b45cf23
--- /dev/null
+++ b/dashboard/src/routes/advanced/+page.svelte
@@ -0,0 +1,81 @@
+<script lang="ts">
+  import { browser } from "$app/environment";
+  import HeaderNav from "$lib/components/HeaderNav.svelte";
+  import PrefillDecodeDisaggregation from "$lib/components/PrefillDecodeDisaggregation.svelte";
+  import { featureFlags, refreshState } from "$lib/stores/app.svelte";
+  import { onMount } from "svelte";
+
+  type TabId = "prefill-decode";
+
+  const tabs: { id: TabId; label: string }[] = [
+    { id: "prefill-decode", label: "Prefill / Decode" },
+  ];
+
+  let activeTab = $state<TabId>(tabs[0].id);
+  let flagsLoaded = $state(false);
+
+  onMount(() => {
+    refreshState().finally(() => {
+      flagsLoaded = true;
+    });
+  });
+
+  const flags = $derived(featureFlags());
+  const enabled = $derived(flags["disaggregation"] === true);
+
+  $effect(() => {
+    if (browser && flagsLoaded && !enabled) {
+      // No advanced features enabled — bounce home.
+      window.location.hash = "/";
+    }
+  });
+</script>
+
+<div class="min-h-screen bg-exo-dark-gray flex flex-col">
+  <HeaderNav />
+
+  <main class="flex-1 max-w-[1100px] mx-auto w-full px-4 md:px-6 py-8">
+    {#if !flagsLoaded}
+      <div class="text-exo-light-gray/60 text-sm">Loading…</div>
+    {:else if !enabled}
+      <div class="text-exo-light-gray/60 text-sm">
+        No advanced features enabled. Set <code
+          class="text-exo-yellow font-mono">ENABLE_DISAGGREGATION=true</code
+        > on the cluster to access prefill/decode disaggregation.
+      </div>
+    {:else}
+      <div class="mb-4">
+        <h1
+          class="text-white text-xl md:text-2xl font-semibold tracking-wide mb-2"
+        >
+          Advanced
+        </h1>
+        <p class="text-exo-light-gray/60 text-sm">
+          Cluster-level configuration. Most users don't need anything here.
+        </p>
+      </div>
+
+      <div
+        class="flex flex-wrap gap-2 mb-6 border-b border-exo-light-gray/10 pb-3"
+      >
+        {#each tabs as tab (tab.id)}
+          <button
+            onclick={() => (activeTab = tab.id)}
+            class="px-3 py-1.5 text-xs rounded-md transition-all cursor-pointer
+              {activeTab === tab.id
+              ? 'bg-exo-yellow/15 text-exo-yellow border border-exo-yellow/30'
+              : 'text-exo-light-gray/60 hover:text-white/80 border border-transparent hover:border-exo-light-gray/20'}"
+          >
+            {tab.label}
+          </button>
+        {/each}
+      </div>
+
+      <div class="space-y-4">
+        {#if activeTab === "prefill-decode"}
+          <PrefillDecodeDisaggregation />
+        {/if}
+      </div>
+    {/if}
+  </main>
+</div>
diff --git a/src/exo/api/main.py b/src/exo/api/main.py
index 565154ea..8fe0cfbe 100644
--- a/src/exo/api/main.py
+++ b/src/exo/api/main.py
@@ -79,6 +79,8 @@ from exo.api.types import (
     ImageListItem,
     ImageListResponse,
     ImageSize,
+    InstanceLinkBody,
+    InstanceLinkResponse,
     ModelList,
     ModelListModel,
     PlaceInstanceParams,
@@ -122,6 +124,7 @@ from exo.master.placement import place_instance as get_instance_placements
 from exo.shared.apply import apply
 from exo.shared.constants import (
     DASHBOARD_DIR,
+    ENABLE_DISAGGREGATION,
     EXO_CACHE_HOME,
     EXO_EVENT_LOG_DIR,
     EXO_IMAGE_CACHE_DIR,
@@ -154,6 +157,7 @@ from exo.shared.types.commands import (
     DeleteCustomModelCard,
     DeleteDownload,
     DeleteInstance,
+    DeleteInstanceLink,
     DownloadCommand,
     ForwarderCommand,
     ForwarderDownloadCommand,
@@ -161,6 +165,7 @@ from exo.shared.types.commands import (
     ImageGeneration,
     PlaceInstance,
     SendInputChunk,
+    SetInstanceLink,
     StartDownload,
     TaskCancelled,
     TaskFinished,
@@ -174,6 +179,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 (
@@ -215,6 +221,17 @@ def _ensure_seed(params: AdvancedImageParams | None) -> AdvancedImageParams:
     return params
 
 
+def _require_disaggregation_enabled() -> None:
+    if not ENABLE_DISAGGREGATION:
+        raise HTTPException(
+            status_code=HTTPStatus.NOT_FOUND,
+            detail=(
+                "Prefill/decode disaggregation is disabled. "
+                "Set ENABLE_DISAGGREGATION=true to enable."
+            ),
+        )
+
+
 class API:
     def __init__(
         self,
@@ -328,6 +345,11 @@ 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("/v1/feature-flags")(self.get_feature_flags)
         self.app.get("/models")(self.get_models)
         self.app.get("/v1/models")(self.get_models)
         self.app.post("/models/add")(self.add_custom_model)
@@ -617,6 +639,49 @@ class API:
             instance_id=instance_id,
         )
 
+    async def get_feature_flags(self) -> dict[str, bool]:
+        return {"disaggregation": ENABLE_DISAGGREGATION}
+
+    async def list_instance_links(self) -> list[InstanceLink]:
+        if not ENABLE_DISAGGREGATION:
+            return []
+        return list(self.state.instance_links.values())
+
+    async def create_instance_link(
+        self, body: InstanceLinkBody
+    ) -> InstanceLinkResponse:
+        _require_disaggregation_enabled()
+        return await self._set_instance_link(InstanceLinkId(), body)
+
+    async def update_instance_link(
+        self, link_id: InstanceLinkId, body: InstanceLinkBody
+    ) -> InstanceLinkResponse:
+        _require_disaggregation_enabled()
+        return await self._set_instance_link(link_id, body)
+
+    async def _set_instance_link(
+        self, link_id: InstanceLinkId, 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:
+        _require_disaggregation_enabled()
+        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(
diff --git a/src/exo/api/types/__init__.py b/src/exo/api/types/__init__.py
index 3a61c1cd..9cb2f834 100644
--- a/src/exo/api/types/__init__.py
+++ b/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
diff --git a/src/exo/api/types/api.py b/src/exo/api/types/api.py
index 859eca77..8cfa10dd 100644
--- a/src/exo/api/types/api.py
+++ b/src/exo/api/types/api.py
@@ -296,6 +296,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",
diff --git a/src/exo/master/main.py b/src/exo/master/main.py
index 35c49390..044dfbf4 100644
--- a/src/exo/master/main.py
+++ b/src/exo/master/main.py
@@ -10,6 +10,7 @@ from exo.master.placement import (
     get_transition_events,
     place_instance,
 )
+from exo.master.placement_utils import find_ip_prioritised
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
 from exo.shared.types.commands import (
@@ -17,6 +18,7 @@ from exo.shared.types.commands import (
     CreateInstance,
     DeleteCustomModelCard,
     DeleteInstance,
+    DeleteInstanceLink,
     ForwarderCommand,
     ForwarderDownloadCommand,
     ImageEdits,
@@ -24,6 +26,7 @@ from exo.shared.types.commands import (
     PlaceInstance,
     RequestEventLog,
     SendInputChunk,
+    SetInstanceLink,
     TaskCancelled,
     TaskFinished,
     TestCommand,
@@ -38,6 +41,8 @@ from exo.shared.types.events import (
     IndexedEvent,
     InputChunkReceived,
     InstanceDeleted,
+    InstanceLinkCreated,
+    InstanceLinkDeleted,
     LocalForwarderEvent,
     NodeGatheredInfo,
     NodeTimedOut,
@@ -48,6 +53,7 @@ from exo.shared.types.events import (
     TracesCollected,
     TracesMerged,
 )
+from exo.shared.types.instance_link import InstanceLink
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
     ImageEdits as ImageEditsTask,
@@ -69,6 +75,46 @@ from exo.utils.event_buffer import MultiSourceBuffer
 from exo.utils.task_group import TaskGroup
 
 
+def _prefill_endpoint_for(state: State, decode_instance_id: InstanceId) -> str | None:
+    decode = state.instances.get(decode_instance_id)
+    if decode is None:
+        return None
+    decode_node = next(iter(decode.shard_assignments.node_to_runner.keys()), None)
+    if decode_node is None:
+        return None
+
+    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)
+
+    in_flight = {TaskStatus.Pending, TaskStatus.Running}
+    task_counts: dict[InstanceId, int] = {
+        src_id: sum(
+            1
+            for task in state.tasks.values()
+            if task.instance_id == src_id and task.task_status in in_flight
+        )
+        for src_id in sources
+    }
+    for src_id in sorted(sources, key=lambda sid: task_counts[sid]):
+        instance = state.instances.get(src_id)
+        if instance is None:
+            continue
+        for node_id, runner_id in instance.shard_assignments.node_to_runner.items():
+            port = state.prefill_server_ports.get(runner_id)
+            if port is None:
+                continue
+            ip = find_ip_prioritised(
+                decode_node, node_id, state.topology, state.node_network, ring=True
+            )
+            if ip is None:
+                continue
+            return f"{ip}:{port}"
+    return None
+
+
 class Master:
     def __init__(
         self,
@@ -128,10 +174,17 @@ class Master:
                         case TestCommand():
                             pass
                         case TextGeneration():
+                            prefill_only: set[InstanceId] = set()
+                            for link in self.state.instance_links.values():
+                                prefill_only.update(link.prefill_instances)
+                            for link in self.state.instance_links.values():
+                                prefill_only.difference_update(link.decode_instances)
+
                             for instance in self.state.instances.values():
                                 if (
                                     instance.shard_assignments.model_id
                                     == command.task_params.model
+                                    and instance.instance_id not in prefill_only
                                 ):
                                     in_flight = {TaskStatus.Pending, TaskStatus.Running}
                                     task_count = sum(
@@ -156,20 +209,27 @@ class Master:
                                 ],
                             )
 
+                            decode_instance_id = available_instance_ids[0]
                             task_id = TaskId()
+                            params = command.task_params.model_copy(
+                                update={
+                                    "prefill_endpoint": _prefill_endpoint_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():
@@ -363,6 +423,21 @@ class Master:
                             generated_events.append(
                                 CustomModelCardDeleted(model_id=command.model_id)
                             )
+                        case SetInstanceLink():
+                            link = InstanceLink(
+                                link_id=command.link_id,
+                                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
diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py
index 65b97831..71593951 100644
--- a/src/exo/master/placement.py
+++ b/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(
diff --git a/src/exo/master/placement_utils.py b/src/exo/master/placement_utils.py
index eb78c8fa..0375e97e 100644
--- a/src/exo/master/placement_utils.py
+++ b/src/exo/master/placement_utils.py
@@ -336,7 +336,7 @@ def _find_connection_ip(
             yield connection.sink_multiaddr.ip_address
 
 
-def _find_ip_prioritised(
+def find_ip_prioritised(
     node_id: NodeId,
     other_node_id: NodeId,
     cycle_digraph: Topology,
@@ -413,7 +413,7 @@ def get_mlx_ring_hosts_by_node(
                 hosts_for_node.append(Host(ip="198.51.100.1", port=0))
                 continue
 
-            connection_ip = _find_ip_prioritised(
+            connection_ip = find_ip_prioritised(
                 node_id, other_node_id, cycle_digraph, node_network, ring=True
             )
             if connection_ip is None:
@@ -445,7 +445,7 @@ def get_mlx_jaccl_coordinators(
         if n == coordinator:
             return "0.0.0.0"
 
-        ip = _find_ip_prioritised(
+        ip = find_ip_prioritised(
             n, coordinator, cycle_digraph, node_network, ring=False
         )
         if ip is not None:
diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py
index fb9f6d9a..959f7765 100644
--- a/src/exo/shared/apply.py
+++ b/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,
@@ -41,7 +44,12 @@ from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.topology import Connection, RDMAConnection
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
-from exo.shared.types.worker.runners import RunnerId, RunnerShutdown, RunnerStatus
+from exo.shared.types.worker.runners import (
+    RunnerId,
+    RunnerReady,
+    RunnerShutdown,
+    RunnerStatus,
+)
 from exo.utils.info_gatherer.info_gatherer import (
     MacmonMetrics,
     MacThunderboltConnections,
@@ -95,6 +103,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 +206,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:
@@ -202,12 +245,28 @@ def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> Sta
         new_runners: Mapping[RunnerId, RunnerStatus] = {
             rid: rs for rid, rs in state.runners.items() if rid != event.runner_id
         }
-        return state.model_copy(update={"runners": new_runners})
+        new_ports: Mapping[RunnerId, int] = {
+            rid: p
+            for rid, p in state.prefill_server_ports.items()
+            if rid != event.runner_id
+        }
+        return state.model_copy(
+            update={"runners": new_runners, "prefill_server_ports": new_ports}
+        )
     new_runners = {
         **state.runners,
         event.runner_id: event.runner_status,
     }
-    return state.model_copy(update={"runners": new_runners})
+    update: dict[str, object] = {"runners": new_runners}
+    if (
+        isinstance(event.runner_status, RunnerReady)
+        and event.runner_status.prefill_server_port is not None
+    ):
+        update["prefill_server_ports"] = {
+            **state.prefill_server_ports,
+            event.runner_id: event.runner_status.prefill_server_port,
+        }
+    return state.model_copy(update=update)
 
 
 def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
diff --git a/src/exo/shared/constants.py b/src/exo/shared/constants.py
index 01f315bc..d7935418 100644
--- a/src/exo/shared/constants.py
+++ b/src/exo/shared/constants.py
@@ -96,6 +96,8 @@ EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true"
 
 EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true"
 
+ENABLE_DISAGGREGATION = os.getenv("ENABLE_DISAGGREGATION", "false").lower() == "true"
+
 EXO_MAX_CONCURRENT_REQUESTS = int(os.getenv("EXO_MAX_CONCURRENT_REQUESTS", "8"))
 
 EXO_MAX_INSTANCE_RETRIES = 5
diff --git a/src/exo/shared/tests/test_apply/test_apply_instance_link.py b/src/exo/shared/tests/test_apply/test_apply_instance_link.py
new file mode 100644
index 00000000..0af2e888
--- /dev/null
+++ b/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
diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py
index b2dc3c89..67d318b2 100644
--- a/src/exo/shared/types/commands.py
+++ b/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
+    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
 )
 
 
diff --git a/src/exo/shared/types/events.py b/src/exo/shared/types/events.py
index b750a2ae..bebaaac5 100644
--- a/src/exo/shared/types/events.py
+++ b/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
 )
 
 
diff --git a/src/exo/shared/types/instance_link.py b/src/exo/shared/types/instance_link.py
new file mode 100644
index 00000000..a3674a92
--- /dev/null
+++ b/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]
diff --git a/src/exo/shared/types/state.py b/src/exo/shared/types/state.py
index 71cffae7..6c976984 100644
--- a/src/exo/shared/types/state.py
+++ b/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,9 @@ 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] = {}
+    prefill_server_ports: Mapping[RunnerId, int] = {}
+
     @field_serializer("topology", mode="plain")
     def _encode_topology(self, value: Topology) -> TopologySnapshot:
         return value.to_snapshot()
diff --git a/src/exo/shared/types/text_generation.py b/src/exo/shared/types/text_generation.py
index 116a5ec2..ccb2512a 100644
--- a/src/exo/shared/types/text_generation.py
+++ b/src/exo/shared/types/text_generation.py
@@ -132,6 +132,8 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     images: list[Base64Image] = Field(default_factory=list)
     image_hashes: dict[int, Base64ImageHash] = Field(default_factory=dict)
 
+    prefill_endpoint: str | None = None
+
     def with_card_sampling_defaults(self) -> "TextGenerationTaskParams":
         from exo.shared.models.model_cards import get_card
 
diff --git a/src/exo/shared/types/worker/runners.py b/src/exo/shared/types/worker/runners.py
index 4875c6e5..c906b6ec 100644
--- a/src/exo/shared/types/worker/runners.py
+++ b/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):
diff --git a/src/exo/utils/ports.py b/src/exo/utils/ports.py
new file mode 100644
index 00000000..f23463df
--- /dev/null
+++ b/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
diff --git a/src/exo/worker/disaggregated/__init__.py b/src/exo/worker/disaggregated/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/exo/worker/disaggregated/protocol.py b/src/exo/worker/disaggregated/protocol.py
new file mode 100644
index 00000000..7e2c0767
--- /dev/null
+++ b/src/exo/worker/disaggregated/protocol.py
@@ -0,0 +1,152 @@
+from typing import BinaryIO, Literal
+
+import msgspec
+
+DType = Literal["bfloat16", "float16", "float32"]
+
+
+class ProtocolError(Exception):
+    pass
+
+
+class Header(msgspec.Struct):
+    request_id: str = ""
+    model_id: str = ""
+    num_layers: int = 0
+    dtype: DType = "bfloat16"
+    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))
diff --git a/src/exo/worker/disaggregated/server.py b/src/exo/worker/disaggregated/server.py
new file mode 100644
index 00000000..67ba89d5
--- /dev/null
+++ b/src/exo/worker/disaggregated/server.py
@@ -0,0 +1,105 @@
+import socket
+import socketserver
+import threading
+from collections.abc import Callable
+from typing import BinaryIO, cast
+
+import msgspec
+from loguru import logger
+
+from exo.worker.disaggregated.protocol import (
+    Header,
+    read_frame,
+    write_error,
+    write_frame,
+    write_header,
+)
+
+
+class PrefillRequest(msgspec.Struct):
+    request_id: str = ""
+    model_id: str = ""
+    token_ids: list[int] = msgspec.field(default_factory=list)
+    start_pos: int = 0
+
+
+_request_encoder = msgspec.msgpack.Encoder()
+_request_decoder: msgspec.msgpack.Decoder[PrefillRequest] = msgspec.msgpack.Decoder(
+    PrefillRequest
+)
+
+
+def write_request(stream: BinaryIO, job: PrefillRequest) -> None:
+    write_frame(stream, _request_encoder.encode(job))
+
+
+def read_request(stream: BinaryIO) -> PrefillRequest:
+    payload = read_frame(stream)
+    if not payload:
+        raise ConnectionError("No request received")
+    return _request_decoder.decode(payload)
+
+
+ResolveHandler = Callable[[PrefillRequest, BinaryIO], bool]
+
+
+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 _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(PrefillServer, self.server)
+        wfile: BinaryIO = cast(BinaryIO, cast(object, self.wfile))
+        rfile: BinaryIO = cast(BinaryIO, cast(object, self.rfile))
+        try:
+            job = read_request(rfile)
+        except ConnectionError:
+            return
+        except (msgspec.DecodeError, ValueError) as exc:
+            _send_error(wfile, 400, f"Bad request: {exc}")
+            return
+        try:
+            picked_up = server.resolve(job, wfile)
+        except Exception as e:
+            logger.opt(exception=e).warning(
+                f"Prefill resolve error for request_id={job.request_id}"
+            )
+            _send_error(wfile, 500, str(e))
+            return
+        if not picked_up:
+            _send_error(
+                wfile, 503, f"Prefill not picked up for request_id={job.request_id!r}"
+            )
+
+
+class PrefillServer(socketserver.ThreadingTCPServer):
+    allow_reuse_address = True
+    daemon_threads = True
+    resolve: ResolveHandler
+
+    def __init__(self, resolve: ResolveHandler, host: str, port: int) -> None:
+        super().__init__((host, port), _PrefillHandler)
+        self.resolve = resolve
+        self._thread = threading.Thread(
+            target=self.serve_forever, name="prefill-server"
+        )
+        self._thread.start()
+        logger.info(f"Prefill server listening on {host}:{port}")
+
+    def stop(self) -> None:
+        self.shutdown()
+        self.server_close()
+        if self._thread is not None:
+            self._thread.join(timeout=5)
+            self._thread = None
diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index 14603f4f..733773f2 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -918,7 +918,7 @@ class DeepseekV4ShardingStrategy(TensorParallelShardingStrategy):
             # Head-parallel attention with interleaved-per-group sharding.
             _shard_v4_attention_heads(layer.attn, self.N, self.group.rank())
             self.sharded_to_all_linear_in_place(layer.attn.wo_a)
-            layer.attn.wo_b = _AllSumLinear(layer.attn.wo_b, self.group)  # type: ignore[assignment]
+            layer.attn.wo_b = _AllSumLinear(layer.attn.wo_b, self.group)  # type: ignore
 
             ffn = layer.ffn
             if getattr(ffn, "shared_experts", None) is not None:
@@ -930,7 +930,7 @@ class DeepseekV4ShardingStrategy(TensorParallelShardingStrategy):
             self.all_to_sharded_linear_in_place(ffn.switch_mlp.up_proj)
             wrapped = ShardedMoEV4(ffn)
             wrapped.sharding_group = self.group
-            layer.ffn = wrapped  # type: ignore[assignment]
+            layer.ffn = wrapped  # type: ignore
 
             mx.eval(layer)
             mx.clear_cache()
diff --git a/src/exo/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py
index 30100db8..7cdcc77f 100644
--- a/src/exo/worker/engines/mlx/cache.py
+++ b/src/exo/worker/engines/mlx/cache.py
@@ -22,8 +22,8 @@ from mlx_lm.models.deepseek_v4 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:
diff --git a/src/exo/worker/engines/mlx/disaggregated/__init__.py b/src/exo/worker/engines/mlx/disaggregated/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/exo/worker/engines/mlx/disaggregated/adapter.py b/src/exo/worker/engines/mlx/disaggregated/adapter.py
new file mode 100644
index 00000000..6753495f
--- /dev/null
+++ b/src/exo/worker/engines/mlx/disaggregated/adapter.py
@@ -0,0 +1,233 @@
+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 mlx_lm.models.deepseek_v4 import DeepseekV4Cache
+
+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
+from exo.worker.runner.bootstrap import logger
+
+_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:
+    match dtype:
+        case "bfloat16":
+            arr = np.frombuffer(data, dtype=np.uint16).reshape(shape).copy()
+            return mx.array(arr).view(mx.bfloat16)
+        case "float16":
+            arr = np.frombuffer(data, dtype=np.float16).reshape(shape).copy()
+            return mx.array(arr)
+        case "float32":
+            arr = np.frombuffer(data, dtype=np.float32).reshape(shape).copy()
+            return mx.array(arr)
+
+
+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):
+        match c:
+            case QuantizedKVCache() | CacheList() | DeepseekV4Cache():
+                raise NotImplementedError
+            case 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),
+                )
+                if tokens_sent != 0 and num_tokens != tokens_sent:
+                    logger.critical(
+                        f"Unexpected number of tokens sent {num_tokens} != {tokens_sent}"
+                    )
+                tokens_sent = num_tokens
+            case ArraysCache():
+                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 write_cache_to_wire(
+    wfile: BinaryIO,
+    cache: KVCacheType,
+    *,
+    request_id: str = "",
+    model_id: str = "",
+    start_pos: int = 0,
+) -> int:
+    dtype = wire_dtype_from_cache(cache)
+    write_header(
+        wfile,
+        Header(
+            request_id=request_id,
+            model_id=model_id,
+            num_layers=len(cache),
+            dtype=dtype,
+            start_pos=start_pos,
+        ),
+    )
+    tokens_sent = send_mlx_kv_cache(wfile, cache, dtype=dtype, start_pos=start_pos)
+    write_done(wfile, tokens_sent)
+    wfile.flush()
+    return tokens_sent
diff --git a/src/exo/worker/engines/mlx/disaggregated/client.py b/src/exo/worker/engines/mlx/disaggregated/client.py
new file mode 100644
index 00000000..7634ac0f
--- /dev/null
+++ b/src/exo/worker/engines/mlx/disaggregated/client.py
@@ -0,0 +1,147 @@
+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.disaggregated.protocol import (
+    ArraysState,
+    Done,
+    Header,
+    KVChunk,
+    TensorBlob,
+    read_header,
+    read_message,
+)
+from exo.worker.disaggregated.server import PrefillRequest, write_request
+from exo.worker.engines.mlx.disaggregated.adapter import (
+    chunk_to_mlx_nhd,
+    inject_arrays_cache,
+    inject_kv_chunk,
+    inject_rotating_kv_chunk,
+)
+
+_SOCKET_TIMEOUT_SECS = 60
+_RECV_BUFFER_BYTES = 4 * 1024 * 1024
+
+
+@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)
+    raise ValueError(f"Invalid endpoint {endpoint}")
+
+
+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:
+        wfile = sock.makefile("wb", buffering=256 * 1024)
+        wstream: BinaryIO = cast(BinaryIO, cast(object, wfile))
+        write_request(wstream, request)
+
+        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
diff --git a/src/exo/worker/engines/mlx/disaggregated/serve.py b/src/exo/worker/engines/mlx/disaggregated/serve.py
new file mode 100644
index 00000000..5220e9d1
--- /dev/null
+++ b/src/exo/worker/engines/mlx/disaggregated/serve.py
@@ -0,0 +1,86 @@
+import time
+
+import mlx.core as mx
+from mlx_lm.sample_utils import make_sampler
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+from exo.worker.disaggregated.server import PrefillRequest
+from exo.worker.engines.mlx.cache import (
+    KVPrefixCache,
+    cache_length,
+    make_kv_cache,
+    snapshot_ssm_states,
+)
+from exo.worker.engines.mlx.generator.generate import prefill as mlx_prefill
+from exo.worker.engines.mlx.types import KVCacheType, Model
+from exo.worker.engines.mlx.utils_mlx import fix_unmatched_think_end_tokens
+from exo.worker.runner.bootstrap import logger
+
+
+def run_prefill_for_request(
+    *,
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    group: mx.distributed.Group | None,
+    kv_prefix_cache: KVPrefixCache | None,
+    request: PrefillRequest,
+) -> KVCacheType:
+    prompt_tokens = mx.array(request.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
+
+    target_offset = max(0, n_tokens - 2)
+    new_tokens = max(0, target_offset - prefix_hit_length)
+    prefill_input = remaining[:new_tokens]
+    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"
+            )
+
+    elapsed = time.perf_counter() - t0
+    final_offset = cache_length(cache)
+    logger.info(
+        f"Prefill: request_id={request.request_id} "
+        f"{n_tokens} tokens (prefix_hit={prefix_hit_length}, "
+        f"final_offset={final_offset}) in {elapsed * 1000:.0f}ms"
+    )
+    return cache
diff --git a/src/exo/worker/engines/mlx/disaggregated/tests/__init__.py b/src/exo/worker/engines/mlx/disaggregated/tests/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/exo/worker/engines/mlx/disaggregated/tests/test_end_to_end.py b/src/exo/worker/engines/mlx/disaggregated/tests/test_end_to_end.py
new file mode 100644
index 00000000..2cfdfc9b
--- /dev/null
+++ b/src/exo/worker/engines/mlx/disaggregated/tests/test_end_to_end.py
@@ -0,0 +1,167 @@
+from typing import BinaryIO
+
+import mlx.core as mx
+import numpy as np
+import pytest
+from mlx_lm.models.cache import KVCache
+
+from exo.worker.disaggregated.protocol import Header, write_done, write_header
+from exo.worker.disaggregated.server import PrefillRequest, PrefillServer
+from exo.worker.engines.mlx.disaggregated.adapter import (
+    send_mlx_kv_cache,
+    wire_dtype_from_cache,
+)
+from exo.worker.engines.mlx.disaggregated.client import (
+    ingest_into_mlx_cache,
+    remote_prefill_fetch,
+)
+
+
+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()
+    with mx.stream(mx.Device(mx.cpu)):
+        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)
+        mx.eval(cache.keys, cache.values)
+    cache.offset = seq_len
+    return cache
+
+
+def _stream_cache(
+    wfile: BinaryIO, cache: KVCache, *, request_id: str, start_pos: int = 0
+) -> None:
+    dtype = wire_dtype_from_cache([cache])
+    write_header(
+        wfile,
+        Header(
+            request_id=request_id,
+            model_id="test-model",
+            num_layers=1,
+            dtype=dtype,
+            start_pos=start_pos,
+        ),
+    )
+    tokens_sent = send_mlx_kv_cache(wfile, [cache], dtype=dtype, start_pos=start_pos)
+    write_done(wfile, tokens_sent)
+    wfile.flush()
+
+
+@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)
+
+    def resolve(job: PrefillRequest, wfile: BinaryIO) -> bool:
+        _stream_cache(wfile, gold, request_id=job.request_id)
+        return True
+
+    server = PrefillServer(resolve=resolve, host="127.0.0.1", port=52417)
+    try:
+        result = remote_prefill_fetch(
+            endpoint="127.0.0.1:52417",
+            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_pickup_failure() -> None:
+    def resolve(_job: PrefillRequest, _wfile: BinaryIO) -> bool:
+        return False
+
+    server = PrefillServer(resolve=resolve, host="127.0.0.1", port=52418)
+    try:
+        with pytest.raises(RuntimeError, match="not picked up"):
+            _ = remote_prefill_fetch(
+                endpoint="127.0.0.1:52418",
+                request=PrefillRequest(
+                    model_id="test-model",
+                    token_ids=[1, 2, 3],
+                    request_id="never-registered",
+                ),
+            )
+    finally:
+        server.stop()
+
+
+@pytest.mark.slow
+def test_server_client_roundtrip_with_start_pos() -> None:
+    seq_len = 8
+    start_pos = 5
+    n_heads = 2
+    head_dim = 4
+    gold = _make_cache(seq_len, n_heads, head_dim)
+
+    def resolve(job: PrefillRequest, wfile: BinaryIO) -> bool:
+        _stream_cache(wfile, gold, request_id=job.request_id, start_pos=start_pos)
+        return True
+
+    server = PrefillServer(resolve=resolve, host="127.0.0.1", port=52419)
+    try:
+        result = remote_prefill_fetch(
+            endpoint="127.0.0.1:52419",
+            request=PrefillRequest(
+                model_id="test-model",
+                token_ids=list(range(seq_len)),
+                request_id="req-1",
+                start_pos=start_pos,
+            ),
+        )
+        assert result.total_tokens == seq_len - start_pos
+        assert result.header.start_pos == start_pos
+
+        dst = KVCache()
+        gold_k = gold.keys
+        gold_v = gold.values
+        assert gold_k is not None and gold_v is not None
+        dst.keys = mx.array(gold_k[:, :, :start_pos, :])
+        dst.values = mx.array(gold_v[:, :, :start_pos, :])
+        dst.offset = start_pos
+
+        final_offset = ingest_into_mlx_cache(result, [dst], start_pos=start_pos)
+        assert final_offset == seq_len
+        assert dst.offset == seq_len
+        dst_k = dst.keys
+        dst_v = dst.values
+        assert dst_k is not None and dst_v is not None
+        assert _equal(dst_k, gold_k)
+        assert _equal(dst_v, gold_v)
+    finally:
+        server.stop()
diff --git a/src/exo/worker/engines/mlx/disaggregated/tests/test_mlx_adapter.py b/src/exo/worker/engines/mlx/disaggregated/tests/test_mlx_adapter.py
new file mode 100644
index 00000000..a706a6b4
--- /dev/null
+++ b/src/exo/worker/engines/mlx/disaggregated/tests/test_mlx_adapter.py
@@ -0,0 +1,270 @@
+import io
+
+import mlx.core as mx
+import numpy as np
+from mlx_lm.models.cache import ArraysCache, KVCache, RotatingKVCache
+
+from exo.worker.disaggregated.protocol import (
+    ArraysState,
+    Done,
+    Header,
+    KVChunk,
+    TensorBlob,
+    read_header,
+    read_message,
+    write_done,
+    write_header,
+)
+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,
+    inject_rotating_kv_chunk,
+    nhd_to_bhsd,
+    send_mlx_kv_cache,
+    wire_dtype_from_cache,
+)
+from exo.worker.engines.mlx.disaggregated.client import (
+    PrefillResult,
+    ingest_into_mlx_cache,
+)
+
+
+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 _make_kv_cache(seq_len: int, n_heads: int, head_dim: int) -> KVCache:
+    cache = KVCache()
+    cache.keys = _rand((1, n_heads, seq_len, head_dim), mx.bfloat16)
+    cache.values = _rand((1, n_heads, seq_len, head_dim), mx.bfloat16)
+    cache.offset = seq_len
+    return cache
+
+
+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
+    src = _make_kv_cache(seq_len, n_heads, head_dim)
+    k_bhsd, v_bhsd = src.keys, src.values
+    assert k_bhsd is not None and v_bhsd is not None
+
+    buf = io.BytesIO()
+    write_header(buf, Header(num_layers=1, dtype="bfloat16"))
+    tokens = send_mlx_kv_cache(buf, [src], dtype="bfloat16")
+    write_done(buf, tokens)
+    buf.seek(0)
+
+    got_hdr = read_header(buf)
+    assert got_hdr.num_layers == 1
+
+    msg = read_message(buf)
+    assert isinstance(msg, KVChunk)
+    assert msg.num_tokens == seq_len
+    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)
+    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
+
+
+def test_send_with_start_pos_only_ships_suffix() -> None:
+    n_heads, head_dim = 2, 4
+    seq_len, start_pos = 6, 4
+    src = _make_kv_cache(seq_len, n_heads, head_dim)
+
+    buf = io.BytesIO()
+    write_header(buf, Header(num_layers=1, dtype="bfloat16", start_pos=start_pos))
+    tokens = send_mlx_kv_cache(buf, [src], dtype="bfloat16", start_pos=start_pos)
+    write_done(buf, tokens)
+    buf.seek(0)
+
+    _ = read_header(buf)
+    msg = read_message(buf)
+    assert isinstance(msg, KVChunk)
+    assert msg.num_tokens == seq_len - start_pos
+
+
+def test_send_skips_layer_when_offset_below_start_pos() -> None:
+    n_heads, head_dim = 2, 4
+    seq_len, start_pos = 3, 5
+    src = _make_kv_cache(seq_len, n_heads, head_dim)
+
+    buf = io.BytesIO()
+    write_header(buf, Header(num_layers=1, dtype="bfloat16", start_pos=start_pos))
+    tokens = send_mlx_kv_cache(buf, [src], dtype="bfloat16", start_pos=start_pos)
+    write_done(buf, tokens)
+    buf.seek(0)
+
+    _ = read_header(buf)
+    msg = read_message(buf)
+    assert isinstance(msg, Done)
+    assert msg.total_tokens == 0
+    assert tokens == 0
+
+
+def test_wire_dtype_from_cache() -> None:
+    src = _make_kv_cache(3, 2, 4)
+    assert wire_dtype_from_cache([src]) == "bfloat16"
+
+    f32 = KVCache()
+    f32.keys = _rand((1, 2, 3, 4), mx.float32)
+    f32.values = _rand((1, 2, 3, 4), mx.float32)
+    f32.offset = 3
+    assert wire_dtype_from_cache([f32]) == "float32"
+
+
+def _decode_payload(payload: bytes) -> PrefillResult:
+    buf = io.BytesIO(payload)
+    hdr = read_header(buf)
+    result = PrefillResult(header=hdr)
+    while True:
+        msg = read_message(buf)
+        if msg is None:
+            break
+        if isinstance(msg, KVChunk):
+            result.kv_chunks.setdefault(msg.layer_idx, []).append(msg)
+        elif isinstance(msg, ArraysState):
+            result.arrays[msg.layer_idx] = msg.arrays
+        elif isinstance(msg, Done):
+            result.total_tokens = msg.total_tokens
+            break
+    return result
+
+
+def test_mixed_cache_roundtrip() -> None:
+    n_heads, head_dim, seq_len = 2, 4, 6
+
+    src_kv = _make_kv_cache(seq_len, n_heads, head_dim)
+
+    src_rot = RotatingKVCache(max_size=16, keep=0)
+    src_rot.keys = _rand((1, n_heads, seq_len, head_dim), mx.bfloat16)
+    src_rot.values = _rand((1, n_heads, seq_len, head_dim), mx.bfloat16)
+    src_rot.offset = seq_len
+    src_rot._idx = seq_len
+
+    src_arr = ArraysCache(size=2)
+    arr_a = _rand((3,), mx.bfloat16)
+    arr_b = _rand((2, 4), mx.bfloat16)
+    src_arr.state = [arr_a, arr_b]
+
+    buf = io.BytesIO()
+    write_header(
+        buf,
+        Header(request_id="req", model_id="m", num_layers=3, dtype="bfloat16"),
+    )
+    tokens_sent = send_mlx_kv_cache(buf, [src_kv, src_rot, src_arr], dtype="bfloat16")
+    write_done(buf, tokens_sent)
+    result = _decode_payload(buf.getvalue())
+
+    assert result.header.num_layers == 3
+    assert result.total_tokens == seq_len
+
+    dst_kv = KVCache()
+    dst_rot = RotatingKVCache(max_size=16, keep=0)
+    dst_arr = ArraysCache(size=2)
+    final_offset = ingest_into_mlx_cache(result, [dst_kv, dst_rot, dst_arr])
+
+    assert final_offset == seq_len
+
+    assert dst_kv.offset == seq_len
+    assert dst_kv.keys is not None and dst_kv.values is not None
+    src_kv_k, src_kv_v = src_kv.keys, src_kv.values
+    assert src_kv_k is not None and src_kv_v is not None
+    assert _equal(dst_kv.keys, src_kv_k)
+    assert _equal(dst_kv.values, src_kv_v)
+
+    assert dst_rot.offset == seq_len
+    assert dst_rot.keys is not None and dst_rot.values is not None
+    src_rot_k, src_rot_v = src_rot.keys, src_rot.values
+    assert src_rot_k is not None and src_rot_v is not None
+    assert _equal(dst_rot.keys, src_rot_k)
+    assert _equal(dst_rot.values, src_rot_v)
+    assert dst_rot._idx == seq_len
+
+    assert len(dst_arr.state) == 2
+    s0, s1 = dst_arr.state[0], dst_arr.state[1]
+    assert s0 is not None and s1 is not None
+    assert _equal(s0, arr_a)
+    assert _equal(s1, arr_b)
+    _ = inject_rotating_kv_chunk
+    _ = nhd_to_bhsd
diff --git a/src/exo/worker/engines/mlx/disaggregated/tests/test_protocol_roundtrip.py b/src/exo/worker/engines/mlx/disaggregated/tests/test_protocol_roundtrip.py
new file mode 100644
index 00000000..7fdfba38
--- /dev/null
+++ b/src/exo/worker/engines/mlx/disaggregated/tests/test_protocol_roundtrip.py
@@ -0,0 +1,154 @@
+import io
+
+import pytest
+
+from exo.worker.disaggregated.protocol import (
+    ArraysState,
+    Done,
+    ErrorMessage,
+    Header,
+    KVChunk,
+    ProtocolError,
+    TensorBlob,
+    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 = Header(
+        request_id="r",
+        model_id="m",
+        num_layers=32,
+        dtype="bfloat16",
+        start_pos=42,
+    )
+    buf = io.BytesIO()
+    write_header(buf, hdr)
+    buf.seek(0)
+    got = read_header(buf)
+    assert got == hdr
+    assert got.dtype == "bfloat16"
+    assert got.num_layers == 32
+    assert got.start_pos == 42
+
+
+def test_kv_chunk_roundtrip() -> None:
+    num_tokens, n_heads, head_dim = 7, 4, 8
+    n_bytes = num_tokens * n_heads * head_dim * 2
+    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,
+        dtype="bfloat16",
+        keys=keys,
+        values=values,
+    )
+    buf.seek(0)
+    msg = read_message(buf)
+    assert isinstance(msg, KVChunk)
+    assert msg.layer_idx == 3
+    assert msg.shape == (num_tokens, n_heads, head_dim)
+    assert msg.dtype == "bfloat16"
+    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)
+    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)
+    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)
+    assert isinstance(msg, ErrorMessage)
+    assert msg.code == 42
+    assert msg.message == "boom"
+
+
+def test_stream_of_messages() -> None:
+    hdr = 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,
+        dtype="float32",
+        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)
+    m2 = read_message(buf)
+    m3 = read_message(buf)
+    m4 = read_message(buf)
+    assert isinstance(m1, KVChunk)
+    assert isinstance(m2, ArraysState)
+    assert isinstance(m3, Done)
+    assert m4 is None
+
+
+def test_corrupt_message_raises() -> None:
+    buf = io.BytesIO()
+    write_header(buf, Header(num_layers=1, dtype="float32"))
+    buf.write((5).to_bytes(4, "big"))
+    buf.write(b"\xff\xff\xff\xff\xff")
+    buf.seek(0)
+    _ = read_header(buf)
+    with pytest.raises(ProtocolError):
+        _ = read_message(buf)
diff --git a/src/exo/worker/engines/mlx/disaggregated/tests/test_server_drain.py b/src/exo/worker/engines/mlx/disaggregated/tests/test_server_drain.py
new file mode 100644
index 00000000..a9ad65e5
--- /dev/null
+++ b/src/exo/worker/engines/mlx/disaggregated/tests/test_server_drain.py
@@ -0,0 +1,120 @@
+"""Server thread receives request, runs resolve in another thread (mimicking
+runner main thread + work queue), streams cache bytes."""
+
+import queue
+import threading
+from typing import BinaryIO
+
+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.disaggregated.protocol import Header, write_done, write_header
+from exo.worker.disaggregated.server import PrefillRequest, PrefillServer
+from exo.worker.engines.mlx.disaggregated.adapter import (
+    send_mlx_kv_cache,
+    wire_dtype_from_cache,
+)
+from exo.worker.engines.mlx.disaggregated.client import (
+    PrefillResult,
+    ingest_into_mlx_cache,
+    remote_prefill_fetch,
+)
+
+
+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)
+
+    request_queue: queue.Queue[tuple[PrefillRequest, BinaryIO, threading.Event]] = (
+        queue.Queue()
+    )
+
+    def resolve(job: PrefillRequest, wfile: BinaryIO) -> bool:
+        done = threading.Event()
+        request_queue.put((job, wfile, done))
+        return done.wait(timeout=5)
+
+    server = PrefillServer(
+        resolve=resolve, host="127.0.0.1", port=(port := random_ephemeral_port())
+    )
+
+    def serve_one(wfile: BinaryIO) -> None:
+        dtype = wire_dtype_from_cache([gold])
+        write_header(
+            wfile,
+            Header(request_id="req-1", model_id="m", num_layers=1, dtype=dtype),
+        )
+        tokens = send_mlx_kv_cache(wfile, [gold], dtype=dtype)
+        write_done(wfile, tokens)
+        wfile.flush()
+
+    drained_job: list[PrefillRequest] = []
+    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:
+        job, wfile, done = request_queue.get(timeout=5)
+        drained_job.append(job)
+        try:
+            serve_one(wfile)
+        finally:
+            done.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()
diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py
index b0b2a537..32f06735 100644
--- a/src/exo/worker/engines/mlx/generator/batch_generate.py
+++ b/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,45 @@ class ExoBatchGenerator:
             if vision is not None
             else contextlib.nullcontext()
         )
+        uncached_count = len(prompt_tokens)
+        use_remote = (
+            uncached_count > REMOTE_PREFILL_MIN_TOKENS
+            and task_params.prefill_endpoint is not None
+        )
+
+        _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 and task_params.prefill_endpoint is not None:
+                try:
+                    _prefill_tps, _prefill_tokens, cache_snapshots = remote_prefill(
+                        prompt_tokens[:-1],
+                        cache,
+                        on_prefill_progress,
+                        endpoint=task_params.prefill_endpoint,
+                        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:
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index 47b32a1e..2e3d0512 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -2,6 +2,7 @@ import contextlib
 import functools
 import math
 import time
+import uuid
 from typing import Callable, Generator, cast, get_args
 
 import mlx.core as mx
@@ -22,7 +23,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 +55,8 @@ from exo.worker.engines.mlx.constants import (
     KV_GROUP_SIZE,
     MAX_TOKENS,
 )
+from exo.worker.engines.mlx.generator.remote_prefill import remote_prefill
+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,
@@ -70,6 +72,8 @@ from exo.worker.engines.mlx.vision import (
 )
 from exo.worker.runner.bootstrap import logger
 
+REMOTE_PREFILL_MIN_TOKENS = 1000
+
 generation_stream = mx.new_stream(mx.default_device())
 
 _MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5
@@ -633,17 +637,42 @@ def mlx_generate(
         if vision is not None
         else contextlib.nullcontext()
     )
+    use_remote = (
+        len(prompt_tokens) > REMOTE_PREFILL_MIN_TOKENS
+        and task.prefill_endpoint is not None
+    )
+    remote_prefilled = False
+    prefill_tps = 0.0
+    prefill_tokens = 0
+    ssm_snapshots_list: list[CacheSnapshot] = []
     with maybe_vision_ctx:
-        prefill_tps, prefill_tokens, ssm_snapshots_list = prefill(
-            model,
-            tokenizer,
-            sampler,
-            prompt_tokens[:-1],
-            caches,
-            group,
-            on_prefill_progress,
-            distributed_prompt_progress_callback,
-        )
+        if use_remote and task.prefill_endpoint is not None:
+            try:
+                prefill_tps, prefill_tokens, ssm_snapshots_list = remote_prefill(
+                    prompt_tokens[:-1],
+                    caches,
+                    on_prefill_progress,
+                    endpoint=task.prefill_endpoint,
+                    request_id=str(uuid.uuid4()),
+                    model_id=str(task.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, ssm_snapshots_list = prefill(
+                model,
+                tokenizer,
+                sampler,
+                prompt_tokens[:-1],
+                caches,
+                group,
+                on_prefill_progress,
+                distributed_prompt_progress_callback,
+            )
     cache_snapshots: list[CacheSnapshot] | None = ssm_snapshots_list or None
 
     if kv_prefix_cache is not None and matched_index is not None and is_exact_hit:
diff --git a/src/exo/worker/engines/mlx/generator/remote_prefill.py b/src/exo/worker/engines/mlx/generator/remote_prefill.py
new file mode 100644
index 00000000..b58ec65b
--- /dev/null
+++ b/src/exo/worker/engines/mlx/generator/remote_prefill.py
@@ -0,0 +1,72 @@
+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 exo.worker.disaggregated.protocol import Header, KVChunk
+from exo.worker.disaggregated.server import PrefillRequest
+from exo.worker.engines.mlx.cache import CacheSnapshot, snapshot_ssm_states
+from exo.worker.engines.mlx.disaggregated.client import (
+    ingest_into_mlx_cache,
+    remote_prefill_fetch,
+)
+from exo.worker.engines.mlx.types import KVCacheType
+from exo.worker.runner.bootstrap import logger
+
+
+def remote_prefill(
+    prompt_tokens: mx.array,
+    cache: KVCacheType,
+    on_prefill_progress: Callable[[int, int], None] | None,
+    *,
+    endpoint: str,
+    request_id: str,
+    model_id: str,
+    start_pos: int = 0,
+) -> tuple[float, int, list[CacheSnapshot]]:
+    t0 = time.perf_counter()
+    total_prompt_tokens = int(prompt_tokens.shape[0])
+    num_layers: int = 0
+
+    def _on_header(header: Header) -> None:
+        nonlocal num_layers
+        num_layers = header.num_layers
+
+    def _on_chunk(_chunk: KVChunk, chunks_received: int) -> None:
+        nonlocal num_layers
+        if on_prefill_progress is None:
+            return
+        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)]
diff --git a/src/exo/worker/engines/mlx/tests/test_batch_generate.py b/src/exo/worker/engines/mlx/tests/test_batch_generate.py
index 3d48590c..6904c18a 100644
--- a/src/exo/worker/engines/mlx/tests/test_batch_generate.py
+++ b/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
 
diff --git a/src/exo/shared/types/mlx.py b/src/exo/worker/engines/mlx/types.py
similarity index 100%
rename from src/exo/shared/types/mlx.py
rename to src/exo/worker/engines/mlx/types.py
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index a016bbf2..7021204a 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/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
 
 
@@ -676,7 +676,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
 
diff --git a/src/exo/worker/engines/mlx/vision.py b/src/exo/worker/engines/mlx/vision.py
index f8af3e97..a79925bc 100644
--- a/src/exo/worker/engines/mlx/vision.py
+++ b/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,
diff --git a/src/exo/worker/runner/llm_inference/batch_generator.py b/src/exo/worker/runner/llm_inference/batch_generator.py
index ee536a99..a58e6486 100644
--- a/src/exo/worker/runner/llm_inference/batch_generator.py
+++ b/src/exo/worker/runner/llm_inference/batch_generator.py
@@ -4,6 +4,7 @@ from abc import ABC, abstractmethod
 from collections import deque
 from collections.abc import Generator, Iterator
 from dataclasses import dataclass, field
+from typing import BinaryIO
 
 import mlx.core as mx
 from mlx_lm.tokenizer_utils import TokenizerWrapper
@@ -12,18 +13,21 @@ 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
 from exo.utils.channels import MpReceiver, MpSender
+from exo.worker.disaggregated.server import PrefillRequest
 from exo.worker.engines.mlx.cache import KVPrefixCache
+from exo.worker.engines.mlx.disaggregated.adapter import write_cache_to_wire
+from exo.worker.engines.mlx.disaggregated.serve import run_prefill_for_request
 from exo.worker.engines.mlx.generator.batch_generate import ExoBatchGenerator
 from exo.worker.engines.mlx.generator.generate import (
     PrefillCancelled,
     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,
@@ -85,6 +89,9 @@ class InferenceGenerator(ABC):
     @abstractmethod
     def close(self) -> None: ...
 
+    @abstractmethod
+    def serve_prefill(self, request: PrefillRequest, wfile: BinaryIO) -> None: ...
+
 
 EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
 EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM"
@@ -308,6 +315,22 @@ class SequentialGenerator(InferenceGenerator):
     def close(self) -> None:
         del self.model, self.tokenizer, self.group
 
+    def serve_prefill(self, request: PrefillRequest, wfile: BinaryIO) -> None:
+        cache = run_prefill_for_request(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            group=self.group,
+            kv_prefix_cache=self.kv_prefix_cache,
+            request=request,
+        )
+        write_cache_to_wire(
+            wfile,
+            cache,
+            request_id=request.request_id,
+            model_id=request.model_id,
+            start_pos=request.start_pos,
+        )
+
 
 @dataclass(eq=False)
 class BatchGenerator(InferenceGenerator):
@@ -534,3 +557,19 @@ class BatchGenerator(InferenceGenerator):
     def close(self) -> None:
         self._mlx_gen.close()
         del self.model, self.tokenizer, self.group
+
+    def serve_prefill(self, request: PrefillRequest, wfile: BinaryIO) -> None:
+        cache = run_prefill_for_request(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            group=self.group,
+            kv_prefix_cache=self.kv_prefix_cache,
+            request=request,
+        )
+        write_cache_to_wire(
+            wfile,
+            cache,
+            request_id=request.request_id,
+            model_id=request.model_id,
+            start_pos=request.start_pos,
+        )
diff --git a/src/exo/worker/runner/llm_inference/model_output_parsers.py b/src/exo/worker/runner/llm_inference/model_output_parsers.py
index 0e75a96d..4952688d 100644
--- a/src/exo/worker/runner/llm_inference/model_output_parsers.py
+++ b/src/exo/worker/runner/llm_inference/model_output_parsers.py
@@ -22,8 +22,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,
 )
diff --git a/src/exo/worker/runner/llm_inference/runner.py b/src/exo/worker/runner/llm_inference/runner.py
index 9cc2506a..20a6a37f 100644
--- a/src/exo/worker/runner/llm_inference/runner.py
+++ b/src/exo/worker/runner/llm_inference/runner.py
@@ -1,13 +1,17 @@
 import os
+import queue
+import threading
 import time
 from collections.abc import Generator
 from dataclasses import dataclass
 from enum import Enum
+from typing import BinaryIO
 
 import mlx.core as mx
-from anyio import WouldBlock
+from anyio import ClosedResourceError, EndOfStream
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
+from exo.shared.constants import ENABLE_DISAGGREGATION
 from exo.shared.models.model_cards import ModelTask
 from exo.shared.types.chunks import GenerationChunk
 from exo.shared.types.common import CommandId, ModelId
@@ -18,7 +22,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,7 +50,13 @@ from exo.shared.types.worker.runners import (
     RunnerWarmingUp,
 )
 from exo.utils.channels import MpReceiver, MpSender
+from exo.utils.ports import random_ephemeral_port
+from exo.worker.disaggregated.server import (
+    PrefillRequest,
+    PrefillServer,
+)
 from exo.worker.engines.mlx.cache import KVPrefixCache
+from exo.worker.engines.mlx.types import Model
 from exo.worker.engines.mlx.utils_mlx import (
     initialize_mlx,
     load_mlx_items,
@@ -63,6 +72,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 PrefillTask:
+    request: PrefillRequest
+    wfile: BinaryIO
+    started: threading.Event
+    done: threading.Event
+
+
+_TaskStreamClosed = object()
+WorkItem = Task | PrefillTask | object
+
 
 class ExitCode(str, Enum):
     AllTasksComplete = "AllTasksComplete"
@@ -110,9 +134,77 @@ 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 not ENABLE_DISAGGREGATION:
+            return None
+        if self.device_rank != 0:
+            return None
+        if self._prefill_server_port is not None:
+            return self._prefill_server_port
+
+        def resolve(request: PrefillRequest, wfile: BinaryIO) -> bool:
+            req = PrefillTask(
+                request=request,
+                wfile=wfile,
+                started=threading.Event(),
+                done=threading.Event(),
+            )
+            self._work_queue.put(req)
+            if not req.started.wait(timeout=PREFILL_PICKUP_TIMEOUT_SECONDS):
+                logger.warning(
+                    f"Prefill request {request.request_id} not picked up within "
+                    f"{PREFILL_PICKUP_TIMEOUT_SECONDS}s — runner busy"
+                )
+                return False
+            if not req.done.wait(timeout=PREFILL_FINISH_TIMEOUT_SECONDS):
+                logger.warning(
+                    f"Prefill request {request.request_id} did not finish within "
+                    f"{PREFILL_FINISH_TIMEOUT_SECONDS}s"
+                )
+            return True
+
+        port = random_ephemeral_port()
+        self._prefill_server = PrefillServer(resolve=resolve, host="0.0.0.0", port=port)
+        self._prefill_server_port = 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")
+        self._task_reader_thread.start()
+
+    def _serve_prefill(self, req: PrefillTask) -> None:
+        req.started.set()
+        try:
+            assert isinstance(self.generator, InferenceGenerator)
+            self.generator.serve_prefill(req.request, req.wfile)
+        except Exception:
+            logger.opt(exception=True).warning(
+                f"Failed to serve prefill request {req.request.request_id}"
+            )
+        finally:
+            req.done.set()
+
     def update_status(self, status: RunnerStatus):
         self.current_status = status
         self.event_sender.send(
@@ -130,8 +222,16 @@ 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:
+        self._start_task_reader()
+        try:
+            while True:
+                item = self._work_queue.get()
+                if item is _TaskStreamClosed:
+                    break
+                if isinstance(item, PrefillTask):
+                    self._serve_prefill(item)
+                    continue
+                task: Task = item  # type: ignore
                 if task.task_id in self.seen:
                     logger.warning("repeat task - potential error")
                     continue
@@ -139,6 +239,14 @@ class Runner:
                 self.handle_first_task(task)
                 if isinstance(self.current_status, RunnerShutdown):
                     break
+        finally:
+            if self._prefill_server is not None:
+                self._prefill_server.stop()
+                self._prefill_server = None
+            self.task_receiver.close()
+            if self._task_reader_thread is not None:
+                self._task_reader_thread.join(timeout=5)
+                self._task_reader_thread = None
 
     def handle_first_task(self, task: Task):
         self.send_task_status(task.task_id, TaskStatus.Running)
@@ -219,8 +327,11 @@ class Runner:
                     f"runner initialized in {time.time() - self.setup_start_time} seconds"
                 )
 
+                self._start_prefill_server()
                 self.send_task_status(task.task_id, TaskStatus.Complete)
-                self.update_status(RunnerReady())
+                self.update_status(
+                    RunnerReady(prefill_server_port=self._prefill_server_port)
+                )
                 logger.info("runner ready")
 
             case TextGeneration() if isinstance(self.current_status, RunnerReady):
@@ -285,29 +396,32 @@ class Runner:
                 self.active_tasks.pop(task_id, None)
 
             try:
-                task = self.task_receiver.receive_nowait()
-
-                if task.task_id in self.seen:
-                    logger.warning("repeat task - potential error")
-                    continue
-                self.seen.add(task.task_id)
-
-                match task:
-                    case TextGeneration():
-                        self.acknowledge_task(task)
-                        self.submit_text_generation(task)
-                    case Shutdown():
-                        self.shutdown(task)
-                        return ExitCode.Shutdown
-                    case _:
-                        raise ValueError(
-                            f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}"
-                        )
-
-            except WouldBlock:
-                pass
+                item = self._work_queue.get_nowait()
+            except queue.Empty:
+                continue
+            if item is _TaskStreamClosed:
+                return ExitCode.Shutdown
+            if isinstance(item, PrefillTask):
+                self._serve_prefill(item)
+                continue
+            task: Task = item  # type: ignore
+            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=}"
+                    )
 
-        self.update_status(RunnerReady())
+        self.update_status(RunnerReady(prefill_server_port=self._prefill_server_port))
         logger.info("runner ready")
 
         return ExitCode.AllTasksComplete
diff --git a/src/exo/worker/tests/unittests/test_mlx/conftest.py b/src/exo/worker/tests/unittests/test_mlx/conftest.py
index a4d71139..bba9f15a 100644
--- a/src/exo/worker/tests/unittests/test_mlx/conftest.py
+++ b/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
 
 
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py b/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py
index 73f8b1fa..3d72d47d 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py
+++ b/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,
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py b/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py
index 0a67f0c3..074a9f94 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py
+++ b/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,
diff --git a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
index a4133b1c..d1f4b65f 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
@@ -152,6 +152,11 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
     )
     monkeypatch.setattr(mlx_batch_generator, "ExoBatchGenerator", FakeExoBatchGenerator)
 
+    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:
diff --git a/src/exo/worker/tests/unittests/test_runner/test_serve_prefill.py b/src/exo/worker/tests/unittests/test_runner/test_serve_prefill.py
new file mode 100644
index 00000000..6c607dfa
--- /dev/null
+++ b/src/exo/worker/tests/unittests/test_runner/test_serve_prefill.py
@@ -0,0 +1,227 @@
+"""Tests for serve_prefill_request — the producer-side path that uses
+KVPrefixCache for cross-request prefix sharing."""
+
+import io
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, cast
+
+import mlx.core as mx
+import pytest
+from mlx_lm.models.cache import KVCache
+
+import exo.worker.engines.mlx.disaggregated.serve as mlx_serve_mod
+from exo.worker.disaggregated.protocol import (
+    Done,
+    Header,
+    KVChunk,
+    read_header,
+    read_message,
+)
+from exo.worker.disaggregated.server import PrefillRequest
+from exo.worker.engines.mlx.cache import KVPrefixCache
+from exo.worker.engines.mlx.disaggregated.adapter import write_cache_to_wire
+
+N_HEADS = 2
+HEAD_DIM = 4
+
+
+@dataclass
+class _FakeTokenizer:
+    has_thinking: bool = False
+    think_start: object | None = None
+    think_end: object | None = None
+
+
+class _FakeModel:
+    def __init__(self) -> None:
+        self.layers = [object()]
+
+    def make_cache(self) -> list[KVCache]:
+        return [KVCache() for _ in self.layers]
+
+
+def _populate_cache_in_place(cache: list[KVCache], n_tokens: int) -> None:
+    mx.random.seed(0)
+    for c in cache:
+        c.keys = (
+            mx.random.uniform(shape=(1, N_HEADS, n_tokens, HEAD_DIM)) * 10
+        ).astype(mx.bfloat16)
+        c.values = (
+            mx.random.uniform(shape=(1, N_HEADS, n_tokens, HEAD_DIM)) * 10
+        ).astype(mx.bfloat16)
+        c.offset = n_tokens
+
+
+def _patch_prefill(monkeypatch: pytest.MonkeyPatch) -> list[mx.array]:
+    """Replace mlx_prefill with a tracker; return the list it appends to."""
+    inputs: list[mx.array] = []
+
+    def fake_prefill(**kwargs: object) -> tuple[float, int, list[object]]:
+        pt = cast(mx.array, kwargs["prompt_tokens"])
+        inputs.append(pt)
+        n = int(pt.shape[0])
+        cache = cast(list[KVCache], kwargs["cache"])
+        existing = int(cache[0].offset) if cache and cache[0].keys is not None else 0
+        _populate_cache_in_place(cache, existing + n)
+        return (0.0, n, [])
+
+    def fake_make_sampler(**_: object) -> Callable[[mx.array], mx.array]:
+        return lambda x: x
+
+    monkeypatch.setattr(mlx_serve_mod, "mlx_prefill", fake_prefill)
+    monkeypatch.setattr(mlx_serve_mod, "make_sampler", fake_make_sampler)
+    return inputs
+
+
+def _decode(payload: bytes) -> tuple[Header, list[KVChunk], int]:
+    buf = io.BytesIO(payload)
+    hdr = read_header(buf)
+    chunks: list[KVChunk] = []
+    total = 0
+    while True:
+        msg = read_message(buf)
+        if msg is None:
+            break
+        if isinstance(msg, KVChunk):
+            chunks.append(msg)
+        elif isinstance(msg, Done):
+            total = msg.total_tokens
+            break
+    return hdr, chunks, total
+
+
+def _serve(request: PrefillRequest, kv_prefix_cache: KVPrefixCache | None) -> bytes:
+    cache = mlx_serve_mod.run_prefill_for_request(
+        model=cast(Any, _FakeModel()),  # pyright: ignore[reportAny]
+        tokenizer=cast(Any, _FakeTokenizer()),  # pyright: ignore[reportAny]
+        group=None,
+        kv_prefix_cache=kv_prefix_cache,
+        request=request,
+    )
+    buf = io.BytesIO()
+    write_cache_to_wire(
+        buf,
+        cache,
+        request_id=request.request_id,
+        model_id=request.model_id,
+        start_pos=request.start_pos,
+    )
+    return buf.getvalue()
+
+
+def test_serve_prefill_runs_full_prefill_when_cache_empty(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    inputs = _patch_prefill(monkeypatch)
+    cache = KVPrefixCache(group=None)
+
+    payload = _serve(
+        PrefillRequest(
+            request_id="r1", model_id="m", token_ids=list(range(20)), start_pos=0
+        ),
+        cache,
+    )
+
+    assert len(inputs) == 1
+    assert int(inputs[0].shape[0]) == 18
+
+    hdr, chunks, total = _decode(payload)
+    assert hdr.start_pos == 0
+    assert total == 18
+    assert len(chunks) == 1
+    assert chunks[0].num_tokens == 18
+
+    assert len(cache.prompts) == 1
+    assert int(cache.prompts[0].shape[0]) == 20
+
+
+def test_serve_prefill_skips_work_on_exact_cache_hit(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    inputs = _patch_prefill(monkeypatch)
+    cache = KVPrefixCache(group=None)
+    tokens = list(range(20))
+
+    _serve(
+        PrefillRequest(request_id="r1", model_id="m", token_ids=tokens, start_pos=0),
+        cache,
+    )
+    inputs.clear()
+
+    _serve(
+        PrefillRequest(request_id="r2", model_id="m", token_ids=tokens, start_pos=0),
+        cache,
+    )
+
+    assert inputs == []
+
+
+def test_serve_prefill_only_runs_suffix_on_partial_match(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    inputs = _patch_prefill(monkeypatch)
+    cache = KVPrefixCache(group=None)
+
+    base = list(range(15))
+    _serve(
+        PrefillRequest(request_id="r1", model_id="m", token_ids=base, start_pos=0),
+        cache,
+    )
+    inputs.clear()
+
+    extended = base + [99, 100, 101, 102, 103]
+    _serve(
+        PrefillRequest(request_id="r2", model_id="m", token_ids=extended, start_pos=0),
+        cache,
+    )
+
+    assert len(inputs) == 1
+    suffix_len = int(inputs[0].shape[0])
+    assert suffix_len < len(extended) - 2
+
+
+def test_serve_prefill_slices_payload_at_client_start_pos(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _patch_prefill(monkeypatch)
+    cache = KVPrefixCache(group=None)
+
+    n_tokens = 20
+    client_has = 12
+
+    payload = _serve(
+        PrefillRequest(
+            request_id="r1",
+            model_id="m",
+            token_ids=list(range(n_tokens)),
+            start_pos=client_has,
+        ),
+        cache,
+    )
+
+    hdr, chunks, total = _decode(payload)
+    assert hdr.start_pos == client_has
+    expected_sent = (n_tokens - 2) - client_has
+    assert total == expected_sent
+    assert len(chunks) == 1
+    assert chunks[0].num_tokens == expected_sent
+
+
+def test_serve_prefill_works_without_prefix_cache(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    inputs = _patch_prefill(monkeypatch)
+
+    payload = _serve(
+        PrefillRequest(
+            request_id="r1", model_id="m", token_ids=list(range(20)), start_pos=0
+        ),
+        None,
+    )
+
+    assert len(inputs) == 1
+    assert int(inputs[0].shape[0]) == 18
+
+    _, _, total = _decode(payload)
+    assert total == 18

← 5d10188d fix: route by in-flight tasks only — completed tasks were sk  ·  back to Exo  ·  fix: uninstall-exo.sh removes both current and legacy bridge 18ffe1df →