From 9a83fa6cdfe467544fca99c62c91cf7a4db0c542 Mon Sep 17 00:00:00 2001 From: Ryuichi Leo Takashige Date: Wed, 11 Mar 2026 17:21:17 +0000 Subject: [PATCH] Patch VLLM to load multiple models dynamically --- flake.nix | 2 +- pyproject.toml | 2 + src/exo/main.py | 83 +++++---- src/exo/vllm_entry.py | 17 ++ src/exo/vllm_patches/__init__.py | 0 src/exo/vllm_patches/growable_cache.py | 246 +++++++++++++++++++++++++ 6 files changed, 311 insertions(+), 39 deletions(-) create mode 100644 src/exo/vllm_entry.py create mode 100644 src/exo/vllm_patches/__init__.py create mode 100644 src/exo/vllm_patches/growable_cache.py diff --git a/flake.nix b/flake.nix index 77942d4a..c25f475c 100644 --- a/flake.nix +++ b/flake.nix @@ -175,7 +175,7 @@ break fi done - export LD_LIBRARY_PATH="${pkgsCuda.cudaPackages.libnvjitlink}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + export LD_LIBRARY_PATH="${pkgsCuda.stdenv.cc.cc.lib}/lib:${pkgsCuda.cudaPackages.libnvjitlink}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" exec exo-cuda "$@" ''; }; diff --git a/pyproject.toml b/pyproject.toml index d6143acc..62ea713d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,6 +100,8 @@ exclude = [ "**/.direnv", "**/rust", "**/.github", + "**/vllm_patches", + "**/vllm_entry.py", ] stubPath = ".mlx_typings" diff --git a/src/exo/main.py b/src/exo/main.py index cbeb09c1..bd44f38c 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -255,40 +255,46 @@ class Node: self.api.unpause(result.won_clock) -def _run_vllm_server(args: "Args") -> None: - from typing import TYPE_CHECKING +def _run_vllm_servers(args: "Args") -> None: + import subprocess + import sys - if not TYPE_CHECKING: - from vllm.entrypoints.openai.api_server import ( - run_server, - ) - from vllm.entrypoints.openai.cli_args import ( - make_arg_parser, - ) - from vllm.utils.argparse_utils import ( - FlexibleArgumentParser, - ) + models = args.models or [args.model or "Qwen/Qwen2.5-0.5B-Instruct"] + procs: list[subprocess.Popen[bytes]] = [] + env = os.environ.copy() + env["VLLM_SERVER_DEV_MODE"] = "1" - vllm_argv = [ + for i, model in enumerate(models): + port = args.api_port + i + cmd = [ + sys.executable, + "-m", + "exo.vllm_entry", "--model", - args.model or "Qwen/Qwen2.5-0.5B-Instruct", + model, "--host", "0.0.0.0", "--port", - str(args.api_port), + str(port), + "--enable-sleep-mode", ] if args.max_model_len is not None: - vllm_argv += ["--max-model-len", str(args.max_model_len)] + cmd += ["--max-model-len", str(args.max_model_len)] + logger.info(f"vLLM [{i}]: {model} on :{port}") + procs.append(subprocess.Popen(cmd, env=env)) - parser = make_arg_parser(FlexibleArgumentParser()) - vllm_args = parser.parse_args(vllm_argv) - logger.info( - f"vLLM detected — starting OpenAI-compatible server on port {args.api_port}" - ) - anyio.run(run_server, vllm_args) + def shutdown(sig: int, frame: object) -> None: + for p in procs: + p.terminate() + + signal.signal(signal.SIGINT, shutdown) + signal.signal(signal.SIGTERM, shutdown) + + for p in procs: + p.wait() -def _check_vllm(args: "Args") -> None: +def _check_vllm_env() -> None: import sys logger.info("Checking PyTorch + CUDA...") @@ -304,7 +310,7 @@ def _check_vllm(args: "Args") -> None: gpu_name: str = torch.cuda.get_device_name(0) compute_cap: tuple[int, int] = torch.cuda.get_device_capability(0) - cuda_version = torch.version.cuda or "unknown" + cuda_version: str = torch.version.cuda or "unknown" logger.info(f" PyTorch {torch.__version__}") logger.info(f" CUDA {cuda_version}") logger.info(f" GPU: {gpu_name} (compute {compute_cap[0]}.{compute_cap[1]})") @@ -316,21 +322,15 @@ def _check_vllm(args: "Args") -> None: logger.critical("FAIL: vLLM not installed") sys.exit(1) - logger.warning(f" vLLM {vllm.__version__}") # type: ignore - logger.warning("All checks passed.") - - try: - _run_vllm_server(args) - return - except ImportError: - pass + logger.info(f" vLLM {vllm.__version__}") # type: ignore + logger.info("All checks passed.") def main(): args = Args.parse() if args.check_vllm: - _check_vllm(args) + _check_vllm_env() soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) target = min(max(soft, 65535), hard) @@ -358,12 +358,12 @@ def main(): logger.info("FAST_SYNCH forced OFF") try: - import vllm # pyright: ignore[reportMissingImports, reportUnusedImport] # noqa: F401 + import vllm # type: ignore except ImportError: - vllm = None # pyright: ignore[reportAssignmentType] + vllm = None if vllm is not None: - _run_vllm_server(args) + _run_vllm_servers(args) return node = anyio.run(Node.create, args) @@ -391,8 +391,9 @@ class Args(CamelCaseModel): no_batch: bool = False fast_synch: bool | None = None # None = auto, True = force on, False = force off check_vllm: bool = False - model: str | None = None # vLLM model name/path - max_model_len: int | None = None # vLLM max sequence length + model: str | None = None + models: list[str] | None = None + max_model_len: int | None = None @classmethod def parse(cls) -> Self: @@ -476,6 +477,12 @@ class Args(CamelCaseModel): default=None, help="Model name or path (used in vLLM mode)", ) + parser.add_argument( + "--models", + type=lambda s: s.split(","), + default=None, + help="Comma-separated list of models to serve (vLLM mode, one per port)", + ) parser.add_argument( "--max-model-len", type=int, diff --git a/src/exo/vllm_entry.py b/src/exo/vllm_entry.py new file mode 100644 index 00000000..a0a33655 --- /dev/null +++ b/src/exo/vllm_entry.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from exo.vllm_patches.growable_cache import patch_vllm + +patch_vllm() + +if __name__ == "__main__": + import sys + + import anyio + from vllm.entrypoints.openai.api_server import run_server + from vllm.entrypoints.openai.cli_args import make_arg_parser + from vllm.utils.argparse_utils import FlexibleArgumentParser + + parser = make_arg_parser(FlexibleArgumentParser()) + args = parser.parse_args(sys.argv[1:]) + anyio.run(run_server, args) diff --git a/src/exo/vllm_patches/__init__.py b/src/exo/vllm_patches/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/exo/vllm_patches/growable_cache.py b/src/exo/vllm_patches/growable_cache.py new file mode 100644 index 00000000..d6d8e0bb --- /dev/null +++ b/src/exo/vllm_patches/growable_cache.py @@ -0,0 +1,246 @@ +from __future__ import annotations + +import torch + +from exo.shared.logging import logger + +INITIAL_FRACTION = 0.05 +GROWTH_HEADROOM_BYTES = 512 * 1024 * 1024 +MIN_GROWTH_BLOCKS = 16 + +_patched = False + + +def patch_vllm() -> None: + global _patched + if _patched: + return + _patched = True + + _patch_determine_available_memory() + _patch_initialize_kv_cache_tensors() + _patch_initialize_from_config() + _patch_kv_cache_manager_init() + _patch_allocate_slots() + 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: + full_available = original(self) + initial = max(int(full_available * INITIAL_FRACTION), 1) + self._growable_max_kv_bytes = full_available + logger.info( + f"Growable KV cache: initial {initial / (1024**3):.2f} GiB " + f"(max {full_available / (1024**3):.2f} GiB)" + ) + return initial + + Worker.determine_available_memory = patched # 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) + _growable_model_runner_ref[0] = self.model_runner + + Worker.initialize_from_config = patched # type: ignore + + +_growable_model_runner_ref: list["object | None"] = [None] + + +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 = _growable_model_runner_ref[0] + + 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 not None: + return result + + if _try_grow_cache(self): + return original(self, request, num_new_tokens, *args, **kwargs) + return 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 + + 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