Fix bench

bench/prefill_decode_bench.py

@@ -10,9 +10,10 @@ 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
+    uv run bench/prefill_decode_bench.py --model <id> --pp 2048,8192 --tg 128
+    uv run bench/prefill_decode_bench.py --model <id> --pp 4096 --tg 128 --repeat 3
+    uv run bench/prefill_decode_bench.py --model <id> --pp 2048 --tg 128 --dry-run
+    uv run python bench/prefill_decode_bench.py --config bench/prefill-decode.toml
 """
 
 from __future__ import annotations

src/exo/api/main.py

@@ -124,7 +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,
+    DISAGGREGATION_MODE,
     EXO_CACHE_HOME,
     EXO_EVENT_LOG_DIR,
     EXO_IMAGE_CACHE_DIR,
@@ -222,12 +222,12 @@ def _ensure_seed(params: AdvancedImageParams | None) -> AdvancedImageParams:
 
 
 def _require_disaggregation_enabled() -> None:
-    if not ENABLE_DISAGGREGATION:
+    if DISAGGREGATION_MODE == 0:
         raise HTTPException(
             status_code=HTTPStatus.NOT_FOUND,
             detail=(
                 "Prefill/decode disaggregation is disabled. "
-                "Set ENABLE_DISAGGREGATION=true to enable."
+                "Set DISAGGREGATION_MODE=1 (non-overlapping) or 2 (overlapping) to enable."
             ),
         )
 
@@ -640,10 +640,10 @@ class API:
         )
 
     async def get_feature_flags(self) -> dict[str, bool]:
-        return {"disaggregation": ENABLE_DISAGGREGATION}
+        return {"disaggregation": DISAGGREGATION_MODE != 0}
 
     async def list_instance_links(self) -> list[InstanceLink]:
-        if not ENABLE_DISAGGREGATION:
+        if DISAGGREGATION_MODE == 0:
             return []
         return list(self.state.instance_links.values())
 

src/exo/shared/constants.py

@@ -96,7 +96,11 @@ 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"
+DISAGGREGATION_MODE = int(os.getenv("DISAGGREGATION_MODE", "0"))
+if DISAGGREGATION_MODE not in (0, 1, 2):
+    raise ValueError(
+        f"DISAGGREGATION_MODE must be 0 (off), 1 (non-overlapping), or 2 (overlapping); got {DISAGGREGATION_MODE}"
+    )
 
 EXO_MAX_CONCURRENT_REQUESTS = int(os.getenv("EXO_MAX_CONCURRENT_REQUESTS", "8"))
 

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

@@ -22,7 +22,6 @@ from exo.worker.disaggregated.protocol import (
     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,
@@ -90,15 +89,22 @@ def nhd_to_bhsd(t: mx.array) -> mx.array:
     return mx.expand_dims(mx.transpose(t, (1, 0, 2)), 0)
 
 
-def send_mlx_kv_cache(
-    stream: BinaryIO,
+def materialize_kv_token_range(
     caches: KVCacheType,
     *,
-    dtype: DType,
-    start_pos: int = 0,
-    max_tokens: int | None = None,
-) -> int:
-    tokens_sent = 0
+    token_start: int,
+    token_end: int,
+) -> list[tuple[int, mx.array, mx.array]]:
+    """Slice [token_start, token_end) from each KV layer and materialize on CPU.
+
+    Returns a list of (layer_idx, k_nhd, v_nhd) tuples ready for serialization.
+    Calls mx.eval to ensure the slice is a snapshot (cache.keys is mutated in
+    place by subsequent prefill steps, so the slice must be evaluated before
+    returning to keep correctness).
+    """
+    if token_end <= token_start:
+        return []
+    out: list[tuple[int, mx.array, mx.array]] = []
     for layer_idx, c in enumerate(caches):
         match c:
             case QuantizedKVCache() | CacheList() | DeepseekV4Cache():
@@ -108,52 +114,104 @@ def send_mlx_kv_cache(
                 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:
+                end = min(token_end, int(c.offset))
+                if end <= token_start:
                     continue
                 with mx.stream(mx.Device(mx.cpu)):
-                    k = mx.array(keys[:, :, start_pos:offset, :])
-                    v = mx.array(values[:, :, start_pos:offset, :])
+                    k = mx.array(keys[:, :, token_start:end, :])
+                    v = mx.array(values[:, :, token_start:end, :])
                     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
+                if int(k_nhd.shape[0]) == 0:
+                    continue
+                out.append((layer_idx, k_nhd, v_nhd))
             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)
+                pass
+    return out
+
+
+def write_materialized_kv_chunk(
+    stream: BinaryIO,
+    *,
+    layer_idx: int,
+    k_nhd: mx.array,
+    v_nhd: mx.array,
+    dtype: DType,
+) -> int:
+    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),
+    )
+    return num_tokens
+
+
+def send_kv_token_range(
+    stream: BinaryIO,
+    caches: KVCacheType,
+    *,
+    dtype: DType,
+    token_start: int,
+    token_end: int,
+) -> int:
+    materialized = materialize_kv_token_range(
+        caches, token_start=token_start, token_end=token_end
+    )
+    tokens_sent = 0
+    for layer_idx, k_nhd, v_nhd in materialized:
+        n = write_materialized_kv_chunk(
+            stream, layer_idx=layer_idx, k_nhd=k_nhd, v_nhd=v_nhd, dtype=dtype
+        )
+        tokens_sent = max(tokens_sent, n)
+    return tokens_sent
+
+
+def send_arrays_states(stream: BinaryIO, caches: KVCacheType) -> None:
+    for layer_idx, c in enumerate(caches):
+        if not isinstance(c, ArraysCache):
+            continue
+        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)
+
+
+def send_mlx_kv_cache(
+    stream: BinaryIO,
+    caches: KVCacheType,
+    *,
+    dtype: DType,
+    start_pos: int = 0,
+    max_tokens: int | None = None,
+) -> int:
+    upper = max((int(c.offset) for c in caches if hasattr(c, "offset")), default=0)
+    if max_tokens is not None:
+        upper = min(upper, max_tokens)
+    tokens_sent = send_kv_token_range(
+        stream, caches, dtype=dtype, token_start=start_pos, token_end=upper
+    )
+    send_arrays_states(stream, caches)
     return tokens_sent
 
 
@@ -208,14 +266,14 @@ 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(
+def write_prefill_header(
     wfile: BinaryIO,
     cache: KVCacheType,
     *,
     request_id: str = "",
     model_id: str = "",
     start_pos: int = 0,
-) -> int:
+) -> DType:
     dtype = wire_dtype_from_cache(cache)
     write_header(
         wfile,
@@ -227,6 +285,47 @@ def write_cache_to_wire(
             start_pos=start_pos,
         ),
     )
+    return dtype
+
+
+def write_prefill_step(
+    wfile: BinaryIO,
+    cache: KVCacheType,
+    *,
+    dtype: DType,
+    token_start: int,
+    token_end: int,
+) -> int:
+    tokens_sent = send_kv_token_range(
+        wfile, cache, dtype=dtype, token_start=token_start, token_end=token_end
+    )
+    if tokens_sent > 0:
+        wfile.flush()
+    return tokens_sent
+
+
+def write_prefill_done(
+    wfile: BinaryIO,
+    cache: KVCacheType,
+    *,
+    total_tokens: int,
+) -> None:
+    send_arrays_states(wfile, cache)
+    write_done(wfile, total_tokens)
+    wfile.flush()
+
+
+def write_cache_to_wire(
+    wfile: BinaryIO,
+    cache: KVCacheType,
+    *,
+    request_id: str = "",
+    model_id: str = "",
+    start_pos: int = 0,
+) -> int:
+    dtype = write_prefill_header(
+        wfile, cache, request_id=request_id, model_id=model_id, 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()

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

@@ -23,6 +23,7 @@ from exo.worker.engines.mlx.disaggregated.adapter import (
     inject_arrays_cache,
     inject_kv_chunk,
     inject_rotating_kv_chunk,
+    nhd_to_bhsd,
 )
 
 _SOCKET_TIMEOUT_SECS = 60
@@ -103,6 +104,121 @@ def remote_prefill_fetch(
         sock.close()
 
 
+def remote_prefill_stream(
+    endpoint: str,
+    request: PrefillRequest,
+    caches: list[KVCache | RotatingKVCache | ArraysCache],
+    *,
+    on_header: Callable[[Header], None] | None = None,
+    on_kv_chunk: Callable[[KVChunk, int], None] | None = None,
+    timeout_secs: float = _SOCKET_TIMEOUT_SECS,
+) -> tuple[Header, int]:
+    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)
+
+        first_seen: set[int] = set()
+        rotating_chunks: dict[int, list[KVChunk]] = defaultdict(list)
+        arrays: dict[int, list[TensorBlob]] = {}
+        per_layer_tokens: dict[int, int] = defaultdict(int)
+        chunks_received = 0
+
+        while True:
+            msg = read_message(stream)
+            if msg is None:
+                break
+            if isinstance(msg, KVChunk):
+                target = caches[msg.layer_idx]
+                if isinstance(target, KVCache):
+                    _stream_inject_kv(
+                        target, msg, first_seen=first_seen, start_pos=request.start_pos
+                    )
+                elif isinstance(target, RotatingKVCache):
+                    rotating_chunks[msg.layer_idx].append(msg)
+                per_layer_tokens[msg.layer_idx] += msg.num_tokens
+                chunks_received += 1
+                if on_kv_chunk is not None:
+                    on_kv_chunk(msg, chunks_received)
+            elif isinstance(msg, ArraysState):
+                arrays[msg.layer_idx] = msg.arrays
+            elif isinstance(msg, Done):
+                break
+            else:
+                raise RuntimeError(f"Prefill server error [{msg.code}]: {msg.message}")
+
+        max_received = max(per_layer_tokens.values(), default=0)
+        final_offset = request.start_pos + max_received
+
+        for i, cache in enumerate(caches):
+            if isinstance(cache, KVCache):
+                if i in first_seen:
+                    cache.offset = final_offset
+            elif isinstance(cache, RotatingKVCache) and i in rotating_chunks:
+                chunks = rotating_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)
+                inject_rotating_kv_chunk(cache, k_nhd, v_nhd, final_offset)
+            if isinstance(cache, ArraysCache) and i in arrays:
+                inject_arrays_cache(cache, arrays[i])
+
+        return header, final_offset
+    finally:
+        sock.close()
+
+
+def _stream_inject_kv(
+    cache: KVCache,
+    chunk: KVChunk,
+    *,
+    first_seen: set[int],
+    start_pos: int,
+) -> None:
+    k_nhd, v_nhd = chunk_to_mlx_nhd(chunk)
+    k_bhsd = nhd_to_bhsd(k_nhd)
+    v_bhsd = nhd_to_bhsd(v_nhd)
+    if chunk.layer_idx not in first_seen:
+        first_seen.add(chunk.layer_idx)
+        existing_k = cache.keys
+        existing_v = cache.values
+        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
+    else:
+        existing_k = cache.keys
+        existing_v = cache.values
+        assert existing_k is not None and existing_v is not None
+        cache.keys = mx.concatenate([existing_k, k_bhsd], axis=2)
+        cache.values = mx.concatenate([existing_v, v_bhsd], axis=2)
+
+
 def ingest_into_mlx_cache(
     result: PrefillResult,
     caches: list[KVCache | RotatingKVCache | ArraysCache],

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

@@ -1,4 +1,5 @@
 import time
+from collections.abc import Callable
 
 import mlx.core as mx
 from mlx_lm.sample_utils import make_sampler
@@ -16,6 +17,8 @@ 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
 
+OnPrefillStep = Callable[[int, KVCacheType], None]
+
 
 def run_prefill_for_request(
     *,
@@ -24,6 +27,7 @@ def run_prefill_for_request(
     group: mx.distributed.Group | None,
     kv_prefix_cache: KVPrefixCache | None,
     request: PrefillRequest,
+    on_step: OnPrefillStep | None = None,
 ) -> KVCacheType:
     prompt_tokens = mx.array(request.token_ids)
     prompt_tokens = fix_unmatched_think_end_tokens(prompt_tokens, tokenizer)
@@ -46,6 +50,18 @@ def run_prefill_for_request(
     prefill_input = remaining[:new_tokens]
     if int(prefill_input.shape[0]) > 0:
         sampler = make_sampler(temp=1.0)
+        on_prefill_progress: Callable[[int, int], None] | None = None
+        if on_step is not None:
+            captured_cache = cache
+            last_seen = [cache_length(cache)]
+
+            def _step(_processed: int, _total: int) -> None:
+                cur = cache_length(captured_cache)
+                if cur > last_seen[0]:
+                    on_step(cur, captured_cache)
+                    last_seen[0] = cur
+
+            on_prefill_progress = _step
         _ = mlx_prefill(
             model=model,
             tokenizer=tokenizer,
@@ -53,7 +69,7 @@ def run_prefill_for_request(
             prompt_tokens=prefill_input,
             cache=cache,
             group=group,
-            on_prefill_progress=None,
+            on_prefill_progress=on_prefill_progress,
             distributed_prompt_progress_callback=None,
         )
 

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

@@ -10,10 +10,14 @@ from exo.worker.disaggregated.server import PrefillRequest, PrefillServer
 from exo.worker.engines.mlx.disaggregated.adapter import (
     send_mlx_kv_cache,
     wire_dtype_from_cache,
+    write_prefill_done,
+    write_prefill_header,
+    write_prefill_step,
 )
 from exo.worker.engines.mlx.disaggregated.client import (
     ingest_into_mlx_cache,
     remote_prefill_fetch,
+    remote_prefill_stream,
 )
 
 
@@ -121,6 +125,116 @@ def test_server_reports_pickup_failure() -> None:
         server.stop()
 
 
+def _stream_cache_in_steps(
+    wfile: BinaryIO,
+    cache: KVCache,
+    *,
+    request_id: str,
+    step_size: int,
+    start_pos: int = 0,
+) -> None:
+    dtype = write_prefill_header(
+        wfile,
+        [cache],
+        request_id=request_id,
+        model_id="test-model",
+        start_pos=start_pos,
+    )
+    cache_offset = int(cache.offset)
+    cur = max(start_pos, 0)
+    total = 0
+    while cur < cache_offset:
+        nxt = min(cur + step_size, cache_offset)
+        sent = write_prefill_step(
+            wfile, [cache], dtype=dtype, token_start=cur, token_end=nxt
+        )
+        total += sent
+        cur = nxt
+    write_prefill_done(wfile, [cache], total_tokens=total)
+
+
+@pytest.mark.slow
+def test_streaming_roundtrip_matches_one_shot() -> None:
+    seq_len = 12
+    n_heads = 2
+    head_dim = 4
+    gold = _make_cache(seq_len, n_heads, head_dim)
+
+    def resolve(job: PrefillRequest, wfile: BinaryIO) -> bool:
+        _stream_cache_in_steps(wfile, gold, request_id=job.request_id, step_size=4)
+        return True
+
+    server = PrefillServer(resolve=resolve, host="127.0.0.1", port=52420)
+    try:
+        dst = KVCache()
+        _, final_offset = remote_prefill_stream(
+            endpoint="127.0.0.1:52420",
+            request=PrefillRequest(
+                model_id="test-model",
+                token_ids=list(range(seq_len)),
+                request_id="req-stream",
+            ),
+            caches=[dst],
+        )
+        assert final_offset == seq_len
+        assert dst.offset == seq_len
+        gold_k = gold.keys
+        gold_v = gold.values
+        dst_k = dst.keys
+        dst_v = dst.values
+        assert gold_k is not None and gold_v is not None
+        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()
+
+
+@pytest.mark.slow
+def test_streaming_roundtrip_with_start_pos() -> None:
+    seq_len = 12
+    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_in_steps(
+            wfile, gold, request_id=job.request_id, step_size=3, start_pos=start_pos
+        )
+        return True
+
+    server = PrefillServer(resolve=resolve, host="127.0.0.1", port=52421)
+    try:
+        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 = remote_prefill_stream(
+            endpoint="127.0.0.1:52421",
+            request=PrefillRequest(
+                model_id="test-model",
+                token_ids=list(range(seq_len)),
+                request_id="req-stream-prefix",
+                start_pos=start_pos,
+            ),
+            caches=[dst],
+        )
+        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()
+
+
 @pytest.mark.slow
 def test_server_client_roundtrip_with_start_pos() -> None:
     seq_len = 8

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

@@ -8,10 +8,7 @@ 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.disaggregated.client import remote_prefill_stream
 from exo.worker.engines.mlx.types import KVCacheType
 from exo.worker.runner.bootstrap import logger
 
@@ -51,13 +48,14 @@ def remote_prefill(
         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)
+    _, final_offset = remote_prefill_stream(
+        endpoint,
+        request,
+        caches,
+        on_header=_on_header,
+        on_kv_chunk=_on_chunk,
+    )
     t_done = time.perf_counter()
 
     num_tokens = final_offset - start_pos
@@ -66,7 +64,6 @@ def remote_prefill(
     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"
+        f"elapsed={(t_done - t0) * 1000:.0f}ms"
     )
     return tps, num_tokens, [snapshot_ssm_states(cache)]

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

@@ -8,7 +8,10 @@ from typing import BinaryIO
 import mlx.core as mx
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
-from exo.shared.constants import EXO_MAX_CONCURRENT_REQUESTS
+from exo.shared.constants import (
+    DISAGGREGATION_MODE,
+    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
@@ -28,7 +31,12 @@ from exo.utils.channels import MpReceiver, MpSender
 from exo.worker.disaggregated.server import PrefillRequest
 from exo.worker.engines.base import Engine
 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.adapter import (
+    write_cache_to_wire,
+    write_prefill_done,
+    write_prefill_header,
+    write_prefill_step,
+)
 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 (
@@ -36,7 +44,7 @@ from exo.worker.engines.mlx.generator.generate import (
     mlx_generate,
     warmup_inference,
 )
-from exo.worker.engines.mlx.types import Model
+from exo.worker.engines.mlx.types import KVCacheType, Model
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
     mx_all_gather_tasks,
@@ -294,14 +302,76 @@ class SequentialGenerator(Engine):
         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(
+        if DISAGGREGATION_MODE == 2:
+            stream_prefill(
+                wfile=wfile,
+                request=request,
+                model=self.model,
+                tokenizer=self.tokenizer,
+                group=self.group,
+                kv_prefix_cache=self.kv_prefix_cache,
+            )
+        else:
+            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,
+            )
+
+
+def stream_prefill(
+    *,
+    wfile: BinaryIO,
+    request: PrefillRequest,
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    group: mx.distributed.Group | None,
+    kv_prefix_cache: KVPrefixCache | None,
+) -> None:
+    streamed = [request.start_pos]
+    header_state: dict[str, object] = {"dtype": None}
+
+    def on_step(cur: int, cache: KVCacheType) -> None:
+        if header_state["dtype"] is None:
+            header_state["dtype"] = write_prefill_header(
+                wfile,
+                cache,
+                request_id=request.request_id,
+                model_id=request.model_id,
+                start_pos=request.start_pos,
+            )
+        token_start = streamed[0]
+        if cur > token_start:
+            written = write_prefill_step(
+                wfile,
+                cache,
+                dtype=header_state["dtype"],  # type: ignore
+                token_start=token_start,
+                token_end=cur,
+            )
+            if written > 0:
+                streamed[0] = cur
+
+    cache = run_prefill_for_request(
+        model=model,
+        tokenizer=tokenizer,
+        group=group,
+        kv_prefix_cache=kv_prefix_cache,
+        request=request,
+        on_step=on_step,
+    )
+
+    if header_state["dtype"] is None:
+        header_state["dtype"] = write_prefill_header(
             wfile,
             cache,
             request_id=request.request_id,
@@ -309,6 +379,24 @@ class SequentialGenerator(Engine):
             start_pos=request.start_pos,
         )
 
+    final_offset = max(
+        (int(c.offset) for c in cache if hasattr(c, "offset")),
+        default=0,
+    )
+    if final_offset > streamed[0]:
+        written = write_prefill_step(
+            wfile,
+            cache,
+            dtype=header_state["dtype"],  # type: ignore
+            token_start=streamed[0],
+            token_end=final_offset,
+        )
+        if written > 0:
+            streamed[0] = final_offset
+
+    total_tokens = max(0, streamed[0] - request.start_pos)
+    write_prefill_done(wfile, cache, total_tokens=total_tokens)
+
 
 @dataclass(eq=False)
 class BatchGenerator(Engine):
@@ -542,17 +630,27 @@ class BatchGenerator(Engine):
         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,
-        )
+        if DISAGGREGATION_MODE == 2:
+            stream_prefill(
+                wfile=wfile,
+                request=request,
+                model=self.model,
+                tokenizer=self.tokenizer,
+                group=self.group,
+                kv_prefix_cache=self.kv_prefix_cache,
+            )
+        else:
+            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,
+            )

src/exo/worker/runner/runner.py

@@ -7,7 +7,7 @@ from typing import BinaryIO
 
 from anyio import ClosedResourceError, EndOfStream
 
-from exo.shared.constants import ENABLE_DISAGGREGATION
+from exo.shared.constants import DISAGGREGATION_MODE
 from exo.shared.types.chunks import Chunk
 from exo.shared.types.common import CommandId
 from exo.shared.types.events import (
@@ -126,7 +126,7 @@ class Runner:
         self.update_status(RunnerIdle())
 
     def _start_prefill_server(self) -> int | None:
-        if not ENABLE_DISAGGREGATION:
+        if DISAGGREGATION_MODE == 0:
             return None
         if self.device_rank != 0:
             return None

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

@@ -208,6 +208,62 @@ def test_serve_prefill_slices_payload_at_client_start_pos(
     assert chunks[0].num_tokens == expected_sent
 
 
+def test_run_prefill_streams_per_step_when_on_step_provided(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Verify on_step fires per progress callback with correct offsets."""
+    step_size = 4
+
+    def fake_prefill(**kwargs: object) -> tuple[float, int, list[object]]:
+        pt = cast(mx.array, kwargs["prompt_tokens"])
+        cache = cast(list[KVCache], kwargs["cache"])
+        existing = int(cache[0].offset) if cache and cache[0].keys is not None else 0
+        n = int(pt.shape[0])
+        cb = cast(Callable[[int, int], None] | None, kwargs.get("on_prefill_progress"))
+        processed = 0
+        while processed < n:
+            advance = min(step_size, n - processed)
+            processed += advance
+            _populate_cache_in_place(cache, existing + processed)
+            if cb is not None:
+                cb(processed, 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)
+
+    from exo.worker.engines.mlx.types import KVCacheType
+
+    seen: list[tuple[int, int]] = []
+
+    def on_step(cur: int, ks: KVCacheType) -> None:
+        seen.append((cur, len(ks)))
+
+    n_tokens = 14
+    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=None,
+        request=PrefillRequest(
+            request_id="r-stream",
+            model_id="m",
+            token_ids=list(range(n_tokens)),
+            start_pos=0,
+        ),
+        on_step=on_step,
+    )
+
+    assert len(seen) >= 2
+    assert seen == sorted(seen)
+    final_offset = seen[-1][0]
+    assert final_offset == n_tokens - 2
+    assert int(cache[0].offset) == n_tokens - 2
+
+
 def test_serve_prefill_works_without_prefix_cache(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
@@ -225,3 +281,53 @@ def test_serve_prefill_works_without_prefix_cache(
 
     _, _, total = _decode(payload)
     assert total == 18
+
+
+def test_stream_prefill_emits_multiple_chunks_when_overlapping(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    import io as _io
+
+    import exo.worker.runner.llm_inference.batch_generator as bg
+
+    step_size = 4
+
+    def fake_prefill(**kwargs: object) -> tuple[float, int, list[object]]:
+        pt = cast(mx.array, kwargs["prompt_tokens"])
+        cache = cast(list[KVCache], kwargs["cache"])
+        existing = int(cache[0].offset) if cache and cache[0].keys is not None else 0
+        n = int(pt.shape[0])
+        cb = cast(Callable[[int, int], None] | None, kwargs.get("on_prefill_progress"))
+        processed = 0
+        while processed < n:
+            advance = min(step_size, n - processed)
+            processed += advance
+            _populate_cache_in_place(cache, existing + processed)
+            if cb is not None:
+                cb(processed, 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)
+
+    n_tokens = 14
+    request = PrefillRequest(
+        request_id="r", model_id="m", token_ids=list(range(n_tokens)), start_pos=0
+    )
+
+    buf = _io.BytesIO()
+    bg.stream_prefill(
+        wfile=buf,
+        request=request,
+        model=cast(Any, _FakeModel()),  # pyright: ignore[reportAny]
+        tokenizer=cast(Any, _FakeTokenizer()),  # pyright: ignore[reportAny]
+        group=None,
+        kv_prefix_cache=None,
+    )
+
+    _, chunks, total = _decode(buf.getvalue())
+    assert total == n_tokens - 2
+    assert len(chunks) >= 2