Bake scale into SDPA pass1 source, remove last scalar input

This reverts commit 97d7f03032.

.mlx_typings/mlx_lm/generate.pyi

@@ -73,9 +73,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,
@@ -252,7 +254,14 @@ class BatchResponse:
 
     texts: List[str]
     stats: BatchStats
+    caches: Optional[List[List[Any]]]
 
+def _left_pad_prompts(prompts: Any, max_length: Optional[int] = ...) -> mx.array: ...
+def _right_pad_prompts(prompts: Any, max_length: Optional[int] = ...) -> mx.array: ...
+def _make_cache(
+    model: Any, left_padding: Any, max_kv_size: Optional[int]
+) -> List[Any]: ...
+def _merge_caches(caches: Any) -> List[Any]: ...
 @dataclass
 class Batch:
     uids: List[int]
@@ -261,39 +270,71 @@ class Batch:
     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:
-        ...
+    samplers: List[Any]
+    logits_processors: List[Any]
+    tokens: List[mx.array]
+    def __len__(self) -> int: ...
+    def filter(self, keep_idx: List[int]) -> None: ...
+    def extend(self, other: "Batch") -> None: ...
+    def extract_cache(self, idx: int) -> List[Any]: ...
 
 class BatchGenerator:
+    model: Any
+    max_kv_size: Optional[int]
+    prefill_step_size: int
+    unprocessed_prompts: List[Any]
+    active_batch: Optional[Batch]
+    prompt_progress_callback: Callable[[List[Tuple[int, int, int]]], None]
+    _stats: BatchStats
+
     @dataclass
     class Response:
         uid: int
         token: int
         logprobs: mx.array
         finish_reason: Optional[str]
+        prompt_cache: Any
 
     def __init__(
         self,
-        model,
+        model: Any,
         max_tokens: int = ...,
-        stop_tokens: Optional[set] = ...,
+        stop_tokens: Optional[set[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 = ...,
+        prompt_progress_callback: Optional[
+            Callable[[List[Tuple[int, int, int]]], None]
+        ] = ...,
+        max_kv_size: Optional[int] = ...,
     ) -> None: ...
+    def close(self) -> None: ...
     def insert(
-        self, prompts, max_tokens: Union[List[int], int, None] = ...
-    ):  # -> list[Any]:
-        ...
-    def stats(self):  # -> BatchStats:
-        ...
-    def next(self):  # -> list[Any]:
-        ...
+        self,
+        prompts: Any,
+        max_tokens: Union[List[int], int, None] = ...,
+        caches: Any = ...,
+        samplers: Optional[List[Any]] = ...,
+        logits_processors: Optional[List[Any]] = ...,
+    ) -> List[int]: ...
+    def remove(
+        self, uids: List[int], return_prompt_caches: bool = ...
+    ) -> Optional[dict[int, List[Any]]]: ...
+    def stats(self) -> BatchStats: ...
+    def next(self) -> List[Response]: ...
+    def _process_prompts(self, prompts: List[Any]) -> Batch: ...
+    def _step(
+        self,
+        input_tokens: mx.array,
+        prompt_cache: List[Any],
+        samplers: Optional[List[Any]],
+        logits_processors: Optional[List[Any]],
+        tokens: List[mx.array],
+    ) -> Tuple[mx.array, List[mx.array]]: ...
 
 def batch_generate(
     model,

.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: ...
 
@@ -92,13 +92,14 @@ class _BaseCache(Cache):
     values: mx.array
     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:
         ...

.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

@@ -7,6 +7,15 @@ import mlx.nn as nn
 
 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

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

CONTRIBUTING.md

@@ -39,6 +39,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
 
@@ -196,6 +198,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 +279,47 @@ 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_MODELS_PATH` | Colon-separated paths to search for pre-downloaded models (e.g., on NFS mounts or shared storage) | None |
+| `EXO_MODELS_DIR` | Directory where exo downloads and stores models | `~/.local/share/exo/models` (Linux) or `~/.exo/models` (macOS) |
+| `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
+EXO_MODELS_PATH=/mnt/nfs/models:/opt/ai-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 +409,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 +546,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.

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

@@ -24,6 +24,7 @@ import json
 import sys
 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
@@ -252,6 +253,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,
@@ -282,6 +289,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 +304,9 @@ 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)
+    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:
@@ -392,52 +405,123 @@ 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]] = []
+                    for r in range(args.repeat):
+                        time.sleep(3)
+
+                        if concurrency <= 1:
+                            # Sequential: single request
+                            try:
+                                row, actual_pp_tokens = run_one_completion(
+                                    client, full_model_id, pp, tg, prompt_sizer
+                                )
+                            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
+                            # Each thread gets its own ExoClient (separate HTTP connection)
+                            batch_results: list[tuple[dict[str, Any], int]] = []
+                            batch_errors = 0
+
+                            def _run_concurrent(
+                                idx: int, *, _pp: int = pp, _tg: int = tg
+                            ) -> tuple[dict[str, Any], int]:
+                                c = ExoClient(
+                                    args.host, args.port, timeout_s=args.timeout
+                                )
+                                return run_one_completion(
+                                    c, full_model_id, _pp, _tg, prompt_sizer
+                                )
+
+                            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
+
+                            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:
+                                agg_gen_tps = sum(
+                                    x["stats"]["generation_tps"]
+                                    for x, _ in batch_results
+                                    if x["stats"]["generation_tps"] > 0
+                                )
+                                logger.info(
+                                    f"[concurrent {concurrency}x]  "
+                                    f"agg_gen_tps={agg_gen_tps:.2f}  "
+                                    f"errors={batch_errors}"
+                                )
+
+                    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
                         )
-                    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)
+                        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)
         finally:
             try:
                 client.request_json("DELETE", f"/instance/{instance_id}")

bench/harness.py

@@ -165,7 +165,9 @@ 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 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,6 +183,16 @@ 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}")
 
 
@@ -445,6 +457,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,

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

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

dashboard/src/lib/components/ModelPickerGroup.svelte

@@ -32,6 +32,7 @@
     group: ModelGroup;
     isExpanded: boolean;
     isFavorite: boolean;
+    isHighlighted?: boolean;
     selectedModelId: string | null;
     canModelFit: (id: string) => boolean;
     getModelFitStatus: (id: string) => ModelFitStatus;
@@ -48,6 +49,7 @@
     group,
     isExpanded,
     isFavorite,
+    isHighlighted = false,
     selectedModelId,
     canModelFit,
     getModelFitStatus,
@@ -150,10 +152,11 @@
 </script>
 
 <div
+  data-model-ids={group.variants.map((v) => v.id).join(" ")}
   class="border-b border-white/5 last:border-b-0 {!anyVariantFits &&
   !anyVariantHasInstance
     ? 'opacity-50'
-    : ''}"
+    : ''} {isHighlighted ? 'model-just-added' : ''}"
 >
   <!-- Main row -->
   <div
@@ -644,3 +647,21 @@
     </div>
   {/if}
 </div>
+
+<style>
+  .model-just-added {
+    animation: highlightFade 4s ease-out forwards;
+  }
+
+  @keyframes highlightFade {
+    0%,
+    40% {
+      background-color: rgba(20, 83, 45, 0.25);
+      box-shadow: inset 0 0 0 1px rgba(74, 222, 128, 0.4);
+    }
+    100% {
+      background-color: transparent;
+      box-shadow: none;
+    }
+  }
+</style>

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -1,4 +1,5 @@
 <script lang="ts">
+  import { tick } from "svelte";
   import { fade, fly } from "svelte/transition";
   import { cubicOut } from "svelte/easing";
   import FamilySidebar from "./FamilySidebar.svelte";
@@ -7,6 +8,7 @@
   import HuggingFaceResultItem from "./HuggingFaceResultItem.svelte";
   import { getNodesWithModelDownloaded } from "$lib/utils/downloads";
   import { getRecentEntries } from "$lib/stores/recents.svelte";
+  import { addToast } from "$lib/stores/toast.svelte";
 
   interface ModelInfo {
     id: string;
@@ -191,6 +193,13 @@
   let hfSearchDebounceTimer: ReturnType<typeof setTimeout> | null = null;
   let manualModelId = $state("");
   let addModelError = $state<string | null>(null);
+  let justAddedModelId = $state<string | null>(null);
+  let justAddedTimer: ReturnType<typeof setTimeout> | null = null;
+
+  // Inline HuggingFace search in main search bar
+  let mainSearchHfResults = $state<HuggingFaceModel[]>([]);
+  let mainSearchHfLoading = $state(false);
+  let mainSearchDebounceTimer: ReturnType<typeof setTimeout> | null = null;
 
   // Reset transient state when modal opens, but preserve tab selection
   $effect(() => {
@@ -200,6 +209,11 @@
       showFilters = false;
       manualModelId = "";
       addModelError = null;
+      justAddedModelId = null;
+      if (justAddedTimer) {
+        clearTimeout(justAddedTimer);
+        justAddedTimer = null;
+      }
     }
   });
 
@@ -214,6 +228,50 @@
     }
   });
 
+  // Inline HuggingFace search when local search returns no results
+  $effect(() => {
+    const query = searchQuery.trim();
+    const noLocalResults = filteredGroups.length === 0;
+
+    if (mainSearchDebounceTimer) {
+      clearTimeout(mainSearchDebounceTimer);
+      mainSearchDebounceTimer = null;
+    }
+
+    if (
+      selectedFamily === "huggingface" ||
+      selectedFamily === "recents" ||
+      selectedFamily === "favorites" ||
+      query.length < 2 ||
+      !noLocalResults
+    ) {
+      mainSearchHfResults = [];
+      mainSearchHfLoading = false;
+      return;
+    }
+
+    mainSearchHfLoading = true;
+    mainSearchDebounceTimer = setTimeout(async () => {
+      try {
+        const response = await fetch(
+          `/models/search?query=${encodeURIComponent(query)}&limit=10`,
+        );
+        if (response.ok) {
+          const results: HuggingFaceModel[] = await response.json();
+          mainSearchHfResults = results.filter(
+            (r) => !existingModelIds.has(r.id),
+          );
+        } else {
+          mainSearchHfResults = [];
+        }
+      } catch {
+        mainSearchHfResults = [];
+      } finally {
+        mainSearchHfLoading = false;
+      }
+    }, 500);
+  });
+
   async function fetchTrendingModels() {
     hfIsLoadingTrending = true;
     try {
@@ -274,6 +332,24 @@
     addModelError = null;
     try {
       await onAddModel(modelId);
+      // Success: show toast, switch to All Models, highlight the model
+      const shortName = modelId.split("/").pop() || modelId;
+      addToast({ type: "success", message: `Added ${shortName}` });
+      justAddedModelId = modelId;
+      selectedFamily = null;
+      searchQuery = "";
+      // Scroll to the newly added model after DOM update
+      await tick();
+      const el = document.querySelector(
+        `[data-model-ids~="${CSS.escape(modelId)}"]`,
+      );
+      el?.scrollIntoView({ behavior: "smooth", block: "center" });
+      // Clear highlight after 4 seconds
+      if (justAddedTimer) clearTimeout(justAddedTimer);
+      justAddedTimer = setTimeout(() => {
+        justAddedModelId = null;
+        justAddedTimer = null;
+      }, 4000);
     } catch (error) {
       addModelError =
         error instanceof Error ? error.message : "Failed to add model";
@@ -841,6 +917,8 @@
                 {group}
                 isExpanded={expandedGroups.has(group.id)}
                 isFavorite={favorites.has(group.id)}
+                isHighlighted={justAddedModelId !== null &&
+                  group.variants.some((v) => v.id === justAddedModelId)}
                 {selectedModelId}
                 {canModelFit}
                 {getModelFitStatus}
@@ -910,6 +988,8 @@
               {group}
               isExpanded={expandedGroups.has(group.id)}
               isFavorite={favorites.has(group.id)}
+              isHighlighted={justAddedModelId !== null &&
+                group.variants.some((v) => v.id === justAddedModelId)}
               {selectedModelId}
               {canModelFit}
               {getModelFitStatus}
@@ -937,6 +1017,8 @@
               {group}
               isExpanded={expandedGroups.has(group.id)}
               isFavorite={favorites.has(group.id)}
+              isHighlighted={justAddedModelId !== null &&
+                group.variants.some((v) => v.id === justAddedModelId)}
               {selectedModelId}
               {canModelFit}
               {getModelFitStatus}
@@ -948,6 +1030,55 @@
               {instanceStatuses}
             />
           {/each}
+          <!-- Inline HuggingFace search results (shown when no local results match) -->
+          {#if filteredGroups.length === 0 && searchQuery.trim().length >= 2 && selectedFamily !== "huggingface" && selectedFamily !== "recents" && selectedFamily !== "favorites"}
+            {#if mainSearchHfLoading}
+              <div
+                class="flex items-center gap-2 px-3 py-2 border-t border-orange-400/20 bg-orange-950/20"
+              >
+                <span
+                  class="w-4 h-4 border-2 border-orange-400 border-t-transparent rounded-full animate-spin"
+                ></span>
+                <span class="text-xs font-mono text-orange-400/60"
+                  >Searching HuggingFace...</span
+                >
+              </div>
+            {:else if mainSearchHfResults.length > 0}
+              <div
+                class="sticky top-0 z-10 flex items-center gap-2 px-3 py-2 bg-orange-950/30 border-y border-orange-400/20 backdrop-blur-sm"
+              >
+                <span
+                  class="text-xs font-mono text-orange-400 tracking-wider uppercase"
+                  >From HuggingFace</span
+                >
+              </div>
+              {#each mainSearchHfResults as model}
+                <HuggingFaceResultItem
+                  {model}
+                  isAdded={existingModelIds.has(model.id)}
+                  isAdding={addingModelId === model.id}
+                  onAdd={() => handleAddModel(model.id)}
+                  onSelect={() => handleSelectHfModel(model.id)}
+                  downloadedOnNodes={downloadsData
+                    ? getNodesWithModelDownloaded(downloadsData, model.id).map(
+                        getNodeName,
+                      )
+                    : []}
+                />
+              {/each}
+              <button
+                type="button"
+                class="w-full px-3 py-2 text-xs font-mono text-orange-400/60 hover:text-orange-400 hover:bg-orange-500/10 transition-colors text-center"
+                onclick={() => {
+                  hfSearchQuery = searchQuery;
+                  searchHuggingFace(searchQuery);
+                  selectedFamily = "huggingface";
+                }}
+              >
+                See all results on Hub
+              </button>
+            {/if}
+          {/if}
         {/if}
       </div>
     </div>

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

@@ -3158,6 +3158,23 @@ class AppStore {
     return (await response.json()) as TraceStatsResponse;
   }
 
+  /**
+   * Delete traces by task IDs
+   */
+  async deleteTraces(
+    taskIds: string[],
+  ): Promise<{ deleted: string[]; notFound: string[] }> {
+    const response = await fetch("/v1/traces/delete", {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify({ taskIds }),
+    });
+    if (!response.ok) {
+      throw new Error(`Failed to delete traces: ${response.status}`);
+    }
+    return await response.json();
+  }
+
   /**
    * Get the URL for the raw trace file (for Perfetto)
    */
@@ -3301,3 +3318,5 @@ export const fetchTraceStats = (taskId: string) =>
   appStore.fetchTraceStats(taskId);
 export const getTraceRawUrl = (taskId: string) =>
   appStore.getTraceRawUrl(taskId);
+export const deleteTraces = (taskIds: string[]) =>
+  appStore.deleteTraces(taskIds);

dashboard/src/routes/+page.svelte

@@ -6050,12 +6050,11 @@
               </div>
 
               <div
-                class="relative aspect-square bg-exo-dark-gray rounded-lg overflow-hidden"
+                class="relative aspect-square bg-exo-dark-gray rounded-lg overflow-hidden pointer-events-none"
               >
                 <TopologyGraph
                   highlightedNodes={highlightedNodes()}
                   filteredNodes={nodeFilter}
-                  onNodeClick={togglePreviewNodeFilter}
                 />
 
                 {@render clusterWarningsCompact()}

dashboard/src/routes/traces/+page.svelte

@@ -3,6 +3,7 @@
   import {
     listTraces,
     getTraceRawUrl,
+    deleteTraces,
     type TraceListItem,
   } from "$lib/stores/app.svelte";
   import HeaderNav from "$lib/components/HeaderNav.svelte";
@@ -10,6 +11,51 @@
   let traces = $state<TraceListItem[]>([]);
   let loading = $state(true);
   let error = $state<string | null>(null);
+  let selectedIds = $state<Set<string>>(new Set());
+  let deleting = $state(false);
+
+  let allSelected = $derived(
+    traces.length > 0 && selectedIds.size === traces.length,
+  );
+
+  function toggleSelect(taskId: string) {
+    const next = new Set(selectedIds);
+    if (next.has(taskId)) {
+      next.delete(taskId);
+    } else {
+      next.add(taskId);
+    }
+    selectedIds = next;
+  }
+
+  function toggleSelectAll() {
+    if (allSelected) {
+      selectedIds = new Set();
+    } else {
+      selectedIds = new Set(traces.map((t) => t.taskId));
+    }
+  }
+
+  async function handleDelete() {
+    if (selectedIds.size === 0) return;
+    const count = selectedIds.size;
+    if (
+      !confirm(
+        `Delete ${count} trace${count === 1 ? "" : "s"}? This cannot be undone.`,
+      )
+    )
+      return;
+    deleting = true;
+    try {
+      await deleteTraces([...selectedIds]);
+      selectedIds = new Set();
+      await refresh();
+    } catch (e) {
+      error = e instanceof Error ? e.message : "Failed to delete traces";
+    } finally {
+      deleting = false;
+    }
+  }
 
   function formatBytes(bytes: number): string {
     if (!bytes || bytes <= 0) return "0B";
@@ -109,6 +155,16 @@
         </h1>
       </div>
       <div class="flex items-center gap-3">
+        {#if selectedIds.size > 0}
+          <button
+            type="button"
+            class="text-xs font-mono text-red-400 hover:text-red-300 transition-colors uppercase border border-red-500/40 px-2 py-1 rounded"
+            onclick={handleDelete}
+            disabled={deleting}
+          >
+            {deleting ? "Deleting..." : `Delete (${selectedIds.size})`}
+          </button>
+        {/if}
         <button
           type="button"
           class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-2 py-1 rounded"
@@ -143,14 +199,41 @@
       </div>
     {:else}
       <div class="space-y-3">
+        <div class="flex items-center gap-2 px-1">
+          <button
+            type="button"
+            class="text-xs font-mono uppercase transition-colors {allSelected
+              ? 'text-exo-yellow'
+              : 'text-exo-light-gray hover:text-exo-yellow'}"
+            onclick={toggleSelectAll}
+          >
+            {allSelected ? "Deselect all" : "Select all"}
+          </button>
+        </div>
         {#each traces as trace}
+          {@const isSelected = selectedIds.has(trace.taskId)}
+          <!-- svelte-ignore a11y_no_static_element_interactions -->
           <div
-            class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-4 flex items-center justify-between gap-4"
+            role="button"
+            tabindex="0"
+            class="w-full text-left rounded border-l-2 border-r border-t border-b transition-all p-4 flex items-center justify-between gap-4 cursor-pointer {isSelected
+              ? 'bg-exo-yellow/10 border-l-exo-yellow border-r-exo-medium-gray/30 border-t-exo-medium-gray/30 border-b-exo-medium-gray/30'
+              : 'bg-exo-black/30 border-l-transparent border-r-exo-medium-gray/30 border-t-exo-medium-gray/30 border-b-exo-medium-gray/30 hover:bg-white/[0.03]'}"
+            onclick={() => toggleSelect(trace.taskId)}
+            onkeydown={(e) => {
+              if (e.key === "Enter" || e.key === " ") {
+                e.preventDefault();
+                toggleSelect(trace.taskId);
+              }
+            }}
           >
             <div class="min-w-0 flex-1">
               <a
                 href="#/traces/{trace.taskId}"
-                class="text-sm font-mono text-white hover:text-exo-yellow transition-colors truncate block"
+                class="text-sm font-mono transition-colors truncate block {isSelected
+                  ? 'text-exo-yellow'
+                  : 'text-white hover:text-exo-yellow'}"
+                onclick={(e) => e.stopPropagation()}
               >
                 {trace.taskId}
               </a>
@@ -160,7 +243,11 @@
                 )}
               </div>
             </div>
-            <div class="flex items-center gap-2 shrink-0">
+            <!-- svelte-ignore a11y_click_events_have_key_events -->
+            <div
+              class="flex items-center gap-2 shrink-0"
+              onclick={(e) => e.stopPropagation()}
+            >
               <a
                 href="#/traces/{trace.taskId}"
                 class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-2 py-1 rounded"

docs/api.md

@@ -4,7 +4,7 @@ This document describes the REST API exposed by the **EXO** service, as implemen
 
 `src/exo/master/api.py`
 
-The API is used to manage model instances in the cluster, inspect cluster state, and perform inference using an OpenAI-compatible interface.
+The API is used to manage model instances in the cluster, inspect cluster state, and perform inference using multiple API-compatible interfaces.
 
 Base URL example:
 
@@ -144,8 +144,59 @@ Placement result.
 
 Returns the list of available models and their metadata.
 
+**Query parameters:**
+
+* `status`: string (optional) - Filter by `downloaded` to show only downloaded models
+
 **Response:**
-Array of model descriptors.
+Array of model descriptors including `is_custom` field for custom HuggingFace models.
+
+### Add Custom Model
+
+**POST** `/models/add`
+
+Add a custom model from HuggingFace hub.
+
+**Request body (example):**
+
+```json
+{
+  "model_id": "mlx-community/my-custom-model"
+}
+```
+
+**Response:**
+Model descriptor for the added model.
+
+**Security note:**
+Models with `trust_remote_code` enabled in their configuration require explicit opt-in (default is false) for security.
+
+### Delete Custom Model
+
+**DELETE** `/models/custom/{model_id}`
+
+Delete a user-added custom model card.
+
+**Path parameters:**
+
+* `model_id`: string, ID of the custom model to delete
+
+**Response:**
+Confirmation JSON with deleted model ID.
+
+### Search Models
+
+**GET** `/models/search`
+
+Search HuggingFace Hub for mlx-community models.
+
+**Query parameters:**
+
+* `query`: string (optional) - Search query
+* `limit`: integer (default: 20) - Maximum number of results
+
+**Response:**
+Array of HuggingFace model search results.
 
 ## 4. Inference / Chat Completions
 
@@ -168,9 +219,123 @@ Executes a chat completion request using an OpenAI-compatible schema. Supports s
 }
 ```
 
+**Request parameters:**
+
+* `model`: string, required - Model ID to use
+* `messages`: array, required - Conversation messages
+* `stream`: boolean (default: false) - Enable streaming responses
+* `max_tokens`: integer (optional) - Maximum tokens to generate
+* `temperature`: float (optional) - Sampling temperature
+* `top_p`: float (optional) - Nucleus sampling parameter
+* `top_k`: integer (optional) - Top-k sampling parameter
+* `stop`: string or array (optional) - Stop sequences
+* `seed`: integer (optional) - Random seed for reproducibility
+* `enable_thinking`: boolean (optional) - Enable thinking mode for capable models (DeepSeek V3.1, Qwen3, GLM-4.7)
+* `tools`: array (optional) - Tool definitions for function calling
+* `logprobs`: boolean (optional) - Return log probabilities
+* `top_logprobs`: integer (optional) - Number of top log probabilities to return
+
 **Response:**
 OpenAI-compatible chat completion response.
 
+**Streaming response format:**
+When `stream=true`, returns Server-Sent Events (SSE) with format:
+
+```
+data: {"id":"...","object":"chat.completion","created":...,"model":"...","choices":[...]}
+
+data: [DONE]
+```
+
+**Non-streaming response includes usage statistics:**
+
+```json
+{
+  "id": "...",
+  "object": "chat.completion",
+  "created": 1234567890,
+  "model": "llama-3.2-1b",
+  "choices": [{
+    "index": 0,
+    "message": {
+      "role": "assistant",
+      "content": "Hello! How can I help you?"
+    },
+    "finish_reason": "stop"
+  }],
+  "usage": {
+    "prompt_tokens": 15,
+    "completion_tokens": 8,
+    "total_tokens": 23
+  }
+}
+```
+
+**Cancellation:**
+You can cancel an active generation by closing the HTTP connection. The server detects the disconnection and stops processing.
+
+### Claude Messages API
+
+**POST** `/v1/messages`
+
+Executes a chat completion request using the Claude Messages API format. Supports streaming and non-streaming modes.
+
+**Request body (example):**
+
+```json
+{
+  "model": "llama-3.2-1b",
+  "messages": [
+    { "role": "user", "content": "Hello" }
+  ],
+  "max_tokens": 1024,
+  "stream": false
+}
+```
+
+**Streaming response format:**
+When `stream=true`, returns Server-Sent Events with Claude-specific event types:
+
+* `message_start` - Message generation started
+* `content_block_start` - Content block started
+* `content_block_delta` - Incremental content chunk
+* `content_block_stop` - Content block completed
+* `message_delta` - Message metadata updates
+* `message_stop` - Message generation completed
+
+**Response:**
+Claude-compatible messages response.
+
+### OpenAI Responses API
+
+**POST** `/v1/responses`
+
+Executes a chat completion request using the OpenAI Responses API format. Supports streaming and non-streaming modes.
+
+**Request body (example):**
+
+```json
+{
+  "model": "llama-3.2-1b",
+  "messages": [
+    { "role": "user", "content": "Hello" }
+  ],
+  "stream": false
+}
+```
+
+**Streaming response format:**
+When `stream=true`, returns Server-Sent Events with response-specific event types:
+
+* `response.created` - Response generation started
+* `response.in_progress` - Response is being generated
+* `response.output_item.added` - New output item added
+* `response.output_item.done` - Output item completed
+* `response.done` - Response generation completed
+
+**Response:**
+OpenAI Responses API-compatible response.
+
 ### Benchmarked Chat Completions
 
 **POST** `/bench/chat/completions`
@@ -181,26 +346,177 @@ Same as `/v1/chat/completions`, but also returns performance and generation stat
 Same schema as `/v1/chat/completions`.
 
 **Response:**
-Chat completion plus benchmarking metrics.
+Chat completion plus benchmarking metrics including:
 
-## 5. Image Generation & Editing
+* `prompt_tps` - Tokens per second during prompt processing
+* `generation_tps` - Tokens per second during generation
+* `prompt_tokens` - Number of prompt tokens
+* `generation_tokens` - Number of generated tokens
+* `peak_memory_usage` - Peak memory used during generation
+
+### Cancel Command
+
+**POST** `/v1/cancel/{command_id}`
+
+Cancels an active generation command (text or image). Notifies workers and closes the stream.
+
+**Path parameters:**
+
+* `command_id`: string, ID of the command to cancel
+
+**Response (example):**
+
+```json
+{
+  "message": "Command cancelled.",
+  "command_id": "cmd-abc-123"
+}
+```
+
+Returns 404 if the command is not found or already completed.
+
+## 5. Ollama API Compatibility
+
+EXO provides Ollama API compatibility for tools like OpenWebUI.
+
+### Ollama Chat
+
+**POST** `/ollama/api/chat`
+**POST** `/ollama/api/api/chat` (alias)
+**POST** `/ollama/api/v1/chat` (alias)
+
+Execute a chat request using Ollama API format.
+
+**Request body (example):**
+
+```json
+{
+  "model": "llama-3.2-1b",
+  "messages": [
+    { "role": "user", "content": "Hello" }
+  ],
+  "stream": false
+}
+```
+
+**Response:**
+Ollama-compatible chat response.
+
+### Ollama Generate
+
+**POST** `/ollama/api/generate`
+
+Execute a text generation request using Ollama API format.
+
+**Request body (example):**
+
+```json
+{
+  "model": "llama-3.2-1b",
+  "prompt": "Hello",
+  "stream": false
+}
+```
+
+**Response:**
+Ollama-compatible generation response.
+
+### Ollama Tags
+
+**GET** `/ollama/api/tags`
+**GET** `/ollama/api/api/tags` (alias)
+**GET** `/ollama/api/v1/tags` (alias)
+
+Returns list of downloaded models in Ollama tags format.
+
+**Response:**
+Array of model tags with metadata.
+
+### Ollama Show
+
+**POST** `/ollama/api/show`
+
+Returns model information in Ollama show format.
+
+**Request body:**
+
+```json
+{
+  "name": "llama-3.2-1b"
+}
+```
+
+**Response:**
+Model details including modelfile and family.
+
+### Ollama PS
+
+**GET** `/ollama/api/ps`
+
+Returns list of running models (active instances).
+
+**Response:**
+Array of active model instances.
+
+### Ollama Version
+
+**GET** `/ollama/api/version`
+**HEAD** `/ollama/` (alias)
+**HEAD** `/ollama/api/version` (alias)
+
+Returns version information for Ollama API compatibility.
+
+**Response:**
+
+```json
+{
+  "version": "exo v1.0"
+}
+```
+
+## 6. Image Generation & Editing
 
 ### Image Generation
 
 **POST** `/v1/images/generations`
 
-Executes an image generation request using an OpenAI-compatible schema with additional advanced_params.
+Executes an image generation request using an OpenAI-compatible schema with additional advanced_params. Supports both streaming and non-streaming modes.
 
 **Request body (example):**
 
 ```json
 {
   "prompt": "a robot playing chess",
-  "model": "flux-dev",
+  "model": "exolabs/FLUX.1-dev",
+  "n": 1,
+  "size": "1024x1024",
   "stream": false,
+  "response_format": "b64_json"
 }
 ```
 
+**Request parameters:**
+
+* `prompt`: string, required - Text description of the image
+* `model`: string, required - Image model ID
+* `n`: integer (default: 1) - Number of images to generate
+* `size`: string (default: "auto") - Image dimensions. Supported sizes:
+  - `512x512`
+  - `768x768`
+  - `1024x768`
+  - `768x1024`
+  - `1024x1024`
+  - `1024x1536`
+  - `1536x1024`
+  - `1024x1365`
+  - `1365x1024`
+* `stream`: boolean (default: false) - Enable streaming for partial images
+* `partial_images`: integer (default: 0) - Number of partial images to stream during generation
+* `response_format`: string (default: "b64_json") - Either `url` or `b64_json`
+* `quality`: string (default: "medium") - Either `high`, `medium`, or `low`
+* `output_format`: string (default: "png") - Either `png`, `jpeg`, or `webp`
+* `advanced_params`: object (optional) - Advanced generation parameters
+
 **Advanced Parameters (`advanced_params`):**
 
 | Parameter | Type | Constraints | Description |
@@ -210,8 +526,54 @@ Executes an image generation request using an OpenAI-compatible schema with addi
 | `guidance` | float | 1.0-20.0 | Classifier-free guidance scale |
 | `negative_prompt` | string | - | Text describing what to avoid in the image |
 
+**Non-streaming response:**
+
+```json
+{
+  "created": 1234567890,
+  "data": [
+    {
+      "b64_json": "iVBORw0KGgoAAAANSUhEUgAA...",
+      "url": null
+    }
+  ]
+}
+```
+
+**Streaming response format:**
+When `stream=true` and `partial_images > 0`, returns Server-Sent Events:
+
+```
+data: {"type":"partial","image_index":0,"partial_index":1,"total_partials":5,"format":"png","data":{"b64_json":"..."}}
+
+data: {"type":"final","image_index":0,"format":"png","data":{"b64_json":"..."}}
+
+data: [DONE]
+```
+
+### Image Editing
+
+**POST** `/v1/images/edits`
+
+Executes an image editing request (img2img) using FLUX.1-Kontext-dev or similar models.
+
+**Request (multipart/form-data):**
+
+* `image`: file, required - Input image to edit
+* `prompt`: string, required - Text description of desired changes
+* `model`: string, required - Image editing model ID (e.g., `exolabs/FLUX.1-Kontext-dev`)
+* `n`: integer (default: 1) - Number of edited images to generate
+* `size`: string (optional) - Output image dimensions
+* `response_format`: string (default: "b64_json") - Either `url` or `b64_json`
+* `input_fidelity`: string (default: "low") - Either `low` or `high` - Controls how closely the output follows the input image
+* `stream`: string (default: "false") - Enable streaming
+* `partial_images`: string (default: "0") - Number of partial images to stream
+* `quality`: string (default: "medium") - Either `high`, `medium`, or `low`
+* `output_format`: string (default: "png") - Either `png`, `jpeg`, or `webp`
+* `advanced_params`: string (optional) - JSON-encoded advanced parameters
+
 **Response:**
-OpenAI-compatible image generation response.
+Same format as `/v1/images/generations`.
 
 ### Benchmarked Image Generation
 
@@ -223,16 +585,15 @@ Same as `/v1/images/generations`, but also returns generation statistics.
 Same schema as `/v1/images/generations`.
 
 **Response:**
-Image generation plus benchmarking metrics.
+Image generation plus benchmarking metrics including:
 
-### Image Editing
-
-**POST** `/v1/images/edits`
-
-Executes an image editing request using an OpenAI-compatible schema with additional advanced_params (same as `/v1/images/generations`).
-
-**Response:**
-Same format as `/v1/images/generations`.
+* `seconds_per_step` - Average time per denoising step
+* `total_generation_time` - Total generation time
+* `num_inference_steps` - Number of inference steps used
+* `num_images` - Number of images generated
+* `image_width` - Output image width
+* `image_height` - Output image height
+* `peak_memory_usage` - Peak memory used during generation
 
 ### Benchmarked Image Editing
 
@@ -246,36 +607,127 @@ Same schema as `/v1/images/edits`.
 **Response:**
 Same format as `/bench/images/generations`, including `generation_stats`.
 
-## 6. Complete Endpoint Summary
+### List Images
+
+**GET** `/images`
+
+List all stored images.
+
+**Response:**
+Array of image metadata including URLs and expiration times.
+
+### Get Image
+
+**GET** `/images/{image_id}`
+
+Retrieve a stored image by ID.
+
+**Path parameters:**
+
+* `image_id`: string, ID of the image
+
+**Response:**
+Image file with appropriate content type.
+
+## 7. Complete Endpoint Summary
 
 ```
+# General
 GET     /node_id
 GET     /state
 GET     /events
 
+# Instance Management
 POST    /instance
 GET     /instance/{instance_id}
 DELETE  /instance/{instance_id}
-
 GET     /instance/previews
 GET     /instance/placement
 POST    /place_instance
 
+# Models
 GET     /models
 GET     /v1/models
+POST    /models/add
+DELETE  /models/custom/{model_id}
+GET     /models/search
 
+# Text Generation (OpenAI Chat Completions)
 POST    /v1/chat/completions
 POST    /bench/chat/completions
 
+# Text Generation (Claude Messages API)
+POST    /v1/messages
+
+# Text Generation (OpenAI Responses API)
+POST    /v1/responses
+
+# Text Generation (Ollama API)
+POST    /ollama/api/chat
+POST    /ollama/api/api/chat
+POST    /ollama/api/v1/chat
+POST    /ollama/api/generate
+GET     /ollama/api/tags
+GET     /ollama/api/api/tags
+GET     /ollama/api/v1/tags
+POST    /ollama/api/show
+GET     /ollama/api/ps
+GET     /ollama/api/version
+HEAD    /ollama/
+HEAD    /ollama/api/version
+
+# Command Control
+POST    /v1/cancel/{command_id}
+
+# Image Generation
 POST    /v1/images/generations
 POST    /bench/images/generations
 POST    /v1/images/edits
 POST    /bench/images/edits
+GET     /images
+GET     /images/{image_id}
 ```
 
-## 7. Notes
+## 8. Notes
 
-* The `/v1/chat/completions` endpoint is compatible with the OpenAI Chat API format, so existing OpenAI clients can be pointed to EXO by changing the base URL.
-* The `/v1/images/generations` and `/v1/images/edits` endpoints are compatible with the OpenAI Images API format.
-* The instance placement endpoints allow you to plan and preview cluster allocations before actually creating instances.
-* The `/events` and `/state` endpoints are primarily intended for operational visibility and debugging.
+### API Compatibility
+
+EXO provides multiple API-compatible interfaces:
+
+* **OpenAI Chat Completions API** - Compatible with OpenAI clients and tools
+* **Claude Messages API** - Compatible with Anthropic's Claude API format
+* **OpenAI Responses API** - Compatible with OpenAI's Responses API format
+* **Ollama API** - Compatible with Ollama and tools like OpenWebUI
+
+Existing OpenAI, Claude, or Ollama clients can be pointed to EXO by changing the base URL.
+
+### Custom Models
+
+You can add custom models from HuggingFace using the `/models/add` endpoint. Custom models are identified by the `is_custom` field in model list responses.
+
+**Security:** Models requiring `trust_remote_code` must be explicitly enabled (default is false) for security. Only enable this if you trust the model's remote code.
+
+### Usage Statistics
+
+Chat completion responses include usage statistics with:
+
+* `prompt_tokens` - Number of tokens in the prompt
+* `completion_tokens` - Number of tokens generated
+* `total_tokens` - Sum of prompt and completion tokens
+
+### Request Cancellation
+
+You can cancel active requests by:
+
+1. Closing the HTTP connection (for streaming requests)
+2. Calling `/v1/cancel/{command_id}` (for any request)
+
+The server detects cancellation and stops processing immediately.
+
+### Instance Placement
+
+The instance placement endpoints allow you to plan and preview cluster allocations before creating instances. This helps optimize resource usage across nodes.
+
+### Observability
+
+The `/events` and `/state` endpoints are primarily intended for operational visibility and debugging.

docs/architecture.md

@@ -28,6 +28,26 @@ There are currently 5 major systems:
     
     Implements a distributed algorithm for master election in unstable networking conditions
 
+## API Layer
+
+The API system uses multiple adapters to support multiple API formats, converting them to a single request / response type.
+
+### Adapter Pattern
+
+Adapters convert between external API formats and EXO's internal types:
+
+```
+Chat Completions → [adapter] → TextGenerationTaskParams → Application
+Claude Messages  → [adapter] → TextGenerationTaskParams → Application
+Responses API    → [adapter] → TextGenerationTaskParams → Application
+Ollama API       → [adapter] → TextGenerationTaskParams → Application
+```
+
+Each adapter implements two key functions:
+1. **Request conversion**: Converts API-specific requests to `TextGenerationTaskParams`
+2. **Response generation**: Converts internal `TokenChunk` streams back to API-specific formats (streaming and non-streaming)
+
+
 ## Topics
 
 There are currently 5 topics:

flake.nix

@@ -108,6 +108,7 @@
                 package = pkgsSwift.swiftPackages.swift-format;
               };
               shfmt.enable = true;
+              taplo.enable = true;
             };
           };
 

pyproject.toml

@@ -5,31 +5,31 @@ description = "Exo"
 readme = "README.md"
 requires-python = ">=3.13"
 dependencies = [
-    "aiofiles>=24.1.0",
-    "aiohttp>=3.12.14",
-    "types-aiofiles>=24.1.0.20250708",
-    "pydantic>=2.11.7",
-    "fastapi>=0.116.1",
-    "filelock>=3.18.0",
-    "rustworkx>=0.17.1",
-    "huggingface-hub>=0.33.4",
-    "psutil>=7.0.0",
-    "loguru>=0.7.3",
-    "exo_pyo3_bindings", # rust bindings
-    "anyio==4.11.0",
-    "mlx; sys_platform == 'darwin'",
-    "mlx[cpu]==0.30.6; sys_platform == 'linux'",
-    "mlx-lm==0.30.7",
-    "tiktoken>=0.12.0", # required for kimi k2 tokenizer
-    "hypercorn>=0.18.0",
-    "openai-harmony>=0.0.8",
-    "httpx>=0.28.1",
-    "tomlkit>=0.14.0",
-    "pillow>=11.0,<12.0", # compatibility with mflux
-    "mflux==0.15.5",
-    "python-multipart>=0.0.21",
-    "msgspec>=0.19.0",
-    "zstandard>=0.23.0",
+  "aiofiles>=24.1.0",
+  "aiohttp>=3.12.14",
+  "types-aiofiles>=24.1.0.20250708",
+  "pydantic>=2.11.7",
+  "fastapi>=0.116.1",
+  "filelock>=3.18.0",
+  "rustworkx>=0.17.1",
+  "huggingface-hub>=0.33.4",
+  "psutil>=7.0.0",
+  "loguru>=0.7.3",
+  "exo_pyo3_bindings",                         # rust bindings
+  "anyio==4.11.0",
+  "mlx; sys_platform == 'darwin'",
+  "mlx[cpu]==0.30.6; sys_platform == 'linux'",
+  "mlx-lm",
+  "tiktoken>=0.12.0",                          # required for kimi k2 tokenizer
+  "hypercorn>=0.18.0",
+  "openai-harmony>=0.0.8",
+  "httpx>=0.28.1",
+  "tomlkit>=0.14.0",
+  "pillow>=11.0,<12.0",                        # compatibility with mflux
+  "mflux==0.15.5",
+  "python-multipart>=0.0.21",
+  "msgspec>=0.19.0",
+  "zstandard>=0.23.0",
 ]
 
 [project.scripts]
@@ -38,12 +38,12 @@ exo = "exo.main:main"
 # dependencies only required for development
 [dependency-groups]
 dev = [
-    "basedpyright>=1.29.0",
-    "pyinstaller>=6.17.0",
-    "pytest>=8.4.0",
-    "pytest-asyncio>=1.0.0",
-    "pytest-env",
-    "ruff>=0.11.13",
+  "basedpyright>=1.29.0",
+  "pyinstaller>=6.17.0",
+  "pytest>=8.4.0",
+  "pytest-asyncio>=1.0.0",
+  "pytest-env",
+  "ruff>=0.11.13",
 ]
 
 # mlx[cuda] requires a newer version of mlx. the ideal on linux is: default to mlx[cpu] unless[cuda] specified.
@@ -57,15 +57,12 @@ dev = [
 ###
 
 [tool.uv.workspace]
-members = [
-    "rust/exo_pyo3_bindings",
-    "bench",
-]
+members = ["rust/exo_pyo3_bindings", "bench"]
 
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
-#mlx-lm = { git = "https://github.com/davidmcc73/mlx-lm", branch = "stable" }
+mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" }
 # Uncomment to use local mlx/mlx-lm development versions:
 # mlx = { path = "/Users/Shared/mlx", editable=true }
 # mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
@@ -95,7 +92,15 @@ reportUnnecessaryTypeIgnoreComment = "error"
 pythonVersion = "3.13"
 pythonPlatform = "Darwin"
 
-exclude = ["**/.venv", "**/venv", "**/__pycache__", "**/exo_scripts", "**/.direnv", "**/rust", "**/.github"]
+exclude = [
+  "**/.venv",
+  "**/venv",
+  "**/__pycache__",
+  "**/exo_scripts",
+  "**/.direnv",
+  "**/rust",
+  "**/.github",
+]
 stubPath = ".mlx_typings"
 
 [[tool.basedpyright.executionEnvironments]]
@@ -109,17 +114,19 @@ root = "src"
 [tool.uv]
 required-version = ">=0.8.6"
 prerelease = "allow"
-environments = [
-    "sys_platform == 'darwin'",
-    "sys_platform == 'linux'",
-]
+environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
 
 ###
 # ruff configuration
 ###
 
 [tool.ruff]
-extend-exclude = ["shared/protobufs/**", "*mlx_typings/**", "rust/exo_pyo3_bindings/**"]
+extend-exclude = [
+  "shared/protobufs/**",
+  "*mlx_typings/**",
+  "rust/exo_pyo3_bindings/**",
+  "bench/vendor/**",
+]
 
 [tool.ruff.lint]
 extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
@@ -127,13 +134,7 @@ extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
 [tool.pytest.ini_options]
 pythonpath = "."
 asyncio_mode = "auto"
-markers = [
-    "slow: marks tests as slow (deselected by default)"
-]
-env = [
-  "EXO_TESTS=1"
-]
+markers = ["slow: marks tests as slow (deselected by default)"]
+env = ["EXO_TESTS=1"]
 addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
-filterwarnings = [
-    "ignore:builtin type Swig:DeprecationWarning",
-]
+filterwarnings = ["ignore:builtin type Swig:DeprecationWarning"]

python/parts.nix

@@ -43,6 +43,27 @@
             final.setuptools
           ];
         });
+        # rouge-score and sacrebleu don't declare setuptools as a build dependency
+        rouge-score = prev.rouge-score.overrideAttrs (old: {
+          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
+            final.setuptools
+          ];
+        });
+        sacrebleu = prev.sacrebleu.overrideAttrs (old: {
+          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
+            final.setuptools
+          ];
+        });
+        sqlitedict = prev.sqlitedict.overrideAttrs (old: {
+          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
+            final.setuptools
+          ];
+        });
+        word2number = prev.word2number.overrideAttrs (old: {
+          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
+            final.setuptools
+          ];
+        });
       } // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
         # Use our pure Nix-built MLX with Metal support (macOS only)
         mlx = self'.packages.mlx;
@@ -96,13 +117,16 @@
         "lib/python3.13/site-packages/nvidia*"
       ];
 
-      exoVenv = (pythonSet.mkVirtualEnv "exo-env" workspace.deps.default).overrideAttrs {
+      # Exclude bench deps from main env (bench has its own benchVenv)
+      exoDeps = removeAttrs workspace.deps.default [ "exo-bench" ];
+
+      exoVenv = (pythonSet.mkVirtualEnv "exo-env" exoDeps).overrideAttrs {
         venvIgnoreCollisions = venvCollisionPaths;
       };
 
       # Virtual environment with dev dependencies for testing
       testVenv = (pythonSet.mkVirtualEnv "exo-test-env" (
-        workspace.deps.default // {
+        exoDeps // {
           exo = [ "dev" ]; # Include pytest, pytest-asyncio, pytest-env
         }
       )).overrideAttrs {
@@ -158,6 +182,7 @@
           exo-test-env = testVenv;
         } // {
         exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
+        exo-eval = mkBenchScript "exo-eval" (inputs.self + /bench/exo_eval.py);
         exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
         exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
       };

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-122B-A10B-4bit"
+n_layers = 48
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 122B A10B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 69593314272

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-6bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-122B-A10B-6bit"
+n_layers = 48
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "6bit"
+base_model = "Qwen3.5 122B A10B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 100120675296

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-122B-A10B-8bit"
+n_layers = 48
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 122B A10B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 130648036320

resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-bf16.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-122B-A10B-bf16"
+n_layers = 48
+hidden_size = 3072
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "bf16"
+base_model = "Qwen3.5 122B A10B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 245125640160

resources/inference_model_cards/mlx-community--Qwen3.5-27B-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-27B-4bit"
+n_layers = 64
+hidden_size = 5120
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 27B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 16054266848

resources/inference_model_cards/mlx-community--Qwen3.5-27B-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-27B-8bit"
+n_layers = 64
+hidden_size = 5120
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 27B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 29500943328

resources/inference_model_cards/mlx-community--Qwen3.5-2B-MLX-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-2B-MLX-8bit"
+n_layers = 24
+hidden_size = 2048
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 2B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 2662787264

resources/inference_model_cards/mlx-community--Qwen3.5-35B-A3B-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-35B-A3B-4bit"
+n_layers = 40
+hidden_size = 2048
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 35B A3B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 20391405152

resources/inference_model_cards/mlx-community--Qwen3.5-35B-A3B-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-35B-A3B-8bit"
+n_layers = 40
+hidden_size = 2048
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 35B A3B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 37721130592

resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-397B-A17B-4bit"
+n_layers = 60
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 397B A17B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 223860768352

resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-6bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-397B-A17B-6bit"
+n_layers = 60
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "6bit"
+base_model = "Qwen3.5 397B A17B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 322946674272

resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-397B-A17B-8bit"
+n_layers = 60
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 397B A17B"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 422032580192

resources/inference_model_cards/mlx-community--Qwen3.5-9B-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-9B-4bit"
+n_layers = 32
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "4bit"
+base_model = "Qwen3.5 9B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 5950062560

resources/inference_model_cards/mlx-community--Qwen3.5-9B-8bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Qwen3.5-9B-8bit"
+n_layers = 32
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "qwen"
+quantization = "8bit"
+base_model = "Qwen3.5 9B"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 10426433504

rust/exo_pyo3_bindings/Cargo.toml

@@ -26,20 +26,24 @@ networking = { workspace = true }
 
 # interop
 pyo3 = { version = "0.27.2", features = [
-    # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
-    # "nightly", # enables better-supported GIL integration
-    "experimental-async", # async support in #[pyfunction] & #[pymethods]
-    #"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
-    #"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
-    # "multiple-pymethods", # allows multiple #[pymethods] sections per class
+  # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
+  # "nightly", # enables better-supported GIL integration
+  "experimental-async", # async support in #[pyfunction] & #[pymethods]
+  #"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
+  #"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
+  # "multiple-pymethods", # allows multiple #[pymethods] sections per class
 
-    # integrations with other libraries
-    # "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
-    # "ordered-float", "rust_decimal", "smallvec",
-    # "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time",  "serde",
+  # integrations with other libraries
+  # "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
+  # "ordered-float", "rust_decimal", "smallvec",
+  # "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time",  "serde",
 ] }
 pyo3-stub-gen = { version = "0.17.2" }
-pyo3-async-runtimes = { version = "0.27.0", features = ["attributes", "tokio-runtime", "testing"] }
+pyo3-async-runtimes = { version = "0.27.0", features = [
+  "attributes",
+  "tokio-runtime",
+  "testing",
+] }
 pyo3-log = "0.13.2"
 
 # macro dependencies

rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi

@@ -2,7 +2,6 @@
 # ruff: noqa: E501, F401
 
 import builtins
-import enum
 import typing
 
 @typing.final
@@ -11,29 +10,6 @@ class AllQueuesFullError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
-@typing.final
-class ConnectionUpdate:
-    @property
-    def update_type(self) -> ConnectionUpdateType:
-        r"""
-        Whether this is a connection or disconnection event
-        """
-    @property
-    def peer_id(self) -> builtins.str:
-        r"""
-        Identity of the peer that we have connected to or disconnected from.
-        """
-    @property
-    def remote_ipv4(self) -> builtins.str:
-        r"""
-        Remote connection's IPv4 address.
-        """
-    @property
-    def remote_tcp_port(self) -> builtins.int:
-        r"""
-        Remote connection's TCP port.
-        """
-
 @typing.final
 class Keypair:
     r"""
@@ -58,21 +34,15 @@ class Keypair:
         Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
         """
 
+@typing.final
+class MessageTooLargeError(builtins.Exception):
+    def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
+    def __repr__(self) -> builtins.str: ...
+    def __str__(self) -> builtins.str: ...
+
 @typing.final
 class NetworkingHandle:
     def __new__(cls, identity: Keypair) -> NetworkingHandle: ...
-    async def connection_update_recv(self) -> ConnectionUpdate:
-        r"""
-        Receives the next `ConnectionUpdate` from networking.
-        """
-    async def connection_update_recv_many(self, limit: builtins.int) -> builtins.list[ConnectionUpdate]:
-        r"""
-        Receives at most `limit` `ConnectionUpdate`s from networking and returns them.
-        
-        For `limit = 0`, an empty collection of `ConnectionUpdate`s will be returned immediately.
-        For `limit > 0`, if there are no `ConnectionUpdate`s in the channel's queue this method
-        will sleep until a `ConnectionUpdate`s is sent.
-        """
     async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Subscribe to a `GossipSub` topic.
@@ -91,24 +61,7 @@ class NetworkingHandle:
         
         If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
         """
-    async def gossipsub_recv(self) -> tuple[builtins.str, bytes]:
-        r"""
-        Receives the next message from the `GossipSub` network.
-        """
-    async def gossipsub_recv_many(self, limit: builtins.int) -> builtins.list[tuple[builtins.str, bytes]]:
-        r"""
-        Receives at most `limit` messages from the `GossipSub` network and returns them.
-        
-        For `limit = 0`, an empty collection of messages will be returned immediately.
-        For `limit > 0`, if there are no messages in the channel's queue this method
-        will sleep until a message is sent.
-        """
-
-@typing.final
-class MessageTooLargeError(builtins.Exception):
-    def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
-    def __repr__(self) -> builtins.str: ...
-    def __str__(self) -> builtins.str: ...
+    async def recv(self) -> PyFromSwarm: ...
 
 @typing.final
 class NoPeersSubscribedToTopicError(builtins.Exception):
@@ -116,11 +69,26 @@ class NoPeersSubscribedToTopicError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
-@typing.final
-class ConnectionUpdateType(enum.Enum):
-    r"""
-    Connection or disconnection event discriminant type.
-    """
-    Connected = ...
-    Disconnected = ...
+class PyFromSwarm:
+    @typing.final
+    class Connection(PyFromSwarm):
+        __match_args__ = ("peer_id", "connected",)
+        @property
+        def peer_id(self) -> builtins.str: ...
+        @property
+        def connected(self) -> builtins.bool: ...
+        def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> PyFromSwarm.Connection: ...
+    
+    @typing.final
+    class Message(PyFromSwarm):
+        __match_args__ = ("origin", "topic", "data",)
+        @property
+        def origin(self) -> builtins.str: ...
+        @property
+        def topic(self) -> builtins.str: ...
+        @property
+        def data(self) -> bytes: ...
+        def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
+    
+    ...
 

rust/exo_pyo3_bindings/pyproject.toml

@@ -4,21 +4,18 @@ build-backend = "maturin"
 
 [project]
 name = "exo_pyo3_bindings"
-version = "0.1.0"
+version = "0.2.1"
 description = "Add your description here"
 readme = "README.md"
 authors = [
-    { name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" }
+  { name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" },
+  { name = "Evan Quiney", email = "evanev7@gmail.com" },
 ]
 requires-python = ">=3.13"
 dependencies = []
 
 [dependency-groups]
-dev = [
-    "exo_pyo3_bindings",
-    "pytest>=8.4.0",
-    "pytest-asyncio>=1.0.0",
-]
+dev = ["exo_pyo3_bindings", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
 
 [tool.maturin]
 #purelib = true

rust/exo_pyo3_bindings/src/lib.rs

@@ -155,6 +155,9 @@ pub(crate) mod ext {
 fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // install logger
     pyo3_log::init();
+    let mut builder = tokio::runtime::Builder::new_multi_thread();
+    builder.enable_all();
+    pyo3_async_runtimes::tokio::init(builder);
 
     // TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
     //       work with maturin, where the types generate correctly, in the right folder, without

rust/exo_pyo3_bindings/src/networking.rs

@@ -1,26 +1,24 @@
-#![allow(
-    clippy::multiple_inherent_impl,
-    clippy::unnecessary_wraps,
-    clippy::unused_self,
-    clippy::needless_pass_by_value
-)]
+use std::pin::Pin;
+use std::sync::Arc;
 
 use crate::r#const::MPSC_CHANNEL_SIZE;
 use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
-use crate::ext::{ResultExt as _, TokioMpscReceiverExt as _, TokioMpscSenderExt as _};
+use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
 use crate::ident::PyKeypair;
+use crate::networking::exception::{
+    PyAllQueuesFullError, PyMessageTooLargeError, PyNoPeersSubscribedToTopicError,
+};
 use crate::pyclass;
-use libp2p::futures::StreamExt as _;
-use libp2p::gossipsub;
-use libp2p::gossipsub::{IdentTopic, Message, MessageId, PublishError};
-use libp2p::swarm::SwarmEvent;
-use networking::discovery;
-use networking::swarm::create_swarm;
+use futures_lite::{Stream, StreamExt as _};
+use libp2p::gossipsub::PublishError;
+use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
+use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::{PyModule, PyModuleMethods as _};
 use pyo3::types::PyBytes;
-use pyo3::{Bound, Py, PyErr, PyResult, PyTraverseError, PyVisit, Python, pymethods};
-use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_enum, gen_stub_pymethods};
-use std::net::IpAddr;
+use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
+use pyo3_stub_gen::derive::{
+    gen_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods,
+};
 use tokio::sync::{Mutex, mpsc, oneshot};
 
 mod exception {
@@ -131,237 +129,45 @@ mod exception {
     }
 }
 
-/// Connection or disconnection event discriminant type.
-#[gen_stub_pyclass_enum]
-#[pyclass(eq, eq_int, name = "ConnectionUpdateType")]
-#[derive(Debug, Clone, PartialEq)]
-enum PyConnectionUpdateType {
-    Connected = 0,
-    Disconnected,
-}
-
-#[gen_stub_pyclass]
-#[pyclass(frozen, name = "ConnectionUpdate")]
-#[derive(Debug, Clone)]
-struct PyConnectionUpdate {
-    /// Whether this is a connection or disconnection event
-    #[pyo3(get)]
-    update_type: PyConnectionUpdateType,
-
-    /// Identity of the peer that we have connected to or disconnected from.
-    #[pyo3(get)]
-    peer_id: String,
-
-    /// Remote connection's IPv4 address.
-    #[pyo3(get)]
-    remote_ipv4: String,
-
-    /// Remote connection's TCP port.
-    #[pyo3(get)]
-    remote_tcp_port: u16,
-}
-
-enum ToTask {
-    GossipsubSubscribe {
-        topic: String,
-        result_tx: oneshot::Sender<PyResult<bool>>,
-    },
-    GossipsubUnsubscribe {
-        topic: String,
-        result_tx: oneshot::Sender<bool>,
-    },
-    GossipsubPublish {
-        topic: String,
-        data: Vec<u8>,
-        result_tx: oneshot::Sender<PyResult<MessageId>>,
-    },
-}
-
-#[allow(clippy::enum_glob_use)]
-async fn networking_task(
-    mut swarm: networking::swarm::Swarm,
-    mut to_task_rx: mpsc::Receiver<ToTask>,
-    connection_update_tx: mpsc::Sender<PyConnectionUpdate>,
-    gossipsub_message_tx: mpsc::Sender<(String, Vec<u8>)>,
-) {
-    use SwarmEvent::*;
-    use ToTask::*;
-    use networking::swarm::BehaviourEvent::*;
-
-    log::info!("RUST: networking task started");
-
-    loop {
-        tokio::select! {
-            message = to_task_rx.recv() => {
-                // handle closed channel
-                let Some(message) = message else {
-                    log::info!("RUST: channel closed");
-                    break;
-                };
-
-                // dispatch incoming messages
-                match message {
-                    GossipsubSubscribe { topic, result_tx } => {
-                        // try to subscribe
-                        let result = swarm.behaviour_mut()
-                            .gossipsub.subscribe(&IdentTopic::new(topic));
-
-                        // send response oneshot
-                        if let Err(e) = result_tx.send(result.pyerr()) {
-                            log::error!("RUST: could not subscribe to gossipsub topic since channel already closed: {e:?}");
-                            continue;
-                        }
-                    }
-                    GossipsubUnsubscribe { topic, result_tx } => {
-                        // try to unsubscribe from the topic
-                        let result = swarm.behaviour_mut()
-                            .gossipsub.unsubscribe(&IdentTopic::new(topic));
-
-                        // send response oneshot (or exit if connection closed)
-                        if let Err(e) = result_tx.send(result) {
-                            log::error!("RUST: could not unsubscribe from gossipsub topic since channel already closed: {e:?}");
-                            continue;
-                        }
-                    }
-                    GossipsubPublish { topic, data, result_tx } => {
-                        // try to publish the data -> catch NoPeersSubscribedToTopic error & convert to correct exception
-                        let result = swarm.behaviour_mut().gossipsub.publish(
-                            IdentTopic::new(topic), data);
-                        let pyresult: PyResult<MessageId> = if let Err(PublishError::NoPeersSubscribedToTopic) = result {
-                            Err(exception::PyNoPeersSubscribedToTopicError::new_err())
-                        } else if let Err(PublishError::AllQueuesFull(_)) = result {
-                            Err(exception::PyAllQueuesFullError::new_err())
-                        } else if let Err(PublishError::MessageTooLarge) = result {
-                            Err(exception::PyMessageTooLargeError::new_err())
-                        } else {
-                            result.pyerr()
-                        };
-
-                        // send response oneshot (or exit if connection closed)
-                        if let Err(e) = result_tx.send(pyresult) {
-                            log::error!("RUST: could not publish gossipsub message since channel already closed: {e:?}");
-                            continue;
-                        }
-                    }
-                }
-            }
-
-            // architectural solution to this problem:
-            // create keep_alive behavior who's job it is to dial peers discovered by mDNS (and drop when expired)
-            //   -> it will emmit TRUE connected/disconnected events consumable elsewhere
-            //
-            // gossipsub will feed off-of dial attempts created by networking, and that will bootstrap its' peers list
-            // then for actual communication it will dial those peers if need-be
-            swarm_event = swarm.select_next_some() => {
-                match swarm_event {
-                    Behaviour(Gossipsub(gossipsub::Event::Message {
-                        message: Message {
-                            topic,
-                            data,
-                            ..
-                        },
-                        ..
-                    })) => {
-                        // topic-ID is just the topic hash!!! (since we used identity hasher)
-                        let message = (topic.into_string(), data);
-
-                        // send incoming message to channel (or exit if connection closed)
-                        if let Err(e) = gossipsub_message_tx.send(message).await {
-                            log::error!("RUST: could not send incoming gossipsub message since channel already closed: {e}");
-                            continue;
-                        }
-                    },
-                    Behaviour(Discovery(discovery::Event::ConnectionEstablished { peer_id, remote_ip, remote_tcp_port, .. })) => {
-                        // grab IPv4 string
-                        let remote_ipv4 = match remote_ip {
-                            IpAddr::V4(ip) => ip.to_string(),
-                            IpAddr::V6(ip) => {
-                                log::warn!("RUST: ignoring connection to IPv6 address: {ip}");
-                                continue;
-                            }
-                        };
-
-                        // send connection event to channel (or exit if connection closed)
-                        if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
-                            update_type: PyConnectionUpdateType::Connected,
-                            peer_id: peer_id.to_base58(),
-                            remote_ipv4,
-                            remote_tcp_port,
-                        }).await {
-                            log::error!("RUST: could not send connection update since channel already closed: {e}");
-                            continue;
-                        }
-                    },
-                    Behaviour(Discovery(discovery::Event::ConnectionClosed { peer_id, remote_ip, remote_tcp_port, .. })) => {
-                        // grab IPv4 string
-                        let remote_ipv4 = match remote_ip {
-                            IpAddr::V4(ip) => ip.to_string(),
-                            IpAddr::V6(ip) => {
-                                log::warn!("RUST: ignoring disconnection from IPv6 address: {ip}");
-                                continue;
-                            }
-                        };
-
-                        // send disconnection event to channel (or exit if connection closed)
-                        if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
-                            update_type: PyConnectionUpdateType::Disconnected,
-                            peer_id: peer_id.to_base58(),
-                            remote_ipv4,
-                            remote_tcp_port,
-                        }).await {
-                            log::error!("RUST: could not send connection update since channel already closed: {e}");
-                            continue;
-                        }
-                    },
-                    e => {
-                        log::info!("RUST: other event {e:?}");
-                    }
-                }
-            }
-        }
-    }
-
-    log::info!("RUST: networking task stopped");
-}
-
 #[gen_stub_pyclass]
 #[pyclass(name = "NetworkingHandle")]
-#[derive(Debug)]
 struct PyNetworkingHandle {
     // channels
-    to_task_tx: Option<mpsc::Sender<ToTask>>,
-    connection_update_rx: Mutex<mpsc::Receiver<PyConnectionUpdate>>,
-    gossipsub_message_rx: Mutex<mpsc::Receiver<(String, Vec<u8>)>>,
+    pub to_swarm: mpsc::Sender<ToSwarm>,
+    pub swarm: Arc<Mutex<Pin<Box<dyn Stream<Item = FromSwarm> + Send>>>>,
 }
 
-impl Drop for PyNetworkingHandle {
-    fn drop(&mut self) {
-        // TODO: may or may not need to await a "kill-signal" oneshot channel message,
-        //       to ensure that the networking task is done BEFORE exiting the clear function...
-        //       but this may require GIL?? and it may not be safe to call GIL here??
-        self.to_task_tx = None; // Using Option<T> as a trick to force channel to be dropped
-    }
+#[gen_stub_pyclass_complex_enum]
+#[pyclass]
+enum PyFromSwarm {
+    Connection {
+        peer_id: String,
+        connected: bool,
+    },
+    Message {
+        origin: String,
+        topic: String,
+        data: Py<PyBytes>,
+    },
 }
-
-#[allow(clippy::expect_used)]
-impl PyNetworkingHandle {
-    fn new(
-        to_task_tx: mpsc::Sender<ToTask>,
-        connection_update_rx: mpsc::Receiver<PyConnectionUpdate>,
-        gossipsub_message_rx: mpsc::Receiver<(String, Vec<u8>)>,
-    ) -> Self {
-        Self {
-            to_task_tx: Some(to_task_tx),
-            connection_update_rx: Mutex::new(connection_update_rx),
-            gossipsub_message_rx: Mutex::new(gossipsub_message_rx),
+impl From<FromSwarm> for PyFromSwarm {
+    fn from(value: FromSwarm) -> Self {
+        match value {
+            FromSwarm::Discovered { peer_id } => Self::Connection {
+                peer_id: peer_id.to_base58(),
+                connected: true,
+            },
+            FromSwarm::Expired { peer_id } => Self::Connection {
+                peer_id: peer_id.to_base58(),
+                connected: false,
+            },
+            FromSwarm::Message { from, topic, data } => Self::Message {
+                origin: from.to_base58(),
+                topic: topic,
+                data: data.pybytes(),
+            },
         }
     }
-
-    const fn to_task_tx(&self) -> &mpsc::Sender<ToTask> {
-        self.to_task_tx
-            .as_ref()
-            .expect("The sender should only be None after de-initialization.")
-    }
 }
 
 #[gen_stub_pymethods]
@@ -375,97 +181,36 @@ impl PyNetworkingHandle {
 
     #[new]
     fn py_new(identity: Bound<'_, PyKeypair>) -> PyResult<Self> {
-        use pyo3_async_runtimes::tokio::get_runtime;
-
         // create communication channels
-        let (to_task_tx, to_task_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
-        let (connection_update_tx, connection_update_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
-        let (gossipsub_message_tx, gossipsub_message_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
+        let (to_swarm, from_client) = mpsc::channel(MPSC_CHANNEL_SIZE);
 
         // get identity
         let identity = identity.borrow().0.clone();
 
         // create networking swarm (within tokio context!! or it crashes)
-        let swarm = get_runtime()
-            .block_on(async { create_swarm(identity) })
-            .pyerr()?;
+        let _guard = pyo3_async_runtimes::tokio::get_runtime().enter();
+        let swarm = create_swarm(identity, from_client).pyerr()?.into_stream();
 
-        // spawn tokio task running the networking logic
-        get_runtime().spawn(async move {
-            networking_task(
-                swarm,
-                to_task_rx,
-                connection_update_tx,
-                gossipsub_message_tx,
-            )
-            .await;
-        });
-        Ok(Self::new(
-            to_task_tx,
-            connection_update_rx,
-            gossipsub_message_rx,
-        ))
+        Ok(Self {
+            swarm: Arc::new(Mutex::new(swarm)),
+            to_swarm,
+        })
     }
 
     #[gen_stub(skip)]
-    const fn __traverse__(&self, _visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
-        Ok(()) // This is needed purely so `__clear__` can work
+    fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let swarm = Arc::clone(&self.swarm);
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            swarm
+                .try_lock()
+                .map_err(|_| PyRuntimeError::new_err("called recv twice concurrently"))?
+                .next()
+                .await
+                .ok_or(PyErr::receiver_channel_closed())
+                .map(PyFromSwarm::from)
+        })
     }
 
-    #[gen_stub(skip)]
-    fn __clear__(&mut self) {
-        // TODO: may or may not need to await a "kill-signal" oneshot channel message,
-        //       to ensure that the networking task is done BEFORE exiting the clear function...
-        //       but this may require GIL?? and it may not be safe to call GIL here??
-        self.to_task_tx = None; // Using Option<T> as a trick to force channel to be dropped
-    }
-
-    // ---- Connection update receiver methods ----
-
-    /// Receives the next `ConnectionUpdate` from networking.
-    async fn connection_update_recv(&self) -> PyResult<PyConnectionUpdate> {
-        self.connection_update_rx
-            .lock()
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-            .recv_py()
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-    }
-
-    /// Receives at most `limit` `ConnectionUpdate`s from networking and returns them.
-    ///
-    /// For `limit = 0`, an empty collection of `ConnectionUpdate`s will be returned immediately.
-    /// For `limit > 0`, if there are no `ConnectionUpdate`s in the channel's queue this method
-    /// will sleep until a `ConnectionUpdate`s is sent.
-    async fn connection_update_recv_many(&self, limit: usize) -> PyResult<Vec<PyConnectionUpdate>> {
-        self.connection_update_rx
-            .lock()
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-            .recv_many_py(limit)
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-    }
-
-    // TODO: rn this blocks main thread if anything else is awaiting the channel (bc its a mutex)
-    //       so its too dangerous to expose just yet. figure out a better semantics for handling this,
-    //       so things don't randomly block
-    // /// Tries to receive the next `ConnectionUpdate` from networking.
-    // fn connection_update_try_recv(&self) -> PyResult<Option<PyConnectionUpdate>> {
-    //     self.connection_update_rx.blocking_lock().try_recv_py()
-    // }
-    //
-    // /// Checks if the `ConnectionUpdate` channel is empty.
-    // fn connection_update_is_empty(&self) -> bool {
-    //     self.connection_update_rx.blocking_lock().is_empty()
-    // }
-    //
-    // /// Returns the number of `ConnectionUpdate`s in the channel.
-    // fn connection_update_len(&self) -> usize {
-    //     self.connection_update_rx.blocking_lock().len()
-    // }
-
     // ---- Gossipsub management methods ----
 
     /// Subscribe to a `GossipSub` topic.
@@ -475,10 +220,10 @@ impl PyNetworkingHandle {
         let (tx, rx) = oneshot::channel();
 
         // send off request to subscribe
-        self.to_task_tx()
-            .send_py(ToTask::GossipsubSubscribe {
+        self.to_swarm
+            .send_py(ToSwarm::Subscribe {
                 topic,
-                result_tx: tx,
+                result_sender: tx,
             })
             .allow_threads_py() // allow-threads-aware async call
             .await?;
@@ -487,6 +232,7 @@ impl PyNetworkingHandle {
         rx.allow_threads_py() // allow-threads-aware async call
             .await
             .map_err(|_| PyErr::receiver_channel_closed())?
+            .pyerr()
     }
 
     /// Unsubscribes from a `GossipSub` topic.
@@ -496,10 +242,10 @@ impl PyNetworkingHandle {
         let (tx, rx) = oneshot::channel();
 
         // send off request to unsubscribe
-        self.to_task_tx()
-            .send_py(ToTask::GossipsubUnsubscribe {
+        self.to_swarm
+            .send_py(ToSwarm::Unsubscribe {
                 topic,
-                result_tx: tx,
+                result_sender: tx,
             })
             .allow_threads_py() // allow-threads-aware async call
             .await?;
@@ -518,11 +264,11 @@ impl PyNetworkingHandle {
 
         // send off request to subscribe
         let data = Python::attach(|py| Vec::from(data.as_bytes(py)));
-        self.to_task_tx()
-            .send_py(ToTask::GossipsubPublish {
+        self.to_swarm
+            .send_py(ToSwarm::Publish {
                 topic,
                 data,
-                result_tx: tx,
+                result_sender: tx,
             })
             .allow_threads_py() // allow-threads-aware async call
             .await?;
@@ -531,64 +277,26 @@ impl PyNetworkingHandle {
         let _ = rx
             .allow_threads_py() // allow-threads-aware async call
             .await
-            .map_err(|_| PyErr::receiver_channel_closed())??;
+            .map_err(|_| PyErr::receiver_channel_closed())?
+            .map_err(|e| match e {
+                PublishError::AllQueuesFull(_) => PyAllQueuesFullError::new_err(),
+                PublishError::MessageTooLarge => PyMessageTooLargeError::new_err(),
+                PublishError::NoPeersSubscribedToTopic => {
+                    PyNoPeersSubscribedToTopicError::new_err()
+                }
+                e => PyRuntimeError::new_err(e.to_string()),
+            })?;
         Ok(())
     }
+}
 
-    // ---- Gossipsub message receiver methods ----
-
-    /// Receives the next message from the `GossipSub` network.
-    async fn gossipsub_recv(&self) -> PyResult<(String, Py<PyBytes>)> {
-        self.gossipsub_message_rx
-            .lock()
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-            .recv_py()
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-            .map(|(t, d)| (t, d.pybytes()))
+pyo3_stub_gen::inventory::submit! {
+    gen_methods_from_python! {
+        r#"
+            class PyNetworkingHandle:
+                async def recv() -> PyFromSwarm: ...
+        "#
     }
-
-    /// Receives at most `limit` messages from the `GossipSub` network and returns them.
-    ///
-    /// For `limit = 0`, an empty collection of messages will be returned immediately.
-    /// For `limit > 0`, if there are no messages in the channel's queue this method
-    /// will sleep until a message is sent.
-    async fn gossipsub_recv_many(&self, limit: usize) -> PyResult<Vec<(String, Py<PyBytes>)>> {
-        Ok(self
-            .gossipsub_message_rx
-            .lock()
-            .allow_threads_py() // allow-threads-aware async call
-            .await
-            .recv_many_py(limit)
-            .allow_threads_py() // allow-threads-aware async call
-            .await?
-            .into_iter()
-            .map(|(t, d)| (t, d.pybytes()))
-            .collect())
-    }
-
-    // TODO: rn this blocks main thread if anything else is awaiting the channel (bc its a mutex)
-    //       so its too dangerous to expose just yet. figure out a better semantics for handling this,
-    //       so things don't randomly block
-    // /// Tries to receive the next message from the `GossipSub` network.
-    // fn gossipsub_try_recv(&self) -> PyResult<Option<(String, Py<PyBytes>)>> {
-    //     Ok(self
-    //         .gossipsub_message_rx
-    //         .blocking_lock()
-    //         .try_recv_py()?
-    //         .map(|(t, d)| (t, d.pybytes())))
-    // }
-    //
-    // /// Checks if the `GossipSub` message channel is empty.
-    // fn gossipsub_is_empty(&self) -> bool {
-    //     self.gossipsub_message_rx.blocking_lock().is_empty()
-    // }
-    //
-    // /// Returns the number of `GossipSub` messages in the channel.
-    // fn gossipsub_len(&self) -> usize {
-    //     self.gossipsub_message_rx.blocking_lock().len()
-    // }
 }
 
 pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
@@ -596,10 +304,8 @@ pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<exception::PyAllQueuesFullError>()?;
     m.add_class::<exception::PyMessageTooLargeError>()?;
 
-    m.add_class::<PyConnectionUpdateType>()?;
-    m.add_class::<PyConnectionUpdate>()?;
-    m.add_class::<PyConnectionUpdateType>()?;
     m.add_class::<PyNetworkingHandle>()?;
+    m.add_class::<PyFromSwarm>()?;
 
     Ok(())
 }

rust/exo_pyo3_bindings/tests/test_python.py

@@ -1,16 +1,20 @@
 import asyncio
 
 import pytest
-from exo_pyo3_bindings import Keypair, NetworkingHandle, NoPeersSubscribedToTopicError
+from exo_pyo3_bindings import (
+    Keypair,
+    NetworkingHandle,
+    NoPeersSubscribedToTopicError,
+    PyFromSwarm,
+)
 
 
 @pytest.mark.asyncio
 async def test_sleep_on_multiple_items() -> None:
     print("PYTHON: starting handle")
-    h = NetworkingHandle(Keypair.generate_ed25519())
+    h = NetworkingHandle(Keypair.generate())
 
-    ct = asyncio.create_task(_await_cons(h))
-    mt = asyncio.create_task(_await_msg(h))
+    rt = asyncio.create_task(_await_recv(h))
 
     # sleep for 4 ticks
     for i in range(4):
@@ -22,13 +26,11 @@ async def test_sleep_on_multiple_items() -> None:
             print("caught it", e)
 
 
-async def _await_cons(h: NetworkingHandle):
+async def _await_recv(h: NetworkingHandle):
     while True:
-        c = await h.connection_update_recv()
-        print(f"PYTHON: connection update: {c}")
-
-
-async def _await_msg(h: NetworkingHandle):
-    while True:
-        m = await h.gossipsub_recv()
-        print(f"PYTHON: message: {m}")
+        event = await h.recv()
+        match event:
+            case PyFromSwarm.Connection() as c:
+                print(f"PYTHON: connection update: {c}")
+            case PyFromSwarm.Message() as m:
+                print(f"PYTHON: message: {m}")

rust/networking/Cargo.toml

@@ -21,13 +21,17 @@ extend = { workspace = true }
 delegate = { workspace = true }
 
 # async
-tokio = { workspace = true, features = ["full"] }
+async-stream = { workspace = true }
 futures-lite = { workspace = true }
 futures-timer = { workspace = true }
+tokio = { workspace = true, features = ["full"] }
 
 # utility dependencies
 util = { workspace = true }
-tracing-subscriber = { version = "0.3.19", features = ["default", "env-filter"] }
+tracing-subscriber = { version = "0.3.19", features = [
+  "default",
+  "env-filter",
+] }
 keccak-const = { workspace = true }
 
 # tracing/logging
@@ -35,3 +39,4 @@ log = { workspace = true }
 
 # networking
 libp2p = { workspace = true, features = ["full"] }
+pin-project = "1.1.10"

rust/networking/examples/chatroom.rs

@@ -1,7 +1,9 @@
 use futures_lite::StreamExt;
-use libp2p::{gossipsub, identity, swarm::SwarmEvent};
-use networking::{discovery, swarm};
-use tokio::{io, io::AsyncBufReadExt as _, select};
+use libp2p::identity;
+use networking::swarm;
+use networking::swarm::{FromSwarm, ToSwarm};
+use tokio::sync::{mpsc, oneshot};
+use tokio::{io, io::AsyncBufReadExt as _};
 use tracing_subscriber::EnvFilter;
 use tracing_subscriber::filter::LevelFilter;
 
@@ -11,64 +13,69 @@ async fn main() {
         .with_env_filter(EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into()))
         .try_init();
 
+    let (to_swarm, from_client) = mpsc::channel(20);
+
     // Configure swarm
-    let mut swarm =
-        swarm::create_swarm(identity::Keypair::generate_ed25519()).expect("Swarm creation failed");
+    let mut swarm = swarm::create_swarm(identity::Keypair::generate_ed25519(), from_client)
+        .expect("Swarm creation failed")
+        .into_stream();
 
     // Create a Gossipsub topic & subscribe
-    let topic = gossipsub::IdentTopic::new("test-net");
-    swarm
-        .behaviour_mut()
-        .gossipsub
-        .subscribe(&topic)
-        .expect("Subscribing to topic failed");
+    let (tx, rx) = oneshot::channel();
+    _ = to_swarm
+        .send(ToSwarm::Subscribe {
+            topic: "test-net".to_string(),
+            result_sender: tx,
+        })
+        .await
+        .expect("should send");
 
     // Read full lines from stdin
     let mut stdin = io::BufReader::new(io::stdin()).lines();
     println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub");
 
+    tokio::task::spawn(async move {
+        rx.await
+            .expect("tx not dropped")
+            .expect("subscribe shouldn't fail");
+        loop {
+            if let Ok(Some(line)) = stdin.next_line().await {
+                let (tx, rx) = oneshot::channel();
+                if let Err(e) = to_swarm
+                    .send(swarm::ToSwarm::Publish {
+                        topic: "test-net".to_string(),
+                        data: line.as_bytes().to_vec(),
+                        result_sender: tx,
+                    })
+                    .await
+                {
+                    println!("Send error: {e:?}");
+                    return;
+                };
+                match rx.await {
+                    Ok(Err(e)) => println!("Publish error: {e:?}"),
+                    Err(e) => println!("Publish error: {e:?}"),
+                    Ok(_) => {}
+                }
+            }
+        }
+    });
+
     // Kick it off
     loop {
-        select! {
-            // on gossipsub outgoing
-            Ok(Some(line)) = stdin.next_line() => {
-                if let Err(e) = swarm
-                    .behaviour_mut().gossipsub
-                    .publish(topic.clone(), line.as_bytes()) {
-                    println!("Publish error: {e:?}");
-                }
+        // on gossipsub outgoing
+        match swarm.next().await {
+            // on gossipsub incoming
+            Some(FromSwarm::Discovered { peer_id }) => {
+                println!("\n\nconnected to {peer_id}\n\n")
             }
-            event = swarm.next() => match event {
-                // on gossipsub incoming
-                Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Gossipsub(gossipsub::Event::Message {
-                    propagation_source: peer_id,
-                    message_id: id,
-                    message,
-                }))) => println!(
-                        "\n\nGot message: '{}' with id: {id} from peer: {peer_id}\n\n",
-                        String::from_utf8_lossy(&message.data),
-                    ),
-
-                // on discovery
-                Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Discovery(e)) )=> match e {
-                    discovery::Event::ConnectionEstablished {
-                        peer_id, connection_id, remote_ip, remote_tcp_port
-                    } => {
-                        println!("\n\nConnected to: {peer_id}; connection ID: {connection_id}; remote IP: {remote_ip}; remote TCP port: {remote_tcp_port}\n\n");
-                    }
-                    discovery::Event::ConnectionClosed {
-                        peer_id, connection_id, remote_ip, remote_tcp_port
-                    } => {
-                        eprintln!("\n\nDisconnected from: {peer_id}; connection ID: {connection_id}; remote IP: {remote_ip}; remote TCP port: {remote_tcp_port}\n\n");
-                    }
-                }
-
-                // ignore outgoing errors: those are normal
-                e@Some(SwarmEvent::OutgoingConnectionError { .. }) => { log::debug!("Outgoing connection error: {e:?}"); }
-
-                // otherwise log any other event
-                e => { log::info!("Other event {e:?}"); }
+            Some(FromSwarm::Expired { peer_id }) => {
+                println!("\n\ndisconnected from {peer_id}\n\n")
             }
+            Some(FromSwarm::Message { from, topic, data }) => {
+                println!("{topic}/{from}:\n{}", String::from_utf8_lossy(&data))
+            }
+            None => {}
         }
     }
 }

rust/networking/src/swarm.rs

@@ -1,9 +1,11 @@
-use crate::alias;
-use crate::swarm::transport::tcp_transport;
-pub use behaviour::{Behaviour, BehaviourEvent};
-use libp2p::{SwarmBuilder, identity};
+use std::pin::Pin;
 
-pub type Swarm = libp2p::Swarm<Behaviour>;
+use crate::swarm::transport::tcp_transport;
+use crate::{alias, discovery};
+pub use behaviour::{Behaviour, BehaviourEvent};
+use futures_lite::{Stream, StreamExt};
+use libp2p::{PeerId, SwarmBuilder, gossipsub, identity, swarm::SwarmEvent};
+use tokio::sync::{mpsc, oneshot};
 
 /// The current version of the network: this prevents devices running different versions of the
 /// software from interacting with each other.
@@ -15,8 +17,136 @@ pub type Swarm = libp2p::Swarm<Behaviour>;
 pub const NETWORK_VERSION: &[u8] = b"v0.0.1";
 pub const OVERRIDE_VERSION_ENV_VAR: &str = "EXO_LIBP2P_NAMESPACE";
 
+// Uses oneshot senders to emulate function calling apis while avoiding requiring unique ownership
+// of the Swarm.
+pub enum ToSwarm {
+    Unsubscribe {
+        topic: String,
+        result_sender: oneshot::Sender<bool>,
+    },
+    Subscribe {
+        topic: String,
+        result_sender: oneshot::Sender<Result<bool, gossipsub::SubscriptionError>>,
+    },
+    Publish {
+        topic: String,
+        data: Vec<u8>,
+        result_sender: oneshot::Sender<Result<gossipsub::MessageId, gossipsub::PublishError>>,
+    },
+}
+pub enum FromSwarm {
+    Message {
+        from: PeerId,
+        topic: String,
+        data: Vec<u8>,
+    },
+    Discovered {
+        peer_id: PeerId,
+    },
+    Expired {
+        peer_id: PeerId,
+    },
+}
+
+pub struct Swarm {
+    swarm: libp2p::Swarm<Behaviour>,
+    from_client: mpsc::Receiver<ToSwarm>,
+}
+
+impl Swarm {
+    pub fn into_stream(self) -> Pin<Box<dyn Stream<Item = FromSwarm> + Send>> {
+        let Swarm {
+            mut swarm,
+            mut from_client,
+        } = self;
+        let stream = async_stream::stream! {
+            loop {
+                tokio::select! {
+                    msg = from_client.recv() => {
+                        let Some(msg) = msg else { break };
+                        on_message(&mut swarm, msg);
+                    }
+                    event = swarm.next() => {
+                        let Some(event) = event else { break };
+                        if let Some(item) = filter_swarm_event(event) {
+                            yield item;
+                        }
+                    }
+                }
+            }
+        };
+        Box::pin(stream)
+    }
+}
+
+fn on_message(swarm: &mut libp2p::Swarm<Behaviour>, message: ToSwarm) {
+    match message {
+        ToSwarm::Subscribe {
+            topic,
+            result_sender,
+        } => {
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .subscribe(&gossipsub::IdentTopic::new(topic));
+            _ = result_sender.send(result);
+        }
+        ToSwarm::Unsubscribe {
+            topic,
+            result_sender,
+        } => {
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .unsubscribe(&gossipsub::IdentTopic::new(topic));
+            _ = result_sender.send(result);
+        }
+        ToSwarm::Publish {
+            topic,
+            data,
+            result_sender,
+        } => {
+            let result = swarm
+                .behaviour_mut()
+                .gossipsub
+                .publish(gossipsub::IdentTopic::new(topic), data);
+            _ = result_sender.send(result);
+        }
+    }
+}
+
+fn filter_swarm_event(event: SwarmEvent<BehaviourEvent>) -> Option<FromSwarm> {
+    match event {
+        SwarmEvent::Behaviour(BehaviourEvent::Gossipsub(gossipsub::Event::Message {
+            message:
+                gossipsub::Message {
+                    source: Some(peer_id),
+                    topic,
+                    data,
+                    ..
+                },
+            ..
+        })) => Some(FromSwarm::Message {
+            from: peer_id,
+            topic: topic.into_string(),
+            data,
+        }),
+        SwarmEvent::Behaviour(BehaviourEvent::Discovery(
+            discovery::Event::ConnectionEstablished { peer_id, .. },
+        )) => Some(FromSwarm::Discovered { peer_id }),
+        SwarmEvent::Behaviour(BehaviourEvent::Discovery(discovery::Event::ConnectionClosed {
+            peer_id,
+            ..
+        })) => Some(FromSwarm::Expired { peer_id }),
+        _ => None,
+    }
+}
+
 /// Create and configure a swarm which listens to all ports on OS
-pub fn create_swarm(keypair: identity::Keypair) -> alias::AnyResult<Swarm> {
+pub fn create_swarm(
+    keypair: identity::Keypair,
+    from_client: mpsc::Receiver<ToSwarm>,
+) -> alias::AnyResult<Swarm> {
     let mut swarm = SwarmBuilder::with_existing_identity(keypair)
         .with_tokio()
         .with_other_transport(tcp_transport)?
@@ -25,7 +155,7 @@ pub fn create_swarm(keypair: identity::Keypair) -> alias::AnyResult<Swarm> {
 
     // Listen on all interfaces and whatever port the OS assigns
     swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
-    Ok(swarm)
+    Ok(Swarm { swarm, from_client })
 }
 
 mod transport {
@@ -133,7 +263,7 @@ mod behaviour {
         gossipsub::Behaviour::new(
             MessageAuthenticity::Signed(keypair.clone()),
             ConfigBuilder::default()
-                .max_transmit_size(1024 * 1024)
+                .max_transmit_size(8 * 1024 * 1024)
                 .validation_mode(ValidationMode::Strict)
                 .build()
                 .expect("the configuration should always be valid"),

src/exo/download/coordinator.py

@@ -1,6 +1,5 @@
 import asyncio
 from dataclasses import dataclass, field
-from random import random
 
 import anyio
 from anyio import current_time
@@ -21,13 +20,9 @@ from exo.shared.types.commands import (
     ForwarderDownloadCommand,
     StartDownload,
 )
-from exo.shared.types.common import NodeId, SessionId, SystemId
+from exo.shared.types.common import NodeId
 from exo.shared.types.events import (
     Event,
-    EventId,
-    # TODO(evan): just for acks, should delete this ASAP
-    GlobalForwarderEvent,
-    LocalForwarderEvent,
     NodeDownloadProgress,
 )
 from exo.shared.types.worker.downloads import (
@@ -38,40 +33,28 @@ from exo.shared.types.worker.downloads import (
     DownloadProgress,
 )
 from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
-from exo.utils.channels import Receiver, Sender, channel
+from exo.utils.channels import Receiver, Sender
 from exo.utils.task_group import TaskGroup
 
 
 @dataclass
 class DownloadCoordinator:
     node_id: NodeId
-    session_id: SessionId
     shard_downloader: ShardDownloader
     download_command_receiver: Receiver[ForwarderDownloadCommand]
-    local_event_sender: Sender[LocalForwarderEvent]
-
-    # ack stuff
-    _global_event_receiver: Receiver[GlobalForwarderEvent]
-    _out_for_delivery: dict[EventId, LocalForwarderEvent] = field(default_factory=dict)
-
+    event_sender: Sender[Event]
     offline: bool = False
 
-    _system_id: SystemId = field(default_factory=SystemId)
-
     # Local state
     download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
     active_downloads: dict[ModelId, asyncio.Task[None]] = field(default_factory=dict)
 
-    # Internal event channel for forwarding (initialized in __post_init__)
-    event_sender: Sender[Event] = field(init=False)
-    event_receiver: Receiver[Event] = field(init=False)
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
 
     # Per-model throttle for download progress events
     _last_progress_time: dict[ModelId, float] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
-        self.event_sender, self.event_receiver = channel[Event]()
         self.shard_downloader.on_progress(self._download_progress_callback)
 
     def _model_dir(self, model_id: ModelId) -> str:
@@ -123,10 +106,7 @@ class DownloadCoordinator:
         try:
             async with self._tg as tg:
                 tg.start_soon(self._command_processor)
-                tg.start_soon(self._forward_events)
                 tg.start_soon(self._emit_existing_download_progress)
-                tg.start_soon(self._resend_out_for_delivery)
-                tg.start_soon(self._clear_ofd)
         finally:
             for task in self.active_downloads.values():
                 task.cancel()
@@ -134,20 +114,6 @@ class DownloadCoordinator:
     def shutdown(self) -> None:
         self._tg.cancel_tasks()
 
-    # directly copied from worker
-    async def _resend_out_for_delivery(self) -> None:
-        # This can also be massively tightened, we should check events are at least a certain age before resending.
-        # Exponential backoff would also certainly help here.
-        while True:
-            await anyio.sleep(1 + random())
-            for event in self._out_for_delivery.copy().values():
-                await self.local_event_sender.send(event)
-
-    async def _clear_ofd(self) -> None:
-        with self._global_event_receiver as events:
-            async for event in events:
-                self._out_for_delivery.pop(event.event.event_id, None)
-
     async def _command_processor(self) -> None:
         with self.download_command_receiver as commands:
             async for cmd in commands:
@@ -167,6 +133,16 @@ class DownloadCoordinator:
         if model_id in self.active_downloads and model_id in self.download_status:
             logger.info(f"Cancelling download for {model_id}")
             self.active_downloads.pop(model_id).cancel()
+            current_status = self.download_status[model_id]
+            pending = DownloadPending(
+                shard_metadata=current_status.shard_metadata,
+                node_id=self.node_id,
+                model_directory=self._model_dir(model_id),
+            )
+            self.download_status[model_id] = pending
+            await self.event_sender.send(
+                NodeDownloadProgress(download_progress=pending)
+            )
 
     async def _start_download(self, shard: ShardMetadata) -> None:
         model_id = shard.model_card.model_id
@@ -320,26 +296,9 @@ class DownloadCoordinator:
             )
             del self.download_status[model_id]
 
-    async def _forward_events(self) -> None:
-        idx = 0
-        with self.event_receiver as events:
-            async for event in events:
-                fe = LocalForwarderEvent(
-                    origin_idx=idx,
-                    origin=self._system_id,
-                    session=self.session_id,
-                    event=event,
-                )
-                idx += 1
-                logger.debug(
-                    f"DownloadCoordinator published event {idx}: {str(event)[:100]}"
-                )
-                await self.local_event_sender.send(fe)
-                self._out_for_delivery[event.event_id] = fe
-
     async def _emit_existing_download_progress(self) -> None:
-        try:
-            while True:
+        while True:
+            try:
                 logger.debug(
                     "DownloadCoordinator: Fetching and emitting existing download progress..."
                 )
@@ -427,8 +386,8 @@ class DownloadCoordinator:
                 logger.debug(
                     "DownloadCoordinator: Done emitting existing download progress."
                 )
-                await anyio.sleep(60)
-        except Exception as e:
-            logger.error(
-                f"DownloadCoordinator: Error emitting existing download progress: {e}"
-            )
+            except Exception as e:
+                logger.error(
+                    f"DownloadCoordinator: Error emitting existing download progress: {e}"
+                )
+            await anyio.sleep(60)

src/exo/download/download_utils.py

@@ -314,9 +314,13 @@ async def fetch_file_list_with_cache(
         _fetched_file_lists_this_session.add(cache_key)
         return file_list
     except Exception as e:
+        logger.opt(exception=e).warning(
+            "Ran into exception when fetching file list from HF."
+        )
+
         if await aios.path.exists(cache_file):
             logger.warning(
-                f"No internet and no cached file list for {model_id} - using local file list"
+                f"No cached file list for {model_id} - using local file list"
             )
             async with aiofiles.open(cache_file, "r") as f:
                 return TypeAdapter(list[FileListEntry]).validate_json(await f.read())

src/exo/download/impl_shard_downloader.py

@@ -19,9 +19,7 @@ def exo_shard_downloader(
     max_parallel_downloads: int = 8, offline: bool = False
 ) -> ShardDownloader:
     return SingletonShardDownloader(
-        CachedShardDownloader(
-            ResumableShardDownloader(max_parallel_downloads, offline=offline)
-        )
+        ResumableShardDownloader(max_parallel_downloads, offline=offline)
     )
 
 
@@ -85,39 +83,6 @@ class SingletonShardDownloader(ShardDownloader):
         return await self.shard_downloader.get_shard_download_status_for_shard(shard)
 
 
-class CachedShardDownloader(ShardDownloader):
-    def __init__(self, shard_downloader: ShardDownloader):
-        self.shard_downloader = shard_downloader
-        self.cache: dict[tuple[str, ShardMetadata], Path] = {}
-
-    def on_progress(
-        self,
-        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
-    ) -> None:
-        self.shard_downloader.on_progress(callback)
-
-    async def ensure_shard(
-        self, shard: ShardMetadata, config_only: bool = False
-    ) -> Path:
-        if (shard.model_card.model_id, shard) in self.cache:
-            return self.cache[(shard.model_card.model_id, shard)]
-
-        target_dir = await self.shard_downloader.ensure_shard(shard, config_only)
-        self.cache[(shard.model_card.model_id, shard)] = target_dir
-        return target_dir
-
-    async def get_shard_download_status(
-        self,
-    ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
-        async for path, status in self.shard_downloader.get_shard_download_status():
-            yield path, status
-
-    async def get_shard_download_status_for_shard(
-        self, shard: ShardMetadata
-    ) -> RepoDownloadProgress:
-        return await self.shard_downloader.get_shard_download_status_for_shard(shard)
-
-
 class ResumableShardDownloader(ShardDownloader):
     def __init__(self, max_parallel_downloads: int = 8, offline: bool = False):
         self.max_parallel_downloads = max_parallel_downloads

src/exo/download/tests/test_coordinator_ack.py

@@ -1,98 +0,0 @@
-from typing import Any
-
-import anyio
-import pytest
-
-from exo.download.coordinator import DownloadCoordinator
-from exo.download.shard_downloader import NoopShardDownloader
-from exo.shared.models.model_cards import ModelCard, ModelTask
-from exo.shared.types.common import ModelId, NodeId, SessionId
-from exo.shared.types.events import (
-    GlobalForwarderEvent,
-    LocalForwarderEvent,
-    NodeDownloadProgress,
-)
-from exo.shared.types.memory import Memory
-from exo.shared.types.worker.downloads import (
-    DownloadPending,
-)
-from exo.shared.types.worker.shards import PipelineShardMetadata
-from exo.utils.channels import channel
-
-# Use the built‑in NoopShardDownloader directly – it already implements the required abstract interface.
-# No additional subclass is needed for this test.
-
-
-@pytest.mark.anyio
-async def test_ack_behaviour():
-    # Create channels (type Any for simplicity)
-    _, command_receiver = channel[Any]()
-    local_sender, _ = channel[Any]()
-    global_sender, global_receiver = channel[Any]()
-
-    # Minimal identifiers
-    node_id = NodeId()
-    session_id = SessionId(master_node_id=node_id, election_clock=0)
-
-    # Create a dummy model card and shard metadata
-    model_id = ModelId("test/model")
-    model_card = ModelCard(
-        model_id=model_id,
-        storage_size=Memory.from_bytes(0),
-        n_layers=1,
-        hidden_size=1,
-        supports_tensor=True,
-        tasks=[ModelTask.TextGeneration],
-    )
-    shard = PipelineShardMetadata(
-        model_card=model_card,
-        device_rank=0,
-        world_size=1,
-        start_layer=0,
-        end_layer=1,
-        n_layers=1,
-    )
-
-    # Instantiate the coordinator with the dummy downloader
-    coord = DownloadCoordinator(
-        node_id=node_id,
-        session_id=session_id,
-        shard_downloader=NoopShardDownloader(),
-        download_command_receiver=command_receiver,
-        local_event_sender=local_sender,
-        _global_event_receiver=global_receiver,
-    )
-
-    async with anyio.create_task_group() as tg:
-        # Start the forwarding and ack‑clearing loops
-        tg.start_soon(coord._forward_events)  # pyright: ignore[reportPrivateUsage]
-        tg.start_soon(coord._clear_ofd)  # pyright: ignore[reportPrivateUsage]
-
-        # Send a pending download progress event via the internal event sender
-        pending = DownloadPending(
-            node_id=node_id,
-            shard_metadata=shard,
-            model_directory="/tmp/model",
-        )
-        await coord.event_sender.send(NodeDownloadProgress(download_progress=pending))
-        # Allow the forwarder to process the event
-        await anyio.sleep(0.1)
-
-        # There should be exactly one entry awaiting ACK
-        assert len(coord._out_for_delivery) == 1  # pyright: ignore[reportPrivateUsage]
-        # Retrieve the stored LocalForwarderEvent
-        stored_fe: LocalForwarderEvent = next(iter(coord._out_for_delivery.values()))  # pyright: ignore[reportPrivateUsage]
-        # Simulate receiving a global ack for this event
-        ack = GlobalForwarderEvent(
-            origin_idx=0,
-            origin=node_id,
-            session=session_id,
-            event=stored_fe.event,
-        )
-        await global_sender.send(ack)
-        # Give the clear‑ofd task a moment to process the ack
-        await anyio.sleep(0.1)
-        # The out‑for‑delivery map should now be empty
-        assert len(coord._out_for_delivery) == 0  # pyright: ignore[reportPrivateUsage]
-        # Cancel background tasks
-        tg.cancel_scope.cancel()

src/exo/download/tests/test_re_download.py

@@ -0,0 +1,211 @@
+"""Tests that re-downloading a previously deleted model completes successfully."""
+
+import asyncio
+import contextlib
+from collections.abc import AsyncIterator, Awaitable
+from datetime import timedelta
+from pathlib import Path
+from typing import Callable
+from unittest.mock import AsyncMock, patch
+
+from exo.download.coordinator import DownloadCoordinator
+from exo.download.download_utils import RepoDownloadProgress
+from exo.download.impl_shard_downloader import SingletonShardDownloader
+from exo.download.shard_downloader import ShardDownloader
+from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
+from exo.shared.types.commands import (
+    DeleteDownload,
+    ForwarderDownloadCommand,
+    StartDownload,
+)
+from exo.shared.types.common import NodeId, SystemId
+from exo.shared.types.events import Event, NodeDownloadProgress
+from exo.shared.types.memory import Memory
+from exo.shared.types.worker.downloads import DownloadCompleted
+from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
+from exo.utils.channels import Receiver, Sender, channel
+
+NODE_ID = NodeId("aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa")
+MODEL_ID = ModelId("test-org/test-model")
+
+
+def _make_shard(model_id: ModelId = MODEL_ID) -> ShardMetadata:
+    return PipelineShardMetadata(
+        model_card=ModelCard(
+            model_id=model_id,
+            storage_size=Memory.from_mb(100),
+            n_layers=28,
+            hidden_size=1024,
+            supports_tensor=False,
+            tasks=[ModelTask.TextGeneration],
+        ),
+        device_rank=0,
+        world_size=1,
+        start_layer=0,
+        end_layer=28,
+        n_layers=28,
+    )
+
+
+class FakeShardDownloader(ShardDownloader):
+    """Fake downloader that simulates a successful download by firing the
+    progress callback with status='complete' when ensure_shard is called."""
+
+    def __init__(self) -> None:
+        self._progress_callbacks: list[
+            Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]]
+        ] = []
+
+    def on_progress(
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
+    ) -> None:
+        self._progress_callbacks.append(callback)
+
+    async def ensure_shard(
+        self,
+        shard: ShardMetadata,
+        config_only: bool = False,  # noqa: ARG002
+    ) -> Path:
+        # Simulate a completed download by firing the progress callback
+        progress = RepoDownloadProgress(
+            repo_id=str(shard.model_card.model_id),
+            repo_revision="main",
+            shard=shard,
+            completed_files=1,
+            total_files=1,
+            downloaded=Memory.from_mb(100),
+            downloaded_this_session=Memory.from_mb(100),
+            total=Memory.from_mb(100),
+            overall_speed=0,
+            overall_eta=timedelta(seconds=0),
+            status="complete",
+        )
+        for cb in self._progress_callbacks:
+            await cb(shard, progress)
+        return Path("/fake/models") / shard.model_card.model_id.normalize()
+
+    async def get_shard_download_status(
+        self,
+    ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
+        if False:  # noqa: SIM108  # empty async generator
+            yield (
+                Path(),
+                RepoDownloadProgress(  # pyright: ignore[reportUnreachable]
+                    repo_id="",
+                    repo_revision="",
+                    shard=_make_shard(),
+                    completed_files=0,
+                    total_files=0,
+                    downloaded=Memory.from_bytes(0),
+                    downloaded_this_session=Memory.from_bytes(0),
+                    total=Memory.from_bytes(0),
+                    overall_speed=0,
+                    overall_eta=timedelta(seconds=0),
+                    status="not_started",
+                ),
+            )
+
+    async def get_shard_download_status_for_shard(
+        self,
+        shard: ShardMetadata,
+    ) -> RepoDownloadProgress:
+        return RepoDownloadProgress(
+            repo_id=str(shard.model_card.model_id),
+            repo_revision="main",
+            shard=shard,
+            completed_files=0,
+            total_files=1,
+            downloaded=Memory.from_bytes(0),
+            downloaded_this_session=Memory.from_bytes(0),
+            total=Memory.from_mb(100),
+            overall_speed=0,
+            overall_eta=timedelta(seconds=0),
+            status="not_started",
+        )
+
+
+async def test_re_download_after_delete_completes() -> None:
+    """A model that was downloaded, deleted, and then re-downloaded should
+    reach DownloadCompleted status. This is an end-to-end test through
+    the DownloadCoordinator."""
+    cmd_send: Sender[ForwarderDownloadCommand]
+    cmd_send, cmd_recv = channel[ForwarderDownloadCommand]()
+    event_send, event_recv = channel[Event]()
+
+    fake_downloader = FakeShardDownloader()
+    wrapped_downloader = SingletonShardDownloader(fake_downloader)
+    coordinator = DownloadCoordinator(
+        node_id=NODE_ID,
+        shard_downloader=wrapped_downloader,
+        download_command_receiver=cmd_recv,
+        event_sender=event_send,
+    )
+
+    shard = _make_shard()
+    origin = SystemId("test")
+
+    with patch("exo.download.coordinator.delete_model", new_callable=AsyncMock):
+        # Run the coordinator in the background
+        coordinator_task = asyncio.create_task(coordinator.run())
+
+        try:
+            # 1. Start first download
+            await cmd_send.send(
+                ForwarderDownloadCommand(
+                    origin=origin,
+                    command=StartDownload(target_node_id=NODE_ID, shard_metadata=shard),
+                )
+            )
+
+            # Wait for DownloadCompleted
+            first_completed = await _wait_for_download_completed(event_recv, MODEL_ID)
+            assert first_completed is not None, "First download should complete"
+
+            # 2. Delete the model
+            await cmd_send.send(
+                ForwarderDownloadCommand(
+                    origin=origin,
+                    command=DeleteDownload(target_node_id=NODE_ID, model_id=MODEL_ID),
+                )
+            )
+            # Give the coordinator time to process the delete
+            await asyncio.sleep(0.05)
+
+            # 3. Re-download the same model
+            await cmd_send.send(
+                ForwarderDownloadCommand(
+                    origin=origin,
+                    command=StartDownload(target_node_id=NODE_ID, shard_metadata=shard),
+                )
+            )
+
+            # Wait for second DownloadCompleted — this is the bug: it never arrives
+            second_completed = await _wait_for_download_completed(event_recv, MODEL_ID)
+            assert second_completed is not None, (
+                "Re-download after deletion should complete"
+            )
+        finally:
+            coordinator.shutdown()
+            coordinator_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await coordinator_task
+
+
+async def _wait_for_download_completed(
+    event_recv: Receiver[Event], model_id: ModelId, timeout: float = 2.0
+) -> DownloadCompleted | None:
+    """Drain events until we see a DownloadCompleted for the given model, or timeout."""
+    try:
+        async with asyncio.timeout(timeout):
+            while True:
+                event = await event_recv.receive()
+                if (
+                    isinstance(event, NodeDownloadProgress)
+                    and isinstance(event.download_progress, DownloadCompleted)
+                    and event.download_progress.shard_metadata.model_card.model_id
+                    == model_id
+                ):
+                    return event.download_progress
+    except TimeoutError:
+        return None

src/exo/main.py

@@ -15,6 +15,7 @@ from exo.download.coordinator import DownloadCoordinator
 from exo.download.impl_shard_downloader import exo_shard_downloader
 from exo.master.api import API  # TODO: should API be in master?
 from exo.master.main import Master
+from exo.routing.event_router import EventRouter
 from exo.routing.router import Router, get_node_id_keypair
 from exo.shared.constants import EXO_LOG
 from exo.shared.election import Election, ElectionResult
@@ -29,6 +30,7 @@ from exo.worker.main import Worker
 @dataclass
 class Node:
     router: Router
+    event_router: EventRouter
     download_coordinator: DownloadCoordinator | None
     worker: Worker | None
     election: Election  # Every node participates in election, as we do want a node to become master even if it isn't a master candidate if no master candidates are present.
@@ -52,6 +54,12 @@ class Node:
         await router.register_topic(topics.ELECTION_MESSAGES)
         await router.register_topic(topics.CONNECTION_MESSAGES)
         await router.register_topic(topics.DOWNLOAD_COMMANDS)
+        event_router = EventRouter(
+            session_id,
+            command_sender=router.sender(topics.COMMANDS),
+            external_outbound=router.sender(topics.LOCAL_EVENTS),
+            external_inbound=router.receiver(topics.GLOBAL_EVENTS),
+        )
 
         logger.info(f"Starting node {node_id}")
 
@@ -59,13 +67,10 @@ class Node:
         if not args.no_downloads:
             download_coordinator = DownloadCoordinator(
                 node_id,
-                session_id,
                 exo_shard_downloader(offline=args.offline),
+                event_sender=event_router.sender(),
                 download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
-                local_event_sender=router.sender(topics.LOCAL_EVENTS),
                 offline=args.offline,
-                # TODO(evan): remove
-                _global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
             )
         else:
             download_coordinator = None
@@ -73,9 +78,8 @@ class Node:
         if args.spawn_api:
             api = API(
                 node_id,
-                session_id,
                 port=args.api_port,
-                global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
+                event_receiver=event_router.receiver(),
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 election_receiver=router.receiver(topics.ELECTION_MESSAGES),
@@ -86,9 +90,8 @@ class Node:
         if not args.no_worker:
             worker = Worker(
                 node_id,
-                session_id,
-                global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
-                local_event_sender=router.sender(topics.LOCAL_EVENTS),
+                event_receiver=event_router.receiver(),
+                event_sender=event_router.sender(),
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
             )
@@ -99,6 +102,7 @@ class Node:
         master = Master(
             node_id,
             session_id,
+            event_sender=event_router.sender(),
             global_event_sender=router.sender(topics.GLOBAL_EVENTS),
             local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
             command_receiver=router.receiver(topics.COMMANDS),
@@ -121,6 +125,7 @@ class Node:
 
         return cls(
             router,
+            event_router,
             download_coordinator,
             worker,
             election,
@@ -136,6 +141,7 @@ class Node:
             signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
             signal.signal(signal.SIGTERM, lambda _, __: self.shutdown())
             tg.start_soon(self.router.run)
+            tg.start_soon(self.event_router.run)
             tg.start_soon(self.election.run)
             if self.download_coordinator:
                 tg.start_soon(self.download_coordinator.run)
@@ -170,6 +176,17 @@ class Node:
                 # - Shutdown and re-create the worker
                 # - Shut down and re-create the API
 
+                if result.is_new_master:
+                    await anyio.sleep(0)
+                    self.event_router.shutdown()
+                    self.event_router = EventRouter(
+                        result.session_id,
+                        self.router.sender(topics.COMMANDS),
+                        self.router.receiver(topics.GLOBAL_EVENTS),
+                        self.router.sender(topics.LOCAL_EVENTS),
+                    )
+                    self._tg.start_soon(self.event_router.run)
+
                 if (
                     result.session_id.master_node_id == self.node_id
                     and self.master is not None
@@ -183,6 +200,7 @@ class Node:
                     self.master = Master(
                         self.node_id,
                         result.session_id,
+                        event_sender=self.event_router.sender(),
                         global_event_sender=self.router.sender(topics.GLOBAL_EVENTS),
                         local_event_receiver=self.router.receiver(topics.LOCAL_EVENTS),
                         command_receiver=self.router.receiver(topics.COMMANDS),
@@ -205,22 +223,16 @@ class Node:
                         f"Node {result.session_id.master_node_id} elected master"
                     )
                 if result.is_new_master:
-                    await anyio.sleep(0)
                     if self.download_coordinator:
                         self.download_coordinator.shutdown()
                         self.download_coordinator = DownloadCoordinator(
                             self.node_id,
-                            result.session_id,
                             exo_shard_downloader(offline=self.offline),
+                            event_sender=self.event_router.sender(),
                             download_command_receiver=self.router.receiver(
                                 topics.DOWNLOAD_COMMANDS
                             ),
-                            local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
                             offline=self.offline,
-                            # TODO(evan): remove
-                            _global_event_receiver=self.router.receiver(
-                                topics.GLOBAL_EVENTS
-                            ),
                         )
                         self._tg.start_soon(self.download_coordinator.run)
                     if self.worker:
@@ -228,11 +240,8 @@ class Node:
                         # TODO: add profiling etc to resource monitor
                         self.worker = Worker(
                             self.node_id,
-                            result.session_id,
-                            global_event_receiver=self.router.receiver(
-                                topics.GLOBAL_EVENTS
-                            ),
-                            local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
+                            event_receiver=self.event_router.receiver(),
+                            event_sender=self.event_router.sender(),
                             command_sender=self.router.sender(topics.COMMANDS),
                             download_command_sender=self.router.sender(
                                 topics.DOWNLOAD_COMMANDS
@@ -240,7 +249,7 @@ class Node:
                         )
                         self._tg.start_soon(self.worker.run)
                     if self.api:
-                        self.api.reset(result.session_id, result.won_clock)
+                        self.api.reset(result.won_clock, self.event_router.receiver())
                 else:
                     if self.api:
                         self.api.unpause(result.won_clock)
@@ -261,6 +270,10 @@ def main():
     if args.offline:
         logger.info("Running in OFFLINE mode — no internet checks, local models only")
 
+    if args.no_batch:
+        os.environ["EXO_NO_BATCH"] = "1"
+        logger.info("Continuous batching disabled (--no-batch)")
+
     # Set FAST_SYNCH override env var for runner subprocesses
     if args.fast_synch is True:
         os.environ["EXO_FAST_SYNCH"] = "on"
@@ -291,6 +304,7 @@ class Args(CamelCaseModel):
     no_worker: bool = False
     no_downloads: bool = False
     offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
+    no_batch: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
 
     @classmethod
@@ -344,6 +358,11 @@ class Args(CamelCaseModel):
             default=os.getenv("EXO_OFFLINE", "false").lower() == "true",
             help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models",
         )
+        parser.add_argument(
+            "--no-batch",
+            action="store_true",
+            help="Disable continuous batching, use sequential generation",
+        )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(
             "--fast-synch",

src/exo/master/adapters/chat_completions.py

@@ -26,7 +26,11 @@ from exo.shared.types.chunks import (
     ToolCallChunk,
 )
 from exo.shared.types.common import CommandId
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+    InputMessage,
+    TextGenerationTaskParams,
+    resolve_reasoning_params,
+)
 
 
 def chat_request_to_text_generation(
@@ -75,6 +79,10 @@ def chat_request_to_text_generation(
             dumped: dict[str, Any] = msg_copy.model_dump(exclude_none=True)
             chat_template_messages.append(dumped)
 
+    resolved_effort, resolved_thinking = resolve_reasoning_params(
+        request.reasoning_effort, request.enable_thinking
+    )
+
     return TextGenerationTaskParams(
         model=request.model,
         input=input_messages
@@ -89,12 +97,16 @@ def chat_request_to_text_generation(
         seed=request.seed,
         stream=request.stream,
         tools=request.tools,
-        enable_thinking=request.enable_thinking,
+        reasoning_effort=resolved_effort,
+        enable_thinking=resolved_thinking,
         chat_template_messages=chat_template_messages
         if chat_template_messages
         else None,
         logprobs=request.logprobs or False,
         top_logprobs=request.top_logprobs,
+        min_p=request.min_p,
+        repetition_penalty=request.repetition_penalty,
+        repetition_context_size=request.repetition_context_size,
     )
 
 

src/exo/master/adapters/responses.py

@@ -42,7 +42,11 @@ from exo.shared.types.openai_responses import (
     ResponseTextDoneEvent,
     ResponseUsage,
 )
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.text_generation import (
+    InputMessage,
+    TextGenerationTaskParams,
+    resolve_reasoning_params,
+)
 
 
 def _format_sse(event: ResponsesStreamEvent) -> str:
@@ -119,6 +123,11 @@ def responses_request_to_text_generation(
         )
         built_chat_template = chat_template_messages if chat_template_messages else None
 
+    effort_from_reasoning = request.reasoning.effort if request.reasoning else None
+    resolved_effort, resolved_thinking = resolve_reasoning_params(
+        effort_from_reasoning, request.enable_thinking
+    )
+
     return TextGenerationTaskParams(
         model=request.model,
         input=input_value,
@@ -132,6 +141,8 @@ def responses_request_to_text_generation(
         stop=request.stop,
         seed=request.seed,
         chat_template_messages=built_chat_template or request.chat_template_messages,
+        reasoning_effort=resolved_effort,
+        enable_thinking=resolved_thinking,
     )
 
 

src/exo/master/api.py

@@ -3,7 +3,7 @@ import contextlib
 import json
 import random
 import time
-from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator
+from collections.abc import AsyncGenerator, Awaitable, Callable, Iterable
 from datetime import datetime, timezone
 from http import HTTPStatus
 from pathlib import Path
@@ -73,6 +73,7 @@ from exo.shared.types.api import (
     BenchChatCompletionResponse,
     BenchImageGenerationResponse,
     BenchImageGenerationTaskParams,
+    CancelCommandResponse,
     ChatCompletionChoice,
     ChatCompletionMessage,
     ChatCompletionRequest,
@@ -81,6 +82,8 @@ from exo.shared.types.api import (
     CreateInstanceResponse,
     DeleteDownloadResponse,
     DeleteInstanceResponse,
+    DeleteTracesRequest,
+    DeleteTracesResponse,
     ErrorInfo,
     ErrorResponse,
     FinishReason,
@@ -140,11 +143,10 @@ from exo.shared.types.commands import (
     TaskFinished,
     TextGeneration,
 )
-from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
+from exo.shared.types.common import CommandId, Id, NodeId, SystemId
 from exo.shared.types.events import (
     ChunkGenerated,
     Event,
-    GlobalForwarderEvent,
     IndexedEvent,
     TracesMerged,
 )
@@ -172,7 +174,6 @@ from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.banner import print_startup_banner
 from exo.utils.channels import Receiver, Sender, channel
-from exo.utils.event_buffer import OrderedBuffer
 from exo.utils.task_group import TaskGroup
 
 _API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
@@ -196,10 +197,9 @@ class API:
     def __init__(
         self,
         node_id: NodeId,
-        session_id: SessionId,
         *,
         port: int,
-        global_event_receiver: Receiver[GlobalForwarderEvent],
+        event_receiver: Receiver[IndexedEvent],
         command_sender: Sender[ForwarderCommand],
         download_command_sender: Sender[ForwarderDownloadCommand],
         # This lets us pause the API if an election is running
@@ -210,11 +210,9 @@ class API:
         self._system_id = SystemId()
         self.command_sender = command_sender
         self.download_command_sender = download_command_sender
-        self.global_event_receiver = global_event_receiver
+        self.event_receiver = event_receiver
         self.election_receiver = election_receiver
-        self.event_buffer: OrderedBuffer[Event] = OrderedBuffer[Event]()
         self.node_id: NodeId = node_id
-        self.session_id: SessionId = session_id
         self.last_completed_election: int = 0
         self.port = port
 
@@ -254,17 +252,18 @@ class API:
         self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
         self._tg: TaskGroup = TaskGroup()
 
-    def reset(self, new_session_id: SessionId, result_clock: int):
+    def reset(self, result_clock: int, event_receiver: Receiver[IndexedEvent]):
         logger.info("Resetting API State")
         self._event_log.close()
         self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
         self.state = State()
         self._system_id = SystemId()
-        self.session_id = new_session_id
-        self.event_buffer = OrderedBuffer[Event]()
         self._text_generation_queues = {}
         self._image_generation_queues = {}
         self.unpause(result_clock)
+        self.event_receiver.close()
+        self.event_receiver = event_receiver
+        self._tg.start_soon(self._apply_state)
 
     def unpause(self, result_clock: int):
         logger.info("Unpausing API")
@@ -324,6 +323,7 @@ class API:
         self.app.get("/images/{image_id}")(self.get_image)
         self.app.post("/v1/messages", response_model=None)(self.claude_messages)
         self.app.post("/v1/responses", response_model=None)(self.openai_responses)
+        self.app.post("/v1/cancel/{command_id}")(self.cancel_command)
 
         # Ollama API
         self.app.head("/ollama/")(self.ollama_version)
@@ -344,6 +344,7 @@ class API:
         self.app.post("/download/start")(self.start_download)
         self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download)
         self.app.get("/v1/traces")(self.list_traces)
+        self.app.post("/v1/traces/delete")(self.delete_traces)
         self.app.get("/v1/traces/{task_id}")(self.get_trace)
         self.app.get("/v1/traces/{task_id}/stats")(self.get_trace_stats)
         self.app.get("/v1/traces/{task_id}/raw")(self.get_trace_raw)
@@ -568,6 +569,25 @@ class API:
             instance_id=instance_id,
         )
 
+    async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
+        """Cancel an active command by closing its stream and notifying workers."""
+        sender = self._text_generation_queues.get(
+            command_id
+        ) or self._image_generation_queues.get(command_id)
+        if sender is None:
+            raise HTTPException(
+                status_code=404,
+                detail="Command not found or already completed",
+            )
+
+        await self._send(TaskCancelled(cancelled_command_id=command_id))
+        sender.close()
+
+        return CancelCommandResponse(
+            message="Command cancelled.",
+            command_id=command_id,
+        )
+
     async def _token_chunk_stream(
         self, command_id: CommandId
     ) -> AsyncGenerator[
@@ -755,7 +775,7 @@ class API:
         return resolved_model
 
     def stream_events(self) -> StreamingResponse:
-        def _generate_json_array(events: Iterator[Event]) -> Iterator[str]:
+        def _generate_json_array(events: Iterable[Event]) -> Iterable[str]:
             yield "["
             first = True
             for event in events:
@@ -1566,26 +1586,42 @@ class API:
     async def search_models(
         self, query: str = "", limit: int = 20
     ) -> list[HuggingFaceSearchResult]:
-        """Search HuggingFace Hub for mlx-community models."""
-        from huggingface_hub import list_models
+        """Search HuggingFace Hub — tries mlx-community first, falls back to all of HuggingFace."""
+        from huggingface_hub import ModelInfo, list_models
 
-        results = list_models(
-            search=query or None,
-            author="mlx-community",
-            sort="downloads",
-            limit=limit,
-        )
-        return [
-            HuggingFaceSearchResult(
-                id=m.id,
-                author=m.author or "",
-                downloads=m.downloads or 0,
-                likes=m.likes or 0,
-                last_modified=str(m.last_modified or ""),
-                tags=list(m.tags or []),
+        def _to_results(models: Iterable[ModelInfo]) -> list[HuggingFaceSearchResult]:
+            return [
+                HuggingFaceSearchResult(
+                    id=m.id,
+                    author=m.author or "",
+                    downloads=m.downloads or 0,
+                    likes=m.likes or 0,
+                    last_modified=str(m.last_modified or ""),
+                    tags=list(m.tags or []),
+                )
+                for m in models
+            ]
+
+        # Search mlx-community first
+        mlx_results = _to_results(
+            list_models(
+                search=query or None,
+                author="mlx-community",
+                sort="downloads",
+                limit=limit,
             )
-            for m in results
-        ]
+        )
+        if mlx_results:
+            return mlx_results
+
+        # Fall back to searching all of HuggingFace
+        return _to_results(
+            list_models(
+                search=query or None,
+                sort="downloads",
+                limit=limit,
+            )
+        )
 
     async def run(self):
         shutdown_ev = anyio.Event()
@@ -1606,7 +1642,7 @@ class API:
         finally:
             self._event_log.close()
             self.command_sender.close()
-            self.global_event_receiver.close()
+            self.event_receiver.close()
 
     async def run_api(self, ev: anyio.Event):
         cfg = Config()
@@ -1623,38 +1659,31 @@ class API:
             )
 
     async def _apply_state(self):
-        with self.global_event_receiver as events:
-            async for f_event in events:
-                if f_event.session != self.session_id:
-                    continue
-                if f_event.origin != self.session_id.master_node_id:
-                    continue
-                self.event_buffer.ingest(f_event.origin_idx, f_event.event)
-                for idx, event in self.event_buffer.drain_indexed():
-                    self._event_log.append(event)
-                    self.state = apply(self.state, IndexedEvent(event=event, idx=idx))
+        with self.event_receiver as events:
+            async for i_event in events:
+                self._event_log.append(i_event.event)
+                self.state = apply(self.state, i_event)
+                event = i_event.event
 
-                    if isinstance(event, ChunkGenerated):
-                        if queue := self._image_generation_queues.get(
-                            event.command_id, None
-                        ):
-                            assert isinstance(event.chunk, ImageChunk)
-                            try:
-                                await queue.send(event.chunk)
-                            except BrokenResourceError:
-                                self._image_generation_queues.pop(
-                                    event.command_id, None
-                                )
-                        if queue := self._text_generation_queues.get(
-                            event.command_id, None
-                        ):
-                            assert not isinstance(event.chunk, ImageChunk)
-                            try:
-                                await queue.send(event.chunk)
-                            except BrokenResourceError:
-                                self._text_generation_queues.pop(event.command_id, None)
-                    if isinstance(event, TracesMerged):
-                        self._save_merged_trace(event)
+                if isinstance(event, ChunkGenerated):
+                    if queue := self._image_generation_queues.get(
+                        event.command_id, None
+                    ):
+                        assert isinstance(event.chunk, ImageChunk)
+                        try:
+                            await queue.send(event.chunk)
+                        except BrokenResourceError:
+                            self._image_generation_queues.pop(event.command_id, None)
+                    if queue := self._text_generation_queues.get(
+                        event.command_id, None
+                    ):
+                        assert not isinstance(event.chunk, ImageChunk)
+                        try:
+                            await queue.send(event.chunk)
+                        except BrokenResourceError:
+                            self._text_generation_queues.pop(event.command_id, None)
+                if isinstance(event, TracesMerged):
+                    self._save_merged_trace(event)
 
     def _save_merged_trace(self, event: TracesMerged) -> None:
         traces = [
@@ -1718,8 +1747,12 @@ class API:
         await self._send_download(command)
         return DeleteDownloadResponse(command_id=command.command_id)
 
-    def _get_trace_path(self, task_id: str) -> Path:
-        return EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
+    @staticmethod
+    def _get_trace_path(task_id: str) -> Path:
+        trace_path = EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
+        if not trace_path.resolve().is_relative_to(EXO_TRACING_CACHE_DIR.resolve()):
+            raise HTTPException(status_code=400, detail=f"Invalid task ID: {task_id}")
+        return trace_path
 
     async def list_traces(self) -> TraceListResponse:
         traces: list[TraceListItem] = []
@@ -1818,6 +1851,18 @@ class API:
             filename=f"trace_{task_id}.json",
         )
 
+    async def delete_traces(self, request: DeleteTracesRequest) -> DeleteTracesResponse:
+        deleted: list[str] = []
+        not_found: list[str] = []
+        for task_id in request.task_ids:
+            trace_path = self._get_trace_path(task_id)
+            if trace_path.exists():
+                trace_path.unlink()
+                deleted.append(task_id)
+            else:
+                not_found.append(task_id)
+        return DeleteTracesResponse(deleted=deleted, not_found=not_found)
+
     async def get_onboarding(self) -> JSONResponse:
         return JSONResponse({"completed": ONBOARDING_COMPLETE_FILE.exists()})
 

src/exo/master/main.py

@@ -60,7 +60,7 @@ from exo.shared.types.tasks import (
     TextGeneration as TextGenerationTask,
 )
 from exo.shared.types.worker.instances import InstanceId
-from exo.utils.channels import Receiver, Sender, channel
+from exo.utils.channels import Receiver, Sender
 from exo.utils.event_buffer import MultiSourceBuffer
 from exo.utils.task_group import TaskGroup
 
@@ -72,25 +72,21 @@ class Master:
         session_id: SessionId,
         *,
         command_receiver: Receiver[ForwarderCommand],
+        event_sender: Sender[Event],
         local_event_receiver: Receiver[LocalForwarderEvent],
         global_event_sender: Sender[GlobalForwarderEvent],
         download_command_sender: Sender[ForwarderDownloadCommand],
     ):
-        self.state = State()
-        self._tg: TaskGroup = TaskGroup()
         self.node_id = node_id
         self.session_id = session_id
+        self.state = State()
+        self._tg: TaskGroup = TaskGroup()
         self.command_task_mapping: dict[CommandId, TaskId] = {}
         self.command_receiver = command_receiver
         self.local_event_receiver = local_event_receiver
         self.global_event_sender = global_event_sender
         self.download_command_sender = download_command_sender
-        send, recv = channel[Event]()
-        self.event_sender: Sender[Event] = send
-        self._loopback_event_receiver: Receiver[Event] = recv
-        self._loopback_event_sender: Sender[LocalForwarderEvent] = (
-            local_event_receiver.clone_sender()
-        )
+        self.event_sender = event_sender
         self._system_id = SystemId()
         self._multi_buffer = MultiSourceBuffer[SystemId, Event]()
         self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master")
@@ -104,15 +100,12 @@ class Master:
             async with self._tg as tg:
                 tg.start_soon(self._event_processor)
                 tg.start_soon(self._command_processor)
-                tg.start_soon(self._loopback_processor)
                 tg.start_soon(self._plan)
         finally:
             self._event_log.close()
             self.global_event_sender.close()
             self.local_event_receiver.close()
             self.command_receiver.close()
-            self._loopback_event_sender.close()
-            self._loopback_event_receiver.close()
 
     async def shutdown(self):
         logger.info("Stopping Master")
@@ -335,17 +328,22 @@ class Master:
                                         task_id=task_id,
                                     )
                                 )
-                        case TaskFinished():
-                            generated_events.append(
-                                TaskDeleted(
-                                    task_id=self.command_task_mapping[
-                                        command.finished_command_id
-                                    ]
+                            else:
+                                logger.warning(
+                                    f"Nonexistent command {command.cancelled_command_id} cancelled"
                                 )
-                            )
-                            self.command_task_mapping.pop(
-                                command.finished_command_id, None
-                            )
+                        case TaskFinished():
+                            if (
+                                task_id := self.command_task_mapping.pop(
+                                    command.finished_command_id, None
+                                )
+                            ) is not None:
+                                generated_events.append(TaskDeleted(task_id=task_id))
+                            else:
+                                logger.warning(
+                                    f"Finished command {command.finished_command_id} finished"
+                                )
+
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             # rate limit to 1000 at a time
@@ -409,22 +407,6 @@ class Master:
                     self._event_log.append(event)
                     await self._send_event(indexed)
 
-    async def _loopback_processor(self) -> None:
-        # this would ideally not be necessary.
-        # this is WAY less hacky than how I was working around this before
-        local_index = 0
-        with self._loopback_event_receiver as events:
-            async for event in events:
-                await self._loopback_event_sender.send(
-                    LocalForwarderEvent(
-                        origin=self._system_id,
-                        origin_idx=local_index,
-                        session=self.session_id,
-                        event=event,
-                    )
-                )
-                local_index += 1
-
     # This function is re-entrant, take care!
     async def _send_event(self, event: IndexedEvent):
         # Convenience method since this line is ugly

src/exo/master/tests/test_cancel_command.py

@@ -0,0 +1,77 @@
+# pyright: reportUnusedFunction=false, reportAny=false
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock
+
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from exo.shared.types.common import CommandId
+
+
+def _make_api() -> Any:
+    """Create a minimal API instance with cancel route and error handler."""
+    from exo.master.api import API
+
+    app = FastAPI()
+    api = object.__new__(API)
+    api.app = app
+    api._text_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._image_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._send = AsyncMock()  # pyright: ignore[reportPrivateUsage]
+    api._setup_exception_handlers()  # pyright: ignore[reportPrivateUsage]
+    app.post("/v1/cancel/{command_id}")(api.cancel_command)
+    return api
+
+
+def test_cancel_nonexistent_command_returns_404() -> None:
+    """Cancel for an unknown command_id returns 404 in OpenAI error format."""
+    api = _make_api()
+    client = TestClient(api.app)
+
+    response = client.post("/v1/cancel/nonexistent-id")
+    assert response.status_code == 404
+    data: dict[str, Any] = response.json()
+    assert "error" in data
+    assert data["error"]["message"] == "Command not found or already completed"
+    assert data["error"]["type"] == "Not Found"
+    assert data["error"]["code"] == 404
+
+
+def test_cancel_active_text_generation() -> None:
+    """Cancel an active text generation command: returns 200, sender.close() called."""
+    api = _make_api()
+    client = TestClient(api.app)
+
+    cid = CommandId("text-cmd-123")
+    sender = MagicMock()
+    api._text_generation_queues[cid] = sender
+
+    response = client.post(f"/v1/cancel/{cid}")
+    assert response.status_code == 200
+    data: dict[str, Any] = response.json()
+    assert data["message"] == "Command cancelled."
+    assert data["command_id"] == str(cid)
+    sender.close.assert_called_once()
+    api._send.assert_called_once()
+    task_cancelled = api._send.call_args[0][0]
+    assert task_cancelled.cancelled_command_id == cid
+
+
+def test_cancel_active_image_generation() -> None:
+    """Cancel an active image generation command: returns 200, sender.close() called."""
+    api = _make_api()
+    client = TestClient(api.app)
+
+    cid = CommandId("img-cmd-456")
+    sender = MagicMock()
+    api._image_generation_queues[cid] = sender
+
+    response = client.post(f"/v1/cancel/{cid}")
+    assert response.status_code == 200
+    data: dict[str, Any] = response.json()
+    assert data["message"] == "Command cancelled."
+    assert data["command_id"] == str(cid)
+    sender.close.assert_called_once()
+    api._send.assert_called_once()
+    task_cancelled = api._send.call_args[0][0]
+    assert task_cancelled.cancelled_command_id == cid

src/exo/master/tests/test_master.py

@@ -17,6 +17,7 @@ from exo.shared.types.commands import (
 )
 from exo.shared.types.common import ModelId, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
+    Event,
     GlobalForwarderEvent,
     IndexedEvent,
     InstanceCreated,
@@ -50,6 +51,22 @@ async def test_master():
     command_sender, co_receiver = channel[ForwarderCommand]()
     local_event_sender, le_receiver = channel[LocalForwarderEvent]()
     fcds, _fcdr = channel[ForwarderDownloadCommand]()
+    ev_send, ev_recv = channel[Event]()
+
+    async def mock_event_router():
+        idx = 0
+        sid = SystemId()
+        with ev_recv as master_events:
+            async for event in master_events:
+                await local_event_sender.send(
+                    LocalForwarderEvent(
+                        origin=sid,
+                        origin_idx=idx,
+                        session=session_id,
+                        event=event,
+                    )
+                )
+                idx += 1
 
     all_events: list[IndexedEvent] = []
 
@@ -67,6 +84,7 @@ async def test_master():
     master = Master(
         node_id,
         session_id,
+        event_sender=ev_send,
         global_event_sender=ge_sender,
         local_event_receiver=le_receiver,
         command_receiver=co_receiver,
@@ -75,6 +93,7 @@ async def test_master():
     logger.info("run the master")
     async with anyio.create_task_group() as tg:
         tg.start_soon(master.run)
+        tg.start_soon(mock_event_router)
 
         # inject a NodeGatheredInfo event
         logger.info("inject a NodeGatheredInfo event")
@@ -197,4 +216,5 @@ async def test_master():
             input=[InputMessage(role="user", content="Hello, how are you?")],
         )
 
+        ev_send.close()
         await master.shutdown()

src/exo/routing/connection_message.py

@@ -1,6 +1,4 @@
-from enum import Enum
-
-from exo_pyo3_bindings import ConnectionUpdate, ConnectionUpdateType
+from exo_pyo3_bindings import PyFromSwarm
 
 from exo.shared.types.common import NodeId
 from exo.utils.pydantic_ext import CamelCaseModel
@@ -8,30 +6,10 @@ from exo.utils.pydantic_ext import CamelCaseModel
 """Serialisable types for Connection Updates/Messages"""
 
 
-class ConnectionMessageType(Enum):
-    Connected = 0
-    Disconnected = 1
-
-    @staticmethod
-    def from_update_type(update_type: ConnectionUpdateType):
-        match update_type:
-            case ConnectionUpdateType.Connected:
-                return ConnectionMessageType.Connected
-            case ConnectionUpdateType.Disconnected:
-                return ConnectionMessageType.Disconnected
-
-
 class ConnectionMessage(CamelCaseModel):
     node_id: NodeId
-    connection_type: ConnectionMessageType
-    remote_ipv4: str
-    remote_tcp_port: int
+    connected: bool
 
     @classmethod
-    def from_update(cls, update: ConnectionUpdate) -> "ConnectionMessage":
-        return cls(
-            node_id=NodeId(update.peer_id),
-            connection_type=ConnectionMessageType.from_update_type(update.update_type),
-            remote_ipv4=update.remote_ipv4,
-            remote_tcp_port=update.remote_tcp_port,
-        )
+    def from_update(cls, update: PyFromSwarm.Connection) -> "ConnectionMessage":
+        return cls(node_id=NodeId(update.peer_id), connected=update.connected)

src/exo/routing/event_router.py

@@ -0,0 +1,161 @@
+from dataclasses import dataclass, field
+from random import random
+
+import anyio
+from anyio import BrokenResourceError, ClosedResourceError
+from anyio.abc import CancelScope
+from loguru import logger
+
+from exo.shared.types.commands import ForwarderCommand, RequestEventLog
+from exo.shared.types.common import SessionId, SystemId
+from exo.shared.types.events import (
+    Event,
+    EventId,
+    GlobalForwarderEvent,
+    IndexedEvent,
+    LocalForwarderEvent,
+)
+from exo.utils.channels import Receiver, Sender, channel
+from exo.utils.event_buffer import OrderedBuffer
+from exo.utils.task_group import TaskGroup
+
+
+@dataclass
+class EventRouter:
+    session_id: SessionId
+    command_sender: Sender[ForwarderCommand]
+    external_inbound: Receiver[GlobalForwarderEvent]
+    external_outbound: Sender[LocalForwarderEvent]
+    _system_id: SystemId = field(init=False, default_factory=SystemId)
+    internal_outbound: list[Sender[IndexedEvent]] = field(
+        init=False, default_factory=list
+    )
+    event_buffer: OrderedBuffer[Event] = field(
+        init=False, default_factory=OrderedBuffer
+    )
+    out_for_delivery: dict[EventId, tuple[float, LocalForwarderEvent]] = field(
+        init=False, default_factory=dict
+    )
+    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
+
+    _nack_cancel_scope: CancelScope | None = field(init=False, default=None)
+    _nack_attempts: int = field(init=False, default=0)
+    _nack_base_seconds: float = field(init=False, default=0.5)
+    _nack_cap_seconds: float = field(init=False, default=10.0)
+
+    async def run(self):
+        try:
+            async with self._tg as tg:
+                tg.start_soon(self._run_ext_in)
+                tg.start_soon(self._simple_retry)
+        finally:
+            self.external_outbound.close()
+            for send in self.internal_outbound:
+                send.close()
+
+    # can make this better in future
+    async def _simple_retry(self):
+        while True:
+            await anyio.sleep(1 + random())
+            # list here is a shallow clone for shared mutation
+            for e_id, (time, event) in list(self.out_for_delivery.items()):
+                if anyio.current_time() > time + 5:
+                    self.out_for_delivery[e_id] = (anyio.current_time(), event)
+                    await self.external_outbound.send(event)
+
+    def sender(self) -> Sender[Event]:
+        send, recv = channel[Event]()
+        if self._tg.is_running():
+            self._tg.start_soon(self._ingest, SystemId(), recv)
+        else:
+            self._tg.queue(self._ingest, SystemId(), recv)
+        return send
+
+    def receiver(self) -> Receiver[IndexedEvent]:
+        send, recv = channel[IndexedEvent]()
+        self.internal_outbound.append(send)
+        return recv
+
+    def shutdown(self) -> None:
+        self._tg.cancel_tasks()
+
+    async def _ingest(self, system_id: SystemId, recv: Receiver[Event]):
+        idx = 0
+        with recv as events:
+            async for event in events:
+                f_ev = LocalForwarderEvent(
+                    origin_idx=idx,
+                    origin=system_id,
+                    session=self.session_id,
+                    event=event,
+                )
+                idx += 1
+                await self.external_outbound.send(f_ev)
+                self.out_for_delivery[event.event_id] = (anyio.current_time(), f_ev)
+
+    async def _run_ext_in(self):
+        buf = OrderedBuffer[Event]()
+        with self.external_inbound as events:
+            async for event in events:
+                if event.session != self.session_id:
+                    continue
+                if event.origin != self.session_id.master_node_id:
+                    continue
+
+                buf.ingest(event.origin_idx, event.event)
+                event_id = event.event.event_id
+                if event_id in self.out_for_delivery:
+                    self.out_for_delivery.pop(event_id)
+
+                drained = buf.drain_indexed()
+                if drained:
+                    self._nack_attempts = 0
+                    if self._nack_cancel_scope:
+                        self._nack_cancel_scope.cancel()
+
+                if not drained and (
+                    self._nack_cancel_scope is None
+                    or self._nack_cancel_scope.cancel_called
+                ):
+                    # Request the next index.
+                    self._tg.start_soon(self._nack_request, buf.next_idx_to_release)
+                    continue
+
+                for idx, event in drained:
+                    to_clear = set[int]()
+                    for i, sender in enumerate(self.internal_outbound):
+                        try:
+                            await sender.send(IndexedEvent(idx=idx, event=event))
+                        except (ClosedResourceError, BrokenResourceError):
+                            to_clear.add(i)
+                    for i in sorted(to_clear, reverse=True):
+                        self.internal_outbound.pop(i)
+
+    async def _nack_request(self, since_idx: int) -> None:
+        # We request all events after (and including) the missing index.
+        # This function is started whenever we receive an event that is out of sequence.
+        # It is cancelled as soon as we receiver an event that is in sequence.
+
+        if since_idx < 0:
+            logger.warning(f"Negative value encountered for nack request {since_idx=}")
+            since_idx = 0
+
+        with CancelScope() as scope:
+            self._nack_cancel_scope = scope
+            delay: float = self._nack_base_seconds * (2.0**self._nack_attempts)
+            delay = min(self._nack_cap_seconds, delay)
+            self._nack_attempts += 1
+            try:
+                await anyio.sleep(delay)
+                logger.info(
+                    f"Nack attempt {self._nack_attempts}: Requesting Event Log from {since_idx}"
+                )
+                await self.command_sender.send(
+                    ForwarderCommand(
+                        origin=self._system_id,
+                        command=RequestEventLog(since_idx=since_idx),
+                    )
+                )
+            finally:
+                if self._nack_cancel_scope is scope:
+                    self._nack_cancel_scope = None

src/exo/routing/router.py

@@ -17,6 +17,7 @@ from exo_pyo3_bindings import (
     MessageTooLargeError,
     NetworkingHandle,
     NoPeersSubscribedToTopicError,
+    PyFromSwarm,
 )
 from filelock import FileLock
 from loguru import logger
@@ -121,7 +122,8 @@ class Router:
             send = self.networking_receiver.clone_sender()
         router = TopicRouter[T](topic, send)
         self.topic_routers[topic.topic] = cast(TopicRouter[CamelCaseModel], router)
-        await self._networking_subscribe(str(topic.topic))
+        if self._tg.is_running():
+            await self._networking_subscribe(topic.topic)
 
     def sender[T: CamelCaseModel](self, topic: TypedTopic[T]) -> Sender[T]:
         router = self.topic_routers.get(topic.topic, None)
@@ -152,8 +154,10 @@ class Router:
                     router = self.topic_routers[topic]
                     tg.start_soon(router.run)
                 tg.start_soon(self._networking_recv)
-                tg.start_soon(self._networking_recv_connection_messages)
                 tg.start_soon(self._networking_publish)
+                # subscribe to pending topics
+                for topic in self.topic_routers:
+                    await self._networking_subscribe(topic)
                 # Router only shuts down if you cancel it.
                 await sleep_forever()
         finally:
@@ -176,46 +180,49 @@ class Router:
     async def _networking_recv(self):
         try:
             while True:
-                topic, data = await self._net.gossipsub_recv()
-                logger.trace(f"Received message on {topic} with payload {data}")
-                if topic not in self.topic_routers:
-                    logger.warning(
-                        f"Received message on unknown or inactive topic {topic}"
-                    )
-                    continue
-
-                router = self.topic_routers[topic]
-                await router.publish_bytes(data)
+                from_swarm = await self._net.recv()
+                logger.debug(from_swarm)
+                match from_swarm:
+                    case PyFromSwarm.Message(origin, topic, data):
+                        logger.trace(
+                            f"Received message on {topic} from {origin} with payload {data}"
+                        )
+                        if topic not in self.topic_routers:
+                            logger.warning(
+                                f"Received message on unknown or inactive topic {topic}"
+                            )
+                            continue
+                        router = self.topic_routers[topic]
+                        await router.publish_bytes(data)
+                    case PyFromSwarm.Connection():
+                        message = ConnectionMessage.from_update(from_swarm)
+                        logger.trace(
+                            f"Received message on connection_messages with payload {message}"
+                        )
+                        if CONNECTION_MESSAGES.topic in self.topic_routers:
+                            router = self.topic_routers[CONNECTION_MESSAGES.topic]
+                            assert router.topic.model_type == ConnectionMessage
+                            router = cast(TopicRouter[ConnectionMessage], router)
+                            await router.publish(message)
+                    case _:
+                        logger.critical(
+                            "failed to exhaustively check FromSwarm messages - logic error"
+                        )
         except Exception as exception:
             logger.opt(exception=exception).error(
                 "Gossipsub receive loop terminated unexpectedly"
             )
             raise
 
-    async def _networking_recv_connection_messages(self):
-        try:
-            while True:
-                update = await self._net.connection_update_recv()
-                message = ConnectionMessage.from_update(update)
-                logger.trace(
-                    f"Received message on connection_messages with payload {message}"
-                )
-                if CONNECTION_MESSAGES.topic in self.topic_routers:
-                    router = self.topic_routers[CONNECTION_MESSAGES.topic]
-                    assert router.topic.model_type == ConnectionMessage
-                    router = cast(TopicRouter[ConnectionMessage], router)
-                    await router.publish(message)
-        except Exception as exception:
-            logger.opt(exception=exception).error(
-                "Connection update receive loop terminated unexpectedly"
-            )
-            raise
-
     async def _networking_publish(self):
         with self.networking_receiver as networked_items:
             async for topic, data in networked_items:
                 try:
                     logger.trace(f"Sending message on {topic} with payload {data}")
+                    if len(data) > 1024 * 1024:
+                        logger.warning(
+                            "Sending overlarge payload, network performance may be temporarily degraded"
+                        )
                     await self._net.gossipsub_publish(topic, data)
                 except NoPeersSubscribedToTopicError:
                     pass
@@ -255,6 +262,6 @@ def get_node_id_keypair(
 
         # if no valid credentials, create new ones and persist
         with open(path, "w+b") as f:
-            keypair = Keypair.generate_ed25519()
+            keypair = Keypair.generate()
             f.write(keypair.to_bytes())
             return keypair

src/exo/shared/constants.py

@@ -81,3 +81,5 @@ EXO_ENABLE_IMAGE_MODELS = (
 EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true"
 
 EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true"
+
+EXO_MAX_CONCURRENT_REQUESTS = int(os.getenv("EXO_MAX_CONCURRENT_REQUESTS", "8"))

src/exo/shared/models/model_cards.py

@@ -190,6 +190,8 @@ class ConfigData(BaseModel):
             ["DeepseekV3ForCausalLM"],
             ["Qwen3NextForCausalLM"],
             ["Qwen3MoeForCausalLM"],
+            ["Qwen3_5MoeForConditionalGeneration"],
+            ["Qwen3_5ForConditionalGeneration"],
             ["MiniMaxM2ForCausalLM"],
             ["LlamaForCausalLM"],
             ["GptOssForCausalLM"],

src/exo/shared/tests/test_election.py

@@ -1,7 +1,7 @@
 import pytest
 from anyio import create_task_group, fail_after, move_on_after
 
-from exo.routing.connection_message import ConnectionMessage, ConnectionMessageType
+from exo.routing.connection_message import ConnectionMessage
 from exo.shared.election import Election, ElectionMessage, ElectionResult
 from exo.shared.types.commands import ForwarderCommand, TestCommand
 from exo.shared.types.common import NodeId, SessionId, SystemId
@@ -327,14 +327,7 @@ async def test_connection_message_triggers_new_round_broadcast() -> None:
             tg.start_soon(election.run)
 
             # Send any connection message object; we close quickly to cancel before result creation
-            await cm_tx.send(
-                ConnectionMessage(
-                    node_id=NodeId(),
-                    connection_type=ConnectionMessageType.Connected,
-                    remote_ipv4="",
-                    remote_tcp_port=0,
-                )
-            )
+            await cm_tx.send(ConnectionMessage(node_id=NodeId(), connected=True))
 
             # Expect a broadcast for the new round at clock=1
             while True:

src/exo/shared/tests/test_resolve_reasoning_params.py

@@ -0,0 +1,30 @@
+import pytest
+
+from exo.shared.types.text_generation import resolve_reasoning_params
+
+
+def test_both_none_returns_none_none() -> None:
+    assert resolve_reasoning_params(None, None) == (None, None)
+
+
+def test_both_set_passes_through_unchanged() -> None:
+    assert resolve_reasoning_params("high", True) == ("high", True)
+    assert resolve_reasoning_params("none", True) == ("none", True)
+    assert resolve_reasoning_params("low", False) == ("low", False)
+
+
+def test_enable_thinking_true_derives_medium() -> None:
+    assert resolve_reasoning_params(None, True) == ("medium", True)
+
+
+def test_enable_thinking_false_derives_none() -> None:
+    assert resolve_reasoning_params(None, False) == ("none", False)
+
+
+def test_reasoning_effort_none_derives_thinking_false() -> None:
+    assert resolve_reasoning_params("none", None) == ("none", False)
+
+
+@pytest.mark.parametrize("effort", ["minimal", "low", "medium", "high", "xhigh"])
+def test_non_none_effort_derives_thinking_true(effort: str) -> None:
+    assert resolve_reasoning_params(effort, None) == (effort, True)  # pyright: ignore[reportArgumentType]

src/exo/shared/types/api.py

@@ -8,6 +8,7 @@ from pydantic import BaseModel, Field, field_validator
 from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.memory import Memory
+from exo.shared.types.text_generation import ReasoningEffort
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel
@@ -198,7 +199,11 @@ class ChatCompletionRequest(BaseModel):
     top_p: float | None = None
     top_k: int | None = None
     tools: list[dict[str, Any]] | None = None
+    reasoning_effort: ReasoningEffort | None = None
     enable_thinking: bool | None = None
+    min_p: float | None = None
+    repetition_penalty: float | None = None
+    repetition_context_size: int | None = None
     tool_choice: str | dict[str, Any] | None = None
     parallel_tool_calls: bool | None = None
     user: str | None = None
@@ -262,6 +267,11 @@ class DeleteInstanceResponse(BaseModel):
     instance_id: InstanceId
 
 
+class CancelCommandResponse(BaseModel):
+    message: str
+    command_id: CommandId
+
+
 ImageSize = Literal[
     "auto",
     "512x512",
@@ -437,3 +447,12 @@ class TraceListItem(CamelCaseModel):
 
 class TraceListResponse(CamelCaseModel):
     traces: list[TraceListItem]
+
+
+class DeleteTracesRequest(CamelCaseModel):
+    task_ids: list[str]
+
+
+class DeleteTracesResponse(CamelCaseModel):
+    deleted: list[str]
+    not_found: list[str]

src/exo/shared/types/mlx.py

@@ -2,6 +2,8 @@
 
 from collections.abc import Sequence
 
+from mlx import core as mx
+from mlx import nn as nn
 from mlx_lm.models.cache import (
     ArraysCache,
     CacheList,
@@ -14,3 +16,16 @@ from mlx_lm.models.cache import (
 KVCacheType = Sequence[
     KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList
 ]
+
+
+# Model is a wrapper function to fix the fact that mlx is not strongly typed in the same way that EXO is.
+# For example - MLX has no guarantee of the interface that nn.Module will expose. But we need a guarantee that it has a __call__() function
+class Model(nn.Module):
+    layers: list[nn.Module]
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache: KVCacheType | None,
+        input_embeddings: mx.array | None = None,
+    ) -> mx.array: ...

src/exo/shared/types/openai_responses.py

@@ -12,6 +12,7 @@ from typing import Any, Literal
 from pydantic import BaseModel, Field
 
 from exo.shared.types.common import ModelId
+from exo.shared.types.text_generation import ReasoningEffort
 
 # Type aliases
 ResponseStatus = Literal["completed", "failed", "in_progress", "incomplete"]
@@ -71,6 +72,13 @@ ResponseInputItem = (
 )
 
 
+class Reasoning(BaseModel, frozen=True):
+    """Reasoning configuration for OpenAI Responses API."""
+
+    effort: ReasoningEffort | None = None
+    summary: Literal["auto", "concise", "detailed"] | None = None
+
+
 class ResponsesRequest(BaseModel, frozen=True):
     """Request body for OpenAI Responses API.
 
@@ -89,8 +97,15 @@ class ResponsesRequest(BaseModel, frozen=True):
     stream: bool = False
     tools: list[dict[str, Any]] | None = None
     metadata: dict[str, str] | None = None
+    reasoning: Reasoning | None = None
 
     # --- exo extensions (not in OpenAI Responses API spec) ---
+    enable_thinking: bool | None = Field(
+        default=None,
+        description="[exo extension] Boolean thinking toggle. Not part of the OpenAI Responses API.",
+        json_schema_extra={"x-exo-extension": True},
+    )
+
     top_k: int | None = Field(
         default=None,
         description="[exo extension] Top-k sampling parameter. Not part of the OpenAI Responses API.",

src/exo/shared/types/tasks.py

@@ -18,6 +18,9 @@ class TaskId(Id):
     pass
 
 
+CANCEL_ALL_TASKS = TaskId("CANCEL_ALL_TASKS")
+
+
 class TaskStatus(str, Enum):
     Pending = "Pending"
     Running = "Running"

src/exo/shared/types/text_generation.py

@@ -11,6 +11,29 @@ from pydantic import BaseModel
 from exo.shared.types.common import ModelId
 
 MessageRole = Literal["user", "assistant", "system", "developer"]
+ReasoningEffort = Literal["none", "minimal", "low", "medium", "high", "xhigh"]
+
+
+def resolve_reasoning_params(
+    reasoning_effort: ReasoningEffort | None,
+    enable_thinking: bool | None,
+) -> tuple[ReasoningEffort | None, bool | None]:
+    """
+    enable_thinking=True  -> reasoning_effort="medium"
+    enable_thinking=False -> reasoning_effort="none"
+    reasoning_effort="none" -> enable_thinking=False
+    reasoning_effort=<anything else> -> enable_thinking=True
+    """
+    resolved_effort: ReasoningEffort | None = reasoning_effort
+    resolved_thinking: bool | None = enable_thinking
+
+    if reasoning_effort is None and enable_thinking is not None:
+        resolved_effort = "medium" if enable_thinking else "none"
+
+    if enable_thinking is None and reasoning_effort is not None:
+        resolved_thinking = reasoning_effort != "none"
+
+    return resolved_effort, resolved_thinking
 
 
 class InputMessage(BaseModel, frozen=True):
@@ -40,6 +63,10 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     stop: str | list[str] | None = None
     seed: int | None = None
     chat_template_messages: list[dict[str, Any]] | None = None
+    reasoning_effort: ReasoningEffort | None = None
     enable_thinking: bool | None = None
     logprobs: bool = False
     top_logprobs: int | None = None
+    min_p: float | None = None
+    repetition_penalty: float | None = None
+    repetition_context_size: int | None = None

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

@@ -1,17 +0,0 @@
-import mlx.core as mx
-import mlx.nn as nn
-from mlx_lm.models.cache import KVCache
-
-# These are wrapper functions to fix the fact that mlx is not strongly typed in the same way that EXO is.
-# For example - MLX has no guarantee of the interface that nn.Module will expose. But we need a guarantee that it has a __call__() function
-
-
-class Model(nn.Module):
-    layers: list[nn.Module]
-
-    def __call__(
-        self,
-        x: mx.array,
-        cache: list[KVCache] | None,
-        input_embeddings: mx.array | None = None,
-    ) -> mx.array: ...

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

@@ -16,6 +16,7 @@ from mlx.nn.layers.distributed import (
 from mlx_lm.models.base import (
     scaled_dot_product_attention,  # pyright: ignore[reportUnknownVariableType]
 )
+from mlx_lm.models.cache import ArraysCache, KVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3MLP
 from mlx_lm.models.deepseek_v3 import Model as DeepseekV3Model
 from mlx_lm.models.deepseek_v32 import DeepseekV32MLP
@@ -31,10 +32,19 @@ from mlx_lm.models.llama import Model as LlamaModel
 from mlx_lm.models.minimax import MiniMaxAttention
 from mlx_lm.models.minimax import Model as MiniMaxModel
 from mlx_lm.models.ministral3 import Model as Ministral3Model
+from mlx_lm.models.qwen3_5 import DecoderLayer as Qwen3_5DecoderLayer
+from mlx_lm.models.qwen3_5 import Model as Qwen3_5TextModel
+from mlx_lm.models.qwen3_5 import Qwen3_5TextModel as Qwen3_5TextModelInner
+from mlx_lm.models.qwen3_5 import SparseMoeBlock as Qwen3_5SparseMoeBlock
+from mlx_lm.models.qwen3_5_moe import Model as Qwen3_5MoeModel
 from mlx_lm.models.qwen3_moe import Model as Qwen3MoeModel
 from mlx_lm.models.qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeSparseMoeBlock
 from mlx_lm.models.qwen3_next import Model as Qwen3NextModel
-from mlx_lm.models.qwen3_next import Qwen3NextDecoderLayer, Qwen3NextSparseMoeBlock
+from mlx_lm.models.qwen3_next import (
+    Qwen3NextDecoderLayer,
+    Qwen3NextGatedDeltaNet,
+    Qwen3NextSparseMoeBlock,
+)
 from mlx_lm.models.step3p5 import Model as Step35Model
 from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP
 from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel
@@ -49,6 +59,21 @@ TimeoutCallback = Callable[[], None]
 LayerLoadedCallback = Callable[[int, int], None]  # (layers_loaded, total_layers)
 
 
+_pending_prefill_sends: list[tuple[mx.array, int, mx.distributed.Group]] = []
+
+
+def flush_prefill_sends() -> None:
+    for output, dst, group in _pending_prefill_sends:
+        sent = mx.distributed.send(output, dst, group=group)
+        mx.async_eval(sent)
+    _pending_prefill_sends.clear()
+
+
+def clear_prefill_sends() -> None:
+    # Discard pending sends (e.g. on cancellation).
+    _pending_prefill_sends.clear()
+
+
 def eval_with_timeout(
     mlx_item: Any,  # pyright: ignore[reportAny]
     timeout_seconds: float = 60.0,
@@ -150,6 +175,7 @@ class PipelineLastLayer(CustomMlxLayer):
         self.group = group
         self.original_layer_signature = signature(self.original_layer.__call__)
         self.is_prefill: bool = False
+        self.queue_sends: bool = False
 
     def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
         cache = self.original_layer_signature.bind_partial(
@@ -163,16 +189,22 @@ class PipelineLastLayer(CustomMlxLayer):
         mx.eval(output)
 
         if self.r != self.s - 1:
-            output = mx.distributed.send(
-                output, (self.r + 1) % self.s, group=self.group
-            )
+            if self.queue_sends:
+                _pending_prefill_sends.append(
+                    (output, (self.r + 1) % self.s, self.group)
+                )
+            else:
+                output = mx.distributed.send(
+                    output, (self.r + 1) % self.s, group=self.group
+                )
             if cache is not None:
                 # CacheList (used by MLA models like DeepSeekV32, GLM MoE DSA)
                 # doesn't have .keys directly; access via first sub-cache.
                 _cache = cache[0] if hasattr(cache, "caches") else cache  # type: ignore
-                _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
+                if hasattr(_cache, "keys"):  # pyright: ignore[reportAny]
+                    _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
             mx.eval(output)
-            if cache is not None:
+            if cache is not None and hasattr(_cache, "keys"):  # type: ignore
                 mx.eval(_cache.keys)  # type: ignore
 
         if not self.is_prefill:
@@ -190,6 +222,12 @@ def set_pipeline_prefill(model: nn.Module, is_prefill: bool) -> None:
             layer.is_prefill = is_prefill
 
 
+def set_pipeline_queue_sends(model: nn.Module, queue_sends: bool) -> None:
+    for layer in model.layers:  # type: ignore
+        if isinstance(layer, PipelineLastLayer):
+            layer.queue_sends = queue_sends
+
+
 def get_inner_model(model: nn.Module) -> nn.Module:
     inner = getattr(model, "model", None)
     if isinstance(inner, nn.Module):
@@ -221,6 +259,32 @@ def get_layers(inner_model_instance: nn.Module) -> list[_LayerCallable]:
     return layers
 
 
+def _patch_qwen35_cache(
+    model: Qwen3_5TextModel,
+    fa_idx: int,
+    has_full_attn: bool,
+    ssm_idx: int,
+    has_linear: bool,
+) -> None:
+    # Hacks to make make_mask happy.
+    original = model.make_cache
+
+    def patched() -> list[ArraysCache | KVCache]:
+        cache: list[ArraysCache | KVCache] = original()
+        if not has_full_attn:
+            entry = cache[fa_idx]
+            orig_make_mask = entry.make_mask
+            entry.make_mask = lambda n, **_kw: orig_make_mask(n)  # type: ignore
+        if not has_linear:
+            orig_ssm_make_mask = cache[ssm_idx].make_mask
+            cache[ssm_idx].make_mask = (  # type: ignore
+                lambda n, **kw: orig_ssm_make_mask(n, **kw) if kw else None  # type: ignore
+            )
+        return cache
+
+    model.make_cache = patched
+
+
 def pipeline_auto_parallel(
     model: nn.Module,
     group: mx.distributed.Group,
@@ -291,6 +355,24 @@ def pipeline_auto_parallel(
         inner_model_instance._swa_idx = 0 if not sliding_layers else sliding_layers[0]
         inner_model_instance._full_idx = 0 if not full_layers else full_layers[0]
 
+    if isinstance(inner_model_instance, Qwen3_5TextModelInner):
+        full_attn_layers = [
+            i for i, layer in enumerate(layers) if not getattr(layer, "is_linear", True)
+        ]
+        linear_layers = [
+            i for i, layer in enumerate(layers) if getattr(layer, "is_linear", False)
+        ]
+        inner_model_instance.fa_idx = full_attn_layers[0] if full_attn_layers else 0
+        inner_model_instance.ssm_idx = linear_layers[0] if linear_layers else 0
+        if not full_attn_layers or not linear_layers:
+            _patch_qwen35_cache(
+                cast(Qwen3_5TextModel, model),
+                fa_idx=inner_model_instance.fa_idx,
+                has_full_attn=bool(full_attn_layers),
+                ssm_idx=inner_model_instance.ssm_idx,
+                has_linear=bool(linear_layers),
+            )
+
     _set_layers(model, layers)
 
     assert isinstance(layers, list), (
@@ -320,7 +402,8 @@ def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
         if cache is not None:
             last = cache[-1]  # type: ignore
             dep_cache = last[0] if hasattr(last, "caches") else last  # type: ignore
-            dep_cache.keys = mx.depends(dep_cache.keys, logits)  # type: ignore
+            if hasattr(dep_cache, "keys") and dep_cache.keys is not None:  # type: ignore
+                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # type: ignore
 
         return logits
 
@@ -443,7 +526,9 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
-    elif isinstance(model, (Qwen3MoeModel, Qwen3NextModel)):
+    elif isinstance(
+        model, (Qwen3MoeModel, Qwen3NextModel, Qwen3_5TextModel, Qwen3_5MoeModel)
+    ):
         tensor_parallel_sharding_strategy = QwenShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -838,7 +923,9 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
         on_timeout: TimeoutCallback | None,
         on_layer_loaded: LayerLoadedCallback | None,
     ) -> nn.Module:
-        model = cast(Qwen3MoeModel | Qwen3NextModel, model)
+        model = cast(
+            Qwen3MoeModel | Qwen3NextModel | Qwen3_5TextModel | Qwen3_5MoeModel, model
+        )
         total = len(model.layers)
         for i, layer in enumerate(model.layers):
             eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
@@ -859,16 +946,39 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 layer.self_attn.n_heads //= self.N
                 layer.self_attn.n_kv_heads //= self.N
             else:
-                assert isinstance(layer, Qwen3NextDecoderLayer)
+                assert isinstance(layer, (Qwen3NextDecoderLayer, Qwen3_5DecoderLayer))
                 if hasattr(layer, "linear_attn"):
                     linear_attn = layer.linear_attn
 
-                    linear_attn.in_proj_qkvz = self.all_to_sharded_linear(
-                        linear_attn.in_proj_qkvz
-                    )
-                    linear_attn.in_proj_ba = self.all_to_sharded_linear(
-                        linear_attn.in_proj_ba
-                    )
+                    if isinstance(linear_attn, Qwen3NextGatedDeltaNet):
+                        # Qwen3-Next: combined projections
+                        linear_attn.in_proj_qkvz = self.all_to_sharded_linear(
+                            linear_attn.in_proj_qkvz
+                        )
+                        linear_attn.in_proj_ba = self.all_to_sharded_linear(
+                            linear_attn.in_proj_ba
+                        )
+                    else:
+                        # Qwen3.5: separate projections
+                        # in_proj_qkv has sections [q(key_dim), k(key_dim), v(value_dim)]
+                        # that must be split section-aware, not as a contiguous block
+                        key_dim = linear_attn.key_dim
+                        value_dim = linear_attn.value_dim
+                        linear_attn.in_proj_qkv = shard_linear(
+                            linear_attn.in_proj_qkv,
+                            "all-to-sharded",
+                            segments=[key_dim, key_dim + key_dim],
+                            group=self.group,
+                        )
+                        linear_attn.in_proj_z = self.all_to_sharded_linear(
+                            linear_attn.in_proj_z
+                        )
+                        linear_attn.in_proj_b = self.all_to_sharded_linear(
+                            linear_attn.in_proj_b
+                        )
+                        linear_attn.in_proj_a = self.all_to_sharded_linear(
+                            linear_attn.in_proj_a
+                        )
                     linear_attn.out_proj = self.sharded_to_all_linear(
                         linear_attn.out_proj
                     )
@@ -930,11 +1040,20 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                     layer.self_attn.num_key_value_heads //= self.N
 
             # Shard the MoE.
-            if isinstance(layer.mlp, (Qwen3MoeSparseMoeBlock, Qwen3NextSparseMoeBlock)):
+            if isinstance(
+                layer.mlp,
+                (
+                    Qwen3MoeSparseMoeBlock,
+                    Qwen3NextSparseMoeBlock,
+                    Qwen3_5SparseMoeBlock,
+                ),
+            ):
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
                 self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
-                if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):
+                if isinstance(
+                    layer.mlp, (Qwen3NextSparseMoeBlock, Qwen3_5SparseMoeBlock)
+                ):
                     self.all_to_sharded_linear_in_place(
                         layer.mlp.shared_expert.gate_proj
                     )

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

@@ -13,8 +13,7 @@ from mlx_lm.models.cache import (
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType
-from exo.worker.engines.mlx import Model
+from exo.shared.types.mlx import KVCacheType, Model
 from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
 from exo.worker.runner.bootstrap import logger
 
@@ -254,9 +253,9 @@ def trim_cache(
             if snapshot is not None and snapshot.states[i] is not None:
                 cache[i] = deepcopy(snapshot.states[i])  # type: ignore
             else:
-                c.state = [None] * len(c.state)  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
+                c.state = [None] * len(c.state)
         else:
-            c.trim(num_tokens)  # pyright: ignore[reportUnknownMemberType]
+            c.trim(num_tokens)
 
 
 def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:

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

@@ -0,0 +1,383 @@
+import time
+from dataclasses import dataclass, field
+from typing import Callable, cast
+
+import mlx.core as mx
+from mlx_lm.generate import (
+    BatchGenerator as MlxBatchGenerator,
+)
+from mlx_lm.models.cache import RotatingKVCache
+from mlx_lm.sample_utils import make_logits_processors, make_sampler
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+from exo.shared.types.api import (
+    CompletionTokensDetails,
+    FinishReason,
+    GenerationStats,
+    PromptTokensDetails,
+    TopLogprobItem,
+    Usage,
+)
+from exo.shared.types.memory import Memory
+from exo.shared.types.mlx import KVCacheType, Model
+from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.worker.runner_response import GenerationResponse
+from exo.worker.engines.mlx.cache import (
+    CacheSnapshot,
+    KVPrefixCache,
+    encode_prompt,
+    make_kv_cache,
+)
+from exo.worker.engines.mlx.constants import DEFAULT_TOP_LOGPROBS, MAX_TOKENS
+from exo.worker.engines.mlx.generator.generate import (
+    ban_token_ids,
+    eos_ids_from_tokenizer,
+    extract_top_logprobs,
+    prefill,
+)
+from exo.worker.engines.mlx.utils_mlx import fix_unmatched_think_end_tokens
+from exo.worker.runner.bootstrap import logger
+
+_MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5
+
+
+def _stop_sequences(task_params: TextGenerationTaskParams) -> list[str]:
+    if task_params.stop is None:
+        return []
+    if isinstance(task_params.stop, str):
+        return [task_params.stop]
+    return task_params.stop
+
+
+@dataclass
+class _EngineTask:
+    uid: int
+    task_params: TextGenerationTaskParams
+    all_prompt_tokens: mx.array
+    prefix_hit_length: int
+    matched_index: int | None
+    cache_snapshots: list[CacheSnapshot] | None
+    on_generation_token: Callable[[], None] | None = None
+    generated_text_parts: list[str] = field(default_factory=list)
+    potential_stop_sequence_text: str = ""
+    completion_tokens: int = 0
+    generation_start_time: float = 0.0
+    in_thinking: bool = False
+    reasoning_tokens: int = 0
+
+
+@dataclass(eq=False)
+class ExoBatchGenerator:
+    model: Model
+    tokenizer: TokenizerWrapper
+    group: mx.distributed.Group | None
+    kv_prefix_cache: KVPrefixCache | None
+
+    _exo_gen: MlxBatchGenerator = field(init=False)
+    _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False)
+
+    def __post_init__(self) -> None:
+        self._exo_gen = MlxBatchGenerator(
+            model=self.model,
+            stop_tokens=set(eos_ids_from_tokenizer(self.tokenizer)),
+            prefill_step_size=4096,
+        )
+
+    @property
+    def has_work(self) -> bool:
+        return (
+            bool(self._active_tasks)
+            or bool(self._exo_gen.unprocessed_prompts)
+            or self._exo_gen.active_batch is not None
+        )
+
+    def submit(
+        self,
+        task_params: TextGenerationTaskParams,
+        prompt: str,
+        on_prefill_progress: Callable[[int, int], None] | None = None,
+        distributed_prompt_progress_callback: Callable[[], None] | None = None,
+        on_generation_token: Callable[[], None] | None = None,
+    ) -> int:
+        all_prompt_tokens = encode_prompt(self.tokenizer, prompt)
+        all_prompt_tokens = fix_unmatched_think_end_tokens(
+            all_prompt_tokens, self.tokenizer
+        )
+
+        is_bench = task_params.bench
+
+        prefix_hit_length = 0
+        matched_index: int | None = None
+        prompt_tokens = all_prompt_tokens
+
+        if self.kv_prefix_cache is not None and not is_bench:
+            cache, remaining_tokens, matched_index = self.kv_prefix_cache.get_kv_cache(
+                self.model, all_prompt_tokens
+            )
+            prefix_hit_length = len(all_prompt_tokens) - len(remaining_tokens)
+            if prefix_hit_length > 0:
+                logger.info(
+                    f"KV cache hit: {prefix_hit_length}/{len(all_prompt_tokens)} tokens "
+                    f"cached ({100 * prefix_hit_length / len(all_prompt_tokens):.1f}%)"
+                )
+                prompt_tokens = remaining_tokens
+            else:
+                cache = make_kv_cache(self.model)
+        else:
+            cache = make_kv_cache(self.model)
+
+        seed = task_params.seed if task_params.seed is not None else 42
+        mx.random.seed(seed)
+
+        sampler = make_sampler(
+            temp=task_params.temperature
+            if task_params.temperature is not None
+            else 0.7,
+            top_p=task_params.top_p if task_params.top_p is not None else 1.0,
+            min_p=task_params.min_p if task_params.min_p is not None else 0.05,
+            top_k=task_params.top_k if task_params.top_k is not None else 0,
+        )
+
+        _prefill_tps, _prefill_tokens, cache_snapshots = prefill(
+            self.model,
+            self.tokenizer,
+            sampler,
+            prompt_tokens[:-1],
+            cache,
+            self.group,
+            on_prefill_progress,
+            distributed_prompt_progress_callback,
+        )
+
+        # We need to clamp rotating kv caches to max size so that mlx lm's _merge_caches behaves
+        for c in cache:
+            if (
+                isinstance(c, RotatingKVCache)
+                and c.keys is not None
+                and c.values is not None
+                and c.keys.shape[2] > c.max_size
+            ):
+                trim_size = c.keys.shape[2] - c.max_size
+                c.keys = c._trim(trim_size, c.keys)
+                c.values = c._trim(trim_size, c.values)
+                c._idx = c.max_size
+
+        if not is_bench:
+            self._save_prefix_cache(
+                all_prompt_tokens,
+                list(cache),
+                cache_snapshots,
+                prefix_hit_length,
+                matched_index,
+            )
+
+        last_tokens = prompt_tokens[-2:]
+
+        logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = (
+            make_logits_processors(
+                repetition_penalty=task_params.repetition_penalty,
+                repetition_context_size=task_params.repetition_context_size,
+            )
+        )
+        if is_bench:
+            # Only sample length eos tokens
+            eos_ids = eos_ids_from_tokenizer(self.tokenizer)
+            logits_processors = [ban_token_ids(eos_ids)] + logits_processors
+
+        max_tokens = task_params.max_output_tokens or MAX_TOKENS
+
+        uids = self._exo_gen.insert(
+            prompts=[last_tokens.tolist()],
+            max_tokens=[max_tokens],
+            caches=[list(cache)],
+            samplers=[sampler],
+            logits_processors=[logits_processors],
+        )
+
+        assert len(uids) == 1
+
+        uid = uids[0]
+
+        self._active_tasks[uid] = _EngineTask(
+            uid=uid,
+            task_params=task_params,
+            all_prompt_tokens=all_prompt_tokens,
+            prefix_hit_length=prefix_hit_length,
+            matched_index=matched_index,
+            cache_snapshots=cache_snapshots or None,
+            on_generation_token=on_generation_token,
+            generation_start_time=time.perf_counter(),
+        )
+
+        return uid
+
+    def step(self) -> list[tuple[int, GenerationResponse]]:
+        if not self.has_work:
+            return []
+
+        responses = self._exo_gen.next()
+
+        results: list[tuple[int, GenerationResponse]] = []
+
+        for response in responses:
+            if response.uid not in self._active_tasks:
+                logger.warning(
+                    f"response uid {response.uid} was not found - should be active"
+                )
+                continue
+
+            state = self._active_tasks[response.uid]
+            if state.on_generation_token is not None:
+                state.on_generation_token()
+            text = (
+                ""
+                if response.finish_reason == "stop"
+                else self.tokenizer.decode([response.token])
+            )
+            state.completion_tokens += 1
+            state.generated_text_parts.append(text)
+            state.potential_stop_sequence_text += text
+
+            think_start = self.tokenizer.think_start
+            think_end = self.tokenizer.think_end
+            if think_start is not None and text == think_start:
+                state.in_thinking = True
+            elif think_end is not None and text == think_end:
+                state.in_thinking = False
+            if state.in_thinking:
+                state.reasoning_tokens += 1
+
+            finish_reason: FinishReason | None = cast(
+                FinishReason | None, response.finish_reason
+            )
+            task_params = state.task_params
+            stop_sequences = _stop_sequences(task_params)
+            max_stop_len = max((len(s) for s in stop_sequences), default=0)
+
+            if stop_sequences:
+                for stop_seq in stop_sequences:
+                    if stop_seq in state.potential_stop_sequence_text:
+                        stop_index = state.potential_stop_sequence_text.find(stop_seq)
+                        text_before_stop = state.potential_stop_sequence_text[
+                            :stop_index
+                        ]
+                        chunk_start = len(state.potential_stop_sequence_text) - len(
+                            text
+                        )
+                        text = text_before_stop[chunk_start:]
+                        finish_reason = "stop"
+                        break
+
+            is_done = finish_reason is not None
+
+            logprob: float | None = None
+            top_logprobs: list[TopLogprobItem] | None = None
+            if task_params.logprobs:
+                logprob, top_logprobs = extract_top_logprobs(
+                    logprobs=response.logprobs,
+                    tokenizer=self.tokenizer,
+                    top_logprobs=task_params.top_logprobs or DEFAULT_TOP_LOGPROBS,
+                    selected_token=response.token,
+                )
+
+            stats: GenerationStats | None = None
+            usage: Usage | None = None
+            if is_done:
+                generation_elapsed = time.perf_counter() - state.generation_start_time
+                generation_tps = (
+                    state.completion_tokens / generation_elapsed
+                    if generation_elapsed > 0
+                    else 0.0
+                )
+                mlx_stats = self._exo_gen.stats()
+                stats = GenerationStats(
+                    prompt_tps=float(mlx_stats.prompt_tps)
+                    if mlx_stats.prompt_time > 0
+                    else 0.0,
+                    generation_tps=float(generation_tps),
+                    prompt_tokens=len(state.all_prompt_tokens),
+                    generation_tokens=state.completion_tokens,
+                    peak_memory_usage=Memory.from_gb(mx.get_peak_memory() / 1e9),
+                )
+                total_prompt_tokens = len(state.all_prompt_tokens)
+                usage = Usage(
+                    prompt_tokens=total_prompt_tokens,
+                    completion_tokens=state.completion_tokens,
+                    total_tokens=total_prompt_tokens + state.completion_tokens,
+                    prompt_tokens_details=PromptTokensDetails(
+                        cached_tokens=state.prefix_hit_length
+                    ),
+                    completion_tokens_details=CompletionTokensDetails(
+                        reasoning_tokens=state.reasoning_tokens
+                    ),
+                )
+
+            results.append(
+                (
+                    response.uid,
+                    GenerationResponse(
+                        text=text,
+                        token=response.token,
+                        logprob=logprob,
+                        top_logprobs=top_logprobs,
+                        finish_reason=finish_reason,
+                        stats=stats,
+                        usage=usage,
+                    ),
+                )
+            )
+
+            if is_done:
+                del self._active_tasks[response.uid]
+            elif (
+                max_stop_len > 0
+                and len(state.potential_stop_sequence_text) > max_stop_len
+            ):
+                state.potential_stop_sequence_text = state.potential_stop_sequence_text[
+                    -max_stop_len:
+                ]
+
+        return results
+
+    def cancel(self, uids: list[int]) -> None:
+        self._exo_gen.remove(uids)
+        for uid in uids:
+            self._active_tasks.pop(uid, None)
+
+    def close(self) -> None:
+        self._exo_gen.close()
+
+    def _save_prefix_cache(
+        self,
+        all_prompt_tokens: mx.array,
+        cache: KVCacheType,
+        cache_snapshots: list[CacheSnapshot] | None,
+        prefix_hit_length: int,
+        matched_index: int | None,
+    ) -> None:
+        if self.kv_prefix_cache is None:
+            return
+
+        try:
+            hit_ratio = (
+                prefix_hit_length / len(all_prompt_tokens)
+                if len(all_prompt_tokens) > 0
+                else 0.0
+            )
+            if (
+                matched_index is not None
+                and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
+            ):
+                self.kv_prefix_cache.update_kv_cache(
+                    matched_index,
+                    all_prompt_tokens,
+                    cache,
+                    cache_snapshots,
+                    restore_pos=prefix_hit_length,
+                )
+            else:
+                self.kv_prefix_cache.add_kv_cache(
+                    all_prompt_tokens, cache, cache_snapshots
+                )
+        except Exception:
+            logger.warning("Failed to save prefix cache", exc_info=True)

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

@@ -1,12 +1,16 @@
+import functools
 import math
 import time
 from copy import deepcopy
 from typing import Callable, Generator, cast, get_args
 
 import mlx.core as mx
-from mlx_lm.generate import stream_generate
+from mlx_lm.generate import (
+    maybe_quantize_kv_cache,
+    stream_generate,
+)
 from mlx_lm.models.cache import ArraysCache, RotatingKVCache
-from mlx_lm.sample_utils import make_sampler
+from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.types.api import (
@@ -19,13 +23,19 @@ from exo.shared.types.api import (
 )
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType
+from exo.shared.types.mlx import KVCacheType, Model
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
 )
-from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.auto_parallel import set_pipeline_prefill
+from exo.worker.engines.mlx.auto_parallel import (
+    PipelineFirstLayer,
+    PipelineLastLayer,
+    clear_prefill_sends,
+    flush_prefill_sends,
+    set_pipeline_prefill,
+    set_pipeline_queue_sends,
+)
 from exo.worker.engines.mlx.cache import (
     CacheSnapshot,
     KVPrefixCache,
@@ -56,6 +66,130 @@ class PrefillCancelled(BaseException):
     """Raised when prefill is cancelled via the progress callback."""
 
 
+def _has_pipeline_communication_layer(model: Model):
+    for layer in model.layers:
+        if isinstance(layer, (PipelineFirstLayer, PipelineLastLayer)):
+            return True
+    return False
+
+
+def pipeline_parallel_prefill(
+    model: Model,
+    prompt: mx.array,
+    prompt_cache: KVCacheType,
+    prefill_step_size: int,
+    kv_group_size: int | None,
+    kv_bits: int | None,
+    prompt_progress_callback: Callable[[int, int], None],
+    distributed_prompt_progress_callback: Callable[[], None] | None,
+    group: mx.distributed.Group,
+) -> None:
+    """Prefill the KV cache for pipeline parallel with overlapping stages.
+
+    Each rank processes the full prompt through its real cache, offset by leading
+    and trailing dummy iterations.
+
+    Total iterations per rank = N_real_chunks + world_size - 1:
+      - rank r leading dummies  (skip_pipeline_io, throwaway cache)
+      - N_real_chunks real      (pipeline IO active, real cache)
+      - (world_size-1-r) trailing dummies (skip_pipeline_io, throwaway cache)
+
+    e.g.
+    Timeline (2 ranks, 3 chunks of 10240 tokens @ step=4096):
+        iter 0: R0 real[0:4096]     R1 dummy
+        iter 1: R0 real[4096:8192]  R1 real[0:4096]
+        iter 2: R0 real[8192:10240] R1 real[4096:8192]
+        iter 3: R0 dummy            R1 real[8192:10240]
+
+    This function is designed to match mlx_lm's stream_generate exactly in terms of
+    side effects (given the same prefill step size)
+    """
+    prefill_step_size = prefill_step_size // min(4, group.size())
+
+    quantize_cache_fn: Callable[..., None] = functools.partial(
+        maybe_quantize_kv_cache,
+        quantized_kv_start=0,
+        kv_group_size=kv_group_size,
+        kv_bits=kv_bits,
+    )
+
+    _prompt_cache: KVCacheType = prompt_cache
+    rank = group.rank()
+    world_size = group.size()
+
+    # Build list of real prompt chunk sizes
+    total = len(prompt)
+    real_chunk_sizes: list[int] = []
+    remaining = total - 1
+    while remaining:
+        n = min(prefill_step_size, remaining)
+        real_chunk_sizes.append(n)
+        remaining -= n
+    n_real = len(real_chunk_sizes)
+
+    # Each rank does: [rank leading dummies] [N real chunks] [world_size-1-rank trailing dummies]
+    n_leading = rank
+    n_trailing = world_size - 1 - rank
+    n_total = n_leading + n_real + n_trailing
+
+    t_start = time.perf_counter()
+    processed = 0
+    logger.info(
+        f"[R{rank}] Pipeline prefill: {n_real} real + {n_leading} leading + {n_trailing} trailing = {n_total} iterations"
+    )
+    clear_prefill_sends()
+
+    # Initial callback matching generate_step
+    prompt_progress_callback(0, total)
+
+    try:
+        with mx.stream(generation_stream):
+            for _ in range(n_leading):
+                if distributed_prompt_progress_callback is not None:
+                    distributed_prompt_progress_callback()
+
+            for i in range(n_real):
+                chunk_size = real_chunk_sizes[i]
+                model(
+                    prompt[processed : processed + chunk_size][None],
+                    cache=_prompt_cache,
+                )
+                quantize_cache_fn(_prompt_cache)
+                processed += chunk_size
+
+                if distributed_prompt_progress_callback is not None:
+                    distributed_prompt_progress_callback()
+
+                flush_prefill_sends()
+
+                prompt_progress_callback(processed, total)
+
+            for _ in range(n_trailing):
+                if distributed_prompt_progress_callback is not None:
+                    distributed_prompt_progress_callback()
+
+    finally:
+        clear_prefill_sends()
+
+    # Post-loop: process remaining 1 token + add +1 entry to match stream_generate.
+    for _ in range(2):
+        with mx.stream(generation_stream):
+            model(prompt[-1:][None], cache=_prompt_cache)
+            quantize_cache_fn(_prompt_cache)
+        flush_prefill_sends()
+
+    assert _prompt_cache is not None
+    mx.eval([c.state for c in _prompt_cache])  # type: ignore
+
+    # Final callback matching generate_step
+    prompt_progress_callback(total, total)
+
+    logger.info(
+        f"[R{rank}] Prefill: {n_real} real + {n_leading}+{n_trailing} dummy iterations, "
+        f"Processed {processed} tokens in {(time.perf_counter() - t_start) * 1000:.1f}ms"
+    )
+
+
 def prefill(
     model: Model,
     tokenizer: TokenizerWrapper,
@@ -64,6 +198,7 @@ def prefill(
     cache: KVCacheType,
     group: mx.distributed.Group | None,
     on_prefill_progress: Callable[[int, int], None] | None,
+    distributed_prompt_progress_callback: Callable[[], None] | None,
 ) -> tuple[float, int, list[CacheSnapshot]]:
     """Prefill the KV cache with prompt tokens.
 
@@ -95,31 +230,57 @@ def prefill(
         if on_prefill_progress is not None:
             on_prefill_progress(processed, total)
 
+    def combined_progress_callback(processed: int, total: int) -> None:
+        if distributed_prompt_progress_callback is not None:
+            distributed_prompt_progress_callback()
+        progress_callback(processed, total)
+
     set_pipeline_prefill(model, is_prefill=True)
 
     mx_barrier(group)
     logger.info("Starting prefill")
 
-    # Use max_tokens=1 because max_tokens=0 does not work.
-    # We just throw away the generated token - we only care about filling the cache
+    is_pipeline = _has_pipeline_communication_layer(model)
+
+    prefill_step_size = 4096
+
     try:
-        for _ in stream_generate(
-            model=model,
-            tokenizer=tokenizer,
-            prompt=prompt_tokens,
-            max_tokens=1,
-            sampler=sampler,
-            prompt_cache=cache,
-            prefill_step_size=4096,
-            kv_group_size=KV_GROUP_SIZE,
-            kv_bits=KV_BITS,
-            prompt_progress_callback=progress_callback,
-        ):
-            break  # Stop after first iteration - cache is now filled
+        if is_pipeline and num_tokens >= prefill_step_size:
+            set_pipeline_queue_sends(model, queue_sends=True)
+            assert group is not None, "Pipeline prefill requires a distributed group"
+            pipeline_parallel_prefill(
+                model=model,
+                prompt=prompt_tokens,
+                prompt_cache=cache,
+                prefill_step_size=prefill_step_size,
+                kv_group_size=KV_GROUP_SIZE,
+                kv_bits=KV_BITS,
+                prompt_progress_callback=progress_callback,
+                distributed_prompt_progress_callback=distributed_prompt_progress_callback,
+                group=group,
+            )
+        else:
+            # Use max_tokens=1 because max_tokens=0 does not work.
+            # We just throw away the generated token - we only care about filling the cache
+            for _ in stream_generate(
+                model=model,
+                tokenizer=tokenizer,
+                prompt=prompt_tokens,
+                max_tokens=1,
+                sampler=sampler,
+                prompt_cache=cache,
+                prefill_step_size=prefill_step_size,
+                kv_group_size=KV_GROUP_SIZE,
+                kv_bits=KV_BITS,
+                prompt_progress_callback=combined_progress_callback,
+            ):
+                break  # Stop after first iteration - cache is now filled
     except PrefillCancelled:
+        set_pipeline_queue_sends(model, queue_sends=False)
         set_pipeline_prefill(model, is_prefill=False)
         raise
 
+    set_pipeline_queue_sends(model, queue_sends=False)
     set_pipeline_prefill(model, is_prefill=False)
 
     # stream_generate added 1 extra generated token to the cache, so we should trim it.
@@ -132,7 +293,7 @@ def prefill(
                 cache[i] = deepcopy(pre_gen.states[i])  # type: ignore
         else:
             assert not isinstance(c, (ArraysCache, RotatingKVCache))
-            c.trim(2)  # pyright: ignore[reportUnknownMemberType]
+            c.trim(2)
 
     elapsed = time.perf_counter() - start_time
     tokens_per_sec = num_tokens / elapsed if elapsed > 0 else 0.0
@@ -148,7 +309,11 @@ def warmup_inference(
     model: Model,
     tokenizer: TokenizerWrapper,
     group: mx.distributed.Group | None,
+    model_id: ModelId,
 ) -> int:
+    logger.info(f"warming up inference for instance: {model_id}")
+    t = time.monotonic()
+
     content = "Prompt to warm up the inference engine. Repeat this."
 
     warmup_prompt = apply_chat_template(
@@ -189,7 +354,25 @@ def warmup_inference(
 
     mx_barrier(group)
 
-    return tokens_generated
+    logger.info(f"warmed up by generating {tokens_generated} tokens")
+    check_for_cancel_every = min(
+        math.ceil(tokens_generated / min(time.monotonic() - t, 0.001)), 100
+    )
+    if group is not None:
+        check_for_cancel_every = int(
+            mx.max(
+                mx.distributed.all_gather(
+                    mx.array([check_for_cancel_every]),
+                    group=group,
+                )
+            ).item()
+        )
+
+    logger.info(
+        f"runner checking for cancellation every {check_for_cancel_every} tokens"
+    )
+
+    return check_for_cancel_every
 
 
 def ban_token_ids(token_ids: list[int]) -> Callable[[mx.array, mx.array], mx.array]:
@@ -275,6 +458,8 @@ def mlx_generate(
     kv_prefix_cache: KVPrefixCache | None,
     group: mx.distributed.Group | None,
     on_prefill_progress: Callable[[int, int], None] | None = None,
+    distributed_prompt_progress_callback: Callable[[], None] | None = None,
+    on_generation_token: Callable[[], None] | None = None,
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
@@ -307,15 +492,21 @@ def mlx_generate(
                 f"KV cache hit: {prefix_hit_length}/{len(all_prompt_tokens)} tokens cached ({100 * prefix_hit_length / len(all_prompt_tokens):.1f}%)"
             )
 
-    logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = []
+    logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = (
+        make_logits_processors(
+            repetition_penalty=task.repetition_penalty,
+            repetition_context_size=task.repetition_context_size,
+        )
+    )
     if is_bench:
         # Only sample length eos tokens
         eos_ids = eos_ids_from_tokenizer(tokenizer)
-        logits_processors = [ban_token_ids(eos_ids)]
+        logits_processors = [ban_token_ids(eos_ids)] + logits_processors
 
     sampler = make_sampler(
         temp=task.temperature if task.temperature is not None else 0.7,
         top_p=task.top_p if task.top_p is not None else 1.0,
+        min_p=task.min_p if task.min_p is not None else 0.05,
         top_k=task.top_k if task.top_k is not None else 0,
     )
 
@@ -336,6 +527,7 @@ def mlx_generate(
         caches,
         group,
         on_prefill_progress,
+        distributed_prompt_progress_callback,
     )
     cache_snapshots: list[CacheSnapshot] | None = ssm_snapshots_list or None
 
@@ -481,6 +673,9 @@ def mlx_generate(
                         full_prompt_tokens, caches, cache_snapshots
                     )
 
+        if on_generation_token is not None:
+            on_generation_token()
+
         yield GenerationResponse(
             text=text,
             token=out.token,

src/exo/worker/engines/mlx/patches/__init__.py

@@ -0,0 +1,47 @@
+"""Model-specific kernel fusion patches for MLX inference.
+
+Detects model type after loading and applies optimized kernel patches
+where available. Currently supports:
+- Qwen3.5 MoE (model_type: qwen3_5_moe):
+    EXO_FUSED_KERNELS=1: oproj mode (4-dispatch MoE fusion)
+    EXO_FUSED_KERNELS=2: fused_gqa_gdn mode (full GDN + GQA + MoE fusion, default)
+
+Set EXO_FUSED_KERNELS=0 to disable all patches (vanilla mode).
+Default: EXO_FUSED_KERNELS=2 (fused_gqa_gdn, best performance).
+"""
+
+import json
+import os
+from pathlib import Path
+
+import mlx.nn as nn
+from loguru import logger
+
+
+def maybe_apply_patches(model: nn.Module, model_path: Path) -> None:
+    """Detect model type and apply kernel fusion patches if available."""
+    fused_mode = os.environ.get("EXO_FUSED_KERNELS", "2")
+    if fused_mode == "0":
+        logger.info("Kernel fusion patches disabled (EXO_FUSED_KERNELS=0)")
+        return
+
+    config_path = model_path / "config.json"
+    if not config_path.exists():
+        return
+
+    with open(config_path) as f:
+        config = json.load(f)
+
+    model_type = config.get("model_type", "")
+
+    if model_type == "qwen3_5_moe":
+        if fused_mode == "1":
+            from .qwen3_5_moe.apply import apply_qwen35_oproj_patches
+
+            logger.info("Detected Qwen3.5 MoE model, applying oproj fusion patches")
+            apply_qwen35_oproj_patches(model)
+        else:
+            from .qwen3_5_moe.apply import apply_qwen35_fused_gqa_gdn_patches
+
+            logger.info("Detected Qwen3.5 MoE model, applying fused GQA+GDN patches")
+            apply_qwen35_fused_gqa_gdn_patches(model)

src/exo/worker/engines/mlx/patches/qwen3_5_moe/apply.py

@@ -0,0 +1,44 @@
+"""Apply kernel fusion patches to Qwen3.5 MoE models.
+
+Entry point called from patches/__init__.py after model type detection.
+Supports oproj mode (MoE only) and fused_gqa_gdn mode (full attention + MoE).
+"""
+
+import time
+
+import mlx.nn as nn
+from loguru import logger
+
+from .common import apply_oproj_fused_patches, apply_fused_gqa_gdn_patches
+
+
+def apply_qwen35_oproj_patches(model: nn.Module) -> None:
+    """Apply oproj 4-dispatch fusion to all layers of a Qwen3.5 MoE model.
+
+    Patches decoder, attention, and MoE __call__ methods to use custom Metal
+    kernels for decode (seq_len=1). Prefill (seq_len>1) falls back to vanilla.
+    """
+    layers = model.layers  # type: ignore[attr-defined]
+    n_layers = len(layers)
+
+    t0 = time.time()
+    apply_oproj_fused_patches(layers, gate_bm=8, free_originals=False)
+    t_patch = time.time() - t0
+
+    logger.info(f"Qwen3.5 oproj fusion: patched {n_layers} layers in {t_patch:.1f}s")
+
+
+def apply_qwen35_fused_gqa_gdn_patches(model: nn.Module) -> None:
+    """Apply full fusion (GDN + GQA attention + oproj MoE) to all layers.
+
+    Fused GDN attention (3/4 layers) + fused GQA attention (1/4 layers)
+    + oproj MoE (all layers). 44% faster than vanilla on Qwen3.5-35B-A3B.
+    """
+    layers = model.layers  # type: ignore[attr-defined]
+    n_layers = len(layers)
+
+    t0 = time.time()
+    apply_fused_gqa_gdn_patches(layers, gate_bm=8, free_originals=False)
+    t_patch = time.time() - t0
+
+    logger.info(f"Qwen3.5 fused GQA+GDN: patched {n_layers} layers in {t_patch:.1f}s")

src/exo/worker/engines/mlx/patches/qwen3_5_moe/common.py

@@ -0,0 +1,346 @@
+"""Weight preparation and patch orchestration for Qwen3.5 kernel fusion.
+
+Supports:
+  apply_oproj_fused_patches    — oproj mode (4-dispatch MoE only)
+  apply_fused_gqa_gdn_patches  — full fusion (fused GDN + GQA attention + oproj MoE)
+
+Adapted from mlx_bench/model_patches/qwen/common.py.
+"""
+
+import mlx.core as mx
+import mlx.nn as nn
+from loguru import logger
+from mlx_lm.models.qwen3_5 import DecoderLayer
+from mlx_lm.models.qwen3_next import Qwen3NextSparseMoeBlock
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _patch_swiglu_weights(moe):
+    """Stack gate+up weights for fused 8-bit SwiGLU kernel."""
+    gate_proj = moe.switch_mlp.gate_proj
+    up_proj = moe.switch_mlp.up_proj
+
+    moe.switch_mlp._fused_w_gu = mx.concatenate(
+        [gate_proj.weight, up_proj.weight], axis=1)
+    moe.switch_mlp._fused_s_gu = mx.concatenate(
+        [gate_proj.scales, up_proj.scales], axis=1)
+    moe.switch_mlp._fused_b_gu = mx.concatenate(
+        [gate_proj.biases, up_proj.biases], axis=1)
+    moe.switch_mlp._fused_n_inter = gate_proj.output_dims
+    moe.switch_mlp._fused_k_hidden = gate_proj.input_dims
+    moe.switch_mlp._fused_group_size = gate_proj.group_size
+
+    mx.eval(moe.switch_mlp._fused_w_gu,
+            moe.switch_mlp._fused_s_gu,
+            moe.switch_mlp._fused_b_gu)
+
+
+def _patch_shared_expert(moe):
+    """Prepare shared expert quantized weights for fused 8-bit path."""
+    shared = moe.shared_expert
+    gp = shared.gate_proj
+    up = shared.up_proj
+    dp = shared.down_proj
+
+    moe._shared_w_gu = mx.concatenate([gp.weight, up.weight], axis=0)
+    moe._shared_s_gu = mx.concatenate([gp.scales, up.scales], axis=0)
+    moe._shared_b_gu = mx.concatenate([gp.biases, up.biases], axis=0)
+
+    moe._shared_down_w = dp.weight
+    moe._shared_down_s = dp.scales
+    moe._shared_down_b = dp.biases
+
+    moe._shared_inter = gp.weight.shape[0]
+    moe._shared_gs = gp.group_size
+
+    mx.eval(moe._shared_w_gu, moe._shared_s_gu, moe._shared_b_gu,
+            moe._shared_down_w, moe._shared_down_s, moe._shared_down_b)
+
+
+def _patch_down_proj(moe):
+    """Extract down_proj weights for merged 8-bit kernel dispatch."""
+    dp = moe.switch_mlp.down_proj
+    moe._down_w = dp.weight
+    moe._down_s = dp.scales
+    moe._down_b = dp.biases
+    moe._down_K = dp.output_dims
+    moe._down_N = dp.input_dims
+    moe._down_gs = dp.group_size
+    mx.eval(moe._down_w, moe._down_s, moe._down_b)
+
+
+def _patch_oproj_gate_rms(layer, gate_bm=8):
+    """Precompute M1/W_fused for fused o_proj + gate GEMV (oproj 4-dispatch mode).
+
+    Gate decomposition:
+      gate_score[e] = W_gate[e,:] @ rms_norm(h)
+      where h = residual + W_oproj @ attn_out
+      rms_norm(h) = h * w_rms * inv_rms
+
+    Expanding:
+      gate_score[e] = (W_fused @ residual + M1 @ attn_out) * inv_rms
+
+    Precomputed offline (per layer, stored on moe block):
+      W_fused = dequant(W_gate) · diag(w_rms)    — (E, K) bf16
+      M1 = W_fused @ dequant(W_oproj)            — (E, K_attn) bf16
+    """
+    moe = layer.mlp
+
+    if layer.is_linear:
+        oproj = layer.linear_attn.out_proj
+    else:
+        oproj = layer.self_attn.o_proj
+
+    # Dequantize gate and o_proj (temporary, for M1 computation)
+    gate = moe.gate
+    W_gate_f32 = mx.dequantize(
+        gate.weight, gate.scales, gate.biases,
+        group_size=gate.group_size, bits=gate.bits,
+    ).astype(mx.float32)
+
+    W_oproj_f32 = mx.dequantize(
+        oproj.weight, oproj.scales, oproj.biases,
+        group_size=oproj.group_size, bits=oproj.bits,
+    ).astype(mx.float32)
+    mx.eval(W_gate_f32, W_oproj_f32)
+
+    rms_weight = layer.post_attention_layernorm.weight.astype(mx.bfloat16)
+
+    w_rms_f32 = rms_weight.astype(mx.float32)
+    W_fused = (W_gate_f32 * w_rms_f32).astype(mx.bfloat16)
+    mx.eval(W_fused)
+    del W_gate_f32
+
+    M1 = (W_fused.astype(mx.float32) @ W_oproj_f32).astype(mx.bfloat16)
+    mx.eval(M1)
+    del W_oproj_f32
+
+    moe._oproj_M1 = M1
+    moe._oproj_W_fused = W_fused
+    moe._oproj_rms_weight = rms_weight
+
+    moe._oproj_w = oproj.weight
+    moe._oproj_s = oproj.scales
+    moe._oproj_b = oproj.biases
+    moe._oproj_K_attn = oproj.weight.shape[1] * 4
+
+    seg = moe.shared_expert_gate
+    moe._seg_w = seg.weight
+    moe._seg_s = seg.scales
+    moe._seg_b = seg.biases
+
+    M = oproj.weight.shape[0]
+    K_hidden = W_fused.shape[1]
+    n_experts = W_fused.shape[0]
+    moe._oproj_M = M
+    moe._oproj_K_hidden = K_hidden
+    moe._oproj_n_experts = n_experts
+    moe._oproj_n_tg = ceil_div(M, 32)
+    moe._oproj_gate_bm = gate_bm
+
+    mx.eval(moe._oproj_rms_weight)
+
+
+def apply_oproj_fused_patches(layers, gate_bm=8, free_originals=False):
+    """Apply oproj-mode fused MoE patches (4-dispatch) to all layers.
+
+    1. Prepare SwiGLU weights (_patch_swiglu_weights)
+    2. Prepare shared expert 8-bit weights (_patch_shared_expert)
+    3. Prepare down_proj weights (_patch_down_proj)
+    4. Precompute M1/W_fused + store o_proj/gate weights (_patch_oproj_gate_rms)
+    5. Replace __call__ methods for decoder + MoE + attention
+    """
+    from .moe import _oproj_moe_call
+    from .decoder import (
+        _oproj_decoder_call,
+        _pre_oproj_attention_call,
+        _pre_oproj_qwen35_linear_attn_call,
+    )
+    from mlx_lm.models.qwen3_next import Qwen3NextAttention
+    from mlx_lm.models.qwen3_5 import GatedDeltaNet
+
+    n_patched = 0
+    for li, layer in enumerate(layers):
+        moe = layer.mlp
+        if isinstance(moe, Qwen3NextSparseMoeBlock):
+            _patch_swiglu_weights(moe)
+            _patch_shared_expert(moe)
+            _patch_down_proj(moe)
+            _patch_oproj_gate_rms(layer, gate_bm=gate_bm)
+
+            if free_originals:
+                for attr in ('weight', 'scales', 'biases'):
+                    for proj in (moe.switch_mlp.gate_proj,
+                                 moe.switch_mlp.up_proj):
+                        try:
+                            delattr(proj, attr)
+                        except AttributeError:
+                            pass
+
+            n_patched += 1
+            if (li + 1) % 10 == 0 or li == 0:
+                logger.info(f"  Patched layer {li+1}/{len(layers)} (oproj mode)")
+
+    Qwen3NextAttention.__call__ = _pre_oproj_attention_call
+    GatedDeltaNet.__call__ = _pre_oproj_qwen35_linear_attn_call
+    Qwen3NextSparseMoeBlock.__call__ = _oproj_moe_call
+    DecoderLayer.__call__ = _oproj_decoder_call
+    logger.info(f"  Patched {n_patched} MoE blocks (oproj mode, 4 dispatches)")
+
+
+def _patch_gdn_proj_weights(attn):
+    """Merge all 4 GDN projection weights into contiguous buffers."""
+    W_merged = mx.concatenate([
+        attn.in_proj_qkv.weight,
+        attn.in_proj_z.weight,
+        attn.in_proj_b.weight,
+        attn.in_proj_a.weight,
+    ], axis=0)
+    S_merged = mx.concatenate([
+        attn.in_proj_qkv.scales,
+        attn.in_proj_z.scales,
+        attn.in_proj_b.scales,
+        attn.in_proj_a.scales,
+    ], axis=0)
+    B_merged = mx.concatenate([
+        attn.in_proj_qkv.biases,
+        attn.in_proj_z.biases,
+        attn.in_proj_b.biases,
+        attn.in_proj_a.biases,
+    ], axis=0)
+    attn._merged_proj_w = W_merged
+    attn._merged_proj_s = S_merged
+    attn._merged_proj_b = B_merged
+    attn._merged_proj_dims = (
+        attn.in_proj_qkv.weight.shape[0],
+        attn.in_proj_z.weight.shape[0],
+        attn.in_proj_b.weight.shape[0],
+        attn.in_proj_a.weight.shape[0],
+    )
+    mx.eval(W_merged, S_merged, B_merged)
+
+
+def _patch_gqa_proj_weights(attn):
+    """Merge GQA q_proj, k_proj, v_proj weights with q_proj row permutation.
+
+    q_proj rows are interleaved [head0_q, head0_gate, head1_q, ...].
+    Permute so queries come first, then gate, then k, then v.
+    Pre-cache constant scalar arrays for kernel dispatch.
+    """
+    q = attn.q_proj
+    k = attn.k_proj
+    v = attn.v_proj
+
+    H_q = attn.num_attention_heads
+    D = attn.head_dim
+
+    W_q = q.weight.reshape(H_q, 2 * D, -1)
+    S_q = q.scales.reshape(H_q, 2 * D, -1)
+    B_q = q.biases.reshape(H_q, 2 * D, -1)
+
+    W_queries = W_q[:, :D, :].reshape(H_q * D, -1)
+    W_gate = W_q[:, D:, :].reshape(H_q * D, -1)
+    S_queries = S_q[:, :D, :].reshape(H_q * D, -1)
+    S_gate = S_q[:, D:, :].reshape(H_q * D, -1)
+    B_queries = B_q[:, :D, :].reshape(H_q * D, -1)
+    B_gate = B_q[:, D:, :].reshape(H_q * D, -1)
+
+    W_merged = mx.contiguous(mx.concatenate([W_queries, W_gate, k.weight, v.weight], axis=0))
+    S_merged = mx.contiguous(mx.concatenate([S_queries, S_gate, k.scales, v.scales], axis=0))
+    B_merged = mx.contiguous(mx.concatenate([B_queries, B_gate, k.biases, v.biases], axis=0))
+
+    attn._merged_proj_w = W_merged
+    attn._merged_proj_s = S_merged
+    attn._merged_proj_b = B_merged
+    N_Q = H_q * D
+    N_GATE = H_q * D
+    N_K = k.weight.shape[0]
+    N_V = v.weight.shape[0]
+    attn._merged_proj_dims = (N_Q, N_GATE, N_K, N_V)
+    mx.eval(W_merged, S_merged, B_merged)
+
+    # Pre-cache constant scalar arrays for kernel dispatch
+    N_TOTAL = N_Q + N_GATE + N_K + N_V
+    K_dim = q.weight.shape[1] * 4  # 8-bit: pack_factor=4
+    attn._kernel_scalars = {
+        'K': mx.array(K_dim, dtype=mx.int32),
+        'N_Q': mx.array(N_Q, dtype=mx.int32),
+        'N_GATE': mx.array(N_GATE, dtype=mx.int32),
+        'N_K': mx.array(N_K, dtype=mx.int32),
+        'N_TOTAL': mx.array(N_TOTAL, dtype=mx.int32),
+        'N_Q_TG': mx.array(ceil_div(N_Q, 8), dtype=mx.int32),
+        'N_GATE_TG': mx.array(ceil_div(N_GATE, 8), dtype=mx.int32),
+        'N_K_TG': mx.array(ceil_div(N_K, 8), dtype=mx.int32),
+        'scale': mx.array(attn.head_dim ** -0.5, dtype=mx.float32),
+        'H_Q': mx.array(attn.num_attention_heads, dtype=mx.int32),
+        'H_KV': mx.array(attn.num_key_value_heads, dtype=mx.int32),
+        'N_blocks': mx.array(128, dtype=mx.int32),
+    }
+    mx.eval(*attn._kernel_scalars.values())
+
+    N_V_TG = ceil_div(N_V, 8)
+    attn._d1_total_tg = ceil_div(N_Q, 8) + ceil_div(N_GATE, 8) + ceil_div(N_K, 8) + N_V_TG
+
+    # Precompute RoPE inv_freq
+    rope_dims = attn.rope.dims
+    half_dims = rope_dims // 2
+    theta = attn.rope.base
+    d_indices = mx.arange(half_dims, dtype=mx.float32)
+    attn._rope_inv_freq = theta ** (-d_indices / half_dims)
+    mx.eval(attn._rope_inv_freq)
+
+
+def apply_fused_gqa_gdn_patches(layers, gate_bm=8, free_originals=False):
+    """Apply all fusions: fused GDN + fused GQA + oproj MoE.
+
+    Combines:
+      - Fused GDN attention (3/4 layers: GatedDeltaNet)
+      - Fused GQA attention (1/4 layers: Qwen3NextAttention)
+      - Oproj MoE (all layers: oproj_gate_gemv + fused MoE dispatches)
+    """
+    from .moe import _oproj_moe_call
+    from .decoder import _fused_gdn_decoder_call
+    from .fused_gdn_attention import _fused_gdn_call
+    from .fused_gqa_attention import _fused_gqa_call
+    from mlx_lm.models.qwen3_next import Qwen3NextAttention
+    from mlx_lm.models.qwen3_5 import GatedDeltaNet
+
+    n_patched = 0
+    n_gdn = 0
+    n_gqa = 0
+    for li, layer in enumerate(layers):
+        moe = layer.mlp
+        if isinstance(moe, Qwen3NextSparseMoeBlock):
+            _patch_swiglu_weights(moe)
+            _patch_shared_expert(moe)
+            _patch_down_proj(moe)
+            _patch_oproj_gate_rms(layer, gate_bm=gate_bm)
+
+            if layer.is_linear:
+                _patch_gdn_proj_weights(layer.linear_attn)
+                n_gdn += 1
+            else:
+                _patch_gqa_proj_weights(layer.self_attn)
+                n_gqa += 1
+
+            if free_originals:
+                for attr in ('weight', 'scales', 'biases'):
+                    for proj in (moe.switch_mlp.gate_proj,
+                                 moe.switch_mlp.up_proj):
+                        try:
+                            delattr(proj, attr)
+                        except AttributeError:
+                            pass
+
+            n_patched += 1
+            if (li + 1) % 10 == 0 or li == 0:
+                logger.info(f"  Patched layer {li+1}/{len(layers)} (fused GQA+GDN mode)")
+
+    GatedDeltaNet.__call__ = _fused_gdn_call
+    Qwen3NextAttention.__call__ = _fused_gqa_call
+    Qwen3NextSparseMoeBlock.__call__ = _oproj_moe_call
+    DecoderLayer.__call__ = _fused_gdn_decoder_call
+    logger.info(f"  Patched {n_patched} MoE blocks ({n_gdn} GDN + {n_gqa} GQA, fused GQA+GDN mode)")

src/exo/worker/engines/mlx/patches/qwen3_5_moe/decoder.py

@@ -0,0 +1,186 @@
+"""Decoder layer __call__ variants for Qwen3.5.
+
+Three modes:
+  _fused_decoder_call: passes residual to fused MoE epilogue (~15 dispatches)
+  _oproj_decoder_call: fuses o_proj + RMSNorm + gate GEMV (4 dispatches)
+  _fused_gdn_decoder_call: fused GDN/GQA attention + oproj MoE (6 dispatches)
+
+Attention patches for oproj mode:
+  _pre_oproj_attention_call: Qwen3NextAttention.__call__ that skips o_proj
+  _pre_oproj_qwen35_linear_attn_call: qwen3_5.GatedDeltaNet.__call__ that skips out_proj
+
+Note: qwen3_5.GatedDeltaNet (used by DecoderLayer) is a DIFFERENT class from
+qwen3_next.Qwen3NextGatedDeltaNet. They have different projection layouts:
+  - qwen3_5.GatedDeltaNet: separate in_proj_qkv, in_proj_z, in_proj_b, in_proj_a
+  - qwen3_next.Qwen3NextGatedDeltaNet: merged in_proj_qkvz, in_proj_ba
+The patch must match qwen3_5.GatedDeltaNet's __call__ structure.
+
+Adapted from mlx_bench/model_patches/qwen/decoder.py.
+"""
+
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.nn.layers.activations import silu as nn_silu
+
+# Map moe block id → parent decoder layer (avoids circular refs in model tree)
+_parent_layer_map = {}
+
+
+def _fused_decoder_call(self, x, mask=None, cache=None):
+    """Decoder layer with residual passed to fused MoE epilogue.
+
+    Replaces:
+      h = x + attn(norm(x))
+      out = h + mlp(norm(h))        # mlp returns MoE output, then adds h
+
+    With:
+      h = x + attn(norm(x))
+      out = mlp(norm(h), _residual=h)  # epilogue fuses: moe_out + h
+    """
+    if self.is_linear:
+        r = self.linear_attn(self.input_layernorm(x), mask, cache)
+    else:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+    h = x + r
+    out = self.mlp(self.post_attention_layernorm(h), _residual=h)
+    return out  # already includes residual add from epilogue
+
+
+def _oproj_decoder_call(self, x, mask=None, cache=None):
+    """Decoder with fused o_proj + RMSNorm + gate GEMV (oproj 4-dispatch mode).
+
+    Skips o_proj, addmm, and post_attention_layernorm — all fused into Dispatch 1.
+    Attention __call__ is patched to return pre-o_proj output.
+
+    Flow:
+      pre_oproj = attn(input_layernorm(x))   # returns BEFORE o_proj
+      MoE receives (pre_oproj, residual=x) and handles o_proj + RMSNorm + gate internally
+    """
+    if self.is_linear:
+        pre_oproj = self.linear_attn(self.input_layernorm(x), mask, cache)
+    else:
+        pre_oproj = self.self_attn(self.input_layernorm(x), mask, cache)
+    _parent_layer_map[id(self.mlp)] = self
+    return self.mlp(pre_oproj, _residual=x)
+
+
+def _fused_gdn_decoder_call(self, x, mask=None, cache=None):
+    """Decoder with fused GDN/GQA attention + oproj MoE (6-dispatch mode).
+
+    GDN layers use fused kernels, GQA layers use fused or vanilla attention.
+    Both return pre-out_proj output. MoE handles oproj_gate_gemv + MoE dispatches.
+
+    Flow is identical to oproj mode — the difference is that GatedDeltaNet.__call__
+    and/or Qwen3NextAttention.__call__ are patched with fused kernel implementations.
+    """
+    if self.is_linear:
+        pre_oproj = self.linear_attn(self.input_layernorm(x), mask, cache)
+    else:
+        pre_oproj = self.self_attn(self.input_layernorm(x), mask, cache)
+    _parent_layer_map[id(self.mlp)] = self
+    return self.mlp(pre_oproj, _residual=x)
+
+
+def _pre_oproj_attention_call(self, x, mask=None, cache=None):
+    """Qwen3NextAttention.__call__ that returns pre-o_proj output.
+
+    Identical to original except final line returns output*sigmoid(gate)
+    instead of self.o_proj(output*sigmoid(gate)).
+    """
+    B, L, D = x.shape
+    q_proj_output = self.q_proj(x)
+    queries, gate = mx.split(
+        q_proj_output.reshape(B, L, self.num_attention_heads, -1), 2, axis=-1
+    )
+    gate = gate.reshape(B, L, -1)
+    keys, values = self.k_proj(x), self.v_proj(x)
+    queries = self.q_norm(queries).transpose(0, 2, 1, 3)
+    keys = self.k_norm(
+        keys.reshape(B, L, self.num_key_value_heads, -1)
+    ).transpose(0, 2, 1, 3)
+    values = values.reshape(B, L, self.num_key_value_heads, -1).transpose(
+        0, 2, 1, 3
+    )
+    if cache is not None:
+        queries = self.rope(queries, offset=cache.offset)
+        keys = self.rope(keys, offset=cache.offset)
+        keys, values = cache.update_and_fetch(keys, values)
+    else:
+        queries = self.rope(queries)
+        keys = self.rope(keys)
+
+    from mlx_lm.models.qwen3_next import scaled_dot_product_attention
+    output = scaled_dot_product_attention(
+        queries, keys, values, cache=cache, scale=self.scale, mask=mask
+    )
+    output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+    return output * mx.sigmoid(gate)  # skip o_proj
+
+
+def _pre_oproj_qwen35_linear_attn_call(
+    self,
+    inputs: mx.array,
+    mask: Optional[mx.array] = None,
+    cache: Optional[Any] = None,
+) -> mx.array:
+    """qwen3_5.GatedDeltaNet.__call__ that returns pre-out_proj output.
+
+    Identical to qwen3_5.GatedDeltaNet.__call__ except final line returns
+    out.reshape(B,S,-1) instead of self.out_proj(out.reshape(B,S,-1)).
+
+    Note: this targets qwen3_5.GatedDeltaNet (separate projections), NOT
+    qwen3_next.Qwen3NextGatedDeltaNet (merged projections). They are
+    different classes with different __call__ bodies.
+    """
+    from mlx_lm.models.gated_delta import gated_delta_update
+
+    B, S, _ = inputs.shape
+
+    qkv = self.in_proj_qkv(inputs)
+    z = self.in_proj_z(inputs).reshape(B, S, self.num_v_heads, self.head_v_dim)
+    b = self.in_proj_b(inputs)
+    a = self.in_proj_a(inputs)
+
+    if cache is not None and cache[0] is not None:
+        conv_state = cache[0]
+    else:
+        conv_state = mx.zeros(
+            (B, self.conv_kernel_size - 1, self.conv_dim),
+            dtype=inputs.dtype,
+        )
+
+    if mask is not None:
+        qkv = mx.where(mask[..., None], qkv, 0)
+    conv_input = mx.concatenate([conv_state, qkv], axis=1)
+    if cache is not None:
+        cache[0] = conv_input[:, -(self.conv_kernel_size - 1) :]
+    conv_out = nn.silu(self.conv1d(conv_input))
+
+    q, k, v = [
+        t.reshape(B, S, h, d)
+        for t, h, d in zip(
+            mx.split(conv_out, [self.key_dim, 2 * self.key_dim], -1),
+            [self.num_k_heads, self.num_k_heads, self.num_v_heads],
+            [self.head_k_dim, self.head_k_dim, self.head_v_dim],
+        )
+    ]
+
+    state = cache[1] if cache else None
+    inv_scale = k.shape[-1] ** -0.5
+    q = (inv_scale**2) * mx.fast.rms_norm(q, None, 1e-6)
+    k = inv_scale * mx.fast.rms_norm(k, None, 1e-6)
+
+    out, state = gated_delta_update(
+        q, k, v, a, b,
+        self.A_log, self.dt_bias,
+        state, mask,
+        use_kernel=not self.training,
+    )
+
+    if cache is not None:
+        cache[1] = state
+
+    out = self.norm(out, z)
+    return out.reshape(B, S, -1)  # skip out_proj

src/exo/worker/engines/mlx/patches/qwen3_5_moe/fused_gdn_attention.py

@@ -0,0 +1,161 @@
+"""Fused GDN attention __call__ for qwen3_5.GatedDeltaNet (Dispatches 2-5).
+
+Replaces the vanilla GatedDeltaNet.__call__ with fused kernel dispatches:
+  Dispatch 2: fused_gdn_projections — merged 8-bit GEMV + conv1d + SiLU(qkv) + SiLU(z)
+              + sigmoid(b)→beta + g=exp(-exp(A_log)*softplus(a+dt_bias))
+  Dispatch 3: fused_qk_rmsnorm — per-head L2-norm on q (×Dk^(-½)) and k
+  Dispatch 4: gated_delta_kernel — GDN recurrence (receives pre-computed g, beta)
+  Dispatch 5: fused_rms_norm_gated — RMSNorm(out, weight) × z_silu
+
+All 4 projection weights are pre-merged into contiguous buffers at patch time
+(_patch_gdn_proj_weights) for better memory locality.
+
+g/beta computation is fused into Dispatch 2 epilogues, eliminating ~8 micro-
+dispatches that gated_delta_update would otherwise generate.
+
+Fused path is decode-only (S=1). For prefill (S>1), falls back to vanilla ops.
+
+Returns pre-out_proj output (same interface as _pre_oproj_qwen35_linear_attn_call).
+Dispatch 1 (input_layernorm) is handled by the decoder.
+Dispatch 6 (oproj_gate_gemv) is handled by the MoE __call__.
+"""
+
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .kernels.fused_gdn_projections_8bit import fused_gdn_projections
+from .kernels.fused_qk_rmsnorm import fused_qk_rmsnorm
+from .kernels.fused_rms_norm_gated import fused_rms_norm_gated
+
+
+def _vanilla_gdn_call(self, inputs, mask, cache):
+    """Vanilla GDN path for prefill (S>1). Returns pre-out_proj output."""
+    from mlx_lm.models.gated_delta import gated_delta_update
+
+    B, S, _ = inputs.shape
+
+    qkv = self.in_proj_qkv(inputs)
+    z = self.in_proj_z(inputs).reshape(B, S, self.num_v_heads, self.head_v_dim)
+    b = self.in_proj_b(inputs)
+    a = self.in_proj_a(inputs)
+
+    if cache is not None and cache[0] is not None:
+        conv_state = cache[0]
+    else:
+        conv_state = mx.zeros(
+            (B, self.conv_kernel_size - 1, self.conv_dim),
+            dtype=inputs.dtype,
+        )
+
+    if mask is not None:
+        qkv = mx.where(mask[..., None], qkv, 0)
+    conv_input = mx.concatenate([conv_state, qkv], axis=1)
+    if cache is not None:
+        cache[0] = conv_input[:, -(self.conv_kernel_size - 1):]
+    conv_out = nn.silu(self.conv1d(conv_input))
+
+    q, k, v = [
+        t.reshape(B, S, h, d)
+        for t, h, d in zip(
+            mx.split(conv_out, [self.key_dim, 2 * self.key_dim], -1),
+            [self.num_k_heads, self.num_k_heads, self.num_v_heads],
+            [self.head_k_dim, self.head_k_dim, self.head_v_dim],
+        )
+    ]
+
+    state = cache[1] if cache else None
+    inv_scale = k.shape[-1] ** -0.5
+    q = inv_scale * q * mx.rsqrt(
+        (q * q).sum(axis=-1, keepdims=True) + 1e-6
+    )
+    k = k * mx.rsqrt(
+        (k * k).sum(axis=-1, keepdims=True) + 1e-6
+    )
+
+    out, state = gated_delta_update(
+        q, k, v, a, b,
+        self.A_log, self.dt_bias,
+        state, mask,
+        use_kernel=True,
+    )
+
+    if cache is not None:
+        cache[1] = state
+
+    out = self.norm(out, z)
+    return out.reshape(B, S, -1)  # skip out_proj
+
+
+def _fused_gdn_call(
+    self,
+    inputs: mx.array,
+    mask: Optional[mx.array] = None,
+    cache: Optional[Any] = None,
+) -> mx.array:
+    """Fused GDN attention: merged projections + existing GDN kernel.
+
+    Decode (S=1): uses fused kernels with merged weight buffers.
+    Prefill (S>1): falls back to vanilla ops.
+
+    Returns pre-out_proj output [B, S, value_dim] for Dispatch 6.
+    """
+    B, S, _ = inputs.shape
+
+    # Prefill fallback: fused kernels are decode-only (S=1)
+    if S > 1:
+        return _vanilla_gdn_call(self, inputs, mask, cache)
+
+    from mlx_lm.models.gated_delta import gated_delta_kernel
+
+    # ── Cache: conv state ──
+    if cache is not None and cache[0] is not None:
+        conv_state = cache[0]
+    else:
+        conv_state = mx.zeros(
+            (B, self.conv_kernel_size - 1, self.conv_dim),
+            dtype=inputs.dtype,
+        )
+
+    # ── Dispatch 2: fused projections (merged GEMV + conv + SiLU + g/beta) ──
+    qkv_conv_silu, z_silu, beta, g, conv_state_out = fused_gdn_projections(
+        inputs,
+        self._merged_proj_w, self._merged_proj_s, self._merged_proj_b,
+        self._merged_proj_dims,
+        conv_state, self.conv1d.weight,
+        self.A_log, self.dt_bias,
+        batch_size=B,
+    )
+
+    if cache is not None:
+        cache[0] = conv_state_out
+
+    # ── Dispatch 3: fused Q/K L2-norm ──
+    qk_normed = fused_qk_rmsnorm(qkv_conv_silu, batch_size=B)
+
+    # ── Split q, k from normed output; v from conv output ──
+    q = qk_normed[:, :, :self.key_dim].reshape(B, S, self.num_k_heads, self.head_k_dim)
+    k = qk_normed[:, :, self.key_dim:].reshape(B, S, self.num_k_heads, self.head_k_dim)
+    v = qkv_conv_silu[:, :, 2 * self.key_dim:].reshape(B, S, self.num_v_heads, self.head_v_dim)
+
+    # ── Dispatch 4: GDN recurrence with pre-computed g/beta ──
+    state = cache[1] if cache else None
+    if state is None:
+        state = mx.zeros(
+            (B, self.num_v_heads, self.head_v_dim, self.head_k_dim),
+            dtype=inputs.dtype,
+        )
+
+    out, state_new = gated_delta_kernel(
+        q, k, v, g, beta, state, mask,
+    )
+
+    if cache is not None:
+        cache[1] = state_new
+
+    # ── Dispatch 5: fused RMSNorm × z_silu ──
+    norm_weight = self.norm.weight
+    result = fused_rms_norm_gated(out, z_silu, norm_weight, batch_size=B)
+
+    return result  # [B, S, value_dim] — skip out_proj (handled by Dispatch 6)

src/exo/worker/engines/mlx/patches/qwen3_5_moe/fused_gqa_attention.py

@@ -0,0 +1,130 @@
+"""Fused GQA attention __call__ for qwen3_next.Qwen3NextAttention.
+
+Replaces vanilla Qwen3NextAttention.__call__ with fused kernel dispatches:
+  Dispatch 1: fused_gqa_projections — merged 8-bit GEMV (q+gate+k+v) + sigmoid(gate)
+  Dispatch 2: fused_qk_norm_rope — RMSNorm + RoPE (TODO: custom kernel)
+  Dispatch 3: KV cache update (MLX built-in)
+  Dispatch 4: custom_sdpa_pass1 (TODO: custom kernel)
+  Dispatch 5: custom_sdpa_pass2_gate (TODO: custom kernel, includes gate multiply)
+  Dispatch 6: oproj_gate_gemv (existing, handled by MoE __call__)
+
+Dispatches 1-2, 4-5 are custom kernels, Dispatch 3 is MLX built-in.
+Returns pre-out_proj output (output * sigmoid(gate)) for Dispatch 6.
+
+Fused path is decode-only (S=1). For prefill (S>1), falls back to vanilla ops.
+"""
+
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .kernels.fused_gqa_projections_8bit import fused_gqa_projections
+from .kernels.fused_qk_norm_rope import fused_qk_norm_rope
+from .kernels.custom_sdpa_pass1 import custom_sdpa_pass1
+from .kernels.custom_sdpa_pass2_gate import custom_sdpa_pass2_gate
+
+
+def _vanilla_gqa_call(self, x, mask, cache):
+    """Vanilla GQA path for prefill (S>1). Returns pre-out_proj output."""
+    B, L, D = x.shape
+    q_proj_output = self.q_proj(x)
+    queries, gate = mx.split(
+        q_proj_output.reshape(B, L, self.num_attention_heads, -1), 2, axis=-1
+    )
+    gate = gate.reshape(B, L, -1)
+    keys, values = self.k_proj(x), self.v_proj(x)
+    queries = self.q_norm(queries).transpose(0, 2, 1, 3)
+    keys = self.k_norm(
+        keys.reshape(B, L, self.num_key_value_heads, -1)
+    ).transpose(0, 2, 1, 3)
+    values = values.reshape(B, L, self.num_key_value_heads, -1).transpose(
+        0, 2, 1, 3
+    )
+    if cache is not None:
+        queries = self.rope(queries, offset=cache.offset)
+        keys = self.rope(keys, offset=cache.offset)
+        keys, values = cache.update_and_fetch(keys, values)
+    else:
+        queries = self.rope(queries)
+        keys = self.rope(keys)
+
+    from mlx_lm.models.qwen3_next import scaled_dot_product_attention
+    output = scaled_dot_product_attention(
+        queries, keys, values, cache=cache, scale=self.scale, mask=mask
+    )
+    output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+    return output * mx.sigmoid(gate)  # skip o_proj
+
+
+def _fused_gqa_call(
+    self,
+    x: mx.array,
+    mask: Optional[mx.array] = None,
+    cache: Optional[Any] = None,
+) -> mx.array:
+    """Fused GQA attention: custom projection + norm/rope kernels.
+
+    Decode (S=1): Dispatch 1 (fused GEMV) + Dispatch 2 (fused norm+rope)
+                  + vanilla cache + SDPA (Dispatches 3-5).
+    Prefill (S>1): falls back to fully vanilla ops.
+
+    Returns pre-out_proj output [B, S, H_q*D] for Dispatch 6.
+    """
+    B, S, _ = x.shape
+
+    # Prefill fallback: fused kernels are decode-only (S=1)
+    if S > 1:
+        return _vanilla_gqa_call(self, x, mask, cache)
+
+    H_q = self.num_attention_heads
+    H_kv = self.num_key_value_heads
+    D = self.head_dim
+
+    # ── Dispatch 1: fused projections (merged GEMV + sigmoid(gate)) ──
+    queries, gate_sigmoid, keys, values = fused_gqa_projections(
+        x,
+        self._merged_proj_w, self._merged_proj_s, self._merged_proj_b,
+        self._merged_proj_dims,
+        batch_size=B,
+        total_tg=getattr(self, '_d1_total_tg', None),
+    )
+
+    # ── Dispatch 2: fused Q/K RMSNorm + RoPE ──
+    queries, keys = fused_qk_norm_rope(
+        queries, keys,
+        self.q_norm.weight, self.k_norm.weight,
+        self._rope_inv_freq, cache.offset,
+        H_q, H_kv, D, batch_size=B,
+    )
+    # queries: [B, H_q, 1, D], keys: [B, H_kv, 1, D]
+    values = values.reshape(B, H_kv, 1, D)
+
+    # ── Dispatch 3: KV cache update ──
+    cache.update_and_fetch(keys, values)
+    N = cache.offset
+    alloc_len = cache.keys.shape[2]
+
+    # ── Dispatch 4: SDPA Pass 1 ──
+    blocks = 128
+    if N < 1024:
+        k_sliced = cache.keys[:, :, :N, :]
+        v_sliced = cache.values[:, :, :N, :]
+        from mlx_lm.models.qwen3_next import scaled_dot_product_attention
+        output = scaled_dot_product_attention(
+            queries, k_sliced, v_sliced, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, S, -1)
+        return output * gate_sigmoid.astype(output.dtype)
+
+    o_partials, sums, maxs = custom_sdpa_pass1(
+        queries, cache.keys, cache.values, self.scale,
+        H_q, H_kv, D, blocks=blocks, batch_size=B,
+        N=N, alloc_len=alloc_len,
+    )
+
+    # ── Dispatch 5: SDPA Pass 2 + gate multiply ──
+    return custom_sdpa_pass2_gate(
+        o_partials, sums, maxs, gate_sigmoid,
+        H_q, D, blocks=blocks, V_SPLIT=4, batch_size=B,
+    )

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/custom_oproj_gate_gemv_8bit.py

@@ -0,0 +1,324 @@
+"""Dispatch 1: Fused o_proj (8-bit) + gate GEMV parts + x² partials for Qwen3.5.
+
+Port of Kimi's custom_oproj_gate_gemv.py adapted for 8-bit quantized o_proj.
+
+Single dispatch with 3 GEMV types sharing 256-thread TGs (8 SGs of 32):
+
+TGs 0 to N_OPROJ_TG-1: o_proj GEMV (8-bit, M=4096, K=8192)
+  8-bit affine dequant: result = scale * Σ(x*w_uint8) + bias * Σ(x)
+  Epilogue: h = oproj_result + residual, h_scaled = h*w_rms, h_out = h, x²_acc
+  TG reduction: 8 SG x² → 1 float per TG.
+
+TGs N_OPROJ_TG to +N_M1_TG-1: M1 GEMV (bf16, E × K_attn)
+  M1 = W_fused @ W_oproj (pre-computed). Input = attn_out.
+  Output: gate_part_a (E,) f32.
+
+TGs +N_M1_TG to end: W_fused GEMV (bf16, E × K)
+  W_fused × residual → gate_part_b (E,) f32.
+"""
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_custom_oproj_8bit_source(n_experts=64, group_size=64, scale_bf16=True):
+    """Generate Metal source for fused 8-bit o_proj + bf16 gate GEMVs."""
+    E = int(n_experts)
+    gs = int(group_size)
+    sc_t = "bfloat16_t" if scale_bf16 else "float"
+    # 8-bit dequant params for o_proj
+    oproj_slid_divisor = gs // 8     # 64/8 = 8
+    oproj_sc_stride = 256 // gs      # 256/64 = 4
+
+    return f"""
+    // ── Constants ──
+    const int TM = 4;       // rows per SG for bf16 GEMVs
+    const int TN = 4;       // elements per thread per iter for bf16 GEMVs
+    const int blockN = 128; // 32 threads × TN=4
+    const int E_CONST = {E};
+
+    int M = M_val;                   // 4096 (hidden_size)
+    int K_attn = K_attn_val;         // 8192 (o_proj input dim)
+    int K_hidden = K_hidden_val;     // 4096 (hidden_size, same as M for Qwen)
+    int N_OPROJ_TG = N_OPROJ_TG_val;
+    int N_M1_TG = N_M1_TG_val;
+    int blockM_gate = BM_GATE_val * TM;  // rows per gate GEMV TG
+
+    uint tg_x = threadgroup_position_in_grid.x;
+    uint sgid = simdgroup_index_in_threadgroup;  // 0..7
+    uint slid = thread_index_in_simdgroup;       // 0..31
+
+    if (tg_x < (uint)N_OPROJ_TG) {{
+        // ════════════════════════════════════════════════════════════════
+        // O_PROJ GEMV: 8-bit quantized, M=4096, K_attn=8192
+        // Each TG handles 32 rows (8 SGs × 4 rows each)
+        // 8-bit affine dequant: result = scale * Σ(x*w) + bias * Σ(x)
+        // ════════════════════════════════════════════════════════════════
+        const int blockM = 32;  // 8 SGs × TM=4
+        const int VPT = 8;     // values per thread per iteration
+        const int BLOCK_SIZE = 256;  // 32 threads × 8 values
+
+        int out_row = int(tg_x) * blockM + int(sgid) * TM;
+        if (out_row >= M) return;
+        out_row = (out_row + TM <= M) ? out_row : (M - TM);
+
+        threadgroup float tgp_x2[8];
+
+        // 8-bit GEMV K-loop: K-outer, tm-inner (input loaded once, reused across TM rows)
+        float acc[TM] = {{0.0f, 0.0f, 0.0f, 0.0f}};
+        float result[TM];
+        int K_groups = K_attn / {gs};
+
+        // Per-row weight/scale/bias pointers
+        const device uint8_t* ws0 = (const device uint8_t*)W_oproj + (long)(out_row + 0) * K_attn + slid * VPT;
+        const device uint8_t* ws1 = (const device uint8_t*)W_oproj + (long)(out_row + 1) * K_attn + slid * VPT;
+        const device uint8_t* ws2 = (const device uint8_t*)W_oproj + (long)(out_row + 2) * K_attn + slid * VPT;
+        const device uint8_t* ws3 = (const device uint8_t*)W_oproj + (long)(out_row + 3) * K_attn + slid * VPT;
+
+        const device {sc_t}* sc0 = (const device {sc_t}*)S_oproj + (long)(out_row + 0) * K_groups + slid / {oproj_slid_divisor};
+        const device {sc_t}* sc1 = (const device {sc_t}*)S_oproj + (long)(out_row + 1) * K_groups + slid / {oproj_slid_divisor};
+        const device {sc_t}* sc2 = (const device {sc_t}*)S_oproj + (long)(out_row + 2) * K_groups + slid / {oproj_slid_divisor};
+        const device {sc_t}* sc3 = (const device {sc_t}*)S_oproj + (long)(out_row + 3) * K_groups + slid / {oproj_slid_divisor};
+
+        const device {sc_t}* bi0 = (const device {sc_t}*)B_oproj + (long)(out_row + 0) * K_groups + slid / {oproj_slid_divisor};
+        const device {sc_t}* bi1 = (const device {sc_t}*)B_oproj + (long)(out_row + 1) * K_groups + slid / {oproj_slid_divisor};
+        const device {sc_t}* bi2 = (const device {sc_t}*)B_oproj + (long)(out_row + 2) * K_groups + slid / {oproj_slid_divisor};
+        const device {sc_t}* bi3 = (const device {sc_t}*)B_oproj + (long)(out_row + 3) * K_groups + slid / {oproj_slid_divisor};
+
+        int xb = slid * VPT;
+
+        for (int k = 0; k < K_attn; k += BLOCK_SIZE) {{
+            // Load input ONCE per K-block
+            float xv[VPT];
+            float xsum = 0.0f;
+            for (int i = 0; i < VPT; i++) {{
+                xv[i] = float(attn_out[xb + i]);
+                xsum += xv[i];
+            }}
+
+            // Row 0
+            {{ float wacc = 0.0f;
+              for (int i = 0; i < VPT; i++) wacc += xv[i] * float(ws0[i]);
+              acc[0] += float(*sc0) * wacc + xsum * float(*bi0);
+              ws0 += BLOCK_SIZE; sc0 += {oproj_sc_stride}; bi0 += {oproj_sc_stride}; }}
+            // Row 1
+            {{ float wacc = 0.0f;
+              for (int i = 0; i < VPT; i++) wacc += xv[i] * float(ws1[i]);
+              acc[1] += float(*sc1) * wacc + xsum * float(*bi1);
+              ws1 += BLOCK_SIZE; sc1 += {oproj_sc_stride}; bi1 += {oproj_sc_stride}; }}
+            // Row 2
+            {{ float wacc = 0.0f;
+              for (int i = 0; i < VPT; i++) wacc += xv[i] * float(ws2[i]);
+              acc[2] += float(*sc2) * wacc + xsum * float(*bi2);
+              ws2 += BLOCK_SIZE; sc2 += {oproj_sc_stride}; bi2 += {oproj_sc_stride}; }}
+            // Row 3
+            {{ float wacc = 0.0f;
+              for (int i = 0; i < VPT; i++) wacc += xv[i] * float(ws3[i]);
+              acc[3] += float(*sc3) * wacc + xsum * float(*bi3);
+              ws3 += BLOCK_SIZE; sc3 += {oproj_sc_stride}; bi3 += {oproj_sc_stride}; }}
+
+            xb += BLOCK_SIZE;
+        }}
+
+        for (int tm = 0; tm < TM; tm++)
+            result[tm] = simd_sum(acc[tm]);
+
+        // Epilogue: addmm + x² + h_scaled + h_out
+        float x2_acc = 0.0f;
+        if (slid == 0) {{
+            for (int tm = 0; tm < TM; tm++) {{
+                int k = out_row + tm;
+                float h = result[tm] + float(residual[k]);
+                x2_acc += h * h;
+                h_scaled[k] = static_cast<bfloat16_t>(h * float(w_rms[k]));
+                h_out[k] = static_cast<bfloat16_t>(h);
+            }}
+        }}
+
+        // TG x² reduction: 8 SGs → 1 value
+        if (slid == 0) tgp_x2[sgid] = x2_acc;
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (sgid == 0 && slid == 0) {{
+            float total = 0.0f;
+            for (int s = 0; s < 8; s++) total += tgp_x2[s];
+            x2_partials[tg_x] = total;
+        }}
+
+    }} else if (tg_x < (uint)(N_OPROJ_TG + N_M1_TG)) {{
+        // ════════════════════════════════════════════════════════════════
+        // M1 GEMV: bf16, M=E, K=K_attn — M1 × attn_out → gate_part_a
+        // ════════════════════════════════════════════════════════════════
+        int local_tg = int(tg_x) - N_OPROJ_TG;
+        int out_row = local_tg * blockM_gate + int(sgid) * TM;
+        if (out_row >= E_CONST) return;
+        out_row = (out_row + TM <= E_CONST) ? out_row : (E_CONST - TM);
+
+        float result[TM] = {{0.0f, 0.0f, 0.0f, 0.0f}};
+        int bn = int(slid) * TN;
+        int n_iter = K_attn / blockN;
+
+        for (int i = 0; i < n_iter; i++) {{
+            float v[TN];
+            for (int tn = 0; tn < TN; tn++)
+                v[tn] = float(attn_out[bn + tn]);
+            for (int tm = 0; tm < TM; tm++) {{
+                float acc = 0.0f;
+                for (int tn = 0; tn < TN; tn++)
+                    acc += float(M1[(out_row + tm) * K_attn + bn + tn]) * v[tn];
+                result[tm] += acc;
+            }}
+            bn += blockN;
+        }}
+
+        // Handle remainder if K_attn not divisible by blockN
+        if (K_attn > n_iter * blockN) {{
+            for (int tm = 0; tm < TM; tm++) {{
+                for (int tn = 0; tn < TN; tn++) {{
+                    if (bn + tn < K_attn)
+                        result[tm] += float(M1[(out_row + tm) * K_attn + bn + tn])
+                                    * float(attn_out[bn + tn]);
+                }}
+            }}
+        }}
+
+        for (int tm = 0; tm < TM; tm++)
+            result[tm] = simd_sum(result[tm]);
+
+        if (slid == 0) {{
+            for (int tm = 0; tm < TM; tm++) {{
+                int e = out_row + tm;
+                if (e < E_CONST)
+                    gate_part_a[e] = result[tm];
+            }}
+        }}
+
+    }} else {{
+        // ════════════════════════════════════════════════════════════════
+        // W_FUSED GEMV: bf16, M=E, K=K_hidden — W_fused × residual → gate_part_b
+        // ════════════════════════════════════════════════════════════════
+        int local_tg = int(tg_x) - N_OPROJ_TG - N_M1_TG;
+        int out_row = local_tg * blockM_gate + int(sgid) * TM;
+        if (out_row >= E_CONST) return;
+        out_row = (out_row + TM <= E_CONST) ? out_row : (E_CONST - TM);
+
+        float result[TM] = {{0.0f, 0.0f, 0.0f, 0.0f}};
+        int bn = int(slid) * TN;
+        int n_iter = K_hidden / blockN;
+
+        for (int i = 0; i < n_iter; i++) {{
+            float v[TN];
+            for (int tn = 0; tn < TN; tn++)
+                v[tn] = float(residual[bn + tn]);
+            for (int tm = 0; tm < TM; tm++) {{
+                float acc = 0.0f;
+                for (int tn = 0; tn < TN; tn++)
+                    acc += float(W_fused[(out_row + tm) * K_hidden + bn + tn]) * v[tn];
+                result[tm] += acc;
+            }}
+            bn += blockN;
+        }}
+
+        if (K_hidden > n_iter * blockN) {{
+            for (int tm = 0; tm < TM; tm++) {{
+                for (int tn = 0; tn < TN; tn++) {{
+                    if (bn + tn < K_hidden)
+                        result[tm] += float(W_fused[(out_row + tm) * K_hidden + bn + tn])
+                                    * float(residual[bn + tn]);
+                }}
+            }}
+        }}
+
+        for (int tm = 0; tm < TM; tm++)
+            result[tm] = simd_sum(result[tm]);
+
+        if (slid == 0) {{
+            for (int tm = 0; tm < TM; tm++) {{
+                int e = out_row + tm;
+                if (e < E_CONST)
+                    gate_part_b[e] = result[tm];
+            }}
+        }}
+    }}
+"""
+
+
+_custom_oproj_8bit_kernels = {}
+
+
+def _get_custom_oproj_8bit_kernel(n_experts=64, group_size=64, scale_bf16=True):
+    key = (n_experts, group_size, scale_bf16)
+    if key not in _custom_oproj_8bit_kernels:
+        sc_tag = "_bf16sc" if scale_bf16 else ""
+        _custom_oproj_8bit_kernels[key] = mx.fast.metal_kernel(
+            name=f"custom_oproj_gate_gemv_8bit_e{n_experts}_gs{group_size}{sc_tag}",
+            input_names=[
+                "W_oproj", "S_oproj", "B_oproj",   # o_proj 8-bit weights
+                "attn_out",                          # (K_attn,) bf16
+                "residual",                          # (K,) bf16
+                "w_rms",                             # (K,) bf16 — RMSNorm weight
+                "M1",                                # (E, K_attn) bf16
+                "W_fused",                           # (E, K) bf16
+                "M_val", "K_attn_val", "K_hidden_val",
+                "N_OPROJ_TG_val", "N_M1_TG_val", "BM_GATE_val",
+            ],
+            output_names=["h_scaled", "h_out", "x2_partials",
+                          "gate_part_a", "gate_part_b"],
+            source=_gen_custom_oproj_8bit_source(n_experts, group_size, scale_bf16),
+        )
+    return _custom_oproj_8bit_kernels[key]
+
+
+def fused_custom_oproj_8bit(W_oproj, S_oproj, B_oproj,
+                             attn_out, residual, w_rms,
+                             M1, W_fused,
+                             M, K_attn, K_hidden=None,
+                             n_experts=64, gate_bm=8,
+                             group_size=64):
+    """Dispatch the fused 8-bit o_proj + bf16 gate GEMVs kernel.
+
+    Three GEMV types in one dispatch:
+    - TGs 0..N_OPROJ_TG-1: 8-bit o_proj GEMV → h_scaled, h_out, x2_partials
+    - TGs N_OPROJ_TG..+N_M1_TG: bf16 M1 GEMV → gate_part_a
+    - TGs +N_M1_TG..end: bf16 W_fused GEMV → gate_part_b
+
+    Args:
+        W_oproj/S_oproj/B_oproj: 8-bit quantized o_proj weights
+        attn_out: (K_attn,) bf16 — pre-o_proj attention output
+        residual: (K,) bf16 — input residual
+        w_rms: (K,) bf16 — post_attention_layernorm weight
+        M1: (E, K_attn) bf16 — precomputed W_fused @ W_oproj
+        W_fused: (E, K) bf16 — precomputed dequant(W_gate) * w_rms
+        M: hidden size (4096)
+        K_attn: attention output dim (8192)
+        K_hidden: hidden size for W_fused (defaults to M)
+        gate_bm: SGs per gate TG (1,2,4,8). Controls TG count.
+
+    Returns:
+        (h_scaled, h_out, x2_partials, gate_part_a, gate_part_b)
+    """
+    M = int(M)
+    K_attn = int(K_attn)
+    K_hidden = int(K_hidden) if K_hidden is not None else M
+    scale_bf16 = (S_oproj.dtype == mx.bfloat16)
+
+    kern = _get_custom_oproj_8bit_kernel(n_experts, group_size, scale_bf16)
+
+    n_oproj_tg = ceil_div(M, 32)  # 128 for M=4096
+    blockM_gate = gate_bm * 4     # rows per gate GEMV TG
+    n_m1_tg = ceil_div(n_experts, blockM_gate)
+    n_wf_tg = ceil_div(n_experts, blockM_gate)
+    total_tg = n_oproj_tg + n_m1_tg + n_wf_tg
+
+    results = kern(
+        inputs=[W_oproj, S_oproj, B_oproj,
+                attn_out, residual, w_rms, M1, W_fused,
+                M, K_attn, K_hidden, n_oproj_tg, n_m1_tg, gate_bm],
+        output_shapes=[(M,), (M,), (n_oproj_tg,), (n_experts,), (n_experts,)],
+        output_dtypes=[mx.bfloat16, mx.bfloat16, mx.float32,
+                       mx.float32, mx.float32],
+        grid=(total_tg * 32, 8, 1),    # total threads: total_tg * 256
+        threadgroup=(32, 8, 1),          # 256 threads per TG (8 SGs of 32)
+    )
+    return results[0], results[1], results[2], results[3], results[4]

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/custom_sdpa_pass1.py

@@ -0,0 +1,136 @@
+"""Custom SDPA Pass 1 for GQA attention (Dispatch 4).
+
+Online softmax over N key positions, strided by blocks.
+Constants baked into Metal source. Per-call varying values (N_keys, alloc_len)
+passed via a params array.
+"""
+
+import mlx.core as mx
+
+
+def _gen_sdpa_pass1_source(D=256, H_q=16, H_kv=2, blocks=128):
+    EPT = D // 32
+    gqa_factor = H_q // H_kv
+    scale = D ** -0.5
+
+    return f"""
+    const int D_DIM = {D};
+    const int EPT = {EPT};
+    const int H_Q = {H_q};
+    const int H_KV = {H_kv};
+    const int N_BLOCKS = {blocks};
+    const int GQA_FACTOR = {gqa_factor};
+    const float SCALE = {scale}f;
+
+    uint simd_lid = thread_index_in_simdgroup;
+    uint kv_head_idx = threadgroup_position_in_grid.x;
+    uint gqa_lane = thread_index_in_threadgroup / 32;
+
+    uint block_idx = threadgroup_position_in_grid.z % (uint)N_BLOCKS;
+    uint batch_idx = threadgroup_position_in_grid.z / (uint)N_BLOCKS;
+
+    int q_head_idx = (int)kv_head_idx * GQA_FACTOR + (int)gqa_lane;
+    int N_val = params[0];
+    int alloc = params[1];
+
+    int q_offset = ((int)batch_idx * H_Q + q_head_idx) * D_DIM;
+
+    int kv_head_offset = ((int)batch_idx * H_KV + (int)kv_head_idx) * alloc * D_DIM;
+    int k_offset = kv_head_offset + (int)block_idx * D_DIM;
+    int v_offset = kv_head_offset + (int)block_idx * D_DIM;
+
+    int out_head_offset = ((int)batch_idx * H_Q + q_head_idx);
+    int o_offset = (out_head_offset * N_BLOCKS + (int)block_idx) * D_DIM;
+    int s_offset = out_head_offset * N_BLOCKS + (int)block_idx;
+
+    float q[{EPT}];
+    for (int i = 0; i < EPT; i++) {{
+        q[i] = SCALE * (float)queries[q_offset + (int)simd_lid * EPT + i];
+    }}
+
+    float max_score = -__FLT_MAX__;
+    float sum_exp = 0.0f;
+    float o[{EPT}] = {{0}};
+
+    int k_stride = N_BLOCKS * D_DIM;
+
+    for (int pos = (int)block_idx; pos < N_val; pos += N_BLOCKS) {{
+        float score = 0.0f;
+        for (int i = 0; i < EPT; i++) {{
+            score += q[i] * (float)keys[k_offset + (int)simd_lid * EPT + i];
+        }}
+        score = simd_sum(score);
+
+        float new_max = metal::max(max_score, score);
+        float factor = metal::fast::exp(max_score - new_max);
+        float exp_score = metal::fast::exp(score - new_max);
+
+        max_score = new_max;
+        sum_exp = sum_exp * factor + exp_score;
+
+        for (int i = 0; i < EPT; i++) {{
+            o[i] = o[i] * factor + exp_score * (float)values[v_offset + (int)simd_lid * EPT + i];
+        }}
+
+        k_offset += k_stride;
+        v_offset += k_stride;
+    }}
+
+    if (simd_lid == 0) {{
+        sums[s_offset] = sum_exp;
+        maxs[s_offset] = max_score;
+    }}
+
+    for (int i = 0; i < EPT; i++) {{
+        o_partials[o_offset + (int)simd_lid * EPT + i] = static_cast<bfloat16_t>(o[i]);
+    }}
+"""
+
+
+_pass1_cache = {}
+
+
+def _get_pass1_kernel(D, H_q, H_kv, blocks, gqa_factor):
+    key = (D, H_q, H_kv, blocks)
+    if key not in _pass1_cache:
+        _pass1_cache[key] = mx.fast.metal_kernel(
+            name=f"sdpa_pass1_D{D}_Hq{H_q}_Hkv{H_kv}_b{blocks}",
+            input_names=["queries", "keys", "values", "params"],
+            output_names=["o_partials", "sums", "maxs"],
+            source=_gen_sdpa_pass1_source(D, H_q, H_kv, blocks),
+        )
+    return _pass1_cache[key]
+
+
+def custom_sdpa_pass1(queries, keys, values, scale, H_q, H_kv, D,
+                      blocks=128, batch_size=1, N=None, alloc_len=None,
+                      scalars=None):
+    B = batch_size
+    gqa_factor = H_q // H_kv
+
+    kern = _get_pass1_kernel(D, H_q, H_kv, blocks, gqa_factor)
+
+    if N is None:
+        N = keys.shape[2]
+    if alloc_len is None:
+        alloc_len = N
+
+    # Pack per-call varying values into a params array (always a pointer)
+    params = mx.array([int(N), int(alloc_len)], dtype=mx.int32)
+
+    results = kern(
+        inputs=[queries, keys, values, params],
+        output_shapes=[
+            (B * H_q * blocks * D,),
+            (B * H_q * blocks,),
+            (B * H_q * blocks,),
+        ],
+        output_dtypes=[mx.bfloat16, mx.float32, mx.float32],
+        grid=(H_kv * 32, gqa_factor, blocks * B),
+        threadgroup=(32, gqa_factor, 1),
+    )
+
+    o_partials = results[0].reshape(B, H_q, blocks, D)
+    sums = results[1].reshape(B, H_q, blocks)
+    maxs = results[2].reshape(B, H_q, blocks)
+    return o_partials, sums, maxs

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/custom_sdpa_pass2_gate.py

@@ -0,0 +1,127 @@
+"""Custom SDPA Pass 2 + gate multiply for GQA attention (Dispatch 5).
+
+32-SG block-parallel reduction + V_SPLIT for high GPU utilization.
+All constants baked into Metal source (no scalar kernel inputs).
+"""
+
+import mlx.core as mx
+
+
+def _gen_sdpa_pass2_gate_source(D=256, V_SPLIT=4, H_q=16, blocks=128):
+    BN = 32
+    V_PER_SPLIT = D // V_SPLIT
+    EPT = V_PER_SPLIT // BN
+
+    return f"""
+    const int D_DIM = {D};
+    const int V_SPLIT = {V_SPLIT};
+    const int V_PER_SPLIT = {V_PER_SPLIT};
+    const int EPT = {EPT};
+    const int BN = {BN};
+    const int H_Q = {H_q};
+    const int N_BLOCKS = {blocks};
+
+    uint tg_idx = threadgroup_position_in_grid.x;
+    uint simd_gid = simdgroup_index_in_threadgroup;
+    uint simd_lid = thread_index_in_simdgroup;
+    uint b_idx = threadgroup_position_in_grid.z;
+
+    int head_idx = (int)tg_idx / V_SPLIT;
+    int split_idx = (int)tg_idx % V_SPLIT;
+    int v_offset = split_idx * V_PER_SPLIT;
+
+    int head_base = ((int)b_idx * H_Q + head_idx);
+
+    int p_base = head_base * N_BLOCKS * D_DIM
+               + (int)simd_gid * D_DIM
+               + v_offset + (int)simd_lid * EPT;
+    int ms_base = head_base * N_BLOCKS;
+
+    // ── Phase 1: Find global max across all blocks ──
+    float local_max = -__FLT_MAX__;
+    for (int b = 0; b < N_BLOCKS / BN; ++b) {{
+        local_max = metal::max(local_max, maxs[ms_base + (int)simd_lid + BN * b]);
+    }}
+    float global_max = simd_max(local_max);
+
+    // ── Phase 2: Compute global sum_exp ──
+    float local_sum = 0.0f;
+    for (int b = 0; b < N_BLOCKS / BN; ++b) {{
+        float factor = metal::fast::exp(maxs[ms_base + (int)simd_lid + BN * b] - global_max);
+        local_sum += factor * sums[ms_base + (int)simd_lid + BN * b];
+    }}
+    float global_sum = simd_sum(local_sum);
+
+    // ── Phase 3: Accumulate V-partials (block-parallel via SGs) ──
+    float o[EPT] = {{0}};
+    for (int b = 0; b < N_BLOCKS / BN; ++b) {{
+        float factor = metal::fast::exp(maxs[ms_base + (int)simd_gid] - global_max);
+        for (int i = 0; i < EPT; i++) {{
+            o[i] += factor * (float)o_partials[p_base + i];
+        }}
+        ms_base += BN;
+        p_base += BN * D_DIM;
+    }}
+
+    // ── Phase 4: Shared memory transpose + reduce across SGs ──
+    threadgroup float tg_mem[BN * BN];
+
+    for (int i = 0; i < EPT; i++) {{
+        tg_mem[(int)simd_lid * BN + (int)simd_gid] = o[i];
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        o[i] = simd_sum(tg_mem[(int)simd_gid * BN + (int)simd_lid]);
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }}
+
+    // ── Phase 5: Normalize + gate multiply + write (simd_lid == 0 only) ──
+    if (simd_lid == 0) {{
+        float inv_sum = (global_sum > 0.0f) ? (1.0f / global_sum) : 0.0f;
+        int out_base = (int)b_idx * H_Q * D_DIM + head_idx * D_DIM
+                     + v_offset + (int)simd_gid * EPT;
+        for (int i = 0; i < EPT; i++) {{
+            float val = o[i] * inv_sum;
+            float gate = gate_sigmoid[out_base + i];
+            attn_output[out_base + i] = static_cast<bfloat16_t>(val * gate);
+        }}
+    }}
+"""
+
+
+_pass2_cache = {}
+
+
+def _get_pass2_kernel(D, V_SPLIT, H_q, blocks):
+    key = (D, V_SPLIT, H_q, blocks)
+    if key not in _pass2_cache:
+        _pass2_cache[key] = mx.fast.metal_kernel(
+            name=f"sdpa_pass2_gate_D{D}_vs{V_SPLIT}_Hq{H_q}_b{blocks}",
+            input_names=["o_partials", "sums", "maxs", "gate_sigmoid"],
+            output_names=["attn_output"],
+            source=_gen_sdpa_pass2_gate_source(D, V_SPLIT, H_q, blocks),
+        )
+    return _pass2_cache[key]
+
+
+def custom_sdpa_pass2_gate(o_partials, sums, maxs, gate_sigmoid,
+                           H_q, D, blocks=128, V_SPLIT=4, batch_size=1,
+                           scalars=None):
+    B = batch_size
+
+    kern = _get_pass2_kernel(D, V_SPLIT, H_q, blocks)
+
+    o_flat = o_partials.reshape(B * H_q * blocks * D)
+    s_flat = sums.reshape(B * H_q * blocks)
+    m_flat = maxs.reshape(B * H_q * blocks)
+    g_flat = gate_sigmoid.reshape(B * H_q * D)
+
+    n_tgs = H_q * V_SPLIT
+    BN = 32
+    results = kern(
+        inputs=[o_flat, s_flat, m_flat, g_flat],
+        output_shapes=[(B * H_q * D,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(n_tgs * BN, BN, B),
+        threadgroup=(BN, BN, 1),
+    )
+
+    return results[0].reshape(B, 1, H_q * D)

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_gdn_projections_8bit.py

@@ -0,0 +1,288 @@
+"""Fused GDN projections for Qwen3.5-35B-A3B (Dispatch 2).
+
+Single dispatch fuses 4 quantized 8-bit GEMVs + depthwise conv1d + activations:
+  - in_proj_qkv (8192×2048): GEMV → conv1d(4-tap) → SiLU → write bf16 + cache update
+  - in_proj_z   (4096×2048): GEMV → SiLU → write f32
+  - in_proj_b   (32×2048):   GEMV → sigmoid → write f32 (beta for GDN kernel)
+  - in_proj_a   (32×2048):   GEMV → g=exp(-exp(A_log)*softplus(a+dt_bias)) → write f32
+
+All 4 projection weight matrices are pre-merged into one contiguous buffer
+(W_merged, S_merged, B_merged) for better memory locality and cache behavior.
+Merging is done offline at patch time by _patch_gdn_proj_weights().
+
+B/A epilogues compute g and beta in-kernel, eliminating ~8 micro-dispatches
+that gated_delta_update would otherwise generate (sigmoid, exp, log, etc.).
+The caller can pass g/beta directly to gated_delta_kernel.
+
+TG-level multiplexing: tgid.y routes to different epilogues.
+Each TG: 64 threads = 2 SGs of 32, produces 8 output rows (4 per SG).
+Standard 8-bit affine GEMV: result = scale * Σ(x[i]*w[i]) + bias * Σ(x[i])
+
+Grid: (32, total_tg * 2, B), TG: (32, 2, 1)
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_fused_gdn_projections_source(K, N_QKV, N_Z, N_B, N_A, group_size=64):
+    """Generate Metal source for fused GDN projections with merged weights.
+
+    All constants baked into Metal source (no scalar kernel inputs).
+    """
+    gs = int(group_size)
+    sc_stride = 256 // gs
+    slid_div = gs // 8
+    N_TOTAL = N_QKV + N_Z + N_B + N_A
+    K_groups = K // gs
+    N_QKV_TG = ceil_div(N_QKV, 8)
+    N_Z_TG = ceil_div(N_Z, 8)
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int GROUP_SIZE = {gs};
+    const int SC_STRIDE = {sc_stride};
+    const int SLID_DIV = {slid_div};
+    const int K = {K};
+    const int K_groups = {K_groups};
+    const int N_QKV = {N_QKV};
+    const int N_Z = {N_Z};
+    const int N_B = {N_B};
+    const int N_TOTAL = {N_TOTAL};
+    const int N_QKV_TG = {N_QKV_TG};
+    const int N_Z_TG = {N_Z_TG};
+    const int N_B_TG = {ceil_div(N_B, 8)};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;  // 0 or 1
+    uint slid = thread_index_in_simdgroup;       // 0..31
+    int b_idx = tgid.z;
+
+    int tg = tgid.y;
+
+    // ─── Determine region and absolute out_row in merged matrix ───
+    int out_row;
+    int region;  // 0=QKV, 1=Z, 2=B, 3=A
+
+    if (tg < N_QKV_TG) {{
+        region = 0;
+        out_row = tg * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_QKV_TG + N_Z_TG) {{
+        region = 1;
+        out_row = N_QKV + (tg - N_QKV_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_QKV_TG + N_Z_TG + N_B_TG) {{
+        region = 2;
+        out_row = N_QKV + N_Z + (tg - N_QKV_TG - N_Z_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else {{
+        region = 3;
+        out_row = N_QKV + N_Z + N_B + (tg - N_QKV_TG - N_Z_TG - N_B_TG) * 8 + sgid * RESULTS_PER_SG;
+    }}
+
+    if (out_row >= N_TOTAL) return;
+
+    // ─── Single pointer into merged weight buffer ───
+    const device uint8_t* ws = (const device uint8_t*)W_merged + (long)out_row * K + slid * VALUES_PER_THREAD;
+    const device bfloat16_t* sc = (const device bfloat16_t*)S_merged + (long)out_row * K_groups + slid / SLID_DIV;
+    const device bfloat16_t* bi = (const device bfloat16_t*)B_merged + (long)out_row * K_groups + slid / SLID_DIV;
+
+    // ─── 8-bit GEMV K-loop (unified for all regions) ───
+    float result[4] = {{0, 0, 0, 0}};
+    int x_base = b_idx * K + slid * VALUES_PER_THREAD;
+
+    for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{
+        float x_thread[8];
+        float xsum = 0;
+        for (int i = 0; i < 8; i++) {{
+            float xi = float(x[x_base + i]);
+            x_thread[i] = xi;
+            xsum += xi;
+        }}
+
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            const device uint8_t* w = ws + row * K;
+            float s_val = float(sc[row * K_groups]);
+            float b_val = float(bi[row * K_groups]);
+            float accum = 0;
+            for (int i = 0; i < 8; i++) {{
+                accum += x_thread[i] * float(w[i]);
+            }}
+            result[row] += s_val * accum + xsum * b_val;
+        }}
+
+        ws += BLOCK_SIZE;
+        sc += SC_STRIDE;
+        bi += SC_STRIDE;
+        x_base += BLOCK_SIZE;
+    }}
+
+    // ─── Reduction ───
+    for (int row = 0; row < RESULTS_PER_SG; row++) {{
+        result[row] = simd_sum(result[row]);
+    }}
+
+    // ─── Region-specific epilogues ───
+    // After simd_sum, all 32 threads have result[0..3].
+    // Threads 0-3 each handle one output row.
+
+    if (region == 0) {{
+        // ═══ QKV: conv1d(4-tap) + SiLU + cache update ═══
+        int c = out_row + (int)slid;  // channel index (= absolute row for QKV)
+        if (slid < (uint)RESULTS_PER_SG && c < N_QKV) {{
+            float qkv_val = result[slid];
+
+            int conv_dim = N_QKV;
+            long cs_base = (long)b_idx * 3 * conv_dim;
+            float s0 = float(conv_state[cs_base + 0 * conv_dim + c]);
+            float s1 = float(conv_state[cs_base + 1 * conv_dim + c]);
+            float s2 = float(conv_state[cs_base + 2 * conv_dim + c]);
+
+            float conv_out = float(conv_w[c * 4 + 0]) * s0
+                           + float(conv_w[c * 4 + 1]) * s1
+                           + float(conv_w[c * 4 + 2]) * s2
+                           + float(conv_w[c * 4 + 3]) * qkv_val;
+
+            float silu_out = conv_out / (1.0f + metal::exp(-conv_out));
+
+            conv_state_out[cs_base + 0 * conv_dim + c] = static_cast<bfloat16_t>(s1);
+            conv_state_out[cs_base + 1 * conv_dim + c] = static_cast<bfloat16_t>(s2);
+            conv_state_out[cs_base + 2 * conv_dim + c] = static_cast<bfloat16_t>(qkv_val);
+
+            qkv_out[b_idx * conv_dim + c] = static_cast<bfloat16_t>(silu_out);
+        }}
+
+    }} else if (region == 1) {{
+        // ═══ Z: SiLU → write f32 ═══
+        int z_row = out_row - N_QKV + (int)slid;
+        if (slid < (uint)RESULTS_PER_SG && z_row < N_Z) {{
+            float val = result[slid];
+            float silu_val = val / (1.0f + metal::exp(-val));
+            z_silu_out[b_idx * N_Z + z_row] = silu_val;
+        }}
+
+    }} else if (region == 2) {{
+        // ═══ B: sigmoid(result) → beta (f32) ═══
+        int b_row = out_row - N_QKV - N_Z + (int)slid;
+        if (slid < (uint)RESULTS_PER_SG && b_row < N_B) {{
+            float val = result[slid];
+            float beta = 1.0f / (1.0f + metal::exp(-val));
+            b_out[b_idx * N_B + b_row] = beta;
+        }}
+
+    }} else {{
+        // ═══ A: g = exp(-exp(A_log) * softplus(a + dt_bias)) → f32 ═══
+        int a_row = out_row - N_QKV - N_Z - N_B + (int)slid;
+        int N_A = N_TOTAL - N_QKV - N_Z - N_B;
+        if (slid < (uint)RESULTS_PER_SG && a_row < N_A) {{
+            float a_val = result[slid];
+            float dt = float(dt_bias_arr[a_row]);
+            float x_g = a_val + dt;
+            // softplus(x) = log(1 + exp(x)), with x>20 shortcut for numerical stability
+            float sp = (x_g > 20.0f) ? x_g : metal::log(1.0f + metal::exp(x_g));
+            float g_val = metal::exp(-metal::exp(float(A_log_arr[a_row])) * sp);
+            a_out[b_idx * N_A + a_row] = g_val;
+        }}
+    }}
+"""
+
+
+_fused_gdn_proj_cache = {}
+
+
+def _get_fused_gdn_proj_kernel(K, N_QKV, N_Z, N_B, N_A, group_size=64):
+    key = (K, N_QKV, N_Z, N_B, N_A, group_size)
+    if key not in _fused_gdn_proj_cache:
+        _fused_gdn_proj_cache[key] = mx.fast.metal_kernel(
+            name=f"fused_gdn_proj_K{K}_NQKV{N_QKV}_NZ{N_Z}_NB{N_B}_NA{N_A}",
+            input_names=[
+                "x",
+                "W_merged", "S_merged", "B_merged",
+                "conv_state", "conv_w",
+                "A_log_arr", "dt_bias_arr",
+            ],
+            output_names=["qkv_out", "z_silu_out", "b_out", "a_out", "conv_state_out"],
+            source=_gen_fused_gdn_projections_source(K, N_QKV, N_Z, N_B, N_A, group_size),
+        )
+    return _fused_gdn_proj_cache[key]
+
+
+def fused_gdn_projections(
+    x,
+    W_merged, S_merged, B_merged,
+    proj_dims,
+    conv_state, conv_weights,
+    A_log, dt_bias,
+    batch_size=1,
+):
+    """Fused GDN projections: 4 GEMVs + conv1d + activations + g/beta.
+
+    Uses pre-merged contiguous weight buffers for all 4 projections.
+    B epilogue computes beta = sigmoid(b) in f32.
+    A epilogue computes g = exp(-exp(A_log) * softplus(a + dt_bias)) in f32.
+    Caller passes g/beta directly to gated_delta_kernel (no micro-dispatches).
+
+    Args:
+        x: [B, 1, K] bf16 — post-RMSNorm hidden state
+        W_merged: [N_TOTAL, K/4] uint32 — merged quantized weights
+        S_merged: [N_TOTAL, K/gs] bf16 — merged scales
+        B_merged: [N_TOTAL, K/gs] bf16 — merged biases
+        proj_dims: (N_QKV, N_Z, N_B, N_A) — per-projection output dims
+        conv_state:  [B, 3, conv_dim] bf16 — previous 3 timesteps
+        conv_weights: [conv_dim, 4, 1] or [conv_dim, 4] bf16 — depthwise conv filters
+        A_log: [Hv] f32 — GDN decay log-parameter
+        dt_bias: [Hv] f32 — GDN time constant bias
+        batch_size: int
+
+    Returns:
+        qkv_conv_silu: [B, 1, N_QKV] bf16 — post-conv, post-SiLU
+        z_silu:        [B, 1, N_Z] f32  — post-SiLU
+        beta:          [B, 1, N_B] f32  — sigmoid(b), ready for GDN kernel
+        g:             [B, 1, N_A] f32  — gating, ready for GDN kernel
+        conv_state_out: [B, 3, N_QKV] bf16
+    """
+    B = batch_size
+
+    N_QKV, N_Z, N_B, N_A = proj_dims
+    K = x.shape[-1]
+
+    kern = _get_fused_gdn_proj_kernel(K, N_QKV, N_Z, N_B, N_A)
+
+    N_QKV_TG = ceil_div(N_QKV, 8)
+    N_Z_TG = ceil_div(N_Z, 8)
+    N_B_TG = ceil_div(N_B, 8)
+    N_A_TG = ceil_div(N_A, 8)
+    total_tg = N_QKV_TG + N_Z_TG + N_B_TG + N_A_TG
+
+    conv_w_flat = conv_weights.reshape(-1, 4) if conv_weights.ndim == 3 else conv_weights
+    x_flat = x.reshape(B, K)
+
+    results = kern(
+        inputs=[
+            x_flat,
+            W_merged, S_merged, B_merged,
+            conv_state, conv_w_flat,
+            A_log, dt_bias,
+        ],
+        output_shapes=[
+            (B * N_QKV,),              # qkv_out
+            (B * N_Z,),                # z_silu_out
+            (B * N_B,),                # beta_out (f32)
+            (B * N_A,),                # g_out (f32)
+            (B * 3 * N_QKV,),          # conv_state_out
+        ],
+        output_dtypes=[mx.bfloat16, mx.float32, mx.float32, mx.float32, mx.bfloat16],
+        grid=(32, total_tg * 2, B),
+        threadgroup=(32, 2, 1),
+    )
+
+    qkv_out = results[0].reshape(B, 1, N_QKV)
+    z_silu = results[1].reshape(B, 1, N_Z)
+    beta = results[2].reshape(B, 1, N_B)
+    g = results[3].reshape(B, 1, N_A)
+    conv_state_out = results[4].reshape(B, 3, N_QKV)
+
+    return qkv_out, z_silu, beta, g, conv_state_out

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_gqa_projections_8bit.py

@@ -0,0 +1,178 @@
+"""Fused GQA projections for Qwen3.5 (Dispatch 1).
+
+Single dispatch fuses 4 quantized 8-bit GEMVs with region-specific epilogues:
+  - q_proj queries: GEMV → raw bf16 write
+  - q_proj gate:    GEMV → sigmoid → f32 write
+  - k_proj:         GEMV → raw bf16 write
+  - v_proj:         GEMV → raw bf16 write
+
+All constants baked into Metal source (no scalar kernel inputs).
+"""
+
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_fused_gqa_projections_source(K, N_Q, N_GATE, N_K, N_V, group_size=64):
+    gs = int(group_size)
+    sc_stride = 256 // gs
+    slid_div = gs // 8
+    N_TOTAL = N_Q + N_GATE + N_K + N_V
+    K_groups = K // gs
+    N_Q_TG = ceil_div(N_Q, 8)
+    N_GATE_TG = ceil_div(N_GATE, 8)
+    N_K_TG = ceil_div(N_K, 8)
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+    const int GROUP_SIZE = {gs};
+    const int SC_STRIDE = {sc_stride};
+    const int SLID_DIV = {slid_div};
+    const int K = {K};
+    const int K_groups = {K_groups};
+    const int N_Q = {N_Q};
+    const int N_GATE = {N_GATE};
+    const int N_K = {N_K};
+    const int N_TOTAL = {N_TOTAL};
+    const int N_Q_TG = {N_Q_TG};
+    const int N_GATE_TG = {N_GATE_TG};
+    const int N_K_TG = {N_K_TG};
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;
+    uint slid = thread_index_in_simdgroup;
+    int b_idx = tgid.z;
+    int tg = tgid.y;
+
+    int out_row;
+    int region;
+
+    if (tg < N_Q_TG) {{
+        region = 0;
+        out_row = tg * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_Q_TG + N_GATE_TG) {{
+        region = 1;
+        out_row = N_Q + (tg - N_Q_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else if (tg < N_Q_TG + N_GATE_TG + N_K_TG) {{
+        region = 2;
+        out_row = N_Q + N_GATE + (tg - N_Q_TG - N_GATE_TG) * 8 + sgid * RESULTS_PER_SG;
+    }} else {{
+        region = 3;
+        out_row = N_Q + N_GATE + N_K + (tg - N_Q_TG - N_GATE_TG - N_K_TG) * 8 + sgid * RESULTS_PER_SG;
+    }}
+
+    if (out_row >= N_TOTAL) return;
+
+    const device uint8_t* ws = (const device uint8_t*)W_merged + (long)out_row * K + slid * VALUES_PER_THREAD;
+    const device bfloat16_t* sc = (const device bfloat16_t*)S_merged + (long)out_row * K_groups + slid / SLID_DIV;
+    const device bfloat16_t* bi = (const device bfloat16_t*)B_merged + (long)out_row * K_groups + slid / SLID_DIV;
+
+    float result[4] = {{0, 0, 0, 0}};
+    int x_base = b_idx * K + slid * VALUES_PER_THREAD;
+
+    for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{
+        float x_thread[8];
+        float xsum = 0;
+        for (int i = 0; i < 8; i++) {{
+            float xi = float(x[x_base + i]);
+            x_thread[i] = xi;
+            xsum += xi;
+        }}
+
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            const device uint8_t* w = ws + row * K;
+            float s_val = float(sc[row * K_groups]);
+            float b_val = float(bi[row * K_groups]);
+            float accum = 0;
+            for (int i = 0; i < 8; i++) {{
+                accum += x_thread[i] * float(w[i]);
+            }}
+            result[row] += s_val * accum + xsum * b_val;
+        }}
+
+        ws += BLOCK_SIZE;
+        sc += SC_STRIDE;
+        bi += SC_STRIDE;
+        x_base += BLOCK_SIZE;
+    }}
+
+    for (int row = 0; row < RESULTS_PER_SG; row++) {{
+        result[row] = simd_sum(result[row]);
+    }}
+
+    if (region == 0) {{
+        int q_row = out_row + (int)slid;
+        if (slid < (uint)RESULTS_PER_SG && q_row < N_Q) {{
+            q_out[b_idx * N_Q + q_row] = static_cast<bfloat16_t>(result[slid]);
+        }}
+    }} else if (region == 1) {{
+        int g_row = out_row - N_Q + (int)slid;
+        if (slid < (uint)RESULTS_PER_SG && g_row < N_GATE) {{
+            float val = result[slid];
+            float sig = 1.0f / (1.0f + metal::exp(-val));
+            gate_out[b_idx * N_GATE + g_row] = sig;
+        }}
+    }} else if (region == 2) {{
+        int k_row = out_row - N_Q - N_GATE + (int)slid;
+        if (slid < (uint)RESULTS_PER_SG && k_row < N_K) {{
+            k_out[b_idx * N_K + k_row] = static_cast<bfloat16_t>(result[slid]);
+        }}
+    }} else {{
+        int v_row = out_row - N_Q - N_GATE - N_K + (int)slid;
+        int N_V = N_TOTAL - N_Q - N_GATE - N_K;
+        if (slid < (uint)RESULTS_PER_SG && v_row < N_V) {{
+            v_out[b_idx * N_V + v_row] = static_cast<bfloat16_t>(result[slid]);
+        }}
+    }}
+"""
+
+
+_fused_gqa_proj_cache = {}
+
+
+def _get_fused_gqa_proj_kernel(K, N_Q, N_GATE, N_K, N_V, group_size=64):
+    key = (K, N_Q, N_GATE, N_K, N_V, group_size)
+    if key not in _fused_gqa_proj_cache:
+        _fused_gqa_proj_cache[key] = mx.fast.metal_kernel(
+            name=f"fused_gqa_proj_K{K}_NQ{N_Q}_NG{N_GATE}_NK{N_K}_NV{N_V}",
+            input_names=["x", "W_merged", "S_merged", "B_merged"],
+            output_names=["q_out", "gate_out", "k_out", "v_out"],
+            source=_gen_fused_gqa_projections_source(K, N_Q, N_GATE, N_K, N_V, group_size),
+        )
+    return _fused_gqa_proj_cache[key]
+
+
+def fused_gqa_projections(
+    x, W_merged, S_merged, B_merged, proj_dims, batch_size=1,
+    scalars=None, total_tg=None,
+):
+    B = batch_size
+    N_Q, N_GATE, N_K, N_V = proj_dims
+    K = x.shape[-1]
+
+    kern = _get_fused_gqa_proj_kernel(K, N_Q, N_GATE, N_K, N_V)
+
+    if total_tg is None:
+        total_tg = ceil_div(N_Q, 8) + ceil_div(N_GATE, 8) + ceil_div(N_K, 8) + ceil_div(N_V, 8)
+
+    x_flat = x.reshape(B, K)
+
+    results = kern(
+        inputs=[x_flat, W_merged, S_merged, B_merged],
+        output_shapes=[
+            (B * N_Q,), (B * N_GATE,), (B * N_K,), (B * N_V,),
+        ],
+        output_dtypes=[mx.bfloat16, mx.float32, mx.bfloat16, mx.bfloat16],
+        grid=(32, total_tg * 2, B),
+        threadgroup=(32, 2, 1),
+    )
+
+    return (results[0].reshape(B, 1, N_Q),
+            results[1].reshape(B, 1, N_GATE),
+            results[2].reshape(B, 1, N_K),
+            results[3].reshape(B, 1, N_V))

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_moe_epilogue_qwen.py

@@ -0,0 +1,97 @@
+"""Fused MoE epilogue for Qwen3.5: weighted sum + shared expert gate + residual.
+
+Computes:
+  Y[j] = bf16( Σ_a(scores[a] * D_routed[a,j]) + gate_shared * D_shared[j] + H[j] )
+
+Two modes:
+  fuse_sigmoid=False: gate_shared is pre-computed sigmoid output (scalar f32)
+  fuse_sigmoid=True:  gate_raw is raw dot product; sigmoid computed internally
+"""
+import mlx.core as mx
+
+
+def _gen_epilogue_qwen_source(K, n_active, fuse_sigmoid=False):
+    """Metal source for fused MoE epilogue with shared expert gate."""
+    if fuse_sigmoid:
+        gate_line = """
+    // Compute sigmoid from raw gate value
+    float gate = 1.0f / (1.0f + metal::exp(-shared_gate_val));"""
+    else:
+        gate_line = """
+    float gate = shared_gate_val;"""
+
+    return f"""
+    const int K_const = {K};
+    const int n_active_const = {n_active};
+
+    uint tid = thread_position_in_grid.x;
+    if (tid >= K_const) return;
+
+    // Weighted sum of routed expert outputs
+    float acc = 0.0f;
+    for (int a = 0; a < n_active_const; a++) {{
+        acc += scores[a] * D_routed[a * K_const + tid];
+    }}
+{gate_line}
+
+    // Shared expert: multiply by gate, add to accumulator
+    float shared_val = float(D_shared[tid]) * gate;
+
+    // Add residual and write
+    Y[tid] = static_cast<bfloat16_t>(acc + shared_val + float(H[tid]));
+"""
+
+
+_epilogue_qwen_kernels = {}
+
+
+def _get_epilogue_qwen_kernel(K, n_active, fuse_sigmoid=False):
+    key = (K, n_active, fuse_sigmoid)
+    if key not in _epilogue_qwen_kernels:
+        sig_tag = "_fsig" if fuse_sigmoid else ""
+        _epilogue_qwen_kernels[key] = mx.fast.metal_kernel(
+            name=f"fused_moe_epilogue_qwen_K{K}_n{n_active}{sig_tag}",
+            input_names=["D_routed", "D_shared", "scores", "H",
+                         "shared_gate_val"],
+            output_names=["Y"],
+            source=_gen_epilogue_qwen_source(K, n_active, fuse_sigmoid),
+        )
+    return _epilogue_qwen_kernels[key]
+
+
+def fused_moe_epilogue_qwen(d_routed, d_shared, scores, h,
+                             shared_gate, k_val, fuse_sigmoid=False):
+    """Fused MoE epilogue with shared expert gate.
+
+    Args:
+        d_routed: routed expert outputs (n_active, K) float32
+        d_shared: shared expert output (K,) float32
+        scores: normalized routing scores (n_active,) float32
+        h: residual hidden state (K,) bfloat16
+        shared_gate: sigmoid output (fuse_sigmoid=False) or raw dot product
+                     (fuse_sigmoid=True), scalar float32
+        k_val: hidden dimension
+        fuse_sigmoid: if True, compute sigmoid(shared_gate) internally
+
+    Returns:
+        Y: (K,) bfloat16 — final layer output
+    """
+    K = int(k_val)
+    n_active = scores.shape[0]
+    kern = _get_epilogue_qwen_kernel(K, n_active, fuse_sigmoid)
+
+    # shared_gate is a scalar — pass as 0-d array
+    if shared_gate.ndim > 0:
+        shared_gate = shared_gate.reshape(())
+
+    tg_size = min(K, 1024)
+    n_tg = (K + tg_size - 1) // tg_size
+
+    Y = kern(
+        inputs=[d_routed, d_shared, scores, h, shared_gate],
+        output_shapes=[(K,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(n_tg * tg_size, 1, 1),
+        threadgroup=(tg_size, 1, 1),
+    )
+    return Y[0]

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_qk_norm_rope.py

@@ -0,0 +1,152 @@
+"""Fused Q/K RMSNorm + RoPE for GQA attention (Dispatch 2).
+
+Performs per-head RMSNorm with learned weight (head_dim=256) then applies
+RoPE on the first 64 dims (partial_rotary_factor=0.25) using non-traditional
+pairing: element p pairs with p+32 (not p+1).
+
+cos/sin values precomputed in Python from inv_freq * position,
+passed as array inputs (no scalar kernel inputs).
+
+All constants baked into Metal source (no scalar kernel inputs).
+"""
+
+import mlx.core as mx
+
+
+def _gen_fused_qk_norm_rope_source(H_q=16, H_kv=2, D=256, rope_dims=64):
+    N_READS = D // 32
+    ROPE_HALF = rope_dims // 2
+    ROPE_THREADS = rope_dims // N_READS
+    FIRST_HALF = ROPE_THREADS // 2
+    PARTNER_XOR = FIRST_HALF
+
+    return f"""
+    const int H_Q = {H_q};
+    const int H_KV = {H_kv};
+    const int D_DIM = {D};
+    const int N_READS = {N_READS};
+    const float EPS = 1e-6f;
+
+    uint head_idx = threadgroup_position_in_grid.x;
+    uint slid = thread_index_in_simdgroup;
+    uint b_idx = thread_position_in_grid.z;
+
+    bool is_q = (head_idx < (uint)H_Q);
+    int head_local = is_q ? (int)head_idx : ((int)head_idx - H_Q);
+
+    // ── Phase 1: Load N_READS elements + partial sum of squares ──
+    int in_base = is_q
+        ? (int)(b_idx * H_Q * D_DIM + (int)head_idx * D_DIM)
+        : (int)(b_idx * H_KV * D_DIM + head_local * D_DIM);
+
+    int elem_base = (int)slid * N_READS;
+    float vals[{N_READS}];
+    float partial_sq = 0.0f;
+
+    for (int i = 0; i < N_READS; i++) {{
+        float xi;
+        if (is_q)
+            xi = (float)queries[in_base + elem_base + i];
+        else
+            xi = (float)keys[in_base + elem_base + i];
+        vals[i] = xi;
+        partial_sq += xi * xi;
+    }}
+
+    // ── Phase 2: RMSNorm reduction ──
+    float sum_sq = simd_sum(partial_sq);
+    float inv_rms = metal::precise::rsqrt(sum_sq / (float)D_DIM + EPS);
+
+    // ── Phase 3: Normalize with learned weight ──
+    for (int i = 0; i < N_READS; i++) {{
+        float w;
+        if (is_q)
+            w = (float)q_norm_w[elem_base + i];
+        else
+            w = (float)k_norm_w[elem_base + i];
+        vals[i] = vals[i] * inv_rms * w;
+    }}
+
+    // ── Phase 4: RoPE on first {rope_dims} dims ──
+    if (slid < {ROPE_THREADS}u) {{
+        ushort partner = (ushort)(slid ^ {PARTNER_XOR}u);
+        int cos_base = (int)(slid & {FIRST_HALF - 1}u) * N_READS;
+
+        // Load precomputed cos/sin (computed in Python from position * inv_freq)
+        float cos_arr[{N_READS}], sin_arr[{N_READS}];
+        for (int i = 0; i < N_READS; i++) {{
+            cos_arr[i] = rope_cos[cos_base + i];
+            sin_arr[i] = rope_sin[cos_base + i];
+        }}
+
+        float partner_vals[{N_READS}];
+        for (int i = 0; i < N_READS; i++)
+            partner_vals[i] = simd_shuffle(vals[i], partner);
+
+        if (slid < {FIRST_HALF}u) {{
+            for (int i = 0; i < N_READS; i++)
+                vals[i] = vals[i] * cos_arr[i] - partner_vals[i] * sin_arr[i];
+        }} else {{
+            for (int i = 0; i < N_READS; i++)
+                vals[i] = partner_vals[i] * sin_arr[i] + vals[i] * cos_arr[i];
+        }}
+    }}
+
+    // ── Phase 5: Write output ──
+    int out_base = is_q
+        ? (int)(b_idx * H_Q * D_DIM + (int)head_idx * D_DIM)
+        : (int)(b_idx * H_KV * D_DIM + head_local * D_DIM);
+
+    for (int i = 0; i < N_READS; i++) {{
+        if (is_q)
+            q_out[out_base + elem_base + i] = static_cast<bfloat16_t>(vals[i]);
+        else
+            k_out[out_base + elem_base + i] = static_cast<bfloat16_t>(vals[i]);
+    }}
+"""
+
+
+_kernel_cache = {}
+
+
+def _get_kernel(H_q, H_kv, D, rope_dims):
+    key = (H_q, H_kv, D, rope_dims)
+    if key not in _kernel_cache:
+        _kernel_cache[key] = mx.fast.metal_kernel(
+            name=f"fused_qk_norm_rope_Hq{H_q}_Hkv{H_kv}_D{D}_rd{rope_dims}",
+            input_names=["queries", "keys", "q_norm_w", "k_norm_w",
+                         "rope_cos", "rope_sin"],
+            output_names=["q_out", "k_out"],
+            source=_gen_fused_qk_norm_rope_source(H_q, H_kv, D, rope_dims),
+        )
+    return _kernel_cache[key]
+
+
+def fused_qk_norm_rope(queries, keys, q_norm_weight, k_norm_weight,
+                        inv_freq, cache_offset, H_q, H_kv, D,
+                        batch_size=1):
+    B = batch_size
+    rope_dims = inv_freq.shape[0] * 2
+
+    kern = _get_kernel(H_q, H_kv, D, rope_dims)
+
+    q_flat = queries.reshape(B, H_q * D)
+    k_flat = keys.reshape(B, H_kv * D)
+
+    # Precompute cos/sin in Python (avoids scalar kernel input)
+    angles = float(cache_offset) * inv_freq
+    rope_cos = mx.cos(angles).astype(mx.float32)
+    rope_sin = mx.sin(angles).astype(mx.float32)
+
+    n_heads = H_q + H_kv
+    results = kern(
+        inputs=[q_flat, k_flat, q_norm_weight, k_norm_weight, rope_cos, rope_sin],
+        output_shapes=[(B * H_q * D,), (B * H_kv * D,)],
+        output_dtypes=[mx.bfloat16, mx.bfloat16],
+        grid=(n_heads * 32, 1, B),
+        threadgroup=(32, 1, 1),
+    )
+
+    q_out = results[0].reshape(B, H_q, 1, D)
+    k_out = results[1].reshape(B, H_kv, 1, D)
+    return q_out, k_out

src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_qk_rmsnorm.py

@@ -0,0 +1,128 @@
+"""Fused Q/K per-head L2-norm for GDN attention (Dispatch 3).
+
+Performs per-head L2 normalization on q and k vectors with different scaling.
+Matches vLLM and latest mlx-lm (qwen3_5.py) which use rsqrt(sum(x²) + eps),
+NOT rms_norm which uses rsqrt(mean(x²) + eps).
+
+From qwen3_5.py (updated to match vLLM):
+    inv_scale = Dk^(-0.5) = 128^(-0.5)
+    q = inv_scale * q * rsqrt(sum(q²) + 1e-6)   → L2-normalize then scale by 1/√Dk
+    k = k * rsqrt(sum(k²) + 1e-6)                → L2-normalize only (no extra scale)
+
+Grid: (32 heads × 32 threads, 1, B).
+Each TG = 32 threads = 1 SG, handles one 128-dim head.
+Dk=128 = 32 threads × 4 elements → exactly 1 SG, no cross-SG reduction.
+"""
+
+import mlx.core as mx
+
+
+def _gen_fused_qk_rmsnorm_source():
+    """Generate Metal source for fused Q/K per-head L2-norm.
+
+    Input: qkv [B, 8192] bf16 (flattened from [B, 1, 8192])
+      - [0, 2048): q = 16 heads × 128
+      - [2048, 4096): k = 16 heads × 128
+      - [4096, 8192): v (untouched)
+
+    Output: qk_out [B, 4096] bf16
+      - [0, 2048): q L2-normalized then scaled by 1/√Dk
+      - [2048, 4096): k L2-normalized (no extra scale)
+
+    Grid: (32 * 32, 1, B), TG: (32, 1, 1)
+      tgid.x 0..15:  q heads → scale = 1/√128
+      tgid.x 16..31: k heads → scale = 1.0
+      tgid.z: batch index
+    """
+    return """
+    const int N_READS = 4;
+    const int DK = 128;
+    const int HK = 16;
+    const float EPS = 1e-6f;
+    const float Q_SCALE = rsqrt(128.0f);        // inv_scale = Dk^(-0.5)
+    const float K_SCALE = 1.0f;                  // no extra scale for k
+
+    uint head_idx = threadgroup_position_in_grid.x;
+    uint slid = thread_index_in_simdgroup;
+    uint b_idx = thread_position_in_grid.z;
+
+    bool is_q = (head_idx < (uint)HK);
+
+    // Input offset: q heads at [0, 2048), k heads at [2048, 4096)
+    int in_base = is_q
+        ? (b_idx * 8192 + head_idx * DK)
+        : (b_idx * 8192 + 2048 + (head_idx - HK) * DK);
+
+    // Output offset: q at [0, 2048), k at [2048, 4096)
+    int out_base = b_idx * 4096 + head_idx * DK;
+
+    // ── Phase 1: Load 4 elements + sum of squares ──
+    float vals[4];
+    float partial_sq = 0.0f;
+    int elem_base = slid * N_READS;
+
+    for (int i = 0; i < N_READS; i++) {
+        float xi = float(qkv[in_base + elem_base + i]);
+        vals[i] = xi;
+        partial_sq += xi * xi;
+    }
+
+    // ── Phase 2: simd reduction (32 threads → full sum of 128 elements) ──
+    float sum_sq = simd_sum(partial_sq);
+
+    // ── Phase 3: compute L2 inv-norm (NOT rms_norm — no /Dk) ──
+    float inv_rms = metal::precise::rsqrt(sum_sq + EPS);
+
+    // ── Phase 4: scale and write ──
+    float scale = is_q ? Q_SCALE : K_SCALE;
+    float combined = inv_rms * scale;
+
+    for (int i = 0; i < N_READS; i++) {
+        qk_out[out_base + elem_base + i] = static_cast<bfloat16_t>(vals[i] * combined);
+    }
+"""
+
+
+_fused_qk_rmsnorm_kernel = None
+
+
+def _get_fused_qk_rmsnorm_kernel():
+    """Get or compile the fused Q/K RMSNorm kernel."""
+    global _fused_qk_rmsnorm_kernel
+    if _fused_qk_rmsnorm_kernel is None:
+        _fused_qk_rmsnorm_kernel = mx.fast.metal_kernel(
+            name="fused_qk_rmsnorm",
+            input_names=["qkv"],
+            output_names=["qk_out"],
+            source=_gen_fused_qk_rmsnorm_source(),
+        )
+    return _fused_qk_rmsnorm_kernel
+
+
+def fused_qk_rmsnorm(qkv_conv_silu, batch_size=1):
+    """Fused Q/K per-head RMSNorm for GDN attention.
+
+    Args:
+        qkv_conv_silu: [B, 1, 8192] bf16 — post-conv, post-SiLU output from Dispatch 2.
+            First 2048 = q (16 heads × 128), next 2048 = k, last 4096 = v.
+        batch_size: int — batch dimension.
+
+    Returns:
+        qk_normed: [B, 1, 4096] bf16 — normalized q (first 2048) and k (next 2048).
+            v is NOT copied; Dispatch 4 reads v directly from qkv_conv_silu[:, :, 4096:].
+    """
+    B = batch_size
+    kern = _get_fused_qk_rmsnorm_kernel()
+
+    # Flatten to [B, 8192] for kernel
+    qkv_flat = qkv_conv_silu.reshape(B, 8192)
+
+    n_heads = 32  # 16 q + 16 k
+    results = kern(
+        inputs=[qkv_flat],
+        output_shapes=[(B * 4096,)],
+        output_dtypes=[mx.bfloat16],
+        grid=(n_heads * 32, 1, B),
+        threadgroup=(32, 1, 1),
+    )
+    return results[0].reshape(B, 1, 4096)