man i gotta commit

python/exo_core/src/exo_core/engine.py

@@ -5,10 +5,10 @@ from typing import Self
 from exo_core.types.tasks import TaskId
 
 
-class Cancelled: ...
+class Cancelled: pass
 
 
-class Finished: ...
+class Finished: pass
 
 
 CANCEL_ALL_TASKS = TaskId("CANCEL_TALL_TASKS")

python/exo_core/src/exo_core/types/chunks.py

@@ -1,7 +1,9 @@
 from collections.abc import Generator
 from typing import Any, Literal
 
-from exo.api.types import (
+from exo_core.model_cards import ModelId
+from exo_core.models import TaggedModel
+from exo_core.types.runner_response import (
     FinishReason,
     GenerationStats,
     ImageGenerationStats,
@@ -9,8 +11,6 @@ from exo.api.types import (
     TopLogprobItem,
     Usage,
 )
-from exo_core.model_cards import ModelId
-from exo_core.models import TaggedModel
 
 from .common import CommandId
 

python/exo_core/src/exo_core/types/image_generation.py

@@ -0,0 +1,108 @@
+from collections.abc import Generator
+from typing import Annotated, Any, Literal, get_args, override
+
+from pydantic import BaseModel, Field, field_validator
+
+ImageSize = Literal[
+    "auto",
+    "512x512",
+    "768x768",
+    "1024x768",
+    "768x1024",
+    "1024x1024",
+    "1024x1536",
+    "1536x1024",
+]
+
+
+def normalize_image_size(v: object) -> ImageSize:
+    """Shared validator for ImageSize fields: maps None → "auto" and rejects invalid values."""
+    if v is None:
+        return "auto"
+    if v not in get_args(ImageSize):
+        raise ValueError(f"Invalid size: {v!r}. Must be one of {get_args(ImageSize)}")
+    return v  # pyright: ignore[reportReturnType]
+
+
+# can we tighten these to CamelCaseModels?
+class AdvancedImageParams(BaseModel):
+    seed: Annotated[int, Field(ge=0)] | None = None
+    num_inference_steps: Annotated[int, Field(ge=1, le=100)] | None = None
+    guidance: Annotated[float, Field(ge=1.0, le=20.0)] | None = None
+    negative_prompt: str | None = None
+    num_sync_steps: Annotated[int, Field(ge=1, le=100)] | None = None
+
+
+class ImageGenerationTaskParams(BaseModel):
+    prompt: str
+    background: str | None = None
+    model: str
+    moderation: str | None = None
+    n: int | None = 1
+    output_compression: int | None = None
+    output_format: Literal["png", "jpeg", "webp"] = "png"
+    partial_images: int | None = 0
+    quality: Literal["high", "medium", "low"] | None = "medium"
+    response_format: Literal["url", "b64_json"] | None = "b64_json"
+    size: ImageSize = "auto"
+    stream: bool | None = False
+    style: str | None = "vivid"
+    user: str | None = None
+    advanced_params: AdvancedImageParams | None = None
+    # Internal flag for benchmark mode - set by API, preserved through serialization
+    bench: bool = False
+
+    @field_validator("size", mode="before")
+    @classmethod
+    def normalize_size(cls, v: object) -> ImageSize:
+        return normalize_image_size(v)
+
+
+class BenchImageGenerationTaskParams(ImageGenerationTaskParams):
+    bench: bool = True
+
+
+class ImageEditsTaskParams(BaseModel):
+    """Internal task params for image-editing requests."""
+
+    image_data: str = ""  # Base64-encoded image (empty when using chunked transfer)
+    total_input_chunks: int = 0
+    prompt: str
+    model: str
+    n: int | None = 1
+    quality: Literal["high", "medium", "low"] | None = "medium"
+    output_format: Literal["png", "jpeg", "webp"] = "png"
+    response_format: Literal["url", "b64_json"] | None = "b64_json"
+    size: ImageSize = "auto"
+    image_strength: float | None = 0.7
+    stream: bool = False
+    partial_images: int | None = 0
+    advanced_params: AdvancedImageParams | None = None
+    bench: bool = False
+
+    @field_validator("size", mode="before")
+    @classmethod
+    def normalize_size(cls, v: object) -> ImageSize:
+        return normalize_image_size(v)
+
+    @override
+    def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
+        for name, value in super().__repr_args__():  # pyright: ignore[reportAny]
+            if name == "image_data":
+                yield name, f"<{len(self.image_data)} chars>"
+            elif name is not None:
+                yield name, value
+
+
+class ImageData(BaseModel):
+    b64_json: str | None = None
+    url: str | None = None
+    revised_prompt: str | None = None
+
+    @override
+    def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
+        for name, value in super().__repr_args__():  # pyright: ignore[reportAny]
+            if name == "b64_json" and self.b64_json is not None:
+                yield name, f"<{len(self.b64_json)} chars>"
+            elif name is not None:
+                yield name, value

python/exo_core/src/exo_core/types/runner_response.py

@@ -1,15 +1,68 @@
 from collections.abc import Generator
-from typing import Any, Literal
+from typing import Any, Literal, override
+from uuid import uuid4
 
-from exo.api.types import (
-    FinishReason,
-    GenerationStats,
-    ImageGenerationStats,
-    ToolCallItem,
-    TopLogprobItem,
-    Usage,
-)
-from exo_core.models import TaggedModel
+from pydantic import Field
+
+from exo_core.models import CamelCaseModel, TaggedModel
+from exo_core.utils.memory import Memory
+
+FinishReason = Literal[
+    "stop", "length", "tool_calls", "content_filter", "function_call", "error"
+]
+
+
+class ToolCallItem(CamelCaseModel):
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    name: str
+    arguments: str
+
+
+class TopLogprobItem(CamelCaseModel):
+    token: str
+    logprob: float
+    bytes: list[int] | None = None
+
+
+class GenerationStats(CamelCaseModel):
+    prompt_tps: float
+    generation_tps: float
+    prompt_tokens: int
+    generation_tokens: int
+    peak_memory_usage: Memory
+
+
+class ImageGenerationStats(CamelCaseModel):
+    seconds_per_step: float
+    total_generation_time: float
+
+    num_inference_steps: int
+    num_images: int
+
+    image_width: int
+    image_height: int
+
+    peak_memory_usage: Memory
+
+
+class PromptTokensDetails(CamelCaseModel):
+    cached_tokens: int = 0
+    audio_tokens: int = 0
+
+
+class CompletionTokensDetails(CamelCaseModel):
+    reasoning_tokens: int = 0
+    audio_tokens: int = 0
+    accepted_prediction_tokens: int = 0
+    rejected_prediction_tokens: int = 0
+
+
+class Usage(CamelCaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    prompt_tokens_details: PromptTokensDetails
+    completion_tokens_details: CompletionTokensDetails
 
 
 class BaseRunnerResponse(TaggedModel):
@@ -37,6 +90,7 @@ class ImageGenerationResponse(BaseRunnerResponse):
     stats: ImageGenerationStats | None = None
     image_index: int = 0
 
+    @override
     def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
         for name, value in super().__repr_args__():  # pyright: ignore[reportAny]
             if name == "image_data":
@@ -52,6 +106,7 @@ class PartialImageResponse(BaseRunnerResponse):
     total_partials: int
     image_index: int = 0
 
+    @override
     def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
         for name, value in super().__repr_args__():  # pyright: ignore[reportAny]
             if name == "image_data":

python/exo_core/src/exo_core/types/tasks.py

@@ -2,16 +2,17 @@ from enum import Enum
 
 from pydantic import Field
 
-from exo.api.types import (
+from exo_core.models import TaggedModel
+
+from .common import CommandId, Id
+from .image_generation import (
     ImageEditsTaskParams,
     ImageGenerationTaskParams,
 )
-from exo_core.models import TaggedModel
-from exo_core.types.common import CommandId, Id
-from exo_core.types.instances import BoundInstance, InstanceId
-from exo_core.types.runners import RunnerId
-from exo_core.types.shards import ShardMetadata
-from exo_core.types.text_generation import TextGenerationTaskParams
+from .instances import BoundInstance, InstanceId
+from .runners import RunnerId
+from .shards import ShardMetadata
+from .text_generation import TextGenerationTaskParams
 
 
 class TaskId(Id):

python/exo_core/src/exo_core/utils/channels.py

@@ -0,0 +1,309 @@
+import contextlib
+import multiprocessing as mp
+from dataclasses import dataclass, field
+from math import inf
+from multiprocessing.synchronize import Event
+from queue import Empty, Full
+from types import TracebackType
+from typing import Any, Self
+
+from anyio import (
+    CapacityLimiter,
+    ClosedResourceError,
+    EndOfStream,
+    WouldBlock,
+    to_thread,
+)
+from anyio.streams.memory import (
+    MemoryObjectReceiveStream as AnyioReceiver,
+)
+from anyio.streams.memory import (
+    MemoryObjectSendStream as AnyioSender,
+)
+from anyio.streams.memory import (
+    MemoryObjectStreamState as AnyioState,
+)
+
+
+class Sender[T](AnyioSender[T]):
+    def clone(self) -> "Sender[T]":
+        if self._closed:
+            raise ClosedResourceError
+        return Sender(_state=self._state)
+
+    def clone_receiver(self) -> "Receiver[T]":
+        """Constructs a Receiver using a Senders shared state - similar to calling Receiver.clone() without needing the receiver"""
+        if self._closed:
+            raise ClosedResourceError
+        return Receiver(_state=self._state)
+
+
+class Receiver[T](AnyioReceiver[T]):
+    def clone(self) -> "Receiver[T]":
+        if self._closed:
+            raise ClosedResourceError
+        return Receiver(_state=self._state)
+
+    def clone_sender(self) -> Sender[T]:
+        """Constructs a Sender using a Receivers shared state - similar to calling Sender.clone() without needing the sender"""
+        if self._closed:
+            raise ClosedResourceError
+        return Sender(_state=self._state)
+
+    def collect(self) -> list[T]:
+        """Collect all currently available items from this receiver"""
+        out: list[T] = []
+        while True:
+            try:
+                item = self.receive_nowait()
+                out.append(item)
+            except WouldBlock:
+                break
+        return out
+
+    async def receive_at_least(self, n: int) -> list[T]:
+        out: list[T] = []
+        out.append(await self.receive())
+        out.extend(self.collect())
+        while len(out) < n:
+            out.append(await self.receive())
+            out.extend(self.collect())
+        return out
+
+    def __enter__(self) -> Self:
+        return self
+
+
+class _MpEndOfStream:
+    pass
+
+
+class MpState[T]:
+    def __init__(self, max_buffer_size: float):
+        if max_buffer_size == inf:
+            max_buffer_size = 0
+        assert isinstance(max_buffer_size, int), (
+            "State should only ever be constructed with an integer or math.inf size."
+        )
+
+        self.max_buffer_size: float = max_buffer_size
+        self.buffer: mp.Queue[T | _MpEndOfStream] = mp.Queue(max_buffer_size)
+        self.closed: Event = mp.Event()
+
+    def __getstate__(self):
+        d = self.__dict__.copy()
+        d.pop("__orig_class__", None)
+        return d
+
+
+@dataclass(eq=False)
+class MpSender[T]:
+    """
+    An interprocess channel, mimicing the Anyio structure.
+    It should be noted that none of the clone methods are implemented for simplicity, for now.
+    """
+
+    _state: MpState[T] = field()
+
+    def send_nowait(self, item: T) -> None:
+        if self._state.closed.is_set():
+            raise ClosedResourceError
+        try:
+            self._state.buffer.put(item, block=False)
+        except Full:
+            raise WouldBlock from None
+        except ValueError as e:
+            print("Unreachable code path - let me know!")
+            raise ClosedResourceError from e
+
+    def send(self, item: T) -> None:
+        if self._state.closed.is_set():
+            raise ClosedResourceError
+        try:
+            self.send_nowait(item)
+        except WouldBlock:
+            # put anyway, blocking
+            self._state.buffer.put(item, block=True)
+
+    async def send_async(self, item: T) -> None:
+        await to_thread.run_sync(
+            self.send, item, limiter=CapacityLimiter(1), abandon_on_cancel=True
+        )
+
+    def close(self) -> None:
+        if not self._state.closed.is_set():
+            self._state.closed.set()
+        with contextlib.suppress(Exception):
+            self._state.buffer.put_nowait(_MpEndOfStream())
+        self._state.buffer.close()
+
+    # == unique to Mp channels ==
+    def join(self) -> None:
+        """Ensure any queued messages are resolved before continuing"""
+        assert self._state.closed.is_set(), (
+            "Mp channels must be closed before being joined"
+        )
+        self._state.buffer.join_thread()
+
+    # == context manager support ==
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        self.close()
+
+    def __getstate__(self) -> dict[str, Any]:
+        d = self.__dict__.copy()
+        d.pop("__orig_class__", None)
+        return d
+
+
+@dataclass(eq=False)
+class MpReceiver[T]:
+    """
+    An interprocess channel, mimicing the Anyio structure.
+    It should be noted that none of the clone methods are implemented for simplicity, for now.
+    """
+
+    _state: MpState[T] = field()
+
+    def receive_nowait(self) -> T:
+        if self._state.closed.is_set():
+            raise ClosedResourceError
+
+        try:
+            item = self._state.buffer.get(block=False)
+            if isinstance(item, _MpEndOfStream):
+                self.close()
+                raise EndOfStream
+            return item
+        except Empty:
+            raise WouldBlock from None
+        except ValueError as e:
+            print("Unreachable code path - let me know!")
+            raise ClosedResourceError from e
+
+    def receive(self) -> T:
+        try:
+            return self.receive_nowait()
+        except WouldBlock:
+            try:
+                item = self._state.buffer.get()
+            except (TypeError, OSError):
+                # Queue pipe can get closed while we are blocked on get().
+                # The underlying connection._handle becomes None, causing
+                # TypeError in read(handle, remaining).
+                raise ClosedResourceError from None
+            if isinstance(item, _MpEndOfStream):
+                self.close()
+                raise EndOfStream from None
+            return item
+
+    async def receive_async(self) -> T:
+        return await to_thread.run_sync(
+            self.receive, limiter=CapacityLimiter(1), abandon_on_cancel=True
+        )
+
+    def close(self) -> None:
+        if not self._state.closed.is_set():
+            self._state.closed.set()
+        with contextlib.suppress(Exception):
+            self._state.buffer.put_nowait(_MpEndOfStream())
+        self._state.buffer.close()
+
+    # == unique to Mp channels ==
+    def join(self) -> None:
+        """Block until all enqueued messages are drained off our side of the buffer"""
+        assert self._state.closed.is_set(), (
+            "Mp channels must be closed before being joined"
+        )
+        self._state.buffer.join_thread()
+
+    # == iterator support ==
+    def __iter__(self) -> Self:
+        return self
+
+    def __next__(self) -> T:
+        try:
+            return self.receive()
+        except EndOfStream:
+            raise StopIteration from None
+
+    # == async iterator support ==
+    def __aiter__(self) -> Self:
+        return self
+
+    async def __anext__(self) -> T:
+        try:
+            return await self.receive_async()
+        except EndOfStream:
+            raise StopAsyncIteration from None
+
+    # == context manager support ==
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        self.close()
+
+    def collect(self) -> list[T]:
+        """Collect all currently available items from this receiver"""
+        out: list[T] = []
+        while True:
+            try:
+                item = self.receive_nowait()
+                out.append(item)
+            except WouldBlock:
+                break
+        return out
+
+    def receive_at_least(self, n: int) -> list[T]:
+        out: list[T] = []
+        out.append(self.receive())
+        out.extend(self.collect())
+        while len(out) < n:
+            out.append(self.receive())
+            out.extend(self.collect())
+        return out
+
+    def __getstate__(self):
+        d = self.__dict__.copy()
+        d.pop("__orig_class__", None)
+        return d
+
+
+class channel[T]:  # noqa: N801
+    """Create a pair of asynchronous channels for communicating within the same process"""
+
+    def __new__(cls, max_buffer_size: float = inf) -> tuple[Sender[T], Receiver[T]]:
+        if max_buffer_size != inf and not isinstance(max_buffer_size, int):
+            raise ValueError("max_buffer_size must be either an integer or math.inf")
+        state = AnyioState[T](max_buffer_size)
+        return Sender(_state=state), Receiver(_state=state)
+
+
+class mp_channel[T]:  # noqa: N801
+    """Create a pair of synchronous channels for interprocess communication"""
+
+    # max buffer size uses math.inf to represent an unbounded queue, and 0 to represent a yet unimplemented "unbuffered" queue.
+    def __new__(cls, max_buffer_size: float = inf) -> tuple[MpSender[T], MpReceiver[T]]:
+        if (
+            max_buffer_size == 0
+            or max_buffer_size != inf
+            and not isinstance(max_buffer_size, int)
+        ):
+            raise ValueError(
+                "max_buffer_size must be either an integer or math.inf. 0-sized buffers are not supported by multiprocessing"
+            )
+        state = MpState[T](max_buffer_size)
+        return MpSender(_state=state), MpReceiver(_state=state)

python/exo_core/src/exo_core/utils/tool_parsers.py

@@ -0,0 +1,242 @@
+import json
+import math
+from dataclasses import dataclass
+from typing import Any, Callable, Final
+
+from exo_core.types.runner_response import ToolCallItem
+
+
+@dataclass
+class ToolParser:
+    start_parsing: str
+    end_parsing: str
+    _inner_parser: Callable[[str], list[ToolCallItem] | None]
+
+    def parse(
+        self, text: str, tools: list[dict[str, Any]] | None
+    ) -> list[ToolCallItem] | None:
+        parsed = self._inner_parser(text)
+        if parsed is None:
+            return None
+        if tools is not None:
+            parsed = _coerce_tool_calls_to_schema(parsed, tools)
+        return parsed
+
+
+def _json_type_matches(value: Any, expected_type: str) -> bool:  # pyright: ignore[reportAny]
+    if expected_type == "object":
+        return isinstance(value, dict)
+    if expected_type == "array":
+        return isinstance(value, list)
+    if expected_type == "string":
+        return isinstance(value, str)
+    if expected_type == "integer":
+        return isinstance(value, int) and not isinstance(value, bool)
+    if expected_type == "number":
+        return (isinstance(value, int) and not isinstance(value, bool)) or isinstance(
+            value, float
+        )
+    if expected_type == "boolean":
+        return isinstance(value, bool)
+    if expected_type == "null":
+        return value is None
+    return False
+
+
+def _coerce_tool_arg_with_schema(value: Any, schema: dict[str, Any]) -> Any:  # pyright: ignore[reportAny]
+    schema_type = schema.get("type")
+
+    if isinstance(schema_type, list):
+        for candidate in schema_type:  # pyright: ignore[reportUnknownVariableType]
+            if not isinstance(candidate, str):
+                continue
+            if candidate == "null" and value is None:
+                return None
+            candidate_schema = {**schema, "type": candidate}
+            coerced = _coerce_tool_arg_with_schema(value, candidate_schema)  # pyright: ignore[reportAny]
+            if _json_type_matches(coerced, candidate):
+                return coerced  # pyright: ignore[reportAny]
+        return value  # pyright: ignore[reportAny]
+
+    if not isinstance(schema_type, str):
+        return value  # pyright: ignore[reportAny]
+
+    if schema_type == "object":
+        parsed = value  # pyright: ignore[reportAny]
+        if isinstance(parsed, str):
+            try:
+                parsed = json.loads(parsed)  # pyright: ignore[reportAny]
+            except Exception:
+                return value  # pyright: ignore[reportAny]
+        if not isinstance(parsed, dict):
+            return value  # pyright: ignore[reportAny]
+        properties = schema.get("properties")
+        if not isinstance(properties, dict):
+            return parsed  # pyright: ignore[reportUnknownVariableType]
+        return {
+            key: (
+                _coerce_tool_arg_with_schema(prop_value, prop_schema)  # pyright: ignore[reportUnknownArgumentType]
+                if isinstance(prop_schema, dict)
+                else prop_value
+            )
+            for key, prop_value in parsed.items()  # pyright: ignore[reportUnknownVariableType]
+            for prop_schema in [properties.get(key)]  # type: ignore
+        }
+
+    if schema_type == "array":
+        parsed = value  # pyright: ignore[reportAny]
+        if isinstance(parsed, str):
+            try:
+                parsed = json.loads(parsed)  # pyright: ignore[reportAny]
+            except Exception:
+                return value  # pyright: ignore[reportAny]
+        if not isinstance(parsed, list):
+            return value  # pyright: ignore[reportAny]
+        item_schema = schema.get("items")
+        if not isinstance(item_schema, dict):
+            return parsed  # pyright: ignore[reportUnknownVariableType]
+        return [_coerce_tool_arg_with_schema(item, item_schema) for item in parsed]  # type: ignore
+
+    if schema_type == "integer":
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, int):
+            return value
+        if isinstance(value, float) and value.is_integer():
+            return int(value)
+        if isinstance(value, str):
+            try:
+                return int(value.strip())
+            except ValueError:
+                return value
+        return value
+
+    if schema_type == "number":
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, (int, float)):
+            return value
+        if isinstance(value, str):
+            try:
+                num = float(value.strip())
+                if math.isfinite(num):
+                    return num
+            except ValueError:
+                return value
+        return value
+
+    if schema_type == "boolean":
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, str):
+            lowered = value.strip().lower()
+            if lowered == "true":
+                return True
+            if lowered == "false":
+                return False
+        return value
+
+    return value  # pyright: ignore[reportAny]
+
+
+def _coerce_tool_calls_to_schema(
+    tool_calls: list[ToolCallItem], tools: list[dict[str, Any]]
+) -> list[ToolCallItem]:
+    schema_by_name: dict[str, dict[str, Any]] = {}
+    for tool in tools:
+        function = tool.get("function")
+        if not isinstance(function, dict):
+            continue
+        name = function.get("name")  # type: ignore
+        parameters = function.get("parameters")  # type: ignore
+        if isinstance(name, str) and isinstance(parameters, dict):
+            schema_by_name[name] = parameters
+
+    if not schema_by_name:
+        return tool_calls
+
+    coerced_calls: list[ToolCallItem] = []
+    for tool_call in tool_calls:
+        schema = schema_by_name.get(tool_call.name)
+        if schema is None:
+            coerced_calls.append(tool_call)
+            continue
+
+        try:
+            parsed_args = json.loads(tool_call.arguments)  # pyright: ignore[reportAny]
+        except Exception:
+            coerced_calls.append(tool_call)
+            continue
+
+        if not isinstance(parsed_args, dict):
+            coerced_calls.append(tool_call)
+            continue
+
+        coerced_args = _coerce_tool_arg_with_schema(parsed_args, schema)  # pyright: ignore[reportAny]
+        if not isinstance(coerced_args, dict):
+            coerced_calls.append(tool_call)
+            continue
+
+        coerced_calls.append(
+            tool_call.model_copy(update={"arguments": json.dumps(coerced_args)})
+        )
+    return coerced_calls
+
+
+def make_mlx_parser(
+    tool_call_start: str,
+    tool_call_end: str,
+    tool_parser: Callable[[str], dict[str, Any] | list[dict[str, Any]]],
+) -> ToolParser:
+    def parse_tool_calls(text: str) -> list[ToolCallItem] | None:
+        try:
+            text = text.removeprefix(tool_call_start)
+            text = text.removesuffix(tool_call_end)
+            parsed = tool_parser(text)
+            if isinstance(parsed, list):
+                return [ToolCallItem.model_validate(_flatten(p)) for p in parsed]
+            else:
+                return [ToolCallItem.model_validate(_flatten(parsed))]
+
+        except Exception:
+            return None
+
+    return ToolParser(
+        start_parsing=tool_call_start,
+        end_parsing=tool_call_end,
+        _inner_parser=parse_tool_calls,
+    )
+
+
+def _parse_json_calls(text: str) -> list[ToolCallItem] | None:
+    try:
+        text = text.removeprefix("<tool_call>")
+        text = text.removesuffix("</tool_call>")
+        top_level = {
+            k: json.dumps(v) if isinstance(v, (dict, list)) else v
+            for k, v in json.loads(text).items()  # pyright: ignore[reportAny]
+        }
+        return [ToolCallItem.model_validate(top_level)]
+    except Exception:
+        return None
+
+
+def _flatten(p: dict[str, Any]) -> dict[str, str]:
+    return {
+        k: json.dumps(v) if isinstance(v, (dict, list)) else str(v)  # pyright: ignore[reportAny]
+        for k, v in p.items()  # pyright: ignore[reportAny]
+    }
+
+
+json_parser: Final[ToolParser] = ToolParser(
+    start_parsing="<tool_call>",
+    end_parsing="</tool_call>",
+    _inner_parser=_parse_json_calls,
+)
+
+
+def infer_tool_parser(chat_template: str) -> ToolParser | None:
+    """Attempt to auto-infer a tool parser from the chat template."""
+    if "<tool_call>" in chat_template and "tool_call.name" in chat_template:
+        return json_parser
+    return None

python/mlx_engine/pyproject.toml

@@ -4,11 +4,27 @@ version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.13"
-dependencies = ["exo_core"]
+dependencies = ["exo_core", "psutil >= 7.0.0", "loguru >= 0.7.3"]
 
 [build-system]
 requires = ["uv_build>=0.8.9,<0.9.0"]
 build-backend = "uv_build"
 
 [tool.uv.sources]
-exo_core = { workspace = true }
+exo_core = { workspace = true, editable = true }
+
+[tool.basedpyright]
+typeCheckingMode = "strict"
+failOnWarnings = true
+
+reportAny = "error"
+reportUnknownVariableType = "error"
+reportUnknownParameterType = "error"
+reportMissingParameterType = "error"
+reportMissingTypeStubs = "error"
+reportInvalidCast = "error"
+reportUnnecessaryCast = "error"
+reportUnnecessaryTypeIgnoreComment = "error"
+
+pythonVersion = "3.13"
+stubPath = "../../.typings"

python/mlx_engine/src/mlx_engine/builder.py

@@ -0,0 +1,132 @@
+from dataclasses import dataclass
+from typing import Self, Callable
+from exo_core.engine import EngineBuilder, Engine
+from exo_core.types.common import ModelId
+from exo_core.types.instances import BoundInstance
+from exo_core.types.tasks import TextGeneration
+from exo_core.types.runner_response import GenerationResponse
+from mlx_engine.utils_mlx import initialize_mlx, load_mlx_items
+from mlx_engine.types import Model
+from mlx_engine.generator.generate import (
+    mlx_generate,
+    warmup_inference,
+)
+from exo_core.utils.tool_parsers import make_mlx_parser
+
+
+@dataclass
+class MlxBuilder(EngineBuilder[BoundInstance, TextGeneration, GenerationResponse]):
+    import mlx.core as mx
+    from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+    model_id: ModelId
+    bound_instance: BoundInstance
+    event_sender: MpSender[Event]
+    cancel_receiver: MpReceiver[TaskId]
+    inference_model: Model | None = None
+    tokenizer: TokenizerWrapper | None = None
+    group: mx.distributed.Group | None = None
+
+    @classmethod
+    def create(
+        cls,
+        bound_instance: BoundInstance,
+        event_sender: MpSender[Event],
+        cancel_receiver: MpReceiver[TaskId],
+    ) -> Self:
+        return cls(
+            bound_instance.instance.shard_assignments.model_id,
+            bound_instance,
+            event_sender,
+            cancel_receiver,
+        )
+
+    def connect(self) -> None:
+        self.group = initialize_mlx(self.bound_instance)
+
+    def load(
+        self,
+        on_timeout: Callable[[], None],
+        on_layer_loaded: Callable[[int, int], None],
+    ) -> None:
+        self.inference_model, self.tokenizer = load_mlx_items(
+            self.bound_instance,
+            self.group,
+            on_timeout=on_timeout,
+            on_layer_loaded=on_layer_loaded,
+        )
+
+    def build(self) -> SequentialGenerator | BatchGenerator:
+        assert self.inference_model
+        assert self.tokenizer
+
+        tool_parser = None
+        logger.info(
+            f"model has_tool_calling={self.tokenizer.has_tool_calling} using tokens {self.tokenizer.tool_call_start}, {self.tokenizer.tool_call_end}"
+        )
+        if (
+            self.tokenizer.tool_call_start
+            and self.tokenizer.tool_call_end
+            and self.tokenizer.tool_parser  # type: ignore
+        ):
+            tool_parser = make_mlx_parser(
+                self.tokenizer.tool_call_start,
+                self.tokenizer.tool_call_end,
+                self.tokenizer.tool_parser,  # type: ignore
+            )
+
+        kv_prefix_cache = KVPrefixCache(self.group)
+
+        from functools import partial
+
+        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(
+                tokenizer=self.tokenizer,
+                group=self.group,
+                tool_parser=tool_parser,
+                kv_prefix_cache=kv_prefix_cache,
+                model_id=self.model_id,
+                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
+
+        logger.info("using BatchGenerator")
+        gen = ExoBatchGenerator(
+            model=self.inference_model,
+            tokenizer=self.tokenizer,
+            group=self.group,
+            kv_prefix_cache=kv_prefix_cache,
+            model_id=self.model_id,
+        )
+        return BatchGenerator(
+            tokenizer=self.tokenizer,
+            group=self.group,
+            tool_parser=tool_parser,
+            kv_prefix_cache=kv_prefix_cache,
+            model_id=self.model_id,
+            device_rank=device_rank,
+            cancel_receiver=self.cancel_receiver,
+            event_sender=self.event_sender,
+            _gen=gen,
+        )
+
+    def close(self):
+        with contextlib.suppress(NameError, AttributeError):
+            del self.inference_model, self.tokenizer

python/mlx_engine/src/mlx_engine/cache.py

@@ -0,0 +1,386 @@
+import os
+from copy import deepcopy
+from typing import TYPE_CHECKING, cast
+
+import mlx.core as mx
+import psutil
+from exo_core.utils.memory import Memory
+from loguru import logger
+from mlx_lm.models.cache import (
+    ArraysCache,
+    CacheList,
+    KVCache,
+    QuantizedKVCache,
+    RotatingKVCache,
+)
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+from mlx_engine.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
+from mlx_engine.types import KVCacheType, MLXCacheType, Model
+
+if TYPE_CHECKING:
+    from vllm_engine.kv_cache import TorchKVCache
+
+
+# Fraction of device memory above which LRU eviction kicks in.
+# Smaller machines need more aggressive eviction.
+def _default_memory_threshold() -> float:
+    total_gb = Memory.from_bytes(psutil.virtual_memory().total).in_gb
+    if total_gb >= 128:
+        return 0.85
+    if total_gb >= 64:
+        return 0.80
+    if total_gb >= 32:
+        return 0.75
+    return 0.70
+
+
+_MEMORY_THRESHOLD = float(
+    os.environ.get("EXO_MEMORY_THRESHOLD", _default_memory_threshold())
+)
+
+
+class CacheSnapshot:
+    """Snapshot of states at a known token position."""
+
+    def __init__(
+        self, states: list[RotatingKVCache | ArraysCache | None], token_count: int
+    ):
+        self.states = states
+        self.token_count = token_count
+
+
+def snapshot_ssm_states(cache: MLXCacheType) -> CacheSnapshot:
+    states: list[ArraysCache | RotatingKVCache | None] = []
+    for c in cache:
+        if isinstance(c, (ArraysCache, RotatingKVCache)):
+            states.append(deepcopy(c))
+        else:
+            states.append(None)
+    token_count = cache_length(cache)
+    return CacheSnapshot(states=states, token_count=token_count)
+
+
+def _find_nearest_snapshot(
+    snapshots: list[CacheSnapshot],
+    target_token_count: int,
+) -> CacheSnapshot | None:
+    best: CacheSnapshot | None = None
+    for snap in snapshots:
+        if snap.token_count <= target_token_count and (
+            best is None or snap.token_count > best.token_count
+        ):
+            best = snap
+    return best
+
+
+def has_non_kv_caches(cache: MLXCacheType) -> bool:
+    """Check if a cache contains any ArraysCache (SSM) entries."""
+    return any(isinstance(c, (ArraysCache, RotatingKVCache)) for c in cache)
+
+
+class KVPrefixCache:
+    def __init__(self, group: mx.distributed.Group | None):
+        self.prompts: list[mx.array] = []  # mx array of tokens (ints)
+        self.caches: list[KVCacheType] = []
+        self._snapshots: list[list[CacheSnapshot] | None] = []
+        self._last_used: list[int] = []  # monotonic counter of last access per entry
+        self._access_counter: int = 0
+        self._group = group
+
+    def clear(self):
+        """Clear all cached prompts and caches."""
+        self.prompts.clear()
+        self.caches.clear()
+        self._snapshots.clear()
+        self._last_used.clear()
+
+    def add_kv_cache(
+        self,
+        prompt_tokens: mx.array,
+        cache: MLXCacheType,
+        ssm_snapshots: list[CacheSnapshot] | None = None,
+    ):
+        """Add a new cache entry. Evicts LRU entries if memory is high."""
+        self._evict_if_needed()
+        self.prompts.append(prompt_tokens)
+        self.caches.append(deepcopy(cache))
+        self._snapshots.append(ssm_snapshots)
+        self._access_counter += 1
+        self._last_used.append(self._access_counter)
+        logger.info(f"KV cache added: {len(prompt_tokens)} tokens")
+
+    def update_kv_cache(
+        self,
+        index: int,
+        prompt_tokens: mx.array,
+        cache: MLXCacheType,
+        snapshots: list[CacheSnapshot] | None,
+        restore_pos: int,
+    ):
+        """Update an existing cache entry in-place."""
+        old_snapshots = self._snapshots[index]
+        merged: list[CacheSnapshot] = []
+        if old_snapshots:
+            merged = [s for s in old_snapshots if s.token_count <= restore_pos]
+        if snapshots:
+            merged.extend(snapshots)
+
+        self.prompts[index] = prompt_tokens
+        self.caches[index] = deepcopy(cache)
+        self._snapshots[index] = merged or None
+        self._access_counter += 1
+        self._last_used[index] = self._access_counter
+        logger.info(f"KV cache updated (index {index}): {len(prompt_tokens)} tokens")
+
+    def _get_mlx_cache(self, index: int) -> MLXCacheType:
+        cached = self.caches[index]
+        return cast(MLXCacheType, cached)
+
+    def _get_snapshot(
+        self, entry_index: int, target_token_count: int
+    ) -> tuple[int, CacheSnapshot | None]:
+        if not has_non_kv_caches(self._get_mlx_cache(entry_index)):
+            return target_token_count, None
+
+        snapshots = self._snapshots[entry_index]
+        if not snapshots:
+            return 0, None
+
+        snap = _find_nearest_snapshot(snapshots, target_token_count)
+        if snap is not None:
+            return snap.token_count, snap
+
+        return 0, None
+
+    def get_kv_cache(
+        self,
+        model: Model,
+        prompt_tokens: mx.array,
+    ) -> tuple[MLXCacheType, mx.array, int | None]:
+        """Get KV cache for prompt, returning remaining tokens to prefill.
+
+        Returns:
+            Tuple of (cache, remaining_tokens, matched_index) where:
+            - cache: KV cache to use for generation
+            - remaining_tokens: tokens that still need prefilling
+            - matched_index: index of the matched entry (None if no match)
+
+        For models with SSM layers (which are ArraysCache in mlx), the cache is trimmed to the
+        nearest SSM snapshot position at or before the match point for correctness.
+        Same for rotating KV Cache.
+        """
+        max_length = len(prompt_tokens)
+
+        best_index: int | None = None
+        best_length = 0
+        is_exact = False
+
+        # Find best cache match
+        for i, cached_prompt in enumerate(self.prompts):
+            length = get_prefix_length(prompt_tokens, cached_prompt)
+            if length >= max_length - 1:
+                best_index, best_length = i, length
+                is_exact = True
+                break
+            if length > best_length:
+                best_index, best_length = i, length
+
+        if best_index is None:
+            return make_kv_cache(model), prompt_tokens, None
+
+        # For exact match: trim to max_length-1 so remaining has the last token
+        # For partial match: trim to best_length, remaining has suffix to prefill
+        # This ensures stream_generate always has at least one token to start with
+        mlx_cache = self._get_mlx_cache(best_index)
+        has_ssm = has_non_kv_caches(mlx_cache)
+        target = (max_length - 1) if is_exact and not has_ssm else best_length
+        restore_pos, restore_snap = self._get_snapshot(best_index, target)
+
+        # No usable snapshot — need fresh cache
+        if restore_snap is None and has_ssm:
+            return make_kv_cache(model), prompt_tokens, None
+
+        prompt_cache = deepcopy(mlx_cache)
+        cached_length = cache_length(mlx_cache)
+        tokens_to_trim = cached_length - restore_pos
+        if tokens_to_trim > 0:
+            trim_cache(prompt_cache, tokens_to_trim, restore_snap)
+            # Reset cache offset to match trimmed length
+            for c in prompt_cache:
+                if hasattr(c, "offset"):
+                    c.offset = restore_pos
+
+        self._access_counter += 1
+        self._last_used[best_index] = self._access_counter
+        remaining = prompt_tokens[restore_pos:]
+
+        return prompt_cache, remaining, best_index
+
+    def lookup(
+        self, prompt_token_ids: list[int]
+    ) -> tuple["TorchKVCache | None", int, int | None]:
+        from exo.worker.engines.vllm.kv_cache import TorchKVCache
+
+        prompt_mx = mx.array(prompt_token_ids)
+        max_length = len(prompt_token_ids)
+        best_index: int | None = None
+        best_length = 0
+
+        for i, cached_prompt in enumerate(self.prompts):
+            length = get_prefix_length(prompt_mx, cached_prompt)
+            if length >= max_length - 1:
+                best_index, best_length = i, length
+                break
+            if length > best_length:
+                best_index, best_length = i, length
+
+        if best_index is None or best_length == 0:
+            return None, 0, None
+
+        best_length = min(best_length, max_length - 1)
+
+        self._access_counter += 1
+        self._last_used[best_index] = self._access_counter
+
+        cached = self.caches[best_index]
+        if isinstance(cached, TorchKVCache):
+            return cached.trim_to(best_length), best_length, best_index
+
+        torch_cache = TorchKVCache.from_mlx_cache(cached)
+        return torch_cache.trim_to(best_length), best_length, best_index
+
+    def add_from_torch(
+        self, prompt_token_ids: list[int], cache: "TorchKVCache"
+    ) -> None:
+        self._evict_if_needed()
+        self.prompts.append(mx.array(prompt_token_ids))
+        self.caches.append(cache.detach_cpu())
+        self._snapshots.append(None)
+        self._access_counter += 1
+        self._last_used.append(self._access_counter)
+        logger.info(f"KV cache added (torch): {len(prompt_token_ids)} tokens")
+
+    def _evict_if_needed(self):
+        """Evict least recently used entries while memory usage is high."""
+        if len(self.caches) == 0:
+            return
+
+        # Evict LRU entries until below threshold
+        while (
+            len(self.caches) > 0
+            and self.get_memory_used_percentage() > _MEMORY_THRESHOLD
+        ):
+            lru_index = self._last_used.index(min(self._last_used))
+            evicted_tokens = len(self.prompts[lru_index])
+            self.prompts.pop(lru_index)
+            self.caches.pop(lru_index)
+            self._snapshots.pop(lru_index)
+            self._last_used.pop(lru_index)
+            logger.info(
+                f"KV cache evicted LRU entry ({evicted_tokens} tokens) due to memory usage"
+            )
+
+    def get_memory_used_percentage(self) -> float:
+        local_pressure: float = get_memory_used_percentage()
+
+        if self._group is None:
+            return local_pressure
+
+        all_pressure = mx.distributed.all_gather(
+            mx.array([local_pressure], dtype=mx.float32),
+            group=self._group,
+        )
+        # .item() evals.
+        max_pressure = float(mx.max(all_pressure).item())
+        return max_pressure
+
+
+def trim_cache(
+    cache: MLXCacheType,
+    num_tokens: int,
+    snapshot: CacheSnapshot | None = None,
+) -> None:
+    for i, c in enumerate(cache):
+        if isinstance(c, (ArraysCache, RotatingKVCache)):
+            if snapshot is not None and snapshot.states[i] is not None:
+                cache[i] = deepcopy(snapshot.states[i])  # type: ignore
+            else:
+                c.state = [None] * len(c.state)
+        else:
+            c.trim(num_tokens)
+
+
+def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
+    """Encode a prompt string to token array.
+
+    For chat-templated prompts (which have their own structure markers like
+    <|im_user|>, <|im_middle|>, etc.), we should NOT add BOS/EOS tokens as
+    that would corrupt the prompt structure.
+    """
+    # Chat templates define their own structure - don't add BOS/EOS
+    prompt_tokens = tokenizer.encode(prompt, add_special_tokens=False)
+    return mx.array(prompt_tokens)
+
+
+def _entry_length(
+    c: KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList,
+) -> int:
+    # Use .offset attribute which KVCache types have (len() not implemented in older QuantizedKVCache).
+    if hasattr(c, "offset"):
+        return c.offset
+    # For CacheList
+    if hasattr(c, "size"):
+        return int(c.size())  # type: ignore
+    return 0
+
+
+def cache_length(cache: MLXCacheType) -> int:
+    """Get the number of tokens in a KV cache."""
+    return max(_entry_length(c) for c in cache)
+
+
+def get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
+    """Find the length of the common prefix between two token arrays."""
+    n = min(int(prompt.shape[0]), int(cached_prompt.shape[0]))
+    if n == 0:
+        return 0
+
+    equal = mx.equal(prompt[:n], cached_prompt[:n]).astype(mx.int32)
+    prefix_mask = mx.cumprod(equal)  # stays 1 until first mismatch, then 0 forever
+    return int(mx.sum(prefix_mask).item())
+
+
+def get_available_memory() -> Memory:
+    mem: int = psutil.virtual_memory().available
+    return Memory.from_bytes(mem)
+
+
+def get_memory_used_percentage() -> float:
+    mem = psutil.virtual_memory()
+    # percent is 0-100
+    return float(mem.percent / 100)
+
+
+def make_kv_cache(
+    model: Model, max_kv_size: int | None = None, keep: int = 0
+) -> MLXCacheType:
+    assert hasattr(model, "layers")
+
+    if hasattr(model, "make_cache"):
+        logger.info("Using MLX LM's make cache")
+        return model.make_cache()  # type: ignore
+
+    if max_kv_size is None:
+        if KV_CACHE_BITS is None:
+            logger.info("Using default KV cache")
+            return [KVCache() for _ in model.layers]
+        else:
+            logger.info("Using quantized KV cache")
+            return [
+                QuantizedKVCache(group_size=CACHE_GROUP_SIZE, bits=KV_CACHE_BITS)
+                for _ in model.layers
+            ]
+    else:
+        logger.info(f"Using rotating KV cache with {max_kv_size=} with {keep=}")
+        return [RotatingKVCache(max_size=max_kv_size, keep=keep) for _ in model.layers]

python/mlx_engine/src/mlx_engine/constants.py

@@ -0,0 +1,17 @@
+# TODO: Do we want so many constants?
+#  I think we want a lot of these as parameters?
+
+KV_GROUP_SIZE: int | None = 32
+KV_BITS: int | None = None
+ATTENTION_KV_BITS: int | None = 4
+MAX_TOKENS: int = 32168
+MAX_KV_SIZE: int | None = 3200
+KEEP_KV_SIZE: int | None = 1600
+QUANTIZE_MODEL_MODE: str | None = "affine"
+CACHE_GROUP_SIZE: int = 64
+KV_CACHE_BITS: int | None = None
+
+DEFAULT_TOP_LOGPROBS: int = 5
+
+# TODO: We should really make this opt-in, but Kimi requires trust_remote_code=True
+TRUST_REMOTE_CODE: bool = True

python/mlx_engine/src/mlx_engine/dsml_encoding.py

@@ -0,0 +1,72 @@
+import json
+import re
+from typing import Any
+
+from mlx_lm.chat_templates import deepseek_v32
+
+from exo_core.types.runner_response import ToolCallItem
+
+BOS_TOKEN: str = deepseek_v32.bos_token
+EOS_TOKEN: str = deepseek_v32.eos_token
+DSML_TOKEN: str = deepseek_v32.dsml_token
+THINKING_START: str = deepseek_v32.thinking_start_token
+THINKING_END: str = deepseek_v32.thinking_end_token
+USER_TOKEN = "<\uff5cUser\uff5c>"
+ASSISTANT_TOKEN = "<\uff5cAssistant\uff5c>"
+TOOL_CALLS_START = f"<{DSML_TOKEN}function_calls>"
+TOOL_CALLS_END = f"</{DSML_TOKEN}function_calls>"
+encode_messages = deepseek_v32.encode_messages
+
+_INVOKE_PATTERN = re.compile(
+    rf"<{re.escape(DSML_TOKEN)}invoke\s+name=\"([^\"]+)\">"
+    rf"(.*?)"
+    rf"</{re.escape(DSML_TOKEN)}invoke>",
+    re.DOTALL,
+)
+
+_PARAM_PATTERN = re.compile(
+    rf"<{re.escape(DSML_TOKEN)}parameter\s+name=\"([^\"]+)\"\s+string=\"(true|false)\">"
+    rf"(.*?)"
+    rf"</{re.escape(DSML_TOKEN)}parameter>",
+    re.DOTALL,
+)
+
+
+def parse_dsml_output(text: str) -> list[ToolCallItem] | None:
+    """Parse DSML function_calls block from model output text.
+
+    Args:
+        text: The text containing the DSML function_calls block
+              (including the start/end markers).
+
+    Returns:
+        List of ToolCallItem, or None if parsing fails.
+    """
+    tool_calls: list[ToolCallItem] = []
+
+    for invoke_match in _INVOKE_PATTERN.finditer(text):
+        func_name = invoke_match.group(1)
+        invoke_body = invoke_match.group(2)
+
+        args: dict[str, Any] = {}
+        for param_match in _PARAM_PATTERN.finditer(invoke_body):
+            param_name = param_match.group(1)
+            is_string = param_match.group(2) == "true"
+            param_value = param_match.group(3)
+
+            if is_string:
+                args[param_name] = param_value
+            else:
+                try:
+                    args[param_name] = json.loads(param_value)
+                except (json.JSONDecodeError, ValueError):
+                    args[param_name] = param_value
+
+        tool_calls.append(
+            ToolCallItem(
+                name=func_name,
+                arguments=json.dumps(args),
+            )
+        )
+
+    return tool_calls if tool_calls else None

python/mlx_engine/src/mlx_engine/generator/batch_generate.py

@@ -0,0 +1,407 @@
+import time
+from dataclasses import dataclass, field
+from typing import Callable, cast
+
+import mlx.core as mx
+from exo_core.types.common import ModelId
+from exo_core.types.runner_response import GenerationResponse
+from exo_core.types.tasks import TaskId
+from exo_core.types.text_generation import TextGenerationTaskParams
+from exo_core.utils.memory import Memory
+from loguru import logger
+from mlx_lm.generate import (
+    BatchGenerator as MlxBatchGenerator,
+)
+from mlx_lm.models.cache import RotatingKVCache
+from mlx_lm.sample_utils import make_logits_processors, make_sampler
+from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper
+
+from exo.api.types import (
+    CompletionTokensDetails,
+    FinishReason,
+    GenerationStats,
+    PromptTokensDetails,
+    TopLogprobItem,
+    Usage,
+)
+from mlx_engine.cache import (
+    CacheSnapshot,
+    KVPrefixCache,
+    encode_prompt,
+    make_kv_cache,
+)
+from mlx_engine.constants import DEFAULT_TOP_LOGPROBS, MAX_TOKENS
+from mlx_engine.generator.generate import (
+    ban_token_ids,
+    eos_ids_from_tokenizer,
+    extract_top_logprobs,
+    prefill,
+    warmup_inference,
+)
+from mlx_engine.types import MLXCacheType, Model
+from mlx_engine.utils_mlx import fix_unmatched_think_end_tokens
+
+_MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5
+
+
+def _stop_sequences(task_params: TextGenerationTaskParams) -> list[str]:
+    if task_params.stop is None:
+        return []
+    if isinstance(task_params.stop, str):
+        return [task_params.stop]
+    return task_params.stop
+
+
+@dataclass
+class _EngineTask:
+    uid: int
+    task_params: TextGenerationTaskParams
+    all_prompt_tokens: mx.array
+    prefix_hit_length: int
+    matched_index: int | None
+    cache_snapshots: list[CacheSnapshot] | None
+    detokenizer: StreamingDetokenizer
+    on_generation_token: Callable[[], None] | None = None
+    generated_text_parts: list[str] = field(default_factory=list)
+    potential_stop_sequence_text: str = ""
+    completion_tokens: int = 0
+    generation_start_time: float = 0.0
+    in_thinking: bool = False
+    reasoning_tokens: int = 0
+    prefill_tps: float = 0.0
+
+
+@dataclass(eq=False)
+class ExoBatchGenerator:
+    model: Model
+    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)
+    _uid_to_task_id: dict[int, TaskId] = field(default_factory=dict, init=False)
+
+    def __post_init__(self) -> None:
+        self._mlx_gen = MlxBatchGenerator(
+            model=self.model,
+            stop_tokens=set(eos_ids_from_tokenizer(self.tokenizer)),
+            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 (
+            bool(self._active_tasks)
+            or bool(self._mlx_gen.unprocessed_prompts)
+            or self._mlx_gen.active_batch is not None
+        )
+
+    def submit(
+        self,
+        task_id: TaskId,
+        task_params: TextGenerationTaskParams,
+        prompt: str,
+        on_prefill_progress: Callable[[int, int], None] | None = None,
+        distributed_prompt_progress_callback: Callable[[], None] | None = None,
+        on_generation_token: Callable[[], None] | None = None,
+    ) -> TaskId:
+        all_prompt_tokens = encode_prompt(self.tokenizer, prompt)
+        all_prompt_tokens = fix_unmatched_think_end_tokens(
+            all_prompt_tokens, self.tokenizer
+        )
+
+        is_bench = task_params.bench
+
+        prefix_hit_length = 0
+        matched_index: int | None = None
+        prompt_tokens = all_prompt_tokens
+
+        if self.kv_prefix_cache is not None and not is_bench:
+            cache, remaining_tokens, matched_index = self.kv_prefix_cache.get_kv_cache(
+                self.model, all_prompt_tokens
+            )
+            prefix_hit_length = len(all_prompt_tokens) - len(remaining_tokens)
+            if prefix_hit_length > 0:
+                logger.info(
+                    f"KV cache hit: {prefix_hit_length}/{len(all_prompt_tokens)} tokens "
+                    f"cached ({100 * prefix_hit_length / len(all_prompt_tokens):.1f}%)"
+                )
+                prompt_tokens = remaining_tokens
+            else:
+                cache = make_kv_cache(self.model)
+        else:
+            cache = make_kv_cache(self.model)
+
+        seed = task_params.seed if task_params.seed is not None else 42
+        mx.random.seed(seed)
+
+        sampler = make_sampler(
+            temp=task_params.temperature
+            if task_params.temperature is not None
+            else 0.7,
+            top_p=task_params.top_p if task_params.top_p is not None else 1.0,
+            min_p=task_params.min_p if task_params.min_p is not None else 0.05,
+            top_k=task_params.top_k if task_params.top_k is not None else 0,
+        )
+
+        _prefill_tps, _prefill_tokens, cache_snapshots = prefill(
+            self.model,
+            self.tokenizer,
+            sampler,
+            prompt_tokens[:-1],
+            cache,
+            self.group,
+            on_prefill_progress,
+            distributed_prompt_progress_callback,
+        )
+
+        # We need to clamp rotating kv caches to max size so that mlx lm's _merge_caches behaves
+        for c in cache:
+            if (
+                isinstance(c, RotatingKVCache)
+                and c.keys is not None
+                and c.values is not None
+                and c.keys.shape[2] > c.max_size
+            ):
+                trim_size = c.keys.shape[2] - c.max_size
+                c.keys = c._trim(trim_size, c.keys)
+                c.values = c._trim(trim_size, c.values)
+                c._idx = c.max_size
+
+        if not is_bench:
+            self._save_prefix_cache(
+                all_prompt_tokens,
+                list(cache),
+                cache_snapshots,
+                prefix_hit_length,
+                matched_index,
+            )
+
+        last_tokens = prompt_tokens[-2:]
+
+        logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = (
+            make_logits_processors(
+                repetition_penalty=task_params.repetition_penalty,
+                repetition_context_size=task_params.repetition_context_size,
+            )
+        )
+        if is_bench:
+            # Only sample length eos tokens
+            eos_ids = eos_ids_from_tokenizer(self.tokenizer)
+            logits_processors = [ban_token_ids(eos_ids)] + logits_processors
+
+        max_tokens = task_params.max_output_tokens or MAX_TOKENS
+
+        uids = self._mlx_gen.insert(
+            prompts=[last_tokens.tolist()],
+            max_tokens=[max_tokens],
+            caches=[list(cache)],
+            samplers=[sampler],
+            logits_processors=[logits_processors],
+        )
+
+        assert len(uids) == 1
+
+        uid = uids[0]
+        self._uid_to_task_id[uid] = task_id
+
+        self._active_tasks[uid] = _EngineTask(
+            uid=uid,
+            task_params=task_params,
+            all_prompt_tokens=all_prompt_tokens,
+            prefix_hit_length=prefix_hit_length,
+            matched_index=matched_index,
+            cache_snapshots=cache_snapshots or None,
+            detokenizer=self.tokenizer.detokenizer,
+            on_generation_token=on_generation_token,
+            generation_start_time=time.perf_counter(),
+            prefill_tps=_prefill_tps,
+        )
+
+        return task_id
+
+    def step(self) -> list[tuple[TaskId, GenerationResponse]]:
+        if not self.has_work:
+            return []
+
+        responses = self._mlx_gen.next()
+
+        results: list[tuple[TaskId, GenerationResponse]] = []
+
+        for response in responses:
+            if response.uid not in self._active_tasks:
+                logger.warning(
+                    f"response uid {response.uid} was not found - should be active"
+                )
+                continue
+
+            state = self._active_tasks[response.uid]
+            if state.on_generation_token is not None:
+                state.on_generation_token()
+            if response.finish_reason != "stop":
+                state.detokenizer.add_token(response.token)
+            if response.finish_reason is not None:
+                state.detokenizer.finalize()
+            text = state.detokenizer.last_segment
+            state.completion_tokens += 1
+            state.generated_text_parts.append(text)
+            state.potential_stop_sequence_text += text
+
+            think_start = self.tokenizer.think_start
+            think_end = self.tokenizer.think_end
+            if think_start is not None and text == think_start:
+                state.in_thinking = True
+            elif think_end is not None and text == think_end:
+                state.in_thinking = False
+            if state.in_thinking:
+                state.reasoning_tokens += 1
+
+            finish_reason: FinishReason | None = cast(
+                FinishReason | None, response.finish_reason
+            )
+            task_params = state.task_params
+            stop_sequences = _stop_sequences(task_params)
+            max_stop_len = max((len(s) for s in stop_sequences), default=0)
+
+            if stop_sequences:
+                for stop_seq in stop_sequences:
+                    if stop_seq in state.potential_stop_sequence_text:
+                        stop_index = state.potential_stop_sequence_text.find(stop_seq)
+                        text_before_stop = state.potential_stop_sequence_text[
+                            :stop_index
+                        ]
+                        chunk_start = len(state.potential_stop_sequence_text) - len(
+                            text
+                        )
+                        text = text_before_stop[chunk_start:]
+                        finish_reason = "stop"
+                        break
+
+            is_done = finish_reason is not None
+
+            logprob: float | None = None
+            top_logprobs: list[TopLogprobItem] | None = None
+            if task_params.logprobs:
+                logprob, top_logprobs = extract_top_logprobs(
+                    logprobs=response.logprobs,
+                    tokenizer=self.tokenizer,
+                    top_logprobs=task_params.top_logprobs or DEFAULT_TOP_LOGPROBS,
+                    selected_token=response.token,
+                )
+
+            stats: GenerationStats | None = None
+            usage: Usage | None = None
+            if is_done:
+                try:
+                    mlx_stats = self._mlx_gen.stats()
+                    generation_tps = mlx_stats.generation_tps
+                except ZeroDivisionError:
+                    generation_elapsed = (
+                        time.perf_counter() - state.generation_start_time
+                    )
+                    generation_tps = (
+                        state.completion_tokens / generation_elapsed
+                        if generation_elapsed > 0
+                        else 0.0
+                    )
+
+                stats = GenerationStats(
+                    prompt_tps=state.prefill_tps,
+                    generation_tps=generation_tps,
+                    prompt_tokens=len(state.all_prompt_tokens),
+                    generation_tokens=state.completion_tokens,
+                    peak_memory_usage=Memory.from_gb(mx.get_peak_memory() / 1e9),
+                )
+                total_prompt_tokens = len(state.all_prompt_tokens)
+                usage = Usage(
+                    prompt_tokens=total_prompt_tokens,
+                    completion_tokens=state.completion_tokens,
+                    total_tokens=total_prompt_tokens + state.completion_tokens,
+                    prompt_tokens_details=PromptTokensDetails(
+                        cached_tokens=state.prefix_hit_length
+                    ),
+                    completion_tokens_details=CompletionTokensDetails(
+                        reasoning_tokens=state.reasoning_tokens
+                    ),
+                )
+
+            results.append(
+                (
+                    self._uid_to_task_id.get(response.uid, TaskId(str(response.uid))),
+                    GenerationResponse(
+                        text=text,
+                        token=response.token,
+                        logprob=logprob,
+                        top_logprobs=top_logprobs,
+                        finish_reason=finish_reason,
+                        stats=stats,
+                        usage=usage,
+                    ),
+                )
+            )
+
+            if is_done:
+                del self._active_tasks[response.uid]
+                self._uid_to_task_id.pop(response.uid, None)
+            elif (
+                max_stop_len > 0
+                and len(state.potential_stop_sequence_text) > max_stop_len
+            ):
+                state.potential_stop_sequence_text = state.potential_stop_sequence_text[
+                    -max_stop_len:
+                ]
+
+        return results
+
+    def cancel(self, task_ids: list[TaskId]) -> None:
+        task_id_set = set(task_ids)
+        uids = [uid for uid, tid in self._uid_to_task_id.items() if tid in task_id_set]
+        if uids:
+            self._mlx_gen.remove(uids)
+        for uid in uids:
+            self._active_tasks.pop(uid, None)
+            self._uid_to_task_id.pop(uid, None)
+
+    def close(self) -> None:
+        self._mlx_gen.close()
+        mx.clear_cache()
+
+    def _save_prefix_cache(
+        self,
+        all_prompt_tokens: mx.array,
+        cache: MLXCacheType,
+        cache_snapshots: list[CacheSnapshot] | None,
+        prefix_hit_length: int,
+        matched_index: int | None,
+    ) -> None:
+        if self.kv_prefix_cache is None:
+            return
+
+        try:
+            hit_ratio = (
+                prefix_hit_length / len(all_prompt_tokens)
+                if len(all_prompt_tokens) > 0
+                else 0.0
+            )
+            if (
+                matched_index is not None
+                and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
+            ):
+                self.kv_prefix_cache.update_kv_cache(
+                    matched_index,
+                    all_prompt_tokens,
+                    cache,
+                    cache_snapshots,
+                    restore_pos=prefix_hit_length,
+                )
+            else:
+                self.kv_prefix_cache.add_kv_cache(
+                    all_prompt_tokens, cache, cache_snapshots
+                )
+        except Exception:
+            logger.warning("Failed to save prefix cache", exc_info=True)

python/mlx_engine/src/mlx_engine/generator/generate.py

@@ -0,0 +1,695 @@
+import functools
+import math
+import time
+from copy import deepcopy
+from typing import Callable, Generator, cast, get_args
+
+import mlx.core as mx
+from exo_core.types.common import ModelId
+from exo_core.types.runner_response import (
+    GenerationResponse,
+)
+from exo_core.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo_core.utils.memory import Memory
+from loguru import logger
+from mlx_lm.generate import (
+    maybe_quantize_kv_cache,
+    stream_generate,
+)
+from mlx_lm.models.cache import ArraysCache, RotatingKVCache
+from mlx_lm.sample_utils import make_logits_processors, make_sampler
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+from exo.api.types import (
+    CompletionTokensDetails,
+    FinishReason,
+    GenerationStats,
+    PromptTokensDetails,
+    TopLogprobItem,
+    Usage,
+)
+from mlx_engine.auto_parallel import (
+    PipelineFirstLayer,
+    PipelineLastLayer,
+    clear_prefill_sends,
+    flush_prefill_sends,
+    set_pipeline_prefill,
+    set_pipeline_queue_sends,
+)
+from mlx_engine.cache import (
+    CacheSnapshot,
+    KVPrefixCache,
+    encode_prompt,
+    has_non_kv_caches,
+    make_kv_cache,
+    snapshot_ssm_states,
+)
+from mlx_engine.constants import (
+    DEFAULT_TOP_LOGPROBS,
+    KV_BITS,
+    KV_GROUP_SIZE,
+    MAX_TOKENS,
+)
+from mlx_engine.types import MLXCacheType, Model
+from mlx_engine.utils_mlx import (
+    apply_chat_template,
+    fix_unmatched_think_end_tokens,
+    mx_barrier,
+)
+
+generation_stream = mx.new_stream(mx.default_device())
+
+_MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5
+
+
+class PrefillCancelled(BaseException):
+    """Raised when prefill is cancelled via the progress callback."""
+
+
+def _has_pipeline_communication_layer(model: Model):
+    for layer in model.layers:
+        if isinstance(layer, (PipelineFirstLayer, PipelineLastLayer)):
+            return True
+    return False
+
+
+def pipeline_parallel_prefill(
+    model: Model,
+    prompt: mx.array,
+    prompt_cache: MLXCacheType,
+    prefill_step_size: int,
+    kv_group_size: int | None,
+    kv_bits: int | None,
+    prompt_progress_callback: Callable[[int, int], None],
+    distributed_prompt_progress_callback: Callable[[], None] | None,
+    group: mx.distributed.Group,
+) -> None:
+    """Prefill the KV cache for pipeline parallel with overlapping stages.
+
+    Each rank processes the full prompt through its real cache, offset by leading
+    and trailing dummy iterations.
+
+    Total iterations per rank = N_real_chunks + world_size - 1:
+      - rank r leading dummies  (skip_pipeline_io, throwaway cache)
+      - N_real_chunks real      (pipeline IO active, real cache)
+      - (world_size-1-r) trailing dummies (skip_pipeline_io, throwaway cache)
+
+    e.g.
+    Timeline (2 ranks, 3 chunks of 10240 tokens @ step=4096):
+        iter 0: R0 real[0:4096]     R1 dummy
+        iter 1: R0 real[4096:8192]  R1 real[0:4096]
+        iter 2: R0 real[8192:10240] R1 real[4096:8192]
+        iter 3: R0 dummy            R1 real[8192:10240]
+
+    This function is designed to match mlx_lm's stream_generate exactly in terms of
+    side effects (given the same prefill step size)
+    """
+    prefill_step_size = prefill_step_size // min(4, group.size())
+
+    quantize_cache_fn: Callable[..., None] = functools.partial(
+        maybe_quantize_kv_cache,
+        quantized_kv_start=0,
+        kv_group_size=kv_group_size,
+        kv_bits=kv_bits,
+    )
+
+    _prompt_cache: MLXCacheType = prompt_cache
+    rank = group.rank()
+    world_size = group.size()
+
+    # Build list of real prompt chunk sizes
+    total = len(prompt)
+    real_chunk_sizes: list[int] = []
+    remaining = total - 1
+    while remaining:
+        n = min(prefill_step_size, remaining)
+        real_chunk_sizes.append(n)
+        remaining -= n
+    n_real = len(real_chunk_sizes)
+
+    # Each rank does: [rank leading dummies] [N real chunks] [world_size-1-rank trailing dummies]
+    n_leading = rank
+    n_trailing = world_size - 1 - rank
+    n_total = n_leading + n_real + n_trailing
+
+    t_start = time.perf_counter()
+    processed = 0
+    logger.info(
+        f"[R{rank}] Pipeline prefill: {n_real} real + {n_leading} leading + {n_trailing} trailing = {n_total} iterations"
+    )
+    clear_prefill_sends()
+
+    # Initial callback matching generate_step
+    prompt_progress_callback(0, total)
+
+    try:
+        with mx.stream(generation_stream):
+            for _ in range(n_leading):
+                if distributed_prompt_progress_callback is not None:
+                    distributed_prompt_progress_callback()
+
+            for i in range(n_real):
+                chunk_size = real_chunk_sizes[i]
+                model(
+                    prompt[processed : processed + chunk_size][None],
+                    cache=_prompt_cache,
+                )
+                quantize_cache_fn(_prompt_cache)
+                processed += chunk_size
+
+                if distributed_prompt_progress_callback is not None:
+                    distributed_prompt_progress_callback()
+
+                flush_prefill_sends()
+
+                prompt_progress_callback(processed, total)
+
+            for _ in range(n_trailing):
+                if distributed_prompt_progress_callback is not None:
+                    distributed_prompt_progress_callback()
+
+    finally:
+        clear_prefill_sends()
+
+    # Post-loop: process remaining 1 token + add +1 entry to match stream_generate.
+    for _ in range(2):
+        with mx.stream(generation_stream):
+            model(prompt[-1:][None], cache=_prompt_cache)
+            quantize_cache_fn(_prompt_cache)
+        flush_prefill_sends()
+
+    assert _prompt_cache is not None
+    mx.eval([c.state for c in _prompt_cache])  # type: ignore
+
+    # Final callback matching generate_step
+    prompt_progress_callback(total, total)
+
+    logger.info(
+        f"[R{rank}] Prefill: {n_real} real + {n_leading}+{n_trailing} dummy iterations, "
+        f"Processed {processed} tokens in {(time.perf_counter() - t_start) * 1000:.1f}ms"
+    )
+
+
+def prefill(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    sampler: Callable[[mx.array], mx.array],
+    prompt_tokens: mx.array,
+    cache: MLXCacheType,
+    group: mx.distributed.Group | None,
+    on_prefill_progress: Callable[[int, int], None] | None,
+    distributed_prompt_progress_callback: Callable[[], None] | None,
+) -> tuple[float, int, list[CacheSnapshot]]:
+    """Prefill the KV cache with prompt tokens.
+
+    This runs the model over the prompt tokens to populate the cache,
+    then trims off the extra generated token.
+
+    Returns:
+        (tokens_per_sec, num_tokens, snapshots)
+    """
+    num_tokens = len(prompt_tokens)
+    if num_tokens == 0:
+        return 0.0, 0, []
+
+    logger.debug(f"Prefilling {num_tokens} tokens...")
+    start_time = time.perf_counter()
+    has_ssm = has_non_kv_caches(cache)
+    snapshots: list[CacheSnapshot] = []
+
+    # TODO(evan): kill the callbacks/runner refactor
+    def progress_callback(processed: int, total: int) -> None:
+        elapsed = time.perf_counter() - start_time
+        tok_per_sec = processed / elapsed if elapsed > 0 else 0
+        logger.debug(
+            f"Prefill progress: {processed}/{total} tokens ({tok_per_sec:.1f} tok/s)"
+        )
+        if has_ssm:
+            snapshots.append(snapshot_ssm_states(cache))
+
+        if on_prefill_progress is not None:
+            on_prefill_progress(processed, total)
+
+    def combined_progress_callback(processed: int, total: int) -> None:
+        if distributed_prompt_progress_callback is not None:
+            distributed_prompt_progress_callback()
+        progress_callback(processed, total)
+
+    set_pipeline_prefill(model, is_prefill=True)
+
+    mx_barrier(group)
+    logger.info("Starting prefill")
+
+    is_pipeline = _has_pipeline_communication_layer(model)
+
+    prefill_step_size = 4096
+
+    try:
+        if is_pipeline and num_tokens >= prefill_step_size:
+            set_pipeline_queue_sends(model, queue_sends=True)
+            assert group is not None, "Pipeline prefill requires a distributed group"
+            pipeline_parallel_prefill(
+                model=model,
+                prompt=prompt_tokens,
+                prompt_cache=cache,
+                prefill_step_size=prefill_step_size,
+                kv_group_size=KV_GROUP_SIZE,
+                kv_bits=KV_BITS,
+                prompt_progress_callback=progress_callback,
+                distributed_prompt_progress_callback=distributed_prompt_progress_callback,
+                group=group,
+            )
+        else:
+            # Use max_tokens=1 because max_tokens=0 does not work.
+            # We just throw away the generated token - we only care about filling the cache
+            for _ in stream_generate(
+                model=model,
+                tokenizer=tokenizer,
+                prompt=prompt_tokens,
+                max_tokens=1,
+                sampler=sampler,
+                prompt_cache=cache,
+                prefill_step_size=prefill_step_size,
+                kv_group_size=KV_GROUP_SIZE,
+                kv_bits=KV_BITS,
+                prompt_progress_callback=combined_progress_callback,
+            ):
+                break  # Stop after first iteration - cache is now filled
+    except PrefillCancelled:
+        set_pipeline_queue_sends(model, queue_sends=False)
+        set_pipeline_prefill(model, is_prefill=False)
+        raise
+
+    set_pipeline_queue_sends(model, queue_sends=False)
+    set_pipeline_prefill(model, is_prefill=False)
+
+    # stream_generate added 1 extra generated token to the cache, so we should trim it.
+    # Because of needing to roll back arrays cache, we will generate on 2 tokens so trim 1 more.
+    pre_gen = deepcopy(snapshots[-2]) if has_ssm else None
+    for i, c in enumerate(cache):
+        if has_ssm and isinstance(c, (ArraysCache, RotatingKVCache)):
+            assert pre_gen is not None
+            if pre_gen.states[i] is not None:
+                cache[i] = deepcopy(pre_gen.states[i])  # type: ignore
+        else:
+            assert not isinstance(c, (ArraysCache, RotatingKVCache))
+            c.trim(2)
+
+    elapsed = time.perf_counter() - start_time
+    tokens_per_sec = num_tokens / elapsed if elapsed > 0 else 0.0
+    logger.debug(
+        f"Prefill complete: {num_tokens} tokens in {elapsed:.2f}s "
+        f"({tokens_per_sec:.1f} tok/s)"
+    )
+    # Exclude the last snapshot
+    return tokens_per_sec, num_tokens, snapshots[:-1] if snapshots else []
+
+
+def warmup_inference(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    group: mx.distributed.Group | None,
+    model_id: ModelId,
+) -> int:
+    logger.info(f"warming up inference for instance: {model_id}")
+    t = time.monotonic()
+
+    content = "Prompt to warm up the inference engine. Repeat this."
+
+    warmup_prompt = apply_chat_template(
+        tokenizer=tokenizer,
+        task_params=TextGenerationTaskParams(
+            model=ModelId(""),
+            input=[InputMessage(role="user", content=content)],
+        ),
+    )
+
+    tokens_generated = 0
+
+    cache = make_kv_cache(
+        model=model,
+    )
+
+    # Use a default sampler for warmup
+    sampler = make_sampler(temp=0.0)
+
+    mx_barrier(group)
+
+    logger.info("Generating warmup tokens")
+    for _r in stream_generate(
+        model=model,
+        tokenizer=tokenizer,
+        prompt=warmup_prompt,
+        max_tokens=50,
+        sampler=sampler,
+        prompt_cache=cache,
+        prefill_step_size=2048,
+        kv_group_size=KV_GROUP_SIZE,
+        kv_bits=KV_BITS,
+    ):
+        logger.info("Generated warmup token: " + str(_r.text))
+        tokens_generated += 1
+
+    logger.info("Generated ALL warmup tokens")
+
+    mx_barrier(group)
+
+    logger.info(f"warmed up by generating {tokens_generated} tokens")
+    check_for_cancel_every = min(
+        math.ceil(tokens_generated / min(time.monotonic() - t, 0.001)), 100
+    )
+    if group is not None:
+        check_for_cancel_every = int(
+            mx.max(
+                mx.distributed.all_gather(
+                    mx.array([check_for_cancel_every]),
+                    group=group,
+                )
+            ).item()
+        )
+
+    logger.info(
+        f"runner checking for cancellation every {check_for_cancel_every} tokens"
+    )
+
+    return check_for_cancel_every
+
+
+def ban_token_ids(token_ids: list[int]) -> Callable[[mx.array, mx.array], mx.array]:
+    token_ids = [int(t) for t in token_ids]
+
+    def proc(_history: mx.array, logits: mx.array) -> mx.array:
+        for tid in token_ids:
+            logits[..., tid] = -1e9
+        return logits
+
+    return proc
+
+
+def eos_ids_from_tokenizer(tokenizer: TokenizerWrapper) -> list[int]:
+    eos: list[int] | None = getattr(tokenizer, "eos_token_ids", None)
+    if eos is None:
+        return []
+    return eos
+
+
+def extract_top_logprobs(
+    logprobs: mx.array,
+    tokenizer: TokenizerWrapper,
+    top_logprobs: int,
+    selected_token: int,
+) -> tuple[float, list[TopLogprobItem]]:
+    """Extract the selected token's logprob and top alternative tokens.
+
+    Args:
+        logprobs: Full vocabulary logprobs array from MLX
+        tokenizer: Tokenizer for decoding token IDs to strings
+        top_logprobs: Number of top alternatives to return
+        selected_token: The token ID that was actually sampled
+
+    Returns:
+        Tuple of (selected_token_logprob, list of TopLogprobItem for top alternatives)
+    """
+    # Get the logprob of the selected token
+    selected_logprob = float(logprobs[selected_token].item())
+
+    # Get top indices (most probable tokens)
+    # mx.argpartition gives indices that would partition the array
+    # We negate logprobs since argpartition finds smallest, and we want largest
+    top_logprobs = min(top_logprobs, logprobs.shape[0])  # Don't exceed vocab size
+    top_indices = mx.argpartition(-logprobs, top_logprobs)[:top_logprobs]
+
+    # Get the actual logprob values for these indices
+    top_values = logprobs[top_indices]
+
+    # Sort by logprob (descending) for consistent ordering
+    sort_order = mx.argsort(-top_values)
+    top_indices = top_indices[sort_order]
+    top_values = top_values[sort_order]
+
+    # Convert to list of TopLogprobItem
+    top_logprob_items: list[TopLogprobItem] = []
+    for i in range(top_logprobs):
+        token_id = int(top_indices[i].item())
+        token_logprob = float(top_values[i].item())
+        if math.isnan(token_logprob):
+            continue
+
+        # Decode token ID to string
+        token_str = tokenizer.decode([token_id])
+        # Get byte representation
+        token_bytes = list(token_str.encode("utf-8"))
+        top_logprob_items.append(
+            TopLogprobItem(
+                token=token_str,
+                logprob=token_logprob,
+                bytes=token_bytes,
+            )
+        )
+
+    return selected_logprob, top_logprob_items
+
+
+def mlx_generate(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    task: TextGenerationTaskParams,
+    prompt: str,
+    kv_prefix_cache: KVPrefixCache | None,
+    group: mx.distributed.Group | None,
+    on_prefill_progress: Callable[[int, int], None] | None = None,
+    distributed_prompt_progress_callback: Callable[[], None] | None = None,
+    on_generation_token: Callable[[], None] | None = None,
+) -> Generator[GenerationResponse]:
+    # Ensure that generation stats only contains peak memory for this generation
+    mx.reset_peak_memory()
+    # TODO: Randomise task seed and set in taskparams, instead of hard coding as 42.
+    seed = task.seed or 42
+    mx.random.seed(seed)
+
+    # Encode prompt once at the top and fix unmatched think tags
+    all_prompt_tokens = encode_prompt(tokenizer, prompt)
+    all_prompt_tokens = fix_unmatched_think_end_tokens(all_prompt_tokens, tokenizer)
+
+    # Do not use the prefix cache if we are trying to do benchmarks.
+    is_bench = task.bench
+    if is_bench:
+        kv_prefix_cache = None
+
+    # Use prefix cache if available, otherwise create fresh cache
+    prefix_hit_length = 0
+    matched_index: int | None = None
+    if kv_prefix_cache is None:
+        caches = make_kv_cache(model=model)
+        prompt_tokens = all_prompt_tokens
+    else:
+        caches, prompt_tokens, matched_index = kv_prefix_cache.get_kv_cache(
+            model, all_prompt_tokens
+        )
+        prefix_hit_length = len(all_prompt_tokens) - len(prompt_tokens)
+        if prefix_hit_length > 0:
+            logger.info(
+                f"KV cache hit: {prefix_hit_length}/{len(all_prompt_tokens)} tokens cached ({100 * prefix_hit_length / len(all_prompt_tokens):.1f}%)"
+            )
+
+    logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = (
+        make_logits_processors(
+            repetition_penalty=task.repetition_penalty,
+            repetition_context_size=task.repetition_context_size,
+        )
+    )
+    if is_bench:
+        # Only sample length eos tokens
+        eos_ids = eos_ids_from_tokenizer(tokenizer)
+        logits_processors = [ban_token_ids(eos_ids)] + logits_processors
+
+    sampler = make_sampler(
+        temp=task.temperature if task.temperature is not None else 0.7,
+        top_p=task.top_p if task.top_p is not None else 1.0,
+        min_p=task.min_p if task.min_p is not None else 0.05,
+        top_k=task.top_k if task.top_k is not None else 0,
+    )
+
+    # Normalize stop sequences to a list
+    stop_sequences: list[str] = (
+        ([task.stop] if isinstance(task.stop, str) else task.stop)
+        if task.stop is not None
+        else []
+    )
+    max_stop_len = max((len(s) for s in stop_sequences), default=0)
+
+    # Prefill cache with all tokens except the last one
+    prefill_tps, prefill_tokens, ssm_snapshots_list = prefill(
+        model,
+        tokenizer,
+        sampler,
+        prompt_tokens[:-1],
+        caches,
+        group,
+        on_prefill_progress,
+        distributed_prompt_progress_callback,
+    )
+    cache_snapshots: list[CacheSnapshot] | None = ssm_snapshots_list or None
+
+    # stream_generate starts from the last token
+    last_token = prompt_tokens[-2:]
+
+    max_tokens = task.max_output_tokens or MAX_TOKENS
+    accumulated_text = ""
+    generated_text_parts: list[str] = []
+    generation_start_time = time.perf_counter()
+    usage: Usage | None = None
+    in_thinking = False
+    reasoning_tokens = 0
+    think_start = tokenizer.think_start
+    think_end = tokenizer.think_end
+
+    logger.info("Starting decode")
+    mx_barrier(group)
+
+    for completion_tokens, out in enumerate(
+        stream_generate(
+            model=model,
+            tokenizer=tokenizer,
+            prompt=last_token,
+            max_tokens=max_tokens,
+            sampler=sampler,
+            logits_processors=logits_processors,
+            prompt_cache=caches,
+            prefill_step_size=1,
+            kv_group_size=KV_GROUP_SIZE,
+            kv_bits=KV_BITS,
+        ),
+        start=1,
+    ):
+        generated_text_parts.append(out.text)
+        accumulated_text += out.text
+
+        if think_start is not None and out.text == think_start:
+            in_thinking = True
+        elif think_end is not None and out.text == think_end:
+            in_thinking = False
+        if in_thinking:
+            reasoning_tokens += 1
+
+        # Check for stop sequences
+        text = out.text
+        finish_reason: FinishReason | None = cast(
+            FinishReason | None, out.finish_reason
+        )
+        stop_matched = False
+
+        if stop_sequences:
+            for stop_seq in stop_sequences:
+                if stop_seq in accumulated_text:
+                    # Trim text to just before the stop sequence
+                    stop_index = accumulated_text.find(stop_seq)
+                    text_before_stop = accumulated_text[:stop_index]
+                    chunk_start = len(accumulated_text) - len(out.text)
+                    text = text_before_stop[chunk_start:]
+                    finish_reason = "stop"
+                    stop_matched = True
+                    break
+
+        is_done = finish_reason is not None
+
+        stats: GenerationStats | None = None
+        if is_done:
+            stats = GenerationStats(
+                prompt_tps=float(prefill_tps or out.prompt_tps),
+                generation_tps=float(out.generation_tps),
+                prompt_tokens=int(prefill_tokens + out.prompt_tokens),
+                generation_tokens=int(out.generation_tokens),
+                peak_memory_usage=Memory.from_gb(out.peak_memory),
+            )
+            if not stop_matched and out.finish_reason not in get_args(FinishReason):
+                logger.warning(
+                    f"Model generated unexpected finish_reason: {out.finish_reason}"
+                )
+
+            total_prompt_tokens = len(all_prompt_tokens)
+            usage = Usage(
+                prompt_tokens=total_prompt_tokens,
+                completion_tokens=completion_tokens,
+                total_tokens=total_prompt_tokens + completion_tokens,
+                prompt_tokens_details=PromptTokensDetails(
+                    cached_tokens=prefix_hit_length
+                ),
+                completion_tokens_details=CompletionTokensDetails(
+                    reasoning_tokens=reasoning_tokens
+                ),
+            )
+
+        # Extract logprobs from the full vocabulary logprobs array
+        logprob: float | None = None
+        top_logprobs: list[TopLogprobItem] | None = None
+        if task.logprobs:
+            logprob, top_logprobs = extract_top_logprobs(
+                logprobs=out.logprobs,
+                tokenizer=tokenizer,
+                top_logprobs=task.top_logprobs or DEFAULT_TOP_LOGPROBS,
+                selected_token=out.token,
+            )
+
+        if is_done:
+            # Log generation stats
+            generation_elapsed = time.perf_counter() - generation_start_time
+            generated_tokens = len(generated_text_parts)
+            generation_tps = (
+                generated_tokens / generation_elapsed if generation_elapsed > 0 else 0.0
+            )
+            logger.debug(
+                f"Generation complete: prefill {prompt_tokens} tokens @ "
+                f"{prefill_tps:.1f} tok/s, generated {generated_tokens} tokens @ "
+                f"{generation_tps:.1f} tok/s"
+            )
+            if kv_prefix_cache is not None:
+                generated_tokens_array = mx.array(
+                    tokenizer.encode(
+                        "".join(generated_text_parts), add_special_tokens=False
+                    )
+                )
+                full_prompt_tokens = mx.concatenate(
+                    [all_prompt_tokens, generated_tokens_array]
+                )
+                hit_ratio = (
+                    prefix_hit_length / len(all_prompt_tokens)
+                    if len(all_prompt_tokens) > 0
+                    else 0.0
+                )
+                if (
+                    matched_index is not None
+                    and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
+                ):
+                    kv_prefix_cache.update_kv_cache(
+                        matched_index,
+                        full_prompt_tokens,
+                        caches,
+                        cache_snapshots,
+                        restore_pos=prefix_hit_length,
+                    )
+                else:
+                    kv_prefix_cache.add_kv_cache(
+                        full_prompt_tokens, caches, cache_snapshots
+                    )
+
+        if on_generation_token is not None:
+            on_generation_token()
+
+        yield GenerationResponse(
+            text=text,
+            token=out.token,
+            logprob=logprob,
+            top_logprobs=top_logprobs,
+            finish_reason=finish_reason,
+            stats=stats,
+            usage=usage,
+        )
+
+        if is_done:
+            mx_barrier(group)
+            break
+
+        # Limit accumulated_text to what's needed for stop sequence detection
+        if max_stop_len > 0 and len(accumulated_text) > max_stop_len:
+            accumulated_text = accumulated_text[-max_stop_len:]

python/mlx_engine/src/mlx_engine/tests/test_batch_generate.py

@@ -0,0 +1,268 @@
+# pyright: reportAny=false, reportUnknownVariableType=false
+# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false
+# pyright: reportUnknownLambdaType=false, reportPrivateUsage=false
+# pyright: reportInvalidCast=false, reportArgumentType=false
+# pyright: reportUnusedImport=false
+"""Test B=1 vs B=2 equivalence for batch generation.
+
+Verifies that running two requests concurrently in a batch (B=2) produces
+identical token selections to running them sequentially (B=1).
+Uses random weights — no model download required.
+"""
+
+from pathlib import Path
+from typing import cast
+
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.utils
+
+# Import batch_generate to activate the right-padding BatchKVCache patch
+import mlx_engine.generator.batch_generate  # noqa: F401
+import pytest
+from mlx_engine.cache import encode_prompt, make_kv_cache
+from mlx_engine.generator.generate import prefill
+from mlx_engine.types import Model
+from mlx_lm.generate import _merge_caches
+from mlx_lm.sample_utils import make_sampler
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+from transformers import AutoTokenizer
+
+NUM_STEPS = 20
+
+
+def _init_random(model: nn.Module) -> None:
+    """Initialize all model parameters with random values."""
+    params = model.parameters()
+    new_params = mlx.utils.tree_map(
+        lambda p: mx.random.normal(shape=p.shape, dtype=p.dtype)
+        if isinstance(p, mx.array)
+        else p,
+        params,
+    )
+    model.update(new_params)
+    mx.eval(model.parameters())
+
+
+def _run_b1_vs_b2(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    tokens_a: mx.array,
+    tokens_b: mx.array,
+) -> tuple[float, int]:
+    """Run B=1 sequential and B=2 batched, return (max_diff, mismatches)."""
+    sampler = make_sampler(temp=0.0)
+
+    # B=1 sequential
+    cache_a1 = make_kv_cache(model)
+    prefill(model, tokenizer, sampler, tokens_a[:-1], cache_a1, None, None, None)
+    merged_a1 = _merge_caches([[c for c in cache_a1]])
+    for c in merged_a1:
+        c.prepare(lengths=[1], right_padding=[0])
+    model(mx.array([[tokens_a[-2].item()]]), cache=merged_a1)
+    mx.eval([c.state for c in merged_a1])
+    for c in merged_a1:
+        c.finalize()
+
+    cache_b1 = make_kv_cache(model)
+    prefill(model, tokenizer, sampler, tokens_b[:-1], cache_b1, None, None, None)
+    merged_b1 = _merge_caches([[c for c in cache_b1]])
+    for c in merged_b1:
+        c.prepare(lengths=[1], right_padding=[0])
+    model(mx.array([[tokens_b[-2].item()]]), cache=merged_b1)
+    mx.eval([c.state for c in merged_b1])
+    for c in merged_b1:
+        c.finalize()
+
+    b1_logits_a: list[mx.array] = []
+    b1_logits_b: list[mx.array] = []
+    next_a, next_b = tokens_a[-1].item(), tokens_b[-1].item()
+    for _ in range(NUM_STEPS):
+        la = model(mx.array([[next_a]]), cache=merged_a1)
+        mx.eval(la)
+        b1_logits_a.append(la[0, -1])
+        next_a = int(mx.argmax(la[0, -1]).item())
+        lb = model(mx.array([[next_b]]), cache=merged_b1)
+        mx.eval(lb)
+        b1_logits_b.append(lb[0, -1])
+        next_b = int(mx.argmax(lb[0, -1]).item())
+
+    # B=2 batched
+    cache_a2 = make_kv_cache(model)
+    cache_b2 = make_kv_cache(model)
+    prefill(model, tokenizer, sampler, tokens_a[:-1], cache_a2, None, None, None)
+    prefill(model, tokenizer, sampler, tokens_b[:-1], cache_b2, None, None, None)
+    merged_b2 = _merge_caches([list(cache_a2), list(cache_b2)])
+    for c in merged_b2:
+        c.prepare(lengths=[1, 1], right_padding=[0, 0])
+    model(
+        mx.array([[tokens_a[-2].item()], [tokens_b[-2].item()]]),
+        cache=merged_b2,
+    )
+    mx.eval([c.state for c in merged_b2])
+    for c in merged_b2:
+        c.finalize()
+
+    b2_logits_a: list[mx.array] = []
+    b2_logits_b: list[mx.array] = []
+    next_a2, next_b2 = tokens_a[-1].item(), tokens_b[-1].item()
+    for _ in range(NUM_STEPS):
+        l2 = model(mx.array([[next_a2], [next_b2]]), cache=merged_b2)
+        mx.eval(l2)
+        b2_logits_a.append(l2[0, -1])
+        b2_logits_b.append(l2[1, -1])
+        next_a2 = int(mx.argmax(l2[0, -1]).item())
+        next_b2 = int(mx.argmax(l2[1, -1]).item())
+
+    # Compare
+    max_diff = 0.0
+    mismatches = 0
+    for step in range(NUM_STEPS):
+        diff_a = float(
+            mx.max(
+                mx.abs(
+                    b1_logits_a[step].astype(mx.float32)
+                    - b2_logits_a[step].astype(mx.float32)
+                )
+            ).item()
+        )
+        diff_b = float(
+            mx.max(
+                mx.abs(
+                    b1_logits_b[step].astype(mx.float32)
+                    - b2_logits_b[step].astype(mx.float32)
+                )
+            ).item()
+        )
+        max_diff = max(max_diff, diff_a, diff_b)
+        if int(mx.argmax(b1_logits_a[step]).item()) != int(
+            mx.argmax(b2_logits_a[step]).item()
+        ):
+            mismatches += 1
+        if int(mx.argmax(b1_logits_b[step]).item()) != int(
+            mx.argmax(b2_logits_b[step]).item()
+        ):
+            mismatches += 1
+
+    return max_diff, mismatches
+
+
+def _make_tokenizer() -> TokenizerWrapper:
+    """Load the Qwen tokenizer (tiny download, shared across Qwen models)."""
+    from huggingface_hub import snapshot_download
+
+    model_path = Path(
+        snapshot_download(
+            "mlx-community/Qwen3.5-35B-A3B-4bit",
+            allow_patterns=["tokenizer*", "*.jinja"],
+        )
+    )
+    hf_tokenizer = AutoTokenizer.from_pretrained(model_path)
+    return TokenizerWrapper(hf_tokenizer)
+
+
+@pytest.mark.slow
+def test_batch_b2_llama() -> None:
+    """Llama-style model (KVCache only) must produce bit-exact logits in B=2.
+
+    Right-padded BatchKVCache keeps data at position 0 for all sequences,
+    so flash attention sees identical data layout as B=1 → bit-exact output.
+    """
+    from mlx_lm.models.llama import Model as LlamaModel
+    from mlx_lm.models.llama import ModelArgs
+
+    mx.random.seed(42)
+    args = ModelArgs(
+        model_type="llama",
+        hidden_size=256,
+        num_hidden_layers=4,
+        intermediate_size=512,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        rms_norm_eps=1e-6,
+        vocab_size=248320,
+        rope_theta=10000.0,
+        tie_word_embeddings=True,
+    )
+    model = LlamaModel(args)
+    _init_random(model)
+
+    tokenizer = _make_tokenizer()
+    tokens_a = encode_prompt(tokenizer, "Write a short essay about AI.")
+    tokens_b = encode_prompt(tokenizer, "Explain evolution briefly.")
+
+    max_diff, mismatches = _run_b1_vs_b2(
+        cast(Model, model), tokenizer, tokens_a, tokens_b
+    )
+    assert mismatches == 0, f"Llama B=2 token mismatches: {mismatches}/{NUM_STEPS * 2}"
+    assert max_diff < 0.002, f"Llama B=2 max logit diff: {max_diff}"
+
+
+@pytest.mark.slow
+def test_batch_b2_qwen35_moe() -> None:
+    """Qwen3.5 MoE model (hybrid SSM+attention+MoE) must produce bit-exact logits in B=2.
+
+    Right-padded BatchKVCache keeps data at position 0 for all sequences,
+    so flash attention sees identical data layout as B=1 → bit-exact output.
+    """
+    from mlx_lm.models.qwen3_5_moe import Model as Qwen35MoeModel
+    from mlx_lm.models.qwen3_5_moe import ModelArgs
+
+    mx.random.seed(42)
+    config = {
+        "model_type": "qwen3_5_moe",
+        "text_config": {
+            "model_type": "qwen3_5_moe_text",
+            "hidden_size": 256,
+            "num_hidden_layers": 8,
+            "intermediate_size": 512,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 2,
+            "rms_norm_eps": 1e-6,
+            "vocab_size": 248320,
+            "head_dim": 64,
+            "max_position_embeddings": 4096,
+            "full_attention_interval": 4,
+            "layer_types": [
+                "linear_attention",
+                "linear_attention",
+                "linear_attention",
+                "full_attention",
+                "linear_attention",
+                "linear_attention",
+                "linear_attention",
+                "full_attention",
+            ],
+            "linear_conv_kernel_dim": 4,
+            "linear_key_head_dim": 64,
+            "linear_num_key_heads": 4,
+            "linear_num_value_heads": 4,
+            "linear_value_head_dim": 64,
+            "mamba_ssm_dtype": "float32",
+            "num_experts": 8,
+            "num_experts_per_tok": 2,
+            "moe_intermediate_size": 256,
+            "shared_expert_intermediate_size": 256,
+            "rope_parameters": {
+                "rope_type": "default",
+                "rope_theta": 10000000,
+            },
+            "attention_bias": False,
+            "attn_output_gate": True,
+        },
+    }
+    args = ModelArgs.from_dict(config)
+    model = Qwen35MoeModel(args)
+    _init_random(model)
+
+    tokenizer = _make_tokenizer()
+    tokens_a = encode_prompt(tokenizer, "Write a short essay about AI.")
+    tokens_b = encode_prompt(tokenizer, "Explain evolution briefly.")
+
+    max_diff, mismatches = _run_b1_vs_b2(
+        cast(Model, model), tokenizer, tokens_a, tokens_b
+    )
+    assert mismatches == 0, (
+        f"Qwen3.5 MoE B=2 token mismatches: {mismatches}/{NUM_STEPS * 2}"
+    )
+    assert max_diff < 0.002, f"Qwen3.5 MoE B=2 max logit diff: {max_diff}"

python/mlx_engine/src/mlx_engine/types.py

@@ -0,0 +1,33 @@
+"""Shared types for MLX-related functionality."""
+
+from collections.abc import Sequence
+
+from mlx import core as mx
+from mlx import nn as nn
+from mlx_lm.models.cache import (
+    ArraysCache,
+    CacheList,
+    KVCache,
+    QuantizedKVCache,
+    RotatingKVCache,
+)
+from vllm_engine.kv_cache import TorchKVCache
+
+MLXCacheType = Sequence[
+    KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList
+]
+
+KVCacheType = MLXCacheType | TorchKVCache
+
+
+# Model is a wrapper function to fix the fact that mlx is not strongly typed in the same way that EXO is.
+# For example - MLX has no guarantee of the interface that nn.Module will expose. But we need a guarantee that it has a __call__() function
+class Model(nn.Module):
+    layers: list[nn.Module]
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache: MLXCacheType | None,
+        input_embeddings: mx.array | None = None,
+    ) -> mx.array: ...

python/mlx_engine/src/mlx_engine/utils_mlx.py

@@ -0,0 +1,810 @@
+import json
+import os
+import re
+import sys
+import tempfile
+import time
+from pathlib import Path
+from typing import Any, cast
+
+# Monkey-patch for transformers 5.x compatibility
+# Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location
+# which was moved in transformers 5.0.0rc2
+try:
+    import transformers.models.gpt2.tokenization_gpt2 as gpt2_tokenization
+    from transformers.convert_slow_tokenizer import bytes_to_unicode
+
+    if not hasattr(gpt2_tokenization, "bytes_to_unicode"):
+        gpt2_tokenization.bytes_to_unicode = bytes_to_unicode  # type: ignore[attr-defined]
+except ImportError:
+    pass  # transformers < 5.0 or bytes_to_unicode not available
+
+from mlx_lm.models.cache import KVCache
+from mlx_lm.models.deepseek_v3 import DeepseekV3Model
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+from mlx_engine.constants import TRUST_REMOTE_CODE
+
+try:
+    from mlx_lm.tokenizer_utils import load_tokenizer
+except ImportError:
+    from mlx_lm.tokenizer_utils import load as load_tokenizer
+import contextlib
+
+import mlx.core as mx
+import mlx.nn as nn
+from exo_core.types.common import Host, ModelId
+from exo_core.types.instances import (
+    BoundInstance,
+    MlxJacclInstance,
+    MlxRingInstance,
+)
+from exo_core.types.shards import (
+    CfgShardMetadata,
+    PipelineShardMetadata,
+    ShardMetadata,
+    TensorShardMetadata,
+)
+from exo_core.types.tasks import TaskId, TextGeneration
+from exo_core.types.text_generation import TextGenerationTaskParams
+from exo_core.utils.downloads import build_model_path
+from exo_core.utils.memory import Memory
+from loguru import logger
+from mlx_lm.utils import load_model
+from pydantic import RootModel
+
+from mlx_engine.auto_parallel import (
+    LayerLoadedCallback,
+    TimeoutCallback,
+    eval_with_timeout,
+    get_inner_model,
+    get_layers,
+    pipeline_auto_parallel,
+    tensor_auto_parallel,
+)
+from mlx_engine.types import Model
+
+Group = mx.distributed.Group
+
+
+def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
+    return Memory.from_float_kb(
+        (model_shard_meta.end_layer - model_shard_meta.start_layer)
+        / model_shard_meta.n_layers
+        * model_shard_meta.model_card.storage_size.in_kb
+        / (
+            1
+            if isinstance(model_shard_meta, PipelineShardMetadata)
+            else model_shard_meta.world_size
+        )
+    )
+
+
+class ModelLoadingTimeoutError(Exception):
+    pass
+
+
+class HostList(RootModel[list[str]]):
+    @classmethod
+    def from_hosts(cls, hosts: list[Host]) -> "HostList":
+        return cls(root=[str(host) for host in hosts])
+
+
+def mlx_distributed_init(
+    bound_instance: BoundInstance,
+) -> Group:
+    """
+    Initialize MLX distributed.
+    """
+    rank = bound_instance.bound_shard.device_rank
+    logger.info(f"Starting initialization for rank {rank}")
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        coordination_file = str(
+            Path(tmpdir) / f"hosts_{bound_instance.instance.instance_id}_{rank}.json"
+        )
+        # TODO: singleton instances
+        match bound_instance.instance:
+            case MlxRingInstance(hosts_by_node=hosts_by_node, ephemeral_port=_):
+                hosts_for_node = hosts_by_node[bound_instance.bound_node_id]
+                hosts_json = HostList.from_hosts(hosts_for_node).model_dump_json()
+
+                with open(coordination_file, "w") as f:
+                    _ = f.write(hosts_json)
+
+                logger.info(
+                    f"rank {rank} hostfile: {coordination_file} hosts: {hosts_json}"
+                )
+
+                os.environ["MLX_HOSTFILE"] = coordination_file
+                os.environ["MLX_RANK"] = str(rank)
+                os.environ["MLX_RING_VERBOSE"] = "1"
+                group = mx.distributed.init(backend="ring", strict=True)
+
+            case MlxJacclInstance(
+                jaccl_devices=jaccl_devices, jaccl_coordinators=jaccl_coordinators
+            ):
+                assert all(
+                    jaccl_devices[i][i] is None for i in range(len(jaccl_devices))
+                )
+                # Use RDMA connectivity matrix
+                jaccl_devices_json = json.dumps(jaccl_devices)
+
+                with open(coordination_file, "w") as f:
+                    _ = f.write(jaccl_devices_json)
+
+                jaccl_coordinator = jaccl_coordinators[bound_instance.bound_node_id]
+
+                logger.info(
+                    f"rank {rank} MLX_IBV_DEVICES: {coordination_file} with devices: {jaccl_devices_json}"
+                )
+                logger.info(f"rank {rank} MLX_JACCL_COORDINATOR: {jaccl_coordinator}")
+                os.environ["MLX_IBV_DEVICES"] = coordination_file
+                os.environ["MLX_RANK"] = str(rank)
+                os.environ["MLX_JACCL_COORDINATOR"] = jaccl_coordinator
+                group = mx.distributed.init(backend="jaccl", strict=True)
+
+            case _:
+                raise ValueError(
+                    f"Unsupported instance type for MLX init: {type(bound_instance.instance)}"
+                )
+
+        logger.info(f"Rank {rank} mlx distributed initialization complete")
+
+        return group
+
+
+def initialize_mlx(
+    bound_instance: BoundInstance,
+) -> Group:
+    # should we unseed it?
+    # TODO: pass in seed from params
+    mx.random.seed(42)
+
+    assert len(bound_instance.instance.shard_assignments.node_to_runner) > 1, (
+        "Tried to initialize mlx for a single node instance"
+    )
+    return mlx_distributed_init(bound_instance)
+
+
+def load_mlx_items(
+    bound_instance: BoundInstance,
+    group: Group | None,
+    on_timeout: TimeoutCallback | None,
+    on_layer_loaded: LayerLoadedCallback | None,
+) -> tuple[Model, TokenizerWrapper]:
+    if group is None:
+        logger.info(f"Single device used for {bound_instance.instance}")
+        model_path = build_model_path(bound_instance.bound_shard.model_card.model_id)
+        start_time = time.perf_counter()
+        model, _ = load_model(model_path, lazy=True, strict=False)
+        # Eval layers one by one for progress reporting
+        try:
+            inner = get_inner_model(model)
+            layers = get_layers(inner)
+            total = len(layers)
+            for i, layer in enumerate(layers):
+                mx.eval(layer)  # type: ignore
+                if on_layer_loaded is not None:
+                    on_layer_loaded(i, total)
+        except ValueError as e:
+            logger.opt(exception=e).debug(
+                "Model architecture doesn't support layer-by-layer progress tracking",
+            )
+        mx.eval(model)
+        end_time = time.perf_counter()
+        logger.info(f"Time taken to load model: {(end_time - start_time):.2f}s")
+        tokenizer = get_tokenizer(model_path, bound_instance.bound_shard)
+
+    else:
+        logger.info("Starting distributed init")
+        start_time = time.perf_counter()
+        model, tokenizer = shard_and_load(
+            bound_instance.bound_shard,
+            group=group,
+            on_timeout=on_timeout,
+            on_layer_loaded=on_layer_loaded,
+        )
+        end_time = time.perf_counter()
+        logger.info(
+            f"Time taken to shard and load model: {(end_time - start_time):.2f}s"
+        )
+
+    set_wired_limit_for_model(get_weights_size(bound_instance.bound_shard))
+
+    mx.clear_cache()
+
+    return cast(Model, model), tokenizer
+
+
+def shard_and_load(
+    shard_metadata: ShardMetadata,
+    group: Group,
+    on_timeout: TimeoutCallback | None,
+    on_layer_loaded: LayerLoadedCallback | None,
+) -> tuple[nn.Module, TokenizerWrapper]:
+    model_path = build_model_path(shard_metadata.model_card.model_id)
+
+    model, _ = load_model(model_path, lazy=True, strict=False)
+    logger.debug(model)
+    if hasattr(model, "model") and isinstance(model.model, DeepseekV3Model):  # type: ignore
+        pass
+        # TODO: See if we should quantize the model.
+        # def is_attention_layer(path: str) -> bool:
+        #     path = path.lower()
+
+        #     return "self_attn" in path and "layernorm" not in path
+
+        # def quant_predicate(path: str, module: nn.Module):
+        #     if not isinstance(module, nn.Linear):
+        #         return False
+
+        #     return is_attention_layer(path)
+        # model, config = quantize_model(
+        #        model, config, group_size=KV_GROUP_SIZE, bits=ATTENTION_KV_BITS, quant_predicate=quant_predicate, mode=QUANTIZE_MODEL_MODE
+        #    )
+
+    assert isinstance(model, nn.Module)
+
+    tokenizer = get_tokenizer(model_path, shard_metadata)
+
+    logger.info(f"Group size: {group.size()}, group rank: {group.rank()}")
+
+    # Estimate timeout based on model size (5x default for large queued workloads)
+    base_timeout = float(os.environ.get("EXO_MODEL_LOAD_TIMEOUT", "300"))
+    model_size = get_weights_size(shard_metadata)
+    timeout_seconds = base_timeout + model_size.in_gb
+    logger.info(
+        f"Evaluating model parameters with timeout of {timeout_seconds:.0f}s "
+        f"(model size: {model_size.in_gb:.1f}GB)"
+    )
+
+    match shard_metadata:
+        case TensorShardMetadata():
+            logger.info(f"loading model from {model_path} with tensor parallelism")
+            model = tensor_auto_parallel(
+                model, group, timeout_seconds, on_timeout, on_layer_loaded
+            )
+        case PipelineShardMetadata():
+            logger.info(f"loading model from {model_path} with pipeline parallelism")
+            model = pipeline_auto_parallel(
+                model, group, shard_metadata, on_layer_loaded=on_layer_loaded
+            )
+            eval_with_timeout(model.parameters(), timeout_seconds, on_timeout)
+        case CfgShardMetadata():
+            raise ValueError(
+                "CfgShardMetadata is not supported for text model loading - "
+                "this metadata type is only for image generation models"
+            )
+
+    # TODO: Do we need this?
+    mx.eval(model)
+
+    logger.debug("SHARDED")
+    logger.debug(model)
+
+    # Synchronize processes before generation to avoid timeout
+    mx_barrier(group)
+
+    return model, tokenizer
+
+
+def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata) -> TokenizerWrapper:
+    """Load tokenizer for a model shard. Delegates to load_tokenizer_for_model_id."""
+    return load_tokenizer_for_model_id(
+        shard_metadata.model_card.model_id,
+        model_path,
+        trust_remote_code=shard_metadata.model_card.trust_remote_code,
+    )
+
+
+def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
+    """
+    Get the EOS token IDs for a model based on its ID.
+
+    Some models require explicit EOS token configuration that isn't in their
+    tokenizer config. This function returns the known EOS token IDs for such models.
+
+    Args:
+        model_id: The HuggingFace model ID
+
+    Returns:
+        List of EOS token IDs, or None if the model uses standard tokenizer config
+    """
+    model_id_lower = model_id.lower()
+    if "kimi-k2" in model_id_lower:
+        return [163586]
+    elif "glm-5" in model_id_lower or "glm-4.7" in model_id_lower:
+        # For GLM-5 and GLM-4.7
+        # 154820: <|endoftext|>, 154827: <|user|>, 154829: <|observation|>
+        return [154820, 154827, 154829]
+    elif "glm" in model_id_lower:
+        # For GLM-4.5 and older
+        return [151336, 151329, 151338]
+    elif "gpt-oss" in model_id_lower:
+        return [200002, 200012]
+    elif "qwen3.5" in model_id_lower or "qwen-3.5" in model_id_lower:
+        # For Qwen3.5: 248046 (<|im_end|>), 248044 (<|endoftext|>)
+        return [248046, 248044]
+    return None
+
+
+def load_tokenizer_for_model_id(
+    model_id: ModelId, model_path: Path, *, trust_remote_code: bool = TRUST_REMOTE_CODE
+) -> TokenizerWrapper:
+    """
+    Load tokenizer for a model given its ID and local path.
+
+    This is the core tokenizer loading logic, handling special cases for different
+    model families (Kimi, GLM, etc.) and transformers 5.x compatibility.
+
+    Args:
+        model_id: The HuggingFace model ID (e.g., "moonshotai/Kimi-K2-Instruct")
+        model_path: Local path where the model/tokenizer files are stored
+
+    Returns:
+        TokenizerWrapper instance configured for the model
+    """
+    model_id_lower = model_id.lower()
+    eos_token_ids = get_eos_token_ids_for_model(model_id)
+
+    # Kimi uses a custom TikTokenTokenizer that transformers 5.x can't load via AutoTokenizer
+    if "kimi-k2" in model_id_lower:
+        import importlib.util
+        import types
+
+        sys.path.insert(0, str(model_path))
+
+        # Load tool_declaration_ts first (tokenization_kimi imports it with relative import)
+        tool_decl_path = model_path / "tool_declaration_ts.py"
+        if tool_decl_path.exists():
+            spec = importlib.util.spec_from_file_location(
+                "tool_declaration_ts", tool_decl_path
+            )
+            if spec and spec.loader:
+                tool_decl_module = importlib.util.module_from_spec(spec)
+                sys.modules["tool_declaration_ts"] = tool_decl_module
+                spec.loader.exec_module(tool_decl_module)
+
+        # Load tokenization_kimi with patched source (convert relative to absolute import)
+        tok_path = model_path / "tokenization_kimi.py"
+        source = tok_path.read_text()
+        source = source.replace("from .tool_declaration_ts", "from tool_declaration_ts")
+        spec = importlib.util.spec_from_file_location("tokenization_kimi", tok_path)
+        if spec:
+            tok_module = types.ModuleType("tokenization_kimi")
+            tok_module.__file__ = str(tok_path)
+            sys.modules["tokenization_kimi"] = tok_module
+            exec(compile(source, tok_path, "exec"), tok_module.__dict__)  # noqa: S102
+            TikTokenTokenizer = tok_module.TikTokenTokenizer  # type: ignore[attr-defined]  # noqa: N806
+        else:
+            from tokenization_kimi import TikTokenTokenizer  # type: ignore[import-not-found]  # noqa: I001
+
+        hf_tokenizer: Any = TikTokenTokenizer.from_pretrained(model_path)  # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType]
+
+        # Patch encode to use internal tiktoken model directly
+        # transformers 5.x has a bug in the encode->pad path for slow tokenizers
+        def _patched_encode(text: str, **_kwargs: object) -> list[int]:
+            # Pass allowed_special="all" to handle special tokens like <|im_user|>
+            return list(hf_tokenizer.model.encode(text, allowed_special="all"))  # pyright: ignore[reportUnknownMemberType,reportUnknownArgumentType]
+
+        hf_tokenizer.encode = _patched_encode
+        return TokenizerWrapper(
+            hf_tokenizer,
+            eos_token_ids=eos_token_ids,
+            tool_call_start="<|tool_calls_section_begin|>",
+            tool_call_end="<|tool_calls_section_end|>",
+            tool_parser=_parse_kimi_tool_calls,
+        )
+
+    tokenizer = load_tokenizer(
+        model_path,
+        tokenizer_config_extra={"trust_remote_code": trust_remote_code},
+        eos_token_ids=eos_token_ids,
+    )
+
+    if "gemma-3" in model_id_lower:
+        gemma_3_eos_id = 1
+        gemma_3_end_of_turn_id = 106
+        if tokenizer.eos_token_ids is not None:
+            if gemma_3_end_of_turn_id not in tokenizer.eos_token_ids:
+                tokenizer.eos_token_ids = list(tokenizer.eos_token_ids) + [
+                    gemma_3_end_of_turn_id
+                ]
+        else:
+            tokenizer.eos_token_ids = [gemma_3_eos_id, gemma_3_end_of_turn_id]
+
+    return tokenizer
+
+
+def _normalize_tool_calls(msg_dict: dict[str, Any]) -> None:
+    """Normalize tool_calls in a message dict.
+
+    OpenAI format has tool_calls[].function.arguments as a JSON string,
+    but some chat templates (e.g., GLM) expect it as a dict.
+    """
+    tool_calls = msg_dict.get("tool_calls")
+    if not tool_calls or not isinstance(tool_calls, list):
+        return
+
+    for tc in tool_calls:  # pyright: ignore[reportUnknownVariableType]
+        if not isinstance(tc, dict):
+            continue
+        func = tc.get("function")  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+        if not isinstance(func, dict):
+            continue
+        args = func.get("arguments")  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+        if isinstance(args, str):
+            with contextlib.suppress(json.JSONDecodeError):
+                func["arguments"] = json.loads(args)
+
+
+def _collect_nested_property_names(schema: dict[str, Any]) -> set[str]:
+    names: set[str] = set()
+    properties: dict[str, Any] = schema.get("properties", {})  # type: ignore[reportAny]
+    for prop_spec in properties.values():  # pyright: ignore[reportAny]
+        if not isinstance(prop_spec, dict):
+            continue
+        if prop_spec.get("type") == "array":  # type: ignore[reportAny]
+            items: dict[str, Any] | None = prop_spec.get("items")  # type: ignore[reportAny]
+            if isinstance(items, dict) and items.get("type") == "object":  # type: ignore[reportAny]
+                inner_props: dict[str, Any] = items.get("properties", {})  # type: ignore[reportAny]
+                for k in inner_props:  # pyright: ignore[reportUnknownVariableType]
+                    names.add(str(k))  # pyright: ignore[reportUnknownArgumentType]
+                names.update(_collect_nested_property_names(items))  # pyright: ignore[reportUnknownArgumentType]
+    return names
+
+
+def _schemas_lost_in_prompt(prompt: str, tools: list[dict[str, Any]]) -> bool:
+    """Return True if nested property names from any tool schema are absent."""
+    for tool in tools:
+        fn: dict[str, Any] = tool.get("function", {})  # type: ignore
+        params: dict[str, Any] = fn.get("parameters", {})  # type: ignore
+        nested = _collect_nested_property_names(params)
+        if nested and not all(name in prompt for name in nested):
+            return True
+    return False
+
+
+_LOSSY_TEMPLATE_PATTERN = re.compile(
+    r"""inner_type\s*==\s*["']object \| object["']\s*or\s*inner_type\|length\s*>\s*\d+""",
+)
+
+
+def _patch_lossy_chat_template(template: str) -> str | None:
+    """Patch chat templates that collapse nested object schemas to ``any[]``.
+
+    Some templates (e.g., GPT-OSS) have a guard like::
+
+        inner_type == "object | object" or inner_type|length > 50
+
+    The length check silently drops complex array-of-object schemas.
+    We remove the length guard, keeping only the object-union check.
+    Returns the patched template, or *None* if no patch was needed.
+    """
+    patched, n = _LOSSY_TEMPLATE_PATTERN.subn(
+        lambda m: m.group(0).split(" or ")[0],  # keep only the object-union check
+        template,
+    )
+    return patched if n > 0 else None
+
+
+def _needs_dsml_encoding(task_params: TextGenerationTaskParams) -> bool:
+    if "deepseek-v3.2" not in task_params.model.lower():
+        return False
+    # Use DSML encoding when tools are provided or tool results are in the conversation
+    if task_params.tools:
+        return True
+    if task_params.chat_template_messages:
+        return any(
+            msg.get("role") == "tool" for msg in task_params.chat_template_messages
+        )
+    return False
+
+
+def apply_chat_template(
+    tokenizer: TokenizerWrapper,
+    task_params: TextGenerationTaskParams,
+) -> str:
+    """Convert TextGenerationTaskParams to a chat template prompt.
+
+    Converts the internal format (input + instructions) to a messages list
+    that can be processed by the tokenizer's chat template.
+
+    When chat_template_messages is available (from Chat Completions API),
+    uses those directly to preserve tool_calls, thinking, and other fields.
+    """
+    formatted_messages: list[dict[str, Any]] = []
+    if task_params.chat_template_messages is not None:
+        # Use pre-formatted messages that preserve tool_calls, thinking, etc.
+        formatted_messages = list(task_params.chat_template_messages)
+        for msg in formatted_messages:
+            _normalize_tool_calls(msg)
+    else:
+        # Add system message (instructions) if present
+        if task_params.instructions:
+            formatted_messages.append(
+                {"role": "system", "content": task_params.instructions}
+            )
+
+        # Convert input to messages
+        for msg in task_params.input:
+            if not msg.content:
+                logger.warning("Received message with empty content, skipping")
+                continue
+            formatted_messages.append({"role": msg.role, "content": msg.content})
+
+    # For assistant prefilling, append content after templating to avoid a closing turn token.
+    partial_assistant_content: str | None = None
+    if formatted_messages and formatted_messages[-1].get("role") == "assistant":
+        partial_assistant_content = cast(str, formatted_messages[-1].get("content", ""))
+        formatted_messages = formatted_messages[:-1]
+
+    if _needs_dsml_encoding(task_params):
+        from mlx_engine.dsml_encoding import encode_messages
+
+        prompt = encode_messages(
+            messages=formatted_messages,
+            thinking_mode="thinking" if task_params.enable_thinking else "chat",
+            tools=task_params.tools,
+        )
+        if partial_assistant_content:
+            prompt += partial_assistant_content
+        logger.info(prompt)
+        return prompt
+
+    extra_kwargs: dict[str, Any] = {}
+    if task_params.enable_thinking is not None:
+        # Qwen3 and GLM use "enable_thinking"; DeepSeek uses "thinking".
+        # Jinja ignores unknown variables, so passing both is safe.
+        extra_kwargs["enable_thinking"] = task_params.enable_thinking
+        extra_kwargs["thinking"] = task_params.enable_thinking
+    if task_params.reasoning_effort is not None:
+        extra_kwargs["reasoning_effort"] = task_params.reasoning_effort
+
+    patched_template: str | None = None
+    if task_params.tools:
+        original_template: str | None = getattr(tokenizer, "chat_template", None)
+        if isinstance(original_template, str):
+            patched_template = _patch_lossy_chat_template(original_template)
+            if patched_template is not None:
+                logger.info(
+                    "Patched lossy chat template (removed inner_type length guard)"
+                )
+
+    prompt: str = tokenizer.apply_chat_template(
+        formatted_messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        tools=task_params.tools,
+        **({"chat_template": patched_template} if patched_template is not None else {}),
+        **extra_kwargs,
+    )
+
+    if task_params.tools and _schemas_lost_in_prompt(prompt, task_params.tools):
+        logger.warning("Chat template lost nested tool schemas even after patching")
+
+    if partial_assistant_content:
+        prompt += partial_assistant_content
+
+    logger.info(prompt)
+
+    return prompt
+
+
+def detect_thinking_prompt_suffix(prompt: str, tokenizer: TokenizerWrapper) -> bool:
+    """
+    Detect if prompt ends with a thinking opening tag that should be
+    prepended to the output stream.
+    """
+    think_token = tokenizer.think_start
+
+    return think_token is not None and prompt.rstrip().endswith(think_token)
+
+
+def fix_unmatched_think_end_tokens(
+    tokens: mx.array, tokenizer: TokenizerWrapper
+) -> mx.array:
+    if not tokenizer.has_thinking:
+        return tokens
+    assert tokenizer.think_start_id
+    assert tokenizer.think_end_id
+    think_start_id: int = tokenizer.think_start_id
+    think_end_id: int = tokenizer.think_end_id
+    token_list: list[int] = cast(list[int], tokens.tolist())
+    result: list[int] = []
+    depth = 0
+    for token in token_list:
+        if token == think_start_id:
+            depth += 1
+        elif token == think_end_id:
+            if depth == 0:
+                result.append(think_start_id)
+            else:
+                depth -= 1
+        result.append(token)
+    return mx.array(result)
+
+
+class NullKVCache(KVCache):
+    """
+    A KVCache that pretends to exist but holds zero tokens.
+    It satisfies .state/.meta_state and never allocates real keys/values.
+    """
+
+    def __init__(self, dtype: mx.Dtype = mx.float16):
+        super().__init__()
+        # zero-length K/V so shapes/dtypes are defined but empty
+        self.keys = mx.zeros((1, 1, 0, 1), dtype=dtype)
+        self.values = mx.zeros((1, 1, 0, 1), dtype=dtype)
+        self.offset = 0
+
+    @property
+    def state(self) -> tuple[mx.array, mx.array]:
+        # matches what mx.save_safetensors / mx.eval expect
+        return self.keys, self.values
+
+    @state.setter
+    def state(self, v: tuple[mx.array, mx.array]) -> None:
+        raise NotImplementedError("We should not be setting a NullKVCache.")
+
+
+def mlx_force_oom(size: int = 200000) -> None:
+    """
+    Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations.
+    """
+    mx.set_default_device(mx.gpu)
+    a = mx.random.uniform(shape=(size, size), dtype=mx.float32)
+    b = mx.random.uniform(shape=(size, size), dtype=mx.float32)
+    mx.eval(a, b)
+    c = mx.matmul(a, b)
+    d = mx.matmul(a, c)
+    e = mx.matmul(b, c)
+    f = mx.sigmoid(d + e)
+    mx.eval(f)
+
+
+def set_wired_limit_for_model(model_size: Memory):
+    """
+    A context manager to temporarily change the wired limit.
+
+    Note, the wired limit should not be changed during an async eval.  If an
+    async eval could be running pass in the streams to synchronize with prior
+    to exiting the context manager.
+    """
+    if not mx.metal.is_available():
+        return
+
+    max_rec_size = Memory.from_bytes(
+        int(mx.device_info()["max_recommended_working_set_size"])
+    )
+    if model_size > 0.9 * max_rec_size:
+        logger.warning(
+            f"Generating with a model that requires {model_size.in_float_mb:.1f} MB "
+            f"which is close to the maximum recommended size of {max_rec_size.in_float_mb:.1f} "
+            "MB. This can be slow. See the documentation for possible work-arounds: "
+            "https://github.com/ml-explore/mlx-lm/tree/main#large-models"
+        )
+    mx.set_wired_limit(max_rec_size.in_bytes)
+    logger.info(f"Wired limit set to {max_rec_size}.")
+
+
+def mlx_cleanup(
+    model: Model | None, tokenizer: TokenizerWrapper | None, group: Group | None
+) -> None:
+    del model, tokenizer, group
+    mx.clear_cache()
+    import gc
+
+    gc.collect()
+
+
+def mx_any(bool_: bool, group: Group | None) -> bool:
+    if group is None:
+        return bool_
+    num_true = mx.distributed.all_sum(
+        mx.array(bool_), group=group, stream=mx.default_stream(mx.Device(mx.cpu))
+    )
+    mx.eval(num_true)
+    return num_true.item() > 0
+
+
+def mx_barrier(group: Group | None):
+    if group is None:
+        return
+    mx.eval(
+        mx.distributed.all_sum(
+            mx.array(1.0), group=group, stream=mx.default_stream(mx.Device(mx.cpu))
+        )
+    )
+
+
+def _parse_kimi_tool_calls(text: str):
+    import regex as re
+
+    # kimi has a fixed function naming scheme, with a json formatted arg
+    #   functions.multiply:0<|tool_call_argument_begin|>{"a": 2, "b": 3}
+    _func_name_regex = re.compile(
+        r"^\s*((?:functions\.)?(.+?):\d+)\s*<\|tool_call_argument_begin\|>", re.DOTALL
+    )
+    _func_arg_regex = re.compile(r"<\|tool_call_argument_begin\|>\s*(.*)\s*", re.DOTALL)
+    _tool_call_split_regex = re.compile(
+        r"<\|tool_call_begin\|>(.*?)<\|tool_call_end\|>", re.DOTALL
+    )
+
+    def _parse_single_tool(text: str) -> dict[str, Any]:
+        func_name_match = _func_name_regex.search(text)
+        if func_name_match is None:
+            raise ValueError("No tool call found.")
+        tool_call_id = func_name_match.group(1)  # e.g. "functions.get_weather:0"
+        func_name = func_name_match.group(2)  # e.g. "get_weather"
+
+        func_args_match = _func_arg_regex.search(text)
+        if func_args_match is None:
+            raise ValueError("No tool call arguments found.")
+        func_args = func_args_match.group(1)
+        try:
+            arg_dct = json.loads(func_args)  # pyright: ignore[reportAny]
+        except Exception:
+            arg_dct = None
+
+        return dict(id=tool_call_id, name=func_name, arguments=arg_dct)
+
+    tool_matches = _tool_call_split_regex.findall(text)
+    if tool_matches:
+        return [_parse_single_tool(match) for match in tool_matches]  # pyright: ignore[reportAny]
+    else:
+        return [_parse_single_tool(text)]
+
+
+def mx_all_gather_tasks(
+    tasks: list[TextGeneration],
+    group: mx.distributed.Group | None,
+) -> tuple[list[TextGeneration], list[TextGeneration]]:
+    def encode_task_id(task_id: TaskId) -> list[int]:
+        utf8_task_id = task_id.encode()
+        return [
+            int.from_bytes(utf8_task_id[i : i + 1]) for i in range(len(utf8_task_id))
+        ]
+
+    def decode_task_id(encoded_task_id: list[int]) -> TaskId:
+        return TaskId(
+            bytes.decode(b"".join((x).to_bytes(length=1) for x in encoded_task_id))
+        )
+
+    uuid_byte_length = 36
+
+    n_tasks = len(tasks)
+    all_counts = cast(
+        list[int],
+        mx.distributed.all_gather(mx.array([n_tasks]), group=group).tolist(),
+    )
+    max_tasks = max(all_counts)
+    world_size: int = 1 if group is None else group.size()
+
+    if max_tasks == 0:
+        return [], []
+
+    padded = [encode_task_id(task.task_id) for task in tasks] + [
+        [0] * uuid_byte_length
+    ] * (max_tasks - n_tasks)
+
+    assert all(len(encoded_task_id) == uuid_byte_length for encoded_task_id in padded)
+
+    gathered = cast(
+        list[list[list[int]]],
+        mx.distributed.all_gather(mx.array(padded), group=group)
+        .reshape(world_size, max_tasks, -1)
+        .tolist(),
+    )
+    all_task_ids: list[list[TaskId]] = [
+        [decode_task_id(encoded_task_id) for encoded_task_id in rank_tasks[:count]]
+        for rank_tasks, count in zip(gathered, all_counts, strict=True)
+    ]
+
+    agreed_ids = set[TaskId].intersection(*(set(tids) for tids in all_task_ids))
+
+    local_tasks = {task.task_id: task for task in tasks}
+    agreed = [local_tasks[tid] for tid in sorted(agreed_ids)]
+    different = [task for task in tasks if task.task_id not in agreed_ids]
+    return agreed, different

python/mlx_engine/uv.lock

@@ -1,123 +0,0 @@
-version = 1
-revision = 3
-requires-python = ">=3.13"
-
-[[package]]
-name = "annotated-types"
-version = "0.7.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/ee/67/531ea369ba64dcff5ec9c3402f9f51bf748cec26dde048a2f973a4eea7f5/annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89", size = 16081, upload-time = "2024-05-20T21:33:25.928Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" },
-]
-
-[[package]]
-name = "exo-core"
-version = "0.1.0"
-source = { directory = "../exo_core" }
-dependencies = [
-    { name = "pydantic" },
-]
-
-[package.metadata]
-requires-dist = [{ name = "pydantic" }]
-
-[[package]]
-name = "mlx-engine"
-version = "0.1.0"
-source = { editable = "." }
-dependencies = [
-    { name = "exo-core" },
-]
-
-[package.metadata]
-requires-dist = [{ name = "exo-core", directory = "../exo_core" }]
-
-[[package]]
-name = "pydantic"
-version = "2.12.5"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "annotated-types" },
-    { name = "pydantic-core" },
-    { name = "typing-extensions" },
-    { name = "typing-inspection" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/69/44/36f1a6e523abc58ae5f928898e4aca2e0ea509b5aa6f6f392a5d882be928/pydantic-2.12.5.tar.gz", hash = "sha256:4d351024c75c0f085a9febbb665ce8c0c6ec5d30e903bdb6394b7ede26aebb49", size = 821591, upload-time = "2025-11-26T15:11:46.471Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/5a/87/b70ad306ebb6f9b585f114d0ac2137d792b48be34d732d60e597c2f8465a/pydantic-2.12.5-py3-none-any.whl", hash = "sha256:e561593fccf61e8a20fc46dfc2dfe075b8be7d0188df33f221ad1f0139180f9d", size = 463580, upload-time = "2025-11-26T15:11:44.605Z" },
-]
-
-[[package]]
-name = "pydantic-core"
-version = "2.41.5"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "typing-extensions" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/71/70/23b021c950c2addd24ec408e9ab05d59b035b39d97cdc1130e1bce647bb6/pydantic_core-2.41.5.tar.gz", hash = "sha256:08daa51ea16ad373ffd5e7606252cc32f07bc72b28284b6bc9c6df804816476e", size = 460952, upload-time = "2025-11-04T13:43:49.098Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/87/06/8806241ff1f70d9939f9af039c6c35f2360cf16e93c2ca76f184e76b1564/pydantic_core-2.41.5-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:941103c9be18ac8daf7b7adca8228f8ed6bb7a1849020f643b3a14d15b1924d9", size = 2120403, upload-time = "2025-11-04T13:40:25.248Z" },
-    { url = "https://files.pythonhosted.org/packages/94/02/abfa0e0bda67faa65fef1c84971c7e45928e108fe24333c81f3bfe35d5f5/pydantic_core-2.41.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:112e305c3314f40c93998e567879e887a3160bb8689ef3d2c04b6cc62c33ac34", size = 1896206, upload-time = "2025-11-04T13:40:27.099Z" },
-    { url = "https://files.pythonhosted.org/packages/15/df/a4c740c0943e93e6500f9eb23f4ca7ec9bf71b19e608ae5b579678c8d02f/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0cbaad15cb0c90aa221d43c00e77bb33c93e8d36e0bf74760cd00e732d10a6a0", size = 1919307, upload-time = "2025-11-04T13:40:29.806Z" },
-    { url = "https://files.pythonhosted.org/packages/9a/e3/6324802931ae1d123528988e0e86587c2072ac2e5394b4bc2bc34b61ff6e/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:03ca43e12fab6023fc79d28ca6b39b05f794ad08ec2feccc59a339b02f2b3d33", size = 2063258, upload-time = "2025-11-04T13:40:33.544Z" },
-    { url = "https://files.pythonhosted.org/packages/c9/d4/2230d7151d4957dd79c3044ea26346c148c98fbf0ee6ebd41056f2d62ab5/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:dc799088c08fa04e43144b164feb0c13f9a0bc40503f8df3e9fde58a3c0c101e", size = 2214917, upload-time = "2025-11-04T13:40:35.479Z" },
-    { url = "https://files.pythonhosted.org/packages/e6/9f/eaac5df17a3672fef0081b6c1bb0b82b33ee89aa5cec0d7b05f52fd4a1fa/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:97aeba56665b4c3235a0e52b2c2f5ae9cd071b8a8310ad27bddb3f7fb30e9aa2", size = 2332186, upload-time = "2025-11-04T13:40:37.436Z" },
-    { url = "https://files.pythonhosted.org/packages/cf/4e/35a80cae583a37cf15604b44240e45c05e04e86f9cfd766623149297e971/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:406bf18d345822d6c21366031003612b9c77b3e29ffdb0f612367352aab7d586", size = 2073164, upload-time = "2025-11-04T13:40:40.289Z" },
-    { url = "https://files.pythonhosted.org/packages/bf/e3/f6e262673c6140dd3305d144d032f7bd5f7497d3871c1428521f19f9efa2/pydantic_core-2.41.5-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:b93590ae81f7010dbe380cdeab6f515902ebcbefe0b9327cc4804d74e93ae69d", size = 2179146, upload-time = "2025-11-04T13:40:42.809Z" },
-    { url = "https://files.pythonhosted.org/packages/75/c7/20bd7fc05f0c6ea2056a4565c6f36f8968c0924f19b7d97bbfea55780e73/pydantic_core-2.41.5-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:01a3d0ab748ee531f4ea6c3e48ad9dac84ddba4b0d82291f87248f2f9de8d740", size = 2137788, upload-time = "2025-11-04T13:40:44.752Z" },
-    { url = "https://files.pythonhosted.org/packages/3a/8d/34318ef985c45196e004bc46c6eab2eda437e744c124ef0dbe1ff2c9d06b/pydantic_core-2.41.5-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:6561e94ba9dacc9c61bce40e2d6bdc3bfaa0259d3ff36ace3b1e6901936d2e3e", size = 2340133, upload-time = "2025-11-04T13:40:46.66Z" },
-    { url = "https://files.pythonhosted.org/packages/9c/59/013626bf8c78a5a5d9350d12e7697d3d4de951a75565496abd40ccd46bee/pydantic_core-2.41.5-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:915c3d10f81bec3a74fbd4faebe8391013ba61e5a1a8d48c4455b923bdda7858", size = 2324852, upload-time = "2025-11-04T13:40:48.575Z" },
-    { url = "https://files.pythonhosted.org/packages/1a/d9/c248c103856f807ef70c18a4f986693a46a8ffe1602e5d361485da502d20/pydantic_core-2.41.5-cp313-cp313-win32.whl", hash = "sha256:650ae77860b45cfa6e2cdafc42618ceafab3a2d9a3811fcfbd3bbf8ac3c40d36", size = 1994679, upload-time = "2025-11-04T13:40:50.619Z" },
-    { url = "https://files.pythonhosted.org/packages/9e/8b/341991b158ddab181cff136acd2552c9f35bd30380422a639c0671e99a91/pydantic_core-2.41.5-cp313-cp313-win_amd64.whl", hash = "sha256:79ec52ec461e99e13791ec6508c722742ad745571f234ea6255bed38c6480f11", size = 2019766, upload-time = "2025-11-04T13:40:52.631Z" },
-    { url = "https://files.pythonhosted.org/packages/73/7d/f2f9db34af103bea3e09735bb40b021788a5e834c81eedb541991badf8f5/pydantic_core-2.41.5-cp313-cp313-win_arm64.whl", hash = "sha256:3f84d5c1b4ab906093bdc1ff10484838aca54ef08de4afa9de0f5f14d69639cd", size = 1981005, upload-time = "2025-11-04T13:40:54.734Z" },
-    { url = "https://files.pythonhosted.org/packages/ea/28/46b7c5c9635ae96ea0fbb779e271a38129df2550f763937659ee6c5dbc65/pydantic_core-2.41.5-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:3f37a19d7ebcdd20b96485056ba9e8b304e27d9904d233d7b1015db320e51f0a", size = 2119622, upload-time = "2025-11-04T13:40:56.68Z" },
-    { url = "https://files.pythonhosted.org/packages/74/1a/145646e5687e8d9a1e8d09acb278c8535ebe9e972e1f162ed338a622f193/pydantic_core-2.41.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1d1d9764366c73f996edd17abb6d9d7649a7eb690006ab6adbda117717099b14", size = 1891725, upload-time = "2025-11-04T13:40:58.807Z" },
-    { url = "https://files.pythonhosted.org/packages/23/04/e89c29e267b8060b40dca97bfc64a19b2a3cf99018167ea1677d96368273/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:25e1c2af0fce638d5f1988b686f3b3ea8cd7de5f244ca147c777769e798a9cd1", size = 1915040, upload-time = "2025-11-04T13:41:00.853Z" },
-    { url = "https://files.pythonhosted.org/packages/84/a3/15a82ac7bd97992a82257f777b3583d3e84bdb06ba6858f745daa2ec8a85/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:506d766a8727beef16b7adaeb8ee6217c64fc813646b424d0804d67c16eddb66", size = 2063691, upload-time = "2025-11-04T13:41:03.504Z" },
-    { url = "https://files.pythonhosted.org/packages/74/9b/0046701313c6ef08c0c1cf0e028c67c770a4e1275ca73131563c5f2a310a/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:4819fa52133c9aa3c387b3328f25c1facc356491e6135b459f1de698ff64d869", size = 2213897, upload-time = "2025-11-04T13:41:05.804Z" },
-    { url = "https://files.pythonhosted.org/packages/8a/cd/6bac76ecd1b27e75a95ca3a9a559c643b3afcd2dd62086d4b7a32a18b169/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2b761d210c9ea91feda40d25b4efe82a1707da2ef62901466a42492c028553a2", size = 2333302, upload-time = "2025-11-04T13:41:07.809Z" },
-    { url = "https://files.pythonhosted.org/packages/4c/d2/ef2074dc020dd6e109611a8be4449b98cd25e1b9b8a303c2f0fca2f2bcf7/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:22f0fb8c1c583a3b6f24df2470833b40207e907b90c928cc8d3594b76f874375", size = 2064877, upload-time = "2025-11-04T13:41:09.827Z" },
-    { url = "https://files.pythonhosted.org/packages/18/66/e9db17a9a763d72f03de903883c057b2592c09509ccfe468187f2a2eef29/pydantic_core-2.41.5-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:2782c870e99878c634505236d81e5443092fba820f0373997ff75f90f68cd553", size = 2180680, upload-time = "2025-11-04T13:41:12.379Z" },
-    { url = "https://files.pythonhosted.org/packages/d3/9e/3ce66cebb929f3ced22be85d4c2399b8e85b622db77dad36b73c5387f8f8/pydantic_core-2.41.5-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:0177272f88ab8312479336e1d777f6b124537d47f2123f89cb37e0accea97f90", size = 2138960, upload-time = "2025-11-04T13:41:14.627Z" },
-    { url = "https://files.pythonhosted.org/packages/a6/62/205a998f4327d2079326b01abee48e502ea739d174f0a89295c481a2272e/pydantic_core-2.41.5-cp314-cp314-musllinux_1_1_armv7l.whl", hash = "sha256:63510af5e38f8955b8ee5687740d6ebf7c2a0886d15a6d65c32814613681bc07", size = 2339102, upload-time = "2025-11-04T13:41:16.868Z" },
-    { url = "https://files.pythonhosted.org/packages/3c/0d/f05e79471e889d74d3d88f5bd20d0ed189ad94c2423d81ff8d0000aab4ff/pydantic_core-2.41.5-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:e56ba91f47764cc14f1daacd723e3e82d1a89d783f0f5afe9c364b8bb491ccdb", size = 2326039, upload-time = "2025-11-04T13:41:18.934Z" },
-    { url = "https://files.pythonhosted.org/packages/ec/e1/e08a6208bb100da7e0c4b288eed624a703f4d129bde2da475721a80cab32/pydantic_core-2.41.5-cp314-cp314-win32.whl", hash = "sha256:aec5cf2fd867b4ff45b9959f8b20ea3993fc93e63c7363fe6851424c8a7e7c23", size = 1995126, upload-time = "2025-11-04T13:41:21.418Z" },
-    { url = "https://files.pythonhosted.org/packages/48/5d/56ba7b24e9557f99c9237e29f5c09913c81eeb2f3217e40e922353668092/pydantic_core-2.41.5-cp314-cp314-win_amd64.whl", hash = "sha256:8e7c86f27c585ef37c35e56a96363ab8de4e549a95512445b85c96d3e2f7c1bf", size = 2015489, upload-time = "2025-11-04T13:41:24.076Z" },
-    { url = "https://files.pythonhosted.org/packages/4e/bb/f7a190991ec9e3e0ba22e4993d8755bbc4a32925c0b5b42775c03e8148f9/pydantic_core-2.41.5-cp314-cp314-win_arm64.whl", hash = "sha256:e672ba74fbc2dc8eea59fb6d4aed6845e6905fc2a8afe93175d94a83ba2a01a0", size = 1977288, upload-time = "2025-11-04T13:41:26.33Z" },
-    { url = "https://files.pythonhosted.org/packages/92/ed/77542d0c51538e32e15afe7899d79efce4b81eee631d99850edc2f5e9349/pydantic_core-2.41.5-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:8566def80554c3faa0e65ac30ab0932b9e3a5cd7f8323764303d468e5c37595a", size = 2120255, upload-time = "2025-11-04T13:41:28.569Z" },
-    { url = "https://files.pythonhosted.org/packages/bb/3d/6913dde84d5be21e284439676168b28d8bbba5600d838b9dca99de0fad71/pydantic_core-2.41.5-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:b80aa5095cd3109962a298ce14110ae16b8c1aece8b72f9dafe81cf597ad80b3", size = 1863760, upload-time = "2025-11-04T13:41:31.055Z" },
-    { url = "https://files.pythonhosted.org/packages/5a/f0/e5e6b99d4191da102f2b0eb9687aaa7f5bea5d9964071a84effc3e40f997/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3006c3dd9ba34b0c094c544c6006cc79e87d8612999f1a5d43b769b89181f23c", size = 1878092, upload-time = "2025-11-04T13:41:33.21Z" },
-    { url = "https://files.pythonhosted.org/packages/71/48/36fb760642d568925953bcc8116455513d6e34c4beaa37544118c36aba6d/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:72f6c8b11857a856bcfa48c86f5368439f74453563f951e473514579d44aa612", size = 2053385, upload-time = "2025-11-04T13:41:35.508Z" },
-    { url = "https://files.pythonhosted.org/packages/20/25/92dc684dd8eb75a234bc1c764b4210cf2646479d54b47bf46061657292a8/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5cb1b2f9742240e4bb26b652a5aeb840aa4b417c7748b6f8387927bc6e45e40d", size = 2218832, upload-time = "2025-11-04T13:41:37.732Z" },
-    { url = "https://files.pythonhosted.org/packages/e2/09/f53e0b05023d3e30357d82eb35835d0f6340ca344720a4599cd663dca599/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:bd3d54f38609ff308209bd43acea66061494157703364ae40c951f83ba99a1a9", size = 2327585, upload-time = "2025-11-04T13:41:40Z" },
-    { url = "https://files.pythonhosted.org/packages/aa/4e/2ae1aa85d6af35a39b236b1b1641de73f5a6ac4d5a7509f77b814885760c/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2ff4321e56e879ee8d2a879501c8e469414d948f4aba74a2d4593184eb326660", size = 2041078, upload-time = "2025-11-04T13:41:42.323Z" },
-    { url = "https://files.pythonhosted.org/packages/cd/13/2e215f17f0ef326fc72afe94776edb77525142c693767fc347ed6288728d/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:d0d2568a8c11bf8225044aa94409e21da0cb09dcdafe9ecd10250b2baad531a9", size = 2173914, upload-time = "2025-11-04T13:41:45.221Z" },
-    { url = "https://files.pythonhosted.org/packages/02/7a/f999a6dcbcd0e5660bc348a3991c8915ce6599f4f2c6ac22f01d7a10816c/pydantic_core-2.41.5-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:a39455728aabd58ceabb03c90e12f71fd30fa69615760a075b9fec596456ccc3", size = 2129560, upload-time = "2025-11-04T13:41:47.474Z" },
-    { url = "https://files.pythonhosted.org/packages/3a/b1/6c990ac65e3b4c079a4fb9f5b05f5b013afa0f4ed6780a3dd236d2cbdc64/pydantic_core-2.41.5-cp314-cp314t-musllinux_1_1_armv7l.whl", hash = "sha256:239edca560d05757817c13dc17c50766136d21f7cd0fac50295499ae24f90fdf", size = 2329244, upload-time = "2025-11-04T13:41:49.992Z" },
-    { url = "https://files.pythonhosted.org/packages/d9/02/3c562f3a51afd4d88fff8dffb1771b30cfdfd79befd9883ee094f5b6c0d8/pydantic_core-2.41.5-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:2a5e06546e19f24c6a96a129142a75cee553cc018ffee48a460059b1185f4470", size = 2331955, upload-time = "2025-11-04T13:41:54.079Z" },
-    { url = "https://files.pythonhosted.org/packages/5c/96/5fb7d8c3c17bc8c62fdb031c47d77a1af698f1d7a406b0f79aaa1338f9ad/pydantic_core-2.41.5-cp314-cp314t-win32.whl", hash = "sha256:b4ececa40ac28afa90871c2cc2b9ffd2ff0bf749380fbdf57d165fd23da353aa", size = 1988906, upload-time = "2025-11-04T13:41:56.606Z" },
-    { url = "https://files.pythonhosted.org/packages/22/ed/182129d83032702912c2e2d8bbe33c036f342cc735737064668585dac28f/pydantic_core-2.41.5-cp314-cp314t-win_amd64.whl", hash = "sha256:80aa89cad80b32a912a65332f64a4450ed00966111b6615ca6816153d3585a8c", size = 1981607, upload-time = "2025-11-04T13:41:58.889Z" },
-    { url = "https://files.pythonhosted.org/packages/9f/ed/068e41660b832bb0b1aa5b58011dea2a3fe0ba7861ff38c4d4904c1c1a99/pydantic_core-2.41.5-cp314-cp314t-win_arm64.whl", hash = "sha256:35b44f37a3199f771c3eaa53051bc8a70cd7b54f333531c59e29fd4db5d15008", size = 1974769, upload-time = "2025-11-04T13:42:01.186Z" },
-]
-
-[[package]]
-name = "typing-extensions"
-version = "4.15.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" },
-]
-
-[[package]]
-name = "typing-inspection"
-version = "0.4.2"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "typing-extensions" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" },
-]

python/parts.nix

@@ -247,27 +247,8 @@
         };
 
       buildSystemsOverlay = final: prev:
-        let
-          addSetupTools = old: {
-            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-              final.setuptools
-            ];
-          };
-
-        in
         {
-          # mlx-lm is a git dependency that needs setuptools
-          mlx-lm = prev.mlx-lm.overrideAttrs addSetupTools;
-          # rouge-score and sacrebleu don't declare setuptools as a build dependency
-          rouge-score = prev.rouge-score.overrideAttrs addSetupTools;
-          sacrebleu = prev.sacrebleu.overrideAttrs addSetupTools;
-          sqlitedict = prev.sqlitedict.overrideAttrs addSetupTools;
-          word2number = prev.word2number.overrideAttrs addSetupTools;
           vllm = prev.vllm.overrideAttrs (old: {
-            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-              final.setuptools
-              final.setuptools-scm
-            ];
             propagatedBuildInputs = (old.propagatedBuildInputs or [ ]) ++ [ final.torch ];
             VLLM_TARGET_DEVICE = "empty";
           });
@@ -320,19 +301,6 @@
           linuxOverlay
         ]
       );
-      editablePythonSet = pythonSet.overrideScope (
-        workspace.mkEditablePyprojectOverlay { root = "$REPO_ROOT"; members = [ "exo" "exo_core" "vllm_engine" "mlx_engine" "exo_bench" ]; }
-      );
-      evenv = editablePythonSet.mkVirtualEnv "exo-dev-env"
-        {
-          exo = lib.optionals isDarwin [ "mlx" ];
-          exo-pyo3-bindings = [ ];
-          exo-bench = [ ];
-          mlx-engine = [ ];
-          vllm-engine = [ ];
-
-        }
-      ;
       cudaPythonSet = (cudaPkgs.callPackage inputs.pyproject-nix.build.packages {
         python = cudaPkgs.python313;
       }).overrideScope (
@@ -348,12 +316,27 @@
       # mlx-cpu and mlx-cuda-13 both ship mlx/ site-packages files; keep first.
       # mlx-cpu/mlx-cuda-13 and nvidia-cudnn-cu12/cu13 ship overlapping files.
       venvCollisionPaths = lib.optionals isLinux [
+        "lib/python3.13/site-packages/cv2*"
         "lib/python3.13/site-packages/mlx*"
         "lib/python3.13/site-packages/nvidia*"
       ];
 
+      editablePythonSet = pythonSet.overrideScope (
+        workspace.mkEditablePyprojectOverlay { root = "$REPO_ROOT"; members = [ "exo" "python/*" "bench" ]; }
+      );
+      evenv = (editablePythonSet.mkVirtualEnv "exo-dev-env"
+        {
+          exo = [ "mlx" "dev" ];
+          exo-pyo3-bindings = [ ];
+          exo-bench = [ ];
+          mlx-engine = [ ];
+          vllm-engine = [ ];
+        }).overrideAttrs {
+        venvIgnoreCollisions = venvCollisionPaths;
+      };
       exoVenv = (pythonSet.mkVirtualEnv "exo-env" {
-        exo = lib.optionals isDarwin [ "mlx" ];
+        exo = [ "mlx" ];
+        mlx-engine = [ ];
         exo-pyo3-bindings = [ ];
       }).overrideAttrs {
         venvIgnoreCollisions = venvCollisionPaths;
@@ -361,6 +344,7 @@
       exoCudaVenv = (cudaPythonSet.mkVirtualEnv "exo-env" {
         exo = lib.optionals cudaPkgs.config.cudaSupport [ "cuda" ];
         exo-pyo3-bindings = [ ];
+        vllm-engine = [ ];
       }).overrideAttrs {
         venvIgnoreCollisions = venvCollisionPaths;
       };

python/vllm_engine/.python-version

@@ -1 +0,0 @@
-3.13

python/vllm_engine/pyproject.toml

@@ -10,6 +10,8 @@ dependencies = [
   "mlx-cuda-13==0.30.6; sys_platform == 'linux'",
   "fastsafetensors>=0.1.10; sys_platform == 'linux'",
   "exo_core",
+  "mlx_engine",                                       # TODO(evan): remove this dependency!
+  "loguru>=0.7.3",
 ]
 
 [tool.uv]
@@ -17,7 +19,8 @@ environments = ["sys_platform == 'linux'"]
 
 
 [tool.uv.sources]
-exo_core = { workspace = true }
+exo_core = { workspace = true, editable = true }
+mlx_engine = { workspace = true, editable = true }
 vllm = { git = "https://github.com/hmellor/vllm.git", rev = "b99bedc737166ae5ca98cb9e3534b96e0c8c69aa" }
 torch = [
   { index = "pytorch-cu130", marker = "platform_machine == 'aarch64'" },

python/vllm_engine/src/vllm_engine/builder.py

@@ -0,0 +1,84 @@
+from dataclasses import dataclass
+from typing import Self, Callable
+from exo_core.constants import EXO_MODELS_DIR
+from exo_core.engine import EngineBuilder, Engine
+from exo_core.types.common import ModelId
+from exo_core.types.instances import BoundInstance
+from exo_core.types.tasks import TextGeneration
+from exo_core.types.runner_response import GenerationResponse
+from vllm_engine.vllm_generator import VllmBatchEngine
+from vllm_engine.vllm_generator import load_vllm_engine
+
+
+@dataclass
+class VllmBuilder(EngineBuilder[BoundInstance, TextGeneration, GenerationResponse]):
+    model_id: ModelId
+    model_path: str
+    trust_remote_code: bool
+    cancel_receiver: MpReceiver[TaskId]
+    event_sender: MpSender[Event]
+    bound_instance: BoundInstance
+
+    @classmethod
+    def create(
+        cls,
+        bound_instance: BoundInstance,
+        event_sender: MpSender[Event],
+        cancel_receiver: MpReceiver[TaskId],
+    ) -> Self:
+        mid = bound_instance.instance.shard_assignments.model_id
+        return cls(
+            mid,
+            str(EXO_MODELS_DIR / mid.normalize()),
+            bound_instance.bound_shard.model_card.trust_remote_code,
+            cancel_receiver,
+            event_sender,
+            bound_instance,
+        )
+
+    def connect(self) -> None:
+        raise NotImplementedError(
+            "Multiple node VLLM instances are not supported at the moment!"
+        )
+
+    def load(
+        self,
+        on_timeout: Callable[[], None],
+        on_layer_loaded: Callable[[int, int], None],
+    ) -> None:
+        self._engine, self._tool_parser, self._prefix_cache = load_vllm_engine(
+            model_path=self.model_path,
+            model_id=self.model_id,
+            trust_remote_code=self.trust_remote_code,
+            n_layers=self.bound_instance.bound_shard.model_card.n_layers,
+            on_layer_loaded=on_layer_loaded,
+        )
+
+    def build(self) -> Engine[TextGeneration, GenerationResponse]:
+        gen = VllmBatchEngine(
+            engine=self._engine,
+            model_id=self.model_id,
+            prefix_cache=self._prefix_cache,
+        )
+        from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+        tokenizer = TokenizerWrapper(self._engine.get_tokenizer())
+        max_concurrent = 1 if os.environ.get("EXO_NO_BATCH") else 8
+
+        logger.info(f"using BatchGenerator (vLLM, max_concurrent={max_concurrent})")
+        return BatchGenerator(
+            tokenizer=tokenizer,
+            group=None,
+            tool_parser=self._tool_parser,
+            kv_prefix_cache=None,
+            model_id=self.model_id,
+            device_rank=0,
+            cancel_receiver=self.cancel_receiver,
+            event_sender=self.event_sender,
+            _gen=gen,
+            max_concurrent_requests=max_concurrent,
+        )
+
+    def close(self) -> None:
+        with contextlib.suppress(NameError, AttributeError):
+            del self._engine, self._prefix_cache, self._tool_parser

python/vllm_engine/src/vllm_engine/growable_cache.py

@@ -1,8 +1,8 @@
 import torch
+from mlx_engine.cache import KVPrefixCache
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
-from exo.shared.logging import logger
-from exo.worker.engines.mlx.cache import KVPrefixCache
+from loguru import logger
 
 INITIAL_FRACTION = 0.05
 GROWTH_HEADROOM_BYTES = 512 * 1024 * 1024

python/vllm_engine/src/vllm_engine/prompt_format.py

@@ -1,14 +1,13 @@
+from exo.worker.engines.mlx.utils_mlx import (
+    apply_chat_template,
+    get_eos_token_ids_for_model,
+)
 from exo_core.types.common import ModelId
 from exo_core.types.text_generation import TextGenerationTaskParams
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.llm_engine import LLMEngine
 
-from exo.worker.engines.mlx.utils_mlx import (
-    apply_chat_template,
-    get_eos_token_ids_for_model,
-)
-
 
 def format_vllm_prompt(
     engine: LLMEngine, params: TextGenerationTaskParams

python/vllm_engine/src/vllm_engine/vllm_generator.py

@@ -8,25 +8,27 @@ from collections.abc import Callable, Generator
 from dataclasses import dataclass, field
 
 import torch
+from mlx_engine.cache import KVPrefixCache
+from mlx_engine.utils_mlx import get_eos_token_ids_for_model
 from exo_core.types.common import ModelId
 from exo_core.types.runner_response import GenerationResponse
 from exo_core.types.tasks import TaskId
 from exo_core.types.text_generation import TextGenerationTaskParams
 from exo_core.utils.memory import Memory
+from exo_core.engine import Engine
+from loguru import logger
 from vllm.engine.arg_utils import EngineArgs
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.llm_engine import LLMEngine
 from vllm.v1.kv_cache_interface import KVCacheConfig
 
-from exo.api.types import (
+from exo_core.types.runner_response import (
     CompletionTokensDetails,
     GenerationStats,
     PromptTokensDetails,
     Usage,
 )
-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.runner.bootstrap import logger
 from exo.worker.runner.llm_inference.tool_parsers import ToolParser, infer_tool_parser
 from vllm_engine.growable_cache import (
     get_model_runner,
@@ -177,9 +179,7 @@ def _build_generation_response(
             else 0.0,
             prompt_tokens=prompt_token_count,
             generation_tokens=completion_tokens,
-            peak_memory_usage=Memory.from_bytes(
-                torch.cuda.max_memory_allocated()  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType, reportAttributeAccessIssue]
-            ),
+            peak_memory_usage=Memory.from_bytes(torch.cuda.max_memory_allocated()),
         )
         mapped_finish_reason = (
             finish_reason
@@ -215,7 +215,7 @@ def vllm_generate(
     stop_ids = _stop_token_ids(tokenizer, model_id)
     max_batch_tokens: int = (
         getattr(engine.model_config, "max_num_batched_tokens", 2048) or 2048
-    )  # type: ignore[reportUnknownMemberType]
+    )
     start_time = time.perf_counter()
     first_token_time: float | None = None
     prev_token_count = 0
@@ -433,7 +433,7 @@ class VllmBatchEngine:
         return results
 
     def cancel(self, task_ids: list[TaskId]) -> None:
-        to_abort = [tid for tid in task_ids if tid in self._active]
+        to_abort = [str(tid) for tid in task_ids if tid in self._active]
         if to_abort:
             self.engine.abort_request(to_abort)
         for tid in task_ids:
@@ -480,18 +480,18 @@ def set_n_layers(n: int) -> None:
 
 
 def _wrap_weights_iterator(
-    original: Callable[..., Generator[tuple[str, "torch.Tensor"], None, None]],
-) -> Callable[..., Generator[tuple[str, "torch.Tensor"], None, None]]:  # pyright: ignore[reportUnknownParameterType]
+    original: Callable[..., Generator[tuple[str, torch.Tensor], None, None]],
+) -> Callable[..., Generator[tuple[str, torch.Tensor], None, None]]:
     def patched(
         hf_weights_files: list[str], *args: object, **kwargs: object
-    ) -> Generator[tuple[str, "torch.Tensor"], None, None]:  # pyright: ignore[reportUnknownParameterType]
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
         callback = get_weight_loading_callback()
         if callback is not None and hf_weights_files:
             total_layers = get_n_layers()
             seen_layers: set[int] = set()
             last_reported = 0
-            for name, tensor in original(hf_weights_files, *args, **kwargs):  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-                yield name, tensor  # pyright: ignore[reportUnknownArgumentType]
+            for name, tensor in original(hf_weights_files, *args, **kwargs):
+                yield name, tensor
                 match = _LAYER_INDEX_PATTERN.search(name)
                 if match:
                     seen_layers.add(int(match.group(1)))
@@ -501,19 +501,19 @@ def _wrap_weights_iterator(
                     last_reported = current
             callback(total_layers, total_layers)
         else:
-            yield from original(hf_weights_files, *args, **kwargs)  # pyright: ignore[reportUnknownMemberType]
+            yield from original(hf_weights_files, *args, **kwargs)
 
     return patched
 
 
-def _monkey_patch_iterator(weight_utils: object, attr_name: str) -> None:  # pyright: ignore[reportUnknownParameterType]
+def _monkey_patch_iterator(weight_utils: object, attr_name: str) -> None:
     original = getattr(weight_utils, attr_name, None)
     if original is None:
         return
-    patched = _wrap_weights_iterator(original)  # pyright: ignore[reportUnknownArgumentType]
+    patched = _wrap_weights_iterator(original)  # pyright: ignore[reportAny]
     setattr(weight_utils, attr_name, patched)
     for mod in list(sys.modules.values()):
-        if mod is None or mod is weight_utils:
+        if mod is weight_utils:
             continue
         for name in list(vars(mod)):
             if vars(mod)[name] is original:
@@ -527,21 +527,21 @@ def _patch_weight_loading_progress() -> None:
     _weight_loading_patched = True
 
     from vllm.model_executor.model_loader import (
-        weight_utils,  # pyright: ignore[reportMissingImports]
+        weight_utils,
     )
 
     _monkey_patch_iterator(weight_utils, "safetensors_weights_iterator")
     _monkey_patch_iterator(weight_utils, "fastsafetensors_weights_iterator")
 
-    import huggingface_hub  # pyright: ignore[reportMissingImports]
+    import huggingface_hub
 
     def _noop_metadata(*_a: object, **_kw: object) -> None:
-        pass  # pyright: ignore[reportUnknownParameterType]
+        pass
 
-    original_metadata = huggingface_hub.get_safetensors_metadata  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-    huggingface_hub.get_safetensors_metadata = _noop_metadata  # pyright: ignore[reportAttributeAccessIssue]
+    original_metadata = huggingface_hub.get_safetensors_metadata
+    huggingface_hub.get_safetensors_metadata = _noop_metadata
     for mod in list(sys.modules.values()):
-        if mod is None or mod is huggingface_hub:
+        if mod is huggingface_hub:
             continue
         for attr in list(vars(mod)):
             if vars(mod)[attr] is original_metadata:
@@ -571,7 +571,7 @@ def load_vllm_engine(
         trust_remote_code=trust_remote_code,
         load_format="fastsafetensors",
         enable_prefix_caching=False,
-        attention_backend="TRITON_ATTN",
+        attention_backend=AttentionBackendEnum.TRITON_ATTN,
         enforce_eager=True,
         disable_log_stats=True,
     )