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

Tidy pass 1

Browse Source
This commit is contained in:
Ryuichi Leo Takashige
2026-03-16 23:31:44 +00:00
parent 04dcdbd127
commit 8cd1308336
16 changed files with 367 additions and 352 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.cuda_typings/vllm/v1/core/kv_cache_manager.pyi
+3 -1
View File
@@ -19,4 +19,6 @@ class KVCacheManager:
self, request: object, num_new_tokens: int, *args: object, **kwargs: object
) -> KVCacheBlocks | None: ...
def get_computed_blocks(self, request: object) -> tuple[KVCacheBlocks, int]: ...
def create_kv_cache_blocks(self, blocks: tuple[list[KVCacheBlock], ...]) -> KVCacheBlocks: ...
def create_kv_cache_blocks(
self, blocks: tuple[list[KVCacheBlock], ...]
) -> KVCacheBlocks: ...
flake.nix
+17 -16
View File
@@ -114,22 +114,23 @@
};
};
packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
let
uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package);
uvLockMlxVersion = mlxPackage.version;
uvLockMlxRev = builtins.elemAt (builtins.split "#" mlxPackage.source.git) 2;
in
{
metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
mlx = pkgs.callPackage ./nix/mlx.nix {
inherit (self'.packages) metal-toolchain;
inherit uvLockMlxVersion uvLockMlxRev;
};
default = self'.packages.exo;
}
) // lib.optionalAttrs (pkgsCuda != null) {
packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin
(
let
uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package);
uvLockMlxVersion = mlxPackage.version;
uvLockMlxRev = builtins.elemAt (builtins.split "#" mlxPackage.source.git) 2;
in
{
metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
mlx = pkgs.callPackage ./nix/mlx.nix {
inherit (self'.packages) metal-toolchain;
inherit uvLockMlxVersion uvLockMlxRev;
};
default = self'.packages.exo;
}
) // lib.optionalAttrs (pkgsCuda != null) {
torch-cuda = pkgsCuda.python313Packages.torch;
vllm-cuda = pkgsCuda.python313Packages.vllm;
src/exo/shared/types/mlx.py
+6 -6
View File
@@ -14,11 +14,11 @@ from mlx_lm.models.cache import (
from exo.worker.engines.kv_cache import TorchKVCache
# This list contains one cache entry per transformer layer
KVCacheType = (
Sequence[KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList]
| 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.
@@ -29,6 +29,6 @@ class Model(nn.Module):
def __call__(
self,
x: mx.array,
cache: KVCacheType | None,
cache: MLXCacheType | None,
input_embeddings: mx.array | None = None,
) -> mx.array: ...
src/exo/worker/engines/kv_cache.py
+101 -74
View File
@@ -3,13 +3,20 @@
# pyright: reportUnknownArgumentType=false
from __future__ import annotations
from collections.abc import Iterator, Sequence
from copy import deepcopy
from dataclasses import dataclass
import mlx.core as mx
import numpy as np
import torch
from mlx_lm.models.cache import ArraysCache, KVCache, RotatingKVCache
from mlx_lm.models.cache import (
ArraysCache,
CacheList,
KVCache,
QuantizedKVCache,
RotatingKVCache,
)
@dataclass
@@ -50,12 +57,11 @@ def _torch_to_mx(t: torch.Tensor) -> mx.array:
return mx.array(t.numpy())
def _split_kv(kv: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Split a vLLM paged KV tensor into K and V views, each [num_blocks, block_size, H, D].
flash_attn backend: (2, num_blocks, block_size, H, D) — K/V at dim 0
triton_attn backend: (num_blocks, 2, block_size, H, D) — K/V at dim 1
"""
def _split_kv(
kv: torch.Tensor | list[torch.Tensor],
) -> tuple[torch.Tensor, torch.Tensor]:
if isinstance(kv, list):
return kv[0], kv[1]
if kv.shape[0] == 2 and kv.shape[1] != 2:
return kv[0], kv[1]
return kv[:, 0], kv[:, 1]
@@ -76,7 +82,9 @@ def _nhd_to_bhsd(kt: torch.Tensor, vt: torch.Tensor) -> tuple[mx.array, mx.array
class TorchKVCache:
def __init__(self, layers: list[LayerState], token_offset_per_group: list[int] | None = None):
def __init__(
self, layers: list[LayerState], token_offset_per_group: list[int] | None = None
):
self.layers = layers
self.token_offset_per_group = token_offset_per_group or []
@@ -101,55 +109,77 @@ class TorchKVCache:
if not layer.keys.is_cuda:
layers.append(layer)
else:
layers.append(KVLayerState(
layers.append(
KVLayerState(
keys=layer.keys.detach().to("cpu", non_blocking=True),
values=layer.values.detach().to("cpu", non_blocking=True),
)
)
elif isinstance(layer, RotatingKVLayerState):
layers.append(
RotatingKVLayerState(
keys=layer.keys.detach().to("cpu", non_blocking=True),
values=layer.values.detach().to("cpu", non_blocking=True),
))
elif isinstance(layer, RotatingKVLayerState):
layers.append(RotatingKVLayerState(
keys=layer.keys.detach().to("cpu", non_blocking=True),
values=layer.values.detach().to("cpu", non_blocking=True),
keep=layer.keep, max_size=layer.max_size,
offset=layer.offset, idx=layer.idx,
))
keep=layer.keep,
max_size=layer.max_size,
offset=layer.offset,
idx=layer.idx,
)
)
else:
layers.append(deepcopy(layer))
if any(layer.keys.is_cuda for layer in self.layers if isinstance(layer, (KVLayerState, RotatingKVLayerState))):
if any(
layer.keys.is_cuda
for layer in self.layers
if isinstance(layer, (KVLayerState, RotatingKVLayerState))
):
torch.cuda.synchronize()
return TorchKVCache(layers, list(self.token_offset_per_group))
def trim_to(self, num_tokens: int) -> TorchKVCache:
layers: list[LayerState] = []
for layer in self.layers:
if isinstance(layer, KVLayerState):
layers.append(KVLayerState(
keys=layer.keys[:num_tokens],
values=layer.values[:num_tokens],
))
else:
layers.append(deepcopy(layer))
return TorchKVCache(layers, list(self.token_offset_per_group))
trimmed = TorchKVCache(list(self.layers), list(self.token_offset_per_group))
trimmed._num_tokens = num_tokens
return trimmed
@property
def num_tokens(self) -> int | None:
return getattr(self, "_num_tokens", None)
@classmethod
def from_mlx_cache(
cls, cache: list[KVCache | RotatingKVCache | ArraysCache],
cls,
cache: Sequence[
KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList
],
) -> TorchKVCache:
layers: list[LayerState] = []
for c in cache:
if isinstance(c, RotatingKVCache):
if c.keys is None:
layers.append(RotatingKVLayerState(
keys=torch.empty(0), values=torch.empty(0),
keep=c.keep, max_size=c.max_size, offset=c.offset, idx=c._idx,
))
layers.append(
RotatingKVLayerState(
keys=torch.empty(0),
values=torch.empty(0),
keep=c.keep,
max_size=c.max_size,
offset=c.offset,
idx=c._idx,
)
)
else:
k, v = c.state
kt, vt = _kv_to_nhd(k, v) # pyright: ignore[reportArgumentType]
keep, max_size, offset, idx = (int(x) for x in c.meta_state)
layers.append(RotatingKVLayerState(
keys=kt, values=vt,
keep=keep, max_size=max_size, offset=offset, idx=idx,
))
layers.append(
RotatingKVLayerState(
keys=kt,
values=vt,
keep=keep,
max_size=max_size,
offset=offset,
idx=idx,
)
)
elif isinstance(c, ArraysCache):
arrays: list[torch.Tensor | None] = []
for arr in c.state:
@@ -157,7 +187,9 @@ class TorchKVCache:
layers.append(ArraysLayerState(arrays=arrays))
else:
if c.keys is None: # pyright: ignore[reportUnnecessaryComparison]
layers.append(KVLayerState(keys=torch.empty(0), values=torch.empty(0)))
layers.append(
KVLayerState(keys=torch.empty(0), values=torch.empty(0))
)
else:
k, v = c.state
kt, vt = _kv_to_nhd(k, v) # pyright: ignore[reportArgumentType]
@@ -173,7 +205,8 @@ class TorchKVCache:
k_mx, v_mx = _nhd_to_bhsd(layer.keys, layer.values)
c.state = (k_mx, v_mx)
c.meta_state = tuple(
str(x) for x in (layer.keep, layer.max_size, layer.offset, layer.idx)
str(x)
for x in (layer.keep, layer.max_size, layer.offset, layer.idx)
)
result.append(c)
elif isinstance(layer, ArraysLayerState):
@@ -194,7 +227,7 @@ class TorchKVCache:
@classmethod
def from_vllm_cache(
cls,
kv_caches: list[torch.Tensor],
kv_caches: list[torch.Tensor | list[torch.Tensor]],
block_ids_per_group: list[list[int]],
layer_to_group: list[int],
num_tokens: int,
@@ -210,33 +243,21 @@ class TorchKVCache:
for layer_idx, kv in enumerate(kv_caches):
gi = layer_to_group[layer_idx]
bt = block_tables[gi]
offset = token_offset_per_group[gi]
k_all, v_all = _split_kv(kv)
if len(bt) == 0:
layers.append(KVLayerState(keys=torch.empty(0), values=torch.empty(0)))
continue
valid_tokens = num_tokens - offset
keys_valid = k_all[bt].flatten(0, 1)[:valid_tokens].to("cpu", non_blocking=True)
values_valid = v_all[bt].flatten(0, 1)[:valid_tokens].to("cpu", non_blocking=True)
keys = k_all[bt].to("cpu", non_blocking=True)
values = v_all[bt].to("cpu", non_blocking=True)
torch.cuda.synchronize()
if offset > 0:
keys = torch.zeros(num_tokens, *keys_valid.shape[1:], dtype=keys_valid.dtype)
values = torch.zeros(num_tokens, *values_valid.shape[1:], dtype=values_valid.dtype)
keys[offset:] = keys_valid
values[offset:] = values_valid
else:
keys = keys_valid[:num_tokens]
values = values_valid[:num_tokens]
layers.append(KVLayerState(keys=keys, values=values))
return cls(layers, list(token_offset_per_group))
def write_to_vllm_blocks(
self,
kv_caches: list[torch.Tensor],
kv_caches: list[torch.Tensor | list[torch.Tensor]],
block_ids_per_group: list[list[int]],
layer_to_group: list[int],
token_offset_per_group: list[int] | None = None,
@@ -244,40 +265,46 @@ class TorchKVCache:
block_tables = [
torch.tensor(ids, dtype=torch.long) for ids in block_ids_per_group
]
if token_offset_per_group is None:
token_offset_per_group = [0] * len(block_ids_per_group)
device = kv_caches[0].device
first = kv_caches[0]
device = first[0].device if isinstance(first, list) else first.device
for layer_idx, layer in enumerate(self.layers):
if not isinstance(layer, KVLayerState):
continue
gi = layer_to_group[layer_idx]
bt = block_tables[gi]
offset = token_offset_per_group[gi]
kv = kv_caches[layer_idx]
k_all, v_all = _split_kv(kv)
block_size = k_all.shape[1]
valid_keys = layer.keys[offset:]
valid_values = layer.values[offset:]
num_valid = valid_keys.shape[0]
n_full = num_valid // block_size
remainder = num_valid % block_size
if n_full > 0:
k_all[bt[:n_full]] = valid_keys[:n_full * block_size].view(n_full, block_size, *valid_keys.shape[1:]).to(device, non_blocking=True)
v_all[bt[:n_full]] = valid_values[:n_full * block_size].view(n_full, block_size, *valid_values.shape[1:]).to(device, non_blocking=True)
if remainder > 0:
k_all[bt[n_full], :remainder] = valid_keys[n_full * block_size:].to(device, non_blocking=True)
v_all[bt[n_full], :remainder] = valid_values[n_full * block_size:].to(device, non_blocking=True)
n_blocks = min(len(bt), layer.keys.shape[0])
if n_blocks > 0:
k_all[bt[:n_blocks]] = layer.keys[:n_blocks].to(
device, non_blocking=True
)
v_all[bt[:n_blocks]] = layer.values[:n_blocks].to(
device, non_blocking=True
)
torch.cuda.synchronize()
def __iter__(self) -> Iterator[LayerState]:
return iter(self.layers)
def __len__(self) -> int:
return len(self.layers)
def __repr__(self) -> str:
parts: list[str] = [f"TorchKVCache({self.num_layers} layers)"]
for i, layer in enumerate(self.layers):
if isinstance(layer, KVLayerState):
parts.append(f" [{i}] KV: keys={list(layer.keys.shape)} values={list(layer.values.shape)} {layer.keys.dtype}")
parts.append(
f" [{i}] KV: keys={list(layer.keys.shape)} values={list(layer.values.shape)} {layer.keys.dtype}"
)
elif isinstance(layer, RotatingKVLayerState):
parts.append(f" [{i}] RotatingKV: keys={list(layer.keys.shape)} keep={layer.keep} max_size={layer.max_size} offset={layer.offset} idx={layer.idx}")
parts.append(
f" [{i}] RotatingKV: keys={list(layer.keys.shape)} keep={layer.keep} max_size={layer.max_size} offset={layer.offset} idx={layer.idx}"
)
else:
shapes = [list(a.shape) if a is not None else None for a in layer.arrays]
shapes = [
list(a.shape) if a is not None else None for a in layer.arrays
]
parts.append(f" [{i}] Arrays: {shapes}")
return "\n".join(parts)
src/exo/worker/engines/mlx/cache.py
+30 -17
View File
@@ -1,5 +1,8 @@
from __future__ import annotations
import os
from copy import deepcopy
from typing import TYPE_CHECKING
import mlx.core as mx
import psutil
@@ -13,10 +16,13 @@ from mlx_lm.models.cache import (
from mlx_lm.tokenizer_utils import TokenizerWrapper
from exo.shared.types.memory import Memory
from exo.shared.types.mlx import KVCacheType, Model
from exo.shared.types.mlx import KVCacheType, MLXCacheType, Model
from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
from exo.worker.runner.bootstrap import logger
if TYPE_CHECKING:
from exo.worker.engines.kv_cache import TorchKVCache
# Fraction of device memory above which LRU eviction kicks in.
# Smaller machines need more aggressive eviction.
@@ -46,7 +52,7 @@ class CacheSnapshot:
self.token_count = token_count
def snapshot_ssm_states(cache: KVCacheType) -> CacheSnapshot:
def snapshot_ssm_states(cache: MLXCacheType) -> CacheSnapshot:
states: list[ArraysCache | RotatingKVCache | None] = []
for c in cache:
if isinstance(c, (ArraysCache, RotatingKVCache)):
@@ -70,7 +76,7 @@ def _find_nearest_snapshot(
return best
def has_non_kv_caches(cache: KVCacheType) -> bool:
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)
@@ -94,7 +100,7 @@ class KVPrefixCache:
def add_kv_cache(
self,
prompt_tokens: mx.array,
cache: KVCacheType,
cache: MLXCacheType,
ssm_snapshots: list[CacheSnapshot] | None = None,
):
"""Add a new cache entry. Evicts LRU entries if memory is high."""
@@ -110,7 +116,7 @@ class KVPrefixCache:
self,
index: int,
prompt_tokens: mx.array,
cache: KVCacheType,
cache: MLXCacheType,
snapshots: list[CacheSnapshot] | None,
restore_pos: int,
):
@@ -129,10 +135,15 @@ class KVPrefixCache:
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]
assert not isinstance(cached, TorchKVCache)
return cached
def _get_snapshot(
self, entry_index: int, target_token_count: int
) -> tuple[int, CacheSnapshot | None]:
if not has_non_kv_caches(self.caches[entry_index]):
if not has_non_kv_caches(self._get_mlx_cache(entry_index)):
return target_token_count, None
snapshots = self._snapshots[entry_index]
@@ -149,7 +160,7 @@ class KVPrefixCache:
self,
model: Model,
prompt_tokens: mx.array,
) -> tuple[KVCacheType, mx.array, int | None]:
) -> tuple[MLXCacheType, mx.array, int | None]:
"""Get KV cache for prompt, returning remaining tokens to prefill.
Returns:
@@ -184,7 +195,8 @@ class KVPrefixCache:
# 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
has_ssm = has_non_kv_caches(self.caches[best_index])
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)
@@ -192,8 +204,8 @@ class KVPrefixCache:
if restore_snap is None and has_ssm:
return make_kv_cache(model), prompt_tokens, None
prompt_cache = deepcopy(self.caches[best_index])
cached_length = cache_length(self.caches[best_index])
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)
@@ -208,7 +220,9 @@ class KVPrefixCache:
return prompt_cache, remaining, best_index
def lookup(self, prompt_token_ids: list[int]) -> tuple["object | None", int, int | None]:
def lookup(
self, prompt_token_ids: list[int]
) -> tuple[TorchKVCache | None, int, int | None]:
from exo.worker.engines.kv_cache import TorchKVCache
prompt_mx = mx.array(prompt_token_ids)
@@ -239,11 +253,10 @@ class KVPrefixCache:
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: "object") -> None:
"""Store a TorchKVCache directly. For vLLM save path — no MLX conversion."""
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()) # type: ignore[reportArgumentType]
self.caches.append(cache.detach_cpu())
self._snapshots.append(None)
self._access_counter += 1
self._last_used.append(self._access_counter)
@@ -285,7 +298,7 @@ class KVPrefixCache:
def trim_cache(
cache: KVCacheType,
cache: MLXCacheType,
num_tokens: int,
snapshot: CacheSnapshot | None = None,
) -> None:
@@ -323,7 +336,7 @@ def _entry_length(
return 0
def cache_length(cache: KVCacheType) -> int:
def cache_length(cache: MLXCacheType) -> int:
"""Get the number of tokens in a KV cache."""
return max(_entry_length(c) for c in cache)
@@ -352,7 +365,7 @@ def get_memory_used_percentage() -> float:
def make_kv_cache(
model: Model, max_kv_size: int | None = None, keep: int = 0
) -> KVCacheType:
) -> MLXCacheType:
assert hasattr(model, "layers")
if hasattr(model, "make_cache"):
src/exo/worker/engines/mlx/generator/batch_generate.py
+2 -2
View File
@@ -19,7 +19,7 @@ from exo.shared.types.api import (
Usage,
)
from exo.shared.types.memory import Memory
from exo.shared.types.mlx import KVCacheType, Model
from exo.shared.types.mlx import MLXCacheType, Model
from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.runner_response import GenerationResponse
from exo.worker.engines.mlx.cache import (
@@ -353,7 +353,7 @@ class ExoBatchGenerator:
def _save_prefix_cache(
self,
all_prompt_tokens: mx.array,
cache: KVCacheType,
cache: MLXCacheType,
cache_snapshots: list[CacheSnapshot] | None,
prefix_hit_length: int,
matched_index: int | None,
src/exo/worker/engines/mlx/generator/generate.py
+4 -4
View File
@@ -23,7 +23,7 @@ from exo.shared.types.api import (
)
from exo.shared.types.common import ModelId
from exo.shared.types.memory import Memory
from exo.shared.types.mlx import KVCacheType, Model
from exo.shared.types.mlx import MLXCacheType, Model
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
from exo.shared.types.worker.runner_response import (
GenerationResponse,
@@ -76,7 +76,7 @@ def _has_pipeline_communication_layer(model: Model):
def pipeline_parallel_prefill(
model: Model,
prompt: mx.array,
prompt_cache: KVCacheType,
prompt_cache: MLXCacheType,
prefill_step_size: int,
kv_group_size: int | None,
kv_bits: int | None,
@@ -113,7 +113,7 @@ def pipeline_parallel_prefill(
kv_bits=kv_bits,
)
_prompt_cache: KVCacheType = prompt_cache
_prompt_cache: MLXCacheType = prompt_cache
rank = group.rank()
world_size = group.size()
@@ -195,7 +195,7 @@ def prefill(
tokenizer: TokenizerWrapper,
sampler: Callable[[mx.array], mx.array],
prompt_tokens: mx.array,
cache: KVCacheType,
cache: MLXCacheType,
group: mx.distributed.Group | None,
on_prefill_progress: Callable[[int, int], None] | None,
distributed_prompt_progress_callback: Callable[[], None] | None,
src/exo/worker/engines/mlx/utils_mlx.py
+11 -103
View File
@@ -168,94 +168,6 @@ def initialize_mlx(
return mlx_distributed_init(bound_instance)
# Quant methods that mlx_lm can handle natively (no conversion needed)
_MLX_NATIVE_QUANT_METHODS = frozenset(
{"awq", "gptq", "bitnet", "mxfp4", "compressed-tensors"}
)
def _needs_mlx_conversion(model_path: Path) -> bool:
"""Check if a model uses a quantization format that mlx_lm cannot load natively.
Returns True for formats like fp8 that need dequantization via mlx_lm.convert.
Returns False for native MLX models, AWQ/GPTQ (without g_idx), etc.
"""
config_file = model_path / "config.json"
if not config_file.exists():
return False
try:
with open(config_file) as f:
config = json.load(f) # pyright: ignore[reportAny]
except (json.JSONDecodeError, OSError):
return False
quant_config: dict[str, object] | None = config.get("quantization_config") # pyright: ignore[reportAny]
if not quant_config:
text_config: dict[str, object] = config.get("text_config", {}) # pyright: ignore[reportAny]
quant_config = text_config.get("quantization_config") # pyright: ignore[reportAssignmentType]
if not quant_config:
return False
quant_method = str(quant_config.get("quant_method", ""))
# GPTQ with g_idx is explicitly unsupported by mlx_lm
if quant_method == "gptq" and quant_config.get("desc_act", False):
return True
# Check for any weight files containing g_idx (GPTQ models that mlx_lm will reject)
if quant_method == "gptq":
try:
from safetensors import safe_open
for st_file in model_path.glob("*.safetensors"):
with safe_open(str(st_file), framework="numpy") as st: # pyright: ignore[reportUnknownVariableType]
keys: list[str] = st.keys() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
if any("g_idx" in str(k) for k in keys): # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType]
return True
break # Only need to check one file for key names
except Exception:
pass
# FP8 and other non-native methods need conversion
return bool(quant_method and quant_method not in _MLX_NATIVE_QUANT_METHODS)
def _convert_to_mlx(model_path: Path) -> Path:
"""Convert a non-MLX model to MLX format using mlx_lm.convert --dequantize.
The converted model is saved alongside the original with a '-mlx' suffix.
Returns the path to the converted model directory.
"""
converted_path = model_path.parent / (model_path.name + "-mlx")
if (
converted_path.exists()
and (converted_path / "config.json").exists()
and list(converted_path.glob("*.safetensors"))
):
logger.info(f"Using previously converted MLX model at {converted_path}")
return converted_path
logger.info(f"Converting model at {model_path} to MLX format (dequantizing)...")
from mlx_lm.convert import convert # pyright: ignore[reportUnknownVariableType]
convert(
hf_path=str(model_path),
mlx_path=str(converted_path),
dequantize=True,
)
logger.info(f"Conversion complete: {converted_path}")
return converted_path
def _maybe_convert_model(model_path: Path) -> Path:
"""If the model needs conversion to MLX format, convert it and return the new path.
Otherwise return the original path unchanged.
"""
if _needs_mlx_conversion(model_path):
return _convert_to_mlx(model_path)
return model_path
def load_mlx_items(
bound_instance: BoundInstance,
group: Group | None,
@@ -264,9 +176,7 @@ def load_mlx_items(
) -> tuple[Model, TokenizerWrapper]:
if group is None:
logger.info(f"Single device used for {bound_instance.instance}")
model_path = _maybe_convert_model(
build_model_path(bound_instance.bound_shard.model_card.model_id)
)
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
@@ -314,9 +224,7 @@ def shard_and_load(
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> tuple[nn.Module, TokenizerWrapper]:
model_path = _maybe_convert_model(
build_model_path(shard_metadata.model_card.model_id)
)
model_path = build_model_path(shard_metadata.model_card.model_id)
model, _ = load_model(model_path, lazy=True, strict=False)
logger.debug(model)
@@ -510,7 +418,7 @@ def load_tokenizer_for_model_id(
return tokenizer
def normalize_tool_calls(msg_dict: dict[str, Any]) -> None:
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,
@@ -532,7 +440,7 @@ def normalize_tool_calls(msg_dict: dict[str, Any]) -> None:
func["arguments"] = json.loads(args)
def collect_nested_property_names(schema: dict[str, Any]) -> set[str]:
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]
@@ -544,16 +452,16 @@ def collect_nested_property_names(schema: dict[str, Any]) -> set[str]:
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]
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:
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)
nested = _collect_nested_property_names(params)
if nested and not all(name in prompt for name in nested):
return True
return False
@@ -564,7 +472,7 @@ _LOSSY_TEMPLATE_PATTERN = re.compile(
)
def patch_lossy_chat_template(template: str) -> str | None:
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::
@@ -612,7 +520,7 @@ def apply_chat_template(
# 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)
_normalize_tool_calls(msg)
else:
# Add system message (instructions) if present
if task_params.instructions:
@@ -659,7 +567,7 @@ def apply_chat_template(
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)
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)"
@@ -674,7 +582,7 @@ def apply_chat_template(
**extra_kwargs,
)
if task_params.tools and schemas_lost_in_prompt(prompt, task_params.tools):
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:
src/exo/worker/engines/vllm/growable_cache.py
+28 -12
View File
@@ -143,7 +143,6 @@ def _patch_allocate_slots() -> None:
KVCacheManager.allocate_slots = patched # type: ignore
def _try_grow_cache(kv_cache_manager: "object") -> bool:
block_pool = kv_cache_manager.block_pool # type: ignore
model_runner = kv_cache_manager._growable_model_runner # type: ignore
@@ -299,7 +298,8 @@ def _patch_get_computed_blocks() -> None:
original = KVCacheManager.get_computed_blocks
def patched(
self: KVCacheManager, request: Request,
self: KVCacheManager,
request: Request,
) -> tuple[KVCacheBlocks, int]:
prefix_cache = _exo_prefix_cache_ref[0]
if prefix_cache is None or request.prompt_token_ids is None:
@@ -310,11 +310,17 @@ def _patch_get_computed_blocks() -> None:
)
try:
torch_cache, num_matched, _ = prefix_cache.lookup(list(request.prompt_token_ids)) # type: ignore[reportUnknownMemberType]
torch_cache, num_matched, _ = prefix_cache.lookup(
list(request.prompt_token_ids)
) # type: ignore[reportUnknownMemberType]
except Exception:
return original(self, request)
if torch_cache is None or not isinstance(torch_cache, _TorchKVCache) or num_matched == 0:
if (
torch_cache is None
or not isinstance(torch_cache, _TorchKVCache)
or num_matched == 0
):
return original(self, request)
from vllm.utils.math_utils import cdiv # type: ignore[reportMissingImports]
@@ -338,7 +344,9 @@ def _patch_get_computed_blocks() -> None:
total_needed = 0
for gi in range(num_groups):
mgr = self.coordinator.single_type_managers[gi] # type: ignore
block_size: int = self.kv_cache_config.kv_cache_groups[gi].kv_cache_spec.block_size # type: ignore
block_size: int = self.kv_cache_config.kv_cache_groups[
gi
].kv_cache_spec.block_size # type: ignore
num_skipped: int = mgr.get_num_skipped_tokens(num_matched) # type: ignore
num_skipped_blocks = num_skipped // block_size
num_real = cdiv(num_matched, block_size) - num_skipped_blocks
@@ -353,11 +361,17 @@ def _patch_get_computed_blocks() -> None:
token_offset_per_group: list[int] = []
for gi in range(num_groups):
mgr = self.coordinator.single_type_managers[gi] # type: ignore
block_size = self.kv_cache_config.kv_cache_groups[gi].kv_cache_spec.block_size # type: ignore
real_blocks: list[KVCacheBlock] = self.block_pool.get_new_blocks(real_block_counts[gi]) # type: ignore
block_size = self.kv_cache_config.kv_cache_groups[
gi
].kv_cache_spec.block_size # type: ignore
real_blocks: list[KVCacheBlock] = self.block_pool.get_new_blocks(
real_block_counts[gi]
) # type: ignore
blocks_per_group.append(real_blocks)
full_block_list = [null_block] * skipped_block_counts[gi] + list(real_blocks)
full_block_list = [null_block] * skipped_block_counts[gi] + list(
real_blocks
)
req_blocks = mgr.req_to_blocks[request.request_id] # type: ignore
req_blocks.extend(full_block_list) # type: ignore
@@ -368,14 +382,16 @@ def _patch_get_computed_blocks() -> None:
model_runner = self._growable_model_runner # type: ignore[reportAttributeAccessIssue]
if model_runner is not None:
torch_cache.write_to_vllm_blocks( # type: ignore
model_runner.kv_caches, block_ids_per_group, layer_to_group, # type: ignore
model_runner.kv_caches,
block_ids_per_group,
layer_to_group, # type: ignore
token_offset_per_group,
)
total_blocks = sum(len(g) for g in blocks_per_group)
logger.info(f"Prefix cache hit: {num_matched} tokens, {total_blocks} blocks ({num_groups} groups)")
logger.info(
f"Prefix cache hit: {num_matched} tokens, {total_blocks} blocks ({num_groups} groups)"
)
return self.empty_kv_cache_blocks, num_matched
KVCacheManager.get_computed_blocks = patched # type: ignore[reportAttributeAccessIssue]
src/exo/worker/engines/vllm/prompt_format.py
+7 -60
View File
@@ -1,81 +1,28 @@
from typing import Any, cast
from mlx_lm.tokenizer_utils import TokenizerWrapper
from vllm.sampling_params import SamplingParams
from vllm.v1.engine.llm_engine import LLMEngine
from exo.shared.types.common import ModelId
from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.worker.engines.mlx.utils_mlx import (
apply_chat_template,
get_eos_token_ids_for_model,
normalize_tool_calls,
patch_lossy_chat_template,
schemas_lost_in_prompt,
)
from exo.worker.runner.bootstrap import logger
def format_vllm_prompt(
engine: LLMEngine, params: TextGenerationTaskParams
) -> tuple[list[int], str, int]:
tokenizer = engine.get_tokenizer()
if params.chat_template_messages is not None:
formatted_messages: list[dict[str, Any]] = list(params.chat_template_messages)
for msg in formatted_messages:
normalize_tool_calls(msg)
else:
formatted_messages = []
if params.instructions:
formatted_messages.append(
{"role": "system", "content": params.instructions}
)
for msg in params.input:
if msg.content:
formatted_messages.append({"role": msg.role, "content": msg.content})
partial_assistant_content: str | None = None
if formatted_messages and formatted_messages[-1].get("role") == "assistant":
last_content = cast(object, formatted_messages[-1].get("content", ""))
partial_assistant_content = str(last_content)
formatted_messages = formatted_messages[:-1]
extra_kwargs: dict[str, bool | str] = {}
if params.enable_thinking is not None:
extra_kwargs["enable_thinking"] = params.enable_thinking
extra_kwargs["thinking"] = params.enable_thinking
if params.reasoning_effort is not None:
extra_kwargs["reasoning_effort"] = params.reasoning_effort
patched_template: str | None = None
chat_template = getattr(tokenizer, "chat_template", None)
if params.tools and isinstance(chat_template, str):
patched_template = patch_lossy_chat_template(chat_template)
if patched_template is not None:
logger.info("Patched lossy chat template (removed inner_type length guard)")
prompt_text: str = tokenizer.apply_chat_template(
formatted_messages,
tokenize=False,
add_generation_prompt=True,
tools=params.tools,
**({"chat_template": patched_template} if patched_template is not None else {}),
**extra_kwargs,
)
assert isinstance(prompt_text, str)
if params.tools and schemas_lost_in_prompt(prompt_text, params.tools):
logger.warning("Chat template lost nested tool schemas even after patching")
if partial_assistant_content:
prompt_text += partial_assistant_content
tokenizer = TokenizerWrapper(engine.get_tokenizer())
prompt_text = apply_chat_template(tokenizer, params)
token_ids: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False) # type: ignore[reportUnknownMemberType]
return token_ids, prompt_text, len(token_ids)
def make_vllm_sampling_params(
engine: LLMEngine, params: TextGenerationTaskParams, model_id: ModelId | None = None
engine: LLMEngine,
params: TextGenerationTaskParams,
model_id: ModelId | None = None,
) -> SamplingParams:
kwargs: dict[str, object] = {}
src/exo/worker/engines/vllm/vllm_generator.py
+120 -32
View File
@@ -1,5 +1,6 @@
import gc
import json
import os
import re
import sys
import time
@@ -24,6 +25,7 @@ from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.runner_response import GenerationResponse
from exo.worker.engines.kv_cache import TorchKVCache
from exo.worker.engines.mlx.cache import KVPrefixCache
from exo.worker.engines.mlx.utils_mlx import get_eos_token_ids_for_model
from exo.worker.engines.vllm.growable_cache import (
_exo_prefix_cache_ref,
_growable_model_runner_ref,
@@ -130,6 +132,17 @@ def _save_prefix_cache(
logger.opt(exception=True).warning("Failed to save prefix cache")
def _stop_token_ids(tokenizer: object, model_id: ModelId) -> set[int]:
ids: set[int] = set()
eos_id = getattr(tokenizer, "eos_token_id", None)
if eos_id is not None:
ids.add(eos_id)
extra = get_eos_token_ids_for_model(model_id)
if extra:
ids.update(extra)
return ids
def _build_generation_response(
tokenizer: object,
token_id: int,
@@ -138,8 +151,9 @@ def _build_generation_response(
completion_tokens: int,
start_time: float,
first_token_time: float | None,
suppress_text: bool = False,
) -> GenerationResponse:
token_text: str = tokenizer.decode([token_id]) # type: ignore[reportUnknownMemberType]
token_text: str = "" if suppress_text else tokenizer.decode([token_id]) # type: ignore[reportUnknownMemberType]
finish_usage: Usage | None = None
finish_stats: GenerationStats | None = None
mapped_finish_reason: str | None = None
@@ -155,15 +169,23 @@ def _build_generation_response(
completion_tokens_details=CompletionTokensDetails(),
)
finish_stats = GenerationStats(
prompt_tps=prompt_token_count / prefill_elapsed if prefill_elapsed > 0 else 0.0,
generation_tps=completion_tokens / decode_elapsed if decode_elapsed > 0 else 0.0,
prompt_tps=prompt_token_count / prefill_elapsed
if prefill_elapsed > 0
else 0.0,
generation_tps=completion_tokens / decode_elapsed
if decode_elapsed > 0
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]
),
)
mapped_finish_reason = finish_reason if finish_reason in ("stop", "length", "content_filter") else "stop"
mapped_finish_reason = (
finish_reason
if finish_reason in ("stop", "length", "content_filter")
else "stop"
)
return GenerationResponse(
text=token_text,
token=token_id,
@@ -190,7 +212,10 @@ def vllm_generate(
engine.add_request(request_id, {"prompt_token_ids": token_ids}, sampling_params)
tokenizer = engine.get_tokenizer()
max_batch_tokens: int = getattr(engine.model_config, "max_num_batched_tokens", 2048) or 2048 # type: ignore[reportUnknownMemberType]
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
@@ -223,24 +248,49 @@ def vllm_generate(
if not prefill_done and new_tokens:
first_token_time = time.perf_counter()
prefill_done = True
_save_prefix_cache(engine, prefix_cache, request_id, token_ids, prompt_token_count)
_save_prefix_cache(
engine, prefix_cache, request_id, token_ids, prompt_token_count
)
for i, token_id in enumerate(new_tokens):
is_last = i == len(new_tokens) - 1
is_final_stop = is_last and finish_reason and token_id in stop_ids
if on_generation_token:
on_generation_token()
yield _build_generation_response(
tokenizer, token_id,
finish_reason if is_last and finish_reason else None,
prompt_token_count, new_token_count,
start_time, first_token_time,
)
if is_final_stop:
yield _build_generation_response(
tokenizer,
token_id,
finish_reason,
prompt_token_count,
new_token_count,
start_time,
first_token_time,
suppress_text=True,
)
else:
yield _build_generation_response(
tokenizer,
token_id,
finish_reason if is_last and finish_reason else None,
prompt_token_count,
new_token_count,
start_time,
first_token_time,
)
def warmup_vllm_engine(engine: LLMEngine) -> None:
tokenizer = engine.get_tokenizer()
messages = [{"role": "user", "content": "Prompt to warm up the inference engine. Repeat this."}]
prompt_text: str = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) # type: ignore
messages = [
{
"role": "user",
"content": "Prompt to warm up the inference engine. Repeat this.",
}
]
prompt_text: str = tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
) # type: ignore
token_ids: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False) # type: ignore
params = SamplingParams(max_tokens=50, detokenize=False)
engine.add_request("warmup", {"prompt_token_ids": token_ids}, params)
@@ -270,13 +320,19 @@ class VllmBatchEngine:
distributed_prompt_progress_callback: Callable[[], None] | None = None,
on_generation_token: Callable[[], None] | None = None,
) -> int:
token_ids, prompt_text, prompt_token_count = format_vllm_prompt(self.engine, task_params)
token_ids, prompt_text, prompt_token_count = format_vllm_prompt(
self.engine, task_params
)
logger.info(prompt_text)
uid = self._next_uid
self._next_uid += 1
request_id = f"vllm-batch-{uid}"
sampling_params = make_vllm_sampling_params(self.engine, task_params, self.model_id)
self.engine.add_request(request_id, {"prompt_token_ids": token_ids}, sampling_params)
sampling_params = make_vllm_sampling_params(
self.engine, task_params, self.model_id
)
self.engine.add_request(
request_id, {"prompt_token_ids": token_ids}, sampling_params
)
self._active[uid] = _EngineRequest(
uid=uid,
request_id=request_id,
@@ -293,7 +349,10 @@ class VllmBatchEngine:
outputs = self.engine.step()
tokenizer = self.engine.get_tokenizer()
max_batch_tokens: int = getattr(self.engine.model_config, "max_num_batched_tokens", 2048) or 2048 # type: ignore[reportUnknownMemberType]
stop_ids = _stop_token_ids(tokenizer, self.model_id)
max_batch_tokens: int = (
getattr(self.engine.model_config, "max_num_batched_tokens", 2048) or 2048
) # type: ignore[reportUnknownMemberType]
results: list[tuple[int, GenerationResponse]] = []
rid_to_uid = {req.request_id: uid for uid, req in self._active.items()}
@@ -305,7 +364,7 @@ class VllmBatchEngine:
req = self._active[uid]
completion = output.outputs[0]
new_token_count = len(completion.token_ids)
new_tokens = completion.token_ids[req.prev_token_count:]
new_tokens = completion.token_ids[req.prev_token_count :]
finish_reason = completion.finish_reason
req.prev_token_count = new_token_count
@@ -313,7 +372,9 @@ class VllmBatchEngine:
req.prefill_steps += 1
if req.on_prefill_progress:
req.on_prefill_progress(
min(req.prefill_steps * max_batch_tokens, req.prompt_token_count),
min(
req.prefill_steps * max_batch_tokens, req.prompt_token_count
),
req.prompt_token_count,
)
continue
@@ -322,20 +383,33 @@ class VllmBatchEngine:
req.first_token_time = time.perf_counter()
req.prefill_done = True
_save_prefix_cache(
self.engine, self.prefix_cache,
req.request_id, req.prompt_token_ids, req.prompt_token_count,
self.engine,
self.prefix_cache,
req.request_id,
req.prompt_token_ids,
req.prompt_token_count,
)
for i, token_id in enumerate(new_tokens):
is_last = i == len(new_tokens) - 1
is_final_stop = is_last and finish_reason and token_id in stop_ids
if req.on_generation_token:
req.on_generation_token()
results.append((uid, _build_generation_response(
tokenizer, token_id,
finish_reason if is_last and finish_reason else None,
req.prompt_token_count, new_token_count,
req.start_time, req.first_token_time,
)))
results.append(
(
uid,
_build_generation_response(
tokenizer,
token_id,
finish_reason if is_last and finish_reason else None,
req.prompt_token_count,
new_token_count,
req.start_time,
req.first_token_time,
suppress_text=bool(is_final_stop),
),
)
)
if finish_reason:
del self._active[uid]
@@ -345,7 +419,9 @@ class VllmBatchEngine:
req.prefill_steps += 1
if req.on_prefill_progress:
req.on_prefill_progress(
min(req.prefill_steps * max_batch_tokens, req.prompt_token_count),
min(
req.prefill_steps * max_batch_tokens, req.prompt_token_count
),
req.prompt_token_count,
)
@@ -359,6 +435,8 @@ class VllmBatchEngine:
self._active.pop(uid, None)
def close(self) -> None:
if not hasattr(self, "engine"):
return
rids = [req.request_id for req in self._active.values()]
if rids:
self.engine.abort_request(rids)
@@ -387,8 +465,12 @@ def _get_total_layers(model_dir: Path) -> int:
return 1
def _wrap_weights_iterator(original: Callable[..., Generator[tuple[str, "torch.Tensor"], None, None]]) -> Callable[..., Generator[tuple[str, "torch.Tensor"], None, None]]: # pyright: ignore[reportUnknownParameterType]
def patched(hf_weights_files: list[str], *args: object, **kwargs: object) -> Generator[tuple[str, "torch.Tensor"], None, None]: # pyright: ignore[reportUnknownParameterType]
def _wrap_weights_iterator(
original: Callable[..., Generator[tuple[str, "torch.Tensor"], None, None]],
) -> Callable[..., Generator[tuple[str, "torch.Tensor"], None, None]]: # pyright: ignore[reportUnknownParameterType]
def patched(
hf_weights_files: list[str], *args: object, **kwargs: object
) -> Generator[tuple[str, "torch.Tensor"], None, None]: # pyright: ignore[reportUnknownParameterType]
callback = _weight_loading_callback
if callback is not None and hf_weights_files:
model_dir = Path(hf_weights_files[0]).parent
@@ -407,6 +489,7 @@ def _wrap_weights_iterator(original: Callable[..., Generator[tuple[str, "torch.T
callback(total_layers, total_layers)
else:
yield from original(hf_weights_files, *args, **kwargs) # pyright: ignore[reportUnknownMemberType]
return patched
@@ -438,7 +521,10 @@ def _patch_weight_loading_progress() -> None:
_monkey_patch_iterator(weight_utils, "fastsafetensors_weights_iterator")
import huggingface_hub # pyright: ignore[reportMissingImports]
def _noop_metadata(*_a: object, **_kw: object) -> None: pass # pyright: ignore[reportUnknownParameterType]
def _noop_metadata(*_a: object, **_kw: object) -> None:
pass # pyright: ignore[reportUnknownParameterType]
original_metadata = huggingface_hub.get_safetensors_metadata # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
huggingface_hub.get_safetensors_metadata = _noop_metadata # pyright: ignore[reportAttributeAccessIssue]
for mod in list(sys.modules.values()):
@@ -459,6 +545,8 @@ def load_vllm_engine(
patch_vllm()
_patch_weight_loading_progress()
os.environ.setdefault("FASTSAFETENSORS_NOGDS", "1")
prefix_cache = KVPrefixCache(group=None)
_exo_prefix_cache_ref[0] = prefix_cache
src/exo/worker/runner/llm_inference/batch_generator.py
+13 -8
View File
@@ -1,9 +1,12 @@
from __future__ import annotations
import itertools
import time
from abc import ABC, abstractmethod
from collections import deque
from collections.abc import Generator, Iterable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
import mlx.core as mx
from mlx_lm.tokenizer_utils import TokenizerWrapper
@@ -19,6 +22,9 @@ from exo.shared.types.worker.runner_response import GenerationResponse, ToolCall
from exo.utils.channels import MpReceiver, MpSender
from exo.worker.engines.mlx.cache import KVPrefixCache
from exo.worker.engines.mlx.generator.batch_generate import ExoBatchGenerator
if TYPE_CHECKING:
from exo.worker.engines.vllm.vllm_generator import VllmBatchEngine
from exo.worker.engines.mlx.generator.generate import (
PrefillCancelled,
mlx_generate,
@@ -141,13 +147,12 @@ class SequentialGenerator(InferenceGenerator):
) = field(default=None, init=False)
def warmup(self) -> None:
if self.model is not None:
self.check_for_cancel_every = warmup_inference(
model=self.model,
tokenizer=self.tokenizer,
group=self.group,
model_id=self.model_id,
)
self.check_for_cancel_every = warmup_inference(
model=self.model,
tokenizer=self.tokenizer,
group=self.group,
model_id=self.model_id,
)
def submit(
self,
@@ -314,7 +319,7 @@ class BatchGenerator(InferenceGenerator):
device_rank: int
cancel_receiver: MpReceiver[TaskId]
event_sender: MpSender[Event]
_gen: ExoBatchGenerator # ExoBatchGenerator or VllmBatchEngine
_gen: ExoBatchGenerator | VllmBatchEngine
max_concurrent_requests: int = EXO_MAX_CONCURRENT_REQUESTS
check_for_cancel_every: int = 50
src/exo/worker/runner/llm_inference/model_output_parsers.py
+13 -7
View File
@@ -35,7 +35,7 @@ def apply_all_parsers(
prompt: str,
tool_parser: ToolParser | None,
tokenizer: TokenizerWrapper,
model_type: type[Model],
model_type: type[Model] | type[None],
model_id: ModelId,
tools: list[dict[str, Any]] | None,
) -> Generator[GenerationResponse | ToolCallResponse | None]:
@@ -83,20 +83,24 @@ def parse_gpt_oss(
try:
stream.process(token_id)
except HarmonyError as e:
logger.error(f"HarmonyError on token_id={response.token} text={response.text!r}: {e}")
logger.error(
f"HarmonyError on token_id={response.token} text={response.text!r}: {e}"
)
return
delta = stream.last_content_delta
ch = stream.current_channel
recipient = stream.current_recipient
effective_recipient = recipient if (recipient is not None and recipient.startswith("functions.")) else None
effective_recipient = (
recipient
if (recipient is not None and recipient.startswith("functions."))
else None
)
if effective_recipient != current_tool_name:
if current_tool_name is not None:
tool_name = current_tool_name.removeprefix("functions.")
logger.info(
f"parse_gpt_oss yielding tool call: name={tool_name!r}"
)
logger.info(f"parse_gpt_oss yielding tool call: name={tool_name!r}")
yield ToolCallResponse(
tool_calls=[
ToolCallItem(
@@ -117,7 +121,9 @@ def parse_gpt_oss(
tool_arg_parts = []
continue
is_suppressed = ch == "analysis" or (recipient is not None and recipient.startswith("!"))
is_suppressed = ch == "analysis" or (
recipient is not None and recipient.startswith("!")
)
if is_suppressed and not thinking:
thinking = True
src/exo/worker/runner/llm_inference/runner.py
+1 -1
View File
@@ -281,8 +281,8 @@ class Runner:
else:
self.generator.shutdown_cleanup()
gc.collect()
self.send_task_status(task.task_id, TaskStatus.Complete)
self.update_status(RunnerShutdown())
self.send_task_status(task.task_id, TaskStatus.Complete)
def submit_text_generation(self, task: TextGeneration):
assert isinstance(self.generator, InferenceGenerator)
src/exo/worker/tests/unittests/test_mlx/test_batch_vs_generate.py
+10 -7
View File
@@ -11,7 +11,7 @@ import pytest
from mlx_lm.tokenizer_utils import TokenizerWrapper
from exo.shared.types.common import ModelId
from exo.shared.types.mlx import KVCacheType, Model
from exo.shared.types.mlx import MLXCacheType, Model
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
from exo.worker.engines.mlx.cache import CacheSnapshot, KVPrefixCache, cache_length
from exo.worker.engines.mlx.generator.batch_generate import ExoBatchGenerator
@@ -128,8 +128,8 @@ def _assert_state_equal(sa: object, sb: object, label: str) -> None:
def _compare_cache_arrays(
cache_a: KVCacheType,
cache_b: KVCacheType,
cache_a: MLXCacheType,
cache_b: MLXCacheType,
label: str = "",
) -> None:
"""Assert two KV caches have identical array values."""
@@ -282,15 +282,18 @@ class TestBatchVsGenerate:
"BatchGenerator didn't save to prefix cache"
)
mlx_cache = kv_mlx._get_mlx_cache(0) # pyright: ignore[reportPrivateUsage]
batch_cache = kv_batch._get_mlx_cache(0) # pyright: ignore[reportPrivateUsage]
_compare_cache_arrays(
kv_mlx.caches[0],
kv_batch.caches[0],
mlx_cache,
batch_cache,
label=f"[{spec.name}] ",
)
# ── Compare cache lengths ──
mlx_len = cache_length(kv_mlx.caches[0])
batch_len = cache_length(kv_batch.caches[0])
mlx_len = cache_length(mlx_cache)
batch_len = cache_length(batch_cache)
assert mlx_len == batch_len, (
f"[{spec.name}] Cache length: mlx={mlx_len} vs batch={batch_len}"
)
src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
+1 -2
View File
@@ -341,10 +341,9 @@ def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
runner_id=RUNNER_1_ID, runner_status=RunnerShuttingDown()
),
TaskAcknowledged(task_id=SHUTDOWN_TASK_ID),
RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerShutdown()),
TaskStatusUpdated(
task_id=SHUTDOWN_TASK_ID, task_status=TaskStatus.Complete
),
# SPECIAL EXCEPTION FOR RUNNER SHUTDOWN
RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerShutdown()),
],
)