From a6cdc93d3c4e2f18339a3a5911af73544cefd698 Mon Sep 17 00:00:00 2001 From: Evan Date: Tue, 24 Mar 2026 11:22:09 +0000 Subject: [PATCH] disable --- python/mlx_engine/src/mlx_engine/builder.py | 4 +++- python/mlx_engine/src/mlx_engine/cache.py | 23 +++++++++++-------- python/mlx_engine/src/mlx_engine/types.py | 4 ++-- python/vllm_engine/src/vllm_engine/builder.py | 8 ++++--- 4 files changed, 23 insertions(+), 16 deletions(-) diff --git a/python/mlx_engine/src/mlx_engine/builder.py b/python/mlx_engine/src/mlx_engine/builder.py index b628db8f..1d7dd695 100644 --- a/python/mlx_engine/src/mlx_engine/builder.py +++ b/python/mlx_engine/src/mlx_engine/builder.py @@ -27,7 +27,9 @@ from mlx_engine.utils_mlx import initialize_mlx, load_mlx_items @dataclass -class MlxBuilder(EngineBuilder[BoundInstance, TextGeneration, GenerationResponse | ToolCallResponse]): +class MlxBuilder( + EngineBuilder[BoundInstance, TextGeneration, GenerationResponse | ToolCallResponse] +): model_id: ModelId bound_instance: BoundInstance event_sender: MpSender[tuple[CommandId, ErrorChunk | PrefillProgressChunk]] diff --git a/python/mlx_engine/src/mlx_engine/cache.py b/python/mlx_engine/src/mlx_engine/cache.py index 0f74ebdd..f6cf49f2 100644 --- a/python/mlx_engine/src/mlx_engine/cache.py +++ b/python/mlx_engine/src/mlx_engine/cache.py @@ -1,6 +1,6 @@ import os from copy import deepcopy -from typing import TYPE_CHECKING, cast +from typing import cast # , TYPE_CHECKING import mlx.core as mx import psutil @@ -18,14 +18,14 @@ from mlx_lm.tokenizer_utils import TokenizerWrapper from mlx_engine.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS from mlx_engine.types import KVCacheType, MLXCacheType, Model -if TYPE_CHECKING: - from vllm_engine.kv_cache import TorchKVCache +# if TYPE_CHECKING: +# from vllm_engine.kv_cache import TorchKVCache # Fraction of device memory above which LRU eviction kicks in. # Smaller machines need more aggressive eviction. def _default_memory_threshold() -> float: - total_gb = Memory.from_bytes(psutil.virtual_memory().total).in_gb # pyright: ignore[reportAny] + total_gb = Memory.from_bytes(psutil.virtual_memory().total).in_gb # pyright: ignore[reportAny] if total_gb >= 128: return 0.85 if total_gb >= 64: @@ -217,9 +217,10 @@ class KVPrefixCache: return prompt_cache, remaining, best_index - def lookup( + def lookup( # type: ignore self, prompt_token_ids: list[int] - ) -> tuple["TorchKVCache | None", int, int | None]: + ) -> tuple["TorchKVCache | None", int, int | None]: # type: ignore + from vllm_engine.kv_cache import TorchKVCache prompt_mx = mx.array(prompt_token_ids) max_length = len(prompt_token_ids) @@ -250,11 +251,13 @@ class KVPrefixCache: return torch_cache.trim_to(best_length), best_length, best_index def add_from_torch( - self, prompt_token_ids: list[int], cache: "TorchKVCache" + self, + prompt_token_ids: list[int], + cache: "TorchKVCache", # type: ignore ) -> None: self._evict_if_needed() self.prompts.append(mx.array(prompt_token_ids)) - self.caches.append(cache.detach_cpu()) + self.caches.append(cache.detach_cpu()) # type: ignore self._snapshots.append(None) self._access_counter += 1 self._last_used.append(self._access_counter) @@ -351,14 +354,14 @@ def get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int: def get_available_memory() -> Memory: - mem: int = psutil.virtual_memory().available # pyright: ignore[reportAny] + mem: int = psutil.virtual_memory().available # pyright: ignore[reportAny] return Memory.from_bytes(mem) def get_memory_used_percentage() -> float: mem = psutil.virtual_memory() # percent is 0-100 - return float(mem.percent / 100) # pyright: ignore[reportAny] + return float(mem.percent / 100) # pyright: ignore[reportAny] def make_kv_cache( diff --git a/python/mlx_engine/src/mlx_engine/types.py b/python/mlx_engine/src/mlx_engine/types.py index 62ec52f5..8140078d 100644 --- a/python/mlx_engine/src/mlx_engine/types.py +++ b/python/mlx_engine/src/mlx_engine/types.py @@ -11,13 +11,13 @@ from mlx_lm.models.cache import ( QuantizedKVCache, RotatingKVCache, ) -from vllm_engine.kv_cache import TorchKVCache +# from vllm_engine.kv_cache import TorchKVCache MLXCacheType = Sequence[ KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList ] -KVCacheType = MLXCacheType | TorchKVCache +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. diff --git a/python/vllm_engine/src/vllm_engine/builder.py b/python/vllm_engine/src/vllm_engine/builder.py index 0b1eb8b6..74666d0a 100644 --- a/python/vllm_engine/src/vllm_engine/builder.py +++ b/python/vllm_engine/src/vllm_engine/builder.py @@ -18,7 +18,9 @@ from vllm_engine.vllm_generator import VllmBatchEngine, load_vllm_engine @dataclass -class VllmBuilder(EngineBuilder[BoundInstance, TextGeneration, GenerationResponse | ToolCallResponse]): +class VllmBuilder( + EngineBuilder[BoundInstance, TextGeneration, GenerationResponse | ToolCallResponse] +): model_id: ModelId model_path: str trust_remote_code: bool @@ -30,8 +32,8 @@ class VllmBuilder(EngineBuilder[BoundInstance, TextGeneration, GenerationRespons def create( cls, bound_instance: BoundInstance, - cancel_receiver: MpReceiver[TaskId], - event_sender: MpSender[tuple[CommandId, ErrorChunk | PrefillProgressChunk]], + cancel_receiver: MpReceiver[TaskId], + event_sender: MpSender[tuple[CommandId, ErrorChunk | PrefillProgressChunk]], ) -> Self: mid = bound_instance.instance.shard_assignments.model_id return cls(