fix: enable MLX CPU inference on x86_64 Linux in Docker

Two issues prevented MLX CPU from working on x86_64 in Docker:

1. Missing BLAS/LAPACK libraries: MLX CPU backend requires libblas-dev,
   liblapack-dev, and liblapacke-dev on Linux. Added to apt-get install.

2. g++ wrapper ordering: The -fpermissive wrapper for GCC 14 was installed
   AFTER uv sync, but MLX may compile extensions during install. Moved
   the wrapper BEFORE uv sync so both build-time and runtime JIT
   compilation benefit from the fix.

MLX publishes manylinux_2_35_x86_64 wheels, so this uses the native
CPU backend — no alternative inference framework needed.

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

.dockerignore

@@ -0,0 +1,15 @@
+.venv/
+.direnv/
+target/
+.git/
+.idea/
+.pytest_cache/
+.ruff_cache/
+dashboard/node_modules/
+dashboard/.svelte-kit/
+dashboard/build/
+dist/
+*.pdb
+**/__pycache__
+**/.DS_Store
+.mlx_typings/

.github/workflows/e2e.yml

@@ -0,0 +1,42 @@
+name: e2e-tests
+
+on:
+  push:
+  pull_request:
+    branches:
+      - staging
+      - main
+
+jobs:
+  e2e:
+    runs-on: ubuntu-latest
+    timeout-minutes: 45
+    steps:
+      - name: Free up disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc \
+            /opt/hostedtoolcache /usr/local/share/boost /usr/share/swift \
+            /opt/microsoft /opt/az
+          docker system prune -af
+          df -h /
+
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          lfs: false
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Build E2E image with cache
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: e2e/Dockerfile
+          tags: exo-e2e:latest
+          load: true
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+
+      - name: Run E2E tests
+        run: python3 e2e/run_all.py

.mlx_typings/mlx_lm/models/glm_moe_dsa.pyi

@@ -1,46 +0,0 @@
-"""Type stubs for mlx_lm.models.glm_moe_dsa"""
-
-from dataclasses import dataclass
-from typing import Any, Dict, Optional
-
-from .base import BaseModelArgs
-from .deepseek_v32 import Model as DSV32Model
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str
-    vocab_size: int
-    hidden_size: int
-    index_head_dim: int
-    index_n_heads: int
-    index_topk: 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: 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_parameters: Dict[str, Any]
-    attention_bias: bool
-    rope_scaling: Dict[str, Any] | None
-    rope_theta: float | None
-
-class Model(DSV32Model):
-    def __init__(self, config: ModelArgs) -> None: ...

.mlx_typings/mlx_lm/models/step3p5.pyi

@@ -0,0 +1,151 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    vocab_size: int
+    num_attention_heads: int
+    num_attention_groups: int
+    head_dim: int
+    intermediate_size: int
+    rms_norm_eps: float
+    rope_theta: float
+    rope_scaling: Optional[Dict[str, Any]]
+    max_position_embeddings: int
+    sliding_window: int
+    layer_types: Optional[List[str]]
+    yarn_only_types: Optional[List[str]]
+    partial_rotary_factors: Optional[List[float]]
+    attention_other_setting: Optional[Dict[str, Any]]
+    use_head_wise_attn_gate: bool
+    moe_num_experts: int
+    moe_top_k: int
+    moe_intermediate_size: int
+    share_expert_dim: int
+    moe_layers_enum: Optional[str]
+    moe_router_scaling_factor: float
+    norm_expert_weight: bool
+    swiglu_limits: Optional[List[float]]
+    swiglu_limits_shared: Optional[List[float]]
+    tie_word_embeddings: bool
+
+class Step3p5MLP(nn.Module):
+    hidden_size: int
+    intermediate_size: int
+    gate_proj: nn.Linear
+    up_proj: nn.Linear
+    down_proj: nn.Linear
+    limit: Optional[float]
+
+    def __init__(
+        self, args: ModelArgs, intermediate_size: int, swiglu_limit: float = 0
+    ) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class Step3p5MoEGate(nn.Module):
+    top_k: int
+    n_routed_experts: int
+    routed_scaling_factor: float
+    norm_topk_prob: bool
+    gate: nn.Linear
+    router_bias: mx.array
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
+
+class Step3p5MoE(nn.Module):
+    gate: Step3p5MoEGate
+    switch_mlp: SwitchGLU
+    share_expert: Step3p5MLP
+    sharding_group: Optional[mx.distributed.Group]
+
+    def __init__(self, args: ModelArgs, layer_idx: int) -> None: ...
+    def __call__(self, x: mx.array) -> mx.array: ...
+
+class Step3p5Attention(nn.Module):
+    is_sliding: bool
+    num_heads: int
+    num_kv_heads: int
+    head_dim: int
+    scale: float
+    q_proj: nn.Linear
+    k_proj: nn.Linear
+    v_proj: nn.Linear
+    o_proj: nn.Linear
+    q_norm: nn.Module
+    k_norm: nn.Module
+    use_head_wise_attn_gate: bool
+    g_proj: nn.Linear
+    rope: nn.Module
+
+    def __init__(self, args: ModelArgs, layer_idx: int) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Step3p5DecoderLayer(nn.Module):
+    self_attn: Step3p5Attention
+    is_sliding: bool
+    is_moe_layer: bool
+    mlp: Step3p5MLP | Step3p5MoE
+    input_layernorm: nn.Module
+    post_attention_layernorm: nn.Module
+
+    def __init__(self, args: ModelArgs, layer_idx: int) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array: ...
+
+class Step3p5Model(nn.Module):
+    args: ModelArgs
+    vocab_size: int
+    num_layers: int
+    embed_tokens: nn.Embedding
+    layers: list[Step3p5DecoderLayer]
+    norm: nn.Module
+    _swa_idx: Optional[int]
+    _full_idx: Optional[int]
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        x: mx.array,
+        cache: Optional[List[Any]] = None,
+    ) -> mx.array: ...
+
+class Model(nn.Module):
+    args: ModelArgs
+    model_type: str
+    model: Step3p5Model
+    lm_head: nn.Linear
+
+    def __init__(self, args: ModelArgs) -> None: ...
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[List[Any]] = None,
+    ) -> mx.array: ...
+    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+    def shard(self, group: Optional[mx.distributed.Group] = None) -> None: ...
+    @property
+    def layers(self) -> list[Step3p5DecoderLayer]: ...
+    def make_cache(self) -> list[Any]: ...
+    @property
+    def cast_predicate(self) -> Any: ...
+    @property
+    def quant_predicate(self) -> Any: ...

conftest.py

@@ -0,0 +1 @@
+collect_ignore = ["tests/start_distributed_test.py"]

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -806,6 +806,7 @@
                 isFavorite={favorites.has(group.id)}
                 {selectedModelId}
                 {canModelFit}
+                {getModelFitStatus}
                 onToggleExpand={() => toggleGroupExpanded(group.id)}
                 onSelectModel={handleSelect}
                 {onToggleFavorite}

download_glm5_shard.sh

@@ -1,22 +0,0 @@
-#!/bin/bash
-# Usage: ./download_glm5_shard.sh <start> <end> [local_dir]
-#
-# Split across 4 Macs:
-#   Mac 1: ./download_glm5_shard.sh 1 71
-#   Mac 2: ./download_glm5_shard.sh 72 141
-#   Mac 3: ./download_glm5_shard.sh 142 212
-#   Mac 4: ./download_glm5_shard.sh 213 282
-
-set -euo pipefail
-
-START=${1:?Usage: $0 <start> <end> [local_dir]}
-END=${2:?Usage: $0 <start> <end> [local_dir]}
-LOCAL_DIR="${3:-GLM-5}"
-
-INCLUDES=()
-for i in $(seq "$START" "$END"); do
-  INCLUDES+=(--include "$(printf 'model-%05d-of-00282.safetensors' "$i")")
-done
-
-echo "Downloading safetensors $START-$END to $LOCAL_DIR"
-hf download zai-org/GLM-5 "${INCLUDES[@]}" --local-dir "$LOCAL_DIR"

e2e/Dockerfile

@@ -0,0 +1,58 @@
+# Stage 1: Build the dashboard
+FROM node:22-slim AS dashboard
+WORKDIR /app/dashboard
+COPY dashboard/package.json dashboard/package-lock.json ./
+RUN npm ci
+COPY dashboard/ .
+RUN npm run build
+
+# Stage 2: Build and run exo
+FROM python:3.13-slim
+
+# Install system dependencies
+# libblas-dev/liblapack-dev/liblapacke-dev are required by MLX CPU backend on Linux
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    pkg-config \
+    libssl-dev \
+    libblas-dev \
+    liblapack-dev \
+    liblapacke-dev \
+    curl \
+    protobuf-compiler \
+    iptables \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Rust nightly
+RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain nightly
+ENV PATH="/root/.cargo/bin:${PATH}"
+
+# Wrap g++ with -fpermissive to fix MLX CPU JIT compilation with GCC 14
+# (GCC 14 treats _Float128/_Float32/_Float64 as built-in types, conflicting with MLX-generated code)
+# Must be done BEFORE uv sync so any source builds also get the fix
+RUN mv /usr/bin/g++ /usr/bin/g++.real && \
+    printf '#!/bin/sh\nexec /usr/bin/g++.real -fpermissive "$@"\n' > /usr/bin/g++ && \
+    chmod +x /usr/bin/g++
+
+# Install uv
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
+
+WORKDIR /app
+
+# Copy dependency files first for better layer caching
+COPY pyproject.toml Cargo.toml uv.lock README.md ./
+COPY rust/ ./rust/
+COPY bench/pyproject.toml ./bench/pyproject.toml
+
+# Copy source and resources
+COPY src/ ./src/
+COPY resources/ ./resources/
+
+# Copy built dashboard from stage 1
+COPY --from=dashboard /app/dashboard/build ./dashboard/build/
+
+# Install Python deps and build Rust bindings, then clean up build artifacts
+# to keep the layer small (Rust target/ and cargo registry can be 1-2 GB)
+RUN uv sync && rm -rf /app/rust/target /root/.cargo/registry /root/.cargo/git
+
+CMD [".venv/bin/exo", "-v"]

e2e/conftest.py

@@ -0,0 +1,195 @@
+"""Shared E2E test infrastructure for exo cluster tests."""
+
+import asyncio
+import json
+import os
+import sys
+from pathlib import Path
+from urllib.error import URLError
+from urllib.request import Request, urlopen
+
+E2E_DIR = Path(__file__).parent.resolve()
+TIMEOUT = int(os.environ.get("E2E_TIMEOUT", "120"))
+
+
+class Cluster:
+    """Async wrapper around a docker compose exo cluster."""
+
+    def __init__(self, name: str, overrides: list[str] | None = None):
+        self.name = name
+        self.project = f"e2e-{name}"
+        compose_files = [str(E2E_DIR / "docker-compose.yml")]
+        for path in overrides or []:
+            compose_files.append(str(E2E_DIR / path))
+        self._compose_base = [
+            "docker",
+            "compose",
+            "-p",
+            self.project,
+            *[arg for f in compose_files for arg in ("-f", f)],
+        ]
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *exc):
+        await self.stop()
+
+    async def _run(self, *args: str, check: bool = True) -> str:
+        proc = await asyncio.create_subprocess_exec(
+            *self._compose_base,
+            *args,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+        )
+        stdout, _ = await proc.communicate()
+        output = stdout.decode()
+        if check and proc.returncode != 0:
+            print(output, file=sys.stderr)
+            raise RuntimeError(
+                f"docker compose {' '.join(args)} failed (rc={proc.returncode})"
+            )
+        return output
+
+    async def build(self):
+        # Skip build if the image was pre-built (e.g. in CI with buildx cache)
+        proc = await asyncio.create_subprocess_exec(
+            "docker",
+            "image",
+            "inspect",
+            "exo-e2e:latest",
+            stdout=asyncio.subprocess.DEVNULL,
+            stderr=asyncio.subprocess.DEVNULL,
+        )
+        await proc.wait()
+        if proc.returncode == 0:
+            print("  Using pre-built image (exo-e2e:latest)")
+            return
+        print("  Building images...")
+        await self._run("build", "--quiet")
+
+    async def start(self):
+        print("  Starting cluster...")
+        await self._run("up", "-d")
+
+    async def stop(self):
+        print("  Cleaning up...")
+        await self._run("down", "--timeout", "5", check=False)
+
+    async def logs(self) -> str:
+        return await self._run("logs", check=False)
+
+    async def exec(
+        self, service: str, *cmd: str, check: bool = True
+    ) -> tuple[int, str]:
+        """Run a command inside a running container. Returns (returncode, output)."""
+        proc = await asyncio.create_subprocess_exec(
+            *self._compose_base,
+            "exec",
+            "-T",
+            service,
+            *cmd,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+        )
+        stdout, _ = await proc.communicate()
+        output = stdout.decode()
+        if check and proc.returncode != 0:
+            raise RuntimeError(
+                f"exec {' '.join(cmd)} in {service} failed (rc={proc.returncode})"
+            )
+        return proc.returncode, output
+
+    async def wait_for(self, description: str, check_fn, timeout: int = TIMEOUT):
+        """Poll check_fn every 2s until it returns True or timeout expires."""
+        print(f"  Waiting for {description}...")
+        deadline = asyncio.get_event_loop().time() + timeout
+        while asyncio.get_event_loop().time() < deadline:
+            if await check_fn():
+                print(f"  {description}")
+                return
+            await asyncio.sleep(2)
+        output = await self.logs()
+        print(f"--- cluster logs ---\n{output}\n---", file=sys.stderr)
+        raise TimeoutError(f"Timed out waiting for {description}")
+
+    async def assert_healthy(self):
+        """Verify the cluster formed correctly: nodes started, discovered each other, elected a master, API responds."""
+
+        async def both_nodes_started():
+            log = await self.logs()
+            return log.count("Starting node") >= 2
+
+        async def nodes_discovered():
+            log = await self.logs()
+            return log.count("ConnectionMessageType.Connected") >= 2
+
+        async def master_elected():
+            log = await self.logs()
+            return "demoting self" in log
+
+        async def api_responding():
+            try:
+                with urlopen("http://localhost:52415/v1/models", timeout=3) as resp:
+                    return resp.status == 200
+            except (URLError, OSError):
+                return False
+
+        await self.wait_for("Both nodes started", both_nodes_started)
+        await self.wait_for("Nodes discovered each other", nodes_discovered)
+        await self.wait_for("Master election resolved", master_elected)
+        await self.wait_for("API responding", api_responding)
+
+    async def _api(
+        self, method: str, path: str, body: dict | None = None, timeout: int = 30
+    ) -> dict:
+        """Make an API request to the cluster. Returns parsed JSON."""
+        url = f"http://localhost:52415{path}"
+        data = json.dumps(body).encode() if body else None
+        req = Request(
+            url, data=data, headers={"Content-Type": "application/json"}, method=method
+        )
+        loop = asyncio.get_event_loop()
+        resp_bytes = await loop.run_in_executor(
+            None, lambda: urlopen(req, timeout=timeout).read()
+        )
+        return json.loads(resp_bytes)
+
+    async def place_model(self, model: str, timeout: int = 600):
+        """Place a model instance on the cluster (triggers download) and wait until it's ready."""
+        await self._api("POST", "/place_instance", {"model_id": model})
+
+        async def model_ready():
+            try:
+                resp = await self._api("GET", "/v1/models")
+                return any(m.get("id") == model for m in resp.get("data", []))
+            except Exception:
+                return False
+
+        await self.wait_for(f"Model {model} ready", model_ready, timeout=timeout)
+
+    async def chat(
+        self, model: str, messages: list[dict], timeout: int = 600, **kwargs
+    ) -> dict:
+        """Send a chat completion request. Retries until model is downloaded and inference completes."""
+        body = json.dumps({"model": model, "messages": messages, **kwargs}).encode()
+        deadline = asyncio.get_event_loop().time() + timeout
+        last_error = None
+
+        while asyncio.get_event_loop().time() < deadline:
+            try:
+                req = Request(
+                    "http://localhost:52415/v1/chat/completions",
+                    data=body,
+                    headers={"Content-Type": "application/json"},
+                )
+                loop = asyncio.get_event_loop()
+                resp_bytes = await loop.run_in_executor(
+                    None, lambda r=req: urlopen(r, timeout=300).read()
+                )
+                return json.loads(resp_bytes)
+            except Exception as e:
+                last_error = e
+                await asyncio.sleep(5)
+
+        raise TimeoutError(f"Chat request failed after {timeout}s: {last_error}")

e2e/docker-compose.yml

@@ -0,0 +1,20 @@
+services:
+  exo-node-1:
+    image: exo-e2e:latest
+    build:
+      context: ..
+      dockerfile: e2e/Dockerfile
+    environment:
+      - EXO_LIBP2P_NAMESPACE=docker-e2e
+    command: [".venv/bin/exo", "-v"]
+    ports:
+      - "52415:52415"
+
+  exo-node-2:
+    image: exo-e2e:latest
+    build:
+      context: ..
+      dockerfile: e2e/Dockerfile
+    environment:
+      - EXO_LIBP2P_NAMESPACE=docker-e2e
+    command: [".venv/bin/exo", "-v"]

e2e/run_all.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+"""Discovers and runs all E2E tests in e2e/test_*.py.
+
+Tests with '# slow' on the first line of their docstring are skipped
+unless --slow is passed or E2E_SLOW=1 is set.
+"""
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+E2E_DIR = Path(__file__).parent.resolve()
+
+
+def is_slow(test_file: Path) -> bool:
+    """Check if the test file is marked as slow (has '# slow' in first 3 lines)."""
+    with open(test_file) as f:
+        for line in f:
+            if line.strip().startswith("#"):
+                continue
+            if line.strip().startswith('"""') or line.strip().startswith("'''"):
+                # Read into the docstring
+                for doc_line in f:
+                    if "slow" in doc_line.lower() and doc_line.strip().startswith(
+                        "slow"
+                    ):
+                        return True
+                    if '"""' in doc_line or "'''" in doc_line:
+                        break
+            break
+    return False
+
+
+def main():
+    run_slow = "--slow" in sys.argv or os.environ.get("E2E_SLOW") == "1"
+    if "--update-snapshots" in sys.argv:
+        os.environ["UPDATE_SNAPSHOTS"] = "1"
+    test_files = sorted(E2E_DIR.glob("test_*.py"))
+    if not test_files:
+        print("No test files found")
+        sys.exit(1)
+
+    passed = 0
+    failed = 0
+    skipped = 0
+    failures = []
+
+    for test_file in test_files:
+        name = test_file.stem
+        if is_slow(test_file) and not run_slow:
+            print(f"=== {name} === SKIPPED (slow, use --slow to run)")
+            skipped += 1
+            continue
+
+        print(f"=== {name} ===")
+        result = subprocess.run([sys.executable, str(test_file)])
+        if result.returncode == 0:
+            passed += 1
+        else:
+            failed += 1
+            failures.append(name)
+        print()
+
+    total = passed + failed + skipped
+    print("================================")
+    print(
+        f"{passed}/{total} tests passed" + (f", {skipped} skipped" if skipped else "")
+    )
+
+    if failed:
+        print(f"Failed: {' '.join(failures)}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

e2e/snapshot.py

@@ -0,0 +1,69 @@
+"""Snapshot testing infrastructure for E2E tests.
+
+Provides deterministic regression testing by comparing inference output
+against saved snapshots. On first run, snapshots are created automatically.
+Set UPDATE_SNAPSHOTS=1 to regenerate snapshots when output intentionally changes.
+
+Snapshots are stored per-architecture (e.g. snapshots/x86_64/, snapshots/arm64/)
+since floating-point results differ between CPU architectures.
+"""
+
+import difflib
+import json
+import os
+import platform
+from pathlib import Path
+
+ARCH = platform.machine()
+SNAPSHOTS_DIR = Path(__file__).parent / "snapshots" / ARCH
+
+
+def assert_snapshot(
+    name: str,
+    content: str,
+    metadata: dict,
+) -> None:
+    """Compare content against a saved snapshot, or create one if missing.
+
+    Args:
+        name: Snapshot identifier (used as filename: snapshots/{arch}/{name}.json).
+        content: The actual inference output to compare.
+        metadata: Additional context stored alongside content (model, seed, etc.).
+                  Not used for comparison -- purely documentary.
+
+    Raises:
+        AssertionError: If content doesn't match the saved snapshot.
+
+    Environment:
+        UPDATE_SNAPSHOTS=1: Overwrite existing snapshot with actual content.
+    """
+    snapshot_file = SNAPSHOTS_DIR / f"{name}.json"
+    update = os.environ.get("UPDATE_SNAPSHOTS") == "1"
+
+    if snapshot_file.exists() and not update:
+        snapshot = json.loads(snapshot_file.read_text())
+        expected = snapshot["content"]
+        if content != expected:
+            diff = "\n".join(
+                difflib.unified_diff(
+                    expected.splitlines(),
+                    content.splitlines(),
+                    fromfile=f"expected ({snapshot_file.relative_to(SNAPSHOTS_DIR.parent.parent)})",
+                    tofile="actual",
+                    lineterm="",
+                )
+            )
+            raise AssertionError(
+                f"Snapshot mismatch for '{name}' on {ARCH}!\n\n"
+                f"{diff}\n\n"
+                f"Expected: {expected!r}\n"
+                f"Actual:   {content!r}\n\n"
+                f"To update: UPDATE_SNAPSHOTS=1 python3 e2e/run_all.py"
+            )
+        print(f"  Output matches snapshot ({ARCH}/{snapshot_file.name})")
+    else:
+        SNAPSHOTS_DIR.mkdir(parents=True, exist_ok=True)
+        snapshot_data = {**metadata, "arch": ARCH, "content": content}
+        snapshot_file.write_text(json.dumps(snapshot_data, indent=2) + "\n")
+        action = "Updated" if update else "Created"
+        print(f"  {action} snapshot: {ARCH}/{snapshot_file.name}")

e2e/test_cluster_formation.py

@@ -0,0 +1,22 @@
+"""Test: Basic cluster formation.
+
+Verifies two nodes discover each other, elect a master, and the API responds.
+"""
+
+import asyncio
+import sys
+
+sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent))
+from conftest import Cluster
+
+
+async def main():
+    async with Cluster("cluster_formation") as cluster:
+        await cluster.build()
+        await cluster.start()
+        await cluster.assert_healthy()
+        print("PASSED: cluster_formation")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

e2e/test_inference_snapshot.py

@@ -0,0 +1,60 @@
+"""Test: Deterministic inference output (snapshot test).
+
+Sends a chat completion request with a fixed seed,
+then verifies the output matches a known-good snapshot. This ensures
+inference produces consistent results across runs.
+
+Uses MLX CPU backend in Docker on x86 Linux.
+"""
+
+import asyncio
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent))
+from snapshot import assert_snapshot
+
+from conftest import Cluster
+
+MODEL = "mlx-community/Qwen3-0.6B-4bit"
+SEED = 42
+PROMPT = "What is 2+2? Reply with just the number."
+MAX_TOKENS = 32
+
+
+async def main():
+    async with Cluster("inference_snapshot") as cluster:
+        await cluster.build()
+        await cluster.start()
+        await cluster.assert_healthy()
+
+        print(f"  Launching model {MODEL}...")
+        await cluster.place_model(MODEL)
+
+        print(f"  Sending chat completion (seed={SEED})...")
+        resp = await cluster.chat(
+            model=MODEL,
+            messages=[{"role": "user", "content": PROMPT}],
+            seed=SEED,
+            max_tokens=MAX_TOKENS,
+        )
+
+        content = resp["choices"][0]["message"]["content"]
+        print(f"  Response: {content!r}")
+
+        assert_snapshot(
+            name="inference_snapshot",
+            content=content,
+            metadata={
+                "model": MODEL,
+                "seed": SEED,
+                "prompt": PROMPT,
+                "max_tokens": MAX_TOKENS,
+            },
+        )
+
+        print("PASSED: inference_snapshot")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

e2e/test_no_internet.py

@@ -0,0 +1,47 @@
+"""Test: Cluster works without internet access.
+
+Verifies exo functions correctly when containers can talk to each other
+but cannot reach the internet. Uses iptables to block all outbound traffic
+except private subnets and multicast (for mDNS discovery).
+"""
+
+import asyncio
+import sys
+
+sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent))
+from conftest import Cluster
+
+
+async def main():
+    async with Cluster(
+        "no_internet",
+        overrides=["tests/no_internet/docker-compose.override.yml"],
+    ) as cluster:
+        await cluster.build()
+        await cluster.start()
+        await cluster.assert_healthy()
+
+        # Verify internet is actually blocked from inside the containers
+        for node in ["exo-node-1", "exo-node-2"]:
+            rc, _ = await cluster.exec(
+                node,
+                "curl",
+                "-sf",
+                "--max-time",
+                "3",
+                "https://huggingface.co",
+                check=False,
+            )
+            assert rc != 0, f"{node} should not be able to reach the internet"
+            print(f"  {node}: internet correctly blocked")
+
+        # Verify exo detected no internet connectivity
+        log = await cluster.logs()
+        assert "Internet connectivity: False" in log, "exo should detect no internet"
+        print("  exo correctly detected no internet connectivity")
+
+        print("PASSED: no_internet")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

e2e/test_snapshot_code_gen.py

@@ -0,0 +1,58 @@
+"""Test: Code generation snapshot.
+
+Verifies deterministic output for a code generation prompt.
+"""
+
+import asyncio
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent))
+from snapshot import assert_snapshot
+
+from conftest import Cluster
+
+MODEL = "mlx-community/Qwen3-0.6B-4bit"
+SEED = 42
+PROMPT = (
+    "Write a Python function to reverse a string. Only output the code, no explanation."
+)
+MAX_TOKENS = 64
+
+
+async def main():
+    async with Cluster("snapshot_code_gen") as cluster:
+        await cluster.build()
+        await cluster.start()
+        await cluster.assert_healthy()
+
+        print(f"  Launching model {MODEL}...")
+        await cluster.place_model(MODEL)
+
+        print(f"  Sending chat completion (seed={SEED})...")
+        resp = await cluster.chat(
+            model=MODEL,
+            messages=[{"role": "user", "content": PROMPT}],
+            seed=SEED,
+            max_tokens=MAX_TOKENS,
+        )
+
+        content = resp["choices"][0]["message"]["content"]
+        print(f"  Response: {content!r}")
+
+        assert_snapshot(
+            name="snapshot_code_gen",
+            content=content,
+            metadata={
+                "model": MODEL,
+                "seed": SEED,
+                "prompt": PROMPT,
+                "max_tokens": MAX_TOKENS,
+            },
+        )
+
+        print("PASSED: snapshot_code_gen")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

e2e/test_snapshot_edge.py

@@ -0,0 +1,63 @@
+"""Test: Edge case snapshots.
+
+Verifies deterministic output for edge-case prompts: single word input,
+special characters, and unicode.
+"""
+
+import asyncio
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent))
+from snapshot import assert_snapshot
+
+from conftest import Cluster
+
+MODEL = "mlx-community/Qwen3-0.6B-4bit"
+SEED = 42
+MAX_TOKENS = 32
+
+CASES = [
+    ("edge_single_word", "Hi"),
+    ("edge_special_chars", "What does 2 * (3 + 4) / 7 - 1 equal? Use <math> tags."),
+    ("edge_unicode", "Translate 'hello' to Japanese, Chinese, and Korean."),
+]
+
+
+async def main():
+    async with Cluster("snapshot_edge") as cluster:
+        await cluster.build()
+        await cluster.start()
+        await cluster.assert_healthy()
+
+        print(f"  Launching model {MODEL}...")
+        await cluster.place_model(MODEL)
+
+        for snapshot_name, prompt in CASES:
+            print(f"  [{snapshot_name}] Sending: {prompt!r}")
+            resp = await cluster.chat(
+                model=MODEL,
+                messages=[{"role": "user", "content": prompt}],
+                seed=SEED,
+                max_tokens=MAX_TOKENS,
+            )
+
+            content = resp["choices"][0]["message"]["content"]
+            print(f"  [{snapshot_name}] Response: {content!r}")
+
+            assert_snapshot(
+                name=snapshot_name,
+                content=content,
+                metadata={
+                    "model": MODEL,
+                    "seed": SEED,
+                    "prompt": prompt,
+                    "max_tokens": MAX_TOKENS,
+                },
+            )
+
+        print("PASSED: snapshot_edge")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

e2e/test_snapshot_long_output.py

@@ -0,0 +1,56 @@
+"""Test: Longer output snapshot.
+
+Verifies deterministic output with a higher max_tokens (128).
+"""
+
+import asyncio
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent))
+from snapshot import assert_snapshot
+
+from conftest import Cluster
+
+MODEL = "mlx-community/Qwen3-0.6B-4bit"
+SEED = 42
+PROMPT = "Explain how a binary search algorithm works."
+MAX_TOKENS = 128
+
+
+async def main():
+    async with Cluster("snapshot_long_output") as cluster:
+        await cluster.build()
+        await cluster.start()
+        await cluster.assert_healthy()
+
+        print(f"  Launching model {MODEL}...")
+        await cluster.place_model(MODEL)
+
+        print(f"  Sending chat completion (seed={SEED}, max_tokens={MAX_TOKENS})...")
+        resp = await cluster.chat(
+            model=MODEL,
+            messages=[{"role": "user", "content": PROMPT}],
+            seed=SEED,
+            max_tokens=MAX_TOKENS,
+        )
+
+        content = resp["choices"][0]["message"]["content"]
+        print(f"  Response: {content!r}")
+
+        assert_snapshot(
+            name="snapshot_long_output",
+            content=content,
+            metadata={
+                "model": MODEL,
+                "seed": SEED,
+                "prompt": PROMPT,
+                "max_tokens": MAX_TOKENS,
+            },
+        )
+
+        print("PASSED: snapshot_long_output")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

e2e/test_snapshot_multi_model.py

@@ -0,0 +1,72 @@
+"""Test: Multi-model snapshot tests.
+slow
+
+Verifies deterministic output across different model architectures to catch
+model-specific regressions. Each model uses its own snapshot file.
+Run with: python3 e2e/run_all.py --slow  or  E2E_SLOW=1 python3 e2e/run_all.py
+"""
+
+import asyncio
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent))
+from snapshot import assert_snapshot
+
+from conftest import Cluster
+
+SEED = 42
+PROMPT = "What is the capital of France?"
+MAX_TOKENS = 32
+
+MODELS = [
+    "mlx-community/SmolLM2-135M-Instruct",
+    "mlx-community/Llama-3.2-1B-Instruct-4bit",
+    "mlx-community/gemma-2-2b-it-4bit",
+]
+
+
+async def main():
+    async with Cluster("snapshot_multi_model") as cluster:
+        await cluster.build()
+        await cluster.start()
+        await cluster.assert_healthy()
+
+        for model in MODELS:
+            short_name = (
+                model.split("/")[-1].lower().replace("-", "_").replace(".", "_")
+            )
+            snapshot_name = f"snapshot_multi_{short_name}"
+
+            print(f"  Launching model {model}...")
+            await cluster.place_model(model)
+
+            print(f"  Sending chat completion (seed={SEED})...")
+            resp = await cluster.chat(
+                model=model,
+                messages=[{"role": "user", "content": PROMPT}],
+                seed=SEED,
+                max_tokens=MAX_TOKENS,
+            )
+
+            content = resp["choices"][0]["message"]["content"]
+            print(f"  [{short_name}] Response: {content!r}")
+
+            assert_snapshot(
+                name=snapshot_name,
+                content=content,
+                metadata={
+                    "model": model,
+                    "seed": SEED,
+                    "prompt": PROMPT,
+                    "max_tokens": MAX_TOKENS,
+                },
+            )
+
+            print(f"  [{short_name}] PASSED")
+
+        print("PASSED: snapshot_multi_model")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

e2e/test_snapshot_reasoning.py

@@ -0,0 +1,56 @@
+"""Test: Reasoning/math snapshot.
+
+Verifies deterministic output for a simple reasoning prompt.
+"""
+
+import asyncio
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent))
+from snapshot import assert_snapshot
+
+from conftest import Cluster
+
+MODEL = "mlx-community/Qwen3-0.6B-4bit"
+SEED = 42
+PROMPT = "If I have 3 apples and give away 1, how many do I have? Think step by step."
+MAX_TOKENS = 64
+
+
+async def main():
+    async with Cluster("snapshot_reasoning") as cluster:
+        await cluster.build()
+        await cluster.start()
+        await cluster.assert_healthy()
+
+        print(f"  Launching model {MODEL}...")
+        await cluster.place_model(MODEL)
+
+        print(f"  Sending chat completion (seed={SEED})...")
+        resp = await cluster.chat(
+            model=MODEL,
+            messages=[{"role": "user", "content": PROMPT}],
+            seed=SEED,
+            max_tokens=MAX_TOKENS,
+        )
+
+        content = resp["choices"][0]["message"]["content"]
+        print(f"  Response: {content!r}")
+
+        assert_snapshot(
+            name="snapshot_reasoning",
+            content=content,
+            metadata={
+                "model": MODEL,
+                "seed": SEED,
+                "prompt": PROMPT,
+                "max_tokens": MAX_TOKENS,
+            },
+        )
+
+        print("PASSED: snapshot_reasoning")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

e2e/tests/no_internet/docker-compose.override.yml

@@ -0,0 +1,32 @@
+# Block all outbound internet traffic using iptables while preserving:
+#   - Multicast (224.0.0.0/4) for mDNS peer discovery
+#   - Private subnets (10/8, 172.16/12, 192.168/16) for inter-container communication
+#   - Loopback (127/8)
+# Requires NET_ADMIN capability for iptables.
+services:
+  exo-node-1:
+    cap_add:
+      - NET_ADMIN
+    entrypoint: ["/bin/sh", "-c"]
+    command:
+      - |
+        iptables -A OUTPUT -d 127.0.0.0/8 -j ACCEPT
+        iptables -A OUTPUT -d 10.0.0.0/8 -j ACCEPT
+        iptables -A OUTPUT -d 172.16.0.0/12 -j ACCEPT
+        iptables -A OUTPUT -d 192.168.0.0/16 -j ACCEPT
+        iptables -A OUTPUT -d 224.0.0.0/4 -j ACCEPT
+        iptables -A OUTPUT -j REJECT
+        exec .venv/bin/exo -v
+  exo-node-2:
+    cap_add:
+      - NET_ADMIN
+    entrypoint: ["/bin/sh", "-c"]
+    command:
+      - |
+        iptables -A OUTPUT -d 127.0.0.0/8 -j ACCEPT
+        iptables -A OUTPUT -d 10.0.0.0/8 -j ACCEPT
+        iptables -A OUTPUT -d 172.16.0.0/12 -j ACCEPT
+        iptables -A OUTPUT -d 192.168.0.0/16 -j ACCEPT
+        iptables -A OUTPUT -d 224.0.0.0/4 -j ACCEPT
+        iptables -A OUTPUT -j REJECT
+        exec .venv/bin/exo -v

pyproject.toml

@@ -17,7 +17,7 @@ dependencies = [
     "loguru>=0.7.3",
     "exo_pyo3_bindings", # rust bindings
     "anyio==4.11.0",
-    "mlx==0.30.6",
+    "mlx==0.30.6; sys_platform == 'darwin'",
     "mlx[cpu]==0.30.6; sys_platform == 'linux'",
     "mlx-lm==0.30.6",
     "tiktoken>=0.12.0", # required for kimi k2 tokenizer
@@ -64,8 +64,6 @@ members = [
 
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
-#mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", marker = "sys_platform == 'darwin'" }
-mlx-lm = { git = "https://github.com/ml-explore/mlx-lm", branch = "main" }
 #mlx-lm = { git = "https://github.com/davidmcc73/mlx-lm", branch = "stable" }
 # Uncomment to use local mlx/mlx-lm development versions:
 # mlx = { path = "/Users/Shared/mlx", editable=true }

resources/inference_model_cards/mlx-community--GLM-5-8bit.toml

@@ -1,12 +0,0 @@
-model_id = "mlx-community/GLM-5-8bit"
-n_layers = 78
-hidden_size = 6144
-supports_tensor = true
-tasks = ["TextGeneration"]
-family = "glm"
-quantization = "8bit"
-base_model = "GLM-5"
-capabilities = ["text", "thinking"]
-
-[storage_size]
-in_bytes = 790517400864

resources/inference_model_cards/mlx-community--GLM-5-MXFP4-Q8.toml

@@ -1,12 +0,0 @@
-model_id = "mlx-community/GLM-5-MXFP4-Q8"
-n_layers = 78
-hidden_size = 6144
-supports_tensor = true
-tasks = ["TextGeneration"]
-family = "glm"
-quantization = "MXFP4-Q8"
-base_model = "GLM-5"
-capabilities = ["text", "thinking"]
-
-[storage_size]
-in_bytes = 405478939008

resources/inference_model_cards/mlx-community--GLM-5-bf16.toml

@@ -1,12 +0,0 @@
-model_id = "mlx-community/GLM-5"
-n_layers = 78
-hidden_size = 6144
-supports_tensor = true
-tasks = ["TextGeneration"]
-family = "glm"
-quantization = "bf16"
-base_model = "GLM-5"
-capabilities = ["text", "thinking"]
-
-[storage_size]
-in_bytes = 1487822475264

resources/inference_model_cards/mlx-community--SmolLM2-135M-Instruct.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/SmolLM2-135M-Instruct"
+n_layers = 30
+hidden_size = 576
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "llama"
+quantization = "bf16"
+base_model = "SmolLM2 135M"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 269060381

resources/inference_model_cards/mlx-community--Step-3.5-Flash-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Step-3.5-Flash-4bit"
+n_layers = 45
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "step"
+quantization = "4bit"
+base_model = "Step 3.5 Flash"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 114572190076

resources/inference_model_cards/mlx-community--Step-3.5-Flash-6bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Step-3.5-Flash-6bit"
+n_layers = 45
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "step"
+quantization = "6bit"
+base_model = "Step 3.5 Flash"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 159039627774

resources/inference_model_cards/mlx-community--Step-3.5-Flash-8Bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/Step-3.5-Flash-8Bit"
+n_layers = 45
+hidden_size = 4096
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "step"
+quantization = "8bit"
+base_model = "Step 3.5 Flash"
+capabilities = ["text", "thinking"]
+
+[storage_size]
+in_bytes = 209082699847

resources/inference_model_cards/mlx-community--gemma-2-2b-it-4bit.toml

@@ -0,0 +1,12 @@
+model_id = "mlx-community/gemma-2-2b-it-4bit"
+n_layers = 26
+hidden_size = 2304
+supports_tensor = false
+tasks = ["TextGeneration"]
+family = "gemma2"
+quantization = "4bit"
+base_model = "Gemma 2 2B"
+capabilities = ["text"]
+
+[storage_size]
+in_bytes = 1492755242

src/exo/shared/constants.py

@@ -7,17 +7,11 @@ from exo.utils.dashboard_path import find_dashboard, find_resources
 _EXO_HOME_ENV = os.environ.get("EXO_HOME", None)
 
 
-def _resolve_env_path(env_value: str) -> Path:
-    """Resolve an environment variable path: absolute paths are used as-is, relative paths are resolved from home."""
-    p = Path(env_value)
-    return p if p.is_absolute() else Path.home() / p
-
-
 def _get_xdg_dir(env_var: str, fallback: str) -> Path:
     """Get XDG directory, prioritising EXO_HOME environment variable if its set. On non-Linux platforms, default to ~/.exo."""
 
     if _EXO_HOME_ENV is not None:
-        return _resolve_env_path(_EXO_HOME_ENV)
+        return Path.home() / _EXO_HOME_ENV
 
     if sys.platform != "linux":
         return Path.home() / ".exo"
@@ -37,19 +31,15 @@ _EXO_MODELS_DIR_ENV = os.environ.get("EXO_MODELS_DIR", None)
 EXO_MODELS_DIR = (
     EXO_DATA_HOME / "models"
     if _EXO_MODELS_DIR_ENV is None
-    else _resolve_env_path(_EXO_MODELS_DIR_ENV)
+    else Path.home() / _EXO_MODELS_DIR_ENV
 )
 _RESOURCES_DIR_ENV = os.environ.get("EXO_RESOURCES_DIR", None)
 RESOURCES_DIR = (
-    find_resources()
-    if _RESOURCES_DIR_ENV is None
-    else _resolve_env_path(_RESOURCES_DIR_ENV)
+    find_resources() if _RESOURCES_DIR_ENV is None else Path.home() / _RESOURCES_DIR_ENV
 )
 _DASHBOARD_DIR_ENV = os.environ.get("EXO_DASHBOARD_DIR", None)
 DASHBOARD_DIR = (
-    find_dashboard()
-    if _DASHBOARD_DIR_ENV is None
-    else _resolve_env_path(_DASHBOARD_DIR_ENV)
+    find_dashboard() if _DASHBOARD_DIR_ENV is None else Path.home() / _DASHBOARD_DIR_ENV
 )
 
 # Log files (data/logs or cache)

src/exo/shared/models/model_cards.py

@@ -182,7 +182,6 @@ class ConfigData(BaseModel):
     def supports_tensor(self) -> bool:
         return self.architectures in [
             ["Glm4MoeLiteForCausalLM"],
-            ["GlmMoeDsaForCausalLM"],
             ["DeepseekV32ForCausalLM"],
             ["DeepseekV3ForCausalLM"],
             ["Qwen3NextForCausalLM"],
@@ -190,6 +189,7 @@ class ConfigData(BaseModel):
             ["MiniMaxM2ForCausalLM"],
             ["LlamaForCausalLM"],
             ["GptOssForCausalLM"],
+            ["Step3p5ForCausalLM"],
         ]
 
     @model_validator(mode="before")

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

@@ -35,6 +35,9 @@ 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 Qwen3NextDecoderLayer, Qwen3NextSparseMoeBlock
+from mlx_lm.models.step3p5 import Model as Step35Model
+from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP
+from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel
 from transformers.models.qwen3.modeling_qwen3 import Qwen3DecoderLayer
 
 from exo.shared.logging import logger
@@ -160,14 +163,11 @@ class PipelineLastLayer(CustomMlxLayer):
                 output, (self.r + 1) % self.s, group=self.group
             )
             if cache is not None:
-                # CacheList (used by MLA models like DeepSeekV32, GLM MoE DSA)
-                # doesn't have .keys directly; access via first sub-cache.
-                dep_cache = cache[0] if hasattr(cache, "caches") else cache  # type: ignore
-                dep_cache.keys = mx.depends(dep_cache.keys, output)  # type: ignore[reportUnknownMemberType]
+                cache.keys = mx.depends(cache.keys, output)  # type: ignore[reportUnknownMemberType]
             if self.is_prefill:
                 mx.eval(output)
                 if cache is not None:
-                    mx.eval(dep_cache.keys)  # type: ignore
+                    mx.eval(cache.keys)  # type: ignore
 
         if not self.is_prefill:
             output = mx.distributed.all_gather(output, group=self.group)[
@@ -267,6 +267,19 @@ def pipeline_auto_parallel(
             )
         )
 
+    if isinstance(inner_model_instance, Step35InnerModel):
+        inner_model_instance.num_layers = len(layers)
+        sliding_layers = [
+            i for i, layer in enumerate(layers) if getattr(layer, "is_sliding", False)
+        ]
+        full_layers = [
+            i
+            for i, layer in enumerate(layers)
+            if not getattr(layer, "is_sliding", True)
+        ]
+        inner_model_instance._swa_idx = 0 if not sliding_layers else sliding_layers[0]
+        inner_model_instance._full_idx = 0 if not full_layers else full_layers[0]
+
     _set_layers(model, layers)
 
     assert isinstance(layers, list), (
@@ -294,9 +307,7 @@ def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
 
         # Add dependency to last cache entry to ensure distributed ops are evaluated
         if cache is not None:
-            last = cache[-1]  # type: ignore
-            dep_cache = last[0] if hasattr(last, "caches") else last  # type: ignore
-            dep_cache.keys = mx.depends(dep_cache.keys, logits)  # type: ignore
+            cache[-1].state = mx.depends(cache[-1].state, logits)  # type: ignore
 
         return logits
 
@@ -322,9 +333,7 @@ def patch_tensor_model[T](model: T) -> T:
 
         # Add dependency to last cache entry to ensure distributed ops are evaluated
         if cache is not None and len(cache) > 0:  # pyright: ignore[reportAny]
-            last = cache[-1]  # pyright: ignore[reportAny]
-            dep_cache = last[0] if hasattr(last, "caches") else last  # pyright: ignore[reportAny]
-            dep_cache.keys = mx.depends(dep_cache.keys, logits)  # pyright: ignore[reportAny,reportUnknownMemberType]
+            cache[-1].state = mx.depends(cache[-1].state, logits)  # pyright: ignore[reportAny,reportUnknownMemberType]
 
         return logits
 
@@ -434,6 +443,14 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
+    elif isinstance(model, Step35Model):
+        tensor_parallel_sharding_strategy = Step35ShardingStrategy(
+            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)}")
 
@@ -530,13 +547,11 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
         on_timeout: TimeoutCallback | None,
     ) -> nn.Module:
         model = cast(DeepseekV3Model, model)
-
         for layer in model.layers:
             eval_with_timeout(
                 layer.parameters(), timeout_seconds / len(model.layers), on_timeout
             )
-
-            # Shard attention heads
+            # Shard the self attention
             if layer.self_attn.q_lora_rank is None:
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
                     layer.self_attn.q_proj
@@ -546,11 +561,10 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
                     layer.self_attn.q_b_proj
                 )
 
-            layer.self_attn.o_proj = self.sharded_to_all_linear(
-                layer.self_attn.o_proj
-            )
+            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
             layer.self_attn.num_heads //= self.N
 
+            # Logic from upstream mlx
             num_heads = layer.self_attn.num_heads
             sh = self.group.rank() * num_heads
             eh = sh + num_heads
@@ -567,17 +581,12 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
                 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)
 
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
             else:
-                if getattr(layer.mlp, "shared_experts", None) is not None:
-                    self.all_to_sharded_linear_in_place(
-                        layer.mlp.shared_experts.gate_proj
-                    )
-                    self.sharded_to_all_linear_in_place(
-                        layer.mlp.shared_experts.down_proj
-                    )
-                    self.all_to_sharded_linear_in_place(
-                        layer.mlp.shared_experts.up_proj
-                    )
+                self.all_to_sharded_linear_in_place(layer.mlp.shared_experts.gate_proj)
+                self.sharded_to_all_linear_in_place(layer.mlp.shared_experts.down_proj)
+                self.all_to_sharded_linear_in_place(layer.mlp.shared_experts.up_proj)
                 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)
@@ -996,3 +1005,46 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
             layer.mlp.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
             mx.eval(layer)
         return model
+
+
+class Step35ShardingStrategy(TensorParallelShardingStrategy):
+    def shard_model(
+        self,
+        model: nn.Module,
+        timeout_seconds: float,
+        on_timeout: TimeoutCallback | None,
+    ) -> nn.Module:
+        model = cast(Step35Model, model)
+
+        for layer in model.layers:
+            eval_with_timeout(
+                layer.parameters(), timeout_seconds / len(model.layers), on_timeout
+            )
+            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_heads //= self.N
+            layer.self_attn.num_kv_heads //= self.N
+
+            if getattr(layer.self_attn, "use_head_wise_attn_gate", False):
+                layer.self_attn.g_proj = self.all_to_sharded_linear(
+                    layer.self_attn.g_proj
+                )
+
+            if isinstance(layer.mlp, Step35MLP):
+                layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj)
+                layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
+                layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
+            else:
+                layer.mlp.sharding_group = self.group
+                self.all_to_sharded_linear_in_place(layer.mlp.share_expert.gate_proj)
+                self.all_to_sharded_linear_in_place(layer.mlp.share_expert.up_proj)
+                self.sharded_to_all_linear_in_place(layer.mlp.share_expert.down_proj)
+                self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
+                self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
+                self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
+
+            mx.eval(layer)
+        return model

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

@@ -311,12 +311,10 @@ def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
     model_id_lower = model_id.lower()
     if "kimi-k2" in model_id_lower:
         return [163586]
-    elif "glm-5" in model_id_lower or "glm-4.7" in model_id_lower:
-        # For GLM-5 and GLM-4.7
+    elif "glm-4.7-flash" in model_id_lower:
         # 154820: <|endoftext|>, 154827: <|user|>, 154829: <|observation|>
         return [154820, 154827, 154829]
     elif "glm" in model_id_lower:
-        # For GLM-4.5 and older
         return [151336, 151329, 151338]
     return None
 

src/exo/worker/runner/runner.py

@@ -295,8 +295,8 @@ def main(
                             patch_kimi_tokenizer(tokenizer)
 
                         # GLM models need patched parser (upstream has bug with None regex match)
-                        elif "glm-4" in shard_metadata.model_card.model_id.lower():
-                            patch_glm4_tokenizer(tokenizer)
+                        elif "glm" in shard_metadata.model_card.model_id.lower():
+                            patch_glm_tokenizer(tokenizer)
 
                         # GPT-OSS specific parsing to match other model formats.
                         elif isinstance(model, GptOssModel):
@@ -863,7 +863,7 @@ def patch_kimi_tokenizer(tokenizer: TokenizerWrapper):
     tokenizer._tool_parser = parse_tool_call
 
 
-def patch_glm4_tokenizer(tokenizer: TokenizerWrapper):
+def patch_glm_tokenizer(tokenizer: TokenizerWrapper):
     """
     Fixed version of mlx_lm's glm47 tool parser that handles regex match failures.
     """

uv.lock

@@ -416,9 +416,9 @@ requires-dist = [
     { name = "hypercorn", specifier = ">=0.18.0" },
     { name = "loguru", specifier = ">=0.7.3" },
     { name = "mflux", specifier = "==0.15.5" },
-    { name = "mlx", specifier = "==0.30.6" },
+    { name = "mlx", marker = "sys_platform == 'darwin'", specifier = "==0.30.6" },
     { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.6" },
-    { name = "mlx-lm", git = "https://github.com/ml-explore/mlx-lm?branch=main" },
+    { name = "mlx-lm", specifier = "==0.30.6" },
     { name = "msgspec", specifier = ">=0.19.0" },
     { name = "openai-harmony", specifier = ">=0.0.8" },
     { name = "pillow", specifier = ">=11.0,<12.0" },
@@ -1098,8 +1098,8 @@ wheels = [
 
 [[package]]
 name = "mlx-lm"
-version = "0.30.7"
-source = { git = "https://github.com/ml-explore/mlx-lm?branch=main#bcf630614ffb5624bcb19870a7bcb0d847e6e98f" }
+version = "0.30.6"
+source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", marker = "sys_platform == 'darwin'" },
@@ -1109,6 +1109,10 @@ dependencies = [
     { name = "sentencepiece", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
 ]
+sdist = { url = "https://files.pythonhosted.org/packages/76/cb/815deddc8699b1f694d7e1f9cbed52934c03a8b49432c8add72932bb2f0b/mlx_lm-0.30.6.tar.gz", hash = "sha256:807e042d7040268f1b19190b7eaefd8b2efbff5590a65460974ad4225b91dda1", size = 271733, upload-time = "2026-02-04T21:27:45.741Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/20/5f/01d281f1fa8a1521d5936659beb4f5ab1f32b463d059263cf9d4cef969d9/mlx_lm-0.30.6-py3-none-any.whl", hash = "sha256:a7405bd581eacc4bf8209d7a6b7f23629585a0d7c6740c2a97e51fee35b3b0e1", size = 379451, upload-time = "2026-02-04T21:27:43.222Z" },
+]
 
 [[package]]
 name = "mlx-metal"