← back to Exo
Cleanup on #1952 (#1960)
f6e418ed237b9fb0a4b32601243cd97889ca2757 · 2026-04-22 17:05:49 +0100 · rltakashige
Files touched
M .mlx_typings/mlx/core/__init__.pyiM .mlx_typings/mlx/nn/layers/base.pyiM .mlx_typings/mlx_lm/models/base.pyiA .mlx_typings/mlx_lm/models/gpt_oss.pyiA .mlx_typings/mlx_lm/models/minimax.pyiM .mlx_typings/mlx_lm/models/nemotron_h.pyiM .mlx_typings/mlx_lm/models/qwen3_next.pyiM bench/eval_tool_calls.pyM bench/harness.pyM src/exo/worker/engines/mlx/auto_parallel.pyM src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.pyM src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.pyA src/exo/worker/tests/unittests/test_mlx/test_tp_bit_exact.py
Diff
commit f6e418ed237b9fb0a4b32601243cd97889ca2757
Author: rltakashige <rl.takashige@gmail.com>
Date: Wed Apr 22 17:05:49 2026 +0100
Cleanup on #1952 (#1960)
---
.mlx_typings/mlx/core/__init__.pyi | 11 +-
.mlx_typings/mlx/nn/layers/base.pyi | 4 +
.mlx_typings/mlx_lm/models/base.pyi | 10 +-
.mlx_typings/mlx_lm/models/gpt_oss.pyi | 103 ++++++
.mlx_typings/mlx_lm/models/minimax.pyi | 94 +++++
.mlx_typings/mlx_lm/models/nemotron_h.pyi | 14 +
.mlx_typings/mlx_lm/models/qwen3_next.pyi | 1 +
bench/eval_tool_calls.py | 119 +++++--
bench/harness.py | 2 +-
src/exo/worker/engines/mlx/auto_parallel.py | 24 +-
.../unittests/test_mlx/test_kv_prefix_cache.py | 12 +-
.../test_mlx/test_prefix_cache_architectures.py | 47 ++-
.../tests/unittests/test_mlx/test_tp_bit_exact.py | 395 +++++++++++++++++++++
13 files changed, 758 insertions(+), 78 deletions(-)
diff --git a/.mlx_typings/mlx/core/__init__.pyi b/.mlx_typings/mlx/core/__init__.pyi
index 421e75db..b0479bf2 100644
--- a/.mlx_typings/mlx/core/__init__.pyi
+++ b/.mlx_typings/mlx/core/__init__.pyi
@@ -1767,12 +1767,12 @@ def clip(
array: The clipped array.
"""
-def compile(
- fun: Callable,
+def compile[F: Callable[..., object]](
+ fun: F,
inputs: object | None = ...,
outputs: object | None = ...,
shapeless: bool = ...,
-) -> Callable:
+) -> F:
"""
Returns a compiled function which produces the same output as ``fun``.
@@ -2915,8 +2915,8 @@ def gather_mm(
a: array,
b: array,
/,
- lhs_indices: array,
- rhs_indices: array,
+ lhs_indices: array | None = ...,
+ rhs_indices: array | None = ...,
*,
sorted_indices: bool = ...,
stream: Stream | Device | None = ...,
@@ -4707,6 +4707,7 @@ def softmax(
/,
axis: int | Sequence[int] | None = ...,
*,
+ precise: bool = ...,
stream: Stream | Device | None = ...,
) -> array:
"""
diff --git a/.mlx_typings/mlx/nn/layers/base.pyi b/.mlx_typings/mlx/nn/layers/base.pyi
index fcd8cf53..e9055679 100644
--- a/.mlx_typings/mlx/nn/layers/base.pyi
+++ b/.mlx_typings/mlx/nn/layers/base.pyi
@@ -57,6 +57,10 @@ class Module(dict):
def __init__(self) -> None:
"""Should be called by the subclasses of ``Module``."""
+ def __getitem__(self, key: str) -> mx.array | Module: ...
+ def get(
+ self, key: str, default: mx.array | Module | None = ...
+ ) -> mx.array | Module | None: ...
@property
def training(self): # -> bool:
"""Boolean indicating if the model is in training mode."""
diff --git a/.mlx_typings/mlx_lm/models/base.pyi b/.mlx_typings/mlx_lm/models/base.pyi
index e549e624..06b16e19 100644
--- a/.mlx_typings/mlx_lm/models/base.pyi
+++ b/.mlx_typings/mlx_lm/models/base.pyi
@@ -3,7 +3,7 @@ This type stub file was generated by pyright.
"""
from dataclasses import dataclass
-from typing import Optional
+from typing import Any, Optional
import mlx.core as mx
@@ -37,10 +37,10 @@ def quantized_scaled_dot_product_attention(
bits: int = ...,
) -> mx.array: ...
def scaled_dot_product_attention(
- queries,
- keys,
- values,
- cache,
+ queries: mx.array,
+ keys: mx.array,
+ values: mx.array,
+ cache: Optional[Any],
scale: float,
mask: Optional[mx.array],
sinks: Optional[mx.array] = ...,
diff --git a/.mlx_typings/mlx_lm/models/gpt_oss.pyi b/.mlx_typings/mlx_lm/models/gpt_oss.pyi
new file mode 100644
index 00000000..dc4ba878
--- /dev/null
+++ b/.mlx_typings/mlx_lm/models/gpt_oss.pyi
@@ -0,0 +1,103 @@
+"""Type stubs for mlx_lm.models.gpt_oss"""
+
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import KVCache
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+ model_type: str
+ hidden_size: int
+ intermediate_size: int
+ num_hidden_layers: int
+ num_attention_heads: int
+ num_key_value_heads: int
+ num_local_experts: int
+ num_experts_per_tok: int
+ vocab_size: int
+ rms_norm_eps: float
+ sliding_window: int
+ layer_types: Optional[List[str]]
+
+def mlx_topk(a: mx.array, k: int, axis: int = -1) -> tuple[mx.array, mx.array]: ...
+
+class AttentionBlock(nn.Module):
+ head_dim: int
+ num_attention_heads: int
+ num_key_value_heads: int
+ num_key_value_groups: int
+ sinks: mx.array
+ q_proj: nn.Linear
+ k_proj: nn.Linear
+ v_proj: nn.Linear
+ o_proj: nn.Linear
+ sm_scale: float
+ rope: nn.Module
+
+ def __init__(self, config: ModelArgs) -> None: ...
+ def __call__(
+ self,
+ x: mx.array,
+ mask: Optional[mx.array] = None,
+ cache: Optional[Any] = None,
+ ) -> mx.array: ...
+
+class TransformerBlock(nn.Module):
+ self_attn: AttentionBlock
+ mlp: MLPBlock
+
+ def __init__(self, config: ModelArgs) -> None: ...
+ def __call__(
+ self,
+ x: mx.array,
+ mask: Optional[mx.array] = None,
+ cache: Optional[Any] = None,
+ ) -> mx.array: ...
+
+class MLPBlock(nn.Module):
+ hidden_size: int
+ num_local_experts: int
+ num_experts_per_tok: int
+ experts: SwitchGLU
+ router: nn.Linear
+ sharding_group: Optional[mx.distributed.Group]
+
+ def __init__(self, config: ModelArgs) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+
+class GptOssMoeModel(nn.Module):
+ embed_tokens: nn.Embedding
+ norm: nn.RMSNorm
+ layer_types: List[str]
+ layers: list[TransformerBlock]
+ window_size: int
+ swa_idx: int
+ ga_idx: int
+
+ def __init__(self, args: ModelArgs) -> None: ...
+ def __call__(
+ self,
+ inputs: mx.array,
+ cache: Optional[Any] = None,
+ ) -> mx.array: ...
+
+class Model(nn.Module):
+ model_type: str
+ model: GptOssMoeModel
+ lm_head: nn.Linear
+
+ def __init__(self, args: ModelArgs) -> None: ...
+ def __call__(
+ self,
+ inputs: mx.array,
+ cache: Optional[Any] = None,
+ ) -> mx.array: ...
+ @property
+ def layers(self) -> list[nn.Module]: ...
+ def make_cache(self) -> list[KVCache]: ...
diff --git a/.mlx_typings/mlx_lm/models/minimax.pyi b/.mlx_typings/mlx_lm/models/minimax.pyi
new file mode 100644
index 00000000..06806d86
--- /dev/null
+++ b/.mlx_typings/mlx_lm/models/minimax.pyi
@@ -0,0 +1,94 @@
+"""Type stubs for mlx_lm.models.minimax"""
+
+from dataclasses import dataclass
+from typing import Any, 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
+ intermediate_size: int
+ num_hidden_layers: int
+ num_attention_heads: int
+ num_key_value_heads: int
+ num_local_experts: int
+ num_experts_per_tok: int
+ max_position_embeddings: int
+
+class MiniMaxAttention(nn.Module):
+ num_heads: int
+ num_attention_heads: int
+ num_key_value_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
+ rope: nn.Module
+
+ def __init__(self, args: ModelArgs) -> None: ...
+ def __call__(
+ self,
+ x: mx.array,
+ mask: Optional[mx.array] = None,
+ cache: Optional[Any] = None,
+ ) -> mx.array: ...
+
+class MiniMaxSparseMoeBlock(nn.Module):
+ num_experts_per_tok: int
+ gate: nn.Linear
+ switch_mlp: SwitchGLU
+ e_score_correction_bias: mx.array
+ sharding_group: Optional[mx.distributed.Group]
+
+ def __init__(self, args: ModelArgs) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+
+class MiniMaxDecoderLayer(nn.Module):
+ self_attn: MiniMaxAttention
+ block_sparse_moe: MiniMaxSparseMoeBlock
+ input_layernorm: nn.RMSNorm
+ post_attention_layernorm: nn.RMSNorm
+
+ def __init__(self, args: ModelArgs) -> None: ...
+ def __call__(
+ self,
+ x: mx.array,
+ mask: Optional[mx.array] = None,
+ cache: Optional[Any] = None,
+ ) -> mx.array: ...
+
+class MiniMaxModel(nn.Module):
+ embed_tokens: nn.Embedding
+ layers: list[MiniMaxDecoderLayer]
+ norm: nn.RMSNorm
+
+ def __init__(self, args: ModelArgs) -> None: ...
+ def __call__(
+ self,
+ inputs: mx.array,
+ cache: Optional[Any] = None,
+ ) -> mx.array: ...
+
+class Model(nn.Module):
+ model_type: str
+ model: MiniMaxModel
+ lm_head: nn.Linear
+
+ def __init__(self, args: ModelArgs) -> None: ...
+ def __call__(
+ self,
+ inputs: mx.array,
+ cache: Optional[Any] = None,
+ ) -> mx.array: ...
+ @property
+ def layers(self) -> list[MiniMaxDecoderLayer]: ...
diff --git a/.mlx_typings/mlx_lm/models/nemotron_h.pyi b/.mlx_typings/mlx_lm/models/nemotron_h.pyi
index 6f38f3ad..32b07271 100644
--- a/.mlx_typings/mlx_lm/models/nemotron_h.pyi
+++ b/.mlx_typings/mlx_lm/models/nemotron_h.pyi
@@ -92,6 +92,15 @@ class NemotronHAttention(nn.Module):
cache: Optional[KVCache] = None,
) -> mx.array: ...
+class MoEGate(nn.Module):
+ config: ModelArgs
+ top_k: int
+ norm_topk_prob: bool
+ weight: mx.array
+
+ def __init__(self, config: ModelArgs) -> None: ...
+ def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
+
class NemotronHMLP(nn.Module):
up_proj: nn.Linear
down_proj: nn.Linear
@@ -102,9 +111,14 @@ class NemotronHMLP(nn.Module):
def __call__(self, x: mx.array) -> mx.array: ...
class NemotronHMoE(nn.Module):
+ config: ModelArgs
num_experts_per_tok: int
+ moe_latent_size: Optional[int]
switch_mlp: SwitchMLP
+ gate: MoEGate
shared_experts: NemotronHMLP
+ fc1_latent_proj: nn.Linear
+ fc2_latent_proj: nn.Linear
def __init__(self, config: ModelArgs) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
diff --git a/.mlx_typings/mlx_lm/models/qwen3_next.pyi b/.mlx_typings/mlx_lm/models/qwen3_next.pyi
index 649669cf..b36ac6b4 100644
--- a/.mlx_typings/mlx_lm/models/qwen3_next.pyi
+++ b/.mlx_typings/mlx_lm/models/qwen3_next.pyi
@@ -71,6 +71,7 @@ class Qwen3NextAttention(nn.Module):
class Qwen3NextSparseMoeBlock(nn.Module):
norm_topk_prob: bool
num_experts: int
+ num_experts_per_tok: int
top_k: int
gate: nn.Linear
switch_mlp: SwitchGLU
diff --git a/bench/eval_tool_calls.py b/bench/eval_tool_calls.py
index 18071473..c2839fcf 100644
--- a/bench/eval_tool_calls.py
+++ b/bench/eval_tool_calls.py
@@ -3,11 +3,13 @@ from __future__ import annotations
import argparse
import contextlib
+import io
import json
import os
import sys
import time
import tomllib
+from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal
@@ -209,7 +211,7 @@ def _openai_build_request(
"model": model,
"messages": messages,
"tools": tools,
- "max_tokens": 16384,
+ "max_tokens": 4096,
"temperature": 0.0,
}
return "/v1/chat/completions", body
@@ -276,7 +278,7 @@ def _openai_build_followup(
"model": model,
"messages": followup_messages,
"tools": tools,
- "max_tokens": 16384,
+ "max_tokens": 4096,
"temperature": 0.0,
}
return "/v1/chat/completions", body
@@ -379,7 +381,7 @@ def _claude_build_request(
"model": model,
"messages": claude_messages,
"tools": claude_tools,
- "max_tokens": 16384,
+ "max_tokens": 4096,
"temperature": 0.0,
}
if system_content is not None:
@@ -489,7 +491,7 @@ def _claude_build_followup(
"model": model,
"messages": claude_messages,
"tools": claude_tools,
- "max_tokens": 16384,
+ "max_tokens": 4096,
"temperature": 0.0,
}
if system_content is not None:
@@ -913,6 +915,12 @@ Examples:
default=1,
help="Repeat each scenario N times (default: 1)",
)
+ parser.add_argument(
+ "--concurrency",
+ type=int,
+ default=1,
+ help="Run up to N scenarios in parallel against the same instance (default: 1)",
+ )
parser.add_argument(
"--scenarios",
nargs="*",
@@ -935,6 +943,13 @@ Examples:
)
args = parser.parse_args()
+ if args.concurrency < 1:
+ print(
+ f"--concurrency must be >= 1 (got {args.concurrency})",
+ file=sys.stderr,
+ )
+ sys.exit(2)
+
all_scenarios = load_scenarios(SCENARIOS_PATH)
if args.scenarios:
scenarios = [s for s in all_scenarios if s.name in args.scenarios]
@@ -1010,42 +1025,72 @@ Examples:
cluster_snapshot = capture_cluster_snapshot(exo)
all_results: list[ScenarioResult] = []
+ tasks: list[tuple[int, Scenario, ApiName]] = [
+ (run_idx, scenario, api_name)
+ for run_idx in range(args.repeat)
+ for scenario in scenarios
+ for api_name in api_names
+ ]
+
+ def _run_one(
+ http_client: httpx.Client,
+ task: tuple[int, Scenario, ApiName],
+ ) -> tuple[tuple[int, Scenario, ApiName], list[ScenarioResult], str]:
+ run_idx, scenario, api_name = task
+ buf = io.StringIO()
+ run_tag = f"[run {run_idx + 1}/{args.repeat}]" if args.repeat > 1 else ""
+ print(
+ f"\n {run_tag}[{api_name:>9}] {scenario.name}: {scenario.description}",
+ file=buf,
+ )
+ scenario_results = run_scenario(
+ http_client,
+ args.host,
+ args.port,
+ full_model_id,
+ scenario,
+ api_name,
+ args.timeout,
+ args.verbose,
+ )
+ for r in scenario_results:
+ status = "PASS" if r.passed else "FAIL"
+ print(
+ f" [{r.phase:>10}] {status} ({r.latency_ms:.0f}ms)",
+ file=buf,
+ )
+ for check_name, check_ok in r.checks.items():
+ mark = "+" if check_ok else "-"
+ print(f" {mark} {check_name}", file=buf)
+ if r.error:
+ print(f" ! {r.error}", file=buf)
+ return task, scenario_results, buf.getvalue()
+
try:
with httpx.Client() as http_client:
- for run_idx in range(args.repeat):
- if args.repeat > 1:
- print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
-
- for scenario in scenarios:
- for api_name in api_names:
- print(
- f"\n [{api_name:>9}] {scenario.name}: {scenario.description}",
- file=log,
- )
-
- scenario_results = run_scenario(
- http_client,
- args.host,
- args.port,
- full_model_id,
- scenario,
- api_name,
- args.timeout,
- args.verbose,
- )
+ if args.concurrency == 1:
+ current_run = -1
+ for task in tasks:
+ run_idx = task[0]
+ if args.repeat > 1 and run_idx != current_run:
+ print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
+ current_run = run_idx
+ _, scenario_results, buffered = _run_one(http_client, task)
+ all_results.extend(scenario_results)
+ log.write(buffered)
+ log.flush()
+ else:
+ print(
+ f"Running {len(tasks)} tasks with concurrency={args.concurrency}",
+ file=log,
+ )
+ with ThreadPoolExecutor(max_workers=args.concurrency) as pool:
+ futures = [pool.submit(_run_one, http_client, t) for t in tasks]
+ for fut in as_completed(futures):
+ _, scenario_results, buffered = fut.result()
all_results.extend(scenario_results)
-
- for r in scenario_results:
- status = "PASS" if r.passed else "FAIL"
- print(
- f" [{r.phase:>10}] {status} ({r.latency_ms:.0f}ms)",
- file=log,
- )
- for check_name, check_ok in r.checks.items():
- mark = "+" if check_ok else "-"
- print(f" {mark} {check_name}", file=log)
- if r.error:
- print(f" ! {r.error}", file=log)
+ log.write(buffered)
+ log.flush()
finally:
try:
exo.request_json("DELETE", f"/instance/{instance_id}")
diff --git a/bench/harness.py b/bench/harness.py
index 440a97cf..9c31fe21 100644
--- a/bench/harness.py
+++ b/bench/harness.py
@@ -564,7 +564,7 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
ap.add_argument(
"--settle-timeout",
type=float,
- default=0,
+ default=60.0,
help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
)
ap.add_argument(
diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index 398595ee..443979f1 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -12,7 +12,7 @@ from mlx.nn.layers.distributed import (
sum_gradients,
)
from mlx_lm.models.base import (
- scaled_dot_product_attention, # pyright: ignore[reportUnknownVariableType]
+ scaled_dot_product_attention,
)
from mlx_lm.models.cache import ArraysCache, KVCache
from mlx_lm.models.deepseek_v3 import DeepseekV3MLP
@@ -306,24 +306,20 @@ def pipeline_auto_parallel(
)
if isinstance(inner_model_instance, GptOssMoeModel):
- inner_model_instance.layer_types = inner_model_instance.layer_types[ # type: ignore
+ inner_model_instance.layer_types = inner_model_instance.layer_types[
start_layer:end_layer
]
# We can assume the model has at least one layer thanks to placement.
# If a layer type doesn't exist, we can set it to 0.
inner_model_instance.swa_idx = (
0
- if "sliding_attention" not in inner_model_instance.layer_types # type: ignore
- else inner_model_instance.layer_types.index( # type: ignore
- "sliding_attention"
- )
+ if "sliding_attention" not in inner_model_instance.layer_types
+ else inner_model_instance.layer_types.index("sliding_attention")
)
inner_model_instance.ga_idx = (
0
- if "full_attention" not in inner_model_instance.layer_types # type: ignore
- else inner_model_instance.layer_types.index( # type: ignore
- "full_attention"
- )
+ if "full_attention" not in inner_model_instance.layer_types
+ else inner_model_instance.layer_types.index("full_attention")
)
if isinstance(inner_model_instance, Step35InnerModel):
@@ -883,7 +879,7 @@ class WrappedMiniMaxAttention(CustomMlxLayer):
keys,
values,
cache=cache,
- scale=self._original_layer.scale, # type: ignore
+ scale=self._original_layer.scale,
mask=mask,
)
@@ -922,8 +918,8 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
self.all_to_sharded_linear_in_place(
layer.block_sparse_moe.switch_mlp.up_proj
)
- layer.block_sparse_moe = ShardedMoE(layer.block_sparse_moe) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
- layer.block_sparse_moe.sharding_group = self.group # pyright: ignore[reportAttributeAccessIssue]
+ layer.block_sparse_moe = ShardedMoE(layer.block_sparse_moe) # type: ignore
+ layer.block_sparse_moe.sharding_group = self.group
mx.eval(layer)
yield ModelLoadingResponse(layers_loaded=i, total=total)
@@ -1172,7 +1168,7 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
self.all_to_sharded_linear_in_place(layer.mlp.experts.up_proj)
layer.mlp = ShardedMoE(layer.mlp) # type: ignore
- layer.mlp.sharding_group = self.group # pyright: ignore[reportAttributeAccessIssue]
+ layer.mlp.sharding_group = self.group
mx.eval(layer)
yield ModelLoadingResponse(layers_loaded=i, total=total)
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py b/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py
index 680fe66b..73f8b1fa 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py
@@ -343,7 +343,7 @@ class TestKVPrefixCacheWithModel:
)
def test_mlx_generate_populates_cache(self, model_and_tokenizer):
- """mlx_generate should save the cache after generation completes."""
+ """mlx_generate should save the post-prefill cache (before the decode loop)."""
model, tokenizer = model_and_tokenizer
kv_prefix_cache = KVPrefixCache(None)
@@ -356,7 +356,6 @@ class TestKVPrefixCacheWithModel:
prompt_tokens = encode_prompt(tokenizer, prompt)
# Consume the entire generator so the cache-saving code after yield runs
- generated_tokens = 0
for _response in mlx_generate(
model=model,
tokenizer=tokenizer,
@@ -365,13 +364,14 @@ class TestKVPrefixCacheWithModel:
kv_prefix_cache=kv_prefix_cache,
group=None,
):
- generated_tokens += 1
+ pass
assert len(kv_prefix_cache.prompts) == 1
assert len(kv_prefix_cache.caches) == 1
- # Cache should contain prompt + generated tokens
- expected_length = len(prompt_tokens) + generated_tokens
- assert cache_length(kv_prefix_cache.caches[0]) == expected_length
+ # add_kv_cache is called before the decode loop and stores a deepcopy of
+ # the cache as it is just after prefill + trim(2). Generation tokens are
+ # never written into the stored entry.
+ assert cache_length(kv_prefix_cache.caches[0]) == len(prompt_tokens) - 2
def test_mlx_generate_second_call_gets_prefix_hit(self, model_and_tokenizer):
"""Second mlx_generate call with same prompt should get a prefix hit from stored cache."""
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py b/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py
index fec5082a..bfaba83b 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py
@@ -14,6 +14,8 @@ import pytest
from mlx.utils import tree_flatten, tree_unflatten
from mlx_lm.tokenizer_utils import TokenizerWrapper
+from exo.download.download_utils import resolve_existing_model
+from exo.shared.constants import EXO_MODELS_DIRS, EXO_MODELS_READ_ONLY_DIRS
from exo.shared.types.common import ModelId
from exo.shared.types.mlx import Model
from exo.shared.types.text_generation import (
@@ -28,8 +30,6 @@ from exo.worker.engines.mlx.utils_mlx import (
load_tokenizer_for_model_id,
)
-HF_CACHE = Path.home() / ".cache" / "huggingface" / "hub"
-
# ── Config reduction ──────────────────────────────────────────────────────── #
_REDUCE = {
@@ -100,12 +100,21 @@ def _reduce_config(cfg: dict[str, Any]) -> dict[str, Any]:
def _find_snapshot(hub_name: str) -> Path | None:
- model_dir = HF_CACHE / f"models--mlx-community--{hub_name}"
- snaps = model_dir / "snapshots"
- if not snaps.exists():
- return None
- children = sorted(snaps.iterdir())
- return children[0] if children else None
+ """Locate a model directory under exo's models dirs.
+
+ Uses resolve_existing_model for fully-downloaded models; falls back to any
+ existing directory (even partial) so that tokenizer-only copies still work.
+ """
+ model_id = ModelId(f"mlx-community/{hub_name}")
+ found = resolve_existing_model(model_id)
+ if found is not None:
+ return found
+ normalized = model_id.normalize()
+ for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
+ candidate = search_dir / normalized
+ if candidate.is_dir():
+ return candidate
+ return None
def _copy_tokenizer(src: Path, dst: Path) -> None:
@@ -192,13 +201,31 @@ ARCHITECTURES: list[ArchSpec] = [
]
+def _has_chat_template(model_dir: Path) -> bool:
+ """Check if a model dir has a usable chat template (inline or separate)."""
+ if (model_dir / "chat_template.jinja").exists():
+ return True
+ cfg = model_dir / "tokenizer_config.json"
+ if not cfg.exists():
+ return False
+ try:
+ data = cast(dict[str, Any], json.loads(cfg.read_text()))
+ except (OSError, json.JSONDecodeError):
+ return False
+ return bool(data.get("chat_template"))
+
+
def _arch_available(spec: ArchSpec) -> bool:
snap = _find_snapshot(spec.hub_name)
if snap is None or not (snap / "config.json").exists():
return False
+ tokenizer_snap = snap
if spec.tokenizer_hub is not None:
- return _find_snapshot(spec.tokenizer_hub) is not None
- return True
+ alt = _find_snapshot(spec.tokenizer_hub)
+ if alt is None:
+ return False
+ tokenizer_snap = alt
+ return _has_chat_template(tokenizer_snap)
def _make_task() -> TextGenerationTaskParams:
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_tp_bit_exact.py b/src/exo/worker/tests/unittests/test_mlx/test_tp_bit_exact.py
new file mode 100644
index 00000000..f0d12847
--- /dev/null
+++ b/src/exo/worker/tests/unittests/test_mlx/test_tp_bit_exact.py
@@ -0,0 +1,395 @@
+# type: ignore
+"""uv run pytest -v -m "" src/exo/worker/tests/unittests/test_mlx/test_tp_bit_exact.py"""
+
+import importlib
+import json
+import multiprocessing as mp
+import os
+import sys
+import tempfile
+import traceback
+
+import numpy as np
+import pytest
+
+MODEL_CONFIGS = {
+ "llama": dict(
+ module="mlx_lm.models.llama",
+ args=dict(
+ model_type="llama",
+ hidden_size=512,
+ intermediate_size=1024,
+ num_hidden_layers=2,
+ num_attention_heads=16,
+ num_key_value_heads=4,
+ rms_norm_eps=1e-6,
+ vocab_size=512,
+ max_position_embeddings=128,
+ head_dim=32,
+ rope_theta=10000.0,
+ ),
+ ),
+ "qwen3_5_moe": dict(
+ module="mlx_lm.models.qwen3_5_moe",
+ args=dict(
+ model_type="qwen3_5_moe",
+ text_config=dict(
+ model_type="qwen3_5_moe",
+ vocab_size=512,
+ hidden_size=512,
+ intermediate_size=1024,
+ num_hidden_layers=4,
+ num_attention_heads=16,
+ num_key_value_heads=4,
+ head_dim=32,
+ max_position_embeddings=128,
+ rms_norm_eps=1e-6,
+ tie_word_embeddings=False,
+ attention_bias=False,
+ full_attention_interval=2,
+ linear_num_value_heads=32,
+ linear_num_key_heads=16,
+ linear_key_head_dim=32,
+ linear_value_head_dim=32,
+ linear_conv_kernel_dim=4,
+ num_experts=16,
+ num_experts_per_tok=2,
+ decoder_sparse_step=1,
+ shared_expert_intermediate_size=256,
+ moe_intermediate_size=256,
+ norm_topk_prob=True,
+ rope_parameters={
+ "type": "default",
+ "rope_theta": 10000.0,
+ "partial_rotary_factor": 0.25,
+ "mrope_section": [11, 11, 10],
+ },
+ ),
+ ),
+ ),
+ "qwen3_next": dict(
+ module="mlx_lm.models.qwen3_next",
+ args=dict(
+ model_type="qwen3_next",
+ hidden_size=512,
+ intermediate_size=1024,
+ num_hidden_layers=4,
+ num_attention_heads=16,
+ num_key_value_heads=4,
+ head_dim=32,
+ max_position_embeddings=128,
+ rms_norm_eps=1e-6,
+ vocab_size=512,
+ attention_bias=False,
+ full_attention_interval=2,
+ linear_num_value_heads=32,
+ linear_num_key_heads=16,
+ linear_key_head_dim=32,
+ linear_value_head_dim=32,
+ linear_conv_kernel_dim=4,
+ num_experts=16,
+ num_experts_per_tok=2,
+ decoder_sparse_step=1,
+ shared_expert_intermediate_size=256,
+ moe_intermediate_size=256,
+ norm_topk_prob=True,
+ mlp_only_layers=[],
+ rope_theta=10000.0,
+ partial_rotary_factor=0.25,
+ ),
+ ),
+ "deepseek_v3": dict(
+ module="mlx_lm.models.deepseek_v3",
+ args=dict(
+ model_type="deepseek_v3",
+ hidden_size=512,
+ intermediate_size=1024,
+ num_hidden_layers=2,
+ num_attention_heads=16,
+ num_key_value_heads=16,
+ vocab_size=512,
+ max_position_embeddings=128,
+ rms_norm_eps=1e-6,
+ n_routed_experts=8,
+ n_shared_experts=1,
+ num_experts_per_tok=2,
+ moe_intermediate_size=256,
+ moe_layer_freq=1,
+ first_k_dense_replace=0,
+ n_group=1,
+ topk_group=1,
+ routed_scaling_factor=1.0,
+ q_lora_rank=None,
+ kv_lora_rank=16,
+ qk_nope_head_dim=16,
+ qk_rope_head_dim=16,
+ v_head_dim=32,
+ rope_theta=10000.0,
+ rope_scaling={},
+ attention_bias=False,
+ norm_topk_prob=True,
+ scoring_func="sigmoid",
+ topk_method="noaux_tc",
+ ),
+ ),
+ "deepseek_v3_q4": dict(
+ module="mlx_lm.models.deepseek_v3",
+ quantize=dict(group_size=32, bits=4, mode="affine"),
+ args=dict(
+ model_type="deepseek_v3",
+ hidden_size=512,
+ intermediate_size=1024,
+ num_hidden_layers=2,
+ num_attention_heads=16,
+ num_key_value_heads=16,
+ vocab_size=512,
+ max_position_embeddings=128,
+ rms_norm_eps=1e-6,
+ n_routed_experts=8,
+ n_shared_experts=1,
+ num_experts_per_tok=2,
+ moe_intermediate_size=256,
+ moe_layer_freq=1,
+ first_k_dense_replace=0,
+ n_group=1,
+ topk_group=1,
+ routed_scaling_factor=1.0,
+ q_lora_rank=None,
+ kv_lora_rank=64,
+ qk_nope_head_dim=32,
+ qk_rope_head_dim=32,
+ v_head_dim=32,
+ rope_theta=10000.0,
+ rope_scaling={},
+ attention_bias=False,
+ norm_topk_prob=True,
+ scoring_func="sigmoid",
+ topk_method="noaux_tc",
+ ),
+ ),
+ "glm4_moe_lite": dict(
+ module="mlx_lm.models.glm4_moe_lite",
+ args=dict(
+ model_type="glm4_moe_lite",
+ hidden_size=512,
+ intermediate_size=1024,
+ num_hidden_layers=2,
+ num_attention_heads=16,
+ num_key_value_heads=16,
+ vocab_size=512,
+ max_position_embeddings=128,
+ rms_norm_eps=1e-6,
+ n_routed_experts=8,
+ n_shared_experts=1,
+ num_experts_per_tok=2,
+ moe_intermediate_size=256,
+ first_k_dense_replace=1,
+ n_group=1,
+ topk_group=1,
+ routed_scaling_factor=1.0,
+ rope_theta=10000.0,
+ attention_bias=False,
+ q_lora_rank=None,
+ kv_lora_rank=16,
+ qk_rope_head_dim=16,
+ qk_nope_head_dim=16,
+ v_head_dim=32,
+ ),
+ ),
+ "minimax": dict(
+ module="mlx_lm.models.minimax",
+ args=dict(
+ model_type="minimax",
+ hidden_size=512,
+ intermediate_size=1024,
+ num_attention_heads=16,
+ num_key_value_heads=4,
+ max_position_embeddings=128,
+ num_experts_per_tok=2,
+ num_local_experts=8,
+ shared_intermediate_size=256,
+ num_hidden_layers=2,
+ rms_norm_eps=1e-6,
+ rope_theta=10000.0,
+ rotary_dim=32,
+ vocab_size=512,
+ ),
+ ),
+ "gpt_oss": dict(
+ module="mlx_lm.models.gpt_oss",
+ args=dict(
+ model_type="gpt_oss",
+ hidden_size=512,
+ intermediate_size=256,
+ num_hidden_layers=2,
+ num_attention_heads=16,
+ num_key_value_heads=4,
+ vocab_size=512,
+ head_dim=32,
+ rms_norm_eps=1e-6,
+ num_local_experts=8,
+ num_experts_per_tok=2,
+ layer_types=["sliding_attention", "full_attention"],
+ sliding_window=64,
+ rope_theta=10000.0,
+ ),
+ ),
+ "gemma4": dict(
+ module="mlx_lm.models.gemma4",
+ args=dict(
+ model_type="gemma4",
+ vocab_size=512,
+ text_config=dict(
+ vocab_size=512,
+ hidden_size=512,
+ intermediate_size=1024,
+ num_hidden_layers=4,
+ num_attention_heads=16,
+ num_key_value_heads=4,
+ head_dim=32,
+ global_head_dim=32,
+ num_kv_shared_layers=0,
+ vocab_size_per_layer_input=512,
+ hidden_size_per_layer_input=512,
+ rms_norm_eps=1e-6,
+ max_position_embeddings=128,
+ sliding_window=64,
+ sliding_window_pattern=2,
+ layer_types=[
+ "sliding_attention",
+ "full_attention",
+ "sliding_attention",
+ "full_attention",
+ ],
+ enable_moe_block=True,
+ num_experts=8,
+ top_k_experts=2,
+ moe_intermediate_size=256,
+ ),
+ ),
+ ),
+}
+
+_PROMPT = [[1, 23, 45, 67, 89, 12, 34, 56]]
+
+
+def _build(name):
+ import mlx.core as mx
+ import mlx.nn as nn
+ from mlx.utils import tree_map_with_path
+
+ import exo.worker.engines.mlx.auto_parallel # noqa: F401
+
+ cfg = MODEL_CONFIGS[name]
+ module = importlib.import_module(cfg["module"])
+ model_cls = module.Model
+ model_args_cls = module.ModelArgs
+
+ mx.random.seed(0)
+ args = model_args_cls(**cfg["args"])
+ m = model_cls(args)
+
+ def _to_bf16(_p, v):
+ if hasattr(v, "dtype") and v.dtype in (mx.float16, mx.float32, mx.bfloat16):
+ return v.astype(mx.bfloat16)
+ return v
+
+ m.update(tree_map_with_path(_to_bf16, m.parameters()))
+ if "quantize" in cfg:
+ nn.quantize(m, **cfg["quantize"])
+ mx.eval(m.parameters())
+ return mx, m
+
+
+def _run(name, out_path, shard):
+ import mlx.core as mx
+
+ if shard:
+ g = mx.distributed.init(backend="ring", strict=True)
+ mx_, m = _build(name)
+ if shard:
+ from exo.worker.engines.mlx.auto_parallel import tensor_auto_parallel
+
+ m = tensor_auto_parallel(m, g, on_layer_loaded=None)
+ mx_.eval(m.parameters())
+ inputs = mx_.array(_PROMPT, dtype=mx_.int32)
+ logits = m(inputs)
+ mx_.eval(logits)
+ np.savez(out_path, logits=np.asarray(logits.astype(mx_.float32)))
+
+
+def _ref_worker(name, out_path, q):
+ try:
+ _run(name, out_path, shard=False)
+ q.put(True)
+ except BaseException as e:
+ q.put(f"{e}\n{traceback.format_exc()}")
+
+
+def _tp_worker(name, rank, hf, out_path, q):
+ os.environ["MLX_HOSTFILE"] = hf
+ os.environ["MLX_RANK"] = str(rank)
+ try:
+ path = out_path if rank == 0 else out_path + f".r{rank}"
+ _run(name, path, shard=True)
+ q.put((rank, True, None))
+ except BaseException as e:
+ q.put((rank, False, f"{e}\n{traceback.format_exc()}"))
+
+
+def _run_compare(name, world_size, port_base):
+ d = tempfile.mkdtemp()
+ ref_path = f"{d}/ref.npz"
+ tp_path = f"{d}/tp.npz"
+ ctx = mp.get_context("spawn")
+ q = ctx.Queue()
+
+ p = ctx.Process(target=_ref_worker, args=(name, ref_path, q))
+ p.start()
+ p.join(300)
+ r = q.get(timeout=10)
+ if r is not True:
+ pytest.fail(f"[{name}] ref FAIL: {str(r)[:500]}")
+
+ hosts = [f"127.0.0.1:{port_base + i}" for i in range(world_size)]
+ with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+ json.dump(hosts, f)
+ hf = f.name
+ ps = [
+ ctx.Process(target=_tp_worker, args=(name, rank, hf, tp_path, q))
+ for rank in range(world_size)
+ ]
+ for pp in ps:
+ pp.start()
+ results = [q.get(timeout=300) for _ in range(world_size)]
+ for pp in ps:
+ pp.join(60)
+ for rank, ok, payload in results:
+ if not ok:
+ pytest.fail(f"[{name}] rank {rank} FAIL: {payload[:500]}")
+
+ ref = np.load(ref_path)["logits"]
+ tp = np.load(tp_path)["logits"]
+ diff = np.abs(ref - tp)
+ max_diff = float(diff.max())
+ mean_diff = float(diff.mean())
+ assert max_diff == 0.0, (
+ f"[{name} TP={world_size}] not bit-exact: max={max_diff} mean={mean_diff}"
+ )
+
+
+pytestmark = [
+ pytest.mark.slow,
+ pytest.mark.skipif(
+ sys.platform != "darwin", reason="MLX distributed requires Metal"
+ ),
+]
+
+
+@pytest.mark.skip("TP=2 is currently very different to TP=1. This test will not pass")
+@pytest.mark.parametrize("world_size", [2, 4])
+@pytest.mark.parametrize("name", list(MODEL_CONFIGS))
+def test_tp_bit_exact(name, world_size):
+ name_idx = list(MODEL_CONFIGS).index(name)
+ port = 32000 + name_idx * 20 + world_size
+ _run_compare(name, world_size, port)
← 7a312a17 Misc fixes: upstream JACCL all_sum, API, etc. + Add Kimi K2.
·
back to Exo
·
Fix event mutation causing indexed vs event mismatch (#1964) 73782ecc →