Use HND layout

Prevents malformed output from Qwen-Image

.cuda_typings/openai_harmony/__init__.pyi

@@ -0,0 +1,20 @@
+from enum import Enum
+
+class HarmonyEncodingName(Enum):
+    HARMONY_GPT_OSS = ...
+
+class HarmonyEncoding: ...
+class HarmonyError(Exception): ...
+
+class Role(Enum):
+    ASSISTANT = ...
+
+class StreamableParser:
+    last_content_delta: str
+    current_channel: str | None
+    current_recipient: str | None
+
+    def __init__(self, encoding: HarmonyEncoding, role: Role = ...) -> None: ...
+    def process(self, token_id: int) -> None: ...
+
+def load_harmony_encoding(name: HarmonyEncodingName) -> HarmonyEncoding: ...

.cuda_typings/pynvml/__init__.pyi

@@ -0,0 +1,17 @@
+class NvmlMemoryInfo:
+    used: int
+    total: int
+    free: int
+
+class NvmlUtilizationRates:
+    gpu: int
+    memory: int
+
+def nvmlInit() -> None: ...
+def nvmlShutdown() -> None: ...
+def nvmlDeviceGetCount() -> int: ...
+def nvmlDeviceGetHandleByIndex(index: int) -> object: ...
+def nvmlDeviceGetUtilizationRates(handle: object) -> NvmlUtilizationRates: ...
+def nvmlDeviceGetTemperature(handle: object, sensor_type: int) -> int: ...
+def nvmlDeviceGetPowerUsage(handle: object) -> int: ...
+def nvmlDeviceGetMemoryInfo(handle: object) -> NvmlMemoryInfo: ...

.cuda_typings/torch/__init__.pyi

@@ -0,0 +1,61 @@
+from typing import Any, Sequence
+
+from torch import backends as backends
+from torch import cuda as cuda
+from torch import distributed as distributed
+
+__version__: str
+
+class version:
+    cuda: str
+
+class dtype: ...
+
+bfloat16: dtype
+float16: dtype
+float32: dtype
+int8: dtype
+int32: dtype
+int64: dtype
+long: dtype
+float8_e4m3fn: dtype
+
+class Tensor:
+    shape: Sequence[int]
+    dtype: dtype
+    def __getitem__(self, key: Any) -> Tensor: ...
+    def __setitem__(self, key: Any, value: Any) -> None: ...
+    def to(self, *args: Any, **kwargs: Any) -> Tensor: ...
+    def cpu(self) -> Tensor: ...
+    def detach(self) -> Tensor: ...
+    def clone(self) -> Tensor: ...
+    def flatten(self, start_dim: int = 0, end_dim: int = -1) -> Tensor: ...
+    def view(self, *shape: Any) -> Tensor: ...
+    def squeeze(self, dim: int = ...) -> Tensor: ...
+    def unsqueeze(self, dim: int) -> Tensor: ...
+    def permute(self, *dims: int) -> Tensor: ...
+    def float(self) -> Tensor: ...
+    def numpy(self) -> Any: ...
+    def numel(self) -> int: ...
+    def nelement(self) -> int: ...
+    @property
+    def is_cuda(self) -> bool: ...
+    @property
+    def device(self) -> device: ...
+    def __len__(self) -> int: ...
+    def data_ptr(self) -> int: ...
+    def tolist(self) -> Any: ...
+    def abs(self) -> Tensor: ...
+    def max(self) -> Tensor: ...
+    def mean(self) -> Tensor: ...
+    def sum(self, dim: int = ...) -> Tensor: ...
+    def item(self) -> float: ...
+
+def tensor(data: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
+def zeros(*size: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
+def empty(*size: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
+def from_numpy(ndarray: Any) -> Tensor: ...
+def inference_mode() -> Any: ...
+
+class device:
+    def __init__(self, type: str, index: int = ...) -> None: ...

.cuda_typings/torch/backends/__init__.pyi

@@ -0,0 +1 @@
+from torch.backends import cuda as cuda

.cuda_typings/torch/backends/cuda/__init__.pyi

@@ -0,0 +1 @@
+def is_built() -> bool: ...

.cuda_typings/torch/cuda/__init__.pyi

@@ -0,0 +1,10 @@
+class _DeviceProperties:
+    total_memory: int
+
+def is_available() -> bool: ...
+def get_device_name(device: int) -> str: ...
+def get_device_properties(device: int) -> _DeviceProperties: ...
+def empty_cache() -> None: ...
+def mem_get_info() -> tuple[int, int]: ...
+def synchronize() -> None: ...
+def max_memory_allocated() -> int: ...

.cuda_typings/torch/distributed/__init__.pyi

@@ -0,0 +1,2 @@
+def is_initialized() -> bool: ...
+def destroy_process_group() -> None: ...

.cuda_typings/vllm/__init__.pyi

@@ -0,0 +1 @@
+__version__: str

.cuda_typings/vllm/config.pyi

@@ -0,0 +1,2 @@
+class ModelConfig:
+    max_model_len: int

.cuda_typings/vllm/engine/arg_utils.pyi

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+
+@dataclass
+class EngineArgs:
+    model: str = ...
+    served_model_name: str | list[str] | None = ...
+    tokenizer: str | None = ...
+    trust_remote_code: bool = ...
+    dtype: str = ...
+    seed: int = ...
+    max_model_len: int | None = ...
+    gpu_memory_utilization: float = ...
+    enforce_eager: bool = ...
+    tensor_parallel_size: int = ...
+    pipeline_parallel_size: int = ...
+    quantization: str | None = ...
+    load_format: str = ...
+    enable_sleep_mode: bool = ...

.cuda_typings/vllm/outputs.pyi

@@ -0,0 +1,17 @@
+class CompletionOutput:
+    index: int
+    text: str
+    token_ids: list[int]
+    cumulative_logprob: float | None
+    logprobs: object | None
+    finish_reason: str | None
+    stop_reason: int | str | None
+
+    def finished(self) -> bool: ...
+
+class RequestOutput:
+    request_id: str
+    prompt: str | None
+    prompt_token_ids: list[int] | None
+    outputs: list[CompletionOutput]
+    finished: bool

.cuda_typings/vllm/sampling_params.pyi

@@ -0,0 +1,11 @@
+class SamplingParams:
+    n: int
+    temperature: float
+    top_p: float
+    top_k: int
+    min_p: float
+    seed: int | None
+    stop: str | list[str] | None
+    max_tokens: int | None
+    logprobs: int | None
+    repetition_penalty: float

.cuda_typings/vllm/tokenizers/__init__.pyi

@@ -0,0 +1,3 @@
+from vllm.tokenizers.protocol import TokenizerLike
+
+__all__ = ["TokenizerLike"]

.cuda_typings/vllm/tokenizers/protocol.pyi

@@ -0,0 +1,15 @@
+from typing import Protocol
+
+class TokenizerLike(Protocol):
+    @property
+    def eos_token_id(self) -> int: ...
+    @property
+    def vocab_size(self) -> int: ...
+    def encode(self, text: str, add_special_tokens: bool = ...) -> list[int]: ...
+    def decode(self, ids: list[int] | int, skip_special_tokens: bool = ...) -> str: ...
+    def apply_chat_template(
+        self,
+        messages: list[dict[str, str]],
+        tools: list[dict[str, object]] | None = ...,
+        **kwargs: object,
+    ) -> str | list[int]: ...

.cuda_typings/vllm/v1/__init__.pyi

@@ -0,0 +1 @@
+

.cuda_typings/vllm/v1/core/__init__.pyi

@@ -0,0 +1 @@
+

.cuda_typings/vllm/v1/core/kv_cache_manager.pyi

@@ -0,0 +1,24 @@
+from collections.abc import Sequence
+
+from vllm.v1.core.kv_cache_utils import BlockPool, KVCacheBlock
+from vllm.v1.kv_cache_interface import KVCacheConfig
+
+class KVCacheBlocks:
+    blocks: tuple[Sequence[KVCacheBlock], ...]
+    def __init__(self, blocks: tuple[Sequence[KVCacheBlock], ...]) -> None: ...
+    def get_block_ids(self) -> tuple[list[int], ...]: ...
+
+class KVCacheManager:
+    block_pool: BlockPool
+    kv_cache_config: KVCacheConfig
+    enable_caching: bool
+    num_kv_cache_groups: int
+    coordinator: object
+    def __init__(self, *args: object, **kwargs: object) -> None: ...
+    def allocate_slots(
+        self, request: object, num_new_tokens: int, *args: object, **kwargs: object
+    ) -> KVCacheBlocks | None: ...
+    def get_computed_blocks(self, request: object) -> tuple[KVCacheBlocks, int]: ...
+    def create_kv_cache_blocks(
+        self, blocks: tuple[list[KVCacheBlock], ...]
+    ) -> KVCacheBlocks: ...

.cuda_typings/vllm/v1/core/kv_cache_utils.pyi

@@ -0,0 +1,16 @@
+class KVCacheBlock:
+    block_id: int
+    ref_cnt: int
+    def __init__(self, block_id: int) -> None: ...
+
+class FreeKVCacheBlockQueue:
+    def append_n(self, blocks: list[KVCacheBlock]) -> None: ...
+    def popleft_n(self, n: int) -> list[KVCacheBlock]: ...
+
+class BlockPool:
+    blocks: list[KVCacheBlock]
+    free_block_queue: FreeKVCacheBlockQueue
+    num_gpu_blocks: int
+    enable_caching: bool
+    def get_num_free_blocks(self) -> int: ...
+    def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]: ...

.cuda_typings/vllm/v1/engine/llm_engine.pyi

@@ -0,0 +1,22 @@
+from vllm.config import ModelConfig
+from vllm.engine.arg_utils import EngineArgs
+from vllm.outputs import RequestOutput
+from vllm.sampling_params import SamplingParams
+from vllm.tokenizers import TokenizerLike
+
+class LLMEngine:
+    tokenizer: TokenizerLike | None
+    model_config: ModelConfig
+
+    @classmethod
+    def from_engine_args(cls, engine_args: EngineArgs) -> LLMEngine: ...
+    def add_request(
+        self,
+        request_id: str,
+        prompt: str,
+        params: SamplingParams,
+        arrival_time: float | None = ...,
+    ) -> None: ...
+    def step(self) -> list[RequestOutput]: ...
+    def has_unfinished_requests(self) -> bool: ...
+    def get_tokenizer(self) -> TokenizerLike: ...

.cuda_typings/vllm/v1/kv_cache_interface.pyi

@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+
+@dataclass
+class KVCacheSpec:
+    block_size: int
+    num_kv_heads: int
+    head_size: int
+
+@dataclass
+class KVCacheGroupSpec:
+    layer_names: list[str]
+    kv_cache_spec: KVCacheSpec
+
+@dataclass
+class KVCacheTensorSpec:
+    shared_by: list[str]
+    size: int
+
+@dataclass
+class KVCacheConfig:
+    num_blocks: int
+    kv_cache_groups: list[KVCacheGroupSpec]
+    kv_cache_tensors: list[KVCacheTensorSpec]

.cuda_typings/vllm/v1/request.pyi

@@ -0,0 +1,6 @@
+class Request:
+    request_id: str
+    prompt_token_ids: list[int] | None
+    num_prompt_tokens: int
+    num_computed_tokens: int
+    num_tokens: int

.cuda_typings/vllm/v1/worker/__init__.pyi

@@ -0,0 +1 @@
+

.cuda_typings/vllm/v1/worker/gpu_model_runner.pyi

@@ -0,0 +1,24 @@
+import torch
+
+class _CompilationConfig:
+    static_forward_context: dict[str, object]
+
+class _ModelConfig:
+    hf_config: object
+
+class GPUModelRunner:
+    kv_caches: list[torch.Tensor]
+    compilation_config: _CompilationConfig
+    model_config: _ModelConfig | None
+    def _allocate_kv_cache_tensors(
+        self, kv_cache_config: object
+    ) -> dict[str, torch.Tensor]: ...
+    def initialize_kv_cache_tensors(
+        self, kv_cache_config: object, kernel_block_sizes: list[int]
+    ) -> dict[str, torch.Tensor]: ...
+    def _reshape_kv_cache_tensors(
+        self,
+        kv_cache_config: object,
+        raw_tensors: dict[str, torch.Tensor],
+        kernel_block_sizes: list[int],
+    ) -> dict[str, torch.Tensor]: ...

.cuda_typings/vllm/v1/worker/gpu_worker.pyi

@@ -0,0 +1,6 @@
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+class Worker:
+    model_runner: GPUModelRunner
+    def determine_available_memory(self) -> int: ...
+    def initialize_from_config(self, kv_cache_config: object) -> None: ...

.cuda_typings/vllm/v1/worker/utils.pyi

@@ -0,0 +1 @@
+def extract_layer_index(layer_name: str, num_attn_module: int) -> int: ...

Cargo.toml

@@ -1,11 +1,6 @@
 [workspace]
 resolver = "3"
-members = [
-  "rust/networking",
-  "rust/exo_pyo3_bindings",
-  "rust/util",
-  "rust/babblerd",
-]
+members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
 
 [workspace.package]
 version = "0.0.1"
@@ -48,6 +43,7 @@ log = "0.4"
 
 # networking
 libp2p = "0.56"
+libp2p-tcp = "0.44"
 
 [workspace.lints.rust]
 static_mut_refs = "warn"      # Or use "warn" instead of deny

bench/eval_tool_calls.py

@@ -17,11 +17,13 @@ from harness import (
     ExoClient,
     ExoHttpError,
     add_common_instance_args,
+    ensure_cuda_available,
     instance_id_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
     run_planning_phase,
     settle_and_fetch_placements,
+    validate_vllm_args,
     wait_for_instance_gone,
     wait_for_instance_ready,
 )
@@ -933,6 +935,7 @@ Examples:
         help="Write JSON results to stdout instead of file",
     )
     args = parser.parse_args()
+    validate_vllm_args(args)
 
     all_scenarios = load_scenarios(SCENARIOS_PATH)
     if args.scenarios:
@@ -952,6 +955,8 @@ Examples:
 
     log = sys.stderr if args.stdout else sys.stdout
     exo = ExoClient(args.host, args.port, timeout_s=args.timeout)
+    if args.ensure_cuda:
+        ensure_cuda_available(exo)
     _short_id, full_model_id = resolve_model_short_id(exo, args.model)
 
     selected = settle_and_fetch_placements(

bench/exo_bench.py

@@ -33,11 +33,13 @@ from harness import (
     ExoClient,
     ExoHttpError,
     add_common_instance_args,
+    ensure_cuda_available,
     instance_id_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
     run_planning_phase,
     settle_and_fetch_placements,
+    validate_vllm_args,
     wait_for_instance_gone,
     wait_for_instance_ready,
 )
@@ -280,6 +282,7 @@ def main() -> int:
         help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
     )
     args = ap.parse_args()
+    validate_vllm_args(args)
 
     pp_list = parse_int_list(args.pp)
     tg_list = parse_int_list(args.tg)
@@ -304,6 +307,8 @@ def main() -> int:
         logger.info(f"pp/tg mode: tandem (zip) - {len(pp_list)} pairs")
 
     client = ExoClient(args.host, args.port, timeout_s=args.timeout)
+    if args.ensure_cuda:
+        ensure_cuda_available(client)
     short_id, full_model_id = resolve_model_short_id(
         client, args.model, force_download=args.force_download
     )

bench/exo_eval.py

@@ -46,11 +46,13 @@ from harness import (
     ExoClient,
     ExoHttpError,
     add_common_instance_args,
+    ensure_cuda_available,
     instance_id_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
     run_planning_phase,
     settle_and_fetch_placements,
+    validate_vllm_args,
     wait_for_instance_gone,
     wait_for_instance_ready,
 )
@@ -1157,6 +1159,7 @@ def main() -> int:
     )
 
     args, _ = ap.parse_known_args()
+    validate_vllm_args(args)
 
     # Resolve tasks
     if args.tasks:
@@ -1173,6 +1176,8 @@ def main() -> int:
 
     # Instance management
     client = ExoClient(args.host, args.port, timeout_s=args.timeout)
+    if args.ensure_cuda:
+        ensure_cuda_available(client)
     instance_id: str | None = None
 
     if not args.skip_instance_setup:

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:
@@ -196,6 +200,31 @@ def resolve_model_short_id(
     raise ValueError(f"Model not found in /models: {model_arg}")
 
 
+def validate_vllm_args(args: argparse.Namespace) -> None:
+    if args.instance_meta != "vllm":
+        return
+    if args.sharding == "tensor":
+        raise SystemExit(
+            "--instance-meta vllm is incompatible with --sharding tensor (vllm is pipeline-only)"
+        )
+    if args.min_nodes > 1:
+        raise SystemExit(
+            "--instance-meta vllm is incompatible with --min-nodes > 1 (vllm is single-node)"
+        )
+    if args.max_nodes > 1:
+        raise SystemExit(
+            "--instance-meta vllm is incompatible with --max-nodes > 1 (vllm is single-node)"
+        )
+
+
+def ensure_cuda_available(client: ExoClient) -> None:
+    capabilities = client.request_json("GET", "/capabilities")
+    if not capabilities or not capabilities.get("vllm_available"):
+        raise SystemExit(
+            "--ensure-cuda: vllm is not available on the exo cluster (no CUDA capability)"
+        )
+
+
 def placement_filter(instance_meta: str, wanted: str) -> bool:
     s = (instance_meta or "").lower()
     if wanted == "both":
@@ -475,7 +504,7 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
         help="Only consider placements using >= this many nodes.",
     )
     ap.add_argument(
-        "--instance-meta", choices=["ring", "jaccl", "both"], default="both"
+        "--instance-meta", choices=["ring", "jaccl", "vllm", "both"], default="both"
     )
     ap.add_argument(
         "--sharding", choices=["pipeline", "tensor", "both"], default="both"
@@ -504,3 +533,8 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
         action="store_true",
         help="Delete existing models from smallest to largest to make room for benchmark model.",
     )
+    ap.add_argument(
+        "--ensure-cuda",
+        action="store_true",
+        help="Verify the exo cluster has CUDA/vllm capability; error if not.",
+    )

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/DeviceIcon.svelte

@@ -9,7 +9,7 @@
    */
 
   interface Props {
-    /** "macbook pro" | "mac studio" | "mac mini" etc. */
+    /** "macbook pro" | "mac studio" | "mac mini" | "dgx spark" | "linux" etc. */
     deviceType: string;
     /** Center X coordinate in SVG space */
     cx: number;
@@ -38,10 +38,43 @@
   const LOGO_NATIVE_WIDTH = 814;
   const LOGO_NATIVE_HEIGHT = 1000;
 
+  // NVIDIA logo SVG path
+  const NVIDIA_LOGO_PATH =
+    "M0.81 0.429V0.299c0.013 -0.001 0.026 -0.002 0.038 -0.002 0.355 -0.011 0.588 0.306 0.588 0.306S1.186 0.952 0.916 0.952c-0.036 0 -0.071 -0.006 -0.105 -0.017V0.542c0.138 0.017 0.166 0.078 0.249 0.216l0.185 -0.155s-0.135 -0.177 -0.362 -0.177c-0.024 -0.001 -0.048 0.001 -0.072 0.003m0 -0.429v0.194l0.038 -0.002c0.494 -0.017 0.816 0.405 0.816 0.405s-0.37 0.45 -0.754 0.45c-0.034 0 -0.066 -0.003 -0.099 -0.009v0.12c0.027 0.003 0.055 0.006 0.082 0.006 0.358 0 0.618 -0.183 0.869 -0.399 0.042 0.034 0.212 0.114 0.247 0.15 -0.238 0.2 -0.794 0.361 -1.11 0.361 -0.03 0 -0.059 -0.002 -0.088 -0.005v0.169h1.362V0zm0 0.935v0.102c-0.331 -0.059 -0.423 -0.404 -0.423 -0.404s0.159 -0.176 0.423 -0.205v0.112h-0.001C0.671 0.524 0.562 0.654 0.562 0.654s0.062 0.218 0.248 0.282m-0.588 -0.316s0.196 -0.29 0.589 -0.32V0.194C0.376 0.229 0 0.597 0 0.597s0.213 0.616 0.81 0.672v-0.112c-0.438 -0.054 -0.588 -0.538 -0.588 -0.538";
+
   const wireColor = "rgba(179,179,179,0.8)";
   const strokeWidth = 1.5;
 
   const modelLower = $derived(deviceType.toLowerCase());
+  const isSpark = $derived(
+    modelLower.includes("dgx") || modelLower.includes("gx10"),
+  );
+  const isLinux = $derived(!isSpark && modelLower.startsWith("linux"));
+  const isLinuxLaptop = $derived(isLinux && modelLower.includes("laptop"));
+
+  // ── DGX Spark dimensions ──
+  const dgxW = $derived(size * 1.55);
+  const dgxH = $derived(size * 0.58);
+  const dgxX = $derived(cx - dgxW / 2);
+  const dgxY = $derived(cy - dgxH / 2);
+  const dgxChassisX = $derived(dgxX - dgxW * 0.03);
+  const dgxChassisW = $derived(dgxW * 1.05);
+  const dgxHandleW = $derived(dgxW * 0.27);
+  const dgxHandleGap = $derived(dgxH * 0.05);
+  const dgxHandleH = $derived(dgxH - dgxHandleGap * 2);
+  const dgxHandleY = $derived(dgxY + dgxHandleGap);
+  const dgxInnerHandleW = $derived(dgxW * 0.12);
+  const dgxInnerHandleH = $derived(dgxHandleH - dgxH * 0.06);
+  const dgxLeftHandleX = $derived(dgxX + 4);
+  const dgxRightHandleX = $derived(dgxX + dgxW - dgxHandleW - 4);
+  const dgxClipId = $derived(`di-dgx-${uid}`);
+  const dgxTextureId = $derived(`di-dgx-tex-${uid}`);
+
+  // ── Linux Desktop dimensions (reuses Mac Studio proportions) ──
+  const linuxDesktopClipId = $derived(`di-linux-desktop-${uid}`);
+
+  // ── Linux Laptop dimensions (reuses MacBook proportions) ──
+  const linuxScreenClipId = $derived(`di-linux-screen-${uid}`);
 
   // ── Mac Studio dimensions (same ratios as TopologyGraph) ──
   const studioW = $derived(size * 1.25);
@@ -114,7 +147,264 @@
   const studioClipId = $derived(`di-studio-${uid}`);
 </script>
 
-{#if modelLower === "mac studio" || modelLower === "mac mini"}
+{#if isSpark}
+  <!-- DGX Spark -->
+  <defs>
+    <clipPath id={dgxClipId}>
+      <rect x={dgxX} y={dgxY} width={dgxW} height={dgxH} rx="3" />
+    </clipPath>
+    <pattern
+      id={dgxTextureId}
+      patternUnits="userSpaceOnUse"
+      width="8"
+      height="8"
+    >
+      <rect width="8" height="8" fill="#6f6248" />
+      <circle cx="2" cy="2" r="1" fill="#5a4f3b" opacity="0.5" />
+      <circle cx="6" cy="6" r="1" fill="#4a4232" opacity="0.45" />
+    </pattern>
+  </defs>
+
+  <!-- Main body -->
+  <rect
+    x={dgxChassisX}
+    y={dgxY}
+    width={dgxChassisW}
+    height={dgxH}
+    rx="3"
+    fill="url(#{dgxTextureId})"
+    stroke={wireColor}
+    stroke-width={strokeWidth}
+  />
+
+  <!-- Side border accents -->
+  <rect
+    x={dgxChassisX}
+    y={dgxY}
+    width={dgxW * 0.02}
+    height={dgxH}
+    fill="#8a7a56"
+  />
+  <rect
+    x={dgxChassisX + dgxChassisW - dgxW * 0.02}
+    y={dgxY}
+    width={dgxW * 0.02}
+    height={dgxH}
+    fill="#8a7a56"
+  />
+
+  <!-- Memory fill -->
+  {#if ramPercent > 0}
+    <rect
+      x={dgxX}
+      y={dgxY + dgxH - (ramPercent / 100) * dgxH}
+      width={dgxW}
+      height={(ramPercent / 100) * dgxH}
+      fill="rgba(255,215,0,0.45)"
+      clip-path="url(#{dgxClipId})"
+    />
+  {/if}
+
+  <!-- Left handle -->
+  <rect
+    x={dgxLeftHandleX}
+    y={dgxHandleY}
+    width={dgxHandleW}
+    height={dgxHandleH}
+    rx="2.4"
+    fill="#b3a170"
+    stroke="#403723"
+    stroke-width="0.7"
+  />
+  <rect
+    x={dgxLeftHandleX + dgxHandleW * 0.06}
+    y={dgxHandleY + dgxH * 0.03}
+    width={dgxInnerHandleW}
+    height={dgxInnerHandleH}
+    rx="1.6"
+    fill="#8a7a56"
+  />
+
+  <!-- Right handle -->
+  <rect
+    x={dgxRightHandleX}
+    y={dgxHandleY}
+    width={dgxHandleW}
+    height={dgxHandleH}
+    rx="2.4"
+    fill="#b3a170"
+    stroke="#403723"
+    stroke-width="0.7"
+  />
+  <rect
+    x={dgxRightHandleX + dgxHandleW - dgxInnerHandleW - dgxHandleW * 0.08}
+    y={dgxHandleY + dgxH * 0.03}
+    width={dgxInnerHandleW}
+    height={dgxInnerHandleH}
+    rx="1.6"
+    fill="#8a7a56"
+  />
+
+  <!-- NVIDIA logo (rotated 90deg on left handle) -->
+  {@const badgeW = dgxW * 0.09}
+  {@const badgeH = dgxHandleH * 0.5}
+  {@const badgeX = dgxLeftHandleX + dgxHandleW - badgeW - dgxHandleW * 0.06}
+  {@const badgeYPos = dgxHandleY + (dgxHandleH - badgeH) / 2}
+  {@const textSz = badgeW * 0.58}
+  {@const logoW = textSz * 1.2}
+  {@const logoH = logoW * (1.438 / 2.174)}
+  {@const ctrX = badgeX + badgeW / 2 - badgeW * 0.03}
+  {@const ctrY = badgeYPos + badgeH / 2}
+  {@const labelGap = badgeW * 0.15}
+  {@const totalW = logoW + labelGap + textSz * 3.6}
+  <g transform="rotate(90 {ctrX} {ctrY})">
+    <svg
+      x={ctrX - totalW / 2}
+      y={ctrY - logoH / 2}
+      width={logoW}
+      height={logoH}
+      viewBox="0 0 2.174 1.438"
+    >
+      <path d={NVIDIA_LOGO_PATH} fill="#76b900" />
+    </svg>
+    <text
+      x={ctrX - totalW / 2 + logoW + labelGap}
+      y={ctrY}
+      text-anchor="start"
+      dominant-baseline="middle"
+      fill="#8a7a56"
+      font-size={textSz}
+      font-family="monospace"
+      font-weight="700">NVIDIA</text
+    >
+  </g>
+{:else if isLinuxLaptop}
+  <!-- Linux Laptop — MacBook shape with Tux logo -->
+  <defs>
+    <clipPath id={linuxScreenClipId}>
+      <rect
+        x={mbScreenX + mbBezel}
+        y={mbY + mbBezel}
+        width={mbScreenW - mbBezel * 2}
+        height={mbScreenH - mbBezel * 2}
+        rx="2"
+      />
+    </clipPath>
+  </defs>
+
+  <rect
+    x={mbScreenX}
+    y={mbY}
+    width={mbScreenW}
+    height={mbScreenH}
+    rx="3"
+    fill="#1a1a1a"
+    stroke={wireColor}
+    stroke-width={strokeWidth}
+  />
+  <rect
+    x={mbScreenX + mbBezel}
+    y={mbY + mbBezel}
+    width={mbScreenW - mbBezel * 2}
+    height={mbScreenH - mbBezel * 2}
+    rx="2"
+    fill="#0a0a12"
+  />
+  {#if ramPercent > 0}
+    <rect
+      x={mbScreenX + mbBezel}
+      y={mbY + mbBezel + (mbMemTotalH - mbMemH)}
+      width={mbScreenW - mbBezel * 2}
+      height={mbMemH}
+      fill="rgba(255,215,0,0.85)"
+      clip-path="url(#{linuxScreenClipId})"
+    />
+  {/if}
+
+  <!-- Terminal prompt on screen -->
+  <text
+    x={cx}
+    y={mbY + mbScreenH / 2}
+    text-anchor="middle"
+    dominant-baseline="middle"
+    fill="#FFFFFF"
+    opacity="0.9"
+    font-size={mbScreenH * 0.25}
+    font-family="SF Mono, Monaco, monospace"
+    font-weight="700">{">_"}</text
+  >
+
+  <path
+    d="M {mbBaseTopX} {mbBaseY} L {mbBaseTopX +
+      mbBaseTopW} {mbBaseY} L {mbBaseBottomX + mbBaseBottomW} {mbBaseY +
+      mbBaseH} L {mbBaseBottomX} {mbBaseY + mbBaseH} Z"
+    fill="#2c2c2c"
+    stroke={wireColor}
+    stroke-width="1"
+  />
+  <rect
+    x={mbKbX}
+    y={mbKbY}
+    width={mbKbW}
+    height={mbKbH}
+    fill="rgba(0,0,0,0.2)"
+    rx="2"
+  />
+  <rect
+    x={mbTpX}
+    y={mbTpY}
+    width={mbTpW}
+    height={mbTpH}
+    fill="rgba(255,255,255,0.08)"
+    rx="2"
+  />
+{:else if isLinux}
+  <!-- Linux Desktop — Mac Studio shape with Tux logo -->
+  <defs>
+    <clipPath id={linuxDesktopClipId}>
+      <rect
+        x={studioX}
+        y={studioY + studioTopH}
+        width={studioW}
+        height={studioH - studioTopH}
+        rx={studioCorner - 1}
+      />
+    </clipPath>
+  </defs>
+
+  <rect
+    x={studioX}
+    y={studioY}
+    width={studioW}
+    height={studioH}
+    rx={studioCorner}
+    fill="#1a1a1a"
+    stroke={wireColor}
+    stroke-width={strokeWidth}
+  />
+  {#if ramPercent > 0}
+    <rect
+      x={studioX}
+      y={studioY + studioTopH + (studioMemTotalH - studioMemH)}
+      width={studioW}
+      height={studioMemH}
+      fill="rgba(255,215,0,0.75)"
+      clip-path="url(#{linuxDesktopClipId})"
+    />
+  {/if}
+
+  <!-- Terminal prompt on front face -->
+  <text
+    x={cx}
+    y={studioY + studioTopH + (studioH - studioTopH) / 2}
+    text-anchor="middle"
+    dominant-baseline="middle"
+    fill="rgba(255,255,255,0.5)"
+    font-size={(studioH - studioTopH) * 0.4}
+    font-family="SF Mono, Monaco, monospace"
+    font-weight="700">{">_"}</text
+  >
+{:else if modelLower === "mac studio" || modelLower === "mac mini"}
   <!-- Mac Studio / Mac Mini -->
   <defs>
     <clipPath id={studioClipId}>

dashboard/src/lib/components/ModelCard.svelte

@@ -169,8 +169,10 @@
 
   function getDeviceType(
     name: string,
-  ): "macbook" | "studio" | "mini" | "unknown" {
+  ): "macbook" | "studio" | "mini" | "dgx" | "linux" | "unknown" {
     const lower = name.toLowerCase();
+    if (lower.includes("dgx") || lower.includes("gx10")) return "dgx";
+    if (lower.includes("linux")) return "linux";
     if (lower.includes("macbook")) return "macbook";
     if (lower.includes("studio")) return "studio";
     if (lower.includes("mini")) return "mini";
@@ -278,7 +280,7 @@
     let placementNodes: Array<{
       id: string;
       deviceName: string;
-      deviceType: "macbook" | "studio" | "mini" | "unknown";
+      deviceType: "macbook" | "studio" | "mini" | "dgx" | "linux" | "unknown";
       totalGB: number;
       currentUsedGB: number;
       modelUsageGB: number;
@@ -968,6 +970,137 @@
                     />
                   {/if}
                 </g>
+              {:else if node.deviceType === "dgx"}
+                <!-- DGX Spark icon -->
+                {@const s = node.iconSize}
+                {@const dgxW = s * 1.4}
+                {@const dgxH = s * 0.52}
+                <g transform="translate({-dgxW / 2}, {-dgxH / 2})">
+                  <!-- Chassis -->
+                  <rect
+                    x="0"
+                    y="0"
+                    width={dgxW}
+                    height={dgxH}
+                    rx="2"
+                    fill="#6f6248"
+                    stroke={node.isUsed ? "#FFD700" : "#4B5563"}
+                    stroke-width="1.5"
+                  />
+                  <!-- Side accents -->
+                  <rect
+                    x="0"
+                    y="0"
+                    width={dgxW * 0.02}
+                    height={dgxH}
+                    fill="#8a7a56"
+                  />
+                  <rect
+                    x={dgxW - dgxW * 0.02}
+                    y="0"
+                    width={dgxW * 0.02}
+                    height={dgxH}
+                    fill="#8a7a56"
+                  />
+                  <!-- Left handle -->
+                  <rect
+                    x={dgxW * 0.04}
+                    y={dgxH * 0.08}
+                    width={dgxW * 0.22}
+                    height={dgxH * 0.84}
+                    rx="2"
+                    fill="#b3a170"
+                    stroke="#403723"
+                    stroke-width="0.5"
+                  />
+                  <!-- Right handle -->
+                  <rect
+                    x={dgxW - dgxW * 0.04 - dgxW * 0.22}
+                    y={dgxH * 0.08}
+                    width={dgxW * 0.22}
+                    height={dgxH * 0.84}
+                    rx="2"
+                    fill="#b3a170"
+                    stroke="#403723"
+                    stroke-width="0.5"
+                  />
+                  <!-- Memory fill -->
+                  <rect
+                    x="2"
+                    y={dgxH - dgxH * (node.currentPercent / 100)}
+                    width={dgxW - 4}
+                    height={dgxH * (node.currentPercent / 100)}
+                    fill="rgba(255,215,0,0.35)"
+                  />
+                  {#if node.modelUsageGB > 0 && node.isUsed}
+                    <rect
+                      x="2"
+                      y={dgxH - dgxH * (node.newPercent / 100)}
+                      width={dgxW - 4}
+                      height={dgxH *
+                        ((node.newPercent - node.currentPercent) / 100)}
+                      fill="#FFD700"
+                      filter="url(#memGlow-{filterId})"
+                      class="animate-pulse-slow"
+                    />
+                  {/if}
+                </g>
+              {:else if node.deviceType === "linux"}
+                <!-- Linux Tux penguin icon -->
+                {@const sz = node.iconSize}
+                {@const sc = sz / 100}
+                <g transform="translate({-sz / 2}, {-sz / 2})">
+                  <!-- Body -->
+                  <path
+                    d="M50 8c-8 0-14 6-14 13 0 4 2 8 5 10-8 4-16 14-16 28v12c0 4 2 7 5 9l-6 4c-2 1-3 3-3 5v3c0 2 2 4 4 4h10l4-4h22l4 4h10c2 0 4-2 4-4v-3c0-2-1-4-3-5l-6-4c3-2 5-5 5-9V59c0-14-8-24-16-28 3-2 5-6 5-10 0-7-6-13-14-13z"
+                    transform="scale({sc})"
+                    fill="#1a1a1a"
+                    stroke={node.isUsed ? "#FFD700" : "#4B5563"}
+                    stroke-width={1.5 / sc}
+                  />
+                  <!-- Belly -->
+                  <path
+                    d="M38 52c0-8 5-15 12-15s12 7 12 15v14c0 4-5 7-12 7s-12-3-12-7V52z"
+                    transform="scale({sc})"
+                    fill="rgba(220,220,220,0.85)"
+                  />
+                  <!-- Eyes -->
+                  <circle cx={44 * sc} cy={16 * sc} r={2.5 * sc} fill="white" />
+                  <circle cx={56 * sc} cy={16 * sc} r={2.5 * sc} fill="white" />
+                  <circle
+                    cx={44 * sc}
+                    cy={16 * sc}
+                    r={1.2 * sc}
+                    fill="#1a1a1a"
+                  />
+                  <circle
+                    cx={56 * sc}
+                    cy={16 * sc}
+                    r={1.2 * sc}
+                    fill="#1a1a1a"
+                  />
+                  <!-- Beak -->
+                  <path
+                    d="M{46 * sc} {22 * sc} L{50 * sc} {27 * sc} L{54 *
+                      sc} {22 * sc} Z"
+                    fill="#E8A317"
+                  />
+                  <!-- Feet -->
+                  <ellipse
+                    cx={42 * sc}
+                    cy={94 * sc}
+                    rx={6 * sc}
+                    ry={2.5 * sc}
+                    fill="#E8A317"
+                  />
+                  <ellipse
+                    cx={58 * sc}
+                    cy={94 * sc}
+                    rx={6 * sc}
+                    ry={2.5 * sc}
+                    fill="#E8A317"
+                  />
+                </g>
               {:else}
                 <!-- Unknown device - hexagon -->
                 <g

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/lib/components/TopologyGraph.svelte

@@ -117,6 +117,10 @@
   const LOGO_NATIVE_WIDTH = 814;
   const LOGO_NATIVE_HEIGHT = 1000;
 
+  // NVIDIA logo SVG path (from exo-nvidia)
+  const NVIDIA_LOGO_PATH =
+    "M0.81 0.429V0.299c0.013 -0.001 0.026 -0.002 0.038 -0.002 0.355 -0.011 0.588 0.306 0.588 0.306S1.186 0.952 0.916 0.952c-0.036 0 -0.071 -0.006 -0.105 -0.017V0.542c0.138 0.017 0.166 0.078 0.249 0.216l0.185 -0.155s-0.135 -0.177 -0.362 -0.177c-0.024 -0.001 -0.048 0.001 -0.072 0.003m0 -0.429v0.194l0.038 -0.002c0.494 -0.017 0.816 0.405 0.816 0.405s-0.37 0.45 -0.754 0.45c-0.034 0 -0.066 -0.003 -0.099 -0.009v0.12c0.027 0.003 0.055 0.006 0.082 0.006 0.358 0 0.618 -0.183 0.869 -0.399 0.042 0.034 0.212 0.114 0.247 0.15 -0.238 0.2 -0.794 0.361 -1.11 0.361 -0.03 0 -0.059 -0.002 -0.088 -0.005v0.169h1.362V0zm0 0.935v0.102c-0.331 -0.059 -0.423 -0.404 -0.423 -0.404s0.159 -0.176 0.423 -0.205v0.112h-0.001C0.671 0.524 0.562 0.654 0.562 0.654s0.062 0.218 0.248 0.282m-0.588 -0.316s0.196 -0.29 0.589 -0.32V0.194C0.376 0.229 0 0.597 0 0.597s0.213 0.616 0.81 0.672v-0.112c-0.438 -0.054 -0.588 -0.538 -0.588 -0.538";
+
   function formatBytes(bytes: number, decimals = 1): string {
     if (!bytes || bytes === 0) return "0B";
     const k = 1024;
@@ -554,6 +558,13 @@
       const clipPathId = `clip-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
 
       const modelLower = modelId.toLowerCase();
+      const identity = identitiesData[nodeInfo.id];
+      const nameLower = (friendlyName || "").toLowerCase();
+      const isSpark = modelLower.includes("dgx") || modelLower.includes("gx10");
+      const isLinux =
+        !isSpark &&
+        (modelLower.startsWith("linux") || identity?.osVersion === "Linux");
+      const isLinuxLaptop = isLinux && modelLower.includes("laptop");
 
       // Check node states for styling
       const isHighlighted = highlightedNodes.has(nodeInfo.id);
@@ -623,7 +634,382 @@
           `${friendlyName}\nID: ${nodeInfo.id.slice(-8)}\nMemory: ${formatBytes(ramUsed)}/${formatBytes(ramTotal)}`,
         );
 
-      if (modelLower === "mac studio") {
+      if (isSpark) {
+        // NVIDIA DGX Spark — gold chassis with textured front, side handles, and NVIDIA badge
+        iconBaseWidth = nodeRadius * 1.55;
+        iconBaseHeight = nodeRadius * 0.58;
+        const x = nodeInfo.x - iconBaseWidth / 2;
+        const y = nodeInfo.y - iconBaseHeight / 2;
+        const chassisX = x - iconBaseWidth * 0.03;
+        const chassisWidth = iconBaseWidth * 1.05;
+        const cornerRadius = 3;
+
+        const dgxClipId = `dgx-clip-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
+        defs
+          .append("clipPath")
+          .attr("id", dgxClipId)
+          .append("rect")
+          .attr("x", x)
+          .attr("y", y)
+          .attr("width", iconBaseWidth)
+          .attr("height", iconBaseHeight)
+          .attr("rx", cornerRadius);
+
+        // Chassis texture pattern
+        const textureId = `chassis-texture-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
+        defs
+          .append("pattern")
+          .attr("id", textureId)
+          .attr("patternUnits", "userSpaceOnUse")
+          .attr("width", 8)
+          .attr("height", 8);
+        const texturePattern = defs.select(`#${textureId}`);
+        texturePattern
+          .append("rect")
+          .attr("width", 8)
+          .attr("height", 8)
+          .attr("fill", "#6f6248");
+        texturePattern
+          .append("circle")
+          .attr("cx", 2)
+          .attr("cy", 2)
+          .attr("r", 1)
+          .attr("fill", "#5a4f3b")
+          .attr("opacity", 0.5);
+        texturePattern
+          .append("circle")
+          .attr("cx", 6)
+          .attr("cy", 6)
+          .attr("r", 1)
+          .attr("fill", "#4a4232")
+          .attr("opacity", 0.45);
+
+        // Main body
+        nodeG
+          .append("rect")
+          .attr("class", "node-outline")
+          .attr("x", chassisX)
+          .attr("y", y)
+          .attr("width", chassisWidth)
+          .attr("height", iconBaseHeight)
+          .attr("rx", cornerRadius)
+          .attr("fill", `url(#${textureId})`)
+          .attr("stroke", wireColor)
+          .attr("stroke-width", strokeWidth);
+
+        // Side border accents
+        const sideThickness = iconBaseWidth * 0.02;
+        nodeG
+          .append("rect")
+          .attr("x", chassisX)
+          .attr("y", y)
+          .attr("width", sideThickness)
+          .attr("height", iconBaseHeight)
+          .attr("fill", "#8a7a56");
+        nodeG
+          .append("rect")
+          .attr("x", chassisX + chassisWidth - sideThickness)
+          .attr("y", y)
+          .attr("width", sideThickness)
+          .attr("height", iconBaseHeight)
+          .attr("fill", "#8a7a56");
+
+        // Memory fill (bottom up)
+        if (ramUsagePercent > 0) {
+          const memFillHeight = (ramUsagePercent / 100) * iconBaseHeight;
+          nodeG
+            .append("rect")
+            .attr("x", x)
+            .attr("y", y + iconBaseHeight - memFillHeight)
+            .attr("width", iconBaseWidth)
+            .attr("height", memFillHeight)
+            .attr("fill", "rgba(255,215,0,0.45)")
+            .attr("clip-path", `url(#${dgxClipId})`);
+        }
+
+        // Side handles with inner recess
+        const handleWidth = iconBaseWidth * 0.27;
+        const handleGap = iconBaseHeight * 0.05;
+        const handleHeight = iconBaseHeight - handleGap * 2;
+        const handleY = y + handleGap;
+        const innerHandleWidth = iconBaseWidth * 0.12;
+        const innerHandleHeight = handleHeight - iconBaseHeight * 0.06;
+        const leftHandleX = x + 4;
+        const rightHandleX = x + iconBaseWidth - handleWidth - 4;
+
+        // Left handle
+        nodeG
+          .append("rect")
+          .attr("x", leftHandleX)
+          .attr("y", handleY)
+          .attr("width", handleWidth)
+          .attr("height", handleHeight)
+          .attr("rx", 2.4)
+          .attr("fill", "#b3a170")
+          .attr("stroke", "#403723")
+          .attr("stroke-width", 0.7);
+        nodeG
+          .append("rect")
+          .attr("x", leftHandleX + handleWidth * 0.06)
+          .attr("y", handleY + iconBaseHeight * 0.03)
+          .attr("width", innerHandleWidth)
+          .attr("height", innerHandleHeight)
+          .attr("rx", 1.6)
+          .attr("fill", "#8a7a56");
+
+        // Right handle
+        nodeG
+          .append("rect")
+          .attr("x", rightHandleX)
+          .attr("y", handleY)
+          .attr("width", handleWidth)
+          .attr("height", handleHeight)
+          .attr("rx", 2.4)
+          .attr("fill", "#b3a170")
+          .attr("stroke", "#403723")
+          .attr("stroke-width", 0.7);
+        nodeG
+          .append("rect")
+          .attr(
+            "x",
+            rightHandleX + handleWidth - innerHandleWidth - handleWidth * 0.08,
+          )
+          .attr("y", handleY + iconBaseHeight * 0.03)
+          .attr("width", innerHandleWidth)
+          .attr("height", innerHandleHeight)
+          .attr("rx", 1.6)
+          .attr("fill", "#8a7a56");
+
+        // NVIDIA logo + text label (rotated 90 deg on left handle)
+        const badgeWidth = iconBaseWidth * 0.09;
+        const badgeHeight = handleHeight * 0.5;
+        const badgeX =
+          leftHandleX + handleWidth - badgeWidth - handleWidth * 0.06;
+        const badgeY = handleY + (handleHeight - badgeHeight) / 2;
+        const textSize = badgeWidth * 0.58;
+        const logoWidth = textSize * 1.2;
+        const logoHeight = logoWidth * (1.438 / 2.174);
+        const centerX = badgeX + badgeWidth / 2 - badgeWidth * 0.03;
+        const centerY = badgeY + badgeHeight / 2;
+        const gap = badgeWidth * 0.15;
+        const totalWidth = logoWidth + gap + textSize * 3.6;
+
+        const labelGroup = nodeG
+          .append("g")
+          .attr("transform", `rotate(90 ${centerX} ${centerY})`);
+
+        labelGroup
+          .append("svg")
+          .attr("x", centerX - totalWidth / 2)
+          .attr("y", centerY - logoHeight / 2)
+          .attr("width", logoWidth)
+          .attr("height", logoHeight)
+          .attr("viewBox", "0 0 2.174 1.438")
+          .append("path")
+          .attr("d", NVIDIA_LOGO_PATH)
+          .attr("fill", "#76b900");
+
+        labelGroup
+          .append("text")
+          .attr("x", centerX - totalWidth / 2 + logoWidth + gap)
+          .attr("y", centerY)
+          .attr("text-anchor", "start")
+          .attr("dominant-baseline", "middle")
+          .attr("fill", "#8a7a56")
+          .attr("font-size", textSize)
+          .attr("font-family", "monospace")
+          .attr("font-weight", "700")
+          .text("NVIDIA");
+      } else if (isLinuxLaptop) {
+        // Linux Laptop — same shape as MacBook but with Tux logo
+        iconBaseWidth = nodeRadius * 1.6;
+        iconBaseHeight = nodeRadius * 1.15;
+        const x = nodeInfo.x - iconBaseWidth / 2;
+        const y = nodeInfo.y - iconBaseHeight / 2;
+
+        const screenHeight = iconBaseHeight * 0.7;
+        const baseHeight = iconBaseHeight * 0.3;
+        const screenWidth = iconBaseWidth * 0.85;
+        const screenX = nodeInfo.x - screenWidth / 2;
+        const screenBezel = 3;
+
+        const linuxScreenClipId = `linux-screen-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
+        defs
+          .append("clipPath")
+          .attr("id", linuxScreenClipId)
+          .append("rect")
+          .attr("x", screenX + screenBezel)
+          .attr("y", y + screenBezel)
+          .attr("width", screenWidth - screenBezel * 2)
+          .attr("height", screenHeight - screenBezel * 2)
+          .attr("rx", 2);
+
+        // Screen outer frame
+        nodeG
+          .append("rect")
+          .attr("class", "node-outline")
+          .attr("x", screenX)
+          .attr("y", y)
+          .attr("width", screenWidth)
+          .attr("height", screenHeight)
+          .attr("rx", 3)
+          .attr("fill", "#1a1a1a")
+          .attr("stroke", wireColor)
+          .attr("stroke-width", strokeWidth);
+
+        // Screen inner
+        nodeG
+          .append("rect")
+          .attr("x", screenX + screenBezel)
+          .attr("y", y + screenBezel)
+          .attr("width", screenWidth - screenBezel * 2)
+          .attr("height", screenHeight - screenBezel * 2)
+          .attr("rx", 2)
+          .attr("fill", "#0a0a12");
+
+        // Memory fill on screen
+        if (ramUsagePercent > 0) {
+          const memFillTotalHeight = screenHeight - screenBezel * 2;
+          const memFillActualHeight =
+            (ramUsagePercent / 100) * memFillTotalHeight;
+          nodeG
+            .append("rect")
+            .attr("x", screenX + screenBezel)
+            .attr(
+              "y",
+              y + screenBezel + (memFillTotalHeight - memFillActualHeight),
+            )
+            .attr("width", screenWidth - screenBezel * 2)
+            .attr("height", memFillActualHeight)
+            .attr("fill", "rgba(255,215,0,0.85)")
+            .attr("clip-path", `url(#${linuxScreenClipId})`);
+        }
+
+        // Terminal prompt on screen
+        nodeG
+          .append("text")
+          .attr("x", nodeInfo.x)
+          .attr("y", y + screenHeight / 2)
+          .attr("text-anchor", "middle")
+          .attr("dominant-baseline", "middle")
+          .attr("fill", "#FFFFFF")
+          .attr("opacity", 0.9)
+          .attr("font-size", screenHeight * 0.25)
+          .attr("font-family", "SF Mono, Monaco, monospace")
+          .attr("font-weight", "700")
+          .text(">_");
+
+        // Keyboard base (trapezoidal)
+        const baseY = y + screenHeight;
+        const baseTopWidth = screenWidth;
+        const baseBottomWidth = iconBaseWidth;
+        const baseTopX = nodeInfo.x - baseTopWidth / 2;
+        const baseBottomX = nodeInfo.x - baseBottomWidth / 2;
+
+        nodeG
+          .append("path")
+          .attr(
+            "d",
+            `M ${baseTopX} ${baseY} L ${baseTopX + baseTopWidth} ${baseY} L ${baseBottomX + baseBottomWidth} ${baseY + baseHeight} L ${baseBottomX} ${baseY + baseHeight} Z`,
+          )
+          .attr("fill", "#2c2c2c")
+          .attr("stroke", wireColor)
+          .attr("stroke-width", 1);
+
+        // Keyboard area
+        const keyboardX = baseTopX + 6;
+        const keyboardY = baseY + 3;
+        const keyboardWidth = baseTopWidth - 12;
+        const keyboardHeight = baseHeight * 0.55;
+        nodeG
+          .append("rect")
+          .attr("x", keyboardX)
+          .attr("y", keyboardY)
+          .attr("width", keyboardWidth)
+          .attr("height", keyboardHeight)
+          .attr("fill", "rgba(0,0,0,0.2)")
+          .attr("rx", 2);
+
+        // Trackpad
+        const trackpadWidth = baseTopWidth * 0.4;
+        const trackpadX = nodeInfo.x - trackpadWidth / 2;
+        const trackpadY = baseY + keyboardHeight + 5;
+        const trackpadHeight = baseHeight * 0.3;
+        nodeG
+          .append("rect")
+          .attr("x", trackpadX)
+          .attr("y", trackpadY)
+          .attr("width", trackpadWidth)
+          .attr("height", trackpadHeight)
+          .attr("fill", "rgba(255,255,255,0.08)")
+          .attr("rx", 2);
+      } else if (isLinux) {
+        // Linux Desktop — same shape as Mac Studio but with Tux logo
+        iconBaseWidth = nodeRadius * 1.25;
+        iconBaseHeight = nodeRadius * 0.85;
+        const x = nodeInfo.x - iconBaseWidth / 2;
+        const y = nodeInfo.y - iconBaseHeight / 2;
+        const cornerRadius = 4;
+        const topSurfaceHeight = iconBaseHeight * 0.15;
+
+        const linuxDesktopClipId = `linux-desktop-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
+        defs
+          .append("clipPath")
+          .attr("id", linuxDesktopClipId)
+          .append("rect")
+          .attr("x", x)
+          .attr("y", y + topSurfaceHeight)
+          .attr("width", iconBaseWidth)
+          .attr("height", iconBaseHeight - topSurfaceHeight)
+          .attr("rx", cornerRadius - 1);
+
+        // Main body
+        nodeG
+          .append("rect")
+          .attr("class", "node-outline")
+          .attr("x", x)
+          .attr("y", y)
+          .attr("width", iconBaseWidth)
+          .attr("height", iconBaseHeight)
+          .attr("rx", cornerRadius)
+          .attr("fill", "#1a1a1a")
+          .attr("stroke", wireColor)
+          .attr("stroke-width", strokeWidth);
+
+        // Memory fill
+        if (ramUsagePercent > 0) {
+          const memFillTotalHeight = iconBaseHeight - topSurfaceHeight;
+          const memFillActualHeight =
+            (ramUsagePercent / 100) * memFillTotalHeight;
+          nodeG
+            .append("rect")
+            .attr("x", x)
+            .attr(
+              "y",
+              y + topSurfaceHeight + (memFillTotalHeight - memFillActualHeight),
+            )
+            .attr("width", iconBaseWidth)
+            .attr("height", memFillActualHeight)
+            .attr("fill", "rgba(255,215,0,0.75)")
+            .attr("clip-path", `url(#${linuxDesktopClipId})`);
+        }
+
+        // Terminal prompt on front face
+        nodeG
+          .append("text")
+          .attr("x", nodeInfo.x)
+          .attr(
+            "y",
+            y + topSurfaceHeight + (iconBaseHeight - topSurfaceHeight) / 2,
+          )
+          .attr("text-anchor", "middle")
+          .attr("dominant-baseline", "middle")
+          .attr("fill", "rgba(255,255,255,0.5)")
+          .attr("font-size", (iconBaseHeight - topSurfaceHeight) * 0.4)
+          .attr("font-family", "SF Mono, Monaco, monospace")
+          .attr("font-weight", "700")
+          .text(">_");
+      } else if (modelLower === "mac studio") {
         // Mac Studio - classic cube with memory fill
         iconBaseWidth = nodeRadius * 1.25;
         iconBaseHeight = nodeRadius * 0.85;
@@ -1182,8 +1568,12 @@
           debugLabelY += debugLineHeight;
         }
 
-        const identity = identitiesData[nodeInfo.id];
-        if (identity?.osVersion) {
+        const dbgIdentity = identitiesData[nodeInfo.id];
+        if (dbgIdentity?.osVersion) {
+          const osLabel =
+            dbgIdentity.osVersion === "Linux"
+              ? "Linux"
+              : `macOS ${dbgIdentity.osVersion}${dbgIdentity.osBuildVersion ? ` (${dbgIdentity.osBuildVersion})` : ""}`;
           nodeG
             .append("text")
             .attr("x", nodeInfo.x)
@@ -1192,9 +1582,7 @@
             .attr("fill", "rgba(179,179,179,0.7)")
             .attr("font-size", debugFontSize)
             .attr("font-family", "SF Mono, Monaco, monospace")
-            .text(
-              `macOS ${identity.osVersion}${identity.osBuildVersion ? ` (${identity.osBuildVersion})` : ""}`,
-            );
+            .text(osLabel);
         }
       }
     });

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

@@ -168,7 +168,7 @@ export interface ModelDownloadStatus {
 export interface PlacementPreview {
   model_id: string;
   sharding: "Pipeline" | "Tensor";
-  instance_meta: "MlxRing" | "MlxJaccl";
+  instance_meta: "MlxRing" | "MlxJaccl" | "Vllm";
   instance: unknown | null;
   memory_delta_by_node: Record<string, number> | null;
   error: string | null;
@@ -547,6 +547,7 @@ class AppStore {
       { total: { inBytes: number }; available: { inBytes: number } }
     >
   >({});
+  vllmAvailable = $state(false);
   placementPreviews = $state<PlacementPreview[]>([]);
   selectedPreviewModelId = $state<string | null>(null);
   isLoadingPreviews = $state(false);
@@ -1331,6 +1332,15 @@ class AppStore {
         this.isConnected = true;
       }
       this.consecutiveFailures = 0;
+
+      fetch("/capabilities")
+        .then((r) => (r.ok ? r.json() : null))
+        .then((data: { vllm_available?: boolean } | null) => {
+          this.vllmAvailable = data?.vllm_available ?? false;
+        })
+        .catch(() => {
+          this.vllmAvailable = false;
+        });
     } catch (error) {
       this.consecutiveFailures++;
       if (
@@ -3334,6 +3344,7 @@ export const nodeThunderbolt = () => appStore.nodeThunderbolt;
 export const nodeRdmaCtl = () => appStore.nodeRdmaCtl;
 export const thunderboltBridgeCycles = () => appStore.thunderboltBridgeCycles;
 export const nodeThunderboltBridge = () => appStore.nodeThunderboltBridge;
+export const vllmAvailable = () => appStore.vllmAvailable;
 
 // Image generation params
 export const imageGenerationParams = () => appStore.getImageGenerationParams();

dashboard/src/routes/+page.svelte

@@ -64,6 +64,7 @@
     nodeThunderboltBridge,
     nodeIdentities,
     isConnected,
+    vllmAvailable,
     type DownloadProgress,
     type PlacementPreview,
   } from "$lib/stores/app.svelte";
@@ -294,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);
@@ -700,7 +701,10 @@
       ? Object.keys(topologyData()!.nodes).length
       : 1;
     const sharding = nodeCount <= 1 ? "Pipeline" : selectedSharding;
-    const instanceType = nodeCount <= 1 ? "MlxRing" : selectedInstanceType;
+    const instanceType =
+      nodeCount <= 1 && selectedInstanceType !== "Vllm"
+        ? "MlxRing"
+        : selectedInstanceType;
     try {
       const placementResponse = await fetch(
         `/instance/placement?model_id=${encodeURIComponent(modelId)}&sharding=${sharding}&instance_meta=${instanceType}&min_nodes=1`,
@@ -884,7 +888,7 @@
   }
 
   let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
-  type InstanceMeta = "MlxRing" | "MlxJaccl";
+  type InstanceMeta = "MlxRing" | "MlxJaccl" | "Vllm";
 
   // Launch defaults persistence
   const LAUNCH_DEFAULTS_KEY = "exo-launch-defaults-v2";
@@ -930,7 +934,11 @@
     // Apply sharding and instance type unconditionally
     selectedSharding = defaults.sharding;
     selectedInstanceType =
-      defaults.instanceType === "MlxRing" ? "MlxRing" : "MlxJaccl";
+      defaults.instanceType === "Vllm"
+        ? "Vllm"
+        : defaults.instanceType === "MlxRing"
+          ? "MlxRing"
+          : "MlxJaccl";
 
     // Apply minNodes if valid (between 1 and maxNodes)
     if (
@@ -1144,9 +1152,7 @@
   }
 
   const matchesSelectedRuntime = (runtime: InstanceMeta): boolean =>
-    selectedInstanceType === "MlxRing"
-      ? runtime === "MlxRing"
-      : runtime === "MlxJaccl";
+    runtime === selectedInstanceType;
 
   // Helper to check if a model can be launched (has valid placement with >= minNodes)
   function canModelFit(modelId: string): boolean {
@@ -2091,6 +2097,7 @@
     let instanceType = "Unknown";
     if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
     else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
+    else if (instanceTag === "VllmInstance") instanceType = "vLLM (CUDA)";
 
     const inst = instance as {
       shardAssignments?: {
@@ -5840,6 +5847,32 @@
                         </span>
                         RDMA (Fast)
                       </button>
+                      {#if vllmAvailable()}
+                        <button
+                          onclick={() => {
+                            selectedInstanceType = "Vllm";
+                            saveLaunchDefaults();
+                          }}
+                          class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
+                          'Vllm'
+                            ? 'bg-transparent text-exo-yellow border-exo-yellow'
+                            : 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
+                        >
+                          <span
+                            class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
+                            'Vllm'
+                              ? 'border-exo-yellow'
+                              : 'border-exo-medium-gray'}"
+                          >
+                            {#if selectedInstanceType === "Vllm"}
+                              <span
+                                class="w-1.5 h-1.5 rounded-full bg-exo-yellow"
+                              ></span>
+                            {/if}
+                          </span>
+                          vLLM (CUDA)
+                        </button>
+                      {/if}
                     </div>
                   </div>
 

flake.nix

@@ -51,8 +51,8 @@
   };
 
   nixConfig = {
-    extra-trusted-public-keys = "exo.cachix.org-1:okq7hl624TBeAR3kV+g39dUFSiaZgLRkLsFBCuJ2NZI=";
-    extra-substituters = "https://exo.cachix.org";
+    extra-trusted-public-keys = "exo.cachix.org-1:okq7hl624TBeAR3kV+g39dUFSiaZgLRkLsFBCuJ2NZI= cache.nixos-cuda.org:74DUi4Ye579gUqzH4ziL9IyiJBlDpMRn9MBN8oNan9M=";
+    extra-substituters = "https://exo.cachix.org https://cache.nixos-cuda.org";
   };
 
   outputs =
@@ -76,6 +76,8 @@
         let
           # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
           pkgsSwift = import inputs.nixpkgs-swift { inherit system; };
+
+          pkgsCuda = import ./nix/cuda-pkgs.nix { nixpkgs = inputs.nixpkgs; inherit system; };
         in
         {
           # Allow unfree for metal-toolchain (needed for Darwin Metal packages)
@@ -112,65 +114,137 @@
             };
           };
 
-          packages = {
-            babeld = pkgs.callPackage ./nix/babeld.nix { };
-          } // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
-            let
-              uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
-              mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package);
-              uvLockMlxVersion = mlxPackage.version;
-              uvLockMlxRev = builtins.elemAt (builtins.split "#" mlxPackage.source.git) 2;
-            in
-            {
-              metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
-              mlx = pkgs.callPackage ./nix/mlx.nix {
-                inherit (self'.packages) metal-toolchain;
-                inherit uvLockMlxVersion uvLockMlxRev;
-              };
-              default = self'.packages.exo;
-            }
-          );
+          packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin
+            (
+              let
+                uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
+                mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package);
+                uvLockMlxVersion = mlxPackage.version;
+                uvLockMlxRev = builtins.elemAt (builtins.split "#" mlxPackage.source.git) 2;
+              in
+              {
+                metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
+                mlx = pkgs.callPackage ./nix/mlx.nix {
+                  inherit (self'.packages) metal-toolchain;
+                  inherit uvLockMlxVersion uvLockMlxRev;
+                };
+                default = self'.packages.exo;
+              }
+            ) // lib.optionalAttrs (pkgsCuda != null) {
+            torch-cuda = pkgsCuda.python313Packages.torch;
+            vllm-cuda = pkgsCuda.python313Packages.vllm;
 
-          devShells.default = with pkgs; pkgs.mkShell {
-            inputsFrom = [ self'.checks.cargo-build ];
-
-            packages =
-              [
-                # FORMATTING
-                config.treefmt.build.wrapper
-
-                # PYTHON
-                python313
-                uv
-                ruff
-                basedpyright
-
-                # RUST
-                config.rust.toolchain
-                maturin
-
-                # NIX
-                nixpkgs-fmt
-
-                # SVELTE
-                nodejs
-
-                # MISC
-                just
-                jq
-              ]
-              ++ lib.optionals stdenv.isDarwin [
-                macmon
+            # Smoke test script for verifying vLLM + CUDA GPU setup
+            vllm-check = pkgs.writeShellApplication {
+              name = "vllm-check";
+              runtimeInputs = [
+                (pkgsCuda.python313.withPackages (ps: [ ps.torch ps.vllm ]))
               ];
+              # On non-NixOS hosts, NVIDIA driver libraries live in /usr/lib and must be
+              # LD_PRELOAD'd individually (adding the whole dir causes SIGILL from conflicts).
+              # These are: CUDA driver, NVML, and the PTX JIT compiler (for flash attention).
+              # libnvJitLink comes from the nix CUDA toolkit via LD_LIBRARY_PATH.
+              text = ''
+                for dir in /usr/lib/aarch64-linux-gnu /usr/lib/x86_64-linux-gnu /usr/lib; do
+                  if [ -e "$dir/libcuda.so.1" ]; then
+                    NVIDIA_LIBS="$dir/libcuda.so.1"
+                    for lib in libnvidia-ml.so.1 libnvidia-ptxjitcompiler.so.1; do
+                      [ -e "$dir/$lib" ] && NVIDIA_LIBS="$NVIDIA_LIBS:$dir/$lib"
+                    done
+                    export LD_PRELOAD="$NVIDIA_LIBS''${LD_PRELOAD:+:$LD_PRELOAD}"
+                    break
+                  fi
+                done
+                export LD_LIBRARY_PATH="${pkgsCuda.cudaPackages.libnvjitlink}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
+                exec python ${inputs.self + /tests/test_vllm_smoke.py}
+              '';
+            };
 
-            OPENSSL_NO_VENDOR = "1";
+            # exo with CUDA torch + vLLM — wraps the uv2nix-built package with host driver libs
+            exo-cuda = pkgs.writeShellApplication {
+              name = "exo-cuda";
+              runtimeInputs = [ self'.packages.exo-cuda-unwrapped ];
+              text = ''
+                for dir in /usr/lib/aarch64-linux-gnu /usr/lib/x86_64-linux-gnu /usr/lib; do
+                  if [ -e "$dir/libcuda.so.1" ]; then
+                    NVIDIA_LIBS="$dir/libcuda.so.1"
+                    for lib in libnvidia-ml.so.1 libnvidia-ptxjitcompiler.so.1; do
+                      [ -e "$dir/$lib" ] && NVIDIA_LIBS="$NVIDIA_LIBS:$dir/$lib"
+                    done
+                    export LD_PRELOAD="$NVIDIA_LIBS''${LD_PRELOAD:+:$LD_PRELOAD}"
+                    break
+                  fi
+                done
+                export LD_LIBRARY_PATH="${pkgsCuda.stdenv.cc.cc.lib}/lib:${pkgsCuda.cudaPackages.libnvjitlink}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
+                exec exo-cuda "$@"
+              '';
+            };
+          };
 
-            shellHook = ''
-              export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib"
-              ${lib.optionalString stdenv.isLinux ''
-                export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH"
-              ''}
-            '';
+          # CUDA development shell with torch + vLLM (aarch64-linux only)
+          devShells = lib.optionalAttrs (pkgsCuda != null)
+            {
+              cuda = pkgs.mkShell {
+                packages = [
+                  (pkgsCuda.python313.withPackages (ps: [
+                    ps.torch
+                    ps.vllm
+                  ]))
+                  pkgs.uv
+                  pkgs.just
+                ];
+
+                shellHook = ''
+                  echo "CUDA dev shell with torch + vLLM"
+                  python -c "import torch; print(f'PyTorch {torch.__version__}, CUDA: {torch.cuda.is_available()}')" 2>/dev/null || true
+                '';
+              };
+            } // {
+
+            default = with pkgs; pkgs.mkShell {
+              inputsFrom = [ self'.checks.cargo-build ];
+
+              packages =
+                [
+                  # FORMATTING
+                  config.treefmt.build.wrapper
+
+                  # PYTHON
+                  python313
+                  uv
+                  ruff
+                  basedpyright
+
+                  # RUST
+                  config.rust.toolchain
+                  maturin
+
+                  # NIX
+                  nixpkgs-fmt
+
+                  # SVELTE
+                  nodejs
+
+                  # MISC
+                  just
+                  jq
+                ]
+                ++ lib.optionals stdenv.isLinux [
+                  unixtools.ifconfig
+                ]
+                ++ lib.optionals stdenv.isDarwin [
+                  macmon
+                ];
+
+              OPENSSL_NO_VENDOR = "1";
+
+              shellHook = ''
+                export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib"
+                ${lib.optionalString stdenv.isLinux ''
+                  export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH"
+                ''}
+              '';
+            };
           };
         };
     };

justfile

@@ -15,6 +15,39 @@ check:
 sync:
     uv sync --all-packages
 
+sync-cuda:
+    #!/usr/bin/env bash
+    set -euo pipefail
+    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
 

nix/babeld.nix

@@ -1,25 +0,0 @@
-{ stdenv
-, lib
-, fetchurl
-}:
-stdenv.mkDerivation rec {
-  pname = "babeld";
-  version = "1.13.1";
-
-  src = fetchurl {
-    url = "https://www.irif.fr/~jch/software/files/${pname}-${version}.tar.gz";
-    hash = "sha256-FfJNJtoMz8Bzq83vAwnygeRoTyqnESb4JlcsTIRejdk=";
-  };
-
-  outputs = [
-    "out"
-    "man"
-  ];
-
-  makeFlags = [
-    "PREFIX=${placeholder "out"}"
-    "ETCDIR=${placeholder "out"}/etc"
-  ]
-  ++ lib.optional stdenv.isDarwin "LDLIBS=''";
-}
-

nix/cuda-pkgs.nix

@@ -0,0 +1,102 @@
+{ nixpkgs, system }:
+let
+  pkgs = import nixpkgs { inherit system; };
+in
+if system == "aarch64-linux" then
+  import nixpkgs
+  {
+    inherit system;
+    config = {
+      allowUnfree = true;
+      allowBroken = true;
+      allowUnsupportedSystem = true;
+      cudaSupport = true;
+      cudaCapabilities = [ "12.1" ];
+    };
+    overlays = [
+      (final: prev:
+        let
+          cudaCompatStub = cfinal: cprev: {
+            cuda_compat = prev.runCommand "cuda13.0-cuda_compat-stub" { } "mkdir -p $out";
+          };
+        in
+        {
+          cudaPackages = prev.cudaPackages_13.overrideScope cudaCompatStub // {
+            override = args:
+              (prev.cudaPackages_13.override args).overrideScope cudaCompatStub;
+          };
+
+          pythonPackagesExtensions = prev.pythonPackagesExtensions ++ [
+            (pyFinal: pyPrev: {
+              fastsafetensors = pyFinal.buildPythonPackage {
+                pname = "fastsafetensors";
+                version = "0.2.2";
+                src = prev.fetchFromGitHub {
+                  owner = "foundation-model-stack";
+                  repo = "fastsafetensors";
+                  rev = "v0.2.2";
+                  hash = "";
+                };
+                pyproject = true;
+                build-system = [
+                  pyFinal.setuptools
+                  pyFinal.pybind11
+                ];
+                buildInputs = [
+                  final.cudaPackages.cuda_cudart
+                  final.cudaPackages.cuda_nvml_dev
+                ];
+                nativeBuildInputs = [
+                  final.cudaPackages.cuda_nvcc
+                ];
+                dependencies = [
+                  pyFinal.typer
+                ];
+                env.CUDA_HOME = "${final.cudaPackages.cuda_nvcc}";
+                pythonImportsCheck = [ "fastsafetensors" ];
+              };
+              cupy = pyPrev.cupy.override {
+                cudaPackages = final.cudaPackages;
+              };
+
+              bitsandbytes = pyPrev.bitsandbytes.overrideAttrs (old: {
+                preConfigure = (old.preConfigure or "") + ''
+                  export CXXFLAGS="''${CXXFLAGS:-} -I${final.cudaPackages.cuda_crt}/include"
+                  export CUDAFLAGS="''${CUDAFLAGS:-} -I${final.cudaPackages.cuda_crt}/include"
+                  export CMAKE_CUDA_FLAGS="''${CMAKE_CUDA_FLAGS:-} -I${final.cudaPackages.cuda_crt}/include"
+                '';
+                buildInputs = (old.buildInputs or [ ]) ++ [ final.cudaPackages.cuda_crt ];
+              });
+
+              vllm = pyPrev.vllm.overrideAttrs (old: {
+                buildInputs = (old.buildInputs or [ ]) ++ [
+                  final.cuda_cccl_with_prefix
+                  final.cudaPackages.cuda_crt
+                ];
+                preConfigure = (old.preConfigure or "") + ''
+                  export CXXFLAGS="''${CXXFLAGS:-} -I${final.cuda_cccl_with_prefix}/include -I${final.cudaPackages.cuda_crt}/include"
+                  export CUDAFLAGS="''${CUDAFLAGS:-} -I${final.cuda_cccl_with_prefix}/include -I${final.cudaPackages.cuda_crt}/include"
+                '';
+              });
+            })
+          ];
+
+          magma-cuda-static = prev.magma-cuda-static.overrideAttrs (old: {
+            postPatch = (old.postPatch or "") + ''
+              sed -i '/err = cudaGetDeviceProperties( &prop, dev );/a\        int clock_khz = 0; cudaDeviceGetAttribute(\&clock_khz, cudaDevAttrClockRate, dev);' interface_cuda/interface.cpp
+              sed -i 's/prop\.clockRate/clock_khz/g' interface_cuda/interface.cpp
+            '';
+          });
+
+          cuda_cccl_with_prefix = prev.runCommand "cuda13.0-cuda_cccl-with-cccl-prefix" { } ''
+            mkdir -p $out/include
+            ln -s ${final.cudaPackages.cuda_cccl}/include $out/include/cccl
+          '';
+
+          opencv = prev.opencv.override { enableCuda = false; };
+          opencv4 = prev.opencv4.override { enableCuda = false; };
+        })
+    ];
+  }
+else
+  null

pyproject.toml

@@ -15,18 +15,17 @@ dependencies = [
   "huggingface-hub>=0.33.4",
   "psutil>=7.0.0",
   "loguru>=0.7.3",
-  "exo_pyo3_bindings",                         # rust bindings
+  "exo_pyo3_bindings",                       # rust bindings
   "anyio==4.11.0",
   "mlx; sys_platform == 'darwin'",
-  "mlx[cpu]==0.30.6; sys_platform == 'linux'",
+  "mlx==0.30.6; sys_platform == 'linux'",
   "mlx-lm",
-  "tiktoken>=0.12.0",                          # required for kimi k2 tokenizer
+  "tiktoken>=0.12.0",                        # required for kimi k2 tokenizer
   "hypercorn>=0.18.0",
   "openai-harmony>=0.0.8",
   "httpx>=0.28.1",
   "tomlkit>=0.14.0",
-  "pillow>=11.0,<12.0",                        # compatibility with mflux
-  "mflux==0.15.5",
+  "mflux==0.16.9; sys_platform == 'darwin'",
   "python-multipart>=0.0.21",
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
@@ -46,11 +45,13 @@ dev = [
   "ruff>=0.11.13",
 ]
 
-# mlx[cuda] requires a newer version of mlx. the ideal on linux is: default to mlx[cpu] unless[cuda] specified.
 [project.optional-dependencies]
-# cuda = [
-#     "mlx[cuda]==0.26.3",
-# ]
+cuda = [
+  "torch>=2.10.0; sys_platform == 'linux'",
+  "vllm>=0.13.0; sys_platform == 'linux'",
+  "mlx-cuda-13==0.30.6; sys_platform == 'linux'",
+  "fastsafetensors>=0.1.10; sys_platform == 'linux'",
+]
 
 ###
 # workspace configuration
@@ -63,10 +64,17 @@ members = ["rust/exo_pyo3_bindings", "bench"]
 exo_pyo3_bindings = { workspace = true }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
 mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" }
+torch = [{ index = "pytorch-cu130", marker = "sys_platform == 'linux'" }]
+vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
 # Uncomment to use local mlx/mlx-lm development versions:
 # mlx = { path = "/Users/Shared/mlx", editable=true }
 # mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
 
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true
+
 [build-system]
 requires = ["uv_build>=0.8.9,<0.9.0"]
 build-backend = "uv_build"
@@ -100,8 +108,16 @@ exclude = [
   "**/.direnv",
   "**/rust",
   "**/.github",
+  "**/vllm_patches",
+  "**/engines/vllm",
 ]
 stubPath = ".mlx_typings"
+extraPaths = [".cuda_typings"]
+
+[[tool.basedpyright.executionEnvironments]]
+root = "src/exo/worker/runner/vllm_inference"
+extraPaths = ["src", ".cuda_typings"]
+reportMissingModuleSource = false
 
 [[tool.basedpyright.executionEnvironments]]
 root = "src"
@@ -114,7 +130,22 @@ root = "src"
 [tool.uv]
 required-version = ">=0.8.6"
 prerelease = "allow"
-environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
+environments = [
+  "sys_platform == 'darwin'",
+  "sys_platform == 'linux' and platform_machine == 'aarch64'",
+]
+no-binary-package = ["vllm", "flashinfer-python"]
+no-build-isolation-package = ["vllm", "flashinfer-python"]
+extra-build-dependencies = { vllm = [
+  "cmake>=3.26.1",
+  "ninja",
+  "packaging>=24.2",
+  "setuptools>=77.0.3,<81.0.0",
+  "setuptools-scm>=8.0",
+  "wheel",
+  "jinja2",
+] }
+conflicts = [[{ package = "exo", extra = "cuda" }, { package = "exo-bench" }]]
 
 ###
 # ruff configuration
@@ -124,6 +155,7 @@ environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
 extend-exclude = [
   "shared/protobufs/**",
   "*mlx_typings/**",
+  "*cuda_typings/**",
   "rust/exo_pyo3_bindings/**",
   "bench/vendor/**",
 ]

python/parts.nix

@@ -3,6 +3,7 @@
   perSystem =
     { config, self', pkgs, lib, system, ... }:
     let
+      pkgsCuda = import ../nix/cuda-pkgs.nix { nixpkgs = inputs.nixpkgs; inherit system; };
       # Load workspace from uv.lock
       workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
         workspaceRoot = inputs.self;
@@ -99,16 +100,18 @@
           }
         );
 
+      baseOverlays = [
+        inputs.pyproject-build-systems.overlays.default
+        overlay
+        exoOverlay
+        buildSystemsOverlay
+        linuxOverlay
+      ];
+
       pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
         inherit python;
       }).overrideScope (
-        lib.composeManyExtensions [
-          inputs.pyproject-build-systems.overlays.default
-          overlay
-          exoOverlay
-          buildSystemsOverlay
-          linuxOverlay
-        ]
+        lib.composeManyExtensions baseOverlays
       );
       # mlx-cpu and mlx-cuda-13 both ship mlx/ site-packages files; keep first.
       # mlx-cpu/mlx-cuda-13 and nvidia-cudnn-cu12/cu13 ship overlapping files.
@@ -172,6 +175,39 @@
             --set EXO_RESOURCES_DIR ${inputs.self + /resources} \
             ${lib.optionalString pkgs.stdenv.hostPlatform.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
         '';
+
+      vllmEnv = pkgsCuda.python313.withPackages (ps: [ ps.vllm ps.fastsafetensors ]);
+
+      vllmSite = pkgs.runCommand "vllm-site-filtered" { } ''
+        mkdir -p $out
+        for pkg in ${vllmEnv}/${python.sitePackages}/*; do
+          name=$(basename "$pkg")
+          case "$name" in
+            anyio*|pydantic*) ;;
+            *) ln -s "$pkg" "$out/$name" ;;
+          esac
+        done
+      '';
+
+      exoCudaDeps = exoDeps // {
+        mlx-cuda-13 = [ ];
+      };
+
+      exoCudaVenv = (pythonSet.mkVirtualEnv "exo-cuda-env" exoCudaDeps).overrideAttrs {
+        venvIgnoreCollisions = venvCollisionPaths;
+      };
+
+      exoCudaPackage = pkgs.runCommand "exo-cuda"
+        {
+          nativeBuildInputs = [ pkgs.makeWrapper ];
+        }
+        ''
+          mkdir -p $out/bin
+          makeWrapper ${exoCudaVenv}/bin/exo $out/bin/exo-cuda \
+            --set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
+            --set EXO_RESOURCES_DIR ${inputs.self + /resources} \
+            --prefix PYTHONPATH : "${vllmSite}"
+        '';
     in
     {
       # Python package only available on macOS (requires MLX/Metal)
@@ -180,12 +216,13 @@
           exo = exoPackage;
           # Test environment for running pytest outside of Nix sandbox (needs GPU access)
           exo-test-env = testVenv;
-        } // {
+        } // lib.optionalAttrs (pkgsCuda != null) {
+        exo-cuda-unwrapped = exoCudaPackage;
+      } // {
         exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
         exo-eval = mkBenchScript "exo-eval" (inputs.self + /bench/exo_eval.py);
         exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
         exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
-        exo-mlx-bandwidth-test = mkPythonScript "exo-mlx-bandwidth-test" (inputs.self + /bench/test_mlx_bandwidth.py);
       };
 
       checks = {

result

@@ -1 +0,0 @@
-/nix/store/41rd4g2p69a2106qscl3xvkzmh6j7nry-babblerd

rust/babblerd/Cargo.toml

@@ -1,19 +0,0 @@
-[package]
-name = "babblerd"
-version.workspace = true
-edition.workspace = true
-
-[dependencies]
-color-eyre = "0.6.5"
-futures-lite.workspace = true
-ipnet = "2.12.0"
-libc = "0.2.183"
-n0-watcher = "0.6.1"
-netwatch = "0.14.0"
-rand = "0.10.0"
-tokio = { workspace = true, features = ["full"] }
-tracing = "0.1.44"
-tracing-subscriber = { version = "0.3.22", features = ["env-filter"] }
-
-[lints]
-workspace = true

rust/babblerd/src/lib.rs

@@ -1,548 +0,0 @@
-pub use babel::{babel, handle_listener};
-pub use error::{BabbleError, Result};
-pub use if_watcher::{PREFIX, watch};
-pub mod error {
-    use std::io;
-
-    pub type Result<T> = core::result::Result<T, BabbleError>;
-    #[derive(Debug)]
-    pub enum BabbleError {
-        Io(io::Error),
-        Unspecified,
-        BabeldCrashed(Option<i32>),
-        FailedToSetIp,
-        Other(String),
-    }
-    impl std::fmt::Display for BabbleError {
-        // use the debug display for now
-        #[inline]
-        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-            write!(f, "{self:?}")
-        }
-    }
-    impl std::error::Error for BabbleError {}
-    impl From<io::Error> for BabbleError {
-        #[inline]
-        fn from(value: io::Error) -> Self {
-            Self::Io(value)
-        }
-    }
-}
-pub mod if_watcher {
-    #[cfg(target_os = "linux")]
-    use std::path::PathBuf;
-    use std::{
-        collections::HashSet,
-        net::{IpAddr, Ipv6Addr},
-    };
-
-    use futures_lite::StreamExt;
-    use ipnet::Ipv6Net;
-    use n0_watcher::Watcher;
-    use netwatch::interfaces::{Interface, IpNet};
-    use tokio::sync::mpsc;
-
-    use crate::{
-        BabbleError, Result,
-        babel::Babble,
-        ip_manager::{add_ip, remove_ip},
-    };
-
-    pub const PREFIX: Ipv6Net = Ipv6Net::new_assert(
-        Ipv6Addr::new(0xfde0, 0x20c6, 0x1fa7, 0xffff, 0, 0, 0, 0),
-        48,
-    );
-
-    trait IfaceExt {
-        fn has_link_local_v6(&self) -> bool;
-        fn is_real_interface(&self) -> bool;
-        fn will_babel(&self) -> bool;
-        fn get_v6_in(&self, prefix: Ipv6Net) -> Option<Ipv6Net>;
-    }
-    impl IfaceExt for Interface {
-        fn will_babel(&self) -> bool {
-            self.has_link_local_v6() && self.is_real_interface() && self.is_up()
-        }
-
-        fn has_link_local_v6(&self) -> bool {
-            let mut has = false;
-            for addr in self.addrs() {
-                let IpAddr::V6(a) = addr.addr() else {
-                    continue;
-                };
-                if a.is_unicast_link_local() {
-                    has = true;
-                    break;
-                }
-            }
-            has
-        }
-        fn is_real_interface(&self) -> bool {
-            // macos is weird. en0 & en1 are ethernet & wifi (varies which is which by device). en3+ is thunderbolt, but at some point becomes usb ethernet.
-            if self.name().strip_prefix("en").is_none()
-            //.and_then(|s| s.parse::<u8>().ok())
-            //.is_none_or(|_n| false)
-            {
-                return false;
-            }
-            #[cfg(target_os = "linux")]
-            {
-                if !PathBuf::from(format!("/sys/class/net/{}/device", self.name())).exists() {
-                    tracing::debug!(
-                        "skipping interface {} as it doesn't correspond to a physical link",
-                        self.name()
-                    );
-                    return false;
-                }
-                let dev_type_path = PathBuf::from(format!("/sys/class/net/{}/type", self.name()));
-                if !dev_type_path.exists() {
-                    tracing::debug!(
-                        "skipping interface {} with no type file at {:?}",
-                        self.name(),
-                        dev_type_path.to_str()
-                    );
-                    return false;
-                }
-                let Ok(dev_type) = std::fs::read_to_string(dev_type_path) else {
-                    return false;
-                };
-                if dev_type.trim() != "1" {
-                    tracing::debug!(
-                        "skipping interface {} with type {:?}",
-                        self.name(),
-                        dev_type
-                    );
-                    return false;
-                }
-            }
-            true
-        }
-        fn get_v6_in(&self, prefix: Ipv6Net) -> Option<Ipv6Net> {
-            for addr in self.addrs() {
-                if let IpNet::V6(v6) = addr
-                    && prefix.contains(&v6.addr())
-                {
-                    return Some(v6);
-                }
-            }
-            None
-        }
-    }
-
-    #[tracing::instrument(skip(send))]
-    pub async fn watch(my_range: Ipv6Net, send: mpsc::Sender<Babble>) -> Result<()> {
-        assert!(PREFIX.contains(&my_range));
-        let mut ready_ifaces = HashSet::new();
-        // 0 is reserved for the first loopback address
-        let mut iface_num: u16 = 1;
-
-        tracing::info!("starting interface monitor");
-        let mon = netwatch::netmon::Monitor::new()
-            .await
-            .map_err(|_| BabbleError::Unspecified)?;
-        let mut mon_stream = mon.interface_state().stream();
-
-        while let Some(s) = mon_stream.next().await {
-            for iface in s.interfaces.values() {
-                if !iface.is_real_interface() {
-                    continue;
-                }
-                for addr in iface.addrs() {
-                    if let IpNet::V6(v6) = addr
-                        && PREFIX.contains(&v6.addr())
-                        && !my_range.contains(&v6.addr())
-                    {
-                        tracing::info!("removing stale ip {v6} from {}", iface.name());
-                        // don't really care if this fails
-                        if let Err(e) = remove_ip(v6, iface).await {
-                            tracing::warn!(%e, "failed to remove ip");
-                        }
-                    }
-                }
-                if !iface.will_babel() {
-                    continue;
-                }
-                if iface.get_v6_in(my_range).is_none() {
-                    let addr = Ipv6Net::new_assert(
-                        Ipv6Addr::from_bits(my_range.addr().to_bits() | u128::from(iface_num)),
-                        128,
-                    );
-                    iface_num += 1;
-                    assert!(iface_num < u16::MAX, "Really? u16::MAX interfaces?");
-                    tracing::info!("adding new ip {addr} to {}", iface.name());
-                    add_ip(addr, iface).await?;
-                }
-
-                if ready_ifaces.insert(iface.name().to_owned()) {
-                    tracing::info!("telling babeld to watch {}", iface.name());
-                    let Ok(()) = send.send(Babble::AddIface(iface.name().to_owned())).await else {
-                        return Ok(());
-                    };
-                }
-            }
-        }
-        tracing::info!("stopping interface monitor");
-
-        Ok(())
-    }
-}
-pub mod babel {
-    #[tracing::instrument(skip_all)]
-    pub async fn handle_listener(sock: UnixStream, mut receiver: broadcast::Receiver<String>) {
-        tracing::info!("new socket conn");
-        let (reader, mut write) = sock.into_split();
-        let mut reader = BufReader::new(reader).lines();
-        loop {
-            tokio::select! {
-                read = reader.next_line() => {
-                    let Ok(Some(_)) = read else { break; };
-                }
-                recv = receiver.recv() => {
-                    let res = match recv {
-                        Err(broadcast::error::RecvError::Lagged(_)) => { tracing::warn!("receiver lagged"); continue; },
-                        Err(broadcast::error::RecvError::Closed) => { tracing::debug!("receiver closed, dropping connection"); break; },
-                        Ok(s) => write.write_all(format!("{s}\n").as_bytes()).await,
-                    };
-                    if let Err(e) = res { tracing::warn!(error=%e, "failed to write to socket"); break; }
-                }
-            };
-        }
-        tracing::info!("closing socket conn");
-        _ = write.shutdown().await;
-    }
-
-    #[cfg(target_os = "macos")]
-    const PRIVATE_SOCK_PATH: &str = "/var/run/babbler/private/babeld.sock";
-    #[cfg(target_os = "linux")]
-    const PRIVATE_SOCK_PATH: &str = "/run/babbler/private/babeld.sock";
-    #[cfg(target_os = "macos")]
-    const PRIVATE_DIR: &str = "/var/run/babbler/private";
-    #[cfg(target_os = "linux")]
-    const PRIVATE_DIR: &str = "/run/babbler/private";
-
-    use ipnet::Ipv6Net;
-    use std::fs::Permissions;
-    use std::io;
-    use std::os::unix::fs::PermissionsExt;
-    use tokio::time::Duration;
-
-    use futures_lite::FutureExt;
-    use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, Lines};
-    use tokio::net::UnixStream;
-    use tokio::net::unix::{OwnedReadHalf, OwnedWriteHalf};
-    use tokio::process::Command;
-    use tokio::sync::{broadcast, mpsc};
-
-    use crate::{BabbleError, Result};
-
-    #[derive(Debug)]
-    pub enum Babble {
-        AddIface(String),
-    }
-
-    pub struct BabeldProcess {
-        proc: tokio::process::Child,
-    }
-    impl Drop for BabeldProcess {
-        #[inline]
-        fn drop(&mut self) {
-            // emergency sigkill babeld process to prevent leakage
-            match self.proc.try_wait() {
-                Ok(None) => {}
-                Ok(Some(sc)) => {
-                    if !sc.success() {
-                        _ = self.proc.start_kill();
-                    }
-                }
-                _ => {
-                    _ = self.proc.start_kill();
-                }
-            }
-        }
-    }
-    impl BabeldProcess {
-        #[tracing::instrument]
-        async fn spawn(my_range: Ipv6Net, iface: String) -> Result<Self> {
-            tokio::fs::create_dir_all(PRIVATE_DIR).await?;
-            tokio::fs::set_permissions(PRIVATE_DIR, Permissions::from_mode(0o0700)).await?;
-            tracing::info!("spawning babeld socket in {PRIVATE_SOCK_PATH}");
-            let res = match Command::new("babeld")
-                .arg("-G")
-                .arg(PRIVATE_SOCK_PATH)
-                .arg("-I")
-                .arg(format!("{PRIVATE_DIR}/babeld.pid"))
-                .arg("-C")
-                .arg(format!("redistribute local ip {my_range}"))
-                .arg("-C")
-                .arg("redistribute local deny")
-                .arg(iface)
-                .spawn()
-            {
-                Ok(proc) => Ok(Self { proc }),
-                Err(e) => {
-                    tracing::warn!(error=%e, "failed to spawn babeld");
-                    Err(e.into())
-                }
-            };
-            tokio::time::sleep(Duration::from_millis(10)).await;
-            while !matches!(tokio::fs::try_exists(PRIVATE_SOCK_PATH).await, Ok(true)) {
-                tracing::info!("where is the sock");
-                tokio::time::sleep(Duration::from_secs(1)).await;
-            }
-            res
-        }
-
-        #[tracing::instrument(skip(read, write, send))]
-        async fn query(
-            read: &mut Lines<BufReader<OwnedReadHalf>>,
-            write: &mut OwnedWriteHalf,
-            send: &broadcast::Sender<String>,
-            cmd: &str,
-        ) -> io::Result<Option<bool>> {
-            write.write_all(cmd.as_bytes()).await?;
-            loop {
-                let Some(line) = read.next_line().await? else {
-                    tracing::warn!("babeld closed unexpectedly");
-                    return Ok(None);
-                };
-                tracing::info!("[babel] {:?}", line);
-                let ret = match line.as_str() {
-                    "ok" => Ok(Some(true)),
-                    "bad" => {
-                        tracing::warn!("malformed message sent to babeld");
-                        Ok(Some(false))
-                    }
-                    _ if line.starts_with("no") => {
-                        tracing::warn!("message rejected");
-                        Ok(Some(false))
-                    }
-                    _ => Ok(None),
-                };
-                let Ok(_) = send.send(line) else {
-                    return Ok(None);
-                };
-                if !matches!(ret, Ok(None)) {
-                    return ret;
-                }
-            }
-        }
-
-        #[tracing::instrument(skip_all)]
-        async fn supervise(
-            &self,
-            mut recv: mpsc::Receiver<Babble>,
-            send: broadcast::Sender<String>,
-        ) -> Result<()> {
-            std::fs::set_permissions(PRIVATE_SOCK_PATH, Permissions::from_mode(0o0600))?;
-            tracing::debug!("connecting to babeld.sock");
-            let (reader, mut writer) = UnixStream::connect(PRIVATE_SOCK_PATH).await?.into_split();
-            let mut babel_lines = BufReader::new(reader).lines();
-            while let Some(s) = babel_lines.next_line().await? {
-                tracing::debug!("[babeld] {}", s);
-                if s == "ok" {
-                    break;
-                }
-            }
-            tracing::info!("babeld ok");
-            /* TODO(evan): push rather than pull
-            if Self::query(&mut babel_lines, &mut writer, "monitor\n")
-                .await?
-                .is_none()
-            {
-                return Ok(());
-            };
-            */
-
-            let mut interval = tokio::time::interval(Duration::from_secs(5));
-            interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
-            loop {
-                tokio::select! {
-                    _ = interval.tick() => {
-                        if Self::query(&mut babel_lines, &mut writer, &send, "dump\n")
-                            .await?
-                            .is_none()
-                        {
-                            return Ok(());
-                        }
-                    },
-                    babble = recv.recv() => {
-                        tracing::debug!("[babble] {:?}", babble);
-                        let Some(babble) = babble else {
-                            break;
-                        };
-                        match babble {
-                            Babble::AddIface(iface) => {
-                                Self::query(&mut babel_lines, &mut writer, &send, format!("interface {iface}\n").as_ref()).await?;
-                            }
-                        }
-                    },
-                    line = babel_lines.next_line() => {
-                        let Ok(Some(line)) = line else {
-                            break;
-                        };
-                        tracing::debug!("[babeld] {}", line);
-                        let Ok(_) = send.send(line) else {
-                            break;
-                        };
-                    },
-                }
-            }
-            Ok(())
-        }
-
-        async fn shutdown(mut self) -> Result<()> {
-            let kill_res = if let Some(pid) = self.proc.id() {
-                let pid: libc::pid_t = pid.try_into().expect("pid overflow");
-                // SAFETY: pid >= 0, freshly checked.
-                let rc = unsafe { libc::kill(pid, libc::SIGINT) };
-                let rc_err = if rc != 0 && rc != libc::ESRCH {
-                    Err(io::Error::last_os_error().into())
-                } else {
-                    Ok(())
-                };
-                let exit_code = async { Some(self.proc.wait().await) }
-                    .or(async {
-                        tokio::time::sleep(Duration::from_secs(5)).await;
-                        None
-                    })
-                    .await;
-                match exit_code {
-                    Some(Ok(code)) => {
-                        if code.success() {
-                            rc_err
-                        } else {
-                            rc_err.and_then(|()| Err(BabbleError::BabeldCrashed(code.code())))
-                        }
-                    }
-                    Some(Err(e)) => Err(e.into()),
-                    None => {
-                        self.proc.kill().await?;
-                        rc_err.and(Err(BabbleError::BabeldCrashed(None)))
-                    }
-                }
-            } else {
-                Ok(())
-            };
-            let rem_res = match std::fs::remove_file(PRIVATE_SOCK_PATH) {
-                Ok(()) => Ok(()),
-                Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
-                Err(e) => Err(e.into()),
-            };
-            kill_res.and(rem_res)
-        }
-    }
-
-    #[tracing::instrument(skip(send, recv))]
-    pub async fn babel(
-        my_range: Ipv6Net,
-        mut recv: mpsc::Receiver<Babble>,
-        send: broadcast::Sender<String>,
-    ) -> Result<()> {
-        let iface = loop {
-            match recv.recv().await {
-                Some(Babble::AddIface(iface)) => {
-                    break iface;
-                }
-                None => return Ok(()),
-            }
-        };
-
-        let babel = BabeldProcess::spawn(my_range, iface).await?;
-        let res1 = babel.supervise(recv, send).await;
-        let res2 = babel.shutdown().await;
-        res1.and(res2)
-    }
-}
-pub(crate) mod ip_manager {
-    pub use sys::add_ip;
-    pub use sys::remove_ip;
-
-    #[cfg(target_os = "linux")]
-    mod sys {
-        use ipnet::Ipv6Net;
-        use netwatch::interfaces::Interface;
-
-        use crate::{BabbleError, Result};
-        use tokio::process::Command;
-
-        #[tracing::instrument]
-        pub async fn add_ip(subnet: Ipv6Net, iface: &Interface) -> Result<()> {
-            let out = Command::new("ip")
-                .arg("addr")
-                .arg("add")
-                .arg(format!("{subnet}"))
-                .arg("dev")
-                .arg(iface.name())
-                .output()
-                .await?;
-            if out.status.success() {
-                Ok(())
-            } else {
-                Err(BabbleError::FailedToSetIp)
-            }
-        }
-
-        #[tracing::instrument]
-        pub async fn remove_ip(v6: Ipv6Net, iface: &Interface) -> Result<()> {
-            let out = Command::new("ip")
-                .arg("addr")
-                .arg("del")
-                .arg(format!("{v6}"))
-                .arg("dev")
-                .arg(iface.name())
-                .output()
-                .await?;
-            if out.status.success() {
-                Ok(())
-            } else {
-                let std_err = String::from_utf8_lossy(&out.stdout);
-                tracing::debug!(%std_err);
-                Err(BabbleError::FailedToSetIp)
-            }
-        }
-    }
-
-    #[cfg(target_os = "macos")]
-    mod sys {
-        use ipnet::Ipv6Net;
-        use netwatch::interfaces::Interface;
-
-        use crate::BabbleError;
-        use crate::Result;
-        use tokio::process::Command;
-
-        #[tracing::instrument]
-        pub async fn add_ip(subnet: Ipv6Net, iface: &Interface) -> Result<()> {
-            let out = Command::new("ifconfig")
-                .arg(iface.name())
-                .arg("inet6")
-                .arg(format!("{subnet}"))
-                .arg("add")
-                .output()
-                .await?;
-            if out.status.success() {
-                Ok(())
-            } else {
-                Err(BabbleError::FailedToSetIp)
-            }
-        }
-
-        #[tracing::instrument]
-        pub async fn remove_ip(v6: Ipv6Net, iface: &Interface) -> Result<()> {
-            let out = Command::new("ifconfig")
-                .arg(iface.name())
-                .arg("inet6")
-                .arg(format!("{v6}"))
-                .arg("delete")
-                .output()
-                .await?;
-            if out.status.success() {
-                Ok(())
-            } else {
-                let std_err = String::from_utf8_lossy(&out.stdout);
-                tracing::debug!(%std_err);
-                Err(BabbleError::FailedToSetIp)
-            }
-        }
-    }
-}

rust/babblerd/src/main.rs

@@ -1,167 +0,0 @@
-use std::{fs::Permissions, io, net::Ipv6Addr, os::unix::fs::PermissionsExt};
-
-use babblerd::babel::handle_listener;
-use color_eyre::eyre::{WrapErr, eyre};
-use ipnet::Ipv6Net;
-use tokio::{
-    net::UnixListener,
-    signal,
-    sync::{broadcast, mpsc},
-    task::{JoinHandle, JoinSet},
-};
-
-#[cfg(target_os = "macos")]
-const PUBLIC_DIR: &str = "/var/run/babbler";
-#[cfg(target_os = "linux")]
-const PUBLIC_DIR: &str = "/run/babbler";
-#[cfg(target_os = "macos")]
-const PUBLIC_SOCK_PATH: &str = "/var/run/babbler/babblerd.sock";
-#[cfg(target_os = "linux")]
-const PUBLIC_SOCK_PATH: &str = "/run/babbler/babblerd.sock";
-
-enum State {
-    Idle,
-    Active {
-        recv: broadcast::Receiver<String>,
-        babel: JoinHandle<babblerd::Result<()>>,
-        watcher: JoinHandle<babblerd::Result<()>>,
-        listeners: JoinSet<()>,
-    },
-}
-
-#[tokio::main]
-async fn main() -> color_eyre::Result<()> {
-    color_eyre::install()?;
-    tracing_subscriber::fmt()
-        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
-        .init();
-    // cleanup old data
-    match std::fs::remove_file(PUBLIC_SOCK_PATH) {
-        Err(e) if e.kind() != io::ErrorKind::NotFound => {
-            return Err(e.into());
-        }
-        Ok(()) => {
-            tracing::info!("cleaned up old file at {PUBLIC_SOCK_PATH}");
-        }
-        _ => {}
-    }
-    std::fs::create_dir_all(PUBLIC_DIR)?;
-    // make our directory world readable
-    if let Err(e) = std::fs::set_permissions(PUBLIC_DIR, Permissions::from_mode(0o0755)) {
-        if e.kind() == io::ErrorKind::PermissionDenied {
-            return Err(eyre!(
-                "Insufficient permissions to run daemon -- did you forget sudo?"
-            ));
-        }
-        return Err(e.into());
-    }
-
-    let res = inner_main().await;
-    _ = std::fs::remove_file(PUBLIC_SOCK_PATH);
-    res
-}
-
-async fn inner_main() -> color_eyre::Result<()> {
-    tracing::info!("creating socket at {PUBLIC_SOCK_PATH}");
-    let public_socket = UnixListener::bind(PUBLIC_SOCK_PATH)?;
-    // make our socket world accessible
-    std::fs::set_permissions(PUBLIC_SOCK_PATH, Permissions::from_mode(0o0666))?;
-    let mut babbler: State = State::Idle;
-
-    loop {
-        babbler = match babbler {
-            State::Idle => {
-                tracing::info!("waiting for connection");
-                tokio::select! {
-                    sig = signal::ctrl_c() => {
-                        sig?;
-                        break;
-                    }
-                    sock = public_socket.accept() => {
-                        let sock = sock?.0;
-                        tracing::info!("starting babeld");
-                        let (br_send, br_recv) = broadcast::channel(1024);
-                        let (mp_send, mp_recv) = mpsc::channel(32);
-                        // node id is a PREFIX/64, NODE_ID/48 and an INTERFACE_ID/16
-                        let ip_node_id = u128::from(rand::random::<u64>() & (u64::MAX>>16)) << 16;
-                        let my_range = Ipv6Net::new_assert(
-                            Ipv6Addr::from_bits(babblerd::PREFIX.addr().to_bits() | ip_node_id),
-                            112,
-                        );
-                        let babel = tokio::spawn(babblerd::babel(my_range, mp_recv, br_send));
-                        let watcher = tokio::spawn(babblerd::watch(my_range, mp_send));
-                        let mut listeners = JoinSet::new();
-                        listeners.spawn(handle_listener(sock, br_recv.resubscribe()));
-                        State::Active { recv: br_recv, babel, watcher, listeners }
-                    }
-                }
-            }
-            State::Active {
-                recv,
-                mut babel,
-                mut watcher,
-                mut listeners,
-            } => {
-                tokio::select! {
-                    sig = signal::ctrl_c() => {
-                        sig?;
-                        drop(recv);
-                        watcher.abort();
-                        if let Ok(e) = watcher.await {
-                            e.wrap_err("while ctrl-c")?;
-                        }
-                        babel.await?.wrap_err("while ctrl-c")?;
-                        while let Some(res) = listeners.join_next().await {
-                            res.wrap_err("while ctrl-c")?;
-                        }
-                        break;
-                    }
-                    next_join_result = listeners.join_next(), if !listeners.is_empty() => {
-                        next_join_result.expect("checked")?;
-                        tracing::info!("dropped a listener");
-                        if listeners.is_empty() {
-                            tracing::info!("stopping babeld");
-                            drop(recv);
-
-                            watcher.abort();
-                            if let Ok(e) = watcher.await {
-                                e.wrap_err("while closing listeners")?;
-                            }
-                            babel.await?.wrap_err("while closing listeners")?;
-
-                            State::Idle
-                        } else {
-                            State::Active { recv, babel, watcher, listeners }
-                        }
-                    }
-                    sock = public_socket.accept() => {
-                        listeners.spawn(handle_listener(sock?.0, recv.resubscribe()));
-                        State::Active { recv, babel, watcher, listeners }
-                    }
-                    res = &mut watcher => {
-                        res??;
-                        drop(recv);
-                        babel.await?.wrap_err("while closing watcher")?;
-                        while let Some(res2) = listeners.join_next().await {
-                            res2.wrap_err("while closing watcher")?;
-                        }
-                        State::Idle
-                    }
-                    res = &mut babel => {
-                        res??;
-                        drop(recv);
-                        watcher.abort();
-                        if let Ok(e) = watcher.await {
-                            e.wrap_err("while closing babeld")?;
-                        }
-                        while let Some(res2) = listeners.join_next().await {
-                            res2.wrap_err("while closing babeld")?;
-                        }
-                        State::Idle
-                    }
-                }
-            }
-        };
-    }
-    Ok(())
-}

rust/parts.nix

@@ -1,7 +1,7 @@
 { inputs, ... }:
 {
   perSystem =
-    { inputs', self', pkgs, lib, ... }:
+    { inputs', pkgs, lib, ... }:
     let
       # Fenix nightly toolchain with all components
       rustToolchain = inputs'.fenix.packages.stable.withComponents [
@@ -79,7 +79,7 @@
       };
 
       config = {
-        packages = rec {
+        packages = {
           # Python bindings wheel via maturin
           exo_pyo3_bindings = craneLib.buildPackage (
             commonArgs
@@ -110,19 +110,6 @@
               '';
             }
           );
-          babblerd-unwrapped = craneLib.buildPackage (
-            commonArgs // {
-              inherit cargoArtifacts;
-              pname = "babblerd-unwrapped";
-            }
-          );
-          babblerd = pkgs.writeShellApplication {
-            name = "babblerd";
-            runtimeInputs = [ self'.packages.babeld ];
-            text = ''
-              exec ${babblerd-unwrapped}/bin/babblerd "$@"
-            '';
-          };
         };
 
         checks = {

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/__main__.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import sys
 from collections.abc import Sequence
 from multiprocessing import freeze_support

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/download/download_utils.py

@@ -753,7 +753,10 @@ async def download_shard(
     )
     filtered_file_list = list(
         filter_repo_objects(
-            file_list, allow_patterns=allow_patterns, key=lambda x: x.path
+            file_list,
+            allow_patterns=allow_patterns,
+            ignore_patterns=["original/*", "metal/*"],
+            key=lambda x: x.path,
         )
     )
 

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/adapters/ollama.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import json
 from collections.abc import AsyncGenerator
 from typing import Any

src/exo/master/api.py

@@ -341,6 +341,7 @@ class API:
         self.app.get("/ollama/api/version")(self.ollama_version)
 
         self.app.get("/state")(lambda: self.state)
+        self.app.get("/capabilities")(self._get_capabilities)
         self.app.get("/events")(self.stream_events)
         self.app.post("/download/start")(self.start_download)
         self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download)
@@ -413,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
@@ -448,18 +450,15 @@ class API:
                 status_code=400, detail=f"Failed to load model card: {exc}"
             ) from exc
         instance_combinations: list[tuple[Sharding, InstanceMeta, int]] = []
+        node_count = len(list(self.state.topology.list_nodes()))
+
         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, len(list(self.state.topology.list_nodes())) + 1
-                        )
-                    ]
+                    [(sharding, instance_meta, i) for i in range(1, node_count + 1)]
                 )
-        # TODO: PDD
-        # instance_combinations.append((Sharding.PrefillDecodeDisaggregation, InstanceMeta.MlxRing, 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:
             try:
@@ -472,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,
@@ -741,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)
@@ -749,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.
@@ -781,6 +788,9 @@ class API:
             )
         return resolved_model
 
+    def _get_capabilities(self) -> dict[str, bool]:
+        return {"vllm_available": any(self.state.node_vllm.values())}
+
     def stream_events(self) -> StreamingResponse:
         def _generate_json_array(events: Iterable[Event]) -> Iterable[str]:
             yield "["

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

@@ -41,6 +41,7 @@ from exo.shared.types.worker.instances import (
     InstanceMeta,
     MlxJacclInstance,
     MlxRingInstance,
+    VllmInstance,
 )
 from exo.shared.types.worker.shards import Sharding
 
@@ -66,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 = [
@@ -78,8 +112,11 @@ def place_instance(
             for cycle in candidate_cycles
             if required_nodes.issubset(cycle.node_ids)
         ]
+    required_memory = command.model_card.storage_size
+    if command.instance_meta == InstanceMeta.Vllm:
+        required_memory = Memory.from_bytes(int(required_memory.in_bytes * 1.3))
     cycles_with_sufficient_memory = filter_cycles_by_memory(
-        candidate_cycles, node_memory, command.model_card.storage_size
+        candidate_cycles, node_memory, required_memory
     )
     if len(cycles_with_sufficient_memory) == 0:
         raise ValueError("No cycles found with sufficient memory")
@@ -139,7 +176,7 @@ def place_instance(
     )
 
     # Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
-    if len(selected_cycle) == 1:
+    if len(selected_cycle) == 1 and command.instance_meta != InstanceMeta.Vllm:
         command.instance_meta = InstanceMeta.MlxRing
         command.sharding = Sharding.Pipeline
 
@@ -199,6 +236,11 @@ def place_instance(
                 hosts_by_node=hosts_by_node,
                 ephemeral_port=ephemeral_port,
             )
+        case InstanceMeta.Vllm:
+            target_instances[instance_id] = VllmInstance(
+                instance_id=instance_id,
+                shard_assignments=shard_assignments,
+            )
 
     return target_instances
 

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/apply.py

@@ -49,9 +49,11 @@ from exo.utils.info_gatherer.info_gatherer import (
     NodeConfig,
     NodeDiskUsage,
     NodeNetworkInterfaces,
+    NvmlMetrics,
     RdmaCtlStatus,
     StaticNodeInformation,
     ThunderboltBridgeInfo,
+    VllmCapability,
 )
 
 
@@ -363,6 +365,16 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 **state.node_rdma_ctl,
                 event.node_id: NodeRdmaCtlStatus(enabled=info.enabled),
             }
+        case NvmlMetrics():
+            update["node_system"] = {
+                **state.node_system,
+                event.node_id: info.system_profile,
+            }
+        case VllmCapability():
+            update["node_vllm"] = {
+                **state.node_vllm,
+                event.node_id: info.available,
+            }
 
     return state.model_copy(update=update)
 

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
@@ -258,21 +302,24 @@ async def fetch_safetensors_size(model_id: ModelId) -> Memory:
 
     target_dir = (await ensure_models_dir()) / model_id.normalize()
     await aios.makedirs(target_dir, exist_ok=True)
-    index_path = await download_file_with_retry(
-        model_id,
-        "main",
-        "model.safetensors.index.json",
-        target_dir,
-        lambda curr_bytes, total_bytes, is_renamed: logger.debug(
-            f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
-        ),
-    )
-    async with aiofiles.open(index_path, "r") as f:
-        index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
+    try:
+        index_path = await download_file_with_retry(
+            model_id,
+            "main",
+            "model.safetensors.index.json",
+            target_dir,
+            lambda curr_bytes, total_bytes, is_renamed: logger.debug(
+                f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
+            ),
+        )
+        async with aiofiles.open(index_path, "r") as f:
+            index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
 
-    metadata = index_data.metadata
-    if metadata is not None:
-        return Memory.from_bytes(metadata.total_size)
+        metadata = index_data.metadata
+        if metadata is not None:
+            return Memory.from_bytes(metadata.total_size)
+    except FileNotFoundError:
+        pass
 
     info = model_info(model_id)
     if info.safetensors is 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/tracing.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import json
 import time
 from collections import defaultdict

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/mlx.py

@@ -12,11 +12,14 @@ from mlx_lm.models.cache import (
     RotatingKVCache,
 )
 
-# This list contains one cache entry per transformer layer
-KVCacheType = Sequence[
+from exo.worker.engines.vllm.kv_cache import TorchKVCache
+
+MLXCacheType = Sequence[
     KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList
 ]
 
+KVCacheType = MLXCacheType | TorchKVCache
+
 
 # Model is a wrapper function to fix the fact that mlx is not strongly typed in the same way that EXO is.
 # For example - MLX has no guarantee of the interface that nn.Module will expose. But we need a guarantee that it has a __call__() function
@@ -26,6 +29,6 @@ class Model(nn.Module):
     def __call__(
         self,
         x: mx.array,
-        cache: KVCacheType | None,
+        cache: MLXCacheType | None,
         input_embeddings: mx.array | None = None,
     ) -> mx.array: ...

src/exo/shared/types/ollama_api.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import time
 from typing import Any, Literal
 

src/exo/shared/types/profiling.py

@@ -61,6 +61,7 @@ class SystemPerformanceProfile(CamelCaseModel):
 
     gpu_usage: float = 0.0
     temp: float = 0.0
+    # Best effort sys power
     sys_power: float = 0.0
     pcpu_usage: float = 0.0
     ecpu_usage: float = 0.0

src/exo/shared/types/state.py

@@ -57,6 +57,7 @@ class State(CamelCaseModel):
     node_thunderbolt: Mapping[NodeId, NodeThunderboltInfo] = {}
     node_thunderbolt_bridge: Mapping[NodeId, ThunderboltBridgeStatus] = {}
     node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] = {}
+    node_vllm: Mapping[NodeId, bool] = {}
 
     # Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes)
     thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = []

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/instances.py

@@ -15,6 +15,7 @@ class InstanceId(Id):
 class InstanceMeta(str, Enum):
     MlxRing = "MlxRing"
     MlxJaccl = "MlxJaccl"
+    Vllm = "Vllm"
 
 
 class BaseInstance(TaggedModel):
@@ -35,8 +36,12 @@ class MlxJacclInstance(BaseInstance):
     jaccl_coordinators: dict[NodeId, str]
 
 
+class VllmInstance(BaseInstance):
+    pass
+
+
 # TODO: Single node instance
-Instance = MlxRingInstance | MlxJacclInstance
+Instance = MlxRingInstance | MlxJacclInstance | VllmInstance
 
 
 class BoundInstance(CamelCaseModel):

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/info_gatherer.py

@@ -32,6 +32,7 @@ from exo.utils.pydantic_ext import TaggedModel
 from exo.utils.task_group import TaskGroup
 
 from .macmon import MacmonMetrics
+from .nvml import NvmlMetrics
 from .system_info import (
     get_friendly_name,
     get_model_and_chip,
@@ -43,6 +44,12 @@ from .system_info import (
 IS_DARWIN = sys.platform == "darwin"
 
 
+def _has_nvml() -> bool:
+    from exo.utils.info_gatherer.nvml import has_nvml
+
+    return has_nvml()
+
+
 async def _get_thunderbolt_devices() -> set[str] | None:
     """Get Thunderbolt interface device names (e.g., en2, en3) from hardware ports.
 
@@ -354,6 +361,21 @@ async def _gather_iface_map() -> dict[str, str] | None:
     return ports
 
 
+class VllmCapability(TaggedModel):
+    available: bool
+    version: str | None = None
+
+    @classmethod
+    async def gather(cls) -> Self:
+        try:
+            import importlib
+
+            vllm = importlib.import_module("vllm")
+            return cls(available=True, version=getattr(vllm, "__version__", None))
+        except ImportError:
+            return cls(available=False)
+
+
 GatheredInfo = (
     MacmonMetrics
     | MemoryUsage
@@ -362,6 +384,8 @@ GatheredInfo = (
     | MacThunderboltConnections
     | RdmaCtlStatus
     | ThunderboltBridgeInfo
+    | VllmCapability
+    | NvmlMetrics
     | NodeConfig
     | MiscData
     | StaticNodeInformation
@@ -381,6 +405,7 @@ class InfoGatherer:
     static_info_poll_interval: float | None = 60
     rdma_ctl_poll_interval: float | None = 10 if IS_DARWIN else None
     disk_poll_interval: float | None = 30
+    vllm_capability_poll_interval: float | None = 60
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
 
     async def run(self):
@@ -389,7 +414,6 @@ class InfoGatherer:
                 if (macmon_path := shutil.which("macmon")) is not None:
                     tg.start_soon(self._monitor_macmon, macmon_path)
                 else:
-                    # macmon not installed — fall back to psutil for memory
                     logger.warning(
                         "macmon not found, falling back to psutil for memory monitoring"
                     )
@@ -397,11 +421,14 @@ class InfoGatherer:
                 tg.start_soon(self._monitor_system_profiler_thunderbolt_data)
                 tg.start_soon(self._monitor_thunderbolt_bridge_status)
                 tg.start_soon(self._monitor_rdma_ctl_status)
+            elif _has_nvml():
+                tg.start_soon(self._monitor_nvml_metrics)
             tg.start_soon(self._watch_system_info)
             tg.start_soon(self._monitor_memory_usage)
             tg.start_soon(self._monitor_misc)
             tg.start_soon(self._monitor_static_info)
             tg.start_soon(self._monitor_disk_usage)
+            tg.start_soon(self._monitor_vllm_capability)
 
             nc = await NodeConfig.gather()
             if nc is not None:
@@ -525,6 +552,26 @@ class InfoGatherer:
                 logger.warning(f"Error gathering disk usage: {e}")
             await anyio.sleep(self.disk_poll_interval)
 
+    async def _monitor_nvml_metrics(self):
+        while True:
+            try:
+                from exo.utils.info_gatherer.nvml import gather_nvidia_metrics
+
+                metrics = gather_nvidia_metrics()
+                if metrics is not None:
+                    await self.info_sender.send(metrics)
+            except Exception as e:
+                logger.warning(f"Error gathering NVML metrics: {e}")
+            await anyio.sleep(1)
+
+    async def _monitor_vllm_capability(self):
+        if self.vllm_capability_poll_interval is None:
+            return
+        try:
+            await self.info_sender.send(await VllmCapability.gather())
+        except Exception as e:
+            logger.warning(f"Error gathering vLLM capability: {e}")
+
     async def _monitor_macmon(self, macmon_path: str):
         if self.macmon_interval is None:
             return

src/exo/utils/info_gatherer/nvml.py

@@ -0,0 +1,85 @@
+# pyright: reportMissingImports=false
+from exo.shared.types.profiling import SystemPerformanceProfile
+from exo.utils.pydantic_ext import TaggedModel
+
+try:
+    from pynvml import (
+        nvmlDeviceGetCount,
+        nvmlDeviceGetHandleByIndex,
+        nvmlDeviceGetPowerUsage,
+        nvmlDeviceGetTemperature,
+        nvmlDeviceGetUtilizationRates,
+        nvmlInit,
+        nvmlShutdown,
+    )
+except ImportError:
+    nvmlDeviceGetCount = None  # noqa: N816
+    nvmlDeviceGetHandleByIndex = None  # noqa: N816
+    nvmlDeviceGetPowerUsage = None  # noqa: N816
+    nvmlDeviceGetTemperature = None  # noqa: N816
+    nvmlDeviceGetUtilizationRates = None  # noqa: N816
+    nvmlInit = None  # noqa: N816
+    nvmlShutdown = None  # noqa: N816
+
+_CPU_POWER_IDLE = 20.0
+_CPU_POWER_MAX = 100.0
+_GPU_POWER_MAX = 120.0
+
+
+class NvmlMetrics(TaggedModel):
+    system_profile: SystemPerformanceProfile
+
+
+def has_nvml() -> bool:
+    if nvmlInit is None:
+        return False
+    try:
+        nvmlInit()
+        count = nvmlDeviceGetCount()  # type: ignore[reportOptionalCall]
+        nvmlShutdown()  # type: ignore[reportOptionalCall]
+        return count > 0
+    except Exception:
+        return False
+
+
+def gather_nvidia_metrics() -> NvmlMetrics | None:
+    if nvmlInit is None or nvmlDeviceGetCount is None or nvmlShutdown is None:
+        return None
+    if nvmlDeviceGetHandleByIndex is None or nvmlDeviceGetUtilizationRates is None:
+        return None
+    if nvmlDeviceGetTemperature is None or nvmlDeviceGetPowerUsage is None:
+        return None
+
+    try:
+        nvmlInit()
+        count = nvmlDeviceGetCount()
+        if count == 0:
+            nvmlShutdown()
+            return None
+
+        total_gpu_util = 0.0
+        total_temp = 0.0
+        total_gpu_power = 0.0
+        for i in range(count):
+            handle = nvmlDeviceGetHandleByIndex(i)
+            util = nvmlDeviceGetUtilizationRates(handle)
+            total_gpu_util += float(util.gpu)
+            total_temp += float(nvmlDeviceGetTemperature(handle, 0))
+            total_gpu_power += float(nvmlDeviceGetPowerUsage(handle)) / 1000.0
+
+        nvmlShutdown()
+
+        gpu_load_fraction = min(total_gpu_power / _GPU_POWER_MAX, 1.0)
+        estimated_cpu_power = (
+            _CPU_POWER_IDLE + (_CPU_POWER_MAX - _CPU_POWER_IDLE) * gpu_load_fraction
+        )
+
+        return NvmlMetrics(
+            system_profile=SystemPerformanceProfile(
+                gpu_usage=total_gpu_util / count / 100.0,
+                temp=total_temp / count,
+                sys_power=total_gpu_power + estimated_cpu_power,
+            ),
+        )
+    except Exception:
+        return None

src/exo/utils/info_gatherer/system_info.py

@@ -1,6 +1,7 @@
 import platform
 import socket
 import sys
+from pathlib import Path
 from subprocess import CalledProcessError
 
 import psutil
@@ -90,39 +91,175 @@ 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:
+    """Read a single DMI sysfs field, returning ``None`` on failure."""
+    try:
+        path = Path(f"/sys/class/dmi/id/{name}")
+        if path.exists():
+            return path.read_text().strip()
+    except (OSError, PermissionError):
+        pass
+    return None
+
+
+async def _get_linux_model_and_chip() -> tuple[str, str]:
+    """Get Linux system information using DMI and /proc/cpuinfo.
+
+    Detects NVIDIA DGX Spark (DMI product_name ``"DGX_Spark"``) and other
+    NVIDIA systems via the ``sys_vendor`` DMI field, falling back to generic
+    Linux identification.
+    """
+    model = "Linux"
+    chip = "Unknown Chip"
+
+    product_name = _read_dmi_field("product_name")
+    sys_vendor = _read_dmi_field("sys_vendor")
+
+    # DGX Spark: DMI product_name may be "DGX_Spark" or "gx10" variant
+    product_lower = (product_name or "").lower()
+    if product_name and ("dgx" in product_lower or "gx10" in product_lower):
+        model = "DGX Spark"
+        try:
+            process = await run_process(
+                ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"]
+            )
+            gpu_name = process.stdout.decode().strip().split("\n")[0]
+            chip = gpu_name if gpu_name and gpu_name != "[N/A]" else "NVIDIA GB10"
+        except (CalledProcessError, FileNotFoundError):
+            chip = "NVIDIA GB10"
+        return (model, chip)
+
+    # Other NVIDIA systems (sys_vendor contains "NVIDIA")
+    if sys_vendor and "NVIDIA" in sys_vendor:
+        model = product_name.replace("_", " ") if product_name else "NVIDIA System"
+        try:
+            process = await run_process(
+                ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"]
+            )
+            gpu_name = process.stdout.decode().strip().split("\n")[0]
+            if gpu_name and gpu_name != "[N/A]":
+                chip = gpu_name
+        except (CalledProcessError, FileNotFoundError):
+            pass
+        return (model, chip)
+
+    # Generic Linux — detect laptop vs desktop via chassis_type
+    # SMBIOS chassis types: 8,9,10,14,31,32 = portable/laptop
+    chassis_type = _read_dmi_field("chassis_type")
+    laptop_chassis_types = {"8", "9", "10", "14", "31", "32"}
+    if chassis_type in laptop_chassis_types:
+        model = "Linux Laptop"
+    elif chassis_type is not None:
+        model = "Linux Desktop"
+
+    # Also check for battery as a fallback laptop indicator
+    if model == "Linux" and Path("/sys/class/power_supply/BAT0").exists():
+        model = "Linux Laptop"
+
+    # Use /proc/cpuinfo for chip
+    cpuinfo_path = Path("/proc/cpuinfo")
+    if cpuinfo_path.exists():
+        try:
+            for line in cpuinfo_path.read_text().splitlines():
+                if line.startswith("model name"):
+                    chip = line.split(":", 1)[1].strip()
+                    break
+        except OSError:
+            pass
+
+    return (model, chip)
 
 
 async def get_model_and_chip() -> tuple[str, str]:
-    """Get Mac system information using system_profiler."""
+    """Get system model and chip information.
+
+    On macOS, uses ``system_profiler``.  On Linux, reads DMI data from
+    sysfs and CPU info from ``/proc/cpuinfo``.
+    """
     model = "Unknown Model"
     chip = "Unknown Chip"
 
-    # TODO: better non mac support
+    if sys.platform == "linux":
+        return await _get_linux_model_and_chip()
+
     if sys.platform != "darwin":
         return (model, chip)
 

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

@@ -1,5 +1,6 @@
 import os
 from copy import deepcopy
+from typing import TYPE_CHECKING, cast
 
 import mlx.core as mx
 import psutil
@@ -13,10 +14,13 @@ from mlx_lm.models.cache import (
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType, Model
+from exo.shared.types.mlx import KVCacheType, MLXCacheType, Model
 from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
 from exo.worker.runner.bootstrap import logger
 
+if TYPE_CHECKING:
+    from exo.worker.engines.vllm.kv_cache import TorchKVCache
+
 
 # Fraction of device memory above which LRU eviction kicks in.
 # Smaller machines need more aggressive eviction.
@@ -46,7 +50,7 @@ class CacheSnapshot:
         self.token_count = token_count
 
 
-def snapshot_ssm_states(cache: KVCacheType) -> CacheSnapshot:
+def snapshot_ssm_states(cache: MLXCacheType) -> CacheSnapshot:
     states: list[ArraysCache | RotatingKVCache | None] = []
     for c in cache:
         if isinstance(c, (ArraysCache, RotatingKVCache)):
@@ -70,7 +74,7 @@ def _find_nearest_snapshot(
     return best
 
 
-def has_non_kv_caches(cache: KVCacheType) -> bool:
+def has_non_kv_caches(cache: MLXCacheType) -> bool:
     """Check if a cache contains any ArraysCache (SSM) entries."""
     return any(isinstance(c, (ArraysCache, RotatingKVCache)) for c in cache)
 
@@ -94,7 +98,7 @@ class KVPrefixCache:
     def add_kv_cache(
         self,
         prompt_tokens: mx.array,
-        cache: KVCacheType,
+        cache: MLXCacheType,
         ssm_snapshots: list[CacheSnapshot] | None = None,
     ):
         """Add a new cache entry. Evicts LRU entries if memory is high."""
@@ -110,7 +114,7 @@ class KVPrefixCache:
         self,
         index: int,
         prompt_tokens: mx.array,
-        cache: KVCacheType,
+        cache: MLXCacheType,
         snapshots: list[CacheSnapshot] | None,
         restore_pos: int,
     ):
@@ -129,10 +133,14 @@ class KVPrefixCache:
         self._last_used[index] = self._access_counter
         logger.info(f"KV cache updated (index {index}): {len(prompt_tokens)} tokens")
 
+    def _get_mlx_cache(self, index: int) -> MLXCacheType:
+        cached = self.caches[index]
+        return cast(MLXCacheType, cached)
+
     def _get_snapshot(
         self, entry_index: int, target_token_count: int
     ) -> tuple[int, CacheSnapshot | None]:
-        if not has_non_kv_caches(self.caches[entry_index]):
+        if not has_non_kv_caches(self._get_mlx_cache(entry_index)):
             return target_token_count, None
 
         snapshots = self._snapshots[entry_index]
@@ -149,7 +157,7 @@ class KVPrefixCache:
         self,
         model: Model,
         prompt_tokens: mx.array,
-    ) -> tuple[KVCacheType, mx.array, int | None]:
+    ) -> tuple[MLXCacheType, mx.array, int | None]:
         """Get KV cache for prompt, returning remaining tokens to prefill.
 
         Returns:
@@ -184,16 +192,25 @@ class KVPrefixCache:
         # For exact match: trim to max_length-1 so remaining has the last token
         # For partial match: trim to best_length, remaining has suffix to prefill
         # This ensures stream_generate always has at least one token to start with
-        has_ssm = has_non_kv_caches(self.caches[best_index])
+        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
 
-        prompt_cache = deepcopy(self.caches[best_index])
-        cached_length = cache_length(self.caches[best_index])
+        prompt_cache = deepcopy(mlx_cache)
+        cached_length = cache_length(mlx_cache)
         tokens_to_trim = cached_length - restore_pos
         if tokens_to_trim > 0:
             trim_cache(prompt_cache, tokens_to_trim, restore_snap)
@@ -208,6 +225,50 @@ class KVPrefixCache:
 
         return prompt_cache, remaining, best_index
 
+    def lookup(
+        self, prompt_token_ids: list[int]
+    ) -> tuple["TorchKVCache | None", int, int | None]:
+        from exo.worker.engines.vllm.kv_cache import TorchKVCache
+
+        prompt_mx = mx.array(prompt_token_ids)
+        max_length = len(prompt_token_ids)
+        best_index: int | None = None
+        best_length = 0
+
+        for i, cached_prompt in enumerate(self.prompts):
+            length = get_prefix_length(prompt_mx, cached_prompt)
+            if length >= max_length - 1:
+                best_index, best_length = i, length
+                break
+            if length > best_length:
+                best_index, best_length = i, length
+
+        if best_index is None or best_length == 0:
+            return None, 0, None
+
+        best_length = min(best_length, max_length - 1)
+
+        self._access_counter += 1
+        self._last_used[best_index] = self._access_counter
+
+        cached = self.caches[best_index]
+        if isinstance(cached, TorchKVCache):
+            return cached.trim_to(best_length), best_length, best_index
+
+        torch_cache = TorchKVCache.from_mlx_cache(cached)
+        return torch_cache.trim_to(best_length), best_length, best_index
+
+    def add_from_torch(
+        self, prompt_token_ids: list[int], cache: "TorchKVCache"
+    ) -> None:
+        self._evict_if_needed()
+        self.prompts.append(mx.array(prompt_token_ids))
+        self.caches.append(cache.detach_cpu())
+        self._snapshots.append(None)
+        self._access_counter += 1
+        self._last_used.append(self._access_counter)
+        logger.info(f"KV cache added (torch): {len(prompt_token_ids)} tokens")
+
     def _evict_if_needed(self):
         """Evict least recently used entries while memory usage is high."""
         if len(self.caches) == 0:
@@ -244,7 +305,7 @@ class KVPrefixCache:
 
 
 def trim_cache(
-    cache: KVCacheType,
+    cache: MLXCacheType,
     num_tokens: int,
     snapshot: CacheSnapshot | None = None,
 ) -> None:
@@ -282,7 +343,7 @@ def _entry_length(
     return 0
 
 
-def cache_length(cache: KVCacheType) -> int:
+def cache_length(cache: MLXCacheType) -> int:
     """Get the number of tokens in a KV cache."""
     return max(_entry_length(c) for c in cache)
 
@@ -311,7 +372,7 @@ def get_memory_used_percentage() -> float:
 
 def make_kv_cache(
     model: Model, max_kv_size: int | None = None, keep: int = 0
-) -> KVCacheType:
+) -> MLXCacheType:
     assert hasattr(model, "layers")
 
     if hasattr(model, "make_cache"):

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
 
@@ -18,13 +18,16 @@ from exo.shared.types.api import (
     TopLogprobItem,
     Usage,
 )
+from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType, Model
+from exo.shared.types.mlx import MLXCacheType, Model
+from exo.shared.types.tasks import TaskId
 from exo.shared.types.text_generation import TextGenerationTaskParams
 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,
 )
@@ -34,6 +37,7 @@ from exo.worker.engines.mlx.generator.generate import (
     eos_ids_from_tokenizer,
     extract_top_logprobs,
     prefill,
+    warmup_inference,
 )
 from exo.worker.engines.mlx.utils_mlx import fix_unmatched_think_end_tokens
 from exo.worker.runner.bootstrap import logger
@@ -74,33 +78,39 @@ class ExoBatchGenerator:
     tokenizer: TokenizerWrapper
     group: mx.distributed.Group | None
     kv_prefix_cache: KVPrefixCache | None
+    model_id: ModelId
 
-    _exo_gen: MlxBatchGenerator = field(init=False)
+    _mlx_gen: MlxBatchGenerator = field(init=False)
     _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False)
+    _uid_to_task_id: dict[int, TaskId] = field(default_factory=dict, init=False)
 
     def __post_init__(self) -> None:
-        self._exo_gen = MlxBatchGenerator(
+        self._mlx_gen = MlxBatchGenerator(
             model=self.model,
             stop_tokens=set(eos_ids_from_tokenizer(self.tokenizer)),
             prefill_step_size=4096,
         )
 
+    def warmup(self) -> int:
+        return warmup_inference(self.model, self.tokenizer, self.group, self.model_id)
+
     @property
     def has_work(self) -> bool:
         return (
             bool(self._active_tasks)
-            or bool(self._exo_gen.unprocessed_prompts)
-            or self._exo_gen.active_batch is not None
+            or bool(self._mlx_gen.unprocessed_prompts)
+            or self._mlx_gen.active_batch is not None
         )
 
     def submit(
         self,
+        task_id: TaskId,
         task_params: TextGenerationTaskParams,
         prompt: str,
         on_prefill_progress: Callable[[int, int], None] | None = None,
         distributed_prompt_progress_callback: Callable[[], None] | None = None,
         on_generation_token: Callable[[], None] | None = None,
-    ) -> int:
+    ) -> TaskId:
         all_prompt_tokens = encode_prompt(self.tokenizer, prompt)
         all_prompt_tokens = fix_unmatched_think_end_tokens(
             all_prompt_tokens, self.tokenizer
@@ -140,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:
@@ -173,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(
@@ -188,7 +231,7 @@ class ExoBatchGenerator:
 
         max_tokens = task_params.max_output_tokens or MAX_TOKENS
 
-        uids = self._exo_gen.insert(
+        uids = self._mlx_gen.insert(
             prompts=[last_tokens.tolist()],
             max_tokens=[max_tokens],
             caches=[list(cache)],
@@ -199,6 +242,7 @@ class ExoBatchGenerator:
         assert len(uids) == 1
 
         uid = uids[0]
+        self._uid_to_task_id[uid] = task_id
 
         self._active_tasks[uid] = _EngineTask(
             uid=uid,
@@ -213,15 +257,15 @@ class ExoBatchGenerator:
             prefill_tps=_prefill_tps,
         )
 
-        return uid
+        return task_id
 
-    def step(self) -> list[tuple[int, GenerationResponse]]:
+    def step(self) -> list[tuple[TaskId, GenerationResponse]]:
         if not self.has_work:
             return []
 
-        responses = self._exo_gen.next()
+        responses = self._mlx_gen.next()
 
-        results: list[tuple[int, GenerationResponse]] = []
+        results: list[tuple[TaskId, GenerationResponse]] = []
 
         for response in responses:
             if response.uid not in self._active_tasks:
@@ -288,7 +332,7 @@ class ExoBatchGenerator:
             usage: Usage | None = None
             if is_done:
                 try:
-                    mlx_stats = self._exo_gen.stats()
+                    mlx_stats = self._mlx_gen.stats()
                     generation_tps = mlx_stats.generation_tps
                 except ZeroDivisionError:
                     generation_elapsed = (
@@ -322,7 +366,7 @@ class ExoBatchGenerator:
 
             results.append(
                 (
-                    response.uid,
+                    self._uid_to_task_id.get(response.uid, TaskId(str(response.uid))),
                     GenerationResponse(
                         text=text,
                         token=response.token,
@@ -337,6 +381,7 @@ class ExoBatchGenerator:
 
             if is_done:
                 del self._active_tasks[response.uid]
+                self._uid_to_task_id.pop(response.uid, None)
             elif (
                 max_stop_len > 0
                 and len(state.potential_stop_sequence_text) > max_stop_len
@@ -347,18 +392,23 @@ class ExoBatchGenerator:
 
         return results
 
-    def cancel(self, uids: list[int]) -> None:
-        self._exo_gen.remove(uids)
+    def cancel(self, task_ids: list[TaskId]) -> None:
+        task_id_set = set(task_ids)
+        uids = [uid for uid, tid in self._uid_to_task_id.items() if tid in task_id_set]
+        if uids:
+            self._mlx_gen.remove(uids)
         for uid in uids:
             self._active_tasks.pop(uid, None)
+            self._uid_to_task_id.pop(uid, None)
 
     def close(self) -> None:
-        self._exo_gen.close()
+        self._mlx_gen.close()
+        mx.clear_cache()
 
     def _save_prefix_cache(
         self,
         all_prompt_tokens: mx.array,
-        cache: KVCacheType,
+        cache: MLXCacheType,
         cache_snapshots: list[CacheSnapshot] | None,
         prefix_hit_length: int,
         matched_index: int | None,

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

@@ -23,7 +23,7 @@ from exo.shared.types.api import (
 )
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType, Model
+from exo.shared.types.mlx import MLXCacheType, Model
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
@@ -76,7 +76,7 @@ def _has_pipeline_communication_layer(model: Model):
 def pipeline_parallel_prefill(
     model: Model,
     prompt: mx.array,
-    prompt_cache: KVCacheType,
+    prompt_cache: MLXCacheType,
     prefill_step_size: int,
     kv_group_size: int | None,
     kv_bits: int | None,
@@ -113,7 +113,7 @@ def pipeline_parallel_prefill(
         kv_bits=kv_bits,
     )
 
-    _prompt_cache: KVCacheType = prompt_cache
+    _prompt_cache: MLXCacheType = prompt_cache
     rank = group.rank()
     world_size = group.size()
 
@@ -195,7 +195,7 @@ def prefill(
     tokenizer: TokenizerWrapper,
     sampler: Callable[[mx.array], mx.array],
     prompt_tokens: mx.array,
-    cache: KVCacheType,
+    cache: MLXCacheType,
     group: mx.distributed.Group | None,
     on_prefill_progress: Callable[[int, int], None] | None,
     distributed_prompt_progress_callback: Callable[[], None] | None,

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

@@ -145,6 +145,11 @@ def mlx_distributed_init(
                 os.environ["MLX_JACCL_COORDINATOR"] = jaccl_coordinator
                 group = mx.distributed.init(backend="jaccl", strict=True)
 
+            case _:
+                raise ValueError(
+                    f"Unsupported instance type for MLX init: {type(bound_instance.instance)}"
+                )
+
         logger.info(f"Rank {rank} mlx distributed initialization complete")
 
         return group
@@ -546,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] = {}
@@ -583,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

@@ -0,0 +1,645 @@
+import os
+import sys
+
+import torch
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+from exo.shared.logging import logger
+from exo.worker.engines.mlx.cache import KVPrefixCache
+
+INITIAL_FRACTION = 0.05
+GROWTH_HEADROOM_BYTES = 512 * 1024 * 1024
+MIN_GROWTH_BLOCKS = 16
+
+_patched = False
+_prefix_cache: KVPrefixCache | None = None
+_model_runner: GPUModelRunner | None = None
+
+
+def get_prefix_cache() -> KVPrefixCache | None:
+    return _prefix_cache
+
+
+def set_prefix_cache(cache: KVPrefixCache | None) -> None:
+    global _prefix_cache
+    _prefix_cache = cache
+
+
+def get_model_runner() -> GPUModelRunner | None:
+    return _model_runner
+
+
+def set_model_runner(runner: GPUModelRunner | None) -> None:
+    global _model_runner
+    _model_runner = runner
+
+
+def patch_vllm() -> None:
+    global _patched
+    if _patched:
+        return
+    _patched = True
+
+    _patch_determine_available_memory()
+    _patch_check_enough_kv_cache_memory()
+    _patch_initialize_kv_cache_tensors()
+    _patch_initialize_from_config()
+    _patch_kv_cache_manager_init()
+    _patch_allocate_slots()
+    _patch_get_computed_blocks()
+    _patch_moe_sum()
+    _patch_marlin_w2_thread_config()
+    _patch_flashinfer_sinks()
+    logger.info("vLLM growable KV cache patch applied")
+
+
+def _patch_determine_available_memory() -> None:
+    from vllm.v1.worker.gpu_worker import Worker
+
+    original = Worker.determine_available_memory
+
+    @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, 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)"
+        )
+        return initial
+
+    Worker.determine_available_memory = patched  # type: ignore
+
+
+def _patch_check_enough_kv_cache_memory() -> None:
+    from vllm.v1.core import kv_cache_utils
+
+    def noop(*_args: "object", **_kwargs: "object") -> None:
+        pass
+
+    kv_cache_utils._check_enough_kv_cache_memory = noop  # type: ignore
+
+
+def _patch_initialize_kv_cache_tensors() -> None:
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+    original_alloc = GPUModelRunner._allocate_kv_cache_tensors
+
+    def patched_alloc(
+        self: "GPUModelRunner", kv_cache_config: "object"
+    ) -> "dict[str, torch.Tensor]":
+        raw_tensors = original_alloc(self, kv_cache_config)
+        self._growable_raw_tensors = {name: t for name, t in raw_tensors.items()}
+        return raw_tensors
+
+    GPUModelRunner._allocate_kv_cache_tensors = patched_alloc  # type: ignore
+
+    original_init_tensors = GPUModelRunner.initialize_kv_cache_tensors
+
+    def patched_init_tensors(
+        self: "GPUModelRunner",
+        kv_cache_config: "object",
+        kernel_block_sizes: "list[int]",
+    ) -> "dict[str, torch.Tensor]":
+        self._growable_kv_cache_config = kv_cache_config
+        self._growable_kernel_block_sizes = kernel_block_sizes
+        return original_init_tensors(self, kv_cache_config, kernel_block_sizes)
+
+    GPUModelRunner.initialize_kv_cache_tensors = patched_init_tensors  # type: ignore
+
+
+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)
+        set_model_runner(self.model_runner)
+
+    Worker.initialize_from_config = patched  # type: ignore
+
+
+def _patch_kv_cache_manager_init() -> None:
+    from vllm.v1.core.kv_cache_manager import KVCacheManager
+
+    original_init = KVCacheManager.__init__
+
+    def patched_init(
+        self: "KVCacheManager", *args: "object", **kwargs: "object"
+    ) -> None:
+        original_init(self, *args, **kwargs)
+        self._growable_model_runner = get_model_runner()
+
+    KVCacheManager.__init__ = patched_init  # type: ignore
+
+
+def _patch_allocate_slots() -> None:
+    from vllm.v1.core.kv_cache_manager import KVCacheManager
+
+    original = KVCacheManager.allocate_slots
+
+    def patched(
+        self: "KVCacheManager",
+        request: "object",
+        num_new_tokens: int,
+        *args: "object",
+        **kwargs: "object",
+    ) -> "object":
+        result = original(self, request, num_new_tokens, *args, **kwargs)
+        if result is None and _try_grow_cache(self):
+            result = original(self, request, num_new_tokens, *args, **kwargs)
+        return result
+
+    KVCacheManager.allocate_slots = patched  # type: ignore
+
+
+def _try_grow_cache(kv_cache_manager: "object") -> bool:
+    block_pool = kv_cache_manager.block_pool  # type: ignore
+    model_runner = kv_cache_manager._growable_model_runner  # type: ignore
+
+    if model_runner is None:
+        return False
+
+    free_bytes, _ = torch.cuda.mem_get_info()
+    if free_bytes < GROWTH_HEADROOM_BYTES:
+        return False
+
+    kv_cache_config = model_runner._growable_kv_cache_config  # type: ignore
+    old_num_blocks: int = kv_cache_config.num_blocks
+
+    total_tensor_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors)
+    per_block_bytes = total_tensor_bytes // old_num_blocks
+
+    usable_bytes = int(free_bytes * 0.8)
+    growth_blocks = min(usable_bytes // per_block_bytes, old_num_blocks)
+
+    if growth_blocks < MIN_GROWTH_BLOCKS:
+        return False
+
+    new_num_blocks = old_num_blocks + growth_blocks
+
+    logger.info(
+        f"Growing KV cache: {old_num_blocks} → {new_num_blocks} blocks "
+        f"(+{growth_blocks * per_block_bytes / (1024**3):.2f} GiB)"
+    )
+
+    try:
+        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:
+        logger.opt(exception=True).error("Failed to grow KV cache")
+        return False
+
+
+
+
+def _grow_tensors(
+    model_runner: "object",
+    kv_cache_config: "object",
+    old_num_blocks: int,
+    new_num_blocks: int,
+) -> None:
+    raw_tensors: dict[str, torch.Tensor] = model_runner._growable_raw_tensors  # type: ignore
+    ratio = new_num_blocks / old_num_blocks
+
+    already_grown: dict[int, torch.Tensor] = {}
+    new_raw_tensors: dict[str, torch.Tensor] = {}
+
+    for layer_name, old_raw in raw_tensors.items():
+        storage_id = old_raw.data_ptr()
+        if storage_id in already_grown:
+            new_raw_tensors[layer_name] = already_grown[storage_id]
+            continue
+
+        old_size = old_raw.numel()
+        new_size = int(old_size * ratio)
+        new_raw = torch.zeros(new_size, dtype=torch.int8, device=old_raw.device)
+        new_raw[:old_size] = old_raw
+        already_grown[storage_id] = new_raw
+        new_raw_tensors[layer_name] = new_raw
+
+    model_runner._growable_raw_tensors = new_raw_tensors  # type: ignore
+
+    kernel_block_sizes: list[int] = model_runner._growable_kernel_block_sizes  # type: ignore
+    new_kv_caches: dict[str, torch.Tensor] = model_runner._reshape_kv_cache_tensors(  # type: ignore
+        kv_cache_config,
+        new_raw_tensors,
+        kernel_block_sizes,
+    )
+
+    forward_context: dict[str, "object"] = (
+        model_runner.compilation_config.static_forward_context
+    )  # type: ignore
+    runner_kv_caches: list[torch.Tensor] = model_runner.kv_caches  # type: ignore
+
+    from collections import defaultdict
+
+    from vllm.v1.worker.utils import extract_layer_index
+
+    num_attn_module = 1
+    hf_config = getattr(getattr(model_runner, "model_config", None), "hf_config", None)  # type: ignore
+    if getattr(hf_config, "model_type", "") == "longcat_flash":
+        num_attn_module = 2
+
+    index2name: dict[int, list[str]] = defaultdict(list)
+    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]:
+            new_ordered.append(new_kv_caches[ln])
+
+    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(
+    block_pool: "object", old_num_blocks: int, new_num_blocks: int
+) -> None:
+    from vllm.v1.core.kv_cache_utils import KVCacheBlock
+
+    new_blocks: list["KVCacheBlock"] = []
+    for idx in range(old_num_blocks, new_num_blocks):
+        block = KVCacheBlock(idx)
+        block_pool.blocks.append(block)  # type: ignore
+        new_blocks.append(block)
+
+    block_pool.free_block_queue.append_n(new_blocks)  # type: ignore
+    block_pool.num_gpu_blocks = new_num_blocks  # type: ignore
+
+
+def _patch_moe_sum() -> None:
+    import vllm._custom_ops as ops  # type: ignore[reportMissingImports]
+
+    def moe_sum_f32(x: "torch.Tensor", output: "torch.Tensor") -> None:
+        output[:] = x.to(torch.float32).sum(dim=1).to(output.dtype)  # type: ignore
+
+    ops.moe_sum = moe_sum_f32  # type: ignore
+
+
+def _patch_marlin_w2_thread_config() -> None:
+    try:
+        import vllm._custom_ops as ops  # type: ignore[reportMissingImports]
+    except ImportError:
+        return
+
+    original_gemm = ops.moe_wna16_marlin_gemm
+
+    def patched_gemm(*args: "object", **kwargs: "object") -> "object":
+        kwargs["thread_k"] = 64
+        kwargs["thread_n"] = 128
+        return original_gemm(*args, **kwargs)
+
+    ops.moe_wna16_marlin_gemm = patched_gemm  # type: ignore
+
+
+def _patch_get_computed_blocks() -> None:
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
+    from vllm.v1.core.kv_cache_utils import KVCacheBlock
+    from vllm.v1.request import Request
+
+    original = KVCacheManager.get_computed_blocks
+
+    def patched(
+        self: KVCacheManager,
+        request: Request,
+    ) -> tuple[KVCacheBlocks, int]:
+        prefix_cache = get_prefix_cache()
+        if prefix_cache is None or request.prompt_token_ids is None:
+            return original(self, request)
+
+        from exo.worker.engines.vllm.kv_cache import (
+            TorchKVCache as _TorchKVCache,  # noqa: F811
+        )
+
+        try:
+            torch_cache, num_matched, _ = prefix_cache.lookup(
+                list(request.prompt_token_ids)
+            )  # type: ignore[reportUnknownMemberType]
+        except Exception:
+            return original(self, request)
+
+        if (
+            torch_cache is None
+            or not isinstance(torch_cache, _TorchKVCache)
+            or num_matched == 0
+        ):
+            return original(self, request)
+
+        from vllm.utils.math_utils import cdiv  # type: ignore[reportMissingImports]
+
+        from exo.worker.engines.vllm.vllm_generator import _build_layer_groups
+
+        num_groups = len(self.kv_cache_config.kv_cache_groups)
+        null_block = self.block_pool.null_block
+        save_offsets = torch_cache.token_offset_per_group or [0] * num_groups
+
+        for gi in range(num_groups):
+            save_off = save_offsets[gi] if gi < len(save_offsets) else 0
+            if save_off > 0:
+                spec = self.kv_cache_config.kv_cache_groups[gi].kv_cache_spec  # type: ignore
+                window = getattr(spec, "sliding_window", 0) or 0
+                if window > 0 and num_matched < save_off + window:
+                    return original(self, request)
+
+        real_block_counts: list[int] = []
+        skipped_block_counts: list[int] = []
+        total_needed = 0
+        for gi in range(num_groups):
+            mgr = self.coordinator.single_type_managers[gi]  # type: ignore
+            block_size: int = self.kv_cache_config.kv_cache_groups[
+                gi
+            ].kv_cache_spec.block_size  # type: ignore
+            num_skipped: int = mgr.get_num_skipped_tokens(num_matched)  # type: ignore
+            num_skipped_blocks = num_skipped // block_size
+            num_real = cdiv(num_matched, block_size) - num_skipped_blocks
+            real_block_counts.append(num_real)
+            skipped_block_counts.append(num_skipped_blocks)
+            total_needed += num_real
+
+        if self.block_pool.get_num_free_blocks() < total_needed:
+            return original(self, request)
+
+        blocks_per_group: list[list[KVCacheBlock]] = []
+        token_offset_per_group: list[int] = []
+        for gi in range(num_groups):
+            mgr = self.coordinator.single_type_managers[gi]  # type: ignore
+            block_size = self.kv_cache_config.kv_cache_groups[
+                gi
+            ].kv_cache_spec.block_size  # type: ignore
+            real_blocks: list[KVCacheBlock] = self.block_pool.get_new_blocks(
+                real_block_counts[gi]
+            )  # type: ignore
+            blocks_per_group.append(real_blocks)
+
+            full_block_list = [null_block] * skipped_block_counts[gi] + list(
+                real_blocks
+            )
+            req_blocks = mgr.req_to_blocks[request.request_id]  # type: ignore
+            req_blocks.extend(full_block_list)  # type: ignore
+
+            token_offset_per_group.append(skipped_block_counts[gi] * block_size)
+
+        block_ids_per_group = [[b.block_id for b in grp] for grp in blocks_per_group]
+        layer_to_group = _build_layer_groups(self.kv_cache_config)
+        model_runner = self._growable_model_runner  # type: ignore[reportAttributeAccessIssue]
+        if model_runner is not None:
+            torch_cache.write_to_vllm_blocks(  # type: ignore
+                model_runner.kv_caches,
+                block_ids_per_group,
+                layer_to_group,  # type: ignore
+                token_offset_per_group,
+            )
+
+        total_blocks = sum(len(g) for g in blocks_per_group)
+        logger.info(
+            f"Prefix cache hit: {num_matched} tokens, {total_blocks} blocks ({num_groups} groups)"
+        )
+        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}")