fix: add distributed sync for tensor parallel operation coordination

Ensures all ranks execute the same operations in the same order to prevent
collective mismatches that would cause deadlocks in distributed inference.

- Add DistributedOp enum and sync_operation() for broadcasting operations
- Add queue_request() and sync_and_insert_pending() to BatchGenerationEngine
- Restructure runner loop to sync operations before execution in distributed mode
- Add tests for distributed sync behavior

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

.github/workflows/pipeline.yml

@@ -94,9 +94,19 @@ jobs:
 
       - uses: ./.github/actions/typecheck
 
-  nix-flake-check:
-    name: Check Nix flake
-    runs-on: ubuntu-latest
+  nix:
+    name: Build and check (${{ matrix.system }})
+    runs-on: ${{ matrix.runner }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - runner: macos-26
+            system: aarch64-darwin
+          - runner: ubuntu-latest
+            system: x86_64-linux
+          - runner: ubuntu-24.04-arm
+            system: aarch64-linux
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -113,83 +123,14 @@ jobs:
           name: exo
           authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
 
-      - name: Run nix flake check
+      - name: Build all Nix outputs
         run: |
-          nix flake check
-        shell: bash
+          nix flake show --json | jq -r '
+            [
+              (.packages."${{ matrix.system }}" // {} | keys[] | ".#packages.${{ matrix.system }}.\(.)"),
+              (.devShells."${{ matrix.system }}" // {} | keys[] | ".#devShells.${{ matrix.system }}.\(.)")
+            ] | .[]
+          ' | xargs nix build
 
-#  ci:
-#    needs: typecheck
-#    runs-on: ubuntu-latest
-#    permissions:
-#      contents: read
-#    env:
-#      GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-#    steps:
-#      - name: Checkout repository
-#        uses: actions/checkout@v4
-#        with:
-#          fetch-depth: 0
-#          token: ${{ secrets.GITHUB_TOKEN }}
-#          lfs: true
-#
-#      - name: Configure git user
-#        run: |
-#          git config --local user.email "github-actions@users.noreply.github.com"
-#          git config --local user.name  "github-actions bot"
-#        shell: bash
-#
-#      - name: Pull LFS files
-#        run: |
-#          echo "Pulling Git LFS files..."
-#          git lfs pull
-#        shell: bash
-#
-#      - name: Setup EXO_HOME and API_PORT
-#        run: |
-#          EXO_HOME=$(mktemp -d -t exo-ci-XXXXXXXX)
-#          # Generate random port (macOS compatible method)
-#          API_PORT=$((49152 + RANDOM % (65535 - 49152 + 1)))
-#          echo "EXO_HOME=$EXO_HOME" >> $GITHUB_ENV
-#          echo "API_PORT=$API_PORT" >> $GITHUB_ENV
-#          echo "Created EXO_HOME: $EXO_HOME"
-#          echo "Generated API_PORT: $API_PORT"
-#        shell: bash
-#
-#      - name: Setup Nix Environment
-#        run: |
-#          echo "Checking for nix installation..."
-#          
-#          # Check if nix binary exists directly
-#          if [ -f /nix/var/nix/profiles/default/bin/nix ]; then
-#            echo "Found nix binary at /nix/var/nix/profiles/default/bin/nix"
-#            export PATH="/nix/var/nix/profiles/default/bin:$PATH"
-#            echo "PATH=$PATH" >> $GITHUB_ENV
-#            nix --version
-#          elif [ -f /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh ]; then
-#            echo "Found nix profile script, sourcing..."
-#            source /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh
-#            nix --version
-#          elif command -v nix >/dev/null 2>&1; then
-#            echo "Nix already in PATH"
-#            nix --version
-#          else
-#            echo "Nix not found. Debugging info:"
-#            echo "Contents of /nix/var/nix/profiles/default/:"
-#            ls -la /nix/var/nix/profiles/default/ 2>/dev/null || echo "Directory not found"
-#            echo "Contents of /nix/var/nix/profiles/default/bin/:"
-#            ls -la /nix/var/nix/profiles/default/bin/ 2>/dev/null || echo "Directory not found"
-#            exit 1
-#          fi
-#        shell: bash
-#
-#      - uses: ./.github/actions/lint-check
-#
-#      - uses: ./.github/actions/unit-test
-#
-#      - name: Cleanup EXO_HOME
-#        run: |
-#          echo "Cleaning up EXO_HOME: $EXO_HOME"
-#          rm -rf "$EXO_HOME"
-#        shell: bash
-#        if: always()
+      - name: Run nix flake check
+        run: nix flake check

.mlx_typings/mlx_lm/generate.pyi

@@ -276,24 +276,23 @@ class BatchGenerator:
         logprobs: mx.array
         finish_reason: Optional[str]
 
+    unprocessed_prompts: List[Any]
+
     def __init__(
         self,
-        model,
+        model: nn.Module,
         max_tokens: int = ...,
-        stop_tokens: Optional[set] = ...,
+        stop_tokens: Optional[set[int]] = ...,
         sampler: Optional[Callable[[mx.array], mx.array]] = ...,
         completion_batch_size: int = ...,
         prefill_batch_size: int = ...,
         prefill_step_size: int = ...,
     ) -> None: ...
     def insert(
-        self, prompts, max_tokens: Union[List[int], int, None] = ...
-    ):  # -> list[Any]:
-        ...
-    def stats(self):  # -> BatchStats:
-        ...
-    def next(self):  # -> list[Any]:
-        ...
+        self, prompts: List[List[int]], max_tokens: Union[List[int], int, None] = ...
+    ) -> List[int]: ...
+    def stats(self) -> BatchStats: ...
+    def next(self) -> List[Response]: ...
 
 def batch_generate(
     model,

.mlx_typings/mlx_lm/models/deepseek_v3.pyi

@@ -0,0 +1,156 @@
+"""Type stubs for mlx_lm.models.deepseek_v3"""
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    moe_intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    n_shared_experts: Optional[int]
+    n_routed_experts: Optional[int]
+    routed_scaling_factor: float
+    kv_lora_rank: int
+    q_lora_rank: Optional[int]
+    qk_rope_head_dim: int
+    v_head_dim: int
+    qk_nope_head_dim: int
+    topk_method: str
+    scoring_func: str
+    norm_topk_prob: bool
+    n_group: int
+    topk_group: int
+    num_experts_per_tok: int
+    moe_layer_freq: int
+    first_k_dense_replace: int
+    max_position_embeddings: int
+    rms_norm_eps: float
+    rope_theta: float
+    rope_scaling: Optional[Dict[str, Any]]
+    attention_bias: bool
+
+class DeepseekV3Attention(nn.Module):
+    config: ModelArgs
+    hidden_size: int
+    num_heads: int
+    max_position_embeddings: int
+    rope_theta: float
+    q_lora_rank: Optional[int]
+    qk_rope_head_dim: int
+    kv_lora_rank: int
+    v_head_dim: int
+    qk_nope_head_dim: int
+    q_head_dim: int
+    scale: float
+    q_proj: nn.Linear
+    q_a_proj: nn.Linear
+    q_a_layernorm: nn.RMSNorm
+    q_b_proj: nn.Linear
+    kv_a_proj_with_mqa: nn.Linear
+    kv_a_layernorm: nn.RMSNorm
+    kv_b_proj: nn.Linear
+    o_proj: nn.Linear
+    rope: Any
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class DeepseekV3MLP(nn.Module):
+    config: ModelArgs
+    hidden_size: int
+    intermediate_size: int
+    gate_proj: nn.Linear
+    up_proj: nn.Linear
+    down_proj: nn.Linear
+
+    def __init__(
+        self,
+        config: ModelArgs,
+        hidden_size: Optional[int] = None,
+        intermediate_size: Optional[int] = None,
+    ) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class MoEGate(nn.Module):
+    config: ModelArgs
+    top_k: int
+    norm_topk_prob: bool
+    n_routed_experts: Optional[int]
+    routed_scaling_factor: float
+    n_group: int
+    topk_group: int
+    weight: mx.array
+    e_score_correction_bias: mx.array
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
+
+class DeepseekV3MoE(nn.Module):
+    config: ModelArgs
+    num_experts_per_tok: int
+    switch_mlp: SwitchGLU
+    gate: MoEGate
+    shared_experts: DeepseekV3MLP
+    sharding_group: Optional[mx.distributed.Group]
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class DeepseekV3DecoderLayer(nn.Module):
+    self_attn: DeepseekV3Attention
+    mlp: DeepseekV3MLP | DeepseekV3MoE
+    input_layernorm: nn.RMSNorm
+    post_attention_layernorm: nn.RMSNorm
+
+    def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class DeepseekV3Model(nn.Module):
+    vocab_size: int
+    embed_tokens: nn.Embedding
+    layers: list[DeepseekV3DecoderLayer]
+    norm: nn.RMSNorm
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    model_type: str
+    model: DeepseekV3Model
+    lm_head: nn.Linear
+
+    def __init__(self, config: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    @property
+    def layers(self) -> list[DeepseekV3DecoderLayer]: ...

.mlx_typings/mlx_lm/models/switch_layers.pyi

@@ -57,6 +57,11 @@ class SwiGLU(nn.Module):
     def __call__(self, x, gate): ...
 
 class SwitchGLU(nn.Module):
+    gate_proj: SwitchLinear
+    up_proj: SwitchLinear
+    down_proj: SwitchLinear
+    activation: SwiGLU
+
     def __init__(
         self,
         input_dims: int,

.mlx_typings/mlx_lm/tokenizer_utils.pyi

@@ -4,6 +4,7 @@ This type stub file was generated by pyright.
 
 from functools import partial
 from pathlib import Path
+from typing import Any
 
 from transformers import PreTrainedTokenizerFast
 
@@ -38,12 +39,18 @@ class StreamingDetokenizer:
     """
 
     __slots__ = ...
-    def reset(self): ...
-    def add_token(self, token): ...
-    def finalize(self): ...
+    tokens: list[int]
+    def reset(self) -> None: ...
+    def add_token(self, token: int) -> None: ...
+    def finalize(self) -> None: ...
     @property
-    def last_segment(self):
+    def text(self) -> str:
+        """The full text decoded so far."""
+        ...
+    @property
+    def last_segment(self) -> str:
         """Return the last segment of readable text since last time this property was accessed."""
+        ...
 
 class NaiveStreamingDetokenizer(StreamingDetokenizer):
     """NaiveStreamingDetokenizer relies on the underlying tokenizer
@@ -103,37 +110,56 @@ class TokenizerWrapper:
     Accessing any attribute other than the ``detokenizer`` is forwarded to the
     huggingface tokenizer.
     """
-    def __init__(self, tokenizer, detokenizer_class=..., eos_token_ids=...) -> None: ...
-    def add_eos_token(self, token: str):  # -> None:
-        ...
-    @property
-    def has_thinking(self):  # -> bool:
-        ...
-    @property
-    def think_start(self):  # -> str | None:
-        ...
-    @property
-    def think_end(self):  # -> str | None:
-        ...
-    @property
-    def has_tool_calling(self):  # -> bool:
-        ...
-    @property
-    def tool_call_start(self):  # -> str | None:
-        ...
-    @property
-    def tool_call_end(self):  # -> str | None:
-        ...
-    @property
-    def detokenizer(self):  # -> NaiveStreamingDetokenizer:
-        """
-        Get a stateful streaming detokenizer.
-        """
 
-    def __getattr__(self, attr):  # -> set[Any] | Any:
-        ...
-    def __setattr__(self, attr, value):  # -> None:
-        ...
+    _tokenizer: PreTrainedTokenizerFast
+    eos_token_id: int | None
+    eos_token: str | None
+    eos_token_ids: list[int] | None
+    bos_token_id: int | None
+    bos_token: str | None
+    vocab_size: int
+    all_special_tokens: list[str]
+
+    def __init__(
+        self,
+        tokenizer: Any,
+        detokenizer_class: Any = ...,
+        eos_token_ids: list[int] | None = ...,
+        chat_template: Any = ...,
+        tool_parser: Any = ...,
+        tool_call_start: str | None = ...,
+        tool_call_end: str | None = ...,
+    ) -> None: ...
+    def encode(self, text: str, **kwargs: Any) -> list[int]: ...
+    def decode(self, token_ids: list[int], **kwargs: Any) -> str: ...
+    def apply_chat_template(
+        self,
+        messages: list[dict[str, Any]],
+        tokenize: bool = False,
+        add_generation_prompt: bool = False,
+        tools: Any = None,
+        **kwargs: Any,
+    ) -> str: ...
+    def get_vocab(self) -> dict[str, int]: ...
+    def add_eos_token(self, token: str) -> None: ...
+    @property
+    def has_thinking(self) -> bool: ...
+    @property
+    def think_start(self) -> str | None: ...
+    @property
+    def think_end(self) -> str | None: ...
+    @property
+    def has_tool_calling(self) -> bool: ...
+    @property
+    def tool_call_start(self) -> str | None: ...
+    @property
+    def tool_call_end(self) -> str | None: ...
+    @property
+    def detokenizer(self) -> NaiveStreamingDetokenizer:
+        """Get a stateful streaming detokenizer."""
+
+    def __getattr__(self, attr: str) -> Any: ...
+    def __setattr__(self, attr: str, value: Any) -> None: ...
 
 class NewlineTokenizer(PreTrainedTokenizerFast):
     """A tokenizer that replaces newlines with <n> and <n> with new line."""
@@ -146,18 +172,11 @@ class NewlineTokenizer(PreTrainedTokenizerFast):
     def batch_decode(self, *args, **kwargs):  # -> list[str]:
         ...
 
-def load_tokenizer(
+def load(
     model_path: Path,
-    tokenizer_config_extra=...,
-    return_tokenizer=...,
-    eos_token_ids=...,
-) -> (
-    TokenizerWrapper
-    | type[SPMStreamingDetokenizer]
-    | partial[SPMStreamingDetokenizer]
-    | type[BPEStreamingDetokenizer]
-    | type[NaiveStreamingDetokenizer]
-):
+    tokenizer_config_extra: dict[str, Any] | None = None,
+    eos_token_ids: list[int] | int | None = None,
+) -> TokenizerWrapper:
     """Load a huggingface tokenizer and try to infer the type of streaming
     detokenizer to use.
 
@@ -165,4 +184,7 @@ def load_tokenizer(
     a Hugging Face repo ID.
     """
 
-def no_bos_or_eos(sequence: list, bos: int, eos: int) -> list: ...
+# Alias for backward compatibility
+load_tokenizer = load
+
+def no_bos_or_eos(sequence: list[int], bos: int, eos: int) -> list[int]: ...

AGENTS.md

@@ -0,0 +1,96 @@
+# AGENTS.md
+
+This file provides guidance to AI coding agents when working with code in this repository.
+
+## Project Overview
+
+exo is a distributed AI inference system that connects multiple devices into a cluster. It enables running large language models across multiple machines using MLX as the inference backend and libp2p for peer-to-peer networking.
+
+## Build & Run Commands
+
+```bash
+# Build the dashboard (required before running exo)
+cd dashboard && npm install && npm run build && cd ..
+
+# Run exo (starts both master and worker with API at http://localhost:52415)
+uv run exo
+
+# Run with verbose logging
+uv run exo -v   # or -vv for more verbose
+
+# Run tests (excludes slow tests by default)
+uv run pytest
+
+# Run all tests including slow tests
+uv run pytest -m ""
+
+# Run a specific test file
+uv run pytest src/exo/shared/tests/test_election.py
+
+# Run a specific test function
+uv run pytest src/exo/shared/tests/test_election.py::test_function_name
+
+# Type checking (strict mode)
+uv run basedpyright
+
+# Linting
+uv run ruff check
+
+# Format code (using nix)
+nix fmt
+```
+
+## Architecture
+
+### Node Composition
+A single exo `Node` (src/exo/main.py) runs multiple components:
+- **Router**: libp2p-based pub/sub messaging via Rust bindings (exo_pyo3_bindings)
+- **Worker**: Handles inference tasks, downloads models, manages runner processes
+- **Master**: Coordinates cluster state, places model instances across nodes
+- **Election**: Bully algorithm for master election
+- **API**: FastAPI server for OpenAI-compatible chat completions
+
+### Message Flow
+Components communicate via typed pub/sub topics (src/exo/routing/topics.py):
+- `GLOBAL_EVENTS`: Master broadcasts indexed events to all workers
+- `LOCAL_EVENTS`: Workers send events to master for indexing
+- `COMMANDS`: Workers/API send commands to master
+- `ELECTION_MESSAGES`: Election protocol messages
+- `CONNECTION_MESSAGES`: libp2p connection updates
+
+### Event Sourcing
+The system uses event sourcing for state management:
+- `State` (src/exo/shared/types/state.py): Immutable state object
+- `apply()` (src/exo/shared/apply.py): Pure function that applies events to state
+- Master indexes events and broadcasts; workers apply indexed events
+
+### Key Type Hierarchy
+- `src/exo/shared/types/`: Pydantic models for all shared types
+  - `events.py`: Event types (discriminated union)
+  - `commands.py`: Command types
+  - `tasks.py`: Task types for worker execution
+  - `state.py`: Cluster state model
+
+### Rust Components
+Rust code in `rust/` provides:
+- `networking`: libp2p networking (gossipsub, peer discovery)
+- `exo_pyo3_bindings`: PyO3 bindings exposing Rust to Python
+- `system_custodian`: System-level operations
+
+### Dashboard
+Svelte 5 + TypeScript frontend in `dashboard/`. Build output goes to `dashboard/build/` and is served by the API.
+
+## Code Style Requirements
+
+From .cursorrules:
+- Strict, exhaustive typing - never bypass the type-checker
+- Use `Literal[...]` for enum-like sets, `typing.NewType` for primitives
+- Pydantic models with `frozen=True` and `strict=True`
+- Pure functions with injectable effect handlers for side-effects
+- Descriptive names - no abbreviations or 3-letter acronyms
+- Catch exceptions only where you can handle them meaningfully
+- Use `@final` and immutability wherever applicable
+
+## Testing
+
+Tests use pytest-asyncio with `asyncio_mode = "auto"`. Tests are in `tests/` subdirectories alongside the code they test. The `EXO_TESTS=1` env var is set during tests.

CLAUDE.md

@@ -0,0 +1 @@
+AGENTS.md

MISSED_THINGS.md

@@ -0,0 +1,41 @@
+# Missed things
+[X] Log EXO_LIBP2P_NAMESPACE on start in exo/main.py
+[X] Ordering of warmup was changed, which is wrong. It was changed to rank < n-1, then rank=n-1. It should be rank!=0 then rank=0 (this matches the auto_parallel implementation. NOTE: we use a different convention to mlx-lm, our terminal rank is rank=n-1 whereas mlx-lm is rank=0 hence i can see why this was changed wrongly).
+[X] Downloads keying by model_id not shard_metadata (worker/plan.py, worker/main.py).
+[X] Fetching download status of all models on start
+[X] Deduplication of tasks in plan_step.
+[X] resolve_allow_patterns should just be wildcard now.
+[] no mx_barrier in genreate.py mlx_generate at the end.
+[] cache assertion not needed in auto_parallel.py PipelineLastLayer.
+[] GPTOSS support dropped in auto_parallel.py.
+[] sharding changed "all-to-sharded" became _all_to_sharded in auto_parallel.py.
+[] same as above with "sharded-to-all" became _sharded_to_all in auto_parallel.py.
+[] Dropped support for Ministral3Model, DeepseekV32Model, Glm4MoeModel, Qwen3NextModel, GptOssMode in auto_parallel.py.
+[] Dropped prefill/decode code in auto_parallel.py and utils_mlx.py.
+[X] KV_CACHE_BITS should be None to disable quantized KV cache.
+[] Dropped _set_nofile_limit in utils_mlx.py.
+[] We have group optional in load_mlx_items in utils_mlx.py.
+[] Dropped add_missing_chat_templates for GptOss in load_mlx_items in utils_mlx.py.
+[] Dropped model.make_cache in make_kv_cache in utils_mlx.py.
+[X] We put cache limit back in utils_mlx.py.
+[] topology.py remove_node removes the connections after checking if node is is in self._node_id_to_rx_id_map. on beta_1 it checks after, so would remove stale connections I guess?
+[] Missing Glm 4.7 model cards (this isn't ready yet but should be picked up, probably create an issue... the blocker is transforemrs version doesn't support the tokenizer for Glm 4.7. rc-1 does but we can't upgrade as it breaks other things.)
+[] try-except in _command_processor only excepts ValueError. This was silently failing leading to un-debuggable errors (we had a KeyError that was happening ). Changed this to catch Exception instead of ValueError. See exo-v2 89ae38405e0052e3c22405daf094b065878aa873 and fb99fea69b5a39017efc90c5dad0072e677455f0.
+[X] In placement.py, place_instance no longer looks at model_meta.supports_tensor and check if this tensor parallel number of nodes is supported by the model's tensor dimensions.
+[X] In placement.py, place_instanec, we no longer have the special case to exclude DeepSeek v3.1 pipeline parallel (it doesn't work).
+[] logger.warning("You have likely selected ibv for a single node instance; falling back to MlxRing") was changed to debug. That will spam this warning since it happens every time we query instance previews.
+[X] In placement_utils.py, get_mlx_jaccl_coordinators, We no longer prioritise Jaccl Coordinator IP. Now it picks the first one, which is unstable (Jaccl coordinator over TB5 is unstable).
+
+
+
+[X] Downloads keying by model_id not shard_metadata (worker/plan.py, worker/main.py).
+[X] Fetching download status of all models on start
+[X] Deduplication of tasks in plan_step.
+[X] resolve_allow_patterns should just be wildcard now.
+[X] KV_CACHE_BITS should be None to disable quantized KV cache.
+[X] We put cache limit back in utils_mlx.py.
+[X] In placement.py, place_instance no longer looks at model_meta.supports_tensor and check if this tensor parallel number of nodes is supported by the model's tensor dimensions.
+[X] In placement.py, place_instanec, we no longer have the special case to exclude DeepSeek v3.1 pipeline parallel (it doesn't work).
+[X] In placement_utils.py, get_mlx_jaccl_coordinators, We no longer prioritise Jaccl Coordinator IP. Now it picks the first one, which is unstable (Jaccl coordinator over TB5 is unstable).
+
+

flake.lock

@@ -21,21 +21,23 @@
         "type": "github"
       }
     },
-    "flake-utils": {
+    "flake-parts": {
       "inputs": {
-        "systems": "systems"
+        "nixpkgs-lib": [
+          "nixpkgs"
+        ]
       },
       "locked": {
-        "lastModified": 1731533236,
-        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
-        "owner": "numtide",
-        "repo": "flake-utils",
-        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "lastModified": 1768135262,
+        "narHash": "sha256-PVvu7OqHBGWN16zSi6tEmPwwHQ4rLPU9Plvs8/1TUBY=",
+        "owner": "hercules-ci",
+        "repo": "flake-parts",
+        "rev": "80daad04eddbbf5a4d883996a73f3f542fa437ac",
         "type": "github"
       },
       "original": {
-        "owner": "numtide",
-        "repo": "flake-utils",
+        "owner": "hercules-ci",
+        "repo": "flake-parts",
         "type": "github"
       }
     },
@@ -58,7 +60,7 @@
     "root": {
       "inputs": {
         "fenix": "fenix",
-        "flake-utils": "flake-utils",
+        "flake-parts": "flake-parts",
         "nixpkgs": "nixpkgs",
         "treefmt-nix": "treefmt-nix"
       }
@@ -80,21 +82,6 @@
         "type": "github"
       }
     },
-    "systems": {
-      "locked": {
-        "lastModified": 1681028828,
-        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
-        "owner": "nix-systems",
-        "repo": "default",
-        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
-        "type": "github"
-      },
-      "original": {
-        "owner": "nix-systems",
-        "repo": "default",
-        "type": "github"
-      }
-    },
     "treefmt-nix": {
       "inputs": {
         "nixpkgs": [

flake.nix

@@ -3,13 +3,17 @@
 
   inputs = {
     nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
-    flake-utils.url = "github:numtide/flake-utils";
-    # Provides Rust dev-env integration:
+
+    flake-parts = {
+      url = "github:hercules-ci/flake-parts";
+      inputs.nixpkgs-lib.follows = "nixpkgs";
+    };
+
     fenix = {
       url = "github:nix-community/fenix";
       inputs.nixpkgs.follows = "nixpkgs";
     };
-    # Provides formatting infrastructure:
+
     treefmt-nix = {
       url = "github:numtide/treefmt-nix";
       inputs.nixpkgs.follows = "nixpkgs";
@@ -17,117 +21,113 @@
   };
 
   nixConfig = {
-    # nix community cachix
     extra-trusted-public-keys = "exo.cachix.org-1:okq7hl624TBeAR3kV+g39dUFSiaZgLRkLsFBCuJ2NZI=";
     extra-substituters = "https://exo.cachix.org";
   };
 
   outputs =
     inputs:
-    let
+    inputs.flake-parts.lib.mkFlake { inherit inputs; } {
       systems = [
         "x86_64-linux"
         "aarch64-darwin"
         "aarch64-linux"
       ];
-      fenixToolchain = system: inputs.fenix.packages.${system}.complete;
-    in
-    inputs.flake-utils.lib.eachSystem systems (
-      system:
-      let
-        pkgs = import inputs.nixpkgs {
-          inherit system;
-          overlays = [ inputs.fenix.overlays.default ];
-        };
-        treefmtEval = inputs.treefmt-nix.lib.evalModule pkgs {
-          projectRootFile = "flake.nix";
-          programs = {
-            nixpkgs-fmt.enable = true;
-            ruff-format = {
-              enable = true;
-              excludes = [ "rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi" ];
+
+      imports = [
+        inputs.treefmt-nix.flakeModule
+      ];
+
+      perSystem =
+        { config, inputs', pkgs, lib, ... }:
+        let
+          fenixToolchain = inputs'.fenix.packages.complete;
+        in
+        {
+          treefmt = {
+            projectRootFile = "flake.nix";
+            programs = {
+              nixpkgs-fmt.enable = true;
+              ruff-format = {
+                enable = true;
+                excludes = [ "rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi" ];
+              };
+              rustfmt = {
+                enable = true;
+                package = fenixToolchain.rustfmt;
+              };
+              prettier = {
+                enable = true;
+                includes = [ "*.ts" ];
+              };
+              swift-format.enable = true;
             };
-            rustfmt = {
-              enable = true;
-              package = (fenixToolchain system).rustfmt;
-            };
-            prettier = {
-              enable = true;
-              includes = [ "*.ts" ];
-            };
-            swift-format.enable = true;
           };
-        };
-      in
-      {
-        formatter = treefmtEval.config.build.wrapper;
-        checks.formatting = treefmtEval.config.build.check inputs.self;
-        checks.lint = pkgs.runCommand "lint-check" { } ''
-          export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
-          ${pkgs.ruff}/bin/ruff check ${inputs.self}/
-          touch $out
-        '';
 
-        devShells.default = pkgs.mkShell {
-          packages =
-            with pkgs;
-            [
-              # FORMATTING
-              treefmtEval.config.build.wrapper
-
-              # PYTHON
-              python313
-              uv
-              ruff
-              basedpyright
-
-              # RUST
-              ((fenixToolchain system).withComponents [
-                "cargo"
-                "rustc"
-                "clippy"
-                "rustfmt"
-                "rust-src"
-              ])
-              rustup # Just here to make RustRover happy
-
-              # NIX
-              nixpkgs-fmt
-
-              # SVELTE
-              nodejs
-
-              # MISC
-              just
-              jq
-            ]
-            ++ (pkgs.lib.optionals pkgs.stdenv.isLinux [
-              # IFCONFIG
-              unixtools.ifconfig
-
-              # Build dependencies for Linux
-              pkg-config
-              openssl
-            ])
-            ++ (pkgs.lib.optionals pkgs.stdenv.isDarwin [
-              # MACMON
-              macmon
-            ]);
-
-          shellHook = ''
-            # PYTHON
-            export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${pkgs.python313}/lib"
-            ${pkgs.lib.optionalString pkgs.stdenv.isLinux ''
-              # Build environment for Linux
-              export PKG_CONFIG_PATH="${pkgs.openssl.dev}/lib/pkgconfig:$PKG_CONFIG_PATH"
-              export LD_LIBRARY_PATH="${pkgs.openssl.out}/lib:$LD_LIBRARY_PATH"
-            ''}
-            echo
-            echo "🍎🍎 Run 'just <recipe>' to get started"
-            just --list
+          checks.lint = pkgs.runCommand "lint-check" { } ''
+            export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
+            ${pkgs.ruff}/bin/ruff check ${inputs.self}/
+            touch $out
           '';
 
+          devShells.default = with pkgs; pkgs.mkShell {
+            packages =
+              [
+                # FORMATTING
+                config.treefmt.build.wrapper
+
+                # PYTHON
+                python313
+                uv
+                ruff
+                basedpyright
+
+                # RUST
+                (fenixToolchain.withComponents [
+                  "cargo"
+                  "rustc"
+                  "clippy"
+                  "rustfmt"
+                  "rust-src"
+                ])
+                rustup # Just here to make RustRover happy
+
+                # NIX
+                nixpkgs-fmt
+
+                # SVELTE
+                nodejs
+
+                # MISC
+                just
+                jq
+              ]
+              ++ (pkgs.lib.optionals pkgs.stdenv.isLinux [
+                # IFCONFIG
+                unixtools.ifconfig
+
+                # Build dependencies for Linux
+                pkg-config
+                openssl
+              ])
+              ++ (pkgs.lib.optionals pkgs.stdenv.isDarwin [
+                # MACMON
+                macmon
+              ]);
+
+            shellHook = ''
+              # PYTHON
+              export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${pkgs.python313}/lib"
+              ${lib.optionalString pkgs.stdenv.isLinux ''
+                # Build environment for Linux
+                export PKG_CONFIG_PATH="${pkgs.openssl.dev}/lib/pkgconfig:$PKG_CONFIG_PATH"
+                export LD_LIBRARY_PATH="${pkgs.openssl.out}/lib:$LD_LIBRARY_PATH"
+              ''}
+              echo
+              echo "🍎🍎 Run 'just <recipe>' to get started"
+              just --list
+            '';
+          };
         };
-      }
-    );
+    };
 }

pyproject.toml

@@ -17,9 +17,9 @@ dependencies = [
     "loguru>=0.7.3",
     "exo_pyo3_bindings", # rust bindings
     "anyio==4.11.0",
-    "mlx>=0.30.1; sys_platform == 'darwin'",
-    "mlx[cpu]>=0.30.1; sys_platform == 'linux'",
-    "mlx-lm>=0.28.3",
+    "mlx==0.30.1; sys_platform == 'darwin'",
+    "mlx[cpu]==0.30.1; sys_platform == 'linux'",
+    "mlx-lm @ git+https://github.com/AlexCheema/mlx-lm.git@fix-transformers-5.0.0rc2",
     "tiktoken>=0.12.0", # required for kimi k2 tokenizer
     "hypercorn>=0.18.0",
     "openai-harmony>=0.0.8",
@@ -33,6 +33,7 @@ exo = "exo.main:main"
 # dependencies only required for development
 [dependency-groups]
 dev = [
+    "basedpyright>=1.29.0",
     "pyinstaller>=6.17.0",
     "pytest>=8.4.0",
     "pytest-asyncio>=1.0.0",
@@ -98,6 +99,7 @@ root = "src"
 
 # supported platforms for this project
 [tool.uv]
+prerelease = "allow"
 environments = [
     "sys_platform == 'darwin'",
     "sys_platform == 'linux'",

src/exo/main.py

@@ -1,6 +1,7 @@
 import argparse
 import multiprocessing as mp
 import os
+import resource
 import signal
 from dataclasses import dataclass, field
 from typing import Self
@@ -195,6 +196,8 @@ class Node:
 
 def main():
     args = Args.parse()
+    soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+    resource.setrlimit(resource.RLIMIT_NOFILE, (max(soft, 65535), hard))
 
     mp.set_start_method("spawn")
     # TODO: Refactor the current verbosity system

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

@@ -50,7 +50,9 @@ class RunnerReady(BaseRunnerStatus):
 
 
 class RunnerRunning(BaseRunnerStatus):
-    pass
+    """Runner is processing requests and can accept more (continuous batching)."""
+
+    active_requests: int = 0
 
 
 class RunnerShuttingDown(BaseRunnerStatus):

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

@@ -10,18 +10,23 @@ from mlx.nn.layers.distributed import (
     shard_linear,
     sum_gradients,
 )
-from mlx_lm.models.cache import (
-    _BaseCache,  # pyright: ignore[reportPrivateUsage]
-)
 from mlx_lm.models.deepseek_v3 import DeepseekV3MLP
 from mlx_lm.models.deepseek_v3 import Model as DeepseekV3Model
+from mlx_lm.models.deepseek_v32 import DeepseekV32MLP
+from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
+from mlx_lm.models.glm4_moe import Model as Glm4MoeModel
+from mlx_lm.models.glm4_moe import MoE
+from mlx_lm.models.gpt_oss import GptOssMoeModel
+from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.models.llama import Model as LlamaModel
+from mlx_lm.models.ministral3 import Model as Ministral3Model
 from mlx_lm.models.qwen3_moe import Model as Qwen3MoeModel
 from mlx_lm.models.qwen3_moe import Qwen3MoeSparseMoeBlock
+from mlx_lm.models.qwen3_next import Model as Qwen3NextModel
+from mlx_lm.models.qwen3_next import Qwen3NextSparseMoeBlock
 
-from exo.shared.types.worker.shards import (
-    PipelineShardMetadata,
-)
+from exo.shared.logging import logger
+from exo.shared.types.worker.shards import PipelineShardMetadata
 
 
 class _LayerCallable(Protocol):
@@ -91,8 +96,6 @@ class PipelineLastLayer(CustomMlxLayer):
             x, *args, **kwargs
         ).arguments.get("cache", None)
 
-        assert cache is None or issubclass(type(cache), _BaseCache)  # type: ignore
-
         output: mx.array = self.original_layer(x, *args, **kwargs)
 
         if self.r != self.s - 1:
@@ -100,7 +103,6 @@ class PipelineLastLayer(CustomMlxLayer):
                 output, (self.r + 1) % self.s, group=self.group
             )
             if cache is not None:
-                # This change happened upstream - check out mlx github somewhere??
                 cache.keys = mx.depends(cache.keys, output)  # type: ignore[reportUnknownMemberType]
 
         output = mx.distributed.all_gather(output, group=self.group)[-output.shape[0] :]
@@ -132,24 +134,6 @@ def _get_layers(inner_model_instance: nn.Module) -> list[_LayerCallable]:
     return layers
 
 
-def _set_layers(model: nn.Module, layers: list[_LayerCallable]) -> None:
-    inner_model_instance = _inner_model(model)
-    if hasattr(inner_model_instance, "layers"):
-        inner_model_instance.layers = layers
-
-        # Update DeepSeek V3 specific parameters when layers are shrunk
-        if isinstance(model, DeepseekV3Model) and hasattr(
-            inner_model_instance, "num_layers"
-        ):
-            inner_model_instance.start_idx = 0
-            inner_model_instance.end_idx = len(layers)
-            inner_model_instance.num_layers = len(layers)
-    elif hasattr(inner_model_instance, "h"):
-        inner_model_instance.h = layers
-    else:
-        raise ValueError("Model must have either a 'layers' or 'h' attribute")
-
-
 def pipeline_auto_parallel(
     model: nn.Module,
     group: mx.distributed.Group,
@@ -165,8 +149,7 @@ def pipeline_auto_parallel(
     """
     inner_model_instance: nn.Module = _inner_model(model)
 
-    # Handle both model.layers and model.h cases
-    layers: list[_LayerCallable] = _get_layers(inner_model_instance)
+    layers = _get_layers(inner_model_instance)
 
     start_layer, end_layer = model_shard_meta.start_layer, model_shard_meta.end_layer
     device_rank, world_size = model_shard_meta.device_rank, model_shard_meta.world_size
@@ -180,6 +163,17 @@ def pipeline_auto_parallel(
         group=group,
     )
 
+    if isinstance(inner_model_instance, GptOssMoeModel):
+        inner_model_instance.layer_types = inner_model_instance.layer_types[  # type: ignore
+            start_layer:end_layer
+        ]
+        inner_model_instance.swa_idx = inner_model_instance.layer_types.index(  # type: ignore
+            "sliding_attention"
+        )
+        inner_model_instance.ga_idx = inner_model_instance.layer_types.index(  # type: ignore
+            "full_attention"
+        )
+
     _set_layers(model, layers)
 
     assert isinstance(layers, list), (
@@ -204,18 +198,44 @@ def tensor_auto_parallel(
         group=group,
     )
 
+    segments: int = 1
+
+    def _all_to_sharded(path: str, weight: mx.array):
+        if path.endswith("bias"):
+            logger.info(f"Sharding bias for {path} - all to sharded")
+            return weight.ndim - 1, segments
+        return max(weight.ndim - 2, 0), segments
+
     all_to_sharded_linear_in_place = partial(
         shard_inplace,
-        sharding="all-to-sharded",
-        group=group,
-    )
-    sharded_to_all_linear_in_place = partial(
-        shard_inplace,
-        sharding="sharded-to-all",
+        sharding=_all_to_sharded,  # type: ignore
         group=group,
     )
 
-    if isinstance(model, LlamaModel):
+    n = group.size()
+
+    def _sharded_to_all(path: str, weight: mx.array):
+        if path.endswith("bias"):
+            logger.info(f"Sharding bias for {path} - sharded to all")
+            weight /= n
+            return None
+        return -1, segments
+
+    sharded_to_all_linear_in_place = partial(
+        shard_inplace,
+        sharding=_sharded_to_all,  # type: ignore
+        group=group,
+    )
+
+    if hasattr(model, "shard"):
+        try:
+            model.shard(group)  # type: ignore
+            return model
+        except (AttributeError, TypeError, NameError):
+            pass
+
+    if isinstance(model, (LlamaModel, Ministral3Model)):
+        logger.warning("shouldn't be hit - upstream sharding exists")
         tensor_parallel_sharding_strategy = LlamaShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -223,7 +243,8 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
-    elif isinstance(model, DeepseekV3Model):
+    elif isinstance(model, (DeepseekV3Model, DeepseekV32Model)):
+        logger.warning("shouldn't be hit - upstream sharding exists")
         tensor_parallel_sharding_strategy = DeepSeekShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -231,7 +252,7 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
-    elif isinstance(model, Qwen3MoeModel):
+    elif isinstance(model, (Qwen3MoeModel, Glm4MoeModel, Qwen3NextModel)):
         tensor_parallel_sharding_strategy = QwenShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -239,6 +260,15 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
+    elif isinstance(model, GptOssModel):
+        tensor_parallel_sharding_strategy = GptOssShardingStrategy(
+            group,
+            all_to_sharded_linear,
+            sharded_to_all_linear,
+            all_to_sharded_linear_in_place,
+            sharded_to_all_linear_in_place,
+        )
+
     else:
         raise ValueError(f"Unsupported model type: {type(model)}")
 
@@ -284,13 +314,38 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
         return model
 
 
+def _set_layers(model: nn.Module, layers: list[_LayerCallable]) -> None:
+    inner_model_instance = _inner_model(model)
+    if hasattr(inner_model_instance, "layers"):
+        inner_model_instance.layers = layers
+
+        # Update DeepSeek V3 specific parameters when layers are shrunk
+        if isinstance(
+            model, (DeepseekV3Model, DeepseekV32Model, Glm4MoeModel)
+        ) and hasattr(inner_model_instance, "num_layers"):
+            logger.info(
+                f"Setting num_layers to {len(layers)} for model {model.model.__class__.__name__}"
+            )
+            inner_model_instance.start_idx = 0
+            inner_model_instance.end_idx = len(layers)
+            inner_model_instance.num_layers = len(layers)
+        elif isinstance(model, Qwen3MoeModel):
+            logger.info(
+                f"Setting num_hidden_layers to {len(layers)} for model {model.model.__class__.__name__}"
+            )
+            inner_model_instance.num_hidden_layers = len(layers)
+    elif hasattr(inner_model_instance, "h"):
+        inner_model_instance.h = layers
+    else:
+        raise ValueError("Model must have either a 'layers' or 'h' attribute")
+
+
 class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
     def shard_model(self, model: nn.Module) -> nn.Module:
         model = cast(DeepseekV3Model, model)
         for layer in model.layers:
             # Shard the self attention
-            if layer.self_attn.q_lora_rank is None:  # pyright: ignore[reportUnnecessaryComparison]
-                # Unfortunately, q_lora_rank can be None despite typing hints.
+            if layer.self_attn.q_lora_rank is None:
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
                     layer.self_attn.q_proj
                 )
@@ -305,7 +360,7 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
             layer.self_attn.num_heads //= self.N
 
             # Shard the MLP
-            if isinstance(layer.mlp, DeepseekV3MLP):
+            if isinstance(layer.mlp, (DeepseekV3MLP, DeepseekV32MLP)):
                 layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj)
                 layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
                 layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
@@ -353,7 +408,9 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
 
             # Shard the MoE. Shard in place since the MoE should be responsible
             # for aggregating the results.
-            if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock):
+            if isinstance(
+                layer.mlp, (Qwen3MoeSparseMoeBlock, MoE, Qwen3NextSparseMoeBlock)
+            ):
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
                 self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
                 self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
@@ -381,3 +438,50 @@ class ShardedQwenMoE(CustomMlxLayer):
         if self.sharding_group is not None:
             y = mx.distributed.all_sum(y, group=self.sharding_group)
         return y
+
+
+class GptOssShardingStrategy(TensorParallelShardingStrategy):
+    def shard_model(self, model: nn.Module) -> nn.Module:
+        model = cast(GptOssMoeModel, model)
+
+        for layer in model.layers:
+            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
+            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
+            layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
+            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
+
+            layer.self_attn.num_attention_heads //= self.N
+            layer.self_attn.num_key_value_heads //= self.N
+            layer.self_attn.num_key_value_groups = (
+                layer.self_attn.num_attention_heads
+                // layer.self_attn.num_key_value_heads
+            )
+
+            layer.self_attn.sinks = layer.self_attn.sinks[
+                layer.self_attn.num_attention_heads
+                * self.group.rank() : layer.self_attn.num_attention_heads
+                * (self.group.rank() + 1)
+            ]
+
+            self.all_to_sharded_linear_in_place(layer.mlp.experts.gate_proj)
+            self.sharded_to_all_linear_in_place(layer.mlp.experts.down_proj)
+            self.all_to_sharded_linear_in_place(layer.mlp.experts.up_proj)
+
+            layer.mlp = ShardedGptOssMoE(layer.mlp)  # type: ignore
+            layer.mlp.sharding_group = self.group
+
+        return model
+
+
+class ShardedGptOssMoE(CustomMlxLayer):
+    def __init__(self, layer: nn.Module):
+        super().__init__(layer)
+        self.sharding_group: mx.distributed.Group | None = None
+
+    def __call__(self, x: mx.array) -> mx.array:
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+        y = self.original_layer(x)
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+        return y

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

@@ -0,0 +1,251 @@
+"""Batch generation engine using mlx_lm's BatchGenerator for continuous batching."""
+
+import time
+from dataclasses import dataclass, field
+
+import mlx.core as mx
+from mlx_lm.generate import BatchGenerator
+from mlx_lm.sample_utils import make_sampler
+from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper
+
+from exo.shared.types.api import FinishReason, GenerationStats
+from exo.shared.types.common import CommandId
+from exo.shared.types.memory import Memory
+from exo.shared.types.tasks import ChatCompletionTaskParams, TaskId
+from exo.shared.types.worker.runner_response import GenerationResponse
+from exo.worker.engines.mlx import Model
+from exo.worker.engines.mlx.constants import MAX_TOKENS
+from exo.worker.engines.mlx.generator.distributed_sync import share_object
+from exo.worker.engines.mlx.utils_mlx import apply_chat_template
+from exo.worker.runner.bootstrap import logger
+
+
+@dataclass
+class ActiveRequest:
+    """Tracks an active request in the batch."""
+
+    command_id: CommandId
+    task_id: TaskId
+    uid: int  # BatchGenerator's internal ID
+    detokenizer: StreamingDetokenizer
+    tokens_generated: int = 0
+    prompt_tokens: int = 0
+    start_time: float = field(default_factory=time.perf_counter)
+
+
+@dataclass
+class BatchedGenerationResponse:
+    """Response from batch engine, tagged with command_id and task_id."""
+
+    command_id: CommandId
+    task_id: TaskId
+    response: GenerationResponse
+
+
+class BatchGenerationEngine:
+    """Manages continuous batching using mlx_lm's BatchGenerator."""
+
+    def __init__(
+        self,
+        model: Model,
+        tokenizer: TokenizerWrapper,
+        group: mx.distributed.Group | None = None,
+        max_tokens: int = MAX_TOKENS,
+        completion_batch_size: int = 32,
+        prefill_batch_size: int = 8,
+        prefill_step_size: int = 2048,
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.max_tokens = max_tokens
+        self.active_requests: dict[int, ActiveRequest] = {}
+        self._pending_inserts: list[tuple[CommandId, TaskId, ChatCompletionTaskParams]] = []
+
+        self.group = group
+        self.rank = group.rank() if group else 0
+        self.is_distributed = group is not None and group.size() > 1
+
+        sampler = make_sampler(temp=0.7, top_p=1.0)
+
+        eos_tokens: set[int] = set(tokenizer.eos_token_ids or [])
+
+        self.batch_gen: BatchGenerator = BatchGenerator(
+            model=model,
+            max_tokens=max_tokens,
+            stop_tokens=eos_tokens,
+            sampler=sampler,
+            completion_batch_size=completion_batch_size,
+            prefill_batch_size=prefill_batch_size,
+            prefill_step_size=prefill_step_size,
+        )
+
+        logger.info(
+            f"BatchGenerationEngine initialized with completion_batch_size={completion_batch_size}, "
+            f"prefill_batch_size={prefill_batch_size}, distributed={self.is_distributed}"
+        )
+
+    def queue_request(
+        self,
+        command_id: CommandId,
+        task_id: TaskId,
+        task_params: ChatCompletionTaskParams,
+    ) -> None:
+        """Queue a request for insertion. Only rank 0 should call this.
+
+        In distributed mode, rank 0 receives tasks from the control plane and
+        queues them here. The actual insertion happens in sync_and_insert_pending()
+        which ensures all ranks insert the same requests together.
+        """
+        assert self.rank == 0, "Only rank 0 should queue requests"
+        self._pending_inserts.append((command_id, task_id, task_params))
+        logger.info(f"Queued request {command_id} for insertion (pending={len(self._pending_inserts)})")
+
+    def sync_and_insert_pending(self) -> list[int]:
+        """Sync pending inserts across ranks and insert them. Returns UIDs.
+
+        This method ensures all ranks insert the same requests in the same order.
+        In non-distributed mode, it simply inserts all pending requests.
+        In distributed mode, it broadcasts pending requests from rank 0 to all ranks.
+        """
+        inserts_to_process: list[tuple[CommandId, TaskId, ChatCompletionTaskParams]]
+
+        if not self.is_distributed:
+            # Non-distributed: just insert directly from pending
+            inserts_to_process = list(self._pending_inserts)
+        else:
+            # Distributed: broadcast pending inserts from rank 0 to all ranks
+            assert self.group is not None
+            pending_data = self._pending_inserts if self.rank == 0 else None
+            synced_data = share_object(pending_data, self.rank, self.group)
+
+            if synced_data is None:
+                self._pending_inserts.clear()
+                return []
+
+            inserts_to_process = synced_data
+
+        uids: list[int] = []
+        for cmd_id, task_id, params in inserts_to_process:
+            uid = self._do_insert(cmd_id, task_id, params)
+            uids.append(uid)
+
+        self._pending_inserts.clear()
+        return uids
+
+    def _do_insert(
+        self,
+        command_id: CommandId,
+        task_id: TaskId,
+        task_params: ChatCompletionTaskParams,
+    ) -> int:
+        """Actually insert a request into BatchGenerator. No sync - called after sync."""
+        prompt_str = apply_chat_template(self.tokenizer, task_params)
+        tokens: list[int] = self.tokenizer.encode(prompt_str, add_special_tokens=False)
+        prompt_tokens = len(tokens)
+        max_tokens = task_params.max_tokens or self.max_tokens
+
+        uids = self.batch_gen.insert([tokens], max_tokens=[max_tokens])
+        uid = uids[0]
+        detokenizer = self.tokenizer.detokenizer
+
+        self.active_requests[uid] = ActiveRequest(
+            command_id=command_id,
+            task_id=task_id,
+            uid=uid,
+            detokenizer=detokenizer,
+            prompt_tokens=prompt_tokens,
+        )
+
+        logger.info(f"Inserted request {command_id} with uid={uid}, prompt_tokens={prompt_tokens}, max_tokens={max_tokens}")
+        return uid
+
+    def step(self) -> list[BatchedGenerationResponse]:
+        """Run one decode step. Syncs completed UIDs across ranks if distributed."""
+        responses = self.batch_gen.next()
+        if not responses:
+            return []
+
+        results: list[BatchedGenerationResponse] = []
+        uids_to_remove: list[int] = []
+
+        for r in responses:
+            uid: int = r.uid
+            req = self.active_requests.get(uid)
+            if req is None:
+                logger.warning(f"Received response for unknown uid={uid}")
+                continue
+
+            req.tokens_generated += 1
+
+            # Decode the token
+            token: int = r.token
+            req.detokenizer.add_token(token)
+            text: str = req.detokenizer.last_segment
+
+            stats: GenerationStats | None = None
+            finish_reason: FinishReason | None = None
+
+            raw_finish_reason: str | None = r.finish_reason
+            if raw_finish_reason is not None:
+                # Finalize to get remaining text
+                req.detokenizer.finalize()
+                text = req.detokenizer.last_segment
+
+                elapsed = time.perf_counter() - req.start_time
+                generation_tps = req.tokens_generated / elapsed if elapsed > 0 else 0.0
+
+                stats = GenerationStats(
+                    prompt_tps=0.0,  # Not tracked per-request in batch mode
+                    generation_tps=generation_tps,
+                    prompt_tokens=req.prompt_tokens,
+                    generation_tokens=req.tokens_generated,
+                    peak_memory_usage=Memory.from_gb(mx.get_peak_memory() / 1e9),
+                )
+
+                if raw_finish_reason == "stop":
+                    finish_reason = "stop"
+                elif raw_finish_reason == "length":
+                    finish_reason = "length"
+                else:
+                    logger.warning(f"Unknown finish_reason: {raw_finish_reason}")
+                    finish_reason = "stop"
+
+                uids_to_remove.append(uid)  # Sync before removal
+                logger.info(f"Request {req.command_id} completed: {req.tokens_generated} tokens, {generation_tps:.2f} tps, reason={finish_reason}")
+
+            results.append(BatchedGenerationResponse(
+                command_id=req.command_id,
+                task_id=req.task_id,
+                response=GenerationResponse(text=text, token=token, finish_reason=finish_reason, stats=stats),
+            ))
+
+        # Sync completed UIDs across ranks before removing
+        if self.is_distributed and uids_to_remove:
+            assert self.group is not None
+            uids_to_remove = share_object(uids_to_remove if self.rank == 0 else None, self.rank, self.group) or []
+
+        for uid in uids_to_remove:
+            if uid in self.active_requests:
+                del self.active_requests[uid]
+
+        return results
+
+    @property
+    def has_active_requests(self) -> bool:
+        return bool(self.active_requests or self.batch_gen.unprocessed_prompts)
+
+    @property
+    def has_pending_inserts(self) -> bool:
+        return bool(self._pending_inserts)
+
+    @property
+    def active_count(self) -> int:
+        return len(self.active_requests)
+
+    @property
+    def pending_count(self) -> int:
+        return len(self.batch_gen.unprocessed_prompts)
+
+    @property
+    def pending_insert_count(self) -> int:
+        return len(self._pending_inserts)

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

@@ -0,0 +1,73 @@
+"""Distributed sync utilities using mx.distributed.all_sum() to broadcast from rank 0."""
+
+# pyright: reportAny=false
+
+import pickle
+from enum import IntEnum
+from typing import TypeVar, cast
+
+import mlx.core as mx
+
+T = TypeVar("T")
+
+
+class DistributedOp(IntEnum):
+    """Operation codes for distributed synchronization.
+
+    Used to ensure all ranks in a distributed setup execute the same
+    operations in the same order, preventing collective mismatches.
+    """
+
+    NOOP = 0  # No operation - rank 0 will block waiting for a task
+    INSERT = 1  # Insert pending requests into the batch
+    STEP = 2  # Run a decode step
+    SHUTDOWN = 3  # All ranks should exit
+
+
+def sync_operation(
+    op: DistributedOp | None,
+    rank: int,
+    group: mx.distributed.Group,
+) -> DistributedOp:
+    """Broadcast operation code from rank 0 to all ranks.
+
+    This ensures all ranks execute the same operation, preventing
+    collective mismatches that would cause deadlocks.
+
+    Args:
+        op: The operation to perform (only rank 0's value is used)
+        rank: This process's rank
+        group: The distributed group
+
+    Returns:
+        The operation that all ranks should execute
+    """
+    if rank == 0:
+        assert op is not None, "Rank 0 must provide an operation"
+        code = mx.array([int(op)], dtype=mx.int32)
+    else:
+        code = mx.array([0], dtype=mx.int32)
+
+    result = mx.distributed.all_sum(code, group=group)
+    mx.eval(result)
+    return DistributedOp(int(result.item()))
+
+
+def share_object(obj: T | None, rank: int, group: mx.distributed.Group) -> T | None:
+    """Broadcast object from rank 0 to all ranks. Two-phase: size then data."""
+    if rank == 0:
+        if obj is None:
+            mx.eval(mx.distributed.all_sum(mx.array([0]), group=group))
+            return None
+        data = mx.array(list(pickle.dumps(obj)), dtype=mx.uint8)
+        mx.eval(mx.distributed.all_sum(mx.array([data.size]), group=group))
+        mx.eval(mx.distributed.all_sum(data, group=group))
+        return obj
+    else:
+        size = int(mx.distributed.all_sum(mx.array([0]), group=group).item())
+        if size == 0:
+            return None
+        data = mx.zeros(size, dtype=mx.uint8)
+        data = mx.distributed.all_sum(data, group=group)
+        mx.eval(data)
+        return cast(T, pickle.loads(bytes(cast(list[int], data.tolist()))))

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

@@ -1,10 +1,23 @@
 import json
 import os
 import resource
+import sys
 import time
 from pathlib import Path
 from typing import Any, cast
 
+# Monkey-patch for transformers 5.x compatibility
+# Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location
+# which was moved in transformers 5.0.0rc2
+try:
+    import transformers.models.gpt2.tokenization_gpt2 as gpt2_tokenization
+    from transformers.convert_slow_tokenizer import bytes_to_unicode
+
+    if not hasattr(gpt2_tokenization, "bytes_to_unicode"):
+        gpt2_tokenization.bytes_to_unicode = bytes_to_unicode  # type: ignore[attr-defined]
+except ImportError:
+    pass  # transformers < 5.0 or bytes_to_unicode not available
+
 from mlx_lm.models.cache import KVCache, QuantizedKVCache, RotatingKVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3Model
 from mlx_lm.tokenizer_utils import TokenizerWrapper
@@ -18,7 +31,7 @@ from exo.worker.engines.mlx.constants import (
 try:
     from mlx_lm.tokenizer_utils import load_tokenizer
 except ImportError:
-    from mlx_lm.tokenizer_utils import load as load_tokenizer  # type: ignore
+    from mlx_lm.tokenizer_utils import load as load_tokenizer
 import contextlib
 
 import mlx.core as mx
@@ -252,26 +265,70 @@ def shard_and_load(
     return model, tokenizer
 
 
-def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata):
-    # TODO: Let's move away from this custom logic to mlx_lm.load()
-    if "kimi-k2" in shard_metadata.model_meta.model_id.lower():
-        eos_token_ids = [163586]
+def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata) -> TokenizerWrapper:
+    """Load tokenizer for a model shard. Delegates to load_tokenizer_for_model_id."""
+    return load_tokenizer_for_model_id(shard_metadata.model_meta.model_id, model_path)
 
-    elif "glm" in shard_metadata.model_meta.model_id.lower():
-        eos_token_ids = [151336, 151329, 151338]
 
-    else:
-        eos_token_ids = None
+def get_eos_token_ids_for_model(model_id: str) -> list[int] | None:
+    """
+    Get the EOS token IDs for a model based on its ID.
 
-    tokenizer = cast(
-        TokenizerWrapper,
-        load_tokenizer(
-            model_path,
-            tokenizer_config_extra={"trust_remote_code": TRUST_REMOTE_CODE},
-            eos_token_ids=eos_token_ids,
-        ),
+    Some models require explicit EOS token configuration that isn't in their
+    tokenizer config. This function returns the known EOS token IDs for such models.
+
+    Args:
+        model_id: The HuggingFace model ID
+
+    Returns:
+        List of EOS token IDs, or None if the model uses standard tokenizer config
+    """
+    model_id_lower = model_id.lower()
+    if "kimi-k2" in model_id_lower:
+        return [163586]
+    elif "glm" in model_id_lower:
+        return [151336, 151329, 151338]
+    return None
+
+
+def load_tokenizer_for_model_id(model_id: str, model_path: Path) -> TokenizerWrapper:
+    """
+    Load tokenizer for a model given its ID and local path.
+
+    This is the core tokenizer loading logic, handling special cases for different
+    model families (Kimi, GLM, etc.) and transformers 5.x compatibility.
+
+    Args:
+        model_id: The HuggingFace model ID (e.g., "moonshotai/Kimi-K2-Instruct")
+        model_path: Local path where the model/tokenizer files are stored
+
+    Returns:
+        TokenizerWrapper instance configured for the model
+    """
+    model_id_lower = model_id.lower()
+    eos_token_ids = get_eos_token_ids_for_model(model_id)
+
+    # Kimi uses a custom TikTokenTokenizer that transformers 5.x can't load via AutoTokenizer
+    if "kimi-k2" in model_id_lower:
+        sys.path.insert(0, str(model_path))
+        from tokenization_kimi import TikTokenTokenizer  # type: ignore[import-not-found]  # noqa: I001
+
+        hf_tokenizer: Any = TikTokenTokenizer.from_pretrained(model_path)  # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType]
+
+        # Patch encode to use internal tiktoken model directly
+        # transformers 5.x has a bug in the encode->pad path for slow tokenizers
+        def _patched_encode(text: str, **_kwargs: object) -> list[int]:
+            # Pass allowed_special="all" to handle special tokens like <|im_user|>
+            return list(hf_tokenizer.model.encode(text, allowed_special="all"))  # pyright: ignore[reportUnknownMemberType,reportUnknownArgumentType]
+
+        hf_tokenizer.encode = _patched_encode
+        return TokenizerWrapper(hf_tokenizer, eos_token_ids=eos_token_ids)
+
+    tokenizer = load_tokenizer(
+        model_path,
+        tokenizer_config_extra={"trust_remote_code": TRUST_REMOTE_CODE},
+        eos_token_ids=eos_token_ids,
     )
-    assert isinstance(tokenizer, TokenizerWrapper)
 
     return tokenizer
 
@@ -301,14 +358,14 @@ def apply_chat_template(
             {k: v for k, v in message.model_dump().items() if v is not None}  # type: ignore
         )
 
-    prompt: str = tokenizer.apply_chat_template(  # type: ignore
+    prompt: str = tokenizer.apply_chat_template(
         formatted_messages,
         tokenize=False,
         add_generation_prompt=True,
         tools=chat_task_data.tools,
     )
 
-    return prompt  # type: ignore
+    return prompt
 
 
 class NullKVCache(KVCache):

src/exo/worker/plan.py

@@ -277,12 +277,14 @@ def _pending_tasks(
             # I have a design point here; this is a state race in disguise as the task status doesn't get updated to completed fast enough
             # however, realistically the task status should be set to completed by the LAST runner, so this is a true race
             # the actual solution is somewhat deeper than this bypass - TODO!
-            if task.task_id in runner.completed:
+            # Also skip tasks in pending to prevent duplicate forwarding with continuous batching
+            if task.task_id in runner.completed or task.task_id in runner.pending:
                 continue
 
             # TODO: Check ordering aligns with MLX distributeds expectations.
 
-            if isinstance(runner.status, RunnerReady) and all(
+            # Allow forwarding tasks when runner is Ready or Running (for continuous batching)
+            if isinstance(runner.status, (RunnerReady, RunnerRunning)) and all(
                 isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning))
                 for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard
             ):

src/exo/worker/runner/runner.py

@@ -1,6 +1,8 @@
+import gc
 import time
 
 import mlx.core as mx
+from anyio import WouldBlock
 
 from exo.shared.types.api import ChatCompletionMessageText
 from exo.shared.types.chunks import TokenChunk
@@ -21,9 +23,6 @@ from exo.shared.types.tasks import (
     TaskStatus,
 )
 from exo.shared.types.worker.instances import BoundInstance
-from exo.shared.types.worker.runner_response import (
-    GenerationResponse,
-)
 from exo.shared.types.worker.runners import (
     RunnerConnected,
     RunnerConnecting,
@@ -39,7 +38,9 @@ from exo.shared.types.worker.runners import (
     RunnerWarmingUp,
 )
 from exo.utils.channels import MpReceiver, MpSender
-from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
+from exo.worker.engines.mlx.generator.batch_engine import BatchGenerationEngine
+from exo.worker.engines.mlx.generator.distributed_sync import DistributedOp, sync_operation
+from exo.worker.engines.mlx.generator.generate import warmup_inference
 from exo.worker.engines.mlx.utils_mlx import (
     initialize_mlx,
     load_mlx_items,
@@ -48,6 +49,26 @@ from exo.worker.engines.mlx.utils_mlx import (
 from exo.worker.runner.bootstrap import logger
 
 
+def _determine_distributed_op(
+    batch_engine: BatchGenerationEngine | None,
+    pending_shutdown: Task | None,
+    should_shutdown: bool,
+) -> DistributedOp:
+    """Determine what operation to perform next in distributed mode.
+
+    Only rank 0's result matters - this gets broadcast to all ranks via sync_operation.
+    """
+    if should_shutdown:
+        return DistributedOp.SHUTDOWN
+    if pending_shutdown is not None and batch_engine is not None and not batch_engine.has_active_requests:
+        return DistributedOp.SHUTDOWN
+    if batch_engine is not None and batch_engine.has_pending_inserts:
+        return DistributedOp.INSERT
+    if batch_engine is not None and batch_engine.has_active_requests:
+        return DistributedOp.STEP
+    return DistributedOp.NOOP
+
+
 def main(
     bound_instance: BoundInstance,
     event_sender: MpSender[Event],
@@ -69,142 +90,247 @@ def main(
     model = None
     tokenizer = None
     group = None
+    batch_engine: BatchGenerationEngine | None = None
+    pending_shutdown: Shutdown | None = None
 
     current_status: RunnerStatus = RunnerIdle()
+
+    def send_status(status: RunnerStatus) -> None:
+        event_sender.send(RunnerStatusUpdated(runner_id=runner_id, runner_status=status))
+
     logger.info("runner created")
-    event_sender.send(
-        RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
-    )
-    with task_receiver as tasks:
-        for task in tasks:
-            event_sender.send(
-                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
-            )
-            event_sender.send(TaskAcknowledged(task_id=task.task_id))
-            match task:
-                case ConnectToGroup() if isinstance(
-                    current_status, (RunnerIdle, RunnerFailed)
-                ):
-                    logger.info("runner connecting")
-                    current_status = RunnerConnecting()
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    group = initialize_mlx(bound_instance)
+    send_status(current_status)
 
-                    logger.info("runner connected")
-                    current_status = RunnerConnected()
+    def handle_task(task: Task, is_deferred: bool = False) -> bool:
+        nonlocal current_status, model, tokenizer, group, batch_engine, pending_shutdown
 
-                # we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
-                case LoadModel() if (
-                    isinstance(current_status, RunnerConnected) and group is not None
-                ) or (isinstance(current_status, RunnerIdle) and group is None):
-                    current_status = RunnerLoading()
-                    logger.info("runner loading")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
+        # For Shutdown, check if we need to defer BEFORE sending Running/Acknowledged
+        if isinstance(task, Shutdown) and not is_deferred:
+            if batch_engine is not None and (batch_engine.has_active_requests or batch_engine.has_pending_inserts):
+                logger.info("deferring shutdown until active requests complete")
+                pending_shutdown = task
+                return True
 
-                    model, tokenizer = load_mlx_items(bound_instance, group)
+        event_sender.send(TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running))
+        event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
-                    current_status = RunnerLoaded()
-                    logger.info("runner loaded")
-                case StartWarmup() if isinstance(current_status, RunnerLoaded):
-                    assert model
-                    assert tokenizer
-                    current_status = RunnerWarmingUp()
-                    logger.info("runner warming up")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
+        match task:
+            case ConnectToGroup() if isinstance(
+                current_status, (RunnerIdle, RunnerFailed)
+            ):
+                logger.info("runner connecting")
+                current_status = RunnerConnecting()
+                send_status(current_status)
+                group = initialize_mlx(bound_instance)
 
-                    logger.info(f"warming up inference for instance: {instance}")
-                    toks = warmup_inference(
-                        model=model,
-                        tokenizer=tokenizer,
-                        # kv_prefix_cache=kv_prefix_cache,  # supply for warmup-time prefix caching
-                    )
-                    logger.info(f"warmed up by generating {toks} tokens")
-                    logger.info(
-                        f"runner initialized in {time.time() - setup_start_time} seconds"
-                    )
-                    current_status = RunnerReady()
-                    logger.info("runner ready")
-                case ChatCompletion(task_params=task_params, command_id=command_id) if (
-                    isinstance(current_status, RunnerReady)
-                ):
-                    assert model
-                    assert tokenizer
-                    logger.info(f"received chat request: {str(task)[:500]}")
-                    current_status = RunnerRunning()
-                    logger.info("runner running")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    assert task_params.messages[0].content is not None
+                logger.info("runner connected")
+                current_status = RunnerConnected()
+                event_sender.send(TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Complete))
+                send_status(current_status)
+
+            case LoadModel() if (
+                isinstance(current_status, RunnerConnected) and group is not None
+            ) or (isinstance(current_status, RunnerIdle) and group is None):
+                current_status = RunnerLoading()
+                logger.info("runner loading")
+                send_status(current_status)
+
+                model, tokenizer = load_mlx_items(bound_instance, group)
+
+                current_status = RunnerLoaded()
+                logger.info("runner loaded")
+                event_sender.send(TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Complete))
+                send_status(current_status)
+
+            case StartWarmup() if isinstance(current_status, RunnerLoaded):
+                assert model is not None
+                assert tokenizer is not None
+                current_status = RunnerWarmingUp()
+                logger.info("runner warming up")
+                send_status(current_status)
+
+                logger.info(f"warming up inference for instance: {instance}")
+                toks = warmup_inference(model=model, tokenizer=tokenizer)
+                logger.info(f"warmed up by generating {toks} tokens")
+                logger.info(f"runner initialized in {time.time() - setup_start_time} seconds")
+
+                batch_engine = BatchGenerationEngine(model=model, tokenizer=tokenizer, group=group)
+
+                current_status = RunnerReady()
+                logger.info("runner ready")
+                event_sender.send(TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Complete))
+                send_status(current_status)
+
+            case ChatCompletion(
+                task_params=task_params, command_id=command_id
+            ) if isinstance(current_status, (RunnerReady, RunnerRunning)):
+                assert batch_engine is not None
+
+                if task_params.messages and task_params.messages[0].content is not None:
                     _check_for_debug_prompts(task_params.messages[0].content)
 
-                    # Generate responses using the actual MLX generation
-                    for response in mlx_generate(
-                        model=model,
-                        tokenizer=tokenizer,
-                        task=task_params,
-                    ):
-                        match response:
-                            case GenerationResponse():
-                                if shard_metadata.device_rank == 0:
-                                    event_sender.send(
-                                        ChunkGenerated(
-                                            command_id=command_id,
-                                            chunk=TokenChunk(
-                                                idx=response.token,
-                                                model=shard_metadata.model_meta.model_id,
-                                                text=response.text,
-                                                token_id=response.token,
-                                                finish_reason=response.finish_reason,
-                                                stats=response.stats,
-                                            ),
-                                        )
-                                    )
-                                # case TokenizedResponse():
-                                # TODO: something here ig
+                # Queue the request - actual insertion happens in sync_and_insert_pending()
+                # In distributed mode, only rank 0 receives tasks from control plane
+                batch_engine.queue_request(command_id=command_id, task_id=task.task_id, task_params=task_params)
 
-                    current_status = RunnerReady()
-                    logger.info("runner ready")
-                case Shutdown():
-                    current_status = RunnerShuttingDown()
-                    logger.info("runner shutting down")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    current_status = RunnerShutdown()
-                case _:
-                    raise ValueError(
-                        f"Received {task.__class__.__name__} outside of state machine in {current_status=}"
-                    )
-            event_sender.send(
-                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Complete)
-            )
-            event_sender.send(
-                RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
-            )
-            if isinstance(current_status, RunnerShutdown):
-                del model, tokenizer, group
-                mx.clear_cache()
-                import gc
+                # Status will be updated after actual insertion in the main loop
+                # For now, set to RunnerRunning to indicate we're processing
+                current_status = RunnerRunning(active_requests=batch_engine.active_count + batch_engine.pending_insert_count)
+                send_status(current_status)
 
-                gc.collect()
-                break
+            case Shutdown():
+                current_status = RunnerShuttingDown()
+                logger.info("runner shutting down")
+                send_status(current_status)
+                event_sender.send(TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Complete))
+                current_status = RunnerShutdown()
+                send_status(current_status)
+                return False
+
+            case _:
+                raise ValueError(
+                    f"Received {task.__class__.__name__} outside of state machine in {current_status=}"
+                )
+
+        return True
+
+    with task_receiver as tasks:
+        running = True
+        is_rank_0 = shard_metadata.device_rank == 0
+        is_distributed = group is not None and group.size() > 1
+
+        while running:
+            if is_distributed:
+                assert group is not None
+                assert batch_engine is not None
+
+                # Distributed mode: synchronize operations across all ranks
+                # Step 1: Only rank 0 checks for tasks and determines operation
+                should_shutdown = False
+                if is_rank_0:
+                    while True:
+                        try:
+                            task = tasks.receive_nowait()
+                            task_result = handle_task(task)
+                            if not task_result:
+                                should_shutdown = True
+                                break
+                        except WouldBlock:
+                            break
+
+                    op = _determine_distributed_op(batch_engine, pending_shutdown, should_shutdown)
+                else:
+                    op = None
+
+                # Step 2: Sync operation across all ranks
+                synced_op = sync_operation(op, shard_metadata.device_rank, group)
+
+                # Step 3: All ranks execute the same operation
+                match synced_op:
+                    case DistributedOp.INSERT:
+                        batch_engine.sync_and_insert_pending()
+                        if is_rank_0:
+                            current_status = RunnerRunning(active_requests=batch_engine.active_count)
+                            send_status(current_status)
+
+                    case DistributedOp.STEP:
+                        for resp in batch_engine.step():
+                            if is_rank_0:
+                                event_sender.send(ChunkGenerated(
+                                    command_id=resp.command_id,
+                                    chunk=TokenChunk(
+                                        idx=resp.response.token,
+                                        model=shard_metadata.model_meta.model_id,
+                                        text=resp.response.text,
+                                        token_id=resp.response.token,
+                                        finish_reason=resp.response.finish_reason,
+                                        stats=resp.response.stats,
+                                    ),
+                                ))
+                            if resp.response.finish_reason is not None:
+                                if is_rank_0:
+                                    event_sender.send(TaskStatusUpdated(task_id=resp.task_id, task_status=TaskStatus.Complete))
+
+                        if is_rank_0:
+                            if batch_engine.has_active_requests:
+                                current_status = RunnerRunning(active_requests=batch_engine.active_count)
+                            else:
+                                current_status = RunnerReady()
+                            send_status(current_status)
+
+                    case DistributedOp.SHUTDOWN:
+                        running = False
+                        if is_rank_0 and pending_shutdown is not None:
+                            handle_task(pending_shutdown, is_deferred=True)
+
+                    case DistributedOp.NOOP:
+                        # No work to do - all ranks poll together
+                        # We can't have rank 0 block while others try to sync
+                        if is_rank_0:
+                            try:
+                                task = tasks.receive_nowait()
+                                task_result = handle_task(task)
+                                if not task_result:
+                                    # Will sync SHUTDOWN on next iteration
+                                    pass
+                            except WouldBlock:
+                                pass
+                        # All ranks: short sleep before looping back to sync
+                        time.sleep(0.001)
+
+            else:
+                # Non-distributed mode: original logic with queue + insert
+                while True:
+                    try:
+                        task = tasks.receive_nowait()
+                        running = handle_task(task)
+                        if not running:
+                            break
+                    except WouldBlock:
+                        break
+
+                if not running:
+                    break
+
+                # Insert any queued requests (non-distributed just inserts directly)
+                # Status was already sent in handle_task when queueing
+                if batch_engine is not None and batch_engine.has_pending_inserts:
+                    batch_engine.sync_and_insert_pending()
+
+                if batch_engine is not None and batch_engine.has_active_requests:
+                    for resp in batch_engine.step():
+                        if shard_metadata.device_rank == 0:
+                            event_sender.send(ChunkGenerated(
+                                command_id=resp.command_id,
+                                chunk=TokenChunk(
+                                    idx=resp.response.token,
+                                    model=shard_metadata.model_meta.model_id,
+                                    text=resp.response.text,
+                                    token_id=resp.response.token,
+                                    finish_reason=resp.response.finish_reason,
+                                    stats=resp.response.stats,
+                                ),
+                            ))
+                        if resp.response.finish_reason is not None:
+                            event_sender.send(TaskStatusUpdated(task_id=resp.task_id, task_status=TaskStatus.Complete))
+
+                    if batch_engine.has_active_requests:
+                        current_status = RunnerRunning(active_requests=batch_engine.active_count)
+                    else:
+                        current_status = RunnerReady()
+                    send_status(current_status)
+
+                    # Process deferred shutdown after all requests complete
+                    if pending_shutdown is not None and not batch_engine.has_active_requests and not batch_engine.has_pending_inserts:
+                        running = handle_task(pending_shutdown, is_deferred=True)
+                else:
+                    task = tasks.receive()
+                    running = handle_task(task)
+
+    # Cleanup
+    del model, tokenizer, group, batch_engine
+    mx.clear_cache()
+    gc.collect()
 
 
 EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"

src/exo/worker/runner/runner_supervisor.py

@@ -105,7 +105,7 @@ class RunnerSupervisor:
             return
 
         # This is overkill but it's not technically bad, just unnecessary.
-        logger.warning("Runner process didn't shutdown succesfully, terminating")
+        logger.warning("Runner process didn't shutdown successfully, terminating")
         self.runner_process.terminate()
         await to_thread.run_sync(self.runner_process.join, 5)
         if not self.runner_process.is_alive():
@@ -128,9 +128,11 @@ class RunnerSupervisor:
 
     async def start_task(self, task: Task):
         if task.task_id in self.completed:
-            logger.info(
-                f"Skipping invalid task {task} as it has already been completed"
-            )
+            logger.info(f"Skipping task {task.task_id} - already completed")
+            return
+        if task.task_id in self.pending:
+            logger.info(f"Skipping task {task.task_id} - already pending")
+            return
         logger.info(f"Starting task {task}")
         event = anyio.Event()
         self.pending[task.task_id] = event
@@ -149,13 +151,17 @@ class RunnerSupervisor:
                     if isinstance(event, RunnerStatusUpdated):
                         self.status = event.runner_status
                     if isinstance(event, TaskAcknowledged):
-                        self.pending.pop(event.task_id).set()
+                        # Just set the event to unblock start_task, but keep in pending
+                        # to prevent duplicate forwarding until completion
+                        if event.task_id in self.pending:
+                            self.pending[event.task_id].set()
                         continue
-                    if (
-                        isinstance(event, TaskStatusUpdated)
-                        and event.task_status == TaskStatus.Complete
+                    if isinstance(event, TaskStatusUpdated) and event.task_status in (
+                        TaskStatus.Complete,
+                        TaskStatus.TimedOut,
+                        TaskStatus.Failed,
                     ):
-                        # If a task has just been completed, we should be working on it.
+                        # If a task has just finished, we should be working on it.
                         assert isinstance(
                             self.status,
                             (
@@ -166,6 +172,8 @@ class RunnerSupervisor:
                                 RunnerShuttingDown,
                             ),
                         )
+                        # Now safe to remove from pending and add to completed
+                        self.pending.pop(event.task_id, None)
                         self.completed.add(event.task_id)
                     await self._event_sender.send(event)
             except (ClosedResourceError, BrokenResourceError) as e:

src/exo/worker/tests/unittests/conftest.py

@@ -20,6 +20,7 @@ class FakeRunnerSupervisor:
     bound_instance: BoundInstance
     status: RunnerStatus
     completed: set[TaskId] = field(default_factory=set)
+    pending: dict[TaskId, object] = field(default_factory=dict)
 
 
 class OtherTask(BaseTask):

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

@@ -0,0 +1,386 @@
+"""
+Unit tests for tokenizer loading and functionality across all supported models.
+
+This test downloads only tokenizer-related files (not full model weights) to verify
+that tokenizers can be loaded and used correctly for encoding/decoding.
+"""
+
+import asyncio
+import contextlib
+from pathlib import Path
+
+import pytest
+
+from exo.shared.models.model_cards import MODEL_CARDS, ModelCard
+from exo.worker.download.download_utils import (
+    download_file_with_retry,
+    ensure_models_dir,
+    fetch_file_list_with_cache,
+)
+from exo.worker.engines.mlx.utils_mlx import (
+    get_eos_token_ids_for_model,
+    load_tokenizer_for_model_id,
+)
+
+# Files needed for tokenizer functionality
+TOKENIZER_FILE_PATTERNS = [
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "special_tokens_map.json",
+    "vocab.json",
+    "vocab.txt",
+    "merges.txt",
+    "tiktoken.model",
+    "added_tokens.json",
+    "tokenizer.model",
+    "tokenization_*.py",  # Custom tokenizer implementations
+]
+
+
+def is_tokenizer_file(filename: str) -> bool:
+    """Check if a file is needed for tokenizer functionality."""
+    for pattern in TOKENIZER_FILE_PATTERNS:
+        if "*" in pattern:
+            prefix = pattern.split("*")[0]
+            suffix = pattern.split("*")[1]
+            if filename.startswith(prefix) and filename.endswith(suffix):
+                return True
+        elif filename == pattern:
+            return True
+    return False
+
+
+async def download_tokenizer_files(model_id: str) -> Path:
+    """Download only the tokenizer-related files for a model."""
+    target_dir = await ensure_models_dir() / model_id.replace("/", "--")
+    target_dir.mkdir(parents=True, exist_ok=True)
+
+    file_list = await fetch_file_list_with_cache(model_id, "main", recursive=True)
+
+    tokenizer_files = [f for f in file_list if is_tokenizer_file(f.path)]
+
+    if not tokenizer_files:
+        pytest.skip(f"No tokenizer files found for {model_id}")
+
+    for file_entry in tokenizer_files:
+        with contextlib.suppress(FileNotFoundError):
+            await download_file_with_retry(
+                model_id, "main", file_entry.path, target_dir
+            )
+
+    return target_dir
+
+
+# Get a sample of models to test (one per family to keep tests fast)
+def get_test_models() -> list[tuple[str, ModelCard]]:
+    """Get a representative sample of models to test."""
+    # Pick one model from each family to test
+    families: dict[str, tuple[str, ModelCard]] = {}
+    for short_id, card in MODEL_CARDS.items():
+        # Extract family name (e.g., "llama-3.1" from "llama-3.1-8b")
+        parts = short_id.split("-")
+        family = "-".join(parts[:2]) if len(parts) >= 2 else parts[0]
+
+        if family not in families:
+            families[family] = (short_id, card)
+
+    return list(families.values())
+
+
+TEST_MODELS: list[tuple[str, ModelCard]] = get_test_models()
+
+
+@pytest.fixture(scope="module")
+def event_loop():
+    """Create event loop for async tests."""
+    loop = asyncio.new_event_loop()
+    yield loop
+    loop.close()
+
+
+@pytest.mark.parametrize(
+    "short_id,model_card",
+    TEST_MODELS,
+    ids=[m[0] for m in TEST_MODELS],
+)
+@pytest.mark.asyncio
+async def test_tokenizer_encode_decode(short_id: str, model_card: ModelCard) -> None:
+    """Test that tokenizer can encode and decode text correctly."""
+    model_id = str(model_card.model_id)
+
+    # Download tokenizer files
+    model_path = await download_tokenizer_files(model_id)
+
+    # Verify required files exist
+    has_tokenizer = (
+        (model_path / "tokenizer.json").exists()
+        or (model_path / "tokenizer_config.json").exists()
+        or (model_path / "tiktoken.model").exists()
+        or (model_path / "tokenizer.model").exists()
+    )
+    if not has_tokenizer:
+        pytest.skip(f"Required tokenizer files not found for {model_id}")
+
+    # Load tokenizer
+    tokenizer = load_tokenizer_for_model_id(model_id, model_path)
+
+    # Test basic encoding
+    test_text = "Hello, world!"
+    encoded = tokenizer.encode(test_text)
+    assert isinstance(encoded, list), f"encode() should return a list for {model_id}"
+    assert len(encoded) > 0, f"encode() should return non-empty list for {model_id}"
+    assert all(isinstance(t, int) for t in encoded), (
+        f"All tokens should be integers for {model_id}"
+    )
+
+    # Test decoding
+    decoded = tokenizer.decode(encoded)
+    assert isinstance(decoded, str), f"decode() should return a string for {model_id}"
+    assert test_text in decoded or decoded.strip() == test_text.strip(), (
+        f"decode(encode(x)) should preserve text for {model_id}: got {decoded!r}"
+    )
+
+    # Test with longer text
+    long_text = "The quick brown fox jumps over the lazy dog. " * 10
+    long_encoded = tokenizer.encode(long_text)
+    assert len(long_encoded) > len(encoded), (
+        f"Longer text should produce more tokens for {model_id}"
+    )
+
+    # Test empty string
+    empty_encoded = tokenizer.encode("")
+    assert isinstance(empty_encoded, list), (
+        f"encode('') should return a list for {model_id}"
+    )
+
+    # Test special characters
+    special_text = 'Hello!\n\tWorld? <test> & "quotes"'
+    special_encoded = tokenizer.encode(special_text)
+    assert len(special_encoded) > 0, f"Special chars should encode for {model_id}"
+
+    # Test unicode
+    unicode_text = "Hello 世界 🌍"
+    unicode_encoded = tokenizer.encode(unicode_text)
+    assert len(unicode_encoded) > 0, f"Unicode should encode for {model_id}"
+
+
+@pytest.mark.parametrize(
+    "short_id,model_card",
+    TEST_MODELS,
+    ids=[m[0] for m in TEST_MODELS],
+)
+@pytest.mark.asyncio
+async def test_tokenizer_has_required_attributes(
+    short_id: str, model_card: ModelCard
+) -> None:
+    """Test that tokenizer has required attributes for inference."""
+    model_id = str(model_card.model_id)
+
+    model_path = await download_tokenizer_files(model_id)
+
+    has_tokenizer = (
+        (model_path / "tokenizer.json").exists()
+        or (model_path / "tokenizer_config.json").exists()
+        or (model_path / "tiktoken.model").exists()
+        or (model_path / "tokenizer.model").exists()
+    )
+    if not has_tokenizer:
+        pytest.skip(f"Required tokenizer files not found for {model_id}")
+
+    tokenizer = load_tokenizer_for_model_id(model_id, model_path)
+    eos_token_ids = get_eos_token_ids_for_model(model_id)
+
+    # Check for vocabulary size
+    empty_vocab: dict[str, int] = {}
+    vocab_size: int = getattr(tokenizer, "vocab_size", None) or len(
+        getattr(tokenizer, "get_vocab", lambda: empty_vocab)()
+    )
+    assert vocab_size > 0, f"Tokenizer should have vocab_size > 0 for {model_id}"
+
+    # Check for EOS token (either from tokenizer or explicitly provided)
+    has_eos = (
+        eos_token_ids is not None
+        or getattr(tokenizer, "eos_token_id", None) is not None
+        or getattr(tokenizer, "eos_token", None) is not None
+    )
+    assert has_eos, f"Tokenizer should have EOS token for {model_id}"
+
+
+@pytest.mark.parametrize(
+    "short_id,model_card",
+    TEST_MODELS,
+    ids=[m[0] for m in TEST_MODELS],
+)
+@pytest.mark.asyncio
+async def test_tokenizer_special_tokens(short_id: str, model_card: ModelCard) -> None:
+    """Test that tokenizer can encode text containing special tokens.
+
+    This is critical because the actual inference path uses prompts with
+    special tokens from chat templates. If special tokens aren't handled
+    correctly, encoding will fail.
+    """
+    model_id = str(model_card.model_id)
+
+    model_path = await download_tokenizer_files(model_id)
+
+    has_tokenizer = (
+        (model_path / "tokenizer.json").exists()
+        or (model_path / "tokenizer_config.json").exists()
+        or (model_path / "tiktoken.model").exists()
+        or (model_path / "tokenizer.model").exists()
+    )
+    assert has_tokenizer, f"Required tokenizer files not found for {model_id}"
+
+    tokenizer = load_tokenizer_for_model_id(model_id, model_path)
+
+    # Get special tokens from the tokenizer
+    special_tokens: list[str] = []
+
+    # Try to get special tokens from various sources
+    if hasattr(tokenizer, "all_special_tokens"):
+        special_tokens.extend(tokenizer.all_special_tokens)
+    elif hasattr(tokenizer, "_tokenizer") and hasattr(
+        tokenizer._tokenizer,
+        "all_special_tokens",
+    ):
+        special_tokens.extend(tokenizer._tokenizer.all_special_tokens)
+
+    # Also check for common special token attributes
+    for attr in [
+        "bos_token",
+        "eos_token",
+        "pad_token",
+        "unk_token",
+        "sep_token",
+        "cls_token",
+    ]:
+        token = getattr(tokenizer, attr, None)
+        if token is None and hasattr(tokenizer, "_tokenizer"):
+            token = getattr(tokenizer._tokenizer, attr, None)
+        if token and isinstance(token, str) and token not in special_tokens:
+            special_tokens.append(token)
+
+    # If we found special tokens, test encoding text that contains them
+    if special_tokens:
+        # Create text with special tokens interspersed
+        test_with_special = f"{special_tokens[0]}Hello world"
+        if len(special_tokens) > 1:
+            test_with_special += f"{special_tokens[1]}"
+
+        encoded = tokenizer.encode(test_with_special)
+        assert isinstance(encoded, list), (
+            f"encode() with special tokens should return list for {model_id}"
+        )
+        assert len(encoded) > 0, (
+            f"encode() with special tokens should return non-empty list for {model_id}"
+        )
+        assert all(isinstance(t, int) for t in encoded), (
+            f"All tokens should be integers for {model_id}"
+        )
+
+        # Verify we can decode
+        decoded = tokenizer.decode(encoded)
+        assert isinstance(decoded, str), f"decode() should return string for {model_id}"
+
+    # Test with angle-bracket tokens (common format for special tokens)
+    # These should not raise errors even if they're not actual special tokens
+    angle_bracket_text = "<|test|>Hello<|end|>"
+    encoded = tokenizer.encode(angle_bracket_text)
+    assert isinstance(encoded, list), (
+        f"encode() with angle brackets should return list for {model_id}"
+    )
+    assert len(encoded) > 0, (
+        f"encode() with angle brackets should be non-empty for {model_id}"
+    )
+
+
+# Specifically test Kimi tokenizer since it has special handling
+@pytest.mark.asyncio
+async def test_kimi_tokenizer_specifically():
+    """Test Kimi tokenizer with its specific patches and quirks."""
+    kimi_models = [
+        (short_id, card)
+        for short_id, card in MODEL_CARDS.items()
+        if "kimi" in short_id.lower()
+    ]
+
+    if not kimi_models:
+        pytest.skip("No Kimi models found in MODEL_CARDS")
+
+    _, model_card = kimi_models[0]
+    model_id = str(model_card.model_id)
+
+    model_path = await download_tokenizer_files(model_id)
+
+    # Ensure the custom tokenizer file exists
+    if not (model_path / "tokenization_kimi.py").exists():
+        pytest.skip("tokenization_kimi.py not found")
+
+    tokenizer = load_tokenizer_for_model_id(model_id, model_path)
+    eos_token_ids = get_eos_token_ids_for_model(model_id)
+
+    # Test encode/decode cycle
+    test_text = "Hello, world!"
+    encoded = tokenizer.encode(test_text)
+    decoded = tokenizer.decode(encoded)
+
+    assert len(encoded) > 0, "Kimi tokenizer should encode text"
+    assert isinstance(decoded, str), "Kimi tokenizer should decode to string"
+
+    # Test that the patched encode works (returns list of ints)
+    assert all(isinstance(t, int) for t in encoded), "Tokens should be integers"
+
+    # Test encoding text with special tokens (like from chat templates)
+    # This is critical - the warmup inference uses prompts with special tokens
+    special_token_text = "<|im_user|>user<|im_middle|>Hello<|im_end|><|im_assistant|>"
+    special_encoded = tokenizer.encode(special_token_text)
+    assert len(special_encoded) > 0, "Kimi tokenizer should handle special tokens"
+    assert all(isinstance(t, int) for t in special_encoded), (
+        "Special token encoding should return integers"
+    )
+
+    # Verify EOS token is set
+    assert eos_token_ids == [163586], "Kimi EOS token should be [163586]"
+
+
+# Test GLM tokenizer since it also has special handling
+@pytest.mark.asyncio
+async def test_glm_tokenizer_specifically():
+    """Test GLM tokenizer with its specific EOS tokens."""
+    glm_models = [
+        (short_id, card)
+        for short_id, card in MODEL_CARDS.items()
+        if "glm" in short_id.lower()
+    ]
+
+    if not glm_models:
+        pytest.skip("No GLM models found in MODEL_CARDS")
+
+    _, model_card = glm_models[0]
+    model_id = str(model_card.model_id)
+
+    model_path = await download_tokenizer_files(model_id)
+
+    has_tokenizer = (model_path / "tokenizer.json").exists() or (
+        model_path / "tokenizer_config.json"
+    ).exists()
+    if not has_tokenizer:
+        pytest.skip("GLM tokenizer files not found")
+
+    tokenizer = load_tokenizer_for_model_id(model_id, model_path)
+    eos_token_ids = get_eos_token_ids_for_model(model_id)
+
+    # Test encode/decode
+    test_text = "Hello, world!"
+    encoded = tokenizer.encode(test_text)
+    decoded = tokenizer.decode(encoded)
+
+    assert len(encoded) > 0, "GLM tokenizer should encode text"
+    assert isinstance(decoded, str), "GLM tokenizer should decode to string"
+
+    # Verify EOS tokens
+    assert eos_token_ids == [
+        151336,
+        151329,
+        151338,
+    ], "GLM EOS tokens should be correct"

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

@@ -0,0 +1,315 @@
+"""
+Tests for continuous batching behavior in the runner.
+
+These tests verify that:
+1. Single requests work through the batch path
+2. Multiple concurrent requests batch together
+3. Tokens are routed to the correct requests
+4. Requests complete at different times appropriately
+"""
+
+# pyright: reportAny=false
+# pyright: reportUnknownArgumentType=false
+# pyright: reportUnknownMemberType=false
+# pyright: reportAttributeAccessIssue=false
+# pyright: reportInvalidTypeVarUse=false
+
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+
+import exo.worker.runner.runner as mlx_runner
+from exo.shared.types.api import ChatCompletionMessage
+from exo.shared.types.common import CommandId, NodeId
+from exo.shared.types.events import (
+    Event,
+    RunnerStatusUpdated,
+    TaskStatusUpdated,
+)
+from exo.shared.types.tasks import (
+    ChatCompletion,
+    ChatCompletionTaskParams,
+    ConnectToGroup,
+    LoadModel,
+    Shutdown,
+    StartWarmup,
+    Task,
+    TaskId,
+    TaskStatus,
+)
+from exo.shared.types.worker.runner_response import GenerationResponse
+from exo.shared.types.worker.runners import RunnerRunning
+from exo.utils.channels import mp_channel
+from exo.worker.engines.mlx.generator.batch_engine import (
+    BatchedGenerationResponse,
+)
+from exo.worker.tests.constants import (
+    INSTANCE_1_ID,
+    MODEL_A_ID,
+    NODE_A,
+    RUNNER_1_ID,
+)
+from exo.worker.tests.unittests.conftest import get_bound_mlx_ring_instance
+
+
+class FakeBatchEngineWithTokens:
+    """
+    Fake batch engine that generates a specified number of tokens per request.
+
+    This simulates realistic batch generation behavior where:
+    - Requests are queued on insert
+    - Each step() call generates one token for all active requests
+    - Requests complete when they've generated all their tokens
+    """
+
+    def __init__(self, *_args: Any, **_kwargs: Any):
+        self._active_requests: dict[int, tuple[CommandId, TaskId, int, int]] = {}
+        self._pending_inserts: list[tuple[CommandId, TaskId, ChatCompletionTaskParams]] = []
+        self._uid_counter = 0
+        self._tokens_per_request = 3  # Default: generate 3 tokens before completing
+        self.rank = 0  # Fake rank for testing
+
+    def queue_request(
+        self,
+        command_id: CommandId,
+        task_id: TaskId,
+        task_params: ChatCompletionTaskParams,
+    ) -> None:
+        """Queue a request for insertion."""
+        self._pending_inserts.append((command_id, task_id, task_params))
+
+    def sync_and_insert_pending(self) -> list[int]:
+        """Insert all pending requests."""
+        uids: list[int] = []
+        for command_id, task_id, task_params in self._pending_inserts:
+            uid = self._do_insert(command_id, task_id, task_params)
+            uids.append(uid)
+        self._pending_inserts.clear()
+        return uids
+
+    @property
+    def has_pending_inserts(self) -> bool:
+        return len(self._pending_inserts) > 0
+
+    def _do_insert(
+        self,
+        command_id: CommandId,
+        task_id: TaskId,
+        task_params: ChatCompletionTaskParams | None,
+    ) -> int:
+        uid = self._uid_counter
+        self._uid_counter += 1
+        # Track: (command_id, task_id, tokens_generated, max_tokens)
+        max_tokens = task_params.max_tokens if task_params else self._tokens_per_request
+        self._active_requests[uid] = (command_id, task_id, 0, max_tokens or 3)
+        return uid
+
+    def step(self) -> list[BatchedGenerationResponse]:
+        results: list[BatchedGenerationResponse] = []
+        uids_to_remove: list[int] = []
+
+        for uid, (command_id, task_id, tokens_gen, max_tokens) in list(
+            self._active_requests.items()
+        ):
+            tokens_gen += 1
+            finish_reason = "stop" if tokens_gen >= max_tokens else None
+            text = f"token{tokens_gen}"
+
+            if finish_reason:
+                uids_to_remove.append(uid)
+            else:
+                self._active_requests[uid] = (
+                    command_id,
+                    task_id,
+                    tokens_gen,
+                    max_tokens,
+                )
+
+            results.append(
+                BatchedGenerationResponse(
+                    command_id=command_id,
+                    task_id=task_id,
+                    response=GenerationResponse(
+                        token=tokens_gen,
+                        text=text,
+                        finish_reason=finish_reason,
+                    ),
+                )
+            )
+
+        for uid in uids_to_remove:
+            del self._active_requests[uid]
+
+        return results
+
+    @property
+    def has_active_requests(self) -> bool:
+        return len(self._active_requests) > 0
+
+    @property
+    def active_count(self) -> int:
+        return len(self._active_requests)
+
+    @property
+    def pending_insert_count(self) -> int:
+        return len(self._pending_inserts)
+
+
+def make_nothin[T, U, V](res: T):
+    def nothin(*_1: U, **_2: V) -> T:
+        return res
+
+    return nothin
+
+
+@pytest.fixture
+def patch_batch_engine(monkeypatch: pytest.MonkeyPatch):
+    """Patch MLX dependencies and use FakeBatchEngineWithTokens."""
+    monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MagicMock()))
+    monkeypatch.setattr(
+        mlx_runner, "load_mlx_items", make_nothin((MagicMock(), MagicMock()))
+    )
+    monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1))
+    monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", make_nothin(None))
+    monkeypatch.setattr(mlx_runner, "BatchGenerationEngine", FakeBatchEngineWithTokens)
+
+
+def _run_with_tasks(tasks: list[Task]) -> list[Event]:
+    """
+    Run tasks through the runner, adding shutdown at the end.
+
+    Tasks are sent in order, with shutdown sent last.
+    The batch engine processes between task handling.
+    """
+    bound_instance = get_bound_mlx_ring_instance(
+        instance_id=INSTANCE_1_ID,
+        model_id=MODEL_A_ID,
+        runner_id=RUNNER_1_ID,
+        node_id=NodeId(NODE_A),
+    )
+
+    task_sender, task_receiver = mp_channel[Task]()
+    event_sender, event_receiver = mp_channel[Event]()
+
+    shutdown_task = Shutdown(
+        task_id=TaskId("shutdown"),
+        instance_id=INSTANCE_1_ID,
+        runner_id=RUNNER_1_ID,
+    )
+
+    with task_sender, event_receiver:
+        # Send all tasks including shutdown
+        for t in tasks:
+            task_sender.send(t)
+        task_sender.send(shutdown_task)
+
+        # Disable cleanup methods to prevent issues
+        event_sender.close = lambda: None
+        event_sender.join = lambda: None
+        task_receiver.close = lambda: None
+        task_receiver.join = lambda: None
+
+        mlx_runner.main(bound_instance, event_sender, task_receiver)
+
+        return event_receiver.collect()
+
+
+INIT_TASK = ConnectToGroup(task_id=TaskId("init"), instance_id=INSTANCE_1_ID)
+LOAD_TASK = LoadModel(task_id=TaskId("load"), instance_id=INSTANCE_1_ID)
+WARMUP_TASK = StartWarmup(task_id=TaskId("warmup"), instance_id=INSTANCE_1_ID)
+
+
+def make_chat_task(
+    task_id: str, command_id: str, max_tokens: int = 3
+) -> ChatCompletion:
+    return ChatCompletion(
+        task_id=TaskId(task_id),
+        command_id=CommandId(command_id),
+        task_params=ChatCompletionTaskParams(
+            model=str(MODEL_A_ID),
+            messages=[ChatCompletionMessage(role="user", content="hello")],
+            stream=True,
+            max_tokens=max_tokens,
+        ),
+        instance_id=INSTANCE_1_ID,
+    )
+
+
+def test_single_request_generates_tokens(patch_batch_engine: None):
+    """
+    Verify a single request generates the expected tokens through the batch path.
+
+    Note: With the current non-blocking design, shutdown is processed before
+    batch steps run when all tasks are queued together. This test verifies
+    the runner status reflects active requests.
+    """
+    chat_task = make_chat_task("chat1", "cmd1", max_tokens=3)
+    events = _run_with_tasks([INIT_TASK, LOAD_TASK, WARMUP_TASK, chat_task])
+
+    # Find RunnerRunning status events - this shows the request was inserted
+    running_events = [
+        e
+        for e in events
+        if isinstance(e, RunnerStatusUpdated)
+        and isinstance(e.runner_status, RunnerRunning)
+    ]
+
+    assert len(running_events) >= 1, "Expected at least one RunnerRunning event"
+    assert running_events[0].runner_status.active_requests == 1
+
+
+def test_runner_status_reflects_active_requests(patch_batch_engine: None):
+    """Verify RunnerRunning status includes active_requests count."""
+    chat_task = make_chat_task("chat1", "cmd1", max_tokens=2)
+    events = _run_with_tasks([INIT_TASK, LOAD_TASK, WARMUP_TASK, chat_task])
+
+    # Find RunnerRunning status events
+    running_events = [
+        e
+        for e in events
+        if isinstance(e, RunnerStatusUpdated)
+        and isinstance(e.runner_status, RunnerRunning)
+    ]
+
+    assert len(running_events) > 0, "Expected at least one RunnerRunning event"
+    assert running_events[0].runner_status.active_requests == 1
+
+
+def test_chat_task_acknowledged(patch_batch_engine: None):
+    """Verify chat completion task is acknowledged with proper status updates."""
+    chat_task = make_chat_task("chat1", "cmd1", max_tokens=2)
+    events = _run_with_tasks([INIT_TASK, LOAD_TASK, WARMUP_TASK, chat_task])
+
+    # Find the chat task status events
+    chat_running = [
+        e
+        for e in events
+        if isinstance(e, TaskStatusUpdated)
+        and e.task_id == TaskId("chat1")
+        and e.task_status == TaskStatus.Running
+    ]
+
+    assert len(chat_running) == 1, "Expected exactly one chat task Running status"
+
+
+def test_multiple_requests_tracked(patch_batch_engine: None):
+    """Verify multiple concurrent requests are tracked in active_requests."""
+    chat1 = make_chat_task("chat1", "cmd1", max_tokens=2)
+    chat2 = make_chat_task("chat2", "cmd2", max_tokens=2)
+    events = _run_with_tasks([INIT_TASK, LOAD_TASK, WARMUP_TASK, chat1, chat2])
+
+    # Find RunnerRunning status events
+    running_events = [
+        e
+        for e in events
+        if isinstance(e, RunnerStatusUpdated)
+        and isinstance(e.runner_status, RunnerRunning)
+    ]
+
+    # Should have at least 2 RunnerRunning events (one per request inserted)
+    assert len(running_events) >= 2, f"Expected at least 2 RunnerRunning events, got {len(running_events)}"
+
+    # First should have 1 active request, second should have 2
+    assert running_events[0].runner_status.active_requests == 1
+    assert running_events[1].runner_status.active_requests == 2

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

@@ -0,0 +1,79 @@
+"""Test for distributed synchronization in batch generation.
+
+These tests verify that all ranks in a distributed setup call the same
+collective operations in the same order, preventing race conditions and deadlocks.
+"""
+
+import pytest
+
+from exo.worker.engines.mlx.generator.distributed_sync import DistributedOp
+from exo.worker.runner.runner import _determine_distributed_op
+
+
+class FakeBatchEngine:
+    """Minimal fake batch engine for testing _determine_distributed_op."""
+
+    def __init__(self, has_pending_inserts: bool = False, has_active_requests: bool = False):
+        self._has_pending_inserts = has_pending_inserts
+        self._has_active_requests = has_active_requests
+
+    @property
+    def has_pending_inserts(self) -> bool:
+        return self._has_pending_inserts
+
+    @property
+    def has_active_requests(self) -> bool:
+        return self._has_active_requests
+
+
+def test_distributed_sync_prevents_race_condition():
+    """
+    Test that the new architecture prevents the race condition.
+
+    In the old code, each rank independently decided what operation to perform
+    based on its local state. This could cause rank 0 to call insert_request()
+    while rank 1 called step(), causing a collective mismatch.
+
+    The fix ensures:
+    1. Only rank 0 determines the operation (via _determine_distributed_op)
+    2. Rank 0's decision is broadcast to all ranks (via sync_operation)
+    3. All ranks then execute the same operation
+    """
+    # Scenario: rank 0 has pending inserts, rank 1 has active requests
+    # In old code: rank 0 would INSERT, rank 1 would STEP -> race condition
+    # In new code: rank 0's decision (INSERT) is broadcast to all ranks
+
+    # Test _determine_distributed_op gives priority to INSERT over STEP
+    batch_engine_with_inserts = FakeBatchEngine(has_pending_inserts=True, has_active_requests=True)
+    op = _determine_distributed_op(batch_engine_with_inserts, pending_shutdown=None, should_shutdown=False)
+    assert op == DistributedOp.INSERT, "INSERT should take priority when there are pending inserts"
+
+    # Test step is used when there are active requests but no pending inserts
+    batch_engine_with_active = FakeBatchEngine(has_pending_inserts=False, has_active_requests=True)
+    op = _determine_distributed_op(batch_engine_with_active, pending_shutdown=None, should_shutdown=False)
+    assert op == DistributedOp.STEP, "STEP should be used when there are active requests"
+
+    # Test NOOP when nothing to do
+    batch_engine_idle = FakeBatchEngine(has_pending_inserts=False, has_active_requests=False)
+    op = _determine_distributed_op(batch_engine_idle, pending_shutdown=None, should_shutdown=False)
+    assert op == DistributedOp.NOOP, "NOOP should be used when nothing to do"
+
+
+def test_distributed_sync_shutdown_handling():
+    """Test that shutdown is properly coordinated across ranks."""
+    batch_engine = FakeBatchEngine(has_pending_inserts=False, has_active_requests=False)
+
+    # Test direct shutdown request
+    op = _determine_distributed_op(batch_engine, pending_shutdown=None, should_shutdown=True)
+    assert op == DistributedOp.SHUTDOWN, "SHUTDOWN should be returned when should_shutdown is True"
+
+    # Test pending shutdown with no active requests
+    from unittest.mock import MagicMock
+    pending_shutdown = MagicMock()
+    op = _determine_distributed_op(batch_engine, pending_shutdown=pending_shutdown, should_shutdown=False)
+    assert op == DistributedOp.SHUTDOWN, "SHUTDOWN should be returned when pending_shutdown and no active requests"
+
+    # Test pending shutdown with active requests - should continue processing
+    batch_engine_active = FakeBatchEngine(has_pending_inserts=False, has_active_requests=True)
+    op = _determine_distributed_op(batch_engine_active, pending_shutdown=pending_shutdown, should_shutdown=False)
+    assert op == DistributedOp.STEP, "Should continue STEP while requests are active even with pending shutdown"

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

@@ -1,11 +1,16 @@
 # Check tasks are complete before runner is ever ready.
+
+# pyright: reportAny=false
+
 from collections.abc import Iterable
-from typing import Callable
+from typing import Any, Callable
+from unittest.mock import MagicMock
 
 import pytest
 
 import exo.worker.runner.runner as mlx_runner
 from exo.shared.types.api import ChatCompletionMessage
+from exo.shared.types.common import CommandId
 from exo.shared.types.chunks import TokenChunk
 from exo.shared.types.events import (
     ChunkGenerated,
@@ -22,6 +27,7 @@ from exo.shared.types.tasks import (
     Shutdown,
     StartWarmup,
     Task,
+    TaskId,
     TaskStatus,
 )
 from exo.shared.types.worker.runner_response import GenerationResponse
@@ -38,6 +44,9 @@ from exo.shared.types.worker.runners import (
     RunnerWarmingUp,
 )
 from exo.utils.channels import mp_channel
+from exo.worker.engines.mlx.generator.batch_engine import (
+    BatchedGenerationResponse,
+)
 
 from ...constants import (
     CHAT_COMPLETION_TASK_ID,
@@ -107,18 +116,85 @@ def assert_events_equal(test_events: Iterable[Event], true_events: Iterable[Even
         assert test_event == true_event, f"{test_event} != {true_event}"
 
 
+class FakeBatchEngine:
+    """
+    Fake batch engine for testing.
+
+    Queues requests on insert, returns one token per step.
+    The runner's non-blocking loop drains all tasks before running batch steps,
+    so this engine queues requests and has_active_requests returns True only
+    after at least one request has been inserted.
+    """
+
+    def __init__(self, *_args: Any, **_kwargs: Any):
+        self._active_requests: dict[int, tuple[CommandId, TaskId]] = {}
+        self._pending_inserts: list[tuple[CommandId, TaskId, ChatCompletionTaskParams]] = []
+        self._uid_counter = 0
+        self.rank = 0  # Fake rank for testing
+
+    def queue_request(
+        self,
+        command_id: CommandId,
+        task_id: TaskId,
+        task_params: ChatCompletionTaskParams,
+    ) -> None:
+        """Queue a request for insertion."""
+        self._pending_inserts.append((command_id, task_id, task_params))
+
+    def sync_and_insert_pending(self) -> list[int]:
+        """Insert all pending requests."""
+        uids: list[int] = []
+        for command_id, task_id, _task_params in self._pending_inserts:
+            uid = self._uid_counter
+            self._uid_counter += 1
+            self._active_requests[uid] = (command_id, task_id)
+            uids.append(uid)
+        self._pending_inserts.clear()
+        return uids
+
+    @property
+    def has_pending_inserts(self) -> bool:
+        return len(self._pending_inserts) > 0
+
+    def step(self) -> list[BatchedGenerationResponse]:
+        results: list[BatchedGenerationResponse] = []
+        # Process all active requests - return one token and complete
+        for uid, (command_id, task_id) in list(self._active_requests.items()):
+            results.append(
+                BatchedGenerationResponse(
+                    command_id=command_id,
+                    task_id=task_id,
+                    response=GenerationResponse(
+                        token=0,
+                        text="hi",
+                        finish_reason="stop",
+                    ),
+                )
+            )
+            del self._active_requests[uid]
+        return results
+
+    @property
+    def has_active_requests(self) -> bool:
+        return len(self._active_requests) > 0
+
+    @property
+    def active_count(self) -> int:
+        return len(self._active_requests)
+
+    @property
+    def pending_insert_count(self) -> int:
+        return len(self._pending_inserts)
+
+
 @pytest.fixture
 def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
-    # initialize_mlx returns a "group" equal to 1
-    monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(1))
-    monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, 1)))
+    # initialize_mlx returns a fake "group" (non-None for state machine)
+    monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MagicMock()))
+    monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((MagicMock(), MagicMock())))
     monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1))
     monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", nothin)
-
-    def fake_generate(*_1: object, **_2: object):
-        yield GenerationResponse(token=0, text="hi", finish_reason="stop")
-
-    monkeypatch.setattr(mlx_runner, "mlx_generate", fake_generate)
+    monkeypatch.setattr(mlx_runner, "BatchGenerationEngine", FakeBatchEngine)
 
 
 def _run(tasks: Iterable[Task]):
@@ -148,7 +224,8 @@ def _run(tasks: Iterable[Task]):
         return event_receiver.collect()
 
 
-def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
+def test_chat_completion_generates_and_completes(patch_out_mlx: pytest.MonkeyPatch):
+    """Verify chat completion generates tokens, completes, and runner returns to Ready."""
     events = _run([INIT_TASK, LOAD_TASK, WARMUP_TASK, CHAT_TASK, SHUTDOWN_TASK])
 
     expected_chunk = ChunkGenerated(
@@ -191,7 +268,9 @@ def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
                 task_id=CHAT_COMPLETION_TASK_ID, task_status=TaskStatus.Running
             ),
             TaskAcknowledged(task_id=CHAT_COMPLETION_TASK_ID),
-            RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerRunning()),
+            RunnerStatusUpdated(
+                runner_id=RUNNER_1_ID, runner_status=RunnerRunning(active_requests=1)
+            ),
             expected_chunk,
             TaskStatusUpdated(
                 task_id=CHAT_COMPLETION_TASK_ID, task_status=TaskStatus.Complete
@@ -206,7 +285,6 @@ def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
             TaskStatusUpdated(
                 task_id=SHUTDOWN_TASK_ID, task_status=TaskStatus.Complete
             ),
-            # SPECIAL EXCEPTION FOR RUNNER SHUTDOWN
             RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerShutdown()),
         ],
     )

tests/headless_runner.py

@@ -49,14 +49,12 @@ class Tests(BaseModel):
     kind: typing.Literal["init", "warmup", "inference"]
 
 
-hn = socket.gethostname()
 mp.set_start_method("spawn", force=True)
 logger_setup(None)
 
 
 async def main():
     logger.info("starting cool server majig")
-    logger.info(hn)
     await assert_downloads()
     cfg = Config()
     cfg.bind = "0.0.0.0:52415"
@@ -81,20 +79,41 @@ async def main():
 async def assert_downloads():
     sd = exo_shard_downloader()
     # await sd.ensure_shard(await build_full_shard(MODEL_CARDS["qwen3-0.6b"].model_id))
-    await sd.ensure_shard(await build_full_shard(MODEL_CARDS["llama-3.2-1b"].model_id))
+    await sd.ensure_shard(
+        await build_full_shard(MODEL_CARDS["llama-3.1-8b-bf16"].model_id)
+    )
+    await sd.ensure_shard(await build_full_shard(MODEL_CARDS["qwen3-30b"].model_id))
+    await sd.ensure_shard(
+        await build_full_shard(MODEL_CARDS["gpt-oss-120b-MXFP4-Q8"].model_id)
+    )
+    await sd.ensure_shard(
+        await build_full_shard(MODEL_CARDS["gpt-oss-20b-4bit"].model_id)
+    )
+    await sd.ensure_shard(
+        await build_full_shard(MODEL_CARDS["glm-4.7-8bit-gs32"].model_id)
+    )
+    await sd.ensure_shard(
+        await build_full_shard(MODEL_CARDS["minimax-m2.1-8bit"].model_id)
+    )
 
 
 async def ring_backend(test: Tests):
     iid = InstanceId(str(hash(str(test.devs))))
-    return await execute_test(test, ring_instance(test, iid))
+    weird_hn = socket.gethostname()
+    for dev in test.devs:
+        if weird_hn.startswith(dev[0]) or dev[0].startswith(weird_hn):
+            hn = dev[0]
+            break
+    else:
+        raise ValueError(f"{weird_hn} not in {test.devs}")
+    return await execute_test(test, ring_instance(test, iid, hn), hn)
 
 
-def ring_instance(test: Tests, iid: InstanceId) -> Instance:
-    global hn
+def ring_instance(test: Tests, iid: InstanceId, hn: str) -> Instance:
     hbn = [Host(ip="i dont care", port=52416) for _ in test.devs]
     world_size = len(test.devs)
     for i in range(world_size):
-        if hn.startswith(test.devs[i][0]):
+        if test.devs[i][0] == hn:
             hn = test.devs[i][0]
             if i - 1 >= 0:
                 hbn[i - 1] = Host(ip=test.devs[i - 1][1], port=52416)
@@ -102,6 +121,8 @@ def ring_instance(test: Tests, iid: InstanceId) -> Instance:
                 hbn[i + 1] = Host(ip=test.devs[i + 1][1], port=52416)
             hbn[i] = Host(ip="0.0.0.0", port=52416)
             break
+    else:
+        raise ValueError(f"{hn} not in {test.devs}")
 
     meta = MODEL_CARDS[test.model_id].metadata
     instance = MlxRingInstance(
@@ -131,10 +152,10 @@ def ring_instance(test: Tests, iid: InstanceId) -> Instance:
     return instance
 
 
-async def execute_test(test: Tests, instance: Instance):
+async def execute_test(test: Tests, instance: Instance, hn: str):
     world_size = len(test.devs)
     iid = InstanceId(str(hash(str(test.devs))))
-    _handle, recv, send = new_runner(instance)
+    _handle, recv, send = new_runner(instance, hn)
     if world_size > 1:
         send.send(ConnectToGroup(instance_id=iid))
     send.send(LoadModel(instance_id=iid))
@@ -181,17 +202,19 @@ async def execute_test(test: Tests, instance: Instance):
 
 async def jaccl_backend(test: Tests):
     iid = InstanceId(str(hash(str(test.devs))))
-    return await execute_test(test, jaccl_instance(test, iid))
+    weird_hn = socket.gethostname()
+    for dev in test.devs:
+        if weird_hn.startswith(dev[0]) or dev[0].startswith(weird_hn):
+            hn = dev[0]
+            break
+    else:
+        raise ValueError(f"{weird_hn} not in {test.devs}")
+    return await execute_test(test, jaccl_instance(test, iid, hn), hn)
 
 
-def jaccl_instance(test: Tests, iid: InstanceId):
-    global hn
+def jaccl_instance(test: Tests, iid: InstanceId, hn: str):
     meta = MODEL_CARDS[test.model_id].metadata
     world_size = len(test.devs)
-    for name, _ in test.devs:
-        if hn.startswith(name):
-            hn = name
-            break
 
     return MlxJacclInstance(
         instance_id=iid,
@@ -220,6 +243,7 @@ def jaccl_instance(test: Tests, iid: InstanceId):
 
 def new_runner(
     instance: Instance,
+    hn: str,
 ) -> tuple[mp.Process, MpReceiver[Event], MpSender[Task]]:
     bound_instance = BoundInstance(
         instance=instance, bound_runner_id=RunnerId(hn), bound_node_id=NodeId(hn)

tests/start_distributed_test.sh

@@ -34,19 +34,23 @@ done
 devs_raw=$(printf "[\"%s\", \"%s\"], " "${weaved[@]}")
 devs="[${devs_raw%, }]"
 
-for i in "${!ips[@]}"; do  
-  { 
-    req="{
-      \"model_id\": \"llama-3.2-1b\",
-      \"devs\": ${devs},
-      \"kind\": \"inference\"
-     }"
-    echo "req $req"
-    curl -sN \
-      -X POST "http://${ips[$i]}:52415/${kind}" \
-      -H "Content-Type: application/json" -d "$req" \
-    2>&1 | sed "s/^/\n${hostnames[$i]}@${ips[$i]}: /" || echo "curl to ${hostnames[$i]} failed"
-  } &
+model_ids=("qwen3-30b" "gpt-oss-120b-MXFP4-Q8" "kimi-k2-thinking")
+
+for model_id in "${model_ids[@]}"; do
+  for i in "${!ips[@]}"; do  
+    { 
+      req="{
+        \"model_id\": \"${model_id}\",
+        \"devs\": ${devs},
+        \"kind\": \"inference\"
+       }"
+      echo "req $req"
+      curl -sN \
+        -X POST "http://${ips[$i]}:52415/${kind}" \
+        -H "Content-Type: application/json" -d "$req" \
+      2>&1 | sed "s/^/\n${hostnames[$i]}@${ips[$i]}: /" || echo "curl to ${hostnames[$i]} failed" && exit 1
+    } &
+  done
+  wait
 done
 
-wait