Merge main

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

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

flake.nix

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

justfile

@@ -15,6 +15,18 @@ check:
 sync:
     uv sync --all-packages
 
+sync-cuda:
+    #!/usr/bin/env bash
+    set -euo pipefail
+    if command -v nvidia-smi &>/dev/null; then
+        arch=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -1 | tr -d ' ')
+        export TORCH_CUDA_ARCH_LIST="$arch"
+        export VLLM_TARGET_DEVICE=cuda
+    fi
+    uv pip install "cmake>=3.26.1" ninja "packaging>=24.2" "setuptools>=77.0.3,<81.0.0" "setuptools-scm>=8.0" wheel jinja2
+    find ~/.cache/uv/git-v0 -name CMakeCache.txt -delete 2>/dev/null || true
+    uv sync --extra cuda
+
 sync-clean:
     uv sync --all-packages --force-reinstall --no-cache
 

nix/cuda-pkgs.nix

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

pyproject.toml

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

python/parts.nix

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

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/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/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)
@@ -448,18 +449,29 @@ class API:
                 status_code=400, detail=f"Failed to load model card: {exc}"
             ) from exc
         instance_combinations: list[tuple[Sharding, InstanceMeta, int]] = []
-        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
-                        )
-                    ]
-                )
-        # TODO: PDD
-        # instance_combinations.append((Sharding.PrefillDecodeDisaggregation, InstanceMeta.MlxRing, 1))
+        node_count = len(list(self.state.topology.list_nodes()))
+
+        # QMM is not available on MLX CUDA. Also, VLLM does not support MLX community models
+        is_mlx_community = str(model_card.model_id).startswith("mlx-community/")
+        is_quantized_mlx = is_mlx_community and model_card.quantization in (
+            "4bit",
+            "8bit",
+        )
+        skip_mlx = any(self.state.node_vllm.values()) and is_quantized_mlx
+        is_vllm_compatible_mlx = is_mlx_community and model_card.quantization in (
+            "",
+            "bf16",
+            "fp16",
+        )
+        skip_vllm = is_mlx_community and not is_vllm_compatible_mlx
+        if not skip_mlx:
+            for sharding in (Sharding.Pipeline, Sharding.Tensor):
+                for instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl):
+                    instance_combinations.extend(
+                        [(sharding, instance_meta, i) for i in range(1, node_count + 1)]
+                    )
+        if any(self.state.node_vllm.values()) and not skip_vllm:
+            instance_combinations.append((Sharding.Pipeline, InstanceMeta.Vllm, 1))
 
         for sharding, instance_meta, min_nodes in instance_combinations:
             try:
@@ -781,6 +793,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/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
 
@@ -78,8 +79,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 +143,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 +203,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/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

@@ -258,21 +258,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/tracing.py

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

src/exo/shared/types/mlx.py

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

src/exo/shared/types/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/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/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
@@ -117,12 +118,97 @@ async def get_network_interfaces() -> list[NetworkInterfaceInfo]:
     return interfaces_info
 
 
+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,7 +192,8 @@ 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)
         target = (max_length - 1) if is_exact and not has_ssm else best_length
         restore_pos, restore_snap = self._get_snapshot(best_index, target)
 
@@ -192,8 +201,8 @@ class KVPrefixCache:
         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 +217,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 +297,7 @@ class KVPrefixCache:
 
 
 def trim_cache(
-    cache: KVCacheType,
+    cache: MLXCacheType,
     num_tokens: int,
     snapshot: CacheSnapshot | None = None,
 ) -> None:
@@ -282,7 +335,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 +364,7 @@ def get_memory_used_percentage() -> float:
 
 def make_kv_cache(
     model: Model, max_kv_size: int | None = None, keep: int = 0
-) -> KVCacheType:
+) -> MLXCacheType:
     assert hasattr(model, "layers")
 
     if hasattr(model, "make_cache"):

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

@@ -18,8 +18,10 @@ 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 (
@@ -34,6 +36,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 +77,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
@@ -188,7 +197,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 +208,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 +223,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 +298,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 +332,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 +347,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 +358,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

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

@@ -0,0 +1,413 @@
+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:
+        try:
+            original(self)
+        except AssertionError:
+            logger.warning(
+                "vLLM memory profiling assertion failed (free memory changed during init, "
+                "likely another process released GPU memory). Continuing with growable cache."
+            )
+        torch.cuda.empty_cache()
+        free_bytes, _ = torch.cuda.mem_get_info()
+        initial = max(int(free_bytes * INITIAL_FRACTION), 1)
+        self._growable_max_kv_bytes = free_bytes
+        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:
+        logger.debug("No model_runner reference — cannot grow cache")
+        return False
+
+    free_bytes, _ = torch.cuda.mem_get_info()
+    if free_bytes < GROWTH_HEADROOM_BYTES:
+        logger.debug(f"Only {free_bytes / (1024**3):.2f} GiB free — not enough to grow")
+        return False
+
+    kv_cache_config = model_runner._growable_kv_cache_config  # type: ignore
+    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:
+        logger.debug(f"Growth too small ({growth_blocks} 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:
+        _grow_tensors(model_runner, kv_cache_config, old_num_blocks, new_num_blocks)
+        _grow_block_pool(block_pool, old_num_blocks, new_num_blocks)
+        kv_cache_config.num_blocks = new_num_blocks
+        for tensor_spec in kv_cache_config.kv_cache_tensors:
+            tensor_spec.size = int(tensor_spec.size * new_num_blocks / old_num_blocks)
+        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
+    runner_kv_caches.clear()
+
+    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)
+
+    for layer_index in sorted(index2name.keys()):
+        for ln in index2name[layer_index]:
+            runner_kv_caches.append(new_kv_caches[ln])
+
+    for layer_name, kv_cache in new_kv_caches.items():
+        forward_context[layer_name].kv_cache = [kv_cache]  # type: ignore
+
+
+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,306 @@
+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,
+    ) -> "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
+
+            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 not isinstance(layer, KVLayerState):
+                continue
+            gi = layer_to_group[layer_idx]
+            bt = block_tables[gi]
+            kv = kv_caches[layer_idx]
+            k_all, v_all = _split_kv(kv)
+            n_blocks = min(len(bt), layer.keys.shape[0])
+            if n_blocks > 0:
+                k_all[bt[:n_blocks]] = layer.keys[:n_blocks].to(
+                    device, non_blocking=True
+                )
+                v_all[bt[:n_blocks]] = layer.values[:n_blocks].to(
+                    device, non_blocking=True
+                )
+        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,594 @@
+import gc
+import math
+import os
+import re
+import sys
+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.info(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)
+    max_batch_tokens: int = (
+        getattr(engine.model_config, "max_num_batched_tokens", 2048) or 2048
+    )  # 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.info(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,
+) -> tuple[LLMEngine, ToolParser | None, KVPrefixCache]:
+    patch_vllm()
+    _patch_weight_loading_progress()
+
+    os.environ.setdefault("FASTSAFETENSORS_NOGDS", "1")
+
+    prefix_cache = KVPrefixCache(group=None)
+    set_prefix_cache(prefix_cache)
+    set_n_layers(n_layers)
+
+    engine_args = EngineArgs(
+        model=model_path,
+        served_model_name=str(model_id),
+        gpu_memory_utilization=0.05,
+        trust_remote_code=trust_remote_code,
+        load_format="fastsafetensors",
+        enable_prefix_caching=False,
+        attention_backend="TRITON_ATTN",
+        enforce_eager=True,
+        disable_log_stats=True,
+    )
+
+    set_weight_loading_callback(on_layer_loaded)
+    engine = LLMEngine.from_engine_args(engine_args)
+
+    tool_parser: ToolParser | None = None
+    tokenizer = engine.get_tokenizer()
+    chat_template = getattr(tokenizer, "chat_template", None)
+    if isinstance(chat_template, str):
+        tool_parser = infer_tool_parser(chat_template)
+        if tool_parser:
+            logger.info(
+                f"inferred tool parser: {tool_parser.start_parsing} / {tool_parser.end_parsing}"
+            )
+
+    logger.info(f"vLLM engine loaded for {model_id}")
+
+    return engine, tool_parser, prefix_cache

src/exo/worker/runner/bootstrap.py

@@ -1,16 +1,48 @@
+import ctypes
 import os
 import resource
+import sys
+from pathlib import Path
 
 import loguru
 
 from exo.shared.types.events import Event, RunnerStatusUpdated
 from exo.shared.types.tasks import Task, TaskId
-from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.instances import BoundInstance, VllmInstance
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
 
 logger: "loguru.Logger" = loguru.logger
 
+_CUDA_HOST_LIBS = ["libcuda.so.1", "libnvidia-ml.so.1", "libnvidia-ptxjitcompiler.so.1"]
+_CUDA_HOST_SEARCH_DIRS = [
+    Path("/usr/lib/aarch64-linux-gnu"),
+    Path("/usr/lib/x86_64-linux-gnu"),
+    Path("/usr/lib64"),
+    Path("/usr/lib"),
+    Path("/usr/local/cuda/lib64"),
+    Path("/usr/local/cuda/compat"),
+]
+
+
+def _ensure_cuda_libs() -> None:
+    if sys.platform != "linux":
+        return
+    for search_dir in _CUDA_HOST_SEARCH_DIRS:
+        driver = search_dir / "libcuda.so.1"
+        if not driver.exists():
+            continue
+        for lib_name in _CUDA_HOST_LIBS:
+            lib_path = search_dir / lib_name
+            if lib_path.exists():
+                try:
+                    ctypes.CDLL(str(lib_path), mode=ctypes.RTLD_GLOBAL)
+                    logger.info(f"Loaded CUDA host lib: {lib_path}")
+                except OSError:
+                    logger.warning(f"Failed to load {lib_path}")
+                    raise
+        return
+
 
 def entrypoint(
     bound_instance: BoundInstance,
@@ -35,7 +67,27 @@ def entrypoint(
 
     # Import main after setting global logger - this lets us just import logger from this module
     try:
-        if bound_instance.is_image_model:
+        if isinstance(bound_instance.instance, VllmInstance):
+            os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+            os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
+            os.environ["VLLM_BATCH_INVARIANT"] = "1"
+            _ensure_cuda_libs()
+            from exo.shared.constants import EXO_MODELS_DIR
+            from exo.worker.runner.llm_inference.runner import Runner, VllmBuilder
+
+            model_id = bound_instance.bound_shard.model_card.model_id
+            builder = VllmBuilder(
+                model_id=model_id,
+                model_path=str(EXO_MODELS_DIR / model_id.normalize()),
+                trust_remote_code=bound_instance.bound_shard.model_card.trust_remote_code,
+                cancel_receiver=cancel_receiver,
+                event_sender=event_sender,
+            )
+            runner = Runner(
+                bound_instance, event_sender, task_receiver, cancel_receiver, builder
+            )
+            runner.main()
+        elif bound_instance.is_image_model:
             from exo.worker.runner.image_models.runner import Runner as ImageRunner
 
             runner = ImageRunner(
@@ -43,10 +95,15 @@ def entrypoint(
             )
             runner.main()
         else:
-            from exo.worker.runner.llm_inference.runner import Runner
+            from exo.worker.runner.llm_inference.runner import MlxBuilder, Runner
 
+            builder = MlxBuilder(
+                model_id=bound_instance.bound_shard.model_card.model_id,
+                event_sender=event_sender,
+                cancel_receiver=cancel_receiver,
+            )
             runner = Runner(
-                bound_instance, event_sender, task_receiver, cancel_receiver
+                bound_instance, event_sender, task_receiver, cancel_receiver, builder
             )
             runner.main()
 

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

@@ -2,8 +2,9 @@ import itertools
 import time
 from abc import ABC, abstractmethod
 from collections import deque
-from collections.abc import Generator, Iterable
+from collections.abc import Callable, Generator, Iterable
 from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
 
 import mlx.core as mx
 from mlx_lm.tokenizer_utils import TokenizerWrapper
@@ -12,17 +13,17 @@ from exo.shared.constants import EXO_MAX_CONCURRENT_REQUESTS
 from exo.shared.types.chunks import ErrorChunk, PrefillProgressChunk
 from exo.shared.types.common import ModelId
 from exo.shared.types.events import ChunkGenerated, Event
-from exo.shared.types.mlx import Model
 from exo.shared.types.tasks import CANCEL_ALL_TASKS, TaskId, TextGeneration
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
 from exo.utils.channels import MpReceiver, MpSender
 from exo.worker.engines.mlx.cache import KVPrefixCache
 from exo.worker.engines.mlx.generator.batch_generate import ExoBatchGenerator
+
+if TYPE_CHECKING:
+    from exo.worker.engines.vllm.vllm_generator import VllmBatchEngine
 from exo.worker.engines.mlx.generator.generate import (
     PrefillCancelled,
-    mlx_generate,
-    warmup_inference,
 )
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
@@ -111,7 +112,6 @@ def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None:
 
 @dataclass(eq=False)
 class SequentialGenerator(InferenceGenerator):
-    model: Model
     tokenizer: TokenizerWrapper
     group: mx.distributed.Group | None
     kv_prefix_cache: KVPrefixCache | None
@@ -120,6 +120,8 @@ class SequentialGenerator(InferenceGenerator):
     device_rank: int
     cancel_receiver: MpReceiver[TaskId]
     event_sender: MpSender[Event]
+    _generate_fn: Callable[..., Generator[GenerationResponse]]
+    _warmup_fn: Callable[[], int]
     check_for_cancel_every: int = 50
 
     _cancelled_tasks: set[TaskId] = field(default_factory=set, init=False)
@@ -140,13 +142,8 @@ class SequentialGenerator(InferenceGenerator):
         | None
     ) = field(default=None, init=False)
 
-    def warmup(self):
-        self.check_for_cancel_every = warmup_inference(
-            model=self.model,
-            tokenizer=self.tokenizer,
-            group=self.group,
-            model_id=self.model_id,
-        )
+    def warmup(self) -> None:
+        self.check_for_cancel_every = self._warmup_fn()
 
     def submit(
         self,
@@ -230,7 +227,6 @@ class SequentialGenerator(InferenceGenerator):
                 apply_chat_template(self.tokenizer, task.task_params),
                 self.tool_parser,
                 self.tokenizer,
-                type(self.model),
                 self.model_id,
                 task.task_params.tools,
             )
@@ -286,9 +282,7 @@ class SequentialGenerator(InferenceGenerator):
 
                 self.agree_on_tasks()
 
-        return mlx_generate(
-            model=self.model,
-            tokenizer=self.tokenizer,
+        return self._generate_fn(
             task=task.task_params,
             prompt=prompt,
             kv_prefix_cache=self.kv_prefix_cache,
@@ -299,12 +293,11 @@ class SequentialGenerator(InferenceGenerator):
         )
 
     def close(self) -> None:
-        del self.model, self.tokenizer, self.group
+        del self.tokenizer, self.group
 
 
 @dataclass(eq=False)
 class BatchGenerator(InferenceGenerator):
-    model: Model
     tokenizer: TokenizerWrapper
     group: mx.distributed.Group | None
     kv_prefix_cache: KVPrefixCache | None
@@ -313,6 +306,8 @@ class BatchGenerator(InferenceGenerator):
     device_rank: int
     cancel_receiver: MpReceiver[TaskId]
     event_sender: MpSender[Event]
+    _gen: "ExoBatchGenerator | VllmBatchEngine"
+    max_concurrent_requests: int = EXO_MAX_CONCURRENT_REQUESTS
     check_for_cancel_every: int = 50
 
     _cancelled_tasks: set[TaskId] = field(default_factory=set, init=False)
@@ -320,9 +315,8 @@ class BatchGenerator(InferenceGenerator):
     _maybe_cancel: list[TextGeneration] = field(default_factory=list, init=False)
     _all_tasks: dict[TaskId, TextGeneration] = field(default_factory=dict, init=False)
     _queue: deque[TextGeneration] = field(default_factory=deque, init=False)
-    _mlx_gen: ExoBatchGenerator = field(init=False)
     _active_tasks: dict[
-        int,
+        TaskId,
         tuple[
             TextGeneration,
             GeneratorQueue[GenerationResponse],
@@ -330,21 +324,8 @@ class BatchGenerator(InferenceGenerator):
         ],
     ] = field(default_factory=dict, init=False)
 
-    def __post_init__(self) -> None:
-        self._mlx_gen = ExoBatchGenerator(
-            model=self.model,
-            tokenizer=self.tokenizer,
-            group=self.group,
-            kv_prefix_cache=self.kv_prefix_cache,
-        )
-
-    def warmup(self):
-        self.check_for_cancel_every = warmup_inference(
-            model=self.model,
-            tokenizer=self.tokenizer,
-            group=self.group,
-            model_id=self.model_id,
-        )
+    def warmup(self) -> None:
+        self.check_for_cancel_every = self._gen.warmup()
 
     def submit(
         self,
@@ -386,10 +367,10 @@ class BatchGenerator(InferenceGenerator):
             self.agree_on_tasks()
 
         # Submit any queued tasks to the engine
-        while self._queue and len(self._active_tasks) < EXO_MAX_CONCURRENT_REQUESTS:
+        while self._queue and len(self._active_tasks) < self.max_concurrent_requests:
             task = self._queue.popleft()
             try:
-                uid = self._start_task(task)
+                task_id = self._start_task(task)
             except PrefillCancelled:
                 continue
             except Exception as e:
@@ -405,16 +386,15 @@ class BatchGenerator(InferenceGenerator):
                     apply_chat_template(self.tokenizer, task.task_params),
                     self.tool_parser,
                     self.tokenizer,
-                    type(self.model),
                     self.model_id,
                     task.task_params.tools,
                 )
-            self._active_tasks[uid] = (task, queue, output_generator)
+            self._active_tasks[task_id] = (task, queue, output_generator)
 
-        if not self._mlx_gen.has_work:
+        if not self._gen.has_work:
             return self._apply_cancellations()
 
-        results = self._mlx_gen.step()
+        results = self._gen.step()
 
         output: list[
             tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
@@ -446,17 +426,17 @@ class BatchGenerator(InferenceGenerator):
 
         cancel_all = CANCEL_ALL_TASKS in self._cancelled_tasks
 
-        uids_to_cancel: list[int] = []
+        ids_to_cancel: list[TaskId] = []
         results: list[tuple[TaskId, Cancelled]] = []
 
-        for uid, (task, _, _) in list(self._active_tasks.items()):
+        for tid, (task, _, _) in list(self._active_tasks.items()):
             if task.task_id in self._cancelled_tasks or cancel_all:
-                uids_to_cancel.append(uid)
+                ids_to_cancel.append(tid)
                 results.append((task.task_id, Cancelled()))
-                del self._active_tasks[uid]
+                del self._active_tasks[tid]
 
-        if uids_to_cancel:
-            self._mlx_gen.cancel(uids_to_cancel)
+        if ids_to_cancel:
+            self._gen.cancel(ids_to_cancel)
 
         already_cancelled = {tid for tid, _ in results}
         for tid in self._cancelled_tasks:
@@ -479,7 +459,7 @@ class BatchGenerator(InferenceGenerator):
                 )
             )
 
-    def _start_task(self, task: TextGeneration) -> int:
+    def _start_task(self, task: TextGeneration) -> TaskId:
         _check_for_debug_prompts(task.task_params)
         prompt = apply_chat_template(self.tokenizer, task.task_params)
 
@@ -516,7 +496,8 @@ class BatchGenerator(InferenceGenerator):
 
                 self.agree_on_tasks()
 
-        return self._mlx_gen.submit(
+        return self._gen.submit(
+            task_id=task.task_id,
             task_params=task.task_params,
             prompt=prompt,
             on_prefill_progress=on_prefill_progress,
@@ -525,5 +506,5 @@ class BatchGenerator(InferenceGenerator):
         )
 
     def close(self) -> None:
-        self._mlx_gen.close()
-        del self.model, self.tokenizer, self.group
+        self._gen.close()
+        del self.tokenizer, self.group

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

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

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

@@ -1,7 +1,12 @@
+import contextlib
+import gc
 import os
 import time
+from abc import ABC, abstractmethod
+from collections.abc import Callable
 from dataclasses import dataclass
 from enum import Enum
+from typing import TYPE_CHECKING
 
 import mlx.core as mx
 from anyio import WouldBlock
@@ -67,12 +72,34 @@ from exo.worker.runner.llm_inference.batch_generator import (
 from .batch_generator import Cancelled, Finished
 from .tool_parsers import make_mlx_parser
 
+if TYPE_CHECKING:
+    pass
+
 
 class ExitCode(str, Enum):
     AllTasksComplete = "AllTasksComplete"
     Shutdown = "Shutdown"
 
 
+class Builder(ABC):
+    @abstractmethod
+    def connect(self, bound_instance: BoundInstance) -> None: ...
+
+    @abstractmethod
+    def load(
+        self,
+        bound_instance: BoundInstance,
+        on_timeout: Callable[[], None],
+        on_layer_loaded: Callable[[int, int], None],
+    ) -> None: ...
+
+    @abstractmethod
+    def build(self) -> InferenceGenerator: ...
+
+    @abstractmethod
+    def close(self) -> None: ...
+
+
 class Runner:
     def __init__(
         self,
@@ -80,6 +107,7 @@ class Runner:
         event_sender: MpSender[Event],
         task_receiver: MpReceiver[Task],
         cancel_receiver: MpReceiver[TaskId],
+        builder: Builder,
     ):
         self.event_sender = event_sender
         self.task_receiver = task_receiver
@@ -102,9 +130,7 @@ class Runner:
 
         self.setup_start_time = time.time()
 
-        self.generator: Builder | InferenceGenerator = Builder(
-            self.model_id, self.event_sender, self.cancel_receiver
-        )
+        self.generator: Builder | InferenceGenerator = builder
 
         self.seen: set[TaskId] = set()
         self.active_tasks: dict[
@@ -132,15 +158,19 @@ class Runner:
         self.event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
     def main(self):
-        with self.task_receiver:
-            for task in self.task_receiver:
-                if task.task_id in self.seen:
-                    logger.warning("repeat task - potential error")
-                    continue
-                self.seen.add(task.task_id)
-                self.handle_first_task(task)
-                if isinstance(self.current_status, RunnerShutdown):
-                    break
+        try:
+            with self.task_receiver:
+                for task in self.task_receiver:
+                    if task.task_id in self.seen:
+                        logger.warning("repeat task - potential error")
+                        continue
+                    self.seen.add(task.task_id)
+                    self.handle_first_task(task)
+                    if isinstance(self.current_status, RunnerShutdown):
+                        break
+        finally:
+            if not isinstance(self.current_status, RunnerShutdown):
+                self.generator.close()
 
     def handle_first_task(self, task: Task):
         self.send_task_status(task.task_id, TaskStatus.Running)
@@ -154,22 +184,28 @@ class Runner:
                 self.update_status(RunnerConnecting())
                 self.acknowledge_task(task)
 
-                self.generator.group = initialize_mlx(self.bound_instance)
+                self.generator.connect(self.bound_instance)
 
                 self.send_task_status(task.task_id, TaskStatus.Complete)
                 self.update_status(RunnerConnected())
                 logger.info("runner connected")
 
             # we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
-            case LoadModel() if isinstance(self.generator, Builder) and (
+            case LoadModel() if (
                 (
-                    isinstance(self.current_status, RunnerConnected)
+                    isinstance(self.generator, MlxBuilder)
+                    and isinstance(self.current_status, RunnerConnected)
                     and self.generator.group is not None
                 )
                 or (
-                    isinstance(self.current_status, RunnerIdle)
+                    isinstance(self.generator, MlxBuilder)
+                    and isinstance(self.current_status, RunnerIdle)
                     and self.generator.group is None
                 )
+                or (
+                    isinstance(self.generator, VllmBuilder)
+                    and isinstance(self.current_status, RunnerIdle)
+                )
             ):
                 total_layers = (
                     self.shard_metadata.end_layer - self.shard_metadata.start_layer
@@ -195,15 +231,12 @@ class Runner:
                 assert (
                     ModelTask.TextGeneration in self.shard_metadata.model_card.tasks
                 ), f"Incorrect model task(s): {self.shard_metadata.model_card.tasks}"
-                self.generator.inference_model, self.generator.tokenizer = (
-                    load_mlx_items(
-                        self.bound_instance,
-                        self.generator.group,
-                        on_timeout=on_model_load_timeout,
-                        on_layer_loaded=on_layer_loaded,
-                    )
-                )
 
+                self.generator.load(
+                    self.bound_instance,
+                    on_timeout=on_model_load_timeout,
+                    on_layer_loaded=on_layer_loaded,
+                )
                 self.generator = self.generator.build()
 
                 self.send_task_status(task.task_id, TaskStatus.Complete)
@@ -245,11 +278,7 @@ class Runner:
         logger.info("runner shutting down")
         self.update_status(RunnerShuttingDown())
         self.acknowledge_task(task)
-        if isinstance(self.generator, InferenceGenerator):
-            self.generator.close()
-        mx.clear_cache()
-        import gc
-
+        self.generator.close()
         gc.collect()
         self.send_task_status(task.task_id, TaskStatus.Complete)
         self.update_status(RunnerShutdown())
@@ -372,7 +401,7 @@ class Runner:
 
 
 @dataclass
-class Builder:
+class MlxBuilder(Builder):
     model_id: ModelId
     event_sender: MpSender[Event]
     cancel_receiver: MpReceiver[TaskId]
@@ -380,9 +409,23 @@ class Builder:
     tokenizer: TokenizerWrapper | None = None
     group: mx.distributed.Group | None = None
 
-    def build(
+    def connect(self, bound_instance: BoundInstance) -> None:
+        self.group = initialize_mlx(bound_instance)
+
+    def load(
         self,
-    ) -> InferenceGenerator:
+        bound_instance: BoundInstance,
+        on_timeout: Callable[[], None],
+        on_layer_loaded: Callable[[int, int], None],
+    ) -> None:
+        self.inference_model, self.tokenizer = load_mlx_items(
+            bound_instance,
+            self.group,
+            on_timeout=on_timeout,
+            on_layer_loaded=on_layer_loaded,
+        )
+
+    def build(self) -> InferenceGenerator:
         assert self.model_id
         assert self.inference_model
         assert self.tokenizer
@@ -404,11 +447,28 @@ class Builder:
 
         kv_prefix_cache = KVPrefixCache(self.group)
 
+        from functools import partial
+
+        from exo.worker.engines.mlx.generator.generate import (
+            mlx_generate,
+            warmup_inference,
+        )
+
         device_rank = 0 if self.group is None else self.group.rank()
+        generate_fn = partial(
+            mlx_generate, model=self.inference_model, tokenizer=self.tokenizer
+        )
+        warmup_fn = partial(
+            warmup_inference,
+            model=self.inference_model,
+            tokenizer=self.tokenizer,
+            group=self.group,
+            model_id=self.model_id,
+        )
+
         if os.environ.get("EXO_NO_BATCH"):
             logger.info("using SequentialGenerator (batching disabled)")
             return SequentialGenerator(
-                model=self.inference_model,
                 tokenizer=self.tokenizer,
                 group=self.group,
                 tool_parser=tool_parser,
@@ -417,10 +477,20 @@ class Builder:
                 device_rank=device_rank,
                 cancel_receiver=self.cancel_receiver,
                 event_sender=self.event_sender,
+                _generate_fn=generate_fn,
+                _warmup_fn=warmup_fn,
             )
+        from exo.worker.runner.llm_inference.batch_generator import ExoBatchGenerator
+
         logger.info("using BatchGenerator")
-        return BatchGenerator(
+        gen = ExoBatchGenerator(
             model=self.inference_model,
+            tokenizer=self.tokenizer,
+            group=self.group,
+            kv_prefix_cache=kv_prefix_cache,
+            model_id=self.model_id,
+        )
+        return BatchGenerator(
             tokenizer=self.tokenizer,
             group=self.group,
             tool_parser=tool_parser,
@@ -429,4 +499,69 @@ class Builder:
             device_rank=device_rank,
             cancel_receiver=self.cancel_receiver,
             event_sender=self.event_sender,
+            _gen=gen,
         )
+
+    def close(self):
+        with contextlib.suppress(NameError, AttributeError):
+            del self.inference_model, self.tokenizer
+
+
+@dataclass
+class VllmBuilder(Builder):
+    model_id: ModelId
+    model_path: str
+    trust_remote_code: bool
+    cancel_receiver: MpReceiver[TaskId]
+    event_sender: MpSender[Event]
+    group: mx.distributed.Group | None = None
+
+    def connect(self, bound_instance: BoundInstance) -> None:
+        raise NotImplementedError(
+            "Multiple node VLLM instances are not supported at the moment!"
+        )
+
+    def load(
+        self,
+        bound_instance: BoundInstance,
+        on_timeout: Callable[[], None],
+        on_layer_loaded: Callable[[int, int], None],
+    ) -> None:
+        from exo.worker.engines.vllm.vllm_generator import load_vllm_engine
+
+        self._engine, self._tool_parser, self._prefix_cache = load_vllm_engine(
+            model_path=self.model_path,
+            model_id=self.model_id,
+            trust_remote_code=self.trust_remote_code,
+            n_layers=bound_instance.bound_shard.model_card.n_layers,
+            on_layer_loaded=on_layer_loaded,
+        )
+
+    def build(self) -> InferenceGenerator:
+        from exo.worker.engines.vllm.vllm_generator import VllmBatchEngine
+
+        gen = VllmBatchEngine(
+            engine=self._engine,
+            model_id=self.model_id,
+            prefix_cache=self._prefix_cache,
+        )
+        tokenizer = TokenizerWrapper(self._engine.get_tokenizer())
+        max_concurrent = 1 if os.environ.get("EXO_NO_BATCH") else 8
+
+        logger.info(f"using BatchGenerator (vLLM, max_concurrent={max_concurrent})")
+        return BatchGenerator(
+            tokenizer=tokenizer,
+            group=None,
+            tool_parser=self._tool_parser,
+            kv_prefix_cache=None,
+            model_id=self.model_id,
+            device_rank=0,
+            cancel_receiver=self.cancel_receiver,
+            event_sender=self.event_sender,
+            _gen=gen,
+            max_concurrent_requests=max_concurrent,
+        )
+
+    def close(self) -> None:
+        with contextlib.suppress(NameError, AttributeError):
+            del self._engine, self._prefix_cache, self._tool_parser

src/exo/worker/tests/unittests/test_mlx/test_batch_vs_generate.py

@@ -1,389 +0,0 @@
-import copy
-import gc
-import json
-import shutil
-import tempfile
-from pathlib import Path
-from typing import Any, cast
-
-import mlx.core as mx
-import pytest
-from mlx_lm.tokenizer_utils import TokenizerWrapper
-
-from exo.shared.types.common import ModelId
-from exo.shared.types.mlx import KVCacheType, Model
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
-from exo.worker.engines.mlx.cache import CacheSnapshot, KVPrefixCache, cache_length
-from exo.worker.engines.mlx.generator.batch_generate import ExoBatchGenerator
-from exo.worker.engines.mlx.generator.generate import mlx_generate
-from exo.worker.engines.mlx.utils_mlx import (
-    apply_chat_template,
-    load_tokenizer_for_model_id,
-)
-
-from .test_prefix_cache_architectures import (
-    ARCHITECTURES,
-    ArchSpec,
-    _arch_available,  # pyright: ignore[reportPrivateUsage]
-    _build_model,  # pyright: ignore[reportPrivateUsage]
-    _copy_tokenizer,  # pyright: ignore[reportPrivateUsage]
-    _find_snapshot,  # pyright: ignore[reportPrivateUsage]
-    _reduce_config,  # pyright: ignore[reportPrivateUsage]
-)
-
-
-def _make_task(
-    content: str = "Hello, what is 2+2?",
-    max_tokens: int = 10,
-    seed: int = 42,
-) -> TextGenerationTaskParams:
-    return TextGenerationTaskParams(
-        model=ModelId("test"),
-        input=[InputMessage(role="user", content=content)],
-        max_output_tokens=max_tokens,
-        temperature=0.7,
-        seed=seed,
-    )
-
-
-# ── Helpers ──────────────────────────────────────────────────────────────── #
-
-
-def _collect_mlx_generate(
-    model: Model,
-    tokenizer: TokenizerWrapper,
-    task: TextGenerationTaskParams,
-    kv_prefix_cache: KVPrefixCache | None,
-) -> list[int]:
-    """Run mlx_generate and collect output token IDs."""
-    prompt = apply_chat_template(tokenizer=tokenizer, task_params=task)
-    tokens: list[int] = []
-    for resp in mlx_generate(
-        model=model,
-        tokenizer=tokenizer,
-        task=task,
-        prompt=prompt,
-        kv_prefix_cache=kv_prefix_cache,
-        group=None,
-    ):
-        tokens.append(resp.token)
-        if resp.finish_reason is not None:
-            break
-    return tokens
-
-
-def _collect_batch_generate(
-    model: Model,
-    tokenizer: TokenizerWrapper,
-    task_params: TextGenerationTaskParams,
-    kv_prefix_cache: KVPrefixCache | None,
-) -> list[int]:
-    """Run ExoBatchGenerator and collect raw output token IDs"""
-    exo_gen = ExoBatchGenerator(
-        model=model,
-        tokenizer=tokenizer,
-        group=None,
-        kv_prefix_cache=kv_prefix_cache,
-    )
-
-    prompt = apply_chat_template(tokenizer=tokenizer, task_params=task_params)
-    exo_gen.submit(task_params=task_params, prompt=prompt)
-
-    tokens: list[int] = []
-    while exo_gen.has_work:
-        results = exo_gen.step()
-        for _uid, response in results:
-            tokens.append(response.token)
-
-    exo_gen.close()
-    return tokens
-
-
-def _assert_state_equal(sa: object, sb: object, label: str) -> None:
-    """Compare two state items, handling both plain arrays and tuples of arrays (CacheList)."""
-    if isinstance(sa, tuple):
-        assert isinstance(sb, tuple), f"{label}: type mismatch"
-        for k, (arr_a, arr_b) in enumerate(
-            zip(
-                cast(tuple[mx.array, ...], sa),
-                cast(tuple[mx.array, ...], sb),
-                strict=True,
-            )
-        ):
-            a_f = mx.array(arr_a).astype(mx.float32)
-            b_f = mx.array(arr_b).astype(mx.float32)
-            if a_f.size == 0:
-                assert b_f.size == 0, f"{label}[{k}]: size mismatch"
-                continue
-            diff = float(mx.max(mx.abs(a_f - b_f)).item())
-            assert diff == 0.0, f"{label}[{k}]: max diff {diff}"
-    else:
-        sa_f = mx.array(cast(mx.array, sa)).astype(mx.float32)
-        sb_f = mx.array(cast(mx.array, sb)).astype(mx.float32)
-        if sa_f.size == 0:
-            assert sb_f.size == 0, f"{label}: size mismatch"
-            return
-        diff = float(mx.max(mx.abs(sa_f - sb_f)).item())
-        assert diff == 0.0, f"{label}: max diff {diff}"
-
-
-def _compare_cache_arrays(
-    cache_a: KVCacheType,
-    cache_b: KVCacheType,
-    label: str = "",
-) -> None:
-    """Assert two KV caches have identical array values."""
-    assert len(cache_a) == len(cache_b), (
-        f"{label}Cache layer count: {len(cache_a)} vs {len(cache_b)}"
-    )
-    for i, (a, b) in enumerate(zip(cache_a, cache_b, strict=True)):
-        assert type(a) is type(b), (
-            f"{label}Layer {i}: type {type(a).__name__} vs {type(b).__name__}"
-        )
-        states_a = a.state
-        states_b = b.state
-        assert len(states_a) == len(states_b), (
-            f"{label}Layer {i}: state count {len(states_a)} vs {len(states_b)}"
-        )
-        for j, (sa, sb) in enumerate(zip(states_a, states_b, strict=True)):
-            if sa is None and sb is None:
-                continue
-            assert sa is not None and sb is not None, (
-                f"{label}Layer {i}, state {j}: one is None"
-            )
-            _assert_state_equal(sa, sb, f"{label}Layer {i}, state {j}")
-
-
-def _safe_state(cache: object) -> list[object]:
-    """Safely access .state on a cache object. Returns [] if uninitialized."""
-    # RotatingKVCache.state crashes when keys is None (uninitialized)
-    if getattr(cache, "keys", _SENTINEL) is None:
-        return []
-    try:
-        return list(cache.state)  # type: ignore[union-attr]
-    except (AttributeError, TypeError):
-        return []
-
-
-_SENTINEL = object()
-
-
-def _compare_snapshots(
-    snaps_a: list[CacheSnapshot] | None,
-    snaps_b: list[CacheSnapshot] | None,
-    label: str = "",
-) -> None:
-    """Assert two snapshot lists are identical."""
-    if snaps_a is None:
-        assert snaps_b is None, f"{label}One side has snapshots, other doesn't"
-        return
-    assert snaps_b is not None, f"{label}One side has snapshots, other doesn't"
-    assert len(snaps_a) == len(snaps_b), (
-        f"{label}Snapshot count: {len(snaps_a)} vs {len(snaps_b)}"
-    )
-    for k, (sa, sb) in enumerate(zip(snaps_a, snaps_b, strict=True)):
-        assert sa.token_count == sb.token_count, (
-            f"{label}Snapshot {k} token_count: {sa.token_count} vs {sb.token_count}"
-        )
-        for layer_i, (s1, s2) in enumerate(zip(sa.states, sb.states, strict=True)):
-            if s1 is None and s2 is None:
-                continue
-            assert s1 is not None and s2 is not None, (
-                f"{label}Snapshot {k}, layer {layer_i}: one state is None"
-            )
-            state_a = _safe_state(s1)
-            state_b = _safe_state(s2)
-            if not state_a and not state_b:
-                continue
-            assert len(state_a) == len(state_b), (
-                f"{label}Snapshot {k}, layer {layer_i}: state length mismatch"
-            )
-            for st_j, (arr_a, arr_b) in enumerate(zip(state_a, state_b, strict=True)):
-                if arr_a is None and arr_b is None:
-                    continue
-                assert arr_a is not None and arr_b is not None
-                _assert_state_equal(
-                    arr_a,
-                    arr_b,
-                    f"{label}Snapshot {k}, layer {layer_i}, state {st_j}",
-                )
-
-
-# ── Test class ────────────────────────────────────────────────────────────── #
-
-
-@pytest.mark.slow
-class TestBatchVsGenerate:
-    """Verify BatchGenerator matches mlx_generate for output tokens and prefix cache."""
-
-    @pytest.fixture(autouse=True)
-    def _cleanup(self):
-        yield
-        mx.clear_cache()
-        gc.collect()
-
-    @pytest.mark.parametrize(
-        "spec",
-        ARCHITECTURES,
-        ids=[a.name for a in ARCHITECTURES],
-    )
-    def test_same_output_and_cache(self, spec: ArchSpec) -> None:
-        if not _arch_available(spec):
-            pytest.skip(f"Model {spec.hub_name} not cached locally")
-
-        snapshot = _find_snapshot(spec.hub_name)
-        assert snapshot is not None
-
-        tmpdir = Path(tempfile.mkdtemp(prefix=f"exo_batchtest_{spec.name}_"))
-        try:
-            # Build reduced config
-            with open(snapshot / "config.json") as f:
-                cfg = cast(dict[str, Any], json.load(f))
-            reduced = _reduce_config(copy.deepcopy(cfg))
-            (tmpdir / "config.json").write_text(json.dumps(reduced))
-
-            # Copy tokenizer
-            tok_src = snapshot
-            if spec.tokenizer_hub is not None:
-                alt = _find_snapshot(spec.tokenizer_hub)
-                if alt is not None:
-                    tok_src = alt
-            _copy_tokenizer(tok_src, tmpdir)
-
-            # Load tokenizer, build model with random weights
-            model_id = ModelId(f"mlx-community/{spec.hub_name}")
-            tokenizer = load_tokenizer_for_model_id(model_id, tmpdir)
-            mx.random.seed(0)
-            model = _build_model(spec.module, reduced)
-
-            task = _make_task()
-
-            # ── Run mlx_generate path ──
-            # Seed is set inside mlx_generate/ExoBatchGenerator.submit from task.seed
-            kv_mlx = KVPrefixCache(None)
-            mlx_tokens = _collect_mlx_generate(model, tokenizer, task, kv_mlx)
-
-            # ── Run batch generator path ──
-            kv_batch = KVPrefixCache(None)
-            batch_tokens = _collect_batch_generate(model, tokenizer, task, kv_batch)
-
-            # ── Compare output tokens ──
-            assert len(mlx_tokens) > 0, "mlx_generate produced no tokens"
-            assert len(batch_tokens) > 0, "BatchGenerator produced no tokens"
-            assert mlx_tokens == batch_tokens, (
-                f"[{spec.name}] Token mismatch:\n"
-                f"  mlx_generate:    {mlx_tokens}\n"
-                f"  BatchGenerator:  {batch_tokens}"
-            )
-
-            # ── Compare prefix cache KV arrays ──
-            assert len(kv_mlx.caches) == 1, "mlx_generate didn't save to prefix cache"
-            assert len(kv_batch.caches) == 1, (
-                "BatchGenerator didn't save to prefix cache"
-            )
-
-            _compare_cache_arrays(
-                kv_mlx.caches[0],
-                kv_batch.caches[0],
-                label=f"[{spec.name}] ",
-            )
-
-            # ── Compare cache lengths ──
-            mlx_len = cache_length(kv_mlx.caches[0])
-            batch_len = cache_length(kv_batch.caches[0])
-            assert mlx_len == batch_len, (
-                f"[{spec.name}] Cache length: mlx={mlx_len} vs batch={batch_len}"
-            )
-
-            # ── Compare snapshots ──
-            _compare_snapshots(
-                kv_mlx._snapshots[0],  # pyright: ignore[reportPrivateUsage]
-                kv_batch._snapshots[0],  # pyright: ignore[reportPrivateUsage]
-                label=f"[{spec.name}] ",
-            )
-
-        finally:
-            shutil.rmtree(tmpdir, ignore_errors=True)
-
-    @pytest.mark.parametrize(
-        "spec",
-        ARCHITECTURES,
-        ids=[a.name for a in ARCHITECTURES],
-    )
-    def test_concurrent_batch_completes(self, spec: ArchSpec) -> None:
-        """Two requests processed concurrently must both complete without
-        crashing and produce non-empty output.
-
-        Note: batch decode logits are NOT bit-exact with sequential because
-        Metal's matmul kernel picks different reduction tiling for B=1 vs B=2
-        when L=1 (decode step). This introduces sub-ULP float16 diffs in
-        gate_proj/down_proj/lm_head which swiglu amplifies by |up_values|.
-        With random weights these accumulate into argmax flips; with trained
-        weights the diffs are absorbed and output matches exactly (verified
-        with real Llama-3.2-1B-Instruct-4bit weights).
-        """
-        if not _arch_available(spec):
-            pytest.skip(f"Model {spec.hub_name} not cached locally")
-
-        snapshot = _find_snapshot(spec.hub_name)
-        assert snapshot is not None
-
-        tmpdir = Path(tempfile.mkdtemp(prefix=f"exo_concurrent_{spec.name}_"))
-        try:
-            with open(snapshot / "config.json") as f:
-                cfg = cast(dict[str, Any], json.load(f))
-            reduced = _reduce_config(copy.deepcopy(cfg))
-            (tmpdir / "config.json").write_text(json.dumps(reduced))
-
-            tok_src = snapshot
-            if spec.tokenizer_hub is not None:
-                alt = _find_snapshot(spec.tokenizer_hub)
-                if alt is not None:
-                    tok_src = alt
-            _copy_tokenizer(tok_src, tmpdir)
-
-            model_id = ModelId(f"mlx-community/{spec.hub_name}")
-            tokenizer = load_tokenizer_for_model_id(model_id, tmpdir)
-            mx.random.seed(0)
-            model = _build_model(spec.module, reduced)
-
-            # Two different prompts → different prompt lengths.
-            task_a = _make_task(content="Hello, what is 2+2?", seed=42)
-            task_a = task_a.model_copy(update={"temperature": 0.0})
-            task_b = _make_task(
-                content="Write a short poem about the ocean and the sky.",
-                seed=99,
-            )
-            task_b = task_b.model_copy(update={"temperature": 0.0})
-
-            # ── Concurrent: submit both to one ExoBatchGenerator ──
-            exo_gen = ExoBatchGenerator(
-                model=model,
-                tokenizer=tokenizer,
-                group=None,
-                kv_prefix_cache=None,
-            )
-
-            prompt_a = apply_chat_template(tokenizer=tokenizer, task_params=task_a)
-            prompt_b = apply_chat_template(tokenizer=tokenizer, task_params=task_b)
-            uid_a = exo_gen.submit(task_params=task_a, prompt=prompt_a)
-            uid_b = exo_gen.submit(task_params=task_b, prompt=prompt_b)
-
-            batch_tokens: dict[int, list[int]] = {uid_a: [], uid_b: []}
-            finished: set[int] = set()
-            while exo_gen.has_work:
-                results = exo_gen.step()
-                for uid, response in results:
-                    batch_tokens[uid].append(response.token)
-                    if response.finish_reason is not None:
-                        finished.add(uid)
-
-            exo_gen.close()
-
-            # ── Verify both completed ──
-            assert len(batch_tokens[uid_a]) > 0, "No tokens for task A"
-            assert len(batch_tokens[uid_b]) > 0, "No tokens for task B"
-            assert uid_a in finished, "Task A never finished"
-            assert uid_b in finished, "Task B never finished"
-        finally:
-            shutil.rmtree(tmpdir, ignore_errors=True)

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

@@ -116,7 +116,6 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
     # initialize_mlx returns a mock group
     monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MockGroup()))
     monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, MockTokenizer)))
-    monkeypatch.setattr(mlx_batch_generator, "warmup_inference", make_nothin(1))
     monkeypatch.setattr(mlx_batch_generator, "_check_for_debug_prompts", nothin)
     monkeypatch.setattr(mlx_batch_generator, "mx_any", make_nothin(False))
 
@@ -139,8 +138,10 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
 
 class FakeExoBatchGenerator:
     def __init__(self, *_args: object, **_kwargs: object) -> None:
-        self._uid_counter = 0
-        self._pending: dict[int, GenerationResponse] = {}
+        self._pending: dict[str, GenerationResponse] = {}
+
+    def warmup(self) -> int:
+        return 50
 
     @property
     def has_work(self) -> bool:
@@ -148,30 +149,29 @@ class FakeExoBatchGenerator:
 
     def submit(
         self,
+        task_id: str = "",
         task_params: object = None,
         prompt: object = None,
         on_prefill_progress: object = None,
         distributed_prompt_progress_callback: object = None,
         on_generation_token: object = None,
-    ) -> int:
-        uid = self._uid_counter
-        self._uid_counter += 1
-        self._pending[uid] = GenerationResponse(
+    ) -> str:
+        self._pending[task_id] = GenerationResponse(
             text="hi",
             token=0,
             finish_reason="stop",
             usage=None,
         )
-        return uid
+        return task_id
 
-    def step(self) -> list[tuple[int, GenerationResponse]]:
+    def step(self) -> list[tuple[str, GenerationResponse]]:
         results = list(self._pending.items())
         self._pending.clear()
         return results
 
-    def cancel(self, uids: list[int]) -> None:
-        for uid in uids:
-            self._pending.pop(uid, None)
+    def cancel(self, task_ids: list[str]) -> None:
+        for tid in task_ids:
+            self._pending.pop(tid, None)
 
     def close(self) -> None:
         pass
@@ -262,11 +262,17 @@ def _run(tasks: Iterable[Task], send_after_ready: list[Task] | None = None):
             "exo.worker.runner.llm_inference.runner.mx.distributed.all_gather",
             make_nothin(mx.array([1])),
         ):
+            builder = mlx_runner.MlxBuilder(
+                model_id=MODEL_A_ID,
+                event_sender=event_sender,  # pyright: ignore[reportArgumentType]
+                cancel_receiver=cancel_receiver,
+            )
             runner = mlx_runner.Runner(
                 bound_instance,
                 event_sender,  # pyright: ignore[reportArgumentType]
                 task_receiver,
                 cancel_receiver,
+                builder,
             )
             runner.main()
 

tmp/spark/mac-usbc-fix.sh

@@ -0,0 +1,52 @@
+#!/usr/bin/env bash
+set -Eeuo pipefail
+
+msg() {
+  printf '\n==> %s\n' "$*"
+}
+
+msg "Bringing down old Apple USB fallback interfaces"
+sudo ifconfig en4 down 2>/dev/null || true
+sudo ifconfig en5 down 2>/dev/null || true
+sudo ifconfig en6 down 2>/dev/null || true
+
+msg "Bringing down current anpi interfaces"
+sudo ifconfig anpi0 down 2>/dev/null || true
+sudo ifconfig anpi1 down 2>/dev/null || true
+sudo ifconfig anpi2 down 2>/dev/null || true
+
+msg "Stopping likely bad exporter services"
+sudo launchctl bootout system/com.apple.usbmuxd 2>/dev/null || true
+sudo launchctl bootout system/com.apple.remoted 2>/dev/null || true
+launchctl bootout "gui/$(id -u)/com.apple.remoted" 2>/dev/null || true
+
+sudo pkill -x usbmuxd 2>/dev/null || true
+sudo pkill -x remoted 2>/dev/null || true
+
+cat <<'EOF'
+
+Now unplug and replug the cable.
+
+After reconnect, press Enter to continue.
+EOF
+
+read -r _
+
+msg "Bringing up anpi interfaces and requesting DHCP"
+for ifn in anpi0 anpi1 anpi2; do
+  sudo ifconfig "$ifn" up 2>/dev/null || true
+  sudo ipconfig set "$ifn" DHCP 2>/dev/null || true
+done
+
+sleep 3
+
+msg "Resulting interface state"
+for ifn in anpi0 anpi1 anpi2; do
+  echo
+  echo "--- $ifn ---"
+  ifconfig "$ifn" 2>/dev/null || true
+  ipconfig getifaddr "$ifn" 2>/dev/null || true
+done
+
+echo
+echo "Use whichever interface got 10.42.0.x or 10.43.0.x"

tmp/spark/spark-build-apple-cdc-ncm.sh

@@ -0,0 +1,180 @@
+#!/usr/bin/env bash
+set -Eeuo pipefail
+
+KVER="$(uname -r)"
+WORKDIR="${HOME}/apple1905-cdcncm"
+SRCVER="6.14.0-1015.15"
+SRCNAME="linux-nvidia-6.14"
+BASE_URL="https://launchpad.net/ubuntu/+archive/primary/+sourcefiles/${SRCNAME}/${SRCVER}"
+
+ORIG="${SRCNAME}_6.14.0.orig.tar.gz"
+DIFF="${SRCNAME}_${SRCVER}.diff.gz"
+DSC="${SRCNAME}_${SRCVER}.dsc"
+
+KBUILD="/lib/modules/${KVER}/build"
+
+msg() { printf '\n==> %s\n' "$*"; }
+die() {
+  echo "ERROR: $*" >&2
+  exit 1
+}
+need_cmd() { command -v "$1" >/dev/null 2>&1 || die "Missing command: $1"; }
+
+[[ $EUID -eq 0 ]] || die "Run with sudo"
+[[ -d $KBUILD ]] || die "Missing kernel headers/build dir: $KBUILD"
+
+need_cmd curl
+need_cmd dpkg-source
+need_cmd make
+need_cmd python3
+need_cmd modprobe
+need_cmd insmod
+need_cmd modinfo
+need_cmd ip
+need_cmd dmesg
+need_cmd find
+
+msg "Installing build deps"
+export DEBIAN_FRONTEND=noninteractive
+apt-get update
+apt-get install -y --no-install-recommends \
+  build-essential \
+  linux-headers-"${KVER}" \
+  dpkg-dev \
+  ca-certificates \
+  curl \
+  kmod \
+  zstd \
+  python3
+
+msg "Preparing workspace"
+rm -rf "$WORKDIR"
+mkdir -p "$WORKDIR"
+cd "$WORKDIR"
+
+msg "Downloading exact source package for ${SRCNAME} ${SRCVER}"
+curl -fL "${BASE_URL}/${ORIG}" -o "${ORIG}"
+curl -fL "${BASE_URL}/${DIFF}" -o "${DIFF}"
+curl -fL "${BASE_URL}/${DSC}" -o "${DSC}"
+
+msg "Extracting source"
+dpkg-source -x "${DSC}"
+
+SRCDIR="$(find . -maxdepth 1 -mindepth 1 -type d -name "${SRCNAME}-*" | head -n1)"
+[[ -n ${SRCDIR} ]] || die "Could not find extracted source directory"
+cd "${SRCDIR}"
+
+[[ -f drivers/net/usb/cdc_ncm.c ]] || die "cdc_ncm.c not found in source tree"
+
+msg "Patching cdc_ncm.c"
+python3 <<'PY'
+from pathlib import Path
+import re
+
+p = Path("drivers/net/usb/cdc_ncm.c")
+s = p.read_text()
+
+if "0x05ac, 0x1905, 0" not in s:
+    block = (
+        '\t/* Apple Mac direct USB-C networking quirk */\n'
+        '\t{ USB_DEVICE_INTERFACE_NUMBER(0x05ac, 0x1905, 0),\n'
+        '\t  .driver_info = (unsigned long)&apple_private_interface_info,\n'
+        '\t},\n'
+        '\t{ USB_DEVICE_INTERFACE_NUMBER(0x05ac, 0x1905, 2),\n'
+        '\t  .driver_info = (unsigned long)&apple_private_interface_info,\n'
+        '\t},\n'
+    )
+
+    pat = re.compile(
+        r'(\{\s*USB_INTERFACE_INFO\s*\(\s*USB_CLASS_COMM\s*,\s*USB_CDC_SUBCLASS_NCM\s*,\s*USB_CDC_PROTO_NONE\s*\)\s*,\s*\.driver_info\s*=\s*\(unsigned long\)&cdc_ncm_info\s*,\s*\},)',
+        re.S,
+    )
+    s2, n = pat.subn(block + r'\1', s, count=1)
+    if n != 1:
+        raise SystemExit("Could not find generic CDC NCM class-match entry to patch")
+    s = s2
+
+# Patch old bind check if needed.
+# Old style:
+#   if (!dev->in || !dev->out || !dev->status)
+# New behavior for this Mac:
+#   allow missing status endpoint for 05ac:1905
+old = re.compile(
+    r'if\s*\(\s*!dev->in\s*\|\|\s*!dev->out\s*\|\|\s*!dev->status\s*\)',
+    re.S,
+)
+repl = (
+    'if (!dev->in || !dev->out || '
+    '(!dev->status && '
+    '!(le16_to_cpu(dev->udev->descriptor.idVendor) == 0x05ac && '
+    'le16_to_cpu(dev->udev->descriptor.idProduct) == 0x1905)))'
+)
+s, n = old.subn(repl, s, count=1)
+
+# If the tree already has the newer FLAG_LINK_INTR logic, leave it alone.
+if n == 0 and "FLAG_LINK_INTR" not in s:
+    raise SystemExit("Could not patch bind_common() missing-status check")
+
+# Add a loud marker so we know the patched module actually loaded.
+if 'Apple 05ac:1905 quirk test module loaded' not in s:
+    m = re.search(r'static\s+int\s+__init\s+cdc_ncm_init\s*\(\s*void\s*\)\s*\{', s)
+    if m:
+        insert_at = m.end()
+        s = s[:insert_at] + '\n\tpr_info("cdc_ncm: Apple 05ac:1905 quirk test module loaded\\n");' + s[insert_at:]
+
+p.write_text(s)
+print("Patched", p)
+PY
+
+msg "Preparing out-of-tree build dir"
+mkdir -p buildmod
+cp drivers/net/usb/cdc_ncm.c buildmod/
+
+cat >buildmod/Makefile <<'EOF'
+obj-m += cdc_ncm.o
+ccflags-y += -Wno-error
+
+all:
+	$(MAKE) -C /lib/modules/$(shell uname -r)/build M=$(PWD) modules
+
+clean:
+	$(MAKE) -C /lib/modules/$(shell uname -r)/build M=$(PWD) clean
+EOF
+
+msg "Building patched module"
+cd buildmod
+make -C "$KBUILD" M="$PWD" modules
+PATCHED_KO="$PWD/cdc_ncm.ko"
+[[ -f $PATCHED_KO ]] || die "Patched cdc_ncm.ko was not built"
+
+msg "Patched module info"
+modinfo "$PATCHED_KO" | sed -n '1,20p'
+
+msg "Reloading module stack"
+modprobe -r cdc_mbim 2>/dev/null || true
+modprobe -r cdc_ncm 2>/dev/null || true
+insmod "$PATCHED_KO"
+
+msg "Recent dmesg"
+dmesg | tail -n 60
+
+cat <<EOF
+
+Done.
+
+Now:
+  1. run: sudo dmesg -w
+  2. unplug and replug the USB-C cable to the Mac
+  3. then run:
+       ip -br link
+       lsusb | grep 05ac:1905
+
+You want to see:
+  - cdc_ncm: Apple 05ac:1905 quirk test module loaded
+  - no bind() failure for 05ac:1905
+  - a new interface appearing
+
+Workspace:
+  $WORKDIR
+
+EOF

tmp/spark/spark-usbc-setup.sh

@@ -0,0 +1,155 @@
+#!/usr/bin/env bash
+set -Eeuo pipefail
+
+PATCHED_CDC_NCM="/root/apple1905-cdcncm/linux-nvidia-6.14-6.14.0/buildmod/cdc_ncm.ko"
+
+need() {
+  command -v "$1" >/dev/null 2>&1 || {
+    echo "Missing command: $1" >&2
+    exit 1
+  }
+}
+
+msg() {
+  printf '\n==> %s\n' "$*"
+}
+
+[[ $EUID -eq 0 ]] || {
+  echo "Run with sudo." >&2
+  exit 1
+}
+
+need modprobe
+need insmod
+need nmcli
+need lsusb
+need ip
+need awk
+need grep
+
+[[ -f $PATCHED_CDC_NCM ]] || {
+  echo "Patched module not found: $PATCHED_CDC_NCM" >&2
+  exit 1
+}
+
+msg "Removing old hard-block config if present"
+rm -f /etc/modprobe.d/zz-kill-mac-usb-net.conf
+depmod -a
+
+msg "Loading USB4 / Type-C / Thunderbolt support"
+modprobe typec || true
+modprobe typec_ucsi || true
+modprobe ucsi_acpi || true
+modprobe thunderbolt || true
+modprobe typec_thunderbolt || true
+modprobe thunderbolt_net || true
+
+msg "Unloading stock USB network stack"
+modprobe -r cdc_mbim cdc_wdm cdc_ncm cdc_ether usbnet 2>/dev/null || true
+
+msg "Loading dependency modules"
+modprobe usbnet
+modprobe cdc_ether
+
+msg "Loading patched cdc_ncm"
+if lsmod | grep -q '^cdc_ncm '; then
+  echo "cdc_ncm already loaded"
+else
+  insmod "$PATCHED_CDC_NCM"
+fi
+
+msg "Current module state"
+lsmod | grep -E 'typec|ucsi|thunderbolt|cdc_ncm|cdc_ether|usbnet' || true
+
+cat <<'EOF'
+
+Now do this on the Mac:
+  1. Run the Mac script below
+  2. Replug the cable
+  3. Wait for the Mac to land on the fast bus
+
+Then run this same script again with:
+  sudo ./spark-usbc-setup.sh finalize
+
+EOF
+
+if [[ ${1:-} != "finalize" ]]; then
+  exit 0
+fi
+
+msg "Waiting for Apple Mac on fast bus (Bus 004 @ 10000M or equivalent)"
+for _ in $(seq 1 60); do
+  if lsusb -t | grep -q '05ac:1905\|Driver=\[none\], 10000M'; then
+    break
+  fi
+  sleep 1
+done
+
+msg "lsusb -t snapshot"
+lsusb -t
+echo
+lsusb | grep '05ac:1905' || true
+
+msg "Waiting for new Apple USB NICs"
+for _ in $(seq 1 30); do
+  mapfile -t APPLE_IFS < <(
+    ip -o link | awk -F': ' '{print $2}' |
+      grep '^enx36be1bab12' || true
+  )
+  if [[ ${#APPLE_IFS[@]} -ge 2 ]]; then
+    break
+  fi
+  sleep 1
+done
+
+if [[ ${#APPLE_IFS[@]} -lt 2 ]]; then
+  echo "Did not find two Apple USB interfaces." >&2
+  ip -br link
+  exit 1
+fi
+
+IF0="${APPLE_IFS[0]}"
+IF1="${APPLE_IFS[1]}"
+
+msg "Found interfaces: $IF0 and $IF1"
+
+msg "Resetting old NetworkManager profiles"
+nmcli connection delete mac-usb-dhcp-0 2>/dev/null || true
+nmcli connection delete mac-usb-dhcp-1 2>/dev/null || true
+
+msg "Bringing up shared DHCP on both interfaces"
+nmcli connection add type ethernet ifname "$IF0" con-name mac-usb-dhcp-0
+nmcli connection modify mac-usb-dhcp-0 \
+  ipv4.method shared \
+  ipv6.method disabled \
+  ipv4.addresses 10.42.0.1/24
+nmcli connection up mac-usb-dhcp-0
+
+nmcli connection add type ethernet ifname "$IF1" con-name mac-usb-dhcp-1
+nmcli connection modify mac-usb-dhcp-1 \
+  ipv4.method shared \
+  ipv6.method disabled \
+  ipv4.addresses 10.43.0.1/24
+nmcli connection up mac-usb-dhcp-1
+
+msg "Current addresses"
+ip -br addr show dev "$IF0"
+ip -br addr show dev "$IF1"
+
+msg "Recent NetworkManager log"
+journalctl -u NetworkManager -n 60 --no-pager || true
+
+cat <<EOF
+
+Done.
+
+DGX shared-DHCP interfaces:
+  $IF0 -> 10.42.0.1/24
+  $IF1 -> 10.43.0.1/24
+
+Next on the Mac:
+  sudo ipconfig set anpi0 DHCP
+  sudo ipconfig set anpi1 DHCP
+  sudo ipconfig set anpi2 DHCP
+
+EOF