← back to Exo
Fix BatchGenerator in line with upstream refactor (and prevent Qwen3.5 memory leak) (#1835)
43b3df45fb4080819f0a29d9db2aa4b76df99501 · 2026-04-07 12:50:12 +0100 · rltakashige
## Motivation
MLX LM has had a massive refactor to their BatchGenerator recently.
Since we'd like new features from MLX LM such as Gemma 4, we need to
update the code to handle this.
Additionally this fixes a significant memory leak in GatedDeltaNet (the
difference is quite substantial, up to 1GB every 1000 tokens, explaining
several memory issues users were facing with Qwen3.5 models)
## Testing
Before
<img width="3146" height="884" alt="image"
src="https://github.com/user-attachments/assets/5af0f55a-393c-4a32-9eed-ae43f1611af4"
/>
After (no memory leak, as one of the changes upstream)
<img width="3190" height="892" alt="image"
src="https://github.com/user-attachments/assets/f0bd128d-fd48-40d4-9bbd-50a564beab14"
/>
Files touched
A .mlx_typings/mlx_lm/_version.pyiM .mlx_typings/mlx_lm/generate.pyiA .mlx_typings/mlx_lm/models/activations.pyiA .mlx_typings/mlx_lm/tuner/dora.pyiA .mlx_typings/mlx_lm/tuner/lora.pyiA .mlx_typings/mlx_lm/tuner/utils.pyiM bench/exo_bench.pyM nix/darwin-build-fixes.patchM nix/mlx.nixM src/exo/worker/engines/mlx/generator/batch_generate.pyM src/exo/worker/engines/mlx/patches/__init__.pyM src/exo/worker/engines/mlx/patches/opt_batch_gen.pyM src/exo/worker/engines/mlx/tests/test_extract_top_logprobs.pyM uv.lock
Diff
commit 43b3df45fb4080819f0a29d9db2aa4b76df99501
Author: rltakashige <rl.takashige@gmail.com>
Date: Tue Apr 7 12:50:12 2026 +0100
Fix BatchGenerator in line with upstream refactor (and prevent Qwen3.5 memory leak) (#1835)
## Motivation
MLX LM has had a massive refactor to their BatchGenerator recently.
Since we'd like new features from MLX LM such as Gemma 4, we need to
update the code to handle this.
Additionally this fixes a significant memory leak in GatedDeltaNet (the
difference is quite substantial, up to 1GB every 1000 tokens, explaining
several memory issues users were facing with Qwen3.5 models)
## Testing
Before
<img width="3146" height="884" alt="image"
src="https://github.com/user-attachments/assets/5af0f55a-393c-4a32-9eed-ae43f1611af4"
/>
After (no memory leak, as one of the changes upstream)
<img width="3190" height="892" alt="image"
src="https://github.com/user-attachments/assets/f0bd128d-fd48-40d4-9bbd-50a564beab14"
/>
---
.mlx_typings/mlx_lm/_version.pyi | 5 +
.mlx_typings/mlx_lm/generate.pyi | 336 ++++++++++++++++-----
.mlx_typings/mlx_lm/models/activations.pyi | 19 ++
.mlx_typings/mlx_lm/tuner/dora.pyi | 50 +++
.mlx_typings/mlx_lm/tuner/lora.pyi | 66 ++++
.mlx_typings/mlx_lm/tuner/utils.pyi | 57 ++++
bench/exo_bench.py | 3 +
nix/darwin-build-fixes.patch | 2 +-
nix/mlx.nix | 2 +-
.../worker/engines/mlx/generator/batch_generate.py | 42 ++-
src/exo/worker/engines/mlx/patches/__init__.py | 1 -
.../worker/engines/mlx/patches/opt_batch_gen.py | 242 ++++++---------
.../engines/mlx/tests/test_extract_top_logprobs.py | 221 --------------
uv.lock | 14 +-
14 files changed, 593 insertions(+), 467 deletions(-)
diff --git a/.mlx_typings/mlx_lm/_version.pyi b/.mlx_typings/mlx_lm/_version.pyi
new file mode 100644
index 00000000..fc797c9a
--- /dev/null
+++ b/.mlx_typings/mlx_lm/_version.pyi
@@ -0,0 +1,5 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+__version__ = ...
diff --git a/.mlx_typings/mlx_lm/generate.pyi b/.mlx_typings/mlx_lm/generate.pyi
index 4006259a..aa38c027 100644
--- a/.mlx_typings/mlx_lm/generate.pyi
+++ b/.mlx_typings/mlx_lm/generate.pyi
@@ -3,13 +3,12 @@ This type stub file was generated by pyright.
"""
import contextlib
-from dataclasses import dataclass
-from typing import Any, Callable, Generator, List, Optional, Tuple, Union
-
import mlx.core as mx
import mlx.nn as nn
+from dataclasses import dataclass
+from collections import deque
+from typing import Any, Callable, Generator, List, Optional, Sequence, Tuple, Union
from transformers import PreTrainedTokenizer
-
from .tokenizer_utils import TokenizerWrapper
DEFAULT_PROMPT = ...
@@ -29,6 +28,7 @@ def str2bool(string): # -> bool:
...
def setup_arg_parser(): # -> ArgumentParser:
"""Set up and return the argument parser."""
+ ...
generation_stream: mx.Stream
@@ -43,6 +43,7 @@ def wired_limit(
async eval could be running pass in the streams to synchronize with prior
to exiting the context manager.
"""
+ ...
@dataclass
class GenerationResponse:
"""
@@ -91,7 +92,7 @@ def generate_step(
kv_bits: Optional[int] = ...,
kv_group_size: int = ...,
quantized_kv_start: int = ...,
- prompt_progress_callback: Optional[Callable[[int], int]] = ...,
+ prompt_progress_callback: Optional[Callable[[int, int], None]] = ...,
input_embeddings: Optional[mx.array] = ...,
) -> Generator[Tuple[mx.array, mx.array], None, None]:
"""
@@ -117,7 +118,7 @@ def generate_step(
kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
quantized_kv_start (int): Step to begin using a quantized KV cache.
when ``kv_bits`` is non-None. Default: ``0``.
- prompt_progress_callback (Callable[[int], int]): A call-back which takes the
+ prompt_progress_callback (Callable[[int, int], None]): A call-back which takes the
prompt tokens processed so far and the total number of prompt tokens.
input_embeddings (mx.array, optional): Input embeddings to use instead of or in
conjunction with prompt tokens. Default: ``None``.
@@ -125,6 +126,7 @@ def generate_step(
Yields:
Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
"""
+ ...
def speculative_generate_step(
prompt: mx.array,
@@ -170,6 +172,7 @@ def speculative_generate_step(
Tuple[mx.array, mx.array, bool]: One token, a vector of log probabilities,
and a bool indicating if the token was generated by the draft model
"""
+ ...
def stream_generate(
model: nn.Module,
@@ -177,7 +180,7 @@ def stream_generate(
prompt: Union[str, mx.array, List[int]],
max_tokens: int = ...,
draft_model: Optional[nn.Module] = ...,
- **kwargs: object,
+ **kwargs: Any,
) -> Generator[GenerationResponse, None, None]:
"""
A generator producing text based on the given prompt from the model.
@@ -199,6 +202,7 @@ def stream_generate(
GenerationResponse: An instance containing the generated text segment and
associated metadata. See :class:`GenerationResponse` for details.
"""
+ ...
def generate(
model: nn.Module,
@@ -219,6 +223,9 @@ def generate(
kwargs: The remaining options get passed to :func:`stream_generate`.
See :func:`stream_generate` for more details.
"""
+ ...
+
+def _merge_caches(caches: List[List[Any]]) -> List[Any]: ...
@dataclass
class BatchStats:
"""
@@ -242,69 +249,196 @@ class BatchStats:
generation_time: float = ...
peak_memory: float = ...
-@dataclass
-class BatchResponse:
+class SequenceStateMachine:
+ """A state machine that uses one Aho-Corasick trie per state to efficiently
+ track state across a generated sequence.
+
+ The transitions are provided as state -> [(sequence, new_state)].
+
+ Example:
+
+ sm = SequenceStateMachine(
+ transitions={
+ "normal": [
+ (think_start_tokens, "reasoning"),
+ (tool_start_tokens, "tool"),
+ (eos, None),
+ ],
+ "reasoning": [
+ (think_end_tokens, "normal"),
+ (eos, None),
+ ],
+ "tool": [
+ (tool_end_tokens, None),
+ (eos, None)
+ ],
+ },
+ initial="normal"
+ )
"""
- An data object to hold a batch generation response.
-
- Args:
- texts: (List[str]): The generated text for each prompt.
- stats (BatchStats): Statistics about the generation.
+ def __init__(self, transitions=..., initial=...) -> None: ...
+ def __deepcopy__(self, memo): # -> SequenceStateMachine:
+ ...
+ def make_state(self): # -> tuple[str, Any, dict[Any, Any]]:
+ ...
+ @staticmethod
+ def match(state, x): # -> tuple[tuple[Any, Any | None, Any], Any | None, Any]:
+ ...
+
+class PromptProcessingBatch:
"""
+ A batch processor for prompt tokens with support for incremental processing.
- texts: List[str]
- stats: BatchStats
- caches: Optional[List[List[Any]]]
+ This class handles batched prompt processing, managing KV caches and preparing
+ tokens for generation. It supports extending, filtering, and splitting batches.
+ """
+ @dataclass
+ class Response:
+ uid: int
+ progress: tuple
+ end_of_segment: bool
+ end_of_prompt: bool
+ ...
-def _left_pad_prompts(prompts: Any, max_length: Optional[int] = ...) -> mx.array: ...
-def _right_pad_prompts(prompts: Any, max_length: Optional[int] = ...) -> mx.array: ...
-def _make_cache(
- model: Any, left_padding: Any, max_kv_size: Optional[int]
-) -> List[Any]: ...
-def _merge_caches(caches: Any) -> List[Any]: ...
-@dataclass
-class Batch:
- uids: List[int]
- y: mx.array
- logprobs: List[mx.array] | mx.array
- max_tokens: List[int]
- num_tokens: List[int]
- cache: 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: ...
- def extend(self, other: "Batch") -> None: ...
+ def __init__(
+ self,
+ model: nn.Module,
+ uids: List[int],
+ caches: List[List[Any]],
+ tokens: Optional[List[List[int]]] = ...,
+ prefill_step_size: int = ...,
+ samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
+ fallback_sampler: Optional[Callable[[mx.array], mx.array]] = ...,
+ logits_processors: Optional[
+ List[List[Callable[[mx.array, mx.array], mx.array]]]
+ ] = ...,
+ state_machines: Optional[List[SequenceStateMachine]] = ...,
+ max_tokens: Optional[List[int]] = ...,
+ ) -> None: ...
+ def __len__(self): # -> int:
+ ...
def extract_cache(self, idx: int) -> List[Any]: ...
+ def extend(self, batch): # -> None:
+ ...
+ def split(self, indices: List[int]): # -> Self:
+ ...
+ def filter(self, keep: List[int]): # -> None:
+ ...
+ def prompt(self, tokens: List[List[int]]): # -> None:
+ """
+ Process prompt tokens through the model.
+
+ Args:
+ tokens: List of token sequences to process.
+ """
+ ...
+
+ def generate(self, tokens: List[List[int]]): # -> GenerationBatch:
+ """
+ Transition from prompt processing to generation.
+
+ Args:
+ tokens: Final tokens for each sequence to start generation.
+
+ Returns:
+ A GenerationBatch ready for token generation.
+ """
+ ...
+
+ @classmethod
+ def empty(
+ cls,
+ model: nn.Module,
+ fallback_sampler: Callable[[mx.array], mx.array],
+ prefill_step_size: int = ...,
+ ): # -> Self:
+ ...
-class BatchGenerator:
- 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
+class GenerationBatch:
+ """
+ A batched token generator that manages multiple sequences in parallel.
+ This class handles the generation phase after prompt processing, managing
+ KV caches, sampling, and stop sequence detection for multiple sequences.
+ """
@dataclass
class Response:
uid: int
token: int
logprobs: mx.array
finish_reason: Optional[str]
- prompt_cache: Any
+ current_state: Optional[str]
+ match_sequence: Optional[List[int]]
+ prompt_cache: Optional[List[Any]]
+ all_tokens: Optional[List[int]]
+ ...
+ model: nn.Module
+ uids: List[int]
+ prompt_cache: List[Any]
+ tokens: List[List[int]]
+ samplers: Optional[List[Callable[[mx.array], mx.array]]]
+ fallback_sampler: Callable[[mx.array], mx.array]
+ logits_processors: Optional[List[List[Callable[[mx.array, mx.array], mx.array]]]]
+ state_machines: List[SequenceStateMachine]
+ max_tokens: List[int]
+ _current_tokens: Optional[mx.array]
+ _current_logprobs: List[mx.array]
+ _next_tokens: mx.array
+ _next_logprobs: List[mx.array]
+ _token_context: List[mx.array]
+ _num_tokens: List[int]
+
+ def __init__(
+ self,
+ model: nn.Module,
+ uids: List[int],
+ inputs: mx.array,
+ prompt_cache: List[Any],
+ tokens: List[List[int]],
+ samplers: Optional[List[Callable[[mx.array], mx.array]]],
+ fallback_sampler: Callable[[mx.array], mx.array],
+ logits_processors: Optional[
+ List[List[Callable[[mx.array, mx.array], mx.array]]]
+ ],
+ state_machines: List[SequenceStateMachine],
+ max_tokens: List[int],
+ ) -> None: ...
+ def __len__(self) -> int: ...
+ def extend(self, batch: GenerationBatch) -> None: ...
+ def extract_cache(self, idx: int) -> List[Any]: ...
+ def filter(self, keep: List[int]) -> None: ...
+ def _step(self) -> Tuple[List[int], List[mx.array]]: ...
+ def next(self) -> List[Response]:
+ """
+ Generate the next batch of tokens.
+
+ Returns:
+ List of Response objects for each sequence in the batch.
+ """
+ ...
+
+ @classmethod
+ def empty(
+ cls, model: nn.Module, fallback_sampler: Callable[[mx.array], mx.array]
+ ): # -> Self:
+ ...
+
+class BatchGenerator:
+ """
+ A batch generator implements continuous batching.
+
+ This class provides automatic management of prompt processing and generation
+ batches, handling the transition between the two.
+
+ It also allows for segmented prompt processing which guarantees that the
+ generator will stop at these boundaries when processing an input.
+ """
def __init__(
self,
- model: Any,
+ model: nn.Module,
max_tokens: int = ...,
- stop_tokens: Optional[set[int]] = ...,
+ stop_tokens: Optional[Sequence[Sequence[int]]] = ...,
sampler: Optional[Callable[[mx.array], mx.array]] = ...,
logits_processors: Optional[
List[Callable[[mx.array, mx.array], mx.array]]
@@ -312,41 +446,85 @@ class BatchGenerator:
completion_batch_size: int = ...,
prefill_batch_size: int = ...,
prefill_step_size: int = ...,
- prompt_progress_callback: Optional[
- Callable[[List[Tuple[int, int, int]]], None]
- ] = ...,
- max_kv_size: Optional[int] = ...,
) -> None: ...
def close(self) -> None: ...
+ def __del__(self): # -> None:
+ ...
+ @contextlib.contextmanager
+ def stats(self, stats=...): # -> Generator[Any | BatchStats, Any, None]:
+ ...
+ _unprocessed_sequences: deque[tuple[Any, ...]]
+ _prompt_batch: PromptProcessingBatch
+ _generation_batch: GenerationBatch
+ _currently_processing: list[Any]
+ _gen_tokens_counter: int
+ _steps_counter: int
+ def _next(
+ self,
+ ) -> tuple[
+ List[PromptProcessingBatch.Response], List[GenerationBatch.Response]
+ ]: ...
def insert(
self,
- prompts: Any,
- max_tokens: Union[List[int], int, None] = ...,
- caches: Any = ...,
- samplers: Optional[List[Any]] = ...,
- logits_processors: Optional[List[Any]] = ...,
+ prompts: List[List[int]],
+ max_tokens: Optional[List[int]] = ...,
+ caches: Optional[List[List[Any]]] = ...,
+ all_tokens: Optional[List[List[int]]] = ...,
+ samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
+ logits_processors: Optional[
+ List[List[Callable[[mx.array, mx.array], mx.array]]]
+ ] = ...,
+ state_machines: Optional[List[SequenceStateMachine]] = ...,
+ ) -> List[int]: ...
+ def insert_segments(
+ self,
+ segments: List[List[List[int]]],
+ max_tokens: Optional[List[int]] = ...,
+ caches: Optional[List[List[Any]]] = ...,
+ all_tokens: Optional[List[List[int]]] = ...,
+ samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
+ logits_processors: Optional[
+ List[List[Callable[[mx.array, mx.array], mx.array]]]
+ ] = ...,
+ state_machines: Optional[List[SequenceStateMachine]] = ...,
) -> List[int]: ...
+ def extract_cache(self, uids: List[int]) -> dict[int, Any]: ...
def remove(
self, uids: List[int], return_prompt_caches: bool = ...
- ) -> Optional[dict[int, List[Any]]]: ...
- def stats(self) -> BatchStats: ...
- def next(self) -> List[Response]: ...
- def _process_prompts(self, prompts: List[Any]) -> Batch: ...
- def _step(
+ ) -> dict[int, Any]: ...
+ @property
+ def prompt_cache_nbytes(self) -> int: ...
+ def next(
self,
- input_tokens: mx.array,
- prompt_cache: List[Any],
- samplers: Optional[List[Any]],
- logits_processors: Optional[List[Any]],
- tokens: List[mx.array],
- ) -> Tuple[mx.array, List[mx.array]]: ...
+ ) -> tuple[
+ List[PromptProcessingBatch.Response], List[GenerationBatch.Response]
+ ]: ...
+ def next_generated(self) -> List[GenerationBatch.Response]: ...
+
+@dataclass
+class BatchResponse:
+ """
+ A data object to hold a batch generation response.
+
+ Args:
+ texts: (List[str]): The generated text for each prompt.
+ stats (BatchStats): Statistics about the generation.
+ """
+
+ texts: List[str]
+ stats: BatchStats
+ caches: Optional[List[List[Any]]]
+ ...
def batch_generate(
model,
tokenizer,
- prompts: List[int],
+ prompts: List[List[int]],
+ prompt_caches: Optional[List[List[Any]]] = ...,
max_tokens: Union[int, List[int]] = ...,
verbose: bool = ...,
+ return_prompt_caches: bool = ...,
+ logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = ...,
**kwargs,
) -> BatchResponse:
"""
@@ -355,14 +533,22 @@ def batch_generate(
Args:
model (nn.Module): The language model.
tokenizer (PreTrainedTokenizer): The tokenizer.
- prompt (List[List[int]]): The input prompts.
+ prompts (List[List[int]]): The input prompts.
+ prompt_caches (List[List[Any]], optional): Pre-computed prompt-caches
+ for each input prompt. Note, unlike ``generate_step``, the caches
+ won't be updated in-place.
verbose (bool): If ``True``, print tokens and timing information.
Default: ``False``.
max_tokens (Union[int, List[int]): Maximum number of output tokens. This
can be per prompt if a list is provided.
+ return_prompt_caches (bool): Return the prompt caches in the batch
+ responses. Default: ``False``.
+ logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
+ A list of functions that take tokens and logits and return the processed logits. Default: ``None``.
kwargs: The remaining options get passed to :obj:`BatchGenerator`.
See :obj:`BatchGenerator` for more details.
"""
+ ...
def main(): # -> None:
...
diff --git a/.mlx_typings/mlx_lm/models/activations.pyi b/.mlx_typings/mlx_lm/models/activations.pyi
new file mode 100644
index 00000000..e7ca000e
--- /dev/null
+++ b/.mlx_typings/mlx_lm/models/activations.pyi
@@ -0,0 +1,19 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.core as mx
+import mlx.nn as nn
+from functools import partial
+
+@partial(mx.compile, shapeless=True)
+def swiglu(gate, x): ...
+@partial(mx.compile, shapeless=True)
+def xielu(x, alpha_p, alpha_n, beta, eps): # -> array:
+ ...
+
+class XieLU(nn.Module):
+ def __init__(
+ self, alpha_p_init=..., alpha_n_init=..., beta=..., eps=...
+ ) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
diff --git a/.mlx_typings/mlx_lm/tuner/dora.pyi b/.mlx_typings/mlx_lm/tuner/dora.pyi
new file mode 100644
index 00000000..1b03544a
--- /dev/null
+++ b/.mlx_typings/mlx_lm/tuner/dora.pyi
@@ -0,0 +1,50 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.nn as nn
+
+class DoRALinear(nn.Module):
+ @staticmethod
+ def from_base(
+ linear: nn.Linear, r: int = ..., dropout: float = ..., scale: float = ...
+ ): # -> DoRALinear:
+ ...
+ def fuse(self, dequantize: bool = ...): # -> QuantizedLinear | Linear:
+ ...
+ def __init__(
+ self,
+ input_dims: int,
+ output_dims: int,
+ r: int = ...,
+ dropout: float = ...,
+ scale: float = ...,
+ bias: bool = ...,
+ ) -> None: ...
+ def set_linear(self, linear): # -> None:
+ """
+ Set the self.linear layer and recompute self.m.
+ """
+ ...
+
+ def __call__(self, x): ...
+
+class DoRAEmbedding(nn.Module):
+ def from_base(
+ embedding: nn.Embedding, r: int = ..., dropout: float = ..., scale: float = ...
+ ): # -> DoRAEmbedding:
+ ...
+ def fuse(self, dequantize: bool = ...): # -> Embedding:
+ ...
+ def __init__(
+ self,
+ num_embeddings: int,
+ dims: int,
+ r: int = ...,
+ dropout: float = ...,
+ scale: float = ...,
+ ) -> None: ...
+ def set_embedding(self, embedding: nn.Module): # -> None:
+ ...
+ def __call__(self, x): ...
+ def as_linear(self, x): ...
diff --git a/.mlx_typings/mlx_lm/tuner/lora.pyi b/.mlx_typings/mlx_lm/tuner/lora.pyi
new file mode 100644
index 00000000..16d356a8
--- /dev/null
+++ b/.mlx_typings/mlx_lm/tuner/lora.pyi
@@ -0,0 +1,66 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.nn as nn
+
+class LoRALinear(nn.Module):
+ @staticmethod
+ def from_base(
+ linear: nn.Linear, r: int = ..., dropout: float = ..., scale: float = ...
+ ): # -> LoRALinear:
+ ...
+ def fuse(self, dequantize: bool = ...): # -> QuantizedLinear | Linear:
+ ...
+ def __init__(
+ self,
+ input_dims: int,
+ output_dims: int,
+ r: int = ...,
+ dropout: float = ...,
+ scale: float = ...,
+ bias: bool = ...,
+ ) -> None: ...
+ def __call__(self, x): # -> array:
+ ...
+
+class LoRASwitchLinear(nn.Module):
+ @staticmethod
+ def from_base(
+ linear: nn.Module, r: int = ..., dropout: float = ..., scale: float = ...
+ ): # -> LoRASwitchLinear:
+ ...
+ def fuse(self, dequantize: bool = ...): # -> QuantizedSwitchLinear | SwitchLinear:
+ ...
+ def __init__(
+ self,
+ input_dims: int,
+ output_dims: int,
+ num_experts: int,
+ r: int = ...,
+ dropout: float = ...,
+ scale: float = ...,
+ bias: bool = ...,
+ ) -> None: ...
+ def __call__(self, x, indices, sorted_indices=...): ...
+
+class LoRAEmbedding(nn.Module):
+ @staticmethod
+ def from_base(
+ embedding: nn.Embedding, r: int = ..., dropout: float = ..., scale: float = ...
+ ): # -> LoRAEmbedding:
+ ...
+ def fuse(self, dequantize: bool = ...): # -> QuantizedEmbedding | Embedding:
+ ...
+ def __init__(
+ self,
+ num_embeddings: int,
+ dims: int,
+ r: int = ...,
+ dropout: float = ...,
+ scale: float = ...,
+ ) -> None: ...
+ def __call__(self, x): # -> array:
+ ...
+ def as_linear(self, x): # -> array:
+ ...
diff --git a/.mlx_typings/mlx_lm/tuner/utils.pyi b/.mlx_typings/mlx_lm/tuner/utils.pyi
new file mode 100644
index 00000000..e579072b
--- /dev/null
+++ b/.mlx_typings/mlx_lm/tuner/utils.pyi
@@ -0,0 +1,57 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.nn as nn
+from typing import Dict
+
+def build_schedule(schedule_config: Dict): # -> Any:
+ """
+ Build a learning rate schedule from the given config.
+ """
+ ...
+
+def linear_to_lora_layers(
+ model: nn.Module, num_layers: int, config: Dict, use_dora: bool = ...
+): # -> None:
+ """
+ Convert some of the models linear layers to lora layers.
+
+ Args:
+ model (nn.Module): The neural network model.
+ num_layers (int): The number of blocks to convert to lora layers
+ starting from the last layer.
+ config (dict): More configuration parameters for LoRA, including the
+ rank, scale, and optional layer keys.
+ use_dora (bool): If True, uses DoRA instead of LoRA.
+ Default: ``False``
+ """
+ ...
+
+def load_adapters(model: nn.Module, adapter_path: str) -> nn.Module:
+ """
+ Load any fine-tuned adapters / layers.
+
+ Args:
+ model (nn.Module): The neural network model.
+ adapter_path (str): Path to the adapter configuration file.
+
+ Returns:
+ nn.Module: The updated model with LoRA layers applied.
+ """
+ ...
+
+def remove_lora_layers(model: nn.Module) -> nn.Module:
+ """
+ Remove the LoRA layers from the model.
+
+ Args:
+ model (nn.Module): The model with LoRA layers.
+
+ Returns:
+ nn.Module: The model without LoRA layers.
+ """
+ ...
+
+def print_trainable_parameters(model): # -> None:
+ ...
diff --git a/bench/exo_bench.py b/bench/exo_bench.py
index d48acdc0..3e16079d 100644
--- a/bench/exo_bench.py
+++ b/bench/exo_bench.py
@@ -238,6 +238,7 @@ def run_one_completion(
"messages": [{"role": "user", "content": content}],
"stream": False,
"max_tokens": tg,
+ "logprobs": False,
}
t0 = time.perf_counter()
@@ -564,6 +565,7 @@ def main() -> int:
"messages": [{"role": "user", "content": content}],
"stream": False,
"max_tokens": tg,
+ "logprobs": False,
}
barrier = threading.Barrier(concurrency)
batch_start = threading.Event()
@@ -602,6 +604,7 @@ def main() -> int:
"stats": stats,
}, _actual_pp
+ inf_t0 = time.monotonic()
with ThreadPoolExecutor(max_workers=concurrency) as pool:
futures = {
pool.submit(_run_concurrent, i): i
diff --git a/nix/darwin-build-fixes.patch b/nix/darwin-build-fixes.patch
index edce60d0..c32dbde8 100644
--- a/nix/darwin-build-fixes.patch
+++ b/nix/darwin-build-fixes.patch
@@ -59,7 +59,7 @@ index 262b0495..5c7446ad 100644
add_custom_command(
OUTPUT ${MLX_METAL_PATH}/mlx.metallib
-- COMMAND xcrun -sdk macosx metallib ${KERNEL_AIR} -o
+- COMMAND xcrun -sdk macosx metal ${KERNEL_AIR} -o
+ COMMAND metallib ${KERNEL_AIR} -o
${MLX_METAL_PATH}/mlx.metallib
DEPENDS ${KERNEL_AIR}
diff --git a/nix/mlx.nix b/nix/mlx.nix
index 39b0ac9f..93c4519b 100644
--- a/nix/mlx.nix
+++ b/nix/mlx.nix
@@ -49,7 +49,7 @@ let
owner = "rltakashige";
repo = "mlx-jaccl-fix-small-recv";
rev = uvLockMlxRev;
- hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ=";
+ hash = "sha256-M9x9QBYxwHv2z47qGZNJ4FgJyqLSIZ/3G1fEFQ421Lo=";
};
patches = [
diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py
index e06cab59..42e02eea 100644
--- a/src/exo/worker/engines/mlx/generator/batch_generate.py
+++ b/src/exo/worker/engines/mlx/generator/batch_generate.py
@@ -40,6 +40,10 @@ from exo.worker.engines.mlx.generator.generate import (
patch_embed_tokens,
prefill,
)
+from exo.worker.engines.mlx.patches.opt_batch_gen import (
+ set_needs_topk,
+ take_ready_topk,
+)
from exo.worker.engines.mlx.utils_mlx import (
detect_thinking_prompt_suffix,
fix_unmatched_think_end_tokens,
@@ -78,7 +82,6 @@ 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
@@ -101,17 +104,18 @@ class ExoBatchGenerator:
def __post_init__(self) -> None:
self._mlx_gen = MlxBatchGenerator(
model=self.model,
- stop_tokens=set(eos_ids_from_tokenizer(self.tokenizer)),
+ stop_tokens=[[t] for t in eos_ids_from_tokenizer(self.tokenizer)],
prefill_step_size=4096,
)
- self._mlx_gen._needs_topk = False # pyright: ignore[reportAttributeAccessIssue]
+ self._step_count = 0
@property
def has_work(self) -> bool:
return (
bool(self._active_tasks)
- or bool(self._mlx_gen.unprocessed_prompts)
- or self._mlx_gen.active_batch is not None
+ or bool(self._mlx_gen._unprocessed_sequences)
+ or len(self._mlx_gen._prompt_batch) > 0
+ or len(self._mlx_gen._generation_batch) > 0
)
def submit(
@@ -246,7 +250,7 @@ class ExoBatchGenerator:
max_tokens = task_params.max_output_tokens or MAX_TOKENS
uids = self._mlx_gen.insert(
- prompts=[last_tokens.tolist()],
+ prompts=[cast(list[int], last_tokens.tolist())],
max_tokens=[max_tokens],
caches=[list(cache)],
samplers=[sampler],
@@ -268,7 +272,6 @@ 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,
media_regions=media_regions,
in_thinking=detect_thinking_prompt_suffix(prompt, self.tokenizer),
)
@@ -279,13 +282,17 @@ class ExoBatchGenerator:
if not self.has_work:
return []
- self._mlx_gen._needs_topk = any( # pyright: ignore[reportAttributeAccessIssue]
- t.task_params.logprobs for t in self._active_tasks.values()
+ gb = self._mlx_gen._generation_batch
+ set_needs_topk(
+ gb,
+ any(t.task_params.logprobs for t in self._active_tasks.values()),
)
_step_tic = time.perf_counter()
- responses = self._mlx_gen.next()
+ _, responses = self._mlx_gen.next()
_next_elapsed = time.perf_counter() - _step_tic
+ topk = take_ready_topk(gb)
+
results: list[tuple[int, GenerationResponse]] = []
for response in responses:
@@ -351,17 +358,19 @@ class ExoBatchGenerator:
logprob: float | None = None
top_logprobs: list[TopLogprobItem] | None = None
if task_params.logprobs:
+ precomputed = topk.for_uid(response.uid)
+ precomputed_indices, precomputed_values, precomputed_selected = (
+ precomputed if precomputed is not None else (None, None, None)
+ )
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
- ),
+ precomputed_indices=precomputed_indices,
+ precomputed_values=precomputed_values,
+ precomputed_selected=precomputed_selected,
)
stats: GenerationStats | None = None
@@ -424,7 +433,8 @@ class ExoBatchGenerator:
_step_elapsed = time.perf_counter() - _step_tic
_overhead = _step_elapsed - _next_elapsed
- if self._mlx_gen._next_count % 64 == 0 and responses:
+ self._step_count += 1
+ if self._step_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)"
)
diff --git a/src/exo/worker/engines/mlx/patches/__init__.py b/src/exo/worker/engines/mlx/patches/__init__.py
index b05a045e..5e86c08c 100644
--- a/src/exo/worker/engines/mlx/patches/__init__.py
+++ b/src/exo/worker/engines/mlx/patches/__init__.py
@@ -10,5 +10,4 @@ def apply_mlx_patches() -> None:
return
_applied = True
patch_yarn_rope()
- # patch_gdn_softplus()
apply_batch_gen_patch()
diff --git a/src/exo/worker/engines/mlx/patches/opt_batch_gen.py b/src/exo/worker/engines/mlx/patches/opt_batch_gen.py
index 015d2e4b..7b07412b 100644
--- a/src/exo/worker/engines/mlx/patches/opt_batch_gen.py
+++ b/src/exo/worker/engines/mlx/patches/opt_batch_gen.py
@@ -1,173 +1,125 @@
-import time
-from typing import Any, cast
+from dataclasses import dataclass, field
+from typing import cast
import mlx.core as mx
-from mlx_lm.generate import BatchGenerator, generation_stream
+from mlx_lm.generate import GenerationBatch
_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
+@dataclass
+class BatchTopKLogprobs:
+ uids: list[int] = field(default_factory=list)
+ indices: mx.array | None = None
+ values: mx.array | None = None
+ selected: mx.array | None = None
+ _uid_to_row: dict[int, int] = field(default_factory=dict, init=False, repr=False)
+
+ def __post_init__(self) -> None:
+ self._uid_to_row = {uid: i for i, uid in enumerate(self.uids)}
+
+ def for_uid(self, uid: int) -> tuple[list[int], list[float], float] | None:
+ if self.indices is None or self.values is None or self.selected is None:
+ return None
+ row = self._uid_to_row.get(uid)
+ if row is None:
+ return None
+ return (
+ cast(list[int], self.indices[row].tolist()),
+ cast(list[float], self.values[row].tolist()),
+ float(self.selected[row].item()),
+ )
+
+
+@dataclass
+class _TopKBuffer:
+ needs_topk: bool = False
+ pending: BatchTopKLogprobs = field(default_factory=BatchTopKLogprobs)
+ ready: BatchTopKLogprobs = field(default_factory=BatchTopKLogprobs)
+
+
+def _get_buffer(batch: GenerationBatch) -> _TopKBuffer:
+ buf = getattr(batch, "_topk_buffer", None)
+ if buf is None:
+ buf = _TopKBuffer()
+ batch._topk_buffer = buf # pyright: ignore[reportAttributeAccessIssue]
+ return buf
+
+def set_needs_topk(batch: GenerationBatch, needed: bool) -> None:
+ _get_buffer(batch).needs_topk = needed
-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
+def take_ready_topk(batch: GenerationBatch) -> BatchTopKLogprobs:
+ return _get_buffer(batch).ready
- 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)
+def _patched_step(self: GenerationBatch) -> tuple[list[int], list[mx.array]]:
+ self._current_tokens = self._next_tokens
+ self._current_logprobs = self._next_logprobs
+ inputs = self._current_tokens
+
+ buf = _get_buffer(self)
+ buf.ready = buf.pending
+ buf.pending = BatchTopKLogprobs()
+
+ logits = self.model(inputs[:, None], cache=self.prompt_cache)
logits = logits[:, -1, :]
- if has_processors:
+ if self.logits_processors is not None and any(self.logits_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)
+ for e in range(len(self.uids)):
+ sample_logits = logits[e : e + 1]
+ for processor in self.logits_processors[e]:
+ sample_logits = processor(mx.array(self.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):
+ if self.samplers is not None and any(self.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)
+ for e in range(len(self.uids)):
+ sample_sampler = self.samplers[e] or self.fallback_sampler
+ all_samples.append(sample_sampler(logprobs[e : e + 1]))
+ sampled = 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:
+ sampled = self.fallback_sampler(logprobs)
+
+ self._next_tokens = sampled
+ self._next_logprobs = list(logprobs)
+
+ if buf.needs_topk:
+ batch_size = len(self.uids)
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]
+ pending_indices = mx.argpartition(-logprobs, k, axis=1)[:, :k]
+ pending_values = mx.take_along_axis(logprobs, pending_indices, axis=1)
+ sort_order = mx.argsort(-pending_values, axis=1)
+ pending_indices = mx.take_along_axis(pending_indices, sort_order, axis=1)
+ pending_values = mx.take_along_axis(pending_values, sort_order, axis=1)
+ pending_selected = logprobs[mx.arange(batch_size), sampled]
+ buf.pending = BatchTopKLogprobs(
+ uids=list(self.uids),
+ indices=pending_indices,
+ values=pending_values,
+ selected=pending_selected,
+ )
mx.async_eval(
- batch.y,
- *batch.logprobs,
- *batch.tokens,
- _pending_topk_idx,
- _pending_topk_val,
- _pending_selected_lps,
+ self._next_tokens,
+ *self._next_logprobs,
+ pending_indices,
+ pending_values,
+ pending_selected,
)
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)
+ mx.async_eval(self._next_tokens, *self._next_logprobs)
+
+ mx.eval(inputs, *self._current_logprobs)
+ token_list = cast(list[int], inputs.tolist())
+ for sti, ti in zip(self.tokens, token_list, strict=True):
+ sti.append(ti)
+ return token_list, self._current_logprobs
def apply_batch_gen_patch() -> None:
- BatchGenerator.next = _patched_public_next
+ GenerationBatch._step = _patched_step
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
index 36e643c2..24c3df8e 100644
--- a/src/exo/worker/engines/mlx/tests/test_extract_top_logprobs.py
+++ b/src/exo/worker/engines/mlx/tests/test_extract_top_logprobs.py
@@ -3,15 +3,9 @@ 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:
@@ -73,218 +67,3 @@ class TestExtractTopLogprobsFallback:
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/uv.lock b/uv.lock
index 133393ef..f40615ff 100644
--- a/uv.lock
+++ b/uv.lock
@@ -558,7 +558,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.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
+ { name = "mlx", version = "0.31.2.dev20260406+90dd61a5", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#90dd61a5f0837f9bbbab4fd3fbfedba1ca5d33e7" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx-vlm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1436,7 +1436,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.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
+ { name = "mlx", version = "0.31.2.dev20260406+90dd61a5", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#90dd61a5f0837f9bbbab4fd3fbfedba1ca5d33e7" }, 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'" },
@@ -1494,8 +1494,8 @@ cuda13 = [
[[package]]
name = "mlx"
-version = "0.31.2.dev20260324+e5e64331"
-source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }
+version = "0.31.2.dev20260406+90dd61a5"
+source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#90dd61a5f0837f9bbbab4fd3fbfedba1ca5d33e7" }
resolution-markers = [
"python_full_version >= '3.14' and sys_platform == 'darwin'",
"python_full_version < '3.14' and sys_platform == 'darwin'",
@@ -1528,10 +1528,10 @@ wheels = [
[[package]]
name = "mlx-lm"
version = "0.31.2"
-source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak#d36e9b661e55a5fc0f77fb6f17ea643aa2dc87aa" }
+source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak#b3540361c2fac915dc0f61ae0ce0de1583bfaa90" }
dependencies = [
{ name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
- { name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
+ { name = "mlx", version = "0.31.2.dev20260406+90dd61a5", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#90dd61a5f0837f9bbbab4fd3fbfedba1ca5d33e7" }, 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'" },
@@ -1548,7 +1548,7 @@ dependencies = [
{ name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "miniaudio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" },
- { name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
+ { name = "mlx", version = "0.31.2.dev20260406+90dd61a5", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#90dd61a5f0837f9bbbab4fd3fbfedba1ca5d33e7" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
← 24420eb1 Fix reasoning_tokens always reported as 0 for thinking model
·
back to Exo
·
Workspace tweaks (#1849) fd5b2328 →