electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve batch performance and stats reporting (#1777)

Browse Source 
## 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"
/>
This commit is contained in:
rltakashige
2026-03-24 14:03:03 +00:00
committed by GitHub
parent 248919c2a8
commit 7df3774ca2
27 changed files with 918 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.mlx_typings/mlx/core/__init__.pyi
+1 -1
View File
@@ -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``.
.mlx_typings/mlx/nn/__init__.pyi
+2 -6
View File
@@ -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
.mlx_typings/mlx/nn/layers/__init__.pyi
+16 -20
View File
@@ -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 *
.mlx_typings/mlx/nn/layers/base.pyi
+1 -1
View File
@@ -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``."""
.mlx_typings/mlx_lm/generate.pyi
+10 -5
View File
@@ -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:
.mlx_typings/mlx_lm/models/cache.pyi
+25 -31
View File
@@ -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,
.mlx_typings/mlx_lm/models/gated_delta.pyi
+35
View File
@@ -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]: ...
.mlx_typings/mlx_lm/models/rope_utils.pyi
+51
View File
@@ -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: ...
bench/exo_bench.py
+5 -7
View File
@@ -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(
dashboard/src/lib/stores/app.svelte.ts
+28 -5
View File
@@ -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)
pyproject.toml
+1 -1
View File
@@ -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 }
src/exo/api/adapters/chat_completions.py
+4
View File
@@ -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"
src/exo/worker/engines/image/models/qwen/wrappers.py
+2 -2
View File
@@ -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)
src/exo/worker/engines/mlx/auto_parallel.py
+1 -1
View File
@@ -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
src/exo/worker/engines/mlx/generator/batch_generate.py
+47 -26
View File
@@ -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,
src/exo/worker/engines/mlx/generator/generate.py
+34 -40
View File
@@ -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
src/exo/worker/engines/mlx/patches/__init__.py
+17
View File
@@ -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()
src/exo/worker/engines/mlx/patches/high_precision_gdn_softplus.py
+27
View File
@@ -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)
src/exo/worker/engines/mlx/patches/opt_batch_gen.py
+173
View File
@@ -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
src/exo/worker/engines/mlx/patches/standard_yarn_rope.py
+118
View File
@@ -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
src/exo/worker/engines/mlx/tests/test_extract_top_logprobs.py
+290
View File
@@ -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)"
src/exo/worker/engines/mlx/utils_mlx.py
+1
View File
@@ -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
src/exo/worker/runner/bootstrap.py
+3
View File
@@ -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
)
src/exo/worker/runner/llm_inference/batch_generator.py
+2 -1
View File
@@ -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))
src/exo/worker/runner/llm_inference/runner.py
+3 -1
View File
@@ -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():
src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py
+2 -2
View File
@@ -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)
uv.lock
Generated
+19 -7
View File
@@ -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'" },