shmovin

python/exo_core/pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "exo-core"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+authors = [
+    { name = "Evan", email = "[email protected]" }
+]
+requires-python = ">=3.13"
+dependencies = []
+
+[build-system]
+requires = ["uv_build>=0.9.24,<0.10.0"]
+build-backend = "uv_build"

python/exo_core/src/exo_core/engine.py

@@ -0,0 +1,65 @@
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable
+from typing import Self
+
+class TaskId(str): ...
+class Cancelled: ...
+class Finished: ...
+
+CANCEL_ALL_TASKS = TaskId("CANCEL_TALL_TASKS")
+
+class Engine[TaskType, ResponseType](ABC):
+    _cancelled_tasks: set[TaskId]
+
+    def should_cancel(self, task_id: TaskId) -> bool:
+        return (
+            task_id in self._cancelled_tasks
+            or CANCEL_ALL_TASKS in self._cancelled_tasks
+        )
+
+    def cancel_task(self, task_id: TaskId):
+        self._cancelled_tasks.add(task_id)
+
+    @abstractmethod
+    def warmup(self) -> None: ...
+
+    @abstractmethod
+    def submit(
+        self,
+        task: TaskType,
+    ) -> None: ...
+
+    @abstractmethod
+    def step(
+        self,
+    ) -> Iterable[
+        tuple[TaskId, ResponseType | Cancelled | Finished]
+    ]: ...
+
+    @abstractmethod
+    def close(self) -> None: ...
+
+class EngineBuilder[SetupType, TaskType, ResponseType](ABC):
+    @classmethod
+    @abstractmethod
+    def create(
+        cls,
+        bound_instance: SetupType,
+    ) -> Self: ...
+
+    @abstractmethod
+    def connect(self) -> None: ...
+
+    @abstractmethod
+    def load(
+        self,
+        on_timeout: Callable[[], None],
+        on_layer_loaded: Callable[[int, int], None],
+    ) -> None: ...
+
+    @abstractmethod
+    def build(self) -> Engine[TaskType, ResponseType]: ...
+
+    @abstractmethod
+    def close(self) -> None: ...
+

python/mlx_engine/pyproject.toml

@@ -1,5 +1,5 @@
 [project]
-name = "mlx-runner"
+name = "mlx-engine"
 version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"

python/mlx_runner/main.py

@@ -1,6 +0,0 @@
-def main():
-    print("Hello from mlx-runner!")
-
-
-if __name__ == "__main__":
-    main()

python/vllm_engine/pyproject.toml

@@ -0,0 +1,38 @@
+[project]
+name = "vllm-runner"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.13"
+dependencies = [
+  "torch>=2.10.0; sys_platform == 'linux'",
+  "vllm>=0.13.0; sys_platform == 'linux' and platform_machine == 'aarch64'",
+  "mlx-cuda-13==0.30.6; sys_platform == 'linux'",
+  "fastsafetensors>=0.1.10; sys_platform == 'linux'",
+]
+
+[tool.uv]
+environments = [
+    "sys_platform == 'linux' and platform_machine == 'aarch64'"
+]
+
+
+[tool.uv.sources]
+vllm = { git = "https://github.com/hmellor/vllm.git", rev = "b99bedc737166ae5ca98cb9e3534b96e0c8c69aa" }
+torch = [{ index = "pytorch-cu130", marker = "platform_machine == 'aarch64'" },
+  { index = "pytorch-cpu", marker = "platform_machine == 'x86_64'" },
+]
+
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[build-system]
+requires = ["uv_build>=0.8.9,<0.9.0"]
+build-backend = "uv_build"

python/vllm_engine/src/vllm_runner/growable_cache.py

@@ -0,0 +1,413 @@
+import torch
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+from exo.shared.logging import logger
+from exo.worker.engines.mlx.cache import KVPrefixCache
+
+INITIAL_FRACTION = 0.05
+GROWTH_HEADROOM_BYTES = 512 * 1024 * 1024
+MIN_GROWTH_BLOCKS = 16
+
+_patched = False
+_prefix_cache: KVPrefixCache | None = None
+_model_runner: GPUModelRunner | None = None
+
+
+def get_prefix_cache() -> KVPrefixCache | None:
+    return _prefix_cache
+
+
+def set_prefix_cache(cache: KVPrefixCache | None) -> None:
+    global _prefix_cache
+    _prefix_cache = cache
+
+
+def get_model_runner() -> GPUModelRunner | None:
+    return _model_runner
+
+
+def set_model_runner(runner: GPUModelRunner | None) -> None:
+    global _model_runner
+    _model_runner = runner
+
+
+def patch_vllm() -> None:
+    global _patched
+    if _patched:
+        return
+    _patched = True
+
+    _patch_determine_available_memory()
+    _patch_check_enough_kv_cache_memory()
+    _patch_initialize_kv_cache_tensors()
+    _patch_initialize_from_config()
+    _patch_kv_cache_manager_init()
+    _patch_allocate_slots()
+    _patch_get_computed_blocks()
+    _patch_moe_sum()
+    _patch_marlin_w2_thread_config()
+    logger.info("vLLM growable KV cache patch applied")
+
+
+def _patch_determine_available_memory() -> None:
+    from vllm.v1.worker.gpu_worker import Worker
+
+    original = Worker.determine_available_memory
+
+    @torch.inference_mode()
+    def patched(self: "Worker") -> int:
+        try:
+            original(self)
+        except AssertionError:
+            logger.warning(
+                "vLLM memory profiling assertion failed (free memory changed during init, "
+                "likely another process released GPU memory). Continuing with growable cache."
+            )
+        torch.cuda.empty_cache()
+        free_bytes, _ = torch.cuda.mem_get_info()
+        initial = max(int(free_bytes * INITIAL_FRACTION), 1)
+        self._growable_max_kv_bytes = free_bytes
+        logger.info(
+            f"Growable KV cache: initial {initial / (1024**3):.2f} GiB "
+            f"(max {free_bytes / (1024**3):.2f} GiB)"
+        )
+        return initial
+
+    Worker.determine_available_memory = patched  # type: ignore
+
+
+def _patch_check_enough_kv_cache_memory() -> None:
+    from vllm.v1.core import kv_cache_utils
+
+    def noop(*_args: "object", **_kwargs: "object") -> None:
+        pass
+
+    kv_cache_utils._check_enough_kv_cache_memory = noop  # type: ignore
+
+
+def _patch_initialize_kv_cache_tensors() -> None:
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+    original_alloc = GPUModelRunner._allocate_kv_cache_tensors
+
+    def patched_alloc(
+        self: "GPUModelRunner", kv_cache_config: "object"
+    ) -> "dict[str, torch.Tensor]":
+        raw_tensors = original_alloc(self, kv_cache_config)
+        self._growable_raw_tensors = {name: t for name, t in raw_tensors.items()}
+        return raw_tensors
+
+    GPUModelRunner._allocate_kv_cache_tensors = patched_alloc  # type: ignore
+
+    original_init_tensors = GPUModelRunner.initialize_kv_cache_tensors
+
+    def patched_init_tensors(
+        self: "GPUModelRunner",
+        kv_cache_config: "object",
+        kernel_block_sizes: "list[int]",
+    ) -> "dict[str, torch.Tensor]":
+        self._growable_kv_cache_config = kv_cache_config
+        self._growable_kernel_block_sizes = kernel_block_sizes
+        return original_init_tensors(self, kv_cache_config, kernel_block_sizes)
+
+    GPUModelRunner.initialize_kv_cache_tensors = patched_init_tensors  # type: ignore
+
+
+def _patch_initialize_from_config() -> None:
+    from vllm.v1.worker.gpu_worker import Worker
+
+    original = Worker.initialize_from_config
+
+    def patched(self: "Worker", kv_cache_config: "object") -> None:
+        original(self, kv_cache_config)
+        set_model_runner(self.model_runner)
+
+    Worker.initialize_from_config = patched  # type: ignore
+
+
+def _patch_kv_cache_manager_init() -> None:
+    from vllm.v1.core.kv_cache_manager import KVCacheManager
+
+    original_init = KVCacheManager.__init__
+
+    def patched_init(
+        self: "KVCacheManager", *args: "object", **kwargs: "object"
+    ) -> None:
+        original_init(self, *args, **kwargs)
+        self._growable_model_runner = get_model_runner()
+
+    KVCacheManager.__init__ = patched_init  # type: ignore
+
+
+def _patch_allocate_slots() -> None:
+    from vllm.v1.core.kv_cache_manager import KVCacheManager
+
+    original = KVCacheManager.allocate_slots
+
+    def patched(
+        self: "KVCacheManager",
+        request: "object",
+        num_new_tokens: int,
+        *args: "object",
+        **kwargs: "object",
+    ) -> "object":
+        result = original(self, request, num_new_tokens, *args, **kwargs)
+        if result is None and _try_grow_cache(self):
+            result = original(self, request, num_new_tokens, *args, **kwargs)
+        return result
+
+    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
+
+    if model_runner is None:
+        logger.debug("No model_runner reference — cannot grow cache")
+        return False
+
+    free_bytes, _ = torch.cuda.mem_get_info()
+    if free_bytes < GROWTH_HEADROOM_BYTES:
+        logger.debug(f"Only {free_bytes / (1024**3):.2f} GiB free — not enough to grow")
+        return False
+
+    kv_cache_config = model_runner._growable_kv_cache_config  # type: ignore
+    old_num_blocks: int = kv_cache_config.num_blocks
+
+    total_tensor_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors)
+    per_block_bytes = total_tensor_bytes // old_num_blocks
+
+    usable_bytes = int(free_bytes * 0.8)
+    growth_blocks = min(usable_bytes // per_block_bytes, old_num_blocks)
+
+    if growth_blocks < MIN_GROWTH_BLOCKS:
+        logger.debug(f"Growth too small ({growth_blocks} blocks)")
+        return False
+
+    new_num_blocks = old_num_blocks + growth_blocks
+
+    logger.info(
+        f"Growing KV cache: {old_num_blocks} → {new_num_blocks} blocks "
+        f"(+{growth_blocks * per_block_bytes / (1024**3):.2f} GiB)"
+    )
+
+    try:
+        _grow_tensors(model_runner, kv_cache_config, old_num_blocks, new_num_blocks)
+        _grow_block_pool(block_pool, old_num_blocks, new_num_blocks)
+        kv_cache_config.num_blocks = new_num_blocks
+        for tensor_spec in kv_cache_config.kv_cache_tensors:
+            tensor_spec.size = int(tensor_spec.size * new_num_blocks / old_num_blocks)
+        logger.info(f"KV cache grown successfully to {new_num_blocks} blocks")
+        return True
+    except Exception:
+        logger.opt(exception=True).error("Failed to grow KV cache")
+        return False
+
+
+def _grow_tensors(
+    model_runner: "object",
+    kv_cache_config: "object",
+    old_num_blocks: int,
+    new_num_blocks: int,
+) -> None:
+    raw_tensors: dict[str, torch.Tensor] = model_runner._growable_raw_tensors  # type: ignore
+    ratio = new_num_blocks / old_num_blocks
+
+    already_grown: dict[int, torch.Tensor] = {}
+    new_raw_tensors: dict[str, torch.Tensor] = {}
+
+    for layer_name, old_raw in raw_tensors.items():
+        storage_id = old_raw.data_ptr()
+        if storage_id in already_grown:
+            new_raw_tensors[layer_name] = already_grown[storage_id]
+            continue
+
+        old_size = old_raw.numel()
+        new_size = int(old_size * ratio)
+        new_raw = torch.zeros(new_size, dtype=torch.int8, device=old_raw.device)
+        new_raw[:old_size] = old_raw
+        already_grown[storage_id] = new_raw
+        new_raw_tensors[layer_name] = new_raw
+
+    model_runner._growable_raw_tensors = new_raw_tensors  # type: ignore
+
+    kernel_block_sizes: list[int] = model_runner._growable_kernel_block_sizes  # type: ignore
+    new_kv_caches: dict[str, torch.Tensor] = model_runner._reshape_kv_cache_tensors(  # type: ignore
+        kv_cache_config,
+        new_raw_tensors,
+        kernel_block_sizes,
+    )
+
+    forward_context: dict[str, "object"] = (
+        model_runner.compilation_config.static_forward_context
+    )  # type: ignore
+    runner_kv_caches: list[torch.Tensor] = model_runner.kv_caches  # type: ignore
+    runner_kv_caches.clear()
+
+    from collections import defaultdict
+
+    from vllm.v1.worker.utils import extract_layer_index
+
+    num_attn_module = 1
+    hf_config = getattr(getattr(model_runner, "model_config", None), "hf_config", None)  # type: ignore
+    if getattr(hf_config, "model_type", "") == "longcat_flash":
+        num_attn_module = 2
+
+    index2name: dict[int, list[str]] = defaultdict(list)
+    for ln in new_kv_caches:
+        index2name[extract_layer_index(ln, num_attn_module)].append(ln)
+
+    for layer_index in sorted(index2name.keys()):
+        for ln in index2name[layer_index]:
+            runner_kv_caches.append(new_kv_caches[ln])
+
+    for layer_name, kv_cache in new_kv_caches.items():
+        forward_context[layer_name].kv_cache = [kv_cache]  # type: ignore
+
+
+def _grow_block_pool(
+    block_pool: "object", old_num_blocks: int, new_num_blocks: int
+) -> None:
+    from vllm.v1.core.kv_cache_utils import KVCacheBlock
+
+    new_blocks: list["KVCacheBlock"] = []
+    for idx in range(old_num_blocks, new_num_blocks):
+        block = KVCacheBlock(idx)
+        block_pool.blocks.append(block)  # type: ignore
+        new_blocks.append(block)
+
+    block_pool.free_block_queue.append_n(new_blocks)  # type: ignore
+    block_pool.num_gpu_blocks = new_num_blocks  # type: ignore
+
+
+def _patch_moe_sum() -> None:
+    import vllm._custom_ops as ops  # type: ignore[reportMissingImports]
+
+    def moe_sum_f32(x: "torch.Tensor", output: "torch.Tensor") -> None:
+        output[:] = x.to(torch.float32).sum(dim=1).to(output.dtype)  # type: ignore
+
+    ops.moe_sum = moe_sum_f32  # type: ignore
+
+
+def _patch_marlin_w2_thread_config() -> None:
+    try:
+        import vllm._custom_ops as ops  # type: ignore[reportMissingImports]
+    except ImportError:
+        return
+
+    original_gemm = ops.moe_wna16_marlin_gemm
+
+    def patched_gemm(*args: "object", **kwargs: "object") -> "object":
+        kwargs["thread_k"] = 64
+        kwargs["thread_n"] = 128
+        return original_gemm(*args, **kwargs)
+
+    ops.moe_wna16_marlin_gemm = patched_gemm  # type: ignore
+
+
+def _patch_get_computed_blocks() -> None:
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
+    from vllm.v1.core.kv_cache_utils import KVCacheBlock
+    from vllm.v1.request import Request
+
+    original = KVCacheManager.get_computed_blocks
+
+    def patched(
+        self: KVCacheManager,
+        request: Request,
+    ) -> tuple[KVCacheBlocks, int]:
+        prefix_cache = get_prefix_cache()
+        if prefix_cache is None or request.prompt_token_ids is None:
+            return original(self, request)
+
+        from exo.worker.engines.vllm.kv_cache import (
+            TorchKVCache as _TorchKVCache,  # noqa: F811
+        )
+
+        try:
+            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
+        ):
+            return original(self, request)
+
+        from vllm.utils.math_utils import cdiv  # type: ignore[reportMissingImports]
+
+        from exo.worker.engines.vllm.vllm_generator import _build_layer_groups
+
+        num_groups = len(self.kv_cache_config.kv_cache_groups)
+        null_block = self.block_pool.null_block
+        save_offsets = torch_cache.token_offset_per_group or [0] * num_groups
+
+        for gi in range(num_groups):
+            save_off = save_offsets[gi] if gi < len(save_offsets) else 0
+            if save_off > 0:
+                spec = self.kv_cache_config.kv_cache_groups[gi].kv_cache_spec  # type: ignore
+                window = getattr(spec, "sliding_window", 0) or 0
+                if window > 0 and num_matched < save_off + window:
+                    return original(self, request)
+
+        real_block_counts: list[int] = []
+        skipped_block_counts: list[int] = []
+        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
+            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
+            real_block_counts.append(num_real)
+            skipped_block_counts.append(num_skipped_blocks)
+            total_needed += num_real
+
+        if self.block_pool.get_num_free_blocks() < total_needed:
+            return original(self, request)
+
+        blocks_per_group: list[list[KVCacheBlock]] = []
+        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
+            blocks_per_group.append(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
+
+            token_offset_per_group.append(skipped_block_counts[gi] * block_size)
+
+        block_ids_per_group = [[b.block_id for b in grp] for grp in blocks_per_group]
+        layer_to_group = _build_layer_groups(self.kv_cache_config)
+        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
+                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)"
+        )
+        return self.empty_kv_cache_blocks, num_matched
+
+    KVCacheManager.get_computed_blocks = patched  # type: ignore[reportAttributeAccessIssue]

python/vllm_engine/src/vllm_runner/kv_cache.py

@@ -0,0 +1,306 @@
+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,
+    CacheList,
+    KVCache,
+    QuantizedKVCache,
+    RotatingKVCache,
+)
+
+
+@dataclass
+class KVLayerState:
+    keys: torch.Tensor  # [seq_len, n_heads, head_dim]
+    values: torch.Tensor  # [seq_len, n_heads, head_dim]
+
+
+@dataclass
+class RotatingKVLayerState:
+    keys: torch.Tensor  # [buffer_len, n_heads, head_dim]
+    values: torch.Tensor  # [buffer_len, n_heads, head_dim]
+    keep: int
+    max_size: int
+    offset: int
+    idx: int
+
+
+@dataclass
+class ArraysLayerState:
+    arrays: list[torch.Tensor | None]
+
+
+LayerState = KVLayerState | RotatingKVLayerState | ArraysLayerState
+
+
+def _mx_to_torch(arr: mx.array) -> torch.Tensor:
+    mx.eval(arr)
+    if arr.dtype == mx.bfloat16:
+        return torch.from_numpy(np.array(arr.astype(mx.float32))).to(torch.bfloat16)
+    return torch.from_numpy(np.array(arr))
+
+
+def _torch_to_mx(t: torch.Tensor) -> mx.array:
+    t = t.detach().cpu()
+    if t.dtype == torch.bfloat16:
+        return mx.array(t.float().numpy()).astype(mx.bfloat16)  # pyright: ignore[reportAny]
+    return mx.array(t.numpy())  # pyright: ignore[reportAny]
+
+
+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]
+
+
+def _kv_to_nhd(k: mx.array, v: mx.array) -> tuple[torch.Tensor, torch.Tensor]:
+    """Convert MLX BHSD [1, H, S, D] to NHD [S, H, D] torch tensors."""
+    kt = _mx_to_torch(k).squeeze(0).permute(1, 0, 2)  # [H,S,D] -> [S,H,D]
+    vt = _mx_to_torch(v).squeeze(0).permute(1, 0, 2)
+    return kt, vt
+
+
+def _nhd_to_bhsd(kt: torch.Tensor, vt: torch.Tensor) -> tuple[mx.array, mx.array]:
+    """Convert NHD [S, H, D] torch tensors to MLX BHSD [1, H, S, D]."""
+    k_mx = _torch_to_mx(kt.permute(1, 0, 2).unsqueeze(0))  # [S,H,D] -> [1,H,S,D]
+    v_mx = _torch_to_mx(vt.permute(1, 0, 2).unsqueeze(0))
+    return k_mx, v_mx
+
+
+class TorchKVCache:
+    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 []
+        self._num_tokens: int | None = None
+
+    @property
+    def num_layers(self) -> int:
+        return len(self.layers)
+
+    def layer(self, idx: int) -> LayerState:
+        return self.layers[idx]
+
+    def kv_layers(self) -> list[tuple[int, KVLayerState | RotatingKVLayerState]]:
+        return [
+            (i, layer)
+            for i, layer in enumerate(self.layers)
+            if isinstance(layer, (KVLayerState, RotatingKVLayerState))
+        ]
+
+    def detach_cpu(self) -> "TorchKVCache":
+        layers: list[LayerState] = []
+        for layer in self.layers:
+            if isinstance(layer, KVLayerState):
+                if not layer.keys.is_cuda:
+                    layers.append(layer)
+                else:
+                    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),
+                        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))
+        ):
+            torch.cuda.synchronize()
+        return TorchKVCache(layers, list(self.token_offset_per_group))
+
+    def trim_to(self, num_tokens: int) -> "TorchKVCache":
+        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: 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,
+                        )
+                    )
+                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)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportUnknownArgumentType]
+                    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:
+                    arrays.append(_mx_to_torch(arr) if arr is not None else None)
+                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))
+                    )
+                else:
+                    k, v = c.state
+                    kt, vt = _kv_to_nhd(k, v)  # pyright: ignore[reportArgumentType]
+                    layers.append(KVLayerState(keys=kt, values=vt))
+        return cls(layers)
+
+    def to_mlx_cache(self) -> list[KVCache | RotatingKVCache | ArraysCache]:
+        result: list[KVCache | RotatingKVCache | ArraysCache] = []
+        for layer in self.layers:
+            if isinstance(layer, RotatingKVLayerState):
+                c = RotatingKVCache(max_size=layer.max_size, keep=layer.keep)
+                if layer.keys.numel() > 0:
+                    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)
+                    )
+                result.append(c)
+            elif isinstance(layer, ArraysLayerState):
+                c = ArraysCache(size=len(layer.arrays))
+                c.state = [
+                    _torch_to_mx(arr) if arr is not None else None
+                    for arr in layer.arrays
+                ]
+                result.append(c)
+            else:
+                c = KVCache()
+                if layer.keys.numel() > 0:
+                    k_mx, v_mx = _nhd_to_bhsd(layer.keys, layer.values)
+                    c.state = (k_mx, v_mx)
+                result.append(c)
+        return result
+
+    @classmethod
+    def from_vllm_cache(
+        cls,
+        kv_caches: list[torch.Tensor | list[torch.Tensor]],
+        block_ids_per_group: list[list[int]],
+        layer_to_group: list[int],
+        num_tokens: int,
+        token_offset_per_group: list[int] | None = None,
+    ) -> "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)
+
+        layers: list[LayerState] = []
+        for layer_idx, kv in enumerate(kv_caches):
+            gi = layer_to_group[layer_idx]
+            bt = block_tables[gi]
+            k_all, v_all = _split_kv(kv)
+
+            if len(bt) == 0:
+                layers.append(KVLayerState(keys=torch.empty(0), values=torch.empty(0)))
+                continue
+
+            keys = k_all[bt].to("cpu", non_blocking=True)
+            values = v_all[bt].to("cpu", non_blocking=True)
+            torch.cuda.synchronize()
+            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 | list[torch.Tensor]],
+        block_ids_per_group: list[list[int]],
+        layer_to_group: list[int],
+        token_offset_per_group: list[int] | None = None,
+    ) -> None:
+        block_tables = [
+            torch.tensor(ids, dtype=torch.long) for ids in block_ids_per_group
+        ]
+
+        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]
+            kv = kv_caches[layer_idx]
+            k_all, v_all = _split_kv(kv)
+            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}"
+                )
+            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}"
+                )
+            else:
+                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)

python/vllm_engine/src/vllm_runner/prompt_format.py

@@ -0,0 +1,57 @@
+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,
+)
+
+
+def format_vllm_prompt(
+    engine: LLMEngine, params: TextGenerationTaskParams
+) -> tuple[list[int], str, int]:
+    # we should have our own wrapper
+    # (instead of abusing mlx's TokenizerWrapper, use tokenizers Tokenizer)
+    tokenizer = TokenizerWrapper(engine.get_tokenizer())
+    prompt_text = apply_chat_template(tokenizer, params)
+    token_ids: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False)  # type: ignore[reportUnknownMemberType]
+    return token_ids, prompt_text, len(token_ids)
+
+
+def make_vllm_sampling_params(
+    engine: LLMEngine,
+    params: TextGenerationTaskParams,
+    model_id: ModelId | None = None,
+) -> SamplingParams:
+    kwargs: dict[str, object] = {}
+
+    if params.max_output_tokens is not None:
+        kwargs["max_tokens"] = params.max_output_tokens
+    else:
+        kwargs["max_tokens"] = min(engine.model_config.max_model_len, 32168)
+    if params.temperature is not None:
+        kwargs["temperature"] = params.temperature
+    if params.top_p is not None:
+        kwargs["top_p"] = params.top_p
+    if params.top_k is not None:
+        kwargs["top_k"] = params.top_k
+    if params.min_p is not None:
+        kwargs["min_p"] = params.min_p
+    if params.stop is not None:
+        kwargs["stop"] = params.stop
+    if params.seed is not None:
+        kwargs["seed"] = params.seed
+    if params.repetition_penalty is not None:
+        kwargs["repetition_penalty"] = params.repetition_penalty
+    if params.logprobs:
+        kwargs["logprobs"] = params.top_logprobs or 1
+
+    if model_id is not None:
+        extra_stop = get_eos_token_ids_for_model(model_id)
+        if extra_stop:
+            kwargs["stop_token_ids"] = extra_stop
+
+    return SamplingParams(**kwargs)

python/vllm_engine/src/vllm_runner/vllm_generator.py

@@ -0,0 +1,594 @@
+import gc
+import math
+import os
+import re
+import sys
+import time
+from collections.abc import Callable, Generator
+from dataclasses import dataclass, field
+
+import torch
+from vllm.engine.arg_utils import EngineArgs
+from vllm.sampling_params import SamplingParams
+from vllm.v1.engine.llm_engine import LLMEngine
+from vllm.v1.kv_cache_interface import KVCacheConfig
+
+from exo.shared.types.api import (
+    CompletionTokensDetails,
+    GenerationStats,
+    PromptTokensDetails,
+    Usage,
+)
+from exo.shared.types.common import ModelId
+from exo.shared.types.memory import Memory
+from exo.shared.types.tasks import TaskId
+from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.worker.runner_response import GenerationResponse
+from exo.worker.engines.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 (
+    get_model_runner,
+    patch_vllm,
+    set_prefix_cache,
+)
+from exo.worker.engines.vllm.kv_cache import TorchKVCache
+from exo.worker.engines.vllm.prompt_format import (
+    format_vllm_prompt,
+    make_vllm_sampling_params,
+)
+from exo.worker.runner.bootstrap import logger
+from exo.worker.runner.llm_inference.tool_parsers import ToolParser, infer_tool_parser
+
+
+def _build_layer_groups(kv_cache_config: KVCacheConfig) -> list[int]:
+    group_lookup: dict[str, int] = {}
+    for group_idx, group_spec in enumerate(kv_cache_config.kv_cache_groups):
+        for layer_name in group_spec.layer_names:
+            group_lookup[layer_name] = group_idx
+
+    layer_to_group: list[int] = []
+    for tensor_spec in kv_cache_config.kv_cache_tensors:
+        for name in tensor_spec.shared_by:
+            layer_to_group.append(group_lookup[name])
+    return layer_to_group
+
+
+@dataclass
+class _EngineRequest:
+    request_id: str
+    prompt_token_count: int
+    prompt_token_ids: list[int]
+    prefill_done: bool = False
+    prefill_steps: int = 0
+    prev_text: str = ""
+    prev_token_count: int = 0
+    start_time: float = field(default_factory=time.perf_counter)
+    first_token_time: float | None = None
+    on_generation_token: Callable[[], None] | None = None
+    on_prefill_progress: Callable[[int, int], None] | None = None
+
+
+def _save_prefix_cache(
+    engine: LLMEngine,
+    prefix_cache: KVPrefixCache,
+    request_id: str,
+    prompt_token_ids: list[int],
+    prompt_token_count: int,
+) -> None:
+    try:
+        coordinator = None
+        model_runner = get_model_runner()
+        kv_cache_config = None
+        try:
+            engine_core = engine.engine_core.engine_core  # type: ignore
+            coordinator = engine_core.scheduler.kv_cache_manager.coordinator  # type: ignore
+            kv_cache_config = engine_core.scheduler.kv_cache_manager.kv_cache_config  # type: ignore
+        except Exception:
+            pass
+        if coordinator is None or model_runner is None or kv_cache_config is None:
+            return
+
+        internal_id: str | None = None
+        for mgr in coordinator.single_type_managers:  # type: ignore
+            for key in mgr.req_to_blocks:  # type: ignore
+                if str(key).startswith(request_id):  # type: ignore
+                    internal_id = str(key)  # type: ignore
+                    break
+            if internal_id:
+                break
+        if internal_id is None:
+            return
+
+        null_block = coordinator.block_pool.null_block  # type: ignore
+        block_ids_per_group: list[list[int]] = []
+        token_offset_per_group: list[int] = []
+        for mgr in coordinator.single_type_managers:  # type: ignore
+            blocks = mgr.req_to_blocks.get(internal_id)  # type: ignore
+            if not blocks:
+                block_ids_per_group.append([])
+                token_offset_per_group.append(0)
+                continue
+            block_size: int = mgr.block_size  # type: ignore
+            num_leading_nulls = 0
+            for b in blocks:  # type: ignore
+                if b is null_block or b.is_null:  # type: ignore
+                    num_leading_nulls += 1
+                else:
+                    break
+            real_blocks = [b for b in blocks if b is not null_block and not b.is_null]  # type: ignore
+            block_ids_per_group.append([b.block_id for b in real_blocks])  # type: ignore
+            token_offset_per_group.append(num_leading_nulls * block_size)
+
+        layer_to_group = _build_layer_groups(kv_cache_config)
+        torch_cache = TorchKVCache.from_vllm_cache(
+            model_runner.kv_caches,  # type: ignore
+            block_ids_per_group,
+            layer_to_group,
+            prompt_token_count,
+            token_offset_per_group,
+        )
+        prefix_cache.add_from_torch(prompt_token_ids, torch_cache)
+    except Exception:
+        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,
+    finish_reason: str | None,
+    prompt_token_count: int,
+    completion_tokens: int,
+    start_time: float,
+    first_token_time: float | None,
+    suppress_text: bool = False,
+) -> GenerationResponse:
+    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
+    if finish_reason:
+        now = time.perf_counter()
+        prefill_elapsed = (first_token_time or now) - start_time
+        decode_elapsed = now - (first_token_time or now)
+        finish_usage = Usage(
+            prompt_tokens=prompt_token_count,
+            completion_tokens=completion_tokens,
+            total_tokens=prompt_token_count + completion_tokens,
+            prompt_tokens_details=PromptTokensDetails(),
+            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_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"
+        )
+    return GenerationResponse(
+        text=token_text,
+        token=token_id,
+        finish_reason=mapped_finish_reason,
+        usage=finish_usage,
+        stats=finish_stats,
+    )
+
+
+def vllm_generate(
+    engine: LLMEngine,
+    model_id: ModelId,
+    task: TextGenerationTaskParams,
+    prompt: str,
+    prefix_cache: KVPrefixCache,
+    on_prefill_progress: Callable[[int, int], None] | None = None,
+    distributed_prompt_progress_callback: Callable[[], None] | None = None,
+    on_generation_token: Callable[[], None] | None = None,
+) -> Generator[GenerationResponse, None, None]:
+    token_ids, prompt_text, prompt_token_count = format_vllm_prompt(engine, task)
+    logger.info(prompt_text)
+    request_id = f"vllm-seq-{time.monotonic_ns()}"
+    sampling_params = make_vllm_sampling_params(engine, task, model_id)
+    engine.add_request(request_id, {"prompt_token_ids": token_ids}, sampling_params)
+
+    tokenizer = engine.get_tokenizer()
+    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
+    prefill_done = False
+    prefill_steps = 0
+
+    while engine.has_unfinished_requests():
+        if distributed_prompt_progress_callback and not prefill_done:
+            distributed_prompt_progress_callback()
+        outputs = engine.step()
+
+        for output in outputs:
+            if output.request_id != request_id:
+                continue
+            completion = output.outputs[0]
+            new_token_count = len(completion.token_ids)
+            new_tokens = completion.token_ids[prev_token_count:]
+            finish_reason = completion.finish_reason
+            prev_token_count = new_token_count
+
+            if not prefill_done and not new_tokens:
+                prefill_steps += 1
+                if on_prefill_progress:
+                    on_prefill_progress(
+                        min(prefill_steps * max_batch_tokens, prompt_token_count),
+                        prompt_token_count,
+                    )
+                continue
+
+            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
+                )
+
+            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()
+                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) -> int:
+    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
+    token_ids: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False)  # type: ignore
+    params = SamplingParams(max_tokens=50, detokenize=False)
+    engine.add_request("warmup", {"prompt_token_ids": token_ids}, params)
+    t = time.monotonic()
+    tokens_generated = 0
+    while engine.has_unfinished_requests():
+        engine.step()
+        tokens_generated += 1
+    elapsed = max(time.monotonic() - t, 0.001)
+    check_for_cancel_every = min(math.ceil(tokens_generated / elapsed), 100)
+    logger.info(
+        f"vLLM warmup complete, check_for_cancel_every={check_for_cancel_every}"
+    )
+    return check_for_cancel_every
+
+
+@dataclass(eq=False)
+class VllmBatchEngine:
+    engine: LLMEngine
+    model_id: ModelId
+    prefix_cache: KVPrefixCache
+
+    _active: dict[TaskId, _EngineRequest] = field(default_factory=dict, init=False)
+
+    def warmup(self) -> int:
+        return warmup_vllm_engine(self.engine)
+
+    @property
+    def has_work(self) -> bool:
+        return bool(self._active) or self.engine.has_unfinished_requests()
+
+    def submit(
+        self,
+        task_id: TaskId,
+        task_params: TextGenerationTaskParams,
+        prompt: str,
+        on_prefill_progress: Callable[[int, int], None] | None = None,
+        distributed_prompt_progress_callback: Callable[[], None] | None = None,
+        on_generation_token: Callable[[], None] | None = None,
+    ) -> TaskId:
+        token_ids, prompt_text, prompt_token_count = format_vllm_prompt(
+            self.engine, task_params
+        )
+        logger.info(prompt_text)
+        sampling_params = make_vllm_sampling_params(
+            self.engine, task_params, self.model_id
+        )
+        self.engine.add_request(
+            task_id, {"prompt_token_ids": token_ids}, sampling_params
+        )
+        self._active[task_id] = _EngineRequest(
+            request_id=task_id,
+            prompt_token_count=prompt_token_count,
+            prompt_token_ids=token_ids,
+            on_generation_token=on_generation_token,
+            on_prefill_progress=on_prefill_progress,
+        )
+        return task_id
+
+    def step(self) -> list[tuple[TaskId, GenerationResponse]]:
+        if not self.has_work:
+            return []
+
+        outputs = self.engine.step()
+        tokenizer = self.engine.get_tokenizer()
+        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[TaskId, GenerationResponse]] = []
+
+        for output in outputs:
+            task_id = TaskId(output.request_id)
+            if task_id not in self._active:
+                continue
+            req = self._active[task_id]
+            completion = output.outputs[0]
+            new_token_count = len(completion.token_ids)
+            new_tokens = completion.token_ids[req.prev_token_count :]
+            finish_reason = completion.finish_reason
+            req.prev_token_count = new_token_count
+
+            if not req.prefill_done and not new_tokens:
+                req.prefill_steps += 1
+                if req.on_prefill_progress:
+                    req.on_prefill_progress(
+                        min(
+                            req.prefill_steps * max_batch_tokens, req.prompt_token_count
+                        ),
+                        req.prompt_token_count,
+                    )
+                continue
+
+            if not req.prefill_done and new_tokens:
+                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,
+                )
+
+            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(
+                    (
+                        task_id,
+                        _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[task_id]
+
+        for req in self._active.values():
+            if not req.prefill_done:
+                req.prefill_steps += 1
+                if req.on_prefill_progress:
+                    req.on_prefill_progress(
+                        min(
+                            req.prefill_steps * max_batch_tokens, req.prompt_token_count
+                        ),
+                        req.prompt_token_count,
+                    )
+
+        return results
+
+    def cancel(self, task_ids: list[TaskId]) -> None:
+        to_abort = [tid for tid in task_ids if tid in self._active]
+        if to_abort:
+            self.engine.abort_request(to_abort)
+        for tid in task_ids:
+            self._active.pop(tid, 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)
+        self._active.clear()
+        del self.engine
+        gc.collect()
+        torch.cuda.empty_cache()
+        if torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group()
+
+
+_weight_loading_callback: Callable[[int, int], None] | None = None
+_weight_loading_patched = False
+
+
+def get_weight_loading_callback() -> Callable[[int, int], None] | None:
+    return _weight_loading_callback
+
+
+def set_weight_loading_callback(cb: Callable[[int, int], None] | None) -> None:
+    global _weight_loading_callback
+    _weight_loading_callback = cb
+
+
+_LAYER_INDEX_PATTERN = re.compile(r"\.layers\.(\d+)\.")
+_n_layers: int = 1
+
+
+def get_n_layers() -> int:
+    return _n_layers
+
+
+def set_n_layers(n: int) -> None:
+    global _n_layers
+    _n_layers = n
+
+
+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 = get_weight_loading_callback()
+        if callback is not None and hf_weights_files:
+            total_layers = get_n_layers()
+            seen_layers: set[int] = set()
+            last_reported = 0
+            for name, tensor in original(hf_weights_files, *args, **kwargs):  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+                yield name, tensor  # pyright: ignore[reportUnknownArgumentType]
+                match = _LAYER_INDEX_PATTERN.search(name)
+                if match:
+                    seen_layers.add(int(match.group(1)))
+                current = len(seen_layers)
+                if current > last_reported:
+                    callback(current, total_layers)
+                    last_reported = current
+            callback(total_layers, total_layers)
+        else:
+            yield from original(hf_weights_files, *args, **kwargs)  # pyright: ignore[reportUnknownMemberType]
+
+    return patched
+
+
+def _monkey_patch_iterator(weight_utils: object, attr_name: str) -> None:  # pyright: ignore[reportUnknownParameterType]
+    original = getattr(weight_utils, attr_name, None)
+    if original is None:
+        return
+    patched = _wrap_weights_iterator(original)  # pyright: ignore[reportUnknownArgumentType]
+    setattr(weight_utils, attr_name, patched)
+    for mod in list(sys.modules.values()):
+        if mod is None or mod is weight_utils:
+            continue
+        for name in list(vars(mod)):
+            if vars(mod)[name] is original:
+                setattr(mod, name, patched)
+
+
+def _patch_weight_loading_progress() -> None:
+    global _weight_loading_patched
+    if _weight_loading_patched:
+        return
+    _weight_loading_patched = True
+
+    from vllm.model_executor.model_loader import (
+        weight_utils,  # pyright: ignore[reportMissingImports]
+    )
+
+    _monkey_patch_iterator(weight_utils, "safetensors_weights_iterator")
+    _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]
+
+    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()):
+        if mod is None or mod is huggingface_hub:
+            continue
+        for attr in list(vars(mod)):
+            if vars(mod)[attr] is original_metadata:
+                setattr(mod, attr, _noop_metadata)
+
+
+def load_vllm_engine(
+    model_path: str,
+    model_id: ModelId,
+    trust_remote_code: bool,
+    n_layers: int = 1,
+    on_layer_loaded: Callable[[int, int], None] | None = None,
+) -> tuple[LLMEngine, ToolParser | None, KVPrefixCache]:
+    patch_vllm()
+    _patch_weight_loading_progress()
+
+    os.environ.setdefault("FASTSAFETENSORS_NOGDS", "1")
+
+    prefix_cache = KVPrefixCache(group=None)
+    set_prefix_cache(prefix_cache)
+    set_n_layers(n_layers)
+
+    engine_args = EngineArgs(
+        model=model_path,
+        served_model_name=str(model_id),
+        gpu_memory_utilization=0.05,
+        trust_remote_code=trust_remote_code,
+        load_format="fastsafetensors",
+        enable_prefix_caching=False,
+        attention_backend="TRITON_ATTN",
+        enforce_eager=True,
+        disable_log_stats=True,
+    )
+
+    set_weight_loading_callback(on_layer_loaded)
+    engine = LLMEngine.from_engine_args(engine_args)
+
+    tool_parser: ToolParser | None = None
+    tokenizer = engine.get_tokenizer()
+    chat_template = getattr(tokenizer, "chat_template", None)
+    if isinstance(chat_template, str):
+        tool_parser = infer_tool_parser(chat_template)
+        if tool_parser:
+            logger.info(
+                f"inferred tool parser: {tool_parser.start_parsing} / {tool_parser.end_parsing}"
+            )
+
+    logger.info(f"vLLM engine loaded for {model_id}")
+
+    return engine, tool_parser, prefix_cache

python/vllm_runner/pyproject.toml

@@ -1,25 +0,0 @@
-[project]
-name = "vllm-runner"
-version = "0.1.0"
-description = "Add your description here"
-readme = "README.md"
-requires-python = ">=3.13"
-dependencies = [
-  "torch>=2.10.0; sys_platform == 'linux'",
-  "vllm>=0.13.0; sys_platform == 'linux'",
-  "mlx-cuda-13==0.30.6; sys_platform == 'linux'",
-  "fastsafetensors>=0.1.10; sys_platform == 'linux'",
-]
-
-[tool.uv.sources]
-vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
-torch = [{ index = "pytorch-cu130", marker = "sys_platform == 'linux'" }]
-
-[[tool.uv.index]]
-name = "pytorch-cu130"
-url = "https://download.pytorch.org/whl/cu130"
-explicit = true
-
-[build-system]
-requires = ["uv_build>=0.8.9,<0.9.0"]
-build-backend = "uv_build"