Address comments 4 - defer to the warmup into the exo batch generator and vllm batch engine and don't store model on the generators.
This commit is contained in:
Ryuichi Leo Takashige
@@ -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 (
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user