vllm support

.envrc

@@ -1,8 +1 @@
 use flake
-
-# creates .venv if doesn't exist and loads its environment
-export VIRTUAL_ENV=".venv"
-if ! [ -d "./$VIRTUAL_ENV" ]; then
-  uv venv
-fi
-layout python

flake.nix

@@ -146,7 +146,7 @@
                 config.treefmt.build.wrapper
 
                 # PYTHON
-                self'.packages.editableVenv
+                self'.packages.exo.passthru.evenv
                 uv
 
                 # RUST

justfile

@@ -40,6 +40,18 @@ build-app: rust-rebuild sync-clean package
     xcodebuild build -project app/EXO/EXO.xcodeproj -scheme EXO -configuration Debug -derivedDataPath app/EXO/build
     @echo "\nBuild complete. Run with:\n  open {{justfile_directory()}}/app/EXO/build/Build/Products/Debug/EXO.app"
 
+sync-cuda:
+    #!/usr/bin/env bash
+    uv sync --extra vllm-cuda13 --extra mlx-cuda13 --no-install-package vllm
+    dest=".venv/lib/python3.13/site-packages"
+    [[ -d $dest/vllm ]] || {
+        nix build .#exo-cuda-13.passthru.evenv
+        # will also grab vllm-0.19.1-distinfo
+        cp -aL result/lib/python3.13/site-packages/vllm* .venv/lib/python3.13/site-packages
+        chmod -R u+rwX .venv/lib/python3.13/site-packages/vllm*
+        rm result
+    }
+
 clean:
     rm -rf **/__pycache__
     rm -rf target/

nix/vllm-setuppy-cmake.patch

@@ -0,0 +1,26 @@
+diff --git a/setup.py b/setup.py
+index 6dc2ed028..bdcc6354a 100644
+--- a/setup.py
++++ b/setup.py
+@@ -18,6 +18,13 @@ from setuptools import Extension, setup
+ from setuptools.command.build_ext import build_ext
+ 
+ 
++if "NIX_ATTRS_JSON_FILE" in os.environ:
++    with open(os.environ["NIX_ATTRS_JSON_FILE"], "r") as f:
++        NIX_ATTRS = json.load(f)
++else:
++    NIX_ATTRS = { "cmakeFlags": os.environ.get("cmakeFlags", "").split() }
++
++
+ def load_module_from_path(module_name, path):
+     spec = importlib.util.spec_from_file_location(module_name, path)
+     module = importlib.util.module_from_spec(spec)
+@@ -184,6 +191,7 @@ class cmake_build_ext(build_ext):
+         cmake_args = [
+             "-DCMAKE_BUILD_TYPE={}".format(cfg),
+             "-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE),
++            *NIX_ATTRS["cmakeFlags"],
+         ]
+ 
+         verbose = envs.VERBOSE

pyproject.toml

@@ -48,27 +48,27 @@ dev = [
 
 [project.optional-dependencies]
 build = ["nanobind"]
-cpu = [
-  "mlx==0.31.1; sys_platform == 'linux'",
-  "mlx-cpu==0.31.1; sys_platform == 'linux'",
+mlx-none = ["anyio"]
+mlx-cpu = [
+  "mlx==0.31.2; sys_platform == 'linux'",
+  "mlx-cpu==0.31.2; sys_platform == 'linux'",
   "mlx-lm; sys_platform == 'linux'",
   "mlx-vlm>=0.3.11; sys_platform== 'linux'",
-  "torch>=2.10.0; sys_platform == 'linux'",
 ]
-cuda12 = [
+mlx-cuda12 = [
   "mlx==0.31.1; sys_platform == 'linux'",
   "mlx-cuda-12==0.31.1; sys_platform == 'linux'",
   "mlx-lm; sys_platform == 'linux'",
   "mlx-vlm>=0.3.11; sys_platform== 'linux'",
-  "torch>=2.10.0; sys_platform == 'linux'",
 ]
-cuda13 = [
+mlx-cuda13 = [
   "mlx==0.31.1; sys_platform == 'linux'",
   "mlx-cuda-13==0.31.1; sys_platform == 'linux'",
   "mlx-lm; sys_platform == 'linux'",
   "mlx-vlm>=0.3.11; sys_platform== 'linux'",
-  "torch>=2.10.0; sys_platform == 'linux'",
 ]
+vllm-none = ["anyio"]
+vllm-cuda13 = ["vllm[cuda13, fastsafetensors]; sys_platform == 'linux'"]
 
 ###
 # workspace configuration
@@ -82,11 +82,12 @@ 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/deepseek-v4" }
 torch = [
-  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'cuda13' and extra != 'cpu' and extra != 'cuda12'" },
-  { index = "pytorch-cu120", marker = "sys_platform == 'linux' and extra == 'cuda12' and extra != 'cpu' and extra != 'cuda13'" },
-  { index = "pytorch-cpu", marker = "(extra != 'cuda12' and extra != 'cuda13' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+  { index = "pytorch-cu130", marker = "(extra == 'mlx-cuda13' or extra == 'vllm-cuda13') and sys_platform == 'linux'" },
+  { index = "pytorch-cu120", marker = "(extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13') and sys_platform == 'linux'" },
+  { index = "pytorch-cpu", marker = "(extra != 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13') and sys_platform == 'linux'" },
+  { index = "pytorch-cpu", marker = "sys_platform == 'darwin'" },
 ]
-vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
+vllm = { git = "http://github.com/evanev7/vllm", branch = "exo" }
 
 [[tool.uv.index]]
 name = "pytorch-cu130"
@@ -156,11 +157,19 @@ root = "src"
 required-version = ">=0.8.6"
 prerelease = "allow"
 environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
-conflicts = [[{ extra = "cuda12" }, { extra = "cuda13" }, { extra = "cpu" }]]
-constraint-dependencies = ["transformers>=5.6.2"]
-override-dependencies = [
-  "mlx==0.31.1; sys_platform=='linux'",
-  "mlx; sys_platform=='darwin'",
+override-dependencies = ["mlx", "opencv-python; python_version < '0'"]
+conflicts = [
+  [
+    { extra = "mlx-cuda13" },
+    { extra = "mlx-cuda12" },
+    { extra = "mlx-cpu" },
+    { extra = "mlx-none" },
+  ],
+  [
+    { extra = "vllm-cuda13" },
+    { extra = "mlx-cuda12" },
+    { extra = "vllm-none" },
+  ],
 ]
 
 [tool.uv.extra-build-dependencies]

python/parts.nix

@@ -10,8 +10,10 @@ let
       inherit (pkgs.stdenv.hostPlatform) isLinux isDarwin isx86_64;
       inherit (pkgs.config) cudaSupport;
       inherit (pkgs) cudaPackages;
-      cuda13Support = cudaSupport && cudaPackages.cudaMajorVersion == "13";
-      libmlx_source = if cuda13Support then "mlx-cuda-13" else if cudaSupport then "mlx-cuda-12" else "mlx-cpu";
+      libmlx_source =
+        if (builtins.elem "mlx-cuda13" members.exo or [ ]) then "mlx-cuda-13"
+        else if (builtins.elem "mlx-cuda12" members.exo or [ ]) then "mlx-cuda-12"
+        else "mlx-cpu";
       python = pkgs.python313;
       cudaLibs = with cudaPackages; [
         cuda_cudart
@@ -113,37 +115,213 @@ let
               });
           } // lib.optionalAttrs isLinux {
           mlx = prev.mlx.overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ lib.optionals cudaSupport [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
-            autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
             postInstall = ''
               cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
             '';
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
         } // lib.optionalAttrs cudaSupport {
           "${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ cudaLibs;
             autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
           nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
-            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
           nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ cudaLibs;
-            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
           nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
-            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
           nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
-            buildInputs = old.buildInputs ++ [ cudaLibs ];
-            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = old.buildInputs ++ cudaLibs;
           });
           torch = prev.torch.overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ cudaLibs;
             autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
+          torchaudio = prev.torchaudio.overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = old.buildInputs ++ [ cudaPackages.cuda_cudart ];
+            preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
+          });
+          torchvision = prev.torchvision.overrideAttrs (old: {
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
+            preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
+          });
+
+          torch-c-dlpack-ext = prev.torch-c-dlpack-ext.overrideAttrs (old: {
+            buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+            preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
+          });
+          # Currently treating vllm as a cuda dep. it obviously exists as a non cuda dep
+          vllm = prev.vllm.overrideAttrs (old:
+            let
+              cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" { } ''
+                mkdir -p $out/include
+                ln -s ${cudaPackages.cuda_cccl}/include $out/include/cccl
+              '';
+
+              cudaRoot = pkgs.symlinkJoin {
+                name = "cuda-merged-exo";
+                paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (cudaLibs ++ [ cudaPackages.cuda_nvcc cuda_cccl_compat ]);
+              };
+
+              cutlass = pkgs.fetchFromGitHub {
+                name = "cutlass-source";
+                owner = "NVIDIA";
+                repo = "cutlass";
+                tag = "v4.2.1";
+                hash = "sha256-iP560D5Vwuj6wX1otJhwbvqe/X4mYVeKTpK533Wr5gY=";
+              };
+              triton-kernels = pkgs.fetchFromGitHub {
+                owner = "triton-lang";
+                repo = "triton";
+                tag = "v3.6.0";
+                hash = "sha256-JFSpQn+WsNnh7CAPlcpOcUp0nyKXNbJEANdXqmkt4Tc=";
+              };
+
+              cutlass-flashmla = pkgs.fetchFromGitHub {
+                owner = "NVIDIA";
+                repo = "cutlass";
+                rev = "147f5673d0c1c3dcf66f78d677fd647e4a020219";
+                hash = "sha256-dHQto08IwTDOIuFUp9jwm1MWkFi8v2YJ/UESrLuG71g=";
+              };
+
+              flashmla = pkgs.stdenv.mkDerivation {
+                pname = "flashmla";
+                version = "1.0.0";
+
+                src = pkgs.fetchFromGitHub {
+                  name = "FlashMLA-source";
+                  owner = "vllm-project";
+                  repo = "FlashMLA";
+                  rev = "c2afa9cb93e674d5a9120a170a6da57b89267208";
+                  hash = "sha256-pKlwxV6G9iHag/jbu3bAyvYvnu5TbrQwUMFV0AlGC3s=";
+                };
+
+                dontConfigure = true;
+
+                buildPhase = ''
+                  rm -rf csrc/cutlass
+                  ln -sf ${cutlass-flashmla} csrc/cutlass
+                '';
+
+                installPhase = ''
+                  cp -rva . $out
+                '';
+              };
+              qutlass = pkgs.fetchFromGitHub {
+                name = "qutlass-source";
+                owner = "IST-DASLab";
+                repo = "qutlass";
+                rev = "830d2c4537c7396e14a02a46fbddd18b5d107c65";
+                hash = "sha256-aG4qd0vlwP+8gudfvHwhtXCFmBOJKQQTvcwahpEqC84=";
+              };
+              vllm-flash-attn = pkgs.stdenv.mkDerivation {
+                pname = "vllm-flash-attn";
+                version = "2.7.2.post1";
+
+                src = pkgs.fetchFromGitHub {
+                  name = "flash-attention-source";
+                  owner = "vllm-project";
+                  repo = "flash-attention";
+                  rev = "188be16520ceefdc625fdf71365585d2ee348fe2";
+                  hash = "sha256-Osec+/IF3+UDtbIhDMBXzUeWJ7hDJNb5FpaVaziPSgM=";
+                };
+
+                patches = [
+                  (pkgs.fetchpatch {
+                    url = "https://github.com/Dao-AILab/flash-attention/commit/dad67c88d4b6122c69d0bed1cebded0cded71cea.patch";
+                    hash = "sha256-JSgXWItOp5KRpFbTQj/cZk+Tqez+4mEz5kmH5EUeQN4=";
+                  })
+                  (pkgs.fetchpatch {
+                    url = "https://github.com/Dao-AILab/flash-attention/commit/e26dd28e487117ee3e6bc4908682f41f31e6f83a.patch";
+                    hash = "sha256-NkCEowXSi+tiWu74Qt+VPKKavx0H9JeteovSJKToK9A=";
+                  })
+                ];
+
+                dontConfigure = true;
+
+                buildPhase = ''
+                  rm -rf csrc/cutlass
+                  ln -sf ${cutlass} csrc/cutlass
+                '';
+
+                installPhase = ''
+                  cp -rva . $out
+                '';
+              };
+            in
+            {
+              patches = (old.patches or [ ]) ++ [ ../nix/vllm-setuppy-cmake.patch ];
+              nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
+                pkgs.cmake
+                pkgs.ninja
+                pkgs.autoAddDriverRunpath
+              ] ++ lib.optionals cudaSupport [
+                cudaPackages.cuda_nvcc
+              ];
+              # TODO: vllm rocm/cpu
+              VLLM_TARGET_DEVICE = "empty";
+              preConfigure = ''
+                export MAX_JOBS="$NIX_BUILD_CORES"
+              '';
+
+              # TODO: vllm non cuda13 support, more arch's, etc.
+            } // lib.optionalAttrs cudaSupport {
+              buildInputs = cudaLibs ++ [ cudaRoot ];
+
+              VLLM_CUDA_VERSION = cudaPackages.cudaMajorMinorVersion;
+              CUDA_HOME = "${cudaRoot}";
+              CUDAToolkit_ROOT = "${cudaRoot}";
+              CUDACXX = "${cudaRoot}/bin/nvcc";
+              VLLM_CUTLASS_SRC_DIR = "${lib.getDev cutlass}";
+              VLLM_TARGET_DEVICE = "cuda";
+              TORCH_CUDA_ARCH_LIST = "12.0;12.1";
+              TRITON_KERNELS_SRC_DIR = "${lib.getDev triton-kernels}/python/triton_kernels/triton_kernels";
+              FLASH_MLA_SRC_DIR = "${lib.getDev flashmla}";
+              QUTLASS_SRC_DIR = "${lib.getDev qutlass}";
+              VLLM_FLASH_ATTN_SRC_DIR = "${lib.getDev vllm-flash-attn}";
+              CAFFE2_USE_CUDNN = "ON";
+              CAFFE2_USE_CUFILE = "ON";
+              CUTLASS_ENABLE_CUBLAS = "ON";
+              CUTLASS_NVCC_ARCHS_ENABLED = "12.0;12.1";
+
+              cmakeFlags = [
+                (lib.cmakeBool "CMAKE_SKIP_INSTALL_RPATH" true)
+                (lib.cmakeBool "CMAKE_BUILD_WITH_INSTALL_RPATH" true)
+                (lib.cmakeFeature "CUDA_HOME" "${cudaRoot}")
+                (lib.cmakeFeature "CUDAToolkit_ROOT" "${cudaRoot}")
+                (lib.cmakeFeature "CMAKE_CUDA_COMPILER" "${cudaRoot}/bin/nvcc")
+                (lib.cmakeFeature "CMAKE_PREFIX_PATH" "${cudaRoot}")
+                (lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_CUTLASS" "${lib.getDev cutlass}")
+                (lib.cmakeFeature "FLASH_MLA_SRC_DIR" "${lib.getDev flashmla}")
+                (lib.cmakeFeature "VLLM_FLASH_ATTN_SRC_DIR" "${lib.getDev vllm-flash-attn}")
+                (lib.cmakeFeature "QUTLASS_SRC_DIR" "${lib.getDev qutlass}")
+                (lib.cmakeFeature "TORCH_CUDA_ARCH_LIST" "12.0;12.1")
+                (lib.cmakeFeature "CUTLASS_NVCC_ARCHS_ENABLED" "${cudaPackages.flags.cmakeCudaArchitecturesString}")
+                (lib.cmakeFeature "CUDA_TOOLKIT_ROOT_DIR" "${cudaRoot}")
+                (lib.cmakeFeature "CAFFE2_USE_CUDNN" "ON")
+                (lib.cmakeFeature "CAFFE2_USE_CUFILE" "ON")
+                (lib.cmakeFeature "CUTLASS_ENABLE_CUBLAS" "ON")
+              ];
+            });
+
+        } // lib.optionalAttrs (cudaSupport && isx86_64) {
+          numba = prev.numba.overrideAttrs (old: {
+            buildInputs = (old.buildInputs or [ ]) ++ [ pkgs.tbb ];
+          });
         };
       pyprojectOverlay = workspace.mkPyprojectOverlay {
         sourcePreference = "wheel";
@@ -164,24 +342,28 @@ let
           buildSystemsOverlay
         ]
       );
-      venv = name: (pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; });
-      mkApp = cmd: name: pkgs.writeShellApplication {
+      # mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
+      venv = name: (pythonSet.mkVirtualEnv "${name}-venv" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ]; });
+      mkApp = text: name: pkgs.writeShellApplication {
         inherit name;
+        text = "exec " + lib.optionalString cudaSupport "nixglhost " + text;
         runtimeEnv = {
           EXO_DASHBOARD_DIR = self'.packages.dashboard;
           EXO_RESOURCES_DIR = inputs.self + /resources;
         };
         runtimeInputs = [
-          # mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
           (venv name)
+          pkgs.nix-gl-host
         ]
         ++ lib.optionals isDarwin [ pkgs.macmon ];
-        text = "exec " + lib.optionalString cudaSupport "${lib.getExe pkgs.nix-gl-host} " + cmd;
+        passthru = {
+          venv = venv name;
+          evenv = ((pythonSet.overrideScope editableOverlay).mkVirtualEnv "${name}-evenv" (members // { exo = (members.exo or [ ]) ++ [ "dev" ]; })).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ]; });
+        };
       };
     in
     {
       inherit venv;
-      editablePythonSet = pythonSet.overrideScope editableOverlay;
       mkPythonScript = path: mkApp ''python ${path} "$@"'';
       mkExo = mkApp ''exo "$@"'';
     };
@@ -191,18 +373,18 @@ in
     { self', pkgs, unfreePkgs, lib, ... }:
     let
       inherit (pkgs.stdenv.hostPlatform) isLinux;
-      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "cpu" ]; }; }) editablePythonSet mkExo;
+      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "mlx-cpu" "vllm-none" ]; }; }) mkExo;
 
       # Virtual environment with dev dependencies for testing
       testVenv = (mkPythonSet {
         inherit self' pkgs lib; members = {
-        exo = [ "dev" "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
+        exo = [ "dev" "mlx-cpu" "vllm-none" ]; # Include pytest, pytest-asyncio, pytest-env
       };
       }).venv "exo-test";
 
       mkBenchScript = (mkPythonSet {
         inherit self' pkgs lib; members = {
-        exo = [ "cpu" ];
+        exo = [ "mlx-cpu" "vllm-none" ];
         exo-bench = [ ]; # Include pytest, pytest-asyncio, pytest-env
       };
       }).mkPythonScript;
@@ -212,12 +394,12 @@ in
         runtimeInputs = [ pkgs.python313 ];
         text = ''exec python ${path} "$@"'';
       };
-
+      cuda12Set = mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "mlx-cuda12" "vllm-none" ]; }; };
+      cuda13Set = mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "mlx-cpu" "vllm-cuda13" ]; }; };
     in
     {
       packages = {
         exo = mkExo "exo";
-        editableVenv = editablePythonSet.mkVirtualEnv "exo-dev-env" { exo = [ "dev" ]; };
         # for running tests in ci
         exo-test-env = testVenv;
         exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
@@ -226,8 +408,8 @@ in
         # used by ./tests/run_exo_on.sh
         exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
       } // lib.optionalAttrs isLinux {
-        exo-cuda-12 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "cuda12" ]; }; }).mkExo "exo-cuda-12";
-        exo-cuda-13 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "cuda13" ]; }; }).mkExo "exo-cuda-13";
+        exo-cuda-12 = cuda12Set.mkExo "exo-cuda-12";
+        exo-cuda-13 = cuda13Set.mkExo "exo-cuda-13";
       };
 
       checks = {

src/exo/master/placement.py

@@ -43,6 +43,7 @@ from exo.shared.types.worker.instances import (
     InstanceMeta,
     MlxJacclInstance,
     MlxRingInstance,
+    VllmInstance,
 )
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.ports import random_ephemeral_port
@@ -202,7 +203,7 @@ def place_instance(
     )
 
     # Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
-    if len(selected_cycle) == 1:
+    if len(selected_cycle) == 1 and command.instance_meta != InstanceMeta.Vllm:
         command = command.model_copy(
             update={
                 "instance_meta": InstanceMeta.MlxRing,
@@ -266,6 +267,11 @@ def place_instance(
                 hosts_by_node=hosts_by_node,
                 ephemeral_port=ephemeral_port,
             )
+        case InstanceMeta.Vllm:
+            target_instances[instance_id] = VllmInstance(
+                instance_id=instance_id,
+                shard_assignments=shard_assignments,
+            )
 
     return target_instances
 

src/exo/shared/types/worker/instances.py

@@ -15,6 +15,7 @@ class InstanceId(Id):
 class InstanceMeta(str, Enum):
     MlxRing = "MlxRing"
     MlxJaccl = "MlxJaccl"
+    Vllm = "Vllm"
 
 
 class BaseInstance(TaggedModel):
@@ -35,8 +36,12 @@ class MlxJacclInstance(BaseInstance):
     jaccl_coordinators: dict[NodeId, str]
 
 
+class VllmInstance(BaseInstance):
+    pass
+
+
 # TODO: Single node instance
-Instance = MlxRingInstance | MlxJacclInstance
+Instance = MlxRingInstance | MlxJacclInstance | VllmInstance
 
 
 class BoundInstance(FrozenModel):

src/exo/worker/engines/mlx/builder.py

@@ -34,7 +34,7 @@ class MlxBuilder(Builder):
     model_id: ModelId
     event_sender: MpSender[Event]
     cancel_receiver: MpReceiver[TaskId]
-    inference_model: Model | None = None
+    model: Model | None = None
     tokenizer: TokenizerWrapper | None = None
     group: mx.distributed.Group | None = None
     vision_processor: VisionProcessor | None = None
@@ -44,7 +44,7 @@ class MlxBuilder(Builder):
 
     def load(self, bound_instance: BoundInstance) -> Generator[ModelLoadingResponse]:
         (
-            self.inference_model,
+            self.model,
             self.tokenizer,
             self.vision_processor,
         ) = yield from load_mlx_items(bound_instance, self.group)
@@ -60,7 +60,7 @@ class MlxBuilder(Builder):
     def build(
         self,
     ) -> Engine:
-        assert self.inference_model
+        assert self.model
         assert self.tokenizer
 
         vision_processor = self.vision_processor
@@ -86,7 +86,7 @@ class MlxBuilder(Builder):
         if os.environ.get("EXO_NO_BATCH"):
             logger.info("using SequentialGenerator (batching disabled)")
             return SequentialGenerator(
-                model=self.inference_model,
+                model=self.model,
                 tokenizer=self.tokenizer,
                 group=self.group,
                 tool_parser=tool_parser,
@@ -100,7 +100,7 @@ class MlxBuilder(Builder):
         else:
             logger.info("using BatchGenerator")
             return BatchGenerator(
-                model=self.inference_model,
+                model=self.model,
                 tokenizer=self.tokenizer,
                 group=self.group,
                 tool_parser=tool_parser,

src/exo/worker/engines/mlx/utils_mlx.py

@@ -50,6 +50,7 @@ from exo.shared.types.worker.instances import (
     BoundInstance,
     MlxJacclInstance,
     MlxRingInstance,
+    VllmInstance,
 )
 from exo.shared.types.worker.runner_response import ModelLoadingResponse
 from exo.shared.types.worker.shards import (
@@ -140,6 +141,8 @@ def mlx_distributed_init(
                 os.environ["MLX_RANK"] = str(rank)
                 os.environ["MLX_JACCL_COORDINATOR"] = jaccl_coordinator
                 group = mx.distributed.init(backend="jaccl", strict=True)
+            case VllmInstance():
+                raise ValueError("loaded VllmInstance in MLX engine")
 
         logger.info(f"Rank {rank} mlx distributed initialization complete")
 

src/exo/worker/engines/vllm/builder.py

@@ -0,0 +1,77 @@
+import contextlib
+import os
+from collections.abc import Generator
+from dataclasses import dataclass
+
+from exo.shared.constants import EXO_MAX_CONCURRENT_REQUESTS
+from exo.shared.types.common import ModelId
+from exo.shared.types.events import Event
+from exo.shared.types.tasks import TaskId
+from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.runner_response import ModelLoadingResponse
+from exo.utils.channels import MpReceiver, MpSender
+from exo.worker.engines.base import Builder, Engine
+from exo.worker.engines.vllm.engine import VllmEngine
+from exo.worker.engines.vllm.generator import VllmBatchEngine, load_vllm_engine
+from exo.worker.runner.bootstrap import logger
+
+
+@dataclass
+class VllmBuilder(Builder):
+    model_id: ModelId
+    event_sender: MpSender[Event]
+    cancel_receiver: MpReceiver[TaskId]
+
+    def connect(self, bound_instance: BoundInstance) -> None:
+        raise NotImplementedError(
+            "Multiple node VLLM instances are not supported at the moment!"
+        )
+
+    def load(
+        self,
+        bound_instance: BoundInstance,
+    ) -> Generator[ModelLoadingResponse]:
+        kv_connector_cls: type[object] | None = None
+        # overlapping = not os.environ.get("EXO_NO_OVERLAPPING_PREFILL_SENDS")
+
+        def on_layer_loaded(loaded: int, total: int) -> None:
+            pass
+
+        self._bound_runner_id = bound_instance.bound_runner_id
+        self._engine, self._tool_parser = load_vllm_engine(
+            model_id=self.model_id,
+            trust_remote_code=bound_instance.bound_shard.model_card.trust_remote_code,
+            n_layers=bound_instance.bound_shard.model_card.n_layers,
+            on_layer_loaded=on_layer_loaded,
+            kv_connector_cls=kv_connector_cls,
+        )
+        return
+        yield
+
+    def build(self) -> Engine:
+        gen = VllmBatchEngine(
+            engine=self._engine,
+            model_id=self.model_id,
+        )
+        try:
+            max_concurrent = (
+                1
+                if bool(os.getenv("EXO_NO_BATCH", False))
+                else EXO_MAX_CONCURRENT_REQUESTS
+            )
+        except Exception:
+            max_concurrent = EXO_MAX_CONCURRENT_REQUESTS
+
+        logger.info(f"using VllmEngine (max_concurrent={max_concurrent})")
+        return VllmEngine(
+            tool_parser=self._tool_parser,
+            model_id=self.model_id,
+            cancel_receiver=self.cancel_receiver,
+            event_sender=self.event_sender,
+            _gen=gen,
+            max_concurrent_requests=max_concurrent,
+        )
+
+    def close(self) -> None:
+        with contextlib.suppress(NameError, AttributeError):
+            del self._engine, self._tool_parser

src/exo/worker/engines/vllm/engine.py

@@ -0,0 +1,285 @@
+import itertools
+import time
+from collections import deque
+from collections.abc import Generator, Iterator
+from dataclasses import dataclass, field
+from typing import BinaryIO
+
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+from exo.shared.constants import EXO_MAX_CONCURRENT_REQUESTS
+from exo.shared.types.chunks import ErrorChunk, GenerationChunk, PrefillProgressChunk
+from exo.shared.types.common import ModelId
+from exo.shared.types.events import ChunkGenerated, Event
+from exo.shared.types.tasks import (
+    CANCEL_ALL_TASKS,
+    GenerationTask,
+    TaskId,
+    TextGeneration,
+)
+from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo.shared.types.worker.runner_response import (
+    CancelledResponse,
+    FinishedResponse,
+    GenerationResponse,
+)
+from exo.utils.channels import MpReceiver, MpSender
+from exo.worker.disaggregated.server import PrefillJob
+from exo.worker.engines.base import Engine
+from exo.worker.engines.vllm.prompt_format import format_vllm_prompt
+from exo.worker.engines.vllm.vllm_generator import VllmBatchEngine
+from exo.worker.runner.bootstrap import logger
+from exo.worker.runner.llm_inference.model_output_parsers import (
+    apply_all_parsers,
+    map_responses_to_chunks,
+)
+from exo.worker.runner.llm_inference.tool_parsers import ToolParser
+
+
+class GeneratorQueue[T]:
+    def __init__(self) -> None:
+        self._q = deque[T]()
+
+    def push(self, t: T) -> None:
+        self._q.append(t)
+
+    def gen(self) -> Generator[T | None]:
+        while True:
+            if len(self._q) == 0:
+                yield None
+            else:
+                yield self._q.popleft()
+
+
+EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
+EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"
+
+
+def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None:
+    """Keep the cheap debug prompt hooks without importing the MLX engine."""
+    if len(task_params.input) == 0:
+        return
+    prompt = task_params.input[0].content
+    if not prompt:
+        return
+    if EXO_RUNNER_MUST_FAIL in prompt:
+        raise Exception("Artificial runner exception - for testing purposes only.")
+    if EXO_RUNNER_MUST_TIMEOUT in prompt:
+        time.sleep(100)
+
+
+@dataclass(eq=False)
+class VllmEngine(Engine):
+    """Single-node vLLM implementation of the exo Engine interface.
+
+    This intentionally duplicates the local orchestration from the MLX
+    BatchGenerator instead of trying to share a batch abstraction too early.
+    The vLLM-specific tokenization/sampling/stepping remains inside
+    VllmBatchEngine.
+    """
+
+    tool_parser: ToolParser | None
+    model_id: ModelId
+    cancel_receiver: MpReceiver[TaskId]
+    event_sender: MpSender[Event]
+    _gen: VllmBatchEngine
+    max_concurrent_requests: int = EXO_MAX_CONCURRENT_REQUESTS
+    check_for_cancel_every: int = 50
+
+    _cancelled_tasks: set[TaskId] = field(default_factory=set, init=False)
+    _all_tasks: dict[TaskId, TextGeneration] = field(default_factory=dict, init=False)
+    _queue: deque[TextGeneration] = field(default_factory=deque, init=False)
+    _active_tasks: dict[
+        TaskId,
+        tuple[
+            TextGeneration,
+            GeneratorQueue[GenerationResponse],
+            Iterator[GenerationChunk | None],
+        ],
+    ] = field(default_factory=dict, init=False)
+
+    def warmup(self) -> None:
+        self.check_for_cancel_every = self._gen.warmup()
+
+    def submit(self, task: GenerationTask) -> None:
+        assert isinstance(task, TextGeneration)
+        self._cancelled_tasks.discard(CANCEL_ALL_TASKS)
+        self._all_tasks[task.task_id] = task
+        self._queue.append(task)
+
+    def step(
+        self,
+    ) -> Iterator[
+        tuple[TaskId, GenerationChunk | CancelledResponse | FinishedResponse]
+    ]:
+        self._collect_cancellations()
+        output: list[
+            tuple[TaskId, GenerationChunk | CancelledResponse | FinishedResponse]
+        ] = list(self._apply_cancellations())
+
+        while self._queue and len(self._active_tasks) < self.max_concurrent_requests:
+            task = self._queue.popleft()
+            if self.should_cancel(task.task_id):
+                output.append((task.task_id, CancelledResponse()))
+                self._all_tasks.pop(task.task_id, None)
+                continue
+
+            try:
+                task_id, queue, output_generator = self._start_task(task)
+            except Exception as e:
+                self._send_error(task, e)
+                self._all_tasks.pop(task.task_id, None)
+                raise
+
+            self._active_tasks[task_id] = (task, queue, output_generator)
+
+        if not self._gen.has_work:
+            return iter(output)
+
+        results = self._gen.step()
+        for task_id, response in results:
+            if task_id not in self._active_tasks:
+                logger.warning(f"{task_id=} not found in active vLLM tasks")
+                continue
+
+            task, queue, output_generator = self._active_tasks[task_id]
+            queue.push(response)
+            while (parsed := next(output_generator, None)) is not None:
+                output.append((task.task_id, parsed))
+
+            if response.finish_reason is not None:
+                output.append((task.task_id, FinishedResponse()))
+                del self._active_tasks[task_id]
+                self._all_tasks.pop(task.task_id, None)
+
+        return itertools.chain(output, self._apply_cancellations())
+
+    def close(self) -> None:
+        self._gen.close()
+
+    def serve_prefill(self, job: PrefillJob, wfile: BinaryIO) -> None:
+        raise NotImplementedError("vLLM serve_prefill is not supported yet")
+
+    def _start_task(
+        self, task: TextGeneration
+    ) -> tuple[
+        TaskId,
+        GeneratorQueue[GenerationResponse],
+        Iterator[GenerationChunk | None],
+    ]:
+        _check_for_debug_prompts(task.task_params)
+
+        token_ids, prompt_text, _ = format_vllm_prompt(
+            self._gen.engine, task.task_params
+        )
+
+        queue = GeneratorQueue[GenerationResponse]()
+        if task.task_params.bench:
+            output_generator: Iterator[GenerationChunk | None] = map(
+                lambda r: map_responses_to_chunks(r, self.model_id), queue.gen()
+            )
+        else:
+            output_generator = apply_all_parsers(
+                queue.gen(),
+                prompt_text,
+                self.tool_parser,
+                TokenizerWrapper(self._gen.engine.get_tokenizer()),
+                self.model_id,
+                task.task_params.tools,
+            )
+
+        check_for_cancel_every = max(self.check_for_cancel_every, 1)
+        tokens_since_cancel_check = check_for_cancel_every
+
+        def on_prefill_progress(processed: int, total: int) -> None:
+            self._collect_cancellations()
+            if self.should_cancel(task.task_id):
+                self._cancelled_tasks.add(task.task_id)
+            self.event_sender.send(
+                ChunkGenerated(
+                    command_id=task.command_id,
+                    chunk=PrefillProgressChunk(
+                        model=self.model_id,
+                        processed_tokens=processed,
+                        total_tokens=total,
+                    ),
+                )
+            )
+
+        def on_generation_token() -> None:
+            nonlocal tokens_since_cancel_check
+            tokens_since_cancel_check += 1
+            if tokens_since_cancel_check >= check_for_cancel_every:
+                tokens_since_cancel_check = 0
+                self._collect_cancellations()
+                if self.should_cancel(task.task_id):
+                    self._cancelled_tasks.add(task.task_id)
+
+        task_id = self._gen.submit(
+            task_id=task.task_id,
+            task_params=task.task_params,
+            on_prefill_progress=on_prefill_progress,
+            on_generation_token=on_generation_token,
+            token_ids=token_ids,
+        )
+        return task_id, queue, output_generator
+
+    def _collect_cancellations(self) -> None:
+        for task_id in self.cancel_receiver.collect():
+            if task_id == CANCEL_ALL_TASKS:
+                self._cancelled_tasks.add(CANCEL_ALL_TASKS)
+            elif task_id in self._all_tasks:
+                self._cancelled_tasks.add(task_id)
+
+    def _apply_cancellations(self) -> Iterator[tuple[TaskId, CancelledResponse]]:
+        if not self._cancelled_tasks:
+            return iter([])
+
+        cancel_all = CANCEL_ALL_TASKS in self._cancelled_tasks
+        results: list[tuple[TaskId, CancelledResponse]] = []
+        task_ids_to_abort: list[TaskId] = []
+
+        for task_id, (task, _, _) in list(self._active_tasks.items()):
+            if cancel_all or task.task_id in self._cancelled_tasks:
+                task_ids_to_abort.append(task_id)
+                results.append((task.task_id, CancelledResponse()))
+                del self._active_tasks[task_id]
+                self._all_tasks.pop(task.task_id, None)
+
+        if self._queue:
+            kept_queue: deque[TextGeneration] = deque()
+            for task in self._queue:
+                if cancel_all or task.task_id in self._cancelled_tasks:
+                    results.append((task.task_id, CancelledResponse()))
+                    self._all_tasks.pop(task.task_id, None)
+                else:
+                    kept_queue.append(task)
+            self._queue = kept_queue
+
+        if task_ids_to_abort:
+            self._gen.cancel(task_ids_to_abort)
+
+        already_cancelled = {task_id for task_id, _ in results}
+        for task_id in self._cancelled_tasks:
+            if (
+                task_id != CANCEL_ALL_TASKS
+                and task_id in self._all_tasks
+                and task_id not in already_cancelled
+            ):
+                results.append((task_id, CancelledResponse()))
+                self._all_tasks.pop(task_id, None)
+
+        self._cancelled_tasks.clear()
+        return iter(results)
+
+    def _send_error(self, task: TextGeneration, e: Exception) -> None:
+        self.event_sender.send(
+            ChunkGenerated(
+                command_id=task.command_id,
+                chunk=ErrorChunk(
+                    model=self.model_id,
+                    finish_reason="error",
+                    error_message=str(e),
+                ),
+            )
+        )

src/exo/worker/engines/vllm/generator.py

@@ -0,0 +1,477 @@
+import gc
+import json
+import math
+import re
+import sys
+import time
+from collections.abc import Callable, Generator
+from dataclasses import dataclass, field
+from typing import cast
+
+import torch
+from vllm.config import CompilationConfig
+from vllm.config.compilation import CUDAGraphMode
+from vllm.config.kv_transfer import KVTransferConfig
+from vllm.engine.arg_utils import EngineArgs
+from vllm.entrypoints.chat_utils import (
+    ChatCompletionMessageParam,
+    CustomChatCompletionMessageParam,
+)
+from vllm.outputs import RequestOutput
+from vllm.sampling_params import SamplingParams
+from vllm.tokenizers import TokenizerLike
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.engine.llm_engine import LLMEngine
+
+from exo.api.types import (
+    CompletionTokensDetails,
+    GenerationStats,
+    PromptTokensDetails,
+    Usage,
+)
+from exo.download.download_utils import build_model_path
+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
+
+# todo; move to different file
+from exo.worker.engines.mlx.utils_mlx import get_eos_token_ids_for_model
+from exo.worker.engines.vllm.prompt_format import (
+    make_vllm_sampling_params,
+)
+from exo.worker.runner.bootstrap import logger
+from exo.worker.runner.llm_inference.tool_parsers import ToolParser, infer_tool_parser
+
+
+@dataclass
+class _EngineRequest:
+    request_id: str
+    prompt_token_count: 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 _stop_token_ids(tokenizer: TokenizerLike, 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)  # pyright: ignore[reportAny]
+    extra = get_eos_token_ids_for_model(model_id)
+    if extra:
+        ids.update(extra)
+    return ids
+
+
+def _build_generation_response(
+    tokenizer: TokenizerLike,
+    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])
+    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()),
+        )
+        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 warmup_vllm_engine(engine: LLMEngine) -> int:
+    tokenizer = engine.get_tokenizer()
+    messages = [
+        cast(
+            ChatCompletionMessageParam,
+            CustomChatCompletionMessageParam(
+                role="user",
+                content="Prompt to warm up the inference engine. Repeat this.",
+            ),
+        )
+    ]
+    prompt_text: str | list[int] = tokenizer.apply_chat_template(  # pyright: ignore[reportUnknownMemberType]
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    if isinstance(prompt_text, list):
+        token_ids = prompt_text
+    else:
+        token_ids: list[int] = tokenizer.encode(prompt_text, add_special_tokens=False)
+
+    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
+
+    _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,
+        token_ids: list[int],
+        on_prefill_progress: Callable[[int, int], None] | None = None,
+        on_generation_token: Callable[[], None] | None = None,
+    ) -> TaskId:
+        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=len(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
+        )
+        results: list[tuple[TaskId, GenerationResponse]] = []
+
+        for output in outputs:
+            # todo: PoolingRequestOutputs
+            assert isinstance(output, RequestOutput)
+            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
+
+            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 = [str(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]]:
+    def patched(
+        hf_weights_files: list[str], *args: object, **kwargs: object
+    ) -> Generator[tuple[str, "torch.Tensor"], None, None]:
+        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):
+                yield name, tensor
+                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)
+
+    return patched
+
+
+def _monkey_patch_iterator(weight_utils: object, attr_name: str) -> None:
+    original = getattr(weight_utils, attr_name, None)
+    if original is None:
+        return
+    patched = _wrap_weights_iterator(original)  # pyright: ignore[reportAny]
+    setattr(weight_utils, attr_name, patched)
+    for mod in list(sys.modules.values()):
+        if 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,
+    )
+
+    _monkey_patch_iterator(weight_utils, "safetensors_weights_iterator")
+    _monkey_patch_iterator(weight_utils, "fastsafetensors_weights_iterator")
+
+    import huggingface_hub
+
+    def _noop_metadata(*_a: object, **_kw: object) -> None:
+        pass
+
+    original_metadata = huggingface_hub.get_safetensors_metadata
+    huggingface_hub.get_safetensors_metadata = _noop_metadata
+    for mod in list(sys.modules.values()):
+        if 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_id: ModelId,
+    trust_remote_code: bool,
+    n_layers: int = 1,
+    on_layer_loaded: Callable[[int, int], None] | None = None,
+    kv_connector_cls: type[object] | None = None,
+) -> tuple[LLMEngine, ToolParser | None]:
+    model_path = build_model_path(model_id)
+    _patch_weight_loading_progress()
+
+    # todo
+    # if kv_connector_cls is not None:
+    # from exo.disaggregated.prefill_server import _patch_vllm_for_connector
+
+    # _patch_vllm_for_connector(kv_connector_cls)
+
+    set_n_layers(n_layers)
+
+    kv_transfer_config: KVTransferConfig | None = None
+    if kv_connector_cls is not None:
+        kv_transfer_config = KVTransferConfig(
+            kv_connector=kv_connector_cls.__name__,
+            kv_connector_module_path=kv_connector_cls.__module__,
+            kv_role="kv_both",
+            kv_load_failure_policy="recompute",
+        )
+
+    has_mamba = False
+    try:
+        with open(model_path / "config.json") as f:
+            model_config = json.load(f)  # pyright: ignore[reportAny]
+        text_config = model_config.get("text_config", model_config)  # pyright: ignore[reportAny]
+        has_mamba = "mamba_ssm_dtype" in text_config or "linear_attention" in (
+            text_config.get("layer_types") or []  # pyright: ignore[reportAny]
+        )
+    except Exception:
+        pass
+    if has_mamba:
+        backends = [AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TRITON_ATTN]
+    else:
+        backends = [
+            AttentionBackendEnum.FLASHINFER,
+            AttentionBackendEnum.FLASH_ATTN,
+            AttentionBackendEnum.TRITON_ATTN,
+        ]
+
+    engine: LLMEngine | None = None
+    for backend in backends:
+        try:
+            engine_args = EngineArgs(
+                model=str(model_path.expanduser().resolve()),
+                served_model_name=str(model_id),
+                gpu_memory_utilization=0.05,
+                trust_remote_code=trust_remote_code,
+                load_format="fastsafetensors",
+                enable_prefix_caching=True,
+                attention_backend=backend,
+                compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE),
+                disable_log_stats=True,
+                max_num_batched_tokens=4096,
+                kv_transfer_config=kv_transfer_config,
+                disable_hybrid_kv_cache_manager=False,
+            )
+
+            set_weight_loading_callback(on_layer_loaded)
+            engine = LLMEngine.from_engine_args(engine_args)
+            logger.info(f"vLLM engine using attention backend: {backend}")
+            break
+        except (ValueError, RuntimeError, NotImplementedError) as e:
+            logger.warning(f"Attention backend {backend} failed: {e}, trying next")
+            engine = None
+
+            gc.collect()
+            torch.cuda.empty_cache()
+            continue
+
+    if engine is None:
+        raise RuntimeError(f"No attention backend worked for {model_id}")
+
+    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

src/exo/worker/engines/vllm/growable_cache.py

@@ -0,0 +1,565 @@
+# pyright: reportPrivateUsage=false, reportAttributeAccessIssue=false
+from typing import TYPE_CHECKING, Any, cast
+
+import torch
+from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.request import Request
+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
+
+if TYPE_CHECKING:
+    from vllm.v1.core.kv_cache_manager import KVCacheManager
+
+
+_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_nogds()
+    _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_nogds() -> None:
+    from vllm.model_executor.model_loader import weight_utils
+
+    original = weight_utils._init_fastsafetensors_loader
+
+    def patched(  # pyright: ignore[reportUnknownParameterType]
+        pg: "torch.distributed.ProcessGroup",
+        device: "torch.device",
+        f_list: list[str],
+        *,
+        nogds: bool = False,
+    ):
+        return original(pg, device, f_list, nogds=True)  # pyright: ignore[reportUnknownVariableType]
+
+    weight_utils._init_fastsafetensors_loader = patched
+
+
+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:
+        import pathlib
+        import shutil
+
+        compile_cache = pathlib.Path.home() / ".cache" / "vllm" / "torch_compile_cache"
+        if compile_cache.exists():
+            shutil.rmtree(compile_cache, ignore_errors=True)
+
+        free_bytes, _ = torch.cuda.mem_get_info()
+        initial = max(int(free_bytes * INITIAL_FRACTION), 1)
+        self._growable_max_kv_bytes = free_bytes
+        self.available_kv_cache_memory_bytes = initial
+        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
+
+
+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
+
+
+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: KVCacheConfig
+    ) -> 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
+
+    original_init_tensors = GPUModelRunner.initialize_kv_cache_tensors
+
+    def patched_init_tensors(
+        self: GPUModelRunner,
+        kv_cache_config: KVCacheConfig,
+        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
+
+
+def _patch_initialize_from_config() -> None:
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+    from vllm.v1.worker.gpu_worker import Worker
+
+    original_init_attn = GPUModelRunner.initialize_attn_backend
+
+    def clear_and_reinit_attn(
+        self: GPUModelRunner,
+        kv_cache_config: KVCacheConfig,
+    ) -> None:
+        self.attn_groups.clear()
+        original_init_attn(self, kv_cache_config)
+
+    GPUModelRunner.initialize_attn_backend = clear_and_reinit_attn
+
+    original = Worker.initialize_from_config
+
+    def patched(self: Worker, kv_cache_config: KVCacheConfig) -> None:
+        original(self, kv_cache_config)
+        set_model_runner(self.model_runner)
+
+    Worker.initialize_from_config = patched
+
+
+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,
+        kv_cache_config: KVCacheConfig,
+        max_model_len: int,
+        hash_block_size: int,
+        enable_caching: bool = True,
+        use_eagle: bool = False,
+        log_stats: bool = False,
+        enable_kv_cache_events: bool = False,
+        dcp_world_size: int = 1,
+        pcp_world_size: int = 1,
+        metrics_collector: KVCacheMetricsCollector | None = None,
+    ) -> None:
+        original_init(
+            self,
+            kv_cache_config,
+            max_model_len,
+            hash_block_size,
+            enable_caching,
+            use_eagle,
+            log_stats,
+            enable_kv_cache_events,
+            dcp_world_size,
+            pcp_world_size,
+            metrics_collector,
+        )
+        self._growable_model_runner = get_model_runner()
+
+    KVCacheManager.__init__ = patched_init
+
+
+def _patch_allocate_slots() -> None:
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
+
+    original = KVCacheManager.allocate_slots
+
+    def patched(
+        self: KVCacheManager,
+        request: Request,
+        num_new_tokens: int,
+        num_new_computed_tokens: int = 0,
+        new_computed_blocks: KVCacheBlocks | None = None,
+        num_lookahead_tokens: int = 0,
+        num_external_computed_tokens: int = 0,
+        delay_cache_blocks: bool = False,
+        num_encoder_tokens: int = 0,
+    ) -> KVCacheBlocks | None:
+        result = original(
+            self,
+            request,
+            num_new_tokens,
+            num_new_computed_tokens,
+            new_computed_blocks,
+            num_lookahead_tokens,
+            num_external_computed_tokens,
+            delay_cache_blocks,
+            num_encoder_tokens,
+        )
+        while result is None and _try_grow_cache(self):
+            result = original(
+                self,
+                request,
+                num_new_tokens,
+                num_new_computed_tokens,
+                new_computed_blocks,
+                num_lookahead_tokens,
+                num_external_computed_tokens,
+                delay_cache_blocks,
+                num_encoder_tokens,
+            )
+
+        return result
+
+    KVCacheManager.allocate_slots = patched
+
+    if hasattr(KVCacheManager, "can_fit_full_sequence"):
+        original_can_fit = KVCacheManager.can_fit_full_sequence
+
+        def patched_can_fit(
+            self: KVCacheManager,
+            request: Request,
+            num_new_computed_tokens: int = 0,
+            new_computed_blocks: KVCacheBlocks | None = None,
+            num_external_computed_tokens: int = 0,
+            num_encoder_tokens: int = 0,
+        ) -> bool:
+            result = original_can_fit(
+                self,
+                request,
+                num_new_computed_tokens,
+                new_computed_blocks,
+                num_external_computed_tokens,
+                num_encoder_tokens,
+            )
+            while not result and _try_grow_cache(self):
+                result = original_can_fit(
+                    self,
+                    request,
+                    num_new_computed_tokens,
+                    new_computed_blocks,
+                    num_external_computed_tokens,
+                    num_encoder_tokens,
+                )
+            return result
+
+        KVCacheManager.can_fit_full_sequence = patched_can_fit
+
+
+def _try_grow_cache(kv_cache_manager: "KVCacheManager") -> bool:
+    block_pool = kv_cache_manager.block_pool
+    model_runner = cast(GPUModelRunner | None, kv_cache_manager._growable_model_runner)
+
+    if model_runner is None:
+        return False
+
+    free_bytes, _ = torch.cuda.mem_get_info()
+    if free_bytes < GROWTH_HEADROOM_BYTES:
+        return False
+
+    kv_cache_config = cast(KVCacheConfig, model_runner._growable_kv_cache_config)
+    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:
+        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:
+        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)
+        _grow_tensors(model_runner, kv_cache_config, old_num_blocks, new_num_blocks)
+        _grow_block_pool(block_pool, old_num_blocks, new_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: GPUModelRunner,
+    kv_cache_config: KVCacheConfig,
+    old_num_blocks: int,
+    new_num_blocks: int,
+) -> None:
+    raw_tensors: dict[str, torch.Tensor] = cast(
+        dict[str, torch.Tensor], model_runner._growable_raw_tensors
+    )
+    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
+
+    kernel_block_sizes: list[int] = cast(
+        list[int], model_runner._growable_kernel_block_sizes
+    )
+    new_kv_caches: dict[str, torch.Tensor] = model_runner._reshape_kv_cache_tensors(
+        kv_cache_config,
+        new_raw_tensors,
+        kernel_block_sizes,
+    )
+
+    forward_context: dict[str, Any] = (
+        model_runner.compilation_config.static_forward_context
+    )
+    runner_kv_caches: list[torch.Tensor] = model_runner.kv_caches
+
+    from collections import defaultdict
+
+    from vllm.model_executor.models.utils import extract_layer_index
+
+    num_attn_module = 1
+    hf_config = getattr(getattr(model_runner, "model_config", None), "hf_config", None)
+    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)
+
+    new_ordered: list[torch.Tensor] = []
+    for layer_index in sorted(index2name.keys()):
+        for ln in index2name[layer_index]:
+            new_ordered.append(new_kv_caches[ln])
+
+    for i, new_kv in enumerate(new_ordered):
+        if i < len(runner_kv_caches):
+            runner_kv_caches[i] = new_kv
+        else:
+            runner_kv_caches.append(new_kv)
+
+    for layer_name, new_kv in new_kv_caches.items():
+        old_kv_list = forward_context[layer_name].kv_cache
+        if old_kv_list is not None and (
+            not isinstance(old_kv_list, torch.Tensor) or old_kv_list.numel() > 0
+        ):
+            old_entry = old_kv_list[0]
+            if isinstance(old_entry, list) and isinstance(new_kv, list):
+                for old_t, new_t in zip(old_entry, new_kv):
+                    old_t.set_(
+                        new_t.storage(),
+                        new_t.storage_offset(),
+                        new_t.shape,
+                        new_t.stride(),
+                    )
+            elif isinstance(old_entry, torch.Tensor) and isinstance(
+                new_kv, torch.Tensor
+            ):
+                old_entry.set_(
+                    new_kv.storage(),
+                    new_kv.storage_offset(),
+                    new_kv.shape,
+                    new_kv.stride(),
+                )
+            else:
+                forward_context[layer_name].kv_cache = [new_kv]
+        else:
+            forward_context[layer_name].kv_cache = [new_kv]
+
+
+def _grow_block_pool(
+    block_pool: BlockPool, 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)
+        new_blocks.append(block)
+
+    block_pool.free_block_queue.append_n(new_blocks)
+    block_pool.num_gpu_blocks = new_num_blocks
+
+
+def _patch_moe_sum() -> None:
+    import vllm._custom_ops as ops
+
+    def moe_sum_f32(x: "torch.Tensor", output: "torch.Tensor") -> None:
+        output[:] = x.to(torch.float32).sum(dim=1).to(output.dtype)
+
+    ops.moe_sum = moe_sum_f32
+
+
+def _patch_marlin_w2_thread_config() -> None:
+    try:
+        import vllm._custom_ops as ops
+    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
+
+
+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)
+            )
+        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
+
+        from exo.worker.engines.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
+                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]
+            block_size: int = self.kv_cache_config.kv_cache_groups[
+                gi
+            ].kv_cache_spec.block_size
+            num_skipped: int = mgr.get_num_skipped_tokens(num_matched)
+            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]
+            block_size = self.kv_cache_config.kv_cache_groups[
+                gi
+            ].kv_cache_spec.block_size
+            real_blocks: list[KVCacheBlock] = self.block_pool.get_new_blocks(
+                real_block_counts[gi]
+            )
+            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]
+            req_blocks.extend(full_block_list)
+
+            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(
+                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

src/exo/worker/engines/vllm/kv_cache.py

@@ -0,0 +1,346 @@
+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,
+        block_sizes_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
+
+            if k_all.dim() >= 4 and len(bt) > 0 and block_sizes_per_group is not None:
+                page_size = k_all.shape[1]
+                sched_block_size = block_sizes_per_group[gi]
+                pages_per_block = sched_block_size // page_size
+                if pages_per_block > 1:
+                    expanded = []
+                    for b in bt.tolist():
+                        start_page = b * pages_per_block
+                        end_page = min(start_page + pages_per_block, k_all.shape[0])
+                        expanded.extend(range(start_page, end_page))
+                    bt = torch.tensor(expanded, dtype=torch.long)
+
+            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 isinstance(layer, ArraysLayerState):
+                gi = layer_to_group[layer_idx]
+                bt = block_tables[gi]
+                kv = kv_caches[layer_idx]
+                if isinstance(kv, list):
+                    for ti, (stored, target) in enumerate(zip(layer.arrays, kv)):
+                        if stored is not None and target is not None:
+                            n = min(len(bt), stored.shape[0])
+                            if n > 0:
+                                target[bt[:n]] = stored[:n].to(
+                                    device, non_blocking=True
+                                )
+                continue
+            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)
+
+            keys = layer.keys
+            values = layer.values
+            block_size = k_all.shape[-3] if k_all.dim() >= 3 else k_all.shape[1]
+            needs_reshape = keys.dim() == 3 and keys.shape[1:] != k_all.shape[1:]
+            if needs_reshape:
+                offset = token_offset_per_group[gi] if token_offset_per_group else 0
+                if offset > 0:
+                    keys = keys[offset:]
+                    values = values[offset:]
+                s, h, d = keys.shape
+                pad = (block_size - s % block_size) % block_size
+                if pad > 0:
+                    keys = torch.nn.functional.pad(keys, (0, 0, 0, 0, 0, pad))
+                    values = torch.nn.functional.pad(values, (0, 0, 0, 0, 0, pad))
+                keys = keys.reshape(-1, block_size, h, d)
+                values = values.reshape(-1, block_size, h, d)
+
+            n_blocks = min(len(bt), keys.shape[0])
+            if n_blocks > 0:
+                k_all[bt[:n_blocks]] = keys[:n_blocks].to(device, non_blocking=True)
+                v_all[bt[:n_blocks]] = 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)

src/exo/worker/engines/vllm/prompt_format.py

@@ -0,0 +1,61 @@
+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)
+    return token_ids, prompt_text, len(token_ids)
+
+
+def make_vllm_sampling_params(
+    engine: LLMEngine,
+    params: TextGenerationTaskParams,
+    model_id: ModelId | None = None,
+) -> SamplingParams:
+    kwargs: SamplingParams = SamplingParams()
+
+    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
+
+    if params.bench:
+        kwargs.ignore_eos = True
+        kwargs.min_tokens = kwargs.max_tokens
+
+    return kwargs

src/exo/worker/runner/bootstrap.py

@@ -5,7 +5,7 @@ import loguru
 
 from exo.shared.types.events import Event, RunnerStatusUpdated
 from exo.shared.types.tasks import Task, TaskId
-from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.instances import BoundInstance, VllmInstance
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
 from exo.worker.engines.base import Builder
@@ -46,6 +46,21 @@ def entrypoint(
             builder = MfluxBuilder(
                 event_sender, cancel_receiver, bound_instance.bound_shard
             )
+        elif isinstance(bound_instance.instance, VllmInstance):
+            from exo.worker.engines.vllm.builder import VllmBuilder
+            from exo.worker.engines.vllm.growable_cache import patch_vllm
+
+            os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+            os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
+            os.environ["VLLM_BATCH_INVARIANT"] = "1"
+            os.environ.setdefault("FASTSAFETENSORS_NOGDS", "1")
+
+            patch_vllm()
+            builder = VllmBuilder(
+                bound_instance.bound_shard.model_card.model_id,
+                event_sender,
+                cancel_receiver,
+            )
         else:
             from exo.worker.engines.mlx.patches import apply_mlx_patches
 

src/exo/worker/runner/llm_inference/batch_generator.py

@@ -221,7 +221,6 @@ class SequentialGenerator(Engine):
                 apply_chat_template(self.tokenizer, task.task_params),
                 self.tool_parser,
                 self.tokenizer,
-                type(self.model),
                 self.model_id,
                 task.task_params.tools,
             )
@@ -418,7 +417,6 @@ class BatchGenerator(Engine):
                     apply_chat_template(self.tokenizer, task.task_params),
                     self.tool_parser,
                     self.tokenizer,
-                    type(self.model),
                     self.model_id,
                     task.task_params.tools,
                 )

src/exo/worker/runner/llm_inference/model_output_parsers.py

@@ -2,9 +2,6 @@ from collections.abc import Callable, Generator, Iterator
 from functools import cache
 from typing import Any
 
-from mlx_lm.models.deepseek_v4 import Model as DeepseekV4Model
-from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
-from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
     HarmonyEncodingName,
@@ -23,7 +20,6 @@ from exo.shared.types.chunks import (
 )
 from exo.shared.types.common import ModelId
 from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
-from exo.worker.engines.mlx.types import Model
 from exo.worker.engines.mlx.utils_mlx import (
     detect_thinking_prompt_suffix,
 )
@@ -69,16 +65,15 @@ def apply_all_parsers(
     prompt: str,
     tool_parser: ToolParser | None,
     tokenizer: TokenizerWrapper,
-    model_type: type[Model],
     model_id: ModelId,
     tools: list[dict[str, Any]] | None,
 ) -> Iterator[GenerationChunk | None]:
     generator = receiver
 
-    normalized_id = model_id.normalize().lower()
-    if issubclass(model_type, GptOssModel):
+    normalized_id = model_id.short().lower()
+    if "gpt-oss" in normalized_id:
         generator = parse_gpt_oss(generator)
-    elif issubclass(model_type, DeepseekV32Model) and "deepseek" in normalized_id:
+    elif "deepseek-v3.2" in normalized_id:
         if tokenizer.has_thinking:
             generator = parse_thinking_models(
                 generator,
@@ -87,7 +82,7 @@ def apply_all_parsers(
                 starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer),
             )
         generator = parse_deepseek_v32(generator)
-    elif issubclass(model_type, DeepseekV4Model) and "deepseek-v4" in normalized_id:
+    elif "deepseek-v4" in normalized_id:
         if tokenizer.has_thinking:
             generator = parse_thinking_models(
                 generator,