more fixes

seems like stablepkgs swiftfmt works now! also bump macmon to 0.7

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

.envrc

@@ -1 +1,8 @@
 use flake
+
+# creates .venv if doesn't exist and loads its environment
+export VIRTUAL_ENV=".venv"
+if ! [ -d "./$VIRTUAL_ENV" ]; then
+  uv venv
+fi
+layout python

.github/workflows/build-app.yml

@@ -159,7 +159,7 @@ jobs:
           fi
 
       - name: Install Homebrew packages
-        run: brew install just awscli macmon
+        run: brew install just awscli
 
       - name: Install UV
         uses: astral-sh/setup-uv@v6
@@ -243,6 +243,14 @@ jobs:
       # Build the bundle
       # ============================================================
 
+      - name: Add pinned macmon to PATH
+        run: |
+          MACMON_DIR=$(nix develop --command sh -c 'dirname $(which macmon)')
+          echo "Using macmon from: $MACMON_DIR"
+          echo "$MACMON_DIR" >> $GITHUB_PATH
+          # Remove any Homebrew macmon so PyInstaller can't accidentally pick it up
+          brew uninstall macmon 2>/dev/null || true
+
       - name: Build PyInstaller bundle
         run: uv run pyinstaller packaging/pyinstaller/exo.spec
 

.gitignore

@@ -38,3 +38,5 @@ bench/**/*.json
 
 # tmp
 tmp/models
+/build/exo
+/.claude/skills

.idea/exo-v2.iml

@@ -1,31 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<module type="EMPTY_MODULE" version="4">
-  <component name="FacetManager">
-    <facet type="Python" name="Python facet">
-      <configuration sdkName="Python 3.13 virtualenv at ~/Desktop/exo/.venv" />
-    </facet>
-  </component>
-  <component name="Go" enabled="true" />
-  <component name="NewModuleRootManager">
-    <content url="file://$MODULE_DIR$">
-      <sourceFolder url="file://$MODULE_DIR$/scripts/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/exo_pyo3_bindings/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/exo_pyo3_bindings/tests" isTestSource="true" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/util/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/networking/examples" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/networking/src" isTestSource="false" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/networking/tests" isTestSource="true" />
-      <sourceFolder url="file://$MODULE_DIR$/rust/system_custodian/src" isTestSource="false" />
-      <excludeFolder url="file://$MODULE_DIR$/.venv" />
-      <excludeFolder url="file://$MODULE_DIR$/.direnv" />
-      <excludeFolder url="file://$MODULE_DIR$/build" />
-      <excludeFolder url="file://$MODULE_DIR$/dist" />
-      <excludeFolder url="file://$MODULE_DIR$/.go_cache" />
-      <excludeFolder url="file://$MODULE_DIR$/rust/target" />
-    </content>
-    <orderEntry type="jdk" jdkName="Python 3.13 (exo)" jdkType="Python SDK" />
-    <orderEntry type="sourceFolder" forTests="false" />
-    <orderEntry type="library" name="Python 3.13 virtualenv at ~/Desktop/exo/.venv interpreter library" level="application" />
-  </component>
-</module>

.idea/externalDependencies.xml

@@ -1,6 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <project version="4">
   <component name="ExternalDependencies">
-    <plugin id="systems.fehn.intellijdirenv" />
+    <plugin id="al.aoli.intellijdirenv" />
   </component>
 </project>

.idea/inspectionProfiles/Project_Default.xml

@@ -1,14 +0,0 @@
-<component name="InspectionProjectProfileManager">
-  <profile version="1.0">
-    <option name="myName" value="Project Default" />
-    <inspection_tool class="PyCompatibilityInspection" enabled="true" level="WARNING" enabled_by_default="true">
-      <option name="ourVersions">
-        <value>
-          <list size="1">
-            <item index="0" class="java.lang.String" itemvalue="3.14" />
-          </list>
-        </value>
-      </option>
-    </inspection_tool>
-  </profile>
-</component>

.idea/misc.xml

@@ -4,7 +4,4 @@
     <option name="sdkName" value="Python 3.13 (exo)" />
   </component>
   <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.13 (exo)" project-jdk-type="Python SDK" />
-  <component name="PythonCompatibilityInspectionAdvertiser">
-    <option name="version" value="3" />
-  </component>
 </project>

.mlx_typings/mlx/core/__init__.pyi

@@ -2396,7 +2396,7 @@ def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array:
         array: The angles in degrees.
     """
 
-def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array]):
+def depends[T](inputs: T, dependencies: array | Sequence[array]) -> T:
     """
     Insert dependencies between arrays in the graph. The outputs are
     identical to ``inputs`` but with dependencies on ``dependencies``.

.mlx_typings/mlx/nn/__init__.pyi

@@ -1,9 +1,5 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from layers import *
-from utils import *
+from .layers import *
+from .utils import *
 
 from . import init as init
 from . import losses as losses

.mlx_typings/mlx/nn/layers/__init__.pyi

@@ -1,20 +1,16 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from activations import *
-from base import *
-from containers import *
-from convolution import *
-from convolution_transpose import *
-from distributed import *
-from dropout import *
-from embedding import *
-from linear import *
-from normalization import *
-from pooling import *
-from positional_encoding import *
-from quantized import *
-from recurrent import *
-from transformer import *
-from upsample import *
+from .activations import *
+from .base import *
+from .containers import *
+from .convolution import *
+from .convolution_transpose import *
+from .distributed import *
+from .dropout import *
+from .embedding import *
+from .linear import *
+from .normalization import *
+from .pooling import *
+from .positional_encoding import *
+from .quantized import *
+from .recurrent import *
+from .transformer import *
+from .upsample import *

.mlx_typings/mlx/nn/layers/base.pyi

@@ -53,7 +53,7 @@ class Module(dict):
         mx.eval(model.parameters())
     """
 
-    __call__: Callable
+    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
     def __init__(self) -> None:
         """Should be called by the subclasses of ``Module``."""
 

.mlx_typings/mlx/nn/layers/convolution.pyi

@@ -32,6 +32,7 @@ class Conv1d(Module):
     """
 
     weight: mx.array
+    bias: mx.array | None
     groups: int
     def __init__(
         self,

.mlx_typings/mlx/nn/layers/linear.pyi

@@ -40,6 +40,10 @@ class Linear(Module):
         bias (bool, optional): If set to ``False`` then the layer will
           not use a bias. Default is ``True``.
     """
+
+    weight: mx.array
+    bias: mx.array | None
+
     def __init__(self, input_dims: int, output_dims: int, bias: bool = ...) -> None: ...
     def __call__(self, x: mx.array) -> mx.array: ...
     def to_quantized(

.mlx_typings/mlx/nn/layers/normalization.pyi

@@ -88,6 +88,9 @@ class RMSNorm(Module):
         dims (int): The feature dimension of the input to normalize over
         eps (float): A small additive constant for numerical stability
     """
+
+    weight: mx.array
+
     def __init__(self, dims: int, eps: float = ...) -> None: ...
     def __call__(self, x) -> mx.array: ...
 

.mlx_typings/mlx/nn/layers/quantized.pyi

@@ -2,7 +2,7 @@
 This type stub file was generated by pyright.
 """
 
-from typing import Callable, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import mlx.core as mx
 from base import Module
@@ -13,8 +13,10 @@ def quantize(
     bits: int = ...,
     *,
     mode: str = ...,
-    class_predicate: Optional[Callable[[str, Module], Union[bool, dict]]] = ...,
-):  # -> None:
+    class_predicate: Optional[
+        Callable[[str, Module], Union[bool, dict[str, Any]]]
+    ] = ...,
+) -> None:
     """Quantize the sub-modules of a module according to a predicate.
 
     By default all layers that define a ``to_quantized(group_size, bits)``

.mlx_typings/mlx_lm/_version.pyi

@@ -0,0 +1,5 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+__version__ = ...

.mlx_typings/mlx_lm/generate.pyi

@@ -3,13 +3,12 @@ This type stub file was generated by pyright.
 """
 
 import contextlib
-from dataclasses import dataclass
-from typing import Any, Callable, Generator, List, Optional, Tuple, Union
-
 import mlx.core as mx
 import mlx.nn as nn
+from dataclasses import dataclass
+from collections import deque
+from typing import Any, Callable, Generator, List, Optional, Sequence, Tuple, Union
 from transformers import PreTrainedTokenizer
-
 from .tokenizer_utils import TokenizerWrapper
 
 DEFAULT_PROMPT = ...
@@ -29,8 +28,9 @@ def str2bool(string):  # -> bool:
     ...
 def setup_arg_parser():  # -> ArgumentParser:
     """Set up and return the argument parser."""
+    ...
 
-generation_stream = ...
+generation_stream: mx.Stream
 
 @contextlib.contextmanager
 def wired_limit(
@@ -43,6 +43,7 @@ def wired_limit(
     async eval could be running pass in the streams to synchronize with prior
     to exiting the context manager.
     """
+    ...
 @dataclass
 class GenerationResponse:
     """
@@ -73,9 +74,11 @@ class GenerationResponse:
     finish_reason: Optional[str] = ...
 
 def maybe_quantize_kv_cache(
-    prompt_cache, quantized_kv_start, kv_group_size, kv_bits
-):  # -> None:
-    ...
+    prompt_cache: Any,
+    quantized_kv_start: int | None,
+    kv_group_size: int | None,
+    kv_bits: int | None,
+) -> None: ...
 def generate_step(
     prompt: mx.array,
     model: nn.Module,
@@ -89,7 +92,7 @@ def generate_step(
     kv_bits: Optional[int] = ...,
     kv_group_size: int = ...,
     quantized_kv_start: int = ...,
-    prompt_progress_callback: Optional[Callable[[int], int]] = ...,
+    prompt_progress_callback: Optional[Callable[[int, int], None]] = ...,
     input_embeddings: Optional[mx.array] = ...,
 ) -> Generator[Tuple[mx.array, mx.array], None, None]:
     """
@@ -115,7 +118,7 @@ def generate_step(
         kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
         quantized_kv_start (int): Step to begin using a quantized KV cache.
            when ``kv_bits`` is non-None. Default: ``0``.
-        prompt_progress_callback (Callable[[int], int]): A call-back which takes the
+        prompt_progress_callback (Callable[[int, int], None]): A call-back which takes the
            prompt tokens processed so far and the total number of prompt tokens.
         input_embeddings (mx.array, optional): Input embeddings to use instead of or in
           conjunction with prompt tokens. Default: ``None``.
@@ -123,6 +126,7 @@ def generate_step(
     Yields:
         Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
     """
+    ...
 
 def speculative_generate_step(
     prompt: mx.array,
@@ -168,6 +172,7 @@ def speculative_generate_step(
         Tuple[mx.array, mx.array, bool]: One token, a vector of log probabilities,
           and a bool indicating if the token was generated by the draft model
     """
+    ...
 
 def stream_generate(
     model: nn.Module,
@@ -175,7 +180,7 @@ def stream_generate(
     prompt: Union[str, mx.array, List[int]],
     max_tokens: int = ...,
     draft_model: Optional[nn.Module] = ...,
-    **kwargs: object,
+    **kwargs: Any,
 ) -> Generator[GenerationResponse, None, None]:
     """
     A generator producing text based on the given prompt from the model.
@@ -197,6 +202,7 @@ def stream_generate(
         GenerationResponse: An instance containing the generated text segment and
             associated metadata. See :class:`GenerationResponse` for details.
     """
+    ...
 
 def generate(
     model: nn.Module,
@@ -217,6 +223,9 @@ def generate(
        kwargs: The remaining options get passed to :func:`stream_generate`.
           See :func:`stream_generate` for more details.
     """
+    ...
+
+def _merge_caches(caches: List[List[Any]]) -> List[Any]: ...
 @dataclass
 class BatchStats:
     """
@@ -240,10 +249,262 @@ class BatchStats:
     generation_time: float = ...
     peak_memory: float = ...
 
+class SequenceStateMachine:
+    """A state machine that uses one Aho-Corasick trie per state to efficiently
+    track state across a generated sequence.
+
+    The transitions are provided as state -> [(sequence, new_state)].
+
+    Example:
+
+        sm = SequenceStateMachine(
+            transitions={
+                "normal": [
+                    (think_start_tokens, "reasoning"),
+                    (tool_start_tokens, "tool"),
+                    (eos, None),
+                ],
+                "reasoning": [
+                    (think_end_tokens, "normal"),
+                    (eos, None),
+                ],
+                "tool": [
+                    (tool_end_tokens, None),
+                    (eos, None)
+                ],
+            },
+            initial="normal"
+        )
+    """
+    def __init__(self, transitions=..., initial=...) -> None: ...
+    def __deepcopy__(self, memo):  # -> SequenceStateMachine:
+        ...
+    def make_state(self):  # -> tuple[str, Any, dict[Any, Any]]:
+        ...
+    @staticmethod
+    def match(state, x):  # -> tuple[tuple[Any, Any | None, Any], Any | None, Any]:
+        ...
+
+class PromptProcessingBatch:
+    """
+    A batch processor for prompt tokens with support for incremental processing.
+
+    This class handles batched prompt processing, managing KV caches and preparing
+    tokens for generation. It supports extending, filtering, and splitting batches.
+    """
+    @dataclass
+    class Response:
+        uid: int
+        progress: tuple
+        end_of_segment: bool
+        end_of_prompt: bool
+        ...
+
+    def __init__(
+        self,
+        model: nn.Module,
+        uids: List[int],
+        caches: List[List[Any]],
+        tokens: Optional[List[List[int]]] = ...,
+        prefill_step_size: int = ...,
+        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
+        fallback_sampler: Optional[Callable[[mx.array], mx.array]] = ...,
+        logits_processors: Optional[
+            List[List[Callable[[mx.array, mx.array], mx.array]]]
+        ] = ...,
+        state_machines: Optional[List[SequenceStateMachine]] = ...,
+        max_tokens: Optional[List[int]] = ...,
+    ) -> None: ...
+    def __len__(self):  # -> int:
+        ...
+    def extract_cache(self, idx: int) -> List[Any]: ...
+    def extend(self, batch):  # -> None:
+        ...
+    def split(self, indices: List[int]):  # -> Self:
+        ...
+    def filter(self, keep: List[int]):  # -> None:
+        ...
+    def prompt(self, tokens: List[List[int]]):  # -> None:
+        """
+        Process prompt tokens through the model.
+
+        Args:
+            tokens: List of token sequences to process.
+        """
+        ...
+
+    def generate(self, tokens: List[List[int]]):  # -> GenerationBatch:
+        """
+        Transition from prompt processing to generation.
+
+        Args:
+            tokens: Final tokens for each sequence to start generation.
+
+        Returns:
+            A GenerationBatch ready for token generation.
+        """
+        ...
+
+    @classmethod
+    def empty(
+        cls,
+        model: nn.Module,
+        fallback_sampler: Callable[[mx.array], mx.array],
+        prefill_step_size: int = ...,
+    ):  # -> Self:
+        ...
+
+class GenerationBatch:
+    """
+    A batched token generator that manages multiple sequences in parallel.
+
+    This class handles the generation phase after prompt processing, managing
+    KV caches, sampling, and stop sequence detection for multiple sequences.
+    """
+    @dataclass
+    class Response:
+        uid: int
+        token: int
+        logprobs: mx.array
+        finish_reason: Optional[str]
+        current_state: Optional[str]
+        match_sequence: Optional[List[int]]
+        prompt_cache: Optional[List[Any]]
+        all_tokens: Optional[List[int]]
+        ...
+
+    model: nn.Module
+    uids: List[int]
+    prompt_cache: List[Any]
+    tokens: List[List[int]]
+    samplers: Optional[List[Callable[[mx.array], mx.array]]]
+    fallback_sampler: Callable[[mx.array], mx.array]
+    logits_processors: Optional[List[List[Callable[[mx.array, mx.array], mx.array]]]]
+    state_machines: List[SequenceStateMachine]
+    max_tokens: List[int]
+    _current_tokens: Optional[mx.array]
+    _current_logprobs: List[mx.array]
+    _next_tokens: mx.array
+    _next_logprobs: List[mx.array]
+    _token_context: List[mx.array]
+    _num_tokens: List[int]
+
+    def __init__(
+        self,
+        model: nn.Module,
+        uids: List[int],
+        inputs: mx.array,
+        prompt_cache: List[Any],
+        tokens: List[List[int]],
+        samplers: Optional[List[Callable[[mx.array], mx.array]]],
+        fallback_sampler: Callable[[mx.array], mx.array],
+        logits_processors: Optional[
+            List[List[Callable[[mx.array, mx.array], mx.array]]]
+        ],
+        state_machines: List[SequenceStateMachine],
+        max_tokens: List[int],
+    ) -> None: ...
+    def __len__(self) -> int: ...
+    def extend(self, batch: GenerationBatch) -> None: ...
+    def extract_cache(self, idx: int) -> List[Any]: ...
+    def filter(self, keep: List[int]) -> None: ...
+    def _step(self) -> Tuple[List[int], List[mx.array]]: ...
+    def next(self) -> List[Response]:
+        """
+        Generate the next batch of tokens.
+
+        Returns:
+            List of Response objects for each sequence in the batch.
+        """
+        ...
+
+    @classmethod
+    def empty(
+        cls, model: nn.Module, fallback_sampler: Callable[[mx.array], mx.array]
+    ):  # -> Self:
+        ...
+
+class BatchGenerator:
+    """
+    A batch generator implements continuous batching.
+
+    This class provides automatic management of prompt processing and generation
+    batches, handling the transition between the two.
+
+    It also allows for segmented prompt processing which guarantees that the
+    generator will stop at these boundaries when processing an input.
+    """
+    def __init__(
+        self,
+        model: nn.Module,
+        max_tokens: int = ...,
+        stop_tokens: Optional[Sequence[Sequence[int]]] = ...,
+        sampler: Optional[Callable[[mx.array], mx.array]] = ...,
+        logits_processors: Optional[
+            List[Callable[[mx.array, mx.array], mx.array]]
+        ] = ...,
+        completion_batch_size: int = ...,
+        prefill_batch_size: int = ...,
+        prefill_step_size: int = ...,
+    ) -> None: ...
+    def close(self) -> None: ...
+    def __del__(self):  # -> None:
+        ...
+    @contextlib.contextmanager
+    def stats(self, stats=...):  # -> Generator[Any | BatchStats, Any, None]:
+        ...
+    _unprocessed_sequences: deque[tuple[Any, ...]]
+    _prompt_batch: PromptProcessingBatch
+    _generation_batch: GenerationBatch
+    _currently_processing: list[Any]
+    _gen_tokens_counter: int
+    _steps_counter: int
+    def _next(
+        self,
+    ) -> tuple[
+        List[PromptProcessingBatch.Response], List[GenerationBatch.Response]
+    ]: ...
+    def insert(
+        self,
+        prompts: List[List[int]],
+        max_tokens: Optional[List[int]] = ...,
+        caches: Optional[List[List[Any]]] = ...,
+        all_tokens: Optional[List[List[int]]] = ...,
+        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
+        logits_processors: Optional[
+            List[List[Callable[[mx.array, mx.array], mx.array]]]
+        ] = ...,
+        state_machines: Optional[List[SequenceStateMachine]] = ...,
+    ) -> List[int]: ...
+    def insert_segments(
+        self,
+        segments: List[List[List[int]]],
+        max_tokens: Optional[List[int]] = ...,
+        caches: Optional[List[List[Any]]] = ...,
+        all_tokens: Optional[List[List[int]]] = ...,
+        samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
+        logits_processors: Optional[
+            List[List[Callable[[mx.array, mx.array], mx.array]]]
+        ] = ...,
+        state_machines: Optional[List[SequenceStateMachine]] = ...,
+    ) -> List[int]: ...
+    def extract_cache(self, uids: List[int]) -> dict[int, Any]: ...
+    def remove(
+        self, uids: List[int], return_prompt_caches: bool = ...
+    ) -> dict[int, Any]: ...
+    @property
+    def prompt_cache_nbytes(self) -> int: ...
+    def next(
+        self,
+    ) -> tuple[
+        List[PromptProcessingBatch.Response], List[GenerationBatch.Response]
+    ]: ...
+    def next_generated(self) -> List[GenerationBatch.Response]: ...
+
 @dataclass
 class BatchResponse:
     """
-    An data object to hold a batch generation response.
+    A data object to hold a batch generation response.
 
     Args:
         texts: (List[str]): The generated text for each prompt.
@@ -252,55 +513,18 @@ class BatchResponse:
 
     texts: List[str]
     stats: BatchStats
-
-@dataclass
-class Batch:
-    uids: List[int]
-    y: mx.array
-    logprobs: mx.array
-    max_tokens: List[int]
-    num_tokens: List[int]
-    cache: List[Any]
-    def __len__(self):  # -> int:
-        ...
-    def filter(self, keep_idx: List[int]):  # -> None:
-        ...
-    def extend(self, other):  # -> None:
-        ...
-
-class BatchGenerator:
-    @dataclass
-    class Response:
-        uid: int
-        token: int
-        logprobs: mx.array
-        finish_reason: Optional[str]
-
-    def __init__(
-        self,
-        model,
-        max_tokens: int = ...,
-        stop_tokens: Optional[set] = ...,
-        sampler: Optional[Callable[[mx.array], mx.array]] = ...,
-        completion_batch_size: int = ...,
-        prefill_batch_size: int = ...,
-        prefill_step_size: int = ...,
-    ) -> None: ...
-    def insert(
-        self, prompts, max_tokens: Union[List[int], int, None] = ...
-    ):  # -> list[Any]:
-        ...
-    def stats(self):  # -> BatchStats:
-        ...
-    def next(self):  # -> list[Any]:
-        ...
+    caches: Optional[List[List[Any]]]
+    ...
 
 def batch_generate(
     model,
     tokenizer,
-    prompts: List[int],
+    prompts: List[List[int]],
+    prompt_caches: Optional[List[List[Any]]] = ...,
     max_tokens: Union[int, List[int]] = ...,
     verbose: bool = ...,
+    return_prompt_caches: bool = ...,
+    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = ...,
     **kwargs,
 ) -> BatchResponse:
     """
@@ -309,14 +533,22 @@ def batch_generate(
     Args:
        model (nn.Module): The language model.
        tokenizer (PreTrainedTokenizer): The tokenizer.
-       prompt (List[List[int]]): The input prompts.
+       prompts (List[List[int]]): The input prompts.
+       prompt_caches (List[List[Any]], optional): Pre-computed prompt-caches
+          for each input prompt. Note, unlike ``generate_step``, the caches
+          won't be updated in-place.
        verbose (bool): If ``True``, print tokens and timing information.
           Default: ``False``.
        max_tokens (Union[int, List[int]): Maximum number of output tokens. This
           can be per prompt if a list is provided.
+       return_prompt_caches (bool): Return the prompt caches in the batch
+          responses. Default: ``False``.
+       logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
+          A list of functions that take tokens and logits and return the processed logits. Default: ``None``.
        kwargs: The remaining options get passed to :obj:`BatchGenerator`.
           See :obj:`BatchGenerator` for more details.
     """
+    ...
 
 def main():  # -> None:
     ...

.mlx_typings/mlx_lm/models/activations.pyi

@@ -0,0 +1,19 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.core as mx
+import mlx.nn as nn
+from functools import partial
+
+@partial(mx.compile, shapeless=True)
+def swiglu(gate, x): ...
+@partial(mx.compile, shapeless=True)
+def xielu(x, alpha_p, alpha_n, beta, eps):  # -> array:
+    ...
+
+class XieLU(nn.Module):
+    def __init__(
+        self, alpha_p_init=..., alpha_n_init=..., beta=..., eps=...
+    ) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...

.mlx_typings/mlx_lm/models/cache.pyi

@@ -16,7 +16,7 @@ class Cache(Protocol):
         self, keys: mx.array, values: mx.array
     ) -> tuple[mx.array, mx.array]: ...
     @property
-    def state(self) -> tuple[mx.array, mx.array]: ...
+    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
     @state.setter
     def state(self, v) -> None: ...
 
@@ -88,17 +88,18 @@ def create_attention_mask(
 ) -> array | Literal["causal"] | None: ...
 
 class _BaseCache(Cache):
-    keys: mx.array
-    values: mx.array
+    keys: mx.array | None
+    values: mx.array | None
     offset: int
     @property
-    def state(self) -> tuple[mx.array, mx.array]: ...
+    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
     @state.setter
     def state(self, v) -> None: ...
     @property
     def meta_state(self) -> Literal[""]: ...
     @meta_state.setter
     def meta_state(self, v) -> None: ...
+    def trim(self, n: int) -> int: ...
     def is_trimmable(self) -> Literal[False]: ...
     @classmethod
     def from_state(cls, state, meta_state) -> Self: ...
@@ -114,15 +115,13 @@ class ConcatenateKVCache(_BaseCache):
     def update_and_fetch(self, keys, values):  # -> tuple[Any | array, Any | array]:
         ...
     @property
-    def state(self):  # -> tuple[Any | array | None, Any | array | None]:
-        ...
+    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
     @state.setter
     def state(self, v):  # -> None:
         ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n):  # -> int:
-        ...
+    def trim(self, n: int) -> int: ...
     def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
         ...
 
@@ -132,10 +131,7 @@ class QuantizedKVCache(_BaseCache):
     def update_and_fetch(self, keys, values):  # -> Any:
         ...
     @property
-    def state(
-        self,
-    ):  # -> tuple[Any | tuple[array, array, array] | None, Any | tuple[array, array, array] | None] | Any:
-        ...
+    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -147,8 +143,7 @@ class QuantizedKVCache(_BaseCache):
         ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n):  # -> int:
-        ...
+    def trim(self, n: int) -> int: ...
     def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
         ...
 
@@ -160,22 +155,30 @@ class KVCache(_BaseCache):
     @property
     def state(
         self,
-    ) -> tuple[array, array]: ...
+    ) -> tuple[mx.array | None, mx.array | None]: ...
     @state.setter
     def state(self, v) -> None: ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n):  # -> int:
-        ...
+    def trim(self, n: int) -> int: ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ...
     ) -> QuantizedKVCache: ...
-    def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
-        ...
+    def make_mask(
+        self, *args: Any, **kwargs: Any
+    ) -> mx.array | Literal["causal"] | None: ...
 
 class RotatingKVCache(_BaseCache):
     step = ...
+    keys: mx.array | None
+    values: mx.array | None
+    keep: int
+    max_size: int
+    _idx: int
     def __init__(self, max_size, keep=...) -> None: ...
+    def _trim(
+        self, trim_size: int, v: mx.array, append: mx.array | None = ...
+    ) -> mx.array: ...
     def update_and_fetch(
         self, keys, values
     ):  # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
@@ -183,8 +186,7 @@ class RotatingKVCache(_BaseCache):
     @property
     def state(
         self,
-    ):  # -> tuple[Any | array, Any | array] | tuple[Any | array | None, Any | array | None]:
-        ...
+    ) -> tuple[mx.array | None, mx.array | None]: ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -196,8 +198,7 @@ class RotatingKVCache(_BaseCache):
         ...
     def is_trimmable(self):  # -> bool:
         ...
-    def trim(self, n):  # -> int:
-        ...
+    def trim(self, n: int) -> int: ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ...
     ) -> QuantizedKVCache: ...
@@ -212,8 +213,7 @@ class ArraysCache(_BaseCache):
         ...
     def __getitem__(self, idx): ...
     @property
-    def state(self):  # -> list[Any | array] | list[array]:
-        ...
+    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -227,8 +227,7 @@ class ArraysCache(_BaseCache):
         In-place extend this cache with the other cache.
         """
 
-    def make_mask(self, N: int):  # -> array | None:
-        ...
+    def make_mask(self, N: int) -> mx.array | None: ...
 
 class MambaCache(ArraysCache):
     def __init__(self, left_padding: Optional[List[int]] = ...) -> None: ...
@@ -239,8 +238,7 @@ class ChunkedKVCache(KVCache):
         ...
     def update_and_fetch(self, keys, values):  # -> tuple[array, array]:
         ...
-    def trim(self, n):  # -> int:
-        ...
+    def trim(self, n: int) -> int: ...
     @property
     def meta_state(self):  # -> tuple[str, ...]:
         ...
@@ -253,10 +251,9 @@ class CacheList(_BaseCache):
     def __getitem__(self, idx): ...
     def is_trimmable(self):  # -> bool:
         ...
-    def trim(self, n): ...
+    def trim(self, n: int) -> int: ...
     @property
-    def state(self):  # -> list[Any]:
-        ...
+    def state(self) -> list[tuple[mx.array | None, mx.array | None]]: ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -271,29 +268,14 @@ class CacheList(_BaseCache):
         """
 
 class BatchKVCache(_BaseCache):
-    step = ...
-    def __init__(self, left_padding: List[int]) -> None:
-        """
-        The BatchKV cache expects inputs to be left-padded.
-
-        E.g. the following prompts:
-
-            [1, 3, 5]
-            [7]
-            [2, 6, 8, 9]
-
-        Should be padded like so:
-
-            [0, 1, 3, 5]
-            [0, 0, 0, 7]
-            [2, 6, 8, 9]
-
-        And ``left_padding`` specifies the amount of padding for each.
-        In this case, ``left_padding = [1, 3, 0]``.
-        """
-
-    def update_and_fetch(self, keys, values):  # -> tuple[array | Any, array | Any]:
-        ...
+    step: int
+    keys: array | None
+    values: array | None
+    offset: array
+    left_padding: array
+    _idx: int
+    def __init__(self, left_padding: List[int]) -> None: ...
+    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
     @property
     def state(
         self,
@@ -319,12 +301,21 @@ class BatchKVCache(_BaseCache):
         """
 
 class BatchRotatingKVCache(_BaseCache):
-    step = ...
-    def __init__(self, max_size, left_padding: List[int]) -> None: ...
-    def update_and_fetch(
-        self, keys, values
-    ):  # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
-        ...
+    step: int
+    keys: array | None
+    values: array | None
+    offset: array
+    left_padding: array
+    max_size: int
+    _idx: int
+    _offset: int
+    rotated: bool
+    _lengths: array | None
+    def __init__(self, max_size: int, left_padding: List[int]) -> None: ...
+    def _trim(self, trim_size: int, v: array, append: array | None = ...) -> array: ...
+    def _update_in_place(self, keys: array, values: array) -> tuple[array, array]: ...
+    def _update_concat(self, keys: array, values: array) -> tuple[array, array]: ...
+    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
     @property
     def state(
         self,

.mlx_typings/mlx_lm/models/gated_delta.pyi

@@ -0,0 +1,35 @@
+from typing import Optional
+
+import mlx.core as mx
+
+def compute_g(A_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array: ...
+def gated_delta_update(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    a: mx.array,
+    b: mx.array,
+    A_log: mx.array,
+    dt_bias: mx.array,
+    state: Optional[mx.array] = ...,
+    mask: Optional[mx.array] = ...,
+    use_kernel: bool = ...,
+) -> tuple[mx.array, mx.array]: ...
+def gated_delta_ops(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: Optional[mx.array] = ...,
+    mask: Optional[mx.array] = ...,
+) -> tuple[mx.array, mx.array]: ...
+def gated_delta_kernel(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: mx.array,
+    mask: Optional[mx.array] = ...,
+) -> tuple[mx.array, mx.array]: ...

.mlx_typings/mlx_lm/models/gemma4.pyi

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from . import gemma4_text
+from .base import BaseModelArgs
+from .cache import KVCache, RotatingKVCache
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: Optional[dict[str, Any]]
+    vocab_size: int
+
+    def __post_init__(self) -> None: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    model_type: str
+    language_model: gemma4_text.Model
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    @property
+    def layers(self) -> list[gemma4_text.DecoderLayer]: ...
+    @property
+    def quant_predicate(self) -> Any: ...
+    def make_cache(self) -> list[KVCache | RotatingKVCache]: ...

.mlx_typings/mlx_lm/models/gemma4_text.pyi

@@ -0,0 +1,179 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import KVCache, RotatingKVCache
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    head_dim: int
+    global_head_dim: int
+    global_partial_rotary_factor: float
+    rms_norm_eps: float
+    vocab_size: int
+    vocab_size_per_layer_input: int
+    num_key_value_heads: int
+    num_global_key_value_heads: Optional[int]
+    num_kv_shared_layers: int
+    pad_token_id: int
+    hidden_size_per_layer_input: int
+    rope_traditional: bool
+    partial_rotary_factor: float
+    rope_parameters: Optional[Dict[str, Any]]
+    sliding_window: int
+    sliding_window_pattern: int
+    max_position_embeddings: int
+    attention_k_eq_v: bool
+    final_logit_softcapping: float
+    use_double_wide_mlp: bool
+    enable_moe_block: bool
+    num_experts: Optional[int]
+    top_k_experts: Optional[int]
+    moe_intermediate_size: Optional[int]
+    layer_types: Optional[List[str]]
+    tie_word_embeddings: bool
+
+    def __post_init__(self) -> None: ...
+
+class MLP(nn.Module):
+    gate_proj: nn.Linear
+    down_proj: nn.Linear
+    up_proj: nn.Linear
+
+    def __init__(self, config: ModelArgs, layer_idx: int = 0) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class Router(nn.Module):
+    proj: nn.Linear
+    scale: mx.array
+    per_expert_scale: mx.array
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
+
+class Experts(nn.Module):
+    switch_glu: SwitchGLU
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self, x: mx.array, top_k_indices: mx.array, top_k_weights: mx.array
+    ) -> mx.array: ...
+
+class Attention(nn.Module):
+    layer_idx: int
+    layer_type: str
+    is_sliding: bool
+    head_dim: int
+    n_heads: int
+    n_kv_heads: int
+    use_k_eq_v: bool
+    scale: float
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    q_norm: nn.Module
+    k_norm: nn.Module
+    v_norm: nn.Module
+    rope: nn.Module
+
+    def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+
+class DecoderLayer(nn.Module):
+    layer_idx: int
+    layer_type: str
+    self_attn: Attention
+    mlp: MLP
+    enable_moe: bool
+    router: Router
+    experts: Experts
+    input_layernorm: nn.Module
+    post_attention_layernorm: nn.Module
+    pre_feedforward_layernorm: nn.Module
+    post_feedforward_layernorm: nn.Module
+    post_feedforward_layernorm_1: nn.Module
+    post_feedforward_layernorm_2: nn.Module
+    pre_feedforward_layernorm_2: nn.Module
+    hidden_size_per_layer_input: int
+    per_layer_input_gate: Optional[nn.Linear]
+    per_layer_projection: Optional[nn.Linear]
+    post_per_layer_input_norm: Optional[nn.Module]
+    layer_scalar: mx.array
+
+    def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = ...,
+        cache: Optional[Any] = ...,
+        per_layer_input: Optional[mx.array] = ...,
+        shared_kv: Optional[tuple[mx.array, mx.array]] = ...,
+        offset: Optional[mx.array] = ...,
+    ) -> tuple[mx.array, tuple[mx.array, mx.array], mx.array]: ...
+
+class Gemma4TextModel(nn.Module):
+    config: ModelArgs
+    vocab_size: int
+    window_size: int
+    sliding_window_pattern: int
+    num_hidden_layers: int
+    embed_tokens: nn.Embedding
+    embed_scale: float
+    layers: list[DecoderLayer]
+    norm: nn.Module
+    hidden_size_per_layer_input: int
+    embed_tokens_per_layer: Optional[nn.Embedding]
+    per_layer_model_projection: Optional[nn.Linear]
+    per_layer_projection_norm: Optional[nn.Module]
+    previous_kvs: list[int]
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: Optional[mx.array] = ...,
+        cache: Optional[list[Any]] = ...,
+        input_embeddings: Optional[mx.array] = ...,
+        per_layer_inputs: Optional[mx.array] = ...,
+    ) -> mx.array: ...
+    def _get_per_layer_inputs(
+        self,
+        input_ids: Optional[mx.array],
+        input_embeddings: Optional[mx.array] = ...,
+    ) -> mx.array: ...
+    def _project_per_layer_inputs(
+        self,
+        input_embeddings: mx.array,
+        per_layer_inputs: Optional[mx.array] = ...,
+    ) -> mx.array: ...
+    def _make_masks(self, h: mx.array, cache: list[Any]) -> list[Any]: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    model_type: str
+    model: Gemma4TextModel
+    final_logit_softcapping: float
+    tie_word_embeddings: bool
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    @property
+    def layers(self) -> list[DecoderLayer]: ...
+    @property
+    def head_dim(self) -> int: ...
+    @property
+    def n_kv_heads(self) -> int: ...
+    @property
+    def quant_predicate(self) -> Any: ...
+    def make_cache(self) -> list[KVCache | RotatingKVCache]: ...

.mlx_typings/mlx_lm/models/nemotron_h.pyi

@@ -0,0 +1,154 @@
+from dataclasses import dataclass
+from typing import Any, List, Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .cache import ArraysCache, KVCache
+from .switch_layers import SwitchMLP
+
+@dataclass
+class ModelArgs:
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    max_position_embeddings: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    attention_bias: bool
+    mamba_num_heads: int
+    mamba_head_dim: int
+    mamba_proj_bias: bool
+    ssm_state_size: int
+    conv_kernel: int
+    n_groups: int
+    mlp_bias: bool
+    layer_norm_epsilon: float
+    use_bias: bool
+    use_conv_bias: bool
+    hybrid_override_pattern: List[str]
+    head_dim: Optional[int]
+    moe_intermediate_size: Optional[int]
+    moe_shared_expert_intermediate_size: Optional[int]
+    n_group: Optional[int]
+    n_routed_experts: Optional[int]
+    n_shared_experts: Optional[int]
+    topk_group: Optional[int]
+    num_experts_per_tok: Optional[int]
+    norm_topk_prob: Optional[bool]
+    routed_scaling_factor: Optional[float]
+    time_step_limit: Optional[Tuple[float, float]]
+    time_step_min: Optional[float]
+    time_step_max: Optional[float]
+
+    @classmethod
+    def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
+    def __post_init__(self) -> None: ...
+
+class NemotronHMamba2Mixer(nn.Module):
+    num_heads: int
+    hidden_size: int
+    ssm_state_size: int
+    conv_kernel_size: int
+    intermediate_size: int
+    n_groups: int
+    head_dim: int
+    conv_dim: int
+    conv1d: nn.Conv1d
+    in_proj: nn.Linear
+    dt_bias: mx.array
+    A_log: mx.array
+    D: mx.array
+    norm: nn.RMSNorm
+    heads_per_group: int
+    out_proj: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        mask: Optional[mx.array],
+        cache: Optional[ArraysCache] = None,
+    ) -> mx.array: ...
+
+class NemotronHAttention(nn.Module):
+    hidden_size: int
+    num_heads: int
+    head_dim: int
+    num_key_value_heads: int
+    scale: float
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[KVCache] = None,
+    ) -> mx.array: ...
+
+class NemotronHMLP(nn.Module):
+    up_proj: nn.Linear
+    down_proj: nn.Linear
+
+    def __init__(
+        self, args: ModelArgs, intermediate_size: Optional[int] = None
+    ) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class NemotronHMoE(nn.Module):
+    num_experts_per_tok: int
+    switch_mlp: SwitchMLP
+    shared_experts: NemotronHMLP
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class NemotronHBlock(nn.Module):
+    block_type: str
+    norm: nn.RMSNorm
+    mixer: NemotronHMamba2Mixer | NemotronHAttention | NemotronHMLP | NemotronHMoE
+
+    def __init__(self, args: ModelArgs, block_type: str) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class NemotronHModel(nn.Module):
+    embeddings: nn.Embedding
+    layers: list[NemotronHBlock]
+    norm_f: nn.RMSNorm
+    fa_idx: int
+    ssm_idx: int
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    backbone: NemotronHModel
+    lm_head: nn.Linear
+    model_type: str
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[NemotronHBlock]: ...
+    def make_cache(self) -> list[ArraysCache | KVCache]: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...

.mlx_typings/mlx_lm/models/qwen3_5.pyi

@@ -0,0 +1,153 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .cache import ArraysCache, KVCache
+from .qwen3_next import (
+    Qwen3NextAttention as Attention,
+    Qwen3NextMLP as MLP,
+    Qwen3NextRMSNormGated as RMSNormGated,
+    Qwen3NextSparseMoeBlock,
+)
+
+SparseMoeBlock = Qwen3NextSparseMoeBlock
+from .switch_layers import SwitchGLU
+
+@dataclass
+class TextModelArgs:
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    max_position_embeddings: int
+    linear_num_value_heads: int
+    linear_num_key_heads: int
+    linear_key_head_dim: int
+    linear_value_head_dim: int
+    linear_conv_kernel_dim: int
+    tie_word_embeddings: bool
+    attention_bias: bool
+    head_dim: Optional[int]
+    full_attention_interval: int
+    num_experts: int
+    num_experts_per_tok: int
+    decoder_sparse_step: int
+    shared_expert_intermediate_size: int
+    moe_intermediate_size: int
+    norm_topk_prob: bool
+    rope_parameters: Optional[dict[str, Any]]
+    partial_rotary_factor: float
+    rope_theta: float
+    rope_scaling: Optional[dict[str, Any]]
+
+    @classmethod
+    def from_dict(cls, params: dict[str, Any]) -> TextModelArgs: ...
+    def __post_init__(self) -> None: ...
+
+class GatedDeltaNet(nn.Module):
+    hidden_size: int
+    num_v_heads: int
+    num_k_heads: int
+    head_k_dim: int
+    head_v_dim: int
+    key_dim: int
+    value_dim: int
+    conv_kernel_size: int
+    conv_dim: int
+    conv1d: nn.Conv1d
+    in_proj_qkv: nn.Linear
+    in_proj_z: nn.Linear
+    in_proj_b: nn.Linear
+    in_proj_a: nn.Linear
+    dt_bias: mx.array
+    A_log: mx.array
+    norm: RMSNormGated
+    out_proj: nn.Linear
+
+    def __init__(self, config: TextModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class DecoderLayer(nn.Module):
+    is_linear: bool
+    linear_attn: GatedDeltaNet
+    self_attn: Attention
+    input_layernorm: nn.RMSNorm
+    post_attention_layernorm: nn.RMSNorm
+    mlp: MLP | SparseMoeBlock
+
+    def __init__(self, args: TextModelArgs, layer_idx: int) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Qwen3_5TextModel(nn.Module):
+    embed_tokens: nn.Embedding
+    layers: list[DecoderLayer]
+    norm: nn.RMSNorm
+    ssm_idx: int
+    fa_idx: int
+
+    def __init__(self, args: TextModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array: ...
+
+class TextModel(nn.Module):
+    args: TextModelArgs
+    model_type: str
+    model: Qwen3_5TextModel
+    lm_head: nn.Linear
+
+    def __init__(self, args: TextModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array: ...
+    @property
+    def layers(self) -> list[DecoderLayer]: ...
+    def make_cache(self) -> list[ArraysCache | KVCache]: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+
+@dataclass
+class ModelArgs:
+    model_type: str
+    text_config: dict[str, Any]
+
+    @classmethod
+    def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    model_type: str
+    language_model: TextModel
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    @property
+    def layers(self) -> list[DecoderLayer]: ...
+    def make_cache(self) -> list[ArraysCache | KVCache]: ...

.mlx_typings/mlx_lm/models/qwen3_5_moe.pyi

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .cache import ArraysCache, KVCache
+from .qwen3_5 import DecoderLayer, Model as Qwen3_5Model, TextModel
+
+@dataclass
+class ModelArgs:
+    model_type: str
+    text_config: dict[str, Any]
+
+    @classmethod
+    def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
+
+class Model(Qwen3_5Model):
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...

.mlx_typings/mlx_lm/models/qwen3_next.pyi

@@ -5,8 +5,18 @@ from typing import Any, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
+from .cache import ArraysCache, KVCache
 from .switch_layers import SwitchGLU
 
+class Qwen3NextRMSNormGated(nn.Module):
+    eps: float
+    weight: mx.array
+
+    def __init__(self, hidden_size: int, eps: float = ...) -> None: ...
+    def __call__(
+        self, hidden_states: mx.array, gate: mx.array | None = None
+    ) -> mx.array: ...
+
 class Qwen3NextMLP(nn.Module):
     gate_proj: nn.Linear
     down_proj: nn.Linear
@@ -90,6 +100,8 @@ class Qwen3NextModel(nn.Module):
     embed_tokens: nn.Embedding
     layers: list[Qwen3NextDecoderLayer]
     norm: nn.RMSNorm
+    ssm_idx: int
+    fa_idx: int
 
     def __init__(self, args: Any) -> None: ...
     def __call__(
@@ -112,3 +124,4 @@ class Model(nn.Module):
     def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
     @property
     def layers(self) -> list[Qwen3NextDecoderLayer]: ...
+    def make_cache(self) -> list[ArraysCache | KVCache]: ...

.mlx_typings/mlx_lm/models/rope_utils.pyi

@@ -0,0 +1,51 @@
+from typing import Any, Optional
+
+import mlx.nn as nn
+
+class YarnRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        scaling_factor: float = ...,
+        original_max_position_embeddings: int = ...,
+        beta_fast: float = ...,
+        beta_slow: float = ...,
+        mscale: float = ...,
+        mscale_all_dim: float = ...,
+    ) -> None: ...
+
+class Llama3RoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        scaling_factor: float = ...,
+        original_max_position_embeddings: int = ...,
+        low_freq_factor: float = ...,
+        high_freq_factor: float = ...,
+    ) -> None: ...
+
+class SuScaledRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        short_factor: Any = ...,
+        long_factor: Any = ...,
+        original_max_position_embeddings: int = ...,
+    ) -> None: ...
+
+def initialize_rope(
+    dims: int,
+    base: float = ...,
+    traditional: bool = ...,
+    scaling_config: Optional[dict[str, Any]] = ...,
+    max_position_embeddings: Optional[int] = ...,
+) -> nn.Module: ...

.mlx_typings/mlx_lm/models/switch_layers.pyi

@@ -73,6 +73,9 @@ class SwitchGLU(nn.Module):
     def __call__(self, x, indices) -> mx.array: ...
 
 class SwitchMLP(nn.Module):
+    fc1: SwitchLinear
+    fc2: SwitchLinear
+
     def __init__(
         self,
         input_dims: int,

.mlx_typings/mlx_lm/sample_utils.pyi

@@ -48,7 +48,7 @@ def make_logits_processors(
     logit_bias: Optional[Dict[int, float]] = ...,
     repetition_penalty: Optional[float] = ...,
     repetition_context_size: Optional[int] = ...,
-):  # -> list[Any]:
+) -> list[Callable[[mx.array, mx.array], mx.array]]:
     """
     Make logits processors for use with ``generate_step``.
 

.mlx_typings/mlx_lm/tokenizer_utils.pyi

@@ -39,11 +39,11 @@ class StreamingDetokenizer:
     """
 
     __slots__ = ...
-    def reset(self): ...
-    def add_token(self, token): ...
-    def finalize(self): ...
+    def reset(self) -> None: ...
+    def add_token(self, token: int) -> None: ...
+    def finalize(self) -> None: ...
     @property
-    def last_segment(self):
+    def last_segment(self) -> str:
         """Return the last segment of readable text since last time this property was accessed."""
 
 class NaiveStreamingDetokenizer(StreamingDetokenizer):
@@ -117,6 +117,8 @@ class TokenizerWrapper:
     think_end: str | None
     think_start_id: int | None
     think_end_id: int | None
+    think_start_tokens: list[int] | None
+    think_end_tokens: list[int] | None
 
     def __init__(
         self,

.mlx_typings/mlx_lm/tuner/dora.pyi

@@ -0,0 +1,50 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.nn as nn
+
+class DoRALinear(nn.Module):
+    @staticmethod
+    def from_base(
+        linear: nn.Linear, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> DoRALinear:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> QuantizedLinear | Linear:
+        ...
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+        bias: bool = ...,
+    ) -> None: ...
+    def set_linear(self, linear):  # -> None:
+        """
+        Set the self.linear layer and recompute self.m.
+        """
+        ...
+
+    def __call__(self, x): ...
+
+class DoRAEmbedding(nn.Module):
+    def from_base(
+        embedding: nn.Embedding, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> DoRAEmbedding:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> Embedding:
+        ...
+    def __init__(
+        self,
+        num_embeddings: int,
+        dims: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+    ) -> None: ...
+    def set_embedding(self, embedding: nn.Module):  # -> None:
+        ...
+    def __call__(self, x): ...
+    def as_linear(self, x): ...

.mlx_typings/mlx_lm/tuner/lora.pyi

@@ -0,0 +1,66 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.nn as nn
+
+class LoRALinear(nn.Module):
+    @staticmethod
+    def from_base(
+        linear: nn.Linear, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> LoRALinear:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> QuantizedLinear | Linear:
+        ...
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+        bias: bool = ...,
+    ) -> None: ...
+    def __call__(self, x):  # -> array:
+        ...
+
+class LoRASwitchLinear(nn.Module):
+    @staticmethod
+    def from_base(
+        linear: nn.Module, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> LoRASwitchLinear:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> QuantizedSwitchLinear | SwitchLinear:
+        ...
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        num_experts: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+        bias: bool = ...,
+    ) -> None: ...
+    def __call__(self, x, indices, sorted_indices=...): ...
+
+class LoRAEmbedding(nn.Module):
+    @staticmethod
+    def from_base(
+        embedding: nn.Embedding, r: int = ..., dropout: float = ..., scale: float = ...
+    ):  # -> LoRAEmbedding:
+        ...
+    def fuse(self, dequantize: bool = ...):  # -> QuantizedEmbedding | Embedding:
+        ...
+    def __init__(
+        self,
+        num_embeddings: int,
+        dims: int,
+        r: int = ...,
+        dropout: float = ...,
+        scale: float = ...,
+    ) -> None: ...
+    def __call__(self, x):  # -> array:
+        ...
+    def as_linear(self, x):  # -> array:
+        ...

.mlx_typings/mlx_lm/tuner/utils.pyi

@@ -0,0 +1,57 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.nn as nn
+from typing import Dict
+
+def build_schedule(schedule_config: Dict):  # -> Any:
+    """
+    Build a learning rate schedule from the given config.
+    """
+    ...
+
+def linear_to_lora_layers(
+    model: nn.Module, num_layers: int, config: Dict, use_dora: bool = ...
+):  # -> None:
+    """
+    Convert some of the models linear layers to lora layers.
+
+    Args:
+        model (nn.Module): The neural network model.
+        num_layers (int): The number of blocks to convert to lora layers
+        starting from the last layer.
+        config (dict): More configuration parameters for LoRA, including the
+          rank, scale, and optional layer keys.
+        use_dora (bool): If True, uses DoRA instead of LoRA.
+          Default: ``False``
+    """
+    ...
+
+def load_adapters(model: nn.Module, adapter_path: str) -> nn.Module:
+    """
+    Load any fine-tuned adapters / layers.
+
+    Args:
+        model (nn.Module): The neural network model.
+        adapter_path (str): Path to the adapter configuration file.
+
+    Returns:
+        nn.Module: The updated model with LoRA layers applied.
+    """
+    ...
+
+def remove_lora_layers(model: nn.Module) -> nn.Module:
+    """
+    Remove the LoRA layers from the model.
+
+    Args:
+        model (nn.Module): The model with LoRA layers.
+
+    Returns:
+        nn.Module: The model without LoRA layers.
+    """
+    ...
+
+def print_trainable_parameters(model):  # -> None:
+    ...

.mlx_typings/mlx_vlm/prompt_utils.pyi

@@ -0,0 +1,12 @@
+from typing import Any
+
+def get_message_json(
+    model_name: str,
+    prompt: str,
+    role: str = "user",
+    skip_image_token: bool = False,
+    skip_audio_token: bool = False,
+    num_images: int = 0,
+    num_audios: int = 0,
+    **kwargs: Any,
+) -> dict[str, Any]: ...

.mlx_typings/mlx_vlm/utils.pyi

@@ -0,0 +1,15 @@
+from pathlib import Path
+from typing import Any
+
+class ImageProcessor:
+    def preprocess(
+        self, images: list[dict[str, Any]], **kwargs: Any
+    ) -> dict[str, Any]: ...
+    def __call__(self, **kwargs: Any) -> dict[str, Any]: ...
+
+def load_image_processor(
+    model_path: str | Path, **kwargs: Any
+) -> ImageProcessor | None: ...
+def load_processor(
+    model_path: str | Path, add_detokenizer: bool = ..., **kwargs: Any
+) -> ImageProcessor: ...

.mlx_typings/safetensors/__init__.pyi

@@ -0,0 +1,8 @@
+from typing import Any, Self
+
+class safe_open:
+    def __init__(self, filename: str, framework: str = "pt") -> None: ...
+    def __enter__(self) -> Self: ...
+    def __exit__(self, *args: Any) -> None: ...
+    def keys(self) -> list[str]: ...
+    def get_tensor(self, name: str) -> Any: ...

.vscode/extensions.json

@@ -1,7 +1,8 @@
 {
     "recommendations": [
         "detachhead.basedpyright",
-        "ms-python.python"
+        "ms-python.python",
+        "jnoortheen.nix-ide"
     ],
     "unwantedRecommendations": [
         "ms-python.vscode-pylance",

.vscode/settings.json

@@ -1,3 +1,32 @@
 {
-    "basedpyright.importStrategy": "fromEnvironment"
+    "files.associations": {
+        "*.nix": "nix",
+    },
+    "nix.enableLanguageServer": true,
+    "nix.serverPath": "nixd",
+    "nix.serverSettings": {
+        "nixd": {
+        "formatting": {
+            "command": ["nixpkgs-fmt"]
+        },
+        "nixpkgs": {
+            "expr": "(builtins.getFlake \"path:${workspaceFolder}\").currentSystem.config._module.args.pkgs"
+        },
+        "options": {
+            "flake-parts": {
+                "expr": "(builtins.getFlake \"path:${workspaceFolder}\").debug.options"
+            },
+            "flake-parts-perSystem": {
+                "expr": "(builtins.getFlake \"path:${workspaceFolder}\").currentSystem.options"
+            }
+        }
+      }
+    },
+    "[nix]": {
+        "editor.defaultFormatter": "jnoortheen.nix-ide"
+    },
+    
+    "python.defaultInterpreterPath": "${workspaceFolder}/.venv/bin/python",
+    "basedpyright.analysis.configFilePath": "${workspaceFolder}/pyproject.toml",
+    "basedpyright.importStrategy": "fromEnvironment",
 }

CONTRIBUTING.md

@@ -11,9 +11,18 @@ To run EXO from source:
   ```bash
   brew install uv
   ```
-- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
+- [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now)
   ```bash
-  brew install macmon
+  curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
+  rustup toolchain install nightly
+  ```
+- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
+  Use the pinned fork revision used by this repo instead of Homebrew `macmon`.
+  ```bash
+  cargo install --git https://github.com/vladkens/macmon \
+    --rev a1cd06b6cc0d5e61db24fd8832e74cd992097a7d \
+    macmon \
+    --force
   ```
 
 ```bash
@@ -39,6 +48,119 @@ Write pure functions where possible. When adding new code, prefer Rust unless th
 
 Run `nix fmt` to auto-format your code before submitting.
 
+## Model Cards
+
+EXO uses TOML-based model cards to define model metadata and capabilities. Model cards are stored in:
+- `resources/inference_model_cards/` for text generation models
+- `resources/image_model_cards/` for image generation models
+- `~/.exo/custom_model_cards/` for user-added custom models
+
+### Adding a Model Card
+
+To add a new model, create a TOML file with the following structure:
+
+```toml
+model_id = "mlx-community/Llama-3.2-1B-Instruct-4bit"
+n_layers = 16
+hidden_size = 2048
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "llama"
+quantization = "4bit"
+base_model = "Llama 3.2 1B"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 729808896
+```
+
+### Required Fields
+
+- `model_id`: Hugging Face model identifier
+- `n_layers`: Number of transformer layers
+- `hidden_size`: Hidden dimension size
+- `supports_tensor`: Whether the model supports tensor parallelism
+- `tasks`: List of supported tasks (`TextGeneration`, `TextToImage`, `ImageToImage`)
+- `family`: Model family (e.g., "llama", "deepseek", "qwen")
+- `quantization`: Quantization level (e.g., "4bit", "8bit", "bf16")
+- `base_model`: Human-readable base model name
+- `capabilities`: List of capabilities (e.g., `["text"]`, `["text", "thinking"]`)
+
+### Optional Fields
+
+- `components`: For multi-component models (like image models with separate text encoders and transformers)
+- `uses_cfg`: Whether the model uses classifier-free guidance (for image models)
+- `trust_remote_code`: Whether to allow remote code execution (defaults to `false` for security)
+
+### Capabilities
+
+The `capabilities` field defines what the model can do:
+- `text`: Standard text generation
+- `thinking`: Model supports chain-of-thought reasoning
+- `thinking_toggle`: Thinking can be enabled/disabled via `enable_thinking` parameter
+- `image_edit`: Model supports image-to-image editing (FLUX.1-Kontext)
+
+### Security Note
+
+By default, `trust_remote_code` is set to `false` for security. Only enable it if the model explicitly requires remote code execution from the Hugging Face hub.
+
+## API Adapters
+
+EXO supports multiple API formats through an adapter pattern. Adapters convert API-specific request formats to the internal `TextGenerationTaskParams` format and convert internal token chunks back to API-specific responses.
+
+### Adapter Architecture
+
+All adapters live in `src/exo/master/adapters/` and follow the same pattern:
+
+1. Convert API-specific requests to `TextGenerationTaskParams`
+2. Handle both streaming and non-streaming response generation
+3. Convert internal `TokenChunk` objects to API-specific formats
+4. Manage error handling and edge cases
+
+### Existing Adapters
+
+- `chat_completions.py`: OpenAI Chat Completions API
+- `claude.py`: Anthropic Claude Messages API
+- `responses.py`: OpenAI Responses API
+- `ollama.py`: Ollama API (for OpenWebUI compatibility)
+
+### Adding a New API Adapter
+
+To add support for a new API format:
+
+1. Create a new adapter file in `src/exo/master/adapters/`
+2. Implement a request conversion function:
+   ```python
+   def your_api_request_to_text_generation(
+       request: YourAPIRequest,
+   ) -> TextGenerationTaskParams:
+       # Convert API request to internal format
+       pass
+   ```
+3. Implement streaming response generation:
+   ```python
+   async def generate_your_api_stream(
+       command_id: CommandId,
+       chunk_stream: AsyncGenerator[TokenChunk | ErrorChunk | ToolCallChunk, None],
+   ) -> AsyncGenerator[str, None]:
+       # Convert internal chunks to API-specific streaming format
+       pass
+   ```
+4. Implement non-streaming response collection:
+   ```python
+   async def collect_your_api_response(
+       command_id: CommandId,
+       chunk_stream: AsyncGenerator[TokenChunk | ErrorChunk | ToolCallChunk, None],
+   ) -> AsyncGenerator[str]:
+       # Collect all chunks and return single response
+       pass
+   ```
+5. Register the adapter endpoints in `src/exo/master/api.py`
+
+The adapter pattern keeps API-specific logic isolated from core inference systems. Internal systems (worker, runner, event sourcing) only see `TextGenerationTaskParams` and `TokenChunk` objects - no API-specific types cross the adapter boundary.
+
+For detailed API documentation, see [docs/api.md](docs/api.md).
+
 ## Testing
 
 EXO relies heavily on manual testing at this point in the project, but this is evolving. Before submitting a change, test both before and after to demonstrate how your change improves behavior. Do the best you can with the hardware you have available - if you need help testing, ask and we'll do our best to assist. Add automated tests where possible - we're actively working to substantially improve our automated testing story.

Cargo.lock

@@ -216,6 +216,28 @@ dependencies = [
  "windows-sys 0.61.2",
 ]
 
+[[package]]
+name = "async-stream"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476"
+dependencies = [
+ "async-stream-impl",
+ "futures-core",
+ "pin-project-lite",
+]
+
+[[package]]
+name = "async-stream-impl"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.111",
+]
+
 [[package]]
 name = "async-trait"
 version = "0.1.89"
@@ -2759,6 +2781,7 @@ dependencies = [
 name = "networking"
 version = "0.0.1"
 dependencies = [
+ "async-stream",
  "delegate",
  "either",
  "extend",
@@ -2767,6 +2790,7 @@ dependencies = [
  "keccak-const",
  "libp2p",
  "log",
+ "pin-project",
  "tokio",
  "tracing-subscriber",
  "util",

Cargo.toml

@@ -1,10 +1,6 @@
 [workspace]
 resolver = "3"
-members = [
-    "rust/networking",
-    "rust/exo_pyo3_bindings",
-    "rust/util",
-]
+members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
 
 [workspace.package]
 version = "0.0.1"
@@ -34,6 +30,7 @@ delegate = "0.13"
 keccak-const = "0.2"
 
 # Async dependencies
+async-stream = "0.3"
 tokio = "1.46"
 futures-lite = "2.6.1"
 futures-timer = "3.0"

README.md

@@ -26,6 +26,8 @@ exo connects all your devices into an AI cluster. Not only does exo enable runni
 - **Topology-Aware Auto Parallel**: exo figures out the best way to split your model across all available devices based on a realtime view of your device topology. It takes into account device resources and network latency/bandwidth between each link.
 - **Tensor Parallelism**: exo supports sharding models, for up to 1.8x speedup on 2 devices and 3.2x speedup on 4 devices.
 - **MLX Support**: exo uses [MLX](https://github.com/ml-explore/mlx) as an inference backend and [MLX distributed](https://ml-explore.github.io/mlx/build/html/usage/distributed.html) for distributed communication.
+- **Multiple API Compatibility**: Compatible with OpenAI Chat Completions API, Claude Messages API, OpenAI Responses API, and Ollama API - use your existing tools and clients.
+- **Custom Model Support**: Load custom models from HuggingFace hub to expand the range of available models.
 
 ## Dashboard
 
@@ -93,11 +95,10 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
   /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
   ```
 - [uv](https://github.com/astral-sh/uv) (for Python dependency management)
-- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
 - [node](https://github.com/nodejs/node) (for building the dashboard)
 
   ```bash
-  brew install uv macmon node
+  brew install uv node
   ```
 - [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now)
 
@@ -105,6 +106,17 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
   curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
   rustup toolchain install nightly
   ```
+- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
+
+  Install the pinned fork revision used by this repo instead of Homebrew `macmon`.
+  Homebrew `macmon 0.6.1` still crashes on Apple M5.
+
+  ```bash
+  cargo install --git https://github.com/vladkens/macmon \
+    --rev a1cd06b6cc0d5e61db24fd8832e74cd992097a7d \
+    macmon \
+    --force
+  ```
 
 Clone the repo, build the dashboard, and run exo:
 
@@ -196,6 +208,8 @@ exo follows the [XDG Base Directory Specification](https://specifications.freede
 - **Configuration files**: `~/.config/exo/` (or `$XDG_CONFIG_HOME/exo/`)
 - **Data files**: `~/.local/share/exo/` (or `$XDG_DATA_HOME/exo/`)
 - **Cache files**: `~/.cache/exo/` (or `$XDG_CACHE_HOME/exo/`)
+- **Log files**: `~/.cache/exo/exo_log/` (with automatic log rotation)
+- **Custom model cards**: `~/.local/share/exo/custom_model_cards/`
 
 You can override these locations by setting the corresponding XDG environment variables.
 
@@ -275,8 +289,51 @@ After that, RDMA will be enabled in macOS and exo will take care of the rest.
 
 ---
 
+## Environment Variables
+
+exo supports several environment variables for configuration:
+
+| Variable | Description | Default |
+|----------|-------------|---------|
+| `EXO_DEFAULT_MODELS_DIR` | Default directory for model downloads and caches. Always first in the writable dirs list. | `~/.local/share/exo/models` (Linux) or `~/.exo/models` (macOS) |
+| `EXO_MODELS_DIRS` | Colon-separated additional writable directories for model downloads. Checked in order after the default; first with enough free space is used. | None |
+| `EXO_MODELS_READ_ONLY_DIRS` | Colon-separated read-only directories to search for pre-downloaded models (e.g., NFS mounts, shared storage). Models here cannot be deleted. | None |
+| `EXO_OFFLINE` | Run without internet connection (uses only local models) | `false` |
+| `EXO_ENABLE_IMAGE_MODELS` | Enable image model support | `false` |
+| `EXO_LIBP2P_NAMESPACE` | Custom namespace for cluster isolation | None |
+| `EXO_FAST_SYNCH` | Control MLX_METAL_FAST_SYNCH behavior (for JACCL backend) | Auto |
+| `EXO_TRACING_ENABLED` | Enable distributed tracing for performance analysis | `false` |
+
+**Example usage:**
+
+```bash
+# Use pre-downloaded models from NFS mount (read-only)
+EXO_MODELS_READ_ONLY_DIRS=/mnt/nfs/models:/opt/ai-models uv run exo
+
+# Download models to an external SSD (falls back to default dir if full)
+EXO_MODELS_DIRS=/Volumes/ExternalSSD/exo-models uv run exo
+
+# Run in offline mode
+EXO_OFFLINE=true uv run exo
+
+# Enable image models
+EXO_ENABLE_IMAGE_MODELS=true uv run exo
+
+# Use custom namespace for cluster isolation
+EXO_LIBP2P_NAMESPACE=my-dev-cluster uv run exo
+```
+
+---
+
 ### Using the API
 
+exo provides multiple API-compatible interfaces for maximum compatibility with existing tools:
+
+- **OpenAI Chat Completions API** - Compatible with OpenAI clients
+- **Claude Messages API** - Compatible with Anthropic's Claude format
+- **OpenAI Responses API** - Compatible with OpenAI's Responses format
+- **Ollama API** - Compatible with Ollama and tools like OpenWebUI
+
 If you prefer to interact with exo via the API, here is an example creating an instance of a small model (`mlx-community/Llama-3.2-1B-Instruct-4bit`), sending a chat completions request and deleting the instance.
 
 ---
@@ -366,14 +423,85 @@ When you're done, delete the instance by its ID (find it via `/state` or `/insta
 curl -X DELETE http://localhost:52415/instance/YOUR_INSTANCE_ID
 ```
 
+### Claude Messages API Compatibility
+
+Use the Claude Messages API format with the `/v1/messages` endpoint:
+
+```bash
+curl -N -X POST http://localhost:52415/v1/messages \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model": "mlx-community/Llama-3.2-1B-Instruct-4bit",
+    "messages": [
+      {"role": "user", "content": "Hello"}
+    ],
+    "max_tokens": 1024,
+    "stream": true
+  }'
+```
+
+### OpenAI Responses API Compatibility
+
+Use the OpenAI Responses API format with the `/v1/responses` endpoint:
+
+```bash
+curl -N -X POST http://localhost:52415/v1/responses \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model": "mlx-community/Llama-3.2-1B-Instruct-4bit",
+    "messages": [
+      {"role": "user", "content": "Hello"}
+    ],
+    "stream": true
+  }'
+```
+
+### Ollama API Compatibility
+
+exo supports Ollama API endpoints for compatibility with tools like OpenWebUI:
+
+```bash
+# Ollama chat
+curl -X POST http://localhost:52415/ollama/api/chat \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model": "mlx-community/Llama-3.2-1B-Instruct-4bit",
+    "messages": [
+      {"role": "user", "content": "Hello"}
+    ],
+    "stream": false
+  }'
+
+# List models (Ollama format)
+curl http://localhost:52415/ollama/api/tags
+```
+
+### Custom Model Loading from HuggingFace
+
+You can add custom models from the HuggingFace hub:
+
+```bash
+curl -X POST http://localhost:52415/models/add \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model_id": "mlx-community/my-custom-model"
+  }'
+```
+
+**Security Note:**
+
+Custom models requiring `trust_remote_code` in their configuration must be explicitly enabled (default is false) for security. Only enable this if you trust the model's remote code execution. Models are fetched from HuggingFace and stored locally as custom model cards.
+
 **Other useful API endpoints*:**
 
 - List all models: `curl http://localhost:52415/models`
+- List downloaded models only: `curl http://localhost:52415/models?status=downloaded`
+- Search HuggingFace: `curl "http://localhost:52415/models/search?query=llama&limit=10"`
 - Inspect instance IDs and deployment state: `curl http://localhost:52415/state`
 
 For further details, see:
 
-- API basic documentation in [docs/api.md](docs/api.md).
+- API documentation in [docs/api.md](docs/api.md).
 - API types and endpoints in [src/exo/master/api.py](src/exo/master/api.py).
 
 ---
@@ -432,4 +560,4 @@ On macOS, exo uses the GPU. On Linux, exo currently runs on CPU. We are working
 
 ## Contributing
 
-See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo.
+See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo.

app/EXO/EXO/ExoProcessController.swift

@@ -4,9 +4,27 @@ import Foundation
 
 private let customNamespaceKey = "EXOCustomNamespace"
 private let hfTokenKey = "EXOHFToken"
+private let hfEndpointKey = "EXOHFEndpoint"
 private let enableImageModelsKey = "EXOEnableImageModels"
 private let offlineModeKey = "EXOOfflineMode"
+private let fastSynchEnabledKey = "EXOFastSynchEnabled"
 private let onboardingCompletedKey = "EXOOnboardingCompleted"
+private let customEnvironmentVariablesKey = "EXOCustomEnvironmentVariables"
+
+/// A user-defined environment variable that is injected into the exo child
+/// process at launch. Used to pass arbitrary key/value settings to exo
+/// without having to add first-class UI for each one.
+struct CustomEnvironmentVariable: Codable, Identifiable, Equatable {
+    var id: UUID
+    var key: String
+    var value: String
+
+    init(id: UUID = UUID(), key: String = "", value: String = "") {
+        self.id = id
+        self.key = key
+        self.value = value
+    }
+}
 
 @MainActor
 final class ExoProcessController: ObservableObject {
@@ -53,6 +71,14 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(hfToken, forKey: hfTokenKey)
         }
     }
+    @Published var hfEndpoint: String = {
+        return UserDefaults.standard.string(forKey: hfEndpointKey) ?? ""
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(hfEndpoint, forKey: hfEndpointKey)
+        }
+    }
     @Published var enableImageModels: Bool = {
         return UserDefaults.standard.bool(forKey: enableImageModelsKey)
     }()
@@ -69,6 +95,36 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(offlineMode, forKey: offlineModeKey)
         }
     }
+    @Published var fastSynchEnabled: Bool = {
+        if UserDefaults.standard.object(forKey: fastSynchEnabledKey) == nil {
+            return true
+        }
+        return UserDefaults.standard.bool(forKey: fastSynchEnabledKey)
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(fastSynchEnabled, forKey: fastSynchEnabledKey)
+        }
+    }
+    @Published var customEnvironmentVariables: [CustomEnvironmentVariable] = {
+        guard
+            let data = UserDefaults.standard.data(forKey: customEnvironmentVariablesKey),
+            let decoded = try? JSONDecoder().decode(
+                [CustomEnvironmentVariable].self, from: data
+            )
+        else {
+            return []
+        }
+        return decoded
+    }()
+    {
+        didSet {
+            guard let data = try? JSONEncoder().encode(customEnvironmentVariables) else {
+                return
+            }
+            UserDefaults.standard.set(data, forKey: customEnvironmentVariablesKey)
+        }
+    }
 
     /// Fires once when EXO transitions to `.running` for the very first time (fresh install).
     @Published private(set) var isFirstLaunchReady = false
@@ -273,12 +329,16 @@ final class ExoProcessController: ObservableObject {
         if !hfToken.isEmpty {
             environment["HF_TOKEN"] = hfToken
         }
+        if !hfEndpoint.isEmpty {
+            environment["HF_ENDPOINT"] = hfEndpoint
+        }
         if enableImageModels {
             environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
         }
         if offlineMode {
             environment["EXO_OFFLINE"] = "true"
         }
+        environment["EXO_FAST_SYNCH"] = fastSynchEnabled ? "true" : "false"
 
         var paths: [String] = []
         if let existing = environment["PATH"], !existing.isEmpty {
@@ -303,6 +363,16 @@ final class ExoProcessController: ObservableObject {
         }
 
         environment["PATH"] = paths.joined(separator: ":")
+
+        // Apply user-defined arbitrary environment variables last so that
+        // power users can override any of the built-in keys above when
+        // necessary. Empty keys are ignored.
+        for variable in customEnvironmentVariables {
+            let trimmedKey = variable.key.trimmingCharacters(in: .whitespaces)
+            guard !trimmedKey.isEmpty else { continue }
+            environment[trimmedKey] = variable.value
+        }
+
         return environment
     }
 

app/EXO/EXO/Views/SettingsView.swift

@@ -12,8 +12,11 @@ struct SettingsView: View {
 
     @State private var pendingNamespace: String = ""
     @State private var pendingHFToken: String = ""
+    @State private var pendingHFEndpoint: String = ""
     @State private var pendingEnableImageModels = false
     @State private var pendingOfflineMode = false
+    @State private var pendingFastSynchEnabled = false
+    @State private var pendingCustomEnvironmentVariables: [CustomEnvironmentVariable] = []
     @State private var needsRestart = false
     @State private var bugReportInFlight = false
     @State private var bugReportMessage: String?
@@ -33,6 +36,10 @@ struct SettingsView: View {
                 .tabItem {
                     Label("Advanced", systemImage: "wrench.and.screwdriver")
                 }
+            environmentTab
+                .tabItem {
+                    Label("Environment", systemImage: "terminal")
+                }
             aboutTab
                 .tabItem {
                     Label("About", systemImage: "info.circle")
@@ -42,8 +49,11 @@ struct SettingsView: View {
         .onAppear {
             pendingNamespace = controller.customNamespace
             pendingHFToken = controller.hfToken
+            pendingHFEndpoint = controller.hfEndpoint
             pendingEnableImageModels = controller.enableImageModels
             pendingOfflineMode = controller.offlineMode
+            pendingFastSynchEnabled = controller.fastSynchEnabled
+            pendingCustomEnvironmentVariables = controller.customEnvironmentVariables
             needsRestart = false
         }
     }
@@ -74,6 +84,17 @@ struct SettingsView: View {
                     .foregroundColor(.secondary)
             }
 
+            Section {
+                LabeledContent("HuggingFace Endpoint") {
+                    TextField("default", text: $pendingHFEndpoint)
+                        .textFieldStyle(.roundedBorder)
+                        .frame(width: 200)
+                }
+                Text("Defaults to huggingface.co. Use a mirror (e.g. hf-mirror.com) for China.")
+                    .font(.caption)
+                    .foregroundColor(.secondary)
+            }
+
             Section {
                 Toggle("Offline Mode", isOn: $pendingOfflineMode)
                 Text("Skip internet checks and use only locally available models.")
@@ -124,6 +145,23 @@ struct SettingsView: View {
 
     private var advancedTab: some View {
         Form {
+            Section("Performance") {
+                Toggle("Fast Synch Enabled", isOn: $pendingFastSynchEnabled)
+                Text(
+                    "Experimental: enables fast CPU to GPU synchronization. Can sometimes cause a \"GPU lock\" where inference hangs for ~10 seconds before starting. Necessary for low latency with RDMA and Tensor Parallelism."
+                )
+                .font(.caption)
+                .foregroundColor(.secondary)
+
+                HStack {
+                    Spacer()
+                    Button("Save & Restart") {
+                        applyAdvancedSettings()
+                    }
+                    .disabled(!hasAdvancedChanges)
+                }
+            }
+
             Section("Onboarding") {
                 HStack {
                     VStack(alignment: .leading) {
@@ -180,6 +218,81 @@ struct SettingsView: View {
         .padding()
     }
 
+    // MARK: - Environment Tab
+
+    private var environmentTab: some View {
+        Form {
+            Section("Custom Environment Variables") {
+                Text("Passed to the exo process at launch. Override built-in defaults here.")
+                    .font(.caption)
+                    .foregroundColor(.secondary)
+
+                if pendingCustomEnvironmentVariables.isEmpty {
+                    Text("No custom variables.")
+                        .font(.caption)
+                        .foregroundColor(.secondary)
+                } else {
+                    ForEach($pendingCustomEnvironmentVariables) { $variable in
+                        HStack(alignment: .center, spacing: 8) {
+                            VStack(spacing: 4) {
+                                TextField("key", text: $variable.key)
+                                    .labelsHidden()
+                                    .textFieldStyle(.roundedBorder)
+                                    .font(.system(.body, design: .monospaced))
+                                TextField("value", text: $variable.value)
+                                    .labelsHidden()
+                                    .textFieldStyle(.roundedBorder)
+                                    .font(.system(.body, design: .monospaced))
+                            }
+                            VStack(spacing: 4) {
+                                Button {
+                                    pendingCustomEnvironmentVariables.removeAll {
+                                        $0.id == variable.id
+                                    }
+                                } label: {
+                                    Image(systemName: "minus.circle")
+                                }
+                                .buttonStyle(.borderless)
+                                .help("Remove variable")
+                                if !isValidEnvironmentVariableName(variable.key) {
+                                    Image(systemName: "exclamationmark.triangle.fill")
+                                        .foregroundColor(.orange)
+                                        .help(
+                                            "Invalid environment variable name. "
+                                                + "Must match [A-Za-z_][A-Za-z0-9_]*."
+                                        )
+                                }
+                            }
+                        }
+                    }
+                }
+
+                HStack {
+                    Button {
+                        pendingCustomEnvironmentVariables.append(
+                            CustomEnvironmentVariable()
+                        )
+                    } label: {
+                        Label("Add Variable", systemImage: "plus")
+                    }
+                    Spacer()
+                }
+            }
+
+            Section {
+                HStack {
+                    Spacer()
+                    Button("Save & Restart") {
+                        applyEnvironmentSettings()
+                    }
+                    .disabled(!hasEnvironmentChanges)
+                }
+            }
+        }
+        .formStyle(.grouped)
+        .padding()
+    }
+
     // MARK: - About Tab
 
     private var aboutTab: some View {
@@ -454,6 +567,7 @@ struct SettingsView: View {
 
     private var hasGeneralChanges: Bool {
         pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
+            || pendingHFEndpoint != controller.hfEndpoint
             || pendingOfflineMode != controller.offlineMode
     }
 
@@ -461,9 +575,18 @@ struct SettingsView: View {
         pendingEnableImageModels != controller.enableImageModels
     }
 
+    private var hasAdvancedChanges: Bool {
+        pendingFastSynchEnabled != controller.fastSynchEnabled
+    }
+
+    private var hasEnvironmentChanges: Bool {
+        pendingCustomEnvironmentVariables != controller.customEnvironmentVariables
+    }
+
     private func applyGeneralSettings() {
         controller.customNamespace = pendingNamespace
         controller.hfToken = pendingHFToken
+        controller.hfEndpoint = pendingHFEndpoint
         controller.offlineMode = pendingOfflineMode
         restartIfRunning()
     }
@@ -473,6 +596,63 @@ struct SettingsView: View {
         restartIfRunning()
     }
 
+    private func applyAdvancedSettings() {
+        controller.fastSynchEnabled = pendingFastSynchEnabled
+        restartIfRunning()
+    }
+
+    private func applyEnvironmentSettings() {
+        // Trim whitespace from keys and drop empty ones so that the stored
+        // form matches what is actually injected into the child process and
+        // hasEnvironmentChanges doesn't show a stale diff after save.
+        let trimmed: [CustomEnvironmentVariable] =
+            pendingCustomEnvironmentVariables.compactMap { variable in
+                let key = variable.key.trimmingCharacters(in: .whitespaces)
+                guard !key.isEmpty else { return nil }
+                return CustomEnvironmentVariable(
+                    id: variable.id, key: key, value: variable.value
+                )
+            }
+
+        // De-duplicate keys, keeping the last occurrence. This matches the
+        // effective semantics of the dictionary assignment in
+        // ExoProcessController.makeEnvironment and avoids silently losing
+        // visible rows after save.
+        var seenKeys = Set<String>()
+        var deduplicatedReversed: [CustomEnvironmentVariable] = []
+        for variable in trimmed.reversed() {
+            if seenKeys.insert(variable.key).inserted {
+                deduplicatedReversed.append(variable)
+            }
+        }
+        let sanitized = Array(deduplicatedReversed.reversed())
+
+        pendingCustomEnvironmentVariables = sanitized
+        controller.customEnvironmentVariables = sanitized
+        restartIfRunning()
+    }
+
+    /// Validates a POSIX-style environment variable name:
+    /// `[A-Za-z_][A-Za-z0-9_]*`. Uses an ASCII-only charset so that
+    /// Unicode letters (e.g. `ñ`, Cyrillic) are rejected in line with what
+    /// the help tooltip advertises. Empty strings are treated as valid
+    /// here so that a freshly added blank row does not immediately look
+    /// broken; the save step filters empty keys out instead.
+    private func isValidEnvironmentVariableName(_ key: String) -> Bool {
+        if key.isEmpty { return true }
+        let headAllowed = CharacterSet(
+            charactersIn: "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz_"
+        )
+        let tailAllowed = headAllowed.union(CharacterSet(charactersIn: "0123456789"))
+        guard let first = key.unicodeScalars.first, headAllowed.contains(first) else {
+            return false
+        }
+        for scalar in key.unicodeScalars.dropFirst() {
+            if !tailAllowed.contains(scalar) { return false }
+        }
+        return true
+    }
+
     private func restartIfRunning() {
         if controller.status == .running || controller.status == .starting {
             controller.restart()

bench/METHODOLOGY.md

@@ -0,0 +1,144 @@
+# Exo-Bench — Methodology
+
+exo bench measures inference throughput and resource consumption of an exo cluster under controlled conditions. It sends prompts to the `/bench/chat/completions` endpoint, collects server-reported timing statistics, and records system-level metrics (power, GPU utilisation, temperature) throughout each run.
+
+The goal is to have accurate, transparent and reproducible numbers to compare speed and scaling across different models and different setups, and to be able to track these results as optimizations and features are added to EXO.
+
+Below is a technical summary of how Exo-Bench works. While the methodology and benchmark may change over time, this document will be kept up to date whenever this happens. If you find an issue with the methodology, or would like a  feature to be added, please open a GitHub issue!
+
+---
+
+## Prompt Construction
+
+Benchmarks need prompts of an exact token length. Unfortunately, we do not have direct access to the model but just the chat completion endpoint. To get around this fact, we create a request that will tokenise to a certain prompt length.
+
+This is achieved by:
+
+1. Tokenising a sample message through the model's `apply_chat_template()` to measure overhead (system tokens, special tokens, chat formatting).
+2. Binary-searching over a repeated atom string (default `"a "`) to find the content length that produces exactly the target number of tokens after template expansion.
+3. Returning both the content string and the verified token count.
+
+The actual token count is recorded in every result row as `pp_tokens`, so downstream analysis can confirm the prompt hit its target.
+
+Chat template formatting means that it may be impossible to attain very small pp benchmarks. e.g. pp=32 may not work. This tradeoff was made because the result of such a small prompt does not seem very interesting or useful for any real-world use cases.
+
+---
+
+## Bench Endpoint
+
+When a request reaches the server via the `/bench/chat/completions` endpoint, three things change compared to a normal chat completion:
+
+- **KV prefix cache is disabled**. Every request starts from a cold cache, ensuring prefill timing is not affected by prior requests.
+- **EOS tokens are banned**. A logits processor suppresses all end-of-sequence tokens, forcing the model to generate exactly `max_tokens` tokens. This guarantees consistent generation length for fair TPS comparison — the model cannot short-circuit a run by stopping early.
+- **No model output parsing**. The bench collection path concatenates raw token text without any model-specific post-processing (thinking tag extraction, structured output handling, etc.). This is to avoid model outputs such as tool parsing or any structural mistakes from breaking the benchmark - we are testing for speed; see Exo-Eval for performance metrics.
+
+---
+
+## Timing
+
+### Prefill TPS
+
+Measured server-side per task.
+
+```
+prefill_tps = num_prompt_tokens / prefill_wall_seconds
+```
+
+### Generation TPS
+
+Measured server-side per task. Each task records wall-clock timestamps as tokens arrive:
+
+- First generated token: timestamp recorded
+- Every subsequent token: timestamp updated
+
+When generation completes:
+
+```
+gen_span = last_token_time - first_token_time
+generation_tps = (completion_tokens - 1) / gen_span
+```
+
+The first token is excluded from the numerator because the rate measures inter-token throughput — the time between the first and last token divided by the number of intervals.
+
+This does mean that tg=1 will not work.
+
+---
+
+## Concurrency
+
+### Single Request
+
+The client records wall-clock `elapsed_s` around the HTTP round-trip (network latency + server prefill + generation + response serialisation). This is a convenience metric for end-to-end latency. The authoritative TPS numbers come from the server-side per-task timing in the `generation_stats` response.
+
+### Concurrent Requests
+
+When `--concurrency N` is set with N > 1, all N requests must hit the server at the same instant. The mechanism:
+
+1. The prompt is built once and shared across all threads.
+2. Each thread gets its own HTTP connection.
+3. A thread barrier blocks all threads until every thread is ready.
+4. The first thread past the barrier records the batch start time and signals the others.
+5. All threads use the same start time as their reference, then fire their HTTP request.
+6. Each thread's `elapsed_s` is measured from the shared start time to its own response completion.
+
+**Batch wall time** is the maximum `elapsed_s` across all N requests — the time until the last request finishes.
+
+### Aggregate TPS
+
+```
+per_req_tps = max(generation_tps across N concurrent requests)
+agg_gen_tps = per_req_tps * concurrency
+```
+
+`max` is used instead of `mean` because all requests run in parallel against the same model. The fastest request's generation rate represents the system's per-stream throughput capacity; multiplying by concurrency gives aggregate throughput.
+
+---
+
+## Warmup
+
+Before measurement begins, `--warmup N` (default: 0) discarded requests are sent using the first pp/tg pair. Warmup results are not included in the output.
+
+---
+
+## System Metrics
+
+A background thread polls each node at 1 Hz, collecting:
+
+- GPU utilisation (%)
+- Temperature (C)
+- System power draw (W)
+- CPU cluster usage (performance and efficiency cores)
+
+**Energy** is computed via trapezoidal integration of the power samples over each inference window (the wall-clock span of each benchmark request or concurrent batch). Average power is `total_joules / total_inference_seconds`.
+
+---
+
+## Output Format
+
+Results are written as JSON with three top-level keys:
+
+- **`runs`**: Array of per-request result objects, each containing:
+  - `elapsed_s`, `output_text_preview` (first 200 chars)
+  - `stats`: `{ prompt_tps, generation_tps, prompt_tokens, generation_tokens, peak_memory_usage }`
+  - Placement metadata: `model_id`, `placement_sharding`, `placement_instance_meta`, `placement_nodes`
+  - Run metadata: `pp_tokens`, `tg`, `repeat_index`, `concurrency`, `concurrent_index`
+  - `download_duration_s` (if model was freshly downloaded)
+- **`cluster`**: Cluster state snapshot at time of benchmarking.
+- **`system_metrics`**: Per-node time-series samples (GPU, power, temperature).
+
+---
+
+## Reproducing Results
+
+```bash
+cd bench && uv run python exo_bench.py \
+  --model "mlx-community/Qwen3.5-27B-4bit" \
+  --instance-meta jaccl \
+  --sharding tensor \
+  --min-nodes 2 --max-nodes 2 \
+  --pp 512 4096 --tg 128 \
+  --repeat 3 \
+  --warmup 1
+```
+
+Run --help for all the available flags.

bench/bench.toml

@@ -2,6 +2,4 @@
 #
 # Lists the suite files to include. Each file defines benchmarks
 # with shared constraints, topology, and default args.
-include = [
-    "single-m3-ultra.toml",
-]
+include = ["single-m3-ultra.toml"]

bench/eval_configs/models.toml

@@ -0,0 +1,292 @@
+# Model evaluation configurations for exo_eval.
+#
+# Each [[model]] entry uses `patterns` — a list of substrings matched
+# against the model_id. First matching entry wins.
+#
+# Required fields:
+#   name, patterns, reasoning
+#
+# Optional per-model overrides (CLI flags take priority over these):
+#   temperature, top_p, max_tokens, reasoning_effort
+#
+# Fallback defaults (when no per-model config):
+#   reasoning:     temperature=1.0, max_tokens=131072, reasoning_effort="high"
+#   non-reasoning: temperature=0.0, max_tokens=16384
+#
+# All per-model values below are sourced from official model cards,
+# generation_config.json files, and vendor documentation.
+
+# ─── Qwen3.5 (Feb 2026) ─────────────────────────────────────────────
+# Source: HuggingFace model cards (Qwen/Qwen3.5-*)
+# 35B-A3B thinking general: temp=1.0, top_p=0.95, top_k=20
+# 397B thinking: temp=0.6, top_p=0.95, top_k=20
+# Non-thinking: temp=0.7, top_p=0.8, top_k=20
+# max_tokens: 32768 general, 81920 for complex math/code
+
+[[model]]
+name = "Qwen3.5 2B"
+patterns = ["Qwen3.5-2B"]
+reasoning = true
+temperature = 0.6
+top_p = 0.95
+max_tokens = 81920
+
+[[model]]
+name = "Qwen3.5 9B"
+patterns = ["Qwen3.5-9B"]
+reasoning = true
+temperature = 0.6
+top_p = 0.95
+max_tokens = 81920
+
+[[model]]
+name = "Qwen3.5 27B"
+patterns = ["Qwen3.5-27B"]
+reasoning = true
+temperature = 0.6
+top_p = 0.95
+max_tokens = 81920
+
+[[model]]
+name = "Qwen3.5 35B A3B"
+patterns = ["Qwen3.5-35B-A3B"]
+reasoning = true
+temperature = 1.0
+top_p = 0.95
+max_tokens = 81920
+
+[[model]]
+name = "Qwen3.5 122B A10B"
+patterns = ["Qwen3.5-122B-A10B"]
+reasoning = true
+temperature = 0.6
+top_p = 0.95
+max_tokens = 81920
+
+[[model]]
+name = "Qwen3.5 397B A17B"
+patterns = ["Qwen3.5-397B-A17B"]
+reasoning = true
+temperature = 0.6
+top_p = 0.95
+max_tokens = 81920
+
+# ─── Qwen3 (Apr 2025) ───────────────────────────────────────────────
+# Source: HuggingFace model cards (Qwen/Qwen3-*)
+# Thinking: temp=0.6, top_p=0.95, top_k=20
+# Non-thinking: temp=0.7, top_p=0.8, top_k=20
+# max_tokens: 32768 general, 38912 for complex math/code
+
+[[model]]
+name = "Qwen3 0.6B"
+patterns = ["Qwen3-0.6B"]
+reasoning = true
+temperature = 0.6
+top_p = 0.95
+max_tokens = 38912
+
+[[model]]
+name = "Qwen3 30B A3B"
+patterns = ["Qwen3-30B-A3B"]
+reasoning = true
+temperature = 0.6
+top_p = 0.95
+max_tokens = 38912
+
+[[model]]
+name = "Qwen3 235B A22B"
+patterns = ["Qwen3-235B-A22B"]
+reasoning = true
+temperature = 0.6
+top_p = 0.95
+max_tokens = 38912
+
+[[model]]
+name = "Qwen3 Next 80B Thinking"
+patterns = ["Qwen3-Next-80B-A3B-Thinking"]
+reasoning = true
+temperature = 0.6
+top_p = 0.95
+max_tokens = 38912
+
+[[model]]
+name = "Qwen3 Next 80B Instruct"
+patterns = ["Qwen3-Next-80B-A3B-Instruct"]
+reasoning = false
+temperature = 0.7
+top_p = 0.8
+max_tokens = 16384
+
+[[model]]
+name = "Qwen3 Coder 480B"
+patterns = ["Qwen3-Coder-480B"]
+reasoning = false
+temperature = 0.7
+top_p = 0.8
+max_tokens = 16384
+
+[[model]]
+name = "Qwen3 Coder Next"
+patterns = ["Qwen3-Coder-Next"]
+reasoning = false
+temperature = 0.7
+top_p = 0.8
+max_tokens = 16384
+
+# ─── GPT-OSS (OpenAI) ───────────────────────────────────────────────
+# Source: OpenAI GitHub README + HuggingFace discussion #21
+# temp=1.0, top_p=1.0, NO top_k, NO repetition_penalty
+# reasoning_effort supported: low/medium/high
+
+[[model]]
+name = "GPT-OSS 20B"
+patterns = ["gpt-oss-20b"]
+reasoning = true
+temperature = 1.0
+top_p = 1.0
+
+[[model]]
+name = "GPT-OSS 120B"
+patterns = ["gpt-oss-120b"]
+reasoning = true
+temperature = 1.0
+top_p = 1.0
+
+# ─── DeepSeek ────────────────────────────────────────────────────────
+# Source: https://api-docs.deepseek.com/quick_start/parameter_settings
+# Coding/Math: temp=0.0, General: temp=1.3, Creative: temp=1.5
+# NOTE: DeepSeek API applies nonlinear temp mapping. These are API values.
+# When running model directly: API temp 1.0 = model temp 0.3
+# We use temp=0.0 for eval (coding/math focus).
+
+[[model]]
+name = "DeepSeek V3.1"
+patterns = ["DeepSeek-V3.1"]
+reasoning = true
+temperature = 0.0
+
+# ─── GLM (ZhipuAI / THUDM) ──────────────────────────────────────────
+# Source: HuggingFace model cards + generation_config.json + docs.z.ai
+# GLM 4.5+: temp=1.0, top_p=0.95
+# Reasoning tasks: 131072 max_tokens; coding/SWE tasks: temp=0.7
+
+[[model]]
+name = "GLM-5"
+patterns = ["GLM-5"]
+reasoning = true
+temperature = 1.0
+top_p = 0.95
+max_tokens = 131072
+
+[[model]]
+name = "GLM 4.5 Air"
+patterns = ["GLM-4.5-Air"]
+reasoning = true
+temperature = 1.0
+top_p = 0.95
+
+[[model]]
+name = "GLM 4.7"
+patterns = ["GLM-4.7-"]
+reasoning = true
+temperature = 1.0
+top_p = 0.95
+max_tokens = 131072
+# Note: matches both GLM-4.7 and GLM-4.7-Flash
+
+# ─── Kimi (Moonshot AI) ─────────────────────────────────────────────
+# Source: HuggingFace model cards (moonshotai/Kimi-K2-*)
+# K2-Instruct: temp=0.6
+# K2-Thinking: temp=1.0, max_length=262144
+# K2.5: thinking temp=1.0, top_p=0.95; instant temp=0.6, top_p=0.95
+
+[[model]]
+name = "Kimi K2 Thinking"
+patterns = ["Kimi-K2-Thinking"]
+reasoning = true
+temperature = 1.0
+max_tokens = 131072
+
+[[model]]
+name = "Kimi K2.5"
+patterns = ["Kimi-K2.5"]
+reasoning = true
+temperature = 1.0
+top_p = 0.95
+max_tokens = 131072
+
+[[model]]
+name = "Kimi K2 Instruct"
+patterns = ["Kimi-K2-Instruct"]
+reasoning = false
+temperature = 0.6
+
+# ─── MiniMax ─────────────────────────────────────────────────────────
+# Source: HuggingFace model cards + generation_config.json
+# All models: temp=1.0, top_p=0.95, top_k=40
+
+[[model]]
+name = "MiniMax M2.5"
+patterns = ["MiniMax-M2.5"]
+reasoning = true
+temperature = 1.0
+top_p = 0.95
+
+[[model]]
+name = "MiniMax M2.1"
+patterns = ["MiniMax-M2.1"]
+reasoning = true
+temperature = 1.0
+top_p = 0.95
+
+# ─── Step (StepFun) ─────────────────────────────────────────────────
+# Source: HuggingFace model card (stepfun-ai/Step-3.5-Flash)
+# Reasoning: temp=1.0, top_p=0.95
+# General chat: temp=0.6, top_p=0.95
+# We use reasoning settings for eval.
+
+[[model]]
+name = "Step 3.5 Flash"
+patterns = ["Step-3.5-Flash"]
+reasoning = true
+temperature = 1.0
+top_p = 0.95
+
+# ─── Llama (Meta) ───────────────────────────────────────────────────
+# Source: generation_config.json + meta-llama/llama-models generation.py
+# All variants: temp=0.6, top_p=0.9
+
+[[model]]
+name = "Llama 3.2 1B"
+patterns = ["Llama-3.2-1B"]
+reasoning = false
+temperature = 0.6
+top_p = 0.9
+
+[[model]]
+name = "Llama 3.2 3B"
+patterns = ["Llama-3.2-3B"]
+reasoning = false
+temperature = 0.6
+top_p = 0.9
+
+[[model]]
+name = "Llama 3.1 8B"
+patterns = ["Llama-3.1-8B", "Meta-Llama-3.1-8B"]
+reasoning = false
+temperature = 0.6
+top_p = 0.9
+
+[[model]]
+name = "Llama 3.1 70B"
+patterns = ["Llama-3.1-70B", "Meta-Llama-3.1-70B"]
+reasoning = false
+temperature = 0.6
+top_p = 0.9
+
+[[model]]
+name = "Llama 3.3 70B"
+patterns = ["Llama-3.3-70B", "llama-3.3-70b"]
+reasoning = false
+temperature = 0.6
+top_p = 0.9

bench/eval_tool_calls.py

@@ -17,11 +17,14 @@ from harness import (
     ExoClient,
     ExoHttpError,
     add_common_instance_args,
+    capture_cluster_snapshot,
+    ensure_cuda_available,
     instance_id_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
     run_planning_phase,
     settle_and_fetch_placements,
+    validate_vllm_args,
     wait_for_instance_gone,
     wait_for_instance_ready,
 )
@@ -933,6 +936,7 @@ Examples:
         help="Write JSON results to stdout instead of file",
     )
     args = parser.parse_args()
+    validate_vllm_args(args)
 
     all_scenarios = load_scenarios(SCENARIOS_PATH)
     if args.scenarios:
@@ -952,6 +956,8 @@ Examples:
 
     log = sys.stderr if args.stdout else sys.stdout
     exo = ExoClient(args.host, args.port, timeout_s=args.timeout)
+    if args.ensure_cuda:
+        ensure_cuda_available(exo)
     _short_id, full_model_id = resolve_model_short_id(exo, args.model)
 
     selected = settle_and_fetch_placements(
@@ -1006,6 +1012,7 @@ Examples:
         sys.exit(1)
 
     time.sleep(1)
+    cluster_snapshot = capture_cluster_snapshot(exo)
     all_results: list[ScenarioResult] = []
 
     try:
@@ -1084,16 +1091,19 @@ Examples:
                 print(f"  - {r.name} [{r.api}/{r.phase}]: {r.error}", file=log)
 
     json_results = [result_to_dict(r) for r in all_results]
+    output: dict[str, Any] = {"results": json_results}
+    if cluster_snapshot:
+        output["cluster"] = cluster_snapshot
 
     if args.stdout:
-        print(json.dumps(json_results, indent=2))
+        print(json.dumps(output, indent=2))
     else:
         json_path = args.json_out
         parent = os.path.dirname(json_path)
         if parent:
             os.makedirs(parent, exist_ok=True)
         with open(json_path, "w") as f:
-            json.dump(json_results, f, indent=2)
+            json.dump(output, f, indent=2)
             f.write("\n")
         print(f"\nJSON results written to {json_path}", file=log)
 

bench/exo_bench.py

@@ -22,8 +22,10 @@ import contextlib
 import itertools
 import json
 import sys
+import threading
 import time
 from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
 from statistics import mean
 from typing import Any
@@ -32,11 +34,15 @@ from harness import (
     ExoClient,
     ExoHttpError,
     add_common_instance_args,
+    capture_cluster_snapshot,
+    ensure_cuda_available,
     instance_id_from_instance,
+    node_ids_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
     run_planning_phase,
     settle_and_fetch_placements,
+    validate_vllm_args,
     wait_for_instance_gone,
     wait_for_instance_ready,
 )
@@ -130,6 +136,91 @@ def format_peak_memory(b: float) -> str:
     raise ValueError("You're using petabytes of memory. Something went wrong...")
 
 
+_SAMPLER_METRICS = ("gpuUsage", "temp", "sysPower", "pcpuUsage", "ecpuUsage")
+
+
+class SystemMetricsSampler:
+    def __init__(self, client: ExoClient, node_ids: list[str], interval_s: float = 1.0):
+        self._client = client
+        self._node_ids = node_ids
+        self._interval_s = interval_s
+        self._samples: dict[str, list[tuple[float, dict[str, float]]]] = {
+            nid: [] for nid in node_ids
+        }
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+
+    def start(self) -> None:
+        self._stop.clear()
+        self._thread = threading.Thread(target=self._poll_loop, daemon=True)
+        self._thread.start()
+
+    def stop(self) -> None:
+        self._stop.set()
+        if self._thread:
+            self._thread.join(timeout=5)
+
+    def _poll_loop(self) -> None:
+        while not self._stop.is_set():
+            t = time.monotonic()
+            for nid in self._node_ids:
+                try:
+                    data = self._client.get_node_system(nid)
+                    if data:
+                        self._samples[nid].append(
+                            (t, {k: data.get(k, 0.0) for k in _SAMPLER_METRICS})
+                        )
+                except Exception:
+                    pass
+            self._stop.wait(self._interval_s)
+
+    def energy_between(self, t0: float, t1: float) -> float:
+        total_joules = 0.0
+        for _nid, samples in self._samples.items():
+            window = [(t, s["sysPower"]) for t, s in samples if t0 <= t <= t1]
+            if len(window) >= 2:
+                for i in range(1, len(window)):
+                    dt = window[i][0] - window[i - 1][0]
+                    avg_power = (window[i][1] + window[i - 1][1]) / 2
+                    total_joules += avg_power * dt
+            elif len(window) == 1:
+                total_joules += window[0][1] * (t1 - t0)
+        return total_joules
+
+    def summarize(self) -> dict[str, dict[str, dict[str, float]]]:
+        result: dict[str, dict[str, dict[str, float]]] = {}
+        for nid, samples in self._samples.items():
+            if not samples:
+                continue
+            metrics: dict[str, dict[str, float]] = {}
+            for key in _SAMPLER_METRICS:
+                values = [s[key] for t, s in samples]
+                metrics[key] = {
+                    "min": round(min(values), 2),
+                    "max": round(max(values), 2),
+                    "mean": round(mean(values), 2),
+                    "samples": len(values),
+                }
+            result[nid] = metrics
+        return result
+
+    def print_summary(self, placement_label: str) -> None:
+        summary = self.summarize()
+        if not summary:
+            return
+        logger.info(f"--- System Metrics ({placement_label}) ---")
+        for nid, metrics in summary.items():
+            gpu = metrics.get("gpuUsage", {})
+            temp = metrics.get("temp", {})
+            power = metrics.get("sysPower", {})
+            logger.info(
+                f"  {nid}:  "
+                f"GPU {gpu.get('mean', 0) * 100:.0f}% avg ({gpu.get('min', 0) * 100:.0f}–{gpu.get('max', 0) * 100:.0f}%)  |  "
+                f"{temp.get('mean', 0):.1f}°C avg  |  "
+                f"{power.get('mean', 0):.1f}W avg"
+            )
+
+
 def parse_int_list(values: list[str]) -> list[int]:
     items: list[int] = []
     for v in values:
@@ -149,6 +240,7 @@ def run_one_completion(
         "messages": [{"role": "user", "content": content}],
         "stream": False,
         "max_tokens": tg,
+        "logprobs": False,
     }
 
     t0 = time.perf_counter()
@@ -252,6 +344,12 @@ def main() -> int:
     ap.add_argument(
         "--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair."
     )
+    ap.add_argument(
+        "--concurrency",
+        nargs="+",
+        default=["1"],
+        help="Concurrency levels (ints). Accepts commas. E.g. --concurrency 1,2,4,8. Default 1.",
+    )
     ap.add_argument(
         "--warmup",
         type=int,
@@ -272,7 +370,19 @@ def main() -> int:
         action="store_true",
         help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
     )
+    ap.add_argument(
+        "--no-system-metrics",
+        action="store_true",
+        help="Disable GPU utilization, temperature, and power collection during inference.",
+    )
+    ap.add_argument(
+        "--metrics-interval",
+        type=float,
+        default=1.0,
+        help="System metrics polling interval in seconds (default: 1.0).",
+    )
     args = ap.parse_args()
+    validate_vllm_args(args)
 
     pp_list = parse_int_list(args.pp)
     tg_list = parse_int_list(args.tg)
@@ -282,6 +392,10 @@ def main() -> int:
     if args.repeat <= 0:
         logger.error("--repeat must be >= 1")
         return 2
+    concurrency_list = parse_int_list(args.concurrency)
+    if not concurrency_list or any(c <= 0 for c in concurrency_list):
+        logger.error("--concurrency values must be >= 1")
+        return 2
 
     # Log pairing mode
     use_combinations = args.all_combinations or len(pp_list) != len(tg_list)
@@ -293,7 +407,11 @@ 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)
-    short_id, full_model_id = resolve_model_short_id(client, args.model)
+    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
+    )
 
     tokenizer = load_tokenizer_for_bench(full_model_id)
     if tokenizer is None:
@@ -351,7 +469,9 @@ def main() -> int:
     else:
         logger.info("Download: model already cached")
 
+    cluster_snapshot = capture_cluster_snapshot(client)
     all_rows: list[dict[str, Any]] = []
+    all_system_metrics: dict[str, dict[str, dict[str, float]]] = {}
 
     for preview in selected:
         instance = preview["instance"]
@@ -377,6 +497,16 @@ def main() -> int:
 
         time.sleep(1)
 
+        sampler: SystemMetricsSampler | None = None
+        if not args.no_system_metrics:
+            nids = node_ids_from_instance(instance)
+            sampler = SystemMetricsSampler(
+                ExoClient(args.host, args.port, timeout_s=30),
+                nids,
+                interval_s=args.metrics_interval,
+            )
+            sampler.start()
+
         try:
             for i in range(args.warmup):
                 run_one_completion(
@@ -392,53 +522,195 @@ def main() -> int:
                 pp_tg_pairs = list(zip(pp_list, tg_list, strict=True))
 
             for pp, tg in pp_tg_pairs:
-                runs: list[dict[str, Any]] = []
-                for r in range(args.repeat):
-                    time.sleep(3)
-                    try:
-                        row, actual_pp_tokens = run_one_completion(
-                            client, full_model_id, pp, tg, prompt_sizer
+                for concurrency in concurrency_list:
+                    logger.info(f"--- pp={pp} tg={tg} concurrency={concurrency} ---")
+                    runs: list[dict[str, Any]] = []
+                    inference_windows: list[tuple[float, float]] = []
+                    for r in range(args.repeat):
+                        time.sleep(3)
+
+                        if concurrency <= 1:
+                            # Sequential: single request
+                            try:
+                                inf_t0 = time.monotonic()
+                                row, actual_pp_tokens = run_one_completion(
+                                    client, full_model_id, pp, tg, prompt_sizer
+                                )
+                                inference_windows.append((inf_t0, time.monotonic()))
+                            except Exception as e:
+                                logger.error(e)
+                                continue
+                            row.update(
+                                {
+                                    "model_short_id": short_id,
+                                    "model_id": full_model_id,
+                                    "placement_sharding": sharding,
+                                    "placement_instance_meta": instance_meta,
+                                    "placement_nodes": n_nodes,
+                                    "instance_id": instance_id,
+                                    "pp_tokens": actual_pp_tokens,
+                                    "tg": tg,
+                                    "repeat_index": r,
+                                    "concurrency": 1,
+                                    **(
+                                        {"download_duration_s": download_duration_s}
+                                        if download_duration_s is not None
+                                        else {}
+                                    ),
+                                }
+                            )
+                            runs.append(row)
+                            all_rows.append(row)
+                        else:
+                            # Concurrent: fire N requests in parallel
+                            # Pre-build prompt once, barrier ensures simultaneous dispatch
+                            content, actual_pp = prompt_sizer.build(pp)
+                            pre_built_payload: dict[str, Any] = {
+                                "model": full_model_id,
+                                "messages": [{"role": "user", "content": content}],
+                                "stream": False,
+                                "max_tokens": tg,
+                                "logprobs": False,
+                            }
+                            barrier = threading.Barrier(concurrency)
+                            batch_start = threading.Event()
+                            batch_t0: float = 0.0
+                            batch_results: list[tuple[dict[str, Any], int]] = []
+                            batch_errors = 0
+
+                            def _run_concurrent(
+                                idx: int,
+                                _barrier: threading.Barrier = barrier,
+                                _batch_start: threading.Event = batch_start,
+                                _payload: dict[str, Any] = pre_built_payload,
+                                _actual_pp: int = actual_pp,
+                            ) -> tuple[dict[str, Any], int]:
+                                nonlocal batch_t0
+                                c = ExoClient(
+                                    args.host, args.port, timeout_s=args.timeout
+                                )
+                                if _barrier.wait() == 0:
+                                    batch_t0 = time.perf_counter()
+                                    _batch_start.set()
+                                else:
+                                    _batch_start.wait()
+                                t0 = batch_t0
+                                out = c.post_bench_chat_completions(_payload)
+                                elapsed = time.perf_counter() - t0
+                                stats = out.get("generation_stats")
+                                choices = out.get("choices") or [{}]
+                                message = (
+                                    choices[0].get("message", {}) if choices else {}
+                                )
+                                text = message.get("content") or ""
+                                return {
+                                    "elapsed_s": elapsed,
+                                    "output_text_preview": text[:200],
+                                    "stats": stats,
+                                }, _actual_pp
+
+                            inf_t0 = time.monotonic()
+                            with ThreadPoolExecutor(max_workers=concurrency) as pool:
+                                futures = {
+                                    pool.submit(_run_concurrent, i): i
+                                    for i in range(concurrency)
+                                }
+                                for fut in as_completed(futures):
+                                    try:
+                                        batch_results.append(fut.result())
+                                    except Exception as e:
+                                        logger.error(f"Concurrent request failed: {e}")
+                                        batch_errors += 1
+                            batch_wall_s = (
+                                max(x["elapsed_s"] for x, _ in batch_results)
+                                if batch_results
+                                else time.perf_counter() - batch_t0
+                            )
+                            inference_windows.append((inf_t0, time.monotonic()))
+
+                            for idx, (row, actual_pp_tokens) in enumerate(
+                                batch_results
+                            ):
+                                row.update(
+                                    {
+                                        "model_short_id": short_id,
+                                        "model_id": full_model_id,
+                                        "placement_sharding": sharding,
+                                        "placement_instance_meta": instance_meta,
+                                        "placement_nodes": n_nodes,
+                                        "instance_id": instance_id,
+                                        "pp_tokens": actual_pp_tokens,
+                                        "tg": tg,
+                                        "repeat_index": r,
+                                        "concurrency": concurrency,
+                                        "concurrent_index": idx,
+                                        **(
+                                            {"download_duration_s": download_duration_s}
+                                            if download_duration_s is not None
+                                            else {}
+                                        ),
+                                    }
+                                )
+                                runs.append(row)
+                                all_rows.append(row)
+
+                            if batch_results:
+                                valid_gen_tps = [
+                                    x["stats"]["generation_tps"]
+                                    for x, _ in batch_results
+                                    if x["stats"]["generation_tps"] > 0
+                                ]
+                                per_req_tps = (
+                                    max(valid_gen_tps) if valid_gen_tps else 0.0
+                                )
+                                agg_gen_tps = per_req_tps * concurrency
+                                logger.info(
+                                    f"[concurrent {concurrency}x]  "
+                                    f"agg_gen_tps={agg_gen_tps:.2f}  "
+                                    f"per_req_tps={per_req_tps:.2f}  "
+                                    f"wall_s={batch_wall_s:.2f}  "
+                                    f"errors={batch_errors}"
+                                )
+
+                    if runs:
+                        prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
+                        valid_gen = [
+                            x["stats"]["generation_tps"]
+                            for x in runs
+                            if x["stats"]["generation_tps"] > 0
+                        ]
+                        per_req_tps = max(valid_gen) if valid_gen else 0.0
+                        gen_tps = per_req_tps * concurrency
+                        ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
+                        gtok = mean(x["stats"]["generation_tokens"] for x in runs)
+                        peak = mean(
+                            x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
                         )
-                    except Exception as e:
-                        logger.error(e)
-                        continue
-                    row.update(
-                        {
-                            "model_short_id": short_id,
-                            "model_id": full_model_id,
-                            "placement_sharding": sharding,
-                            "placement_instance_meta": instance_meta,
-                            "placement_nodes": n_nodes,
-                            "instance_id": instance_id,
-                            "pp_tokens": actual_pp_tokens,
-                            "tg": tg,
-                            "repeat_index": r,
-                            **(
-                                {"download_duration_s": download_duration_s}
-                                if download_duration_s is not None
-                                else {}
-                            ),
-                        }
-                    )
-                    runs.append(row)
-                    all_rows.append(row)
 
-                if runs:
-                    prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
-                    gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
-                    ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
-                    gtok = mean(x["stats"]["generation_tokens"] for x in runs)
-                    peak = mean(
-                        x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
-                    )
-
-                    logger.info(
-                        f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
-                        f"prompt_tokens={ptok} gen_tokens={gtok}    "
-                        f"peak_memory={format_peak_memory(peak)}\n"
-                    )
-                time.sleep(2)
+                        summary = (
+                            f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
+                            f"prompt_tokens={ptok} gen_tokens={gtok}    "
+                            f"peak_memory={format_peak_memory(peak)}"
+                        )
+                        if sampler and inference_windows:
+                            joules = sum(
+                                sampler.energy_between(t0, t1)
+                                for t0, t1 in inference_windows
+                            )
+                            inf_seconds = sum(t1 - t0 for t0, t1 in inference_windows)
+                            avg_watts = joules / inf_seconds if inf_seconds > 0 else 0
+                            summary += f"    energy={joules:.1f}J ({avg_watts:.1f}W avg over {inf_seconds:.1f}s inference)"
+                        logger.info(f"{summary}\n")
+                    time.sleep(2)
         finally:
+            if sampler:
+                sampler.stop()
+                placement_label = f"{sharding}/{instance_meta}/{n_nodes} nodes"
+                sampler.print_summary(placement_label)
+                placement_metrics = sampler.summarize()
+                if placement_metrics:
+                    all_system_metrics.update(placement_metrics)
+
             try:
                 client.request_json("DELETE", f"/instance/{instance_id}")
             except ExoHttpError as e:
@@ -449,11 +721,17 @@ def main() -> int:
 
             time.sleep(5)
 
+    output: dict[str, Any] = {"runs": all_rows}
+    if cluster_snapshot:
+        output["cluster"] = cluster_snapshot
+    if all_system_metrics:
+        output["system_metrics"] = all_system_metrics
+
     if args.stdout:
-        json.dump(all_rows, sys.stdout, indent=2, ensure_ascii=False)
+        json.dump(output, sys.stdout, indent=2, ensure_ascii=False)
     elif args.json_out:
         with open(args.json_out, "w", encoding="utf-8") as f:
-            json.dump(all_rows, f, indent=2, ensure_ascii=False)
+            json.dump(output, f, indent=2, ensure_ascii=False)
         logger.debug(f"\nWrote results JSON: {args.json_out}")
 
     return 0

bench/harness.py

@@ -69,6 +69,39 @@ class ExoClient:
     def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
         return self.request_json("POST", "/bench/chat/completions", body=payload)
 
+    def post_bench_disaggregated(self, payload: dict[str, Any]) -> dict[str, Any]:
+        payload["disaggregated"] = True
+        return self.request_json("POST", "/bench/chat/completions", body=payload)
+
+    def get_state_path(self, path: str) -> Any:
+        try:
+            return self.request_json("GET", f"/state/{path}")
+        except ExoHttpError as e:
+            if e.status == 404:
+                return None
+            raise
+
+    def get_instance(self, instance_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"instances/{instance_id}")
+
+    def get_runner(self, runner_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"runners/{runner_id}")
+
+    def get_node_downloads(self, node_id: str) -> list[dict[str, Any]] | None:
+        return self.get_state_path(f"downloads/{node_id}")
+
+    def get_node_disk(self, node_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"nodeDisk/{node_id}")
+
+    def get_node_system(self, node_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"nodeSystem/{node_id}")
+
+    def get_node_identities(self) -> dict[str, Any] | None:
+        return self.get_state_path("nodeIdentities")
+
+    def get_topology(self) -> dict[str, Any] | None:
+        return self.get_state_path("topology")
+
 
 def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]:
     if len(instance) != 1:
@@ -97,6 +130,11 @@ def runner_ids_from_instance(instance: dict[str, Any]) -> list[str]:
     return list(runner_to_shard.keys())
 
 
+def node_ids_from_instance(instance: dict[str, Any]) -> list[str]:
+    inner = unwrap_instance(instance)
+    return list(inner["shardAssignments"]["nodeToRunner"].keys())
+
+
 def runner_ready(runner: dict[str, Any]) -> bool:
     return "RunnerReady" in runner
 
@@ -116,13 +154,12 @@ def wait_for_instance_ready(
 ) -> None:
     start_time = time.time()
     instance_existed = False
+    last_loaded: dict[str, int] = {}
     while time.time() - start_time < timeout:
-        state = client.request_json("GET", "/state")
-        instances = state.get("instances", {})
+        instance = client.get_instance(instance_id)
 
-        if instance_id not in instances:
+        if instance is None:
             if instance_existed:
-                # Instance was deleted after being created - likely due to runner failure
                 raise RuntimeError(
                     f"Instance {instance_id} was deleted (runner may have failed)"
                 )
@@ -130,18 +167,25 @@ def wait_for_instance_ready(
             continue
 
         instance_existed = True
-        instance = instances[instance_id]
-        runner_ids = runner_ids_from_instance(instance)
-        runners = state.get("runners", {})
+        rids = runner_ids_from_instance(instance)
 
-        # Check for failed runners first
-        for rid in runner_ids:
-            runner = runners.get(rid, {})
+        all_ready = True
+        for rid in rids:
+            runner = client.get_runner(rid) or {}
             if runner_failed(runner):
                 error_msg = get_runner_failed_message(runner) or "Unknown error"
                 raise RuntimeError(f"Runner {rid} failed: {error_msg}")
+            if "RunnerLoading" in runner:
+                loading = runner["RunnerLoading"]
+                loaded = loading.get("layersLoaded", 0)
+                total = loading.get("totalLayers", 0)
+                if total > 0 and last_loaded.get(rid) != loaded:
+                    last_loaded[rid] = loaded
+                    logger.debug(f"Runner {rid}: loading layers {loaded}/{total}")
+            if not runner_ready(runner):
+                all_ready = False
 
-        if all(runner_ready(runners.get(rid, {})) for rid in runner_ids):
+        if all_ready:
             return
 
         time.sleep(0.1)
@@ -165,7 +209,26 @@ def wait_for_instance_gone(
     raise TimeoutError(f"Instance {instance_id} did not get deleted within {timeout=}")
 
 
-def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]:
+def capture_cluster_snapshot(client: ExoClient) -> dict[str, Any]:
+    snapshot: dict[str, Any] = {}
+    identities = client.get_node_identities()
+    if identities:
+        snapshot["nodeIdentities"] = identities
+    topology = client.get_topology()
+    if topology:
+        snapshot["topology"] = topology
+    node_memory = client.get_state_path("nodeMemory")
+    if node_memory:
+        snapshot["nodeMemory"] = node_memory
+    node_system = client.get_state_path("nodeSystem")
+    if node_system:
+        snapshot["nodeSystem"] = node_system
+    return snapshot
+
+
+def resolve_model_short_id(
+    client: ExoClient, model_arg: str, *, force_download: bool = False
+) -> tuple[str, str]:
     models = client.request_json("GET", "/models") or {}
     data = models.get("data") or []
 
@@ -181,9 +244,44 @@ def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]
             full_id = str(m["hugging_face_id"])
             return short_id, full_id
 
+    if force_download and "/" in model_arg:
+        logger.info(f"Model not in /models, adding from HuggingFace: {model_arg}")
+        result = client.request_json(
+            "POST", "/models/add", body={"model_id": model_arg}
+        )
+        if result:
+            short_id = str(result.get("name") or model_arg.rsplit("/", 1)[-1])
+            full_id = str(result.get("hugging_face_id") or model_arg)
+            return short_id, full_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":
@@ -314,16 +412,11 @@ def run_planning_phase(
     node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys())
     runner_to_shard = inner["shardAssignments"]["runnerToShard"]
 
-    state = client.request_json("GET", "/state")
-    downloads = state.get("downloads", {})
-    node_disk = state.get("nodeDisk", {})
-
     needs_download = False
 
     for node_id in node_ids:
-        node_downloads = downloads.get(node_id, [])
+        node_downloads = client.get_node_downloads(node_id) or []
 
-        # Check if model already downloaded on this node
         already_downloaded = any(
             "DownloadCompleted" in p
             and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
@@ -337,8 +430,7 @@ def run_planning_phase(
 
         needs_download = True
 
-        # Wait for disk info if settle_deadline is set
-        disk_info = node_disk.get(node_id, {})
+        disk_info = client.get_node_disk(node_id) or {}
         backoff = _SETTLE_INITIAL_BACKOFF_S
         while not disk_info and settle_deadline and time.monotonic() < settle_deadline:
             remaining = settle_deadline - time.monotonic()
@@ -347,9 +439,7 @@ def run_planning_phase(
             )
             time.sleep(min(backoff, remaining))
             backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
-            state = client.request_json("GET", "/state")
-            node_disk = state.get("nodeDisk", {})
-            disk_info = node_disk.get(node_id, {})
+            disk_info = client.get_node_disk(node_id) or {}
 
         if not disk_info:
             logger.warning(f"No disk info for {node_id}, skipping space check")
@@ -365,7 +455,6 @@ def run_planning_phase(
                 f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
             )
 
-        # Delete from smallest to largest (skip read-only models from EXO_MODELS_PATH)
         completed = [
             (
                 unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
@@ -405,21 +494,20 @@ def run_planning_phase(
     # Wait for downloads
     start = time.time()
     while time.time() - start < timeout:
-        state = client.request_json("GET", "/state")
-        downloads = state.get("downloads", {})
         all_done = True
         for node_id in node_ids:
+            node_downloads = client.get_node_downloads(node_id) or []
             done = any(
                 "DownloadCompleted" in p
                 and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[
                     "modelCard"
                 ]["modelId"]
                 == full_model_id
-                for p in downloads.get(node_id, [])
+                for p in node_downloads
             )
             failed = [
                 p["DownloadFailed"]["errorMessage"]
-                for p in downloads.get(node_id, [])
+                for p in node_downloads
                 if "DownloadFailed" in p
                 and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][
                     "modelId"
@@ -430,6 +518,27 @@ def run_planning_phase(
                 raise RuntimeError(f"Download failed on {node_id}: {failed[0]}")
             if not done:
                 all_done = False
+                ongoing = [
+                    p
+                    for p in node_downloads
+                    if "DownloadOngoing" in p
+                    and unwrap_instance(p["DownloadOngoing"]["shardMetadata"])[
+                        "modelCard"
+                    ]["modelId"]
+                    == full_model_id
+                ]
+                if ongoing:
+                    prog = ongoing[0]["DownloadOngoing"]["downloadProgress"]
+                    speed_mb = prog.get("speed", 0) / (1024 * 1024)
+                    eta_s = prog.get("etaMs", 0) / 1000
+                    dl_bytes = prog.get("downloaded", {}).get("inBytes", 0)
+                    total_bytes = prog.get("total", {}).get("inBytes", 0)
+                    pct = (dl_bytes / total_bytes * 100) if total_bytes else 0
+                    logger.info(
+                        f"Downloading on {node_id}: {pct:.1f}% @ {speed_mb:.1f} MB/s, "
+                        f"ETA {eta_s:.0f}s "
+                        f"({prog.get('completedFiles', 0)}/{prog.get('totalFiles', 0)} files)"
+                    )
         if all_done:
             if download_t0 is not None:
                 return time.perf_counter() - download_t0
@@ -445,6 +554,11 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
         "--port", type=int, default=int(os.environ.get("EXO_PORT", "52415"))
     )
     ap.add_argument("--model", required=True, help="Model short id or huggingface id")
+    ap.add_argument(
+        "--force-download",
+        action="store_true",
+        help="If model not in /models, add it from HuggingFace via exo and download.",
+    )
     ap.add_argument(
         "--max-nodes",
         type=int,
@@ -458,7 +572,7 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
         help="Only consider placements using >= this many nodes.",
     )
     ap.add_argument(
-        "--instance-meta", choices=["ring", "jaccl", "both"], default="both"
+        "--instance-meta", choices=["ring", "jaccl", "vllm", "both"], default="both"
     )
     ap.add_argument(
         "--sharding", choices=["pipeline", "tensor", "both"], default="both"
@@ -487,3 +601,8 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
         action="store_true",
         help="Delete existing models from smallest to largest to make room for benchmark model.",
     )
+    ap.add_argument(
+        "--ensure-cuda",
+        action="store_true",
+        help="Verify the exo cluster has CUDA/vllm capability; error if not.",
+    )

bench/parallel_requests.py

@@ -0,0 +1,255 @@
+# type: ignore
+import argparse
+import asyncio
+import sys
+import termios
+import time
+import tty
+
+import aiohttp
+
+NUM_REQUESTS = 10
+BASE_URL = ""
+
+QUESTIONS = [
+    "What is the capital of Australia?",
+    "How many bones are in the human body?",
+    "What year did World War II end?",
+    "What is the speed of light in meters per second?",
+    "Who wrote Romeo and Juliet?",
+    "What is the chemical formula for water?",
+    "How many planets are in our solar system?",
+    "What is the largest ocean on Earth?",
+    "Who painted the Mona Lisa?",
+    "What is the boiling point of water in Celsius?",
+]
+
+
+def write(s: str) -> None:
+    sys.stdout.write(s)
+
+
+# ---------------------------------------------------------------------------
+# Model picker (same style as exo_eval)
+# ---------------------------------------------------------------------------
+
+
+def fetch_models() -> list[str]:
+    import json
+    import urllib.request
+
+    with urllib.request.urlopen(f"{BASE_URL}/state") as resp:
+        data = json.loads(resp.read())
+    model_ids: set[str] = set()
+    for instance in data.get("instances", {}).values():
+        for variant in instance.values():
+            sa = variant.get("shardAssignments", {})
+            model_id = sa.get("modelId")
+            if model_id:
+                model_ids.add(model_id)
+    return sorted(model_ids)
+
+
+def pick_model() -> str | None:
+    models = fetch_models()
+    if not models:
+        print("No models found.")
+        return None
+
+    cursor = 0
+    total_lines = len(models) + 4
+
+    def render(first: bool = False) -> None:
+        if not first:
+            write(f"\033[{total_lines}A")
+        write("\033[J")
+        write("\033[1mSelect model\033[0m (up/down, enter confirm, q quit)\r\n\r\n")
+        for i, model in enumerate(models):
+            line = f"  {'>' if i == cursor else ' '} {model}"
+            write(f"\033[7m{line}\033[0m\r\n" if i == cursor else f"{line}\r\n")
+        write("\r\n")
+        sys.stdout.flush()
+
+    fd = sys.stdin.fileno()
+    old = termios.tcgetattr(fd)
+    try:
+        tty.setraw(fd)
+        write("\033[?25l")
+        render(first=True)
+        while True:
+            ch = sys.stdin.read(1)
+            if ch in ("q", "\x03"):
+                write("\033[?25h\033[0m\r\n")
+                return None
+            elif ch in ("\r", "\n"):
+                break
+            elif ch == "\x1b":
+                seq = sys.stdin.read(2)
+                if seq == "[A":
+                    cursor = (cursor - 1) % len(models)
+                elif seq == "[B":
+                    cursor = (cursor + 1) % len(models)
+            render()
+    finally:
+        termios.tcsetattr(fd, termios.TCSADRAIN, old)
+        write(f"\033[{total_lines}A\033[J")  # clear picker UI
+        write("\033[?25h\033[0m")
+        sys.stdout.flush()
+
+    return models[cursor]
+
+
+# ---------------------------------------------------------------------------
+# Parallel requests
+# ---------------------------------------------------------------------------
+
+statuses: list[str] = []
+times: list[str] = []
+previews: list[str] = []
+tokens: list[str] = []
+full_responses: list[dict | None] = []
+total_lines = 0
+start_time: float = 0
+selected_model: str = ""
+
+
+def render_progress(first: bool = False) -> None:
+    if not first:
+        write(f"\033[{total_lines}A")
+    write("\033[J")
+    elapsed = time.monotonic() - start_time if start_time else 0
+    done = sum(1 for s in statuses if s == "done")
+    write(
+        f"\033[1m{selected_model}\033[0m  [{done}/{NUM_REQUESTS}]  {elapsed:.1f}s\r\n\r\n"
+    )
+
+    for i in range(NUM_REQUESTS):
+        q = QUESTIONS[i % len(QUESTIONS)]
+        status = statuses[i]
+        if status == "pending":
+            color = "\033[33m"  # yellow
+        elif status == "running":
+            color = "\033[36m"  # cyan
+        elif status == "done":
+            color = "\033[32m"  # green
+        else:
+            color = "\033[31m"  # red
+        write(
+            f"  {i:>2}  {color}{status:<8}\033[0m  {times[i]:>6}  {tokens[i]:>5}tok  {q[:40]:<40}  {previews[i][:50]}\r\n"
+        )
+
+    write("\r\n")
+    sys.stdout.flush()
+
+
+async def send_request(
+    session: aiohttp.ClientSession, i: int, lock: asyncio.Lock
+) -> None:
+    payload = {
+        "model": selected_model,
+        "messages": [{"role": "user", "content": QUESTIONS[i % len(QUESTIONS)]}],
+        "max_tokens": 1024,
+    }
+    statuses[i] = "running"
+    async with lock:
+        render_progress()
+    t0 = time.monotonic()
+    try:
+        async with session.post(
+            f"{BASE_URL}/v1/chat/completions", json=payload
+        ) as resp:
+            data = await resp.json()
+            elapsed = time.monotonic() - t0
+            full_responses[i] = data
+            times[i] = f"{elapsed:.1f}s"
+            if resp.status == 200:
+                choice = data["choices"][0]
+                msg = choice["message"]
+                content = msg.get("content", "")
+                previews[i] = content[:50].replace("\n", " ") or "(empty)"
+                if "usage" in data:
+                    tokens[i] = str(data["usage"].get("total_tokens", ""))
+                statuses[i] = "done"
+            else:
+                statuses[i] = f"err:{resp.status}"
+                previews[i] = str(data.get("error", {}).get("message", ""))[:50]
+    except Exception as e:
+        elapsed = time.monotonic() - t0
+        times[i] = f"{elapsed:.1f}s"
+        statuses[i] = "error"
+        previews[i] = str(e)[:50]
+    async with lock:
+        render_progress()
+
+
+async def run_requests(print_stdout: bool = False) -> None:
+    global start_time, total_lines, statuses, times, previews, tokens, full_responses
+
+    statuses = ["pending"] * NUM_REQUESTS
+    times = ["-"] * NUM_REQUESTS
+    previews = ["-"] * NUM_REQUESTS
+    tokens = ["-"] * NUM_REQUESTS
+    full_responses = [None] * NUM_REQUESTS
+    total_lines = NUM_REQUESTS + 4
+
+    write("\033[?25l")  # hide cursor
+    start_time = time.monotonic()
+    render_progress(first=True)
+    lock = asyncio.Lock()
+    try:
+        async with aiohttp.ClientSession() as session:
+            tasks = [send_request(session, i, lock) for i in range(NUM_REQUESTS)]
+            await asyncio.gather(*tasks)
+        total = time.monotonic() - start_time
+        write(
+            f"\033[1m=== All {NUM_REQUESTS} requests done in {total:.1f}s ===\033[0m\r\n\r\n"
+        )
+
+        if print_stdout:
+            for i in range(NUM_REQUESTS):
+                data = full_responses[i]
+                if not data or "choices" not in data:
+                    continue
+                choice = data["choices"][0]
+                msg = choice["message"]
+                q = QUESTIONS[i % len(QUESTIONS)]
+                write(f"\033[1m--- #{i}: {q} ---\033[0m\r\n")
+                if msg.get("reasoning_content"):
+                    write(f"\033[2m[Thinking]: {msg['reasoning_content']}\033[0m\r\n")
+                write(f"{msg.get('content', '')}\r\n")
+                if "usage" in data:
+                    u = data["usage"]
+                    write(
+                        f"\033[2m[Usage: prompt={u.get('prompt_tokens')}, "
+                        f"completion={u.get('completion_tokens')}, "
+                        f"total={u.get('total_tokens')}]\033[0m\r\n"
+                    )
+                write("\r\n")
+    finally:
+        write("\033[?25h")  # show cursor
+        sys.stdout.flush()
+
+
+def main() -> None:
+    global selected_model, BASE_URL
+    parser = argparse.ArgumentParser(
+        description="Send parallel requests to an exo cluster"
+    )
+    parser.add_argument(
+        "--host", required=True, help="Hostname of the exo node (e.g. s1)"
+    )
+    parser.add_argument("--port", type=int, default=52415, help="Port (default: 52415)")
+    parser.add_argument(
+        "--stdout", action="store_true", help="Print full responses after completion"
+    )
+    args = parser.parse_args()
+    BASE_URL = f"http://{args.host}:{args.port}"
+    model = pick_model()
+    if not model:
+        return
+    selected_model = model
+    asyncio.run(run_requests(print_stdout=args.stdout))
+
+
+if __name__ == "__main__":
+    main()

bench/pyproject.toml

@@ -4,13 +4,18 @@ version = "0.1.0"
 description = "Benchmarking tool for exo distributed inference"
 requires-python = ">=3.13"
 dependencies = [
-    "httpx>=0.27.0",
-    "loguru>=0.7.3",
-    "transformers>=5.0.0",
-    "huggingface-hub>=0.33.4",
-    "tiktoken>=0.12.0",
-    "jinja2>=3.1.0",
-    "protobuf>=5.29.0",
+  "httpx>=0.27.0",
+  "loguru>=0.7.3",
+  "transformers>=5.0.0",
+  "huggingface-hub>=0.33.4",
+  "tiktoken>=0.12.0",
+  "jinja2>=3.1.0",
+  "protobuf>=5.29.0",
+  "datasets>=2.0.0",
+  "math-verify>=0.7.0",
+  "lm-eval[api,math]>=0.4.0",
+  "human-eval>=1.0.3",
+  "numpy>=1.24.0",
 ]
 
 [build-system]

bench/single-m3-ultra.toml

@@ -2,10 +2,10 @@
 #
 # Shared constraints applied to ALL benchmarks in this file.
 constraints = [
-    "All(MacOsBuild(=25D125))",
-    "Hosts(=1)",
-    "All(Chip(m3_ultra))",
-    "All(GpuCores(=80))",
+  "All(MacOsBuild(=25D125))",
+  "Hosts(=1)",
+  "All(Chip(m3_ultra))",
+  "All(GpuCores(=80))",
 ]
 
 [topology]

bench/test_mlx_bandwidth.py

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

bench/vendor/lcb_testing_util.py

@@ -0,0 +1,593 @@
+# type: ignore
+# Vendored from LiveCodeBench (https://github.com/LiveCodeBench/LiveCodeBench)
+# File: lcb_runner/evaluation/testing_util.py
+# License: MIT
+# Vendored 2026-03-07 — do not modify without updating from upstream.
+
+import ast
+import faulthandler
+import json
+import platform
+
+# to run the solution files we're using a timing based approach
+import signal
+import sys
+import time
+
+# used for debugging to time steps
+from datetime import datetime
+from decimal import Decimal
+from enum import Enum
+from io import StringIO
+
+# from pyext import RuntimeModule
+from types import ModuleType
+
+# used for testing the code that reads from input
+from unittest.mock import mock_open, patch
+
+import_string = "from string import *\nfrom re import *\nfrom datetime import *\nfrom collections import *\nfrom heapq import *\nfrom bisect import *\nfrom copy import *\nfrom math import *\nfrom random import *\nfrom statistics import *\nfrom itertools import *\nfrom functools import *\nfrom operator import *\nfrom io import *\nfrom sys import *\nfrom json import *\nfrom builtins import *\nfrom typing import *\nimport string\nimport re\nimport datetime\nimport collections\nimport heapq\nimport bisect\nimport copy\nimport math\nimport random\nimport statistics\nimport itertools\nimport functools\nimport operator\nimport io\nimport sys\nimport json\nsys.setrecursionlimit(50000)\n"
+
+
+def truncatefn(s, length=300):
+    if isinstance(s, str):
+        pass
+    else:
+        s = str(s)
+    if len(s) <= length:
+        return s
+
+    return s[: length // 2] + "...(truncated) ..." + s[-length // 2 :]
+
+
+class CODE_TYPE(Enum):
+    call_based = 0
+    standard_input = 1
+
+
+# stuff for setting up signal timer
+class TimeoutException(Exception):
+    pass
+
+
+def timeout_handler(signum, frame):
+    print("timeout occured: alarm went off")
+    raise TimeoutException
+
+
+# used to capture stdout as a list
+# from https://stackoverflow.com/a/16571630/6416660
+# alternative use redirect_stdout() from contextlib
+class Capturing(list):
+    def __enter__(self):
+        self._stdout = sys.stdout
+        sys.stdout = self._stringio = StringIO()
+        # Make closing the StringIO a no-op
+        self._stringio.close = lambda x: 1
+        return self
+
+    def __exit__(self, *args):
+        self.append(self._stringio.getvalue())
+        del self._stringio  # free up some memory
+        sys.stdout = self._stdout
+
+
+# Custom mock for sys.stdin that supports buffer attribute
+class MockStdinWithBuffer:
+    def __init__(self, inputs: str):
+        self.inputs = inputs
+        self._stringio = StringIO(inputs)
+        self.buffer = MockBuffer(inputs)
+
+    def read(self, *args):
+        return self.inputs
+
+    def readline(self, *args):
+        return self._stringio.readline(*args)
+
+    def readlines(self, *args):
+        return self.inputs.split("\n")
+
+    def __getattr__(self, name):
+        # Delegate other attributes to StringIO
+        return getattr(self._stringio, name)
+
+
+class MockBuffer:
+    def __init__(self, inputs: str):
+        self.inputs = inputs.encode("utf-8")  # Convert to bytes
+
+    def read(self, *args):
+        # Return as byte strings that can be split
+        return self.inputs
+
+    def readline(self, *args):
+        return self.inputs.split(b"\n")[0] + b"\n"
+
+
+def clean_if_name(code: str) -> str:
+    try:
+        astree = ast.parse(code)
+        last_block = astree.body[-1]
+        if isinstance(last_block, ast.If):
+            condition = last_block.test
+            if ast.unparse(condition).strip() == "__name__ == '__main__'":
+                code = (
+                    ast.unparse(astree.body[:-1]) + "\n" + ast.unparse(last_block.body)  # type: ignore
+                )
+    except:
+        pass
+
+    return code
+
+
+def make_function(code: str) -> str:
+    try:
+        import_stmts = []
+        all_other_stmts = []
+        astree = ast.parse(code)
+        for stmt in astree.body:
+            if isinstance(stmt, (ast.Import, ast.ImportFrom)):
+                import_stmts.append(stmt)
+            else:
+                all_other_stmts.append(stmt)
+
+        function_ast = ast.FunctionDef(
+            name="wrapped_function",
+            args=ast.arguments(
+                posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]
+            ),
+            body=all_other_stmts,
+            decorator_list=[],
+            lineno=-1,
+        )
+        main_code = (
+            import_string
+            + "\n"
+            + ast.unparse(import_stmts)
+            + "\n"
+            + ast.unparse(function_ast)
+        )
+        return main_code
+    except Exception:
+        return code
+
+
+def call_method(method, inputs):
+    if isinstance(inputs, list):
+        inputs = "\n".join(inputs)
+
+    inputs_line_iterator = iter(inputs.split("\n"))
+
+    # Create custom stdin mock with buffer support
+    mock_stdin = MockStdinWithBuffer(inputs)
+
+    # sys.setrecursionlimit(10000)
+
+    # @patch('builtins.input', side_effect=inputs.split("\n"))
+    @patch("builtins.open", mock_open(read_data=inputs))
+    @patch("sys.stdin", mock_stdin)  # Use our custom mock instead of StringIO
+    @patch("sys.stdin.readline", lambda *args: next(inputs_line_iterator))
+    @patch("sys.stdin.readlines", lambda *args: inputs.split("\n"))
+    @patch("sys.stdin.read", lambda *args: inputs)
+    # @patch('sys.stdout.write', print)
+    def _inner_call_method(_method):
+        try:
+            return _method()
+        except SystemExit:
+            pass
+        finally:
+            pass
+
+    return _inner_call_method(method)
+
+
+def get_function(compiled_sol, fn_name: str):  # type: ignore
+    try:
+        assert hasattr(compiled_sol, fn_name)
+        return getattr(compiled_sol, fn_name)
+    except Exception:
+        return
+
+
+def compile_code(code: str, timeout: int):
+    signal.alarm(timeout)
+    try:
+        tmp_sol = ModuleType("tmp_sol", "")
+        exec(code, tmp_sol.__dict__)
+        if "class Solution" in code:
+            # leetcode wraps solutions in `Solution`
+            # this is a hack to check if it is leetcode solution or not
+            # currently livecodebench only supports LeetCode but
+            # else condition allows future extensibility to other platforms
+            compiled_sol = tmp_sol.Solution()
+        else:
+            # do nothing in the other case since function is accesible
+            compiled_sol = tmp_sol
+
+        assert compiled_sol is not None
+    finally:
+        signal.alarm(0)
+
+    return compiled_sol
+
+
+def convert_line_to_decimals(line: str) -> tuple[bool, list[Decimal]]:
+    try:
+        decimal_line = [Decimal(elem) for elem in line.split()]
+    except:
+        return False, []
+    return True, decimal_line
+
+
+def get_stripped_lines(val: str):
+    ## you don't want empty lines to add empty list after splitlines!
+    val = val.strip()
+
+    return [val_line.strip() for val_line in val.split("\n")]
+
+
+def grade_call_based(
+    code: str, all_inputs: list, all_outputs: list, fn_name: str, timeout: int
+):
+    # call-based clean up logic
+    # need to wrap in try-catch logic after to catch the correct errors, but for now this is fine.
+    code = import_string + "\n\n" + code
+    compiled_sol = compile_code(code, timeout)
+
+    if compiled_sol is None:
+        return
+
+    method = get_function(compiled_sol, fn_name)
+
+    if method is None:
+        return
+
+    all_inputs = [
+        [json.loads(line) for line in inputs.split("\n")] for inputs in all_inputs
+    ]
+
+    all_outputs = [json.loads(output) for output in all_outputs]
+
+    total_execution = 0
+    all_results = []
+    for idx, (gt_inp, gt_out) in enumerate(zip(all_inputs, all_outputs)):
+        signal.alarm(timeout)
+        faulthandler.enable()
+        try:
+            # can lock here so time is useful
+            start = time.time()
+            prediction = method(*gt_inp)
+            total_execution += time.time() - start
+            signal.alarm(0)
+
+            # don't penalize model if it produces tuples instead of lists
+            # ground truth sequences are not tuples
+            if isinstance(prediction, tuple):
+                prediction = list(prediction)
+
+            tmp_result = prediction == gt_out
+
+            # handle floating point comparisons
+
+            all_results.append(tmp_result)
+
+            if not tmp_result:
+                return all_results, {
+                    "output": truncatefn(prediction),
+                    "inputs": truncatefn(gt_inp),
+                    "expected": truncatefn(gt_out),
+                    "error_code": -2,
+                    "error_message": "Wrong Answer",
+                }
+        except Exception as e:
+            signal.alarm(0)
+            if "timeoutexception" in repr(e).lower():
+                all_results.append(-3)
+                return all_results, {
+                    "error": repr(e),
+                    "error_code": -3,
+                    "error_message": "Time Limit Exceeded",
+                    "inputs": truncatefn(gt_inp),
+                    "expected": truncatefn(gt_out),
+                }
+            else:
+                all_results.append(-4)
+                return all_results, {
+                    "error": repr(e),
+                    "error_code": -4,
+                    "error_message": "Runtime Error",
+                    "inputs": truncatefn(gt_inp),
+                    "expected": truncatefn(gt_out),
+                }
+
+        finally:
+            signal.alarm(0)
+            faulthandler.disable()
+
+    return all_results, {"execution time": total_execution}
+
+
+def grade_stdio(
+    code: str,
+    all_inputs: list,
+    all_outputs: list,
+    timeout: int,
+):
+    ## runtime doesn't interact well with __name__ == '__main__'
+    code = clean_if_name(code)
+
+    ## we wrap the given code inside another function
+    code = make_function(code)
+
+    compiled_sol = compile_code(code, timeout)
+    if compiled_sol is None:
+        return
+
+    method = get_function(compiled_sol, "wrapped_function")
+
+    if method is None:
+        return
+
+    all_results = []
+    total_execution_time = 0
+    for idx, (gt_inp, gt_out) in enumerate(zip(all_inputs, all_outputs)):
+        signal.alarm(timeout)
+        faulthandler.enable()
+
+        signal.alarm(timeout)
+        with Capturing() as captured_output:
+            try:
+                start = time.time()
+                call_method(method, gt_inp)
+                total_execution_time += time.time() - start
+                # reset the alarm
+                signal.alarm(0)
+            except Exception as e:
+                signal.alarm(0)
+                if "timeoutexception" in repr(e).lower():
+                    all_results.append(-3)
+                    return all_results, {
+                        "error": repr(e),
+                        "error_code": -3,
+                        "error_message": "Time Limit Exceeded",
+                        "inputs": truncatefn(gt_inp),
+                        "expected": truncatefn(gt_out),
+                    }
+                else:
+                    all_results.append(-4)
+                    return all_results, {
+                        "error": repr(e),
+                        "error_code": -4,
+                        "error_message": "Runtime Error",
+                        "inputs": truncatefn(gt_inp),
+                        "expected": truncatefn(gt_out),
+                    }
+
+            finally:
+                signal.alarm(0)
+                faulthandler.disable()
+
+        prediction = captured_output[0]
+
+        stripped_prediction_lines = get_stripped_lines(prediction)
+        stripped_gt_out_lines = get_stripped_lines(gt_out)
+
+        ## WA happens in multiple circumstances
+        ## so cache the return to make it clean!
+        WA_send_args = {
+            "output": truncatefn(prediction),
+            "inputs": truncatefn(gt_inp),
+            "expected": truncatefn(gt_out),
+            "error_code": -2,
+        }
+
+        if len(stripped_prediction_lines) != len(stripped_gt_out_lines):
+            all_results.append(-2)
+            WA_send_args["error_message"] = "Wrong answer: mismatched output length"
+            return all_results, WA_send_args
+
+        for output_line_idx, (
+            stripped_prediction_line,
+            stripped_gt_out_line,
+        ) in enumerate(zip(stripped_prediction_lines, stripped_gt_out_lines)):
+            WA_send_args["error_message"] = (
+                f"Wrong answer at {output_line_idx=}: {truncatefn(stripped_prediction_line)} != {truncatefn(stripped_gt_out_line)}"
+            )
+
+            ## CASE 1: exact match
+            if stripped_prediction_line == stripped_gt_out_line:
+                continue
+
+            ## CASE 2: element-wise comparision
+            ## if there are floating elements
+            ## use `decimal` library for good floating point comparision
+            ## otherwise gotcha: np.isclose(50000000000000000, 50000000000000001) = True
+            ## note that we should always be able to convert to decimals
+
+            success, decimal_prediction_line = convert_line_to_decimals(
+                stripped_prediction_line
+            )
+            if not success:
+                all_results.append(-2)
+                return all_results, WA_send_args
+            success, decimal_gtout_line = convert_line_to_decimals(stripped_gt_out_line)
+            if not success:
+                all_results.append(-2)
+                return all_results, WA_send_args
+
+            if decimal_prediction_line == decimal_gtout_line:
+                continue
+
+            all_results.append(-2)
+            return all_results, WA_send_args
+        all_results.append(True)
+
+    return all_results, {"execution time": total_execution_time}
+
+
+def run_test(sample, test=None, debug=False, timeout=6):
+    """
+    if test(generated_code) is not None it'll try to run the code.
+    otherwise it'll just return an input and output pair.
+    """
+    signal.signal(signal.SIGALRM, timeout_handler)
+
+    # Disable functionalities that can make destructive changes to the test.
+    # max memory is set to 4GB
+    reliability_guard()
+
+    if debug:
+        print(f"start = {datetime.now().time()}")
+
+    try:
+        in_outs = json.loads(sample["input_output"])
+    except ValueError as e:
+        raise e
+        in_outs = None
+
+    if in_outs:
+        if in_outs.get("fn_name") is None:
+            which_type = CODE_TYPE.standard_input  # Standard input
+            method_name = None
+
+        else:
+            which_type = CODE_TYPE.call_based  # Call-based
+            method_name = in_outs["fn_name"]
+
+    if debug:
+        print(f"loaded input_output = {datetime.now().time()}")
+
+    if test is None:
+        assert False, "should not happen: test code is none"
+        return in_outs, {"error": "no test code provided"}
+    elif test is not None:
+        results = []
+        sol = import_string
+        if debug:
+            print(f"loading test code = {datetime.now().time()}")
+
+        if which_type == CODE_TYPE.call_based:
+            signal.alarm(timeout)
+            try:
+                results, metadata = grade_call_based(
+                    code=test,
+                    all_inputs=in_outs["inputs"],
+                    all_outputs=in_outs["outputs"],
+                    fn_name=method_name,
+                    timeout=timeout,
+                )
+                return results, metadata
+            except Exception as e:
+                return [-4], {
+                    "error_code": -4,
+                    "error_message": f"Error during testing: {e}",
+                }
+            finally:
+                signal.alarm(0)
+        elif which_type == CODE_TYPE.standard_input:
+            # sol
+            # if code has if __name__ == "__main__": then remove it
+
+            signal.alarm(timeout)
+            try:
+                results, metadata = grade_stdio(
+                    code=test,
+                    all_inputs=in_outs["inputs"],
+                    all_outputs=in_outs["outputs"],
+                    timeout=timeout,
+                )
+                return results, metadata
+            except Exception as e:
+                return [-4], {
+                    "error_code": -4,
+                    "error_message": f"Error during testing: {e}",
+                }
+            finally:
+                signal.alarm(0)
+
+
+def reliability_guard(maximum_memory_bytes=None):
+    """
+    This disables various destructive functions and prevents the generated code
+    from interfering with the test (e.g. fork bomb, killing other processes,
+    removing filesystem files, etc.)
+    WARNING
+    This function is NOT a security sandbox. Untrusted code, including, model-
+    generated code, should not be blindly executed outside of one. See the
+    Codex paper for more information about OpenAI's code sandbox, and proceed
+    with caution.
+    """
+
+    if maximum_memory_bytes is not None:
+        import resource
+
+        resource.setrlimit(
+            resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes)
+        )
+        resource.setrlimit(
+            resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes)
+        )
+        if not platform.uname().system == "Darwin":
+            resource.setrlimit(
+                resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes)
+            )
+
+    faulthandler.disable()
+
+    import builtins
+
+    # builtins.exit = None
+    builtins.quit = None
+
+    import os
+
+    os.environ["OMP_NUM_THREADS"] = "1"
+
+    os.kill = None
+    os.system = None
+    os.putenv = None
+    os.remove = None
+    os.removedirs = None
+    os.rmdir = None
+    os.fchdir = None
+    os.setuid = None
+    os.fork = None
+    os.forkpty = None
+    os.killpg = None
+    os.rename = None
+    os.renames = None
+    os.truncate = None
+    os.replace = None
+    os.unlink = None
+    os.fchmod = None
+    os.fchown = None
+    os.chmod = None
+    os.chown = None
+    os.chroot = None
+    os.fchdir = None
+    os.lchflags = None
+    os.lchmod = None
+    os.lchown = None
+    os.getcwd = None
+    os.chdir = None
+
+    import shutil
+
+    shutil.rmtree = None
+    shutil.move = None
+    shutil.chown = None
+
+    import subprocess
+
+    subprocess.Popen = None
+
+    __builtins__["help"] = None
+
+    import sys
+
+    sys.modules["ipdb"] = None
+    sys.modules["joblib"] = None
+    sys.modules["resource"] = None
+    sys.modules["psutil"] = None
+    sys.modules["tkinter"] = None

bench_overnight.sh

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

dashboard/package-lock.json

@@ -11,7 +11,8 @@
 				"highlight.js": "^11.11.1",
 				"katex": "^0.16.27",
 				"marked": "^17.0.1",
-				"mode-watcher": "^1.1.0"
+				"mode-watcher": "^1.1.0",
+				"pdfjs-dist": "^5.6.205"
 			},
 			"devDependencies": {
 				"@sveltejs/adapter-static": "^3.0.10",
@@ -518,6 +519,256 @@
 				"@jridgewell/sourcemap-codec": "^1.4.14"
 			}
 		},
+		"node_modules/@napi-rs/canvas": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas/-/canvas-0.1.97.tgz",
+			"integrity": "sha512-8cFniXvrIEnVwuNSRCW9wirRZbHvrD3JVujdS2P5n5xiJZNZMOZcfOvJ1pb66c7jXMKHHglJEDVJGbm8XWFcXQ==",
+			"license": "MIT",
+			"optional": true,
+			"workspaces": [
+				"e2e/*"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			},
+			"optionalDependencies": {
+				"@napi-rs/canvas-android-arm64": "0.1.97",
+				"@napi-rs/canvas-darwin-arm64": "0.1.97",
+				"@napi-rs/canvas-darwin-x64": "0.1.97",
+				"@napi-rs/canvas-linux-arm-gnueabihf": "0.1.97",
+				"@napi-rs/canvas-linux-arm64-gnu": "0.1.97",
+				"@napi-rs/canvas-linux-arm64-musl": "0.1.97",
+				"@napi-rs/canvas-linux-riscv64-gnu": "0.1.97",
+				"@napi-rs/canvas-linux-x64-gnu": "0.1.97",
+				"@napi-rs/canvas-linux-x64-musl": "0.1.97",
+				"@napi-rs/canvas-win32-arm64-msvc": "0.1.97",
+				"@napi-rs/canvas-win32-x64-msvc": "0.1.97"
+			}
+		},
+		"node_modules/@napi-rs/canvas-android-arm64": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-android-arm64/-/canvas-android-arm64-0.1.97.tgz",
+			"integrity": "sha512-V1c/WVw+NzH8vk7ZK/O8/nyBSCQimU8sfMsB/9qeSvdkGKNU7+mxy/bIF0gTgeBFmHpj30S4E9WHMSrxXGQuVQ==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"android"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-darwin-arm64": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-darwin-arm64/-/canvas-darwin-arm64-0.1.97.tgz",
+			"integrity": "sha512-ok+SCEF4YejcxuJ9Rm+WWunHHpf2HmiPxfz6z1a/NFQECGXtsY7A4B8XocK1LmT1D7P174MzwPF9Wy3AUAwEPw==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"darwin"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-darwin-x64": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-darwin-x64/-/canvas-darwin-x64-0.1.97.tgz",
+			"integrity": "sha512-PUP6e6/UGlclUvAQNnuXCcnkpdUou6VYZfQOQxExLp86epOylmiwLkqXIvpFmjoTEDmPmXrI+coL/9EFU1gKPA==",
+			"cpu": [
+				"x64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"darwin"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-arm-gnueabihf": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm-gnueabihf/-/canvas-linux-arm-gnueabihf-0.1.97.tgz",
+			"integrity": "sha512-XyXH2L/cic8eTNtbrXCcvqHtMX/nEOxN18+7rMrAM2XtLYC/EB5s0wnO1FsLMWmK+04ZSLN9FBGipo7kpIkcOw==",
+			"cpu": [
+				"arm"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-arm64-gnu": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm64-gnu/-/canvas-linux-arm64-gnu-0.1.97.tgz",
+			"integrity": "sha512-Kuq/M3djq0K8ktgz6nPlK7Ne5d4uWeDxPpyKWOjWDK2RIOhHVtLtyLiJw2fuldw7Vn4mhw05EZXCEr4Q76rs9w==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-arm64-musl": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-arm64-musl/-/canvas-linux-arm64-musl-0.1.97.tgz",
+			"integrity": "sha512-kKmSkQVnWeqg7qdsiXvYxKhAFuHz3tkBjW/zyQv5YKUPhotpaVhpBGv5LqCngzyuRV85SXoe+OFj+Tv0a0QXkQ==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-riscv64-gnu": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-riscv64-gnu/-/canvas-linux-riscv64-gnu-0.1.97.tgz",
+			"integrity": "sha512-Jc7I3A51jnEOIAXeLsN/M/+Z28LUeakcsXs07FLq9prXc0eYOtVwsDEv913Gr+06IRo34gJJVgT0TXvmz+N2VA==",
+			"cpu": [
+				"riscv64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-x64-gnu": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-x64-gnu/-/canvas-linux-x64-gnu-0.1.97.tgz",
+			"integrity": "sha512-iDUBe7AilfuBSRbSa8/IGX38Mf+iCSBqoVKLSQ5XaY2JLOaqz1TVyPFEyIck7wT6mRQhQt5sN6ogfjIDfi74tg==",
+			"cpu": [
+				"x64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-linux-x64-musl": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-linux-x64-musl/-/canvas-linux-x64-musl-0.1.97.tgz",
+			"integrity": "sha512-AKLFd/v0Z5fvgqBDqhvqtAdx+fHMJ5t9JcUNKq4FIZ5WH+iegGm8HPdj00NFlCSnm83Fp3Ln8I2f7uq1aIiWaA==",
+			"cpu": [
+				"x64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"linux"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-win32-arm64-msvc": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-win32-arm64-msvc/-/canvas-win32-arm64-msvc-0.1.97.tgz",
+			"integrity": "sha512-u883Yr6A6fO7Vpsy9YE4FVCIxzzo5sO+7pIUjjoDLjS3vQaNMkVzx5bdIpEL+ob+gU88WDK4VcxYMZ6nmnoX9A==",
+			"cpu": [
+				"arm64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"win32"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
+		"node_modules/@napi-rs/canvas-win32-x64-msvc": {
+			"version": "0.1.97",
+			"resolved": "https://registry.npmjs.org/@napi-rs/canvas-win32-x64-msvc/-/canvas-win32-x64-msvc-0.1.97.tgz",
+			"integrity": "sha512-sWtD2EE3fV0IzN+iiQUqr/Q1SwqWhs2O1FKItFlxtdDkikpEj5g7DKQpY3x55H/MAOnL8iomnlk3mcEeGiUMoQ==",
+			"cpu": [
+				"x64"
+			],
+			"license": "MIT",
+			"optional": true,
+			"os": [
+				"win32"
+			],
+			"engines": {
+				"node": ">= 10"
+			},
+			"funding": {
+				"type": "github",
+				"url": "https://github.com/sponsors/Brooooooklyn"
+			}
+		},
 		"node_modules/@polka/url": {
 			"version": "1.0.0-next.29",
 			"resolved": "https://registry.npmjs.org/@polka/url/-/url-1.0.0-next.29.tgz",
@@ -2635,6 +2886,26 @@
 				"node": "^10 || ^12 || ^13.7 || ^14 || >=15.0.1"
 			}
 		},
+		"node_modules/node-readable-to-web-readable-stream": {
+			"version": "0.4.2",
+			"resolved": "https://registry.npmjs.org/node-readable-to-web-readable-stream/-/node-readable-to-web-readable-stream-0.4.2.tgz",
+			"integrity": "sha512-/cMZNI34v//jUTrI+UIo4ieHAB5EZRY/+7OmXZgBxaWBMcW2tGdceIw06RFxWxrKZ5Jp3sI2i5TsRo+CBhtVLQ==",
+			"license": "MIT",
+			"optional": true
+		},
+		"node_modules/pdfjs-dist": {
+			"version": "5.6.205",
+			"resolved": "https://registry.npmjs.org/pdfjs-dist/-/pdfjs-dist-5.6.205.tgz",
+			"integrity": "sha512-tlUj+2IDa7G1SbvBNN74UHRLJybZDWYom+k6p5KIZl7huBvsA4APi6mKL+zCxd3tLjN5hOOEE9Tv7VdzO88pfg==",
+			"license": "Apache-2.0",
+			"engines": {
+				"node": ">=20.19.0 || >=22.13.0 || >=24"
+			},
+			"optionalDependencies": {
+				"@napi-rs/canvas": "^0.1.96",
+				"node-readable-to-web-readable-stream": "^0.4.2"
+			}
+		},
 		"node_modules/picocolors": {
 			"version": "1.1.1",
 			"resolved": "https://registry.npmjs.org/picocolors/-/picocolors-1.1.1.tgz",

dashboard/package.json

@@ -31,6 +31,7 @@
 		"highlight.js": "^11.11.1",
 		"katex": "^0.16.27",
 		"marked": "^17.0.1",
-		"mode-watcher": "^1.1.0"
+		"mode-watcher": "^1.1.0",
+		"pdfjs-dist": "^5.6.205"
 	}
 }

dashboard/src/lib/components/ChatForm.svelte

@@ -1,9 +1,6 @@
 <script lang="ts">
   import {
     isLoading,
-    sendMessage,
-    generateImage,
-    editImage,
     editingImage,
     clearEditingImage,
     selectedChatModel,
@@ -28,7 +25,7 @@
     modelTasks?: Record<string, string[]>;
     modelCapabilities?: Record<string, string[]>;
     onSend?: () => void;
-    onAutoSend?: (
+    onAutoSend: (
       content: string,
       files?: {
         id: string;
@@ -216,49 +213,10 @@
     uploadedFiles = [];
     resetTextareaHeight();
 
-    // When onAutoSend is provided, the parent controls all send logic
-    // (including launching non-running models before sending)
-    if (onAutoSend) {
-      onAutoSend(content, files);
-      onSend?.();
-      setTimeout(() => textareaRef?.focus(), 10);
-      return;
-    }
-
-    // Use image editing if in edit mode
-    if (isEditMode && currentEditingImage && content) {
-      editImage(content, currentEditingImage.imageDataUrl);
-    }
-    // If user attached an image with an ImageToImage model, use edit endpoint
-    else if (
-      currentModel &&
-      modelSupportsImageEditing(currentModel) &&
-      files.length > 0 &&
-      content
-    ) {
-      // Use the first attached image for editing
-      const imageFile = files[0];
-      if (imageFile.preview) {
-        editImage(content, imageFile.preview);
-      }
-    } else if (
-      currentModel &&
-      modelSupportsTextToImage(currentModel) &&
-      content
-    ) {
-      // Use image generation for text-to-image models
-      generateImage(content);
-    } else {
-      sendMessage(
-        content,
-        files,
-        modelSupportsThinking() ? thinkingEnabled : null,
-      );
-    }
-
+    // Parent controls all send logic (including image routing,
+    // launching non-running models before sending, etc.)
+    onAutoSend(content, files);
     onSend?.();
-
-    // Refocus the textarea after sending
     setTimeout(() => textareaRef?.focus(), 10);
   }
 

dashboard/src/lib/components/ChatMessages.svelte

@@ -139,6 +139,8 @@
         return "🖼";
       case "text":
         return "📄";
+      case "pdf":
+        return "📑";
       default:
         return "📎";
     }

dashboard/src/lib/components/ChatModelSelector.svelte

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

dashboard/src/lib/components/ChatSidebar.svelte

@@ -18,12 +18,18 @@
     class?: string;
     onNewChat?: () => void;
     onSelectConversation?: () => void;
+    isMobileDrawer?: boolean;
+    isOpen?: boolean;
+    onClose?: () => void;
   }
 
   let {
     class: className = "",
     onNewChat,
     onSelectConversation,
+    isMobileDrawer = false,
+    isOpen = false,
+    onClose,
   }: Props = $props();
 
   const conversationList = $derived(conversations());
@@ -53,6 +59,10 @@
   function handleSelectConversation(id: string) {
     onSelectConversation?.();
     loadConversation(id);
+    // Close mobile drawer when selecting a conversation
+    if (isMobileDrawer && isOpen) {
+      onClose?.();
+    }
   }
 
   function handleStartEdit(id: string, name: string, event: MouseEvent) {
@@ -252,9 +262,7 @@
   }
 </script>
 
-<aside
-  class="flex flex-col h-full bg-exo-dark-gray border-r border-exo-yellow/10 {className}"
->
+{#snippet sidebarContent()}
   <!-- Header -->
   <div class="p-4">
     <button
@@ -591,4 +599,30 @@
       </button>
     </div>
   </div>
-</aside>
+{/snippet}
+
+{#if isMobileDrawer}
+  <!-- Mobile drawer with overlay -->
+  {#if isOpen}
+    <!-- Overlay backdrop -->
+    <button
+      type="button"
+      class="fixed inset-0 bg-black/60 backdrop-blur-sm z-40 md:hidden"
+      onclick={() => onClose?.()}
+      aria-label="Close sidebar"
+    ></button>
+    <!-- Drawer panel -->
+    <aside
+      class="fixed left-0 top-0 bottom-0 w-72 bg-exo-dark-gray border-r border-exo-yellow/10 z-50 flex flex-col md:hidden"
+    >
+      {@render sidebarContent()}
+    </aside>
+  {/if}
+{:else}
+  <!-- Desktop sidebar -->
+  <aside
+    class="flex flex-col h-full bg-exo-dark-gray border-r border-exo-yellow/10 {className}"
+  >
+    {@render sidebarContent()}
+  </aside>
+{/if}

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

@@ -82,6 +82,24 @@
       d="M12.025 1.13c-5.77 0-10.449 4.647-10.449 10.378 0 1.112.178 2.181.503 3.185.064-.222.203-.444.416-.577a.96.96 0 0 1 .524-.15c.293 0 .584.124.84.284.278.173.48.408.71.694.226.282.458.611.684.951v-.014c.017-.324.106-.622.264-.874s.403-.487.762-.543c.3-.047.596.06.787.203s.31.313.4.467c.15.257.212.468.233.542.01.026.653 1.552 1.657 2.54.616.605 1.01 1.223 1.082 1.912.055.537-.096 1.059-.38 1.572.637.121 1.294.187 1.967.187.657 0 1.298-.063 1.921-.178-.287-.517-.44-1.041-.384-1.581.07-.69.465-1.307 1.081-1.913 1.004-.987 1.647-2.513 1.657-2.539.021-.074.083-.285.233-.542.09-.154.208-.323.4-.467a1.08 1.08 0 0 1 .787-.203c.359.056.604.29.762.543s.247.55.265.874v.015c.225-.34.457-.67.683-.952.23-.286.432-.52.71-.694.257-.16.547-.284.84-.285a.97.97 0 0 1 .524.151c.228.143.373.388.43.625l.006.04a10.3 10.3 0 0 0 .534-3.273c0-5.731-4.678-10.378-10.449-10.378M8.327 6.583a1.5 1.5 0 0 1 .713.174 1.487 1.487 0 0 1 .617 2.013c-.183.343-.762-.214-1.102-.094-.38.134-.532.914-.917.71a1.487 1.487 0 0 1 .69-2.803m7.486 0a1.487 1.487 0 0 1 .689 2.803c-.385.204-.536-.576-.916-.71-.34-.12-.92.437-1.103.094a1.487 1.487 0 0 1 .617-2.013 1.5 1.5 0 0 1 .713-.174m-10.68 1.55a.96.96 0 1 1 0 1.921.96.96 0 0 1 0-1.92m13.838 0a.96.96 0 1 1 0 1.92.96.96 0 0 1 0-1.92M8.489 11.458c.588.01 1.965 1.157 3.572 1.164 1.607-.007 2.984-1.155 3.572-1.164.196-.003.305.12.305.454 0 .886-.424 2.328-1.563 3.202-.22-.756-1.396-1.366-1.63-1.32q-.011.001-.02.006l-.044.026-.01.008-.03.024q-.018.017-.035.036l-.032.04a1 1 0 0 0-.058.09l-.014.025q-.049.088-.11.19a1 1 0 0 1-.083.116 1.2 1.2 0 0 1-.173.18q-.035.029-.075.058a1.3 1.3 0 0 1-.251-.243 1 1 0 0 1-.076-.107c-.124-.193-.177-.363-.337-.444-.034-.016-.104-.008-.2.022q-.094.03-.216.087-.06.028-.125.063l-.13.074q-.067.04-.136.086a3 3 0 0 0-.135.096 3 3 0 0 0-.26.219 2 2 0 0 0-.12.121 2 2 0 0 0-.106.128l-.002.002a2 2 0 0 0-.09.132l-.001.001a1.2 1.2 0 0 0-.105.212q-.013.036-.024.073c-1.139-.875-1.563-2.317-1.563-3.203 0-.334.109-.457.305-.454m.836 10.354c.824-1.19.766-2.082-.365-3.194-1.13-1.112-1.789-2.738-1.789-2.738s-.246-.945-.806-.858-.97 1.499.202 2.362c1.173.864-.233 1.45-.685.64-.45-.812-1.683-2.896-2.322-3.295s-1.089-.175-.938.647 2.822 2.813 2.562 3.244-1.176-.506-1.176-.506-2.866-2.567-3.49-1.898.473 1.23 2.037 2.16c1.564.932 1.686 1.178 1.464 1.53s-3.675-2.511-4-1.297c-.323 1.214 3.524 1.567 3.287 2.405-.238.839-2.71-1.587-3.216-.642-.506.946 3.49 2.056 3.522 2.064 1.29.33 4.568 1.028 5.713-.624m5.349 0c-.824-1.19-.766-2.082.365-3.194 1.13-1.112 1.789-2.738 1.789-2.738s.246-.945.806-.858.97 1.499-.202 2.362c-1.173.864.233 1.45.685.64.451-.812 1.683-2.896 2.322-3.295s1.089-.175.938.647-2.822 2.813-2.562 3.244 1.176-.506 1.176-.506 2.866-2.567 3.49-1.898-.473 1.23-2.037 2.16c-1.564.932-1.686 1.178-1.464 1.53s3.675-2.511 4-1.297c.323 1.214-3.524 1.567-3.287 2.405.238.839 2.71-1.587 3.216-.642.506.946-3.49 2.056-3.522 2.064-1.29.33-4.568 1.028-5.713-.624"
     />
   </svg>
+{:else if family === "step"}
+  <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
+    <path
+      d="M22.012 0h1.032v.927H24v.968h-.956V3.78h-1.032V1.896h-1.878v-.97h1.878V0zM2.6 12.371V1.87h.969v10.502h-.97zm10.423.66h10.95v.918h-6.208v9.579h-4.742V13.03zM5.629 3.333v12.356H0v4.51h10.386V8L20.859 8l-.003-4.668-15.227.001z"
+    />
+  </svg>
+{:else if family === "nemotron"}
+  <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
+    <path
+      d="M8.948 8.798v-1.43a6.7 6.7 0 0 1 .424-.018c3.922-.124 6.493 3.374 6.493 3.374s-2.774 3.851-5.75 3.851c-.398 0-.787-.062-1.158-.185v-4.346c1.528.185 1.837.857 2.747 2.385l2.04-1.714s-1.492-1.952-4-1.952a6.016 6.016 0 0 0-.796.035m0-4.735v2.138l.424-.027c5.45-.185 9.01 4.47 9.01 4.47s-4.08 4.964-8.33 4.964c-.37 0-.733-.035-1.095-.097v1.325c.3.035.61.062.91.062 3.957 0 6.82-2.023 9.593-4.408.459.371 2.34 1.263 2.73 1.652-2.633 2.208-8.772 3.984-12.253 3.984-.335 0-.653-.018-.971-.053v1.864H24V4.063zm0 10.326v1.131c-3.657-.654-4.673-4.46-4.673-4.46s1.758-1.944 4.673-2.262v1.237H8.94c-1.528-.186-2.73 1.245-2.73 1.245s.68 2.412 2.739 3.11M2.456 10.9s2.164-3.197 6.5-3.533V6.201C4.153 6.59 0 10.653 0 10.653s2.35 6.802 8.948 7.42v-1.237c-4.84-.6-6.492-5.936-6.492-5.936z"
+    />
+  </svg>
+{:else if family === "gemma"}
+  <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
+    <path
+      d="M12.48 10.92v3.28h7.84c-.24 1.84-.853 3.187-1.787 4.133-1.147 1.147-2.933 2.4-6.053 2.4-4.827 0-8.6-3.893-8.6-8.72s3.773-8.72 8.6-8.72c2.6 0 4.507 1.027 5.907 2.347l2.307-2.307C18.747 1.44 16.133 0 12.48 0 5.867 0 .307 5.387.307 12s5.56 12 12.173 12c3.573 0 6.267-1.173 8.373-3.36 2.16-2.16 2.84-5.213 2.84-7.667 0-.76-.053-1.467-.173-2.053H12.48z"
+    />
+  </svg>
 {:else}
   <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
     <path

dashboard/src/lib/components/FamilySidebar.svelte

@@ -31,6 +31,8 @@
     kimi: "Kimi",
     flux: "FLUX",
     "qwen-image": "Qwen Img",
+    nemotron: "NVIDIA",
+    gemma: "Google",
   };
 
   function getFamilyName(family: string): string {
@@ -41,31 +43,20 @@
 </script>
 
 <div
-  class="flex flex-col gap-1 py-2 px-1 border-r border-exo-yellow/10 bg-exo-medium-gray/30 min-w-[64px] overflow-y-auto scrollbar-hide"
+  class="flex flex-col gap-1 py-2 px-1 border-r border-exo-yellow/10 bg-exo-medium-gray/30 min-w-[80px] sm:min-w-[72px] overflow-y-auto scrollbar-hide"
 >
   <!-- All models (no filter) -->
   <button
     type="button"
     onclick={() => onSelect(null)}
-    class="group flex flex-col items-center justify-center p-2 rounded transition-all duration-200 cursor-pointer {selectedFamily ===
+    class="group flex items-center justify-center px-3 py-2.5 rounded transition-all duration-200 cursor-pointer min-h-[44px] sm:min-h-0 {selectedFamily ===
     null
       ? 'bg-exo-yellow/20 border-l-2 border-exo-yellow'
       : 'hover:bg-white/5 border-l-2 border-transparent'}"
     title="All models"
   >
-    <svg
-      class="w-5 h-5 {selectedFamily === null
-        ? 'text-exo-yellow'
-        : 'text-white/50 group-hover:text-white/70'}"
-      viewBox="0 0 24 24"
-      fill="currentColor"
-    >
-      <path
-        d="M4 8h4V4H4v4zm6 12h4v-4h-4v4zm-6 0h4v-4H4v4zm0-6h4v-4H4v4zm6 0h4v-4h-4v4zm6-10v4h4V4h-4zm-6 4h4V4h-4v4zm6 6h4v-4h-4v4zm0 6h4v-4h-4v4z"
-      />
-    </svg>
     <span
-      class="text-[9px] font-mono mt-0.5 {selectedFamily === null
+      class="text-[12px] font-mono font-medium {selectedFamily === null
         ? 'text-exo-yellow'
         : 'text-white/40 group-hover:text-white/60'}">All</span
     >
@@ -89,7 +80,7 @@
           : "text-white/50 group-hover:text-amber-400/70"}
       />
       <span
-        class="text-[9px] font-mono mt-0.5 {selectedFamily === 'favorites'
+        class="text-[11px] font-mono mt-0.5 {selectedFamily === 'favorites'
           ? 'text-amber-400'
           : 'text-white/40 group-hover:text-white/60'}">Faves</span
       >
@@ -114,7 +105,7 @@
           : "text-white/50 group-hover:text-white/70"}
       />
       <span
-        class="text-[9px] font-mono mt-0.5 {selectedFamily === 'recents'
+        class="text-[11px] font-mono mt-0.5 {selectedFamily === 'recents'
           ? 'text-exo-yellow'
           : 'text-white/40 group-hover:text-white/60'}">Recent</span
       >
@@ -138,7 +129,7 @@
         : "text-white/50 group-hover:text-orange-400/70"}
     />
     <span
-      class="text-[9px] font-mono mt-0.5 {selectedFamily === 'huggingface'
+      class="text-[11px] font-mono mt-0.5 {selectedFamily === 'huggingface'
         ? 'text-orange-400'
         : 'text-white/40 group-hover:text-white/60'}">Hub</span
     >
@@ -164,7 +155,7 @@
           : "text-white/50 group-hover:text-white/70"}
       />
       <span
-        class="text-[9px] font-mono mt-0.5 truncate max-w-full {selectedFamily ===
+        class="text-[11px] font-mono mt-0.5 truncate max-w-full {selectedFamily ===
         family
           ? 'text-exo-yellow'
           : 'text-white/40 group-hover:text-white/60'}"

dashboard/src/lib/components/HeaderNav.svelte

@@ -1,15 +1,38 @@
 <script lang="ts">
   import { browser } from "$app/environment";
 
-  export let showHome = true;
-  export let onHome: (() => void) | null = null;
-  export let showSidebarToggle = false;
-  export let sidebarVisible = true;
-  export let onToggleSidebar: (() => void) | null = null;
-  export let downloadProgress: {
-    count: number;
-    percentage: number;
-  } | null = null;
+  interface Props {
+    showHome?: boolean;
+    onHome?: (() => void) | null;
+    showSidebarToggle?: boolean;
+    sidebarVisible?: boolean;
+    onToggleSidebar?: (() => void) | null;
+    showMobileMenuToggle?: boolean;
+    mobileMenuOpen?: boolean;
+    onToggleMobileMenu?: (() => void) | null;
+    showMobileRightToggle?: boolean;
+    mobileRightOpen?: boolean;
+    onToggleMobileRight?: (() => void) | null;
+    downloadProgress?: {
+      count: number;
+      percentage: number;
+    } | null;
+  }
+
+  let {
+    showHome = true,
+    onHome = null,
+    showSidebarToggle = false,
+    sidebarVisible = true,
+    onToggleSidebar = null,
+    showMobileMenuToggle = false,
+    mobileMenuOpen = false,
+    onToggleMobileMenu = null,
+    showMobileRightToggle = false,
+    mobileRightOpen = false,
+    onToggleMobileRight = null,
+    downloadProgress = null,
+  }: Props = $props();
 
   function handleHome(): void {
     if (onHome) {
@@ -27,49 +50,96 @@
       onToggleSidebar();
     }
   }
+
+  function handleToggleMobileMenu(): void {
+    if (onToggleMobileMenu) {
+      onToggleMobileMenu();
+    }
+  }
+
+  function handleToggleMobileRight(): void {
+    if (onToggleMobileRight) {
+      onToggleMobileRight();
+    }
+  }
 </script>
 
 <header
-  class="relative z-20 flex items-center justify-center px-6 pt-8 pb-4 bg-exo-dark-gray"
+  class="relative z-20 flex items-center justify-center px-4 md:px-6 pt-4 md:pt-8 pb-3 md:pb-4 bg-exo-dark-gray"
 >
-  <!-- Left: Sidebar Toggle -->
-  {#if showSidebarToggle}
-    <div class="absolute left-6 top-1/2 -translate-y-1/2">
-      <button
-        onclick={handleToggleSidebar}
-        class="p-2 rounded border border-exo-light-gray/30 hover:border-exo-yellow/50 hover:bg-exo-medium-gray/30 transition-colors cursor-pointer"
-        title={sidebarVisible ? "Hide sidebar" : "Show sidebar"}
-        aria-label={sidebarVisible
-          ? "Hide conversation sidebar"
-          : "Show conversation sidebar"}
-        aria-pressed={sidebarVisible}
+  <!-- Left: Sidebar Toggle (desktop) or Mobile Sidebar Toggle (mobile) -->
+  <div
+    class="absolute left-4 md:left-6 top-1/2 -translate-y-1/2 flex items-center gap-2"
+  >
+    <!-- Mobile sidebar toggle -->
+    <button
+      onclick={handleToggleMobileMenu}
+      class="p-2 rounded border border-exo-light-gray/30 hover:border-exo-yellow/50 hover:bg-exo-medium-gray/30 transition-colors cursor-pointer md:hidden"
+      title={mobileMenuOpen ? "Hide sidebar" : "Show sidebar"}
+      aria-label={mobileMenuOpen
+        ? "Hide conversation sidebar"
+        : "Show conversation sidebar"}
+      aria-pressed={mobileMenuOpen}
+    >
+      <svg
+        fill="none"
+        viewBox="0 0 24 24"
+        stroke="currentColor"
+        stroke-width="2"
+        class="w-5 h-5 {mobileMenuOpen
+          ? 'text-exo-yellow'
+          : 'text-exo-light-gray'}"
       >
-        <svg
-          class="w-5 h-5 {sidebarVisible
-            ? 'text-exo-yellow'
-            : 'text-exo-light-gray'}"
-          fill="none"
-          viewBox="0 0 24 24"
-          stroke="currentColor"
-          stroke-width="2"
-        >
-          {#if sidebarVisible}
-            <path
-              stroke-linecap="round"
-              stroke-linejoin="round"
-              d="M11 19l-7-7 7-7m8 14l-7-7 7-7"
-            />
-          {:else}
-            <path
-              stroke-linecap="round"
-              stroke-linejoin="round"
-              d="M13 5l7 7-7 7M5 5l7 7-7 7"
-            />
-          {/if}
-        </svg>
-      </button>
-    </div>
-  {/if}
+        {#if mobileMenuOpen}
+          <path
+            stroke-linecap="round"
+            stroke-linejoin="round"
+            d="M11 19l-7-7 7-7m8 14l-7-7 7-7"
+          ></path>
+        {:else}
+          <path
+            stroke-linecap="round"
+            stroke-linejoin="round"
+            d="M13 5l7 7-7 7M5 5l7 7-7 7"
+          ></path>
+        {/if}
+      </svg>
+    </button>
+    <!-- Desktop sidebar toggle -->
+    <button
+      onclick={handleToggleSidebar}
+      class="p-2 rounded border border-exo-light-gray/30 hover:border-exo-yellow/50 hover:bg-exo-medium-gray/30 transition-colors cursor-pointer hidden md:block"
+      title={sidebarVisible ? "Hide sidebar" : "Show sidebar"}
+      aria-label={sidebarVisible
+        ? "Hide conversation sidebar"
+        : "Show conversation sidebar"}
+      aria-pressed={sidebarVisible}
+    >
+      <svg
+        fill="none"
+        viewBox="0 0 24 24"
+        stroke="currentColor"
+        stroke-width="2"
+        class="w-5 h-5 {sidebarVisible
+          ? 'text-exo-yellow'
+          : 'text-exo-light-gray'}"
+      >
+        {#if sidebarVisible}
+          <path
+            stroke-linecap="round"
+            stroke-linejoin="round"
+            d="M11 19l-7-7 7-7m8 14l-7-7 7-7"
+          ></path>
+        {:else}
+          <path
+            stroke-linecap="round"
+            stroke-linejoin="round"
+            d="M13 5l7 7-7 7M5 5l7 7-7 7"
+          ></path>
+        {/if}
+      </svg>
+    </button>
+  </div>
 
   <!-- Center: Logo (clickable to go home) -->
   <button
@@ -83,19 +153,55 @@
     <img
       src="/exo-logo.png"
       alt="EXO"
-      class="h-18 drop-shadow-[0_0_4px_rgba(255,215,0,0.3)]"
+      class="h-12 md:h-18 drop-shadow-[0_0_4px_rgba(255,215,0,0.3)]"
     />
   </button>
 
-  <!-- Right: Home + Downloads -->
+  <!-- Right: Home + Downloads + Mobile Right Toggle -->
   <nav
-    class="absolute right-6 top-1/2 -translate-y-1/2 flex items-center gap-4"
+    class="absolute right-4 md:right-6 top-1/2 -translate-y-1/2 flex items-center gap-2 md:gap-4"
     aria-label="Main navigation"
   >
+    <!-- Mobile right sidebar toggle (instances/models) - only show when not in chat mode -->
+    {#if showMobileRightToggle}
+      <button
+        onclick={handleToggleMobileRight}
+        class="p-2 rounded border border-exo-light-gray/30 hover:border-exo-yellow/50 hover:bg-exo-medium-gray/30 transition-colors cursor-pointer md:hidden"
+        title={mobileRightOpen ? "Hide instances" : "Show instances"}
+        aria-label={mobileRightOpen
+          ? "Hide instances panel"
+          : "Show instances panel"}
+        aria-pressed={mobileRightOpen}
+      >
+        <svg
+          fill="none"
+          viewBox="0 0 24 24"
+          stroke="currentColor"
+          stroke-width="2"
+          class="w-5 h-5 {mobileRightOpen
+            ? 'text-exo-yellow'
+            : 'text-exo-light-gray'}"
+        >
+          {#if mobileRightOpen}
+            <path
+              stroke-linecap="round"
+              stroke-linejoin="round"
+              d="M13 5l7 7-7 7M5 5l7 7-7 7"
+            ></path>
+          {:else}
+            <path
+              stroke-linecap="round"
+              stroke-linejoin="round"
+              d="M11 19l-7-7 7-7m8 14l-7-7 7-7"
+            ></path>
+          {/if}
+        </svg>
+      </button>
+    {/if}
     {#if showHome}
       <button
         onclick={handleHome}
-        class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
+        class="flex text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase items-center gap-2 cursor-pointer"
         title="Back to topology view"
       >
         <svg
@@ -111,12 +217,12 @@
             d="M3 12l2-2m0 0l7-7 7 7M5 10v10a1 1 0 001 1h3m10-11l2 2m-2-2v10a1 1 0 01-1 1h-3m-6 0a1 1 0 001-1v-4a1 1 0 011-1h2a1 1 0 011 1v4a1 1 0 001 1m-6 0h6"
           />
         </svg>
-        Home
+        <span class="hidden sm:inline">Home</span>
       </button>
     {/if}
     <a
       href="/#/downloads"
-      class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
+      class="text-xs md:text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-1.5 md:gap-2 cursor-pointer"
       title="View downloads overview"
     >
       {#if downloadProgress}
@@ -168,7 +274,28 @@
           <path d="M5 21h14" />
         </svg>
       {/if}
-      Downloads
+      <span class="hidden sm:inline">Downloads</span>
+    </a>
+    <a
+      href="/#/integrations"
+      class="text-xs md:text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-1.5 md:gap-2 cursor-pointer"
+      title="Integration configs for external tools"
+    >
+      <svg
+        class="w-4 h-4"
+        viewBox="0 0 24 24"
+        fill="none"
+        stroke="currentColor"
+        stroke-width="2"
+        stroke-linecap="round"
+        stroke-linejoin="round"
+      >
+        <path d="M10 13a5 5 0 0 0 7.54.54l3-3a5 5 0 0 0-7.07-7.07l-1.72 1.71" />
+        <path
+          d="M14 11a5 5 0 0 0-7.54-.54l-3 3a5 5 0 0 0 7.07 7.07l1.71-1.71"
+        />
+      </svg>
+      <span class="hidden sm:inline">Integrations</span>
     </a>
   </nav>
 </header>

dashboard/src/lib/components/HuggingFaceResultItem.svelte

@@ -36,8 +36,12 @@
     return num.toString();
   }
 
-  // Extract model name from full ID (e.g., "mlx-community/Llama-3.2-1B" -> "Llama-3.2-1B")
-  const modelName = $derived(model.id.split("/").pop() || model.id);
+  // Show short name for mlx-community models, full ID for everything else
+  const modelName = $derived(
+    model.author === "mlx-community"
+      ? model.id.split("/").pop() || model.id
+      : model.id,
+  );
 </script>
 
 <div

dashboard/src/lib/components/IntegrationCard.svelte

@@ -0,0 +1,52 @@
+<script lang="ts">
+  interface Props {
+    title: string;
+    subtitle: string;
+    config: string;
+    description?: string;
+    language?: "json" | "bash";
+  }
+
+  let {
+    title,
+    subtitle,
+    config,
+    description = "",
+    language = "json",
+  }: Props = $props();
+
+  let copied = $state(false);
+
+  async function copyToClipboard() {
+    await navigator.clipboard.writeText(config);
+    copied = true;
+    setTimeout(() => (copied = false), 2000);
+  }
+</script>
+
+<div
+  class="border border-exo-light-gray/20 rounded-lg bg-exo-medium-gray/20 overflow-hidden"
+>
+  <div class="flex items-center justify-between px-5 py-4">
+    <div>
+      <h3 class="text-white text-sm font-semibold tracking-wide">{title}</h3>
+      <p class="text-exo-light-gray/60 text-xs mt-0.5 font-mono">{subtitle}</p>
+    </div>
+    <button
+      onclick={copyToClipboard}
+      class="px-3 py-1.5 text-xs rounded border transition-all duration-200 cursor-pointer
+        {copied
+        ? 'border-green-500/50 text-green-400 bg-green-500/10'
+        : 'border-exo-light-gray/30 text-exo-light-gray hover:border-exo-yellow/50 hover:text-exo-yellow'}"
+    >
+      {copied ? "Copied!" : "Copy"}
+    </button>
+  </div>
+  {#if description}
+    <p class="text-exo-light-gray/70 text-xs px-5 pb-3">{description}</p>
+  {/if}
+  <div class="bg-black/30 border-t border-exo-light-gray/10">
+    <pre
+      class="text-xs text-exo-light-gray/90 font-mono p-4 overflow-x-auto whitespace-pre">{config}</pre>
+  </div>
+</div>

dashboard/src/lib/components/MarkdownContent.svelte

@@ -507,9 +507,29 @@
   });
 
   $effect(() => {
-    if (containerRef && processedHtml) {
-      setupCopyButtons();
+    if (!containerRef || !browser) return;
+
+    function handleDelegatedClick(event: MouseEvent) {
+      const codeBtn = (event.target as HTMLElement).closest(
+        ".copy-code-btn",
+      ) as HTMLButtonElement | null;
+      if (codeBtn) {
+        handleCopyClick({ currentTarget: codeBtn } as unknown as Event);
+        return;
+      }
+      const mathBtn = (event.target as HTMLElement).closest(
+        ".copy-math-btn",
+      ) as HTMLButtonElement | null;
+      if (mathBtn) {
+        handleMathCopyClick({ currentTarget: mathBtn } as unknown as Event);
+        return;
+      }
     }
+
+    containerRef.addEventListener("click", handleDelegatedClick);
+    return () => {
+      containerRef?.removeEventListener("click", handleDelegatedClick);
+    };
   });
 </script>