more fixes

.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: ...

bench/eval_tool_calls.py

@@ -18,11 +18,13 @@ from harness import (
     ExoHttpError,
     add_common_instance_args,
     capture_cluster_snapshot,
+    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,
 )
@@ -934,6 +936,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:
@@ -953,6 +956,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

@@ -35,12 +35,14 @@ from harness import (
     ExoHttpError,
     add_common_instance_args,
     capture_cluster_snapshot,
+    ensure_cuda_available,
     instance_id_from_instance,
     node_ids_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,
 )
@@ -380,6 +382,7 @@ def main() -> int:
         help="System metrics polling interval in seconds (default: 1.0).",
     )
     args = ap.parse_args()
+    validate_vllm_args(args)
 
     pp_list = parse_int_list(args.pp)
     tg_list = parse_int_list(args.tg)
@@ -404,6 +407,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

@@ -47,11 +47,13 @@ from harness import (
     ExoHttpError,
     add_common_instance_args,
     capture_cluster_snapshot,
+    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,
 )
@@ -1160,6 +1162,7 @@ def main() -> int:
     )
 
     args, _ = ap.parse_known_args()
+    validate_vllm_args(args)
 
     # Resolve tasks
     if args.tasks:
@@ -1176,6 +1179,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 get_state_path(self, path: str) -> Any:
         try:
             return self.request_json("GET", f"/state/{path}")
@@ -253,6 +257,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":
@@ -543,7 +572,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"
@@ -572,3 +601,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/test_mlx_bandwidth.py

@@ -0,0 +1,377 @@
+# type: ignore
+import argparse
+import json
+import os
+import statistics
+import sys
+import tempfile
+import time
+
+import mlx.core as mx
+
+DTYPE_MAP = {
+    "float32": (mx.float32, 4),
+    "float16": (mx.float16, 2),
+    "bfloat16": (mx.bfloat16, 2),
+}
+
+SIZES = [
+    1 * 1024,
+    4 * 1024,
+    16 * 1024,
+    64 * 1024,
+    256 * 1024,
+    1 * 1024 * 1024,
+    4 * 1024 * 1024,
+    16 * 1024 * 1024,
+    64 * 1024 * 1024,
+    256 * 1024 * 1024,
+    1 * 1024 * 1024 * 1024,
+    2 * 1024 * 1024 * 1024,
+    4 * 1024 * 1024 * 1024,
+    8 * 1024 * 1024 * 1024,
+]
+
+
+def format_bytes(n: int) -> str:
+    if n >= 1024 * 1024 * 1024:
+        return f"{n / (1024 * 1024 * 1024):.0f} GB"
+    if n >= 1024 * 1024:
+        return f"{n / (1024 * 1024):.0f} MB"
+    if n >= 1024:
+        return f"{n / 1024:.0f} KB"
+    return f"{n} B"
+
+
+def format_time(seconds: float) -> str:
+    if seconds >= 1.0:
+        return f"{seconds:.3f} s"
+    if seconds >= 0.001:
+        return f"{seconds * 1000:.2f} ms"
+    return f"{seconds * 1_000_000:.1f} us"
+
+
+def format_bandwidth(bytes_per_sec: float) -> str:
+    if bytes_per_sec >= 1024 * 1024 * 1024:
+        return f"{bytes_per_sec / (1024 * 1024 * 1024):.2f} GB/s"
+    if bytes_per_sec >= 1024 * 1024:
+        return f"{bytes_per_sec / (1024 * 1024):.1f} MB/s"
+    return f"{bytes_per_sec / 1024:.1f} KB/s"
+
+
+def barrier(group: mx.distributed.Group) -> None:
+    mx.eval(mx.distributed.all_sum(mx.array(1.0), group=group))
+
+
+def init_ring(
+    rank: int, self_ip: str, peer_ip: str, port: int, tmpdir: str
+) -> mx.distributed.Group:
+    if rank == 0:
+        hosts = [f"{self_ip}:{port}", f"{peer_ip}:{port}"]
+    else:
+        hosts = [f"{peer_ip}:{port}", f"{self_ip}:{port}"]
+
+    hostfile = os.path.join(tmpdir, "hosts.json")
+    with open(hostfile, "w") as f:
+        json.dump(hosts, f)
+
+    for var in ("MLX_HOSTFILE", "MLX_RANK", "MLX_IBV_DEVICES", "MLX_JACCL_COORDINATOR"):
+        os.environ.pop(var, None)
+
+    os.environ["MLX_HOSTFILE"] = hostfile
+    os.environ["MLX_RANK"] = str(rank)
+    return mx.distributed.init(backend="ring", strict=True)
+
+
+def init_jaccl(
+    rank: int, interface: str, coordinator: str, port: int, tmpdir: str
+) -> mx.distributed.Group:
+    devices = [[None, interface], [interface, None]]
+    devfile = os.path.join(tmpdir, "devices.json")
+    with open(devfile, "w") as f:
+        json.dump(devices, f)
+
+    for var in ("MLX_HOSTFILE", "MLX_RANK", "MLX_IBV_DEVICES", "MLX_JACCL_COORDINATOR"):
+        os.environ.pop(var, None)
+
+    os.environ["MLX_IBV_DEVICES"] = devfile
+    os.environ["MLX_RANK"] = str(rank)
+    if rank == 0:
+        os.environ["MLX_JACCL_COORDINATOR"] = f"0.0.0.0:{port}"
+    else:
+        os.environ["MLX_JACCL_COORDINATOR"] = coordinator
+
+    return mx.distributed.init(backend="jaccl", strict=True)
+
+
+def bench_unidirectional(
+    group: mx.distributed.Group,
+    rank: int,
+    size_bytes: int,
+    dtype: mx.Dtype,
+    element_size: int,
+    warmup: int,
+    iterations: int,
+) -> list[float]:
+    n_elements = size_bytes // element_size
+    tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
+    mx.eval(tensor)
+
+    for _ in range(warmup):
+        if rank == 0:
+            sent = mx.distributed.send(tensor, dst=1, group=group)
+            mx.eval(sent)
+        else:
+            received = mx.distributed.recv_like(tensor, src=0, group=group)
+            mx.eval(received)
+        barrier(group)
+
+    times: list[float] = []
+    for _ in range(iterations):
+        barrier(group)
+        t0 = time.perf_counter()
+        if rank == 0:
+            sent = mx.distributed.send(tensor, dst=1, group=group)
+            mx.eval(sent)
+        else:
+            received = mx.distributed.recv_like(tensor, src=0, group=group)
+            mx.eval(received)
+        barrier(group)
+        t1 = time.perf_counter()
+        times.append(t1 - t0)
+
+    return times
+
+
+def bench_rtt(
+    group: mx.distributed.Group,
+    rank: int,
+    size_bytes: int,
+    dtype: mx.Dtype,
+    element_size: int,
+    warmup: int,
+    iterations: int,
+) -> list[float]:
+    n_elements = size_bytes // element_size
+    tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
+    mx.eval(tensor)
+
+    for _ in range(warmup):
+        if rank == 0:
+            sent = mx.distributed.send(tensor, dst=1, group=group)
+            mx.eval(sent)
+            received = mx.distributed.recv_like(tensor, src=1, group=group)
+            mx.eval(received)
+        else:
+            received = mx.distributed.recv_like(tensor, src=0, group=group)
+            mx.eval(received)
+            sent = mx.distributed.send(received, dst=0, group=group)
+            mx.eval(sent)
+        barrier(group)
+
+    times: list[float] = []
+    for _ in range(iterations):
+        barrier(group)
+        t0 = time.perf_counter()
+        if rank == 0:
+            sent = mx.distributed.send(tensor, dst=1, group=group)
+            mx.eval(sent)
+            received = mx.distributed.recv_like(tensor, src=1, group=group)
+            mx.eval(received)
+        else:
+            received = mx.distributed.recv_like(tensor, src=0, group=group)
+            mx.eval(received)
+            sent = mx.distributed.send(received, dst=0, group=group)
+            mx.eval(sent)
+        barrier(group)
+        t1 = time.perf_counter()
+        times.append(t1 - t0)
+
+    return times
+
+
+def bench_all_gather(
+    group: mx.distributed.Group,
+    rank: int,
+    size_bytes: int,
+    dtype: mx.Dtype,
+    element_size: int,
+    warmup: int,
+    iterations: int,
+) -> list[float]:
+    n_elements = (size_bytes // 2) // element_size
+    tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
+    mx.eval(tensor)
+
+    for _ in range(warmup):
+        gathered = mx.distributed.all_gather(tensor, group=group)
+        mx.eval(gathered)
+        barrier(group)
+
+    times: list[float] = []
+    for _ in range(iterations):
+        barrier(group)
+        t0 = time.perf_counter()
+        gathered = mx.distributed.all_gather(tensor, group=group)
+        mx.eval(gathered)
+        t1 = time.perf_counter()
+        times.append(t1 - t0)
+
+    return times
+
+
+def print_table(title: str, rows: list[dict[str, str]]) -> None:
+    print(f"\n=== {title} ===")
+    headers = ["Size", "Median", "Min", "Max", "Bandwidth"]
+    widths = [
+        max(len(h), max((len(r[h]) for r in rows), default=0)) + 2 for h in headers
+    ]
+    header_line = "".join(h.ljust(w) for h, w in zip(headers, widths, strict=True))
+    print(header_line)
+    print("-" * len(header_line))
+    for row in rows:
+        print("".join(row[h].ljust(w) for h, w in zip(headers, widths, strict=True)))
+
+
+def run_bench(
+    name: str,
+    bench_fn,
+    group: mx.distributed.Group,
+    rank: int,
+    dtype: mx.Dtype,
+    element_size: int,
+    warmup: int,
+    iterations: int,
+    bw_multiplier: int = 1,
+) -> None:
+    rows: list[dict[str, str]] = []
+    for size in SIZES:
+        if rank == 0:
+            print(f"  {name}: {format_bytes(size)}...", end="", flush=True)
+        times = bench_fn(group, rank, size, dtype, element_size, warmup, iterations)
+        if rank == 0:
+            med = statistics.median(times)
+            mn = min(times)
+            mx_ = max(times)
+            bw = (size * bw_multiplier) / med
+            rows.append(
+                {
+                    "Size": format_bytes(size),
+                    "Median": format_time(med),
+                    "Min": format_time(mn),
+                    "Max": format_time(mx_),
+                    "Bandwidth": format_bandwidth(bw),
+                }
+            )
+            print(f" {format_bandwidth(bw)}")
+    if rank == 0:
+        print_table(name, rows)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="MLX Distributed Communication Benchmark"
+    )
+    subparsers = parser.add_subparsers(dest="backend", required=True)
+
+    ring_parser = subparsers.add_parser("ring")
+    ring_parser.add_argument("--rank", type=int, required=True, choices=[0, 1])
+    ring_parser.add_argument("--self-ip", required=True)
+    ring_parser.add_argument("--peer-ip", required=True)
+    ring_parser.add_argument("--port", type=int, default=5555)
+
+    jaccl_parser = subparsers.add_parser("jaccl")
+    jaccl_parser.add_argument("--rank", type=int, required=True, choices=[0, 1])
+    jaccl_parser.add_argument("--interface", required=True)
+    jaccl_parser.add_argument(
+        "--coordinator",
+        type=str,
+        default=None,
+        help="IP:PORT of rank 0 (required for rank 1)",
+    )
+    jaccl_parser.add_argument(
+        "--port", type=int, default=9999, help="Coordinator port (rank 0 only)"
+    )
+
+    for p in [ring_parser, jaccl_parser]:
+        p.add_argument("--warmup", type=int, default=3)
+        p.add_argument("--iterations", type=int, default=10)
+        p.add_argument("--dtype", choices=list(DTYPE_MAP.keys()), default="float32")
+
+    args = parser.parse_args()
+
+    if args.backend == "jaccl" and args.rank == 1 and args.coordinator is None:
+        jaccl_parser.error("--coordinator is required for rank 1")
+
+    return args
+
+
+def main() -> int:
+    args = parse_args()
+    dtype, element_size = DTYPE_MAP[args.dtype]
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        if args.backend == "ring":
+            print(f"Initializing ring backend (rank {args.rank})...")
+            group = init_ring(args.rank, args.self_ip, args.peer_ip, args.port, tmpdir)
+        else:
+            print(f"Initializing jaccl backend (rank {args.rank})...")
+            group = init_jaccl(
+                args.rank, args.interface, args.coordinator or "", args.port, tmpdir
+            )
+
+        print(f"Rank {group.rank()} of {group.size()} initialized")
+        barrier(group)
+
+        if args.rank == 0:
+            print("\nMLX Distributed Communication Benchmark")
+            print(
+                f"Backend: {args.backend} | Dtype: {args.dtype} | Warmup: {args.warmup} | Iterations: {args.iterations}"
+            )
+
+        run_bench(
+            "Unidirectional (rank 0 -> rank 1)",
+            bench_unidirectional,
+            group,
+            args.rank,
+            dtype,
+            element_size,
+            args.warmup,
+            args.iterations,
+        )
+        run_bench(
+            "Round-Trip (ping-pong)",
+            bench_rtt,
+            group,
+            args.rank,
+            dtype,
+            element_size,
+            args.warmup,
+            args.iterations,
+            bw_multiplier=2,
+        )
+        run_bench(
+            "All-Gather",
+            bench_all_gather,
+            group,
+            args.rank,
+            dtype,
+            element_size,
+            args.warmup,
+            args.iterations,
+        )
+
+        if args.rank == 0:
+            print("\nDone.")
+        else:
+            print("Rank 1 complete.")
+
+    return 0
+
+
+if __name__ == "__main__":
+    try:
+        sys.exit(main())
+    except KeyboardInterrupt:
+        print("\nInterrupted.")
+        sys.exit(1)

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

@@ -168,8 +168,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";
@@ -277,7 +279,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;
@@ -990,6 +992,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)) {
@@ -580,7 +584,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;
@@ -548,6 +548,7 @@ class AppStore {
       { total: { inBytes: number }; available: { inBytes: number } }
     >
   >({});
+  vllmAvailable = $state(false);
   placementPreviews = $state<PlacementPreview[]>([]);
   selectedPreviewModelId = $state<string | null>(null);
   isLoadingPreviews = $state(false);
@@ -1332,6 +1333,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 (
@@ -3490,6 +3500,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

@@ -65,6 +65,7 @@
     nodeThunderboltBridge,
     nodeIdentities,
     isConnected,
+    vllmAvailable,
     type DownloadProgress,
     type PlacementPreview,
   } from "$lib/stores/app.svelte";
@@ -296,7 +297,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);
@@ -702,7 +703,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`,
@@ -886,7 +890,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";
@@ -932,7 +936,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 (
@@ -1146,9 +1154,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 {
@@ -2063,6 +2069,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?: {
@@ -5814,6 +5821,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

@@ -48,8 +48,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 = inputs:
@@ -70,6 +70,9 @@
       debug = true; # Enable options autocompletion
 
       perSystem = { config, self', pkgs, lib, system, ... }:
+        let
+          pkgsCuda = import ./nix/cuda-pkgs.nix { inherit (inputs) nixpkgs; inherit system; };
+        in
         {
           # Allow unfree for metal-toolchain (needed for Darwin Metal packages)
           _module.args.pkgs = import inputs.nixpkgs {
@@ -118,28 +121,51 @@
             };
           };
 
-          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;
-            }
-          );
+          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;
+            # 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 "$@"
+              '';
+            };
+          };
 
-          devShells.default = with pkgs; pkgs.mkShell {
+          devShells.default = pkgs.mkShell {
             inputsFrom = [ self'.checks.cargo-build ];
 
             packages =
-              [
+              with pkgs; [
                 # FORMATTING
                 config.treefmt.build.wrapper
 
@@ -171,12 +197,11 @@
                 macmon
               ];
 
-            OPENSSL_NO_VENDOR = "1";
 
             shellHook = ''
               export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${self'.packages.python}/lib"
-              ${lib.optionalString stdenv.isLinux ''
-                export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH"
+              ${lib.optionalString pkgs.stdenv.isLinux ''
+                export LD_LIBRARY_PATH="${pkgs.openssl.out}/lib:$LD_LIBRARY_PATH"
               ''}
             '';
           };

justfile

@@ -18,6 +18,18 @@ 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 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
+
 sync-clean:
     uv sync --all-packages --force-reinstall --no-cache
 

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,17 +15,17 @@ dependencies = [
   "huggingface-hub>=1.8.0",
   "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",
-  "mflux==0.17.2",
+  "mflux==0.17.2; sys_platform == 'darwin'",
   "python-multipart>=0.0.21",
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
@@ -47,11 +47,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 +65,22 @@ members = ["rust/exo_pyo3_bindings", "bench"]
 [tool.uv.sources]
 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'" }
+torch = [
+  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and platform_machine!='x86_64'" },
+  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and platform_machine=='x86_64'" }
+]
+vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
 mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/fix-arrayscache-leak" }
-# 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
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
 
 [build-system]
 requires = ["uv_build>=0.8.9,<0.9.0"]
@@ -101,8 +115,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"
@@ -115,8 +137,14 @@ root = "src"
 [tool.uv]
 required-version = ">=0.8.6"
 prerelease = "allow"
-environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
-extra-build-dependencies = { "miniaudio" = ["setuptools", "cffi", "pycparser"] }
+environments = [
+  "sys_platform == 'linux' and platform_machine == 'x86_64'",
+  "sys_platform == 'linux' and platform_machine == 'aarch64'",
+  "sys_platform == 'darwin'",
+]
+
+[tool.uv.extra-build-dependencies]
+miniaudio = ["setuptools", "cffi", "pycparser"]
 
 ###
 # ruff configuration
@@ -126,6 +154,7 @@ extra-build-dependencies = { "miniaudio" = ["setuptools", "cffi", "pycparser"] }
 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,10 +216,11 @@
           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;
+      } // {
 
         inherit python;
-
         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);

scripts/check_kv_compat.py

@@ -0,0 +1,25 @@
+import sys
+
+sys.path.insert(0, "src")
+import mlx.core as mx
+from mlx_lm import load
+from mlx_lm.models.cache import KVCache, RotatingKVCache
+
+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,196 @@
+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 json
+import socket
+from collections import defaultdict
+from pathlib import Path
+
+import mlx.core as mx
+import torch
+from mlx_lm import load
+from mlx_lm.models.cache import ArraysCache, RotatingKVCache
+
+from exo.disaggregated.prefill_client import _torch_to_mx
+from exo.disaggregated.protocol import (
+    ArraysState,
+    Done,
+    KVChunk,
+    read_header,
+    read_message,
+)
+
+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_pipeline_prefill.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+"""Test prefill throughput through the full exo pipeline.
+
+Connects to a running exo instance, creates a vLLM placement for the given model,
+sends a prompt of the specified token count, and reports prompt_tps.
+
+Usage:
+    uv run python scripts/test_pipeline_prefill.py --model openai/gpt-oss-20b --pp 8192 --host gx10-de89
+    uv run python scripts/test_pipeline_prefill.py --model openai/gpt-oss-20b --pp 4096,8192 --host gx10-de89 --repeat 3
+"""
+from __future__ import annotations
+
+import argparse
+import contextlib
+import sys
+import time
+from typing import Any
+
+sys.path.insert(0, "bench")
+
+from exo_bench import PromptSizer, load_tokenizer_for_bench
+from harness import (
+    ExoClient,
+    ExoHttpError,
+    instance_id_from_instance,
+    resolve_model_short_id,
+    run_planning_phase,
+    settle_and_fetch_placements,
+    wait_for_instance_gone,
+    wait_for_instance_ready,
+)
+
+
+def run_one(client: ExoClient, model_id: str, pp: int, tg: int, sizer: PromptSizer) -> dict[str, Any]:
+    content, actual_pp = sizer.build(pp)
+    payload = {
+        "model": model_id,
+        "messages": [{"role": "user", "content": content}],
+        "stream": False,
+        "max_tokens": tg,
+    }
+    t0 = time.perf_counter()
+    out = client.post_bench_chat_completions(payload)
+    elapsed = time.perf_counter() - t0
+    stats = out.get("generation_stats", {})
+    return {
+        "elapsed_s": elapsed,
+        "prompt_tps": stats.get("prompt_tps", 0),
+        "prompt_tokens": stats.get("prompt_tokens", 0),
+        "generation_tps": stats.get("generation_tps", 0),
+        "generation_tokens": stats.get("generation_tokens", 0),
+    }
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description="Test prefill throughput through exo pipeline")
+    ap.add_argument("--host", default="localhost")
+    ap.add_argument("--port", type=int, default=52415)
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--pp", required=True, help="Prompt token counts, comma-separated")
+    ap.add_argument("--tg", type=int, default=2, help="Generation tokens (default 2)")
+    ap.add_argument("--repeat", type=int, default=1)
+    ap.add_argument("--timeout", type=float, default=1200.0)
+    ap.add_argument("--force-download", action="store_true")
+    args = ap.parse_args()
+
+    pp_list = [int(x.strip()) for x in args.pp.split(",")]
+    client = ExoClient(args.host, args.port, timeout_s=args.timeout)
+
+    short_id, full_model_id = resolve_model_short_id(client, args.model, force_download=args.force_download)
+    print(f"Model: {short_id} ({full_model_id})")
+
+    tokenizer = load_tokenizer_for_bench(full_model_id)
+    sizer = PromptSizer(tokenizer)
+
+    ns = argparse.Namespace(
+        instance_meta="vllm",
+        sharding="both",
+        max_nodes=1,
+        min_nodes=1,
+        skip_pipeline_jaccl=False,
+        skip_tensor_ring=False,
+        settle_timeout=60,
+        danger_delete_downloads=False,
+        timeout=args.timeout,
+    )
+
+    selected = settle_and_fetch_placements(client, full_model_id, ns, settle_timeout=60)
+    if not selected:
+        print("ERROR: No vLLM placements available")
+        return 1
+
+    preview = selected[0]
+    instance = preview["instance"]
+    instance_id = instance_id_from_instance(instance)
+    print(f"Placement: {preview.get('sharding')} / {preview.get('instance_meta')} / id={instance_id}")
+
+    run_planning_phase(client, full_model_id, preview, False, args.timeout, None)
+
+    client.request_json("POST", "/instance", body={"instance": instance})
+    try:
+        wait_for_instance_ready(client, instance_id)
+    except (RuntimeError, TimeoutError) as e:
+        print(f"ERROR: Failed to initialize: {e}")
+        with contextlib.suppress(ExoHttpError):
+            client.request_json("DELETE", f"/instance/{instance_id}")
+        return 1
+
+    print("Instance ready. Running prefill tests...")
+    time.sleep(2)
+
+    try:
+        for pp in pp_list:
+            for r in range(args.repeat):
+                time.sleep(1)
+                try:
+                    result = run_one(client, full_model_id, pp, args.tg, sizer)
+                except Exception as e:
+                    print(f"  pp={pp} repeat={r+1}: FAILED - {e}")
+                    continue
+                prompt_tps = result["prompt_tps"]
+                marker = "OK" if prompt_tps >= 3500 else "BELOW TARGET"
+                print(
+                    f"  pp={pp} repeat={r+1}: "
+                    f"prompt_tps={prompt_tps:.1f} "
+                    f"prompt_tokens={result['prompt_tokens']} "
+                    f"elapsed={result['elapsed_s']:.2f}s "
+                    f"[{marker}]"
+                )
+    finally:
+        try:
+            client.request_json("DELETE", f"/instance/{instance_id}")
+        except ExoHttpError as e:
+            if e.status != 404:
+                raise
+        wait_for_instance_gone(client, instance_id)
+        print("Instance cleaned up.")
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(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/api/adapters/chat_completions.py

@@ -168,6 +168,7 @@ async 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,
         images=images,
     )
 

src/exo/api/adapters/ollama.py

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

src/exo/api/main.py

@@ -358,6 +358,7 @@ class API:
         self.app.get("/ollama/api/ps")(self.ollama_ps)
         self.app.get("/ollama/api/version")(self.ollama_version)
 
+        self.app.get("/capabilities")(self._get_capabilities)
         self.app.get("/state")(self.get_state)
         self.app.get("/state/{path:path}")(self.get_state)
         self.app.get("/events")(self.stream_events)
@@ -451,6 +452,7 @@ class API:
                 node_network=self.state.node_network,
                 topology=self.state.topology,
                 current_instances=self.state.instances,
+                node_vllm=self.state.node_vllm,
                 download_status=self.state.downloads,
             )
         except ValueError as exc:
@@ -487,18 +489,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:
@@ -511,6 +510,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,
@@ -846,27 +846,34 @@ 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 = await self._send_text_generation_with_images(task_params)
 
         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.
@@ -885,6 +892,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/api/types/api.py

@@ -229,10 +229,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/api/types/ollama_api.py

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

src/exo/disaggregated/batch_connector.py

@@ -0,0 +1,111 @@
+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).to("cpu", non_blocking=True) for t in kv_layer]  # pyright: ignore[reportAny]
+            return
+
+        if slot_mapping is not None:
+            if kv_layer.shape[0] == 2:  # pyright: ignore[reportAny]
+                k_all = kv_layer[0]  # pyright: ignore[reportAny]
+                v_all = kv_layer[1]  # pyright: ignore[reportAny]
+            else:
+                k_all = kv_layer[:, 0]  # pyright: ignore[reportAny]
+                v_all = 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.to("cpu", non_blocking=True)], dim=0),  # type: ignore
+                    "values": torch.cat([prev["values"], values.to("cpu", non_blocking=True)], dim=0),  # type: ignore
+                }
+            else:
+                self.captured_layers[layer_idx] = {
+                    "keys": keys.to("cpu", non_blocking=True),  # pyright: ignore[reportAny]
+                    "values": values.to("cpu", non_blocking=True),  # 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,664 @@
+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.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,130 @@
+from __future__ import annotations
+
+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_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).to("cpu", non_blocking=True) 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 = kv_layer[0]  # pyright: ignore[reportAny]
+                v_all = kv_layer[1]  # pyright: ignore[reportAny]
+            else:
+                k_all = kv_layer[:, 0]  # pyright: ignore[reportAny]
+                v_all = 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.to("cpu", non_blocking=True), values.to("cpu", non_blocking=True)))  # pyright: ignore[reportAny]
+        else:
+            self._queue.put((layer_idx, kv_layer.to("cpu", non_blocking=True).clone(), kv_layer.to("cpu", non_blocking=True).clone()))  # pyright: ignore[reportAny]
+
+    def wait_for_save(self) -> None:
+        pass
+
+    def finish(self) -> None:
+        self._queue.put(None)
+
+    def request_finished_all_groups(self, request: Any, block_ids: tuple[list[int], ...]) -> tuple[bool, dict[str, Any] | None]:  # pyright: ignore[reportAny]
+        return False, None
+
+    def get_num_new_matched_tokens(self, request: Any, num_computed_tokens: int) -> tuple[int, bool]:  # pyright: ignore[reportAny]
+        return 0, False
+
+    def update_state_after_alloc(self, request: Any, blocks: Any, num_external_tokens: int) -> None:  # pyright: ignore[reportAny]
+        pass
+
+    def build_connector_meta(self, scheduler_output: Any) -> StreamingConnectorMetadata:  # pyright: ignore[reportAny]
+        return StreamingConnectorMetadata()

src/exo/main.py

@@ -287,6 +287,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"] = "true"
@@ -318,6 +322,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
     bootstrap_peers: list[str] = []
     libp2p_port: int
@@ -378,6 +383,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",
+        )
         parser.add_argument(
             "--bootstrap-peers",
             type=lambda s: [p for p in s.split(",") if p],

src/exo/master/main.py

@@ -62,7 +62,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.disk_event_log import DiskEventLog
 from exo.utils.event_buffer import MultiSourceBuffer
@@ -97,8 +98,72 @@ 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
+
+
+            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:
@@ -112,14 +177,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
@@ -128,19 +193,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(
@@ -149,12 +217,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,
@@ -163,7 +245,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,
                                     ),
                                 )
                             )
@@ -298,6 +380,7 @@ class Master:
                                 self.state.instances,
                                 self.state.node_memory,
                                 self.state.node_network,
+                                self.state.node_vllm,
                                 download_status=self.state.downloads,
                             )
                             transition_events = get_transition_events(
@@ -388,7 +471,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

@@ -44,6 +44,7 @@ from exo.shared.types.worker.instances import (
     InstanceMeta,
     MlxJacclInstance,
     MlxRingInstance,
+    VllmInstance,
 )
 from exo.shared.types.worker.shards import Sharding
 
@@ -108,12 +109,45 @@ 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,
     download_status: Mapping[NodeId, Sequence[DownloadProgress]] | 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) \
+        and 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) \
+        and 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 = [
@@ -121,8 +155,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")
@@ -192,7 +229,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
 
@@ -252,6 +289,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

@@ -147,7 +147,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
@@ -189,7 +189,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]
@@ -216,7 +216,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]
@@ -245,7 +245,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):
@@ -344,7 +344,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
@@ -432,7 +432,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

@@ -51,9 +51,11 @@ from exo.utils.info_gatherer.info_gatherer import (
     NodeConfig,
     NodeDiskUsage,
     NodeNetworkInterfaces,
+    NvmlMetrics,
     RdmaCtlStatus,
     StaticNodeInformation,
     ThunderboltBridgeInfo,
+    VllmCapability,
 )
 
 
@@ -373,6 +375,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

@@ -1,4 +1,5 @@
 import json
+import re
 from enum import Enum
 from typing import Annotated, Any
 
@@ -40,6 +41,32 @@ _BUILTIN_CARD_DIRS = [
 
 _card_cache: dict[ModelId, "ModelCard"] = {}
 
+_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
+
 
 def _detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
     normalized = model_id.normalize()
@@ -150,6 +177,15 @@ class ModelCard(CamelCaseModel):
             )
         return self
 
+    @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]:
@@ -196,6 +232,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("_", "-")
+
         return ModelCard(
             model_id=ModelId(model_id),
             storage_size=mem_size_bytes,
@@ -206,6 +245,8 @@ class ModelCard(CamelCaseModel):
             context_length=config_data.max_position_embeddings,
             tasks=[ModelTask.TextGeneration],
             trust_remote_code=False,
+            base_model=base_model,
+            family=family,
             is_custom=True,
             vision=config_data.vision,
         )
@@ -229,6 +270,7 @@ async def delete_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
@@ -333,21 +375,24 @@ async def fetch_safetensors_size(model_id: ModelId) -> Memory:
     from exo.shared.types.worker.downloads import ModelSafetensorsIndex
 
     target_dir = await resolve_model_dir(model_id)
-    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/tracing.py

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

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/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

@@ -108,6 +108,8 @@ 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
     images: list[Base64Image] = Field(default_factory=list)
     image_hashes: dict[int, Base64ImageHash] = Field(default_factory=dict)
     total_input_chunks: int = 0

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

@@ -31,6 +31,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,
@@ -42,6 +43,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.
 
@@ -353,6 +360,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
@@ -361,6 +383,8 @@ GatheredInfo = (
     | MacThunderboltConnections
     | RdmaCtlStatus
     | ThunderboltBridgeInfo
+    | VllmCapability
+    | NvmlMetrics
     | NodeConfig
     | MiscData
     | StaticNodeInformation
@@ -371,6 +395,7 @@ GatheredInfo = (
 @dataclass
 class InfoGatherer:
     info_sender: Sender[GatheredInfo]
+    vllm_capability_poll_interval: float | None = 60
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
     _psutil_enabled: bool = field(init=False, default=False)
 
@@ -417,12 +442,15 @@ class InfoGatherer:
                 tg.start_soon(self._monitor_system_profiler_thunderbolt_data, 5)
                 tg.start_soon(self._monitor_thunderbolt_bridge_status, 10)
                 tg.start_soon(self._monitor_rdma_ctl_status, 10)
-            if not IS_DARWIN:
+            else:
                 tg.start_soon(self._monitor_memory_usage, 1)
+                if _has_nvml():
+                    tg.start_soon(self._monitor_nvml_metrics, 1)
             tg.start_soon(self._watch_system_info, 10)
             tg.start_soon(self._monitor_misc, 60)
             tg.start_soon(self._monitor_static_info, 60)
             tg.start_soon(self._monitor_disk_usage, 30)
+            tg.start_soon(self._monitor_vllm_capability, 60)
 
             nc = await NodeConfig.gather()
             if nc is not None:
@@ -538,6 +566,26 @@ class InfoGatherer:
                 logger.opt(exception=e).warning("Error gathering disk usage")
             await anyio.sleep(disk_poll_interval)
 
+    async def _monitor_nvml_metrics(self, nvml_interval: float):
+        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(nvml_interval)
+
+    async def _monitor_vllm_capability(self, vllm_capability_poll_interval: int):
+        try:
+            while True:
+                await self.info_sender.send(await VllmCapability.gather())
+                await anyio.sleep(vllm_capability_poll_interval)
+        except Exception as e:
+            logger.warning(f"Error gathering vLLM capability: {e}")
+
     async def _monitor_macmon(self, macmon_interval: float):
         if (
             macmon_path := os.getenv("EXO_MACMON_PATH") or shutil.which("macmon")

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,6 +1,6 @@
 import os
 from copy import deepcopy
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import mlx.core as mx
 import psutil
@@ -14,12 +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.mlx.vision import MediaRegion
+    from exo.worker.engines.vllm.kv_cache import TorchKVCache
 
 
 # Fraction of device memory above which LRU eviction kicks in.
@@ -50,7 +51,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)):
@@ -74,7 +75,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)
 
@@ -100,7 +101,7 @@ class KVPrefixCache:
     def add_kv_cache(
         self,
         prompt_tokens: mx.array,
-        cache: KVCacheType,
+        cache: MLXCacheType,
         ssm_snapshots: list[CacheSnapshot] | None = None,
         media_regions: list["MediaRegion"] | None = None,
     ):
@@ -118,7 +119,7 @@ class KVPrefixCache:
         self,
         index: int,
         prompt_tokens: mx.array,
-        cache: KVCacheType,
+        cache: MLXCacheType,
         snapshots: list[CacheSnapshot] | None,
         restore_pos: int,
         media_regions: list["MediaRegion"] | None = None,
@@ -139,10 +140,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]
@@ -206,16 +211,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)
@@ -230,6 +244,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")
+
     @staticmethod
     def _validate_media_match(
         match_length: int,
@@ -298,7 +356,7 @@ class KVPrefixCache:
 
 
 def trim_cache(
-    cache: KVCacheType,
+    cache: MLXCacheType,
     num_tokens: int,
     snapshot: CacheSnapshot | None = None,
 ) -> None:
@@ -336,7 +394,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), default=0)
 
@@ -365,7 +423,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/generator/batch_generate.py

@@ -22,13 +22,16 @@ from exo.api.types 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,
 )
@@ -39,6 +42,7 @@ from exo.worker.engines.mlx.generator.generate import (
     extract_top_logprobs,
     patch_embed_tokens,
     prefill,
+    warmup_inference,
 )
 from exo.worker.engines.mlx.patches.opt_batch_gen import (
     set_needs_topk,
@@ -93,10 +97,12 @@ class ExoBatchGenerator:
     tokenizer: TokenizerWrapper
     group: mx.distributed.Group | None
     kv_prefix_cache: KVPrefixCache | None
+    model_id: ModelId
     vision_processor: VisionProcessor | None = None
 
     _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._mlx_gen = MlxBatchGenerator(
@@ -106,6 +112,9 @@ class ExoBatchGenerator:
         )
         self._step_count = 0
 
+    def warmup(self) -> int:
+        return warmup_inference(self.model, self.tokenizer, self.group, self.model_id)
+
     @property
     def has_work(self) -> bool:
         return (
@@ -117,12 +126,13 @@ class ExoBatchGenerator:
 
     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
@@ -185,28 +195,61 @@ class ExoBatchGenerator:
             top_k=task_params.top_k if task_params.top_k is not None else 0,
         )
 
-        vision_ctx = (
-            patch_embed_tokens(
-                self.model,
-                vision.embeddings,
-                prefix_hit_length,
-                len(prompt_tokens) - 1,
-                image_token_id=vision.image_token_id,
-            )
-            if vision is not None
-            else contextlib.nullcontext()
-        )
-        with vision_ctx:
-            _prefill_tps, _prefill_tokens, cache_snapshots = prefill(
-                self.model,
-                self.tokenizer,
-                sampler,
-                prompt_tokens[:-1],
-                cache,
-                self.group,
-                on_prefill_progress,
-                distributed_prompt_progress_callback,
+        _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:
+            vision_ctx = (
+                patch_embed_tokens(
+                    self.model,
+                    vision.embeddings,
+                    prefix_hit_length,
+                    len(prompt_tokens) - 1,
+                    image_token_id=vision.image_token_id,
+                )
+                if vision is not None
+                else contextlib.nullcontext()
             )
+            with vision_ctx:
+                _prefill_tps, _prefill_tokens, cache_snapshots = prefill(
+                    self.model,
+                    self.tokenizer,
+                    sampler,
+                    prompt_tokens[:-1],
+                    cache,
+                    self.group,
+                    on_prefill_progress,
+                    distributed_prompt_progress_callback,
+                )
 
         # We need to clamp rotating kv caches to max size so that mlx lm's _merge_caches behaves
         for c in cache:
@@ -235,7 +278,7 @@ class ExoBatchGenerator:
                 media_regions,
             )
 
-        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(
@@ -251,7 +294,7 @@ class ExoBatchGenerator:
         max_tokens = task_params.max_output_tokens or MAX_TOKENS
 
         uids = self._mlx_gen.insert(
-            prompts=[cast(list[int], last_tokens.tolist())],
+            prompts=[last_tokens.tolist()],
             max_tokens=[max_tokens],
             caches=[list(cache)],
             samplers=[sampler],
@@ -261,6 +304,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,
@@ -276,9 +320,9 @@ class ExoBatchGenerator:
             media_regions=media_regions,
         )
 
-        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 []
 
@@ -293,7 +337,7 @@ class ExoBatchGenerator:
 
         topk = take_ready_topk(gb)
 
-        results: list[tuple[int, GenerationResponse]] = []
+        results: list[tuple[TaskId, GenerationResponse]] = []
 
         for response in responses:
             if response.uid not in self._active_tasks:
@@ -367,8 +411,13 @@ class ExoBatchGenerator:
             stats: GenerationStats | None = None
             usage: Usage | None = None
             if is_done:
-                if state.completion_tokens > 1:
-                    gen_span = state.last_gen_token_time - state.first_gen_token_time
+                try:
+                    mlx_stats = self._mlx_gen.stats()
+                    generation_tps = mlx_stats.generation_tps
+                except ZeroDivisionError:
+                    generation_elapsed = (
+                        time.perf_counter() - state.generation_start_time
+                    )
                     generation_tps = (
                         (state.completion_tokens - 1) / gen_span
                         if gen_span > 0
@@ -399,7 +448,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,
@@ -414,6 +463,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
@@ -432,18 +482,23 @@ class ExoBatchGenerator:
 
         return results
 
-    def cancel(self, uids: list[int]) -> None:
-        self._mlx_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._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

@@ -24,7 +24,7 @@ from exo.api.types import (
 )
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType, Model
+from exo.shared.types.mlx import MLXCacheType, Model
 from exo.shared.types.text_generation import (
     InputMessage,
     InputMessageContent,
@@ -160,7 +160,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,
@@ -197,7 +197,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()
 
@@ -280,7 +280,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

@@ -148,6 +148,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
@@ -556,7 +561,7 @@ def render_chat_template(
         )
         if partial_assistant_content:
             prompt += partial_assistant_content
-        logger.info(prompt)
+        logger.debug(prompt)
         return prompt
 
     for msg in formatted_messages:

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

@@ -0,0 +1,457 @@
+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()
+    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)
+
+        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_model_runner import GPUModelRunner
+    from vllm.v1.worker.gpu_worker import Worker
+
+    original = Worker.initialize_from_config
+    original_init_attn = GPUModelRunner.initialize_attn_backend
+
+    def clear_and_reinit_attn(self: "GPUModelRunner", *args: object, **kwargs: object) -> None:
+        self.attn_groups.clear()
+        original_init_attn(self, *args, **kwargs)
+
+    GPUModelRunner.initialize_attn_backend = clear_and_reinit_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)
+        while 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
+
+    if hasattr(KVCacheManager, "can_fit_full_sequence"):
+        original_can_fit = KVCacheManager.can_fit_full_sequence
+
+        def patched_can_fit(self: "KVCacheManager", *args: "object", **kwargs: "object") -> bool:
+            result = original_can_fit(self, *args, **kwargs)
+            while not result and _try_grow_cache(self):
+                result = original_can_fit(self, *args, **kwargs)
+            return result
+
+        KVCacheManager.can_fit_full_sequence = patched_can_fit  # 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 is not None and (not isinstance(old_kv_list, torch.Tensor) or old_kv_list.numel() > 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]
+
+

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

@@ -0,0 +1,344 @@
+from collections.abc import Iterator, Sequence
+from copy import deepcopy
+from dataclasses import dataclass
+
+import mlx.core as mx
+import numpy as np
+import torch
+from mlx_lm.models.cache import (
+    ArraysCache,
+    CacheList,
+    KVCache,
+    QuantizedKVCache,
+    RotatingKVCache,
+)
+
+
+@dataclass
+class KVLayerState:
+    keys: torch.Tensor  # [seq_len, n_heads, head_dim]
+    values: torch.Tensor  # [seq_len, n_heads, head_dim]
+
+
+@dataclass
+class RotatingKVLayerState:
+    keys: torch.Tensor  # [buffer_len, n_heads, head_dim]
+    values: torch.Tensor  # [buffer_len, n_heads, head_dim]
+    keep: int
+    max_size: int
+    offset: int
+    idx: int
+
+
+@dataclass
+class ArraysLayerState:
+    arrays: list[torch.Tensor | None]
+
+
+LayerState = KVLayerState | RotatingKVLayerState | ArraysLayerState
+
+
+def _mx_to_torch(arr: mx.array) -> torch.Tensor:
+    mx.eval(arr)
+    if arr.dtype == mx.bfloat16:
+        return torch.from_numpy(np.array(arr.astype(mx.float32))).to(torch.bfloat16)
+    return torch.from_numpy(np.array(arr))
+
+
+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)  # pyright: ignore[reportAny]
+    return mx.array(t.numpy())  # pyright: ignore[reportAny]
+
+
+def _split_kv(
+    kv: torch.Tensor | list[torch.Tensor],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if isinstance(kv, list):
+        return kv[0], kv[1]
+    if kv.shape[0] == 2 and kv.shape[1] != 2:
+        return kv[0], kv[1]
+    return kv[:, 0], kv[:, 1]
+
+
+def _kv_to_nhd(k: mx.array, v: mx.array) -> tuple[torch.Tensor, torch.Tensor]:
+    """Convert MLX BHSD [1, H, S, D] to NHD [S, H, D] torch tensors."""
+    kt = _mx_to_torch(k).squeeze(0).permute(1, 0, 2)  # [H,S,D] -> [S,H,D]
+    vt = _mx_to_torch(v).squeeze(0).permute(1, 0, 2)
+    return kt, vt
+
+
+def _nhd_to_bhsd(kt: torch.Tensor, vt: torch.Tensor) -> tuple[mx.array, mx.array]:
+    """Convert NHD [S, H, D] torch tensors to MLX BHSD [1, H, S, D]."""
+    k_mx = _torch_to_mx(kt.permute(1, 0, 2).unsqueeze(0))  # [S,H,D] -> [1,H,S,D]
+    v_mx = _torch_to_mx(vt.permute(1, 0, 2).unsqueeze(0))
+    return k_mx, v_mx
+
+
+class TorchKVCache:
+    def __init__(
+        self, layers: list[LayerState], token_offset_per_group: list[int] | None = None
+    ):
+        self.layers = layers
+        self.token_offset_per_group = token_offset_per_group or []
+        self._num_tokens: int | None = None
+
+    @property
+    def num_layers(self) -> int:
+        return len(self.layers)
+
+    def layer(self, idx: int) -> LayerState:
+        return self.layers[idx]
+
+    def kv_layers(self) -> list[tuple[int, KVLayerState | RotatingKVLayerState]]:
+        return [
+            (i, layer)
+            for i, layer in enumerate(self.layers)
+            if isinstance(layer, (KVLayerState, RotatingKVLayerState))
+        ]
+
+    def detach_cpu(self) -> "TorchKVCache":
+        layers: list[LayerState] = []
+        for layer in self.layers:
+            if isinstance(layer, KVLayerState):
+                if not layer.keys.is_cuda:
+                    layers.append(layer)
+                else:
+                    layers.append(
+                        KVLayerState(
+                            keys=layer.keys.detach().to("cpu", non_blocking=True),
+                            values=layer.values.detach().to("cpu", non_blocking=True),
+                        )
+                    )
+            elif isinstance(layer, RotatingKVLayerState):
+                layers.append(
+                    RotatingKVLayerState(
+                        keys=layer.keys.detach().to("cpu", non_blocking=True),
+                        values=layer.values.detach().to("cpu", non_blocking=True),
+                        keep=layer.keep,
+                        max_size=layer.max_size,
+                        offset=layer.offset,
+                        idx=layer.idx,
+                    )
+                )
+            else:
+                layers.append(deepcopy(layer))
+        if any(
+            layer.keys.is_cuda
+            for layer in self.layers
+            if isinstance(layer, (KVLayerState, RotatingKVLayerState))
+        ):
+            torch.cuda.synchronize()
+        return TorchKVCache(layers, list(self.token_offset_per_group))
+
+    def trim_to(self, num_tokens: int) -> "TorchKVCache":
+        trimmed = TorchKVCache(list(self.layers), list(self.token_offset_per_group))
+        trimmed._num_tokens = num_tokens
+        return trimmed
+
+    @property
+    def num_tokens(self) -> int | None:
+        return getattr(self, "_num_tokens", None)
+
+    @classmethod
+    def from_mlx_cache(
+        cls,
+        cache: Sequence[
+            KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList
+        ],
+    ) -> "TorchKVCache":
+        layers: list[LayerState] = []
+        for c in cache:
+            if isinstance(c, RotatingKVCache):
+                if c.keys is None:
+                    layers.append(
+                        RotatingKVLayerState(
+                            keys=torch.empty(0),
+                            values=torch.empty(0),
+                            keep=c.keep,
+                            max_size=c.max_size,
+                            offset=c.offset,
+                            idx=c._idx,
+                        )
+                    )
+                else:
+                    k, v = c.state
+                    kt, vt = _kv_to_nhd(k, v)  # pyright: ignore[reportArgumentType]
+                    keep, max_size, offset, idx = (int(x) for x in c.meta_state)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportUnknownArgumentType]
+                    layers.append(
+                        RotatingKVLayerState(
+                            keys=kt,
+                            values=vt,
+                            keep=keep,
+                            max_size=max_size,
+                            offset=offset,
+                            idx=idx,
+                        )
+                    )
+            elif isinstance(c, ArraysCache):
+                arrays: list[torch.Tensor | None] = []
+                for arr in c.state:
+                    arrays.append(_mx_to_torch(arr) if arr is not None else None)
+                layers.append(ArraysLayerState(arrays=arrays))
+            else:
+                if c.keys is None:  # pyright: ignore[reportUnnecessaryComparison]
+                    layers.append(
+                        KVLayerState(keys=torch.empty(0), values=torch.empty(0))
+                    )
+                else:
+                    k, v = c.state
+                    kt, vt = _kv_to_nhd(k, v)  # pyright: ignore[reportArgumentType]
+                    layers.append(KVLayerState(keys=kt, values=vt))
+        return cls(layers)
+
+    def to_mlx_cache(self) -> list[KVCache | RotatingKVCache | ArraysCache]:
+        result: list[KVCache | RotatingKVCache | ArraysCache] = []
+        for layer in self.layers:
+            if isinstance(layer, RotatingKVLayerState):
+                c = RotatingKVCache(max_size=layer.max_size, keep=layer.keep)
+                if layer.keys.numel() > 0:
+                    k_mx, v_mx = _nhd_to_bhsd(layer.keys, layer.values)
+                    c.state = (k_mx, v_mx)
+                    c.meta_state = tuple(
+                        str(x)
+                        for x in (layer.keep, layer.max_size, layer.offset, layer.idx)
+                    )
+                result.append(c)
+            elif isinstance(layer, ArraysLayerState):
+                c = ArraysCache(size=len(layer.arrays))
+                c.state = [
+                    _torch_to_mx(arr) if arr is not None else None
+                    for arr in layer.arrays
+                ]
+                result.append(c)
+            else:
+                c = KVCache()
+                if layer.keys.numel() > 0:
+                    k_mx, v_mx = _nhd_to_bhsd(layer.keys, layer.values)
+                    c.state = (k_mx, v_mx)
+                result.append(c)
+        return result
+
+    @classmethod
+    def from_vllm_cache(
+        cls,
+        kv_caches: list[torch.Tensor | list[torch.Tensor]],
+        block_ids_per_group: list[list[int]],
+        layer_to_group: list[int],
+        num_tokens: int,
+        token_offset_per_group: list[int] | None = None,
+        block_sizes_per_group: list[int] | None = None,
+    ) -> "TorchKVCache":
+        block_tables = [
+            torch.tensor(ids, dtype=torch.long) for ids in block_ids_per_group
+        ]
+        if token_offset_per_group is None:
+            token_offset_per_group = [0] * len(block_ids_per_group)
+
+        layers: list[LayerState] = []
+        for layer_idx, kv in enumerate(kv_caches):
+            gi = layer_to_group[layer_idx]
+            bt = block_tables[gi]
+            k_all, v_all = _split_kv(kv)
+
+            if len(bt) == 0:
+                layers.append(KVLayerState(keys=torch.empty(0), values=torch.empty(0)))
+                continue
+
+            if k_all.dim() >= 4 and len(bt) > 0 and block_sizes_per_group is not None:
+                page_size = k_all.shape[1]
+                sched_block_size = block_sizes_per_group[gi]
+                pages_per_block = sched_block_size // page_size
+                if pages_per_block > 1:
+                    expanded = []
+                    for b in bt.tolist():
+                        start_page = b * pages_per_block
+                        end_page = min(start_page + pages_per_block, k_all.shape[0])
+                        expanded.extend(range(start_page, end_page))
+                    bt = torch.tensor(expanded, dtype=torch.long)
+
+            keys = k_all[bt].to("cpu", non_blocking=True)
+            values = v_all[bt].to("cpu", non_blocking=True)
+            torch.cuda.synchronize()
+            layers.append(KVLayerState(keys=keys, values=values))
+        return cls(layers, list(token_offset_per_group))
+
+    def write_to_vllm_blocks(
+        self,
+        kv_caches: list[torch.Tensor | list[torch.Tensor]],
+        block_ids_per_group: list[list[int]],
+        layer_to_group: list[int],
+        token_offset_per_group: list[int] | None = None,
+    ) -> None:
+        block_tables = [
+            torch.tensor(ids, dtype=torch.long) for ids in block_ids_per_group
+        ]
+
+        first = kv_caches[0]
+        device = first[0].device if isinstance(first, list) else first.device
+        for layer_idx, layer in enumerate(self.layers):
+            if isinstance(layer, ArraysLayerState):
+                gi = layer_to_group[layer_idx]
+                bt = block_tables[gi]
+                kv = kv_caches[layer_idx]
+                if isinstance(kv, list):
+                    for ti, (stored, target) in enumerate(zip(layer.arrays, kv)):
+                        if stored is not None and target is not None:
+                            n = min(len(bt), stored.shape[0])
+                            if n > 0:
+                                target[bt[:n]] = stored[:n].to(device, non_blocking=True)
+                continue
+            if not isinstance(layer, KVLayerState):
+                continue
+            gi = layer_to_group[layer_idx]
+            bt = block_tables[gi]
+            kv = kv_caches[layer_idx]
+            k_all, v_all = _split_kv(kv)
+
+            keys = layer.keys
+            values = layer.values
+            block_size = k_all.shape[-3] if k_all.dim() >= 3 else k_all.shape[1]
+            needs_reshape = keys.dim() == 3 and keys.shape[1:] != k_all.shape[1:]
+            if needs_reshape:
+                offset = token_offset_per_group[gi] if token_offset_per_group else 0
+                if offset > 0:
+                    keys = keys[offset:]
+                    values = values[offset:]
+                s, h, d = keys.shape
+                pad = (block_size - s % block_size) % block_size
+                if pad > 0:
+                    keys = torch.nn.functional.pad(keys, (0, 0, 0, 0, 0, pad))
+                    values = torch.nn.functional.pad(values, (0, 0, 0, 0, 0, pad))
+                keys = keys.reshape(-1, block_size, h, d)
+                values = values.reshape(-1, block_size, h, d)
+
+            n_blocks = min(len(bt), keys.shape[0])
+            if n_blocks > 0:
+                k_all[bt[:n_blocks]] = keys[:n_blocks].to(device, non_blocking=True)
+                v_all[bt[:n_blocks]] = values[:n_blocks].to(device, non_blocking=True)
+        torch.cuda.synchronize()
+
+    def __iter__(self) -> Iterator[LayerState]:
+        return iter(self.layers)
+
+    def __len__(self) -> int:
+        return len(self.layers)
+
+    def __repr__(self) -> str:
+        parts: list[str] = [f"TorchKVCache({self.num_layers} layers)"]
+        for i, layer in enumerate(self.layers):
+            if isinstance(layer, KVLayerState):
+                parts.append(
+                    f"  [{i}] KV: keys={list(layer.keys.shape)} values={list(layer.values.shape)} {layer.keys.dtype}"
+                )
+            elif isinstance(layer, RotatingKVLayerState):
+                parts.append(
+                    f"  [{i}] RotatingKV: keys={list(layer.keys.shape)} keep={layer.keep} max_size={layer.max_size} offset={layer.offset} idx={layer.idx}"
+                )
+            else:
+                shapes = [
+                    list(a.shape) if a is not None else None for a in layer.arrays
+                ]
+                parts.append(f"  [{i}] Arrays: {shapes}")
+        return "\n".join(parts)

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

@@ -0,0 +1,61 @@
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+from vllm.sampling_params import SamplingParams
+from vllm.v1.engine.llm_engine import LLMEngine
+
+from exo.shared.types.common import ModelId
+from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.worker.engines.mlx.utils_mlx import (
+    apply_chat_template,
+    get_eos_token_ids_for_model,
+)
+
+
+def format_vllm_prompt(
+    engine: LLMEngine, params: TextGenerationTaskParams
+) -> tuple[list[int], str, int]:
+    # we should have our own wrapper
+    # (instead of abusing mlx's TokenizerWrapper, use tokenizers Tokenizer)
+    tokenizer = TokenizerWrapper(engine.get_tokenizer())
+    prompt_text = apply_chat_template(tokenizer, params)
+    token_ids: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False)  # type: ignore[reportUnknownMemberType]
+    return token_ids, prompt_text, len(token_ids)
+
+
+def make_vllm_sampling_params(
+    engine: LLMEngine,
+    params: TextGenerationTaskParams,
+    model_id: ModelId | None = None,
+) -> SamplingParams:
+    kwargs: dict[str, object] = {}
+
+    if params.max_output_tokens is not None:
+        kwargs["max_tokens"] = params.max_output_tokens
+    else:
+        kwargs["max_tokens"] = min(engine.model_config.max_model_len, 32168)
+    if params.temperature is not None:
+        kwargs["temperature"] = params.temperature
+    if params.top_p is not None:
+        kwargs["top_p"] = params.top_p
+    if params.top_k is not None:
+        kwargs["top_k"] = params.top_k
+    if params.min_p is not None:
+        kwargs["min_p"] = params.min_p
+    if params.stop is not None:
+        kwargs["stop"] = params.stop
+    if params.seed is not None:
+        kwargs["seed"] = params.seed
+    if params.repetition_penalty is not None:
+        kwargs["repetition_penalty"] = params.repetition_penalty
+    if params.logprobs:
+        kwargs["logprobs"] = params.top_logprobs or 1
+
+    if model_id is not None:
+        extra_stop = get_eos_token_ids_for_model(model_id)
+        if extra_stop:
+            kwargs["stop_token_ids"] = extra_stop
+
+    if params.bench:
+        kwargs["ignore_eos"] = True
+        kwargs["min_tokens"] = kwargs["max_tokens"]
+
+    return SamplingParams(**kwargs)

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

@@ -0,0 +1,647 @@
+import gc
+import math
+import os
+import re
+import sys
+import time
+from collections.abc import Callable, Generator
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import torch
+from vllm.engine.arg_utils import EngineArgs
+from vllm.sampling_params import SamplingParams
+from vllm.v1.engine.llm_engine import LLMEngine
+from vllm.v1.kv_cache_interface import KVCacheConfig
+
+from exo.api.types import (
+    CompletionTokensDetails,
+    GenerationStats,
+    PromptTokensDetails,
+    Usage,
+)
+from exo.shared.types.common import ModelId
+from exo.shared.types.memory import Memory
+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 KVPrefixCache
+from exo.worker.engines.mlx.utils_mlx import get_eos_token_ids_for_model
+from exo.worker.engines.vllm.growable_cache import (
+    get_model_runner,
+    patch_vllm,
+    set_prefix_cache,
+)
+from exo.worker.engines.vllm.kv_cache import TorchKVCache
+from exo.worker.engines.vllm.prompt_format import (
+    format_vllm_prompt,
+    make_vllm_sampling_params,
+)
+from exo.worker.runner.bootstrap import logger
+from exo.worker.runner.llm_inference.tool_parsers import ToolParser, infer_tool_parser
+
+
+def _build_layer_groups(kv_cache_config: KVCacheConfig) -> list[int]:
+    group_lookup: dict[str, int] = {}
+    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: list[int] = []
+    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
+
+
+@dataclass
+class _EngineRequest:
+    request_id: str
+    prompt_token_count: int
+    prompt_token_ids: list[int]
+    prefill_done: bool = False
+    prefill_steps: int = 0
+    prev_text: str = ""
+    prev_token_count: int = 0
+    start_time: float = field(default_factory=time.perf_counter)
+    first_token_time: float | None = None
+    on_generation_token: Callable[[], None] | None = None
+    on_prefill_progress: Callable[[int, int], None] | None = None
+
+
+def _save_prefix_cache(
+    engine: LLMEngine,
+    prefix_cache: KVPrefixCache,
+    request_id: str,
+    prompt_token_ids: list[int],
+    prompt_token_count: int,
+) -> None:
+    try:
+        coordinator = None
+        model_runner = get_model_runner()
+        kv_cache_config = None
+        try:
+            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
+        except Exception:
+            pass
+        if coordinator is None or model_runner is None or kv_cache_config is None:
+            return
+
+        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
+
+        null_block = coordinator.block_pool.null_block  # type: ignore
+        block_ids_per_group: list[list[int]] = []
+        token_offset_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)
+                continue
+            block_size: int = mgr.block_size  # type: ignore
+            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)
+        torch_cache = TorchKVCache.from_vllm_cache(
+            model_runner.kv_caches,  # type: ignore
+            block_ids_per_group,
+            layer_to_group,
+            prompt_token_count,
+            token_offset_per_group,
+        )
+        prefix_cache.add_from_torch(prompt_token_ids, torch_cache)
+    except Exception:
+        logger.opt(exception=True).warning("Failed to save prefix cache")
+
+
+def _stop_token_ids(tokenizer: object, model_id: ModelId) -> set[int]:
+    ids: set[int] = set()
+    eos_id = getattr(tokenizer, "eos_token_id", None)
+    if eos_id is not None:
+        ids.add(eos_id)
+    extra = get_eos_token_ids_for_model(model_id)
+    if extra:
+        ids.update(extra)
+    return ids
+
+
+def _build_generation_response(
+    tokenizer: object,
+    token_id: int,
+    finish_reason: str | None,
+    prompt_token_count: int,
+    completion_tokens: int,
+    start_time: float,
+    first_token_time: float | None,
+    suppress_text: bool = False,
+) -> GenerationResponse:
+    token_text: str = "" if suppress_text else tokenizer.decode([token_id])  # type: ignore[reportUnknownMemberType]
+    finish_usage: Usage | None = None
+    finish_stats: GenerationStats | None = None
+    mapped_finish_reason: str | None = None
+    if finish_reason:
+        now = time.perf_counter()
+        prefill_elapsed = (first_token_time or now) - start_time
+        decode_elapsed = now - (first_token_time or now)
+        finish_usage = Usage(
+            prompt_tokens=prompt_token_count,
+            completion_tokens=completion_tokens,
+            total_tokens=prompt_token_count + completion_tokens,
+            prompt_tokens_details=PromptTokensDetails(),
+            completion_tokens_details=CompletionTokensDetails(),
+        )
+        finish_stats = GenerationStats(
+            prompt_tps=prompt_token_count / prefill_elapsed
+            if prefill_elapsed > 0
+            else 0.0,
+            generation_tps=completion_tokens / decode_elapsed
+            if decode_elapsed > 0
+            else 0.0,
+            prompt_tokens=prompt_token_count,
+            generation_tokens=completion_tokens,
+            peak_memory_usage=Memory.from_bytes(
+                torch.cuda.max_memory_allocated()  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType, reportAttributeAccessIssue]
+            ),
+        )
+        mapped_finish_reason = (
+            finish_reason
+            if finish_reason in ("stop", "length", "content_filter")
+            else "stop"
+        )
+    return GenerationResponse(
+        text=token_text,
+        token=token_id,
+        finish_reason=mapped_finish_reason,
+        usage=finish_usage,
+        stats=finish_stats,
+    )
+
+
+def vllm_generate(
+    engine: LLMEngine,
+    model_id: ModelId,
+    task: TextGenerationTaskParams,
+    prompt: str,
+    prefix_cache: KVPrefixCache,
+    on_prefill_progress: Callable[[int, int], None] | None = None,
+    distributed_prompt_progress_callback: Callable[[], None] | None = None,
+    on_generation_token: Callable[[], None] | None = None,
+) -> Generator[GenerationResponse, None, None]:
+    token_ids, prompt_text, prompt_token_count = format_vllm_prompt(engine, task)
+    logger.debug(prompt_text)
+    request_id = f"vllm-seq-{time.monotonic_ns()}"
+    sampling_params = make_vllm_sampling_params(engine, task, model_id)
+    engine.add_request(request_id, {"prompt_token_ids": token_ids}, sampling_params)
+
+    tokenizer = engine.get_tokenizer()
+    stop_ids = _stop_token_ids(tokenizer, model_id)
+    DEFAULT_PREFILL_STEP_SIZE = 8192
+    max_batch_tokens: int = (
+        getattr(engine.model_config, "max_num_batched_tokens", DEFAULT_PREFILL_STEP_SIZE) or DEFAULT_PREFILL_STEP_SIZE
+    )  # type: ignore[reportUnknownMemberType]
+    start_time = time.perf_counter()
+    first_token_time: float | None = None
+    prev_token_count = 0
+    prefill_done = False
+    prefill_steps = 0
+
+    while engine.has_unfinished_requests():
+        if distributed_prompt_progress_callback and not prefill_done:
+            distributed_prompt_progress_callback()
+        outputs = engine.step()
+
+        for output in outputs:
+            if output.request_id != request_id:
+                continue
+            completion = output.outputs[0]
+            new_token_count = len(completion.token_ids)
+            new_tokens = completion.token_ids[prev_token_count:]
+            finish_reason = completion.finish_reason
+            prev_token_count = new_token_count
+
+            if not prefill_done and not new_tokens:
+                prefill_steps += 1
+                if on_prefill_progress:
+                    on_prefill_progress(
+                        min(prefill_steps * max_batch_tokens, prompt_token_count),
+                        prompt_token_count,
+                    )
+                continue
+
+            if not prefill_done and new_tokens:
+                first_token_time = time.perf_counter()
+                prefill_done = True
+                _save_prefix_cache(
+                    engine, prefix_cache, request_id, token_ids, prompt_token_count
+                )
+
+            for i, token_id in enumerate(new_tokens):
+                is_last = i == len(new_tokens) - 1
+                is_final_stop = is_last and finish_reason and token_id in stop_ids
+                if on_generation_token:
+                    on_generation_token()
+                if is_final_stop:
+                    yield _build_generation_response(
+                        tokenizer,
+                        token_id,
+                        finish_reason,
+                        prompt_token_count,
+                        new_token_count,
+                        start_time,
+                        first_token_time,
+                        suppress_text=True,
+                    )
+                else:
+                    yield _build_generation_response(
+                        tokenizer,
+                        token_id,
+                        finish_reason if is_last and finish_reason else None,
+                        prompt_token_count,
+                        new_token_count,
+                        start_time,
+                        first_token_time,
+                    )
+
+
+def warmup_vllm_engine(engine: LLMEngine) -> int:
+    tokenizer = engine.get_tokenizer()
+    messages = [
+        {
+            "role": "user",
+            "content": "Prompt to warm up the inference engine. Repeat this.",
+        }
+    ]
+    prompt_text: str = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )  # type: ignore
+    token_ids: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False)  # type: ignore
+    params = SamplingParams(max_tokens=50, detokenize=False)
+    engine.add_request("warmup", {"prompt_token_ids": token_ids}, params)
+    t = time.monotonic()
+    tokens_generated = 0
+    while engine.has_unfinished_requests():
+        engine.step()
+        tokens_generated += 1
+    elapsed = max(time.monotonic() - t, 0.001)
+    check_for_cancel_every = min(math.ceil(tokens_generated / elapsed), 100)
+    logger.info(
+        f"vLLM warmup complete, check_for_cancel_every={check_for_cancel_every}"
+    )
+    return check_for_cancel_every
+
+
+@dataclass(eq=False)
+class VllmBatchEngine:
+    engine: LLMEngine
+    model_id: ModelId
+    prefix_cache: KVPrefixCache
+
+    _active: dict[TaskId, _EngineRequest] = field(default_factory=dict, init=False)
+
+    def warmup(self) -> int:
+        return warmup_vllm_engine(self.engine)
+
+    @property
+    def has_work(self) -> bool:
+        return bool(self._active) or self.engine.has_unfinished_requests()
+
+    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,
+    ) -> TaskId:
+        token_ids, prompt_text, prompt_token_count = format_vllm_prompt(
+            self.engine, task_params
+        )
+        logger.debug(prompt_text)
+        sampling_params = make_vllm_sampling_params(
+            self.engine, task_params, self.model_id
+        )
+        self.engine.add_request(
+            task_id, {"prompt_token_ids": token_ids}, sampling_params
+        )
+        self._active[task_id] = _EngineRequest(
+            request_id=task_id,
+            prompt_token_count=prompt_token_count,
+            prompt_token_ids=token_ids,
+            on_generation_token=on_generation_token,
+            on_prefill_progress=on_prefill_progress,
+        )
+        return task_id
+
+    def step(self) -> list[tuple[TaskId, GenerationResponse]]:
+        if not self.has_work:
+            return []
+
+        outputs = self.engine.step()
+        tokenizer = self.engine.get_tokenizer()
+        stop_ids = _stop_token_ids(tokenizer, self.model_id)
+        max_batch_tokens: int = (
+            getattr(self.engine.model_config, "max_num_batched_tokens", 2048) or 2048
+        )  # type: ignore[reportUnknownMemberType]
+        results: list[tuple[TaskId, GenerationResponse]] = []
+
+        for output in outputs:
+            task_id = TaskId(output.request_id)
+            if task_id not in self._active:
+                continue
+            req = self._active[task_id]
+            completion = output.outputs[0]
+            new_token_count = len(completion.token_ids)
+            new_tokens = completion.token_ids[req.prev_token_count :]
+            finish_reason = completion.finish_reason
+            req.prev_token_count = new_token_count
+
+            if not req.prefill_done and not new_tokens:
+                req.prefill_steps += 1
+                if req.on_prefill_progress:
+                    req.on_prefill_progress(
+                        min(
+                            req.prefill_steps * max_batch_tokens, req.prompt_token_count
+                        ),
+                        req.prompt_token_count,
+                    )
+                continue
+
+            if not req.prefill_done and new_tokens:
+                req.first_token_time = time.perf_counter()
+                req.prefill_done = True
+                _save_prefix_cache(
+                    self.engine,
+                    self.prefix_cache,
+                    req.request_id,
+                    req.prompt_token_ids,
+                    req.prompt_token_count,
+                )
+
+            for i, token_id in enumerate(new_tokens):
+                is_last = i == len(new_tokens) - 1
+                is_final_stop = is_last and finish_reason and token_id in stop_ids
+                if req.on_generation_token:
+                    req.on_generation_token()
+                results.append(
+                    (
+                        task_id,
+                        _build_generation_response(
+                            tokenizer,
+                            token_id,
+                            finish_reason if is_last and finish_reason else None,
+                            req.prompt_token_count,
+                            new_token_count,
+                            req.start_time,
+                            req.first_token_time,
+                            suppress_text=bool(is_final_stop),
+                        ),
+                    )
+                )
+
+            if finish_reason:
+                del self._active[task_id]
+
+        for req in self._active.values():
+            if not req.prefill_done:
+                req.prefill_steps += 1
+                if req.on_prefill_progress:
+                    req.on_prefill_progress(
+                        min(
+                            req.prefill_steps * max_batch_tokens, req.prompt_token_count
+                        ),
+                        req.prompt_token_count,
+                    )
+
+        return results
+
+    def cancel(self, task_ids: list[TaskId]) -> None:
+        to_abort = [tid for tid in task_ids if tid in self._active]
+        if to_abort:
+            self.engine.abort_request(to_abort)
+        for tid in task_ids:
+            self._active.pop(tid, None)
+
+    def close(self) -> None:
+        if not hasattr(self, "engine"):
+            return
+        rids = [req.request_id for req in self._active.values()]
+        if rids:
+            self.engine.abort_request(rids)
+        self._active.clear()
+        del self.engine
+        gc.collect()
+        torch.cuda.empty_cache()
+        if torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group()
+
+
+_weight_loading_callback: Callable[[int, int], None] | None = None
+_weight_loading_patched = False
+
+
+def get_weight_loading_callback() -> Callable[[int, int], None] | None:
+    return _weight_loading_callback
+
+
+def set_weight_loading_callback(cb: Callable[[int, int], None] | None) -> None:
+    global _weight_loading_callback
+    _weight_loading_callback = cb
+
+
+_LAYER_INDEX_PATTERN = re.compile(r"\.layers\.(\d+)\.")
+_n_layers: int = 1
+
+
+def get_n_layers() -> int:
+    return _n_layers
+
+
+def set_n_layers(n: int) -> None:
+    global _n_layers
+    _n_layers = n
+
+
+def _wrap_weights_iterator(
+    original: Callable[..., Generator[tuple[str, "torch.Tensor"], None, None]],
+) -> Callable[..., Generator[tuple[str, "torch.Tensor"], None, None]]:  # pyright: ignore[reportUnknownParameterType]
+    def patched(
+        hf_weights_files: list[str], *args: object, **kwargs: object
+    ) -> Generator[tuple[str, "torch.Tensor"], None, None]:  # pyright: ignore[reportUnknownParameterType]
+        callback = get_weight_loading_callback()
+        if callback is not None and hf_weights_files:
+            total_layers = get_n_layers()
+            seen_layers: set[int] = set()
+            last_reported = 0
+            for name, tensor in original(hf_weights_files, *args, **kwargs):  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+                yield name, tensor  # pyright: ignore[reportUnknownArgumentType]
+                match = _LAYER_INDEX_PATTERN.search(name)
+                if match:
+                    seen_layers.add(int(match.group(1)))
+                current = len(seen_layers)
+                if current > last_reported:
+                    callback(current, total_layers)
+                    last_reported = current
+            callback(total_layers, total_layers)
+        else:
+            yield from original(hf_weights_files, *args, **kwargs)  # pyright: ignore[reportUnknownMemberType]
+
+    return patched
+
+
+def _monkey_patch_iterator(weight_utils: object, attr_name: str) -> None:  # pyright: ignore[reportUnknownParameterType]
+    original = getattr(weight_utils, attr_name, None)
+    if original is None:
+        return
+    patched = _wrap_weights_iterator(original)  # pyright: ignore[reportUnknownArgumentType]
+    setattr(weight_utils, attr_name, patched)
+    for mod in list(sys.modules.values()):
+        if mod is None or mod is weight_utils:
+            continue
+        for name in list(vars(mod)):
+            if vars(mod)[name] is original:
+                setattr(mod, name, patched)
+
+
+def _patch_weight_loading_progress() -> None:
+    global _weight_loading_patched
+    if _weight_loading_patched:
+        return
+    _weight_loading_patched = True
+
+    from vllm.model_executor.model_loader import (
+        weight_utils,  # pyright: ignore[reportMissingImports]
+    )
+
+    _monkey_patch_iterator(weight_utils, "safetensors_weights_iterator")
+    _monkey_patch_iterator(weight_utils, "fastsafetensors_weights_iterator")
+
+    import huggingface_hub  # pyright: ignore[reportMissingImports]
+
+    def _noop_metadata(*_a: object, **_kw: object) -> None:
+        pass  # pyright: ignore[reportUnknownParameterType]
+
+    original_metadata = huggingface_hub.get_safetensors_metadata  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+    huggingface_hub.get_safetensors_metadata = _noop_metadata  # pyright: ignore[reportAttributeAccessIssue]
+    for mod in list(sys.modules.values()):
+        if mod is None or mod is huggingface_hub:
+            continue
+        for attr in list(vars(mod)):
+            if vars(mod)[attr] is original_metadata:
+                setattr(mod, attr, _noop_metadata)
+
+
+def load_vllm_engine(
+    model_path: Path,
+    model_id: ModelId,
+    trust_remote_code: bool,
+    n_layers: int = 1,
+    on_layer_loaded: Callable[[int, int], None] | None = None,
+    kv_connector_cls: type[object] | None = None,
+) -> tuple[LLMEngine, ToolParser | None, KVPrefixCache]:
+    patch_vllm()
+    _patch_weight_loading_progress()
+
+    if kv_connector_cls is not None:
+        from exo.disaggregated.prefill_server import _patch_vllm_for_connector
+
+        _patch_vllm_for_connector(kv_connector_cls)
+
+    os.environ.setdefault("FASTSAFETENSORS_NOGDS", "1")
+
+    prefix_cache = KVPrefixCache(group=None)
+    set_prefix_cache(prefix_cache)
+    set_n_layers(n_layers)
+
+    kv_transfer_config: dict[str, str] | None = None
+    if kv_connector_cls is not None:
+        kv_transfer_config = {
+            "kv_connector": f"{kv_connector_cls.__module__}:{kv_connector_cls.__name__}",
+            "kv_role": "kv_both",
+        }
+
+    import json
+    from pathlib import Path
+
+    is_nvfp4 = "nvfp4" in model_path.lower() or "nvfp4" in str(model_id).lower()
+    has_mamba = False
+    is_mxfp4 = False
+    config_path = Path(model_path) / "config.json"
+    if config_path.exists():
+        with open(config_path) as f:
+            model_config = json.load(f)
+        text_config = model_config.get("text_config", model_config)
+        has_mamba = "mamba_ssm_dtype" in text_config or "linear_attention" in (text_config.get("layer_types") or [])
+        quant_config = model_config.get("quantization_config") or text_config.get("quantization_config")
+        if quant_config and quant_config.get("quant_method") == "mxfp4":
+            is_mxfp4 = True
+    if has_mamba:
+        backends = ["FLASH_ATTN", "TRITON_ATTN"]
+    else:
+        backends = ["FLASHINFER", "FLASH_ATTN", "TRITON_ATTN"]
+
+    engine: LLMEngine | None = None
+    for backend in backends:
+        try:
+            engine_args = EngineArgs(
+                model=model_path,
+                served_model_name=str(model_id),
+                gpu_memory_utilization=0.05,
+                trust_remote_code=trust_remote_code,
+                load_format="fastsafetensors",
+                enable_prefix_caching=False,
+                attention_backend=backend,
+                compilation_config={"cudagraph_mode": "none"},
+                disable_log_stats=True,
+                max_num_batched_tokens=4096,
+                kv_transfer_config=kv_transfer_config,  # type: ignore
+                disable_hybrid_kv_cache_manager=False,
+            )
+
+            set_weight_loading_callback(on_layer_loaded)
+            engine = LLMEngine.from_engine_args(engine_args)
+            logger.info(f"vLLM engine using attention backend: {backend}")
+            break
+        except (ValueError, RuntimeError, NotImplementedError) as e:
+            logger.warning(f"Attention backend {backend} failed: {e}, trying next")
+            engine = None
+            import gc
+            gc.collect()
+            torch.cuda.empty_cache()
+            continue
+
+    if engine is None:
+        raise RuntimeError(f"No attention backend worked for {model_id}")
+
+    tool_parser: ToolParser | None = None
+    tokenizer = engine.get_tokenizer()
+    chat_template = getattr(tokenizer, "chat_template", None)
+    if isinstance(chat_template, str):
+        tool_parser = infer_tool_parser(chat_template)
+        if tool_parser:
+            logger.info(
+                f"inferred tool parser: {tool_parser.start_parsing} / {tool_parser.end_parsing}"
+            )
+
+    logger.info(f"vLLM engine loaded for {model_id}")
+
+    return engine, tool_parser, prefix_cache

src/exo/worker/main.py

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

src/exo/worker/runner/bootstrap.py

@@ -1,17 +1,75 @@
+import ctypes
 import os
 import resource
+import sys
+import urllib.request
+from pathlib import Path
 
 import loguru
 
 from exo.shared.types.events import Event, RunnerStatusUpdated
 from exo.shared.types.tasks import Task, TaskId
-from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.instances import BoundInstance, VllmInstance
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
 from exo.worker.engines.mlx.patches import apply_mlx_patches
 
 logger: "loguru.Logger" = loguru.logger
 
+_TIKTOKEN_BASE_URL = "https://openaipublic.blob.core.windows.net/encodings"
+_TIKTOKEN_FILES = ["o200k_base.tiktoken", "cl100k_base.tiktoken"]
+
+_CUDA_HOST_LIBS = ["libcuda.so.1", "libnvidia-ml.so.1", "libnvidia-ptxjitcompiler.so.1"]
+_CUDA_HOST_SEARCH_DIRS = [
+    Path("/usr/lib/aarch64-linux-gnu"),
+    Path("/usr/lib/x86_64-linux-gnu"),
+    Path("/usr/lib64"),
+    Path("/usr/lib"),
+    Path("/usr/local/cuda/lib64"),
+    Path("/usr/local/cuda/compat"),
+]
+
+
+def _ensure_tiktoken_encodings() -> None:
+    if os.environ.get("TIKTOKEN_ENCODINGS_BASE"):
+        return
+    from exo.shared.constants import EXO_CACHE_HOME
+
+    enc_dir = EXO_CACHE_HOME / "encodings"
+    enc_dir.mkdir(parents=True, exist_ok=True)
+    for fname in _TIKTOKEN_FILES:
+        dest = enc_dir / fname
+        if dest.exists():
+            continue
+        url = f"{_TIKTOKEN_BASE_URL}/{fname}"
+        logger.info(f"Downloading {url} -> {dest}")
+        try:
+            urllib.request.urlretrieve(url, dest)
+        except Exception:
+            logger.warning(f"Failed to download {fname}, harmony encoding may fail")
+            return
+    os.environ["TIKTOKEN_ENCODINGS_BASE"] = str(enc_dir)
+    logger.info(f"Set TIKTOKEN_ENCODINGS_BASE={enc_dir}")
+
+
+def _ensure_cuda_libs() -> None:
+    if sys.platform != "linux":
+        return
+    for search_dir in _CUDA_HOST_SEARCH_DIRS:
+        driver = search_dir / "libcuda.so.1"
+        if not driver.exists():
+            continue
+        for lib_name in _CUDA_HOST_LIBS:
+            lib_path = search_dir / lib_name
+            if lib_path.exists():
+                try:
+                    ctypes.CDLL(str(lib_path), mode=ctypes.RTLD_GLOBAL)
+                    logger.info(f"Loaded CUDA host lib: {lib_path}")
+                except OSError:
+                    logger.warning(f"Failed to load {lib_path}")
+                    raise
+        return
+
 
 def entrypoint(
     bound_instance: BoundInstance,
@@ -36,7 +94,26 @@ def entrypoint(
 
     # Import main after setting global logger - this lets us just import logger from this module
     try:
-        if bound_instance.is_image_model:
+        if isinstance(bound_instance.instance, VllmInstance):
+            os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+            os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
+            if "/usr/local/cuda-13.2/compat" not in os.environ.get("LD_LIBRARY_PATH", ""):
+                _ensure_cuda_libs()
+            _ensure_tiktoken_encodings()
+            from exo.worker.runner.llm_inference.runner import Runner, VllmBuilder
+
+            model_id = bound_instance.bound_shard.model_card.model_id
+            builder = VllmBuilder(
+                model_id=model_id,
+                trust_remote_code=bound_instance.bound_shard.model_card.trust_remote_code,
+                cancel_receiver=cancel_receiver,
+                event_sender=event_sender,
+            )
+            runner = Runner(
+                bound_instance, event_sender, task_receiver, cancel_receiver, builder
+            )
+            runner.main()
+        elif bound_instance.is_image_model:
             from exo.worker.runner.image_models.runner import Runner as ImageRunner
 
             runner = ImageRunner(
@@ -44,12 +121,17 @@ def entrypoint(
             )
             runner.main()
         else:
-            from exo.worker.runner.llm_inference.runner import Runner
-
+            from exo.worker.runner.llm_inference.runner import MlxBuilder, Runner
             apply_mlx_patches()
 
+            builder = MlxBuilder(
+                model_id=bound_instance.bound_shard.model_card.model_id,
+                event_sender=event_sender,
+                cancel_receiver=cancel_receiver,
+            )
+
             runner = Runner(
-                bound_instance, event_sender, task_receiver, cancel_receiver
+                bound_instance, event_sender, task_receiver, cancel_receiver, builder
             )
             runner.main()
 

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

@@ -2,8 +2,9 @@ import itertools
 import time
 from abc import ABC, abstractmethod
 from collections import deque
-from collections.abc import Generator, Iterable
+from collections.abc import Callable, Generator, Iterable
 from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
 
 import mlx.core as mx
 from mlx_lm.tokenizer_utils import TokenizerWrapper
@@ -12,17 +13,17 @@ from exo.shared.constants import EXO_MAX_CONCURRENT_REQUESTS
 from exo.shared.types.chunks import ErrorChunk, PrefillProgressChunk
 from exo.shared.types.common import ModelId
 from exo.shared.types.events import ChunkGenerated, Event
-from exo.shared.types.mlx import Model
 from exo.shared.types.tasks import CANCEL_ALL_TASKS, TaskId, TextGeneration
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
 from exo.utils.channels import MpReceiver, MpSender
 from exo.worker.engines.mlx.cache import KVPrefixCache
 from exo.worker.engines.mlx.generator.batch_generate import ExoBatchGenerator
+
+if TYPE_CHECKING:
+    from exo.worker.engines.vllm.vllm_generator import VllmBatchEngine
 from exo.worker.engines.mlx.generator.generate import (
     PrefillCancelled,
-    mlx_generate,
-    warmup_inference,
 )
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
@@ -112,7 +113,6 @@ def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None:
 
 @dataclass(eq=False)
 class SequentialGenerator(InferenceGenerator):
-    model: Model
     tokenizer: TokenizerWrapper
     group: mx.distributed.Group | None
     kv_prefix_cache: KVPrefixCache | None
@@ -121,6 +121,8 @@ class SequentialGenerator(InferenceGenerator):
     device_rank: int
     cancel_receiver: MpReceiver[TaskId]
     event_sender: MpSender[Event]
+    _generate_fn: Callable[..., Generator[GenerationResponse]]
+    _warmup_fn: Callable[[], int]
     vision_processor: VisionProcessor | None = None
     check_for_cancel_every: int = 50
 
@@ -142,13 +144,8 @@ class SequentialGenerator(InferenceGenerator):
         | None
     ) = field(default=None, init=False)
 
-    def warmup(self):
-        self.check_for_cancel_every = warmup_inference(
-            model=self.model,
-            tokenizer=self.tokenizer,
-            group=self.group,
-            model_id=self.model_id,
-        )
+    def warmup(self) -> None:
+        self.check_for_cancel_every = self._warmup_fn()
 
     def submit(
         self,
@@ -240,7 +237,6 @@ class SequentialGenerator(InferenceGenerator):
                 apply_chat_template(self.tokenizer, task.task_params),
                 self.tool_parser,
                 self.tokenizer,
-                type(self.model),
                 self.model_id,
                 task.task_params.tools,
             )
@@ -296,9 +292,7 @@ class SequentialGenerator(InferenceGenerator):
 
                 self.agree_on_tasks()
 
-        return mlx_generate(
-            model=self.model,
-            tokenizer=self.tokenizer,
+        return self._generate_fn(
             task=task.task_params,
             prompt=prompt,
             kv_prefix_cache=self.kv_prefix_cache,
@@ -310,12 +304,14 @@ class SequentialGenerator(InferenceGenerator):
         )
 
     def close(self) -> None:
-        del self.model, self.tokenizer, self.group
+        if hasattr(self, "tokenizer"):
+            del self.tokenizer
+        if hasattr(self, "group"):
+            del self.group
 
 
 @dataclass(eq=False)
 class BatchGenerator(InferenceGenerator):
-    model: Model
     tokenizer: TokenizerWrapper
     group: mx.distributed.Group | None
     kv_prefix_cache: KVPrefixCache | None
@@ -324,6 +320,8 @@ class BatchGenerator(InferenceGenerator):
     device_rank: int
     cancel_receiver: MpReceiver[TaskId]
     event_sender: MpSender[Event]
+    _gen: "ExoBatchGenerator | VllmBatchEngine"
+    max_concurrent_requests: int = EXO_MAX_CONCURRENT_REQUESTS
     check_for_cancel_every: int = 50
     vision_processor: VisionProcessor | None = None
 
@@ -332,9 +330,8 @@ class BatchGenerator(InferenceGenerator):
     _maybe_cancel: list[TextGeneration] = field(default_factory=list, init=False)
     _all_tasks: dict[TaskId, TextGeneration] = field(default_factory=dict, init=False)
     _queue: deque[TextGeneration] = field(default_factory=deque, init=False)
-    _mlx_gen: ExoBatchGenerator = field(init=False)
     _active_tasks: dict[
-        int,
+        TaskId,
         tuple[
             TextGeneration,
             GeneratorQueue[GenerationResponse],
@@ -342,23 +339,6 @@ class BatchGenerator(InferenceGenerator):
         ],
     ] = field(default_factory=dict, init=False)
 
-    def __post_init__(self) -> None:
-        self._mlx_gen = ExoBatchGenerator(
-            model=self.model,
-            tokenizer=self.tokenizer,
-            group=self.group,
-            kv_prefix_cache=self.kv_prefix_cache,
-            vision_processor=self.vision_processor,
-        )
-
-    def warmup(self):
-        self.check_for_cancel_every = warmup_inference(
-            model=self.model,
-            tokenizer=self.tokenizer,
-            group=self.group,
-            model_id=self.model_id,
-        )
-
     def submit(
         self,
         task: TextGeneration,
@@ -399,10 +379,10 @@ class BatchGenerator(InferenceGenerator):
             self.agree_on_tasks()
 
         # Submit any queued tasks to the engine
-        while self._queue and len(self._active_tasks) < EXO_MAX_CONCURRENT_REQUESTS:
+        while self._queue and len(self._active_tasks) < self.max_concurrent_requests:
             task = self._queue.popleft()
             try:
-                uid = self._start_task(task)
+                task_id = self._start_task(task)
             except PrefillCancelled:
                 continue
             except Exception as e:
@@ -418,16 +398,15 @@ class BatchGenerator(InferenceGenerator):
                     apply_chat_template(self.tokenizer, task.task_params),
                     self.tool_parser,
                     self.tokenizer,
-                    type(self.model),
                     self.model_id,
                     task.task_params.tools,
                 )
-            self._active_tasks[uid] = (task, queue, output_generator)
+            self._active_tasks[task_id] = (task, queue, output_generator)
 
-        if not self._mlx_gen.has_work:
+        if not self._gen.has_work:
             return self._apply_cancellations()
 
-        results = self._mlx_gen.step()
+        results = self._gen.step()
 
         output: list[
             tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
@@ -459,17 +438,17 @@ class BatchGenerator(InferenceGenerator):
 
         cancel_all = CANCEL_ALL_TASKS in self._cancelled_tasks
 
-        uids_to_cancel: list[int] = []
+        ids_to_cancel: list[TaskId] = []
         results: list[tuple[TaskId, Cancelled]] = []
 
-        for uid, (task, _, _) in list(self._active_tasks.items()):
+        for tid, (task, _, _) in list(self._active_tasks.items()):
             if task.task_id in self._cancelled_tasks or cancel_all:
-                uids_to_cancel.append(uid)
+                ids_to_cancel.append(tid)
                 results.append((task.task_id, Cancelled()))
-                del self._active_tasks[uid]
+                del self._active_tasks[tid]
 
-        if uids_to_cancel:
-            self._mlx_gen.cancel(uids_to_cancel)
+        if ids_to_cancel:
+            self._gen.cancel(ids_to_cancel)
 
         already_cancelled = {tid for tid, _ in results}
         for tid in self._cancelled_tasks:
@@ -492,7 +471,7 @@ class BatchGenerator(InferenceGenerator):
                 )
             )
 
-    def _start_task(self, task: TextGeneration) -> int:
+    def _start_task(self, task: TextGeneration) -> TaskId:
         _check_for_debug_prompts(task.task_params)
         prompt = apply_chat_template(self.tokenizer, task.task_params)
 
@@ -529,7 +508,8 @@ class BatchGenerator(InferenceGenerator):
 
                 self.agree_on_tasks()
 
-        return self._mlx_gen.submit(
+        return self._gen.submit(
+            task_id=task.task_id,
             task_params=task.task_params,
             prompt=prompt,
             on_prefill_progress=on_prefill_progress,
@@ -538,5 +518,8 @@ class BatchGenerator(InferenceGenerator):
         )
 
     def close(self) -> None:
-        self._mlx_gen.close()
-        del self.model, self.tokenizer, self.group
+        self._gen.close()
+        if hasattr(self, "tokenizer"):
+            del self.tokenizer
+        if hasattr(self, "group"):
+            del self.group

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

@@ -2,12 +2,10 @@ from collections.abc import Generator
 from functools import cache
 from typing import Any
 
-from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
-from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.tokenizer_utils import TokenizerWrapper
-from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
+from openai_harmony import (
     HarmonyEncodingName,
-    HarmonyError,  # pyright: ignore[reportUnknownVariableType]
+    HarmonyError,
     Role,
     StreamableParser,
     load_harmony_encoding,
@@ -15,7 +13,6 @@ from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
 
 from exo.api.types import ToolCallItem
 from exo.shared.types.common import ModelId
-from exo.shared.types.mlx import Model
 from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
 from exo.worker.engines.mlx.utils_mlx import (
     detect_thinking_prompt_suffix,
@@ -61,18 +58,14 @@ def apply_all_parsers(
     prompt: str,
     tool_parser: ToolParser | None,
     tokenizer: TokenizerWrapper,
-    model_type: type[Model],
     model_id: ModelId,
     tools: list[dict[str, Any]] | None,
 ) -> Generator[GenerationResponse | ToolCallResponse | None]:
     mlx_generator = receiver
 
-    if issubclass(model_type, GptOssModel):
+    if "gpt-oss" in model_id.short().lower():
         mlx_generator = parse_gpt_oss(mlx_generator)
-    elif (
-        issubclass(model_type, DeepseekV32Model)
-        and "deepseek" in model_id.normalize().lower()
-    ):
+    elif "deepseek-v3.2" in model_id.short().lower():
         mlx_generator = parse_deepseek_v32(mlx_generator)
     else:
         if tokenizer.has_thinking:
@@ -88,6 +81,9 @@ def apply_all_parsers(
 
     return count_reasoning_tokens(mlx_generator)
 
+_GPT_OSS_CHANNEL_TOKEN = 200005
+_GPT_OSS_MESSAGE_TOKEN = 200008
+
 
 def parse_gpt_oss(
     responses: Generator[GenerationResponse | None],
@@ -101,44 +97,42 @@ def parse_gpt_oss(
         if response is None:
             yield None
             continue
+
+        token_id = response.token
+
         try:
-            stream.process(response.token)
-        except HarmonyError:
-            logger.error("Encountered critical Harmony Error, returning early")
+            stream.process(token_id)
+        except HarmonyError as e:
+            logger.error(
+                f"HarmonyError on token_id={response.token} text={response.text!r}: {e}"
+            )
             return
 
         delta = stream.last_content_delta
         ch = stream.current_channel
         recipient = stream.current_recipient
 
-        # Debug: log every token with state
-        logger.debug(
-            f"parse_gpt_oss token={response.token} text={response.text!r} "
-            f"recipient={recipient!r} ch={ch!r} delta={delta!r} "
-            f"state={stream.state} current_tool={current_tool_name!r}"
+        effective_recipient = (
+            recipient
+            if (recipient is not None and recipient.startswith("functions."))
+            else None
         )
-
-        if recipient != current_tool_name:
+        if effective_recipient != current_tool_name:
             if current_tool_name is not None:
-                prefix = "functions."
-                if current_tool_name.startswith(prefix):
-                    current_tool_name = current_tool_name[len(prefix) :]
-                logger.info(
-                    f"parse_gpt_oss yielding tool call: name={current_tool_name!r}"
-                )
+                tool_name = current_tool_name.removeprefix("functions.")
+                logger.info(f"parse_gpt_oss yielding tool call: name={tool_name!r}")
                 yield ToolCallResponse(
                     tool_calls=[
                         ToolCallItem(
-                            name=current_tool_name,
+                            name=tool_name,
                             arguments="".join(tool_arg_parts).strip(),
                         )
                     ],
                     usage=response.usage,
                 )
                 tool_arg_parts = []
-            current_tool_name = recipient
+            current_tool_name = effective_recipient
 
-        # If inside a tool call, accumulate arguments
         if current_tool_name is not None:
             if delta:
                 tool_arg_parts.append(delta)
@@ -147,13 +141,17 @@ def parse_gpt_oss(
                 tool_arg_parts = []
             continue
 
+        is_suppressed = ch == "analysis" or (
+            recipient is not None and recipient.startswith("!")
+        )
+
         if delta:
             yield response.model_copy(
-                update={"text": delta, "is_thinking": ch == "analysis"}
+                update={"text": delta, "is_thinking": is_suppressed}
             )
 
         if response.finish_reason is not None:
-            yield response
+            yield response.model_copy(update={"text": ""})
 
 
 def parse_deepseek_v32(

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

@@ -1,12 +1,19 @@
+import contextlib
+import socketserver
+import gc
 import os
 import time
+from abc import ABC, abstractmethod
+from collections.abc import Callable
 from dataclasses import dataclass
 from enum import Enum
+from typing import TYPE_CHECKING
 
 import mlx.core as mx
 from anyio import WouldBlock
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
+from exo.download.download_utils import build_model_path
 from exo.shared.models.model_cards import ModelTask
 from exo.shared.types.chunks import (
     ErrorChunk,
@@ -68,12 +75,34 @@ from exo.worker.runner.llm_inference.batch_generator import (
 from .batch_generator import Cancelled, Finished
 from .tool_parsers import make_mlx_parser
 
+if TYPE_CHECKING:
+    pass
+
 
 class ExitCode(str, Enum):
     AllTasksComplete = "AllTasksComplete"
     Shutdown = "Shutdown"
 
 
+class Builder(ABC):
+    @abstractmethod
+    def connect(self, bound_instance: BoundInstance) -> None: ...
+
+    @abstractmethod
+    def load(
+        self,
+        bound_instance: BoundInstance,
+        on_timeout: Callable[[], None],
+        on_layer_loaded: Callable[[int, int], None],
+    ) -> None: ...
+
+    @abstractmethod
+    def build(self) -> InferenceGenerator: ...
+
+    @abstractmethod
+    def close(self) -> None: ...
+
+
 class Runner:
     def __init__(
         self,
@@ -81,6 +110,7 @@ class Runner:
         event_sender: MpSender[Event],
         task_receiver: MpReceiver[Task],
         cancel_receiver: MpReceiver[TaskId],
+        builder: Builder,
     ):
         self.event_sender = event_sender
         self.task_receiver = task_receiver
@@ -103,17 +133,14 @@ class Runner:
 
         self.setup_start_time = time.time()
 
-        self.generator: Builder | InferenceGenerator = Builder(
-            self.model_id,
-            self.event_sender,
-            self.cancel_receiver,
-        )
+        self.generator: Builder | InferenceGenerator = builder
 
         self.seen: set[TaskId] = set()
         self.active_tasks: dict[
             TaskId,
             TextGeneration,
         ] = {}
+        self.prefill_server_port: int | None = None
 
         logger.info("runner created")
         self.update_status(RunnerIdle())
@@ -135,15 +162,19 @@ class Runner:
         self.event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
     def main(self):
-        with self.task_receiver:
-            for task in self.task_receiver:
-                if task.task_id in self.seen:
-                    logger.warning("repeat task - potential error")
-                    continue
-                self.seen.add(task.task_id)
-                self.handle_first_task(task)
-                if isinstance(self.current_status, RunnerShutdown):
-                    break
+        try:
+            with self.task_receiver:
+                for task in self.task_receiver:
+                    if task.task_id in self.seen:
+                        logger.warning("repeat task - potential error")
+                        continue
+                    self.seen.add(task.task_id)
+                    self.handle_first_task(task)
+                    if isinstance(self.current_status, RunnerShutdown):
+                        break
+        finally:
+            if not isinstance(self.current_status, RunnerShutdown):
+                self.generator.close()
 
     def handle_first_task(self, task: Task):
         self.send_task_status(task.task_id, TaskStatus.Running)
@@ -155,22 +186,28 @@ class Runner:
                 self.update_status(RunnerConnecting())
                 self.acknowledge_task(task)
 
-                self.generator.group = initialize_mlx(self.bound_instance)
+                self.generator.connect(self.bound_instance)
 
                 self.send_task_status(task.task_id, TaskStatus.Complete)
                 self.update_status(RunnerConnected())
                 logger.info("runner connected")
 
             # we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
-            case LoadModel() if isinstance(self.generator, Builder) and (
+            case LoadModel() if (
                 (
-                    isinstance(self.current_status, RunnerConnected)
+                    isinstance(self.generator, MlxBuilder)
+                    and isinstance(self.current_status, RunnerConnected)
                     and self.generator.group is not None
                 )
                 or (
-                    isinstance(self.current_status, RunnerIdle)
+                    isinstance(self.generator, MlxBuilder)
+                    and isinstance(self.current_status, RunnerIdle)
                     and self.generator.group is None
                 )
+                or (
+                    isinstance(self.generator, VllmBuilder)
+                    and isinstance(self.current_status, RunnerIdle)
+                )
             ):
                 total_layers = (
                     self.shard_metadata.end_layer - self.shard_metadata.start_layer
@@ -196,18 +233,15 @@ class Runner:
                 assert (
                     ModelTask.TextGeneration in self.shard_metadata.model_card.tasks
                 ), f"Incorrect model task(s): {self.shard_metadata.model_card.tasks}"
-                (
-                    self.generator.inference_model,
-                    self.generator.tokenizer,
-                    self.generator.vision_processor,
-                ) = load_mlx_items(
+
+                self.generator.load(
                     self.bound_instance,
-                    self.generator.group,
                     on_timeout=on_model_load_timeout,
                     on_layer_loaded=on_layer_loaded,
                 )
-
+                builder_ref = self.generator
                 self.generator = self.generator.build()
+                self.prefill_server_port = getattr(builder_ref, "_prefill_server_port", None)
 
                 self.send_task_status(task.task_id, TaskStatus.Complete)
                 self.update_status(RunnerLoaded())
@@ -227,7 +261,7 @@ class Runner:
                 )
 
                 self.send_task_status(task.task_id, TaskStatus.Complete)
-                self.update_status(RunnerReady())
+                self.update_status(RunnerReady(prefill_server_port=self.prefill_server_port))
                 logger.info("runner ready")
 
             case TextGeneration() if isinstance(self.current_status, RunnerReady):
@@ -248,11 +282,7 @@ class Runner:
         logger.info("runner shutting down")
         self.update_status(RunnerShuttingDown())
         self.acknowledge_task(task)
-        if isinstance(self.generator, InferenceGenerator):
-            self.generator.close()
-        mx.clear_cache()
-        import gc
-
+        self.generator.close()
         gc.collect()
         self.send_task_status(task.task_id, TaskStatus.Complete)
         self.update_status(RunnerShutdown())
@@ -316,7 +346,7 @@ class Runner:
             except WouldBlock:
                 pass
 
-        self.update_status(RunnerReady())
+        self.update_status(RunnerReady(prefill_server_port=self.prefill_server_port))
         logger.info("runner ready")
 
         return ExitCode.AllTasksComplete
@@ -377,7 +407,7 @@ class Runner:
 
 
 @dataclass
-class Builder:
+class MlxBuilder(Builder):
     model_id: ModelId
     event_sender: MpSender[Event]
     cancel_receiver: MpReceiver[TaskId]
@@ -386,9 +416,23 @@ class Builder:
     group: mx.distributed.Group | None = None
     vision_processor: VisionProcessor | None = None
 
-    def build(
+    def connect(self, bound_instance: BoundInstance) -> None:
+        self.group = initialize_mlx(bound_instance)
+
+    def load(
         self,
-    ) -> InferenceGenerator:
+        bound_instance: BoundInstance,
+        on_timeout: Callable[[], None],
+        on_layer_loaded: Callable[[int, int], None],
+    ) -> None:
+        self.inference_model, self.tokenizer, self.vision_processor = load_mlx_items(
+            bound_instance,
+            self.group,
+            on_timeout=on_timeout,
+            on_layer_loaded=on_layer_loaded,
+        )
+
+    def build(self) -> InferenceGenerator:
         assert self.model_id
         assert self.inference_model
         assert self.tokenizer
@@ -412,11 +456,28 @@ class Builder:
 
         kv_prefix_cache = KVPrefixCache(self.group)
 
+        from functools import partial
+
+        from exo.worker.engines.mlx.generator.generate import (
+            mlx_generate,
+            warmup_inference,
+        )
+
         device_rank = 0 if self.group is None else self.group.rank()
+        generate_fn = partial(
+            mlx_generate, model=self.inference_model, tokenizer=self.tokenizer
+        )
+        warmup_fn = partial(
+            warmup_inference,
+            model=self.inference_model,
+            tokenizer=self.tokenizer,
+            group=self.group,
+            model_id=self.model_id,
+        )
+
         if os.environ.get("EXO_NO_BATCH"):
             logger.info("using SequentialGenerator (batching disabled)")
             return SequentialGenerator(
-                model=self.inference_model,
                 tokenizer=self.tokenizer,
                 group=self.group,
                 tool_parser=tool_parser,
@@ -426,10 +487,20 @@ class Builder:
                 cancel_receiver=self.cancel_receiver,
                 event_sender=self.event_sender,
                 vision_processor=vision_processor,
+                _generate_fn=generate_fn,
+                _warmup_fn=warmup_fn,
             )
+        from exo.worker.runner.llm_inference.batch_generator import ExoBatchGenerator
+
         logger.info("using BatchGenerator")
-        return BatchGenerator(
+        gen = ExoBatchGenerator(
             model=self.inference_model,
+            tokenizer=self.tokenizer,
+            group=self.group,
+            kv_prefix_cache=kv_prefix_cache,
+            model_id=self.model_id,
+        )
+        return BatchGenerator(
             tokenizer=self.tokenizer,
             group=self.group,
             tool_parser=tool_parser,
@@ -438,5 +509,114 @@ class Builder:
             device_rank=device_rank,
             cancel_receiver=self.cancel_receiver,
             event_sender=self.event_sender,
+            _gen=gen,
             vision_processor=vision_processor,
         )
+
+    def close(self):
+        with contextlib.suppress(NameError, AttributeError):
+            del self.inference_model, self.tokenizer
+
+
+@dataclass
+class VllmBuilder(Builder):
+    model_id: ModelId
+    trust_remote_code: bool
+    cancel_receiver: MpReceiver[TaskId]
+    event_sender: MpSender[Event]
+    group: mx.distributed.Group | None = None
+    _prefill_server: socketserver.ThreadingTCPServer | None = None
+    _prefill_server_port: int | None = None
+
+    def connect(self, bound_instance: BoundInstance) -> None:
+        raise NotImplementedError(
+            "Multiple node VLLM instances are not supported at the moment!"
+        )
+
+    def load(
+        self,
+        bound_instance: BoundInstance,
+        on_timeout: Callable[[], None],
+        on_layer_loaded: Callable[[int, int], None],
+    ) -> None:
+        from exo.worker.engines.vllm.vllm_generator import load_vllm_engine
+
+        kv_connector_cls: type[object] | None = None
+        overlapping = not os.environ.get("EXO_NO_OVERLAPPING_PREFILL_SENDS")
+        if overlapping:
+            from exo.disaggregated.streaming_connector import StreamingConnector
+            kv_connector_cls = StreamingConnector
+        else:
+            from exo.disaggregated.batch_connector import BatchConnector
+            kv_connector_cls = BatchConnector
+
+        self._bound_runner_id = bound_instance.bound_runner_id
+        self._engine, self._tool_parser, self._prefix_cache = load_vllm_engine(
+            model_path=build_model_path(self.model_id),
+            model_id=self.model_id,
+            trust_remote_code=self.trust_remote_code,
+            n_layers=bound_instance.bound_shard.model_card.n_layers,
+            on_layer_loaded=on_layer_loaded,
+            kv_connector_cls=kv_connector_cls,
+        )
+
+    def build(self) -> InferenceGenerator:
+        from exo.worker.engines.vllm.vllm_generator import VllmBatchEngine
+
+        gen = VllmBatchEngine(
+            engine=self._engine,
+            model_id=self.model_id,
+            prefix_cache=self._prefix_cache,
+        )
+        tokenizer = TokenizerWrapper(self._engine.get_tokenizer())
+        max_concurrent = 1 if os.environ.get("EXO_NO_BATCH") else 8
+
+        from exo.master.placement import random_ephemeral_port
+
+        prefill_port = random_ephemeral_port()
+        overlapping = not os.environ.get("EXO_NO_OVERLAPPING_PREFILL_SENDS")
+        try:
+            from exo.disaggregated.prefill_server import start_prefill_server
+            from exo.shared.types.events import RunnerStatusUpdated
+            from exo.shared.types.worker.runners import RunnerReady, RunnerRunning
+
+            runner_id = self._bound_runner_id
+
+            def _on_prefill_status(running: bool) -> None:
+                port = prefill_port
+                if running:
+                    self.event_sender.send(RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerRunning(prefill_server_port=port)))
+                else:
+                    self.event_sender.send(RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerReady(prefill_server_port=port)))
+
+            self._prefill_server = start_prefill_server(
+                engine=self._engine,
+                bind_address="0.0.0.0",
+                port=prefill_port,
+                overlapping=overlapping,
+                prefix_cache=self._prefix_cache,
+                on_status_change=_on_prefill_status,
+            )
+            self._prefill_server_port = prefill_port
+        except Exception as e:
+            logger.opt(exception=e).warning("Failed to start prefill server")
+
+        logger.info(f"using BatchGenerator (vLLM, max_concurrent={max_concurrent})")
+        return BatchGenerator(
+            tokenizer=tokenizer,
+            group=None,
+            tool_parser=self._tool_parser,
+            kv_prefix_cache=None,
+            model_id=self.model_id,
+            device_rank=0,
+            cancel_receiver=self.cancel_receiver,
+            event_sender=self.event_sender,
+            _gen=gen,
+            max_concurrent_requests=max_concurrent,
+        )
+
+    def close(self) -> None:
+        with contextlib.suppress(NameError, AttributeError):
+            if hasattr(self, "_prefill_server") and self._prefill_server is not None:
+                self._prefill_server.shutdown()
+            del self._engine, self._prefix_cache, self._tool_parser

src/exo/worker/runner/runner_supervisor.py

@@ -1,5 +1,6 @@
 import contextlib
 import multiprocessing as mp
+import os
 import signal
 from dataclasses import dataclass, field
 from typing import Self
@@ -81,6 +82,21 @@ class RunnerSupervisor:
         task_sender, task_recv = mp_channel[Task]()
         cancel_sender, cancel_recv = mp_channel[TaskId]()
 
+        from exo.shared.types.worker.instances import VllmInstance
+
+        if isinstance(bound_instance.instance, VllmInstance):
+            if os.path.isdir("/usr/local/cuda-13.2/compat"):
+                os.environ["LD_LIBRARY_PATH"] = "/usr/local/cuda-13.2/compat:" + os.environ.get("LD_LIBRARY_PATH", "")
+            cublas_preload: list[str] = []
+            for lib in ["libcublasLt.so.13", "libcublas.so.13"]:
+                path = f"/usr/local/cuda-13.2/lib64/{lib}"
+                if os.path.exists(path):
+                    cublas_preload.append(path)
+            if cublas_preload:
+                existing = os.environ.get("LD_PRELOAD", "")
+                os.environ["LD_PRELOAD"] = ":".join(cublas_preload) + (":" + existing if existing else "")
+            os.environ["FLASHINFER_DISABLE_VERSION_CHECK"] = "1"
+        # uses spawn
         runner_process = mp.Process(
             target=entrypoint,
             args=(

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

@@ -145,8 +145,10 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
 
 class FakeExoBatchGenerator:
     def __init__(self, *_args: object, **_kwargs: object) -> None:
-        self._uid_counter = 0
-        self._pending: dict[int, GenerationResponse] = {}
+        self._pending: dict[str, GenerationResponse] = {}
+
+    def warmup(self) -> int:
+        return 50
 
     @property
     def has_work(self) -> bool:
@@ -154,30 +156,29 @@ class FakeExoBatchGenerator:
 
     def submit(
         self,
+        task_id: str = "",
         task_params: object = None,
         prompt: object = None,
         on_prefill_progress: object = None,
         distributed_prompt_progress_callback: object = None,
         on_generation_token: object = None,
-    ) -> int:
-        uid = self._uid_counter
-        self._uid_counter += 1
-        self._pending[uid] = GenerationResponse(
+    ) -> str:
+        self._pending[task_id] = GenerationResponse(
             text="hi",
             token=0,
             finish_reason="stop",
             usage=None,
         )
-        return uid
+        return task_id
 
-    def step(self) -> list[tuple[int, GenerationResponse]]:
+    def step(self) -> list[tuple[str, GenerationResponse]]:
         results = list(self._pending.items())
         self._pending.clear()
         return results
 
-    def cancel(self, uids: list[int]) -> None:
-        for uid in uids:
-            self._pending.pop(uid, None)
+    def cancel(self, task_ids: list[str]) -> None:
+        for tid in task_ids:
+            self._pending.pop(tid, None)
 
     def close(self) -> None:
         pass
@@ -268,11 +269,17 @@ def _run(tasks: Iterable[Task], send_after_ready: list[Task] | None = None):
             "exo.worker.runner.llm_inference.runner.mx.distributed.all_gather",
             make_nothin(mx.array([1])),
         ):
+            builder = mlx_runner.MlxBuilder(
+                model_id=MODEL_A_ID,
+                event_sender=event_sender,  # pyright: ignore[reportArgumentType]
+                cancel_receiver=cancel_receiver,
+            )
             runner = mlx_runner.Runner(
                 bound_instance,
                 event_sender,  # pyright: ignore[reportArgumentType]
                 task_receiver,
                 cancel_receiver,
+                builder,
             )
             runner.main()
 

test_growable_compile.py

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