diff --git a/pyproject.toml b/pyproject.toml index 9ce9a8d8..2023b0ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,7 +66,7 @@ exo_pyo3_bindings = { workspace = true } mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" } mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" } torch = [{ index = "pytorch-cu130", marker = "sys_platform == 'linux'" }] -vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" } +vllm = { git = "https://github.com/rltakashige/vllm-transformers-sm121.git" } # Uncomment to use local mlx/mlx-lm development versions: # mlx = { path = "/Users/Shared/mlx", editable=true } # mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true } diff --git a/src/exo/worker/engines/vllm/growable_cache.py b/src/exo/worker/engines/vllm/growable_cache.py index 33689a26..e0e8e52f 100644 --- a/src/exo/worker/engines/vllm/growable_cache.py +++ b/src/exo/worker/engines/vllm/growable_cache.py @@ -26,7 +26,6 @@ def patch_vllm() -> None: _patch_allocate_slots() _patch_get_computed_blocks() _patch_moe_sum() - _patch_marlin_w2_thread_config() logger.info("vLLM growable KV cache patch applied") @@ -275,22 +274,6 @@ def _patch_moe_sum() -> None: 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