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,

.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,18 +155,18 @@ 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 = ...
@@ -183,8 +178,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 +190,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 +205,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 +219,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 +230,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 +243,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/pyproject.toml

@@ -4,13 +4,13 @@ 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",
 ]
 
 [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]

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/ml-explore/mlx-lm", rev = "834fac934c4e04de9b3d723e2b9287a2c60cfd4a" }
 # 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,18 @@ 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/**",
+]
 
 [tool.ruff.lint]
 extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
@@ -127,13 +133,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"]

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)

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,15 @@ 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,
+        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)
@@ -524,15 +525,15 @@ class API:
 
             if (
                 model_card.model_id,
-                instance.sharding(),
-                instance.instance_meta(),
+                sharding,
+                instance_meta,
                 len(placement_node_ids),
             ) not in seen:
                 previews.append(
                     PlacementPreview(
                         model_id=model_card.model_id,
-                        sharding=instance.sharding(),
-                        instance_meta=instance.instance_meta(),
+                        sharding=sharding,
+                        instance_meta=instance_meta,
                         instance=instance,
                         memory_delta_by_node=memory_delta_by_node or None,
                         error=None,
@@ -541,8 +542,8 @@ class API:
             seen.add(
                 (
                     model_card.model_id,
-                    instance.sharding(),
-                    instance.instance_meta(),
+                    sharding,
+                    instance_meta,
                     len(placement_node_ids),
                 )
             )
@@ -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/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,10 @@ 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
+    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 +266,11 @@ class DeleteInstanceResponse(BaseModel):
     instance_id: InstanceId
 
 
+class CancelCommandResponse(BaseModel):
+    message: str
+    command_id: CommandId
+
+
 ImageSize = Literal[
     "auto",
     "512x512",
@@ -437,3 +446,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/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,9 @@ 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
+    repetition_penalty: float | None = None
+    repetition_context_size: int | None = None

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

@@ -4,13 +4,7 @@ from pydantic import model_validator
 
 from exo.shared.models.model_cards import ModelTask
 from exo.shared.types.common import Host, Id, NodeId
-from exo.shared.types.worker.runners import RunnerId, ShardAssignments
-from exo.shared.types.worker.shards import (
-    PipelineShardMetadata,
-    Sharding,
-    ShardMetadata,
-    TensorShardMetadata,
-)
+from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
 
@@ -30,40 +24,16 @@ class BaseInstance(TaggedModel):
     def shard(self, runner_id: RunnerId) -> ShardMetadata | None:
         return self.shard_assignments.runner_to_shard.get(runner_id, None)
 
-    @staticmethod
-    def instance_meta() -> InstanceMeta: ...
-
-    def sharding(self) -> Sharding:
-        if all(
-            isinstance(sm, PipelineShardMetadata)
-            for sm in self.shard_assignments.runner_to_shard.values()
-        ):
-            return Sharding.Pipeline
-        if all(
-            isinstance(sm, TensorShardMetadata)
-            for sm in self.shard_assignments.runner_to_shard.values()
-        ):
-            return Sharding.Tensor
-        raise ValueError("shard metadata malformed")
-
 
 class MlxRingInstance(BaseInstance):
     hosts_by_node: dict[NodeId, list[Host]]
     ephemeral_port: int
 
-    @staticmethod
-    def instance_meta() -> InstanceMeta:
-        return InstanceMeta.MlxRing
-
 
 class MlxJacclInstance(BaseInstance):
     jaccl_devices: list[list[str | None]]
     jaccl_coordinators: dict[NodeId, str]
 
-    @staticmethod
-    def instance_meta() -> InstanceMeta:
-        return InstanceMeta.MlxJaccl
-
 
 # TODO: Single node instance
 Instance = MlxRingInstance | MlxJacclInstance

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/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
@@ -275,6 +436,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,11 +470,16 @@ 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,
@@ -336,6 +504,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 +650,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)

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

@@ -0,0 +1,117 @@
+"""Fused RMSNormGated for GDN attention (Dispatch 5).
+
+Fuses RMSNorm(out, weight) × z_silu into one kernel.
+SiLU on z was already applied in Dispatch 2, so z_silu arrives as f32.
+
+From qwen3_next.py (Qwen3NextRMSNormGated):
+    x = rms_norm(hidden_states, weight, eps)     # weight: [Dv=128]
+    gate = silu(z.float())                        # already done in Dispatch 2
+    return (gate * x).to(hidden_states.dtype)
+
+Grid: (32 heads × 32 threads, 1, B).
+Each TG = 32 threads = 1 SG, handles one 128-dim head.
+Dv=128 = 32 threads × 4 elements → exactly 1 SG.
+"""
+
+import mlx.core as mx
+
+
+def _gen_fused_rms_norm_gated_source():
+    """Generate Metal source for fused RMSNormGated.
+
+    Inputs:
+      gdn_out: [B, Hv*Dv] bf16  — GDN output, flattened (Hv=32, Dv=128)
+      z_silu:  [B, Hv*Dv] f32   — post-SiLU z from Dispatch 2
+      weight:  [Dv] f32         — RMSNormGated learned weight (128 elements)
+
+    Output:
+      out: [B, Hv*Dv] bf16      — result = z_silu * rms_norm(gdn_out, weight)
+
+    Grid: (32 * 32, 1, B), TG: (32, 1, 1)
+      tgid.x: head index (0..31)
+      tgid.z: batch index
+    """
+    return """
+    const int N_READS = 4;
+    const int DV = 128;
+    const int HV = 32;
+    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;
+
+    int base = b_idx * HV * DV + head_idx * DV;
+    int elem_base = slid * N_READS;
+
+    // ── Phase 1: Load gdn_out elements + sum of squares ──
+    float gdn_vals[4];
+    float partial_sq = 0.0f;
+
+    for (int i = 0; i < N_READS; i++) {
+        float xi = float(gdn_out[base + elem_base + i]);
+        gdn_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 inv_rms ──
+    float inv_rms = metal::precise::rsqrt(sum_sq / float(DV) + EPS);
+
+    // ── Phase 4: RMSNorm × z_silu, write bf16 ──
+    for (int i = 0; i < N_READS; i++) {
+        int idx = elem_base + i;
+        float w = float(weight[idx]);                           // learned weight[Dv]
+        float normed = gdn_vals[i] * inv_rms * w;              // RMSNorm
+        float z_val = z_silu[base + idx];                       // already f32, post-SiLU
+        out[base + idx] = static_cast<bfloat16_t>(z_val * normed);
+    }
+"""
+
+
+_fused_rms_norm_gated_kernel = None
+
+
+def _get_fused_rms_norm_gated_kernel():
+    """Get or compile the fused RMSNormGated kernel."""
+    global _fused_rms_norm_gated_kernel
+    if _fused_rms_norm_gated_kernel is None:
+        _fused_rms_norm_gated_kernel = mx.fast.metal_kernel(
+            name="fused_rms_norm_gated",
+            input_names=["gdn_out", "z_silu", "weight"],
+            output_names=["out"],
+            source=_gen_fused_rms_norm_gated_source(),
+        )
+    return _fused_rms_norm_gated_kernel
+
+
+def fused_rms_norm_gated(gdn_out, z_silu, weight, batch_size=1):
+    """Fused RMSNormGated: RMSNorm(out, weight) × z_silu.
+
+    Args:
+        gdn_out: [B, 1, Hv, Dv] bf16 — GDN recurrence output (Hv=32, Dv=128).
+        z_silu:  [B, 1, 4096] f32 — post-SiLU z from Dispatch 2.
+        weight:  [128] f32 — RMSNormGated learned weight (Dv elements).
+        batch_size: int.
+
+    Returns:
+        out: [B, 1, 4096] bf16 — ready for out_proj in Dispatch 6.
+    """
+    B = batch_size
+    kern = _get_fused_rms_norm_gated_kernel()
+
+    # Flatten to [B, 4096]
+    gdn_flat = gdn_out.reshape(B, 4096)
+    z_flat = z_silu.reshape(B, 4096)
+
+    n_heads = 32  # Hv
+    results = kern(
+        inputs=[gdn_flat, z_flat, weight],
+        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)

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

@@ -0,0 +1,177 @@
+"""GDN recurrence with pre-computed g and beta (Dispatch 4).
+
+Modified version of gated_delta_step from mlx-lm-fork/mlx_lm/models/gated_delta.py.
+Instead of computing g = exp(-exp(A_log) * softplus(a + dt_bias)) and beta = sigmoid(b)
+inside the kernel, accepts them as pre-computed f32 inputs from Dispatch 2.
+
+Non-vectorized only (Qwen3.5-35B-A3B uses scalar gating per head).
+
+Grid: (32, Dv, B*Hv) = (32, 128, B*32), TG: (32, 4, 1)
+"""
+
+from typing import Optional, Tuple
+
+import mlx.core as mx
+
+
+def _make_gdn_precomputed_kernel(has_mask=False):
+    """Build the GDN kernel with pre-computed g and beta."""
+    if not mx.metal.is_available():
+        return None
+
+    mask_source = "mask[b_idx * T + t]" if has_mask else "true"
+
+    source = f"""
+        auto n = thread_position_in_grid.z;
+        auto b_idx = n / Hv;
+        auto hv_idx = n % Hv;
+        auto hk_idx = hv_idx / (Hv / Hk);
+        constexpr int n_per_t = Dk / 32;
+
+        // q, k: [B, T, Hk, Dk]
+        auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
+        auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
+
+        // v, y: [B, T, Hv, Dv]
+        auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
+        y += b_idx * T * Hv * Dv + hv_idx * Dv;
+
+        auto dk_idx = thread_position_in_threadgroup.x;
+        auto dv_idx = thread_position_in_grid.y;
+
+        // state_in, state_out: [B, Hv, Dv, Dk]
+        auto i_state = state_in + (n * Dv + dv_idx) * Dk;
+        auto o_state = state_out + (n * Dv + dv_idx) * Dk;
+
+        float state[n_per_t];
+        for (int i = 0; i < n_per_t; ++i) {{
+          auto s_idx = n_per_t * dk_idx + i;
+          state[i] = static_cast<float>(i_state[s_idx]);
+        }}
+
+        // g: [B, T, Hv] f32 — pre-computed decay gate
+        auto g_ = g + b_idx * T * Hv;
+        // beta: [B, T, Hv] f32 — pre-computed sigmoid(b)
+        auto beta_ = beta + b_idx * T * Hv;
+
+        for (int t = 0; t < T; ++t) {{
+          if ({mask_source}) {{
+            // Pre-computed g and beta (no softplus/exp/sigmoid needed)
+            float g_val = g_[hv_idx];
+            float beta_val = beta_[hv_idx];
+
+            float kv_mem = 0.0f;
+            for (int i = 0; i < n_per_t; ++i) {{
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = state[i] * g_val;
+              kv_mem += state[i] * k_[s_idx];
+            }}
+            kv_mem = simd_sum(kv_mem);
+
+            auto delta = (v_[dv_idx] - kv_mem) * beta_val;
+
+            float out = 0.0f;
+            for (int i = 0; i < n_per_t; ++i) {{
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = state[i] + k_[s_idx] * delta;
+              out += state[i] * q_[s_idx];
+            }}
+            out = simd_sum(out);
+            if (thread_index_in_simdgroup == 0) {{
+              y[dv_idx] = static_cast<InT>(out);
+            }}
+          }}
+          // Increment data pointers to next time step
+          q_ += Hk * Dk;
+          k_ += Hk * Dk;
+          v_ += Hv * Dv;
+          y += Hv * Dv;
+          g_ += Hv;
+          beta_ += Hv;
+        }}
+        for (int i = 0; i < n_per_t; ++i) {{
+          auto s_idx = n_per_t * dk_idx + i;
+          o_state[s_idx] = static_cast<InT>(state[i]);
+        }}
+    """
+
+    inputs = ["q", "k", "v", "g", "beta", "state_in", "T"]
+    if has_mask:
+        inputs.append("mask")
+
+    suffix = "_precomputed"
+    if has_mask:
+        suffix += "_mask"
+
+    return mx.fast.metal_kernel(
+        name=f"gated_delta_step{suffix}",
+        input_names=inputs,
+        output_names=["y", "state_out"],
+        source=source,
+    )
+
+
+_gdn_precomputed_kernel = None
+_gdn_precomputed_kernel_masked = None
+
+
+def _get_gdn_precomputed_kernel(has_mask=False):
+    """Get or compile the pre-computed GDN kernel."""
+    global _gdn_precomputed_kernel, _gdn_precomputed_kernel_masked
+    if has_mask:
+        if _gdn_precomputed_kernel_masked is None:
+            _gdn_precomputed_kernel_masked = _make_gdn_precomputed_kernel(has_mask=True)
+        return _gdn_precomputed_kernel_masked
+    else:
+        if _gdn_precomputed_kernel is None:
+            _gdn_precomputed_kernel = _make_gdn_precomputed_kernel(has_mask=False)
+        return _gdn_precomputed_kernel
+
+
+def gated_delta_update_precomputed(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: mx.array,
+    mask: Optional[mx.array] = None,
+) -> Tuple[mx.array, mx.array]:
+    """GDN recurrence with pre-computed g and beta.
+
+    Args:
+        q: [B, T, Hk, Dk] bf16 — normalized q from Dispatch 3
+        k: [B, T, Hk, Dk] bf16 — normalized k from Dispatch 3
+        v: [B, T, Hv, Dv] bf16 — v from Dispatch 2 (qkv_conv_silu[:, :, 4096:])
+        g: [B, T, Hv] f32 — pre-computed decay gate from Dispatch 2
+        beta: [B, T, Hv] f32 — pre-computed sigmoid(b) from Dispatch 2
+        state: [B, Hv, Dv, Dk] bf16 — recurrent state from cache
+        mask: [B, T] optional
+
+    Returns:
+        y: [B, T, Hv, Dv] bf16
+        new_state: [B, Hv, Dv, Dk] bf16
+    """
+    B, T, Hk, Dk = k.shape
+    Hv, Dv = v.shape[2:]
+    input_type = q.dtype
+
+    kernel = _get_gdn_precomputed_kernel(has_mask=mask is not None)
+    inputs = [q, k, v, g, beta, state, T]
+    if mask is not None:
+        inputs.append(mask)
+
+    return kernel(
+        inputs=inputs,
+        template=[
+            ("InT", input_type),
+            ("Dk", Dk),
+            ("Dv", Dv),
+            ("Hk", Hk),
+            ("Hv", Hv),
+        ],
+        grid=(32, Dv, B * Hv),
+        threadgroup=(32, 4, 1),
+        output_shapes=[(B, T, Hv, Dv), state.shape],
+        output_dtypes=[input_type, input_type],
+    )

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

@@ -0,0 +1,218 @@
+"""Merged 8-bit down_proj GEMV for Qwen3.5 routed + shared experts.
+
+Port of the 4-bit down_proj kernel, adapted for 8-bit quantization (gs=64).
+Maps intermediate → hidden: (n_active, N_IN) → (n_active, K_OUT).
+
+Two paths via tgid.z:
+- 8-bit routed experts (z < n_active): uint8 dequant with expert index lookup, f32 output
+- 8-bit shared expert (z == n_active): uint8 dequant (no expert index), f32 output
+"""
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_merged_down_8bit_source(group_size=64, scale_bf16=True):
+    """Metal source for merged 8-bit down_proj GEMV.
+
+    Same structure as gate GEMV: 2 SGs × 4 rows/SG = 8 output rows per TG.
+    8-bit dequant: result = scale * Σ(x[i]*w[i]) + bias * Σ(x[i])
+    Both routed and shared paths use 8-bit dequantization.
+    """
+    gs = int(group_size)
+    sc_stride = 256 // gs  # groups per K-block (256/64 = 4)
+    slid_divisor = gs // 8  # threads per group (64/8 = 8)
+    sc_t = "bfloat16_t" if scale_bf16 else "float"
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int VALUES_PER_THREAD = 8;
+    const int BLOCK_SIZE = 256;
+
+    int K_OUT = K_OUT_val;
+    int N_IN = N_IN_val;
+    int SHARED_N_IN = SHARED_N_IN_val;
+    int n_active = n_active_val;
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;  // 0 or 1
+    uint slid = thread_index_in_simdgroup;       // 0..31
+
+    int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG;
+    if (out_row >= K_OUT) return;
+
+    if (tgid.z < (uint)n_active) {{
+        // ═══════ 8-BIT ROUTED EXPERT PATH ═══════
+        int N_groups = N_IN / {gs};
+
+        int expert = inds[tgid.z];
+
+        const device uint8_t* ws = (const device uint8_t*)W
+            + (long)expert * K_OUT * N_IN + out_row * N_IN + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc = (const device {sc_t}*)S
+            + (long)expert * K_OUT * N_groups + out_row * N_groups + slid / {slid_divisor};
+        const device {sc_t}* bi = (const device {sc_t}*)B_q
+            + (long)expert * K_OUT * N_groups + out_row * N_groups + slid / {slid_divisor};
+
+        const device float* x_ptr = (const device float*)X_routed
+            + tgid.z * N_IN;
+
+        int x_base = slid * VALUES_PER_THREAD;
+        float result[4] = {{0, 0, 0, 0}};
+
+        for (int k = 0; k < N_IN; k += BLOCK_SIZE) {{
+            float x_thread[8];
+            float xsum = 0;
+            for (int i = 0; i < 8; i++) {{
+                float xi = x_ptr[x_base + i];
+                x_thread[i] = xi;
+                xsum += xi;
+            }}
+
+            for (int row = 0; row < RESULTS_PER_SG; row++) {{
+                const device uint8_t* wl = ws + row * N_IN;
+                float s = float(sc[row * N_groups]);
+                float b = float(bi[row * N_groups]);
+                float accum = 0;
+                for (int i = 0; i < 8; i++) {{
+                    accum += x_thread[i] * float(wl[i]);
+                }}
+                result[row] += s * accum + xsum * b;
+            }}
+
+            ws += BLOCK_SIZE;
+            sc += {sc_stride};
+            bi += {sc_stride};
+            x_base += BLOCK_SIZE;
+        }}
+
+        device float* yp = Y_routed + tgid.z * K_OUT + out_row;
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            float r = simd_sum(result[row]);
+            if (slid == 0) {{
+                yp[row] = r;
+            }}
+        }}
+
+    }} else {{
+        // ═══════ 8-BIT SHARED EXPERT PATH ═══════
+        // Same dequant as routed path, but no expert index lookup.
+        // W_shared_down is (K_OUT, SHARED_N_IN/4) uint32 → (K_OUT, SHARED_N_IN) uint8
+        int N_groups = SHARED_N_IN / {gs};
+
+        const device uint8_t* ws = (const device uint8_t*)W_shared_down
+            + (long)out_row * SHARED_N_IN + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc = (const device {sc_t}*)S_shared_down
+            + (long)out_row * N_groups + slid / {slid_divisor};
+        const device {sc_t}* bi = (const device {sc_t}*)B_shared_down
+            + (long)out_row * N_groups + slid / {slid_divisor};
+
+        // X_shared is float32 (output of Kernel 1 shared path)
+        const device float* x_ptr = (const device float*)X_shared;
+
+        int x_base = slid * VALUES_PER_THREAD;
+        float result[4] = {{0, 0, 0, 0}};
+
+        for (int k = 0; k < SHARED_N_IN; k += BLOCK_SIZE) {{
+            float x_thread[8];
+            float xsum = 0;
+            for (int i = 0; i < 8; i++) {{
+                float xi = x_ptr[x_base + i];
+                x_thread[i] = xi;
+                xsum += xi;
+            }}
+
+            for (int row = 0; row < RESULTS_PER_SG; row++) {{
+                const device uint8_t* wl = ws + row * SHARED_N_IN;
+                float s = float(sc[row * N_groups]);
+                float b = float(bi[row * N_groups]);
+                float accum = 0;
+                for (int i = 0; i < 8; i++) {{
+                    accum += x_thread[i] * float(wl[i]);
+                }}
+                result[row] += s * accum + xsum * b;
+            }}
+
+            ws += BLOCK_SIZE;
+            sc += {sc_stride};
+            bi += {sc_stride};
+            x_base += BLOCK_SIZE;
+        }}
+
+        device float* yp = Y_shared + out_row;
+        for (int tm = 0; tm < RESULTS_PER_SG; tm++) {{
+            float r = simd_sum(result[tm]);
+            if (slid == 0) {{
+                yp[tm] = r;
+            }}
+        }}
+    }}
+"""
+
+
+_merged_down_8bit_kernels = {}
+
+
+def _get_merged_down_8bit_kernel(group_size=64, scale_bf16=True):
+    key = (group_size, scale_bf16)
+    if key not in _merged_down_8bit_kernels:
+        sc_tag = "_bf16sc" if scale_bf16 else ""
+        _merged_down_8bit_kernels[key] = mx.fast.metal_kernel(
+            name=f"merged_down_proj_8bit_gs{group_size}{sc_tag}",
+            input_names=["W", "S", "B_q",
+                         "W_shared_down", "S_shared_down", "B_shared_down",
+                         "X_routed", "X_shared", "inds",
+                         "K_OUT_val", "N_IN_val", "SHARED_N_IN_val", "n_active_val"],
+            output_names=["Y_routed", "Y_shared"],
+            source=_gen_merged_down_8bit_source(group_size, scale_bf16),
+        )
+    return _merged_down_8bit_kernels[key]
+
+
+def fused_merged_down_proj_8bit(w_q, s, b_q,
+                                 w_shared_down, s_shared_down, b_shared_down,
+                                 x_routed, x_shared, inds,
+                                 k_out, n_in, group_size=64,
+                                 shared_n_in=None):
+    """Single-dispatch merged down_proj for 8-bit routed + 8-bit shared experts.
+
+    Args:
+        w_q: routed quantized down weights (E, K_OUT, N_IN/4) uint32
+        s: routed scales (E, K_OUT, N_IN/gs) bfloat16
+        b_q: routed biases (E, K_OUT, N_IN/gs) bfloat16
+        w_shared_down: shared expert down weight (K_OUT, SHARED_N_IN/4) uint32
+        s_shared_down: shared expert down scales (K_OUT, SHARED_N_IN/gs) bfloat16
+        b_shared_down: shared expert down biases (K_OUT, SHARED_N_IN/gs) bfloat16
+        x_routed: routed SwiGLU output (n_active, N_IN) float32
+        x_shared: shared SwiGLU output (SHARED_N_IN,) float32
+        inds: selected expert indices (n_active,) uint32
+        k_out: output dimension (4096 for Qwen3.5)
+        n_in: routed expert input dimension (1024 for Qwen3.5)
+        group_size: quantization group size (64)
+        shared_n_in: shared expert input dimension (defaults to n_in)
+
+    Returns:
+        (Y_routed, Y_shared):
+            Y_routed: (n_active, k_out) float32
+            Y_shared: (k_out,) float32
+    """
+    scale_bf16 = (s.dtype == mx.bfloat16)
+    kern = _get_merged_down_8bit_kernel(group_size, scale_bf16)
+    n_active = inds.shape[0]
+    k_out_val = int(k_out)
+    n_in_val = int(n_in)
+    shared_n_in_val = int(shared_n_in) if shared_n_in is not None else n_in_val
+    y_groups = ceil_div(k_out_val, 8)
+    Y = kern(
+        inputs=[w_q, s, b_q,
+                w_shared_down, s_shared_down, b_shared_down,
+                x_routed, x_shared, inds,
+                k_out_val, n_in_val, shared_n_in_val, n_active],
+        output_shapes=[(n_active, k_out_val), (k_out_val,)],
+        output_dtypes=[mx.float32, mx.float32],
+        grid=(32, y_groups * 2, n_active + 1),
+        threadgroup=(32, 2, 1),
+    )
+    return Y[0], Y[1]

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

@@ -0,0 +1,274 @@
+"""Merged 8-bit fused gate+up+SwiGLU for Qwen3.5 routed + shared experts.
+
+Port of the 4-bit kernel for Kimi K2.5, adapted for:
+- 8-bit quantization (direct uint8 byte reads, no nibble extraction)
+- group_size=64 (one scale+bias per 64 elements)
+- Qwen3.5 dimensions (K=4096, N_INTER=1024, E=512, top_k=10)
+- No score normalization in kernel (handled by gate dispatch)
+
+Two paths via tgid.z:
+- 8-bit routed experts (z < n_active): uint8 dequant with expert index lookup, f32 output
+- 8-bit shared expert (z == n_active): uint8 dequant (no expert index), f32 output
+"""
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_merged_8bit_source(group_size=64, scale_bf16=True):
+    """Generate Metal source for merged 8-bit fused gate+up+SwiGLU.
+
+    Both routed and shared expert paths use 8-bit dequant:
+        result = scale * Σ(x[i]*w[i]) + bias * Σ(x[i])
+
+    Each thread processes VALUES_PER_THREAD=8 elements of K per iteration.
+    BLOCK_SIZE = 32 * 8 = 256 elements per K-block.
+    """
+    gs = int(group_size)
+    sc_stride = 256 // gs  # groups consumed per K-block (256/64 = 4)
+    slid_divisor = gs // 8  # threads per group (64/8 = 8)
+    sc_t = "bfloat16_t" if scale_bf16 else "float"
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+
+    int N_INTER = N_INTER_val;
+    int SHARED_INTER = SHARED_INTER_val;
+    int K = K_val;
+    int n_active = n_active_val;
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;  // 0 or 1
+    uint slid = thread_index_in_simdgroup;       // 0..31
+
+    int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG;
+
+    // Routed path uses N_INTER rows, shared path uses SHARED_INTER rows
+    int row_limit = (tgid.z < (uint)n_active) ? N_INTER : SHARED_INTER;
+    if (out_row >= row_limit) return;
+
+    float gate_result[4] = {{0, 0, 0, 0}};
+    float up_result[4] = {{0, 0, 0, 0}};
+
+    if (tgid.z < (uint)n_active) {{
+        // ═══════ 8-BIT ROUTED EXPERT PATH ═══════
+        const int VALUES_PER_THREAD = 8;
+        const int BLOCK_SIZE = 256;  // 32 * 8
+
+        int N_TOTAL = 2 * N_INTER;  // gate + up stacked
+        int K_groups = K / {gs};
+
+        int expert = inds[tgid.z];
+
+        // W is stored as uint32 (4 bytes per uint32), cast to uint8_t
+        // Layout: (E, N_TOTAL, K/4) uint32 → (E, N_TOTAL, K) uint8
+        const device uint8_t* ws_gate = (const device uint8_t*)W
+            + (long)expert * N_TOTAL * K + out_row * K + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_gate = (const device {sc_t}*)S
+            + (long)expert * N_TOTAL * K_groups + out_row * K_groups + slid / {slid_divisor};
+        const device {sc_t}* bi_gate = (const device {sc_t}*)B_q
+            + (long)expert * N_TOTAL * K_groups + out_row * K_groups + slid / {slid_divisor};
+
+        const device uint8_t* ws_up = (const device uint8_t*)W
+            + (long)expert * N_TOTAL * K + (out_row + N_INTER) * K + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_up = (const device {sc_t}*)S
+            + (long)expert * N_TOTAL * K_groups + (out_row + N_INTER) * K_groups + slid / {slid_divisor};
+        const device {sc_t}* bi_up = (const device {sc_t}*)B_q
+            + (long)expert * N_TOTAL * K_groups + (out_row + N_INTER) * K_groups + slid / {slid_divisor};
+
+        int x_base = slid * VALUES_PER_THREAD;
+
+        for (int k = 0; k < K; k += BLOCK_SIZE) {{
+            // Load 8 x values and compute sum
+            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++) {{
+                // Gate projection
+                const device uint8_t* wg = ws_gate + row * K;
+                float sg = float(sc_gate[row * K_groups]);
+                float bg = float(bi_gate[row * K_groups]);
+                float accum_g = 0;
+                for (int i = 0; i < 8; i++) {{
+                    accum_g += x_thread[i] * float(wg[i]);
+                }}
+                gate_result[row] += sg * accum_g + xsum * bg;
+
+                // Up projection
+                const device uint8_t* wu = ws_up + row * K;
+                float su = float(sc_up[row * K_groups]);
+                float bu = float(bi_up[row * K_groups]);
+                float accum_u = 0;
+                for (int i = 0; i < 8; i++) {{
+                    accum_u += x_thread[i] * float(wu[i]);
+                }}
+                up_result[row] += su * accum_u + xsum * bu;
+            }}
+
+            ws_gate += BLOCK_SIZE;
+            ws_up += BLOCK_SIZE;
+            sc_gate += {sc_stride};
+            sc_up += {sc_stride};
+            bi_gate += {sc_stride};
+            bi_up += {sc_stride};
+            x_base += BLOCK_SIZE;
+        }}
+
+        // Epilogue: SwiGLU + write f32 to Y_routed
+        device float* yp = Y_routed + tgid.z * N_INTER + out_row;
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            float g = simd_sum(gate_result[row]);
+            float u = simd_sum(up_result[row]);
+            if (slid == 0) {{
+                float silu_g = g / (1.0f + metal::exp(-g));
+                yp[row] = silu_g * u;
+            }}
+        }}
+
+    }} else {{
+        // ═══════ 8-BIT SHARED EXPERT PATH ═══════
+        // Same dequant as routed path, but no expert index lookup.
+        // W_shared is (2*SHARED_INTER, K/4) uint32 → (2*SHARED_INTER, K) uint8
+        const int VALUES_PER_THREAD = 8;
+        const int BLOCK_SIZE = 256;  // 32 * 8
+
+        int K_groups = K / {gs};
+
+        const device uint8_t* ws_gate = (const device uint8_t*)W_shared
+            + (long)out_row * K + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_gate = (const device {sc_t}*)S_shared
+            + (long)out_row * K_groups + slid / {slid_divisor};
+        const device {sc_t}* bi_gate = (const device {sc_t}*)B_shared
+            + (long)out_row * K_groups + slid / {slid_divisor};
+
+        const device uint8_t* ws_up = (const device uint8_t*)W_shared
+            + (long)(out_row + SHARED_INTER) * K + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_up = (const device {sc_t}*)S_shared
+            + (long)(out_row + SHARED_INTER) * K_groups + slid / {slid_divisor};
+        const device {sc_t}* bi_up = (const device {sc_t}*)B_shared
+            + (long)(out_row + SHARED_INTER) * K_groups + slid / {slid_divisor};
+
+        int x_base = slid * VALUES_PER_THREAD;
+
+        for (int k = 0; k < K; k += 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++) {{
+                // Gate projection
+                const device uint8_t* wg = ws_gate + row * K;
+                float sg = float(sc_gate[row * K_groups]);
+                float bg = float(bi_gate[row * K_groups]);
+                float accum_g = 0;
+                for (int i = 0; i < 8; i++) {{
+                    accum_g += x_thread[i] * float(wg[i]);
+                }}
+                gate_result[row] += sg * accum_g + xsum * bg;
+
+                // Up projection
+                const device uint8_t* wu = ws_up + row * K;
+                float su = float(sc_up[row * K_groups]);
+                float bu = float(bi_up[row * K_groups]);
+                float accum_u = 0;
+                for (int i = 0; i < 8; i++) {{
+                    accum_u += x_thread[i] * float(wu[i]);
+                }}
+                up_result[row] += su * accum_u + xsum * bu;
+            }}
+
+            ws_gate += BLOCK_SIZE;
+            ws_up += BLOCK_SIZE;
+            sc_gate += {sc_stride};
+            sc_up += {sc_stride};
+            bi_gate += {sc_stride};
+            bi_up += {sc_stride};
+            x_base += BLOCK_SIZE;
+        }}
+
+        // Epilogue: SwiGLU + write f32 to Y_shared
+        device float* yp = Y_shared + out_row;
+        for (int tm = 0; tm < RESULTS_PER_SG; tm++) {{
+            float g = simd_sum(gate_result[tm]);
+            float u = simd_sum(up_result[tm]);
+            if (slid == 0) {{
+                float silu_g = g / (1.0f + metal::exp(-g));
+                yp[tm] = silu_g * u;
+            }}
+        }}
+    }}
+"""
+
+
+_merged_8bit_kernels = {}
+
+
+def _get_merged_8bit_kernel(group_size=64, scale_bf16=True):
+    key = (group_size, scale_bf16)
+    if key not in _merged_8bit_kernels:
+        sc_tag = "_bf16sc" if scale_bf16 else ""
+        _merged_8bit_kernels[key] = mx.fast.metal_kernel(
+            name=f"merged_routed_shared_swiglu_8bit_gs{group_size}{sc_tag}",
+            input_names=["W", "S", "B_q", "W_shared", "S_shared", "B_shared",
+                         "X", "inds",
+                         "N_INTER_val", "SHARED_INTER_val", "K_val", "n_active_val"],
+            output_names=["Y_routed", "Y_shared"],
+            source=_gen_merged_8bit_source(group_size, scale_bf16),
+        )
+    return _merged_8bit_kernels[key]
+
+
+def fused_merged_gate_up_swiglu_8bit(w_q, s, b_q,
+                                      w_shared, s_shared, b_shared,
+                                      x, inds,
+                                      n_inter, k_hidden, group_size=64,
+                                      shared_inter=None):
+    """Single-dispatch merged gate+up+SwiGLU for 8-bit routed + 8-bit shared experts.
+
+    Args:
+        w_q: stacked routed quantized weights (E, 2*N_INTER, K/4) uint32
+        s: routed scales (E, 2*N_INTER, K/gs) bfloat16
+        b_q: routed biases (E, 2*N_INTER, K/gs) bfloat16
+        w_shared: shared expert gate+up stacked (2*SHARED_INTER, K/4) uint32
+        s_shared: shared expert scales (2*SHARED_INTER, K/gs) bfloat16
+        b_shared: shared expert biases (2*SHARED_INTER, K/gs) bfloat16
+        x: input vector (K,) bfloat16
+        inds: selected expert indices (n_active,) uint32
+        n_inter: routed expert intermediate size (1024 for Qwen3.5)
+        k_hidden: hidden size (4096 for Qwen3.5)
+        group_size: quantization group size (64)
+        shared_inter: shared expert intermediate size (defaults to n_inter)
+
+    Returns:
+        (Y_routed, Y_shared):
+            Y_routed: (n_active, n_inter) float32
+            Y_shared: (shared_inter,) float32
+    """
+    scale_bf16 = (s.dtype == mx.bfloat16)
+    kern = _get_merged_8bit_kernel(group_size, scale_bf16)
+    n_active = inds.shape[0]
+    n_inter_val = int(n_inter)
+    shared_inter_val = int(shared_inter) if shared_inter is not None else n_inter_val
+    # Grid y must cover max(n_inter, shared_inter)
+    max_inter = max(n_inter_val, shared_inter_val)
+    Y = kern(
+        inputs=[w_q, s, b_q, w_shared, s_shared, b_shared,
+                x, inds,
+                n_inter_val, shared_inter_val, int(k_hidden), n_active],
+        output_shapes=[(n_active, n_inter_val), (shared_inter_val,)],
+        output_dtypes=[mx.float32, mx.float32],
+        grid=(32, ceil_div(max_inter, 8) * 2, n_active + 1),
+        threadgroup=(32, 2, 1),
+    )
+    return Y[0], Y[1]

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

@@ -0,0 +1,493 @@
+"""Dispatch 2: SwiGLU with softmax prologue for Qwen3.5 oproj fusion.
+
+Port of Kimi's oproj_topk_v2_swiglu.py adapted for:
+- 8-bit quantized weights with gs=64 (Kimi uses 4-bit gs=32)
+- Softmax routing (Kimi uses sigmoid)
+- E up to 512 with multiple scores per thread (SPT = ceil(E/64))
+- Shared expert gate GEMV hidden in TG(0,0,0) SG 0
+- Both routed and shared paths use 8-bit quantized weights
+
+Prologue (all TGs):
+  Phase 1: distributed x² sum → inv_rms
+  Phase 2 (routed TGs only): gate scores → softmax → parallel top-k → norm_topk_prob
+  Phase 3 (TG(0,0,0) SG 0): shared_expert_gate 8-bit GEMV → gate_raw
+
+Main SwiGLU (after prologue):
+  Routed (z < n_active): 8-bit gate+up+SwiGLU with h_scaled input (inv_rms factored out)
+  Shared (z == n_active): 8-bit SwiGLU for shared expert
+"""
+import mlx.core as mx
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+def _gen_oproj_softmax_topk_swiglu_8bit_source(group_size=64, scale_bf16=True,
+                                                 n_experts=64, top_k=10,
+                                                 norm_topk=True):
+    """Generate Metal source for softmax + top-k + SwiGLU with oproj prologue."""
+    gs = int(group_size)
+    sc_stride = 256 // gs       # groups per K-block (256/64 = 4)
+    slid_divisor = gs // 8      # threads per group (64/8 = 8)
+    sc_t = "bfloat16_t" if scale_bf16 else "float"
+    E = int(n_experts)
+    K_TOP = int(top_k)
+    SPT = (E + 63) // 64       # scores per thread
+
+    # Score normalization block (TG(0,0,0) thread 0 only)
+    if norm_topk:
+        score_norm_block = f"""
+    // ── norm_topk_prob + write indices (TG(0,0,0) thread 0) ──
+    if (tgid.y == 0 && tgid.z == 0 && tid == 0) {{
+        float total = 0.0f;
+        for (int a = 0; a < {K_TOP}; a++) total += tg_selected_scores[a];
+        float inv_total = 1.0f / total;
+        for (int a = 0; a < {K_TOP}; a++) {{
+            norm_scores[a] = tg_selected_scores[a] * inv_total;
+            out_inds[a] = (uint)tg_inds[a];
+        }}
+    }}"""
+    else:
+        score_norm_block = f"""
+    // ── Write indices and raw scores (TG(0,0,0) thread 0) ──
+    if (tgid.y == 0 && tgid.z == 0 && tid == 0) {{
+        for (int a = 0; a < {K_TOP}; a++) {{
+            norm_scores[a] = tg_selected_scores[a];
+            out_inds[a] = (uint)tg_inds[a];
+        }}
+    }}"""
+
+    return f"""
+    const int RESULTS_PER_SG = 4;
+    const int E_CONST = {E};
+    const int K_TOP_CONST = {K_TOP};
+    const int SPT = {SPT};
+
+    int N_INTER = N_INTER_val;
+    int SHARED_INTER = SHARED_INTER_val;
+    int K = K_val;
+    int n_active = n_active_val;
+    int N_OPROJ_TG = N_OPROJ_TG_val;
+
+    uint3 tgid = threadgroup_position_in_grid;
+    uint sgid = simdgroup_index_in_threadgroup;  // 0 or 1
+    uint slid = thread_index_in_simdgroup;       // 0..31
+    int tid = int(sgid) * 32 + int(slid);        // 0..63
+
+    // ═══════════════════════════════════════════════════════════════
+    // PROLOGUE PHASE 1: distributed x² sum → inv_rms (ALL TGs)
+    // ═══════════════════════════════════════════════════════════════
+    int chunk = (N_OPROJ_TG + 63) / 64;
+    int x2_start = tid * chunk;
+    int x2_end = min(x2_start + chunk, N_OPROJ_TG);
+    float local_x2 = 0.0f;
+    for (int i = x2_start; i < x2_end; i++) local_x2 += x2_partials[i];
+    float sg_x2_sum = simd_sum(local_x2);
+
+    threadgroup float tg_x2_sg[2];
+    if (slid == 0) tg_x2_sg[sgid] = sg_x2_sum;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    float total_x2 = tg_x2_sg[0] + tg_x2_sg[1];
+    float inv_rms = metal::precise::rsqrt(total_x2 / (float)K + 1e-6f);
+
+    // ═══════════════════════════════════════════════════════════════
+    // PROLOGUE PHASE 2: Softmax + Top-k (routed TGs only)
+    // ═══════════════════════════════════════════════════════════════
+    threadgroup int tg_inds[{K_TOP}];
+    threadgroup float tg_selected_scores[{K_TOP}];
+
+    if (tgid.z < (uint)n_active) {{
+        // Load gate scores and apply inv_rms
+        float my_scores[SPT];
+        for (int j = 0; j < SPT; j++) {{
+            int e = tid * SPT + j;
+            if (e < E_CONST)
+                my_scores[j] = (gate_part_a[e] + gate_part_b[e]) * inv_rms;
+            else
+                my_scores[j] = -1e30f;
+        }}
+
+        // ── Softmax: distributed max ──
+        float local_max = -1e30f;
+        for (int j = 0; j < SPT; j++)
+            local_max = max(local_max, my_scores[j]);
+        float sg_max_val = simd_max(local_max);
+        threadgroup float tg_softmax_sg[2];
+        if (slid == 0) tg_softmax_sg[sgid] = sg_max_val;
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        float tg_max = max(tg_softmax_sg[0], tg_softmax_sg[1]);
+
+        // ── Softmax: exp + distributed sum ──
+        float local_sum = 0.0f;
+        for (int j = 0; j < SPT; j++) {{
+            float e_val = metal::exp(my_scores[j] - tg_max);
+            my_scores[j] = e_val;
+            local_sum += e_val;
+        }}
+        float sg_sum_val = simd_sum(local_sum);
+        if (slid == 0) tg_softmax_sg[sgid] = sg_sum_val;
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        float tg_sum = tg_softmax_sg[0] + tg_softmax_sg[1];
+
+        // ── Softmax: normalize ──
+        float inv_sum = 1.0f / tg_sum;
+        for (int j = 0; j < SPT; j++)
+            my_scores[j] *= inv_sum;
+
+        // ── Parallel top-k: K_TOP rounds ──
+        threadgroup float tg_tk_val[2];
+        threadgroup int tg_tk_info[2];
+
+        for (int round = 0; round < K_TOP_CONST; round++) {{
+            // Find local best among SPT scores
+            float best = -1.0f;
+            int best_e = -1;
+            for (int j = 0; j < SPT; j++) {{
+                int e = tid * SPT + j;
+                if (e < E_CONST && my_scores[j] > best) {{
+                    best = my_scores[j];
+                    best_e = e;
+                }}
+            }}
+
+            // SG-level max + winner identification
+            float sg_best = simd_max(best);
+            int candidate = (best == sg_best && best > 0.0f) ? int(slid) : 999;
+            int sg_winner = simd_min(candidate);
+
+            if (slid == 0) {{
+                tg_tk_val[sgid] = sg_best;
+                tg_tk_info[sgid] = sg_winner;
+            }}
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+
+            // Determine global winner (SG with higher max; tie-break: SG 0)
+            int winner_sg = (tg_tk_val[0] >= tg_tk_val[1]) ? 0 : 1;
+            int winner_lane = tg_tk_info[winner_sg];
+            int winner_tid = winner_sg * 32 + winner_lane;
+
+            // Winner writes expert index and score to TG memory
+            if (tid == winner_tid) {{
+                tg_inds[round] = best_e;
+                tg_selected_scores[round] = best;
+                // Disable the winning score in register
+                for (int j = 0; j < SPT; j++) {{
+                    if (tid * SPT + j == best_e) {{
+                        my_scores[j] = -1.0f;
+                        break;
+                    }}
+                }}
+            }}
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+        }}
+    }}
+{score_norm_block}
+
+    // ═══════════════════════════════════════════════════════════════
+    // PROLOGUE PHASE 3: Shared expert gate (TG(0,0,0) SG 0 only)
+    // 8-bit GEMV: W_seg (1, K) × X (K,) → gate_raw scalar
+    // Input is h_scaled; multiply result by inv_rms for true gate value
+    // ═══════════════════════════════════════════════════════════════
+    if (tgid.y == 0 && tgid.z == 0 && sgid == 0) {{
+        const int VPT = 8;
+        const int BLOCK = 256;  // 32 * VPT
+        int K_groups_seg = K / {gs};
+
+        const device uint8_t* wg_seg = (const device uint8_t*)W_seg
+            + slid * VPT;
+        const device {sc_t}* sc_seg = (const device {sc_t}*)S_seg
+            + slid / {slid_divisor};
+        const device {sc_t}* bi_seg = (const device {sc_t}*)B_seg
+            + slid / {slid_divisor};
+        int xb = slid * VPT;
+
+        float gate_acc = 0.0f;
+        for (int k = 0; k < K; k += BLOCK) {{
+            float xsum = 0.0f, wacc = 0.0f;
+            for (int i = 0; i < VPT; i++) {{
+                float xi = float(X[xb + i]);
+                xsum += xi;
+                wacc += xi * float(wg_seg[i]);
+            }}
+            gate_acc += float(*sc_seg) * wacc + xsum * float(*bi_seg);
+            wg_seg += BLOCK;
+            sc_seg += {sc_stride};
+            bi_seg += {sc_stride};
+            xb += BLOCK;
+        }}
+        gate_acc = simd_sum(gate_acc);
+        if (slid == 0) gate_raw[0] = gate_acc * inv_rms;
+    }}
+
+    // ═══════════════════════════════════════════════════════════════
+    // MAIN SwiGLU BODY
+    // ═══════════════════════════════════════════════════════════════
+    int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG;
+    int row_limit = (tgid.z < (uint)n_active) ? N_INTER : SHARED_INTER;
+    if (out_row >= row_limit) return;
+
+    float gate_result[4] = {{0, 0, 0, 0}};
+    float up_result[4] = {{0, 0, 0, 0}};
+
+    if (tgid.z < (uint)n_active) {{
+        // ═══════ 8-BIT ROUTED EXPERT PATH ═══════
+        const int VALUES_PER_THREAD = 8;
+        const int BLOCK_SIZE = 256;
+
+        int N_TOTAL = 2 * N_INTER;
+        int K_groups = K / {gs};
+
+        int expert = tg_inds[tgid.z];
+
+        const device uint8_t* ws_gate = (const device uint8_t*)W
+            + (long)expert * N_TOTAL * K + out_row * K + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_gate = (const device {sc_t}*)S
+            + (long)expert * N_TOTAL * K_groups + out_row * K_groups
+            + slid / {slid_divisor};
+        const device {sc_t}* bi_gate = (const device {sc_t}*)B_q
+            + (long)expert * N_TOTAL * K_groups + out_row * K_groups
+            + slid / {slid_divisor};
+
+        const device uint8_t* ws_up = (const device uint8_t*)W
+            + (long)expert * N_TOTAL * K + (out_row + N_INTER) * K
+            + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_up = (const device {sc_t}*)S
+            + (long)expert * N_TOTAL * K_groups + (out_row + N_INTER) * K_groups
+            + slid / {slid_divisor};
+        const device {sc_t}* bi_up = (const device {sc_t}*)B_q
+            + (long)expert * N_TOTAL * K_groups + (out_row + N_INTER) * K_groups
+            + slid / {slid_divisor};
+
+        int x_base = slid * VALUES_PER_THREAD;
+
+        for (int k = 0; k < K; k += 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++) {{
+                // Gate projection
+                const device uint8_t* wg = ws_gate + row * K;
+                float sg = float(sc_gate[row * K_groups]);
+                float bg = float(bi_gate[row * K_groups]);
+                float accum_g = 0;
+                for (int i = 0; i < 8; i++)
+                    accum_g += x_thread[i] * float(wg[i]);
+                gate_result[row] += sg * accum_g + xsum * bg;
+
+                // Up projection
+                const device uint8_t* wu = ws_up + row * K;
+                float su = float(sc_up[row * K_groups]);
+                float bu = float(bi_up[row * K_groups]);
+                float accum_u = 0;
+                for (int i = 0; i < 8; i++)
+                    accum_u += x_thread[i] * float(wu[i]);
+                up_result[row] += su * accum_u + xsum * bu;
+            }}
+
+            ws_gate += BLOCK_SIZE;
+            ws_up += BLOCK_SIZE;
+            sc_gate += {sc_stride};
+            sc_up += {sc_stride};
+            bi_gate += {sc_stride};
+            bi_up += {sc_stride};
+            x_base += BLOCK_SIZE;
+        }}
+
+        // Epilogue: apply inv_rms (factored), SwiGLU, write f32
+        device float* yp = Y_routed + tgid.z * N_INTER + out_row;
+        for (int row = 0; row < RESULTS_PER_SG; row++) {{
+            float g = simd_sum(gate_result[row]) * inv_rms;
+            float u = simd_sum(up_result[row]) * inv_rms;
+            if (slid == 0) {{
+                float silu_g = g / (1.0f + metal::exp(-g));
+                yp[row] = silu_g * u;
+            }}
+        }}
+
+    }} else {{
+        // ═══════ 8-BIT SHARED EXPERT PATH ═══════
+        const int VALUES_PER_THREAD = 8;
+        const int BLOCK_SIZE = 256;
+
+        int K_groups = K / {gs};
+
+        const device uint8_t* ws_gate = (const device uint8_t*)W_shared
+            + (long)out_row * K + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_gate = (const device {sc_t}*)S_shared
+            + (long)out_row * K_groups + slid / {slid_divisor};
+        const device {sc_t}* bi_gate = (const device {sc_t}*)B_shared
+            + (long)out_row * K_groups + slid / {slid_divisor};
+
+        const device uint8_t* ws_up = (const device uint8_t*)W_shared
+            + (long)(out_row + SHARED_INTER) * K + slid * VALUES_PER_THREAD;
+        const device {sc_t}* sc_up = (const device {sc_t}*)S_shared
+            + (long)(out_row + SHARED_INTER) * K_groups + slid / {slid_divisor};
+        const device {sc_t}* bi_up = (const device {sc_t}*)B_shared
+            + (long)(out_row + SHARED_INTER) * K_groups + slid / {slid_divisor};
+
+        int x_base = slid * VALUES_PER_THREAD;
+
+        for (int k = 0; k < K; k += 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* wg_s = ws_gate + row * K;
+                float sg_s = float(sc_gate[row * K_groups]);
+                float bg_s = float(bi_gate[row * K_groups]);
+                float accum_g = 0;
+                for (int i = 0; i < 8; i++)
+                    accum_g += x_thread[i] * float(wg_s[i]);
+                gate_result[row] += sg_s * accum_g + xsum * bg_s;
+
+                const device uint8_t* wu_s = ws_up + row * K;
+                float su_s = float(sc_up[row * K_groups]);
+                float bu_s = float(bi_up[row * K_groups]);
+                float accum_u = 0;
+                for (int i = 0; i < 8; i++)
+                    accum_u += x_thread[i] * float(wu_s[i]);
+                up_result[row] += su_s * accum_u + xsum * bu_s;
+            }}
+
+            ws_gate += BLOCK_SIZE;
+            ws_up += BLOCK_SIZE;
+            sc_gate += {sc_stride};
+            sc_up += {sc_stride};
+            bi_gate += {sc_stride};
+            bi_up += {sc_stride};
+            x_base += BLOCK_SIZE;
+        }}
+
+        // Epilogue: apply inv_rms (factored), SwiGLU, write f32
+        device float* yp = Y_shared + out_row;
+        for (int tm = 0; tm < RESULTS_PER_SG; tm++) {{
+            float g = simd_sum(gate_result[tm]) * inv_rms;
+            float u = simd_sum(up_result[tm]) * inv_rms;
+            if (slid == 0) {{
+                float silu_g = g / (1.0f + metal::exp(-g));
+                yp[tm] = silu_g * u;
+            }}
+        }}
+    }}
+"""
+
+
+_oproj_softmax_swiglu_8bit_kernels = {}
+
+
+def _get_oproj_softmax_swiglu_8bit_kernel(group_size=64, scale_bf16=True,
+                                            n_experts=64, top_k=10,
+                                            norm_topk=True):
+    key = (group_size, scale_bf16, n_experts, top_k, norm_topk)
+    if key not in _oproj_softmax_swiglu_8bit_kernels:
+        sc_tag = "_bf16sc" if scale_bf16 else ""
+        nt_tag = "_nt" if norm_topk else ""
+        _oproj_softmax_swiglu_8bit_kernels[key] = mx.fast.metal_kernel(
+            name=(f"oproj_softmax_topk_swiglu_8bit_gs{group_size}"
+                  f"_e{n_experts}_k{top_k}{sc_tag}{nt_tag}"),
+            input_names=[
+                "W", "S", "B_q",                    # routed expert weights
+                "W_shared", "S_shared", "B_shared",  # shared expert weights
+                "X",                                  # h_scaled (K,) bf16
+                "gate_part_a", "gate_part_b",         # (E,) f32
+                "x2_partials",                        # (N_OPROJ_TG,) f32
+                "W_seg", "S_seg", "B_seg",            # shared_expert_gate weights
+                "N_INTER_val", "SHARED_INTER_val",
+                "K_val", "n_active_val", "N_OPROJ_TG_val",
+            ],
+            output_names=["Y_routed", "Y_shared", "out_inds",
+                          "norm_scores", "gate_raw"],
+            source=_gen_oproj_softmax_topk_swiglu_8bit_source(
+                group_size, scale_bf16, n_experts, top_k, norm_topk),
+        )
+    return _oproj_softmax_swiglu_8bit_kernels[key]
+
+
+def fused_oproj_softmax_topk_swiglu_8bit(
+    w_q, s, b_q,                         # routed expert weights
+    w_shared, s_shared, b_shared,         # shared expert weights
+    h_scaled,                             # (K,) bf16 — h * w_rms
+    gate_part_a, gate_part_b,             # (E,) f32 — gate decomposition
+    x2_partials,                          # (N_OPROJ_TG,) f32 — per-TG x²
+    w_seg, s_seg, b_seg,                  # shared_expert_gate 8-bit weights
+    n_inter, k_hidden,                    # MoE dimensions
+    n_experts, top_k,                     # routing params
+    n_oproj_tg,                           # number of o_proj TGs (for x²)
+    group_size=64,
+    shared_inter=None,
+    norm_topk=True,
+):
+    """Single-dispatch softmax + top-k + merged 8-bit SwiGLU with oproj prologue.
+
+    Prologue:
+      Phase 1: distributed x² → inv_rms (all TGs)
+      Phase 2: softmax(gate_part_a + gate_part_b) → top-k → norm (routed TGs)
+      Phase 3: shared_expert_gate 8-bit GEMV → gate_raw (TG(0,0,0) SG 0)
+
+    Main: 8-bit gate+up+SwiGLU for routed + shared experts.
+    inv_rms is factored out of the inner loop and applied once in epilogue.
+
+    Args:
+        h_scaled: (K,) bf16 — h * w_rms from Dispatch 1
+        gate_part_a: (E,) f32 — M1 @ attn_out from Dispatch 1
+        gate_part_b: (E,) f32 — W_fused @ residual from Dispatch 1
+        x2_partials: (N_OPROJ_TG,) f32 — per-TG Σh² from Dispatch 1
+        w_seg/s_seg/b_seg: shared_expert_gate 8-bit weights (1, K/4) uint32
+
+    Returns:
+        (Y_routed, Y_shared, out_inds, norm_scores, gate_raw):
+            Y_routed: (top_k, n_inter) f32
+            Y_shared: (shared_inter,) f32
+            out_inds: (top_k,) uint32
+            norm_scores: (top_k,) f32
+            gate_raw: (1,) f32 — raw shared expert gate value (sigmoid in epilogue)
+    """
+    scale_bf16 = (s.dtype == mx.bfloat16)
+    kern = _get_oproj_softmax_swiglu_8bit_kernel(
+        group_size, scale_bf16, n_experts, top_k, norm_topk)
+
+    n_inter_val = int(n_inter)
+    shared_inter_val = int(shared_inter) if shared_inter is not None else n_inter_val
+    n_active = int(top_k)
+    max_inter = max(n_inter_val, shared_inter_val)
+
+    results = kern(
+        inputs=[
+            w_q, s, b_q,
+            w_shared, s_shared, b_shared,
+            h_scaled,
+            gate_part_a, gate_part_b,
+            x2_partials,
+            w_seg, s_seg, b_seg,
+            n_inter_val, shared_inter_val,
+            int(k_hidden), n_active, int(n_oproj_tg),
+        ],
+        output_shapes=[
+            (n_active, n_inter_val),     # Y_routed
+            (shared_inter_val,),         # Y_shared
+            (n_active,),                 # out_inds
+            (n_active,),                 # norm_scores
+            (1,),                        # gate_raw
+        ],
+        output_dtypes=[
+            mx.float32,    # Y_routed
+            mx.float32,    # Y_shared
+            mx.uint32,     # out_inds
+            mx.float32,    # norm_scores
+            mx.float32,    # gate_raw
+        ],
+        grid=(32, ceil_div(max_inter, 8) * 2, n_active + 1),
+        threadgroup=(32, 2, 1),
+    )
+    return results[0], results[1], results[2], results[3], results[4]

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

@@ -0,0 +1,245 @@
+"""MoE __call__ variants for Qwen3.5.
+
+Two modes:
+  _fused_moe_call: gate→SwiGLU→down_proj→epilogue (~15 dispatches, fused to ~4+MLX)
+  _oproj_moe_call: o_proj+gate→SwiGLU(w/softmax prologue)→down_proj→epilogue (4 dispatches)
+
+Kernels used:
+  merged_routed_shared_swiglu_8bit.py    — 8-bit gate+up+SwiGLU (_fused_moe_call)
+  merged_down_proj_8bit.py               — 8-bit down_proj (both modes)
+  fused_moe_epilogue_qwen.py             — weighted sum + sigmoid gate + residual (both modes)
+  custom_oproj_gate_gemv_8bit.py         — 8-bit o_proj + bf16 gate GEMVs (_oproj_moe_call)
+  oproj_softmax_topk_swiglu_8bit.py      — SwiGLU with softmax prologue (_oproj_moe_call)
+
+Adapted from mlx_bench/model_patches/qwen/moe.py.
+"""
+
+import mlx.core as mx
+
+from .kernels.merged_routed_shared_swiglu_8bit import fused_merged_gate_up_swiglu_8bit
+from .kernels.merged_down_proj_8bit import fused_merged_down_proj_8bit
+from .kernels.fused_moe_epilogue_qwen import fused_moe_epilogue_qwen
+from .kernels.custom_oproj_gate_gemv_8bit import fused_custom_oproj_8bit
+from .kernels.oproj_softmax_topk_swiglu_8bit import fused_oproj_softmax_topk_swiglu_8bit
+
+
+def _vanilla_moe_call(self, x):
+    """Original Qwen3NextSparseMoeBlock.__call__ (for prefill fallback)."""
+    gates = self.gate(x)
+    gates = mx.softmax(gates, axis=-1, precise=True)
+
+    k = self.top_k
+    inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]
+    scores = mx.take_along_axis(gates, inds, axis=-1)
+    if self.norm_topk_prob:
+        scores = scores / scores.sum(axis=-1, keepdims=True)
+
+    y = self.switch_mlp(x, inds)
+    y = (y * scores[..., None]).sum(axis=-2)
+
+    shared_y = self.shared_expert(x)
+    shared_y = mx.sigmoid(self.shared_expert_gate(x)) * shared_y
+
+    return y + shared_y
+
+
+def _fused_moe_call(self, x, _residual=None):
+    """Qwen3.5 MoE with fused kernels (4 custom dispatches).
+
+    Falls back to vanilla for prefill (seq_len > 1).
+
+    Args:
+        x: (B, S, K) bf16 — post-layernorm hidden state
+        _residual: (B, S, K) bf16 — pre-layernorm hidden state for residual add.
+                   If None, epilogue skips residual (returns MoE output only).
+    """
+    # Fused kernels are decode-only (seq_len=1). Fall back for prefill.
+    if x.shape[-2] > 1:
+        # Try vanilla switch_mlp path if weights still exist
+        has_switch_weights = hasattr(self.switch_mlp, 'gate_proj') and \
+            hasattr(self.switch_mlp.gate_proj, 'weight')
+        if has_switch_weights:
+            out = _vanilla_moe_call(self, x)
+        else:
+            # Weights were freed — process tokens one by one through fused path
+            outs = []
+            for t in range(x.shape[-2]):
+                xt = x[:, t:t+1, :]
+                res_t = _residual[:, t:t+1, :] if _residual is not None else None
+                outs.append(_fused_moe_call(self, xt, _residual=res_t))
+            return mx.concatenate(outs, axis=1)
+        if _residual is not None:
+            out = out + _residual
+        return out
+
+    # ── Gate routing (vanilla MLX ops) ──
+    gates = self.gate(x)
+    gates = mx.softmax(gates, axis=-1, precise=True)
+
+    k = self.top_k
+    inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]
+    scores = mx.take_along_axis(gates, inds, axis=-1)
+    if self.norm_topk_prob:
+        scores = scores / scores.sum(axis=-1, keepdims=True)
+
+    x_flat = x.reshape(-1).astype(mx.bfloat16)  # (K,)
+    inds_flat = inds.reshape(-1).astype(mx.uint32)
+    scores_flat = scores.reshape(-1).astype(mx.float32)
+
+    # ── Dispatch 1: Merged gate+up+SwiGLU (8-bit routed + 8-bit shared) ──
+    y_routed, y_shared = fused_merged_gate_up_swiglu_8bit(
+        self.switch_mlp._fused_w_gu,
+        self.switch_mlp._fused_s_gu,
+        self.switch_mlp._fused_b_gu,
+        self._shared_w_gu,
+        self._shared_s_gu,
+        self._shared_b_gu,
+        x_flat,
+        inds_flat,
+        self.switch_mlp._fused_n_inter,
+        self.switch_mlp._fused_k_hidden,
+        group_size=self.switch_mlp._fused_group_size,
+        shared_inter=self._shared_inter,
+    )
+
+    # ── Dispatch 2: Merged down_proj (8-bit routed + 8-bit shared) ──
+    y_down_routed, y_down_shared = fused_merged_down_proj_8bit(
+        self._down_w,
+        self._down_s,
+        self._down_b,
+        self._shared_down_w,
+        self._shared_down_s,
+        self._shared_down_b,
+        y_routed,
+        y_shared,
+        inds_flat,
+        self._down_K,
+        self._down_N,
+        group_size=self._down_gs,
+        shared_n_in=self._shared_inter,
+    )
+
+    # ── Dispatch 3: Shared expert gate (small GEMV) ──
+    shared_gate_out = self.shared_expert_gate(x.reshape(1, 1, -1))
+    shared_gate_val = mx.sigmoid(shared_gate_out.reshape(()))
+
+    # ── Dispatch 4: Fused epilogue ──
+    if _residual is not None:
+        h_flat = _residual.reshape(-1).astype(mx.bfloat16)
+    else:
+        h_flat = mx.zeros((self.switch_mlp._fused_k_hidden,), dtype=mx.bfloat16)
+
+    y = fused_moe_epilogue_qwen(
+        y_down_routed,    # (n_active, K) f32
+        y_down_shared,    # (K,) f32
+        scores_flat,      # (n_active,) f32
+        h_flat,           # (K,) bf16
+        shared_gate_val,  # scalar f32
+        self._down_K,     # K
+    )
+    return y.reshape(1, 1, -1)
+
+
+def _oproj_moe_call(self, attn_out_3d, _residual=None):
+    """Qwen3.5 MoE with fused o_proj + gate GEMVs (4 custom dispatches).
+
+    Receives raw attention output (pre-o_proj) and residual.
+    Fuses o_proj, RMSNorm, gate softmax, top-k, score norm, shared_expert_gate,
+    SwiGLU, down_proj, and epilogue into 4 dispatches.
+
+    Dispatch 1: 8-bit o_proj + bf16 M1/W_fused GEMVs → h_scaled, h_out, x2_partials,
+                gate_part_a, gate_part_b
+    Dispatch 2: SwiGLU with softmax prologue (inv_rms + softmax + top-k + score norm
+                + shared_expert_gate in TG(0,0,0)) → y_routed, y_shared, inds, scores,
+                gate_raw
+    Dispatch 3: Merged down_proj (unchanged)
+    Dispatch 4: Fused epilogue with sigmoid(gate_raw)
+
+    Args:
+        attn_out_3d: (1, 1, K_attn) bf16 — raw attention output (pre-o_proj)
+        _residual: (1, 1, K) bf16 — input residual (before attention)
+    """
+    # Prefill fallback (S > 1): restore o_proj + vanilla MoE
+    if attn_out_3d.shape[-2] > 1:
+        from .decoder import _parent_layer_map
+        parent = _parent_layer_map[id(self)]
+        if parent.is_linear:
+            oproj_mod = parent.linear_attn.out_proj
+        else:
+            oproj_mod = parent.self_attn.o_proj
+        projected = oproj_mod(attn_out_3d)
+        h = _residual + projected
+        out = _vanilla_moe_call(self, parent.post_attention_layernorm(h))
+        return out + h
+
+    attn_out = attn_out_3d.reshape(-1).astype(mx.bfloat16)  # (K_attn,)
+    residual = _residual.reshape(-1).astype(mx.bfloat16)     # (K,)
+
+    K = self._oproj_M
+    K_attn = self._oproj_K_attn
+
+    # ── Dispatch 1: fused o_proj + M1 GEMV + W_fused GEMV ──
+    h_scaled, h_out, x2_partials, gate_part_a, gate_part_b = \
+        fused_custom_oproj_8bit(
+            self._oproj_w, self._oproj_s, self._oproj_b,
+            attn_out, residual, self._oproj_rms_weight,
+            self._oproj_M1, self._oproj_W_fused,
+            M=K, K_attn=K_attn, K_hidden=self._oproj_K_hidden,
+            n_experts=self._oproj_n_experts,
+            gate_bm=self._oproj_gate_bm,
+        )
+
+    # ── Dispatch 2: SwiGLU with softmax prologue ──
+    y_routed, y_shared, inds, scores, gate_raw = \
+        fused_oproj_softmax_topk_swiglu_8bit(
+            self.switch_mlp._fused_w_gu,
+            self.switch_mlp._fused_s_gu,
+            self.switch_mlp._fused_b_gu,
+            self._shared_w_gu,
+            self._shared_s_gu,
+            self._shared_b_gu,
+            h_scaled,
+            gate_part_a,
+            gate_part_b,
+            x2_partials,
+            self._seg_w,
+            self._seg_s,
+            self._seg_b,
+            n_inter=self.switch_mlp._fused_n_inter,
+            k_hidden=K,
+            n_experts=self._oproj_n_experts,
+            top_k=self.top_k,
+            n_oproj_tg=self._oproj_n_tg,
+            group_size=self.switch_mlp._fused_group_size,
+            shared_inter=self._shared_inter,
+            norm_topk=self.norm_topk_prob,
+        )
+
+    # ── Dispatch 3: Merged down_proj (8-bit routed + 8-bit shared) ──
+    y_down_routed, y_down_shared = fused_merged_down_proj_8bit(
+        self._down_w,
+        self._down_s,
+        self._down_b,
+        self._shared_down_w,
+        self._shared_down_s,
+        self._shared_down_b,
+        y_routed,
+        y_shared,
+        inds,
+        self._down_K,
+        self._down_N,
+        group_size=self._down_gs,
+        shared_n_in=self._shared_inter,
+    )
+
+    # ── Dispatch 4: Fused epilogue with sigmoid(gate_raw) ──
+    y = fused_moe_epilogue_qwen(
+        y_down_routed,    # (n_active, K) f32
+        y_down_shared,    # (K,) f32
+        scores,           # (n_active,) f32 — norm_topk_prob from prologue
+        h_out,            # (K,) bf16 — post-o_proj hidden for residual add
+        gate_raw,         # scalar f32 — raw dot product, sigmoid fused
+        K,
+        fuse_sigmoid=True,
+    )
+    return y.reshape(1, 1, -1)

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

@@ -40,6 +40,7 @@ from pydantic import RootModel
 from exo.download.download_utils import build_model_path
 from exo.shared.types.common import Host
 from exo.shared.types.memory import Memory
+from exo.shared.types.mlx import Model
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import (
     BoundInstance,
@@ -52,7 +53,6 @@ from exo.shared.types.worker.shards import (
     ShardMetadata,
     TensorShardMetadata,
 )
-from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.auto_parallel import (
     LayerLoadedCallback,
     TimeoutCallback,
@@ -189,6 +189,10 @@ def load_mlx_items(
         mx.eval(model)
         end_time = time.perf_counter()
         logger.info(f"Time taken to load model: {(end_time - start_time):.2f}s")
+
+        from exo.worker.engines.mlx.patches import maybe_apply_patches
+        maybe_apply_patches(model, model_path)
+
         tokenizer = get_tokenizer(model_path, bound_instance.bound_shard)
 
     else:
@@ -318,6 +322,9 @@ def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
         return [151336, 151329, 151338]
     elif "gpt-oss" in model_id_lower:
         return [200002, 200012]
+    elif "qwen3.5" in model_id_lower or "qwen-3.5" in model_id_lower:
+        # For Qwen3.5: 248046 (<|im_end|>), 248044 (<|endoftext|>)
+        return [248046, 248044]
     return None
 
 
@@ -551,6 +558,8 @@ def apply_chat_template(
         # Jinja ignores unknown variables, so passing both is safe.
         extra_kwargs["enable_thinking"] = task_params.enable_thinking
         extra_kwargs["thinking"] = task_params.enable_thinking
+    if task_params.reasoning_effort is not None:
+        extra_kwargs["reasoning_effort"] = task_params.reasoning_effort
 
     patched_template: str | None = None
     if task_params.tools:

src/exo/worker/main.py

@@ -1,9 +1,8 @@
 from collections import defaultdict
 from datetime import datetime, timezone
-from random import random
 
 import anyio
-from anyio import CancelScope, fail_after
+from anyio import fail_after
 from loguru import logger
 
 from exo.download.download_utils import resolve_model_in_path
@@ -13,17 +12,13 @@ from exo.shared.types.api import ImageEditsTaskParams
 from exo.shared.types.commands import (
     ForwarderCommand,
     ForwarderDownloadCommand,
-    RequestEventLog,
     StartDownload,
 )
-from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
+from exo.shared.types.common import CommandId, NodeId, SystemId
 from exo.shared.types.events import (
     Event,
-    EventId,
-    GlobalForwarderEvent,
     IndexedEvent,
     InputChunkReceived,
-    LocalForwarderEvent,
     NodeDownloadProgress,
     NodeGatheredInfo,
     TaskCreated,
@@ -46,7 +41,6 @@ from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.runners import RunnerId
 from exo.utils.channels import Receiver, Sender, channel
-from exo.utils.event_buffer import OrderedBuffer
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
 from exo.utils.info_gatherer.net_profile import check_reachable
 from exo.utils.keyed_backoff import KeyedBackoff
@@ -59,38 +53,26 @@ class Worker:
     def __init__(
         self,
         node_id: NodeId,
-        session_id: SessionId,
         *,
-        global_event_receiver: Receiver[GlobalForwarderEvent],
-        local_event_sender: Sender[LocalForwarderEvent],
+        event_receiver: Receiver[IndexedEvent],
+        event_sender: Sender[Event],
         # This is for requesting updates. It doesn't need to be a general command sender right now,
         # but I think it's the correct way to be thinking about commands
         command_sender: Sender[ForwarderCommand],
         download_command_sender: Sender[ForwarderDownloadCommand],
     ):
         self.node_id: NodeId = node_id
-        self.session_id: SessionId = session_id
-
-        self.global_event_receiver = global_event_receiver
-        self.local_event_sender = local_event_sender
+        self.event_receiver = event_receiver
+        self.event_sender = event_sender
         self.command_sender = command_sender
         self.download_command_sender = download_command_sender
-        self.event_buffer = OrderedBuffer[Event]()
-        self.out_for_delivery: dict[EventId, LocalForwarderEvent] = {}
 
         self.state: State = State()
         self.runners: dict[RunnerId, RunnerSupervisor] = {}
         self._tg: TaskGroup = TaskGroup()
 
-        self._nack_cancel_scope: CancelScope | None = None
-        self._nack_attempts: int = 0
-        self._nack_base_seconds: float = 0.5
-        self._nack_cap_seconds: float = 10.0
-
         self._system_id = SystemId()
 
-        self.event_sender, self.event_receiver = channel[Event]()
-
         # Buffer for input image chunks (for image editing)
         self.input_chunk_buffer: dict[CommandId, dict[int, str]] = {}
         self.input_chunk_counts: dict[CommandId, int] = {}
@@ -108,14 +90,12 @@ class Worker:
                 tg.start_soon(info_gatherer.run)
                 tg.start_soon(self._forward_info, info_recv)
                 tg.start_soon(self.plan_step)
-                tg.start_soon(self._resend_out_for_delivery)
                 tg.start_soon(self._event_applier)
-                tg.start_soon(self._forward_events)
                 tg.start_soon(self._poll_connection_updates)
         finally:
             # Actual shutdown code - waits for all tasks to complete before executing.
             logger.info("Stopping Worker")
-            self.local_event_sender.close()
+            self.event_sender.close()
             self.command_sender.close()
             self.download_command_sender.close()
             for runner in self.runners.values():
@@ -133,47 +113,22 @@ class Worker:
                 )
 
     async def _event_applier(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)
-                event_id = f_event.event.event_id
-                if event_id in self.out_for_delivery:
-                    del self.out_for_delivery[event_id]
-
+        with self.event_receiver as events:
+            async for event in events:
                 # 2. for each event, apply it to the state
-                indexed_events = self.event_buffer.drain_indexed()
-                if indexed_events:
-                    self._nack_attempts = 0
+                self.state = apply(self.state, event=event)
+                event = event.event
 
-                if not indexed_events 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, self.state.last_event_applied_idx + 1
+                # Buffer input image chunks for image editing
+                if isinstance(event, InputChunkReceived):
+                    cmd_id = event.command_id
+                    if cmd_id not in self.input_chunk_buffer:
+                        self.input_chunk_buffer[cmd_id] = {}
+                        self.input_chunk_counts[cmd_id] = event.chunk.total_chunks
+
+                    self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
+                        event.chunk.data
                     )
-                    continue
-                elif indexed_events and self._nack_cancel_scope:
-                    self._nack_cancel_scope.cancel()
-
-                for idx, event in indexed_events:
-                    self.state = apply(self.state, IndexedEvent(idx=idx, event=event))
-
-                    # Buffer input image chunks for image editing
-                    if isinstance(event, InputChunkReceived):
-                        cmd_id = event.command_id
-                        if cmd_id not in self.input_chunk_buffer:
-                            self.input_chunk_buffer[cmd_id] = {}
-                            self.input_chunk_counts[cmd_id] = event.chunk.total_chunks
-
-                        self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
-                            event.chunk.data
-                        )
 
     async def plan_step(self):
         while True:
@@ -325,43 +280,6 @@ class Worker:
                 instance.shard_assignments.node_to_runner[self.node_id]
             ].start_task(task)
 
-    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
-
-    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)
-
     def _create_supervisor(self, task: CreateRunner) -> RunnerSupervisor:
         """Creates and stores a new AssignedRunner with initial downloading status."""
         runner = RunnerSupervisor.create(
@@ -372,21 +290,6 @@ class Worker:
         self._tg.start_soon(runner.run)
         return runner
 
-    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"Worker published event {idx}: {str(event)[:100]}")
-                await self.local_event_sender.send(fe)
-                self.out_for_delivery[event.event_id] = fe
-
     async def _poll_connection_updates(self):
         while True:
             edges = set(

src/exo/worker/plan.py

@@ -297,10 +297,10 @@ def _pending_tasks(
             # the task status _should_ be set to completed by the LAST runner
             # it is currently set by the first
             # this is definitely a hack
-            if task.task_id in runner.completed:
+            if task.task_id in runner.completed or task.task_id in runner.in_progress:
                 continue
 
-            if isinstance(runner.status, RunnerReady) and all(
+            if isinstance(runner.status, (RunnerReady, RunnerRunning)) and all(
                 isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning))
                 for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard
             ):

src/exo/worker/runner/bootstrap.py

@@ -32,11 +32,19 @@ def entrypoint(
     # Import main after setting global logger - this lets us just import logger from this module
     try:
         if bound_instance.is_image_model:
-            from exo.worker.runner.image_models.runner import main
-        else:
-            from exo.worker.runner.llm_inference.runner import main
+            from exo.worker.runner.image_models.runner import Runner as ImageRunner
 
-        main(bound_instance, event_sender, task_receiver, cancel_receiver)
+            runner = ImageRunner(
+                bound_instance, event_sender, task_receiver, cancel_receiver
+            )
+            runner.main()
+        else:
+            from exo.worker.runner.llm_inference.runner import Runner
+
+            runner = Runner(
+                bound_instance, event_sender, task_receiver, cancel_receiver
+            )
+            runner.main()
 
     except ClosedResourceError:
         logger.warning("Runner communication closed unexpectedly")