Use HND layout

bench/harness.py

@@ -69,6 +69,10 @@ class ExoClient:
     def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
         return self.request_json("POST", "/bench/chat/completions", body=payload)
 
+    def post_bench_disaggregated(self, payload: dict[str, Any]) -> dict[str, Any]:
+        payload["disaggregated"] = True
+        return self.request_json("POST", "/bench/chat/completions", body=payload)
+
 
 def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]:
     if len(instance) != 1:

bench/prefill_decode_bench.py

@@ -0,0 +1,399 @@
+# type: ignore
+#!/usr/bin/env python3
+"""Disaggregated prefill-decode benchmark for exo.
+
+Measures throughput when a vLLM node handles prefill (KV generation + transfer)
+and an MLX node handles decode (token generation).
+
+Requires a cluster with at least one CUDA/vLLM node and one Apple Silicon/MLX node.
+
+Usage:
+    uv run python prefill_decode_bench.py --model <decode-model> --pp 2048,8192 --tg 128
+    uv run python prefill_decode_bench.py --model <decode-model> --prefill-model <vllm-model> --pp 4096 --tg 128
+    uv run python prefill_decode_bench.py --model <model-id> --pp 2048 --tg 128 --no-overlapping
+    uv run python prefill_decode_bench.py --model <model-id> --pp 2048 --tg 128 --dry-run
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import itertools
+import json
+import sys
+import time
+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,
+    nodes_used_in_instance,
+    resolve_model_short_id,
+    run_planning_phase,
+    settle_and_fetch_placements,
+    wait_for_instance_gone,
+    wait_for_instance_ready,
+)
+from loguru import logger
+
+
+def find_vllm_placement(placements: list[dict[str, Any]]) -> dict[str, Any] | None:
+    for p in placements:
+        if "vllm" in str(p.get("instance_meta", "")).lower():
+            return p
+    return None
+
+
+def find_mlx_placement(
+    placements: list[dict[str, Any]], decode_meta: str
+) -> dict[str, Any] | None:
+    for p in placements:
+        meta = str(p.get("instance_meta", "")).lower()
+        if decode_meta == "both":
+            if "ring" in meta or "jaccl" in meta:
+                return p
+        elif decode_meta in meta:
+            return p
+    return None
+
+
+def run_one_disaggregated(
+    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_disaggregated(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 main() -> int:
+    ap = argparse.ArgumentParser(
+        prog="prefill-decode-bench",
+        description="Benchmark disaggregated prefill-decode (vLLM prefill → MLX decode).",
+    )
+    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 per placement (uses first pp/tg).",
+    )
+    ap.add_argument(
+        "--prefill-model",
+        default=None,
+        help="Model for the vLLM prefill node (defaults to --model if not set).",
+    )
+    ap.add_argument(
+        "--no-overlapping",
+        action="store_true",
+        help="Use batch KV transfer instead of streaming (overlapping).",
+    )
+    ap.add_argument(
+        "--decode-instance-meta",
+        choices=["ring", "jaccl", "both"],
+        default="ring",
+        help="Instance meta for the decode (MLX) node.",
+    )
+    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 (cartesian product) even when lists have equal length.",
+    )
+    args = ap.parse_args()
+
+    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 requires >1000 uncached tokens)")
+            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")
+
+    overlapping = not args.no_overlapping
+    logger.info(f"KV transfer mode: {'overlapping (streaming)' if overlapping else 'batch'}")
+
+    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
+    )
+
+    prefill_model_arg = args.prefill_model or args.model
+    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] tokenizer load failed")
+
+    try:
+        prompt_sizer = PromptSizer(tokenizer)
+    except Exception:
+        logger.error("[prefill-decode-bench] tokenizer usable but prompt sizing failed")
+        raise
+
+    original_instance_meta = args.instance_meta
+    original_max_nodes = args.max_nodes
+
+    args.instance_meta = "vllm"
+    args.min_nodes = 1
+    args.max_nodes = 1
+    vllm_placements = settle_and_fetch_placements(
+        client, prefill_full_id, args, settle_timeout=args.settle_timeout
+    )
+
+    args.instance_meta = args.decode_instance_meta
+    args.min_nodes = 1
+    args.max_nodes = original_max_nodes
+    mlx_placements = settle_and_fetch_placements(
+        client, decode_full_id, args, settle_timeout=args.settle_timeout
+    )
+
+    args.instance_meta = original_instance_meta
+
+    vllm_placement = find_vllm_placement(vllm_placements)
+    mlx_placement = find_mlx_placement(mlx_placements, args.decode_instance_meta)
+
+    if not vllm_placement:
+        logger.error("No vLLM placement found. Need a CUDA node for prefill.")
+        return 1
+    if not mlx_placement:
+        logger.error(f"No MLX ({args.decode_instance_meta}) placement found for decode.")
+        return 1
+
+    vllm_instance = vllm_placement["instance"]
+    mlx_instance = mlx_placement["instance"]
+    vllm_instance_id = instance_id_from_instance(vllm_instance)
+    mlx_instance_id = instance_id_from_instance(mlx_instance)
+    vllm_meta = str(vllm_placement.get("instance_meta", ""))
+    mlx_meta = str(mlx_placement.get("instance_meta", ""))
+    mlx_sharding = str(mlx_placement.get("sharding", ""))
+    mlx_nodes = nodes_used_in_instance(mlx_instance)
+
+    logger.info("=" * 80)
+    logger.info(f"PREFILL (vLLM):  {vllm_meta} / {prefill_short_id} ({prefill_full_id}) / instance_id={vllm_instance_id}")
+    logger.info(f"DECODE  (MLX):   {mlx_meta} / {mlx_sharding} / nodes={mlx_nodes} / {decode_short_id} ({decode_full_id}) / instance_id={mlx_instance_id}")
+    logger.info(f"Overlapping: {overlapping}")
+
+    if args.dry_run:
+        return 0
+
+    settle_deadline = (
+        time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
+    )
+
+    logger.info("Planning phase: checking downloads...")
+    download_duration_s = run_planning_phase(
+        client,
+        prefill_full_id,
+        vllm_placement,
+        args.danger_delete_downloads,
+        args.timeout,
+        settle_deadline,
+    )
+    download_duration_mlx = run_planning_phase(
+        client,
+        decode_full_id,
+        mlx_placement,
+        args.danger_delete_downloads,
+        args.timeout,
+        settle_deadline,
+    )
+    if download_duration_s is not None:
+        logger.info(f"Download (vLLM): {download_duration_s:.1f}s")
+    if download_duration_mlx is not None:
+        logger.info(f"Download (MLX): {download_duration_mlx:.1f}s")
+
+    # Create vLLM instance first (starts prefill server)
+    logger.info("Creating vLLM (prefill) instance...")
+    client.request_json("POST", "/instance", body={"instance": vllm_instance})
+    try:
+        wait_for_instance_ready(client, vllm_instance_id)
+    except (RuntimeError, TimeoutError) as e:
+        logger.error(f"Failed to initialize vLLM instance: {e}")
+        with contextlib.suppress(ExoHttpError):
+            client.request_json("DELETE", f"/instance/{vllm_instance_id}")
+        return 1
+    logger.info("vLLM (prefill) instance ready")
+
+    # Create MLX instance (decode)
+    logger.info("Creating MLX (decode) instance...")
+    client.request_json("POST", "/instance", body={"instance": mlx_instance})
+    try:
+        wait_for_instance_ready(client, mlx_instance_id)
+    except (RuntimeError, TimeoutError) as e:
+        logger.error(f"Failed to initialize MLX instance: {e}")
+        with contextlib.suppress(ExoHttpError):
+            client.request_json("DELETE", f"/instance/{mlx_instance_id}")
+        with contextlib.suppress(ExoHttpError):
+            client.request_json("DELETE", f"/instance/{vllm_instance_id}")
+        return 1
+    logger.info("MLX (decode) instance ready")
+
+    time.sleep(2)
+
+    all_rows: list[dict[str, Any]] = []
+
+    try:
+        for i in range(args.warmup):
+            run_one_disaggregated(
+                client, decode_full_id, pp_list[0], tg_list[0], prompt_sizer
+            )
+            logger.debug(f"  warmup {i + 1}/{args.warmup} done")
+
+        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))
+
+        for pp, tg in pp_tg_pairs:
+            logger.info(f"--- pp={pp} tg={tg} ---")
+            runs: list[dict[str, Any]] = []
+            for r in range(args.repeat):
+                time.sleep(3)
+
+                try:
+                    row, actual_pp_tokens = run_one_disaggregated(
+                        client, decode_full_id, pp, tg, prompt_sizer
+                    )
+                except Exception as e:
+                    logger.error(e)
+                    continue
+
+                row.update(
+                    {
+                        "prefill_model_short_id": prefill_short_id,
+                        "prefill_model_id": prefill_full_id,
+                        "prefill_instance_meta": vllm_meta,
+                        "prefill_instance_id": vllm_instance_id,
+                        "decode_model_short_id": decode_short_id,
+                        "decode_model_id": decode_full_id,
+                        "decode_instance_meta": mlx_meta,
+                        "decode_sharding": mlx_sharding,
+                        "decode_nodes": mlx_nodes,
+                        "decode_instance_id": mlx_instance_id,
+                        "overlapping": overlapping,
+                        "pp_tokens": actual_pp_tokens,
+                        "tg": tg,
+                        "repeat_index": r,
+                    }
+                )
+                runs.append(row)
+                all_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"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\n"
+                )
+            time.sleep(2)
+    finally:
+        try:
+            client.request_json("DELETE", f"/instance/{mlx_instance_id}")
+        except ExoHttpError as e:
+            if e.status != 404:
+                raise
+        try:
+            client.request_json("DELETE", f"/instance/{vllm_instance_id}")
+        except ExoHttpError as e:
+            if e.status != 404:
+                raise
+        wait_for_instance_gone(client, mlx_instance_id)
+        wait_for_instance_gone(client, vllm_instance_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())

bench_overnight.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+set -e
+
+export PATH="/opt/homebrew/bin:$PATH"
+
+echo "=== Starting overnight bench runs at $(date) ==="
+
+echo "--- [4/8] Qwen3.5-122B-A10B-GPTQ-Int4 ---"
+echo "Skipping because Int 4"
+#uv run bench/exo_bench.py --force-download --model "Qwen/Qwen3.5-122B-A10B-GPTQ-Int4" --pp 700 --tg 36000 --repeat 1
+
+echo "--- [5/8] Qwen3.5-27B-FP8 ---"
+#uv run bench/exo_bench.py --force-download --model "Qwen/Qwen3.5-27B-FP8" --pp 700 --tg 35133 --repeat 1
+
+echo "--- [6/8] GLM-4.7-Flash-bf16 ---"
+uv run bench/exo_bench.py --force-download --model "mlx-community/GLM-4.7-Flash-bf16" --pp 700 --tg 29000 --repeat 1
+
+echo "--- [7/8] NVIDIA-Nemotron-3-Nano-30B-A3B (23000,1200) ---"
+uv run bench/exo_bench.py --force-download --model "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-BF16" --pp 700 --tg 23000,1200 --repeat 1
+
+echo "--- [8/8] Qwen3.5-27B-bf16 ---"
+uv run bench/exo_bench.py --force-download --model "mlx-community/Qwen3.5-27B-bf16" --pp 700 --tg 35400 --repeat 1
+
+echo "=== All bench runs complete at $(date) ==="

dashboard/src/lib/components/ChatModelSelector.svelte

@@ -9,6 +9,10 @@
     quantization: string;
   }
 
+  function normalizeBaseModel(s: string): string {
+    return s.toLowerCase().replace(/[-_]/g, " ").trim();
+  }
+
   // Auto mode tier list (for when user just starts typing)
   export const AUTO_TIERS: string[][] = [
     // Tier 1 (frontier)
@@ -43,8 +47,9 @@
 
   /** Return the tier index (0 = best) for a base_model name. */
   export function getAutoTierIndex(baseModel: string): number {
+    const norm = normalizeBaseModel(baseModel);
     for (let i = 0; i < AUTO_TIERS.length; i++) {
-      if (AUTO_TIERS[i].includes(baseModel)) return i;
+      if (AUTO_TIERS[i].some((t) => normalizeBaseModel(t) === norm)) return i;
     }
     return AUTO_TIERS.length; // not in any tier → lowest priority
   }
@@ -60,7 +65,7 @@
         const variants = modelList
           .filter(
             (m) =>
-              m.base_model === baseModel &&
+              normalizeBaseModel(m.base_model) === normalizeBaseModel(baseModel) &&
               (m.storage_size_megabytes || 0) / 1024 <= memoryGB &&
               (m.storage_size_megabytes || 0) > 0,
           )
@@ -162,7 +167,7 @@
   /** For a given base_model name, find the biggest quant variant that fits in memory. */
   function pickBestVariant(baseModel: string): ChatModelInfo | null {
     const variants = models
-      .filter((m) => m.base_model === baseModel && fitsInMemory(m))
+      .filter((m) => normalizeBaseModel(m.base_model) === normalizeBaseModel(baseModel) && fitsInMemory(m))
       .sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
     return variants[0] ?? null;
   }

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -379,12 +379,16 @@
     return hfTrendingModels;
   });
 
+  function normalizeBaseModel(s: string): string {
+    return s.toLowerCase().replace(/[-_]/g, " ").trim();
+  }
+
   // Group models by base_model
   const groupedModels = $derived.by((): ModelGroup[] => {
     const groups = new Map<string, ModelGroup>();
 
     for (const model of models) {
-      const groupId = model.base_model || model.id;
+      const groupId = normalizeBaseModel(model.base_model || model.id);
       const groupName = model.base_model || model.name || model.id;
 
       if (!groups.has(groupId)) {
@@ -578,7 +582,7 @@
       const model = models.find((m) => m.id === id);
       if (model) {
         result.push({
-          id: model.base_model || model.id,
+          id: normalizeBaseModel(model.base_model || model.id),
           name: model.name || model.id,
           capabilities: model.capabilities || ["text"],
           family: model.family || "",

dashboard/src/routes/+page.svelte

@@ -295,7 +295,7 @@
     const seen = new Set<string>();
     const deduped: typeof candidates = [];
     for (const m of candidates) {
-      const key = m.base_model || m.family || m.id;
+      const key = (m.base_model || m.family || m.id).toLowerCase().replace(/[-_]/g, " ");
       if (seen.has(key)) continue;
       seen.add(key);
       deduped.push(m);

justfile

@@ -21,12 +21,33 @@ sync-cuda:
     if command -v nvidia-smi &>/dev/null; then
         arch=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -1 | tr -d ' ')
         export TORCH_CUDA_ARCH_LIST="$arch"
+        export FLASHINFER_CUDA_ARCH_LIST="${arch}a"
         export VLLM_TARGET_DEVICE=cuda
     fi
     uv pip install "cmake>=3.26.1" ninja "packaging>=24.2" "setuptools>=77.0.3,<81.0.0" "setuptools-scm>=8.0" wheel jinja2
     find ~/.cache/uv/git-v0 -name CMakeCache.txt -delete 2>/dev/null || true
     uv sync --extra cuda
 
+build-flashinfer-sm121:
+    #!/usr/bin/env bash
+    set -euo pipefail
+    FLASHINFER_VERSION="v0.6.6"
+    BUILD_DIR="/tmp/flashinfer-build"
+    EXO_VENV="$(pwd)/.venv"
+    rm -rf "$BUILD_DIR"
+    git clone --depth 1 --branch "$FLASHINFER_VERSION" --recurse-submodules --shallow-submodules https://github.com/flashinfer-ai/flashinfer.git "$BUILD_DIR"
+    cd "$BUILD_DIR"
+    export FLASHINFER_CUDA_ARCH_LIST="12.1a"
+    export TORCH_CUDA_ARCH_LIST="12.1"
+    export FLASHINFER_ENABLE_AOT=1
+    UV="$(command -v uv || echo "$HOME/.local/bin/uv")"
+    VIRTUAL_ENV="$EXO_VENV" "$UV" pip install . --no-build-isolation
+    echo "Building AOT cubin wheel..."
+    cd flashinfer-cubin && VIRTUAL_ENV="$EXO_VENV" "$UV" pip install . --no-build-isolation && cd ..
+    echo "Building JIT cache wheel..."
+    cd flashinfer-jit-cache && VIRTUAL_ENV="$EXO_VENV" "$UV" pip install . --no-build-isolation && cd ..
+    echo "FlashInfer built from source with SM121 AOT kernels"
+
 sync-clean:
     uv sync --all-packages --force-reinstall --no-cache
 

pyproject.toml

@@ -134,8 +134,8 @@ environments = [
   "sys_platform == 'darwin'",
   "sys_platform == 'linux' and platform_machine == 'aarch64'",
 ]
-no-binary-package = ["vllm"]
-no-build-isolation-package = ["vllm"]
+no-binary-package = ["vllm", "flashinfer-python"]
+no-build-isolation-package = ["vllm", "flashinfer-python"]
 extra-build-dependencies = { vllm = [
   "cmake>=3.26.1",
   "ninja",

scripts/check_kv_compat.py

@@ -0,0 +1,24 @@
+import sys
+sys.path.insert(0, "src")
+import mlx.core as mx
+from mlx_lm import load
+from mlx_lm.models.cache import RotatingKVCache, KVCache
+
+model, tok = load("mlx-community/gpt-oss-20b-MXFP4-Q8")
+
+prompt = "Hello " * 2000
+tokens = tok.encode(prompt)
+print(f"Tokens: {len(tokens)}")
+
+cache = model.make_cache()
+token_arr = mx.array([tokens])
+logits = model(token_arr, cache=cache)
+mx.eval(logits)
+
+for i, c in enumerate(cache[:6]):
+    if isinstance(c, KVCache) and not isinstance(c, RotatingKVCache) and c.keys is not None:
+        k = c.keys.astype(mx.float32)
+        print(f"Layer {i} KVCache: shape={c.keys.shape} offset={c.offset} first=[{float(k[0,0,0,0]):.6f}, {float(k[0,0,0,1]):.6f}] last=[{float(k[0,0,-1,-2]):.6f}, {float(k[0,0,-1,-1]):.6f}]")
+    elif isinstance(c, RotatingKVCache) and c.keys is not None:
+        k = c.keys.astype(mx.float32)
+        print(f"Layer {i} RotatingKV: shape={c.keys.shape} _idx={c._idx} offset={c.offset} first=[{float(k[0,0,0,0]):.6f}, {float(k[0,0,0,1]):.6f}]")

scripts/check_rotating_cache.py

@@ -0,0 +1,12 @@
+import mlx.core as mx
+from mlx_lm import load
+from mlx_lm.models.cache import RotatingKVCache
+
+model, tok = load("mlx-community/gpt-oss-20b-MXFP4-Q8")
+cache = model.make_cache()
+tokens = mx.ones((1, 5000), dtype=mx.int32)
+model(tokens, cache=cache)
+mx.eval([c.keys for c in cache if c.keys is not None])
+for i, c in enumerate(cache[:4]):
+    if isinstance(c, RotatingKVCache):
+        print(f"Layer {i}: _idx={c._idx} offset={c.offset} keep={c.keep} max_size={c.max_size} keys={c.keys.shape}")

scripts/compare_kv.py

@@ -0,0 +1,186 @@
+import sys
+sys.path.insert(0, "src")
+from exo.worker.engines.mlx.gdn_softplus_patch import patch_gdn_softplus
+from exo.worker.engines.mlx.yarn_rope_patch import patch_yarn_rope
+patch_gdn_softplus()
+patch_yarn_rope()
+import mlx.core as mx
+import torch
+import socket
+from pathlib import Path
+import json
+from collections import defaultdict
+from mlx_lm import load
+from mlx_lm.models.cache import ArraysCache, RotatingKVCache, KVCache
+from exo.disaggregated.protocol import read_header, read_message, ArraysState, KVChunk, Done
+from exo.disaggregated.prefill_client import _nhd_to_bhsd, _torch_to_mx
+
+ENDPOINT = sys.argv[1] if len(sys.argv) > 1 else "10.43.0.1:62988"
+MODEL = sys.argv[2] if len(sys.argv) > 2 else "mlx-community/Llama-3.2-1B-Instruct-bf16"
+MODEL_PATH = sys.argv[3] if len(sys.argv) > 3 else None
+
+model, tok = load(MODEL_PATH or str(Path.home() / ".exo/models" / MODEL.replace("/", "--")))
+prompt = "The quick brown fox jumps over the lazy dog. " * 3000
+tokens = tok.encode(prompt)
+print(f"Tokens: {len(tokens)}")
+
+host, port = ENDPOINT.rsplit(":", 1)
+sock = socket.create_connection((host, int(port)), timeout=60)
+request = json.dumps({"model": MODEL, "token_ids": tokens, "start_pos": 0}).encode() + b"\n"
+sock.sendall(request)
+stream = sock.makefile("rb", buffering=65536)
+header = read_header(stream)
+
+vllm_kv = defaultdict(list)
+vllm_arrays: dict[int, list[torch.Tensor]] = {}
+while True:
+    msg = read_message(stream, header)
+    if msg is None or isinstance(msg, Done):
+        break
+    if isinstance(msg, KVChunk):
+        vllm_kv[msg.layer_idx].append((msg.keys, msg.values))
+    elif isinstance(msg, ArraysState):
+        vllm_arrays[msg.layer_idx] = msg.arrays
+sock.close()
+
+print(f"Received {len(vllm_kv)} KV layers, {len(vllm_arrays)} arrays layers from vLLM")
+
+if hasattr(model, "make_cache"):
+    mlx_cache = model.make_cache()
+else:
+    from mlx_lm.models.cache import make_prompt_cache
+    mlx_cache = make_prompt_cache(model)
+token_arr = mx.array([tokens[:-2]])
+mlx_logits = model(token_arr, cache=mlx_cache)
+mx.eval(mlx_logits)
+
+for i in range(min(6, len(mlx_cache))):
+    c = mlx_cache[i]
+    if isinstance(c, ArraysCache):
+        if i in vllm_arrays:
+            vllm_arrs = vllm_arrays[i]
+            mlx_state = c.state
+            print(f"Layer {i} (Arrays): mlx_state={len(mlx_state)} arrays, vllm={len(vllm_arrs)} arrays")
+            for ai, (m_arr, v_arr) in enumerate(zip(mlx_state, vllm_arrs)):
+                if m_arr is None:
+                    continue
+                v_mx = _torch_to_mx(v_arr).astype(mx.float32)
+                m_f = m_arr.astype(mx.float32)
+                if m_f.shape != v_mx.shape:
+                    print(f"  [{ai}] SHAPE MISMATCH mlx={m_f.shape} vllm={v_mx.shape}")
+                else:
+                    d = mx.abs(m_f - v_mx)
+                    a = m_f.reshape(-1)
+                    b = v_mx.reshape(-1)
+                    cos = float(mx.sum(a * b).item()) / (float(mx.sqrt(mx.sum(a * a)).item()) * float(mx.sqrt(mx.sum(b * b)).item()) + 1e-8)
+                    print(f"  [{ai}] cosine_sim={cos:.6f} max_diff={mx.max(d).item():.6f} mean_diff={mx.mean(d).item():.6f} shape={m_f.shape}")
+        else:
+            print(f"Layer {i} (Arrays): no vLLM data")
+        continue
+    if c.keys is None:
+        continue
+    mlx_k = c.keys.astype(mx.float32)
+
+    if i not in vllm_kv:
+        print(f"Layer {i}: no vLLM data")
+        continue
+    chunks = vllm_kv[i]
+    vk = torch.cat([k for k, v in chunks], dim=0) if len(chunks) > 1 else chunks[0][0]
+    vk_mx = _torch_to_mx(vk.permute(1, 0, 2).unsqueeze(0)).astype(mx.float32)
+
+    n = min(mlx_k.shape[2], vk_mx.shape[2])
+    diff = mx.abs(mlx_k[:, :, :n, :] - vk_mx[:, :, :n, :])
+    max_diff = mx.max(diff).item()
+    mean_diff = mx.mean(diff).item()
+    cache_type = "RotatingKV" if isinstance(c, RotatingKVCache) else "KV"
+    print(f"Layer {i} ({cache_type}): mlx={mlx_k.shape} vllm={vk_mx.shape} max_diff={max_diff:.6f} mean_diff={mean_diff:.6f}")
+
+    a = mlx_k[:, :, :n, :].reshape(-1)
+    b_vec = vk_mx[:, :, :n, :].reshape(-1)
+    cos_sim = float(mx.sum(a * b_vec).item()) / (float(mx.sqrt(mx.sum(a * a)).item()) * float(mx.sqrt(mx.sum(b_vec * b_vec)).item()) + 1e-8)
+    diff_tensor = mx.abs(mlx_k[:, :, :n, :] - vk_mx[:, :, :n, :])
+    max_idx = mx.argmax(diff_tensor.reshape(-1)).item()
+    total_elems = diff_tensor.shape[1] * n * diff_tensor.shape[3]
+    h_idx = (max_idx // (n * diff_tensor.shape[3])) % diff_tensor.shape[1]
+    s_idx = (max_idx // diff_tensor.shape[3]) % n
+    d_idx = max_idx % diff_tensor.shape[3]
+    print(f"  cosine_sim={cos_sim:.6f} max_diff={max_diff:.4f} at h={h_idx} pos={s_idx} dim={d_idx}: mlx={float(mlx_k[0,h_idx,s_idx,d_idx].item()):.6f} vllm={float(vk_mx[0,h_idx,s_idx,d_idx].item()):.6f}")
+
+    D = mlx_k.shape[3]
+    for pos in [0, 100, n-1]:
+        mlx_row = [float(mlx_k[0, 0, pos, d].item()) for d in range(D)]
+        vllm_row = [float(vk_mx[0, 0, pos, d].item()) for d in range(D)]
+        diffs = [abs(mlx_row[d] - vllm_row[d]) for d in range(D)]
+        top5 = sorted(range(D), key=lambda d: -diffs[d])[:5]
+        print(f"  pos={pos} top5 diff dims: {[(d, f'{diffs[d]:.3f}', f'mlx={mlx_row[d]:.3f}', f'vllm={vllm_row[d]:.3f}') for d in top5]}")
+
+print("\n--- Run 2: cached request ---")
+sock2 = socket.create_connection((host, int(port)), timeout=60)
+request2 = json.dumps({"model": MODEL, "token_ids": tokens, "start_pos": 0}).encode() + b"\n"
+sock2.sendall(request2)
+stream2 = sock2.makefile("rb", buffering=65536)
+
+first_byte = stream2.peek(1)[:1]
+if first_byte == b"{":
+    line2 = stream2.readline()
+    print(f"Server error: {json.loads(line2.decode())}")
+    sys.exit(1)
+
+header2 = read_header(stream2)
+vllm_kv2 = defaultdict(list)
+vllm_arrays2: dict[int, list[torch.Tensor]] = {}
+total_tokens2 = 0
+while True:
+    msg = read_message(stream2, header2)
+    if msg is None:
+        break
+    if isinstance(msg, KVChunk):
+        vllm_kv2[msg.layer_idx].append((msg.keys, msg.values))
+    elif isinstance(msg, ArraysState):
+        vllm_arrays2[msg.layer_idx] = msg.arrays
+    elif isinstance(msg, Done):
+        total_tokens2 = msg.total_tokens
+        break
+sock2.close()
+
+kv_tokens2 = 0
+if vllm_kv2:
+    first_layer = next(iter(vllm_kv2.values()))
+    kv_tokens2 = sum(k.shape[0] for k, v in first_layer)
+print(f"Received {len(vllm_kv2)} KV layers ({kv_tokens2} tokens), {len(vllm_arrays2)} arrays layers, total_tokens={total_tokens2}")
+
+for i in range(min(6, len(mlx_cache))):
+    c = mlx_cache[i]
+    if isinstance(c, ArraysCache):
+        if i in vllm_arrays2:
+            vllm_arrs = vllm_arrays2[i]
+            mlx_state = c.state
+            for ai, (m_arr, v_arr) in enumerate(zip(mlx_state, vllm_arrs)):
+                if m_arr is None:
+                    continue
+                v_mx = _torch_to_mx(v_arr).astype(mx.float32)
+                m_f = m_arr.astype(mx.float32)
+                if m_f.shape != v_mx.shape:
+                    print(f"Layer {i} [{ai}] SHAPE MISMATCH mlx={m_f.shape} vllm={v_mx.shape}")
+                else:
+                    a2 = m_f.reshape(-1)
+                    b2 = v_mx.reshape(-1)
+                    cos2 = float(mx.sum(a2 * b2).item()) / (float(mx.sqrt(mx.sum(a2 * a2)).item()) * float(mx.sqrt(mx.sum(b2 * b2)).item()) + 1e-8)
+                    print(f"Layer {i} (Arrays) [{ai}] cosine_sim={cos2:.6f} shape={m_f.shape}")
+        continue
+    if c.keys is None or i not in vllm_kv2:
+        continue
+    mlx_k = c.keys.astype(mx.float32)
+    chunks = vllm_kv2[i]
+    vk = torch.cat([k for k, v in chunks], dim=0) if len(chunks) > 1 else chunks[0][0]
+    vk_mx = _torch_to_mx(vk.permute(1, 0, 2).unsqueeze(0)).astype(mx.float32)
+    n = min(mlx_k.shape[2], vk_mx.shape[2])
+    a2 = mlx_k[:, :, :n, :].reshape(-1)
+    b2 = vk_mx[:, :, :n, :].reshape(-1)
+    cos2 = float(mx.sum(a2 * b2).item()) / (float(mx.sqrt(mx.sum(a2 * a2)).item()) * float(mx.sqrt(mx.sum(b2 * b2)).item()) + 1e-8)
+    print(f"Layer {i} (KV) cosine_sim={cos2:.6f} mlx={mlx_k.shape} vllm={vk_mx.shape}")
+
+if len(vllm_kv2) > 0:
+    print("PASS")
+else:
+    print("FAIL")

scripts/disaggregated/capture_connector.py

@@ -0,0 +1,87 @@
+"""Minimal KVConnector that captures per-layer cache data."""
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+)
+
+captured_layers: dict[str, Any] = {}
+
+
+@dataclass
+class CaptureMetadata(KVConnectorMetadata):
+    pass
+
+
+class CaptureConnector(KVConnectorBase_V1):
+    def __init__(self, vllm_config, role, kv_cache_config=None):
+        super().__init__(vllm_config, role, kv_cache_config)
+
+    def start_load_kv(self, forward_context, **kwargs):
+        pass
+
+    def wait_for_layer_load(self, layer_name):
+        pass
+
+    def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs):
+        import time
+        slot_mapping = getattr(attn_metadata, 'slot_mapping', None)
+        if slot_mapping is not None and slot_mapping.shape[0] <= 100:
+            return
+        t0 = time.perf_counter()
+        torch.cuda.synchronize()
+        t_sync = time.perf_counter() - t0
+        if isinstance(kv_layer, (list, tuple)):
+            captured_layers[layer_name] = [t.cpu().clone() for t in kv_layer]
+        else:
+            slot_mapping = getattr(attn_metadata, 'slot_mapping', None)
+            if slot_mapping is not None:
+                if kv_layer.shape[0] == 2:
+                    k_all = kv_layer[0]
+                    v_all = kv_layer[1]
+                else:
+                    k_all = kv_layer[:, 0]
+                    v_all = kv_layer[:, 1]
+                k_flat = k_all.reshape(-1, *k_all.shape[-2:])
+                v_flat = v_all.reshape(-1, *v_all.shape[-2:])
+                valid = slot_mapping >= 0
+                safe_sm = slot_mapping.clamp(min=0)
+                keys = k_flat[safe_sm]
+                values = v_flat[safe_sm]
+                keys[~valid] = 0
+                values[~valid] = 0
+                prev = captured_layers.get(layer_name)
+                if isinstance(prev, dict) and "keys" in prev:
+                    t1 = time.perf_counter()
+                    captured_layers[layer_name] = {
+                        "keys": torch.cat([prev["keys"], keys.cpu()], dim=0),
+                        "values": torch.cat([prev["values"], values.cpu()], dim=0),
+                    }
+                    t_copy = time.perf_counter() - t1
+                else:
+                    t1 = time.perf_counter()
+                    captured_layers[layer_name] = {
+                        "keys": keys.cpu(),
+                        "values": values.cpu(),
+                    }
+                    t_copy = time.perf_counter() - t1
+                if "layers.3." in layer_name:
+                    print(f"    [attn save] sync={t_sync*1000:.1f}ms copy={t_copy*1000:.1f}ms tokens={keys.shape[0]}")
+            else:
+                captured_layers[layer_name] = kv_layer.cpu().clone()
+
+    def wait_for_save(self):
+        pass
+
+    def get_num_new_matched_tokens(self, request, num_computed_tokens):
+        return 0, False
+
+    def update_state_after_alloc(self, request, blocks, num_external_tokens):
+        pass
+
+    def build_connector_meta(self, scheduler_output):
+        return CaptureMetadata()

scripts/disaggregated/inspect_vllm_kv.py

@@ -0,0 +1,156 @@
+"""Inspect vLLM KV cache structure per-layer after prefill.
+
+Runs on DGX Spark. Prints per-layer shapes, dtypes, kv_cache_config,
+and layer_to_group mapping to understand what vLLM stores for each
+model architecture (standard attention, sliding window, GatedDeltaNet).
+
+Usage:
+  uv run python scripts/disaggregated/inspect_vllm_kv.py --model ~/.local/share/exo/models/openai--gpt-oss-20b
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src"))
+
+os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
+os.environ["VLLM_BATCH_INVARIANT"] = "1"
+
+from exo.worker.runner.bootstrap import _ensure_cuda_libs
+
+_ensure_cuda_libs()
+
+import torch
+
+
+def _build_layer_groups(kv_cache_config):
+    group_lookup = {}
+    for group_idx, group_spec in enumerate(kv_cache_config.kv_cache_groups):
+        for layer_name in group_spec.layer_names:
+            group_lookup[layer_name] = group_idx
+
+    layer_to_group = []
+    for tensor_spec in kv_cache_config.kv_cache_tensors:
+        for name in tensor_spec.shared_by:
+            layer_to_group.append(group_lookup[name])
+    return layer_to_group
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True, help="Path to model")
+    parser.add_argument("--prompt", default="Hello, world! How are you today?", help="Prompt to prefill")
+    args = parser.parse_args()
+
+    from exo.worker.engines.vllm.growable_cache import get_model_runner
+    from exo.worker.engines.vllm.vllm_generator import load_vllm_engine
+
+    print(f"Loading vLLM engine from {args.model}...")
+    engine, _, prefix_cache = load_vllm_engine(
+        model_path=args.model,
+        model_id=args.model,
+        trust_remote_code=True,
+    )
+    print("Engine loaded.\n")
+
+    from vllm import SamplingParams
+
+    tokenizer = engine.get_tokenizer()
+    token_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
+    print(f"Prompt: {args.prompt!r}")
+    print(f"Token IDs: {len(token_ids)} tokens\n")
+
+    request_id = "inspect-test"
+    params = SamplingParams(max_tokens=1, detokenize=False)
+    engine.add_request(request_id, {"prompt_token_ids": token_ids}, params)
+
+    while engine.has_unfinished_requests():
+        engine.step()
+
+    model_runner = get_model_runner()
+    if model_runner is None:
+        print("ERROR: model_runner is None")
+        return
+
+    print("=" * 70)
+    print("PER-LAYER KV CACHE TENSORS (model_runner.kv_caches)")
+    print("=" * 70)
+    kv_caches = model_runner.kv_caches
+    for i, kv in enumerate(kv_caches):
+        if isinstance(kv, list):
+            shapes = [t.shape for t in kv]
+            dtypes = [t.dtype for t in kv]
+            print(f"  Layer {i:3d}: list of {len(kv)} tensors — shapes={shapes}, dtypes={dtypes}")
+        elif isinstance(kv, torch.Tensor):
+            print(f"  Layer {i:3d}: shape={tuple(kv.shape)}, dtype={kv.dtype}, device={kv.device}")
+        else:
+            print(f"  Layer {i:3d}: type={type(kv).__name__}")
+    print(f"\n  Total layers with KV: {len(kv_caches)}\n")
+
+    engine_core = engine.engine_core.engine_core
+    kv_cache_config = engine_core.scheduler.kv_cache_manager.kv_cache_config
+
+    print("=" * 70)
+    print("KV CACHE CONFIG")
+    print("=" * 70)
+
+    print(f"\n  Number of KV cache groups: {len(kv_cache_config.kv_cache_groups)}")
+    for gi, group in enumerate(kv_cache_config.kv_cache_groups):
+        print(f"\n  Group {gi}:")
+        print(f"    Layer names ({len(group.layer_names)}):")
+        for name in group.layer_names[:5]:
+            print(f"      {name}")
+        if len(group.layer_names) > 5:
+            print(f"      ... and {len(group.layer_names) - 5} more")
+
+    print(f"\n  Number of KV cache tensors: {len(kv_cache_config.kv_cache_tensors)}")
+    for ti, tensor_spec in enumerate(kv_cache_config.kv_cache_tensors):
+        shared = tensor_spec.shared_by[:3]
+        extra = f" ... +{len(tensor_spec.shared_by)-3}" if len(tensor_spec.shared_by) > 3 else ""
+        print(f"  Tensor {ti}: shared_by={shared}{extra}")
+
+    layer_to_group = _build_layer_groups(kv_cache_config)
+    print(f"\n  layer_to_group ({len(layer_to_group)} entries): {layer_to_group[:10]}{'...' if len(layer_to_group) > 10 else ''}")
+
+    coordinator = engine_core.scheduler.kv_cache_manager.coordinator
+    null_block = coordinator.block_pool.null_block
+
+    internal_id = None
+    for mgr in coordinator.single_type_managers:
+        for key in mgr.req_to_blocks:
+            if str(key).startswith(request_id):
+                internal_id = str(key)
+                break
+        if internal_id:
+            break
+
+    if internal_id:
+        print(f"\n  Request internal_id: {internal_id}")
+        for gi, mgr in enumerate(coordinator.single_type_managers):
+            blocks = mgr.req_to_blocks.get(internal_id)
+            if blocks:
+                real_blocks = [b for b in blocks if b is not null_block and not b.is_null]
+                null_count = len(blocks) - len(real_blocks)
+                print(f"  Group {gi}: {len(real_blocks)} real blocks, {null_count} null blocks, block_size={mgr.block_size}")
+            else:
+                print(f"  Group {gi}: no blocks")
+
+    print("\n" + "=" * 70)
+    print("SUMMARY")
+    print("=" * 70)
+    print(f"  Model: {args.model}")
+    print(f"  KV cache layers: {len(kv_caches)}")
+    print(f"  KV cache groups: {len(kv_cache_config.kv_cache_groups)}")
+    print(f"  Layer-to-group mapping entries: {len(layer_to_group)}")
+    unique_shapes = set()
+    for kv in kv_caches:
+        if isinstance(kv, torch.Tensor):
+            unique_shapes.add(tuple(kv.shape))
+    print(f"  Unique tensor shapes: {unique_shapes}")
+
+
+if __name__ == "__main__":
+    main()

scripts/disaggregated/test_kv_extract.py

@@ -0,0 +1,261 @@
+"""Extract KV cache per-layer from vLLM using a real KVConnector.
+
+Patches vLLM to allow KVConnector on hybrid models (attention + GDN).
+
+Usage:
+  uv run python scripts/disaggregated/test_kv_extract.py --model ~/.local/share/exo/models/Qwen--Qwen3.5-2B --output /tmp/kv_cache_qwen35/
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src"))
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
+os.environ["VLLM_BATCH_INVARIANT"] = "1"
+
+from exo.worker.runner.bootstrap import _ensure_cuda_libs
+
+_ensure_cuda_libs()
+
+import torch
+
+
+def _patch_vllm_for_connector():
+    """Patch vLLM to allow KVConnector on hybrid models."""
+    from vllm.v1.core import kv_cache_utils
+
+    original_unify = kv_cache_utils.unify_hybrid_kv_cache_specs
+
+    def patched_unify(kv_cache_spec):
+        try:
+            original_unify(kv_cache_spec)
+        except ValueError:
+            pass
+
+    kv_cache_utils.unify_hybrid_kv_cache_specs = patched_unify
+
+    from vllm.v1.core.sched import scheduler as sched_mod
+    original_connector_finished = sched_mod.Scheduler._connector_finished
+
+    def patched_connector_finished(self, request):
+        return False, None
+
+    sched_mod.Scheduler._connector_finished = patched_connector_finished
+
+    from capture_connector import CaptureConnector
+    from vllm.distributed.kv_transfer.kv_connector import factory
+
+    original_get = factory.KVConnectorFactory._get_connector_class_with_compat
+
+    @classmethod
+    def patched_get(cls, kv_transfer_config):
+        if "capture_connector" in (kv_transfer_config.kv_connector or ""):
+            return CaptureConnector, None
+        return original_get.__func__(cls, kv_transfer_config)
+
+    factory.KVConnectorFactory._get_connector_class_with_compat = patched_get
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--output", required=True)
+    _lorem = "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Curabitur pretium tincidunt lacus. Nulla gravida orci a odio. Nullam varius, turpis et commodo pharetra, est eros bibendum elit, nec luctus magna felis sollicitudin mauris. Integer in mauris eu nibh euismod gravida. Duis ac tellus et risus vulputate vehicula. Donec lobortis risus a elit. Etiam tempor. Ut ullamcorper, ligula ut dictum pharetra, nisi nunc fringilla magna, in commodo elit erat nec turpis. Ut pharetra augue nec augue. Nam elit agna, endrerit sit amet, tincidunt ac, viverra sed, nulla. Donec porta diam eu massa. Quisque diam lorem, interdum vitae, dapibus ac, scelerisque vitae, pede. Donec eget tellus non erat lacinia fermentum. Donec in velit vel ipsum auctor pulvinar. Vestibulum iaculis lacinia est. Proin dictum elementum velit. Fusce euismod consequat ante. Lorem ipsum dolor sit amet, consectetuer adipiscing elit. Pellentesque sed dolor. Aliquam congue fermentum nisl. Mauris accumsan nulla vel diam. Sed in lacus ut enim adipiscing aliquet. Nulla venenatis. In pede mi, aliquet sit amet, euismod in, auctor ut, ligula. Aliquam dapibus tincidunt metus. Praesent justo dolor, lobortis quis, lobortis dignissim, pulvinar ac, lorem. "
+    parser.add_argument("--prompt", default=_lorem * 21 + "Now answer this question: What is the capital of France and why is it historically significant? Give a detailed answer.")
+    args = parser.parse_args()
+
+    out_dir = Path(args.output)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    _patch_vllm_for_connector()
+
+    from vllm.engine.arg_utils import EngineArgs
+    from vllm.v1.engine.llm_engine import LLMEngine
+
+    from exo.worker.engines.mlx.cache import KVPrefixCache
+    from exo.worker.engines.vllm.growable_cache import patch_vllm, set_prefix_cache
+
+    patch_vllm()
+
+    prefix_cache = KVPrefixCache(group=None)
+    set_prefix_cache(prefix_cache)
+
+    engine_args = EngineArgs(
+        model=args.model,
+        served_model_name=args.model,
+        gpu_memory_utilization=0.05,
+        trust_remote_code=True,
+        load_format="fastsafetensors",
+        enable_prefix_caching=False,
+        attention_backend="TRITON_ATTN",
+        enforce_eager=True,
+        disable_log_stats=True,
+        kv_transfer_config={
+            "kv_connector": "capture_connector:CaptureConnector",
+            "kv_role": "kv_both",
+        },
+    )
+
+    print("Loading engine with KVConnector...")
+    engine = LLMEngine.from_engine_args(engine_args)
+    print("Engine loaded.")
+
+
+    from exo.worker.engines.vllm.growable_cache import get_model_runner
+
+    model_runner = get_model_runner()
+
+    import vllm.model_executor.layers.mamba.ops.causal_conv1d as cc_mod
+    from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
+        causal_conv1d_fn as orig_causal_conv1d_fn,
+    )
+
+    gdn_states: dict[int, dict[str, torch.Tensor]] = {}
+    gdn_call_idx = [0]
+    gdn_layer_order = [0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14, 16, 17, 18, 20, 21, 22]
+
+    def patched_causal_conv1d_fn(*args, conv_states=None, cache_indices=None, **kwargs):
+        result = orig_causal_conv1d_fn(*args, conv_states=conv_states, cache_indices=cache_indices, **kwargs)
+        if conv_states is not None and cache_indices is not None:
+            x = args[0] if args else None
+            if x is not None and x.shape[0] <= 100:
+                return result
+            import time as _time
+            t0 = _time.perf_counter()
+            torch.cuda.synchronize()
+            t_sync = _time.perf_counter() - t0
+            ci = cache_indices[0].item() if cache_indices.numel() > 0 else 0
+            idx = gdn_call_idx[0]
+            layer_idx = gdn_layer_order[idx % len(gdn_layer_order)]
+            t1 = _time.perf_counter()
+            conv_at_ci = conv_states[ci:ci+1].transpose(-1, -2).contiguous().cpu()
+            t_copy = _time.perf_counter() - t1
+            gdn_states.setdefault(layer_idx, {})["conv"] = conv_at_ci
+            gdn_states[layer_idx]["ci"] = ci
+            if gdn_call_idx[0] < 3:
+                print(f"    [gdn save] sync={t_sync*1000:.1f}ms copy={t_copy*1000:.1f}ms layer={layer_idx}")
+            gdn_call_idx[0] += 1
+        return result
+
+    cc_mod.causal_conv1d_fn = patched_causal_conv1d_fn
+    for mod in list(sys.modules.values()):
+        if mod is None or mod is cc_mod:
+            continue
+        if hasattr(mod, 'causal_conv1d_fn') and mod.causal_conv1d_fn is orig_causal_conv1d_fn:
+            mod.causal_conv1d_fn = patched_causal_conv1d_fn
+    print("  Patched causal_conv1d_fn")
+
+    from exo.shared.types.tasks import TaskId
+    from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+    from exo.worker.engines.vllm.vllm_generator import VllmBatchEngine
+
+    batch_engine = VllmBatchEngine(engine=engine, model_id=args.model, prefix_cache=prefix_cache)
+
+    task = TextGenerationTaskParams(
+        model=args.model,
+        input=[InputMessage(role="user", content=args.prompt)],
+        max_completion_tokens=1,
+    )
+
+    task_id = batch_engine.submit(task_id=TaskId("extract"), task_params=task, prompt=args.prompt)
+
+    print("Running prefill via VllmBatchEngine...")
+    t0 = time.perf_counter()
+    while batch_engine.has_work:
+        results = batch_engine.step()
+        for tid, resp in results:
+            print(f"  Prefill done in {(time.perf_counter()-t0)*1000:.0f}ms")
+            batch_engine.cancel([tid])
+            break
+        if results:
+            break
+    t1 = time.perf_counter()
+    print(f"Total: {(t1-t0)*1000:.0f}ms")
+
+    prompt_mx = prefix_cache.prompts[0] if prefix_cache.prompts else None
+    token_ids = [int(x) for x in prompt_mx.tolist()] if prompt_mx is not None else []
+
+    from capture_connector import captured_layers
+
+    print(f"\nCaptured {len(captured_layers)} layers via save_kv_layer:")
+    for name in sorted(captured_layers.keys()):
+        v = captured_layers[name]
+        if isinstance(v, list):
+            print(f"  {name}: {[tuple(t.shape) for t in v]}")
+        elif isinstance(v, torch.Tensor):
+            print(f"  {name}: {tuple(v.shape)}")
+        else:
+            print(f"  {name}: {type(v).__name__}")
+
+    num_tokens = len(token_ids)
+    print(f"  Chat-templated prompt: {num_tokens} tokens")
+    total_layers = 24
+
+    for f_old in out_dir.glob("layer_*"):
+        f_old.unlink()
+
+    metadata = {
+        "model": args.model,
+        "prompt": args.prompt,
+        "num_tokens": num_tokens,
+        "token_ids": token_ids,
+        "num_layers": total_layers,
+        "layers": [],
+    }
+
+    print(f"\nSaving {total_layers} layers...")
+
+    torch.cuda.synchronize()
+    for layer_idx in sorted(gdn_states.keys()):
+        ci = gdn_states[layer_idx]["ci"]
+        kv = model_runner.kv_caches[layer_idx]
+        if isinstance(kv, (list, tuple)) and len(kv) > 1:
+            rec_pool = kv[1]
+            rec = rec_pool[ci:ci+1].cpu().clone()
+            gdn_states[layer_idx]["rec"] = rec
+
+    for li in range(total_layers):
+        if li in gdn_states:
+            s = gdn_states[li]
+            conv = s.get("conv")
+            rec = s.get("rec")
+            torch.save(conv, out_dir / f"layer_{li:03d}_conv.pt")
+            if rec is not None:
+                torch.save(rec, out_dir / f"layer_{li:03d}_rec.pt")
+            metadata["layers"].append({"type": "gdn", "conv": list(conv.shape), "rec": list(rec.shape) if rec is not None else None})
+            print(f"  Layer {li}: GDN conv={tuple(conv.shape)}, rec={tuple(rec.shape) if rec is not None else 'None'}")
+        else:
+            attn_name = None
+            for n in captured_layers:
+                parts = n.split(".")
+                for pi, p in enumerate(parts):
+                    if p == "layers" and pi + 1 < len(parts) and parts[pi + 1] == str(li):
+                        attn_name = n
+                        break
+            if attn_name and isinstance(captured_layers[attn_name], dict):
+                kv = captured_layers[attn_name]
+                torch.save(kv["keys"], out_dir / f"layer_{li:03d}_keys.pt")
+                torch.save(kv["values"], out_dir / f"layer_{li:03d}_values.pt")
+                if "last_chunk_keys" in kv:
+                    torch.save(kv["last_chunk_keys"], out_dir / f"layer_{li:03d}_keys_last.pt")
+                    torch.save(kv["last_chunk_values"], out_dir / f"layer_{li:03d}_values_last.pt")
+                metadata["layers"].append({"type": "kv", "keys_shape": list(kv["keys"].shape), "values_shape": list(kv["values"].shape)})
+                print(f"  Layer {li}: KV keys={tuple(kv['keys'].shape)}, values={tuple(kv['values'].shape)}")
+            else:
+                metadata["layers"].append({"type": "missing"})
+                print(f"  Layer {li}: MISSING")
+
+    with open(out_dir / "metadata.json", "w") as f:
+        json.dump(metadata, f, indent=2)
+    print(f"\nSaved metadata to {out_dir}/metadata.json")
+
+
+if __name__ == "__main__":
+    main()

scripts/disaggregated/test_kv_inject.py

@@ -0,0 +1,265 @@
+"""Inject extracted vLLM KV cache into MLX model caches and test decode.
+
+Runs on Mac (Apple Silicon). Loads per-layer KV tensors saved by
+test_kv_extract.py, converts to MLX format, injects into MLX caches,
+and generates tokens to verify correctness.
+
+Usage:
+  uv run python scripts/disaggregated/test_kv_inject.py \
+    --model mlx-community/gpt-oss-20b-MXFP4-Q8 \
+    --kv-dir /path/to/extracted/kv_cache/ \
+    --num-tokens 20
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src"))
+
+import mlx.core as mx
+import torch
+from mlx_lm import load
+from mlx_lm.models.cache import ArraysCache, KVCache, RotatingKVCache
+
+
+def _torch_to_mx(t: torch.Tensor) -> mx.array:
+    t = t.detach().cpu()
+    if t.dtype == torch.bfloat16:
+        return mx.array(t.float().numpy()).astype(mx.bfloat16)
+    return mx.array(t.numpy())
+
+
+def _to_bhsd(keys: torch.Tensor, values: torch.Tensor, num_tokens: int) -> tuple[mx.array, mx.array]:
+    """Convert vLLM block format to MLX BHSD [1, H, S, D].
+
+    Input can be:
+      - 4D [blocks, block_size, H, D] — flatten to [blocks*block_size, H, D], trim to num_tokens
+      - 3D [S, H, D] — use directly
+    """
+    if keys.dim() == 4:
+        keys = keys.reshape(-1, keys.shape[2], keys.shape[3])[:num_tokens]
+        values = values.reshape(-1, values.shape[2], values.shape[3])[:num_tokens]
+    elif keys.dim() == 3:
+        keys = keys[:num_tokens]
+        values = values[:num_tokens]
+
+    k_mx = _torch_to_mx(keys.permute(1, 0, 2).unsqueeze(0))
+    v_mx = _torch_to_mx(values.permute(1, 0, 2).unsqueeze(0))
+    return k_mx, v_mx
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True, help="MLX model path/ID")
+    parser.add_argument("--kv-dir", required=True, help="Directory with extracted KV tensors")
+    parser.add_argument("--num-tokens", type=int, default=500, help="Tokens to generate")
+    parser.add_argument("--prompt", default=None, help="Override prompt (must match extraction prompt)")
+    args = parser.parse_args()
+
+    kv_dir = Path(args.kv_dir)
+    with open(kv_dir / "metadata.json") as f:
+        metadata = json.load(f)
+
+    num_extracted_layers = metadata["num_layers"]
+    num_tokens = metadata["num_tokens"]
+    vllm_token_ids = metadata.get("token_ids", [])
+
+    print(f"Extracted KV: {num_extracted_layers} layers, {num_tokens} tokens")
+    if vllm_token_ids:
+        print(f"  Using vLLM token_ids ({len(vllm_token_ids)} tokens)")
+    else:
+        print("  WARNING: No token_ids in metadata")
+
+    print(f"\nLoading MLX model: {args.model}")
+    model, tokenizer = load(args.model)
+
+    caches = model.make_cache()
+    num_model_layers = len(caches)
+    print(f"\nMLX model expects {num_model_layers} cache layers:")
+    for i, c in enumerate(caches):
+        print(f"  Layer {i:3d}: {type(c).__name__}", end="")
+        if isinstance(c, RotatingKVCache):
+            print(f" (max_size={c.max_size}, keep={c.keep})", end="")
+        elif isinstance(c, ArraysCache):
+            print(f" (size={len(c.state)})", end="")
+        print()
+
+    layer_info = metadata.get("layers", [])
+    print(f"\nExtracted {num_extracted_layers} layers from vLLM")
+
+    print("\nInjecting KV cache into MLX caches...")
+    injected = 0
+    skipped = 0
+    for i in range(num_model_layers):
+        cache = caches[i]
+
+        if isinstance(cache, ArraysCache):
+            conv_path = kv_dir / f"layer_{i:03d}_conv.pt"
+            rec_path = kv_dir / f"layer_{i:03d}_rec.pt"
+            keys_path = kv_dir / f"layer_{i:03d}_keys.pt"
+            values_path = kv_dir / f"layer_{i:03d}_values.pt"
+            if conv_path.exists():
+                conv = torch.load(conv_path, weights_only=True)
+                rec = torch.load(rec_path, weights_only=True) if rec_path.exists() else None
+                states = [_torch_to_mx(conv)]
+                states.append(_torch_to_mx(rec) if rec is not None else None)
+                cache.state = states
+                injected += 1
+                print(f"  Layer {i}: ArraysCache conv={tuple(conv.shape)}, rec={tuple(rec.shape) if rec is not None else 'None'}")
+            elif keys_path.exists():
+                conv = torch.load(keys_path, weights_only=True)
+                rec = torch.load(values_path, weights_only=True)
+                cache.state = [_torch_to_mx(conv), _torch_to_mx(rec)]
+                injected += 1
+                print(f"  Layer {i}: ArraysCache (legacy) conv={tuple(conv.shape)}, rec={tuple(rec.shape)}")
+            else:
+                print(f"  Layer {i}: SKIP — ArraysCache, no files")
+                skipped += 1
+            continue
+
+        keys_path = kv_dir / f"layer_{i:03d}_keys.pt"
+        values_path = kv_dir / f"layer_{i:03d}_values.pt"
+        if not keys_path.exists():
+            skipped += 1
+            continue
+
+        keys_torch = torch.load(keys_path, weights_only=True)
+        values_torch = torch.load(values_path, weights_only=True)
+        k_mx, v_mx = _to_bhsd(keys_torch, values_torch, num_tokens)
+        seq_len = int(k_mx.shape[2])
+
+        if isinstance(cache, KVCache) and not isinstance(cache, RotatingKVCache):
+            cache.keys = k_mx
+            cache.values = v_mx
+            cache.offset = seq_len
+            injected += 1
+        elif isinstance(cache, RotatingKVCache):
+            if seq_len <= cache.max_size:
+                cache.keys = k_mx
+                cache.values = v_mx
+                cache.offset = seq_len
+                cache._idx = seq_len
+            else:
+                keep = cache.keep
+                window = cache.max_size
+                sink_keys = k_mx[:, :, :keep, :]
+                sink_values = v_mx[:, :, :keep, :]
+                recent_keys = k_mx[:, :, -(window - keep):, :]
+                recent_values = v_mx[:, :, -(window - keep):, :]
+                cache.keys = mx.concatenate([sink_keys, recent_keys], axis=2)
+                cache.values = mx.concatenate([sink_values, recent_values], axis=2)
+                cache.offset = seq_len
+                cache._idx = keep
+            injected += 1
+            print(f"  Layer {i}: RotatingKVCache (seq_len={seq_len}, max_size={cache.max_size})")
+        else:
+            print(f"  Layer {i}: SKIP — {type(cache).__name__}")
+            skipped += 1
+
+    print(f"\n  Injected: {injected} layers, Skipped: {skipped} layers")
+
+    from exo.worker.engines.vllm.kv_cache import TorchKVCache as TKV
+    print("\nRound-trip test (MLX → torch → MLX)...")
+    rt_caches = model.make_cache()
+    rt_tokens = mx.array(vllm_token_ids)
+    rt_logits = model(rt_tokens[None], cache=rt_caches)
+    mx.eval(rt_logits)
+    torch_rt = TKV.from_mlx_cache(rt_caches)
+    back_rt = torch_rt.to_mlx_cache()
+    rt_max_diff = 0.0
+    for i in range(len(rt_caches)):
+        nc = rt_caches[i]
+        bc = back_rt[i]
+        if isinstance(nc, ArraysCache):
+            for ai in range(len(nc.state)):
+                if nc.state[ai] is not None and bc.state[ai] is not None:
+                    d = mx.max(mx.abs(nc.state[ai].astype(mx.float32) - bc.state[ai].astype(mx.float32))).item()
+                    rt_max_diff = max(rt_max_diff, d)
+        elif isinstance(nc, (KVCache, RotatingKVCache)) and nc.keys is not None:
+            nk, nv = nc.state
+            bk, bv = bc.state
+            d = mx.max(mx.abs(nk.astype(mx.float32) - bk.astype(mx.float32))).item()
+            rt_max_diff = max(rt_max_diff, d)
+    print(f"  Round-trip max diff: {rt_max_diff:.4e} ({'PASS' if rt_max_diff < 0.01 else 'FAIL'})")
+
+    print("\nComparing with MLX-native prefill...")
+    native_caches = rt_caches
+
+    for i in range(num_model_layers):
+        nc = native_caches[i]
+        ic = caches[i]
+        if isinstance(nc, KVCache) and not isinstance(nc, RotatingKVCache) and nc.keys is not None and ic.keys is not None:
+            s = min(nc.offset, ic.offset)
+            nk = nc.keys[:, :, :s, :].astype(mx.float32)
+            ik = ic.keys[:, :, :s, :].astype(mx.float32)
+            nv = nc.values[:, :, :s, :].astype(mx.float32)
+            iv = ic.values[:, :, :s, :].astype(mx.float32)
+            k_diff = mx.max(mx.abs(nk - ik)).item()
+            v_diff = mx.max(mx.abs(nv - iv)).item()
+            if k_diff > 0.01 or i < 4 or i == num_model_layers - 1:
+                print(f"  Layer {i:3d} KVCache: k_diff={k_diff:.4e}, v_diff={v_diff:.4e}, offset native={nc.offset} injected={ic.offset}")
+        elif isinstance(nc, RotatingKVCache):
+            pass
+        elif isinstance(nc, ArraysCache):
+            for ai in range(len(nc.state)):
+                na = nc.state[ai]
+                ia = ic.state[ai]
+                if na is not None and ia is not None:
+                    diff = mx.max(mx.abs(na.astype(mx.float32) - ia.astype(mx.float32))).item()
+                    if diff > 0.01 or i < 4 or i == num_model_layers - 1:
+                        print(f"  Layer {i:3d} Arrays[{ai}]: diff={diff:.4e}, native_shape={na.shape}, injected_shape={ia.shape}")
+
+    native_last = mx.array([vllm_token_ids[-1]])
+    native_decode_logits = model(native_last[None], cache=native_caches)
+    mx.eval(native_decode_logits)
+    native_first = mx.argmax(native_decode_logits[:, -1, :], axis=-1)
+    print(f"  Native decode first token: {native_first.item()}, text: {tokenizer.decode([native_first.item()])!r}")
+
+    print(f"\nDecoding {args.num_tokens} tokens with injected cache...")
+    last_tokens = mx.array(vllm_token_ids[-2:])
+    logits = model(last_tokens[None], cache=caches)
+    mx.eval(logits)
+
+    generated_tokens = []
+    token = mx.argmax(logits[:, -1, :], axis=-1)
+    mx.eval(token)
+    generated_tokens.append(token.item())
+
+    for _ in range(args.num_tokens - 1):
+        logits = model(token[None], cache=caches)
+        mx.eval(logits)
+        token = mx.argmax(logits[:, -1, :], axis=-1)
+        mx.eval(token)
+        generated_tokens.append(token.item())
+
+    generated_text = tokenizer.decode(generated_tokens)
+
+    print(f"\n{'='*70}")
+    print("RESULTS")
+    print(f"{'='*70}")
+    print(f"  Model (vLLM): {metadata['model']}")
+    print(f"  Model (MLX):  {args.model}")
+    print(f"  Prompt tokens: {num_tokens}")
+    print(f"  Layers injected: {injected}/{num_model_layers}")
+    print("  Type mismatches: 0")
+    print(f"  Generated {len(generated_tokens)} tokens")
+    print(f"  Text: {generated_text!r}")
+
+    if False:
+        print("\n  GAPS FOUND:")
+        for idx, got, expected in type_mismatches:
+            print(f"    Layer {idx}: vLLM gives KV tensors, MLX wants {expected}")
+        arrays_layers = [i for i, c in enumerate(caches) if isinstance(c, ArraysCache)]
+        if arrays_layers:
+            print(f"    ArraysCache layers (not populated): {arrays_layers[:10]}{'...' if len(arrays_layers) > 10 else ''}")
+
+    if generated_tokens and not all(t == generated_tokens[0] for t in generated_tokens):
+        print("\n  COHERENT OUTPUT: YES (varied tokens)")
+    else:
+        print("\n  COHERENT OUTPUT: POSSIBLY NOT (all same token)")
+
+
+if __name__ == "__main__":
+    main()

scripts/test_vllm_multi.sh

@@ -0,0 +1,80 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+HOST="${1:-gx10-de89}"
+PORT="${2:-52415}"
+NUM_REQUESTS="${3:-4}"
+MODEL="${4:-Qwen/Qwen2.5-0.5B-Instruct}"
+
+echo "Sending $NUM_REQUESTS parallel requests to $HOST:$PORT ($MODEL) with ~32k token prompts..."
+echo
+
+tmpdir=$(mktemp -d)
+pids=()
+for i in $(seq 1 "$NUM_REQUESTS"); do
+  (
+    python3 -c "
+import json, sys, time, urllib.request
+
+import random
+random.seed($i * 9999)
+topics = [
+    'mathematics', 'philosophy', 'religion', 'culture', 'astronomy',
+    'biology', 'music', 'architecture', 'literature', 'physics',
+    'chemistry', 'geology', 'psychology', 'economics', 'linguistics',
+]
+random.shuffle(topics)
+sentences = []
+for j in range(95):
+    t1, t2, t3 = topics[j % len(topics)], topics[(j+3) % len(topics)], topics[(j+7) % len(topics)]
+    sentences.append(
+        f'In the field of {t1}, the number {$i * 1000 + j} holds particular significance '
+        f'when examining its relationship to {t2} and {t3}. Scholars have long debated '
+        f'whether the patterns observed in iteration {j} of this analysis reveal deeper '
+        f'structural connections between seemingly unrelated disciplines. The evidence '
+        f'from experiment {$i * 7 + j * 13} suggests that cross-domain numerical '
+        f'correlations emerge at scale {j * $i}, challenging conventional assumptions '
+        f'about the independence of these fields. ')
+prompt = ' '.join(sentences) + f' Summarize the key finding about the number {$i}.'
+payload = json.dumps({
+    'model': '$MODEL',
+    'messages': [{'role': 'user', 'content': prompt}],
+    'max_tokens': 1,
+    'stream': True,
+}).encode()
+req = urllib.request.Request(
+    'http://$HOST:$PORT/v1/chat/completions',
+    data=payload,
+    headers={'Content-Type': 'application/json'},
+)
+t0 = time.perf_counter()
+try:
+    resp = urllib.request.urlopen(req, timeout=300)
+    first_byte = None
+    for line in resp:
+        if first_byte is None:
+            first_byte = time.perf_counter()
+        line = line.decode().strip()
+        if line.startswith('data: ') and line != 'data: [DONE]':
+            break
+    ttft = (first_byte or time.perf_counter()) - t0
+    prompt_tokens = len(prompt.split()) * 1.3  # rough estimate
+    tps = prompt_tokens / ttft
+    print(f'request $i: TTFT={ttft:.2f}s  ~{int(prompt_tokens)} prompt tokens  ~{int(tps)} tok/s prefill')
+except Exception as e:
+    elapsed = time.perf_counter() - t0
+    print(f'request $i: FAILED after {elapsed:.2f}s — {e}', file=sys.stderr)
+    sys.exit(1)
+" >"$tmpdir/$i" 2>&1
+  ) &
+  pids+=($!)
+done
+
+for pid in "${pids[@]}"; do
+  wait "$pid"
+done
+
+for i in $(seq 1 "$NUM_REQUESTS"); do
+  cat "$tmpdir/$i"
+done
+rm -rf "$tmpdir"

src/exo/disaggregated/batch_connector.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (  # pyright: ignore[reportMissingImports]
+    KVConnectorBase_V1,  # pyright: ignore[reportUnknownVariableType]
+    KVConnectorMetadata,  # pyright: ignore[reportUnknownVariableType]
+    KVConnectorRole,  # pyright: ignore[reportUnknownVariableType]
+    SupportsHMA,  # pyright: ignore[reportUnknownVariableType]
+)
+
+_LAYER_RE = re.compile(r"layers\.(\d+)\.")
+
+_shared_captured_layers: dict[int, dict[str, torch.Tensor]] = {}
+_shared_captured_arrays: dict[int, list[torch.Tensor]] = {}
+
+
+def get_shared_captured_layers() -> dict[int, dict[str, torch.Tensor]]:
+    return _shared_captured_layers
+
+
+def get_shared_captured_arrays() -> dict[int, list[torch.Tensor]]:
+    return _shared_captured_arrays
+
+
+def clear_shared_captured_layers() -> None:
+    _shared_captured_layers.clear()
+    _shared_captured_arrays.clear()
+
+
+@dataclass
+class BatchConnectorMetadata(KVConnectorMetadata):  # pyright: ignore[reportUntypedBaseClass]
+    pass
+
+
+class BatchConnector(KVConnectorBase_V1, SupportsHMA):  # pyright: ignore[reportUntypedBaseClass]
+    captured_layers: dict[int, dict[str, torch.Tensor]]
+
+    def __init__(self, vllm_config: Any, role: KVConnectorRole, kv_cache_config: Any = None) -> None:  # type: ignore
+        super().__init__(vllm_config, role, kv_cache_config)  # pyright: ignore[reportUnknownMemberType]
+        self.captured_layers = _shared_captured_layers
+
+    def start_load_kv(self, forward_context: Any, **kwargs: Any) -> None:  # pyright: ignore[reportAny]
+        pass
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        pass
+
+    def save_kv_layer(self, layer_name: str, kv_layer: Any, attn_metadata: Any, **kwargs: Any) -> None:  # pyright: ignore[reportAny]
+        slot_mapping = getattr(attn_metadata, "slot_mapping", None)  # pyright: ignore[reportAny]
+        if slot_mapping is not None and slot_mapping.shape[0] <= 100:  # pyright: ignore[reportAny]
+            return
+
+        m = _LAYER_RE.search(layer_name)
+        if m is None:
+            return
+        layer_idx = int(m.group(1))
+
+        if isinstance(kv_layer, (list, tuple)):
+            from exo.disaggregated.streaming_connector import _to_bf16
+
+            _shared_captured_arrays[layer_idx] = [_to_bf16(t).cpu() for t in kv_layer]  # pyright: ignore[reportAny]
+            return
+
+        if slot_mapping is not None:
+            from exo.disaggregated.streaming_connector import _to_nhd
+
+            if kv_layer.shape[0] == 2:  # pyright: ignore[reportAny]
+                k_all = _to_nhd(kv_layer[0])  # pyright: ignore[reportAny]
+                v_all = _to_nhd(kv_layer[1])  # pyright: ignore[reportAny]
+            else:
+                k_all = _to_nhd(kv_layer[:, 0])  # pyright: ignore[reportAny]
+                v_all = _to_nhd(kv_layer[:, 1])  # pyright: ignore[reportAny]
+            k_flat = k_all.reshape(-1, *k_all.shape[-2:])  # pyright: ignore[reportAny]
+            v_flat = v_all.reshape(-1, *v_all.shape[-2:])  # pyright: ignore[reportAny]
+            valid = slot_mapping >= 0  # pyright: ignore[reportAny]
+            safe_sm = slot_mapping.clamp(min=0)  # pyright: ignore[reportAny]
+            keys = k_flat[safe_sm][valid]  # pyright: ignore[reportAny]
+            values = v_flat[safe_sm][valid]  # pyright: ignore[reportAny]
+            from exo.disaggregated.streaming_connector import _to_bf16
+
+            keys = _to_bf16(keys)  # pyright: ignore[reportAny]
+            values = _to_bf16(values)  # pyright: ignore[reportAny]
+
+            prev = self.captured_layers.get(layer_idx)
+            if prev is not None:
+                self.captured_layers[layer_idx] = {
+                    "keys": torch.cat([prev["keys"], keys.cpu()], dim=0),  # type: ignore
+                    "values": torch.cat([prev["values"], values.cpu()], dim=0),  # type: ignore
+                }
+            else:
+                self.captured_layers[layer_idx] = {
+                    "keys": keys.cpu(),  # pyright: ignore[reportAny]
+                    "values": values.cpu(),  # pyright: ignore[reportAny]
+                }
+
+    def wait_for_save(self) -> None:
+        pass
+
+    def request_finished_all_groups(self, request: Any, block_ids: tuple[list[int], ...]) -> tuple[bool, dict[str, Any] | None]:  # pyright: ignore[reportAny]
+        return False, None
+
+    def get_num_new_matched_tokens(self, request: Any, num_computed_tokens: int) -> tuple[int, bool]:  # pyright: ignore[reportAny]
+        return 0, False
+
+    def update_state_after_alloc(self, request: Any, blocks: Any, num_external_tokens: int) -> None:  # pyright: ignore[reportAny]
+        pass
+
+    def build_connector_meta(self, scheduler_output: Any) -> BatchConnectorMetadata:  # pyright: ignore[reportAny]
+        return BatchConnectorMetadata()

src/exo/disaggregated/prefill_client.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import json
+import socket
+import time
+from collections import defaultdict
+from collections.abc import Callable
+from typing import TYPE_CHECKING, BinaryIO, cast
+
+import mlx.core as mx
+import torch
+from mlx_lm.models.cache import ArraysCache, KVCache, RotatingKVCache
+
+from exo.disaggregated.protocol import (
+    ArraysState,
+    Done,
+    KVChunk,
+    read_header,
+    read_message,
+)
+
+if TYPE_CHECKING:
+    from exo.shared.types.mlx import Model
+
+from exo.worker.runner.bootstrap import logger
+
+
+def _torch_to_mx(t: torch.Tensor) -> mx.array:
+    t_cpu: torch.Tensor = t.detach().cpu()
+    if t_cpu.dtype == torch.bfloat16:
+        return mx.array(t_cpu.float().numpy()).astype(mx.bfloat16)  # pyright: ignore[reportAny]
+    return mx.array(t_cpu.numpy())  # pyright: ignore[reportAny]
+
+
+def _nhd_to_bhsd(keys: torch.Tensor, values: torch.Tensor) -> tuple[mx.array, mx.array]:
+    k_mx = _torch_to_mx(keys.permute(1, 0, 2).unsqueeze(0))
+    v_mx = _torch_to_mx(values.permute(1, 0, 2).unsqueeze(0))
+    return k_mx, v_mx
+
+
+def _inject_kv_cache(cache: KVCache, keys: torch.Tensor, values: torch.Tensor, num_tokens: int) -> None:
+    k_mx, v_mx = _nhd_to_bhsd(keys, values)
+    cache.keys = k_mx
+    cache.values = v_mx
+    cache.offset = num_tokens
+
+
+def _inject_rotating_kv_cache(cache: RotatingKVCache, keys: torch.Tensor, values: torch.Tensor, num_tokens: int) -> None:
+    k_mx, v_mx = _nhd_to_bhsd(keys, values)
+    seq_len = int(k_mx.shape[2])
+    cache.keys = k_mx
+    cache.values = v_mx
+    cache.offset = num_tokens
+    cache._idx = seq_len
+
+
+def _inject_arrays_cache(cache: ArraysCache, arrays: list[torch.Tensor]) -> None:
+    cache.state = [_torch_to_mx(arr) for arr in arrays]
+
+
+def remote_prefill(
+    endpoint: str,
+    token_ids: list[int],
+    model_id: str,
+    mlx_model: Model,
+    on_prefill_progress: Callable[[int, int], None] | None = None,
+    existing_cache: list[KVCache | RotatingKVCache | ArraysCache] | None = None,
+    start_pos: int = 0,
+) -> tuple[list[KVCache | RotatingKVCache | ArraysCache], int]:
+    if ":" in endpoint:
+        host, port_str = endpoint.rsplit(":", 1)
+        port = int(port_str)
+    else:
+        host = endpoint
+        port = 8900
+
+    logger.info(f"Connecting to prefill server at {host}:{port} ({len(token_ids)} tokens, start_pos={start_pos})")
+    t0 = time.perf_counter()
+
+    sock = socket.create_connection((host, port), timeout=60)
+    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 4 * 1024 * 1024)
+    try:
+        request = json.dumps({"model": model_id, "token_ids": token_ids, "start_pos": start_pos}).encode("utf-8") + b"\n"
+        sock.sendall(request)
+
+        raw_stream = sock.makefile("rb", buffering=256 * 1024)
+        stream: BinaryIO = raw_stream  # pyright: ignore[reportAssignmentType]
+
+        first_byte: bytes = raw_stream.peek(1)[:1]  # type: ignore
+        if first_byte == b"{":
+            line = stream.readline()
+            error_resp: dict[str, object] = json.loads(line.decode("utf-8"))  # pyright: ignore[reportAny]
+            raise RuntimeError(f"Prefill server error: {error_resp.get('error', 'unknown')}")
+
+        header = read_header(stream)
+        num_layers: int = header["num_layers"]  # pyright: ignore[reportAssignmentType]
+        total_prompt_tokens = len(token_ids)
+
+        kv_buffers: dict[int, list[tuple[torch.Tensor, torch.Tensor]]] = defaultdict(list)
+        arrays_buffers: dict[int, list[torch.Tensor]] = {}
+        total_tokens = 0
+        layers_seen: set[int] = set()
+
+        tokens_received = 0
+        chunks_received = 0
+        t_first_chunk = None
+        while True:
+            msg = read_message(stream, header)
+            if msg is None:
+                break
+
+            if isinstance(msg, KVChunk):
+                if t_first_chunk is None:
+                    t_first_chunk = time.perf_counter()
+                kv_buffers[msg.layer_idx].append((msg.keys, msg.values))
+                chunks_received += 1
+                layers_seen.add(msg.layer_idx)
+                tokens_received += msg.num_tokens
+                if on_prefill_progress and num_layers > 0 and chunks_received % num_layers == 0:
+                    on_prefill_progress(
+                        min(tokens_received // num_layers, total_prompt_tokens - start_pos),
+                        total_prompt_tokens - start_pos,
+                    )
+            elif isinstance(msg, ArraysState):
+                arrays_buffers[msg.layer_idx] = msg.arrays
+            elif isinstance(msg, Done):  # pyright: ignore[reportUnnecessaryIsInstance]
+                total_tokens = msg.total_tokens
+                break
+
+        t_received = time.perf_counter()
+    finally:
+        sock.close()
+
+    if existing_cache is not None and start_pos > 0:
+        caches = existing_cache
+    else:
+        if hasattr(mlx_model, "make_cache"):
+            caches = cast(list[KVCache | RotatingKVCache | ArraysCache], mlx_model.make_cache())  # pyright: ignore[reportUnknownMemberType]
+        else:
+            from mlx_lm.models.cache import make_prompt_cache
+            caches = cast(list[KVCache | RotatingKVCache | ArraysCache], make_prompt_cache(mlx_model))  # pyright: ignore[reportUnknownMemberType]
+
+    max_received = max((sum(k.shape[0] for k, _v in chunks) for chunks in kv_buffers.values()), default=0)
+    final_offset = start_pos + max_received
+
+    for i, cache in enumerate(caches):
+        if i in kv_buffers:
+            chunks = kv_buffers[i]
+            all_keys: torch.Tensor
+            all_values: torch.Tensor
+            if len(chunks) == 1:
+                all_keys, all_values = chunks[0]
+            else:
+                all_keys = torch.cat([k for k, _v in chunks], dim=0)  # type: ignore
+                all_values = torch.cat([v for _k, v in chunks], dim=0)  # type: ignore
+
+            if isinstance(cache, RotatingKVCache):
+                _inject_rotating_kv_cache(cache, all_keys, all_values, final_offset)  # pyright: ignore[reportUnknownArgumentType]
+            elif isinstance(cache, KVCache):
+                if start_pos > 0 and cache.keys is not None:
+                    k_new, v_new = _nhd_to_bhsd(all_keys, all_values)  # pyright: ignore[reportUnknownArgumentType]
+                    cache.keys = mx.concatenate([cache.keys[:, :, :start_pos, :], k_new], axis=2)
+                    cache.values = mx.concatenate([cache.values[:, :, :start_pos, :], v_new], axis=2)
+                    cache.offset = final_offset
+                else:
+                    _inject_kv_cache(cache, all_keys, all_values, final_offset)  # pyright: ignore[reportUnknownArgumentType]
+
+        if i in arrays_buffers and isinstance(cache, ArraysCache):
+            _inject_arrays_cache(cache, arrays_buffers[i])
+
+    t_injected = time.perf_counter()
+    logger.info(
+        f"Remote prefill: {total_tokens} new tokens (start_pos={start_pos}, final_offset={final_offset}), "
+        f"transfer={((t_received - t0) * 1000):.0f}ms, "
+        f"inject={((t_injected - t_received) * 1000):.0f}ms, "
+        f"total={((t_injected - t0) * 1000):.0f}ms"
+    )
+
+    return caches, final_offset

src/exo/disaggregated/prefill_server.py

@@ -0,0 +1,665 @@
+from __future__ import annotations
+
+import contextlib
+import json
+import socket
+import socketserver
+import threading
+import time
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from exo.disaggregated.protocol import (
+    write_arrays_state,
+    write_done,
+    write_header,
+    write_kv_chunk,
+)
+
+if TYPE_CHECKING:
+    from vllm.v1.engine.llm_engine import LLMEngine
+
+from exo.worker.engines.mlx.cache import KVPrefixCache
+from exo.worker.engines.vllm.kv_cache import KVLayerState, TorchKVCache
+from exo.worker.runner.bootstrap import logger
+
+_engine_ref: LLMEngine | None = None
+_prefix_cache_ref: KVPrefixCache | None = None
+_overlapping: bool = True
+_on_status_change: Callable[[bool], None] | None = None
+_connector_patched: bool = False
+_gdn_patched: bool = False
+_gdn_states: dict[int, dict[str, torch.Tensor]] = {}
+_gdn_layer_order: list[int] = []
+_gdn_call_idx: list[int] = [0]
+_ssm_call_idx: list[int] = [0]
+
+
+def _patch_vllm_for_connector(connector_class: type[Any]) -> None:  # pyright: ignore[reportUnusedFunction]
+    global _connector_patched
+    if _connector_patched:
+        return
+    _connector_patched = True
+
+    from vllm.v1.core import kv_cache_utils
+
+    original_unify = kv_cache_utils.unify_hybrid_kv_cache_specs  # type: ignore
+
+    def patched_unify(kv_cache_spec: Any) -> None:  # pyright: ignore[reportAny]
+        with contextlib.suppress(ValueError):
+            original_unify(kv_cache_spec)
+
+    kv_cache_utils.unify_hybrid_kv_cache_specs = patched_unify  # pyright: ignore[reportAttributeAccessIssue]
+
+    from vllm.v1.core.sched import (  # pyright: ignore[reportMissingImports]
+        scheduler as sched_mod,  # pyright: ignore[reportUnknownVariableType]
+    )
+
+    def patched_connector_finished(_self: Any, _request: Any) -> tuple[bool, Any]:  # pyright: ignore[reportAny]
+        return False, None
+
+    sched_mod.Scheduler._connector_finished = patched_connector_finished  # pyright: ignore[reportUnknownMemberType]
+
+    from vllm.distributed.kv_transfer.kv_connector import (  # pyright: ignore[reportMissingImports]
+        factory,  # pyright: ignore[reportUnknownVariableType]
+    )
+
+    original_get = factory.KVConnectorFactory._get_connector_class_with_compat  # type: ignore
+
+    @classmethod
+    def patched_get(cls: Any, kv_transfer_config: Any) -> tuple[Any, Any]:  # pyright: ignore[reportAny]
+        kv_conn = getattr(kv_transfer_config, "kv_connector", None) or ""  # pyright: ignore[reportAny]
+        if "streaming_connector" in kv_conn or "batch_connector" in kv_conn:
+            return connector_class, None
+        return original_get.__func__(cls, kv_transfer_config)  # type: ignore
+
+    factory.KVConnectorFactory._get_connector_class_with_compat = patched_get  # pyright: ignore[reportUnknownMemberType]
+
+
+def _patch_gdn_capture() -> None:
+    global _gdn_patched
+    if _gdn_patched:
+        return
+    _gdn_patched = True
+
+    try:
+        import vllm.model_executor.layers.mamba.ops.causal_conv1d as cc_mod  # type: ignore
+        from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
+            causal_conv1d_fn as orig_fn,  # type: ignore
+        )
+    except ImportError:
+        return
+
+    def patched_fn(*args: Any, conv_states: Any = None, cache_indices: Any = None, **kwargs: Any) -> Any:
+        result = orig_fn(*args, conv_states=conv_states, cache_indices=cache_indices, **kwargs)  # type: ignore
+        if conv_states is not None and cache_indices is not None:
+            x = args[0] if args else None
+            if x is not None and x.shape[0] <= 100:  # type: ignore
+                return result
+            ci: int = cache_indices[0].item() if cache_indices.numel() > 0 else 0  # type: ignore
+            idx = _gdn_call_idx[0]
+            if _gdn_layer_order and idx < len(_gdn_layer_order) * 100:
+                layer_idx = _gdn_layer_order[idx % len(_gdn_layer_order)]
+                conv_at_ci = conv_states[ci : ci + 1].transpose(-1, -2).contiguous().cpu()  # type: ignore
+                _gdn_states.setdefault(layer_idx, {})["conv"] = conv_at_ci
+                _gdn_states[layer_idx]["ci"] = ci  # type: ignore
+            _gdn_call_idx[0] += 1
+        return result
+
+    cc_mod.causal_conv1d_fn = patched_fn  # type: ignore
+    import sys
+    for mod in list(sys.modules.values()):
+        if mod is None or mod is cc_mod:
+            continue
+        if hasattr(mod, "causal_conv1d_fn") and mod.causal_conv1d_fn is orig_fn:
+            mod.causal_conv1d_fn = patched_fn
+    logger.info("Patched causal_conv1d_fn for GDN state capture")
+
+    try:
+        from vllm.model_executor.models import qwen3_next as qn_mod  # type: ignore
+
+        orig_chunk = getattr(qn_mod, "fi_chunk_gated_delta_rule", None)  # type: ignore
+        if orig_chunk is None:
+            return
+
+        def patched_chunk(*args: Any, **kwargs: Any) -> Any:
+            result = orig_chunk(*args, **kwargs)
+            output_final_state = kwargs.get("output_final_state", False)
+            if output_final_state and isinstance(result, tuple) and len(result) == 2:
+                _, ssm_state = result
+                idx = _ssm_call_idx[0]
+                if _gdn_layer_order and idx < len(_gdn_layer_order) * 100:
+                    layer_idx = _gdn_layer_order[idx % len(_gdn_layer_order)]
+                    _gdn_states.setdefault(layer_idx, {})["ssm"] = ssm_state.cpu()  # type: ignore
+                _ssm_call_idx[0] += 1
+            return result
+
+        qn_mod.fi_chunk_gated_delta_rule = patched_chunk  # type: ignore
+
+        orig_fla_chunk = getattr(qn_mod, "fla_chunk_gated_delta_rule", None)  # type: ignore
+        if orig_fla_chunk is not None:
+            def patched_fla_chunk(*args: Any, **kwargs: Any) -> Any:
+                result = orig_fla_chunk(*args, **kwargs)
+                output_final_state = kwargs.get("output_final_state", False)
+                if output_final_state and isinstance(result, tuple) and len(result) == 2:
+                    _, ssm_state = result
+                    idx = _ssm_call_idx[0]
+                    if _gdn_layer_order and idx < len(_gdn_layer_order) * 100:
+                        layer_idx = _gdn_layer_order[idx % len(_gdn_layer_order)]
+                        _gdn_states.setdefault(layer_idx, {})["ssm"] = ssm_state.cpu()  # type: ignore
+                    _ssm_call_idx[0] += 1
+                return result
+            qn_mod.fla_chunk_gated_delta_rule = patched_fla_chunk  # type: ignore
+
+        logger.info("Patched chunk_gated_delta_rule for SSM state capture")
+    except ImportError:
+        pass
+
+
+def _init_gdn_layer_order() -> None:
+    from exo.worker.engines.vllm.growable_cache import get_model_runner
+
+    model_runner = get_model_runner()
+    if model_runner is None:
+        return
+    kv_caches = model_runner.kv_caches  # type: ignore
+    _gdn_layer_order.clear()
+    for li in range(len(kv_caches)):  # type: ignore
+        kv = kv_caches[li]  # type: ignore
+        if isinstance(kv, (list, tuple)) and len(kv) > 1:
+            _gdn_layer_order.append(li)
+    if _gdn_layer_order:
+        logger.info(f"GDN layer order: {_gdn_layer_order} ({len(_gdn_layer_order)} layers)")
+
+
+def _get_layer_info(engine: LLMEngine) -> tuple[int, str, list[dict[str, Any]]]:
+    from exo.worker.engines.vllm.growable_cache import get_model_runner
+
+    model_runner = get_model_runner()
+    assert model_runner is not None
+    kv_caches = model_runner.kv_caches
+    num_layers: int = len(kv_caches)
+
+    layers_info: list[dict[str, Any]] = []
+    for li in range(num_layers):
+        kv = kv_caches[li]
+        if isinstance(kv, (list, tuple)) and len(kv) > 1:
+            layers_info.append({"type": "arrays", "sizes": [2]})
+        else:
+            sample = kv[0] if isinstance(kv, (list, tuple)) else kv
+            n_heads: int = sample.shape[-2]
+            head_dim: int = sample.shape[-1]
+            layers_info.append({"type": "kv", "n_heads": n_heads, "head_dim": head_dim})
+
+    dtype_str = "bfloat16"
+    return num_layers, dtype_str, layers_info
+
+
+def _run_prefill_overlapping(engine: LLMEngine, token_ids: list[int], start_pos: int, wfile: Any) -> None:  # pyright: ignore[reportAny]
+    from exo.worker.engines.vllm.growable_cache import get_model_runner
+
+    model_runner = get_model_runner()
+    assert model_runner is not None
+
+    from exo.disaggregated.streaming_connector import (
+        get_shared_arrays_queue,
+        get_shared_queue,
+        reset_shared_queue,
+    )
+
+    reset_shared_queue()
+    _gdn_states.clear()
+    _gdn_call_idx[0] = 0
+    _ssm_call_idx[0] = 0
+    layer_queue = get_shared_queue()
+    arrays_queue = get_shared_arrays_queue()
+
+    server_cached = 0
+    cached_data: TorchKVCache | None = None
+    if _prefix_cache_ref is not None:
+        cached_data, server_cached, _ = _prefix_cache_ref.lookup(token_ids)
+        if not isinstance(cached_data, TorchKVCache):
+            cached_data = None
+            server_cached = 0
+    skip_tokens = max(0, start_pos - server_cached)
+
+    num_layers, dtype_str, layers_info = _get_layer_info(engine)
+    write_header(wfile, {"num_layers": num_layers, "dtype": dtype_str, "layers": layers_info})  # pyright: ignore[reportAny]
+
+    if cached_data is not None and start_pos < server_cached:
+        from exo.worker.engines.vllm.kv_cache import ArraysLayerState
+        kv_sent = 0
+        arr_sent = 0
+        for i, layer in enumerate(cached_data.layers):
+            if isinstance(layer, KVLayerState) and layer.keys.numel() > 0:
+                keys = layer.keys
+                values = layer.values
+                if keys.shape != values.shape:
+                    logger.warning(f"Skipping layer {i}: keys={list(keys.shape)} != values={list(values.shape)}")
+                    continue
+                if keys.dim() == 4:
+                    keys = keys.reshape(-1, keys.shape[-2], keys.shape[-1])
+                    values = values.reshape(-1, values.shape[-2], values.shape[-1])
+                keys = keys[start_pos:server_cached]
+                values = values[start_pos:server_cached]
+                if keys.numel() > 0:
+                    write_kv_chunk(wfile, i, keys, values)  # pyright: ignore[reportAny]
+                    kv_sent += 1
+
+            elif isinstance(layer, ArraysLayerState):
+                arrays = [a for a in layer.arrays if a is not None]
+                if arrays:
+                    write_arrays_state(wfile, i, arrays)  # pyright: ignore[reportAny]
+                    arr_sent += 1
+        logger.info(f"Sent cached: {kv_sent} KV, {arr_sent} arrays for positions {start_pos}-{server_cached}")
+
+    from vllm.sampling_params import (
+        SamplingParams,
+    )
+
+    prefill_token_ids = token_ids[:-2] if len(token_ids) > 2 else token_ids
+    request_id = f"prefill-{time.monotonic_ns()}"
+    params = SamplingParams(max_tokens=2, detokenize=False)  # pyright: ignore[reportCallIssue]
+    engine.add_request(request_id, {"prompt_token_ids": prefill_token_ids}, params)  # pyright: ignore[reportArgumentType]
+
+    chunks_sent = [0]
+    layer_token_counts: dict[int, int] = {}
+    all_kv_chunks: list[tuple[int, torch.Tensor, torch.Tensor]] = []
+
+    def writer_loop() -> None:
+        while True:
+            item = layer_queue.get()
+            if item is None:
+                break
+            layer_idx, keys, values = item
+            all_kv_chunks.append((layer_idx, keys, values))
+
+            prev = layer_token_counts.get(layer_idx, 0)
+            n = keys.shape[0]
+            new_total = prev + n
+            layer_token_counts[layer_idx] = new_total
+
+            if new_total <= skip_tokens:
+                continue
+            if prev < skip_tokens:
+                trim = skip_tokens - prev
+                keys = keys[trim:]
+                values = values[trim:]
+            if chunks_sent[0] == 0:
+                logger.info(f"First KV chunk: layer={layer_idx} keys={keys.shape} keys.dtype={keys.dtype} values.dtype={values.dtype}")
+            write_kv_chunk(wfile, layer_idx, keys, values)  # pyright: ignore[reportAny]
+            chunks_sent[0] += 1
+
+    writer_thread = threading.Thread(target=writer_loop, daemon=True)
+    writer_thread.start()
+
+    while engine.has_unfinished_requests():
+        outputs = engine.step()
+        for output in outputs:
+            if output.request_id == request_id and output.outputs[0].token_ids:
+                engine.abort_request([request_id])  # type: ignore
+                break
+        else:
+            continue
+        break
+
+    layer_queue.put(None)
+    writer_thread.join()
+    actual_per_layer = max(layer_token_counts.values()) if layer_token_counts else 0
+    cached_tokens_sent = max(0, server_cached - start_pos) if cached_data is not None and start_pos < server_cached else 0
+    tokens_sent = cached_tokens_sent + max(0, actual_per_layer - skip_tokens)
+    logger.info(f"Overlapping prefill: sent {chunks_sent[0]} chunks, {tokens_sent} tokens (server_cached={server_cached}, skip={skip_tokens})")
+
+    while not arrays_queue.empty():
+        item = arrays_queue.get_nowait()
+        if item is not None:
+            layer_idx, arrays = item
+            write_arrays_state(wfile, layer_idx, arrays)  # pyright: ignore[reportAny]
+
+    gdn_snapshot: list[tuple[int, list[torch.Tensor]]] = []
+    for layer_idx in sorted(_gdn_states.keys()):
+        state = _gdn_states[layer_idx]
+        arrs: list[torch.Tensor] = []
+        if "conv" in state:
+            arrs.append(state["conv"])
+        if "ssm" in state:
+            arrs.append(state["ssm"])
+        if arrs:
+            gdn_snapshot.append((layer_idx, arrs))
+
+    cached_arrays: list[tuple[int, list[torch.Tensor]]] = []
+    _stream_gdn_states_and_collect(engine, wfile, num_layers, layers_info, cached_arrays)
+    write_done(wfile, tokens_sent)  # pyright: ignore[reportAny]
+
+    connector_cache = _build_torch_cache(all_kv_chunks, gdn_snapshot, num_layers)
+    threading.Thread(target=_store_prefix_cache, args=(prefill_token_ids, connector_cache), daemon=True).start()
+
+
+def _run_prefill_batch(engine: LLMEngine, token_ids: list[int], start_pos: int, wfile: Any) -> None:  # pyright: ignore[reportAny]
+    from exo.worker.engines.vllm.growable_cache import get_model_runner
+
+    num_layers, dtype_str, layers_info = _get_layer_info(engine)
+
+    model_runner = get_model_runner()
+    assert model_runner is not None
+
+    from exo.disaggregated.batch_connector import (
+        clear_shared_captured_layers,
+        get_shared_captured_arrays,
+        get_shared_captured_layers,
+    )
+
+    _gdn_states.clear()
+    _gdn_call_idx[0] = 0
+    clear_shared_captured_layers()
+    captured_layers = get_shared_captured_layers()
+    captured_arrays = get_shared_captured_arrays()
+
+    server_cached = 0
+    if _prefix_cache_ref is not None:
+        _, server_cached, _ = _prefix_cache_ref.lookup(token_ids)
+    skip_tokens = max(0, start_pos - server_cached)
+
+    from vllm.sampling_params import (
+        SamplingParams,
+    )
+
+    prefill_token_ids = token_ids[:-2] if len(token_ids) > 2 else token_ids
+    request_id = f"prefill-{time.monotonic_ns()}"
+    params = SamplingParams(max_tokens=2, detokenize=False)  # pyright: ignore[reportCallIssue]
+    engine.add_request(request_id, {"prompt_token_ids": prefill_token_ids}, params)  # pyright: ignore[reportArgumentType]
+
+    while engine.has_unfinished_requests():
+        outputs = engine.step()
+        for output in outputs:
+            if output.request_id == request_id and output.outputs[0].token_ids:
+                engine.abort_request([request_id])  # type: ignore
+                break
+        else:
+            continue
+        break
+
+    write_header(wfile, {"num_layers": num_layers, "dtype": dtype_str, "layers": layers_info})  # pyright: ignore[reportAny]
+
+    all_kv: list[tuple[int, torch.Tensor, torch.Tensor]] = []
+    for layer_idx in sorted(captured_layers.keys()):
+        layer_data = captured_layers[layer_idx]
+        keys = layer_data["keys"]
+        values = layer_data["values"]
+        all_kv.append((layer_idx, keys, values))
+        if keys.shape[0] > skip_tokens:
+            write_kv_chunk(wfile, layer_idx, keys[skip_tokens:], values[skip_tokens:])  # pyright: ignore[reportAny]
+
+    actual_per_layer = max((k.shape[0] for _, k, _ in all_kv), default=0)
+    tokens_sent = max(0, actual_per_layer - skip_tokens)
+    logger.info(f"Batch prefill: {len(all_kv)} layers, {tokens_sent} tokens sent (server_cached={server_cached}, skip={skip_tokens}, captured={actual_per_layer})")
+
+    batch_arrays: list[tuple[int, list[torch.Tensor]]] = list(captured_arrays.items())
+    for layer_idx, arrs in batch_arrays:
+        write_arrays_state(wfile, layer_idx, arrs)  # pyright: ignore[reportAny]
+    clear_shared_captured_layers()
+
+    cached_arrays: list[tuple[int, list[torch.Tensor]]] = []
+    _stream_gdn_states_and_collect(engine, wfile, num_layers, layers_info, cached_arrays)
+    write_done(wfile, tokens_sent)  # pyright: ignore[reportAny]
+
+    connector_cache = _build_torch_cache(all_kv, batch_arrays, num_layers)
+    threading.Thread(target=_store_prefix_cache, args=(prefill_token_ids, connector_cache), daemon=True).start()
+
+
+def _stream_gdn_states_and_collect(
+    _engine: LLMEngine,
+    wfile: Any,
+    num_layers: int,
+    layers_info: list[dict[str, Any]],
+    out_arrays: list[tuple[int, list[torch.Tensor]]],
+) -> None:  # type: ignore
+    from exo.worker.engines.vllm.growable_cache import get_model_runner
+
+    if not _gdn_states:
+        return
+
+    model_runner = get_model_runner()
+    if model_runner is None:
+        return
+
+    kv_caches = model_runner.kv_caches  # type: ignore
+    torch.cuda.synchronize()
+
+    for layer_idx in sorted(_gdn_states.keys()):
+        try:
+            state = _gdn_states[layer_idx]
+            conv = state.get("conv")
+            ssm = state.get("ssm")
+
+            arrays: list[torch.Tensor] = []
+            if conv is not None:
+                arrays.append(conv)
+            if ssm is not None:
+                arrays.append(ssm)
+            if arrays:
+                write_arrays_state(wfile, layer_idx, arrays)  # type: ignore
+                out_arrays.append((layer_idx, arrays))
+        except Exception:
+            logger.opt(exception=True).warning(f"Failed to capture GDN state for layer {layer_idx}")
+
+    _gdn_states.clear()
+    _gdn_call_idx[0] = 0
+
+
+def _build_torch_cache(kv_chunks: list[tuple[int, torch.Tensor, torch.Tensor]], arrays_chunks: list[tuple[int, list[torch.Tensor]]], num_layers: int) -> TorchKVCache:
+    from exo.worker.engines.vllm.kv_cache import ArraysLayerState
+
+    layers_by_idx: dict[int, KVLayerState | ArraysLayerState] = {}
+    for layer_idx, keys, values in kv_chunks:
+        if layer_idx in layers_by_idx:
+            prev = layers_by_idx[layer_idx]
+            if isinstance(prev, KVLayerState):
+                layers_by_idx[layer_idx] = KVLayerState(
+                    keys=torch.cat([prev.keys, keys], dim=0),  # type: ignore
+                    values=torch.cat([prev.values, values], dim=0),  # type: ignore
+                )
+        else:
+            layers_by_idx[layer_idx] = KVLayerState(keys=keys, values=values)
+    for layer_idx, arrays in arrays_chunks:
+        layers_by_idx[layer_idx] = ArraysLayerState(arrays=[a if isinstance(a, torch.Tensor) else None for a in arrays])
+
+    ordered: list[KVLayerState | ArraysLayerState] = []
+    for i in range(num_layers):
+        if i in layers_by_idx:
+            ordered.append(layers_by_idx[i])
+        else:
+            ordered.append(KVLayerState(keys=torch.empty(0), values=torch.empty(0)))
+    return TorchKVCache(ordered)
+
+
+def _extract_vllm_cache(engine: LLMEngine, request_id: str, num_tokens: int) -> TorchKVCache | None:
+    try:
+        from exo.worker.engines.vllm.vllm_generator import _save_prefix_cache
+        from exo.worker.engines.vllm.growable_cache import get_model_runner
+        from exo.worker.engines.vllm.vllm_generator import _build_layer_groups
+        model_runner = get_model_runner()
+        if model_runner is None:
+            return None
+        engine_core = engine.engine_core.engine_core  # type: ignore
+        coordinator = engine_core.scheduler.kv_cache_manager.coordinator  # type: ignore
+        kv_cache_config = engine_core.scheduler.kv_cache_manager.kv_cache_config  # type: ignore
+
+        internal_id: str | None = None
+        for mgr in coordinator.single_type_managers:  # type: ignore
+            for key in mgr.req_to_blocks:  # type: ignore
+                if str(key).startswith(request_id):  # type: ignore
+                    internal_id = str(key)  # type: ignore
+                    break
+            if internal_id:
+                break
+        if internal_id is None:
+            return None
+
+        null_block = coordinator.block_pool.null_block  # type: ignore
+        block_ids_per_group: list[list[int]] = []
+        token_offset_per_group: list[int] = []
+        block_sizes_per_group: list[int] = []
+        for mgr in coordinator.single_type_managers:  # type: ignore
+            blocks = mgr.req_to_blocks.get(internal_id)  # type: ignore
+            if not blocks:
+                block_ids_per_group.append([])
+                token_offset_per_group.append(0)
+                block_sizes_per_group.append(0)
+                continue
+            block_size: int = mgr.block_size  # type: ignore
+            block_sizes_per_group.append(block_size)
+            num_leading_nulls = 0
+            for b in blocks:  # type: ignore
+                if b is null_block or b.is_null:  # type: ignore
+                    num_leading_nulls += 1
+                else:
+                    break
+            real_blocks = [b for b in blocks if b is not null_block and not b.is_null]  # type: ignore
+            block_ids_per_group.append([b.block_id for b in real_blocks])  # type: ignore
+            token_offset_per_group.append(num_leading_nulls * block_size)
+
+        layer_to_group = _build_layer_groups(kv_cache_config)
+        return TorchKVCache.from_vllm_cache(
+            model_runner.kv_caches,  # type: ignore
+            block_ids_per_group,
+            layer_to_group,
+            num_tokens,
+            token_offset_per_group,
+            block_sizes_per_group,
+        )
+    except Exception:
+        logger.opt(exception=True).warning("Failed to extract vLLM cache")
+        return None
+
+
+def _store_prefix_cache(token_ids: list[int], torch_cache: TorchKVCache) -> None:
+    if _prefix_cache_ref is None:
+        return
+    try:
+        before = len(_prefix_cache_ref.prompts)
+        _prefix_cache_ref.add_from_torch(token_ids, torch_cache)
+        after = len(_prefix_cache_ref.prompts)
+        if after > before:
+            logger.info(f"Server prefix cache: saved {len(token_ids)} tokens (entries: {before} → {after})")
+    except Exception:
+        logger.opt(exception=True).warning("Failed to store prefix cache")
+
+
+def _check_cache(token_ids: list[int]) -> TorchKVCache | None:
+    if _prefix_cache_ref is None:
+        return None
+    import mlx.core as mx
+
+    prompt_arr = mx.array(token_ids)
+    best_index: int | None = None
+    best_length = 0
+    for i, cached_prompt in enumerate(_prefix_cache_ref.prompts):
+        prefix_len = min(len(cached_prompt), len(prompt_arr))
+        if prefix_len == 0:
+            continue
+        match_len = int(mx.sum(cached_prompt[:prefix_len] == prompt_arr[:prefix_len]).item())  # pyright: ignore[reportAny]
+        if match_len == len(token_ids) and match_len == len(cached_prompt) and match_len > best_length:
+            best_index = i
+            best_length = match_len
+
+    if best_index is None:
+        return None
+
+    cached = _prefix_cache_ref.caches[best_index]
+    if isinstance(cached, TorchKVCache):
+        return cached
+    return None
+
+
+def _send_cached(torch_cache: TorchKVCache, token_ids: list[int], wfile: Any, engine: LLMEngine) -> None:
+    num_layers, dtype_str, layers_info = _get_layer_info(engine)
+    write_header(wfile, {"num_layers": num_layers, "dtype": dtype_str, "layers": layers_info})  # type: ignore
+    from exo.worker.engines.vllm.kv_cache import ArraysLayerState
+
+    kv_sent = 0
+    arr_sent = 0
+    for i, layer in enumerate(torch_cache.layers):
+        if isinstance(layer, KVLayerState) and layer.keys.numel() > 0:
+            write_kv_chunk(wfile, i, layer.keys, layer.values)  # type: ignore
+            kv_sent += 1
+        elif isinstance(layer, ArraysLayerState):
+            arrays = [a for a in layer.arrays if a is not None]
+            if arrays:
+                write_arrays_state(wfile, i, arrays)  # type: ignore
+                arr_sent += 1
+    logger.info(f"_send_cached: sent {kv_sent} KV layers, {arr_sent} arrays layers")
+    write_done(wfile, len(token_ids))  # type: ignore
+
+
+class _PrefillHandler(socketserver.StreamRequestHandler):
+    def setup(self) -> None:
+        super().setup()
+        self.request.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)  # type: ignore
+        self.request.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 4 * 1024 * 1024)  # type: ignore
+
+    def handle(self) -> None:
+        try:
+            line = self.rfile.readline()
+            if not line:
+                return
+            request: dict[str, Any] = json.loads(line.decode("utf-8"))  # pyright: ignore[reportAny]
+            token_ids: list[int] = request["token_ids"]  # pyright: ignore[reportAny]
+            start_pos: int = request.get("start_pos", 0)  # pyright: ignore[reportAny]
+
+            engine = _engine_ref
+            if engine is None:
+                error = json.dumps({"error": "No engine loaded"}).encode("utf-8") + b"\n"
+                self.wfile.write(error)
+                return
+
+            if engine.has_unfinished_requests():
+                error = json.dumps({"error": "Engine busy"}).encode("utf-8") + b"\n"
+                self.wfile.write(error)
+                return
+
+            logger.info(f"Prefill request: {len(token_ids)} tokens, start_pos={start_pos}, overlapping={_overlapping}")
+            t0 = time.perf_counter()
+
+            if _on_status_change:
+                _on_status_change(True)
+            try:
+                if _overlapping:
+                    _run_prefill_overlapping(engine, token_ids, start_pos, self.wfile)
+                else:
+                    _run_prefill_batch(engine, token_ids, start_pos, self.wfile)
+            finally:
+                if _on_status_change:
+                    _on_status_change(False)
+
+            elapsed = time.perf_counter() - t0
+            logger.info(f"Prefill complete: {len(token_ids)} tokens in {elapsed*1000:.0f}ms ({len(token_ids)/elapsed:.0f} tok/s)")
+        except Exception:
+            logger.opt(exception=True).error("Prefill handler error")
+
+
+def start_prefill_server(
+    engine: LLMEngine,
+    bind_address: str,
+    port: int,
+    overlapping: bool = True,
+    prefix_cache: KVPrefixCache | None = None,
+    on_status_change: Callable[[bool], None] | None = None,
+) -> socketserver.ThreadingTCPServer:
+    global _engine_ref, _overlapping, _prefix_cache_ref, _on_status_change
+    _engine_ref = engine
+    _overlapping = overlapping
+    _prefix_cache_ref = prefix_cache
+    _on_status_change = on_status_change
+
+    _patch_gdn_capture()
+    _init_gdn_layer_order()
+
+    server = socketserver.ThreadingTCPServer((bind_address, port), _PrefillHandler)
+    server.daemon_threads = True
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    logger.info(f"Prefill TCP server started on {bind_address}:{port} (overlapping={overlapping})")
+    return server

src/exo/disaggregated/protocol.py

@@ -0,0 +1,186 @@
+from __future__ import annotations
+
+import io
+import json
+import struct
+from dataclasses import dataclass
+from typing import BinaryIO
+
+import torch
+
+MSG_KV_CHUNK: int = 0x01
+MSG_ARRAYS_STATE: int = 0x02
+MSG_DONE: int = 0x03
+
+
+@dataclass
+class KVChunk:
+    layer_idx: int
+    num_tokens: int
+    keys: torch.Tensor
+    values: torch.Tensor
+
+
+@dataclass
+class ArraysState:
+    layer_idx: int
+    arrays: list[torch.Tensor]
+
+
+@dataclass
+class Done:
+    total_tokens: int
+
+
+Message = KVChunk | ArraysState | Done
+
+
+def _write_exactly(stream: BinaryIO, data: bytes) -> None:
+    stream.write(data)
+    stream.flush()
+
+
+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 _str_to_dtype(s: str) -> torch.dtype:
+    return {
+        "float16": torch.float16,
+        "bfloat16": torch.bfloat16,
+        "float32": torch.float32,
+    }[s]
+
+
+def _dtype_size(dtype: torch.dtype) -> int:
+    return {torch.float16: 2, torch.bfloat16: 2, torch.float32: 4}[dtype]
+
+
+def write_header(stream: BinaryIO, header: dict[str, object]) -> None:
+    payload = json.dumps(header).encode("utf-8")
+    _write_exactly(stream, struct.pack(">I", len(payload)))
+    _write_exactly(stream, payload)
+
+
+def _tensor_to_bytes(t: torch.Tensor) -> bytes:
+    if t.dtype == torch.bfloat16:
+        return t.contiguous().view(torch.int16).numpy().tobytes()  # type: ignore
+    return t.contiguous().numpy().tobytes()  # type: ignore
+
+
+def write_kv_chunk(stream: BinaryIO, layer_idx: int, keys: torch.Tensor, values: torch.Tensor) -> None:
+    if keys.dim() == 4:
+        keys = keys.reshape(-1, keys.shape[-2], keys.shape[-1])
+        values = values.reshape(-1, values.shape[-2], values.shape[-1])
+    keys_bytes = _tensor_to_bytes(keys)
+    values_bytes = _tensor_to_bytes(values)
+    num_tokens: int = keys.shape[0]
+    n_heads: int = keys.shape[1]
+    head_dim: int = keys.shape[2]
+    header = struct.pack(">BIIII", MSG_KV_CHUNK, layer_idx, num_tokens, n_heads, head_dim)
+    _write_exactly(stream, header + keys_bytes + values_bytes)
+
+
+def _dtype_to_str(dtype: torch.dtype) -> str:
+    return {torch.float16: "float16", torch.bfloat16: "bfloat16", torch.float32: "float32"}[dtype]
+
+
+def write_arrays_state(stream: BinaryIO, layer_idx: int, arrays: list[torch.Tensor]) -> None:
+    buf = io.BytesIO()
+    buf.write(struct.pack(">BI", MSG_ARRAYS_STATE, layer_idx))
+    buf.write(struct.pack(">I", len(arrays)))
+    for arr in arrays:
+        dtype_str = _dtype_to_str(arr.dtype).encode("utf-8")
+        buf.write(struct.pack(">I", len(dtype_str)))
+        buf.write(dtype_str)
+        shape: tuple[int, ...] = tuple(arr.shape)
+        buf.write(struct.pack(">I", len(shape)))
+        for dim in shape:
+            buf.write(struct.pack(">I", dim))
+        buf.write(_tensor_to_bytes(arr))
+    _write_exactly(stream, buf.getvalue())
+
+
+def write_done(stream: BinaryIO, total_tokens: int) -> None:
+    _write_exactly(stream, struct.pack(">BI", MSG_DONE, total_tokens))
+
+
+def read_header(stream: BinaryIO) -> dict[str, object]:
+    raw = _read_exactly(stream, 4)
+    if not raw:
+        raise ConnectionError("No header received")
+    length: int = struct.unpack(">I", raw)[0]  # pyright: ignore[reportAny]
+    payload = _read_exactly(stream, length)
+    return json.loads(payload.decode("utf-8"))  # pyright: ignore[reportAny]
+
+
+def read_message(stream: BinaryIO, header: dict[str, object]) -> Message | None:
+    type_byte = _read_exactly(stream, 1)
+    if not type_byte:
+        return None
+    msg_type = type_byte[0]
+
+    if msg_type == MSG_KV_CHUNK:
+        layer_idx: int
+        num_tokens: int
+        n_heads: int
+        head_dim: int
+        layer_idx, num_tokens, n_heads, head_dim = struct.unpack(">IIII", _read_exactly(stream, 16))  # pyright: ignore[reportAny]
+        dtype = _str_to_dtype(str(header["dtype"]))
+        elem_size = _dtype_size(dtype)
+        tensor_bytes: int = num_tokens * n_heads * head_dim * elem_size
+        keys_raw = _read_exactly(stream, tensor_bytes)
+        values_raw = _read_exactly(stream, tensor_bytes)
+        shape = (num_tokens, n_heads, head_dim)
+        if dtype == torch.bfloat16:
+            keys: torch.Tensor = torch.frombuffer(bytearray(keys_raw), dtype=torch.int16).view(torch.bfloat16).reshape(shape).clone()  # type: ignore
+            values: torch.Tensor = torch.frombuffer(bytearray(values_raw), dtype=torch.int16).view(torch.bfloat16).reshape(shape).clone()  # type: ignore
+        else:
+            keys = torch.frombuffer(bytearray(keys_raw), dtype=dtype).reshape(shape).clone()  # type: ignore
+            values = torch.frombuffer(bytearray(values_raw), dtype=dtype).reshape(shape).clone()  # type: ignore
+        return KVChunk(layer_idx=layer_idx, num_tokens=num_tokens, keys=keys, values=values)  # pyright: ignore[reportUnknownArgumentType]
+
+    if msg_type == MSG_ARRAYS_STATE:
+        arr_layer_idx: int
+        num_arrays: int
+        arr_layer_idx, = struct.unpack(">I", _read_exactly(stream, 4))  # pyright: ignore[reportAny]
+        num_arrays, = struct.unpack(">I", _read_exactly(stream, 4))  # pyright: ignore[reportAny]
+        fallback_dtype = _str_to_dtype(str(header["dtype"]))
+        arrays: list[torch.Tensor] = []
+        for _ in range(num_arrays):
+            dtype_len_raw = _read_exactly(stream, 4)
+            dtype_len: int = struct.unpack(">I", dtype_len_raw)[0]  # pyright: ignore[reportAny]
+            if dtype_len > 0 and dtype_len < 20:
+                dtype_str_bytes = _read_exactly(stream, dtype_len)
+                arr_dtype = _str_to_dtype(dtype_str_bytes.decode("utf-8"))
+            else:
+                arr_dtype = fallback_dtype
+            elem_size = _dtype_size(arr_dtype)
+            ndim: int
+            ndim, = struct.unpack(">I", _read_exactly(stream, 4))  # pyright: ignore[reportAny]
+            shape_arr = struct.unpack(f">{ndim}I", _read_exactly(stream, ndim * 4))
+            total_elems = 1
+            for d in shape_arr:  # pyright: ignore[reportAny]
+                total_elems *= d  # pyright: ignore[reportAny]
+            raw = _read_exactly(stream, total_elems * elem_size)
+            if arr_dtype == torch.bfloat16:
+                t: torch.Tensor = torch.frombuffer(bytearray(raw), dtype=torch.int16).view(torch.bfloat16).reshape(shape_arr).clone()  # type: ignore
+            else:
+                t = torch.frombuffer(bytearray(raw), dtype=arr_dtype).reshape(shape_arr).clone()  # type: ignore
+            arrays.append(t)  # pyright: ignore[reportUnknownArgumentType]
+        return ArraysState(layer_idx=arr_layer_idx, arrays=arrays)
+
+    if msg_type == MSG_DONE:
+        total_tokens: int
+        total_tokens, = struct.unpack(">I", _read_exactly(stream, 4))  # pyright: ignore[reportAny]
+        return Done(total_tokens=total_tokens)
+
+    raise ValueError(f"Unknown message type: {msg_type:#x}")

src/exo/disaggregated/streaming_connector.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+import os
+import queue
+import re
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (  # pyright: ignore[reportMissingImports]
+    KVConnectorBase_V1,  # pyright: ignore[reportUnknownVariableType]
+    KVConnectorMetadata,  # pyright: ignore[reportUnknownVariableType]
+    KVConnectorRole,  # pyright: ignore[reportUnknownVariableType]
+    SupportsHMA,  # pyright: ignore[reportUnknownVariableType]
+)
+
+_LAYER_RE = re.compile(r"layers\.(\d+)\.")
+
+
+def _to_nhd(t: torch.Tensor) -> torch.Tensor:
+    if os.environ.get("VLLM_KV_CACHE_LAYOUT", "HND") == "HND" and t.dim() == 3:
+        return t.permute(1, 0, 2)
+    return t
+
+
+def _to_bf16(t: torch.Tensor) -> torch.Tensor:
+    if t.dtype == torch.uint8:
+        t = t.view(torch.float8_e4m3fn)  # type: ignore
+    if t.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):  # type: ignore
+        return t.to(torch.float32).to(torch.bfloat16)
+    if t.dtype in (torch.bfloat16, torch.float16, torch.float32):
+        return t
+    return t.to(torch.bfloat16)
+
+_shared_queue: queue.Queue[tuple[int, torch.Tensor, torch.Tensor] | None] = queue.Queue()
+_shared_arrays_queue: queue.Queue[tuple[int, list[torch.Tensor]] | None] = queue.Queue()
+
+
+def get_shared_queue() -> queue.Queue[tuple[int, torch.Tensor, torch.Tensor] | None]:
+    return _shared_queue
+
+
+def get_shared_arrays_queue() -> queue.Queue[tuple[int, list[torch.Tensor]] | None]:
+    return _shared_arrays_queue
+
+
+def reset_shared_queue() -> None:
+    while not _shared_queue.empty():
+        try:
+            _shared_queue.get_nowait()
+        except queue.Empty:
+            break
+    while not _shared_arrays_queue.empty():
+        try:
+            _shared_arrays_queue.get_nowait()
+        except queue.Empty:
+            break
+
+
+@dataclass
+class StreamingConnectorMetadata(KVConnectorMetadata):  # pyright: ignore[reportUntypedBaseClass]
+    pass
+
+
+class StreamingConnector(KVConnectorBase_V1, SupportsHMA):  # pyright: ignore[reportUntypedBaseClass]
+    _queue: queue.Queue[tuple[int, torch.Tensor, torch.Tensor] | None]
+
+    _save_count: int = 0
+
+    def __init__(self, vllm_config: Any, role: KVConnectorRole, kv_cache_config: Any = None) -> None:  # type: ignore
+        super().__init__(vllm_config, role, kv_cache_config)  # pyright: ignore[reportUnknownMemberType]
+        self._queue = _shared_queue
+
+    @property
+    def layer_queue(self) -> queue.Queue[tuple[int, torch.Tensor, torch.Tensor] | None]:
+        return self._queue
+
+    def start_load_kv(self, forward_context: Any, **kwargs: Any) -> None:  # pyright: ignore[reportAny]
+        pass
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        pass
+
+    def save_kv_layer(self, layer_name: str, kv_layer: Any, attn_metadata: Any, **kwargs: Any) -> None:  # pyright: ignore[reportAny]
+        slot_mapping = getattr(attn_metadata, "slot_mapping", None)  # pyright: ignore[reportAny]
+        if slot_mapping is not None and slot_mapping.shape[0] <= 100:  # pyright: ignore[reportAny]
+            return
+
+        m = _LAYER_RE.search(layer_name)
+        if m is None:
+            return
+        layer_idx = int(m.group(1))
+
+        if isinstance(kv_layer, (list, tuple)):
+            arrays = [_to_bf16(t).cpu() for t in kv_layer]  # pyright: ignore[reportAny]
+            _shared_arrays_queue.put((layer_idx, arrays))
+            return
+
+        if self._save_count < 1:
+            self._save_count += 1
+
+        if slot_mapping is not None:
+            if kv_layer.shape[0] == 2:  # pyright: ignore[reportAny]
+                k_all = _to_nhd(kv_layer[0])  # pyright: ignore[reportAny]
+                v_all = _to_nhd(kv_layer[1])  # pyright: ignore[reportAny]
+            else:
+                k_all = _to_nhd(kv_layer[:, 0])  # pyright: ignore[reportAny]
+                v_all = _to_nhd(kv_layer[:, 1])  # pyright: ignore[reportAny]
+            k_flat = k_all.reshape(-1, *k_all.shape[-2:])  # pyright: ignore[reportAny]
+            v_flat = v_all.reshape(-1, *v_all.shape[-2:])  # pyright: ignore[reportAny]
+            valid = slot_mapping >= 0  # pyright: ignore[reportAny]
+            safe_sm = slot_mapping.clamp(min=0)  # pyright: ignore[reportAny]
+            keys = k_flat[safe_sm][valid]  # pyright: ignore[reportAny]
+            values = v_flat[safe_sm][valid]  # pyright: ignore[reportAny]
+            keys = _to_bf16(keys)  # pyright: ignore[reportAny]
+            values = _to_bf16(values)  # pyright: ignore[reportAny]
+            self._queue.put((layer_idx, keys.cpu(), values.cpu()))  # pyright: ignore[reportAny]
+        else:
+            self._queue.put((layer_idx, kv_layer.cpu().clone(), kv_layer.cpu().clone()))  # pyright: ignore[reportAny]
+
+    def wait_for_save(self) -> None:
+        pass
+
+    def finish(self) -> None:
+        self._queue.put(None)
+
+    def request_finished_all_groups(self, request: Any, block_ids: tuple[list[int], ...]) -> tuple[bool, dict[str, Any] | None]:  # pyright: ignore[reportAny]
+        return False, None
+
+    def get_num_new_matched_tokens(self, request: Any, num_computed_tokens: int) -> tuple[int, bool]:  # pyright: ignore[reportAny]
+        return 0, False
+
+    def update_state_after_alloc(self, request: Any, blocks: Any, num_external_tokens: int) -> None:  # pyright: ignore[reportAny]
+        pass
+
+    def build_connector_meta(self, scheduler_output: Any) -> StreamingConnectorMetadata:  # pyright: ignore[reportAny]
+        return StreamingConnectorMetadata()

src/exo/main.py

@@ -274,6 +274,10 @@ def main():
         os.environ["EXO_NO_BATCH"] = "1"
         logger.info("Continuous batching disabled (--no-batch)")
 
+    if args.no_overlapping_prefill_sends:
+        os.environ["EXO_NO_OVERLAPPING_PREFILL_SENDS"] = "1"
+        logger.info("Overlapping prefill sends disabled (--no-overlapping-prefill-sends)")
+
     # Set FAST_SYNCH override env var for runner subprocesses
     if args.fast_synch is True:
         os.environ["EXO_FAST_SYNCH"] = "on"
@@ -305,6 +309,7 @@ class Args(CamelCaseModel):
     no_downloads: bool = False
     offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
     no_batch: bool = False
+    no_overlapping_prefill_sends: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
 
     @classmethod
@@ -363,6 +368,11 @@ class Args(CamelCaseModel):
             action="store_true",
             help="Disable continuous batching, use sequential generation",
         )
+        parser.add_argument(
+            "--no-overlapping-prefill-sends",
+            action="store_true",
+            help="Disable overlapping KV transfer during disaggregated prefill",
+        )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(
             "--fast-synch",

src/exo/master/adapters/chat_completions.py

@@ -107,6 +107,7 @@ def chat_request_to_text_generation(
         min_p=request.min_p,
         repetition_penalty=request.repetition_penalty,
         repetition_context_size=request.repetition_context_size,
+        prefill_endpoints=request.prefill_endpoints,
     )
 
 

src/exo/master/api.py

@@ -414,6 +414,7 @@ class API:
                 node_network=self.state.node_network,
                 topology=self.state.topology,
                 current_instances=self.state.instances,
+                node_vllm=self.state.node_vllm,
             )
         except ValueError as exc:
             raise HTTPException(status_code=400, detail=str(exc)) from exc
@@ -451,26 +452,12 @@ class API:
         instance_combinations: list[tuple[Sharding, InstanceMeta, int]] = []
         node_count = len(list(self.state.topology.list_nodes()))
 
-        # QMM is not available on MLX CUDA. Also, VLLM does not support MLX community models
-        is_mlx_community = str(model_card.model_id).startswith("mlx-community/")
-        is_quantized_mlx = is_mlx_community and model_card.quantization in (
-            "4bit",
-            "8bit",
-        )
-        skip_mlx = any(self.state.node_vllm.values()) and is_quantized_mlx
-        is_vllm_compatible_mlx = is_mlx_community and model_card.quantization in (
-            "",
-            "bf16",
-            "fp16",
-        )
-        skip_vllm = is_mlx_community and not is_vllm_compatible_mlx
-        if not skip_mlx:
-            for sharding in (Sharding.Pipeline, Sharding.Tensor):
-                for instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl):
-                    instance_combinations.extend(
-                        [(sharding, instance_meta, i) for i in range(1, node_count + 1)]
-                    )
-        if any(self.state.node_vllm.values()) and not skip_vllm:
+        for sharding in (Sharding.Pipeline, Sharding.Tensor):
+            for instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl):
+                instance_combinations.extend(
+                    [(sharding, instance_meta, i) for i in range(1, node_count + 1)]
+                )
+        if any(self.state.node_vllm.values()):
             instance_combinations.append((Sharding.Pipeline, InstanceMeta.Vllm, 1))
 
         for sharding, instance_meta, min_nodes in instance_combinations:
@@ -484,6 +471,7 @@ class API:
                     ),
                     node_memory=self.state.node_memory,
                     node_network=self.state.node_network,
+                    node_vllm=self.state.node_vllm,
                     topology=self.state.topology,
                     current_instances=self.state.instances,
                     required_nodes=required_nodes,
@@ -753,7 +741,9 @@ class API:
         )
         task_params = task_params.model_copy(update={"model": resolved_model})
 
-        task_params = task_params.model_copy(update={"stream": False, "bench": True})
+        task_params = task_params.model_copy(
+            update={"stream": False, "bench": True, **({"disaggregated_bench": True} if payload.disaggregated else {})}
+        )
 
         command = TextGeneration(task_params=task_params)
         await self._send(command)
@@ -761,20 +751,25 @@ class API:
         return await self._collect_text_generation_with_stats(command.command_id)
 
     async def _resolve_and_validate_text_model(self, model_id: ModelId) -> ModelId:
-        """Validate a text model exists and return the resolved model ID.
+        from exo.shared.models.model_cards import derive_base_model
 
-        Raises HTTPException 404 if no instance is found for the model.
-        """
-        if not any(
+        if any(
             instance.shard_assignments.model_id == model_id
             for instance in self.state.instances.values()
         ):
-            await self._trigger_notify_user_to_download_model(model_id)
-            raise HTTPException(
-                status_code=404,
-                detail=f"No instance found for model {model_id}",
-            )
-        return model_id
+            return model_id
+
+        request_base = derive_base_model(str(model_id))
+        for instance in self.state.instances.values():
+            first_shard = next(iter(instance.shard_assignments.runner_to_shard.values()), None)
+            if first_shard is not None and first_shard.model_card.base_model.lower() == request_base.lower():
+                return instance.shard_assignments.model_id
+
+        await self._trigger_notify_user_to_download_model(model_id)
+        raise HTTPException(
+            status_code=404,
+            detail=f"No instance found for model {model_id}",
+        )
 
     async def _validate_image_model(self, model: ModelId) -> ModelId:
         """Validate model exists and return resolved model ID.

src/exo/master/main.py

@@ -59,7 +59,8 @@ from exo.shared.types.tasks import (
 from exo.shared.types.tasks import (
     TextGeneration as TextGenerationTask,
 )
-from exo.shared.types.worker.instances import InstanceId
+from exo.shared.types.worker.instances import Instance, InstanceId, VllmInstance
+from exo.shared.types.worker.runners import RunnerReady, RunnerRunning
 from exo.utils.channels import Receiver, Sender
 from exo.utils.event_buffer import MultiSourceBuffer
 from exo.utils.task_group import TaskGroup
@@ -93,8 +94,73 @@ class Master:
         self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {}
         self._expected_ranks: dict[TaskId, set[int]] = {}
 
+    def _find_prefill_endpoints(self, decode_instance: Instance, decode_model_base: str) -> list[str]:
+        from exo.master.placement_utils import (
+            _find_ip_prioritised as find_ip_prioritised,  # pyright: ignore[reportPrivateUsage]
+        )
+        from exo.shared.models.model_cards import derive_base_model
+
+        endpoints: list[tuple[int, str]] = []
+        vllm_instance_count = 0
+        for instance in self.state.instances.values():
+            if not isinstance(instance, VllmInstance):
+                continue
+            if instance.instance_id == decode_instance.instance_id:
+                continue
+            vllm_instance_count += 1
+            first_shard = next(iter(instance.shard_assignments.runner_to_shard.values()), None)
+            if first_shard is None:
+                logger.info(f"Prefill routing: VllmInstance {instance.instance_id} has no shards")
+                continue
+            if derive_base_model(first_shard.model_card.base_model).lower() != decode_model_base.lower():
+                logger.info(
+                    f"Prefill routing: VllmInstance {instance.instance_id} base_model "
+                    f"{first_shard.model_card.base_model!r} != decode {decode_model_base!r}"
+                )
+                continue
+
+            pass
+
+            for node_id, runner_id in instance.shard_assignments.node_to_runner.items():
+                runner_status = self.state.runners.get(runner_id)
+                if not isinstance(runner_status, (RunnerReady, RunnerRunning)):
+                    logger.info(f"Prefill routing: runner {runner_id} not ready ({type(runner_status).__name__})")
+                    continue
+                port = runner_status.prefill_server_port
+                if port is None:
+                    logger.info(f"Prefill routing: runner {runner_id} has no prefill_server_port")
+                    continue
+
+                decode_node = next(iter(decode_instance.shard_assignments.node_to_runner.keys()), None)
+                if decode_node is None:
+                    continue
+
+                ip = find_ip_prioritised(decode_node, node_id, self.state.topology, self.state.node_network, ring=True)
+                if ip is None:
+                    logger.info(f"Prefill routing: no IP route from {decode_node} to {node_id}")
+                    continue
+
+                ip_type = "unknown"
+                node_net = self.state.node_network.get(node_id)
+                if node_net:
+                    for iface in node_net.interfaces:
+                        if iface.ip_address == ip:
+                            ip_type = iface.interface_type
+                            break
+                priority = {"thunderbolt": 0, "maybe_ethernet": 1, "ethernet": 2, "wifi": 3, "unknown": 4}.get(ip_type, 4)
+                endpoints.append((priority, f"{ip}:{port}"))
+
+        if not endpoints:
+            logger.info(
+                f"Prefill routing: no endpoints found for base_model={decode_model_base!r} "
+                f"(total VllmInstances in cluster: {vllm_instance_count})"
+            )
+
+        endpoints.sort(key=lambda x: x[0])
+        return [ep for _, ep in endpoints]
+
     async def run(self):
-        logger.info("Starting Master")
+        logger.debug("Starting Master")
 
         try:
             async with self._tg as tg:
@@ -108,14 +174,14 @@ class Master:
             self.command_receiver.close()
 
     async def shutdown(self):
-        logger.info("Stopping Master")
+        logger.debug("Stopping Master")
         self._tg.cancel_tasks()
 
     async def _command_processor(self) -> None:
         with self.command_receiver as commands:
             async for forwarder_command in commands:
                 try:
-                    logger.info(f"Executing command: {forwarder_command.command}")
+                    logger.debug(f"Executing command: {forwarder_command.command}")
 
                     generated_events: list[Event] = []
                     command = forwarder_command.command
@@ -124,19 +190,22 @@ class Master:
                         case TestCommand():
                             pass
                         case TextGeneration():
+                            from exo.shared.models.model_cards import derive_base_model
+
+                            request_base = derive_base_model(str(command.task_params.model))
+
                             for instance in self.state.instances.values():
-                                if (
-                                    instance.shard_assignments.model_id
-                                    == command.task_params.model
-                                ):
-                                    task_count = sum(
-                                        1
-                                        for task in self.state.tasks.values()
-                                        if task.instance_id == instance.instance_id
-                                    )
-                                    instance_task_counts[instance.instance_id] = (
-                                        task_count
-                                    )
+                                exact_match = instance.shard_assignments.model_id == command.task_params.model
+                                first_shard = next(iter(instance.shard_assignments.runner_to_shard.values()), None)
+                                base_match = first_shard is not None and first_shard.model_card.base_model.lower() == request_base.lower()
+                                if not (exact_match or base_match):
+                                    continue
+                                task_count = sum(
+                                    1
+                                    for task in self.state.tasks.values()
+                                    if task.instance_id == instance.instance_id
+                                )
+                                instance_task_counts[instance.instance_id] = task_count
 
                             if not instance_task_counts:
                                 raise ValueError(
@@ -145,12 +214,26 @@ class Master:
 
                             available_instance_ids = sorted(
                                 instance_task_counts.keys(),
-                                key=lambda instance_id: instance_task_counts[
-                                    instance_id
-                                ],
+                                key=lambda instance_id: (
+                                    0 if not isinstance(self.state.instances[instance_id], VllmInstance) else 1,
+                                    instance_task_counts[instance_id],
+                                ),
                             )
 
                             task_id = TaskId()
+                            decode_instance = self.state.instances[available_instance_ids[0]]
+                            logger.info(
+                                f"Decode routing: model={command.task_params.model} base={request_base} "
+                                f"instance={available_instance_ids[0]} type={type(decode_instance).__name__} "
+                                f"candidates={len(instance_task_counts)}"
+                            )
+                            task_params = command.task_params
+                            if not task_params.prefill_endpoints:
+                                prefill_eps = self._find_prefill_endpoints(decode_instance, request_base)
+                                logger.info(f"Prefill endpoints resolved: {prefill_eps}")
+                                if prefill_eps:
+                                    task_params = task_params.model_copy(update={"prefill_endpoints": prefill_eps})
+
                             generated_events.append(
                                 TaskCreated(
                                     task_id=task_id,
@@ -159,7 +242,7 @@ class Master:
                                         command_id=command.command_id,
                                         instance_id=available_instance_ids[0],
                                         task_status=TaskStatus.Pending,
-                                        task_params=command.task_params,
+                                        task_params=task_params,
                                     ),
                                 )
                             )
@@ -294,6 +377,7 @@ class Master:
                                 self.state.instances,
                                 self.state.node_memory,
                                 self.state.node_network,
+                                self.state.node_vllm,
                             )
                             transition_events = get_transition_events(
                                 self.state.instances, placement, self.state.tasks
@@ -375,7 +459,7 @@ class Master:
             for node_id, time in self.state.last_seen.items():
                 now = datetime.now(tz=timezone.utc)
                 if now - time > timedelta(seconds=30):
-                    logger.info(f"Manually removing node {node_id} due to inactivity")
+                    logger.debug(f"Manually removing node {node_id} due to inactivity")
                     await self.event_sender.send(NodeTimedOut(node_id=node_id))
 
             await anyio.sleep(10)

src/exo/master/placement.py

@@ -67,11 +67,44 @@ def place_instance(
     current_instances: Mapping[InstanceId, Instance],
     node_memory: Mapping[NodeId, MemoryUsage],
     node_network: Mapping[NodeId, NodeNetworkInfo],
+    node_vllm: Mapping[NodeId, bool],
     required_nodes: set[NodeId] | None = None,
 ) -> dict[InstanceId, Instance]:
     cycles = topology.get_cycles()
     candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
 
+    # vLLM instances can only be placed on nodes that have vLLM available.
+    # vLLM does not support quantized mlx-community models (only bf16 or unquantized).
+    if command.instance_meta == InstanceMeta.Vllm:
+        is_mlx_community = str(command.model_card.model_id).startswith("mlx-community/")
+        if is_mlx_community and command.model_card.quantization not in ("", "bf16"):
+            raise ValueError("vLLM does not support quantized mlx-community models")
+        candidate_cycles = [
+            cycle
+            for cycle in candidate_cycles
+            if all(node_vllm.get(nid, False) for nid in cycle.node_ids)
+        ]
+
+    # QMM/quantized ops are not available on MLX CUDA — exclude CUDA nodes for quantized MLX models.
+    if command.instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl):
+        if command.model_card.quantization not in ("", "bf16"):
+            candidate_cycles = [
+                cycle
+                for cycle in candidate_cycles
+                if not any(node_vllm.get(nid, False) for nid in cycle.node_ids)
+            ]
+
+    # mlx-community models should prefer Apple Silicon nodes over CUDA nodes.
+    if command.instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl):
+        if str(command.model_card.model_id).startswith("mlx-community/"):
+            apple_silicon_cycles = [
+                cycle
+                for cycle in candidate_cycles
+                if not any(node_vllm.get(nid, False) for nid in cycle.node_ids)
+            ]
+            if apple_silicon_cycles:
+                candidate_cycles = apple_silicon_cycles
+
     # Filter to cycles containing all required nodes (subset matching)
     if required_nodes:
         candidate_cycles = [

src/exo/master/tests/test_placement.py

@@ -137,7 +137,7 @@ def test_get_instance_placements_create_instance(
     topology.add_connection(conn_b_a)
 
     # act
-    placements = place_instance(cic, topology, {}, node_memory, node_network)
+    placements = place_instance(cic, topology, {}, node_memory, node_network, {})
 
     # assert
     assert len(placements) == 1
@@ -179,7 +179,7 @@ def test_get_instance_placements_one_node_exact_fit() -> None:
             tasks=[ModelTask.TextGeneration],
         ),
     )
-    placements = place_instance(cic, topology, {}, node_memory, node_network)
+    placements = place_instance(cic, topology, {}, node_memory, node_network, {})
 
     assert len(placements) == 1
     instance_id = list(placements.keys())[0]
@@ -206,7 +206,7 @@ def test_get_instance_placements_one_node_fits_with_extra_memory() -> None:
             tasks=[ModelTask.TextGeneration],
         ),
     )
-    placements = place_instance(cic, topology, {}, node_memory, node_network)
+    placements = place_instance(cic, topology, {}, node_memory, node_network, {})
 
     assert len(placements) == 1
     instance_id = list(placements.keys())[0]
@@ -235,7 +235,7 @@ def test_get_instance_placements_one_node_not_fit() -> None:
     )
 
     with pytest.raises(ValueError, match="No cycles found with sufficient memory"):
-        place_instance(cic, topology, {}, node_memory, node_network)
+        place_instance(cic, topology, {}, node_memory, node_network, {})
 
 
 def test_get_transition_events_no_change(instance: Instance):
@@ -334,7 +334,7 @@ def test_placement_selects_leaf_nodes(
     cic = place_instance_command(model_card=model_card)
 
     # act
-    placements = place_instance(cic, topology, {}, node_memory, node_network)
+    placements = place_instance(cic, topology, {}, node_memory, node_network, {})
 
     # assert
     assert len(placements) == 1
@@ -422,7 +422,7 @@ def test_tensor_rdma_backend_connectivity_matrix(
     )
 
     # act
-    placements = place_instance(cic, topology, {}, node_memory, node_network)
+    placements = place_instance(cic, topology, {}, node_memory, node_network, {})
 
     # assert
     assert len(placements) == 1

src/exo/shared/models/model_cards.py

@@ -38,6 +38,35 @@ CARD_SEARCH_PATH = [
 
 _card_cache: dict[ModelId, "ModelCard"] = {}
 
+import re
+
+_QUANT_SUFFIXES = re.compile(
+    r"[-_ ](?: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]+))*$",
+    re.IGNORECASE,
+)
+
+
+def _normalize_base_model(s: str) -> str:
+    return s.replace("-", " ").replace("_", " ").replace("  ", " ").strip()
+
+
+def derive_base_model(model_id: str) -> str:
+    short = model_id.split("/")[-1] if "/" in model_id else model_id
+    base = _QUANT_SUFFIXES.sub("", short)
+    return _normalize_base_model(base)
+
+
+def derive_family(model_id: str) -> str:
+    short = model_id.split("/")[-1] if "/" in model_id else model_id
+    short = _QUANT_SUFFIXES.sub("", short).lower().replace("_", "-")
+    parts = re.split(r"[-.]", short)
+    family_parts: list[str] = []
+    for p in parts:
+        if p.isdigit() or re.match(r"^\d+[bm]?$", p, re.IGNORECASE):
+            break
+        family_parts.append(p)
+    return "-".join(family_parts) if family_parts else short
+
 
 async def _refresh_card_cache():
     for path in CARD_SEARCH_PATH:
@@ -93,6 +122,15 @@ class ModelCard(CamelCaseModel):
     uses_cfg: bool = False
     trust_remote_code: bool = True
 
+    @model_validator(mode="after")
+    def _ensure_derived_fields(self) -> "ModelCard":
+        if not self.base_model:
+            self.base_model = derive_base_model(self.model_id)
+        else:
+            stripped = _QUANT_SUFFIXES.sub("", self.base_model)
+            self.base_model = _normalize_base_model(stripped)
+        return self
+
     @field_validator("tasks", mode="before")
     @classmethod
     def _validate_tasks(cls, v: list[str | ModelTask]) -> list[ModelTask]:
@@ -132,6 +170,9 @@ class ModelCard(CamelCaseModel):
         num_layers = config_data.layer_count
         mem_size_bytes = await fetch_safetensors_size(model_id)
 
+        base_model = derive_base_model(model_id)
+        family = (config_data.model_type or "").replace("_", "-")
+
         mc = ModelCard(
             model_id=ModelId(model_id),
             storage_size=mem_size_bytes,
@@ -141,6 +182,8 @@ class ModelCard(CamelCaseModel):
             num_key_value_heads=config_data.num_key_value_heads,
             tasks=[ModelTask.TextGeneration],
             trust_remote_code=False,
+            base_model=base_model,
+            family=family,
         )
         await mc.save_to_custom_dir()
         _card_cache[model_id] = mc
@@ -170,6 +213,7 @@ def is_custom_card(model_id: ModelId) -> bool:
 class ConfigData(BaseModel):
     model_config = {"extra": "ignore"}  # Allow unknown fields
 
+    model_type: str | None = None
     architectures: list[str] | None = None
     hidden_size: Annotated[int, Field(ge=0)] | None = None
     num_key_value_heads: PositiveInt | None = None

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

@@ -8,7 +8,7 @@ def test_apply_runner_shutdown_removes_runner():
     runner_id = RunnerId()
     state = State(runners={runner_id: RunnerIdle()})
 
-    new_state = apply_runner_status_updated(
+    new_state = appprefilly_runner_status_updated(
         RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerShutdown()), state
     )
 

src/exo/shared/types/api.py

@@ -221,10 +221,11 @@ class ChatCompletionRequest(BaseModel):
     tool_choice: str | dict[str, Any] | None = None
     parallel_tool_calls: bool | None = None
     user: str | None = None
+    prefill_endpoints: list[str] | None = None
 
 
 class BenchChatCompletionRequest(ChatCompletionRequest):
-    pass
+    disaggregated: bool = False
 
 
 class AddCustomModelParams(BaseModel):

src/exo/shared/types/text_generation.py

@@ -70,3 +70,5 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     min_p: float | None = None
     repetition_penalty: float | None = None
     repetition_context_size: int | None = None
+    prefill_endpoints: list[str] | None = None
+    disaggregated_bench: bool = False

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

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

src/exo/utils/info_gatherer/system_info.py

@@ -91,31 +91,82 @@ async def _get_interface_types_from_networksetup() -> dict[str, InterfaceType]:
     return types
 
 
+def _classify_unknown_darwin_interface(name: str) -> InterfaceType:
+    if name.lower().startswith("anpi"):
+        return "thunderbolt"
+    return "unknown"
+
+
+async def _get_linux_network_interfaces() -> list[NetworkInterfaceInfo]:
+    import json as _json
+
+    try:
+        result = await run_process(["ip", "-j", "addr", "show"])
+    except (CalledProcessError, FileNotFoundError):
+        return []
+
+    data: list[dict[str, object]] = _json.loads(result.stdout)  # pyright: ignore[reportAny]
+    interfaces: list[NetworkInterfaceInfo] = []
+    for iface in data:
+        name: str = iface.get("ifname", "")  # pyright: ignore[reportAssignmentType, reportAny]
+        link_type: str = iface.get("link_type", "")  # pyright: ignore[reportAssignmentType, reportAny]
+
+        iface_type: InterfaceType
+        if link_type == "loopback":
+            continue
+        elif link_type == "ether":
+            if name.startswith(("wl", "wlan")):
+                iface_type = "wifi"
+            elif name.startswith(("docker", "br-", "veth")):
+                iface_type = "unknown"
+            elif name.startswith(("thunderbolt", "tb", "enx")):
+                iface_type = "thunderbolt"
+            else:
+                iface_type = "ethernet"
+        elif link_type in ("none", "tun"):
+            iface_type = "unknown"
+        else:
+            iface_type = "unknown"
+
+        for addr_info in iface.get("addr_info", []):  # pyright: ignore[reportAny]
+            family: str = addr_info.get("family", "")  # pyright: ignore[reportAny]
+            ip: str = addr_info.get("local", "")  # pyright: ignore[reportAny]
+            if family in ("inet", "inet6") and ip:
+                interfaces.append(NetworkInterfaceInfo(name=name, ip_address=ip, interface_type=iface_type))
+
+    return interfaces
+
+
 async def get_network_interfaces() -> list[NetworkInterfaceInfo]:
     """
-    Retrieves detailed network interface information on macOS.
-    Parses output from 'networksetup -listallhardwareports' and 'ifconfig'
+    Retrieves detailed network interface information on macOS or Linux.
+    On MacOS: parses output from 'networksetup -listallhardwareports' and 'ifconfig'
     to determine interface names, IP addresses, and types (ethernet, wifi, vpn, other).
+    Falls back to using ip -j addr show on other platforms.
     Returns a list of NetworkInterfaceInfo objects.
     """
-    interfaces_info: list[NetworkInterfaceInfo] = []
-    interface_types = await _get_interface_types_from_networksetup()
-
-    for iface, services in psutil.net_if_addrs().items():
-        for service in services:
-            match service.family:
-                case socket.AF_INET | socket.AF_INET6:
-                    interfaces_info.append(
-                        NetworkInterfaceInfo(
-                            name=iface,
-                            ip_address=service.address,
-                            interface_type=interface_types.get(iface, "unknown"),
+    if sys.platform == "darwin":
+        interfaces_info: list[NetworkInterfaceInfo] = []
+        interface_types = await _get_interface_types_from_networksetup()
+        for iface, services in psutil.net_if_addrs().items():
+            for service in services:
+                match service.family:
+                    case socket.AF_INET | socket.AF_INET6:
+                        iface_type = interface_types.get(iface, "unknown")
+                        if iface_type == "unknown":
+                            iface_type = _classify_unknown_darwin_interface(iface)
+                        interfaces_info.append(
+                            NetworkInterfaceInfo(
+                                name=iface,
+                                ip_address=service.address,
+                                interface_type=iface_type,
+                            )
                         )
-                    )
-                case _:
-                    pass
+                    case _:
+                        pass
+        return interfaces_info
 
-    return interfaces_info
+    return await _get_linux_network_interfaces()
 
 
 def _read_dmi_field(name: str) -> str | None:

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

@@ -194,10 +194,18 @@ class KVPrefixCache:
         # This ensures stream_generate always has at least one token to start with
         mlx_cache = self._get_mlx_cache(best_index)
         has_ssm = has_non_kv_caches(mlx_cache)
+        snapshots_available = self._snapshots[best_index] is not None
+
+        if is_exact and has_ssm and not snapshots_available:
+            prompt_cache = deepcopy(mlx_cache)
+            self._access_counter += 1
+            self._last_used[best_index] = self._access_counter
+            remaining = prompt_tokens[best_length:]
+            return prompt_cache, remaining, best_index
+
         target = (max_length - 1) if is_exact and not has_ssm else best_length
         restore_pos, restore_snap = self._get_snapshot(best_index, target)
 
-        # No usable snapshot — need fresh cache
         if restore_snap is None and has_ssm:
             return make_kv_cache(model), prompt_tokens, None
 

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

@@ -0,0 +1,58 @@
+"""Patch mlx_lm's GDN gated_delta_update to match vLLM's float32 precision.
+
+vLLM computes both softplus (gating) and sigmoid (beta) in float32.
+mlx_lm computes them in bfloat16. The precision difference compounds
+through the SSM recurrence over thousands of tokens.
+"""
+
+import sys
+from functools import partial
+from typing import Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+@partial(mx.compile, shapeless=True)
+def _compute_g_f32(A_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array:
+    return mx.exp(
+        -mx.exp(A_log.astype(mx.float32))
+        * nn.softplus((a + dt_bias).astype(mx.float32))
+    )
+
+
+def patch_gdn_softplus() -> None:
+    from mlx_lm.models import gated_delta
+
+    orig_update = gated_delta.gated_delta_update
+    orig_ops = gated_delta.gated_delta_ops
+    orig_kernel = gated_delta.gated_delta_kernel
+
+    def patched_gated_delta_update(
+        q: mx.array,
+        k: mx.array,
+        v: mx.array,
+        a: mx.array,
+        b: mx.array,
+        A_log: mx.array,
+        dt_bias: mx.array,
+        state: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+        use_kernel: bool = True,
+    ) -> Tuple[mx.array, mx.array]:
+        beta = mx.sigmoid(b.astype(mx.float32)).astype(b.dtype)
+        g = _compute_g_f32(A_log, a, dt_bias)
+        if state is None:
+            B, _, Hk, Dk = q.shape
+            Hv, Dv = v.shape[-2:]
+            state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
+
+        return orig_ops(q, k, v, g, beta, state, mask)
+
+    gated_delta.gated_delta_update = patched_gated_delta_update
+
+    for mod in list(sys.modules.values()):
+        if mod is None or mod is gated_delta:
+            continue
+        if getattr(mod, "gated_delta_update", None) is orig_update:
+            mod.gated_delta_update = patched_gated_delta_update

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

@@ -6,7 +6,7 @@ import mlx.core as mx
 from mlx_lm.generate import (
     BatchGenerator as MlxBatchGenerator,
 )
-from mlx_lm.models.cache import RotatingKVCache
+from mlx_lm.models.cache import KVCache, RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper
 
@@ -27,6 +27,7 @@ from exo.shared.types.worker.runner_response import GenerationResponse
 from exo.worker.engines.mlx.cache import (
     CacheSnapshot,
     KVPrefixCache,
+    cache_length,
     encode_prompt,
     make_kv_cache,
 )
@@ -149,16 +150,49 @@ class ExoBatchGenerator:
             top_k=task_params.top_k if task_params.top_k is not None else 0,
         )
 
-        _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,
-        )
+        _prefill_tps: float = 0.0
+        cache_snapshots: list[CacheSnapshot] | None = None
+        used_remote_prefill = False
+        uncached_count = len(prompt_tokens)
+        if (
+            uncached_count > 1000
+            and task_params.prefill_endpoints
+            and (not is_bench or task_params.disaggregated_bench)
+        ):
+            from exo.disaggregated.prefill_client import remote_prefill
+
+            t0 = time.perf_counter()
+            try:
+                injected_cache, total_tokens = remote_prefill(
+                    endpoint=task_params.prefill_endpoints[0],
+                    token_ids=[int(t) for t in all_prompt_tokens.tolist()],  # type: ignore
+                    model_id=str(task_params.model),
+                    mlx_model=self.model,
+                    on_prefill_progress=on_prefill_progress,
+                    existing_cache=list(cache) if prefix_hit_length > 0 else None,
+                    start_pos=cache_length(cache) if prefix_hit_length > 0 else 0,
+                )
+                cache = injected_cache
+                from exo.worker.engines.mlx.cache import snapshot_ssm_states
+
+                cache_snapshots = [snapshot_ssm_states(cache)]
+                _prefill_tps = total_tokens / max(time.perf_counter() - t0, 0.001)
+                used_remote_prefill = True
+                logger.info(f"Remote prefill: {total_tokens} tokens at {_prefill_tps:.0f} tok/s")
+            except Exception:
+                logger.opt(exception=True).warning("Remote prefill failed, falling back to local")
+
+        if not used_remote_prefill:
+            _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,
+            )
 
         # We need to clamp rotating kv caches to max size so that mlx lm's _merge_caches behaves
         for c in cache:
@@ -182,7 +216,7 @@ class ExoBatchGenerator:
                 matched_index,
             )
 
-        last_tokens = prompt_tokens[-2:]
+        last_tokens = mx.array(all_prompt_tokens[-2:]) if used_remote_prefill else prompt_tokens[-2:]
 
         logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = (
             make_logits_processors(

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

@@ -551,7 +551,7 @@ def apply_chat_template(
         )
         if partial_assistant_content:
             prompt += partial_assistant_content
-        logger.info(prompt)
+        logger.debug(prompt)
         return prompt
 
     extra_kwargs: dict[str, Any] = {}
@@ -588,7 +588,7 @@ def apply_chat_template(
     if partial_assistant_content:
         prompt += partial_assistant_content
 
-    logger.info(prompt)
+    logger.debug(prompt)
 
     return prompt
 

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

@@ -0,0 +1,104 @@
+"""Patch mlx_lm's YarnRoPE to match vLLM's inverse-frequency blending formula.
+
+mlx_lm's YarnRoPE uses a harmonic blend of frequencies. vLLM uses a linear blend
+of inverse frequencies. These produce different rotation angles, causing KV cache
+mismatch in disaggregated prefill. This patch replaces the frequency computation
+to match vLLM exactly, including support for the `truncate` parameter.
+"""
+
+import math
+
+import mlx.core as mx
+from mlx_lm.models import rope_utils
+
+
+_original_YarnRoPE_init = rope_utils.YarnRoPE.__init__
+
+
+def _patched_yarn_init(
+    self,  # type: ignore
+    dims,  # type: ignore
+    traditional=False,
+    max_position_embeddings=2048,
+    base=10000,
+    scaling_factor=1.0,
+    original_max_position_embeddings=4096,
+    beta_fast=32,
+    beta_slow=1,
+    mscale=1,
+    mscale_all_dim=0,
+    truncate=True,
+) -> None:
+    super(rope_utils.YarnRoPE, self).__init__()
+
+    def yarn_find_correction_dim(num_rotations: float) -> float:
+        return (dims * math.log(original_max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base))
+
+    def yarn_find_correction_range() -> tuple[float, float]:
+        low: float = yarn_find_correction_dim(beta_fast)
+        high: float = yarn_find_correction_dim(beta_slow)
+        if truncate:
+            low = math.floor(low)
+            high = math.ceil(high)
+        return max(low, 0), min(high, dims - 1)
+
+    def yarn_get_mscale(scale: float = 1, ms: float = 1) -> float:
+        if scale <= 1:
+            return 1.0
+        return 0.1 * ms * math.log(scale) + 1.0
+
+    def yarn_linear_ramp_mask(min_val: float, max_val: float, dim: int) -> mx.array:
+        if min_val == max_val:
+            max_val += 0.001
+        linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (max_val - min_val)
+        return mx.clip(linear_func, 0, 1)
+
+    self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(scaling_factor, mscale_all_dim)
+    pos_freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
+    inv_freq_extrapolation = 1.0 / pos_freqs
+    inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
+    low, high = yarn_find_correction_range()
+    inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dims // 2)
+    inv_freq = inv_freq_interpolation * (1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask
+    self._freqs = 1.0 / inv_freq
+    self.dims = dims
+    self.traditional = traditional
+
+
+def _patched_initialize_rope(
+    dims: int,
+    base: float,
+    traditional: bool,
+    scaling_config: dict | None = None,
+    max_position_embeddings: int | None = None,
+) -> object:  # type: ignore
+    if scaling_config is not None:
+        rope_type = scaling_config.get("type") or scaling_config.get("rope_type", "default")
+    else:
+        rope_type = "default"
+
+    if rope_type in ("yarn", "deepseek_yarn", "telechat3-yarn"):
+        scaling_factor = scaling_config["factor"]  # type: ignore
+        rope_kwargs = {
+            key: scaling_config[key]  # type: ignore
+            for key in ["original_max_position_embeddings", "beta_fast", "beta_slow", "mscale", "mscale_all_dim", "truncate"]
+            if key in scaling_config  # type: ignore
+        }
+        return rope_utils.YarnRoPE(
+            dims=dims,
+            max_position_embeddings=max_position_embeddings,
+            traditional=traditional,
+            scaling_factor=scaling_factor,
+            base=base,
+            **rope_kwargs,
+        )
+
+    return _original_initialize_rope(dims, base, traditional, scaling_config, max_position_embeddings)
+
+
+_original_initialize_rope = rope_utils.initialize_rope
+
+
+def patch_yarn_rope() -> None:
+    rope_utils.YarnRoPE.__init__ = _patched_yarn_init  # type: ignore
+    rope_utils.initialize_rope = _patched_initialize_rope  # type: ignore

src/exo/worker/engines/vllm/growable_cache.py

@@ -1,3 +1,6 @@
+import os
+import sys
+
 import torch
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
@@ -46,6 +49,7 @@ def patch_vllm() -> None:
     _patch_get_computed_blocks()
     _patch_moe_sum()
     _patch_marlin_w2_thread_config()
+    _patch_flashinfer_sinks()
     logger.info("vLLM growable KV cache patch applied")
 
 
@@ -56,17 +60,25 @@ def _patch_determine_available_memory() -> None:
 
     @torch.inference_mode()
     def patched(self: "Worker") -> int:
+        import pathlib
+        import shutil
+
+        compile_cache = pathlib.Path.home() / ".cache" / "vllm" / "torch_compile_cache"
+        if compile_cache.exists():
+            shutil.rmtree(compile_cache, ignore_errors=True)
+
+        real_empty_cache = torch.cuda.empty_cache
+        torch.cuda.empty_cache = lambda: None  # type: ignore
         try:
             original(self)
-        except AssertionError:
-            logger.warning(
-                "vLLM memory profiling assertion failed (free memory changed during init, "
-                "likely another process released GPU memory). Continuing with growable cache."
-            )
-        torch.cuda.empty_cache()
+        except (AssertionError, Exception):
+            pass
+        finally:
+            torch.cuda.empty_cache = real_empty_cache  # type: ignore
         free_bytes, _ = torch.cuda.mem_get_info()
         initial = max(int(free_bytes * INITIAL_FRACTION), 1)
         self._growable_max_kv_bytes = free_bytes
+        self.available_kv_cache_memory_bytes = initial
         logger.info(
             f"Growable KV cache: initial {initial / (1024**3):.2f} GiB "
             f"(max {free_bytes / (1024**3):.2f} GiB)"
@@ -115,8 +127,17 @@ def _patch_initialize_kv_cache_tensors() -> None:
 
 def _patch_initialize_from_config() -> None:
     from vllm.v1.worker.gpu_worker import Worker
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
     original = Worker.initialize_from_config
+    original_init_attn = GPUModelRunner.initialize_attn_backend
+
+    def idempotent_init_attn(self: "GPUModelRunner", *args: object, **kwargs: object) -> None:
+        if len(self.attn_groups) > 0:
+            return
+        original_init_attn(self, *args, **kwargs)
+
+    GPUModelRunner.initialize_attn_backend = idempotent_init_attn  # type: ignore[reportAttributeAccessIssue]
 
     def patched(self: "Worker", kv_cache_config: "object") -> None:
         original(self, kv_cache_config)
@@ -164,12 +185,10 @@ def _try_grow_cache(kv_cache_manager: "object") -> bool:
     model_runner = kv_cache_manager._growable_model_runner  # type: ignore
 
     if model_runner is None:
-        logger.debug("No model_runner reference — cannot grow cache")
         return False
 
     free_bytes, _ = torch.cuda.mem_get_info()
     if free_bytes < GROWTH_HEADROOM_BYTES:
-        logger.debug(f"Only {free_bytes / (1024**3):.2f} GiB free — not enough to grow")
         return False
 
     kv_cache_config = model_runner._growable_kv_cache_config  # type: ignore
@@ -182,7 +201,6 @@ def _try_grow_cache(kv_cache_manager: "object") -> bool:
     growth_blocks = min(usable_bytes // per_block_bytes, old_num_blocks)
 
     if growth_blocks < MIN_GROWTH_BLOCKS:
-        logger.debug(f"Growth too small ({growth_blocks} blocks)")
         return False
 
     new_num_blocks = old_num_blocks + growth_blocks
@@ -193,11 +211,11 @@ def _try_grow_cache(kv_cache_manager: "object") -> bool:
     )
 
     try:
-        _grow_tensors(model_runner, kv_cache_config, old_num_blocks, new_num_blocks)
-        _grow_block_pool(block_pool, old_num_blocks, new_num_blocks)
         kv_cache_config.num_blocks = new_num_blocks
         for tensor_spec in kv_cache_config.kv_cache_tensors:
             tensor_spec.size = int(tensor_spec.size * new_num_blocks / old_num_blocks)
+        _grow_tensors(model_runner, kv_cache_config, old_num_blocks, new_num_blocks)
+        _grow_block_pool(block_pool, old_num_blocks, new_num_blocks)
         logger.info(f"KV cache grown successfully to {new_num_blocks} blocks")
         return True
     except Exception:
@@ -205,6 +223,8 @@ def _try_grow_cache(kv_cache_manager: "object") -> bool:
         return False
 
 
+
+
 def _grow_tensors(
     model_runner: "object",
     kv_cache_config: "object",
@@ -243,7 +263,6 @@ def _grow_tensors(
         model_runner.compilation_config.static_forward_context
     )  # type: ignore
     runner_kv_caches: list[torch.Tensor] = model_runner.kv_caches  # type: ignore
-    runner_kv_caches.clear()
 
     from collections import defaultdict
 
@@ -258,12 +277,37 @@ def _grow_tensors(
     for ln in new_kv_caches:
         index2name[extract_layer_index(ln, num_attn_module)].append(ln)
 
+    new_ordered: list[torch.Tensor | list[torch.Tensor]] = []
     for layer_index in sorted(index2name.keys()):
         for ln in index2name[layer_index]:
-            runner_kv_caches.append(new_kv_caches[ln])
+            new_ordered.append(new_kv_caches[ln])
 
-    for layer_name, kv_cache in new_kv_caches.items():
-        forward_context[layer_name].kv_cache = [kv_cache]  # type: ignore
+    for i, new_kv in enumerate(new_ordered):
+        if i < len(runner_kv_caches):
+            old_kv = runner_kv_caches[i]
+            if isinstance(old_kv, list) and isinstance(new_kv, list):
+                for j, (old_t, new_t) in enumerate(zip(old_kv, new_kv)):
+                    old_t.set_(new_t.storage(), new_t.storage_offset(), new_t.shape, new_t.stride())  # type: ignore
+            elif isinstance(old_kv, torch.Tensor) and isinstance(new_kv, torch.Tensor):
+                old_kv.set_(new_kv.storage(), new_kv.storage_offset(), new_kv.shape, new_kv.stride())  # type: ignore
+            else:
+                runner_kv_caches[i] = new_kv
+        else:
+            runner_kv_caches.append(new_kv)
+
+    for layer_name, new_kv in new_kv_caches.items():
+        old_kv_list = forward_context[layer_name].kv_cache  # type: ignore
+        if old_kv_list and len(old_kv_list) > 0:
+            old_entry = old_kv_list[0]
+            if isinstance(old_entry, list) and isinstance(new_kv, list):
+                for j, (old_t, new_t) in enumerate(zip(old_entry, new_kv)):
+                    old_t.set_(new_t.storage(), new_t.storage_offset(), new_t.shape, new_t.stride())  # type: ignore
+            elif isinstance(old_entry, torch.Tensor) and isinstance(new_kv, torch.Tensor):
+                old_entry.set_(new_kv.storage(), new_kv.storage_offset(), new_kv.shape, new_kv.stride())  # type: ignore
+            else:
+                forward_context[layer_name].kv_cache = [new_kv]  # type: ignore
+        else:
+            forward_context[layer_name].kv_cache = [new_kv]  # type: ignore
 
 
 def _grow_block_pool(
@@ -411,3 +455,191 @@ def _patch_get_computed_blocks() -> None:
         return self.empty_kv_cache_blocks, num_matched
 
     KVCacheManager.get_computed_blocks = patched  # type: ignore[reportAttributeAccessIssue]
+
+
+def _patch_flashinfer_sinks() -> None:
+    cap = torch.cuda.get_device_capability()
+    if cap[0] < 12:
+        return
+
+    _patch_is_blackwell_class()
+    _patch_flashinfer_compilation_context()
+    _patch_flashinfer_fmha_sm12x()
+    _patch_triton_ptxas_path()
+    logger.info("FlashInfer SM12x Blackwell-class patches applied")
+
+
+def _patch_is_blackwell_class() -> None:
+    from vllm.platforms.cuda import CudaPlatform
+
+    original_is_family = CudaPlatform.is_device_capability_family.__func__  # type: ignore[reportAttributeAccessIssue]
+
+    @classmethod  # type: ignore[reportArgumentType]
+    def _patched_is_family(cls: type, capability: int, device_id: int = 0) -> bool:
+        if capability == 100:
+            current = cls.get_device_capability(device_id=device_id)
+            if current is not None and current.major in (10, 11, 12):
+                return True
+        return original_is_family(cls, capability, device_id)
+
+    CudaPlatform.is_device_capability_family = _patched_is_family  # type: ignore[reportAttributeAccessIssue]
+
+    try:
+        from vllm.v1.attention.backends.flashinfer import FlashInferBackend, FlashInferMetadataBuilder
+    except ImportError:
+        return
+
+    import vllm.utils.flashinfer as fi_mod
+
+    def _supports_sink_sm12x(cls: type) -> bool:
+        return True
+    FlashInferBackend.supports_sink = classmethod(_supports_sink_sm12x)  # type: ignore[reportAttributeAccessIssue]
+
+    original_init = FlashInferMetadataBuilder.__init__
+
+    def _patched_init(self: object, *args: object, **kwargs: object) -> None:
+        try:
+            original_init(self, *args, **kwargs)
+        except NotImplementedError as e:
+            if "attention sinks" not in str(e):
+                raise
+            logger.info("Bypassing FlashInfer sinks check for SM12x (native compiled attention)")
+            _orig = fi_mod.can_use_trtllm_attention
+            fi_mod.can_use_trtllm_attention = lambda *a, **k: True  # type: ignore[reportAttributeAccessIssue]
+            for mod in list(sys.modules.values()):
+                if mod is not None and mod is not fi_mod:
+                    for attr in list(vars(mod)):
+                        if vars(mod)[attr] is _orig:
+                            setattr(mod, attr, fi_mod.can_use_trtllm_attention)
+            original_init(self, *args, **kwargs)
+            fi_mod.can_use_trtllm_attention = _orig  # type: ignore[reportAttributeAccessIssue]
+            for mod in list(sys.modules.values()):
+                if mod is not None and mod is not fi_mod:
+                    for attr in list(vars(mod)):
+                        if vars(mod)[attr] is fi_mod.can_use_trtllm_attention:
+                            setattr(mod, attr, _orig)
+            setattr(self, 'use_trtllm_decode_attention', False)
+
+    FlashInferMetadataBuilder.__init__ = _patched_init  # type: ignore[reportAttributeAccessIssue]
+
+    original_build = FlashInferMetadataBuilder.build
+
+    def _patched_build(self: object, *args: object, **kwargs: object) -> object:
+        real_sinks = getattr(self, 'has_sinks', False)
+        if real_sinks:
+            setattr(self, 'has_sinks', False)
+        try:
+            return original_build(self, *args, **kwargs)
+        finally:
+            if real_sinks:
+                setattr(self, 'has_sinks', True)
+
+    FlashInferMetadataBuilder.build = _patched_build  # type: ignore[reportAttributeAccessIssue]
+
+
+def _patch_flashinfer_compilation_context() -> None:
+    try:
+        from flashinfer.compilation_context import CompilationContext
+    except ImportError:
+        return
+
+    original_get_flags = CompilationContext.get_nvcc_flags_list
+
+    def _patched_get_flags(self: object, supported_major_versions: list[int] | None = None) -> list[str]:
+        if supported_major_versions and 10 in supported_major_versions and 12 not in supported_major_versions:
+            supported_major_versions = list(supported_major_versions) + [12]
+        return original_get_flags(self, supported_major_versions)
+
+    CompilationContext.get_nvcc_flags_list = _patched_get_flags  # type: ignore[reportAttributeAccessIssue]
+    logger.info("Patched FlashInfer CompilationContext: SM12x added to SM10x version filters")
+
+
+def _patch_flashinfer_fmha_sm12x() -> None:
+    try:
+        import flashinfer
+        fi_include = os.path.join(os.path.dirname(flashinfer.__file__), "data", "include", "flashinfer", "trtllm")
+    except ImportError:
+        return
+
+    common_h = os.path.join(fi_include, "common.h")
+    runner_cuh = os.path.join(fi_include, "fmha", "fmhaRunner.cuh")
+    kernels_cuh = os.path.join(fi_include, "fmha", "fmhaKernels.cuh")
+
+    if not os.path.exists(common_h):
+        return
+
+    with open(common_h) as f:
+        content = f.read()
+    needs_patch = False
+    if "kSM_120" not in content:
+        content = content.replace(
+            "constexpr int32_t kSM_103 = 103;",
+            "constexpr int32_t kSM_103 = 103;\nconstexpr int32_t kSM_120 = 120;",
+        )
+        needs_patch = True
+    if "sm > 120 && sm < 130" not in content:
+        content = content.replace(
+            "return sm_major * 10 + sm_minor;",
+            "int sm = sm_major * 10 + sm_minor;\n  if (sm > 120 && sm < 130) sm = 120;\n  return sm;",
+        )
+        needs_patch = True
+    if needs_patch:
+        with open(common_h, "w") as f:
+            f.write(content)
+        logger.info("Patched flashinfer common.h: kSM_120 + SM121->120 normalization")
+
+    with open(runner_cuh) as f:
+        content = f.read()
+    if "kSM_120" not in content:
+        content = content.replace(
+            "mSM == kSM_100 || mSM == kSM_103",
+            "mSM == kSM_100 || mSM == kSM_103 || mSM == kSM_120",
+        )
+        with open(runner_cuh, "w") as f:
+            f.write(content)
+        logger.info("Patched flashinfer fmhaRunner.cuh: added kSM_120")
+
+    with open(kernels_cuh) as f:
+        content = f.read()
+    if "kSM_120" not in content:
+        content = content.replace(
+            "if (gpuSM == kSM_103) {\n    return kernelSM == kSM_100f || kernelSM == kSM_103;",
+            "if (gpuSM == kSM_103 || gpuSM == kSM_120) {\n    return kernelSM == kSM_100f || kernelSM == kSM_103 || kernelSM == kSM_120;",
+        )
+        with open(kernels_cuh, "w") as f:
+            f.write(content)
+        logger.info("Patched flashinfer fmhaKernels.cuh: added kSM_120 kernel dispatch")
+
+    fi_csrc = os.path.join(os.path.dirname(flashinfer.__file__), "data", "csrc")
+    moe_launcher = os.path.join(fi_csrc, "trtllm_fused_moe_kernel_launcher.cu")
+    if os.path.exists(moe_launcher):
+        with open(moe_launcher) as f:
+            content = f.read()
+        if "major, 10) <<" in content and "major, 10) || std::get<0>(device_props) == 12)" not in content:
+            content = content.replace(
+                'TVM_FFI_ICHECK_EQ(major, 10) << "MoE kernel requires 10.x',
+                'TVM_FFI_ICHECK(major == 10 || major == 12) << "MoE kernel requires 10.x/12.x',
+            )
+            content = content.replace(
+                'std::get<0>(device_props) == 10)',
+                'std::get<0>(device_props) == 10 || std::get<0>(device_props) == 12)',
+            )
+            with open(moe_launcher, "w") as f:
+                f.write(content)
+            cache_dir = os.path.expanduser("~/.cache/flashinfer/0.6.6/121a/cached_ops/fused_moe_trtllm_sm100")
+            if os.path.isdir(cache_dir):
+                import shutil
+                shutil.rmtree(cache_dir)
+            logger.info("Patched flashinfer trtllm_fused_moe_kernel_launcher.cu: added SM12x support")
+
+
+def _patch_triton_ptxas_path() -> None:
+    import os
+    import shutil
+
+    if os.environ.get("TRITON_PTXAS_PATH"):
+        return
+    ptxas = shutil.which("ptxas")
+    if ptxas:
+        os.environ["TRITON_PTXAS_PATH"] = ptxas
+        logger.info(f"Set TRITON_PTXAS_PATH={ptxas}")

src/exo/worker/engines/vllm/kv_cache.py

@@ -228,6 +228,7 @@ class TorchKVCache:
         layer_to_group: list[int],
         num_tokens: int,
         token_offset_per_group: list[int] | None = None,
+        block_sizes_per_group: list[int] | None = None,
     ) -> "TorchKVCache":
         block_tables = [
             torch.tensor(ids, dtype=torch.long) for ids in block_ids_per_group
@@ -245,6 +246,18 @@ class TorchKVCache:
                 layers.append(KVLayerState(keys=torch.empty(0), values=torch.empty(0)))
                 continue
 
+            if k_all.dim() >= 4 and len(bt) > 0 and block_sizes_per_group is not None:
+                page_size = k_all.shape[1]
+                sched_block_size = block_sizes_per_group[gi]
+                pages_per_block = sched_block_size // page_size
+                if pages_per_block > 1:
+                    expanded = []
+                    for b in bt.tolist():
+                        start_page = b * pages_per_block
+                        end_page = min(start_page + pages_per_block, k_all.shape[0])
+                        expanded.extend(range(start_page, end_page))
+                    bt = torch.tensor(expanded, dtype=torch.long)
+
             keys = k_all[bt].to("cpu", non_blocking=True)
             values = v_all[bt].to("cpu", non_blocking=True)
             torch.cuda.synchronize()
@@ -265,20 +278,45 @@ class TorchKVCache:
         first = kv_caches[0]
         device = first[0].device if isinstance(first, list) else first.device
         for layer_idx, layer in enumerate(self.layers):
+            if isinstance(layer, ArraysLayerState):
+                gi = layer_to_group[layer_idx]
+                bt = block_tables[gi]
+                kv = kv_caches[layer_idx]
+                if isinstance(kv, list):
+                    for ti, (stored, target) in enumerate(zip(layer.arrays, kv)):
+                        if stored is not None and target is not None:
+                            n = min(len(bt), stored.shape[0])
+                            if n > 0:
+                                target[bt[:n]] = stored[:n].to(device, non_blocking=True)
+                continue
             if not isinstance(layer, KVLayerState):
                 continue
             gi = layer_to_group[layer_idx]
             bt = block_tables[gi]
             kv = kv_caches[layer_idx]
             k_all, v_all = _split_kv(kv)
-            n_blocks = min(len(bt), layer.keys.shape[0])
+
+            keys = layer.keys
+            values = layer.values
+            block_size = k_all.shape[-3] if k_all.dim() >= 3 else k_all.shape[1]
+            needs_reshape = keys.dim() == 3 and keys.shape[1:] != k_all.shape[1:]
+            if needs_reshape:
+                offset = token_offset_per_group[gi] if token_offset_per_group else 0
+                if offset > 0:
+                    keys = keys[offset:]
+                    values = values[offset:]
+                s, h, d = keys.shape
+                pad = (block_size - s % block_size) % block_size
+                if pad > 0:
+                    keys = torch.nn.functional.pad(keys, (0, 0, 0, 0, 0, pad))
+                    values = torch.nn.functional.pad(values, (0, 0, 0, 0, 0, pad))
+                keys = keys.reshape(-1, block_size, h, d)
+                values = values.reshape(-1, block_size, h, d)
+
+            n_blocks = min(len(bt), keys.shape[0])
             if n_blocks > 0:
-                k_all[bt[:n_blocks]] = layer.keys[:n_blocks].to(
-                    device, non_blocking=True
-                )
-                v_all[bt[:n_blocks]] = layer.values[:n_blocks].to(
-                    device, non_blocking=True
-                )
+                k_all[bt[:n_blocks]] = keys[:n_blocks].to(device, non_blocking=True)
+                v_all[bt[:n_blocks]] = values[:n_blocks].to(device, non_blocking=True)
         torch.cuda.synchronize()
 
     def __iter__(self) -> Iterator[LayerState]:

src/exo/worker/engines/vllm/vllm_generator.py

@@ -3,6 +3,7 @@ import math
 import os
 import re
 import sys
+import tempfile
 import time
 from collections.abc import Callable, Generator
 from dataclasses import dataclass, field
@@ -206,15 +207,16 @@ def vllm_generate(
     on_generation_token: Callable[[], None] | None = None,
 ) -> Generator[GenerationResponse, None, None]:
     token_ids, prompt_text, prompt_token_count = format_vllm_prompt(engine, task)
-    logger.info(prompt_text)
+    logger.debug(prompt_text)
     request_id = f"vllm-seq-{time.monotonic_ns()}"
     sampling_params = make_vllm_sampling_params(engine, task, model_id)
     engine.add_request(request_id, {"prompt_token_ids": token_ids}, sampling_params)
 
     tokenizer = engine.get_tokenizer()
     stop_ids = _stop_token_ids(tokenizer, model_id)
+    DEFAULT_PREFILL_STEP_SIZE = 8192
     max_batch_tokens: int = (
-        getattr(engine.model_config, "max_num_batched_tokens", 2048) or 2048
+        getattr(engine.model_config, "max_num_batched_tokens", DEFAULT_PREFILL_STEP_SIZE) or DEFAULT_PREFILL_STEP_SIZE
     )  # type: ignore[reportUnknownMemberType]
     start_time = time.perf_counter()
     first_token_time: float | None = None
@@ -334,7 +336,7 @@ class VllmBatchEngine:
         token_ids, prompt_text, prompt_token_count = format_vllm_prompt(
             self.engine, task_params
         )
-        logger.info(prompt_text)
+        logger.debug(prompt_text)
         sampling_params = make_vllm_sampling_params(
             self.engine, task_params, self.model_id
         )
@@ -554,30 +556,82 @@ def load_vllm_engine(
     trust_remote_code: bool,
     n_layers: int = 1,
     on_layer_loaded: Callable[[int, int], None] | None = None,
+    kv_connector_cls: type[object] | None = None,
 ) -> tuple[LLMEngine, ToolParser | None, KVPrefixCache]:
     patch_vllm()
     _patch_weight_loading_progress()
 
+    if kv_connector_cls is not None:
+        from exo.disaggregated.prefill_server import _patch_vllm_for_connector
+
+        _patch_vllm_for_connector(kv_connector_cls)
+
     os.environ.setdefault("FASTSAFETENSORS_NOGDS", "1")
 
     prefix_cache = KVPrefixCache(group=None)
     set_prefix_cache(prefix_cache)
     set_n_layers(n_layers)
 
-    engine_args = EngineArgs(
-        model=model_path,
-        served_model_name=str(model_id),
-        gpu_memory_utilization=0.05,
-        trust_remote_code=trust_remote_code,
-        load_format="fastsafetensors",
-        enable_prefix_caching=False,
-        attention_backend="TRITON_ATTN",
-        enforce_eager=True,
-        disable_log_stats=True,
-    )
+    kv_transfer_config: dict[str, str] | None = None
+    if kv_connector_cls is not None:
+        kv_transfer_config = {
+            "kv_connector": f"{kv_connector_cls.__module__}:{kv_connector_cls.__name__}",
+            "kv_role": "kv_both",
+        }
 
-    set_weight_loading_callback(on_layer_loaded)
-    engine = LLMEngine.from_engine_args(engine_args)
+    import json
+    from pathlib import Path
+
+    is_nvfp4 = "nvfp4" in model_path.lower() or "nvfp4" in str(model_id).lower()
+    has_mamba = False
+    is_mxfp4 = False
+    config_path = Path(model_path) / "config.json"
+    if config_path.exists():
+        with open(config_path) as f:
+            model_config = json.load(f)
+        text_config = model_config.get("text_config", model_config)
+        has_mamba = "mamba_ssm_dtype" in text_config or "linear_attention" in (text_config.get("layer_types") or [])
+        quant_config = model_config.get("quantization_config") or text_config.get("quantization_config")
+        if quant_config and quant_config.get("quant_method") == "mxfp4":
+            is_mxfp4 = True
+    if is_mxfp4:
+        os.environ.setdefault("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "1")
+    if has_mamba:
+        backends = ["FLASH_ATTN", "TRITON_ATTN"]
+    else:
+        # backends = ["FLASH_ATTN", "TRITON_ATTN"]
+        backends = ["FLASHINFER", "FLASH_ATTN", "TRITON_ATTN"]
+
+    engine: LLMEngine | None = None
+    for backend in backends:
+        try:
+            engine_args = EngineArgs(
+                model=model_path,
+                served_model_name=str(model_id),
+                gpu_memory_utilization=0.05,
+                trust_remote_code=trust_remote_code,
+                load_format="fastsafetensors",
+                enable_prefix_caching=False,
+                attention_backend=backend,
+                compilation_config={},
+                disable_log_stats=True,
+                max_num_batched_tokens=4096,
+                kv_transfer_config=kv_transfer_config,  # type: ignore
+                disable_hybrid_kv_cache_manager=False,
+                kv_cache_dtype="auto",
+                **({"moe_backend": "flashinfer_cutlass"} if is_mxfp4 else {}),
+            )
+
+            set_weight_loading_callback(on_layer_loaded)
+            engine = LLMEngine.from_engine_args(engine_args)
+            logger.info(f"vLLM engine using attention backend: {backend}")
+            break
+        except (ValueError, RuntimeError) as e:
+            logger.warning(f"Attention backend {backend} failed: {e}, trying next")
+            continue
+
+    if engine is None:
+        raise RuntimeError(f"No attention backend worked for {model_id}")
 
     tool_parser: ToolParser | None = None
     tokenizer = engine.get_tokenizer()

src/exo/worker/main.py

@@ -1,3 +1,4 @@
+import json
 from collections import defaultdict
 from datetime import datetime, timezone
 
@@ -7,7 +8,7 @@ from loguru import logger
 
 from exo.download.download_utils import resolve_model_in_path
 from exo.shared.apply import apply
-from exo.shared.models.model_cards import ModelId
+from exo.shared.models.model_cards import ModelId, derive_base_model
 from exo.shared.types.api import ImageEditsTaskParams
 from exo.shared.types.commands import (
     ForwarderCommand,
@@ -92,6 +93,7 @@ class Worker:
                 tg.start_soon(self.plan_step)
                 tg.start_soon(self._event_applier)
                 tg.start_soon(self._poll_connection_updates)
+                tg.start_soon(self._update_prefill_endpoints)
         finally:
             # Actual shutdown code - waits for all tasks to complete before executing.
             logger.info("Stopping Worker")
@@ -130,6 +132,57 @@ class Worker:
                         event.chunk.data
                     )
 
+    _IFACE_PRIORITY = {"ethernet": 0, "maybe_ethernet": 1, "wifi": 2, "unknown": 3, "thunderbolt": 4}
+
+    def _best_ip_for_node(self, node_id: NodeId) -> str | None:
+        net = self.state.node_network.get(node_id)
+        if not net or not net.interfaces:
+            return None
+        candidates = [
+            iface for iface in net.interfaces
+            if iface.ip_address not in ("127.0.0.1", "::1") and not iface.ip_address.startswith("fe80:")
+        ]
+        if not candidates:
+            return None
+        candidates.sort(key=lambda i: self._IFACE_PRIORITY.get(i.interface_type, 3))
+        return candidates[0].ip_address
+
+    async def _update_prefill_endpoints(self) -> None:
+        while True:
+            await anyio.sleep(5)
+            try:
+                for runner_sup in self.runners.values():
+                    instance = runner_sup.bound_instance.instance
+                    my_model_id = instance.shard_assignments.model_id
+                    my_runner_id = runner_sup.bound_instance.bound_runner_id
+
+                    endpoints: list[dict[str, object]] = []
+                    for rid, status in self.state.runners.items():
+                        if rid == my_runner_id:
+                            continue
+                        port = getattr(status, "prefill_server_port", None)
+                        if not port:
+                            continue
+                        for other_inst in self.state.instances.values():
+                            if rid not in other_inst.shard_assignments.runner_to_shard:
+                                continue
+                            other_base = derive_base_model(other_inst.shard_assignments.model_id)
+                            my_base = derive_base_model(my_model_id)
+                            if other_base != my_base:
+                                continue
+                            for node_id in other_inst.shard_assignments.node_to_runner:
+                                ip = self._best_ip_for_node(node_id)
+                                if ip:
+                                    endpoints.append({"host": ip, "port": port})
+
+                    safe_model = str(my_model_id).replace("/", "--")
+                    # TODO: Change this to be in the task with a list of optional prefill endpoints.
+                    path = f"/tmp/exo_prefill_endpoints_{safe_model}.json"
+                    with open(path, "w") as f:
+                        json.dump(endpoints, f)
+            except:
+                logger.warning("Updating prefill endpoints failed")
+
     async def plan_step(self):
         while True:
             await anyio.sleep(0.1)

src/exo/worker/runner/bootstrap.py

@@ -2,6 +2,7 @@ import ctypes
 import os
 import resource
 import sys
+import urllib.request
 from pathlib import Path
 
 import loguru
@@ -14,6 +15,9 @@ from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
 
 logger: "loguru.Logger" = loguru.logger
 
+_TIKTOKEN_BASE_URL = "https://openaipublic.blob.core.windows.net/encodings"
+_TIKTOKEN_FILES = ["o200k_base.tiktoken", "cl100k_base.tiktoken"]
+
 _CUDA_HOST_LIBS = ["libcuda.so.1", "libnvidia-ml.so.1", "libnvidia-ptxjitcompiler.so.1"]
 _CUDA_HOST_SEARCH_DIRS = [
     Path("/usr/lib/aarch64-linux-gnu"),
@@ -25,6 +29,28 @@ _CUDA_HOST_SEARCH_DIRS = [
 ]
 
 
+def _ensure_tiktoken_encodings() -> None:
+    if os.environ.get("TIKTOKEN_ENCODINGS_BASE"):
+        return
+    from exo.shared.constants import EXO_CACHE_HOME
+
+    enc_dir = EXO_CACHE_HOME / "encodings"
+    enc_dir.mkdir(parents=True, exist_ok=True)
+    for fname in _TIKTOKEN_FILES:
+        dest = enc_dir / fname
+        if dest.exists():
+            continue
+        url = f"{_TIKTOKEN_BASE_URL}/{fname}"
+        logger.info(f"Downloading {url} -> {dest}")
+        try:
+            urllib.request.urlretrieve(url, dest)
+        except Exception:
+            logger.warning(f"Failed to download {fname}, harmony encoding may fail")
+            return
+    os.environ["TIKTOKEN_ENCODINGS_BASE"] = str(enc_dir)
+    logger.info(f"Set TIKTOKEN_ENCODINGS_BASE={enc_dir}")
+
+
 def _ensure_cuda_libs() -> None:
     if sys.platform != "linux":
         return
@@ -65,13 +91,25 @@ def entrypoint(
 
     logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")
 
+    from exo.worker.engines.mlx.yarn_rope_patch import patch_yarn_rope
+
+    patch_yarn_rope()
+
+    from exo.worker.engines.mlx.gdn_softplus_patch import patch_gdn_softplus
+
+    patch_gdn_softplus()
+
     # Import main after setting global logger - this lets us just import logger from this module
     try:
         if isinstance(bound_instance.instance, VllmInstance):
             os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
-            os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
-            os.environ["VLLM_BATCH_INVARIANT"] = "1"
+            import torch
+            cap = torch.cuda.get_device_capability()
+            layout = "HND" if cap[0] >= 12 else "NHD"
+            os.environ["VLLM_KV_CACHE_LAYOUT"] = layout
+            # os.environ["VLLM_BATCH_INVARIANT"] = "1"
             _ensure_cuda_libs()
+            _ensure_tiktoken_encodings()
             from exo.shared.constants import EXO_MODELS_DIR
             from exo.worker.runner.llm_inference.runner import Runner, VllmBuilder
 

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

@@ -293,7 +293,10 @@ class SequentialGenerator(InferenceGenerator):
         )
 
     def close(self) -> None:
-        del self.tokenizer, self.group
+        if hasattr(self, "tokenizer"):
+            del self.tokenizer
+        if hasattr(self, "group"):
+            del self.group
 
 
 @dataclass(eq=False)
@@ -507,4 +510,7 @@ class BatchGenerator(InferenceGenerator):
 
     def close(self) -> None:
         self._gen.close()
-        del self.tokenizer, self.group
+        if hasattr(self, "tokenizer"):
+            del self.tokenizer
+        if hasattr(self, "group"):
+            del self.group

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

@@ -137,6 +137,7 @@ class Runner:
             TaskId,
             TextGeneration,
         ] = {}
+        self.prefill_server_port: int | None = None
 
         logger.info("runner created")
         self.update_status(RunnerIdle())
@@ -237,7 +238,9 @@ class Runner:
                     on_timeout=on_model_load_timeout,
                     on_layer_loaded=on_layer_loaded,
                 )
+                builder_ref = self.generator
                 self.generator = self.generator.build()
+                self.prefill_server_port = getattr(builder_ref, "_prefill_server_port", None)
 
                 self.send_task_status(task.task_id, TaskStatus.Complete)
                 self.update_status(RunnerLoaded())
@@ -257,7 +260,7 @@ class Runner:
                 )
 
                 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):
@@ -342,7 +345,7 @@ class Runner:
             except WouldBlock:
                 pass
 
-        self.update_status(RunnerReady())
+        self.update_status(RunnerReady(prefill_server_port=self.prefill_server_port))
         logger.info("runner ready")
 
         return ExitCode.AllTasksComplete
@@ -529,12 +532,23 @@ class VllmBuilder(Builder):
     ) -> None:
         from exo.worker.engines.vllm.vllm_generator import load_vllm_engine
 
+        kv_connector_cls: type[object] | None = None
+        overlapping = not os.environ.get("EXO_NO_OVERLAPPING_PREFILL_SENDS")
+        if overlapping:
+            from exo.disaggregated.streaming_connector import StreamingConnector
+            kv_connector_cls = StreamingConnector
+        else:
+            from exo.disaggregated.batch_connector import BatchConnector
+            kv_connector_cls = BatchConnector
+
+        self._bound_runner_id = bound_instance.bound_runner_id
         self._engine, self._tool_parser, self._prefix_cache = load_vllm_engine(
             model_path=self.model_path,
             model_id=self.model_id,
             trust_remote_code=self.trust_remote_code,
             n_layers=bound_instance.bound_shard.model_card.n_layers,
             on_layer_loaded=on_layer_loaded,
+            kv_connector_cls=kv_connector_cls,
         )
 
     def build(self) -> InferenceGenerator:
@@ -548,6 +562,39 @@ class VllmBuilder(Builder):
         tokenizer = TokenizerWrapper(self._engine.get_tokenizer())
         max_concurrent = 1 if os.environ.get("EXO_NO_BATCH") else 8
 
+        from exo.master.placement import random_ephemeral_port
+
+        prefill_port = random_ephemeral_port()
+        overlapping = not os.environ.get("EXO_NO_OVERLAPPING_PREFILL_SENDS")
+        try:
+            from exo.disaggregated.prefill_server import start_prefill_server
+
+            from exo.shared.types.events import RunnerStatusUpdated
+            from exo.shared.types.worker.runners import RunnerReady, RunnerRunning
+
+            runner_id = self._bound_runner_id
+
+            def _on_prefill_status(running: bool) -> None:
+                port = prefill_port
+                if running:
+                    self.event_sender.send(RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerRunning(prefill_server_port=port)))
+                else:
+                    self.event_sender.send(RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerReady(prefill_server_port=port)))
+
+            self._prefill_server = start_prefill_server(
+                engine=self._engine,
+                bind_address="0.0.0.0",
+                port=prefill_port,
+                overlapping=overlapping,
+                prefix_cache=self._prefix_cache,
+                on_status_change=_on_prefill_status,
+            )
+            self._prefill_server_port = prefill_port
+        except Exception:
+            logger.opt(exception=True).warning("Failed to start prefill server")
+            self._prefill_server = None
+            self._prefill_server_port = None
+
         logger.info(f"using BatchGenerator (vLLM, max_concurrent={max_concurrent})")
         return BatchGenerator(
             tokenizer=tokenizer,
@@ -564,4 +611,6 @@ class VllmBuilder(Builder):
 
     def close(self) -> None:
         with contextlib.suppress(NameError, AttributeError):
+            if hasattr(self, "_prefill_server") and self._prefill_server is not None:
+                self._prefill_server.shutdown()
             del self._engine, self._prefix_cache, self._tool_parser

src/exo/worker/runner/runner_supervisor.py

@@ -81,7 +81,14 @@ class RunnerSupervisor:
         task_sender, task_recv = mp_channel[Task]()
         cancel_sender, cancel_recv = mp_channel[TaskId]()
 
-        runner_process = mp.Process(
+        from exo.shared.types.worker.instances import VllmInstance
+
+        # vLLM runners use "spawn" to avoid inheriting the parent's CUDA state.
+        # With "fork", the parent's partial CUDA init (from device detection) is
+        # inherited by the child, which conflicts with torch.compile's inductor
+        # backend (cudagraph_mode=none) and causes CUDA illegal instruction errors.
+        ctx = mp.get_context("spawn") if isinstance(bound_instance.instance, VllmInstance) else mp
+        runner_process = ctx.Process(
             target=entrypoint,
             args=(
                 bound_instance,

test_growable_compile.py

@@ -0,0 +1,98 @@
+"""Test hybrid prefix cache: _extract_vllm_cache for attn + captured SSM for mamba."""
+import os, time
+os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
+from exo.worker.engines.vllm.growable_cache import patch_vllm, set_prefix_cache, get_model_runner
+patch_vllm()
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+KVConnectorFactory.register_connector("StreamingConnector", "exo.disaggregated.streaming_connector", "StreamingConnector")
+from vllm.engine.arg_utils import EngineArgs
+from vllm.sampling_params import SamplingParams
+from vllm.v1.engine.llm_engine import LLMEngine
+
+MODEL = os.path.expanduser("~/.local/share/exo/models/Sehyo--Qwen3.5-35B-A3B-NVFP4")
+GEN = 600
+ea = EngineArgs(model=MODEL, served_model_name="test", gpu_memory_utilization=0.05, trust_remote_code=False,
+    load_format="fastsafetensors", enable_prefix_caching=True, attention_backend="FLASH_ATTN",
+    compilation_config={"cudagraph_mode": "none"}, disable_log_stats=True, max_num_batched_tokens=4096,
+    kv_transfer_config={"kv_connector": "StreamingConnector", "kv_role": "kv_both"},
+    disable_hybrid_kv_cache_manager=False)
+engine = LLMEngine.from_engine_args(ea)
+tok = engine.get_tokenizer()
+from exo.worker.engines.mlx.cache import KVPrefixCache
+pc = KVPrefixCache(group=None)
+set_prefix_cache(pc)
+
+from exo.disaggregated.prefill_server import (
+    _patch_gdn_capture, _init_gdn_layer_order, _gdn_states, _gdn_call_idx, _ssm_call_idx,
+    _extract_vllm_cache,
+)
+from exo.disaggregated.streaming_connector import reset_shared_queue
+_patch_gdn_capture()
+_init_gdn_layer_order()
+print(f"Engine loaded")
+
+article = ("The European Union announced sweeping new regulations on artificial intelligence. " * 500)
+tids = tok.encode(article)[:22000]
+msgs = [{"role": "user", "content": tok.decode(tids) + "\nSummarize the key points of this article."}]
+tids = tok.encode(tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True))
+ptids = tids[:-2]
+print(f"Prompt: {len(ptids)} tokens")
+
+reset_shared_queue()
+_gdn_states.clear()
+_gdn_call_idx[0] = 0
+_ssm_call_idx[0] = 0
+
+engine.add_request("r1", {"prompt_token_ids": ptids}, SamplingParams(max_tokens=2, temperature=0.7))
+done = False
+tc = None
+while engine.has_unfinished_requests() and not done:
+    for out in engine.step():
+        if out.outputs and out.outputs[0].token_ids:
+            tc = _extract_vllm_cache(engine, "r1", len(ptids))
+            engine.abort_request(["r1"])
+            done = True; break
+print(f"Extracted: {tc.num_layers if tc else 'NONE'} layers")
+
+if tc and _gdn_states:
+    from exo.worker.engines.vllm.kv_cache import ArraysLayerState, KVLayerState
+    replaced = 0
+    for layer_idx in sorted(_gdn_states.keys()):
+        state = _gdn_states[layer_idx]
+        arrays = []
+        if "conv" in state: arrays.append(state["conv"])
+        if "ssm" in state: arrays.append(state["ssm"])
+        if arrays and layer_idx < len(tc.layers):
+            tc.layers[layer_idx] = ArraysLayerState(arrays=arrays)
+            replaced += 1
+    print(f"Replaced {replaced} GDN layers with clean prefill state")
+    kv_c = sum(1 for l in tc.layers if isinstance(l, KVLayerState) and l.keys.numel() > 0)
+    arr_c = sum(1 for l in tc.layers if isinstance(l, ArraysLayerState))
+    print(f"Final cache: {kv_c} KV layers, {arr_c} Arrays layers")
+
+import mlx.core as mx
+pc.add_kv_cache(mx.array(ptids), tc, None)
+print("Stored hybrid cache")
+
+engine.add_request("r2", {"prompt_token_ids": ptids}, SamplingParams(max_tokens=GEN, temperature=0.7))
+t2 = time.perf_counter()
+prev = 0; text2 = ""; done2 = False
+while engine.has_unfinished_requests() and not done2:
+    for out in engine.step():
+        if out.outputs:
+            prev = len(out.outputs[0].token_ids)
+            if out.outputs[0].text: text2 = out.outputs[0].text
+        if out.finished: done2 = True; break
+e2 = time.perf_counter() - t2
+print(f"\nRequest 2: {prev} tokens in {e2:.1f}s ({prev/max(e2,0.01):.1f} tok/s)")
+print(f"Output: {text2[:500]}")
+
+keywords = ["regulation", "AI", "high-risk", "compliance", "transparency", "ban", "EU", "framework"]
+hits = sum(1 for kw in keywords if kw.lower() in text2.lower())
+print(f"\nKeyword hits: {hits}/{len(keywords)}")
+if hits >= 2:
+    print("PASS")
+else:
+    print(f"FAIL ({hits} hits)")
+    exit(1)

tmp/bench_overnight.log

@@ -0,0 +1,82 @@
+=== Starting overnight bench runs at Tue Mar 10 22:27:20 GMT 2026 ===
+--- [1/8] Qwen3.5-27B-GPTQ-Int4 ---
+2026-03-10 22:27:22.370 | INFO     | __main__:main:300 - pp/tg mode: combinations (product) - 2 pairs
+You are using a model of type qwen3_5 to instantiate a model of type . This is not supported for all configurations of models and can yield errors.
+2026-03-10 22:27:23.659 | DEBUG    | __main__:main:317 - [exo-bench] loaded tokenizer: mlx-community/Qwen3.5-27B-GPTQ-Int4 for prompt sizer
+2026-03-10 22:27:23.661 | DEBUG    | __main__:main:339 - exo-bench model: short_id=Qwen3.5-27B-GPTQ-Int4 full_id=mlx-community/Qwen3.5-27B-GPTQ-Int4
+2026-03-10 22:27:23.661 | INFO     | __main__:main:340 - placements: 1
+2026-03-10 22:27:23.661 | INFO     | __main__:main:342 -   - Pipeline / MlxRing / nodes=1
+2026-03-10 22:27:23.661 | INFO     | __main__:main:353 - Planning phase: checking downloads...
+2026-03-10 22:27:23.670 | INFO     | harness:run_planning_phase:415 - Started download on 12D3KooWGXXhpS3kzjfDVuBGX8AeARLjVdAFaDouYJtXDVXkyq7f
+2026-03-10 22:27:23.674 | INFO     | __main__:main:365 - Download: model already cached
+2026-03-10 22:27:23.675 | INFO     | __main__:main:377 - ================================================================================
+2026-03-10 22:27:23.675 | INFO     | __main__:main:378 - PLACEMENT: Pipeline / MlxRing / nodes=1 / instance_id=725d8b5f-cc83-4dd9-8a4e-c6e2bc5b0607
+2026-03-10 22:27:31.871 | INFO     | __main__:main:409 - --- pp=700 tg=32067 concurrency=1 ---
+2026-03-10 22:27:34.896 | INFO     | __main__:build:224 - tok=700
+2026-03-10 23:00:45.925 | INFO     | __main__:main:519 - prompt_tps=11.76 gen_tps=16.13    prompt_tokens=700 gen_tokens=32067    peak_memory=31.81GB
+
+2026-03-10 23:00:47.935 | INFO     | __main__:main:409 - --- pp=700 tg=33085 concurrency=1 ---
+2026-03-10 23:00:50.948 | INFO     | __main__:build:224 - tok=700
+2026-03-10 23:35:06.460 | INFO     | __main__:main:519 - prompt_tps=11.87 gen_tps=16.12    prompt_tokens=700 gen_tokens=33085    peak_memory=31.89GB
+
+2026-03-10 23:35:08.978 | DEBUG    | __main__:main:532 - Deleted instance 725d8b5f-cc83-4dd9-8a4e-c6e2bc5b0607
+2026-03-10 23:35:13.985 | DEBUG    | __main__:main:541 - 
+Wrote results JSON: bench/results.json
+--- [2/8] NVIDIA-Nemotron-3-Nano-30B-A3B (1120,1330,23100) ---
+2026-03-10 23:35:17.156 | INFO     | __main__:main:300 - pp/tg mode: combinations (product) - 3 pairs
+2026-03-10 23:35:18.698 | DEBUG    | __main__:main:317 - [exo-bench] loaded tokenizer: mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-BF16 for prompt sizer
+2026-03-10 23:35:18.701 | DEBUG    | __main__:main:339 - exo-bench model: short_id=NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-BF16 full_id=mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-BF16
+2026-03-10 23:35:18.701 | INFO     | __main__:main:340 - placements: 1
+2026-03-10 23:35:18.701 | INFO     | __main__:main:342 -   - Pipeline / MlxRing / nodes=1
+2026-03-10 23:35:18.701 | INFO     | __main__:main:353 - Planning phase: checking downloads...
+2026-03-10 23:35:18.710 | INFO     | harness:run_planning_phase:415 - Started download on 12D3KooWGXXhpS3kzjfDVuBGX8AeARLjVdAFaDouYJtXDVXkyq7f
+2026-03-10 23:35:18.716 | INFO     | __main__:main:365 - Download: model already cached
+2026-03-10 23:35:18.716 | INFO     | __main__:main:377 - ================================================================================
+2026-03-10 23:35:18.716 | INFO     | __main__:main:378 - PLACEMENT: Pipeline / MlxRing / nodes=1 / instance_id=fb4d3883-2570-42e9-b3ee-aa1e4a1f8121
+2026-03-10 23:35:43.422 | INFO     | __main__:main:409 - --- pp=700 tg=1120 concurrency=1 ---
+2026-03-10 23:35:46.449 | INFO     | __main__:build:224 - tok=700
+2026-03-10 23:36:05.631 | INFO     | __main__:main:519 - prompt_tps=42.56 gen_tps=62.91    prompt_tokens=700 gen_tokens=1120    peak_memory=64.47GB
+
+2026-03-10 23:36:07.636 | INFO     | __main__:main:409 - --- pp=700 tg=1330 concurrency=1 ---
+2026-03-10 23:36:10.654 | INFO     | __main__:build:224 - tok=700
+2026-03-10 23:36:33.088 | INFO     | __main__:main:519 - prompt_tps=42.28 gen_tps=62.93    prompt_tokens=700 gen_tokens=1330    peak_memory=64.47GB
+
+2026-03-10 23:36:35.097 | INFO     | __main__:main:409 - --- pp=700 tg=23100 concurrency=1 ---
+2026-03-10 23:36:38.107 | INFO     | __main__:build:224 - tok=700
+2026-03-10 23:43:00.943 | INFO     | __main__:main:519 - prompt_tps=42.50 gen_tps=60.57    prompt_tokens=700 gen_tokens=23100    peak_memory=64.47GB
+
+2026-03-10 23:43:03.952 | DEBUG    | __main__:main:532 - Deleted instance fb4d3883-2570-42e9-b3ee-aa1e4a1f8121
+2026-03-10 23:43:08.954 | DEBUG    | __main__:main:541 - 
+Wrote results JSON: bench/results.json
+--- [3/8] Qwen3.5-35B-A3B-bf16 ---
+2026-03-11 10:33:18.576 | INFO     | __main__:main:300 - pp/tg mode: combinations (product) - 4 pairs
+2026-03-11 10:33:18.581 | INFO     | harness:resolve_model_short_id:187 - Model not in /models, adding from HuggingFace: mlx-community/Qwen3.5-35B-A3B-bf16
+You are using a model of type qwen3_5_moe to instantiate a model of type . This is not supported for all configurations of models and can yield errors.
+2026-03-11 10:33:22.909 | DEBUG    | __main__:main:317 - [exo-bench] loaded tokenizer: mlx-community/Qwen3.5-35B-A3B-bf16 for prompt sizer
+2026-03-11 10:33:22.913 | DEBUG    | __main__:main:339 - exo-bench model: short_id=Qwen3.5-35B-A3B-bf16 full_id=mlx-community/Qwen3.5-35B-A3B-bf16
+2026-03-11 10:33:22.913 | INFO     | __main__:main:340 - placements: 1
+2026-03-11 10:33:22.913 | INFO     | __main__:main:342 -   - Pipeline / MlxRing / nodes=1
+2026-03-11 10:33:22.913 | INFO     | __main__:main:353 - Planning phase: checking downloads...
+2026-03-11 10:33:22.923 | INFO     | harness:run_planning_phase:415 - Started download on 12D3KooWGXXhpS3kzjfDVuBGX8AeARLjVdAFaDouYJtXDVXkyq7f
+2026-03-11 10:44:01.613 | INFO     | __main__:main:363 - Download: 638.7s (freshly downloaded)
+2026-03-11 10:44:01.613 | INFO     | __main__:main:377 - ================================================================================
+2026-03-11 10:44:01.613 | INFO     | __main__:main:378 - PLACEMENT: Pipeline / MlxRing / nodes=1 / instance_id=3a5aa42f-a0e1-4273-b61e-56b4e34e0c1d
+2026-03-11 10:44:27.758 | INFO     | __main__:main:409 - --- pp=700 tg=6200 concurrency=1 ---
+2026-03-11 10:44:30.785 | INFO     | __main__:build:224 - tok=700
+2026-03-11 10:46:16.814 | INFO     | __main__:main:519 - prompt_tps=39.17 gen_tps=59.35    prompt_tokens=700 gen_tokens=6200    peak_memory=70.10GB
+
+2026-03-11 10:46:18.816 | INFO     | __main__:main:409 - --- pp=700 tg=6450 concurrency=1 ---
+2026-03-11 10:46:21.834 | INFO     | __main__:build:224 - tok=700
+2026-03-11 10:48:11.923 | INFO     | __main__:main:519 - prompt_tps=38.95 gen_tps=59.50    prompt_tokens=700 gen_tokens=6450    peak_memory=70.10GB
+
+2026-03-11 10:48:13.927 | INFO     | __main__:main:409 - --- pp=700 tg=25600 concurrency=1 ---
+2026-03-11 10:48:16.940 | INFO     | __main__:build:224 - tok=700
+2026-03-11 10:55:53.583 | INFO     | __main__:main:519 - prompt_tps=39.36 gen_tps=56.27    prompt_tokens=700 gen_tokens=25600    peak_memory=70.34GB
+
+2026-03-11 10:55:55.585 | INFO     | __main__:main:409 - --- pp=700 tg=38000 concurrency=1 ---
+2026-03-11 10:55:58.608 | INFO     | __main__:build:224 - tok=700
+2026-03-11 11:07:35.629 | INFO     | __main__:main:519 - prompt_tps=39.43 gen_tps=54.66    prompt_tokens=700 gen_tokens=38000    peak_memory=70.61GB
+
+2026-03-11 11:07:38.598 | DEBUG    | __main__:main:532 - Deleted instance 3a5aa42f-a0e1-4273-b61e-56b4e34e0c1d
+2026-03-11 11:07:43.607 | DEBUG    | __main__:main:541 -
+Wrote results JSON: bench/results.json