← back to Exo
Improve batch performance and stats reporting (#1777)
7df3774ca24014af27fddcb905529000b0277e3d · 2026-03-24 14:03:03 +0000 · rltakashige
## Motivation
Batch generation reports incorrect statistics, as mlx lm never clears
the original stats, meaning they get polluted over time.
The dashboard also seems considerably slower than bench statistics.
We also have a large discrepancy between B=1 batch generating and
mlx_generate.
Extracting logprobs is massively expensive, causing up to a 25% slowdown
compared to pure batching.
```
[ 12:02:01.1240AM | INFO ] step overhead: 3.49ms (next=12.49ms total=15.99ms)
[ 12:02:02.1600AM | INFO ] step overhead: 3.23ms (next=13.01ms total=16.24ms)
[ 12:02:03.2228AM | INFO ] step overhead: 3.28ms (next=13.38ms total=16.66ms)
[ 12:02:04.2798AM | INFO ] step overhead: 3.25ms (next=12.84ms total=16.10ms)
[ 12:02:05.3152AM | INFO ] step overhead: 3.18ms (next=12.61ms total=15.79ms)
[ 12:02:06.3522AM | INFO ] step overhead: 3.41ms (next=12.83ms total=16.25ms)
[ 12:02:07.3987AM | INFO ] step overhead: 3.38ms (next=13.14ms total=16.52ms)
[ 12:02:08.4537AM | INFO ] step overhead: 1.84ms (next=19.44ms total=21.28ms)
```
## Changes
1. Report stats ourselves instead of using mlx lm's stats for batch
generation (they use perf_counter anyway).
2. Adjust exo bench to match
3. Improve logprobs extraction speed by 10x, improving tps for dashboard
& any requests for logprobs
4. Use an SSE comment to align the speed to the real numbers at the end
of generation
5. Patch mlx for several optimizations given our assumptions and use
cases (e.g. use vllm style RoPE).
6. Switch MLX LM version to latest main, including support for Nemotron
Super and some Qwen3.5 fixes.
## Why It Works
1. Exo bench no longer reports polluted stats
2. Exo bench now handles the reported per-request stats rather than the
aggregate stats
3. The decode speed now jumps back to a real number at the end of the
generation
4. Large batch speedup for rotating KV cache models + 1:1 matching cache
with vllm
## Test Plan
### Manual Testing
Needs testing on OpenCode and CC
Needs eval testing
### Automated Testing
Only going to show the performance optimization difference after the
accurate reporting:
**GPT OSS 20B MXFP4 Q8 (large change)**
Before:
<img width="2466" height="1534" alt="image"
src="https://github.com/user-attachments/assets/88b50637-fca2-4db4-9413-b9eee6e2057e"
/>
<img width="2410" height="1240" alt="image"
src="https://github.com/user-attachments/assets/21e5c76a-2f5f-44d2-8953-121b3ebdbd68"
/>
After:
<img width="2476" height="1472" alt="image"
src="https://github.com/user-attachments/assets/fec5cfbd-fff8-430a-b12e-a329410107a2"
/>
<img width="2454" height="1236" alt="image"
src="https://github.com/user-attachments/assets/0400344b-a4a6-42c0-a9dd-4ee91ade714a"
/>
**Qwen 3.5 35B A3B 8bit (No change)**
Before:
<img width="2414" height="1396" alt="image"
src="https://github.com/user-attachments/assets/e75f0b38-df5d-49fd-ab90-bc1667d981b3"
/>
After:
<img width="2346" height="1234" alt="image"
src="https://github.com/user-attachments/assets/eabfb59c-851f-4d88-b927-e1e699a75cc6"
/>
**Llama 3.2 1B Instruct 4bit (small change)**
Before:
<img width="2516" height="1220" alt="image"
src="https://github.com/user-attachments/assets/c2873655-acff-4536-8263-fb8aea33db80"
/>
After:
<img width="2566" height="1370" alt="image"
src="https://github.com/user-attachments/assets/15f95c75-1c2f-4474-85a2-88c4d0a32543"
/>
Files touched
M .mlx_typings/mlx/core/__init__.pyiM .mlx_typings/mlx/nn/__init__.pyiM .mlx_typings/mlx/nn/layers/__init__.pyiM .mlx_typings/mlx/nn/layers/base.pyiM .mlx_typings/mlx_lm/generate.pyiM .mlx_typings/mlx_lm/models/cache.pyiA .mlx_typings/mlx_lm/models/gated_delta.pyiA .mlx_typings/mlx_lm/models/rope_utils.pyiM bench/exo_bench.pyM dashboard/src/lib/stores/app.svelte.tsM pyproject.tomlM src/exo/api/adapters/chat_completions.pyM src/exo/worker/engines/image/models/qwen/wrappers.pyM src/exo/worker/engines/mlx/auto_parallel.pyM src/exo/worker/engines/mlx/generator/batch_generate.pyM src/exo/worker/engines/mlx/generator/generate.pyA src/exo/worker/engines/mlx/patches/__init__.pyA src/exo/worker/engines/mlx/patches/high_precision_gdn_softplus.pyA src/exo/worker/engines/mlx/patches/opt_batch_gen.pyA src/exo/worker/engines/mlx/patches/standard_yarn_rope.pyA src/exo/worker/engines/mlx/tests/test_extract_top_logprobs.pyM src/exo/worker/engines/mlx/utils_mlx.pyM src/exo/worker/runner/bootstrap.pyM src/exo/worker/runner/llm_inference/batch_generator.pyM src/exo/worker/runner/llm_inference/runner.pyM src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.pyM uv.lock
Diff
commit 7df3774ca24014af27fddcb905529000b0277e3d
Author: rltakashige <rl.takashige@gmail.com>
Date: Tue Mar 24 14:03:03 2026 +0000
Improve batch performance and stats reporting (#1777)
## Motivation
Batch generation reports incorrect statistics, as mlx lm never clears
the original stats, meaning they get polluted over time.
The dashboard also seems considerably slower than bench statistics.
We also have a large discrepancy between B=1 batch generating and
mlx_generate.
Extracting logprobs is massively expensive, causing up to a 25% slowdown
compared to pure batching.
```
[ 12:02:01.1240AM | INFO ] step overhead: 3.49ms (next=12.49ms total=15.99ms)
[ 12:02:02.1600AM | INFO ] step overhead: 3.23ms (next=13.01ms total=16.24ms)
[ 12:02:03.2228AM | INFO ] step overhead: 3.28ms (next=13.38ms total=16.66ms)
[ 12:02:04.2798AM | INFO ] step overhead: 3.25ms (next=12.84ms total=16.10ms)
[ 12:02:05.3152AM | INFO ] step overhead: 3.18ms (next=12.61ms total=15.79ms)
[ 12:02:06.3522AM | INFO ] step overhead: 3.41ms (next=12.83ms total=16.25ms)
[ 12:02:07.3987AM | INFO ] step overhead: 3.38ms (next=13.14ms total=16.52ms)
[ 12:02:08.4537AM | INFO ] step overhead: 1.84ms (next=19.44ms total=21.28ms)
```
## Changes
1. Report stats ourselves instead of using mlx lm's stats for batch
generation (they use perf_counter anyway).
2. Adjust exo bench to match
3. Improve logprobs extraction speed by 10x, improving tps for dashboard
& any requests for logprobs
4. Use an SSE comment to align the speed to the real numbers at the end
of generation
5. Patch mlx for several optimizations given our assumptions and use
cases (e.g. use vllm style RoPE).
6. Switch MLX LM version to latest main, including support for Nemotron
Super and some Qwen3.5 fixes.
## Why It Works
1. Exo bench no longer reports polluted stats
2. Exo bench now handles the reported per-request stats rather than the
aggregate stats
3. The decode speed now jumps back to a real number at the end of the
generation
4. Large batch speedup for rotating KV cache models + 1:1 matching cache
with vllm
## Test Plan
### Manual Testing
Needs testing on OpenCode and CC
Needs eval testing
### Automated Testing
Only going to show the performance optimization difference after the
accurate reporting:
**GPT OSS 20B MXFP4 Q8 (large change)**
Before:
<img width="2466" height="1534" alt="image"
src="https://github.com/user-attachments/assets/88b50637-fca2-4db4-9413-b9eee6e2057e"
/>
<img width="2410" height="1240" alt="image"
src="https://github.com/user-attachments/assets/21e5c76a-2f5f-44d2-8953-121b3ebdbd68"
/>
After:
<img width="2476" height="1472" alt="image"
src="https://github.com/user-attachments/assets/fec5cfbd-fff8-430a-b12e-a329410107a2"
/>
<img width="2454" height="1236" alt="image"
src="https://github.com/user-attachments/assets/0400344b-a4a6-42c0-a9dd-4ee91ade714a"
/>
**Qwen 3.5 35B A3B 8bit (No change)**
Before:
<img width="2414" height="1396" alt="image"
src="https://github.com/user-attachments/assets/e75f0b38-df5d-49fd-ab90-bc1667d981b3"
/>
After:
<img width="2346" height="1234" alt="image"
src="https://github.com/user-attachments/assets/eabfb59c-851f-4d88-b927-e1e699a75cc6"
/>
**Llama 3.2 1B Instruct 4bit (small change)**
Before:
<img width="2516" height="1220" alt="image"
src="https://github.com/user-attachments/assets/c2873655-acff-4536-8263-fb8aea33db80"
/>
After:
<img width="2566" height="1370" alt="image"
src="https://github.com/user-attachments/assets/15f95c75-1c2f-4474-85a2-88c4d0a32543"
/>
---
.mlx_typings/mlx/core/__init__.pyi | 2 +-
.mlx_typings/mlx/nn/__init__.pyi | 8 +-
.mlx_typings/mlx/nn/layers/__init__.pyi | 36 ++-
.mlx_typings/mlx/nn/layers/base.pyi | 2 +-
.mlx_typings/mlx_lm/generate.pyi | 15 +-
.mlx_typings/mlx_lm/models/cache.pyi | 56 ++--
.mlx_typings/mlx_lm/models/gated_delta.pyi | 35 +++
.mlx_typings/mlx_lm/models/rope_utils.pyi | 51 ++++
bench/exo_bench.py | 12 +-
dashboard/src/lib/stores/app.svelte.ts | 33 ++-
pyproject.toml | 2 +-
src/exo/api/adapters/chat_completions.py | 4 +
.../worker/engines/image/models/qwen/wrappers.py | 4 +-
src/exo/worker/engines/mlx/auto_parallel.py | 2 +-
.../worker/engines/mlx/generator/batch_generate.py | 73 ++++--
src/exo/worker/engines/mlx/generator/generate.py | 74 +++---
src/exo/worker/engines/mlx/patches/__init__.py | 17 ++
.../mlx/patches/high_precision_gdn_softplus.py | 27 ++
.../worker/engines/mlx/patches/opt_batch_gen.py | 173 ++++++++++++
.../engines/mlx/patches/standard_yarn_rope.py | 118 +++++++++
.../engines/mlx/tests/test_extract_top_logprobs.py | 290 +++++++++++++++++++++
src/exo/worker/engines/mlx/utils_mlx.py | 1 +
src/exo/worker/runner/bootstrap.py | 3 +
.../worker/runner/llm_inference/batch_generator.py | 3 +-
src/exo/worker/runner/llm_inference/runner.py | 4 +-
.../tests/unittests/test_mlx/test_auto_parallel.py | 4 +-
uv.lock | 26 +-
27 files changed, 918 insertions(+), 157 deletions(-)
diff --git a/.mlx_typings/mlx/core/__init__.pyi b/.mlx_typings/mlx/core/__init__.pyi
index cab1d1ce..421e75db 100644
--- a/.mlx_typings/mlx/core/__init__.pyi
+++ b/.mlx_typings/mlx/core/__init__.pyi
@@ -2396,7 +2396,7 @@ def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array:
array: The angles in degrees.
"""
-def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array]):
+def depends[T](inputs: T, dependencies: array | Sequence[array]) -> T:
"""
Insert dependencies between arrays in the graph. The outputs are
identical to ``inputs`` but with dependencies on ``dependencies``.
diff --git a/.mlx_typings/mlx/nn/__init__.pyi b/.mlx_typings/mlx/nn/__init__.pyi
index 4c999379..d332c0d9 100644
--- a/.mlx_typings/mlx/nn/__init__.pyi
+++ b/.mlx_typings/mlx/nn/__init__.pyi
@@ -1,9 +1,5 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from layers import *
-from utils import *
+from .layers import *
+from .utils import *
from . import init as init
from . import losses as losses
diff --git a/.mlx_typings/mlx/nn/layers/__init__.pyi b/.mlx_typings/mlx/nn/layers/__init__.pyi
index f22856cd..0f09d7d9 100644
--- a/.mlx_typings/mlx/nn/layers/__init__.pyi
+++ b/.mlx_typings/mlx/nn/layers/__init__.pyi
@@ -1,20 +1,16 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from activations import *
-from base import *
-from containers import *
-from convolution import *
-from convolution_transpose import *
-from distributed import *
-from dropout import *
-from embedding import *
-from linear import *
-from normalization import *
-from pooling import *
-from positional_encoding import *
-from quantized import *
-from recurrent import *
-from transformer import *
-from upsample import *
+from .activations import *
+from .base import *
+from .containers import *
+from .convolution import *
+from .convolution_transpose import *
+from .distributed import *
+from .dropout import *
+from .embedding import *
+from .linear import *
+from .normalization import *
+from .pooling import *
+from .positional_encoding import *
+from .quantized import *
+from .recurrent import *
+from .transformer import *
+from .upsample import *
diff --git a/.mlx_typings/mlx/nn/layers/base.pyi b/.mlx_typings/mlx/nn/layers/base.pyi
index 70557bac..fcd8cf53 100644
--- a/.mlx_typings/mlx/nn/layers/base.pyi
+++ b/.mlx_typings/mlx/nn/layers/base.pyi
@@ -53,7 +53,7 @@ class Module(dict):
mx.eval(model.parameters())
"""
- __call__: Callable
+ def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
def __init__(self) -> None:
"""Should be called by the subclasses of ``Module``."""
diff --git a/.mlx_typings/mlx_lm/generate.pyi b/.mlx_typings/mlx_lm/generate.pyi
index 02904d81..4006259a 100644
--- a/.mlx_typings/mlx_lm/generate.pyi
+++ b/.mlx_typings/mlx_lm/generate.pyi
@@ -30,7 +30,7 @@ def str2bool(string): # -> bool:
def setup_arg_parser(): # -> ArgumentParser:
"""Set up and return the argument parser."""
-generation_stream = ...
+generation_stream: mx.Stream
@contextlib.contextmanager
def wired_limit(
@@ -266,12 +266,12 @@ def _merge_caches(caches: Any) -> List[Any]: ...
class Batch:
uids: List[int]
y: mx.array
- logprobs: mx.array
+ logprobs: List[mx.array] | mx.array
max_tokens: List[int]
num_tokens: List[int]
cache: List[Any]
- samplers: List[Any]
- logits_processors: List[Any]
+ samplers: List[Callable[[mx.array], mx.array] | None]
+ logits_processors: List[List[Callable[[mx.array, mx.array], mx.array]]]
tokens: List[mx.array]
def __len__(self) -> int: ...
def filter(self, keep_idx: List[int]) -> None: ...
@@ -279,13 +279,18 @@ class Batch:
def extract_cache(self, idx: int) -> List[Any]: ...
class BatchGenerator:
- model: Any
+ model: nn.Module
+ sampler: Callable[[mx.array], mx.array]
+ stop_tokens: set[int]
max_kv_size: Optional[int]
prefill_step_size: int
+ completion_batch_size: int
+ prefill_batch_size: int
unprocessed_prompts: List[Any]
active_batch: Optional[Batch]
prompt_progress_callback: Callable[[List[Tuple[int, int, int]]], None]
_stats: BatchStats
+ _next_count: int
@dataclass
class Response:
diff --git a/.mlx_typings/mlx_lm/models/cache.pyi b/.mlx_typings/mlx_lm/models/cache.pyi
index 1a934f40..8641815e 100644
--- a/.mlx_typings/mlx_lm/models/cache.pyi
+++ b/.mlx_typings/mlx_lm/models/cache.pyi
@@ -88,8 +88,8 @@ def create_attention_mask(
) -> array | Literal["causal"] | None: ...
class _BaseCache(Cache):
- keys: mx.array
- values: mx.array
+ keys: mx.array | None
+ values: mx.array | None
offset: int
@property
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
@@ -268,29 +268,14 @@ class CacheList(_BaseCache):
"""
class BatchKVCache(_BaseCache):
- step = ...
- def __init__(self, left_padding: List[int]) -> None:
- """
- The BatchKV cache expects inputs to be left-padded.
-
- E.g. the following prompts:
-
- [1, 3, 5]
- [7]
- [2, 6, 8, 9]
-
- Should be padded like so:
-
- [0, 1, 3, 5]
- [0, 0, 0, 7]
- [2, 6, 8, 9]
-
- And ``left_padding`` specifies the amount of padding for each.
- In this case, ``left_padding = [1, 3, 0]``.
- """
-
- def update_and_fetch(self, keys, values): # -> tuple[array | Any, array | Any]:
- ...
+ step: int
+ keys: array | None
+ values: array | None
+ offset: array
+ left_padding: array
+ _idx: int
+ def __init__(self, left_padding: List[int]) -> None: ...
+ def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
@property
def state(
self,
@@ -316,12 +301,21 @@ class BatchKVCache(_BaseCache):
"""
class BatchRotatingKVCache(_BaseCache):
- step = ...
- def __init__(self, max_size, left_padding: List[int]) -> None: ...
- def update_and_fetch(
- self, keys, values
- ): # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
- ...
+ step: int
+ keys: array | None
+ values: array | None
+ offset: array
+ left_padding: array
+ max_size: int
+ _idx: int
+ _offset: int
+ rotated: bool
+ _lengths: array | None
+ def __init__(self, max_size: int, left_padding: List[int]) -> None: ...
+ def _trim(self, trim_size: int, v: array, append: array | None = ...) -> array: ...
+ def _update_in_place(self, keys: array, values: array) -> tuple[array, array]: ...
+ def _update_concat(self, keys: array, values: array) -> tuple[array, array]: ...
+ def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
@property
def state(
self,
diff --git a/.mlx_typings/mlx_lm/models/gated_delta.pyi b/.mlx_typings/mlx_lm/models/gated_delta.pyi
new file mode 100644
index 00000000..a539b66a
--- /dev/null
+++ b/.mlx_typings/mlx_lm/models/gated_delta.pyi
@@ -0,0 +1,35 @@
+from typing import Optional
+
+import mlx.core as mx
+
+def compute_g(A_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array: ...
+def gated_delta_update(
+ q: mx.array,
+ k: mx.array,
+ v: mx.array,
+ a: mx.array,
+ b: mx.array,
+ A_log: mx.array,
+ dt_bias: mx.array,
+ state: Optional[mx.array] = ...,
+ mask: Optional[mx.array] = ...,
+ use_kernel: bool = ...,
+) -> tuple[mx.array, mx.array]: ...
+def gated_delta_ops(
+ q: mx.array,
+ k: mx.array,
+ v: mx.array,
+ g: mx.array,
+ beta: mx.array,
+ state: Optional[mx.array] = ...,
+ mask: Optional[mx.array] = ...,
+) -> tuple[mx.array, mx.array]: ...
+def gated_delta_kernel(
+ q: mx.array,
+ k: mx.array,
+ v: mx.array,
+ g: mx.array,
+ beta: mx.array,
+ state: mx.array,
+ mask: Optional[mx.array] = ...,
+) -> tuple[mx.array, mx.array]: ...
diff --git a/.mlx_typings/mlx_lm/models/rope_utils.pyi b/.mlx_typings/mlx_lm/models/rope_utils.pyi
new file mode 100644
index 00000000..498608fc
--- /dev/null
+++ b/.mlx_typings/mlx_lm/models/rope_utils.pyi
@@ -0,0 +1,51 @@
+from typing import Any, Optional
+
+import mlx.nn as nn
+
+class YarnRoPE(nn.Module):
+ def __init__(
+ self,
+ dims: int,
+ traditional: bool = ...,
+ max_position_embeddings: int = ...,
+ base: float = ...,
+ scaling_factor: float = ...,
+ original_max_position_embeddings: int = ...,
+ beta_fast: float = ...,
+ beta_slow: float = ...,
+ mscale: float = ...,
+ mscale_all_dim: float = ...,
+ ) -> None: ...
+
+class Llama3RoPE(nn.Module):
+ def __init__(
+ self,
+ dims: int,
+ traditional: bool = ...,
+ max_position_embeddings: int = ...,
+ base: float = ...,
+ scaling_factor: float = ...,
+ original_max_position_embeddings: int = ...,
+ low_freq_factor: float = ...,
+ high_freq_factor: float = ...,
+ ) -> None: ...
+
+class SuScaledRoPE(nn.Module):
+ def __init__(
+ self,
+ dims: int,
+ traditional: bool = ...,
+ max_position_embeddings: int = ...,
+ base: float = ...,
+ short_factor: Any = ...,
+ long_factor: Any = ...,
+ original_max_position_embeddings: int = ...,
+ ) -> None: ...
+
+def initialize_rope(
+ dims: int,
+ base: float = ...,
+ traditional: bool = ...,
+ scaling_config: Optional[dict[str, Any]] = ...,
+ max_position_embeddings: Optional[int] = ...,
+) -> nn.Module: ...
diff --git a/bench/exo_bench.py b/bench/exo_bench.py
index 20b0a546..2daa4927 100644
--- a/bench/exo_bench.py
+++ b/bench/exo_bench.py
@@ -501,23 +501,21 @@ def main() -> int:
for x, _ in batch_results
if x["stats"]["generation_tps"] > 0
]
- agg_gen_tps = (
+ per_req_tps = (
mean(valid_gen_tps) if valid_gen_tps else 0.0
)
- gen_tps = agg_gen_tps / concurrency
+ agg_gen_tps = per_req_tps * concurrency
logger.info(
f"[concurrent {concurrency}x] "
f"agg_gen_tps={agg_gen_tps:.2f} "
- f"gen_tps={gen_tps:.2f} "
+ f"per_req_tps={per_req_tps:.2f} "
f"errors={batch_errors}"
)
if runs:
prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
- gen_tps = mean(
- x["stats"]["generation_tps"] / x["concurrency"]
- for x in runs
- )
+ per_req_tps = mean(x["stats"]["generation_tps"] for x in runs)
+ gen_tps = per_req_tps * concurrency
ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
gtok = mean(x["stats"]["generation_tokens"] for x in runs)
peak = mean(
diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts
index 44fd2962..12ca231f 100644
--- a/dashboard/src/lib/stores/app.svelte.ts
+++ b/dashboard/src/lib/stores/app.svelte.ts
@@ -1793,6 +1793,14 @@ class AppStore {
this.persistConversation(targetConversationId);
}
},
+ {
+ generation_stats: (data) => {
+ const stats = data as { generation_tps: number };
+ if (stats.generation_tps > 0) {
+ this.tps = stats.generation_tps;
+ }
+ },
+ },
);
// Final update
@@ -1990,6 +1998,14 @@ class AppStore {
this.persistConversation(targetConversationId);
}
},
+ {
+ generation_stats: (data) => {
+ const stats = data as { generation_tps: number };
+ if (stats.generation_tps > 0) {
+ this.tps = stats.generation_tps;
+ }
+ },
+ },
);
// Final cleanup of the message (if conversation still exists)
@@ -2397,7 +2413,7 @@ class AppStore {
let streamedContent = "";
let streamedThinking = "";
-
+ let serverTpsReceived = false;
interface ChatCompletionChunk {
choices?: Array<{
delta?: { content?: string; reasoning_content?: string };
@@ -2462,7 +2478,6 @@ class AppStore {
tokenCount += 1;
this.totalTokens = tokenCount;
- // Update real-time TPS during streaming
if (firstTokenTime !== null && tokenCount > 1) {
const elapsed = performance.now() - firstTokenTime;
this.tps = (tokenCount / elapsed) * 1000;
@@ -2513,16 +2528,24 @@ class AppStore {
startedAt: this.prefillProgress?.startedAt ?? performance.now(),
};
},
+ generation_stats: (data) => {
+ const stats = data as { generation_tps: number };
+
+ if (stats.generation_tps > 0) {
+ this.tps = stats.generation_tps;
+ serverTpsReceived = true;
+ }
+ },
},
);
// Clear prefill progress after stream ends
this.prefillProgress = null;
- // Calculate final TPS
- if (firstTokenTime !== null && tokenCount > 1) {
+ // Use server-side TPS if available, otherwise fall back to client-side
+ if (!serverTpsReceived && firstTokenTime !== null && tokenCount > 1) {
const totalGenerationTime = performance.now() - firstTokenTime;
- this.tps = (tokenCount / totalGenerationTime) * 1000; // tokens per second
+ this.tps = (tokenCount / totalGenerationTime) * 1000;
}
// Final cleanup of the message (if conversation still exists)
diff --git a/pyproject.toml b/pyproject.toml
index b42cec58..bab185d8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -61,7 +61,7 @@ members = ["rust/exo_pyo3_bindings", "bench"]
[tool.uv.sources]
exo_pyo3_bindings = { workspace = true }
mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
-mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" }
+mlx-lm = { git = "https://github.com/ml-explore/mlx-lm", branch = "main" }
# Uncomment to use local mlx/mlx-lm development versions:
# mlx = { path = "/Users/Shared/mlx", editable=true }
# mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
diff --git a/src/exo/api/adapters/chat_completions.py b/src/exo/api/adapters/chat_completions.py
index 2272cedc..f0aad021 100644
--- a/src/exo/api/adapters/chat_completions.py
+++ b/src/exo/api/adapters/chat_completions.py
@@ -202,6 +202,8 @@ async def generate_chat_stream(
usage=last_usage,
)
yield f"data: {tool_response.model_dump_json()}\n\n"
+ if chunk.stats is not None:
+ yield f": generation_stats {chunk.stats.model_dump_json()}\n\n"
yield "data: [DONE]\n\n"
return
@@ -216,6 +218,8 @@ async def generate_chat_stream(
yield f"data: {chunk_response.model_dump_json()}\n\n"
if chunk.finish_reason is not None:
+ if chunk.stats is not None:
+ yield f": generation_stats {chunk.stats.model_dump_json()}\n\n"
yield "data: [DONE]\n\n"
diff --git a/src/exo/worker/engines/image/models/qwen/wrappers.py b/src/exo/worker/engines/image/models/qwen/wrappers.py
index 106f8fec..f0733070 100644
--- a/src/exo/worker/engines/image/models/qwen/wrappers.py
+++ b/src/exo/worker/engines/image/models/qwen/wrappers.py
@@ -56,10 +56,10 @@ class QwenJointBlockWrapper(JointBlockWrapper[QwenTransformerBlock]):
attn = self.block.attn
img_mod_params = self.block.img_mod_linear(
- self.block.img_mod_silu(text_embeddings) # pyright: ignore[reportUnknownArgumentType]
+ self.block.img_mod_silu(text_embeddings)
)
txt_mod_params = self.block.txt_mod_linear(
- self.block.txt_mod_silu(text_embeddings) # pyright: ignore[reportUnknownArgumentType]
+ self.block.txt_mod_silu(text_embeddings)
)
img_mod1, img_mod2 = mx.split(img_mod_params, 2, axis=-1)
diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index 031fa4a1..c1169a5e 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -480,7 +480,7 @@ def patch_tensor_model[T](model: T) -> T:
last = cache[-1] # pyright: ignore[reportAny]
dep_cache = last[0] if hasattr(last, "caches") else last # pyright: ignore[reportAny]
if hasattr(dep_cache, "keys"): # type: ignore
- dep_cache.keys = mx.depends(dep_cache.keys, logits) # pyright: ignore[reportAny,reportUnknownMemberType]
+ dep_cache.keys = mx.depends(dep_cache.keys, logits) # pyright: ignore[reportAny]
return logits
diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py
index 3c8ef562..d9410c18 100644
--- a/src/exo/worker/engines/mlx/generator/batch_generate.py
+++ b/src/exo/worker/engines/mlx/generator/batch_generate.py
@@ -6,6 +6,9 @@ import mlx.core as mx
from mlx_lm.generate import (
BatchGenerator as MlxBatchGenerator,
)
+from mlx_lm.generate import (
+ generation_stream,
+)
from mlx_lm.models.cache import RotatingKVCache
from mlx_lm.sample_utils import make_logits_processors, make_sampler
from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper
@@ -63,6 +66,7 @@ class _EngineTask:
potential_stop_sequence_text: str = ""
completion_tokens: int = 0
generation_start_time: float = 0.0
+ generation_time_at_start: float = 0.0
in_thinking: bool = False
reasoning_tokens: int = 0
prefill_tps: float = 0.0
@@ -75,22 +79,23 @@ class ExoBatchGenerator:
group: mx.distributed.Group | None
kv_prefix_cache: KVPrefixCache | None
- _exo_gen: MlxBatchGenerator = field(init=False)
+ _mlx_gen: MlxBatchGenerator = field(init=False)
_active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False)
def __post_init__(self) -> None:
- self._exo_gen = MlxBatchGenerator(
+ self._mlx_gen = MlxBatchGenerator(
model=self.model,
stop_tokens=set(eos_ids_from_tokenizer(self.tokenizer)),
prefill_step_size=4096,
)
+ self._mlx_gen._needs_topk = False # pyright: ignore[reportAttributeAccessIssue]
@property
def has_work(self) -> bool:
return (
bool(self._active_tasks)
- or bool(self._exo_gen.unprocessed_prompts)
- or self._exo_gen.active_batch is not None
+ or bool(self._mlx_gen.unprocessed_prompts)
+ or self._mlx_gen.active_batch is not None
)
def submit(
@@ -188,7 +193,7 @@ class ExoBatchGenerator:
max_tokens = task_params.max_output_tokens or MAX_TOKENS
- uids = self._exo_gen.insert(
+ uids = self._mlx_gen.insert(
prompts=[last_tokens.tolist()],
max_tokens=[max_tokens],
caches=[list(cache)],
@@ -211,6 +216,7 @@ class ExoBatchGenerator:
on_generation_token=on_generation_token,
generation_start_time=time.perf_counter(),
prefill_tps=_prefill_tps,
+ generation_time_at_start=self._mlx_gen._stats.generation_time,
)
return uid
@@ -219,7 +225,12 @@ class ExoBatchGenerator:
if not self.has_work:
return []
- responses = self._exo_gen.next()
+ self._mlx_gen._needs_topk = any( # pyright: ignore[reportAttributeAccessIssue]
+ t.task_params.logprobs for t in self._active_tasks.values()
+ )
+ _step_tic = time.perf_counter()
+ responses = self._mlx_gen.next()
+ _next_elapsed = time.perf_counter() - _step_tic
results: list[tuple[int, GenerationResponse]] = []
@@ -277,28 +288,31 @@ class ExoBatchGenerator:
logprob: float | None = None
top_logprobs: list[TopLogprobItem] | None = None
if task_params.logprobs:
- logprob, top_logprobs = extract_top_logprobs(
- logprobs=response.logprobs,
- tokenizer=self.tokenizer,
- top_logprobs=task_params.top_logprobs or DEFAULT_TOP_LOGPROBS,
- selected_token=response.token,
- )
+ with mx.stream(generation_stream):
+ logprob, top_logprobs = extract_top_logprobs(
+ logprobs=response.logprobs,
+ tokenizer=self.tokenizer,
+ top_logprobs=task_params.top_logprobs or DEFAULT_TOP_LOGPROBS,
+ selected_token=response.token,
+ precomputed_indices=getattr(response, "_topk_indices", None),
+ precomputed_values=getattr(response, "_topk_values", None),
+ precomputed_selected=getattr(
+ response, "_selected_logprob", None
+ ),
+ )
stats: GenerationStats | None = None
usage: Usage | None = None
if is_done:
- try:
- mlx_stats = self._exo_gen.stats()
- generation_tps = mlx_stats.generation_tps
- except ZeroDivisionError:
- generation_elapsed = (
- time.perf_counter() - state.generation_start_time
- )
- generation_tps = (
- state.completion_tokens / generation_elapsed
- if generation_elapsed > 0
- else 0.0
- )
+ gen_time_delta = (
+ self._mlx_gen._stats.generation_time
+ - state.generation_time_at_start
+ )
+ generation_tps = (
+ state.completion_tokens / gen_time_delta
+ if gen_time_delta > 0
+ else 0.0
+ )
stats = GenerationStats(
prompt_tps=state.prefill_tps,
@@ -345,15 +359,22 @@ class ExoBatchGenerator:
-max_stop_len:
]
+ _step_elapsed = time.perf_counter() - _step_tic
+ _overhead = _step_elapsed - _next_elapsed
+ if self._mlx_gen._next_count % 64 == 0 and responses:
+ logger.debug(
+ f"step overhead: {_overhead * 1000:.2f}ms (next={_next_elapsed * 1000:.2f}ms total={_step_elapsed * 1000:.2f}ms)"
+ )
+
return results
def cancel(self, uids: list[int]) -> None:
- self._exo_gen.remove(uids)
+ self._mlx_gen.remove(uids)
for uid in uids:
self._active_tasks.pop(uid, None)
def close(self) -> None:
- self._exo_gen.close()
+ self._mlx_gen.close()
def _save_prefix_cache(
self,
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index f5470014..16de44d8 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -179,7 +179,8 @@ def pipeline_parallel_prefill(
flush_prefill_sends()
assert _prompt_cache is not None
- mx.eval([c.state for c in _prompt_cache]) # type: ignore
+ with mx.stream(generation_stream):
+ mx.eval([c.state for c in _prompt_cache]) # type: ignore
# Final callback matching generate_step
prompt_progress_callback(total, total)
@@ -398,52 +399,44 @@ def extract_top_logprobs(
tokenizer: TokenizerWrapper,
top_logprobs: int,
selected_token: int,
+ precomputed_indices: list[int] | None = None,
+ precomputed_values: list[float] | None = None,
+ precomputed_selected: float | None = None,
) -> tuple[float, list[TopLogprobItem]]:
- """Extract the selected token's logprob and top alternative tokens.
-
- Args:
- logprobs: Full vocabulary logprobs array from MLX
- tokenizer: Tokenizer for decoding token IDs to strings
- top_logprobs: Number of top alternatives to return
- selected_token: The token ID that was actually sampled
-
- Returns:
- Tuple of (selected_token_logprob, list of TopLogprobItem for top alternatives)
- """
- # Get the logprob of the selected token
- selected_logprob = float(logprobs[selected_token].item())
-
- # Get top indices (most probable tokens)
- # mx.argpartition gives indices that would partition the array
- # We negate logprobs since argpartition finds smallest, and we want largest
- top_logprobs = min(top_logprobs, logprobs.shape[0]) # Don't exceed vocab size
- top_indices = mx.argpartition(-logprobs, top_logprobs)[:top_logprobs]
-
- # Get the actual logprob values for these indices
- top_values = logprobs[top_indices]
-
- # Sort by logprob (descending) for consistent ordering
- sort_order = mx.argsort(-top_values)
- top_indices = top_indices[sort_order]
- top_values = top_values[sort_order]
+ if (
+ precomputed_indices is not None
+ and precomputed_values is not None
+ and precomputed_selected is not None
+ ):
+ top_indices_list: list[int] = precomputed_indices[:top_logprobs]
+ top_values_list: list[float] = precomputed_values[:top_logprobs]
+ selected_logprob = precomputed_selected
+ else:
+ selected_logprob_arr = logprobs[selected_token]
+ top_logprobs = min(top_logprobs, logprobs.shape[0] - 1)
+ top_indices = mx.argpartition(-logprobs, top_logprobs)[:top_logprobs]
+ top_values = logprobs[top_indices]
+ sort_order = mx.argsort(-top_values)
+ top_indices = top_indices[sort_order]
+ top_values = top_values[sort_order]
+ mx.eval(selected_logprob_arr, top_indices, top_values)
+ selected_logprob = float(selected_logprob_arr.item())
+ top_indices_list = top_indices.tolist() # type: ignore
+ top_values_list = top_values.tolist() # type: ignore
# Convert to list of TopLogprobItem
top_logprob_items: list[TopLogprobItem] = []
- for i in range(top_logprobs):
- token_id = int(top_indices[i].item())
- token_logprob = float(top_values[i].item())
+ for token_id, token_logprob in zip(top_indices_list, top_values_list, strict=True):
if math.isnan(token_logprob):
continue
# Decode token ID to string
token_str = tokenizer.decode([token_id])
- # Get byte representation
- token_bytes = list(token_str.encode("utf-8"))
top_logprob_items.append(
TopLogprobItem(
token=token_str,
logprob=token_logprob,
- bytes=token_bytes,
+ bytes=list(token_str.encode("utf-8")),
)
)
@@ -624,12 +617,13 @@ def mlx_generate(
logprob: float | None = None
top_logprobs: list[TopLogprobItem] | None = None
if task.logprobs:
- logprob, top_logprobs = extract_top_logprobs(
- logprobs=out.logprobs,
- tokenizer=tokenizer,
- top_logprobs=task.top_logprobs or DEFAULT_TOP_LOGPROBS,
- selected_token=out.token,
- )
+ with mx.stream(generation_stream):
+ logprob, top_logprobs = extract_top_logprobs(
+ logprobs=out.logprobs,
+ tokenizer=tokenizer,
+ top_logprobs=task.top_logprobs or DEFAULT_TOP_LOGPROBS,
+ selected_token=out.token,
+ )
if is_done:
# Log generation stats
diff --git a/src/exo/worker/engines/mlx/patches/__init__.py b/src/exo/worker/engines/mlx/patches/__init__.py
new file mode 100644
index 00000000..0a14b262
--- /dev/null
+++ b/src/exo/worker/engines/mlx/patches/__init__.py
@@ -0,0 +1,17 @@
+from exo.worker.engines.mlx.patches.high_precision_gdn_softplus import (
+ patch_gdn_softplus,
+)
+from exo.worker.engines.mlx.patches.opt_batch_gen import apply_batch_gen_patch
+from exo.worker.engines.mlx.patches.standard_yarn_rope import patch_yarn_rope
+
+_applied = False
+
+
+def apply_mlx_patches() -> None:
+ global _applied
+ if _applied:
+ return
+ _applied = True
+ patch_yarn_rope()
+ patch_gdn_softplus()
+ apply_batch_gen_patch()
diff --git a/src/exo/worker/engines/mlx/patches/high_precision_gdn_softplus.py b/src/exo/worker/engines/mlx/patches/high_precision_gdn_softplus.py
new file mode 100644
index 00000000..a1961e11
--- /dev/null
+++ b/src/exo/worker/engines/mlx/patches/high_precision_gdn_softplus.py
@@ -0,0 +1,27 @@
+import sys
+
+import mlx.core as mx
+from mlx_lm.models.gated_delta import compute_g
+
+
+def _compute_g_f32(a_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array:
+ return mx.exp(
+ -mx.exp(a_log.astype(mx.float32))
+ * mx.where(
+ (a + dt_bias).astype(mx.float32) > 20,
+ (a + dt_bias).astype(mx.float32),
+ mx.log1p(mx.exp((a + dt_bias).astype(mx.float32))),
+ )
+ ).astype(a.dtype)
+
+
+def patch_gdn_softplus() -> None:
+ from mlx_lm.models import gated_delta
+
+ gated_delta.compute_g = _compute_g_f32
+
+ for mod in list(sys.modules.values()):
+ if mod is gated_delta:
+ continue
+ if getattr(mod, "compute_g", None) is compute_g:
+ object.__setattr__(mod, "compute_g", _compute_g_f32)
diff --git a/src/exo/worker/engines/mlx/patches/opt_batch_gen.py b/src/exo/worker/engines/mlx/patches/opt_batch_gen.py
new file mode 100644
index 00000000..015d2e4b
--- /dev/null
+++ b/src/exo/worker/engines/mlx/patches/opt_batch_gen.py
@@ -0,0 +1,173 @@
+import time
+from typing import Any, cast
+
+import mlx.core as mx
+from mlx_lm.generate import BatchGenerator, generation_stream
+
+_PRECOMPUTE_TOP_K = 20
+
+_original_public_next = BatchGenerator.next
+
+_pending_topk_idx: mx.array | None = None
+_pending_topk_val: mx.array | None = None
+_pending_selected_lps: mx.array | None = None
+
+
+def _fast_next(self: BatchGenerator) -> list[BatchGenerator.Response]:
+ tic = time.perf_counter()
+ batch = self.active_batch
+ assert batch is not None
+ batch_size = len(batch)
+
+ prev_tokens = batch.y
+ prev_logprobs = batch.logprobs
+
+ has_processors = any(p for ps in batch.logits_processors for p in ps)
+ if has_processors:
+ for i, toks in enumerate(batch.tokens):
+ batch.tokens[i] = mx.concatenate([toks, prev_tokens[i : i + 1]])
+
+ logits = self.model(prev_tokens[:, None], cache=batch.cache)
+ logits = logits[:, -1, :]
+
+ if has_processors:
+ processed_logits: list[mx.array] = []
+ for e in range(batch_size):
+ sample_logits: mx.array = logits[e : e + 1]
+ for processor in batch.logits_processors[e]:
+ sample_logits = processor(batch.tokens[e], sample_logits)
+ processed_logits.append(sample_logits)
+ logits = mx.concatenate(processed_logits, axis=0)
+
+ logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+
+ if (
+ batch_size == 1
+ or any(batch.samplers)
+ and all(s is batch.samplers[0] for s in batch.samplers)
+ ):
+ sampler = batch.samplers[0] or self.sampler
+ batch.y = sampler(logprobs)
+ elif any(batch.samplers):
+ all_samples: list[mx.array] = []
+ for e in range(batch_size):
+ s = batch.samplers[e] or self.sampler
+ all_samples.append(s(logprobs[e : e + 1]))
+ batch.y = mx.concatenate(all_samples, axis=0)
+ else:
+ batch.y = self.sampler(logprobs)
+ batch.logprobs = list(logprobs)
+
+ global _pending_topk_idx, _pending_topk_val, _pending_selected_lps
+
+ emit_topk_indices: list[list[int]] = (
+ cast(list[list[int]], _pending_topk_idx.tolist())
+ if _pending_topk_idx is not None
+ else []
+ )
+ emit_topk_values: list[list[float]] = (
+ cast(list[list[float]], _pending_topk_val.tolist())
+ if _pending_topk_val is not None
+ else []
+ )
+ emit_selected_lps: list[float] = (
+ cast(list[float], _pending_selected_lps.tolist())
+ if _pending_selected_lps is not None
+ else []
+ )
+
+ needs_topk: bool = getattr(self, "_needs_topk", False)
+ if needs_topk:
+ k = min(_PRECOMPUTE_TOP_K, logprobs.shape[1])
+ _pending_topk_idx = mx.argpartition(-logprobs, k, axis=1)[:, :k]
+ _pending_topk_val = mx.take_along_axis(logprobs, _pending_topk_idx, axis=1)
+ sort_order = mx.argsort(-_pending_topk_val, axis=1)
+ _pending_topk_idx = mx.take_along_axis(_pending_topk_idx, sort_order, axis=1)
+ _pending_topk_val = mx.take_along_axis(_pending_topk_val, sort_order, axis=1)
+ _pending_selected_lps = logprobs[mx.arange(batch_size), batch.y]
+ mx.async_eval(
+ batch.y,
+ *batch.logprobs,
+ *batch.tokens,
+ _pending_topk_idx,
+ _pending_topk_val,
+ _pending_selected_lps,
+ )
+ else:
+ _pending_topk_idx = None
+ _pending_topk_val = None
+ _pending_selected_lps = None
+ mx.async_eval(batch.y, *batch.logprobs, *batch.tokens)
+
+ prev_token_list: list[int] = cast(list[int], prev_tokens.tolist())
+
+ toc = time.perf_counter()
+ self._stats.generation_time += toc - tic
+
+ keep_idx: list[int] = []
+ end_idx: list[int] = []
+ responses: list[Any] = []
+ stop_tokens = self.stop_tokens
+
+ for e in range(batch_size):
+ t = prev_token_list[e]
+ uid = batch.uids[e]
+ num_tok = batch.num_tokens[e] + 1
+ batch.num_tokens[e] = num_tok
+
+ if t in stop_tokens:
+ finish_reason = "stop"
+ end_idx.append(e)
+ elif num_tok >= batch.max_tokens[e]:
+ finish_reason = "length"
+ end_idx.append(e)
+ else:
+ finish_reason = None
+ keep_idx.append(e)
+
+ cache = None
+ if finish_reason is not None:
+ cache = batch.extract_cache(e)
+ response = self.Response(uid, t, prev_logprobs[e], finish_reason, cache)
+ if emit_topk_indices and e < len(emit_topk_indices):
+ response._topk_indices = emit_topk_indices[e] # pyright: ignore[reportAttributeAccessIssue]
+ response._topk_values = emit_topk_values[e] # pyright: ignore[reportAttributeAccessIssue]
+ response._selected_logprob = emit_selected_lps[e] # pyright: ignore[reportAttributeAccessIssue]
+ responses.append(response)
+
+ if end_idx:
+ if keep_idx:
+ batch.filter(keep_idx)
+ if (
+ _pending_topk_idx is not None
+ and _pending_topk_val is not None
+ and _pending_selected_lps is not None
+ ):
+ ki = mx.array(keep_idx)
+ _pending_topk_idx = _pending_topk_idx[ki]
+ _pending_topk_val = _pending_topk_val[ki]
+ _pending_selected_lps = _pending_selected_lps[ki]
+ else:
+ self.active_batch = None
+ _pending_topk_idx = None
+ _pending_topk_val = None
+ _pending_selected_lps = None
+
+ self._next_count += 1
+ if self._next_count % 512 == 0:
+ mx.clear_cache()
+ self._stats.generation_tokens += len(responses)
+ return responses
+
+
+def _patched_public_next(self: BatchGenerator) -> list[BatchGenerator.Response]:
+ batch = self.active_batch
+ # Only do decode with fast_next
+ if batch is not None and not self.unprocessed_prompts:
+ with mx.stream(generation_stream):
+ return _fast_next(self)
+ return _original_public_next(self)
+
+
+def apply_batch_gen_patch() -> None:
+ BatchGenerator.next = _patched_public_next
diff --git a/src/exo/worker/engines/mlx/patches/standard_yarn_rope.py b/src/exo/worker/engines/mlx/patches/standard_yarn_rope.py
new file mode 100644
index 00000000..2dbee8ee
--- /dev/null
+++ b/src/exo/worker/engines/mlx/patches/standard_yarn_rope.py
@@ -0,0 +1,118 @@
+import math
+
+import mlx.core as mx
+from mlx_lm.models import rope_utils
+
+_original_YarnRoPE_init = rope_utils.YarnRoPE.__init__ # noqa: N816
+_original_initialize_rope = rope_utils.initialize_rope
+
+
+def _patched_yarn_init(
+ self: rope_utils.YarnRoPE,
+ dims: int,
+ traditional: bool = False,
+ max_position_embeddings: int = 2048,
+ base: float = 10000,
+ scaling_factor: float = 1.0,
+ original_max_position_embeddings: int = 4096,
+ beta_fast: float = 32,
+ beta_slow: float = 1,
+ mscale: float = 1,
+ mscale_all_dim: float = 0,
+ truncate: bool = True,
+) -> None:
+ """Patch mlx_lm's YarnRoPE to match vLLM's inverse-frequency blending formula for compatability."""
+
+ super(rope_utils.YarnRoPE, self).__init__()
+
+ def yarn_find_correction_dim(num_rotations: float) -> float:
+ return (
+ dims
+ * math.log(original_max_position_embeddings / (num_rotations * 2 * math.pi))
+ ) / (2 * math.log(base))
+
+ def yarn_find_correction_range() -> tuple[float, float]:
+ low: float = yarn_find_correction_dim(beta_fast)
+ high: float = yarn_find_correction_dim(beta_slow)
+ if truncate:
+ low = math.floor(low)
+ high = math.ceil(high)
+ return max(low, 0), min(high, dims - 1)
+
+ def yarn_get_mscale(scale: float = 1, ms: float = 1) -> float:
+ if scale <= 1:
+ return 1.0
+ return 0.1 * ms * math.log(scale) + 1.0
+
+ def yarn_linear_ramp_mask(min_val: float, max_val: float, dim: int) -> mx.array:
+ if min_val == max_val:
+ max_val += 0.001
+ linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (max_val - min_val)
+ return mx.clip(linear_func, 0, 1)
+
+ self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
+ scaling_factor, mscale_all_dim
+ )
+ pos_freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
+ inv_freq_extrapolation = 1.0 / pos_freqs
+ inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
+ low, high = yarn_find_correction_range()
+ inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dims // 2)
+ inv_freq = (
+ inv_freq_interpolation * (1 - inv_freq_mask)
+ + inv_freq_extrapolation * inv_freq_mask
+ )
+ self._freqs = 1.0 / inv_freq
+ self.dims = dims
+ self.traditional = traditional
+
+
+def _patched_initialize_rope(
+ dims: int,
+ base: float,
+ traditional: bool,
+ scaling_config: dict[str, str | int | float | bool] | None = None,
+ max_position_embeddings: int | None = None,
+) -> object:
+ rope_type = "default"
+ if scaling_config is not None:
+ rope_type = str(
+ scaling_config.get("type") or scaling_config.get("rope_type", "default")
+ )
+
+ # All the yarn rope types supported in mlx lm
+ if rope_type in ("yarn", "deepseek_yarn"):
+ assert scaling_config is not None
+ cfg = scaling_config
+
+ def _float(key: str, default: float) -> float:
+ v = cfg.get(key)
+ return float(v) if v is not None else default
+
+ def _int(key: str, default: int) -> int:
+ v = cfg.get(key)
+ return int(v) if v is not None else default
+
+ return rope_utils.YarnRoPE(
+ dims=dims,
+ max_position_embeddings=max_position_embeddings or 2048,
+ traditional=traditional,
+ scaling_factor=_float("factor", 1.0),
+ base=base,
+ original_max_position_embeddings=_int(
+ "original_max_position_embeddings", 4096
+ ),
+ beta_fast=_float("beta_fast", 32),
+ beta_slow=_float("beta_slow", 1),
+ mscale=_float("mscale", 1),
+ mscale_all_dim=_float("mscale_all_dim", 0),
+ )
+
+ return _original_initialize_rope(
+ dims, base, traditional, scaling_config, max_position_embeddings
+ )
+
+
+def patch_yarn_rope() -> None:
+ rope_utils.YarnRoPE.__init__ = _patched_yarn_init
+ rope_utils.initialize_rope = _patched_initialize_rope
diff --git a/src/exo/worker/engines/mlx/tests/test_extract_top_logprobs.py b/src/exo/worker/engines/mlx/tests/test_extract_top_logprobs.py
new file mode 100644
index 00000000..36e643c2
--- /dev/null
+++ b/src/exo/worker/engines/mlx/tests/test_extract_top_logprobs.py
@@ -0,0 +1,290 @@
+# type: ignore
+import math
+from unittest.mock import MagicMock
+
+import mlx.core as mx
+import mlx.nn as nn
+import pytest
+from mlx_lm.generate import BatchGenerator
+
+from exo.worker.engines.mlx.generator.generate import extract_top_logprobs
+from exo.worker.engines.mlx.patches.opt_batch_gen import (
+ _PRECOMPUTE_TOP_K,
+ apply_batch_gen_patch,
+)
+
+
+def _mock_tokenizer() -> MagicMock:
+ tok = MagicMock()
+ tok.decode = lambda ids: f"tok_{ids[0]}"
+ return tok
+
+
+def _make_logprobs(values: list[float]) -> mx.array:
+ arr = mx.array(values, dtype=mx.float32)
+ mx.eval(arr)
+ return arr
+
+
+class TestExtractTopLogprobsFallback:
+ def test_returns_correct_selected_logprob(self) -> None:
+ lp = _make_logprobs([-1.0, -2.0, -0.5, -3.0, -4.0])
+ selected, _ = extract_top_logprobs(
+ lp, _mock_tokenizer(), top_logprobs=3, selected_token=2
+ )
+ assert selected == pytest.approx(-0.5)
+
+ def test_returns_top_k_sorted_descending(self) -> None:
+ lp = _make_logprobs([-1.0, -2.0, -0.5, -3.0, -4.0])
+ _, items = extract_top_logprobs(
+ lp, _mock_tokenizer(), top_logprobs=3, selected_token=0
+ )
+ logprob_values = [item.logprob for item in items]
+ assert logprob_values == sorted(logprob_values, reverse=True)
+ assert len(items) == 3
+
+ def test_top_tokens_are_most_probable(self) -> None:
+ lp = _make_logprobs([-5.0, -1.0, -3.0, -0.1, -2.0])
+ _, items = extract_top_logprobs(
+ lp, _mock_tokenizer(), top_logprobs=2, selected_token=0
+ )
+ token_ids = [int(item.token.split("_")[1]) for item in items]
+ assert 3 in token_ids
+ assert 1 in token_ids
+
+ def test_top_logprobs_clamped_to_vocab_size(self) -> None:
+ lp = _make_logprobs([-1.0, -2.0, -3.0, -4.0, -5.0])
+ _, items = extract_top_logprobs(
+ lp, _mock_tokenizer(), top_logprobs=10, selected_token=0
+ )
+ assert len(items) == 4
+
+ def test_nan_logprobs_filtered(self) -> None:
+ lp = _make_logprobs([-1.0, float("nan"), -0.5])
+ _, items = extract_top_logprobs(
+ lp, _mock_tokenizer(), top_logprobs=3, selected_token=0
+ )
+ for item in items:
+ assert not math.isnan(item.logprob)
+
+ def test_token_bytes_correct(self) -> None:
+ tok = MagicMock()
+ tok.decode = lambda ids: "hello"
+ lp = _make_logprobs([-1.0, -2.0])
+ _, items = extract_top_logprobs(lp, tok, top_logprobs=2, selected_token=0)
+ assert items[0].bytes == list("hello".encode("utf-8"))
+
+
+class TestExtractTopLogprobsPrecomputed:
+ def test_uses_precomputed_data(self) -> None:
+ lp = _make_logprobs([-99.0])
+ selected, items = extract_top_logprobs(
+ lp,
+ _mock_tokenizer(),
+ top_logprobs=2,
+ selected_token=0,
+ precomputed_indices=[3, 1, 0],
+ precomputed_values=[-0.1, -1.0, -5.0],
+ precomputed_selected=-0.1,
+ )
+ assert selected == pytest.approx(-0.1)
+ assert len(items) == 2
+ assert items[0].token == "tok_3"
+ assert items[0].logprob == pytest.approx(-0.1)
+ assert items[1].token == "tok_1"
+ assert items[1].logprob == pytest.approx(-1.0)
+
+ def test_slices_precomputed_to_requested_k(self) -> None:
+ lp = _make_logprobs([-99.0])
+ _, items = extract_top_logprobs(
+ lp,
+ _mock_tokenizer(),
+ top_logprobs=1,
+ selected_token=0,
+ precomputed_indices=[3, 1, 0, 2, 4],
+ precomputed_values=[-0.1, -1.0, -2.0, -3.0, -4.0],
+ precomputed_selected=-0.1,
+ )
+ assert len(items) == 1
+ assert items[0].token == "tok_3"
+
+ def test_falls_back_when_precomputed_partial(self) -> None:
+ lp = _make_logprobs([-1.0, -2.0, -0.5])
+ selected, items = extract_top_logprobs(
+ lp,
+ _mock_tokenizer(),
+ top_logprobs=2,
+ selected_token=2,
+ precomputed_indices=[0, 2],
+ precomputed_values=None,
+ precomputed_selected=None,
+ )
+ assert selected == pytest.approx(-0.5)
+ assert len(items) == 2
+
+ def test_precomputed_matches_fallback(self) -> None:
+ lp = _make_logprobs([-1.0, -0.3, -2.5, -0.1, -4.0, -0.8, -3.0, -1.5])
+ tok = _mock_tokenizer()
+
+ selected_fb, items_fb = extract_top_logprobs(
+ lp, tok, top_logprobs=5, selected_token=1
+ )
+
+ pre_indices = [item.token.split("_")[1] for item in items_fb]
+ pre_indices_int = [int(x) for x in pre_indices]
+ pre_values = [item.logprob for item in items_fb]
+
+ selected_pc, items_pc = extract_top_logprobs(
+ lp,
+ tok,
+ top_logprobs=5,
+ selected_token=1,
+ precomputed_indices=pre_indices_int,
+ precomputed_values=pre_values,
+ precomputed_selected=selected_fb,
+ )
+
+ assert selected_pc == pytest.approx(selected_fb)
+ assert len(items_pc) == len(items_fb)
+ for a, b in zip(items_pc, items_fb, strict=True):
+ assert a.token == b.token
+ assert a.logprob == pytest.approx(b.logprob)
+
+
+def _tiny_model() -> nn.Module:
+ from mlx_lm.models.llama import Model, ModelArgs
+
+ mx.random.seed(42)
+ args = ModelArgs(
+ model_type="llama",
+ hidden_size=64,
+ num_hidden_layers=2,
+ intermediate_size=128,
+ num_attention_heads=2,
+ num_key_value_heads=1,
+ rms_norm_eps=1e-6,
+ vocab_size=256,
+ rope_theta=10000.0,
+ tie_word_embeddings=True,
+ )
+ model = Model(args)
+ mx.eval(model.parameters())
+ return model
+
+
+@pytest.mark.slow
+class TestBatchedTopKPrecompute:
+ @pytest.fixture(autouse=True)
+ def _reset_globals(self) -> None:
+ import exo.worker.engines.mlx.patches.opt_batch_gen as _mod
+
+ _mod._pending_topk_idx = None
+ _mod._pending_topk_val = None
+ _mod._pending_selected_lps = None
+
+ def _run_generator(
+ self, model: nn.Module, prompts: list[list[int]], steps: int, needs_topk: bool
+ ) -> list[list[BatchGenerator.Response]]:
+ apply_batch_gen_patch()
+ gen = BatchGenerator(model=model, stop_tokens=set(), prefill_step_size=512)
+ gen._needs_topk = needs_topk
+ gen.insert(prompts)
+ all_responses: list[list[BatchGenerator.Response]] = []
+ for _ in range(steps + len(prompts)):
+ responses = gen.next()
+ if responses:
+ all_responses.append(responses)
+ if gen.active_batch is None and not gen.unprocessed_prompts:
+ break
+ gen.close()
+ return all_responses
+
+ def test_precomputed_topk_attached_to_responses(self) -> None:
+ model = _tiny_model()
+ steps = self._run_generator(model, [[1, 2, 3]], 5, needs_topk=True)
+ found_precomputed = False
+ for step_responses in steps:
+ for resp in step_responses:
+ if hasattr(resp, "_topk_indices"):
+ found_precomputed = True
+ assert hasattr(resp, "_topk_values"), (
+ "Response missing _topk_values"
+ )
+ assert hasattr(resp, "_selected_logprob"), (
+ "Response missing _selected_logprob"
+ )
+ assert len(resp._topk_indices) == _PRECOMPUTE_TOP_K
+ assert len(resp._topk_values) == _PRECOMPUTE_TOP_K
+ assert found_precomputed, "No responses had precomputed topk"
+
+ def test_no_topk_when_not_needed(self) -> None:
+ model = _tiny_model()
+ steps = self._run_generator(model, [[1, 2, 3]], 5, needs_topk=False)
+ for step_responses in steps:
+ for resp in step_responses:
+ assert not hasattr(resp, "_topk_indices")
+
+ def test_precomputed_matches_fallback_in_batch(self) -> None:
+ model = _tiny_model()
+ tok = _mock_tokenizer()
+ steps = self._run_generator(model, [[1, 2, 3]], 10, needs_topk=True)
+ for step_responses in steps[1:]:
+ for resp in step_responses:
+ if not hasattr(resp, "_topk_indices"):
+ continue
+ selected_fb, items_fb = extract_top_logprobs(
+ resp.logprobs, tok, top_logprobs=5, selected_token=resp.token
+ )
+ selected_pc, items_pc = extract_top_logprobs(
+ resp.logprobs,
+ tok,
+ top_logprobs=5,
+ selected_token=resp.token,
+ precomputed_indices=resp._topk_indices,
+ precomputed_values=resp._topk_values,
+ precomputed_selected=resp._selected_logprob,
+ )
+ assert selected_pc == pytest.approx(selected_fb, abs=1e-5)
+ for a, b in zip(items_pc, items_fb, strict=True):
+ assert a.token == b.token
+ assert a.logprob == pytest.approx(b.logprob, abs=1e-5)
+
+ def test_topk_correct_after_batch_shrink(self) -> None:
+ model = _tiny_model()
+ tok = _mock_tokenizer()
+ apply_batch_gen_patch()
+ gen = BatchGenerator(
+ model=model, stop_tokens={0}, prefill_step_size=512, max_tokens=3
+ )
+ gen._needs_topk = True
+ gen.insert([[1, 2, 3], [4, 5, 6]], max_tokens=[3, 20])
+
+ seen_shrink = False
+ for _ in range(30):
+ responses = gen.next()
+ for resp in responses:
+ if resp.finish_reason is not None:
+ seen_shrink = True
+ continue
+ if not hasattr(resp, "_topk_indices"):
+ continue
+ selected_fb, items_fb = extract_top_logprobs(
+ resp.logprobs, tok, top_logprobs=5, selected_token=resp.token
+ )
+ selected_pc, _ = extract_top_logprobs(
+ resp.logprobs,
+ tok,
+ top_logprobs=5,
+ selected_token=resp.token,
+ precomputed_indices=resp._topk_indices,
+ precomputed_values=resp._topk_values,
+ precomputed_selected=resp._selected_logprob,
+ )
+ assert selected_pc == pytest.approx(selected_fb, abs=1e-5), (
+ f"Mismatch after batch shrink: precomputed={selected_pc}, fallback={selected_fb}"
+ )
+ if gen.active_batch is None and not gen.unprocessed_prompts:
+ break
+
+ gen.close()
+ assert seen_shrink, "Expected at least one request to finish (batch shrink)"
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 8b22e6b3..91e6fd60 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -638,6 +638,7 @@ class NullKVCache(KVCache):
@property
def state(self) -> tuple[mx.array, mx.array]:
# matches what mx.save_safetensors / mx.eval expect
+ assert self.keys is not None and self.values is not None
return self.keys, self.values
@state.setter
diff --git a/src/exo/worker/runner/bootstrap.py b/src/exo/worker/runner/bootstrap.py
index 35d4250e..644b1a51 100644
--- a/src/exo/worker/runner/bootstrap.py
+++ b/src/exo/worker/runner/bootstrap.py
@@ -8,6 +8,7 @@ from exo.shared.types.tasks import Task, TaskId
from exo.shared.types.worker.instances import BoundInstance
from exo.shared.types.worker.runners import RunnerFailed
from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
+from exo.worker.engines.mlx.patches import apply_mlx_patches
logger: "loguru.Logger" = loguru.logger
@@ -45,6 +46,8 @@ def entrypoint(
else:
from exo.worker.runner.llm_inference.runner import Runner
+ apply_mlx_patches()
+
runner = Runner(
bound_instance, event_sender, task_receiver, cancel_receiver
)
diff --git a/src/exo/worker/runner/llm_inference/batch_generator.py b/src/exo/worker/runner/llm_inference/batch_generator.py
index e0382ca4..33bbb3a6 100644
--- a/src/exo/worker/runner/llm_inference/batch_generator.py
+++ b/src/exo/worker/runner/llm_inference/batch_generator.py
@@ -427,7 +427,8 @@ class BatchGenerator(InferenceGenerator):
task, queue, output_generator = self._active_tasks[uid]
queue.push(response)
- parsed = next(output_generator)
+ # If a generator fails to parse for some reason and returns early, we should not crash
+ parsed = next(output_generator, None)
if parsed is not None:
output.append((task.task_id, parsed))
diff --git a/src/exo/worker/runner/llm_inference/runner.py b/src/exo/worker/runner/llm_inference/runner.py
index 1f9f1e05..4f2f3d30 100644
--- a/src/exo/worker/runner/llm_inference/runner.py
+++ b/src/exo/worker/runner/llm_inference/runner.py
@@ -319,7 +319,9 @@ class Runner:
return ExitCode.AllTasksComplete
def send_response(
- self, response: GenerationResponse | ToolCallResponse, command_id: CommandId
+ self,
+ response: GenerationResponse | ToolCallResponse,
+ command_id: CommandId,
):
match response:
case GenerationResponse():
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py b/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py
index 850db320..e13eb0f7 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py
@@ -43,8 +43,8 @@ def run_pipeline_device(
def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
for layer in self.layers:
- x = layer(x, *args, **kwargs) # pyright: ignore[reportUnknownVariableType]
- return x # pyright: ignore[reportUnknownVariableType]
+ x = layer(x, *args, **kwargs)
+ return x
try:
group = mx.distributed.init(backend="ring", strict=True)
diff --git a/uv.lock b/uv.lock
index ba630360..82411d80 100644
--- a/uv.lock
+++ b/uv.lock
@@ -213,14 +213,20 @@ sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8
wheels = [
{ url = "https://files.pythonhosted.org/packages/b0/1e/d22cc63332bd59b06481ceaac49d6c507598642e2230f201649058a7e704/cffi-2.0.0-cp313-cp313-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:07b271772c100085dd28b74fa0cd81c8fb1a3ba18b21e03d7c27f3436a10606b", size = 212446, upload-time = "2025-09-08T23:23:03.472Z" },
{ url = "https://files.pythonhosted.org/packages/a9/f5/a2c23eb03b61a0b8747f211eb716446c826ad66818ddc7810cc2cc19b3f2/cffi-2.0.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d48a880098c96020b02d5a1f7d9251308510ce8858940e6fa99ece33f610838b", size = 220101, upload-time = "2025-09-08T23:23:04.792Z" },
+ { url = "https://files.pythonhosted.org/packages/f2/7f/e6647792fc5850d634695bc0e6ab4111ae88e89981d35ac269956605feba/cffi-2.0.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:f93fd8e5c8c0a4aa1f424d6173f14a892044054871c771f8566e4008eaa359d2", size = 207948, upload-time = "2025-09-08T23:23:06.127Z" },
+ { url = "https://files.pythonhosted.org/packages/cb/1e/a5a1bd6f1fb30f22573f76533de12a00bf274abcdc55c8edab639078abb6/cffi-2.0.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:dd4f05f54a52fb558f1ba9f528228066954fee3ebe629fc1660d874d040ae5a3", size = 206422, upload-time = "2025-09-08T23:23:07.753Z" },
{ url = "https://files.pythonhosted.org/packages/98/df/0a1755e750013a2081e863e7cd37e0cdd02664372c754e5560099eb7aa44/cffi-2.0.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c8d3b5532fc71b7a77c09192b4a5a200ea992702734a2e9279a37f2478236f26", size = 219499, upload-time = "2025-09-08T23:23:09.648Z" },
{ url = "https://files.pythonhosted.org/packages/50/e1/a969e687fcf9ea58e6e2a928ad5e2dd88cc12f6f0ab477e9971f2309b57c/cffi-2.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d9b29c1f0ae438d5ee9acb31cadee00a58c46cc9c0b2f9038c6b0b3470877a8c", size = 222928, upload-time = "2025-09-08T23:23:10.928Z" },
{ url = "https://files.pythonhosted.org/packages/36/54/0362578dd2c9e557a28ac77698ed67323ed5b9775ca9d3fe73fe191bb5d8/cffi-2.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6d50360be4546678fc1b79ffe7a66265e28667840010348dd69a314145807a1b", size = 221302, upload-time = "2025-09-08T23:23:12.42Z" },
{ url = "https://files.pythonhosted.org/packages/d6/43/0e822876f87ea8a4ef95442c3d766a06a51fc5298823f884ef87aaad168c/cffi-2.0.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:24b6f81f1983e6df8db3adc38562c83f7d4a0c36162885ec7f7b77c7dcbec97b", size = 220049, upload-time = "2025-09-08T23:23:20.853Z" },
+ { url = "https://files.pythonhosted.org/packages/b4/89/76799151d9c2d2d1ead63c2429da9ea9d7aac304603de0c6e8764e6e8e70/cffi-2.0.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:12873ca6cb9b0f0d3a0da705d6086fe911591737a59f28b7936bdfed27c0d47c", size = 207793, upload-time = "2025-09-08T23:23:22.08Z" },
+ { url = "https://files.pythonhosted.org/packages/bb/dd/3465b14bb9e24ee24cb88c9e3730f6de63111fffe513492bf8c808a3547e/cffi-2.0.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:d9b97165e8aed9272a6bb17c01e3cc5871a594a446ebedc996e2397a1c1ea8ef", size = 206300, upload-time = "2025-09-08T23:23:23.314Z" },
{ url = "https://files.pythonhosted.org/packages/47/d9/d83e293854571c877a92da46fdec39158f8d7e68da75bf73581225d28e90/cffi-2.0.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:afb8db5439b81cf9c9d0c80404b60c3cc9c3add93e114dcae767f1477cb53775", size = 219244, upload-time = "2025-09-08T23:23:24.541Z" },
{ url = "https://files.pythonhosted.org/packages/2b/0f/1f177e3683aead2bb00f7679a16451d302c436b5cbf2505f0ea8146ef59e/cffi-2.0.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:737fe7d37e1a1bffe70bd5754ea763a62a066dc5913ca57e957824b72a85e205", size = 222828, upload-time = "2025-09-08T23:23:26.143Z" },
{ url = "https://files.pythonhosted.org/packages/c6/0f/cafacebd4b040e3119dcb32fed8bdef8dfe94da653155f9d0b9dc660166e/cffi-2.0.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:38100abb9d1b1435bc4cc340bb4489635dc2f0da7456590877030c9b3d40b0c1", size = 220926, upload-time = "2025-09-08T23:23:27.873Z" },
{ url = "https://files.pythonhosted.org/packages/be/b4/c56878d0d1755cf9caa54ba71e5d049479c52f9e4afc230f06822162ab2f/cffi-2.0.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7cc09976e8b56f8cebd752f7113ad07752461f48a58cbba644139015ac24954c", size = 221593, upload-time = "2025-09-08T23:23:31.91Z" },
+ { url = "https://files.pythonhosted.org/packages/e0/0d/eb704606dfe8033e7128df5e90fee946bbcb64a04fcdaa97321309004000/cffi-2.0.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:92b68146a71df78564e4ef48af17551a5ddd142e5190cdf2c5624d0c3ff5b2e8", size = 209354, upload-time = "2025-09-08T23:23:33.214Z" },
+ { url = "https://files.pythonhosted.org/packages/d8/19/3c435d727b368ca475fb8742ab97c9cb13a0de600ce86f62eab7fa3eea60/cffi-2.0.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:b1e74d11748e7e98e2f426ab176d4ed720a64412b6a15054378afdb71e0f37dc", size = 208480, upload-time = "2025-09-08T23:23:34.495Z" },
{ url = "https://files.pythonhosted.org/packages/d0/44/681604464ed9541673e486521497406fadcc15b5217c3e326b061696899a/cffi-2.0.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:28a3a209b96630bca57cce802da70c266eb08c6e97e5afd61a75611ee6c64592", size = 221584, upload-time = "2025-09-08T23:23:36.096Z" },
{ url = "https://files.pythonhosted.org/packages/25/8e/342a504ff018a2825d395d44d63a767dd8ebc927ebda557fecdaca3ac33a/cffi-2.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:7553fb2090d71822f02c629afe6042c299edf91ba1bf94951165613553984512", size = 224443, upload-time = "2025-09-08T23:23:37.328Z" },
{ url = "https://files.pythonhosted.org/packages/e1/5e/b666bacbbc60fbf415ba9988324a132c9a7a0448a9a8f125074671c0f2c3/cffi-2.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6c6c373cfc5c83a975506110d17457138c8c63016b563cc9ed6e056a82f13ce4", size = 223437, upload-time = "2025-09-08T23:23:38.945Z" },
@@ -344,8 +350,10 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/5c/49/498c86566a1d80e978b42f0d702795f69887005548c041636df6ae1ca64c/cryptography-46.0.3-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:01ca9ff2885f3acc98c29f1860552e37f6d7c7d013d7334ff2a9de43a449315d", size = 4450807, upload-time = "2025-10-15T23:16:56.414Z" },
{ url = "https://files.pythonhosted.org/packages/4b/0a/863a3604112174c8624a2ac3c038662d9e59970c7f926acdcfaed8d61142/cryptography-46.0.3-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:6eae65d4c3d33da080cff9c4ab1f711b15c1d9760809dad6ea763f3812d254cb", size = 4299615, upload-time = "2025-10-15T23:16:58.442Z" },
{ url = "https://files.pythonhosted.org/packages/64/02/b73a533f6b64a69f3cd3872acb6ebc12aef924d8d103133bb3ea750dc703/cryptography-46.0.3-cp311-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:e5bf0ed4490068a2e72ac03d786693adeb909981cc596425d09032d372bcc849", size = 4016800, upload-time = "2025-10-15T23:17:00.378Z" },
+ { url = "https://files.pythonhosted.org/packages/25/d5/16e41afbfa450cde85a3b7ec599bebefaef16b5c6ba4ec49a3532336ed72/cryptography-46.0.3-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:5ecfccd2329e37e9b7112a888e76d9feca2347f12f37918facbb893d7bb88ee8", size = 4984707, upload-time = "2025-10-15T23:17:01.98Z" },
{ url = "https://files.pythonhosted.org/packages/c9/56/e7e69b427c3878352c2fb9b450bd0e19ed552753491d39d7d0a2f5226d41/cryptography-46.0.3-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:a2c0cd47381a3229c403062f764160d57d4d175e022c1df84e168c6251a22eec", size = 4482541, upload-time = "2025-10-15T23:17:04.078Z" },
{ url = "https://files.pythonhosted.org/packages/78/f6/50736d40d97e8483172f1bb6e698895b92a223dba513b0ca6f06b2365339/cryptography-46.0.3-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:549e234ff32571b1f4076ac269fcce7a808d3bf98b76c8dd560e42dbc66d7d91", size = 4299464, upload-time = "2025-10-15T23:17:05.483Z" },
+ { url = "https://files.pythonhosted.org/packages/00/de/d8e26b1a855f19d9994a19c702fa2e93b0456beccbcfe437eda00e0701f2/cryptography-46.0.3-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:c0a7bb1a68a5d3471880e264621346c48665b3bf1c3759d682fc0864c540bd9e", size = 4950838, upload-time = "2025-10-15T23:17:07.425Z" },
{ url = "https://files.pythonhosted.org/packages/8f/29/798fc4ec461a1c9e9f735f2fc58741b0daae30688f41b2497dcbc9ed1355/cryptography-46.0.3-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:10b01676fc208c3e6feeb25a8b83d81767e8059e1fe86e1dc62d10a3018fa926", size = 4481596, upload-time = "2025-10-15T23:17:09.343Z" },
{ url = "https://files.pythonhosted.org/packages/15/8d/03cd48b20a573adfff7652b76271078e3045b9f49387920e7f1f631d125e/cryptography-46.0.3-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0abf1ffd6e57c67e92af68330d05760b7b7efb243aab8377e583284dbab72c71", size = 4426782, upload-time = "2025-10-15T23:17:11.22Z" },
{ url = "https://files.pythonhosted.org/packages/fa/b1/ebacbfe53317d55cf33165bda24c86523497a6881f339f9aae5c2e13e57b/cryptography-46.0.3-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a04bee9ab6a4da801eb9b51f1b708a1b5b5c9eb48c03f74198464c66f0d344ac", size = 4698381, upload-time = "2025-10-15T23:17:12.829Z" },
@@ -353,8 +361,10 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/c5/fd/bc1daf8230eaa075184cbbf5f8cd00ba9db4fd32d63fb83da4671b72ed8a/cryptography-46.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:39b6755623145ad5eff1dab323f4eae2a32a77a7abef2c5089a04a3d04366715", size = 4435078, upload-time = "2025-10-15T23:17:23.042Z" },
{ url = "https://files.pythonhosted.org/packages/82/98/d3bd5407ce4c60017f8ff9e63ffee4200ab3e23fe05b765cab805a7db008/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:db391fa7c66df6762ee3f00c95a89e6d428f4d60e7abc8328f4fe155b5ac6e54", size = 4293460, upload-time = "2025-10-15T23:17:24.885Z" },
{ url = "https://files.pythonhosted.org/packages/26/e9/e23e7900983c2b8af7a08098db406cf989d7f09caea7897e347598d4cd5b/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:78a97cf6a8839a48c49271cdcbd5cf37ca2c1d6b7fdd86cc864f302b5e9bf459", size = 3995237, upload-time = "2025-10-15T23:17:26.449Z" },
+ { url = "https://files.pythonhosted.org/packages/91/15/af68c509d4a138cfe299d0d7ddb14afba15233223ebd933b4bbdbc7155d3/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:dfb781ff7eaa91a6f7fd41776ec37c5853c795d3b358d4896fdbb5df168af422", size = 4967344, upload-time = "2025-10-15T23:17:28.06Z" },
{ url = "https://files.pythonhosted.org/packages/ca/e3/8643d077c53868b681af077edf6b3cb58288b5423610f21c62aadcbe99f4/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:6f61efb26e76c45c4a227835ddeae96d83624fb0d29eb5df5b96e14ed1a0afb7", size = 4466564, upload-time = "2025-10-15T23:17:29.665Z" },
{ url = "https://files.pythonhosted.org/packages/0e/43/c1e8726fa59c236ff477ff2b5dc071e54b21e5a1e51aa2cee1676f1c986f/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:23b1a8f26e43f47ceb6d6a43115f33a5a37d57df4ea0ca295b780ae8546e8044", size = 4292415, upload-time = "2025-10-15T23:17:31.686Z" },
+ { url = "https://files.pythonhosted.org/packages/42/f9/2f8fefdb1aee8a8e3256a0568cffc4e6d517b256a2fe97a029b3f1b9fe7e/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:b419ae593c86b87014b9be7396b385491ad7f320bde96826d0dd174459e54665", size = 4931457, upload-time = "2025-10-15T23:17:33.478Z" },
{ url = "https://files.pythonhosted.org/packages/79/30/9b54127a9a778ccd6d27c3da7563e9f2d341826075ceab89ae3b41bf5be2/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:50fc3343ac490c6b08c0cf0d704e881d0d660be923fd3076db3e932007e726e3", size = 4466074, upload-time = "2025-10-15T23:17:35.158Z" },
{ url = "https://files.pythonhosted.org/packages/ac/68/b4f4a10928e26c941b1b6a179143af9f4d27d88fe84a6a3c53592d2e76bf/cryptography-46.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:22d7e97932f511d6b0b04f2bfd818d73dcd5928db509460aaf48384778eb6d20", size = 4420569, upload-time = "2025-10-15T23:17:37.188Z" },
{ url = "https://files.pythonhosted.org/packages/a3/49/3746dab4c0d1979888f125226357d3262a6dd40e114ac29e3d2abdf1ec55/cryptography-46.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:d55f3dffadd674514ad19451161118fd010988540cee43d8bc20675e775925de", size = 4681941, upload-time = "2025-10-15T23:17:39.236Z" },
@@ -362,8 +372,10 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/26/42/fa8389d4478368743e24e61eea78846a0006caffaf72ea24a15159215a14/cryptography-46.0.3-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:15ab9b093e8f09daab0f2159bb7e47532596075139dd74365da52ecc9cb46c5d", size = 4440029, upload-time = "2025-10-15T23:17:49.837Z" },
{ url = "https://files.pythonhosted.org/packages/5f/eb/f483db0ec5ac040824f269e93dd2bd8a21ecd1027e77ad7bdf6914f2fd80/cryptography-46.0.3-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:46acf53b40ea38f9c6c229599a4a13f0d46a6c3fa9ef19fc1a124d62e338dfa0", size = 4297222, upload-time = "2025-10-15T23:17:51.357Z" },
{ url = "https://files.pythonhosted.org/packages/fd/cf/da9502c4e1912cb1da3807ea3618a6829bee8207456fbbeebc361ec38ba3/cryptography-46.0.3-cp38-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:10ca84c4668d066a9878890047f03546f3ae0a6b8b39b697457b7757aaf18dbc", size = 4012280, upload-time = "2025-10-15T23:17:52.964Z" },
+ { url = "https://files.pythonhosted.org/packages/6b/8f/9adb86b93330e0df8b3dcf03eae67c33ba89958fc2e03862ef1ac2b42465/cryptography-46.0.3-cp38-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:36e627112085bb3b81b19fed209c05ce2a52ee8b15d161b7c643a7d5a88491f3", size = 4978958, upload-time = "2025-10-15T23:17:54.965Z" },
{ url = "https://files.pythonhosted.org/packages/d1/a0/5fa77988289c34bdb9f913f5606ecc9ada1adb5ae870bd0d1054a7021cc4/cryptography-46.0.3-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:1000713389b75c449a6e979ffc7dcc8ac90b437048766cef052d4d30b8220971", size = 4473714, upload-time = "2025-10-15T23:17:56.754Z" },
{ url = "https://files.pythonhosted.org/packages/14/e5/fc82d72a58d41c393697aa18c9abe5ae1214ff6f2a5c18ac470f92777895/cryptography-46.0.3-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:b02cf04496f6576afffef5ddd04a0cb7d49cf6be16a9059d793a30b035f6b6ac", size = 4296970, upload-time = "2025-10-15T23:17:58.588Z" },
+ { url = "https://files.pythonhosted.org/packages/78/06/5663ed35438d0b09056973994f1aec467492b33bd31da36e468b01ec1097/cryptography-46.0.3-cp38-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:71e842ec9bc7abf543b47cf86b9a743baa95f4677d22baa4c7d5c69e49e9bc04", size = 4940236, upload-time = "2025-10-15T23:18:00.897Z" },
{ url = "https://files.pythonhosted.org/packages/fc/59/873633f3f2dcd8a053b8dd1d38f783043b5fce589c0f6988bf55ef57e43e/cryptography-46.0.3-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:402b58fc32614f00980b66d6e56a5b4118e6cb362ae8f3fda141ba4689bd4506", size = 4472642, upload-time = "2025-10-15T23:18:02.749Z" },
{ url = "https://files.pythonhosted.org/packages/3d/39/8e71f3930e40f6877737d6f69248cf74d4e34b886a3967d32f919cc50d3b/cryptography-46.0.3-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ef639cb3372f69ec44915fafcd6698b6cc78fbe0c2ea41be867f6ed612811963", size = 4423126, upload-time = "2025-10-15T23:18:04.85Z" },
{ url = "https://files.pythonhosted.org/packages/cd/c7/f65027c2810e14c3e7268353b1681932b87e5a48e65505d8cc17c99e36ae/cryptography-46.0.3-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:3b51b8ca4f1c6453d8829e1eb7299499ca7f313900dd4d89a24b8b87c0a780d4", size = 4686573, upload-time = "2025-10-15T23:18:06.908Z" },
@@ -473,7 +485,7 @@ dependencies = [
{ name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" },
- { name = "mlx", version = "0.30.7.dev20260303+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+ { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "openai-harmony", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -512,7 +524,7 @@ requires-dist = [
{ name = "mflux", specifier = "==0.16.9" },
{ name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" },
{ name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.6" },
- { name = "mlx-lm", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Feval-left-padding-in-batched-rotation" },
+ { name = "mlx-lm", git = "https://github.com/ml-explore/mlx-lm?branch=main" },
{ name = "msgspec", specifier = ">=0.19.0" },
{ name = "openai-harmony", specifier = ">=0.0.8" },
{ name = "psutil", specifier = ">=7.0.0" },
@@ -1351,7 +1363,7 @@ dependencies = [
{ name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" },
- { name = "mlx", version = "0.30.7.dev20260303+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+ { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1399,7 +1411,7 @@ cuda13 = [
[[package]]
name = "mlx"
-version = "0.30.7.dev20260303+257d5692"
+version = "0.30.7.dev20260225+257d5692"
source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }
resolution-markers = [
"python_full_version >= '3.14' and sys_platform == 'darwin'",
@@ -1432,11 +1444,11 @@ wheels = [
[[package]]
name = "mlx-lm"
-version = "0.31.0"
-source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Feval-left-padding-in-batched-rotation#5e0c484cfb5c68d281a71409927ee1bb75adaae2" }
+version = "0.31.2"
+source = { git = "https://github.com/ml-explore/mlx-lm?branch=main#ed7884cb80968e0e77fce6cde5d1597952bbd524" }
dependencies = [
{ name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
- { name = "mlx", version = "0.30.7.dev20260303+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+ { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
← 248919c2 Fix first start in offline mode crash (#1782)
·
back to Exo
·
fix warmup consistency for slow machines (#1748) fac6832e →