bump versions correctly

.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

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

bench/exo_bench.py

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

bench/exo_eval.py

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

bench/harness.py

@@ -196,6 +196,31 @@ def resolve_model_short_id(
     raise ValueError(f"Model not found in /models: {model_arg}")
 
 
+def validate_vllm_args(args: argparse.Namespace) -> None:
+    if args.instance_meta != "vllm":
+        return
+    if args.sharding == "tensor":
+        raise SystemExit(
+            "--instance-meta vllm is incompatible with --sharding tensor (vllm is pipeline-only)"
+        )
+    if args.min_nodes > 1:
+        raise SystemExit(
+            "--instance-meta vllm is incompatible with --min-nodes > 1 (vllm is single-node)"
+        )
+    if args.max_nodes > 1:
+        raise SystemExit(
+            "--instance-meta vllm is incompatible with --max-nodes > 1 (vllm is single-node)"
+        )
+
+
+def ensure_cuda_available(client: ExoClient) -> None:
+    capabilities = client.request_json("GET", "/capabilities")
+    if not capabilities or not capabilities.get("vllm_available"):
+        raise SystemExit(
+            "--ensure-cuda: vllm is not available on the exo cluster (no CUDA capability)"
+        )
+
+
 def placement_filter(instance_meta: str, wanted: str) -> bool:
     s = (instance_meta or "").lower()
     if wanted == "both":
@@ -475,7 +500,7 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
         help="Only consider placements using >= this many nodes.",
     )
     ap.add_argument(
-        "--instance-meta", choices=["ring", "jaccl", "both"], default="both"
+        "--instance-meta", choices=["ring", "jaccl", "vllm", "both"], default="both"
     )
     ap.add_argument(
         "--sharding", choices=["pipeline", "tensor", "both"], default="both"
@@ -504,3 +529,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,8 @@
         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 +168,11 @@
   /** For a given base_model name, find the biggest quant variant that fits in memory. */
   function pickBestVariant(baseModel: string): ChatModelInfo | null {
     const variants = models
-      .filter((m) => m.base_model === baseModel && fitsInMemory(m))
+      .filter(
+        (m) =>
+          normalizeBaseModel(m.base_model) === normalizeBaseModel(baseModel) &&
+          fitsInMemory(m),
+      )
       .sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
     return variants[0] ?? null;
   }

dashboard/src/lib/components/DeviceIcon.svelte

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

dashboard/src/lib/components/ModelCard.svelte

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

dashboard/src/lib/components/ModelPickerModal.svelte

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

dashboard/src/lib/components/TopologyGraph.svelte

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

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

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

dashboard/src/routes/+page.svelte

@@ -64,6 +64,7 @@
     nodeThunderboltBridge,
     nodeIdentities,
     isConnected,
+    vllmAvailable,
     type DownloadProgress,
     type PlacementPreview,
   } from "$lib/stores/app.svelte";
@@ -294,7 +295,9 @@
     const seen = new Set<string>();
     const deduped: typeof candidates = [];
     for (const m of candidates) {
-      const key = m.base_model || m.family || m.id;
+      const key = (m.base_model || m.family || m.id)
+        .toLowerCase()
+        .replace(/[-_]/g, " ");
       if (seen.has(key)) continue;
       seen.add(key);
       deduped.push(m);
@@ -700,7 +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`,
@@ -884,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";
@@ -930,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 (
@@ -1144,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 {
@@ -2091,6 +2099,7 @@
     let instanceType = "Unknown";
     if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
     else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
+    else if (instanceTag === "VllmInstance") instanceType = "vLLM (CUDA)";
 
     const inst = instance as {
       shardAssignments?: {
@@ -5840,6 +5849,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.lock

@@ -2,11 +2,11 @@
   "nodes": {
     "crane": {
       "locked": {
-        "lastModified": 1767744144,
-        "narHash": "sha256-9/9ntI0D+HbN4G0TrK3KmHbTvwgswz7p8IEJsWyef8Q=",
+        "lastModified": 1774313767,
+        "narHash": "sha256-hy0XTQND6avzGEUFrJtYBBpFa/POiiaGBr2vpU6Y9tY=",
         "owner": "ipetkov",
         "repo": "crane",
-        "rev": "2fb033290bf6b23f226d4c8b32f7f7a16b043d7e",
+        "rev": "3d9df76e29656c679c744968b17fbaf28f0e923d",
         "type": "github"
       },
       "original": {
@@ -47,11 +47,11 @@
         "rust-analyzer-src": "rust-analyzer-src"
       },
       "locked": {
-        "lastModified": 1768287139,
-        "narHash": "sha256-nsXFt0OzUi6K7dUzzJD5/v9e0Ic+fvclfIW936/43ZM=",
+        "lastModified": 1774596377,
+        "narHash": "sha256-DiSLMxyTwIUAlhOe34r6kKNQRv6PTF+vf0MG45mAyn4=",
         "owner": "nix-community",
         "repo": "fenix",
-        "rev": "a4a3aa956931f90f35453cb519e4545e9ad7f773",
+        "rev": "a88a1c8cf2f094da6347fcec54089f4bcb518409",
         "type": "github"
       },
       "original": {
@@ -83,11 +83,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1768135262,
-        "narHash": "sha256-PVvu7OqHBGWN16zSi6tEmPwwHQ4rLPU9Plvs8/1TUBY=",
+        "lastModified": 1772408722,
+        "narHash": "sha256-rHuJtdcOjK7rAHpHphUb1iCvgkU3GpfvicLMwwnfMT0=",
         "owner": "hercules-ci",
         "repo": "flake-parts",
-        "rev": "80daad04eddbbf5a4d883996a73f3f542fa437ac",
+        "rev": "f20dc5d9b8027381c474144ecabc9034d6a839a3",
         "type": "github"
       },
       "original": {
@@ -164,11 +164,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1763662255,
-        "narHash": "sha256-4bocaOyLa3AfiS8KrWjZQYu+IAta05u3gYZzZ6zXbT0=",
+        "lastModified": 1773870109,
+        "narHash": "sha256-ZoTdqZP03DcdoyxvpFHCAek4bkPUTUPUF3oCCgc3dP4=",
         "owner": "pyproject-nix",
         "repo": "build-system-pkgs",
-        "rev": "042904167604c681a090c07eb6967b4dd4dae88c",
+        "rev": "b6e74f433b02fa4b8a7965ee24680f4867e2926f",
         "type": "github"
       },
       "original": {
@@ -184,11 +184,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1764134915,
-        "narHash": "sha256-xaKvtPx6YAnA3HQVp5LwyYG1MaN4LLehpQI8xEdBvBY=",
+        "lastModified": 1774498001,
+        "narHash": "sha256-wTfdyzzrmpuqt4TQQNqilF91v0m5Mh1stNy9h7a/WK4=",
         "owner": "pyproject-nix",
         "repo": "pyproject.nix",
-        "rev": "2c8df1383b32e5443c921f61224b198a2282a657",
+        "rev": "794afa6eb588b498344f2eaa36ab1ceb7e6b0b09",
         "type": "github"
       },
       "original": {
@@ -214,11 +214,11 @@
     "rust-analyzer-src": {
       "flake": false,
       "locked": {
-        "lastModified": 1768224240,
-        "narHash": "sha256-Pp1dDrXKPBUJReZnnDElFyHYn67XTd48zRhToheLjtk=",
+        "lastModified": 1774569884,
+        "narHash": "sha256-E8iWEPzg7OnE0XXXjo75CX7xFauqzJuGZ5wSO9KS8Ek=",
         "owner": "rust-lang",
         "repo": "rust-analyzer",
-        "rev": "725349602e525df37f377701e001fe8aab807878",
+        "rev": "443ddcddd0c73b07b799d052f5ef3b448c2f3508",
         "type": "github"
       },
       "original": {
@@ -257,11 +257,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1768158989,
-        "narHash": "sha256-67vyT1+xClLldnumAzCTBvU0jLZ1YBcf4vANRWP3+Ak=",
+        "lastModified": 1773297127,
+        "narHash": "sha256-6E/yhXP7Oy/NbXtf1ktzmU8SdVqJQ09HC/48ebEGBpk=",
         "owner": "numtide",
         "repo": "treefmt-nix",
-        "rev": "e96d59dff5c0d7fddb9d113ba108f03c3ef99eca",
+        "rev": "71b125cd05fbfd78cab3e070b73544abe24c5016",
         "type": "github"
       },
       "original": {
@@ -280,11 +280,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1767701098,
-        "narHash": "sha256-CJhKZnWb3gumR9oTRjFvCg/6lYTGbZRU7xtvcyWIRwU=",
+        "lastModified": 1774490495,
+        "narHash": "sha256-a9WmQWj8fF7BctZGCoyzpUjP6GJw8H+lxl+zxpGnETk=",
         "owner": "pyproject-nix",
         "repo": "uv2nix",
-        "rev": "9d357f0d2ce6f5f35ec7959d7e704452352eb4da",
+        "rev": "18ae62fc5e389e3069854a7c66455c22e31708fc",
         "type": "github"
       },
       "original": {

flake.nix

@@ -51,8 +51,8 @@
   };
 
   nixConfig = {
-    extra-trusted-public-keys = "exo.cachix.org-1:okq7hl624TBeAR3kV+g39dUFSiaZgLRkLsFBCuJ2NZI=";
-    extra-substituters = "https://exo.cachix.org";
+    extra-trusted-public-keys = "exo.cachix.org-1:okq7hl624TBeAR3kV+g39dUFSiaZgLRkLsFBCuJ2NZI= cache.nixos-cuda.org:74DUi4Ye579gUqzH4ziL9IyiJBlDpMRn9MBN8oNan9M=";
+    extra-substituters = "https://exo.cachix.org https://cache.nixos-cuda.org";
   };
 
   outputs =
@@ -72,17 +72,26 @@
       ];
 
       perSystem =
-        { config, self', inputs', pkgs, lib, system, ... }:
+        { config, self', pkgs, lib, system, ... }:
         let
           # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
           pkgsSwift = import inputs.nixpkgs-swift { inherit system; };
         in
         {
+          _module.args.cudaPkgs = import inputs.nixpkgs
+            {
+              inherit system;
+              config = {
+                allowUnfreePredicate = pkg: builtins.elem (lib.getName pkg) [ "cuda-merged" "cuda_cuobjdump" "cuda_gdb" "cuda_nvcc" "cuda_nvdisasm" "cuda_nvprune" "cuda_cccl" "cuda_cudart" "cuda_cupti" "cuda_cuxxfilt" "cuda_nvml_dev" "cuda_nvrtc" "cuda_nvtx" "cuda_profiler_api" "cuda_sanitizer_api" "libcublas" "libcufft" "libcurand" "libcusolver" "libnvjitlink" "libcusparse" "libnpp" "cudnn" "libcusparse_lt" "libcufile" "libnvshmem" "libnvvm" "cuda_crt" ];
+                cudaSupport = true;
+                cudaCapabilities = [ "12.1" ];
+              };
+            };
           # Allow unfree for metal-toolchain (needed for Darwin Metal packages)
           _module.args.pkgs = import inputs.nixpkgs {
             inherit system;
-            config.allowUnfreePredicate = pkg: (pkg.pname or "") == "metal-toolchain";
-            overlays = [
+            config.allowUnfreePredicate = pkg: builtins.elem (lib.getName pkg) [ "metal-toolchain" ];
+            overlays = lib.optionals (system == "aarch64-darwin") [
               (import ./nix/apple-sdk-overlay.nix)
             ];
           };
@@ -112,67 +121,68 @@
             };
           };
 
-          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
+          packages = (lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin
+
             {
               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;
             }
-          );
-
-          devShells.default = with pkgs; pkgs.mkShell {
-            inputsFrom = [ self'.checks.cargo-build ];
-
-            packages =
-              [
-                # FORMATTING
-                config.treefmt.build.wrapper
-
-                # PYTHON
-                python313
-                uv
-                ruff
-                basedpyright
-
-                # RUST
-                config.rust.toolchain
-                maturin
-
-                # NIX
-                nixpkgs-fmt
-
-                # SVELTE
-                nodejs
-
-                # MISC
-                just
-                jq
-              ]
-              ++ lib.optionals stdenv.isLinux [
-                unixtools.ifconfig
-              ]
-              ++ lib.optionals stdenv.isDarwin [
-                macmon
-              ];
-
-            OPENSSL_NO_VENDOR = "1";
-
-            shellHook = ''
-              export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib"
-              ${lib.optionalString stdenv.isLinux ''
-                export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH"
-              ''}
-            '';
+          ) // {
+            default = self'.packages.exo;
           };
+          devShells =
+            {
+              default = with pkgs; mkShell {
+                inputsFrom = [ self'.checks.cargo-build ];
+
+                packages =
+                  [
+                    # FORMATTING
+                    config.treefmt.build.wrapper
+
+                    # PYTHON
+                    self'.packages.editable-venv
+                    uv
+
+                    # RUST
+                    config.rust.toolchain
+                    maturin
+
+                    # NIX
+                    nixpkgs-fmt
+
+                    # SVELTE
+                    nodejs
+
+                    # MISC
+                    just
+                    jq
+                  ]
+                  ++ lib.optionals stdenv.isLinux [
+                    unixtools.ifconfig
+                  ]
+                  ++ lib.optionals stdenv.isDarwin [
+                    macmon
+                  ];
+
+                env = {
+                  UV_NO_SYNC = "1";
+                  UV_PYTHON = "${self'.packages.editable-venv}/bin/python";
+                  UV_PYTHON_DOWNLOADS = "never";
+                  UV_PROJECT_ENVIRONMENT = self'.packages.editable-venv;
+                  VIRTUAL_ENV = self'.packages.editable-venv;
+                  OPENSSL_NO_VENDOR = "1";
+                };
+
+                shellHook = ''
+                  unset PYTHONPATH
+                  export REPO_ROOT=$(git rev-parse --show-toplevel)
+                  export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${self'.packages.editable-venv}/lib"
+                  ${lib.optionalString stdenv.isLinux ''
+                    export LD_LIBRARY_PATH="${openssl.out}/lib:${lib.getLib pkgs.util-linux}/lib:${lib.getLib pkgs.systemd}/lib:${lib.getLib pkgs.numactl}/lib:${lib.getLib pkgs.stdenv.cc.cc.lib}/lib:$LD_LIBRARY_PATH"
+                  ''}
+                '';
+              };
+            };
         };
     };
 }

justfile

@@ -15,6 +15,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

nix/mlx_patch_fmod.patch

@@ -0,0 +1,20 @@
+diff --git a/mlx/backend/cuda/device/binary_ops.cuh b/mlx/backend/cuda/device/binary_ops.cuh
+index b0b79628..bbc377be 100644
+--- a/mlx/backend/cuda/device/binary_ops.cuh
++++ b/mlx/backend/cuda/device/binary_ops.cuh
+@@ -46,6 +46,15 @@ struct Remainder {
+       }
+     } else if constexpr (is_complex_v<T>) {
+       return x % y;
++    } else if constexpr (
++        cuda::std::is_same_v<T, __half> ||
++        cuda::std::is_same_v<T, __nv_bfloat16>) {
++      float yf = static_cast<float>(y);
++      float r = cuda::std::fmod(static_cast<float>(x), y);
++      if (r != 0.0f && ((r < 0.0f) != (yf < 0.0f))) {
++        r += yf;
++      }
++      return static_cast<T>(r);
+     } else {
+       T r = cuda::std::fmod(x, y);
+       if (r != 0 && (r < 0 != y < 0)) {

nix/vllm_uv2nix_cmake.patch

@@ -0,0 +1,12 @@
+diff --git a/setup.py b/setup.py
+index 68861fe4b..fd4738089 100644
+--- a/setup.py
++++ b/setup.py
+@@ -150,6 +150,7 @@ class cmake_build_ext(build_ext):
+         cmake_args = [
+             "-DCMAKE_BUILD_TYPE={}".format(cfg),
+             "-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE),
++            *json.loads(os.environ.get("UV2NIX_CMAKE_FLAGS_JSON", "[]"))
+         ]
+ 
+         verbose = envs.VERBOSE

nix/xgrammar_cmake.patch

@@ -0,0 +1,25 @@
+diff --git a/cpp/nanobind/CMakeLists.txt b/cpp/nanobind/CMakeLists.txt
+index ebfd3da..6ef06e3 100644
+--- a/cpp/nanobind/CMakeLists.txt
++++ b/cpp/nanobind/CMakeLists.txt
+@@ -18,8 +18,18 @@ target_sources(python_methods PRIVATE python_methods.cc)
+ target_link_libraries(python_methods PUBLIC xgrammar)
+ 
+ # Any code that uses nanobind directly lives here
+-nanobind_add_module(xgrammar_bindings LTO nanobind.cc)
+-target_link_libraries(xgrammar_bindings PRIVATE python_methods)
++nanobind_build_library(nanobind)
++add_library(xgrammar_bindings MODULE nanobind.cc)
++target_link_libraries(xgrammar_bindings PRIVATE
++  python_methods
++  nanobind
++)
++nanobind_opt_size(xgrammar_bindings)
++nanobind_lto(xgrammar_bindings)
++nanobind_set_visibility(xgrammar_bindings)
++nanobind_extension(xgrammar_bindings)
++nanobind_compile_options(xgrammar_bindings)
++nanobind_link_options(xgrammar_bindings)
+ 
+ if(DEFINED SKBUILD_PROJECT_NAME)
+   # Building wheel through scikit-build-core

nix/xgrammar_cmake_no_LTO.patch

@@ -0,0 +1,13 @@
+diff --git a/cpp/nanobind/CMakeLists.txt b/cpp/nanobind/CMakeLists.txt
+index 03cf0dc..3c9a90a 100644
+--- a/cpp/nanobind/CMakeLists.txt
++++ b/cpp/nanobind/CMakeLists.txt
+@@ -18,7 +18,7 @@ target_sources(python_methods PRIVATE python_methods.cc)
+ target_link_libraries(python_methods PUBLIC xgrammar)
+ 
+ # Any code that uses nanobind directly lives here
+-nanobind_add_module(xgrammar_bindings LTO nanobind.cc)
++nanobind_add_module(xgrammar_bindings nanobind.cc)
+ target_link_libraries(xgrammar_bindings PRIVATE python_methods)
+ 
+ if(DEFINED SKBUILD_PROJECT_NAME)

pyproject.toml

@@ -15,17 +15,17 @@ dependencies = [
   "huggingface-hub>=0.33.4",
   "psutil>=7.0.0",
   "loguru>=0.7.3",
-  "exo_pyo3_bindings",                         # rust bindings
+  "exo_pyo3_bindings",                       # rust bindings
   "anyio==4.11.0",
-  "mlx; sys_platform == 'darwin'",
-  "mlx[cpu]==0.30.6; sys_platform == 'linux'",
+  "mlx",
+  "mlx[cpu]; 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.16.9",
+  "mflux==0.17.2; sys_platform == 'darwin'",
   "python-multipart>=0.0.21",
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
@@ -45,11 +45,14 @@ 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",
-# ]
+build = ["nanobind"]
+cuda = [
+  "torch>=2.10.0; sys_platform == 'linux'",
+  "vllm>=0.13.0; sys_platform == 'linux'",
+  "fastsafetensors>=0.1.10; sys_platform == 'linux'",
+  "mlx[cuda13]; sys_platform == 'linux'",
+]
 
 ###
 # workspace configuration
@@ -60,11 +63,15 @@ 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'" }
+mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks" }
 mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" }
-# 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 }
+torch = [{ index = "pytorch-cu130", marker = "sys_platform == 'linux'" }]
+vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
+
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true
 
 [build-system]
 requires = ["uv_build>=0.8.9,<0.9.0"]
@@ -99,8 +106,16 @@ exclude = [
   "**/.direnv",
   "**/rust",
   "**/.github",
+  "**/vllm_patches",
+  "**/engines/vllm",
 ]
 stubPath = ".mlx_typings"
+extraPaths = [".cuda_typings"]
+
+[[tool.basedpyright.executionEnvironments]]
+root = "src/exo/worker/runner/vllm_inference"
+extraPaths = ["src", ".cuda_typings"]
+reportMissingModuleSource = false
 
 [[tool.basedpyright.executionEnvironments]]
 root = "src"
@@ -114,6 +129,46 @@ root = "src"
 required-version = ">=0.8.6"
 prerelease = "allow"
 environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
+conflicts = [[{ package = "exo", extra = "cuda" }, { package = "exo-bench" }]]
+override-dependencies = ["compressed-tensors==0.13.0", "torch==2.10.0"]
+constraint-dependencies = ["transformers>=5.0.0,<5.4.0"]
+
+[tool.uv.extra-build-dependencies]
+mlx = [
+  "setuptools",
+  "typing-extensions",
+  "nanobind",
+  "pybind11",
+  "wheel",
+  "cmake",
+  "ninja",
+]
+mlx-lm = ["setuptools"]
+xgrammar = [
+  "nanobind",
+  "setuptools",
+  "scikit-build-core",
+  "packaging",
+  "pathspec",
+]
+rouge-score = ["setuptools"]
+sacrebleu = ["setuptools"]
+sqlitedict = ["setuptools"]
+word2number = ["setuptools"]
+vllm = [
+  "setuptools",
+  "setuptools-scm",
+  "scikit-build-core",
+  "jinja2",
+  "wheel",
+  "markupsafe",
+  "typing-extensions",
+  "torch",
+]
+fastsafetensors = ["setuptools", "pybind11"]
+torch = ["typing-extensions"]
+torchvision = ["torch"]
+torchaudio = ["torch"]
 
 ###
 # ruff configuration
@@ -121,8 +176,8 @@ environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
 
 [tool.ruff]
 extend-exclude = [
-  "shared/protobufs/**",
   "*mlx_typings/**",
+  "*cuda_typings/**",
   "rust/exo_pyo3_bindings/**",
   "bench/vendor/**",
 ]

python/parts.nix

@@ -1,146 +1,434 @@
 { inputs, ... }:
-{
-  perSystem =
-    { config, self', pkgs, lib, system, ... }:
+let
+  mkPythonSet = { self', pkgs, lib, apple-sdk, editable ? false }:
     let
+      inherit (pkgs.stdenv.hostPlatform) isDarwin isLinux isx86_64;
+      inherit (pkgs.config) cudaSupport;
+      inherit (pkgs) cudaPackages;
+      cudaLibs = with cudaPackages; [
+        cuda_cudart
+        cuda_cccl
+        cuda_cupti
+        cuda_nvrtc
+        cuda_nvtx
+        cudnn
+        libcufile
+        libcublas
+        libcufft
+        libcurand
+        libcusolver
+        libcusparse
+        libcusparse_lt
+        libnvjitlink
+        libnvshmem
+        nccl
+      ];
+      cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" { } ''
+        mkdir -p $out/include
+        ln -s ${cudaPackages.cuda_cccl}/include $out/include/cccl
+      '';
+
+      cudaRoot = pkgs.symlinkJoin {
+        name = "cuda-merged-exo";
+        paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (cudaLibs ++ [ cudaPackages.cuda_nvcc cuda_cccl_compat ]);
+      };
+      exoOverlay = final: prev:
+        {
+          mlx = prev.mlx.overrideAttrs (old:
+            let
+              # Static dependencies included directly during compilation
+              gguf-tools = pkgs.fetchFromGitHub {
+                owner = "antirez";
+                repo = "gguf-tools";
+                rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
+                hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
+              };
+
+              metal_cpp = pkgs.fetchzip {
+                url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip";
+                hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
+              };
+
+              nanobind = pkgs.fetchFromGitHub {
+                owner = "wjakob";
+                repo = "nanobind";
+                rev = "v2.10.2";
+                hash = "sha256-io44YhN+VpfHFWyvvLWSanRgbzA0whK8WlDNRi3hahU=";
+                fetchSubmodules = true;
+              };
+
+              nvtx = pkgs.fetchFromGitHub {
+                name = "nvtx3";
+                owner = "NVIDIA";
+                repo = "NVTX";
+                rev = "v3.1.1";
+                hash = "sha256-sx72N+Gskg9Vtqc3sXsWoE/2PHFI2Hq08lEaw0sll5Y=";
+              };
+              cudnn = pkgs.fetchFromGitHub {
+                name = "cudnn_frontend";
+                owner = "NVIDIA";
+                repo = "cudnn-frontend";
+                rev = "v1.16.0";
+                hash = "sha256-+8aBl9dKd2Uz50XoOr91NRyJ4OGJtzfDNNNYGQJ9b94=";
+              };
+              mlx_cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" { } ''
+                mkdir -p $out/include
+                exit 1
+                ln -s ${cudaPackages.cuda_cccl}/include/cuda $out/include/cuda
+                ln -s ${cudaPackages.cuda_cccl}/include/nv $out/include/nv
+              '';
+            in
+            {
+              nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake ] ++ lib.optionals isDarwin [ self'.packages.metal-toolchain ] ++ lib.optionals cudaSupport [
+                cudaPackages.cuda_nvcc
+                pkgs.autoAddDriverRunpath
+                pkgs.autoPatchelfHook
+              ];
+              # TODO: non-sdk_26 support
+              buildInputs = (old.buildInputs or [ ])
+              ++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.openblas ]
+              ++ lib.optionals isDarwin [ apple-sdk ]
+              ++ lib.optionals cudaSupport (cudaLibs ++ [ cudaPackages.cudnn ]);
+              patches = (old.patches or [ ])
+              ++ lib.optionals cudaSupport [ ../nix/mlx_patch_fmod.patch ]
+              ++ lib.optionals isDarwin [
+                (pkgs.replaceVars ../nix/darwin-build-fixes.patch {
+                  sdkVersion = apple-sdk.version;
+                  inherit (self'.packages.metal-toolchain) metalVersion;
+                })
+              ];
+              postPatch = ''
+                substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
+                  --replace-fail "g++" "${lib.getExe' pkgs.stdenv.cc "c++"}"
+              '';
+              autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ lib.optionals cudaSupport [ "libcuda.so.1" ];
+
+              DEV_RELEASE = 1;
+              CMAKE_ARGS = toString ([
+                (lib.cmakeBool "USE_SYSTEM_FMT" true)
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${pkgs.nlohmann_json.src}")
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}")
+                (lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
+                (lib.cmakeBool "MLX_BUILD_CPU" true)
+                (lib.cmakeBool "MLX_BUILD_METAL" isDarwin)
+                (lib.cmakeBool "MLX_BUILD_CUDA" false)
+                (lib.cmakeOptionType "string" "CMAKE_INSTALL_LIBDIR" "lib")
+              ] ++ lib.optionals cudaSupport [
+                (lib.cmakeOptionType "filepath" "CUDAToolkit_ROOT" "${cudaRoot}")
+                (lib.cmakeOptionType "string" "MLX_CUDA_ARCHITECTURES" "121")
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_CUTLASS" "${cudaPackages.cutlass}")
+                # TODO: replace with cudaPackages.cudnn
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_CUDNN" "${cudnn}")
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_CCCL" "${cudaPackages.cuda_cccl}")
+                # TODO: replace with cudaPackages.nvtx
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NVTX3" "${nvtx}")
+              ] ++ lib.optionals isDarwin [
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
+                (lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${apple-sdk.version}")
+                (lib.cmakeOptionType "filepath" "CMAKE_OSX_SYSROOT" "${apple-sdk.passthru.sdkroot}")
+              ] ++ lib.optionals (isDarwin && isx86_64) [
+                (lib.cmakeBool "MLX_ENABLE_X64_MAC" true)
+              ]);
+            } // lib.optionalAttrs isDarwin {
+              SDKROOT = apple-sdk.passthru.sdkroot;
+              MACOSX_DEPLOYMENT_TARGET = apple-sdk.version;
+            });
+          exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
+            pname = "exo-pyo3-bindings";
+            version = "0.1.0";
+            src = self'.packages.exo_pyo3_bindings;
+            # Install from pre-built wheel
+            nativeBuildInputs = [ final.pyprojectWheelHook ];
+            dontStrip = true;
+            passthru = prev.exo-pyo3-bindings.passthru or { };
+            postInstall = ''
+              local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings
+              cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/
+              touch $siteDir/py.typed
+            '';
+          };
+          torch = prev.torch.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = (old.buildInputs or [ ]) ++ lib.optionals cudaSupport cudaLibs;
+            autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ lib.optionals cudaSupport [ "libcuda.so.1" ];
+          });
+          torchaudio = prev.torchaudio.overrideAttrs (old:
+            {
+              nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+              buildInputs = (old.buildInputs or [ ]) ++ lib.optionals cudaSupport [
+                cudaPackages.cuda_cudart
+              ];
+              preFixup = (old.preFixup or "") + ''
+                addAutoPatchelfSearchPath "${final.torch}"
+              '';
+              autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ lib.optionals cudaSupport [ "libcuda.so.1" ];
+            });
+          torchvision = prev.torchvision.overrideAttrs (old:
+            {
+              nativebuildInputs = (old.nativebuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+              preFixup = (old.preFixup or "") + ''
+                addAutoPatchelfSearchPath "${final.torch}"
+              '';
+              autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ lib.optionals cudaSupport [ "libcuda.so.1" ];
+            });
+          torch-c-dlpack-ext = prev.torch-c-dlpack-ext.overrideAttrs (old:
+            {
+              nativebuildInputs = (old.nativebuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+              preFixup = (old.preFixup or "") + ''
+                addAutoPatchelfSearchPath "${final.torch}"
+	      '';
+              autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ lib.optionals cudaSupport [ "libcuda.so.1" ];
+            });
+          xgrammar = prev.xgrammar.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake pkgs.autoPatchelfHook ];
+          });
+          # Currently treating vllm as a cuda dep. it obviously exists as a non cuda dep
+          vllm = prev.vllm.overrideAttrs (old:
+            let
+              cutlass = pkgs.fetchFromGitHub {
+                name = "cutlass-source";
+                owner = "NVIDIA";
+                repo = "cutlass";
+                tag = "v4.2.1";
+                hash = "sha256-iP560D5Vwuj6wX1otJhwbvqe/X4mYVeKTpK533Wr5gY=";
+              };
+              triton-kernels = pkgs.fetchFromGitHub {
+                owner = "triton-lang";
+                repo = "triton";
+                tag = "v3.6.0";
+                hash = "sha256-JFSpQn+WsNnh7CAPlcpOcUp0nyKXNbJEANdXqmkt4Tc=";
+              };
+
+              cutlass-flashmla = pkgs.fetchFromGitHub {
+                owner = "NVIDIA";
+                repo = "cutlass";
+                rev = "147f5673d0c1c3dcf66f78d677fd647e4a020219";
+                hash = "sha256-dHQto08IwTDOIuFUp9jwm1MWkFi8v2YJ/UESrLuG71g=";
+              };
+
+              flashmla = pkgs.stdenv.mkDerivation {
+                pname = "flashmla";
+                version = "1.0.0";
+
+                src = pkgs.fetchFromGitHub {
+                  name = "FlashMLA-source";
+                  owner = "vllm-project";
+                  repo = "FlashMLA";
+                  rev = "c2afa9cb93e674d5a9120a170a6da57b89267208";
+                  hash = "sha256-pKlwxV6G9iHag/jbu3bAyvYvnu5TbrQwUMFV0AlGC3s=";
+                };
+
+                dontConfigure = true;
+
+                buildPhase = ''
+                  rm -rf csrc/cutlass
+                  ln -sf ${cutlass-flashmla} csrc/cutlass
+                '';
+
+                installPhase = ''
+                  cp -rva . $out
+                '';
+              };
+              qutlass = pkgs.fetchFromGitHub {
+                name = "qutlass-source";
+                owner = "IST-DASLab";
+                repo = "qutlass";
+                rev = "830d2c4537c7396e14a02a46fbddd18b5d107c65";
+                hash = "sha256-aG4qd0vlwP+8gudfvHwhtXCFmBOJKQQTvcwahpEqC84=";
+              };
+              vllm-flash-attn = pkgs.stdenv.mkDerivation {
+                pname = "vllm-flash-attn";
+                version = "2.7.2.post1";
+
+                src = pkgs.fetchFromGitHub {
+                  name = "flash-attention-source";
+                  owner = "vllm-project";
+                  repo = "flash-attention";
+                  rev = "188be16520ceefdc625fdf71365585d2ee348fe2";
+                  hash = "sha256-Osec+/IF3+UDtbIhDMBXzUeWJ7hDJNb5FpaVaziPSgM=";
+                };
+
+                patches = [
+                  (pkgs.fetchpatch {
+                    url = "https://github.com/Dao-AILab/flash-attention/commit/dad67c88d4b6122c69d0bed1cebded0cded71cea.patch";
+                    hash = "sha256-JSgXWItOp5KRpFbTQj/cZk+Tqez+4mEz5kmH5EUeQN4=";
+                  })
+                  (pkgs.fetchpatch {
+                    url = "https://github.com/Dao-AILab/flash-attention/commit/e26dd28e487117ee3e6bc4908682f41f31e6f83a.patch";
+                    hash = "sha256-NkCEowXSi+tiWu74Qt+VPKKavx0H9JeteovSJKToK9A=";
+                  })
+                ];
+
+                dontConfigure = true;
+
+                buildPhase = ''
+                  rm -rf csrc/cutlass
+                  ln -sf ${cutlass} csrc/cutlass
+                '';
+
+                installPhase = ''
+                  cp -rva . $out
+                '';
+              };
+            in
+            {
+              patches = (old.patches or [ ]) ++ [ ../nix/vllm_uv2nix_cmake.patch ];
+              nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
+                pkgs.cmake
+                pkgs.ninja
+                pkgs.autoAddDriverRunpath
+              ] ++ lib.optionals cudaSupport [
+                cudaPackages.cuda_nvcc
+              ];
+              # TODO: vllm rocm/cpu
+              VLLM_TARGET_DEVICE = "empty";
+              # TODO: vllm non cuda13 support, more arch's, etc.
+            } // lib.optionalAttrs cudaSupport {
+              buildInputs = (old.buildInputs or [ ]) ++ cudaLibs;
+
+              CUDA_HOME = "${cudaRoot}";
+              CUDAToolkit_ROOT = "${cudaRoot}";
+              CUDACXX = "${cudaRoot}/bin/nvcc";
+              VLLM_CUTLASS_SRC_DIR = "${lib.getDev cutlass}";
+              VLLM_TARGET_DEVICE = "cuda";
+              TORCH_CUDA_ARCH_LIST = "12.0;12.1";
+              TRITON_KERNELS_SRC_DIR = "${lib.getDev triton-kernels}/python/triton_kernels/triton_kernels";
+              FLASH_MLA_SRC_DIR = "${lib.getDev flashmla}";
+              QUTLASS_SRC_DIR = "${lib.getDev qutlass}";
+              VLLM_FLASH_ATTN_SRC_DIR = "${lib.getDev vllm-flash-attn}";
+              CAFFE2_USE_CUDNN = "ON";
+              CAFFE2_USE_CUFILE = "ON";
+              CUTLASS_ENABLE_CUBLAS = "ON";
+              CUTLASS_NVCC_ARCHS_ENABLED = "12.0;12.1";
+
+              UV2NIX_CMAKE_FLAGS_JSON = builtins.toJSON [
+                "-DCUDAToolkit_ROOT=${cudaRoot}"
+                "-DCMAKE_CUDA_COMPILER=${cudaRoot}/bin/nvcc"
+                "-DCMAKE_PREFIX_PATH=${cudaRoot}"
+                "-DFETCHCONTENT_SOURCE_DIR_CUTLASS=${lib.getDev cutlass}"
+                "-DFLASH_MLA_SRC_DIR=${lib.getDev flashmla}"
+                "-DVLLM_FLASH_ATTN_SRC_DIR=${lib.getDev vllm-flash-attn}"
+                "-DQUTLASS_SRC_DIR=${lib.getDev qutlass}"
+                "-DTORCH_CUDA_ARCH_LIST=12.0;12.1"
+                "-DCUTLASS_NVCC_ARCHS_ENABLED=${cudaPackages.flags.cmakeCudaArchitecturesString}"
+                "-DCAFFE2_USE_CUDNN=ON"
+                "-DCAFFE2_USE_CUFILE=ON"
+                "-DCUTLASS_ENABLE_CUBLAS=ON"
+              ];
+            });
+        } // lib.optionalAttrs cudaSupport {
+          nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = (old.buildInputs or [ ]) ++ [ pkgs.rdma-core ];
+            propagatedBuildInputs = (old.propagatedBuildInputs or [ ]) ++ [ pkgs.util-linux ];
+          });
+          nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = (old.buildInputs or [ ]) ++ (with cudaPackages; [ libnvjitlink libcublas libcusparse ]);
+          });
+          nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = (old.buildInputs or [ ]) ++ [ cudaPackages.libnvjitlink ];
+          });
+          nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = (old.buildInputs or [ ]) ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
+          });
+          nvidia-cutlass-dsl-libs-base = prev.nvidia-cutlass-dsl-libs-base.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+            autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ [ "libcuda.so.1" ];
+          });
+        } // lib.optionalAttrs (cudaSupport && isx86_64) {
+          numba = prev.numba.overrideAttrs (old: {
+            buildInputs = (old.buildInputs or [ ]) ++ [ pkgs.tbb ];
+          });
+          intel-openmp = prev.intel-openmp.overrideAttrs (_old: {
+            postFixup = ''
+              rm -f $out/lib/libarcher.so
+              rm -f $out/lib/libomptarget.so
+              rm -f $out/lib/libomptarget.rtl.*.so*
+              rm -f $out/lib/libomptarget.sycl.wrap.so
+            '';
+          });
+        };
+
       # Load workspace from uv.lock
       workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
-        workspaceRoot = inputs.self;
+        workspaceRoot = ../.;
       };
 
       # Create overlay from workspace
       # Use wheels from PyPI for most packages; we override mlx with our pure Nix Metal build
       overlay = workspace.mkPyprojectOverlay { sourcePreference = "wheel"; };
-
-      # Override overlay to inject Nix-built components
-      exoOverlay = final: prev: {
-        # Replace workspace exo_pyo3_bindings with Nix-built wheel.
-        # Preserve passthru so mkVirtualEnv can resolve dependency groups.
-        # Copy .pyi stub + py.typed marker so basedpyright can find the types.
-        exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
-          pname = "exo-pyo3-bindings";
-          version = "0.1.0";
-          src = self'.packages.exo_pyo3_bindings;
-          # Install from pre-built wheel
-          nativeBuildInputs = [ final.pyprojectWheelHook ];
-          dontStrip = true;
-          passthru = prev.exo-pyo3-bindings.passthru or { };
-          postInstall = ''
-            local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings
-            cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/
-            touch $siteDir/py.typed
-          '';
-        };
-      };
-
+    in
+    (pkgs.callPackage inputs.pyproject-nix.build.packages {
       python = pkgs.python313;
+    }).overrideScope (
+      lib.composeManyExtensions ([
+        inputs.pyproject-build-systems.overlays.default
+        overlay
+        exoOverlay
+      ] ++ lib.optionals editable [
+        (workspace.mkEditablePyprojectOverlay { root = "$REPO_ROOT"; members = [ "exo" "bench" ]; })
+      ])
+    );
 
-      # Overlay to provide build systems and custom packages
-      buildSystemsOverlay = final: prev: {
-        # mlx-lm is a git dependency that needs setuptools
-        mlx-lm = prev.mlx-lm.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-        # rouge-score and sacrebleu don't declare setuptools as a build dependency
-        rouge-score = prev.rouge-score.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-        sacrebleu = prev.sacrebleu.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-        sqlitedict = prev.sqlitedict.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-        word2number = prev.word2number.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-      } // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
-        # Use our pure Nix-built MLX with Metal support (macOS only)
-        mlx = self'.packages.mlx;
+  mkExo = args@{ self', pkgs, lib, ... }:
+    let
+      venv = ((mkPythonSet args).mkVirtualEnv "exo-env" {
+        exo = lib.optionals pkgs.config.cudaSupport [ "cuda" ];
+      }).overrideAttrs {
+        venvSkip = [ "lib/python3.13/site-packages/build_backend.py" ];
       };
+    in
+    pkgs.runCommand "exo"
+      {
+        nativeBuildInputs = [ pkgs.makeWrapper ];
+      }
+      ''
+        mkdir -p $out/bin
 
-      # Additional overlay for Linux-specific fixes (type checking env).
-      # Native wheels have shared lib dependencies we don't need at type-check time.
-      linuxOverlay = final: prev:
-        let
-          ignoreMissing = drv: drv.overrideAttrs { autoPatchelfIgnoreMissingDeps = [ "*" ]; };
-          nvidiaPackages = lib.filterAttrs (name: _: lib.hasPrefix "nvidia-" name) prev;
-        in
-        lib.optionalAttrs pkgs.stdenv.hostPlatform.isLinux (
-          (lib.mapAttrs (_: ignoreMissing) nvidiaPackages) // {
-            mlx = ignoreMissing prev.mlx;
-            mlx-cuda-13 = prev.mlx-cuda-13.overrideAttrs (old: {
-              buildInputs = (old.buildInputs or [ ]) ++ [
-                final.nvidia-cublas
-                final.nvidia-cuda-nvrtc
-                final.nvidia-cudnn-cu13
-                final.nvidia-nccl-cu13
-              ];
-              preFixup = ''
-                addAutoPatchelfSearchPath ${final.nvidia-cublas}
-                addAutoPatchelfSearchPath ${final.nvidia-cuda-nvrtc}
-                addAutoPatchelfSearchPath ${final.nvidia-cudnn-cu13}
-                addAutoPatchelfSearchPath ${final.nvidia-nccl-cu13}
-              '';
-              autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-            });
-            torch = ignoreMissing prev.torch;
-            triton = ignoreMissing prev.triton;
-          }
-        );
+        # Create wrapper script
+        makeWrapper ${venv}/bin/exo $out/bin/exo \
+          --set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
+          --set EXO_RESOURCES_DIR ${inputs.self + /resources} \
+          ${lib.optionalString pkgs.stdenv.hostPlatform.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
+      '';
+in
+{
+  perSystem =
+    { self', pkgs, cudaPkgs, lib, ... }:
+    let
+      inherit (pkgs.stdenv.hostPlatform) isDarwin;
+      pythonSet = mkPythonSet { inherit self' pkgs lib; apple-sdk = pkgs.apple-sdk_26; };
+      # taking cudaPkgs.cudaPackages_13.pkgs creates a new nixpkgs that defaults to cuda 13
+      cudaPythonSet = mkPythonSet { inherit self' lib; inherit (cudaPkgs.cudaPackages_13) pkgs; apple-sdk = pkgs.apple-sdk_26; };
 
-      pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
-        inherit python;
-      }).overrideScope (
-        lib.composeManyExtensions [
-          inputs.pyproject-build-systems.overlays.default
-          overlay
-          exoOverlay
-          buildSystemsOverlay
-          linuxOverlay
-        ]
-      );
-      # 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.
-      venvCollisionPaths = lib.optionals pkgs.stdenv.hostPlatform.isLinux [
-        "lib/python3.13/site-packages/mlx*"
-        "lib/python3.13/site-packages/nvidia*"
-      ];
-
-      # Exclude bench deps from main env (bench has its own benchVenv)
-      exoDeps = removeAttrs workspace.deps.default [ "exo-bench" ];
-
-      exoVenv = (pythonSet.mkVirtualEnv "exo-env" exoDeps).overrideAttrs {
-        venvIgnoreCollisions = venvCollisionPaths;
+      editablePythonSet = mkPythonSet { inherit self' lib pkgs; apple-sdk = pkgs.apple-sdk_26; editable = true; };
+      evenv = (editablePythonSet.mkVirtualEnv "exo-dev-env"
+        {
+          exo = [ "dev" ];
+          exo-pyo3-bindings = [ ];
+          exo-bench = [ ];
+        }).overrideAttrs {
+        venvSkip = [ "lib/python3.13/site-packages/build_backend.py" ];
+      };
+      exoCudaVenv = (cudaPythonSet.mkVirtualEnv "exo-env" {
+        exo = [ "cuda" ];
+        exo-pyo3-bindings = [ ];
+      }).overrideAttrs {
+        venvSkip = [ "lib/python3.13/site-packages/build_backend.py" ];
       };
 
       # Virtual environment with dev dependencies for testing
-      testVenv = (pythonSet.mkVirtualEnv "exo-test-env" (
-        exoDeps // {
+      testVenv = (pythonSet.mkVirtualEnv "exo-test-env"
+        {
           exo = [ "dev" ]; # Include pytest, pytest-asyncio, pytest-env
+          exo-pyo3-bindings = [ ];
         }
-      )).overrideAttrs {
-        venvIgnoreCollisions = venvCollisionPaths;
-      };
-
-      mkPythonScript = name: path: pkgs.writeShellApplication {
-        inherit name;
-        runtimeInputs = [ exoVenv ];
-        runtimeEnv = {
-          EXO_DASHBOARD_DIR = self'.packages.dashboard;
-          EXO_RESOURCES_DIR = inputs.self + /resources;
-        };
-        text = ''exec python ${path} "$@"'';
+      ).overrideAttrs {
+        # venvIgnoreCollisions = venvCollisionPaths;
       };
 
       benchVenv = pythonSet.mkVirtualEnv "exo-bench-env" {
@@ -159,32 +447,34 @@
         text = ''exec python ${path} "$@"'';
       };
 
-      exoPackage = pkgs.runCommand "exo"
+      exoCudaPackage = cudaPkgs.runCommand "exo"
         {
-          nativeBuildInputs = [ pkgs.makeWrapper ];
+          nativeBuildInputs = [ cudaPkgs.makeWrapper ];
         }
         ''
           mkdir -p $out/bin
 
           # Create wrapper script
-          makeWrapper ${exoVenv}/bin/exo $out/bin/exo \
+          makeWrapper ${exoCudaVenv}/bin/exo $out/bin/exo \
             --set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
             --set EXO_RESOURCES_DIR ${inputs.self + /resources} \
-            ${lib.optionalString pkgs.stdenv.hostPlatform.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
+            --prefix LD_LIBRARY_PATH : /run/opengl-driver/lib:${lib.getLib pkgs.util-linux}/lib:${lib.getLib pkgs.systemd}/lib:${lib.getLib pkgs.numactl}/lib:${lib.getLib pkgs.stdenv.cc.cc.lib}/lib \
+            ${lib.optionalString isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
         '';
     in
     {
-      # Python package only available on macOS (requires MLX/Metal)
-      packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin
-        {
-          exo = exoPackage;
-          # Test environment for running pytest outside of Nix sandbox (needs GPU access)
-          exo-test-env = testVenv;
-        } // {
+      packages = {
+        exo = mkExo { inherit self' lib pkgs; apple-sdk = pkgs.apple-sdk_26; };
+        exo-cuda = exoCudaPackage;
         exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
         exo-eval = mkBenchScript "exo-eval" (inputs.self + /bench/exo_eval.py);
         exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
         exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
+        editable-venv = evenv;
+      } // lib.optionalAttrs isDarwin {
+        # Test environment for running pytest outside of Nix sandbox (needs GPU access)
+        exo-test-env = testVenv;
+        exo-osx14 = mkExo { inherit self' lib pkgs; apple-sdk = pkgs.apple-sdk_14; };
       };
 
       checks = {

scripts/check_kv_compat.py

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

scripts/check_rotating_cache.py

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

scripts/disaggregated/capture_connector.py

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

scripts/test_vllm_multi.sh

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

src/exo/__main__.py

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

src/exo/disaggregated/batch_connector.py

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

src/exo/disaggregated/prefill_client.py

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

src/exo/disaggregated/protocol.py

@@ -0,0 +1,227 @@
+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,143 @@
+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).cpu() for t in kv_layer]  # pyright: ignore[reportAny]
+            _shared_arrays_queue.put((layer_idx, arrays))
+            return
+
+        if self._save_count < 1:
+            self._save_count += 1
+
+        if slot_mapping is not None:
+            if kv_layer.shape[0] == 2:  # pyright: ignore[reportAny]
+                k_all = 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.cpu(), values.cpu()))  # pyright: ignore[reportAny]
+        else:
+            self._queue.put((layer_idx, kv_layer.cpu().clone(), kv_layer.cpu().clone()))  # pyright: ignore[reportAny]
+
+    def wait_for_save(self) -> None:
+        pass
+
+    def finish(self) -> None:
+        self._queue.put(None)
+
+    def request_finished_all_groups(
+        self, request: Any, block_ids: tuple[list[int], ...]
+    ) -> tuple[bool, dict[str, Any] | None]:  # pyright: ignore[reportAny]
+        return False, None
+
+    def get_num_new_matched_tokens(
+        self, request: Any, num_computed_tokens: int
+    ) -> tuple[int, bool]:  # pyright: ignore[reportAny]
+        return 0, False
+
+    def update_state_after_alloc(
+        self, request: Any, blocks: Any, num_external_tokens: int
+    ) -> None:  # pyright: ignore[reportAny]
+        pass
+
+    def build_connector_meta(self, scheduler_output: Any) -> StreamingConnectorMetadata:  # pyright: ignore[reportAny]
+        return StreamingConnectorMetadata()

src/exo/download/download_utils.py

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

src/exo/main.py

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

src/exo/master/adapters/chat_completions.py

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

src/exo/master/adapters/ollama.py

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

src/exo/master/api.py

@@ -341,6 +341,7 @@ class API:
         self.app.get("/ollama/api/version")(self.ollama_version)
 
         self.app.get("/state")(lambda: self.state)
+        self.app.get("/capabilities")(self._get_capabilities)
         self.app.get("/events")(self.stream_events)
         self.app.post("/download/start")(self.start_download)
         self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download)
@@ -413,6 +414,7 @@ class API:
                 node_network=self.state.node_network,
                 topology=self.state.topology,
                 current_instances=self.state.instances,
+                node_vllm=self.state.node_vllm,
             )
         except ValueError as exc:
             raise HTTPException(status_code=400, detail=str(exc)) from exc
@@ -448,18 +450,15 @@ class API:
                 status_code=400, detail=f"Failed to load model card: {exc}"
             ) from exc
         instance_combinations: list[tuple[Sharding, InstanceMeta, int]] = []
+        node_count = len(list(self.state.topology.list_nodes()))
+
         for sharding in (Sharding.Pipeline, Sharding.Tensor):
             for instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl):
                 instance_combinations.extend(
-                    [
-                        (sharding, instance_meta, i)
-                        for i in range(
-                            1, len(list(self.state.topology.list_nodes())) + 1
-                        )
-                    ]
+                    [(sharding, instance_meta, i) for i in range(1, node_count + 1)]
                 )
-        # TODO: PDD
-        # instance_combinations.append((Sharding.PrefillDecodeDisaggregation, InstanceMeta.MlxRing, 1))
+        if any(self.state.node_vllm.values()):
+            instance_combinations.append((Sharding.Pipeline, InstanceMeta.Vllm, 1))
 
         for sharding, instance_meta, min_nodes in instance_combinations:
             try:
@@ -472,6 +471,7 @@ class API:
                     ),
                     node_memory=self.state.node_memory,
                     node_network=self.state.node_network,
+                    node_vllm=self.state.node_vllm,
                     topology=self.state.topology,
                     current_instances=self.state.instances,
                     required_nodes=required_nodes,
@@ -749,20 +749,30 @@ class API:
         return await self._collect_text_generation_with_stats(command.command_id)
 
     async def _resolve_and_validate_text_model(self, model_id: ModelId) -> ModelId:
-        """Validate a text model exists and return the resolved model ID.
+        from exo.shared.models.model_cards import derive_base_model
 
-        Raises HTTPException 404 if no instance is found for the model.
-        """
-        if not any(
+        if any(
             instance.shard_assignments.model_id == model_id
             for instance in self.state.instances.values()
         ):
-            await self._trigger_notify_user_to_download_model(model_id)
-            raise HTTPException(
-                status_code=404,
-                detail=f"No instance found for model {model_id}",
+            return model_id
+
+        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
             )
-        return model_id
+            if (
+                first_shard is not None
+                and first_shard.model_card.base_model.lower() == request_base.lower()
+            ):
+                return instance.shard_assignments.model_id
+
+        await self._trigger_notify_user_to_download_model(model_id)
+        raise HTTPException(
+            status_code=404,
+            detail=f"No instance found for model {model_id}",
+        )
 
     async def _validate_image_model(self, model: ModelId) -> ModelId:
         """Validate model exists and return resolved model ID.
@@ -781,6 +791,9 @@ class API:
             )
         return resolved_model
 
+    def _get_capabilities(self) -> dict[str, bool]:
+        return {"vllm_available": any(self.state.node_vllm.values())}
+
     def stream_events(self) -> StreamingResponse:
         def _generate_json_array(events: Iterable[Event]) -> Iterable[str]:
             yield "["

src/exo/master/main.py

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

src/exo/master/placement.py

@@ -41,6 +41,7 @@ from exo.shared.types.worker.instances import (
     InstanceMeta,
     MlxJacclInstance,
     MlxRingInstance,
+    VllmInstance,
 )
 from exo.shared.types.worker.shards import Sharding
 
@@ -66,11 +67,44 @@ def place_instance(
     current_instances: Mapping[InstanceId, Instance],
     node_memory: Mapping[NodeId, MemoryUsage],
     node_network: Mapping[NodeId, NodeNetworkInfo],
+    node_vllm: Mapping[NodeId, bool],
     required_nodes: set[NodeId] | None = None,
 ) -> dict[InstanceId, Instance]:
     cycles = topology.get_cycles()
     candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
 
+    # vLLM instances can only be placed on nodes that have vLLM available.
+    # vLLM does not support quantized mlx-community models (only bf16 or unquantized).
+    if command.instance_meta == InstanceMeta.Vllm:
+        is_mlx_community = str(command.model_card.model_id).startswith("mlx-community/")
+        if is_mlx_community and command.model_card.quantization not in ("", "bf16"):
+            raise ValueError("vLLM does not support quantized mlx-community models")
+        candidate_cycles = [
+            cycle
+            for cycle in candidate_cycles
+            if all(node_vllm.get(nid, False) for nid in cycle.node_ids)
+        ]
+
+    # QMM/quantized ops are not available on MLX CUDA — exclude CUDA nodes for quantized MLX models.
+    if command.instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl):
+        if command.model_card.quantization not in ("", "bf16"):
+            candidate_cycles = [
+                cycle
+                for cycle in candidate_cycles
+                if not any(node_vllm.get(nid, False) for nid in cycle.node_ids)
+            ]
+
+    # mlx-community models should prefer Apple Silicon nodes over CUDA nodes.
+    if command.instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl):
+        if str(command.model_card.model_id).startswith("mlx-community/"):
+            apple_silicon_cycles = [
+                cycle
+                for cycle in candidate_cycles
+                if not any(node_vllm.get(nid, False) for nid in cycle.node_ids)
+            ]
+            if apple_silicon_cycles:
+                candidate_cycles = apple_silicon_cycles
+
     # Filter to cycles containing all required nodes (subset matching)
     if required_nodes:
         candidate_cycles = [
@@ -78,8 +112,11 @@ def place_instance(
             for cycle in candidate_cycles
             if required_nodes.issubset(cycle.node_ids)
         ]
+    required_memory = command.model_card.storage_size
+    if command.instance_meta == InstanceMeta.Vllm:
+        required_memory = Memory.from_bytes(int(required_memory.in_bytes * 1.3))
     cycles_with_sufficient_memory = filter_cycles_by_memory(
-        candidate_cycles, node_memory, command.model_card.storage_size
+        candidate_cycles, node_memory, required_memory
     )
     if len(cycles_with_sufficient_memory) == 0:
         raise ValueError("No cycles found with sufficient memory")
@@ -139,7 +176,7 @@ def place_instance(
     )
 
     # Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
-    if len(selected_cycle) == 1:
+    if len(selected_cycle) == 1 and command.instance_meta != InstanceMeta.Vllm:
         command.instance_meta = InstanceMeta.MlxRing
         command.sharding = Sharding.Pipeline
 
@@ -199,6 +236,11 @@ def place_instance(
                 hosts_by_node=hosts_by_node,
                 ephemeral_port=ephemeral_port,
             )
+        case InstanceMeta.Vllm:
+            target_instances[instance_id] = VllmInstance(
+                instance_id=instance_id,
+                shard_assignments=shard_assignments,
+            )
 
     return target_instances
 

src/exo/master/tests/test_placement.py

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

src/exo/shared/apply.py

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

src/exo/shared/models/model_cards.py

@@ -38,6 +38,35 @@ CARD_SEARCH_PATH = [
 
 _card_cache: dict[ModelId, "ModelCard"] = {}
 
+import re
+
+_QUANT_SUFFIXES = re.compile(
+    r"[-_ ](?:MLX|MXFP[0-9]+|NVFP[0-9]+|GPTQ|AWQ|GGUF|fp16|bf16|fp8|int[0-9]+|[0-9]+(?:\.[0-9]+)?bit|Q[0-9]+(?:_[A-Z0-9]+)?|gs[0-9]+)(?:[-_ ](?:MLX|Q[0-9]+|Int[0-9]+|[A-Z0-9]+|gs[0-9]+))*$",
+    re.IGNORECASE,
+)
+
+
+def _normalize_base_model(s: str) -> str:
+    return s.replace("-", " ").replace("_", " ").replace("  ", " ").strip()
+
+
+def derive_base_model(model_id: str) -> str:
+    short = model_id.split("/")[-1] if "/" in model_id else model_id
+    base = _QUANT_SUFFIXES.sub("", short)
+    return _normalize_base_model(base)
+
+
+def derive_family(model_id: str) -> str:
+    short = model_id.split("/")[-1] if "/" in model_id else model_id
+    short = _QUANT_SUFFIXES.sub("", short).lower().replace("_", "-")
+    parts = re.split(r"[-.]", short)
+    family_parts: list[str] = []
+    for p in parts:
+        if p.isdigit() or re.match(r"^\d+[bm]?$", p, re.IGNORECASE):
+            break
+        family_parts.append(p)
+    return "-".join(family_parts) if family_parts else short
+
 
 async def _refresh_card_cache():
     for path in CARD_SEARCH_PATH:
@@ -93,6 +122,15 @@ class ModelCard(CamelCaseModel):
     uses_cfg: bool = False
     trust_remote_code: bool = True
 
+    @model_validator(mode="after")
+    def _ensure_derived_fields(self) -> "ModelCard":
+        if not self.base_model:
+            self.base_model = derive_base_model(self.model_id)
+        else:
+            stripped = _QUANT_SUFFIXES.sub("", self.base_model)
+            self.base_model = _normalize_base_model(stripped)
+        return self
+
     @field_validator("tasks", mode="before")
     @classmethod
     def _validate_tasks(cls, v: list[str | ModelTask]) -> list[ModelTask]:
@@ -132,6 +170,9 @@ class ModelCard(CamelCaseModel):
         num_layers = config_data.layer_count
         mem_size_bytes = await fetch_safetensors_size(model_id)
 
+        base_model = derive_base_model(model_id)
+        family = (config_data.model_type or "").replace("_", "-")
+
         mc = ModelCard(
             model_id=ModelId(model_id),
             storage_size=mem_size_bytes,
@@ -141,6 +182,8 @@ class ModelCard(CamelCaseModel):
             num_key_value_heads=config_data.num_key_value_heads,
             tasks=[ModelTask.TextGeneration],
             trust_remote_code=False,
+            base_model=base_model,
+            family=family,
         )
         await mc.save_to_custom_dir()
         _card_cache[model_id] = mc
@@ -170,6 +213,7 @@ def is_custom_card(model_id: ModelId) -> bool:
 class ConfigData(BaseModel):
     model_config = {"extra": "ignore"}  # Allow unknown fields
 
+    model_type: str | None = None
     architectures: list[str] | None = None
     hidden_size: Annotated[int, Field(ge=0)] | None = None
     num_key_value_heads: PositiveInt | None = None
@@ -258,21 +302,24 @@ async def fetch_safetensors_size(model_id: ModelId) -> Memory:
 
     target_dir = (await ensure_models_dir()) / model_id.normalize()
     await aios.makedirs(target_dir, exist_ok=True)
-    index_path = await download_file_with_retry(
-        model_id,
-        "main",
-        "model.safetensors.index.json",
-        target_dir,
-        lambda curr_bytes, total_bytes, is_renamed: logger.debug(
-            f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
-        ),
-    )
-    async with aiofiles.open(index_path, "r") as f:
-        index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
+    try:
+        index_path = await download_file_with_retry(
+            model_id,
+            "main",
+            "model.safetensors.index.json",
+            target_dir,
+            lambda curr_bytes, total_bytes, is_renamed: logger.debug(
+                f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
+            ),
+        )
+        async with aiofiles.open(index_path, "r") as f:
+            index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
 
-    metadata = index_data.metadata
-    if metadata is not None:
-        return Memory.from_bytes(metadata.total_size)
+        metadata = index_data.metadata
+        if metadata is not None:
+            return Memory.from_bytes(metadata.total_size)
+    except FileNotFoundError:
+        pass
 
     info = model_info(model_id)
     if info.safetensors is None:

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

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

src/exo/shared/tracing.py

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

src/exo/shared/types/api.py

@@ -221,6 +221,7 @@ 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):

src/exo/shared/types/mlx.py

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

src/exo/shared/types/ollama_api.py

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

src/exo/shared/types/profiling.py

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

src/exo/shared/types/state.py

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

src/exo/shared/types/text_generation.py

@@ -70,3 +70,4 @@ 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

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

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

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

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

src/exo/utils/info_gatherer/info_gatherer.py

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

src/exo/utils/info_gatherer/nvml.py

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

src/exo/utils/info_gatherer/system_info.py

@@ -1,6 +1,7 @@
 import platform
 import socket
 import sys
+from pathlib import Path
 from subprocess import CalledProcessError
 
 import psutil
@@ -90,39 +91,179 @@ async def _get_interface_types_from_networksetup() -> dict[str, InterfaceType]:
     return types
 
 
+def _classify_unknown_darwin_interface(name: str) -> InterfaceType:
+    if name.lower().startswith("anpi"):
+        return "thunderbolt"
+    return "unknown"
+
+
+async def _get_linux_network_interfaces() -> list[NetworkInterfaceInfo]:
+    import json as _json
+
+    try:
+        result = await run_process(["ip", "-j", "addr", "show"])
+    except (CalledProcessError, FileNotFoundError):
+        return []
+
+    data: list[dict[str, object]] = _json.loads(result.stdout)  # pyright: ignore[reportAny]
+    interfaces: list[NetworkInterfaceInfo] = []
+    for iface in data:
+        name: str = iface.get("ifname", "")  # pyright: ignore[reportAssignmentType, reportAny]
+        link_type: str = iface.get("link_type", "")  # pyright: ignore[reportAssignmentType, reportAny]
+
+        iface_type: InterfaceType
+        if link_type == "loopback":
+            continue
+        elif link_type == "ether":
+            if name.startswith(("wl", "wlan")):
+                iface_type = "wifi"
+            elif name.startswith(("docker", "br-", "veth")):
+                iface_type = "unknown"
+            elif name.startswith(("thunderbolt", "tb", "enx")):
+                iface_type = "thunderbolt"
+            else:
+                iface_type = "ethernet"
+        elif link_type in ("none", "tun"):
+            iface_type = "unknown"
+        else:
+            iface_type = "unknown"
+
+        for addr_info in iface.get("addr_info", []):  # pyright: ignore[reportAny]
+            family: str = addr_info.get("family", "")  # pyright: ignore[reportAny]
+            ip: str = addr_info.get("local", "")  # pyright: ignore[reportAny]
+            if family in ("inet", "inet6") and ip:
+                interfaces.append(
+                    NetworkInterfaceInfo(
+                        name=name, ip_address=ip, interface_type=iface_type
+                    )
+                )
+
+    return interfaces
+
+
 async def get_network_interfaces() -> list[NetworkInterfaceInfo]:
     """
-    Retrieves detailed network interface information on macOS.
-    Parses output from 'networksetup -listallhardwareports' and 'ifconfig'
+    Retrieves detailed network interface information on macOS or Linux.
+    On MacOS: parses output from 'networksetup -listallhardwareports' and 'ifconfig'
     to determine interface names, IP addresses, and types (ethernet, wifi, vpn, other).
+    Falls back to using ip -j addr show on other platforms.
     Returns a list of NetworkInterfaceInfo objects.
     """
-    interfaces_info: list[NetworkInterfaceInfo] = []
-    interface_types = await _get_interface_types_from_networksetup()
-
-    for iface, services in psutil.net_if_addrs().items():
-        for service in services:
-            match service.family:
-                case socket.AF_INET | socket.AF_INET6:
-                    interfaces_info.append(
-                        NetworkInterfaceInfo(
-                            name=iface,
-                            ip_address=service.address,
-                            interface_type=interface_types.get(iface, "unknown"),
+    if sys.platform == "darwin":
+        interfaces_info: list[NetworkInterfaceInfo] = []
+        interface_types = await _get_interface_types_from_networksetup()
+        for iface, services in psutil.net_if_addrs().items():
+            for service in services:
+                match service.family:
+                    case socket.AF_INET | socket.AF_INET6:
+                        iface_type = interface_types.get(iface, "unknown")
+                        if iface_type == "unknown":
+                            iface_type = _classify_unknown_darwin_interface(iface)
+                        interfaces_info.append(
+                            NetworkInterfaceInfo(
+                                name=iface,
+                                ip_address=service.address,
+                                interface_type=iface_type,
+                            )
                         )
-                    )
-                case _:
-                    pass
+                    case _:
+                        pass
+        return interfaces_info
 
-    return interfaces_info
+    return await _get_linux_network_interfaces()
+
+
+def _read_dmi_field(name: str) -> str | None:
+    """Read a single DMI sysfs field, returning ``None`` on failure."""
+    try:
+        path = Path(f"/sys/class/dmi/id/{name}")
+        if path.exists():
+            return path.read_text().strip()
+    except (OSError, PermissionError):
+        pass
+    return None
+
+
+async def _get_linux_model_and_chip() -> tuple[str, str]:
+    """Get Linux system information using DMI and /proc/cpuinfo.
+
+    Detects NVIDIA DGX Spark (DMI product_name ``"DGX_Spark"``) and other
+    NVIDIA systems via the ``sys_vendor`` DMI field, falling back to generic
+    Linux identification.
+    """
+    model = "Linux"
+    chip = "Unknown Chip"
+
+    product_name = _read_dmi_field("product_name")
+    sys_vendor = _read_dmi_field("sys_vendor")
+
+    # DGX Spark: DMI product_name may be "DGX_Spark" or "gx10" variant
+    product_lower = (product_name or "").lower()
+    if product_name and ("dgx" in product_lower or "gx10" in product_lower):
+        model = "DGX Spark"
+        try:
+            process = await run_process(
+                ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"]
+            )
+            gpu_name = process.stdout.decode().strip().split("\n")[0]
+            chip = gpu_name if gpu_name and gpu_name != "[N/A]" else "NVIDIA GB10"
+        except (CalledProcessError, FileNotFoundError):
+            chip = "NVIDIA GB10"
+        return (model, chip)
+
+    # Other NVIDIA systems (sys_vendor contains "NVIDIA")
+    if sys_vendor and "NVIDIA" in sys_vendor:
+        model = product_name.replace("_", " ") if product_name else "NVIDIA System"
+        try:
+            process = await run_process(
+                ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"]
+            )
+            gpu_name = process.stdout.decode().strip().split("\n")[0]
+            if gpu_name and gpu_name != "[N/A]":
+                chip = gpu_name
+        except (CalledProcessError, FileNotFoundError):
+            pass
+        return (model, chip)
+
+    # Generic Linux — detect laptop vs desktop via chassis_type
+    # SMBIOS chassis types: 8,9,10,14,31,32 = portable/laptop
+    chassis_type = _read_dmi_field("chassis_type")
+    laptop_chassis_types = {"8", "9", "10", "14", "31", "32"}
+    if chassis_type in laptop_chassis_types:
+        model = "Linux Laptop"
+    elif chassis_type is not None:
+        model = "Linux Desktop"
+
+    # Also check for battery as a fallback laptop indicator
+    if model == "Linux" and Path("/sys/class/power_supply/BAT0").exists():
+        model = "Linux Laptop"
+
+    # Use /proc/cpuinfo for chip
+    cpuinfo_path = Path("/proc/cpuinfo")
+    if cpuinfo_path.exists():
+        try:
+            for line in cpuinfo_path.read_text().splitlines():
+                if line.startswith("model name"):
+                    chip = line.split(":", 1)[1].strip()
+                    break
+        except OSError:
+            pass
+
+    return (model, chip)
 
 
 async def get_model_and_chip() -> tuple[str, str]:
-    """Get Mac system information using system_profiler."""
+    """Get system model and chip information.
+
+    On macOS, uses ``system_profiler``.  On Linux, reads DMI data from
+    sysfs and CPU info from ``/proc/cpuinfo``.
+    """
     model = "Unknown Model"
     chip = "Unknown Chip"
 
-    # TODO: better non mac support
+    if sys.platform == "linux":
+        return await _get_linux_model_and_chip()
+
     if sys.platform != "darwin":
         return (model, chip)
 

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

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

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

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

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

@@ -6,7 +6,7 @@ import mlx.core as mx
 from mlx_lm.generate import (
     BatchGenerator as MlxBatchGenerator,
 )
-from mlx_lm.models.cache import RotatingKVCache
+from mlx_lm.models.cache import KVCache, RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper
 
@@ -18,13 +18,16 @@ from exo.shared.types.api import (
     TopLogprobItem,
     Usage,
 )
+from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType, Model
+from exo.shared.types.mlx import MLXCacheType, Model
+from exo.shared.types.tasks import TaskId
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.runner_response import GenerationResponse
 from exo.worker.engines.mlx.cache import (
     CacheSnapshot,
     KVPrefixCache,
+    cache_length,
     encode_prompt,
     make_kv_cache,
 )
@@ -34,6 +37,7 @@ from exo.worker.engines.mlx.generator.generate import (
     eos_ids_from_tokenizer,
     extract_top_logprobs,
     prefill,
+    warmup_inference,
 )
 from exo.worker.engines.mlx.utils_mlx import fix_unmatched_think_end_tokens
 from exo.worker.runner.bootstrap import logger
@@ -74,33 +78,39 @@ class ExoBatchGenerator:
     tokenizer: TokenizerWrapper
     group: mx.distributed.Group | None
     kv_prefix_cache: KVPrefixCache | None
+    model_id: ModelId
 
-    _exo_gen: MlxBatchGenerator = field(init=False)
+    _mlx_gen: MlxBatchGenerator = field(init=False)
     _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False)
+    _uid_to_task_id: dict[int, TaskId] = field(default_factory=dict, init=False)
 
     def __post_init__(self) -> None:
-        self._exo_gen = MlxBatchGenerator(
+        self._mlx_gen = MlxBatchGenerator(
             model=self.model,
             stop_tokens=set(eos_ids_from_tokenizer(self.tokenizer)),
             prefill_step_size=4096,
         )
 
+    def warmup(self) -> int:
+        return warmup_inference(self.model, self.tokenizer, self.group, self.model_id)
+
     @property
     def has_work(self) -> bool:
         return (
             bool(self._active_tasks)
-            or bool(self._exo_gen.unprocessed_prompts)
-            or self._exo_gen.active_batch is not None
+            or bool(self._mlx_gen.unprocessed_prompts)
+            or self._mlx_gen.active_batch is not None
         )
 
     def submit(
         self,
+        task_id: TaskId,
         task_params: TextGenerationTaskParams,
         prompt: str,
         on_prefill_progress: Callable[[int, int], None] | None = None,
         distributed_prompt_progress_callback: Callable[[], None] | None = None,
         on_generation_token: Callable[[], None] | None = None,
-    ) -> int:
+    ) -> TaskId:
         all_prompt_tokens = encode_prompt(self.tokenizer, prompt)
         all_prompt_tokens = fix_unmatched_think_end_tokens(
             all_prompt_tokens, self.tokenizer
@@ -140,16 +150,49 @@ class ExoBatchGenerator:
             top_k=task_params.top_k if task_params.top_k is not None else 0,
         )
 
-        _prefill_tps, _prefill_tokens, cache_snapshots = prefill(
-            self.model,
-            self.tokenizer,
-            sampler,
-            prompt_tokens[:-1],
-            cache,
-            self.group,
-            on_prefill_progress,
-            distributed_prompt_progress_callback,
-        )
+        _prefill_tps: float = 0.0
+        cache_snapshots: list[CacheSnapshot] | None = None
+        used_remote_prefill = False
+        uncached_count = len(prompt_tokens)
+        if uncached_count > 1000 and task_params.prefill_endpoints and not is_bench:
+            from exo.disaggregated.prefill_client import remote_prefill
+
+            t0 = time.perf_counter()
+            try:
+                injected_cache, total_tokens = remote_prefill(
+                    endpoint=task_params.prefill_endpoints[0],
+                    token_ids=[int(t) for t in all_prompt_tokens.tolist()],  # type: ignore
+                    model_id=str(task_params.model),
+                    mlx_model=self.model,
+                    on_prefill_progress=on_prefill_progress,
+                    existing_cache=list(cache) if prefix_hit_length > 0 else None,
+                    start_pos=cache_length(cache) if prefix_hit_length > 0 else 0,
+                )
+                cache = injected_cache
+                from exo.worker.engines.mlx.cache import snapshot_ssm_states
+
+                cache_snapshots = [snapshot_ssm_states(cache)]
+                _prefill_tps = total_tokens / max(time.perf_counter() - t0, 0.001)
+                used_remote_prefill = True
+                logger.info(
+                    f"Remote prefill: {total_tokens} tokens at {_prefill_tps:.0f} tok/s"
+                )
+            except Exception:
+                logger.opt(exception=True).warning(
+                    "Remote prefill failed, falling back to local"
+                )
+
+        if not used_remote_prefill:
+            _prefill_tps, _prefill_tokens, cache_snapshots = prefill(
+                self.model,
+                self.tokenizer,
+                sampler,
+                prompt_tokens[:-1],
+                cache,
+                self.group,
+                on_prefill_progress,
+                distributed_prompt_progress_callback,
+            )
 
         # We need to clamp rotating kv caches to max size so that mlx lm's _merge_caches behaves
         for c in cache:
@@ -173,7 +216,11 @@ class ExoBatchGenerator:
                 matched_index,
             )
 
-        last_tokens = prompt_tokens[-2:]
+        last_tokens = (
+            mx.array(all_prompt_tokens[-2:])
+            if used_remote_prefill
+            else prompt_tokens[-2:]
+        )
 
         logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = (
             make_logits_processors(
@@ -188,7 +235,7 @@ class ExoBatchGenerator:
 
         max_tokens = task_params.max_output_tokens or MAX_TOKENS
 
-        uids = self._exo_gen.insert(
+        uids = self._mlx_gen.insert(
             prompts=[last_tokens.tolist()],
             max_tokens=[max_tokens],
             caches=[list(cache)],
@@ -199,6 +246,7 @@ class ExoBatchGenerator:
         assert len(uids) == 1
 
         uid = uids[0]
+        self._uid_to_task_id[uid] = task_id
 
         self._active_tasks[uid] = _EngineTask(
             uid=uid,
@@ -213,15 +261,15 @@ class ExoBatchGenerator:
             prefill_tps=_prefill_tps,
         )
 
-        return uid
+        return task_id
 
-    def step(self) -> list[tuple[int, GenerationResponse]]:
+    def step(self) -> list[tuple[TaskId, GenerationResponse]]:
         if not self.has_work:
             return []
 
-        responses = self._exo_gen.next()
+        responses = self._mlx_gen.next()
 
-        results: list[tuple[int, GenerationResponse]] = []
+        results: list[tuple[TaskId, GenerationResponse]] = []
 
         for response in responses:
             if response.uid not in self._active_tasks:
@@ -288,7 +336,7 @@ class ExoBatchGenerator:
             usage: Usage | None = None
             if is_done:
                 try:
-                    mlx_stats = self._exo_gen.stats()
+                    mlx_stats = self._mlx_gen.stats()
                     generation_tps = mlx_stats.generation_tps
                 except ZeroDivisionError:
                     generation_elapsed = (
@@ -322,7 +370,7 @@ class ExoBatchGenerator:
 
             results.append(
                 (
-                    response.uid,
+                    self._uid_to_task_id.get(response.uid, TaskId(str(response.uid))),
                     GenerationResponse(
                         text=text,
                         token=response.token,
@@ -337,6 +385,7 @@ class ExoBatchGenerator:
 
             if is_done:
                 del self._active_tasks[response.uid]
+                self._uid_to_task_id.pop(response.uid, None)
             elif (
                 max_stop_len > 0
                 and len(state.potential_stop_sequence_text) > max_stop_len
@@ -347,18 +396,23 @@ class ExoBatchGenerator:
 
         return results
 
-    def cancel(self, uids: list[int]) -> None:
-        self._exo_gen.remove(uids)
+    def cancel(self, task_ids: list[TaskId]) -> None:
+        task_id_set = set(task_ids)
+        uids = [uid for uid, tid in self._uid_to_task_id.items() if tid in task_id_set]
+        if uids:
+            self._mlx_gen.remove(uids)
         for uid in uids:
             self._active_tasks.pop(uid, None)
+            self._uid_to_task_id.pop(uid, None)
 
     def close(self) -> None:
-        self._exo_gen.close()
+        self._mlx_gen.close()
+        mx.clear_cache()
 
     def _save_prefix_cache(
         self,
         all_prompt_tokens: mx.array,
-        cache: KVCacheType,
+        cache: MLXCacheType,
         cache_snapshots: list[CacheSnapshot] | None,
         prefix_hit_length: int,
         matched_index: int | None,

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

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

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

@@ -145,6 +145,11 @@ def mlx_distributed_init(
                 os.environ["MLX_JACCL_COORDINATOR"] = jaccl_coordinator
                 group = mx.distributed.init(backend="jaccl", strict=True)
 
+            case _:
+                raise ValueError(
+                    f"Unsupported instance type for MLX init: {type(bound_instance.instance)}"
+                )
+
         logger.info(f"Rank {rank} mlx distributed initialization complete")
 
         return group
@@ -546,7 +551,7 @@ def apply_chat_template(
         )
         if partial_assistant_content:
             prompt += partial_assistant_content
-        logger.info(prompt)
+        logger.debug(prompt)
         return prompt
 
     extra_kwargs: dict[str, Any] = {}
@@ -583,7 +588,7 @@ def apply_chat_template(
     if partial_assistant_content:
         prompt += partial_assistant_content
 
-    logger.info(prompt)
+    logger.debug(prompt)
 
     return prompt
 

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

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

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

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

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

@@ -0,0 +1,346 @@
+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,57 @@
+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
+
+    return SamplingParams(**kwargs)

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

@@ -0,0 +1,644 @@
+import gc
+import math
+import os
+import re
+import sys
+import tempfile
+import time
+from collections.abc import Callable, Generator
+from dataclasses import dataclass, field
+
+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.shared.types.api 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: str,
+    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
+    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 []
+        )
+    if is_nvfp4 and not has_mamba:
+        backends = ["FLASHINFER", "FLASH_ATTN", "TRITON_ATTN"]
+    else:
+        backends = ["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) as e:
+            logger.warning(f"Attention backend {backend} failed: {e}, trying next")
+            continue
+
+    if engine is None:
+        raise RuntimeError(f"No attention backend worked for {model_id}")
+
+    tool_parser: ToolParser | None = None
+    tokenizer = engine.get_tokenizer()
+    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