Bake scale into SDPA pass1 source, remove last scalar input

This reverts commit 97d7f03032.

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

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.

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>

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,7 +346,13 @@ 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:
+
+* `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
 
@@ -204,24 +375,148 @@ Cancels an active generation command (text or image). Notifies workers and close
 
 Returns 404 if the command is not found or already completed.
 
-## 5. Image Generation & Editing
+## 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 |
@@ -231,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
 
@@ -244,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
 
@@ -267,37 +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:

rust/exo_pyo3_bindings/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "exo_pyo3_bindings"
-version = "0.2.0"
+version = "0.2.1"
 description = "Add your description here"
 readme = "README.md"
 authors = [

rust/networking/src/swarm.rs

@@ -263,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

@@ -133,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
@@ -287,8 +297,8 @@ class DownloadCoordinator:
             del self.download_status[model_id]
 
     async def _emit_existing_download_progress(self) -> None:
-        try:
-            while True:
+        while True:
+            try:
                 logger.debug(
                     "DownloadCoordinator: Fetching and emitting existing download progress..."
                 )
@@ -376,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/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_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/master/adapters/chat_completions.py

@@ -104,6 +104,8 @@ def chat_request_to_text_generation(
         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/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
@@ -775,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:
@@ -1586,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()

src/exo/routing/router.py

@@ -219,6 +219,10 @@ class Router:
             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

src/exo/shared/types/api.py

@@ -201,6 +201,8 @@ class ChatCompletionRequest(BaseModel):
     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

src/exo/shared/types/text_generation.py

@@ -67,3 +67,5 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     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/worker/engines/mlx/generator/generate.py

@@ -10,7 +10,7 @@ from mlx_lm.generate import (
     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 (
@@ -470,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,

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

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

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

@@ -116,6 +116,9 @@ def parse_gpt_oss(
         if current_tool_name is not None:
             if delta:
                 tool_arg_parts.append(delta)
+            if response.finish_reason is not None:
+                yield response.model_copy(update={"text": "".join(tool_arg_parts)})
+                tool_arg_parts = []
             continue
 
         if ch == "analysis" and not thinking:

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

@@ -216,6 +216,8 @@ class Runner:
                         tok.tool_parser,  # pyright: ignore[reportAny]
                     )
 
+                self.generator.kv_prefix_cache = KVPrefixCache(self.generator.group)
+
                 self.generator = self.generator.build()
 
                 self.send_task_status(task.task_id, TaskStatus.Complete)

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

@@ -1,5 +1,6 @@
 from collections.abc import Generator
 
+from exo.shared.types.api import FinishReason
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
     ToolCallResponse,
@@ -83,6 +84,7 @@ THINKING_THEN_TOOL_TOKENS: list[tuple[int, str]] = [
 
 def _make_gen_responses(
     tokens: list[tuple[int, str]],
+    last_finish_reason: FinishReason = "stop",
 ) -> list[GenerationResponse]:
     """Build GenerationResponse list from (token_id, text) pairs."""
     responses: list[GenerationResponse] = []
@@ -92,7 +94,7 @@ def _make_gen_responses(
             GenerationResponse(
                 text=text,
                 token=tid,
-                finish_reason="stop" if is_last else None,
+                finish_reason=last_finish_reason if is_last else None,
                 usage=None,
             )
         )
@@ -101,11 +103,12 @@ def _make_gen_responses(
 
 def _collect(
     tokens: list[tuple[int, str]],
+    last_finish_reason: FinishReason = "stop",
 ) -> list[GenerationResponse | ToolCallResponse]:
     """Feed tokens through parse_gpt_oss and collect all yielded responses."""
 
     def _gen() -> Generator[GenerationResponse, None, None]:
-        yield from _make_gen_responses(tokens)
+        yield from _make_gen_responses(tokens, last_finish_reason)
 
     return list(x for x in parse_gpt_oss(_gen()) if x is not None)
 
@@ -171,3 +174,80 @@ class TestParseGptOssThinkingThenToolCall:
         tc = _get_tool_call(results)
         assert tc.tool_calls[0].name == "get_current_weather"
         assert "Tokyo" in tc.tool_calls[0].arguments
+
+
+# fmt: off
+# Truncated tool call: recipient + channel + message + partial args, no <|call|>
+TRUNCATED_TOOL_CALL_TOKENS: list[tuple[int, str]] = [
+    (316,    " to"),
+    (28,     "="),
+    (44580,  "functions"),
+    (775,    ".get"),
+    (23981,  "_current"),
+    (170154, "_weather"),
+    (_CHANNEL, "<|channel|>"),
+    (12606,  "comment"),
+    (815,    "ary"),
+    (5701,   " json"),
+    (_MESSAGE, "<|message|>"),
+    (10848,  '{"'),
+    (7693,   "location"),
+    (1243,   '":'),
+    (392,    ' "'),
+    (173844, "Tokyo"),
+    # No <|call|> — generation truncated here
+]
+
+# Plain text tokens (no tool call)
+PLAIN_TEXT_TOKENS: list[tuple[int, str]] = [
+    (_CHANNEL, "<|channel|>"),
+    (35644,  "analysis"),
+    (_MESSAGE, "<|message|>"),
+    (12845,  "Let"),
+    (668,    " me"),
+    (2411,   " think"),
+    (1078,   " about"),
+    (495,    " this"),
+    (13,     "."),
+    (_END,   "<|end|>"),
+    (_START, "<|start|>"),
+    (_ASSISTANT, "assistant"),
+    (_CHANNEL, "<|channel|>"),
+    (12606,  "comment"),
+    (815,    "ary"),
+    (_MESSAGE, "<|message|>"),
+    (9906,   "Hello"),
+    (14,     ","),
+    (2989,   " world"),
+]
+# fmt: on
+
+
+class TestParseGptOssMaxTokensTruncation:
+    """Truncated tool calls must still yield finish_reason."""
+
+    def test_truncated_tool_call_yields_finish_reason(self):
+        results = _collect(TRUNCATED_TOOL_CALL_TOKENS, last_finish_reason="length")
+        gen_responses = [r for r in results if isinstance(r, GenerationResponse)]
+        finish_reasons = [
+            r.finish_reason for r in gen_responses if r.finish_reason is not None
+        ]
+        assert "length" in finish_reasons
+
+    def test_truncated_tool_call_emits_partial_args(self):
+        results = _collect(TRUNCATED_TOOL_CALL_TOKENS, last_finish_reason="length")
+        gen_responses = [r for r in results if isinstance(r, GenerationResponse)]
+        last = [r for r in gen_responses if r.finish_reason is not None][-1]
+        assert len(last.text) > 0
+
+    def test_truncated_plain_text_still_works(self):
+        results = _collect(PLAIN_TEXT_TOKENS, last_finish_reason="length")
+        gen_responses = [r for r in results if isinstance(r, GenerationResponse)]
+        finish_reasons = [
+            r.finish_reason for r in gen_responses if r.finish_reason is not None
+        ]
+        assert "length" in finish_reasons
+        # Verify non-empty text was yielded (delta text differs from raw token text
+        # due to Harmony encoding, so we just check something was emitted)
+        all_text = "".join(r.text for r in gen_responses)
+        assert len(all_text) > 0