diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py index f965367a..cf7f4762 100644 --- a/src/exo/worker/engines/mlx/generator/batch_generate.py +++ b/src/exo/worker/engines/mlx/generator/batch_generate.py @@ -18,6 +18,7 @@ from exo.shared.types.api import ( TopLogprobItem, Usage, ) +from exo.shared.types.common import ModelId from exo.shared.types.memory import Memory from exo.shared.types.mlx import MLXCacheType, Model from exo.shared.types.tasks import TaskId @@ -35,6 +36,7 @@ from exo.worker.engines.mlx.generator.generate import ( eos_ids_from_tokenizer, extract_top_logprobs, prefill, + warmup_inference, ) from exo.worker.engines.mlx.utils_mlx import fix_unmatched_think_end_tokens from exo.worker.runner.bootstrap import logger @@ -75,6 +77,7 @@ class ExoBatchGenerator: tokenizer: TokenizerWrapper group: mx.distributed.Group | None kv_prefix_cache: KVPrefixCache | None + model_id: ModelId _mlx_gen: MlxBatchGenerator = field(init=False) _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False) @@ -87,6 +90,9 @@ class ExoBatchGenerator: prefill_step_size=4096, ) + def warmup(self) -> int: + return warmup_inference(self.model, self.tokenizer, self.group, self.model_id) + @property def has_work(self) -> bool: return ( diff --git a/src/exo/worker/engines/vllm/prompt_format.py b/src/exo/worker/engines/vllm/prompt_format.py index 330a1060..272e5372 100644 --- a/src/exo/worker/engines/vllm/prompt_format.py +++ b/src/exo/worker/engines/vllm/prompt_format.py @@ -13,6 +13,8 @@ from exo.worker.engines.mlx.utils_mlx import ( def format_vllm_prompt( engine: LLMEngine, params: TextGenerationTaskParams ) -> tuple[list[int], str, int]: + # we should have our own wrapper + # (instead of abusing mlx's TokenizerWrapper, use tokenizers Tokenizer) tokenizer = TokenizerWrapper(engine.get_tokenizer()) prompt_text = apply_chat_template(tokenizer, params) token_ids: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False) # type: ignore[reportUnknownMemberType] diff --git a/src/exo/worker/engines/vllm/vllm_generator.py b/src/exo/worker/engines/vllm/vllm_generator.py index fda8e66b..812b4401 100644 --- a/src/exo/worker/engines/vllm/vllm_generator.py +++ b/src/exo/worker/engines/vllm/vllm_generator.py @@ -1,4 +1,5 @@ import gc +import math import os import re import sys @@ -23,7 +24,6 @@ from exo.shared.types.memory import Memory from exo.shared.types.tasks import TaskId from exo.shared.types.text_generation import TextGenerationTaskParams from exo.shared.types.worker.runner_response import GenerationResponse -from exo.worker.engines.vllm.kv_cache import TorchKVCache from exo.worker.engines.mlx.cache import KVPrefixCache from exo.worker.engines.mlx.utils_mlx import get_eos_token_ids_for_model from exo.worker.engines.vllm.growable_cache import ( @@ -31,6 +31,7 @@ from exo.worker.engines.vllm.growable_cache import ( patch_vllm, set_prefix_cache, ) +from exo.worker.engines.vllm.kv_cache import TorchKVCache from exo.worker.engines.vllm.prompt_format import ( format_vllm_prompt, make_vllm_sampling_params, @@ -279,7 +280,7 @@ def vllm_generate( ) -def warmup_vllm_engine(engine: LLMEngine) -> None: +def warmup_vllm_engine(engine: LLMEngine) -> int: tokenizer = engine.get_tokenizer() messages = [ { @@ -293,9 +294,17 @@ def warmup_vllm_engine(engine: LLMEngine) -> None: token_ids: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False) # type: ignore params = SamplingParams(max_tokens=50, detokenize=False) engine.add_request("warmup", {"prompt_token_ids": token_ids}, params) + t = time.monotonic() + tokens_generated = 0 while engine.has_unfinished_requests(): engine.step() - logger.info("vLLM warmup complete") + tokens_generated += 1 + elapsed = max(time.monotonic() - t, 0.001) + check_for_cancel_every = min(math.ceil(tokens_generated / elapsed), 100) + logger.info( + f"vLLM warmup complete, check_for_cancel_every={check_for_cancel_every}" + ) + return check_for_cancel_every @dataclass(eq=False) @@ -306,6 +315,9 @@ class VllmBatchEngine: _active: dict[TaskId, _EngineRequest] = field(default_factory=dict, init=False) + def warmup(self) -> int: + return warmup_vllm_engine(self.engine) + @property def has_work(self) -> bool: return bool(self._active) or self.engine.has_unfinished_requests() diff --git a/src/exo/worker/runner/llm_inference/batch_generator.py b/src/exo/worker/runner/llm_inference/batch_generator.py index 8eba561d..a37e637d 100644 --- a/src/exo/worker/runner/llm_inference/batch_generator.py +++ b/src/exo/worker/runner/llm_inference/batch_generator.py @@ -2,7 +2,7 @@ import itertools import time from abc import ABC, abstractmethod from collections import deque -from collections.abc import Generator, Iterable +from collections.abc import Callable, Generator, Iterable from dataclasses import dataclass, field from typing import TYPE_CHECKING @@ -13,7 +13,6 @@ from exo.shared.constants import EXO_MAX_CONCURRENT_REQUESTS from exo.shared.types.chunks import ErrorChunk, PrefillProgressChunk from exo.shared.types.common import ModelId from exo.shared.types.events import ChunkGenerated, Event -from exo.shared.types.mlx import Model from exo.shared.types.tasks import CANCEL_ALL_TASKS, TaskId, TextGeneration from exo.shared.types.text_generation import TextGenerationTaskParams from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse @@ -25,8 +24,6 @@ if TYPE_CHECKING: from exo.worker.engines.vllm.vllm_generator import VllmBatchEngine from exo.worker.engines.mlx.generator.generate import ( PrefillCancelled, - mlx_generate, - warmup_inference, ) from exo.worker.engines.mlx.utils_mlx import ( apply_chat_template, @@ -115,7 +112,6 @@ def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None: @dataclass(eq=False) class SequentialGenerator(InferenceGenerator): - model: Model tokenizer: TokenizerWrapper group: mx.distributed.Group | None kv_prefix_cache: KVPrefixCache | None @@ -124,6 +120,8 @@ class SequentialGenerator(InferenceGenerator): device_rank: int cancel_receiver: MpReceiver[TaskId] event_sender: MpSender[Event] + _generate_fn: Callable[..., Generator[GenerationResponse]] + _warmup_fn: Callable[[], int] check_for_cancel_every: int = 50 _cancelled_tasks: set[TaskId] = field(default_factory=set, init=False) @@ -145,12 +143,7 @@ class SequentialGenerator(InferenceGenerator): ) = field(default=None, init=False) def warmup(self) -> None: - self.check_for_cancel_every = warmup_inference( - model=self.model, - tokenizer=self.tokenizer, - group=self.group, - model_id=self.model_id, - ) + self.check_for_cancel_every = self._warmup_fn() def submit( self, @@ -234,7 +227,6 @@ class SequentialGenerator(InferenceGenerator): apply_chat_template(self.tokenizer, task.task_params), self.tool_parser, self.tokenizer, - type(self.model), self.model_id, task.task_params.tools, ) @@ -290,9 +282,7 @@ class SequentialGenerator(InferenceGenerator): self.agree_on_tasks() - return mlx_generate( - model=self.model, - tokenizer=self.tokenizer, + return self._generate_fn( task=task.task_params, prompt=prompt, kv_prefix_cache=self.kv_prefix_cache, @@ -303,12 +293,11 @@ class SequentialGenerator(InferenceGenerator): ) def close(self) -> None: - del self.model, self.tokenizer, self.group + del self.tokenizer, self.group @dataclass(eq=False) class BatchGenerator(InferenceGenerator): - model: Model | None tokenizer: TokenizerWrapper group: mx.distributed.Group | None kv_prefix_cache: KVPrefixCache | None @@ -336,13 +325,7 @@ class BatchGenerator(InferenceGenerator): ] = field(default_factory=dict, init=False) def warmup(self) -> None: - if self.model is not None: - self.check_for_cancel_every = warmup_inference( - model=self.model, - tokenizer=self.tokenizer, - group=self.group, - model_id=self.model_id, - ) + self.check_for_cancel_every = self._gen.warmup() def submit( self, @@ -403,7 +386,6 @@ class BatchGenerator(InferenceGenerator): apply_chat_template(self.tokenizer, task.task_params), self.tool_parser, self.tokenizer, - type(self.model), self.model_id, task.task_params.tools, ) @@ -525,6 +507,4 @@ class BatchGenerator(InferenceGenerator): def close(self) -> None: self._gen.close() - if self.model is not None: - del self.model del self.tokenizer, self.group diff --git a/src/exo/worker/runner/llm_inference/model_output_parsers.py b/src/exo/worker/runner/llm_inference/model_output_parsers.py index 4d8c50d4..7c6b677a 100644 --- a/src/exo/worker/runner/llm_inference/model_output_parsers.py +++ b/src/exo/worker/runner/llm_inference/model_output_parsers.py @@ -2,8 +2,6 @@ from collections.abc import Generator from functools import cache from typing import Any -from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model -from mlx_lm.models.gpt_oss import Model as GptOssModel from mlx_lm.tokenizer_utils import TokenizerWrapper from openai_harmony import ( HarmonyEncodingName, @@ -15,7 +13,6 @@ from openai_harmony import ( from exo.shared.types.api import ToolCallItem from exo.shared.types.common import ModelId -from exo.shared.types.mlx import Model from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse from exo.worker.engines.mlx.utils_mlx import ( detect_thinking_prompt_suffix, @@ -35,29 +32,28 @@ def apply_all_parsers( prompt: str, tool_parser: ToolParser | None, tokenizer: TokenizerWrapper, - model_type: type[Model] | type[None], model_id: ModelId, tools: list[dict[str, Any]] | None, ) -> Generator[GenerationResponse | ToolCallResponse | None]: - mlx_generator = receiver + gen = receiver if tokenizer.has_thinking: - mlx_generator = parse_thinking_models( - mlx_generator, + gen = parse_thinking_models( + gen, tokenizer.think_start, tokenizer.think_end, starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer), ) lower = model_id.normalize().lower() - if issubclass(model_type, GptOssModel) or "gpt-oss" in lower or "gpt_oss" in lower: - mlx_generator = parse_gpt_oss(mlx_generator) - elif issubclass(model_type, DeepseekV32Model) or "deepseek" in lower: - mlx_generator = parse_deepseek_v32(mlx_generator) + if "gpt-oss" in lower or "gpt_oss" in lower: + gen = parse_gpt_oss(gen) + elif "deepseek" in lower: + gen = parse_deepseek_v32(gen) elif tool_parser: - mlx_generator = parse_tool_calls(mlx_generator, tool_parser, tools) + gen = parse_tool_calls(gen, tool_parser, tools) - return mlx_generator + return gen _GPT_OSS_CHANNEL_TOKEN = 200005 diff --git a/src/exo/worker/runner/llm_inference/runner.py b/src/exo/worker/runner/llm_inference/runner.py index 431acde6..52b26f6b 100644 --- a/src/exo/worker/runner/llm_inference/runner.py +++ b/src/exo/worker/runner/llm_inference/runner.py @@ -447,11 +447,28 @@ class MlxBuilder(Builder): kv_prefix_cache = KVPrefixCache(self.group) + from functools import partial + + from exo.worker.engines.mlx.generator.generate import ( + mlx_generate, + warmup_inference, + ) + device_rank = 0 if self.group is None else self.group.rank() + generate_fn = partial( + mlx_generate, model=self.inference_model, tokenizer=self.tokenizer + ) + warmup_fn = partial( + warmup_inference, + model=self.inference_model, + tokenizer=self.tokenizer, + group=self.group, + model_id=self.model_id, + ) + if os.environ.get("EXO_NO_BATCH"): logger.info("using SequentialGenerator (batching disabled)") return SequentialGenerator( - model=self.inference_model, tokenizer=self.tokenizer, group=self.group, tool_parser=tool_parser, @@ -460,6 +477,8 @@ class MlxBuilder(Builder): device_rank=device_rank, cancel_receiver=self.cancel_receiver, event_sender=self.event_sender, + _generate_fn=generate_fn, + _warmup_fn=warmup_fn, ) from exo.worker.runner.llm_inference.batch_generator import ExoBatchGenerator @@ -469,9 +488,9 @@ class MlxBuilder(Builder): tokenizer=self.tokenizer, group=self.group, kv_prefix_cache=kv_prefix_cache, + model_id=self.model_id, ) return BatchGenerator( - model=self.inference_model, tokenizer=self.tokenizer, group=self.group, tool_parser=tool_parser, @@ -519,12 +538,8 @@ class VllmBuilder(Builder): ) def build(self) -> InferenceGenerator: - from exo.worker.engines.vllm.vllm_generator import ( - VllmBatchEngine, - warmup_vllm_engine, - ) + from exo.worker.engines.vllm.vllm_generator import VllmBatchEngine - warmup_vllm_engine(self._engine) gen = VllmBatchEngine( engine=self._engine, model_id=self.model_id, @@ -535,7 +550,6 @@ class VllmBuilder(Builder): logger.info(f"using BatchGenerator (vLLM, max_concurrent={max_concurrent})") return BatchGenerator( - model=None, tokenizer=tokenizer, group=None, tool_parser=self._tool_parser, diff --git a/src/exo/worker/tests/unittests/test_mlx/test_batch_vs_generate.py b/src/exo/worker/tests/unittests/test_mlx/test_batch_vs_generate.py deleted file mode 100644 index 836da6c2..00000000 --- a/src/exo/worker/tests/unittests/test_mlx/test_batch_vs_generate.py +++ /dev/null @@ -1,399 +0,0 @@ -import copy -import gc -import json -import shutil -import tempfile -from pathlib import Path -from typing import Any, cast - -import mlx.core as mx -import pytest -from mlx_lm.tokenizer_utils import TokenizerWrapper - -from exo.shared.types.common import ModelId -from exo.shared.types.mlx import MLXCacheType, Model -from exo.shared.types.tasks import TaskId -from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams -from exo.worker.engines.mlx.cache import CacheSnapshot, KVPrefixCache, cache_length -from exo.worker.engines.mlx.generator.batch_generate import ExoBatchGenerator -from exo.worker.engines.mlx.generator.generate import mlx_generate -from exo.worker.engines.mlx.utils_mlx import ( - apply_chat_template, - load_tokenizer_for_model_id, -) - -from .test_prefix_cache_architectures import ( - ARCHITECTURES, - ArchSpec, - _arch_available, # pyright: ignore[reportPrivateUsage] - _build_model, # pyright: ignore[reportPrivateUsage] - _copy_tokenizer, # pyright: ignore[reportPrivateUsage] - _find_snapshot, # pyright: ignore[reportPrivateUsage] - _reduce_config, # pyright: ignore[reportPrivateUsage] -) - - -def _make_task( - content: str = "Hello, what is 2+2?", - max_tokens: int = 10, - seed: int = 42, -) -> TextGenerationTaskParams: - return TextGenerationTaskParams( - model=ModelId("test"), - input=[InputMessage(role="user", content=content)], - max_output_tokens=max_tokens, - temperature=0.7, - seed=seed, - ) - - -# ── Helpers ──────────────────────────────────────────────────────────────── # - - -def _collect_mlx_generate( - model: Model, - tokenizer: TokenizerWrapper, - task: TextGenerationTaskParams, - kv_prefix_cache: KVPrefixCache | None, -) -> list[int]: - """Run mlx_generate and collect output token IDs.""" - prompt = apply_chat_template(tokenizer=tokenizer, task_params=task) - tokens: list[int] = [] - for resp in mlx_generate( - model=model, - tokenizer=tokenizer, - task=task, - prompt=prompt, - kv_prefix_cache=kv_prefix_cache, - group=None, - ): - tokens.append(resp.token) - if resp.finish_reason is not None: - break - return tokens - - -def _collect_batch_generate( - model: Model, - tokenizer: TokenizerWrapper, - task_params: TextGenerationTaskParams, - kv_prefix_cache: KVPrefixCache | None, -) -> list[int]: - """Run ExoBatchGenerator and collect raw output token IDs""" - exo_gen = ExoBatchGenerator( - model=model, - tokenizer=tokenizer, - group=None, - kv_prefix_cache=kv_prefix_cache, - ) - - prompt = apply_chat_template(tokenizer=tokenizer, task_params=task_params) - exo_gen.submit( - task_id=TaskId("test-single"), task_params=task_params, prompt=prompt - ) - - tokens: list[int] = [] - while exo_gen.has_work: - results = exo_gen.step() - for _uid, response in results: - tokens.append(response.token) - - exo_gen.close() - return tokens - - -def _assert_state_equal(sa: object, sb: object, label: str) -> None: - """Compare two state items, handling both plain arrays and tuples of arrays (CacheList).""" - if isinstance(sa, tuple): - assert isinstance(sb, tuple), f"{label}: type mismatch" - for k, (arr_a, arr_b) in enumerate( - zip( - cast(tuple[mx.array, ...], sa), - cast(tuple[mx.array, ...], sb), - strict=True, - ) - ): - a_f = mx.array(arr_a).astype(mx.float32) - b_f = mx.array(arr_b).astype(mx.float32) - if a_f.size == 0: - assert b_f.size == 0, f"{label}[{k}]: size mismatch" - continue - diff = float(mx.max(mx.abs(a_f - b_f)).item()) - assert diff == 0.0, f"{label}[{k}]: max diff {diff}" - else: - sa_f = mx.array(cast(mx.array, sa)).astype(mx.float32) - sb_f = mx.array(cast(mx.array, sb)).astype(mx.float32) - if sa_f.size == 0: - assert sb_f.size == 0, f"{label}: size mismatch" - return - diff = float(mx.max(mx.abs(sa_f - sb_f)).item()) - assert diff == 0.0, f"{label}: max diff {diff}" - - -def _compare_cache_arrays( - cache_a: MLXCacheType, - cache_b: MLXCacheType, - label: str = "", -) -> None: - """Assert two KV caches have identical array values.""" - assert len(cache_a) == len(cache_b), ( - f"{label}Cache layer count: {len(cache_a)} vs {len(cache_b)}" - ) - for i, (a, b) in enumerate(zip(cache_a, cache_b, strict=True)): - assert type(a) is type(b), ( - f"{label}Layer {i}: type {type(a).__name__} vs {type(b).__name__}" - ) - states_a = a.state - states_b = b.state - assert len(states_a) == len(states_b), ( - f"{label}Layer {i}: state count {len(states_a)} vs {len(states_b)}" - ) - for j, (sa, sb) in enumerate(zip(states_a, states_b, strict=True)): - if sa is None and sb is None: - continue - assert sa is not None and sb is not None, ( - f"{label}Layer {i}, state {j}: one is None" - ) - _assert_state_equal(sa, sb, f"{label}Layer {i}, state {j}") - - -def _safe_state(cache: object) -> list[object]: - """Safely access .state on a cache object. Returns [] if uninitialized.""" - # RotatingKVCache.state crashes when keys is None (uninitialized) - if getattr(cache, "keys", _SENTINEL) is None: - return [] - try: - return list(cache.state) # type: ignore[union-attr] - except (AttributeError, TypeError): - return [] - - -_SENTINEL = object() - - -def _compare_snapshots( - snaps_a: list[CacheSnapshot] | None, - snaps_b: list[CacheSnapshot] | None, - label: str = "", -) -> None: - """Assert two snapshot lists are identical.""" - if snaps_a is None: - assert snaps_b is None, f"{label}One side has snapshots, other doesn't" - return - assert snaps_b is not None, f"{label}One side has snapshots, other doesn't" - assert len(snaps_a) == len(snaps_b), ( - f"{label}Snapshot count: {len(snaps_a)} vs {len(snaps_b)}" - ) - for k, (sa, sb) in enumerate(zip(snaps_a, snaps_b, strict=True)): - assert sa.token_count == sb.token_count, ( - f"{label}Snapshot {k} token_count: {sa.token_count} vs {sb.token_count}" - ) - for layer_i, (s1, s2) in enumerate(zip(sa.states, sb.states, strict=True)): - if s1 is None and s2 is None: - continue - assert s1 is not None and s2 is not None, ( - f"{label}Snapshot {k}, layer {layer_i}: one state is None" - ) - state_a = _safe_state(s1) - state_b = _safe_state(s2) - if not state_a and not state_b: - continue - assert len(state_a) == len(state_b), ( - f"{label}Snapshot {k}, layer {layer_i}: state length mismatch" - ) - for st_j, (arr_a, arr_b) in enumerate(zip(state_a, state_b, strict=True)): - if arr_a is None and arr_b is None: - continue - assert arr_a is not None and arr_b is not None - _assert_state_equal( - arr_a, - arr_b, - f"{label}Snapshot {k}, layer {layer_i}, state {st_j}", - ) - - -# ── Test class ────────────────────────────────────────────────────────────── # - - -@pytest.mark.slow -class TestBatchVsGenerate: - """Verify BatchGenerator matches mlx_generate for output tokens and prefix cache.""" - - @pytest.fixture(autouse=True) - def _cleanup(self): - yield - mx.clear_cache() - gc.collect() - - @pytest.mark.parametrize( - "spec", - ARCHITECTURES, - ids=[a.name for a in ARCHITECTURES], - ) - def test_same_output_and_cache(self, spec: ArchSpec) -> None: - if not _arch_available(spec): - pytest.skip(f"Model {spec.hub_name} not cached locally") - - snapshot = _find_snapshot(spec.hub_name) - assert snapshot is not None - - tmpdir = Path(tempfile.mkdtemp(prefix=f"exo_batchtest_{spec.name}_")) - try: - # Build reduced config - with open(snapshot / "config.json") as f: - cfg = cast(dict[str, Any], json.load(f)) - reduced = _reduce_config(copy.deepcopy(cfg)) - (tmpdir / "config.json").write_text(json.dumps(reduced)) - - # Copy tokenizer - tok_src = snapshot - if spec.tokenizer_hub is not None: - alt = _find_snapshot(spec.tokenizer_hub) - if alt is not None: - tok_src = alt - _copy_tokenizer(tok_src, tmpdir) - - # Load tokenizer, build model with random weights - model_id = ModelId(f"mlx-community/{spec.hub_name}") - tokenizer = load_tokenizer_for_model_id(model_id, tmpdir) - mx.random.seed(0) - model = _build_model(spec.module, reduced) - - task = _make_task() - - # ── Run mlx_generate path ── - # Seed is set inside mlx_generate/ExoBatchGenerator.submit from task.seed - kv_mlx = KVPrefixCache(None) - mlx_tokens = _collect_mlx_generate(model, tokenizer, task, kv_mlx) - - # ── Run batch generator path ── - kv_batch = KVPrefixCache(None) - batch_tokens = _collect_batch_generate(model, tokenizer, task, kv_batch) - - # ── Compare output tokens ── - assert len(mlx_tokens) > 0, "mlx_generate produced no tokens" - assert len(batch_tokens) > 0, "BatchGenerator produced no tokens" - assert mlx_tokens == batch_tokens, ( - f"[{spec.name}] Token mismatch:\n" - f" mlx_generate: {mlx_tokens}\n" - f" BatchGenerator: {batch_tokens}" - ) - - # ── Compare prefix cache KV arrays ── - assert len(kv_mlx.caches) == 1, "mlx_generate didn't save to prefix cache" - assert len(kv_batch.caches) == 1, ( - "BatchGenerator didn't save to prefix cache" - ) - - mlx_cache = kv_mlx._get_mlx_cache(0) # pyright: ignore[reportPrivateUsage] - batch_cache = kv_batch._get_mlx_cache(0) # pyright: ignore[reportPrivateUsage] - - _compare_cache_arrays( - mlx_cache, - batch_cache, - label=f"[{spec.name}] ", - ) - - # ── Compare cache lengths ── - mlx_len = cache_length(mlx_cache) - batch_len = cache_length(batch_cache) - assert mlx_len == batch_len, ( - f"[{spec.name}] Cache length: mlx={mlx_len} vs batch={batch_len}" - ) - - # ── Compare snapshots ── - _compare_snapshots( - kv_mlx._snapshots[0], # pyright: ignore[reportPrivateUsage] - kv_batch._snapshots[0], # pyright: ignore[reportPrivateUsage] - label=f"[{spec.name}] ", - ) - - finally: - shutil.rmtree(tmpdir, ignore_errors=True) - - @pytest.mark.parametrize( - "spec", - ARCHITECTURES, - ids=[a.name for a in ARCHITECTURES], - ) - def test_concurrent_batch_completes(self, spec: ArchSpec) -> None: - """Two requests processed concurrently must both complete without - crashing and produce non-empty output. - - Note: batch decode logits are NOT bit-exact with sequential because - Metal's matmul kernel picks different reduction tiling for B=1 vs B=2 - when L=1 (decode step). This introduces sub-ULP float16 diffs in - gate_proj/down_proj/lm_head which swiglu amplifies by |up_values|. - With random weights these accumulate into argmax flips; with trained - weights the diffs are absorbed and output matches exactly (verified - with real Llama-3.2-1B-Instruct-4bit weights). - """ - if not _arch_available(spec): - pytest.skip(f"Model {spec.hub_name} not cached locally") - - snapshot = _find_snapshot(spec.hub_name) - assert snapshot is not None - - tmpdir = Path(tempfile.mkdtemp(prefix=f"exo_concurrent_{spec.name}_")) - try: - with open(snapshot / "config.json") as f: - cfg = cast(dict[str, Any], json.load(f)) - reduced = _reduce_config(copy.deepcopy(cfg)) - (tmpdir / "config.json").write_text(json.dumps(reduced)) - - tok_src = snapshot - if spec.tokenizer_hub is not None: - alt = _find_snapshot(spec.tokenizer_hub) - if alt is not None: - tok_src = alt - _copy_tokenizer(tok_src, tmpdir) - - model_id = ModelId(f"mlx-community/{spec.hub_name}") - tokenizer = load_tokenizer_for_model_id(model_id, tmpdir) - mx.random.seed(0) - model = _build_model(spec.module, reduced) - - # Two different prompts → different prompt lengths. - task_a = _make_task(content="Hello, what is 2+2?", seed=42) - task_a = task_a.model_copy(update={"temperature": 0.0}) - task_b = _make_task( - content="Write a short poem about the ocean and the sky.", - seed=99, - ) - task_b = task_b.model_copy(update={"temperature": 0.0}) - - # ── Concurrent: submit both to one ExoBatchGenerator ── - exo_gen = ExoBatchGenerator( - model=model, - tokenizer=tokenizer, - group=None, - kv_prefix_cache=None, - ) - - prompt_a = apply_chat_template(tokenizer=tokenizer, task_params=task_a) - prompt_b = apply_chat_template(tokenizer=tokenizer, task_params=task_b) - tid_a = exo_gen.submit( - task_id=TaskId("batch-a"), task_params=task_a, prompt=prompt_a - ) - tid_b = exo_gen.submit( - task_id=TaskId("batch-b"), task_params=task_b, prompt=prompt_b - ) - - batch_tokens: dict[str, list[int]] = {tid_a: [], tid_b: []} - finished: set[str] = set() - while exo_gen.has_work: - results = exo_gen.step() - for tid, response in results: - batch_tokens[tid].append(response.token) - if response.finish_reason is not None: - finished.add(tid) - - exo_gen.close() - - # ── Verify both completed ── - assert len(batch_tokens[tid_a]) > 0, "No tokens for task A" - assert len(batch_tokens[tid_b]) > 0, "No tokens for task B" - assert tid_a in finished, "Task A never finished" - assert tid_b in finished, "Task B never finished" - finally: - shutil.rmtree(tmpdir, ignore_errors=True) diff --git a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py index 7a883005..3600290e 100644 --- a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py +++ b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py @@ -116,7 +116,6 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch): # initialize_mlx returns a mock group monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MockGroup())) monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, MockTokenizer))) - monkeypatch.setattr(mlx_batch_generator, "warmup_inference", make_nothin(1)) monkeypatch.setattr(mlx_batch_generator, "_check_for_debug_prompts", nothin) monkeypatch.setattr(mlx_batch_generator, "mx_any", make_nothin(False)) @@ -141,6 +140,9 @@ class FakeExoBatchGenerator: def __init__(self, *_args: object, **_kwargs: object) -> None: self._pending: dict[str, GenerationResponse] = {} + def warmup(self) -> int: + return 50 + @property def has_work(self) -> bool: return bool(self._pending)