diff --git a/.cuda_typings/vllm/v1/core/kv_cache_manager.pyi b/.cuda_typings/vllm/v1/core/kv_cache_manager.pyi index d6d6a838..a36a944b 100644 --- a/.cuda_typings/vllm/v1/core/kv_cache_manager.pyi +++ b/.cuda_typings/vllm/v1/core/kv_cache_manager.pyi @@ -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: ... diff --git a/flake.nix b/flake.nix index c83c71ec..886a505d 100644 --- a/flake.nix +++ b/flake.nix @@ -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; diff --git a/src/exo/shared/types/mlx.py b/src/exo/shared/types/mlx.py index a7e5faeb..f8e5b174 100644 --- a/src/exo/shared/types/mlx.py +++ b/src/exo/shared/types/mlx.py @@ -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: ... diff --git a/src/exo/worker/engines/kv_cache.py b/src/exo/worker/engines/kv_cache.py index be9ab040..c5074ed5 100644 --- a/src/exo/worker/engines/kv_cache.py +++ b/src/exo/worker/engines/kv_cache.py @@ -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) diff --git a/src/exo/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py index 5cefe075..6ba698a7 100644 --- a/src/exo/worker/engines/mlx/cache.py +++ b/src/exo/worker/engines/mlx/cache.py @@ -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"): diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py index 4be53b0d..2fb01e91 100644 --- a/src/exo/worker/engines/mlx/generator/batch_generate.py +++ b/src/exo/worker/engines/mlx/generator/batch_generate.py @@ -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, diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index eda6a7e5..e8b3cddc 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -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, diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index a4c42998..516c073c 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -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: diff --git a/src/exo/worker/engines/vllm/growable_cache.py b/src/exo/worker/engines/vllm/growable_cache.py index 33689a26..d75e908e 100644 --- a/src/exo/worker/engines/vllm/growable_cache.py +++ b/src/exo/worker/engines/vllm/growable_cache.py @@ -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] - - diff --git a/src/exo/worker/engines/vllm/prompt_format.py b/src/exo/worker/engines/vllm/prompt_format.py index 60e75bc7..330a1060 100644 --- a/src/exo/worker/engines/vllm/prompt_format.py +++ b/src/exo/worker/engines/vllm/prompt_format.py @@ -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] = {} diff --git a/src/exo/worker/engines/vllm/vllm_generator.py b/src/exo/worker/engines/vllm/vllm_generator.py index ab9bb947..b4a69bf2 100644 --- a/src/exo/worker/engines/vllm/vllm_generator.py +++ b/src/exo/worker/engines/vllm/vllm_generator.py @@ -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 diff --git a/src/exo/worker/runner/llm_inference/batch_generator.py b/src/exo/worker/runner/llm_inference/batch_generator.py index 5dea6da4..0f0327cd 100644 --- a/src/exo/worker/runner/llm_inference/batch_generator.py +++ b/src/exo/worker/runner/llm_inference/batch_generator.py @@ -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 diff --git a/src/exo/worker/runner/llm_inference/model_output_parsers.py b/src/exo/worker/runner/llm_inference/model_output_parsers.py index 0e5a4cf2..4d8c50d4 100644 --- a/src/exo/worker/runner/llm_inference/model_output_parsers.py +++ b/src/exo/worker/runner/llm_inference/model_output_parsers.py @@ -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 diff --git a/src/exo/worker/runner/llm_inference/runner.py b/src/exo/worker/runner/llm_inference/runner.py index 3de42709..c36d42fb 100644 --- a/src/exo/worker/runner/llm_inference/runner.py +++ b/src/exo/worker/runner/llm_inference/runner.py @@ -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) diff --git a/src/exo/worker/tests/unittests/test_mlx/test_batch_vs_generate.py b/src/exo/worker/tests/unittests/test_mlx/test_batch_vs_generate.py index 5d1a792b..3fc1c58c 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_batch_vs_generate.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_batch_vs_generate.py @@ -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}" ) diff --git a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py index 00f45921..2f4d69cb 100644 --- a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py +++ b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py @@ -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()), ], )