bump versions correctly

bench_overnight.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+set -e
+
+export PATH="/opt/homebrew/bin:$PATH"
+
+echo "=== Starting overnight bench runs at $(date) ==="
+
+echo "--- [4/8] Qwen3.5-122B-A10B-GPTQ-Int4 ---"
+echo "Skipping because Int 4"
+#uv run bench/exo_bench.py --force-download --model "Qwen/Qwen3.5-122B-A10B-GPTQ-Int4" --pp 700 --tg 36000 --repeat 1
+
+echo "--- [5/8] Qwen3.5-27B-FP8 ---"
+#uv run bench/exo_bench.py --force-download --model "Qwen/Qwen3.5-27B-FP8" --pp 700 --tg 35133 --repeat 1
+
+echo "--- [6/8] GLM-4.7-Flash-bf16 ---"
+uv run bench/exo_bench.py --force-download --model "mlx-community/GLM-4.7-Flash-bf16" --pp 700 --tg 29000 --repeat 1
+
+echo "--- [7/8] NVIDIA-Nemotron-3-Nano-30B-A3B (23000,1200) ---"
+uv run bench/exo_bench.py --force-download --model "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-BF16" --pp 700 --tg 23000,1200 --repeat 1
+
+echo "--- [8/8] Qwen3.5-27B-bf16 ---"
+uv run bench/exo_bench.py --force-download --model "mlx-community/Qwen3.5-27B-bf16" --pp 700 --tg 35400 --repeat 1
+
+echo "=== All bench runs complete at $(date) ==="

dashboard/src/lib/components/ChatModelSelector.svelte

@@ -9,6 +9,10 @@
     quantization: string;
   }
 
+  function normalizeBaseModel(s: string): string {
+    return s.toLowerCase().replace(/[-_]/g, " ").trim();
+  }
+
   // Auto mode tier list (for when user just starts typing)
   export const AUTO_TIERS: string[][] = [
     // Tier 1 (frontier)
@@ -43,8 +47,9 @@
 
   /** Return the tier index (0 = best) for a base_model name. */
   export function getAutoTierIndex(baseModel: string): number {
+    const norm = normalizeBaseModel(baseModel);
     for (let i = 0; i < AUTO_TIERS.length; i++) {
-      if (AUTO_TIERS[i].includes(baseModel)) return i;
+      if (AUTO_TIERS[i].some((t) => normalizeBaseModel(t) === norm)) return i;
     }
     return AUTO_TIERS.length; // not in any tier → lowest priority
   }
@@ -60,7 +65,8 @@
         const variants = modelList
           .filter(
             (m) =>
-              m.base_model === baseModel &&
+              normalizeBaseModel(m.base_model) ===
+                normalizeBaseModel(baseModel) &&
               (m.storage_size_megabytes || 0) / 1024 <= memoryGB &&
               (m.storage_size_megabytes || 0) > 0,
           )
@@ -162,7 +168,11 @@
   /** For a given base_model name, find the biggest quant variant that fits in memory. */
   function pickBestVariant(baseModel: string): ChatModelInfo | null {
     const variants = models
-      .filter((m) => m.base_model === baseModel && fitsInMemory(m))
+      .filter(
+        (m) =>
+          normalizeBaseModel(m.base_model) === normalizeBaseModel(baseModel) &&
+          fitsInMemory(m),
+      )
       .sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
     return variants[0] ?? null;
   }

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -379,12 +379,16 @@
     return hfTrendingModels;
   });
 
+  function normalizeBaseModel(s: string): string {
+    return s.toLowerCase().replace(/[-_]/g, " ").trim();
+  }
+
   // Group models by base_model
   const groupedModels = $derived.by((): ModelGroup[] => {
     const groups = new Map<string, ModelGroup>();
 
     for (const model of models) {
-      const groupId = model.base_model || model.id;
+      const groupId = normalizeBaseModel(model.base_model || model.id);
       const groupName = model.base_model || model.name || model.id;
 
       if (!groups.has(groupId)) {
@@ -578,7 +582,7 @@
       const model = models.find((m) => m.id === id);
       if (model) {
         result.push({
-          id: model.base_model || model.id,
+          id: normalizeBaseModel(model.base_model || model.id),
           name: model.name || model.id,
           capabilities: model.capabilities || ["text"],
           family: model.family || "",

dashboard/src/routes/+page.svelte

@@ -295,7 +295,9 @@
     const seen = new Set<string>();
     const deduped: typeof candidates = [];
     for (const m of candidates) {
-      const key = m.base_model || m.family || m.id;
+      const key = (m.base_model || m.family || m.id)
+        .toLowerCase()
+        .replace(/[-_]/g, " ");
       if (seen.has(key)) continue;
       seen.add(key);
       deduped.push(m);

flake.lock

@@ -2,11 +2,11 @@
   "nodes": {
     "crane": {
       "locked": {
-        "lastModified": 1767744144,
-        "narHash": "sha256-9/9ntI0D+HbN4G0TrK3KmHbTvwgswz7p8IEJsWyef8Q=",
+        "lastModified": 1774313767,
+        "narHash": "sha256-hy0XTQND6avzGEUFrJtYBBpFa/POiiaGBr2vpU6Y9tY=",
         "owner": "ipetkov",
         "repo": "crane",
-        "rev": "2fb033290bf6b23f226d4c8b32f7f7a16b043d7e",
+        "rev": "3d9df76e29656c679c744968b17fbaf28f0e923d",
         "type": "github"
       },
       "original": {
@@ -47,11 +47,11 @@
         "rust-analyzer-src": "rust-analyzer-src"
       },
       "locked": {
-        "lastModified": 1768287139,
-        "narHash": "sha256-nsXFt0OzUi6K7dUzzJD5/v9e0Ic+fvclfIW936/43ZM=",
+        "lastModified": 1774596377,
+        "narHash": "sha256-DiSLMxyTwIUAlhOe34r6kKNQRv6PTF+vf0MG45mAyn4=",
         "owner": "nix-community",
         "repo": "fenix",
-        "rev": "a4a3aa956931f90f35453cb519e4545e9ad7f773",
+        "rev": "a88a1c8cf2f094da6347fcec54089f4bcb518409",
         "type": "github"
       },
       "original": {
@@ -83,11 +83,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1768135262,
-        "narHash": "sha256-PVvu7OqHBGWN16zSi6tEmPwwHQ4rLPU9Plvs8/1TUBY=",
+        "lastModified": 1772408722,
+        "narHash": "sha256-rHuJtdcOjK7rAHpHphUb1iCvgkU3GpfvicLMwwnfMT0=",
         "owner": "hercules-ci",
         "repo": "flake-parts",
-        "rev": "80daad04eddbbf5a4d883996a73f3f542fa437ac",
+        "rev": "f20dc5d9b8027381c474144ecabc9034d6a839a3",
         "type": "github"
       },
       "original": {
@@ -164,11 +164,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1763662255,
-        "narHash": "sha256-4bocaOyLa3AfiS8KrWjZQYu+IAta05u3gYZzZ6zXbT0=",
+        "lastModified": 1773870109,
+        "narHash": "sha256-ZoTdqZP03DcdoyxvpFHCAek4bkPUTUPUF3oCCgc3dP4=",
         "owner": "pyproject-nix",
         "repo": "build-system-pkgs",
-        "rev": "042904167604c681a090c07eb6967b4dd4dae88c",
+        "rev": "b6e74f433b02fa4b8a7965ee24680f4867e2926f",
         "type": "github"
       },
       "original": {
@@ -184,11 +184,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1764134915,
-        "narHash": "sha256-xaKvtPx6YAnA3HQVp5LwyYG1MaN4LLehpQI8xEdBvBY=",
+        "lastModified": 1774498001,
+        "narHash": "sha256-wTfdyzzrmpuqt4TQQNqilF91v0m5Mh1stNy9h7a/WK4=",
         "owner": "pyproject-nix",
         "repo": "pyproject.nix",
-        "rev": "2c8df1383b32e5443c921f61224b198a2282a657",
+        "rev": "794afa6eb588b498344f2eaa36ab1ceb7e6b0b09",
         "type": "github"
       },
       "original": {
@@ -214,11 +214,11 @@
     "rust-analyzer-src": {
       "flake": false,
       "locked": {
-        "lastModified": 1768224240,
-        "narHash": "sha256-Pp1dDrXKPBUJReZnnDElFyHYn67XTd48zRhToheLjtk=",
+        "lastModified": 1774569884,
+        "narHash": "sha256-E8iWEPzg7OnE0XXXjo75CX7xFauqzJuGZ5wSO9KS8Ek=",
         "owner": "rust-lang",
         "repo": "rust-analyzer",
-        "rev": "725349602e525df37f377701e001fe8aab807878",
+        "rev": "443ddcddd0c73b07b799d052f5ef3b448c2f3508",
         "type": "github"
       },
       "original": {
@@ -257,11 +257,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1768158989,
-        "narHash": "sha256-67vyT1+xClLldnumAzCTBvU0jLZ1YBcf4vANRWP3+Ak=",
+        "lastModified": 1773297127,
+        "narHash": "sha256-6E/yhXP7Oy/NbXtf1ktzmU8SdVqJQ09HC/48ebEGBpk=",
         "owner": "numtide",
         "repo": "treefmt-nix",
-        "rev": "e96d59dff5c0d7fddb9d113ba108f03c3ef99eca",
+        "rev": "71b125cd05fbfd78cab3e070b73544abe24c5016",
         "type": "github"
       },
       "original": {
@@ -280,11 +280,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1767701098,
-        "narHash": "sha256-CJhKZnWb3gumR9oTRjFvCg/6lYTGbZRU7xtvcyWIRwU=",
+        "lastModified": 1774490495,
+        "narHash": "sha256-a9WmQWj8fF7BctZGCoyzpUjP6GJw8H+lxl+zxpGnETk=",
         "owner": "pyproject-nix",
         "repo": "uv2nix",
-        "rev": "9d357f0d2ce6f5f35ec7959d7e704452352eb4da",
+        "rev": "18ae62fc5e389e3069854a7c66455c22e31708fc",
         "type": "github"
       },
       "original": {

flake.nix

@@ -72,19 +72,26 @@
       ];
 
       perSystem =
-        { config, self', inputs', pkgs, lib, system, ... }:
+        { config, self', pkgs, lib, system, ... }:
         let
           # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
           pkgsSwift = import inputs.nixpkgs-swift { inherit system; };
-
-          pkgsCuda = import ./nix/cuda-pkgs.nix { nixpkgs = inputs.nixpkgs; inherit system; };
         in
         {
+          _module.args.cudaPkgs = import inputs.nixpkgs
+            {
+              inherit system;
+              config = {
+                allowUnfreePredicate = pkg: builtins.elem (lib.getName pkg) [ "cuda-merged" "cuda_cuobjdump" "cuda_gdb" "cuda_nvcc" "cuda_nvdisasm" "cuda_nvprune" "cuda_cccl" "cuda_cudart" "cuda_cupti" "cuda_cuxxfilt" "cuda_nvml_dev" "cuda_nvrtc" "cuda_nvtx" "cuda_profiler_api" "cuda_sanitizer_api" "libcublas" "libcufft" "libcurand" "libcusolver" "libnvjitlink" "libcusparse" "libnpp" "cudnn" "libcusparse_lt" "libcufile" "libnvshmem" "libnvvm" "cuda_crt" ];
+                cudaSupport = true;
+                cudaCapabilities = [ "12.1" ];
+              };
+            };
           # Allow unfree for metal-toolchain (needed for Darwin Metal packages)
           _module.args.pkgs = import inputs.nixpkgs {
             inherit system;
-            config.allowUnfreePredicate = pkg: (pkg.pname or "") == "metal-toolchain";
-            overlays = [
+            config.allowUnfreePredicate = pkg: builtins.elem (lib.getName pkg) [ "metal-toolchain" ];
+            overlays = lib.optionals (system == "aarch64-darwin") [
               (import ./nix/apple-sdk-overlay.nix)
             ];
           };
@@ -114,138 +121,68 @@
             };
           };
 
-          packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin
-            (
-              let
-                uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
-                mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package);
-                uvLockMlxVersion = mlxPackage.version;
-                uvLockMlxRev = builtins.elemAt (builtins.split "#" mlxPackage.source.git) 2;
-              in
-              {
-                metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
-                mlx = pkgs.callPackage ./nix/mlx.nix {
-                  inherit (self'.packages) metal-toolchain;
-                  inherit uvLockMlxVersion uvLockMlxRev;
-                };
-                default = self'.packages.exo;
-              }
-            ) // lib.optionalAttrs (pkgsCuda != null) {
-            torch-cuda = pkgsCuda.python313Packages.torch;
-            vllm-cuda = pkgsCuda.python313Packages.vllm;
+          packages = (lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin
 
-            # Smoke test script for verifying vLLM + CUDA GPU setup
-            vllm-check = pkgs.writeShellApplication {
-              name = "vllm-check";
-              runtimeInputs = [
-                (pkgsCuda.python313.withPackages (ps: [ ps.torch ps.vllm ]))
-              ];
-              # On non-NixOS hosts, NVIDIA driver libraries live in /usr/lib and must be
-              # LD_PRELOAD'd individually (adding the whole dir causes SIGILL from conflicts).
-              # These are: CUDA driver, NVML, and the PTX JIT compiler (for flash attention).
-              # libnvJitLink comes from the nix CUDA toolkit via LD_LIBRARY_PATH.
-              text = ''
-                for dir in /usr/lib/aarch64-linux-gnu /usr/lib/x86_64-linux-gnu /usr/lib; do
-                  if [ -e "$dir/libcuda.so.1" ]; then
-                    NVIDIA_LIBS="$dir/libcuda.so.1"
-                    for lib in libnvidia-ml.so.1 libnvidia-ptxjitcompiler.so.1; do
-                      [ -e "$dir/$lib" ] && NVIDIA_LIBS="$NVIDIA_LIBS:$dir/$lib"
-                    done
-                    export LD_PRELOAD="$NVIDIA_LIBS''${LD_PRELOAD:+:$LD_PRELOAD}"
-                    break
-                  fi
-                done
-                export LD_LIBRARY_PATH="${pkgsCuda.cudaPackages.libnvjitlink}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
-                exec python ${inputs.self + /tests/test_vllm_smoke.py}
-              '';
-            };
-
-            # exo with CUDA torch + vLLM — wraps the uv2nix-built package with host driver libs
-            exo-cuda = pkgs.writeShellApplication {
-              name = "exo-cuda";
-              runtimeInputs = [ self'.packages.exo-cuda-unwrapped ];
-              text = ''
-                for dir in /usr/lib/aarch64-linux-gnu /usr/lib/x86_64-linux-gnu /usr/lib; do
-                  if [ -e "$dir/libcuda.so.1" ]; then
-                    NVIDIA_LIBS="$dir/libcuda.so.1"
-                    for lib in libnvidia-ml.so.1 libnvidia-ptxjitcompiler.so.1; do
-                      [ -e "$dir/$lib" ] && NVIDIA_LIBS="$NVIDIA_LIBS:$dir/$lib"
-                    done
-                    export LD_PRELOAD="$NVIDIA_LIBS''${LD_PRELOAD:+:$LD_PRELOAD}"
-                    break
-                  fi
-                done
-                export LD_LIBRARY_PATH="${pkgsCuda.stdenv.cc.cc.lib}/lib:${pkgsCuda.cudaPackages.libnvjitlink}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
-                exec exo-cuda "$@"
-              '';
-            };
-          };
-
-          # CUDA development shell with torch + vLLM (aarch64-linux only)
-          devShells = lib.optionalAttrs (pkgsCuda != null)
             {
-              cuda = pkgs.mkShell {
-                packages = [
-                  (pkgsCuda.python313.withPackages (ps: [
-                    ps.torch
-                    ps.vllm
-                  ]))
-                  pkgs.uv
-                  pkgs.just
-                ];
+              metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
+            }
+          ) // {
+            default = self'.packages.exo;
+          };
+          devShells =
+            {
+              default = with pkgs; mkShell {
+                inputsFrom = [ self'.checks.cargo-build ];
+
+                packages =
+                  [
+                    # FORMATTING
+                    config.treefmt.build.wrapper
+
+                    # PYTHON
+                    self'.packages.editable-venv
+                    uv
+
+                    # RUST
+                    config.rust.toolchain
+                    maturin
+
+                    # NIX
+                    nixpkgs-fmt
+
+                    # SVELTE
+                    nodejs
+
+                    # MISC
+                    just
+                    jq
+                  ]
+                  ++ lib.optionals stdenv.isLinux [
+                    unixtools.ifconfig
+                  ]
+                  ++ lib.optionals stdenv.isDarwin [
+                    macmon
+                  ];
+
+                env = {
+                  UV_NO_SYNC = "1";
+                  UV_PYTHON = "${self'.packages.editable-venv}/bin/python";
+                  UV_PYTHON_DOWNLOADS = "never";
+                  UV_PROJECT_ENVIRONMENT = self'.packages.editable-venv;
+                  VIRTUAL_ENV = self'.packages.editable-venv;
+                  OPENSSL_NO_VENDOR = "1";
+                };
 
                 shellHook = ''
-                  echo "CUDA dev shell with torch + vLLM"
-                  python -c "import torch; print(f'PyTorch {torch.__version__}, CUDA: {torch.cuda.is_available()}')" 2>/dev/null || true
+                  unset PYTHONPATH
+                  export REPO_ROOT=$(git rev-parse --show-toplevel)
+                  export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${self'.packages.editable-venv}/lib"
+                  ${lib.optionalString stdenv.isLinux ''
+                    export LD_LIBRARY_PATH="${openssl.out}/lib:${lib.getLib pkgs.util-linux}/lib:${lib.getLib pkgs.systemd}/lib:${lib.getLib pkgs.numactl}/lib:${lib.getLib pkgs.stdenv.cc.cc.lib}/lib:$LD_LIBRARY_PATH"
+                  ''}
                 '';
               };
-            } // {
-
-            default = with pkgs; pkgs.mkShell {
-              inputsFrom = [ self'.checks.cargo-build ];
-
-              packages =
-                [
-                  # FORMATTING
-                  config.treefmt.build.wrapper
-
-                  # PYTHON
-                  python313
-                  uv
-                  ruff
-                  basedpyright
-
-                  # RUST
-                  config.rust.toolchain
-                  maturin
-
-                  # NIX
-                  nixpkgs-fmt
-
-                  # SVELTE
-                  nodejs
-
-                  # MISC
-                  just
-                  jq
-                ]
-                ++ lib.optionals stdenv.isLinux [
-                  unixtools.ifconfig
-                ]
-                ++ lib.optionals stdenv.isDarwin [
-                  macmon
-                ];
-
-              OPENSSL_NO_VENDOR = "1";
-
-              shellHook = ''
-                export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib"
-                ${lib.optionalString stdenv.isLinux ''
-                  export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH"
-                ''}
-              '';
             };
-          };
         };
     };
 }

nix/mlx_patch_fmod.patch

@@ -0,0 +1,20 @@
+diff --git a/mlx/backend/cuda/device/binary_ops.cuh b/mlx/backend/cuda/device/binary_ops.cuh
+index b0b79628..bbc377be 100644
+--- a/mlx/backend/cuda/device/binary_ops.cuh
++++ b/mlx/backend/cuda/device/binary_ops.cuh
+@@ -46,6 +46,15 @@ struct Remainder {
+       }
+     } else if constexpr (is_complex_v<T>) {
+       return x % y;
++    } else if constexpr (
++        cuda::std::is_same_v<T, __half> ||
++        cuda::std::is_same_v<T, __nv_bfloat16>) {
++      float yf = static_cast<float>(y);
++      float r = cuda::std::fmod(static_cast<float>(x), y);
++      if (r != 0.0f && ((r < 0.0f) != (yf < 0.0f))) {
++        r += yf;
++      }
++      return static_cast<T>(r);
+     } else {
+       T r = cuda::std::fmod(x, y);
+       if (r != 0 && (r < 0 != y < 0)) {

nix/vllm_uv2nix_cmake.patch

@@ -0,0 +1,12 @@
+diff --git a/setup.py b/setup.py
+index 68861fe4b..fd4738089 100644
+--- a/setup.py
++++ b/setup.py
+@@ -150,6 +150,7 @@ class cmake_build_ext(build_ext):
+         cmake_args = [
+             "-DCMAKE_BUILD_TYPE={}".format(cfg),
+             "-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE),
++            *json.loads(os.environ.get("UV2NIX_CMAKE_FLAGS_JSON", "[]"))
+         ]
+ 
+         verbose = envs.VERBOSE

nix/xgrammar_cmake.patch

@@ -0,0 +1,25 @@
+diff --git a/cpp/nanobind/CMakeLists.txt b/cpp/nanobind/CMakeLists.txt
+index ebfd3da..6ef06e3 100644
+--- a/cpp/nanobind/CMakeLists.txt
++++ b/cpp/nanobind/CMakeLists.txt
+@@ -18,8 +18,18 @@ target_sources(python_methods PRIVATE python_methods.cc)
+ target_link_libraries(python_methods PUBLIC xgrammar)
+ 
+ # Any code that uses nanobind directly lives here
+-nanobind_add_module(xgrammar_bindings LTO nanobind.cc)
+-target_link_libraries(xgrammar_bindings PRIVATE python_methods)
++nanobind_build_library(nanobind)
++add_library(xgrammar_bindings MODULE nanobind.cc)
++target_link_libraries(xgrammar_bindings PRIVATE
++  python_methods
++  nanobind
++)
++nanobind_opt_size(xgrammar_bindings)
++nanobind_lto(xgrammar_bindings)
++nanobind_set_visibility(xgrammar_bindings)
++nanobind_extension(xgrammar_bindings)
++nanobind_compile_options(xgrammar_bindings)
++nanobind_link_options(xgrammar_bindings)
+ 
+ if(DEFINED SKBUILD_PROJECT_NAME)
+   # Building wheel through scikit-build-core

nix/xgrammar_cmake_no_LTO.patch

@@ -0,0 +1,13 @@
+diff --git a/cpp/nanobind/CMakeLists.txt b/cpp/nanobind/CMakeLists.txt
+index 03cf0dc..3c9a90a 100644
+--- a/cpp/nanobind/CMakeLists.txt
++++ b/cpp/nanobind/CMakeLists.txt
+@@ -18,7 +18,7 @@ target_sources(python_methods PRIVATE python_methods.cc)
+ target_link_libraries(python_methods PUBLIC xgrammar)
+ 
+ # Any code that uses nanobind directly lives here
+-nanobind_add_module(xgrammar_bindings LTO nanobind.cc)
++nanobind_add_module(xgrammar_bindings nanobind.cc)
+ target_link_libraries(xgrammar_bindings PRIVATE python_methods)
+ 
+ if(DEFINED SKBUILD_PROJECT_NAME)

pyproject.toml

@@ -17,15 +17,15 @@ dependencies = [
   "loguru>=0.7.3",
   "exo_pyo3_bindings",                       # rust bindings
   "anyio==4.11.0",
-  "mlx; sys_platform == 'darwin'",
-  "mlx==0.30.6; sys_platform == 'linux'",
+  "mlx",
+  "mlx[cpu]; sys_platform == 'linux'",
   "mlx-lm",
   "tiktoken>=0.12.0",                        # required for kimi k2 tokenizer
   "hypercorn>=0.18.0",
   "openai-harmony>=0.0.8",
   "httpx>=0.28.1",
   "tomlkit>=0.14.0",
-  "mflux==0.16.9; sys_platform == 'darwin'",
+  "mflux==0.17.2; sys_platform == 'darwin'",
   "python-multipart>=0.0.21",
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
@@ -46,11 +46,12 @@ dev = [
 ]
 
 [project.optional-dependencies]
+build = ["nanobind"]
 cuda = [
   "torch>=2.10.0; sys_platform == 'linux'",
   "vllm>=0.13.0; sys_platform == 'linux'",
-  "mlx-cuda-13==0.30.6; sys_platform == 'linux'",
   "fastsafetensors>=0.1.10; sys_platform == 'linux'",
+  "mlx[cuda13]; sys_platform == 'linux'",
 ]
 
 ###
@@ -62,13 +63,10 @@ members = ["rust/exo_pyo3_bindings", "bench"]
 
 [tool.uv.sources]
 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 = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks" }
 mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" }
 torch = [{ index = "pytorch-cu130", marker = "sys_platform == 'linux'" }]
 vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
-# Uncomment to use local mlx/mlx-lm development versions:
-# mlx = { path = "/Users/Shared/mlx", editable=true }
-# mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
 
 [[tool.uv.index]]
 name = "pytorch-cu130"
@@ -130,22 +128,47 @@ root = "src"
 [tool.uv]
 required-version = ">=0.8.6"
 prerelease = "allow"
-environments = [
-  "sys_platform == 'darwin'",
-  "sys_platform == 'linux' and platform_machine == 'aarch64'",
-]
-no-binary-package = ["vllm"]
-no-build-isolation-package = ["vllm"]
-extra-build-dependencies = { vllm = [
-  "cmake>=3.26.1",
-  "ninja",
-  "packaging>=24.2",
-  "setuptools>=77.0.3,<81.0.0",
-  "setuptools-scm>=8.0",
-  "wheel",
-  "jinja2",
-] }
+environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
 conflicts = [[{ package = "exo", extra = "cuda" }, { package = "exo-bench" }]]
+override-dependencies = ["compressed-tensors==0.13.0", "torch==2.10.0"]
+constraint-dependencies = ["transformers>=5.0.0,<5.4.0"]
+
+[tool.uv.extra-build-dependencies]
+mlx = [
+  "setuptools",
+  "typing-extensions",
+  "nanobind",
+  "pybind11",
+  "wheel",
+  "cmake",
+  "ninja",
+]
+mlx-lm = ["setuptools"]
+xgrammar = [
+  "nanobind",
+  "setuptools",
+  "scikit-build-core",
+  "packaging",
+  "pathspec",
+]
+rouge-score = ["setuptools"]
+sacrebleu = ["setuptools"]
+sqlitedict = ["setuptools"]
+word2number = ["setuptools"]
+vllm = [
+  "setuptools",
+  "setuptools-scm",
+  "scikit-build-core",
+  "jinja2",
+  "wheel",
+  "markupsafe",
+  "typing-extensions",
+  "torch",
+]
+fastsafetensors = ["setuptools", "pybind11"]
+torch = ["typing-extensions"]
+torchvision = ["torch"]
+torchaudio = ["torch"]
 
 ###
 # ruff configuration
@@ -153,7 +176,6 @@ conflicts = [[{ package = "exo", extra = "cuda" }, { package = "exo-bench" }]]
 
 [tool.ruff]
 extend-exclude = [
-  "shared/protobufs/**",
   "*mlx_typings/**",
   "*cuda_typings/**",
   "rust/exo_pyo3_bindings/**",

python/parts.nix

@@ -1,149 +1,434 @@
 { inputs, ... }:
-{
-  perSystem =
-    { config, self', pkgs, lib, system, ... }:
+let
+  mkPythonSet = { self', pkgs, lib, apple-sdk, editable ? false }:
     let
-      pkgsCuda = import ../nix/cuda-pkgs.nix { nixpkgs = inputs.nixpkgs; inherit system; };
+      inherit (pkgs.stdenv.hostPlatform) isDarwin isLinux isx86_64;
+      inherit (pkgs.config) cudaSupport;
+      inherit (pkgs) cudaPackages;
+      cudaLibs = with cudaPackages; [
+        cuda_cudart
+        cuda_cccl
+        cuda_cupti
+        cuda_nvrtc
+        cuda_nvtx
+        cudnn
+        libcufile
+        libcublas
+        libcufft
+        libcurand
+        libcusolver
+        libcusparse
+        libcusparse_lt
+        libnvjitlink
+        libnvshmem
+        nccl
+      ];
+      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 ]);
+      };
+      exoOverlay = final: prev:
+        {
+          mlx = prev.mlx.overrideAttrs (old:
+            let
+              # Static dependencies included directly during compilation
+              gguf-tools = pkgs.fetchFromGitHub {
+                owner = "antirez";
+                repo = "gguf-tools";
+                rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
+                hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
+              };
+
+              metal_cpp = pkgs.fetchzip {
+                url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip";
+                hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
+              };
+
+              nanobind = pkgs.fetchFromGitHub {
+                owner = "wjakob";
+                repo = "nanobind";
+                rev = "v2.10.2";
+                hash = "sha256-io44YhN+VpfHFWyvvLWSanRgbzA0whK8WlDNRi3hahU=";
+                fetchSubmodules = true;
+              };
+
+              nvtx = pkgs.fetchFromGitHub {
+                name = "nvtx3";
+                owner = "NVIDIA";
+                repo = "NVTX";
+                rev = "v3.1.1";
+                hash = "sha256-sx72N+Gskg9Vtqc3sXsWoE/2PHFI2Hq08lEaw0sll5Y=";
+              };
+              cudnn = pkgs.fetchFromGitHub {
+                name = "cudnn_frontend";
+                owner = "NVIDIA";
+                repo = "cudnn-frontend";
+                rev = "v1.16.0";
+                hash = "sha256-+8aBl9dKd2Uz50XoOr91NRyJ4OGJtzfDNNNYGQJ9b94=";
+              };
+              mlx_cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" { } ''
+                mkdir -p $out/include
+                exit 1
+                ln -s ${cudaPackages.cuda_cccl}/include/cuda $out/include/cuda
+                ln -s ${cudaPackages.cuda_cccl}/include/nv $out/include/nv
+              '';
+            in
+            {
+              nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake ] ++ lib.optionals isDarwin [ self'.packages.metal-toolchain ] ++ lib.optionals cudaSupport [
+                cudaPackages.cuda_nvcc
+                pkgs.autoAddDriverRunpath
+                pkgs.autoPatchelfHook
+              ];
+              # TODO: non-sdk_26 support
+              buildInputs = (old.buildInputs or [ ])
+              ++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.openblas ]
+              ++ lib.optionals isDarwin [ apple-sdk ]
+              ++ lib.optionals cudaSupport (cudaLibs ++ [ cudaPackages.cudnn ]);
+              patches = (old.patches or [ ])
+              ++ lib.optionals cudaSupport [ ../nix/mlx_patch_fmod.patch ]
+              ++ lib.optionals isDarwin [
+                (pkgs.replaceVars ../nix/darwin-build-fixes.patch {
+                  sdkVersion = apple-sdk.version;
+                  inherit (self'.packages.metal-toolchain) metalVersion;
+                })
+              ];
+              postPatch = ''
+                substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
+                  --replace-fail "g++" "${lib.getExe' pkgs.stdenv.cc "c++"}"
+              '';
+              autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ lib.optionals cudaSupport [ "libcuda.so.1" ];
+
+              DEV_RELEASE = 1;
+              CMAKE_ARGS = toString ([
+                (lib.cmakeBool "USE_SYSTEM_FMT" true)
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${pkgs.nlohmann_json.src}")
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}")
+                (lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
+                (lib.cmakeBool "MLX_BUILD_CPU" true)
+                (lib.cmakeBool "MLX_BUILD_METAL" isDarwin)
+                (lib.cmakeBool "MLX_BUILD_CUDA" false)
+                (lib.cmakeOptionType "string" "CMAKE_INSTALL_LIBDIR" "lib")
+              ] ++ lib.optionals cudaSupport [
+                (lib.cmakeOptionType "filepath" "CUDAToolkit_ROOT" "${cudaRoot}")
+                (lib.cmakeOptionType "string" "MLX_CUDA_ARCHITECTURES" "121")
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_CUTLASS" "${cudaPackages.cutlass}")
+                # TODO: replace with cudaPackages.cudnn
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_CUDNN" "${cudnn}")
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_CCCL" "${cudaPackages.cuda_cccl}")
+                # TODO: replace with cudaPackages.nvtx
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NVTX3" "${nvtx}")
+              ] ++ lib.optionals isDarwin [
+                (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
+                (lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${apple-sdk.version}")
+                (lib.cmakeOptionType "filepath" "CMAKE_OSX_SYSROOT" "${apple-sdk.passthru.sdkroot}")
+              ] ++ lib.optionals (isDarwin && isx86_64) [
+                (lib.cmakeBool "MLX_ENABLE_X64_MAC" true)
+              ]);
+            } // lib.optionalAttrs isDarwin {
+              SDKROOT = apple-sdk.passthru.sdkroot;
+              MACOSX_DEPLOYMENT_TARGET = apple-sdk.version;
+            });
+          exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
+            pname = "exo-pyo3-bindings";
+            version = "0.1.0";
+            src = self'.packages.exo_pyo3_bindings;
+            # Install from pre-built wheel
+            nativeBuildInputs = [ final.pyprojectWheelHook ];
+            dontStrip = true;
+            passthru = prev.exo-pyo3-bindings.passthru or { };
+            postInstall = ''
+              local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings
+              cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/
+              touch $siteDir/py.typed
+            '';
+          };
+          torch = prev.torch.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = (old.buildInputs or [ ]) ++ lib.optionals cudaSupport cudaLibs;
+            autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ lib.optionals cudaSupport [ "libcuda.so.1" ];
+          });
+          torchaudio = prev.torchaudio.overrideAttrs (old:
+            {
+              nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+              buildInputs = (old.buildInputs or [ ]) ++ lib.optionals cudaSupport [
+                cudaPackages.cuda_cudart
+              ];
+              preFixup = (old.preFixup or "") + ''
+                addAutoPatchelfSearchPath "${final.torch}"
+              '';
+              autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ lib.optionals cudaSupport [ "libcuda.so.1" ];
+            });
+          torchvision = prev.torchvision.overrideAttrs (old:
+            {
+              nativebuildInputs = (old.nativebuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+              preFixup = (old.preFixup or "") + ''
+                addAutoPatchelfSearchPath "${final.torch}"
+              '';
+              autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ lib.optionals cudaSupport [ "libcuda.so.1" ];
+            });
+          torch-c-dlpack-ext = prev.torch-c-dlpack-ext.overrideAttrs (old:
+            {
+              nativebuildInputs = (old.nativebuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+              preFixup = (old.preFixup or "") + ''
+                addAutoPatchelfSearchPath "${final.torch}"
+	      '';
+              autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ lib.optionals cudaSupport [ "libcuda.so.1" ];
+            });
+          xgrammar = prev.xgrammar.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake pkgs.autoPatchelfHook ];
+          });
+          # Currently treating vllm as a cuda dep. it obviously exists as a non cuda dep
+          vllm = prev.vllm.overrideAttrs (old:
+            let
+              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_uv2nix_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";
+              # TODO: vllm non cuda13 support, more arch's, etc.
+            } // lib.optionalAttrs cudaSupport {
+              buildInputs = (old.buildInputs or [ ]) ++ cudaLibs;
+
+              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";
+
+              UV2NIX_CMAKE_FLAGS_JSON = builtins.toJSON [
+                "-DCUDAToolkit_ROOT=${cudaRoot}"
+                "-DCMAKE_CUDA_COMPILER=${cudaRoot}/bin/nvcc"
+                "-DCMAKE_PREFIX_PATH=${cudaRoot}"
+                "-DFETCHCONTENT_SOURCE_DIR_CUTLASS=${lib.getDev cutlass}"
+                "-DFLASH_MLA_SRC_DIR=${lib.getDev flashmla}"
+                "-DVLLM_FLASH_ATTN_SRC_DIR=${lib.getDev vllm-flash-attn}"
+                "-DQUTLASS_SRC_DIR=${lib.getDev qutlass}"
+                "-DTORCH_CUDA_ARCH_LIST=12.0;12.1"
+                "-DCUTLASS_NVCC_ARCHS_ENABLED=${cudaPackages.flags.cmakeCudaArchitecturesString}"
+                "-DCAFFE2_USE_CUDNN=ON"
+                "-DCAFFE2_USE_CUFILE=ON"
+                "-DCUTLASS_ENABLE_CUBLAS=ON"
+              ];
+            });
+        } // lib.optionalAttrs cudaSupport {
+          nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = (old.buildInputs or [ ]) ++ [ pkgs.rdma-core ];
+            propagatedBuildInputs = (old.propagatedBuildInputs or [ ]) ++ [ pkgs.util-linux ];
+          });
+          nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = (old.buildInputs or [ ]) ++ (with cudaPackages; [ libnvjitlink libcublas libcusparse ]);
+          });
+          nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = (old.buildInputs or [ ]) ++ [ cudaPackages.libnvjitlink ];
+          });
+          nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = (old.buildInputs or [ ]) ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
+          });
+          nvidia-cutlass-dsl-libs-base = prev.nvidia-cutlass-dsl-libs-base.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.autoAddDriverRunpath ];
+            autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ [ "libcuda.so.1" ];
+          });
+        } // lib.optionalAttrs (cudaSupport && isx86_64) {
+          numba = prev.numba.overrideAttrs (old: {
+            buildInputs = (old.buildInputs or [ ]) ++ [ pkgs.tbb ];
+          });
+          intel-openmp = prev.intel-openmp.overrideAttrs (_old: {
+            postFixup = ''
+              rm -f $out/lib/libarcher.so
+              rm -f $out/lib/libomptarget.so
+              rm -f $out/lib/libomptarget.rtl.*.so*
+              rm -f $out/lib/libomptarget.sycl.wrap.so
+            '';
+          });
+        };
+
       # Load workspace from uv.lock
       workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
-        workspaceRoot = inputs.self;
+        workspaceRoot = ../.;
       };
 
       # Create overlay from workspace
       # Use wheels from PyPI for most packages; we override mlx with our pure Nix Metal build
       overlay = workspace.mkPyprojectOverlay { sourcePreference = "wheel"; };
-
-      # Override overlay to inject Nix-built components
-      exoOverlay = final: prev: {
-        # Replace workspace exo_pyo3_bindings with Nix-built wheel.
-        # Preserve passthru so mkVirtualEnv can resolve dependency groups.
-        # Copy .pyi stub + py.typed marker so basedpyright can find the types.
-        exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
-          pname = "exo-pyo3-bindings";
-          version = "0.1.0";
-          src = self'.packages.exo_pyo3_bindings;
-          # Install from pre-built wheel
-          nativeBuildInputs = [ final.pyprojectWheelHook ];
-          dontStrip = true;
-          passthru = prev.exo-pyo3-bindings.passthru or { };
-          postInstall = ''
-            local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings
-            cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/
-            touch $siteDir/py.typed
-          '';
-        };
-      };
-
+    in
+    (pkgs.callPackage inputs.pyproject-nix.build.packages {
       python = pkgs.python313;
-
-      # Overlay to provide build systems and custom packages
-      buildSystemsOverlay = final: prev: {
-        # mlx-lm is a git dependency that needs setuptools
-        mlx-lm = prev.mlx-lm.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-        # rouge-score and sacrebleu don't declare setuptools as a build dependency
-        rouge-score = prev.rouge-score.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-        sacrebleu = prev.sacrebleu.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-        sqlitedict = prev.sqlitedict.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-        word2number = prev.word2number.overrideAttrs (old: {
-          nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-            final.setuptools
-          ];
-        });
-      } // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
-        # Use our pure Nix-built MLX with Metal support (macOS only)
-        mlx = self'.packages.mlx;
-      };
-
-      # Additional overlay for Linux-specific fixes (type checking env).
-      # Native wheels have shared lib dependencies we don't need at type-check time.
-      linuxOverlay = final: prev:
-        let
-          ignoreMissing = drv: drv.overrideAttrs { autoPatchelfIgnoreMissingDeps = [ "*" ]; };
-          nvidiaPackages = lib.filterAttrs (name: _: lib.hasPrefix "nvidia-" name) prev;
-        in
-        lib.optionalAttrs pkgs.stdenv.hostPlatform.isLinux (
-          (lib.mapAttrs (_: ignoreMissing) nvidiaPackages) // {
-            mlx = ignoreMissing prev.mlx;
-            mlx-cuda-13 = prev.mlx-cuda-13.overrideAttrs (old: {
-              buildInputs = (old.buildInputs or [ ]) ++ [
-                final.nvidia-cublas
-                final.nvidia-cuda-nvrtc
-                final.nvidia-cudnn-cu13
-                final.nvidia-nccl-cu13
-              ];
-              preFixup = ''
-                addAutoPatchelfSearchPath ${final.nvidia-cublas}
-                addAutoPatchelfSearchPath ${final.nvidia-cuda-nvrtc}
-                addAutoPatchelfSearchPath ${final.nvidia-cudnn-cu13}
-                addAutoPatchelfSearchPath ${final.nvidia-nccl-cu13}
-              '';
-              autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-            });
-            torch = ignoreMissing prev.torch;
-            triton = ignoreMissing prev.triton;
-          }
-        );
-
-      baseOverlays = [
+    }).overrideScope (
+      lib.composeManyExtensions ([
         inputs.pyproject-build-systems.overlays.default
         overlay
         exoOverlay
-        buildSystemsOverlay
-        linuxOverlay
-      ];
+      ] ++ lib.optionals editable [
+        (workspace.mkEditablePyprojectOverlay { root = "$REPO_ROOT"; members = [ "exo" "bench" ]; })
+      ])
+    );
 
-      pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
-        inherit python;
-      }).overrideScope (
-        lib.composeManyExtensions baseOverlays
-      );
-      # mlx-cpu and mlx-cuda-13 both ship mlx/ site-packages files; keep first.
-      # mlx-cpu/mlx-cuda-13 and nvidia-cudnn-cu12/cu13 ship overlapping files.
-      venvCollisionPaths = lib.optionals pkgs.stdenv.hostPlatform.isLinux [
-        "lib/python3.13/site-packages/mlx*"
-        "lib/python3.13/site-packages/nvidia*"
-      ];
+  mkExo = args@{ self', pkgs, lib, ... }:
+    let
+      venv = ((mkPythonSet args).mkVirtualEnv "exo-env" {
+        exo = lib.optionals pkgs.config.cudaSupport [ "cuda" ];
+      }).overrideAttrs {
+        venvSkip = [ "lib/python3.13/site-packages/build_backend.py" ];
+      };
+    in
+    pkgs.runCommand "exo"
+      {
+        nativeBuildInputs = [ pkgs.makeWrapper ];
+      }
+      ''
+        mkdir -p $out/bin
 
-      # Exclude bench deps from main env (bench has its own benchVenv)
-      exoDeps = removeAttrs workspace.deps.default [ "exo-bench" ];
+        # Create wrapper script
+        makeWrapper ${venv}/bin/exo $out/bin/exo \
+          --set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
+          --set EXO_RESOURCES_DIR ${inputs.self + /resources} \
+          ${lib.optionalString pkgs.stdenv.hostPlatform.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
+      '';
+in
+{
+  perSystem =
+    { self', pkgs, cudaPkgs, lib, ... }:
+    let
+      inherit (pkgs.stdenv.hostPlatform) isDarwin;
+      pythonSet = mkPythonSet { inherit self' pkgs lib; apple-sdk = pkgs.apple-sdk_26; };
+      # taking cudaPkgs.cudaPackages_13.pkgs creates a new nixpkgs that defaults to cuda 13
+      cudaPythonSet = mkPythonSet { inherit self' lib; inherit (cudaPkgs.cudaPackages_13) pkgs; apple-sdk = pkgs.apple-sdk_26; };
 
-      exoVenv = (pythonSet.mkVirtualEnv "exo-env" exoDeps).overrideAttrs {
-        venvIgnoreCollisions = venvCollisionPaths;
+      editablePythonSet = mkPythonSet { inherit self' lib pkgs; apple-sdk = pkgs.apple-sdk_26; editable = true; };
+      evenv = (editablePythonSet.mkVirtualEnv "exo-dev-env"
+        {
+          exo = [ "dev" ];
+          exo-pyo3-bindings = [ ];
+          exo-bench = [ ];
+        }).overrideAttrs {
+        venvSkip = [ "lib/python3.13/site-packages/build_backend.py" ];
+      };
+      exoCudaVenv = (cudaPythonSet.mkVirtualEnv "exo-env" {
+        exo = [ "cuda" ];
+        exo-pyo3-bindings = [ ];
+      }).overrideAttrs {
+        venvSkip = [ "lib/python3.13/site-packages/build_backend.py" ];
       };
 
       # Virtual environment with dev dependencies for testing
-      testVenv = (pythonSet.mkVirtualEnv "exo-test-env" (
-        exoDeps // {
+      testVenv = (pythonSet.mkVirtualEnv "exo-test-env"
+        {
           exo = [ "dev" ]; # Include pytest, pytest-asyncio, pytest-env
+          exo-pyo3-bindings = [ ];
         }
-      )).overrideAttrs {
-        venvIgnoreCollisions = venvCollisionPaths;
-      };
-
-      mkPythonScript = name: path: pkgs.writeShellApplication {
-        inherit name;
-        runtimeInputs = [ exoVenv ];
-        runtimeEnv = {
-          EXO_DASHBOARD_DIR = self'.packages.dashboard;
-          EXO_RESOURCES_DIR = inputs.self + /resources;
-        };
-        text = ''exec python ${path} "$@"'';
+      ).overrideAttrs {
+        # venvIgnoreCollisions = venvCollisionPaths;
       };
 
       benchVenv = pythonSet.mkVirtualEnv "exo-bench-env" {
@@ -162,67 +447,34 @@
         text = ''exec python ${path} "$@"'';
       };
 
-      exoPackage = pkgs.runCommand "exo"
+      exoCudaPackage = cudaPkgs.runCommand "exo"
         {
-          nativeBuildInputs = [ pkgs.makeWrapper ];
+          nativeBuildInputs = [ cudaPkgs.makeWrapper ];
         }
         ''
           mkdir -p $out/bin
 
           # Create wrapper script
-          makeWrapper ${exoVenv}/bin/exo $out/bin/exo \
+          makeWrapper ${exoCudaVenv}/bin/exo $out/bin/exo \
             --set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
             --set EXO_RESOURCES_DIR ${inputs.self + /resources} \
-            ${lib.optionalString pkgs.stdenv.hostPlatform.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
-        '';
-
-      vllmEnv = pkgsCuda.python313.withPackages (ps: [ ps.vllm ps.fastsafetensors ]);
-
-      vllmSite = pkgs.runCommand "vllm-site-filtered" { } ''
-        mkdir -p $out
-        for pkg in ${vllmEnv}/${python.sitePackages}/*; do
-          name=$(basename "$pkg")
-          case "$name" in
-            anyio*|pydantic*) ;;
-            *) ln -s "$pkg" "$out/$name" ;;
-          esac
-        done
-      '';
-
-      exoCudaDeps = exoDeps // {
-        mlx-cuda-13 = [ ];
-      };
-
-      exoCudaVenv = (pythonSet.mkVirtualEnv "exo-cuda-env" exoCudaDeps).overrideAttrs {
-        venvIgnoreCollisions = venvCollisionPaths;
-      };
-
-      exoCudaPackage = pkgs.runCommand "exo-cuda"
-        {
-          nativeBuildInputs = [ pkgs.makeWrapper ];
-        }
-        ''
-          mkdir -p $out/bin
-          makeWrapper ${exoCudaVenv}/bin/exo $out/bin/exo-cuda \
-            --set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
-            --set EXO_RESOURCES_DIR ${inputs.self + /resources} \
-            --prefix PYTHONPATH : "${vllmSite}"
+            --prefix LD_LIBRARY_PATH : /run/opengl-driver/lib:${lib.getLib pkgs.util-linux}/lib:${lib.getLib pkgs.systemd}/lib:${lib.getLib pkgs.numactl}/lib:${lib.getLib pkgs.stdenv.cc.cc.lib}/lib \
+            ${lib.optionalString isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
         '';
     in
     {
-      # Python package only available on macOS (requires MLX/Metal)
-      packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin
-        {
-          exo = exoPackage;
-          # Test environment for running pytest outside of Nix sandbox (needs GPU access)
-          exo-test-env = testVenv;
-        } // lib.optionalAttrs (pkgsCuda != null) {
-        exo-cuda-unwrapped = exoCudaPackage;
-      } // {
+      packages = {
+        exo = mkExo { inherit self' lib pkgs; apple-sdk = pkgs.apple-sdk_26; };
+        exo-cuda = exoCudaPackage;
         exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
         exo-eval = mkBenchScript "exo-eval" (inputs.self + /bench/exo_eval.py);
         exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
         exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
+        editable-venv = evenv;
+      } // lib.optionalAttrs isDarwin {
+        # Test environment for running pytest outside of Nix sandbox (needs GPU access)
+        exo-test-env = testVenv;
+        exo-osx14 = mkExo { inherit self' lib pkgs; apple-sdk = pkgs.apple-sdk_14; };
       };
 
       checks = {

scripts/check_kv_compat.py

@@ -0,0 +1,33 @@
+import sys
+
+sys.path.insert(0, "src")
+import mlx.core as mx
+from mlx_lm import load
+from mlx_lm.models.cache import RotatingKVCache, KVCache
+
+model, tok = load("mlx-community/gpt-oss-20b-MXFP4-Q8")
+
+prompt = "Hello " * 2000
+tokens = tok.encode(prompt)
+print(f"Tokens: {len(tokens)}")
+
+cache = model.make_cache()
+token_arr = mx.array([tokens])
+logits = model(token_arr, cache=cache)
+mx.eval(logits)
+
+for i, c in enumerate(cache[:6]):
+    if (
+        isinstance(c, KVCache)
+        and not isinstance(c, RotatingKVCache)
+        and c.keys is not None
+    ):
+        k = c.keys.astype(mx.float32)
+        print(
+            f"Layer {i} KVCache: shape={c.keys.shape} offset={c.offset} first=[{float(k[0, 0, 0, 0]):.6f}, {float(k[0, 0, 0, 1]):.6f}] last=[{float(k[0, 0, -1, -2]):.6f}, {float(k[0, 0, -1, -1]):.6f}]"
+        )
+    elif isinstance(c, RotatingKVCache) and c.keys is not None:
+        k = c.keys.astype(mx.float32)
+        print(
+            f"Layer {i} RotatingKV: shape={c.keys.shape} _idx={c._idx} offset={c.offset} first=[{float(k[0, 0, 0, 0]):.6f}, {float(k[0, 0, 0, 1]):.6f}]"
+        )

scripts/check_rotating_cache.py

@@ -0,0 +1,14 @@
+import mlx.core as mx
+from mlx_lm import load
+from mlx_lm.models.cache import RotatingKVCache
+
+model, tok = load("mlx-community/gpt-oss-20b-MXFP4-Q8")
+cache = model.make_cache()
+tokens = mx.ones((1, 5000), dtype=mx.int32)
+model(tokens, cache=cache)
+mx.eval([c.keys for c in cache if c.keys is not None])
+for i, c in enumerate(cache[:4]):
+    if isinstance(c, RotatingKVCache):
+        print(
+            f"Layer {i}: _idx={c._idx} offset={c.offset} keep={c.keep} max_size={c.max_size} keys={c.keys.shape}"
+        )

scripts/compare_kv.py

@@ -0,0 +1,232 @@
+import sys
+
+sys.path.insert(0, "src")
+from exo.worker.engines.mlx.gdn_softplus_patch import patch_gdn_softplus
+from exo.worker.engines.mlx.yarn_rope_patch import patch_yarn_rope
+
+patch_gdn_softplus()
+patch_yarn_rope()
+import mlx.core as mx
+import torch
+import socket
+from pathlib import Path
+import json
+from collections import defaultdict
+from mlx_lm import load
+from mlx_lm.models.cache import ArraysCache, RotatingKVCache, KVCache
+from exo.disaggregated.protocol import (
+    read_header,
+    read_message,
+    ArraysState,
+    KVChunk,
+    Done,
+)
+from exo.disaggregated.prefill_client import _nhd_to_bhsd, _torch_to_mx
+
+ENDPOINT = sys.argv[1] if len(sys.argv) > 1 else "10.43.0.1:62988"
+MODEL = sys.argv[2] if len(sys.argv) > 2 else "mlx-community/Llama-3.2-1B-Instruct-bf16"
+MODEL_PATH = sys.argv[3] if len(sys.argv) > 3 else None
+
+model, tok = load(
+    MODEL_PATH or str(Path.home() / ".exo/models" / MODEL.replace("/", "--"))
+)
+prompt = "The quick brown fox jumps over the lazy dog. " * 3000
+tokens = tok.encode(prompt)
+print(f"Tokens: {len(tokens)}")
+
+host, port = ENDPOINT.rsplit(":", 1)
+sock = socket.create_connection((host, int(port)), timeout=60)
+request = (
+    json.dumps({"model": MODEL, "token_ids": tokens, "start_pos": 0}).encode() + b"\n"
+)
+sock.sendall(request)
+stream = sock.makefile("rb", buffering=65536)
+header = read_header(stream)
+
+vllm_kv = defaultdict(list)
+vllm_arrays: dict[int, list[torch.Tensor]] = {}
+while True:
+    msg = read_message(stream, header)
+    if msg is None or isinstance(msg, Done):
+        break
+    if isinstance(msg, KVChunk):
+        vllm_kv[msg.layer_idx].append((msg.keys, msg.values))
+    elif isinstance(msg, ArraysState):
+        vllm_arrays[msg.layer_idx] = msg.arrays
+sock.close()
+
+print(f"Received {len(vllm_kv)} KV layers, {len(vllm_arrays)} arrays layers from vLLM")
+
+if hasattr(model, "make_cache"):
+    mlx_cache = model.make_cache()
+else:
+    from mlx_lm.models.cache import make_prompt_cache
+
+    mlx_cache = make_prompt_cache(model)
+token_arr = mx.array([tokens[:-2]])
+mlx_logits = model(token_arr, cache=mlx_cache)
+mx.eval(mlx_logits)
+
+for i in range(min(6, len(mlx_cache))):
+    c = mlx_cache[i]
+    if isinstance(c, ArraysCache):
+        if i in vllm_arrays:
+            vllm_arrs = vllm_arrays[i]
+            mlx_state = c.state
+            print(
+                f"Layer {i} (Arrays): mlx_state={len(mlx_state)} arrays, vllm={len(vllm_arrs)} arrays"
+            )
+            for ai, (m_arr, v_arr) in enumerate(zip(mlx_state, vllm_arrs)):
+                if m_arr is None:
+                    continue
+                v_mx = _torch_to_mx(v_arr).astype(mx.float32)
+                m_f = m_arr.astype(mx.float32)
+                if m_f.shape != v_mx.shape:
+                    print(f"  [{ai}] SHAPE MISMATCH mlx={m_f.shape} vllm={v_mx.shape}")
+                else:
+                    d = mx.abs(m_f - v_mx)
+                    a = m_f.reshape(-1)
+                    b = v_mx.reshape(-1)
+                    cos = float(mx.sum(a * b).item()) / (
+                        float(mx.sqrt(mx.sum(a * a)).item())
+                        * float(mx.sqrt(mx.sum(b * b)).item())
+                        + 1e-8
+                    )
+                    print(
+                        f"  [{ai}] cosine_sim={cos:.6f} max_diff={mx.max(d).item():.6f} mean_diff={mx.mean(d).item():.6f} shape={m_f.shape}"
+                    )
+        else:
+            print(f"Layer {i} (Arrays): no vLLM data")
+        continue
+    if c.keys is None:
+        continue
+    mlx_k = c.keys.astype(mx.float32)
+
+    if i not in vllm_kv:
+        print(f"Layer {i}: no vLLM data")
+        continue
+    chunks = vllm_kv[i]
+    vk = torch.cat([k for k, v in chunks], dim=0) if len(chunks) > 1 else chunks[0][0]
+    vk_mx = _torch_to_mx(vk.permute(1, 0, 2).unsqueeze(0)).astype(mx.float32)
+
+    n = min(mlx_k.shape[2], vk_mx.shape[2])
+    diff = mx.abs(mlx_k[:, :, :n, :] - vk_mx[:, :, :n, :])
+    max_diff = mx.max(diff).item()
+    mean_diff = mx.mean(diff).item()
+    cache_type = "RotatingKV" if isinstance(c, RotatingKVCache) else "KV"
+    print(
+        f"Layer {i} ({cache_type}): mlx={mlx_k.shape} vllm={vk_mx.shape} max_diff={max_diff:.6f} mean_diff={mean_diff:.6f}"
+    )
+
+    a = mlx_k[:, :, :n, :].reshape(-1)
+    b_vec = vk_mx[:, :, :n, :].reshape(-1)
+    cos_sim = float(mx.sum(a * b_vec).item()) / (
+        float(mx.sqrt(mx.sum(a * a)).item())
+        * float(mx.sqrt(mx.sum(b_vec * b_vec)).item())
+        + 1e-8
+    )
+    diff_tensor = mx.abs(mlx_k[:, :, :n, :] - vk_mx[:, :, :n, :])
+    max_idx = mx.argmax(diff_tensor.reshape(-1)).item()
+    total_elems = diff_tensor.shape[1] * n * diff_tensor.shape[3]
+    h_idx = (max_idx // (n * diff_tensor.shape[3])) % diff_tensor.shape[1]
+    s_idx = (max_idx // diff_tensor.shape[3]) % n
+    d_idx = max_idx % diff_tensor.shape[3]
+    print(
+        f"  cosine_sim={cos_sim:.6f} max_diff={max_diff:.4f} at h={h_idx} pos={s_idx} dim={d_idx}: mlx={float(mlx_k[0, h_idx, s_idx, d_idx].item()):.6f} vllm={float(vk_mx[0, h_idx, s_idx, d_idx].item()):.6f}"
+    )
+
+    D = mlx_k.shape[3]
+    for pos in [0, 100, n - 1]:
+        mlx_row = [float(mlx_k[0, 0, pos, d].item()) for d in range(D)]
+        vllm_row = [float(vk_mx[0, 0, pos, d].item()) for d in range(D)]
+        diffs = [abs(mlx_row[d] - vllm_row[d]) for d in range(D)]
+        top5 = sorted(range(D), key=lambda d: -diffs[d])[:5]
+        print(
+            f"  pos={pos} top5 diff dims: {[(d, f'{diffs[d]:.3f}', f'mlx={mlx_row[d]:.3f}', f'vllm={vllm_row[d]:.3f}') for d in top5]}"
+        )
+
+print("\n--- Run 2: cached request ---")
+sock2 = socket.create_connection((host, int(port)), timeout=60)
+request2 = (
+    json.dumps({"model": MODEL, "token_ids": tokens, "start_pos": 0}).encode() + b"\n"
+)
+sock2.sendall(request2)
+stream2 = sock2.makefile("rb", buffering=65536)
+
+first_byte = stream2.peek(1)[:1]
+if first_byte == b"{":
+    line2 = stream2.readline()
+    print(f"Server error: {json.loads(line2.decode())}")
+    sys.exit(1)
+
+header2 = read_header(stream2)
+vllm_kv2 = defaultdict(list)
+vllm_arrays2: dict[int, list[torch.Tensor]] = {}
+total_tokens2 = 0
+while True:
+    msg = read_message(stream2, header2)
+    if msg is None:
+        break
+    if isinstance(msg, KVChunk):
+        vllm_kv2[msg.layer_idx].append((msg.keys, msg.values))
+    elif isinstance(msg, ArraysState):
+        vllm_arrays2[msg.layer_idx] = msg.arrays
+    elif isinstance(msg, Done):
+        total_tokens2 = msg.total_tokens
+        break
+sock2.close()
+
+kv_tokens2 = 0
+if vllm_kv2:
+    first_layer = next(iter(vllm_kv2.values()))
+    kv_tokens2 = sum(k.shape[0] for k, v in first_layer)
+print(
+    f"Received {len(vllm_kv2)} KV layers ({kv_tokens2} tokens), {len(vllm_arrays2)} arrays layers, total_tokens={total_tokens2}"
+)
+
+for i in range(min(6, len(mlx_cache))):
+    c = mlx_cache[i]
+    if isinstance(c, ArraysCache):
+        if i in vllm_arrays2:
+            vllm_arrs = vllm_arrays2[i]
+            mlx_state = c.state
+            for ai, (m_arr, v_arr) in enumerate(zip(mlx_state, vllm_arrs)):
+                if m_arr is None:
+                    continue
+                v_mx = _torch_to_mx(v_arr).astype(mx.float32)
+                m_f = m_arr.astype(mx.float32)
+                if m_f.shape != v_mx.shape:
+                    print(
+                        f"Layer {i} [{ai}] SHAPE MISMATCH mlx={m_f.shape} vllm={v_mx.shape}"
+                    )
+                else:
+                    a2 = m_f.reshape(-1)
+                    b2 = v_mx.reshape(-1)
+                    cos2 = float(mx.sum(a2 * b2).item()) / (
+                        float(mx.sqrt(mx.sum(a2 * a2)).item())
+                        * float(mx.sqrt(mx.sum(b2 * b2)).item())
+                        + 1e-8
+                    )
+                    print(
+                        f"Layer {i} (Arrays) [{ai}] cosine_sim={cos2:.6f} shape={m_f.shape}"
+                    )
+        continue
+    if c.keys is None or i not in vllm_kv2:
+        continue
+    mlx_k = c.keys.astype(mx.float32)
+    chunks = vllm_kv2[i]
+    vk = torch.cat([k for k, v in chunks], dim=0) if len(chunks) > 1 else chunks[0][0]
+    vk_mx = _torch_to_mx(vk.permute(1, 0, 2).unsqueeze(0)).astype(mx.float32)
+    n = min(mlx_k.shape[2], vk_mx.shape[2])
+    a2 = mlx_k[:, :, :n, :].reshape(-1)
+    b2 = vk_mx[:, :, :n, :].reshape(-1)
+    cos2 = float(mx.sum(a2 * b2).item()) / (
+        float(mx.sqrt(mx.sum(a2 * a2)).item()) * float(mx.sqrt(mx.sum(b2 * b2)).item())
+        + 1e-8
+    )
+    print(f"Layer {i} (KV) cosine_sim={cos2:.6f} mlx={mlx_k.shape} vllm={vk_mx.shape}")
+
+if len(vllm_kv2) > 0:
+    print("PASS")
+else:
+    print("FAIL")

scripts/disaggregated/capture_connector.py

@@ -0,0 +1,90 @@
+"""Minimal KVConnector that captures per-layer cache data."""
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+)
+
+captured_layers: dict[str, Any] = {}
+
+
+@dataclass
+class CaptureMetadata(KVConnectorMetadata):
+    pass
+
+
+class CaptureConnector(KVConnectorBase_V1):
+    def __init__(self, vllm_config, role, kv_cache_config=None):
+        super().__init__(vllm_config, role, kv_cache_config)
+
+    def start_load_kv(self, forward_context, **kwargs):
+        pass
+
+    def wait_for_layer_load(self, layer_name):
+        pass
+
+    def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs):
+        import time
+
+        slot_mapping = getattr(attn_metadata, "slot_mapping", None)
+        if slot_mapping is not None and slot_mapping.shape[0] <= 100:
+            return
+        t0 = time.perf_counter()
+        torch.cuda.synchronize()
+        t_sync = time.perf_counter() - t0
+        if isinstance(kv_layer, (list, tuple)):
+            captured_layers[layer_name] = [t.cpu().clone() for t in kv_layer]
+        else:
+            slot_mapping = getattr(attn_metadata, "slot_mapping", None)
+            if slot_mapping is not None:
+                if kv_layer.shape[0] == 2:
+                    k_all = kv_layer[0]
+                    v_all = kv_layer[1]
+                else:
+                    k_all = kv_layer[:, 0]
+                    v_all = kv_layer[:, 1]
+                k_flat = k_all.reshape(-1, *k_all.shape[-2:])
+                v_flat = v_all.reshape(-1, *v_all.shape[-2:])
+                valid = slot_mapping >= 0
+                safe_sm = slot_mapping.clamp(min=0)
+                keys = k_flat[safe_sm]
+                values = v_flat[safe_sm]
+                keys[~valid] = 0
+                values[~valid] = 0
+                prev = captured_layers.get(layer_name)
+                if isinstance(prev, dict) and "keys" in prev:
+                    t1 = time.perf_counter()
+                    captured_layers[layer_name] = {
+                        "keys": torch.cat([prev["keys"], keys.cpu()], dim=0),
+                        "values": torch.cat([prev["values"], values.cpu()], dim=0),
+                    }
+                    t_copy = time.perf_counter() - t1
+                else:
+                    t1 = time.perf_counter()
+                    captured_layers[layer_name] = {
+                        "keys": keys.cpu(),
+                        "values": values.cpu(),
+                    }
+                    t_copy = time.perf_counter() - t1
+                if "layers.3." in layer_name:
+                    print(
+                        f"    [attn save] sync={t_sync * 1000:.1f}ms copy={t_copy * 1000:.1f}ms tokens={keys.shape[0]}"
+                    )
+            else:
+                captured_layers[layer_name] = kv_layer.cpu().clone()
+
+    def wait_for_save(self):
+        pass
+
+    def get_num_new_matched_tokens(self, request, num_computed_tokens):
+        return 0, False
+
+    def update_state_after_alloc(self, request, blocks, num_external_tokens):
+        pass
+
+    def build_connector_meta(self, scheduler_output):
+        return CaptureMetadata()

scripts/disaggregated/inspect_vllm_kv.py

@@ -0,0 +1,172 @@
+"""Inspect vLLM KV cache structure per-layer after prefill.
+
+Runs on DGX Spark. Prints per-layer shapes, dtypes, kv_cache_config,
+and layer_to_group mapping to understand what vLLM stores for each
+model architecture (standard attention, sliding window, GatedDeltaNet).
+
+Usage:
+  uv run python scripts/disaggregated/inspect_vllm_kv.py --model ~/.local/share/exo/models/openai--gpt-oss-20b
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src"))
+
+os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
+os.environ["VLLM_BATCH_INVARIANT"] = "1"
+
+from exo.worker.runner.bootstrap import _ensure_cuda_libs
+
+_ensure_cuda_libs()
+
+import torch
+
+
+def _build_layer_groups(kv_cache_config):
+    group_lookup = {}
+    for group_idx, group_spec in enumerate(kv_cache_config.kv_cache_groups):
+        for layer_name in group_spec.layer_names:
+            group_lookup[layer_name] = group_idx
+
+    layer_to_group = []
+    for tensor_spec in kv_cache_config.kv_cache_tensors:
+        for name in tensor_spec.shared_by:
+            layer_to_group.append(group_lookup[name])
+    return layer_to_group
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True, help="Path to model")
+    parser.add_argument(
+        "--prompt", default="Hello, world! How are you today?", help="Prompt to prefill"
+    )
+    args = parser.parse_args()
+
+    from exo.worker.engines.vllm.growable_cache import get_model_runner
+    from exo.worker.engines.vllm.vllm_generator import load_vllm_engine
+
+    print(f"Loading vLLM engine from {args.model}...")
+    engine, _, prefix_cache = load_vllm_engine(
+        model_path=args.model,
+        model_id=args.model,
+        trust_remote_code=True,
+    )
+    print("Engine loaded.\n")
+
+    from vllm import SamplingParams
+
+    tokenizer = engine.get_tokenizer()
+    token_ids = tokenizer.encode(args.prompt, add_special_tokens=False)
+    print(f"Prompt: {args.prompt!r}")
+    print(f"Token IDs: {len(token_ids)} tokens\n")
+
+    request_id = "inspect-test"
+    params = SamplingParams(max_tokens=1, detokenize=False)
+    engine.add_request(request_id, {"prompt_token_ids": token_ids}, params)
+
+    while engine.has_unfinished_requests():
+        engine.step()
+
+    model_runner = get_model_runner()
+    if model_runner is None:
+        print("ERROR: model_runner is None")
+        return
+
+    print("=" * 70)
+    print("PER-LAYER KV CACHE TENSORS (model_runner.kv_caches)")
+    print("=" * 70)
+    kv_caches = model_runner.kv_caches
+    for i, kv in enumerate(kv_caches):
+        if isinstance(kv, list):
+            shapes = [t.shape for t in kv]
+            dtypes = [t.dtype for t in kv]
+            print(
+                f"  Layer {i:3d}: list of {len(kv)} tensors — shapes={shapes}, dtypes={dtypes}"
+            )
+        elif isinstance(kv, torch.Tensor):
+            print(
+                f"  Layer {i:3d}: shape={tuple(kv.shape)}, dtype={kv.dtype}, device={kv.device}"
+            )
+        else:
+            print(f"  Layer {i:3d}: type={type(kv).__name__}")
+    print(f"\n  Total layers with KV: {len(kv_caches)}\n")
+
+    engine_core = engine.engine_core.engine_core
+    kv_cache_config = engine_core.scheduler.kv_cache_manager.kv_cache_config
+
+    print("=" * 70)
+    print("KV CACHE CONFIG")
+    print("=" * 70)
+
+    print(f"\n  Number of KV cache groups: {len(kv_cache_config.kv_cache_groups)}")
+    for gi, group in enumerate(kv_cache_config.kv_cache_groups):
+        print(f"\n  Group {gi}:")
+        print(f"    Layer names ({len(group.layer_names)}):")
+        for name in group.layer_names[:5]:
+            print(f"      {name}")
+        if len(group.layer_names) > 5:
+            print(f"      ... and {len(group.layer_names) - 5} more")
+
+    print(f"\n  Number of KV cache tensors: {len(kv_cache_config.kv_cache_tensors)}")
+    for ti, tensor_spec in enumerate(kv_cache_config.kv_cache_tensors):
+        shared = tensor_spec.shared_by[:3]
+        extra = (
+            f" ... +{len(tensor_spec.shared_by) - 3}"
+            if len(tensor_spec.shared_by) > 3
+            else ""
+        )
+        print(f"  Tensor {ti}: shared_by={shared}{extra}")
+
+    layer_to_group = _build_layer_groups(kv_cache_config)
+    print(
+        f"\n  layer_to_group ({len(layer_to_group)} entries): {layer_to_group[:10]}{'...' if len(layer_to_group) > 10 else ''}"
+    )
+
+    coordinator = engine_core.scheduler.kv_cache_manager.coordinator
+    null_block = coordinator.block_pool.null_block
+
+    internal_id = None
+    for mgr in coordinator.single_type_managers:
+        for key in mgr.req_to_blocks:
+            if str(key).startswith(request_id):
+                internal_id = str(key)
+                break
+        if internal_id:
+            break
+
+    if internal_id:
+        print(f"\n  Request internal_id: {internal_id}")
+        for gi, mgr in enumerate(coordinator.single_type_managers):
+            blocks = mgr.req_to_blocks.get(internal_id)
+            if blocks:
+                real_blocks = [
+                    b for b in blocks if b is not null_block and not b.is_null
+                ]
+                null_count = len(blocks) - len(real_blocks)
+                print(
+                    f"  Group {gi}: {len(real_blocks)} real blocks, {null_count} null blocks, block_size={mgr.block_size}"
+                )
+            else:
+                print(f"  Group {gi}: no blocks")
+
+    print("\n" + "=" * 70)
+    print("SUMMARY")
+    print("=" * 70)
+    print(f"  Model: {args.model}")
+    print(f"  KV cache layers: {len(kv_caches)}")
+    print(f"  KV cache groups: {len(kv_cache_config.kv_cache_groups)}")
+    print(f"  Layer-to-group mapping entries: {len(layer_to_group)}")
+    unique_shapes = set()
+    for kv in kv_caches:
+        if isinstance(kv, torch.Tensor):
+            unique_shapes.add(tuple(kv.shape))
+    print(f"  Unique tensor shapes: {unique_shapes}")
+
+
+if __name__ == "__main__":
+    main()

scripts/disaggregated/test_kv_extract.py

@@ -0,0 +1,302 @@
+"""Extract KV cache per-layer from vLLM using a real KVConnector.
+
+Patches vLLM to allow KVConnector on hybrid models (attention + GDN).
+
+Usage:
+  uv run python scripts/disaggregated/test_kv_extract.py --model ~/.local/share/exo/models/Qwen--Qwen3.5-2B --output /tmp/kv_cache_qwen35/
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src"))
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
+os.environ["VLLM_BATCH_INVARIANT"] = "1"
+
+from exo.worker.runner.bootstrap import _ensure_cuda_libs
+
+_ensure_cuda_libs()
+
+import torch
+
+
+def _patch_vllm_for_connector():
+    """Patch vLLM to allow KVConnector on hybrid models."""
+    from vllm.v1.core import kv_cache_utils
+
+    original_unify = kv_cache_utils.unify_hybrid_kv_cache_specs
+
+    def patched_unify(kv_cache_spec):
+        try:
+            original_unify(kv_cache_spec)
+        except ValueError:
+            pass
+
+    kv_cache_utils.unify_hybrid_kv_cache_specs = patched_unify
+
+    from vllm.v1.core.sched import scheduler as sched_mod
+
+    original_connector_finished = sched_mod.Scheduler._connector_finished
+
+    def patched_connector_finished(self, request):
+        return False, None
+
+    sched_mod.Scheduler._connector_finished = patched_connector_finished
+
+    from capture_connector import CaptureConnector
+    from vllm.distributed.kv_transfer.kv_connector import factory
+
+    original_get = factory.KVConnectorFactory._get_connector_class_with_compat
+
+    @classmethod
+    def patched_get(cls, kv_transfer_config):
+        if "capture_connector" in (kv_transfer_config.kv_connector or ""):
+            return CaptureConnector, None
+        return original_get.__func__(cls, kv_transfer_config)
+
+    factory.KVConnectorFactory._get_connector_class_with_compat = patched_get
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--output", required=True)
+    _lorem = "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Curabitur pretium tincidunt lacus. Nulla gravida orci a odio. Nullam varius, turpis et commodo pharetra, est eros bibendum elit, nec luctus magna felis sollicitudin mauris. Integer in mauris eu nibh euismod gravida. Duis ac tellus et risus vulputate vehicula. Donec lobortis risus a elit. Etiam tempor. Ut ullamcorper, ligula ut dictum pharetra, nisi nunc fringilla magna, in commodo elit erat nec turpis. Ut pharetra augue nec augue. Nam elit agna, endrerit sit amet, tincidunt ac, viverra sed, nulla. Donec porta diam eu massa. Quisque diam lorem, interdum vitae, dapibus ac, scelerisque vitae, pede. Donec eget tellus non erat lacinia fermentum. Donec in velit vel ipsum auctor pulvinar. Vestibulum iaculis lacinia est. Proin dictum elementum velit. Fusce euismod consequat ante. Lorem ipsum dolor sit amet, consectetuer adipiscing elit. Pellentesque sed dolor. Aliquam congue fermentum nisl. Mauris accumsan nulla vel diam. Sed in lacus ut enim adipiscing aliquet. Nulla venenatis. In pede mi, aliquet sit amet, euismod in, auctor ut, ligula. Aliquam dapibus tincidunt metus. Praesent justo dolor, lobortis quis, lobortis dignissim, pulvinar ac, lorem. "
+    parser.add_argument(
+        "--prompt",
+        default=_lorem * 21
+        + "Now answer this question: What is the capital of France and why is it historically significant? Give a detailed answer.",
+    )
+    args = parser.parse_args()
+
+    out_dir = Path(args.output)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    _patch_vllm_for_connector()
+
+    from vllm.engine.arg_utils import EngineArgs
+    from vllm.v1.engine.llm_engine import LLMEngine
+
+    from exo.worker.engines.mlx.cache import KVPrefixCache
+    from exo.worker.engines.vllm.growable_cache import patch_vllm, set_prefix_cache
+
+    patch_vllm()
+
+    prefix_cache = KVPrefixCache(group=None)
+    set_prefix_cache(prefix_cache)
+
+    engine_args = EngineArgs(
+        model=args.model,
+        served_model_name=args.model,
+        gpu_memory_utilization=0.05,
+        trust_remote_code=True,
+        load_format="fastsafetensors",
+        enable_prefix_caching=False,
+        attention_backend="TRITON_ATTN",
+        enforce_eager=True,
+        disable_log_stats=True,
+        kv_transfer_config={
+            "kv_connector": "capture_connector:CaptureConnector",
+            "kv_role": "kv_both",
+        },
+    )
+
+    print("Loading engine with KVConnector...")
+    engine = LLMEngine.from_engine_args(engine_args)
+    print("Engine loaded.")
+
+    from exo.worker.engines.vllm.growable_cache import get_model_runner
+
+    model_runner = get_model_runner()
+
+    import vllm.model_executor.layers.mamba.ops.causal_conv1d as cc_mod
+    from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
+        causal_conv1d_fn as orig_causal_conv1d_fn,
+    )
+
+    gdn_states: dict[int, dict[str, torch.Tensor]] = {}
+    gdn_call_idx = [0]
+    gdn_layer_order = [0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14, 16, 17, 18, 20, 21, 22]
+
+    def patched_causal_conv1d_fn(*args, conv_states=None, cache_indices=None, **kwargs):
+        result = orig_causal_conv1d_fn(
+            *args, conv_states=conv_states, cache_indices=cache_indices, **kwargs
+        )
+        if conv_states is not None and cache_indices is not None:
+            x = args[0] if args else None
+            if x is not None and x.shape[0] <= 100:
+                return result
+            import time as _time
+
+            t0 = _time.perf_counter()
+            torch.cuda.synchronize()
+            t_sync = _time.perf_counter() - t0
+            ci = cache_indices[0].item() if cache_indices.numel() > 0 else 0
+            idx = gdn_call_idx[0]
+            layer_idx = gdn_layer_order[idx % len(gdn_layer_order)]
+            t1 = _time.perf_counter()
+            conv_at_ci = conv_states[ci : ci + 1].transpose(-1, -2).contiguous().cpu()
+            t_copy = _time.perf_counter() - t1
+            gdn_states.setdefault(layer_idx, {})["conv"] = conv_at_ci
+            gdn_states[layer_idx]["ci"] = ci
+            if gdn_call_idx[0] < 3:
+                print(
+                    f"    [gdn save] sync={t_sync * 1000:.1f}ms copy={t_copy * 1000:.1f}ms layer={layer_idx}"
+                )
+            gdn_call_idx[0] += 1
+        return result
+
+    cc_mod.causal_conv1d_fn = patched_causal_conv1d_fn
+    for mod in list(sys.modules.values()):
+        if mod is None or mod is cc_mod:
+            continue
+        if (
+            hasattr(mod, "causal_conv1d_fn")
+            and mod.causal_conv1d_fn is orig_causal_conv1d_fn
+        ):
+            mod.causal_conv1d_fn = patched_causal_conv1d_fn
+    print("  Patched causal_conv1d_fn")
+
+    from exo.shared.types.tasks import TaskId
+    from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+    from exo.worker.engines.vllm.vllm_generator import VllmBatchEngine
+
+    batch_engine = VllmBatchEngine(
+        engine=engine, model_id=args.model, prefix_cache=prefix_cache
+    )
+
+    task = TextGenerationTaskParams(
+        model=args.model,
+        input=[InputMessage(role="user", content=args.prompt)],
+        max_completion_tokens=1,
+    )
+
+    task_id = batch_engine.submit(
+        task_id=TaskId("extract"), task_params=task, prompt=args.prompt
+    )
+
+    print("Running prefill via VllmBatchEngine...")
+    t0 = time.perf_counter()
+    while batch_engine.has_work:
+        results = batch_engine.step()
+        for tid, resp in results:
+            print(f"  Prefill done in {(time.perf_counter() - t0) * 1000:.0f}ms")
+            batch_engine.cancel([tid])
+            break
+        if results:
+            break
+    t1 = time.perf_counter()
+    print(f"Total: {(t1 - t0) * 1000:.0f}ms")
+
+    prompt_mx = prefix_cache.prompts[0] if prefix_cache.prompts else None
+    token_ids = [int(x) for x in prompt_mx.tolist()] if prompt_mx is not None else []
+
+    from capture_connector import captured_layers
+
+    print(f"\nCaptured {len(captured_layers)} layers via save_kv_layer:")
+    for name in sorted(captured_layers.keys()):
+        v = captured_layers[name]
+        if isinstance(v, list):
+            print(f"  {name}: {[tuple(t.shape) for t in v]}")
+        elif isinstance(v, torch.Tensor):
+            print(f"  {name}: {tuple(v.shape)}")
+        else:
+            print(f"  {name}: {type(v).__name__}")
+
+    num_tokens = len(token_ids)
+    print(f"  Chat-templated prompt: {num_tokens} tokens")
+    total_layers = 24
+
+    for f_old in out_dir.glob("layer_*"):
+        f_old.unlink()
+
+    metadata = {
+        "model": args.model,
+        "prompt": args.prompt,
+        "num_tokens": num_tokens,
+        "token_ids": token_ids,
+        "num_layers": total_layers,
+        "layers": [],
+    }
+
+    print(f"\nSaving {total_layers} layers...")
+
+    torch.cuda.synchronize()
+    for layer_idx in sorted(gdn_states.keys()):
+        ci = gdn_states[layer_idx]["ci"]
+        kv = model_runner.kv_caches[layer_idx]
+        if isinstance(kv, (list, tuple)) and len(kv) > 1:
+            rec_pool = kv[1]
+            rec = rec_pool[ci : ci + 1].cpu().clone()
+            gdn_states[layer_idx]["rec"] = rec
+
+    for li in range(total_layers):
+        if li in gdn_states:
+            s = gdn_states[li]
+            conv = s.get("conv")
+            rec = s.get("rec")
+            torch.save(conv, out_dir / f"layer_{li:03d}_conv.pt")
+            if rec is not None:
+                torch.save(rec, out_dir / f"layer_{li:03d}_rec.pt")
+            metadata["layers"].append(
+                {
+                    "type": "gdn",
+                    "conv": list(conv.shape),
+                    "rec": list(rec.shape) if rec is not None else None,
+                }
+            )
+            print(
+                f"  Layer {li}: GDN conv={tuple(conv.shape)}, rec={tuple(rec.shape) if rec is not None else 'None'}"
+            )
+        else:
+            attn_name = None
+            for n in captured_layers:
+                parts = n.split(".")
+                for pi, p in enumerate(parts):
+                    if (
+                        p == "layers"
+                        and pi + 1 < len(parts)
+                        and parts[pi + 1] == str(li)
+                    ):
+                        attn_name = n
+                        break
+            if attn_name and isinstance(captured_layers[attn_name], dict):
+                kv = captured_layers[attn_name]
+                torch.save(kv["keys"], out_dir / f"layer_{li:03d}_keys.pt")
+                torch.save(kv["values"], out_dir / f"layer_{li:03d}_values.pt")
+                if "last_chunk_keys" in kv:
+                    torch.save(
+                        kv["last_chunk_keys"], out_dir / f"layer_{li:03d}_keys_last.pt"
+                    )
+                    torch.save(
+                        kv["last_chunk_values"],
+                        out_dir / f"layer_{li:03d}_values_last.pt",
+                    )
+                metadata["layers"].append(
+                    {
+                        "type": "kv",
+                        "keys_shape": list(kv["keys"].shape),
+                        "values_shape": list(kv["values"].shape),
+                    }
+                )
+                print(
+                    f"  Layer {li}: KV keys={tuple(kv['keys'].shape)}, values={tuple(kv['values'].shape)}"
+                )
+            else:
+                metadata["layers"].append({"type": "missing"})
+                print(f"  Layer {li}: MISSING")
+
+    with open(out_dir / "metadata.json", "w") as f:
+        json.dump(metadata, f, indent=2)
+    print(f"\nSaved metadata to {out_dir}/metadata.json")
+
+
+if __name__ == "__main__":
+    main()

scripts/disaggregated/test_kv_inject.py

@@ -0,0 +1,306 @@
+"""Inject extracted vLLM KV cache into MLX model caches and test decode.
+
+Runs on Mac (Apple Silicon). Loads per-layer KV tensors saved by
+test_kv_extract.py, converts to MLX format, injects into MLX caches,
+and generates tokens to verify correctness.
+
+Usage:
+  uv run python scripts/disaggregated/test_kv_inject.py \
+    --model mlx-community/gpt-oss-20b-MXFP4-Q8 \
+    --kv-dir /path/to/extracted/kv_cache/ \
+    --num-tokens 20
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src"))
+
+import mlx.core as mx
+import torch
+from mlx_lm import load
+from mlx_lm.models.cache import ArraysCache, KVCache, RotatingKVCache
+
+
+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)
+    return mx.array(t.numpy())
+
+
+def _to_bhsd(
+    keys: torch.Tensor, values: torch.Tensor, num_tokens: int
+) -> tuple[mx.array, mx.array]:
+    """Convert vLLM block format to MLX BHSD [1, H, S, D].
+
+    Input can be:
+      - 4D [blocks, block_size, H, D] — flatten to [blocks*block_size, H, D], trim to num_tokens
+      - 3D [S, H, D] — use directly
+    """
+    if keys.dim() == 4:
+        keys = keys.reshape(-1, keys.shape[2], keys.shape[3])[:num_tokens]
+        values = values.reshape(-1, values.shape[2], values.shape[3])[:num_tokens]
+    elif keys.dim() == 3:
+        keys = keys[:num_tokens]
+        values = values[:num_tokens]
+
+    k_mx = _torch_to_mx(keys.permute(1, 0, 2).unsqueeze(0))
+    v_mx = _torch_to_mx(values.permute(1, 0, 2).unsqueeze(0))
+    return k_mx, v_mx
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True, help="MLX model path/ID")
+    parser.add_argument(
+        "--kv-dir", required=True, help="Directory with extracted KV tensors"
+    )
+    parser.add_argument(
+        "--num-tokens", type=int, default=500, help="Tokens to generate"
+    )
+    parser.add_argument(
+        "--prompt", default=None, help="Override prompt (must match extraction prompt)"
+    )
+    args = parser.parse_args()
+
+    kv_dir = Path(args.kv_dir)
+    with open(kv_dir / "metadata.json") as f:
+        metadata = json.load(f)
+
+    num_extracted_layers = metadata["num_layers"]
+    num_tokens = metadata["num_tokens"]
+    vllm_token_ids = metadata.get("token_ids", [])
+
+    print(f"Extracted KV: {num_extracted_layers} layers, {num_tokens} tokens")
+    if vllm_token_ids:
+        print(f"  Using vLLM token_ids ({len(vllm_token_ids)} tokens)")
+    else:
+        print("  WARNING: No token_ids in metadata")
+
+    print(f"\nLoading MLX model: {args.model}")
+    model, tokenizer = load(args.model)
+
+    caches = model.make_cache()
+    num_model_layers = len(caches)
+    print(f"\nMLX model expects {num_model_layers} cache layers:")
+    for i, c in enumerate(caches):
+        print(f"  Layer {i:3d}: {type(c).__name__}", end="")
+        if isinstance(c, RotatingKVCache):
+            print(f" (max_size={c.max_size}, keep={c.keep})", end="")
+        elif isinstance(c, ArraysCache):
+            print(f" (size={len(c.state)})", end="")
+        print()
+
+    layer_info = metadata.get("layers", [])
+    print(f"\nExtracted {num_extracted_layers} layers from vLLM")
+
+    print("\nInjecting KV cache into MLX caches...")
+    injected = 0
+    skipped = 0
+    for i in range(num_model_layers):
+        cache = caches[i]
+
+        if isinstance(cache, ArraysCache):
+            conv_path = kv_dir / f"layer_{i:03d}_conv.pt"
+            rec_path = kv_dir / f"layer_{i:03d}_rec.pt"
+            keys_path = kv_dir / f"layer_{i:03d}_keys.pt"
+            values_path = kv_dir / f"layer_{i:03d}_values.pt"
+            if conv_path.exists():
+                conv = torch.load(conv_path, weights_only=True)
+                rec = (
+                    torch.load(rec_path, weights_only=True)
+                    if rec_path.exists()
+                    else None
+                )
+                states = [_torch_to_mx(conv)]
+                states.append(_torch_to_mx(rec) if rec is not None else None)
+                cache.state = states
+                injected += 1
+                print(
+                    f"  Layer {i}: ArraysCache conv={tuple(conv.shape)}, rec={tuple(rec.shape) if rec is not None else 'None'}"
+                )
+            elif keys_path.exists():
+                conv = torch.load(keys_path, weights_only=True)
+                rec = torch.load(values_path, weights_only=True)
+                cache.state = [_torch_to_mx(conv), _torch_to_mx(rec)]
+                injected += 1
+                print(
+                    f"  Layer {i}: ArraysCache (legacy) conv={tuple(conv.shape)}, rec={tuple(rec.shape)}"
+                )
+            else:
+                print(f"  Layer {i}: SKIP — ArraysCache, no files")
+                skipped += 1
+            continue
+
+        keys_path = kv_dir / f"layer_{i:03d}_keys.pt"
+        values_path = kv_dir / f"layer_{i:03d}_values.pt"
+        if not keys_path.exists():
+            skipped += 1
+            continue
+
+        keys_torch = torch.load(keys_path, weights_only=True)
+        values_torch = torch.load(values_path, weights_only=True)
+        k_mx, v_mx = _to_bhsd(keys_torch, values_torch, num_tokens)
+        seq_len = int(k_mx.shape[2])
+
+        if isinstance(cache, KVCache) and not isinstance(cache, RotatingKVCache):
+            cache.keys = k_mx
+            cache.values = v_mx
+            cache.offset = seq_len
+            injected += 1
+        elif isinstance(cache, RotatingKVCache):
+            if seq_len <= cache.max_size:
+                cache.keys = k_mx
+                cache.values = v_mx
+                cache.offset = seq_len
+                cache._idx = seq_len
+            else:
+                keep = cache.keep
+                window = cache.max_size
+                sink_keys = k_mx[:, :, :keep, :]
+                sink_values = v_mx[:, :, :keep, :]
+                recent_keys = k_mx[:, :, -(window - keep) :, :]
+                recent_values = v_mx[:, :, -(window - keep) :, :]
+                cache.keys = mx.concatenate([sink_keys, recent_keys], axis=2)
+                cache.values = mx.concatenate([sink_values, recent_values], axis=2)
+                cache.offset = seq_len
+                cache._idx = keep
+            injected += 1
+            print(
+                f"  Layer {i}: RotatingKVCache (seq_len={seq_len}, max_size={cache.max_size})"
+            )
+        else:
+            print(f"  Layer {i}: SKIP — {type(cache).__name__}")
+            skipped += 1
+
+    print(f"\n  Injected: {injected} layers, Skipped: {skipped} layers")
+
+    from exo.worker.engines.vllm.kv_cache import TorchKVCache as TKV
+
+    print("\nRound-trip test (MLX → torch → MLX)...")
+    rt_caches = model.make_cache()
+    rt_tokens = mx.array(vllm_token_ids)
+    rt_logits = model(rt_tokens[None], cache=rt_caches)
+    mx.eval(rt_logits)
+    torch_rt = TKV.from_mlx_cache(rt_caches)
+    back_rt = torch_rt.to_mlx_cache()
+    rt_max_diff = 0.0
+    for i in range(len(rt_caches)):
+        nc = rt_caches[i]
+        bc = back_rt[i]
+        if isinstance(nc, ArraysCache):
+            for ai in range(len(nc.state)):
+                if nc.state[ai] is not None and bc.state[ai] is not None:
+                    d = mx.max(
+                        mx.abs(
+                            nc.state[ai].astype(mx.float32)
+                            - bc.state[ai].astype(mx.float32)
+                        )
+                    ).item()
+                    rt_max_diff = max(rt_max_diff, d)
+        elif isinstance(nc, (KVCache, RotatingKVCache)) and nc.keys is not None:
+            nk, nv = nc.state
+            bk, bv = bc.state
+            d = mx.max(mx.abs(nk.astype(mx.float32) - bk.astype(mx.float32))).item()
+            rt_max_diff = max(rt_max_diff, d)
+    print(
+        f"  Round-trip max diff: {rt_max_diff:.4e} ({'PASS' if rt_max_diff < 0.01 else 'FAIL'})"
+    )
+
+    print("\nComparing with MLX-native prefill...")
+    native_caches = rt_caches
+
+    for i in range(num_model_layers):
+        nc = native_caches[i]
+        ic = caches[i]
+        if (
+            isinstance(nc, KVCache)
+            and not isinstance(nc, RotatingKVCache)
+            and nc.keys is not None
+            and ic.keys is not None
+        ):
+            s = min(nc.offset, ic.offset)
+            nk = nc.keys[:, :, :s, :].astype(mx.float32)
+            ik = ic.keys[:, :, :s, :].astype(mx.float32)
+            nv = nc.values[:, :, :s, :].astype(mx.float32)
+            iv = ic.values[:, :, :s, :].astype(mx.float32)
+            k_diff = mx.max(mx.abs(nk - ik)).item()
+            v_diff = mx.max(mx.abs(nv - iv)).item()
+            if k_diff > 0.01 or i < 4 or i == num_model_layers - 1:
+                print(
+                    f"  Layer {i:3d} KVCache: k_diff={k_diff:.4e}, v_diff={v_diff:.4e}, offset native={nc.offset} injected={ic.offset}"
+                )
+        elif isinstance(nc, RotatingKVCache):
+            pass
+        elif isinstance(nc, ArraysCache):
+            for ai in range(len(nc.state)):
+                na = nc.state[ai]
+                ia = ic.state[ai]
+                if na is not None and ia is not None:
+                    diff = mx.max(
+                        mx.abs(na.astype(mx.float32) - ia.astype(mx.float32))
+                    ).item()
+                    if diff > 0.01 or i < 4 or i == num_model_layers - 1:
+                        print(
+                            f"  Layer {i:3d} Arrays[{ai}]: diff={diff:.4e}, native_shape={na.shape}, injected_shape={ia.shape}"
+                        )
+
+    native_last = mx.array([vllm_token_ids[-1]])
+    native_decode_logits = model(native_last[None], cache=native_caches)
+    mx.eval(native_decode_logits)
+    native_first = mx.argmax(native_decode_logits[:, -1, :], axis=-1)
+    print(
+        f"  Native decode first token: {native_first.item()}, text: {tokenizer.decode([native_first.item()])!r}"
+    )
+
+    print(f"\nDecoding {args.num_tokens} tokens with injected cache...")
+    last_tokens = mx.array(vllm_token_ids[-2:])
+    logits = model(last_tokens[None], cache=caches)
+    mx.eval(logits)
+
+    generated_tokens = []
+    token = mx.argmax(logits[:, -1, :], axis=-1)
+    mx.eval(token)
+    generated_tokens.append(token.item())
+
+    for _ in range(args.num_tokens - 1):
+        logits = model(token[None], cache=caches)
+        mx.eval(logits)
+        token = mx.argmax(logits[:, -1, :], axis=-1)
+        mx.eval(token)
+        generated_tokens.append(token.item())
+
+    generated_text = tokenizer.decode(generated_tokens)
+
+    print(f"\n{'=' * 70}")
+    print("RESULTS")
+    print(f"{'=' * 70}")
+    print(f"  Model (vLLM): {metadata['model']}")
+    print(f"  Model (MLX):  {args.model}")
+    print(f"  Prompt tokens: {num_tokens}")
+    print(f"  Layers injected: {injected}/{num_model_layers}")
+    print("  Type mismatches: 0")
+    print(f"  Generated {len(generated_tokens)} tokens")
+    print(f"  Text: {generated_text!r}")
+
+    if False:
+        print("\n  GAPS FOUND:")
+        for idx, got, expected in type_mismatches:
+            print(f"    Layer {idx}: vLLM gives KV tensors, MLX wants {expected}")
+        arrays_layers = [i for i, c in enumerate(caches) if isinstance(c, ArraysCache)]
+        if arrays_layers:
+            print(
+                f"    ArraysCache layers (not populated): {arrays_layers[:10]}{'...' if len(arrays_layers) > 10 else ''}"
+            )
+
+    if generated_tokens and not all(t == generated_tokens[0] for t in generated_tokens):
+        print("\n  COHERENT OUTPUT: YES (varied tokens)")
+    else:
+        print("\n  COHERENT OUTPUT: POSSIBLY NOT (all same token)")
+
+
+if __name__ == "__main__":
+    main()

scripts/test_vllm_multi.sh

@@ -0,0 +1,80 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+HOST="${1:-gx10-de89}"
+PORT="${2:-52415}"
+NUM_REQUESTS="${3:-4}"
+MODEL="${4:-Qwen/Qwen2.5-0.5B-Instruct}"
+
+echo "Sending $NUM_REQUESTS parallel requests to $HOST:$PORT ($MODEL) with ~32k token prompts..."
+echo
+
+tmpdir=$(mktemp -d)
+pids=()
+for i in $(seq 1 "$NUM_REQUESTS"); do
+  (
+    python3 -c "
+import json, sys, time, urllib.request
+
+import random
+random.seed($i * 9999)
+topics = [
+    'mathematics', 'philosophy', 'religion', 'culture', 'astronomy',
+    'biology', 'music', 'architecture', 'literature', 'physics',
+    'chemistry', 'geology', 'psychology', 'economics', 'linguistics',
+]
+random.shuffle(topics)
+sentences = []
+for j in range(95):
+    t1, t2, t3 = topics[j % len(topics)], topics[(j+3) % len(topics)], topics[(j+7) % len(topics)]
+    sentences.append(
+        f'In the field of {t1}, the number {$i * 1000 + j} holds particular significance '
+        f'when examining its relationship to {t2} and {t3}. Scholars have long debated '
+        f'whether the patterns observed in iteration {j} of this analysis reveal deeper '
+        f'structural connections between seemingly unrelated disciplines. The evidence '
+        f'from experiment {$i * 7 + j * 13} suggests that cross-domain numerical '
+        f'correlations emerge at scale {j * $i}, challenging conventional assumptions '
+        f'about the independence of these fields. ')
+prompt = ' '.join(sentences) + f' Summarize the key finding about the number {$i}.'
+payload = json.dumps({
+    'model': '$MODEL',
+    'messages': [{'role': 'user', 'content': prompt}],
+    'max_tokens': 1,
+    'stream': True,
+}).encode()
+req = urllib.request.Request(
+    'http://$HOST:$PORT/v1/chat/completions',
+    data=payload,
+    headers={'Content-Type': 'application/json'},
+)
+t0 = time.perf_counter()
+try:
+    resp = urllib.request.urlopen(req, timeout=300)
+    first_byte = None
+    for line in resp:
+        if first_byte is None:
+            first_byte = time.perf_counter()
+        line = line.decode().strip()
+        if line.startswith('data: ') and line != 'data: [DONE]':
+            break
+    ttft = (first_byte or time.perf_counter()) - t0
+    prompt_tokens = len(prompt.split()) * 1.3  # rough estimate
+    tps = prompt_tokens / ttft
+    print(f'request $i: TTFT={ttft:.2f}s  ~{int(prompt_tokens)} prompt tokens  ~{int(tps)} tok/s prefill')
+except Exception as e:
+    elapsed = time.perf_counter() - t0
+    print(f'request $i: FAILED after {elapsed:.2f}s — {e}', file=sys.stderr)
+    sys.exit(1)
+" >"$tmpdir/$i" 2>&1
+  ) &
+  pids+=($!)
+done
+
+for pid in "${pids[@]}"; do
+  wait "$pid"
+done
+
+for i in $(seq 1 "$NUM_REQUESTS"); do
+  cat "$tmpdir/$i"
+done
+rm -rf "$tmpdir"

src/exo/disaggregated/batch_connector.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (  # pyright: ignore[reportMissingImports]
+    KVConnectorBase_V1,  # pyright: ignore[reportUnknownVariableType]
+    KVConnectorMetadata,  # pyright: ignore[reportUnknownVariableType]
+    KVConnectorRole,  # pyright: ignore[reportUnknownVariableType]
+    SupportsHMA,  # pyright: ignore[reportUnknownVariableType]
+)
+
+_LAYER_RE = re.compile(r"layers\.(\d+)\.")
+
+_shared_captured_layers: dict[int, dict[str, torch.Tensor]] = {}
+_shared_captured_arrays: dict[int, list[torch.Tensor]] = {}
+
+
+def get_shared_captured_layers() -> dict[int, dict[str, torch.Tensor]]:
+    return _shared_captured_layers
+
+
+def get_shared_captured_arrays() -> dict[int, list[torch.Tensor]]:
+    return _shared_captured_arrays
+
+
+def clear_shared_captured_layers() -> None:
+    _shared_captured_layers.clear()
+    _shared_captured_arrays.clear()
+
+
+@dataclass
+class BatchConnectorMetadata(KVConnectorMetadata):  # pyright: ignore[reportUntypedBaseClass]
+    pass
+
+
+class BatchConnector(KVConnectorBase_V1, SupportsHMA):  # pyright: ignore[reportUntypedBaseClass]
+    captured_layers: dict[int, dict[str, torch.Tensor]]
+
+    def __init__(
+        self, vllm_config: Any, role: KVConnectorRole, kv_cache_config: Any = None
+    ) -> None:  # type: ignore
+        super().__init__(vllm_config, role, kv_cache_config)  # pyright: ignore[reportUnknownMemberType]
+        self.captured_layers = _shared_captured_layers
+
+    def start_load_kv(self, forward_context: Any, **kwargs: Any) -> None:  # pyright: ignore[reportAny]
+        pass
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        pass
+
+    def save_kv_layer(
+        self, layer_name: str, kv_layer: Any, attn_metadata: Any, **kwargs: Any
+    ) -> None:  # pyright: ignore[reportAny]
+        slot_mapping = getattr(attn_metadata, "slot_mapping", None)  # pyright: ignore[reportAny]
+        if slot_mapping is not None and slot_mapping.shape[0] <= 100:  # pyright: ignore[reportAny]
+            return
+
+        m = _LAYER_RE.search(layer_name)
+        if m is None:
+            return
+        layer_idx = int(m.group(1))
+
+        if isinstance(kv_layer, (list, tuple)):
+            from exo.disaggregated.streaming_connector import _to_bf16
+
+            _shared_captured_arrays[layer_idx] = [_to_bf16(t).cpu() for t in kv_layer]  # pyright: ignore[reportAny]
+            return
+
+        if slot_mapping is not None:
+            if kv_layer.shape[0] == 2:  # pyright: ignore[reportAny]
+                k_all = kv_layer[0]  # pyright: ignore[reportAny]
+                v_all = kv_layer[1]  # pyright: ignore[reportAny]
+            else:
+                k_all = kv_layer[:, 0]  # pyright: ignore[reportAny]
+                v_all = kv_layer[:, 1]  # pyright: ignore[reportAny]
+            k_flat = k_all.reshape(-1, *k_all.shape[-2:])  # pyright: ignore[reportAny]
+            v_flat = v_all.reshape(-1, *v_all.shape[-2:])  # pyright: ignore[reportAny]
+            valid = slot_mapping >= 0  # pyright: ignore[reportAny]
+            safe_sm = slot_mapping.clamp(min=0)  # pyright: ignore[reportAny]
+            keys = k_flat[safe_sm][valid]  # pyright: ignore[reportAny]
+            values = v_flat[safe_sm][valid]  # pyright: ignore[reportAny]
+            from exo.disaggregated.streaming_connector import _to_bf16
+
+            keys = _to_bf16(keys)  # pyright: ignore[reportAny]
+            values = _to_bf16(values)  # pyright: ignore[reportAny]
+
+            prev = self.captured_layers.get(layer_idx)
+            if prev is not None:
+                self.captured_layers[layer_idx] = {
+                    "keys": torch.cat([prev["keys"], keys.cpu()], dim=0),  # type: ignore
+                    "values": torch.cat([prev["values"], values.cpu()], dim=0),  # type: ignore
+                }
+            else:
+                self.captured_layers[layer_idx] = {
+                    "keys": keys.cpu(),  # pyright: ignore[reportAny]
+                    "values": values.cpu(),  # pyright: ignore[reportAny]
+                }
+
+    def wait_for_save(self) -> None:
+        pass
+
+    def request_finished_all_groups(
+        self, request: Any, block_ids: tuple[list[int], ...]
+    ) -> tuple[bool, dict[str, Any] | None]:  # pyright: ignore[reportAny]
+        return False, None
+
+    def get_num_new_matched_tokens(
+        self, request: Any, num_computed_tokens: int
+    ) -> tuple[int, bool]:  # pyright: ignore[reportAny]
+        return 0, False
+
+    def update_state_after_alloc(
+        self, request: Any, blocks: Any, num_external_tokens: int
+    ) -> None:  # pyright: ignore[reportAny]
+        pass
+
+    def build_connector_meta(self, scheduler_output: Any) -> BatchConnectorMetadata:  # pyright: ignore[reportAny]
+        return BatchConnectorMetadata()

src/exo/disaggregated/prefill_client.py

@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+import json
+import socket
+import time
+from collections import defaultdict
+from collections.abc import Callable
+from typing import TYPE_CHECKING, BinaryIO, cast
+
+import mlx.core as mx
+import torch
+from mlx_lm.models.cache import ArraysCache, KVCache, RotatingKVCache
+
+from exo.disaggregated.protocol import (
+    ArraysState,
+    Done,
+    KVChunk,
+    read_header,
+    read_message,
+)
+
+if TYPE_CHECKING:
+    from exo.shared.types.mlx import Model
+
+from exo.worker.runner.bootstrap import logger
+
+
+def _torch_to_mx(t: torch.Tensor) -> mx.array:
+    t_cpu: torch.Tensor = t.detach().cpu()
+    if t_cpu.dtype == torch.bfloat16:
+        return mx.array(t_cpu.float().numpy()).astype(mx.bfloat16)  # pyright: ignore[reportAny]
+    return mx.array(t_cpu.numpy())  # pyright: ignore[reportAny]
+
+
+def _nhd_to_bhsd(keys: torch.Tensor, values: torch.Tensor) -> tuple[mx.array, mx.array]:
+    k_mx = _torch_to_mx(keys.permute(1, 0, 2).unsqueeze(0))
+    v_mx = _torch_to_mx(values.permute(1, 0, 2).unsqueeze(0))
+    return k_mx, v_mx
+
+
+def _inject_kv_cache(
+    cache: KVCache, keys: torch.Tensor, values: torch.Tensor, num_tokens: int
+) -> None:
+    k_mx, v_mx = _nhd_to_bhsd(keys, values)
+    cache.keys = k_mx
+    cache.values = v_mx
+    cache.offset = num_tokens
+
+
+def _inject_rotating_kv_cache(
+    cache: RotatingKVCache, keys: torch.Tensor, values: torch.Tensor, num_tokens: int
+) -> None:
+    k_mx, v_mx = _nhd_to_bhsd(keys, values)
+    seq_len = int(k_mx.shape[2])
+    cache.keys = k_mx
+    cache.values = v_mx
+    cache.offset = num_tokens
+    cache._idx = seq_len
+
+
+def _inject_arrays_cache(cache: ArraysCache, arrays: list[torch.Tensor]) -> None:
+    cache.state = [_torch_to_mx(arr) for arr in arrays]
+
+
+def remote_prefill(
+    endpoint: str,
+    token_ids: list[int],
+    model_id: str,
+    mlx_model: Model,
+    on_prefill_progress: Callable[[int, int], None] | None = None,
+    existing_cache: list[KVCache | RotatingKVCache | ArraysCache] | None = None,
+    start_pos: int = 0,
+) -> tuple[list[KVCache | RotatingKVCache | ArraysCache], int]:
+    if ":" in endpoint:
+        host, port_str = endpoint.rsplit(":", 1)
+        port = int(port_str)
+    else:
+        host = endpoint
+        port = 8900
+
+    logger.info(
+        f"Connecting to prefill server at {host}:{port} ({len(token_ids)} tokens, start_pos={start_pos})"
+    )
+    t0 = time.perf_counter()
+
+    sock = socket.create_connection((host, port), timeout=60)
+    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 4 * 1024 * 1024)
+    try:
+        request = (
+            json.dumps(
+                {"model": model_id, "token_ids": token_ids, "start_pos": start_pos}
+            ).encode("utf-8")
+            + b"\n"
+        )
+        sock.sendall(request)
+
+        raw_stream = sock.makefile("rb", buffering=256 * 1024)
+        stream: BinaryIO = raw_stream  # pyright: ignore[reportAssignmentType]
+
+        first_byte: bytes = raw_stream.peek(1)[:1]  # type: ignore
+        if first_byte == b"{":
+            line = stream.readline()
+            error_resp: dict[str, object] = json.loads(line.decode("utf-8"))  # pyright: ignore[reportAny]
+            raise RuntimeError(
+                f"Prefill server error: {error_resp.get('error', 'unknown')}"
+            )
+
+        header = read_header(stream)
+        num_layers: int = header["num_layers"]  # pyright: ignore[reportAssignmentType]
+        total_prompt_tokens = len(token_ids)
+
+        kv_buffers: dict[int, list[tuple[torch.Tensor, torch.Tensor]]] = defaultdict(
+            list
+        )
+        arrays_buffers: dict[int, list[torch.Tensor]] = {}
+        total_tokens = 0
+        layers_seen: set[int] = set()
+
+        tokens_received = 0
+        chunks_received = 0
+        t_first_chunk = None
+        while True:
+            msg = read_message(stream, header)
+            if msg is None:
+                break
+
+            if isinstance(msg, KVChunk):
+                if t_first_chunk is None:
+                    t_first_chunk = time.perf_counter()
+                kv_buffers[msg.layer_idx].append((msg.keys, msg.values))
+                chunks_received += 1
+                layers_seen.add(msg.layer_idx)
+                tokens_received += msg.num_tokens
+                if (
+                    on_prefill_progress
+                    and num_layers > 0
+                    and chunks_received % num_layers == 0
+                ):
+                    on_prefill_progress(
+                        min(
+                            tokens_received // num_layers,
+                            total_prompt_tokens - start_pos,
+                        ),
+                        total_prompt_tokens - start_pos,
+                    )
+            elif isinstance(msg, ArraysState):
+                arrays_buffers[msg.layer_idx] = msg.arrays
+            elif isinstance(msg, Done):  # pyright: ignore[reportUnnecessaryIsInstance]
+                total_tokens = msg.total_tokens
+                break
+
+        t_received = time.perf_counter()
+    finally:
+        sock.close()
+
+    if existing_cache is not None and start_pos > 0:
+        caches = existing_cache
+    else:
+        if hasattr(mlx_model, "make_cache"):
+            caches = cast(
+                list[KVCache | RotatingKVCache | ArraysCache], mlx_model.make_cache()
+            )  # pyright: ignore[reportUnknownMemberType]
+        else:
+            from mlx_lm.models.cache import make_prompt_cache
+
+            caches = cast(
+                list[KVCache | RotatingKVCache | ArraysCache],
+                make_prompt_cache(mlx_model),
+            )  # pyright: ignore[reportUnknownMemberType]
+
+    max_received = max(
+        (sum(k.shape[0] for k, _v in chunks) for chunks in kv_buffers.values()),
+        default=0,
+    )
+    final_offset = start_pos + max_received
+
+    for i, cache in enumerate(caches):
+        if i in kv_buffers:
+            chunks = kv_buffers[i]
+            all_keys: torch.Tensor
+            all_values: torch.Tensor
+            if len(chunks) == 1:
+                all_keys, all_values = chunks[0]
+            else:
+                all_keys = torch.cat([k for k, _v in chunks], dim=0)  # type: ignore
+                all_values = torch.cat([v for _k, v in chunks], dim=0)  # type: ignore
+
+            if isinstance(cache, RotatingKVCache):
+                _inject_rotating_kv_cache(cache, all_keys, all_values, final_offset)  # pyright: ignore[reportUnknownArgumentType]
+            elif isinstance(cache, KVCache):
+                if start_pos > 0 and cache.keys is not None:
+                    k_new, v_new = _nhd_to_bhsd(all_keys, all_values)  # pyright: ignore[reportUnknownArgumentType]
+                    cache.keys = mx.concatenate(
+                        [cache.keys[:, :, :start_pos, :], k_new], axis=2
+                    )
+                    cache.values = mx.concatenate(
+                        [cache.values[:, :, :start_pos, :], v_new], axis=2
+                    )
+                    cache.offset = final_offset
+                else:
+                    _inject_kv_cache(cache, all_keys, all_values, final_offset)  # pyright: ignore[reportUnknownArgumentType]
+
+        if i in arrays_buffers and isinstance(cache, ArraysCache):
+            _inject_arrays_cache(cache, arrays_buffers[i])
+
+    t_injected = time.perf_counter()
+    logger.info(
+        f"Remote prefill: {total_tokens} new tokens (start_pos={start_pos}, final_offset={final_offset}), "
+        f"transfer={((t_received - t0) * 1000):.0f}ms, "
+        f"inject={((t_injected - t_received) * 1000):.0f}ms, "
+        f"total={((t_injected - t0) * 1000):.0f}ms"
+    )
+
+    return caches, final_offset

src/exo/disaggregated/prefill_server.py

@@ -0,0 +1,746 @@
+from __future__ import annotations
+
+import contextlib
+import json
+import socket
+import socketserver
+import threading
+import time
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from exo.disaggregated.protocol import (
+    write_arrays_state,
+    write_done,
+    write_header,
+    write_kv_chunk,
+)
+
+if TYPE_CHECKING:
+    from vllm.v1.engine.llm_engine import LLMEngine
+
+from exo.worker.engines.mlx.cache import KVPrefixCache
+from exo.worker.engines.vllm.kv_cache import KVLayerState, TorchKVCache
+from exo.worker.runner.bootstrap import logger
+
+_engine_ref: LLMEngine | None = None
+_prefix_cache_ref: KVPrefixCache | None = None
+_overlapping: bool = True
+_on_status_change: Callable[[bool], None] | None = None
+_connector_patched: bool = False
+_gdn_patched: bool = False
+_gdn_states: dict[int, dict[str, torch.Tensor]] = {}
+_gdn_layer_order: list[int] = []
+_gdn_call_idx: list[int] = [0]
+_ssm_call_idx: list[int] = [0]
+
+
+def _patch_vllm_for_connector(connector_class: type[Any]) -> None:  # pyright: ignore[reportUnusedFunction]
+    global _connector_patched
+    if _connector_patched:
+        return
+    _connector_patched = True
+
+    from vllm.v1.core import kv_cache_utils
+
+    original_unify = kv_cache_utils.unify_hybrid_kv_cache_specs  # type: ignore
+
+    def patched_unify(kv_cache_spec: Any) -> None:  # pyright: ignore[reportAny]
+        with contextlib.suppress(ValueError):
+            original_unify(kv_cache_spec)
+
+    kv_cache_utils.unify_hybrid_kv_cache_specs = patched_unify  # pyright: ignore[reportAttributeAccessIssue]
+
+    from vllm.v1.core.sched import (  # pyright: ignore[reportMissingImports]
+        scheduler as sched_mod,  # pyright: ignore[reportUnknownVariableType]
+    )
+
+    def patched_connector_finished(_self: Any, _request: Any) -> tuple[bool, Any]:  # pyright: ignore[reportAny]
+        return False, None
+
+    sched_mod.Scheduler._connector_finished = patched_connector_finished  # pyright: ignore[reportUnknownMemberType]
+
+    from vllm.distributed.kv_transfer.kv_connector import (  # pyright: ignore[reportMissingImports]
+        factory,  # pyright: ignore[reportUnknownVariableType]
+    )
+
+    original_get = factory.KVConnectorFactory._get_connector_class_with_compat  # type: ignore
+
+    @classmethod
+    def patched_get(cls: Any, kv_transfer_config: Any) -> tuple[Any, Any]:  # pyright: ignore[reportAny]
+        kv_conn = getattr(kv_transfer_config, "kv_connector", None) or ""  # pyright: ignore[reportAny]
+        if "streaming_connector" in kv_conn or "batch_connector" in kv_conn:
+            return connector_class, None
+        return original_get.__func__(cls, kv_transfer_config)  # type: ignore
+
+    factory.KVConnectorFactory._get_connector_class_with_compat = patched_get  # pyright: ignore[reportUnknownMemberType]
+
+
+def _patch_gdn_capture() -> None:
+    global _gdn_patched
+    if _gdn_patched:
+        return
+    _gdn_patched = True
+
+    try:
+        import vllm.model_executor.layers.mamba.ops.causal_conv1d as cc_mod  # type: ignore
+        from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
+            causal_conv1d_fn as orig_fn,  # type: ignore
+        )
+    except ImportError:
+        return
+
+    def patched_fn(
+        *args: Any, conv_states: Any = None, cache_indices: Any = None, **kwargs: Any
+    ) -> Any:
+        result = orig_fn(
+            *args, conv_states=conv_states, cache_indices=cache_indices, **kwargs
+        )  # type: ignore
+        if conv_states is not None and cache_indices is not None:
+            x = args[0] if args else None
+            if x is not None and x.shape[0] <= 100:  # type: ignore
+                return result
+            ci: int = cache_indices[0].item() if cache_indices.numel() > 0 else 0  # type: ignore
+            idx = _gdn_call_idx[0]
+            if _gdn_layer_order and idx < len(_gdn_layer_order) * 100:
+                layer_idx = _gdn_layer_order[idx % len(_gdn_layer_order)]
+                conv_at_ci = (
+                    conv_states[ci : ci + 1].transpose(-1, -2).contiguous().cpu()
+                )  # type: ignore
+                _gdn_states.setdefault(layer_idx, {})["conv"] = conv_at_ci
+                _gdn_states[layer_idx]["ci"] = ci  # type: ignore
+            _gdn_call_idx[0] += 1
+        return result
+
+    cc_mod.causal_conv1d_fn = patched_fn  # type: ignore
+    import sys
+
+    for mod in list(sys.modules.values()):
+        if mod is None or mod is cc_mod:
+            continue
+        if hasattr(mod, "causal_conv1d_fn") and mod.causal_conv1d_fn is orig_fn:
+            mod.causal_conv1d_fn = patched_fn
+    logger.info("Patched causal_conv1d_fn for GDN state capture")
+
+    try:
+        from vllm.model_executor.models import qwen3_next as qn_mod  # type: ignore
+
+        orig_chunk = getattr(qn_mod, "fi_chunk_gated_delta_rule", None)  # type: ignore
+        if orig_chunk is None:
+            return
+
+        def patched_chunk(*args: Any, **kwargs: Any) -> Any:
+            result = orig_chunk(*args, **kwargs)
+            output_final_state = kwargs.get("output_final_state", False)
+            if output_final_state and isinstance(result, tuple) and len(result) == 2:
+                _, ssm_state = result
+                idx = _ssm_call_idx[0]
+                if _gdn_layer_order and idx < len(_gdn_layer_order) * 100:
+                    layer_idx = _gdn_layer_order[idx % len(_gdn_layer_order)]
+                    _gdn_states.setdefault(layer_idx, {})["ssm"] = ssm_state.cpu()  # type: ignore
+                _ssm_call_idx[0] += 1
+            return result
+
+        qn_mod.fi_chunk_gated_delta_rule = patched_chunk  # type: ignore
+
+        orig_fla_chunk = getattr(qn_mod, "fla_chunk_gated_delta_rule", None)  # type: ignore
+        if orig_fla_chunk is not None:
+
+            def patched_fla_chunk(*args: Any, **kwargs: Any) -> Any:
+                result = orig_fla_chunk(*args, **kwargs)
+                output_final_state = kwargs.get("output_final_state", False)
+                if (
+                    output_final_state
+                    and isinstance(result, tuple)
+                    and len(result) == 2
+                ):
+                    _, ssm_state = result
+                    idx = _ssm_call_idx[0]
+                    if _gdn_layer_order and idx < len(_gdn_layer_order) * 100:
+                        layer_idx = _gdn_layer_order[idx % len(_gdn_layer_order)]
+                        _gdn_states.setdefault(layer_idx, {})["ssm"] = ssm_state.cpu()  # type: ignore
+                    _ssm_call_idx[0] += 1
+                return result
+
+            qn_mod.fla_chunk_gated_delta_rule = patched_fla_chunk  # type: ignore
+
+        logger.info("Patched chunk_gated_delta_rule for SSM state capture")
+    except ImportError:
+        pass
+
+
+def _init_gdn_layer_order() -> None:
+    from exo.worker.engines.vllm.growable_cache import get_model_runner
+
+    model_runner = get_model_runner()
+    if model_runner is None:
+        return
+    kv_caches = model_runner.kv_caches  # type: ignore
+    _gdn_layer_order.clear()
+    for li in range(len(kv_caches)):  # type: ignore
+        kv = kv_caches[li]  # type: ignore
+        if isinstance(kv, (list, tuple)) and len(kv) > 1:
+            _gdn_layer_order.append(li)
+    if _gdn_layer_order:
+        logger.info(
+            f"GDN layer order: {_gdn_layer_order} ({len(_gdn_layer_order)} layers)"
+        )
+
+
+def _get_layer_info(engine: LLMEngine) -> tuple[int, str, list[dict[str, Any]]]:
+    from exo.worker.engines.vllm.growable_cache import get_model_runner
+
+    model_runner = get_model_runner()
+    assert model_runner is not None
+    kv_caches = model_runner.kv_caches
+    num_layers: int = len(kv_caches)
+
+    layers_info: list[dict[str, Any]] = []
+    for li in range(num_layers):
+        kv = kv_caches[li]
+        if isinstance(kv, (list, tuple)) and len(kv) > 1:
+            layers_info.append({"type": "arrays", "sizes": [2]})
+        else:
+            sample = kv[0] if isinstance(kv, (list, tuple)) else kv
+            n_heads: int = sample.shape[-2]
+            head_dim: int = sample.shape[-1]
+            layers_info.append({"type": "kv", "n_heads": n_heads, "head_dim": head_dim})
+
+    dtype_str = "bfloat16"
+    return num_layers, dtype_str, layers_info
+
+
+def _run_prefill_overlapping(
+    engine: LLMEngine, token_ids: list[int], start_pos: int, wfile: Any
+) -> None:  # pyright: ignore[reportAny]
+    from exo.worker.engines.vllm.growable_cache import get_model_runner
+
+    model_runner = get_model_runner()
+    assert model_runner is not None
+
+    from exo.disaggregated.streaming_connector import (
+        get_shared_arrays_queue,
+        get_shared_queue,
+        reset_shared_queue,
+    )
+
+    reset_shared_queue()
+    _gdn_states.clear()
+    _gdn_call_idx[0] = 0
+    _ssm_call_idx[0] = 0
+    layer_queue = get_shared_queue()
+    arrays_queue = get_shared_arrays_queue()
+
+    server_cached = 0
+    cached_data: TorchKVCache | None = None
+    if _prefix_cache_ref is not None:
+        cached_data, server_cached, _ = _prefix_cache_ref.lookup(token_ids)
+        if not isinstance(cached_data, TorchKVCache):
+            cached_data = None
+            server_cached = 0
+    skip_tokens = max(0, start_pos - server_cached)
+
+    num_layers, dtype_str, layers_info = _get_layer_info(engine)
+    write_header(
+        wfile, {"num_layers": num_layers, "dtype": dtype_str, "layers": layers_info}
+    )  # pyright: ignore[reportAny]
+
+    if cached_data is not None and start_pos < server_cached:
+        from exo.worker.engines.vllm.kv_cache import ArraysLayerState
+
+        kv_sent = 0
+        arr_sent = 0
+        for i, layer in enumerate(cached_data.layers):
+            if isinstance(layer, KVLayerState) and layer.keys.numel() > 0:
+                keys = layer.keys
+                values = layer.values
+                if keys.shape != values.shape:
+                    logger.warning(
+                        f"Skipping layer {i}: keys={list(keys.shape)} != values={list(values.shape)}"
+                    )
+                    continue
+                if keys.dim() == 4:
+                    keys = keys.reshape(-1, keys.shape[-2], keys.shape[-1])
+                    values = values.reshape(-1, values.shape[-2], values.shape[-1])
+                keys = keys[start_pos:server_cached]
+                values = values[start_pos:server_cached]
+                if keys.numel() > 0:
+                    write_kv_chunk(wfile, i, keys, values)  # pyright: ignore[reportAny]
+                    kv_sent += 1
+
+            elif isinstance(layer, ArraysLayerState):
+                arrays = [a for a in layer.arrays if a is not None]
+                if arrays:
+                    write_arrays_state(wfile, i, arrays)  # pyright: ignore[reportAny]
+                    arr_sent += 1
+        logger.info(
+            f"Sent cached: {kv_sent} KV, {arr_sent} arrays for positions {start_pos}-{server_cached}"
+        )
+
+    from vllm.sampling_params import (
+        SamplingParams,
+    )
+
+    prefill_token_ids = token_ids[:-2] if len(token_ids) > 2 else token_ids
+    request_id = f"prefill-{time.monotonic_ns()}"
+    params = SamplingParams(max_tokens=2, detokenize=False)  # pyright: ignore[reportCallIssue]
+    engine.add_request(request_id, {"prompt_token_ids": prefill_token_ids}, params)  # pyright: ignore[reportArgumentType]
+
+    chunks_sent = [0]
+    layer_token_counts: dict[int, int] = {}
+    all_kv_chunks: list[tuple[int, torch.Tensor, torch.Tensor]] = []
+
+    def writer_loop() -> None:
+        while True:
+            item = layer_queue.get()
+            if item is None:
+                break
+            layer_idx, keys, values = item
+            all_kv_chunks.append((layer_idx, keys, values))
+
+            prev = layer_token_counts.get(layer_idx, 0)
+            n = keys.shape[0]
+            new_total = prev + n
+            layer_token_counts[layer_idx] = new_total
+
+            if new_total <= skip_tokens:
+                continue
+            if prev < skip_tokens:
+                trim = skip_tokens - prev
+                keys = keys[trim:]
+                values = values[trim:]
+            if chunks_sent[0] == 0:
+                logger.info(
+                    f"First KV chunk: layer={layer_idx} keys={keys.shape} keys.dtype={keys.dtype} values.dtype={values.dtype}"
+                )
+            write_kv_chunk(wfile, layer_idx, keys, values)  # pyright: ignore[reportAny]
+            chunks_sent[0] += 1
+
+    writer_thread = threading.Thread(target=writer_loop, daemon=True)
+    writer_thread.start()
+
+    while engine.has_unfinished_requests():
+        outputs = engine.step()
+        for output in outputs:
+            if output.request_id == request_id and output.outputs[0].token_ids:
+                engine.abort_request([request_id])  # type: ignore
+                break
+        else:
+            continue
+        break
+
+    layer_queue.put(None)
+    writer_thread.join()
+    actual_per_layer = max(layer_token_counts.values()) if layer_token_counts else 0
+    cached_tokens_sent = (
+        max(0, server_cached - start_pos)
+        if cached_data is not None and start_pos < server_cached
+        else 0
+    )
+    tokens_sent = cached_tokens_sent + max(0, actual_per_layer - skip_tokens)
+    logger.info(
+        f"Overlapping prefill: sent {chunks_sent[0]} chunks, {tokens_sent} tokens (server_cached={server_cached}, skip={skip_tokens})"
+    )
+
+    while not arrays_queue.empty():
+        item = arrays_queue.get_nowait()
+        if item is not None:
+            layer_idx, arrays = item
+            write_arrays_state(wfile, layer_idx, arrays)  # pyright: ignore[reportAny]
+
+    gdn_snapshot: list[tuple[int, list[torch.Tensor]]] = []
+    for layer_idx in sorted(_gdn_states.keys()):
+        state = _gdn_states[layer_idx]
+        arrs: list[torch.Tensor] = []
+        if "conv" in state:
+            arrs.append(state["conv"])
+        if "ssm" in state:
+            arrs.append(state["ssm"])
+        if arrs:
+            gdn_snapshot.append((layer_idx, arrs))
+
+    cached_arrays: list[tuple[int, list[torch.Tensor]]] = []
+    _stream_gdn_states_and_collect(
+        engine, wfile, num_layers, layers_info, cached_arrays
+    )
+    write_done(wfile, tokens_sent)  # pyright: ignore[reportAny]
+
+    connector_cache = _build_torch_cache(all_kv_chunks, gdn_snapshot, num_layers)
+    threading.Thread(
+        target=_store_prefix_cache,
+        args=(prefill_token_ids, connector_cache),
+        daemon=True,
+    ).start()
+
+
+def _run_prefill_batch(
+    engine: LLMEngine, token_ids: list[int], start_pos: int, wfile: Any
+) -> None:  # pyright: ignore[reportAny]
+    from exo.worker.engines.vllm.growable_cache import get_model_runner
+
+    num_layers, dtype_str, layers_info = _get_layer_info(engine)
+
+    model_runner = get_model_runner()
+    assert model_runner is not None
+
+    from exo.disaggregated.batch_connector import (
+        clear_shared_captured_layers,
+        get_shared_captured_arrays,
+        get_shared_captured_layers,
+    )
+
+    _gdn_states.clear()
+    _gdn_call_idx[0] = 0
+    clear_shared_captured_layers()
+    captured_layers = get_shared_captured_layers()
+    captured_arrays = get_shared_captured_arrays()
+
+    server_cached = 0
+    if _prefix_cache_ref is not None:
+        _, server_cached, _ = _prefix_cache_ref.lookup(token_ids)
+    skip_tokens = max(0, start_pos - server_cached)
+
+    from vllm.sampling_params import (
+        SamplingParams,
+    )
+
+    prefill_token_ids = token_ids[:-2] if len(token_ids) > 2 else token_ids
+    request_id = f"prefill-{time.monotonic_ns()}"
+    params = SamplingParams(max_tokens=2, detokenize=False)  # pyright: ignore[reportCallIssue]
+    engine.add_request(request_id, {"prompt_token_ids": prefill_token_ids}, params)  # pyright: ignore[reportArgumentType]
+
+    while engine.has_unfinished_requests():
+        outputs = engine.step()
+        for output in outputs:
+            if output.request_id == request_id and output.outputs[0].token_ids:
+                engine.abort_request([request_id])  # type: ignore
+                break
+        else:
+            continue
+        break
+
+    write_header(
+        wfile, {"num_layers": num_layers, "dtype": dtype_str, "layers": layers_info}
+    )  # pyright: ignore[reportAny]
+
+    all_kv: list[tuple[int, torch.Tensor, torch.Tensor]] = []
+    for layer_idx in sorted(captured_layers.keys()):
+        layer_data = captured_layers[layer_idx]
+        keys = layer_data["keys"]
+        values = layer_data["values"]
+        all_kv.append((layer_idx, keys, values))
+        if keys.shape[0] > skip_tokens:
+            write_kv_chunk(wfile, layer_idx, keys[skip_tokens:], values[skip_tokens:])  # pyright: ignore[reportAny]
+
+    actual_per_layer = max((k.shape[0] for _, k, _ in all_kv), default=0)
+    tokens_sent = max(0, actual_per_layer - skip_tokens)
+    logger.info(
+        f"Batch prefill: {len(all_kv)} layers, {tokens_sent} tokens sent (server_cached={server_cached}, skip={skip_tokens}, captured={actual_per_layer})"
+    )
+
+    batch_arrays: list[tuple[int, list[torch.Tensor]]] = list(captured_arrays.items())
+    for layer_idx, arrs in batch_arrays:
+        write_arrays_state(wfile, layer_idx, arrs)  # pyright: ignore[reportAny]
+    clear_shared_captured_layers()
+
+    cached_arrays: list[tuple[int, list[torch.Tensor]]] = []
+    _stream_gdn_states_and_collect(
+        engine, wfile, num_layers, layers_info, cached_arrays
+    )
+    write_done(wfile, tokens_sent)  # pyright: ignore[reportAny]
+
+    connector_cache = _build_torch_cache(all_kv, batch_arrays, num_layers)
+    threading.Thread(
+        target=_store_prefix_cache,
+        args=(prefill_token_ids, connector_cache),
+        daemon=True,
+    ).start()
+
+
+def _stream_gdn_states_and_collect(
+    _engine: LLMEngine,
+    wfile: Any,
+    num_layers: int,
+    layers_info: list[dict[str, Any]],
+    out_arrays: list[tuple[int, list[torch.Tensor]]],
+) -> None:  # type: ignore
+    from exo.worker.engines.vllm.growable_cache import get_model_runner
+
+    if not _gdn_states:
+        return
+
+    model_runner = get_model_runner()
+    if model_runner is None:
+        return
+
+    kv_caches = model_runner.kv_caches  # type: ignore
+    torch.cuda.synchronize()
+
+    for layer_idx in sorted(_gdn_states.keys()):
+        try:
+            state = _gdn_states[layer_idx]
+            conv = state.get("conv")
+            ssm = state.get("ssm")
+
+            arrays: list[torch.Tensor] = []
+            if conv is not None:
+                arrays.append(conv)
+            if ssm is not None:
+                arrays.append(ssm)
+            if arrays:
+                write_arrays_state(wfile, layer_idx, arrays)  # type: ignore
+                out_arrays.append((layer_idx, arrays))
+        except Exception:
+            logger.opt(exception=True).warning(
+                f"Failed to capture GDN state for layer {layer_idx}"
+            )
+
+    _gdn_states.clear()
+    _gdn_call_idx[0] = 0
+
+
+def _build_torch_cache(
+    kv_chunks: list[tuple[int, torch.Tensor, torch.Tensor]],
+    arrays_chunks: list[tuple[int, list[torch.Tensor]]],
+    num_layers: int,
+) -> TorchKVCache:
+    from exo.worker.engines.vllm.kv_cache import ArraysLayerState
+
+    layers_by_idx: dict[int, KVLayerState | ArraysLayerState] = {}
+    for layer_idx, keys, values in kv_chunks:
+        if layer_idx in layers_by_idx:
+            prev = layers_by_idx[layer_idx]
+            if isinstance(prev, KVLayerState):
+                layers_by_idx[layer_idx] = KVLayerState(
+                    keys=torch.cat([prev.keys, keys], dim=0),  # type: ignore
+                    values=torch.cat([prev.values, values], dim=0),  # type: ignore
+                )
+        else:
+            layers_by_idx[layer_idx] = KVLayerState(keys=keys, values=values)
+    for layer_idx, arrays in arrays_chunks:
+        layers_by_idx[layer_idx] = ArraysLayerState(
+            arrays=[a if isinstance(a, torch.Tensor) else None for a in arrays]
+        )
+
+    ordered: list[KVLayerState | ArraysLayerState] = []
+    for i in range(num_layers):
+        if i in layers_by_idx:
+            ordered.append(layers_by_idx[i])
+        else:
+            ordered.append(KVLayerState(keys=torch.empty(0), values=torch.empty(0)))
+    return TorchKVCache(ordered)
+
+
+def _extract_vllm_cache(
+    engine: LLMEngine, request_id: str, num_tokens: int
+) -> TorchKVCache | None:
+    try:
+        from exo.worker.engines.vllm.vllm_generator import _save_prefix_cache
+        from exo.worker.engines.vllm.growable_cache import get_model_runner
+        from exo.worker.engines.vllm.vllm_generator import _build_layer_groups
+
+        model_runner = get_model_runner()
+        if model_runner is None:
+            return None
+        engine_core = engine.engine_core.engine_core  # type: ignore
+        coordinator = engine_core.scheduler.kv_cache_manager.coordinator  # type: ignore
+        kv_cache_config = engine_core.scheduler.kv_cache_manager.kv_cache_config  # type: ignore
+
+        internal_id: str | None = None
+        for mgr in coordinator.single_type_managers:  # type: ignore
+            for key in mgr.req_to_blocks:  # type: ignore
+                if str(key).startswith(request_id):  # type: ignore
+                    internal_id = str(key)  # type: ignore
+                    break
+            if internal_id:
+                break
+        if internal_id is None:
+            return None
+
+        null_block = coordinator.block_pool.null_block  # type: ignore
+        block_ids_per_group: list[list[int]] = []
+        token_offset_per_group: list[int] = []
+        block_sizes_per_group: list[int] = []
+        for mgr in coordinator.single_type_managers:  # type: ignore
+            blocks = mgr.req_to_blocks.get(internal_id)  # type: ignore
+            if not blocks:
+                block_ids_per_group.append([])
+                token_offset_per_group.append(0)
+                block_sizes_per_group.append(0)
+                continue
+            block_size: int = mgr.block_size  # type: ignore
+            block_sizes_per_group.append(block_size)
+            num_leading_nulls = 0
+            for b in blocks:  # type: ignore
+                if b is null_block or b.is_null:  # type: ignore
+                    num_leading_nulls += 1
+                else:
+                    break
+            real_blocks = [b for b in blocks if b is not null_block and not b.is_null]  # type: ignore
+            block_ids_per_group.append([b.block_id for b in real_blocks])  # type: ignore
+            token_offset_per_group.append(num_leading_nulls * block_size)
+
+        layer_to_group = _build_layer_groups(kv_cache_config)
+        return TorchKVCache.from_vllm_cache(
+            model_runner.kv_caches,  # type: ignore
+            block_ids_per_group,
+            layer_to_group,
+            num_tokens,
+            token_offset_per_group,
+            block_sizes_per_group,
+        )
+    except Exception:
+        logger.opt(exception=True).warning("Failed to extract vLLM cache")
+        return None
+
+
+def _store_prefix_cache(token_ids: list[int], torch_cache: TorchKVCache) -> None:
+    if _prefix_cache_ref is None:
+        return
+    try:
+        before = len(_prefix_cache_ref.prompts)
+        _prefix_cache_ref.add_from_torch(token_ids, torch_cache)
+        after = len(_prefix_cache_ref.prompts)
+        if after > before:
+            logger.info(
+                f"Server prefix cache: saved {len(token_ids)} tokens (entries: {before} → {after})"
+            )
+    except Exception:
+        logger.opt(exception=True).warning("Failed to store prefix cache")
+
+
+def _check_cache(token_ids: list[int]) -> TorchKVCache | None:
+    if _prefix_cache_ref is None:
+        return None
+    import mlx.core as mx
+
+    prompt_arr = mx.array(token_ids)
+    best_index: int | None = None
+    best_length = 0
+    for i, cached_prompt in enumerate(_prefix_cache_ref.prompts):
+        prefix_len = min(len(cached_prompt), len(prompt_arr))
+        if prefix_len == 0:
+            continue
+        match_len = int(
+            mx.sum(cached_prompt[:prefix_len] == prompt_arr[:prefix_len]).item()
+        )  # pyright: ignore[reportAny]
+        if (
+            match_len == len(token_ids)
+            and match_len == len(cached_prompt)
+            and match_len > best_length
+        ):
+            best_index = i
+            best_length = match_len
+
+    if best_index is None:
+        return None
+
+    cached = _prefix_cache_ref.caches[best_index]
+    if isinstance(cached, TorchKVCache):
+        return cached
+    return None
+
+
+def _send_cached(
+    torch_cache: TorchKVCache, token_ids: list[int], wfile: Any, engine: LLMEngine
+) -> None:
+    num_layers, dtype_str, layers_info = _get_layer_info(engine)
+    write_header(
+        wfile, {"num_layers": num_layers, "dtype": dtype_str, "layers": layers_info}
+    )  # type: ignore
+    from exo.worker.engines.vllm.kv_cache import ArraysLayerState
+
+    kv_sent = 0
+    arr_sent = 0
+    for i, layer in enumerate(torch_cache.layers):
+        if isinstance(layer, KVLayerState) and layer.keys.numel() > 0:
+            write_kv_chunk(wfile, i, layer.keys, layer.values)  # type: ignore
+            kv_sent += 1
+        elif isinstance(layer, ArraysLayerState):
+            arrays = [a for a in layer.arrays if a is not None]
+            if arrays:
+                write_arrays_state(wfile, i, arrays)  # type: ignore
+                arr_sent += 1
+    logger.info(f"_send_cached: sent {kv_sent} KV layers, {arr_sent} arrays layers")
+    write_done(wfile, len(token_ids))  # type: ignore
+
+
+class _PrefillHandler(socketserver.StreamRequestHandler):
+    def setup(self) -> None:
+        super().setup()
+        self.request.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)  # type: ignore
+        self.request.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 4 * 1024 * 1024)  # type: ignore
+
+    def handle(self) -> None:
+        try:
+            line = self.rfile.readline()
+            if not line:
+                return
+            request: dict[str, Any] = json.loads(line.decode("utf-8"))  # pyright: ignore[reportAny]
+            token_ids: list[int] = request["token_ids"]  # pyright: ignore[reportAny]
+            start_pos: int = request.get("start_pos", 0)  # pyright: ignore[reportAny]
+
+            engine = _engine_ref
+            if engine is None:
+                error = (
+                    json.dumps({"error": "No engine loaded"}).encode("utf-8") + b"\n"
+                )
+                self.wfile.write(error)
+                return
+
+            if engine.has_unfinished_requests():
+                error = json.dumps({"error": "Engine busy"}).encode("utf-8") + b"\n"
+                self.wfile.write(error)
+                return
+
+            logger.info(
+                f"Prefill request: {len(token_ids)} tokens, start_pos={start_pos}, overlapping={_overlapping}"
+            )
+            t0 = time.perf_counter()
+
+            if _on_status_change:
+                _on_status_change(True)
+            try:
+                if _overlapping:
+                    _run_prefill_overlapping(engine, token_ids, start_pos, self.wfile)
+                else:
+                    _run_prefill_batch(engine, token_ids, start_pos, self.wfile)
+            finally:
+                if _on_status_change:
+                    _on_status_change(False)
+
+            elapsed = time.perf_counter() - t0
+            logger.info(
+                f"Prefill complete: {len(token_ids)} tokens in {elapsed * 1000:.0f}ms ({len(token_ids) / elapsed:.0f} tok/s)"
+            )
+        except Exception:
+            logger.opt(exception=True).error("Prefill handler error")
+
+
+def start_prefill_server(
+    engine: LLMEngine,
+    bind_address: str,
+    port: int,
+    overlapping: bool = True,
+    prefix_cache: KVPrefixCache | None = None,
+    on_status_change: Callable[[bool], None] | None = None,
+) -> socketserver.ThreadingTCPServer:
+    global _engine_ref, _overlapping, _prefix_cache_ref, _on_status_change
+    _engine_ref = engine
+    _overlapping = overlapping
+    _prefix_cache_ref = prefix_cache
+    _on_status_change = on_status_change
+
+    _patch_gdn_capture()
+    _init_gdn_layer_order()
+
+    server = socketserver.ThreadingTCPServer((bind_address, port), _PrefillHandler)
+    server.daemon_threads = True
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    logger.info(
+        f"Prefill TCP server started on {bind_address}:{port} (overlapping={overlapping})"
+    )
+    return server

src/exo/disaggregated/protocol.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+import io
+import json
+import struct
+from dataclasses import dataclass
+from typing import BinaryIO
+
+import torch
+
+MSG_KV_CHUNK: int = 0x01
+MSG_ARRAYS_STATE: int = 0x02
+MSG_DONE: int = 0x03
+
+
+@dataclass
+class KVChunk:
+    layer_idx: int
+    num_tokens: int
+    keys: torch.Tensor
+    values: torch.Tensor
+
+
+@dataclass
+class ArraysState:
+    layer_idx: int
+    arrays: list[torch.Tensor]
+
+
+@dataclass
+class Done:
+    total_tokens: int
+
+
+Message = KVChunk | ArraysState | Done
+
+
+def _write_exactly(stream: BinaryIO, data: bytes) -> None:
+    stream.write(data)
+    stream.flush()
+
+
+def _read_exactly(stream: BinaryIO, n: int) -> bytes:
+    buf = bytearray()
+    while len(buf) < n:
+        chunk = stream.read(n - len(buf))
+        if not chunk:
+            if len(buf) == 0:
+                return b""
+            raise ConnectionError(f"Connection closed after {len(buf)}/{n} bytes")
+        buf.extend(chunk)
+    return bytes(buf)
+
+
+def _str_to_dtype(s: str) -> torch.dtype:
+    return {
+        "float16": torch.float16,
+        "bfloat16": torch.bfloat16,
+        "float32": torch.float32,
+    }[s]
+
+
+def _dtype_size(dtype: torch.dtype) -> int:
+    return {torch.float16: 2, torch.bfloat16: 2, torch.float32: 4}[dtype]
+
+
+def write_header(stream: BinaryIO, header: dict[str, object]) -> None:
+    payload = json.dumps(header).encode("utf-8")
+    _write_exactly(stream, struct.pack(">I", len(payload)))
+    _write_exactly(stream, payload)
+
+
+def _tensor_to_bytes(t: torch.Tensor) -> bytes:
+    if t.dtype == torch.bfloat16:
+        return t.contiguous().view(torch.int16).numpy().tobytes()  # type: ignore
+    return t.contiguous().numpy().tobytes()  # type: ignore
+
+
+def write_kv_chunk(
+    stream: BinaryIO, layer_idx: int, keys: torch.Tensor, values: torch.Tensor
+) -> None:
+    if keys.dim() == 4:
+        keys = keys.reshape(-1, keys.shape[-2], keys.shape[-1])
+        values = values.reshape(-1, values.shape[-2], values.shape[-1])
+    keys_bytes = _tensor_to_bytes(keys)
+    values_bytes = _tensor_to_bytes(values)
+    num_tokens: int = keys.shape[0]
+    n_heads: int = keys.shape[1]
+    head_dim: int = keys.shape[2]
+    header = struct.pack(
+        ">BIIII", MSG_KV_CHUNK, layer_idx, num_tokens, n_heads, head_dim
+    )
+    _write_exactly(stream, header + keys_bytes + values_bytes)
+
+
+def _dtype_to_str(dtype: torch.dtype) -> str:
+    return {
+        torch.float16: "float16",
+        torch.bfloat16: "bfloat16",
+        torch.float32: "float32",
+    }[dtype]
+
+
+def write_arrays_state(
+    stream: BinaryIO, layer_idx: int, arrays: list[torch.Tensor]
+) -> None:
+    buf = io.BytesIO()
+    buf.write(struct.pack(">BI", MSG_ARRAYS_STATE, layer_idx))
+    buf.write(struct.pack(">I", len(arrays)))
+    for arr in arrays:
+        dtype_str = _dtype_to_str(arr.dtype).encode("utf-8")
+        buf.write(struct.pack(">I", len(dtype_str)))
+        buf.write(dtype_str)
+        shape: tuple[int, ...] = tuple(arr.shape)
+        buf.write(struct.pack(">I", len(shape)))
+        for dim in shape:
+            buf.write(struct.pack(">I", dim))
+        buf.write(_tensor_to_bytes(arr))
+    _write_exactly(stream, buf.getvalue())
+
+
+def write_done(stream: BinaryIO, total_tokens: int) -> None:
+    _write_exactly(stream, struct.pack(">BI", MSG_DONE, total_tokens))
+
+
+def read_header(stream: BinaryIO) -> dict[str, object]:
+    raw = _read_exactly(stream, 4)
+    if not raw:
+        raise ConnectionError("No header received")
+    length: int = struct.unpack(">I", raw)[0]  # pyright: ignore[reportAny]
+    payload = _read_exactly(stream, length)
+    return json.loads(payload.decode("utf-8"))  # pyright: ignore[reportAny]
+
+
+def read_message(stream: BinaryIO, header: dict[str, object]) -> Message | None:
+    type_byte = _read_exactly(stream, 1)
+    if not type_byte:
+        return None
+    msg_type = type_byte[0]
+
+    if msg_type == MSG_KV_CHUNK:
+        layer_idx: int
+        num_tokens: int
+        n_heads: int
+        head_dim: int
+        layer_idx, num_tokens, n_heads, head_dim = struct.unpack(
+            ">IIII", _read_exactly(stream, 16)
+        )  # pyright: ignore[reportAny]
+        dtype = _str_to_dtype(str(header["dtype"]))
+        elem_size = _dtype_size(dtype)
+        tensor_bytes: int = num_tokens * n_heads * head_dim * elem_size
+        keys_raw = _read_exactly(stream, tensor_bytes)
+        values_raw = _read_exactly(stream, tensor_bytes)
+        shape = (num_tokens, n_heads, head_dim)
+        if dtype == torch.bfloat16:
+            keys: torch.Tensor = (
+                torch.frombuffer(bytearray(keys_raw), dtype=torch.int16)
+                .view(torch.bfloat16)
+                .reshape(shape)
+                .clone()
+            )  # type: ignore
+            values: torch.Tensor = (
+                torch.frombuffer(bytearray(values_raw), dtype=torch.int16)
+                .view(torch.bfloat16)
+                .reshape(shape)
+                .clone()
+            )  # type: ignore
+        else:
+            keys = (
+                torch.frombuffer(bytearray(keys_raw), dtype=dtype)
+                .reshape(shape)
+                .clone()
+            )  # type: ignore
+            values = (
+                torch.frombuffer(bytearray(values_raw), dtype=dtype)
+                .reshape(shape)
+                .clone()
+            )  # type: ignore
+        return KVChunk(
+            layer_idx=layer_idx, num_tokens=num_tokens, keys=keys, values=values
+        )  # pyright: ignore[reportUnknownArgumentType]
+
+    if msg_type == MSG_ARRAYS_STATE:
+        arr_layer_idx: int
+        num_arrays: int
+        (arr_layer_idx,) = struct.unpack(">I", _read_exactly(stream, 4))  # pyright: ignore[reportAny]
+        (num_arrays,) = struct.unpack(">I", _read_exactly(stream, 4))  # pyright: ignore[reportAny]
+        fallback_dtype = _str_to_dtype(str(header["dtype"]))
+        arrays: list[torch.Tensor] = []
+        for _ in range(num_arrays):
+            dtype_len_raw = _read_exactly(stream, 4)
+            dtype_len: int = struct.unpack(">I", dtype_len_raw)[0]  # pyright: ignore[reportAny]
+            if dtype_len > 0 and dtype_len < 20:
+                dtype_str_bytes = _read_exactly(stream, dtype_len)
+                arr_dtype = _str_to_dtype(dtype_str_bytes.decode("utf-8"))
+            else:
+                arr_dtype = fallback_dtype
+            elem_size = _dtype_size(arr_dtype)
+            ndim: int
+            (ndim,) = struct.unpack(">I", _read_exactly(stream, 4))  # pyright: ignore[reportAny]
+            shape_arr = struct.unpack(f">{ndim}I", _read_exactly(stream, ndim * 4))
+            total_elems = 1
+            for d in shape_arr:  # pyright: ignore[reportAny]
+                total_elems *= d  # pyright: ignore[reportAny]
+            raw = _read_exactly(stream, total_elems * elem_size)
+            if arr_dtype == torch.bfloat16:
+                t: torch.Tensor = (
+                    torch.frombuffer(bytearray(raw), dtype=torch.int16)
+                    .view(torch.bfloat16)
+                    .reshape(shape_arr)
+                    .clone()
+                )  # type: ignore
+            else:
+                t = (
+                    torch.frombuffer(bytearray(raw), dtype=arr_dtype)
+                    .reshape(shape_arr)
+                    .clone()
+                )  # type: ignore
+            arrays.append(t)  # pyright: ignore[reportUnknownArgumentType]
+        return ArraysState(layer_idx=arr_layer_idx, arrays=arrays)
+
+    if msg_type == MSG_DONE:
+        total_tokens: int
+        (total_tokens,) = struct.unpack(">I", _read_exactly(stream, 4))  # pyright: ignore[reportAny]
+        return Done(total_tokens=total_tokens)
+
+    raise ValueError(f"Unknown message type: {msg_type:#x}")

src/exo/disaggregated/streaming_connector.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import queue
+import re
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (  # pyright: ignore[reportMissingImports]
+    KVConnectorBase_V1,  # pyright: ignore[reportUnknownVariableType]
+    KVConnectorMetadata,  # pyright: ignore[reportUnknownVariableType]
+    KVConnectorRole,  # pyright: ignore[reportUnknownVariableType]
+    SupportsHMA,  # pyright: ignore[reportUnknownVariableType]
+)
+
+_LAYER_RE = re.compile(r"layers\.(\d+)\.")
+
+
+def _to_bf16(t: torch.Tensor) -> torch.Tensor:
+    if t.dtype == torch.uint8:
+        t = t.view(torch.float8_e4m3fn)  # type: ignore
+    if t.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):  # type: ignore
+        return t.to(torch.float32).to(torch.bfloat16)
+    if t.dtype in (torch.bfloat16, torch.float16, torch.float32):
+        return t
+    return t.to(torch.bfloat16)
+
+
+_shared_queue: queue.Queue[tuple[int, torch.Tensor, torch.Tensor] | None] = (
+    queue.Queue()
+)
+_shared_arrays_queue: queue.Queue[tuple[int, list[torch.Tensor]] | None] = queue.Queue()
+
+
+def get_shared_queue() -> queue.Queue[tuple[int, torch.Tensor, torch.Tensor] | None]:
+    return _shared_queue
+
+
+def get_shared_arrays_queue() -> queue.Queue[tuple[int, list[torch.Tensor]] | None]:
+    return _shared_arrays_queue
+
+
+def reset_shared_queue() -> None:
+    while not _shared_queue.empty():
+        try:
+            _shared_queue.get_nowait()
+        except queue.Empty:
+            break
+    while not _shared_arrays_queue.empty():
+        try:
+            _shared_arrays_queue.get_nowait()
+        except queue.Empty:
+            break
+
+
+@dataclass
+class StreamingConnectorMetadata(KVConnectorMetadata):  # pyright: ignore[reportUntypedBaseClass]
+    pass
+
+
+class StreamingConnector(KVConnectorBase_V1, SupportsHMA):  # pyright: ignore[reportUntypedBaseClass]
+    _queue: queue.Queue[tuple[int, torch.Tensor, torch.Tensor] | None]
+
+    _save_count: int = 0
+
+    def __init__(
+        self, vllm_config: Any, role: KVConnectorRole, kv_cache_config: Any = None
+    ) -> None:  # type: ignore
+        super().__init__(vllm_config, role, kv_cache_config)  # pyright: ignore[reportUnknownMemberType]
+        self._queue = _shared_queue
+
+    @property
+    def layer_queue(self) -> queue.Queue[tuple[int, torch.Tensor, torch.Tensor] | None]:
+        return self._queue
+
+    def start_load_kv(self, forward_context: Any, **kwargs: Any) -> None:  # pyright: ignore[reportAny]
+        pass
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        pass
+
+    def save_kv_layer(
+        self, layer_name: str, kv_layer: Any, attn_metadata: Any, **kwargs: Any
+    ) -> None:  # pyright: ignore[reportAny]
+        slot_mapping = getattr(attn_metadata, "slot_mapping", None)  # pyright: ignore[reportAny]
+        if slot_mapping is not None and slot_mapping.shape[0] <= 100:  # pyright: ignore[reportAny]
+            return
+
+        m = _LAYER_RE.search(layer_name)
+        if m is None:
+            return
+        layer_idx = int(m.group(1))
+
+        if isinstance(kv_layer, (list, tuple)):
+            arrays = [_to_bf16(t).cpu() for t in kv_layer]  # pyright: ignore[reportAny]
+            _shared_arrays_queue.put((layer_idx, arrays))
+            return
+
+        if self._save_count < 1:
+            self._save_count += 1
+
+        if slot_mapping is not None:
+            if kv_layer.shape[0] == 2:  # pyright: ignore[reportAny]
+                k_all = kv_layer[0]  # pyright: ignore[reportAny]
+                v_all = kv_layer[1]  # pyright: ignore[reportAny]
+            else:
+                k_all = kv_layer[:, 0]  # pyright: ignore[reportAny]
+                v_all = kv_layer[:, 1]  # pyright: ignore[reportAny]
+            k_flat = k_all.reshape(-1, *k_all.shape[-2:])  # pyright: ignore[reportAny]
+            v_flat = v_all.reshape(-1, *v_all.shape[-2:])  # pyright: ignore[reportAny]
+            valid = slot_mapping >= 0  # pyright: ignore[reportAny]
+            safe_sm = slot_mapping.clamp(min=0)  # pyright: ignore[reportAny]
+            keys = k_flat[safe_sm][valid]  # pyright: ignore[reportAny]
+            values = v_flat[safe_sm][valid]  # pyright: ignore[reportAny]
+            keys = _to_bf16(keys)  # pyright: ignore[reportAny]
+            values = _to_bf16(values)  # pyright: ignore[reportAny]
+            self._queue.put((layer_idx, keys.cpu(), values.cpu()))  # pyright: ignore[reportAny]
+        else:
+            self._queue.put((layer_idx, kv_layer.cpu().clone(), kv_layer.cpu().clone()))  # pyright: ignore[reportAny]
+
+    def wait_for_save(self) -> None:
+        pass
+
+    def finish(self) -> None:
+        self._queue.put(None)
+
+    def request_finished_all_groups(
+        self, request: Any, block_ids: tuple[list[int], ...]
+    ) -> tuple[bool, dict[str, Any] | None]:  # pyright: ignore[reportAny]
+        return False, None
+
+    def get_num_new_matched_tokens(
+        self, request: Any, num_computed_tokens: int
+    ) -> tuple[int, bool]:  # pyright: ignore[reportAny]
+        return 0, False
+
+    def update_state_after_alloc(
+        self, request: Any, blocks: Any, num_external_tokens: int
+    ) -> None:  # pyright: ignore[reportAny]
+        pass
+
+    def build_connector_meta(self, scheduler_output: Any) -> StreamingConnectorMetadata:  # pyright: ignore[reportAny]
+        return StreamingConnectorMetadata()

src/exo/main.py

@@ -274,6 +274,12 @@ def main():
         os.environ["EXO_NO_BATCH"] = "1"
         logger.info("Continuous batching disabled (--no-batch)")
 
+    if args.no_overlapping_prefill_sends:
+        os.environ["EXO_NO_OVERLAPPING_PREFILL_SENDS"] = "1"
+        logger.info(
+            "Overlapping prefill sends disabled (--no-overlapping-prefill-sends)"
+        )
+
     # Set FAST_SYNCH override env var for runner subprocesses
     if args.fast_synch is True:
         os.environ["EXO_FAST_SYNCH"] = "on"
@@ -305,6 +311,7 @@ class Args(CamelCaseModel):
     no_downloads: bool = False
     offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
     no_batch: bool = False
+    no_overlapping_prefill_sends: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
 
     @classmethod
@@ -363,6 +370,11 @@ class Args(CamelCaseModel):
             action="store_true",
             help="Disable continuous batching, use sequential generation",
         )
+        parser.add_argument(
+            "--no-overlapping-prefill-sends",
+            action="store_true",
+            help="Disable overlapping KV transfer during disaggregated prefill",
+        )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(
             "--fast-synch",

src/exo/master/adapters/chat_completions.py

@@ -107,6 +107,7 @@ def chat_request_to_text_generation(
         min_p=request.min_p,
         repetition_penalty=request.repetition_penalty,
         repetition_context_size=request.repetition_context_size,
+        prefill_endpoints=request.prefill_endpoints,
     )
 
 

src/exo/master/api.py

@@ -414,6 +414,7 @@ class API:
                 node_network=self.state.node_network,
                 topology=self.state.topology,
                 current_instances=self.state.instances,
+                node_vllm=self.state.node_vllm,
             )
         except ValueError as exc:
             raise HTTPException(status_code=400, detail=str(exc)) from exc
@@ -451,26 +452,12 @@ class API:
         instance_combinations: list[tuple[Sharding, InstanceMeta, int]] = []
         node_count = len(list(self.state.topology.list_nodes()))
 
-        # QMM is not available on MLX CUDA. Also, VLLM does not support MLX community models
-        is_mlx_community = str(model_card.model_id).startswith("mlx-community/")
-        is_quantized_mlx = is_mlx_community and model_card.quantization in (
-            "4bit",
-            "8bit",
-        )
-        skip_mlx = any(self.state.node_vllm.values()) and is_quantized_mlx
-        is_vllm_compatible_mlx = is_mlx_community and model_card.quantization in (
-            "",
-            "bf16",
-            "fp16",
-        )
-        skip_vllm = is_mlx_community and not is_vllm_compatible_mlx
-        if not skip_mlx:
-            for sharding in (Sharding.Pipeline, Sharding.Tensor):
-                for instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl):
-                    instance_combinations.extend(
-                        [(sharding, instance_meta, i) for i in range(1, node_count + 1)]
-                    )
-        if any(self.state.node_vllm.values()) and not skip_vllm:
+        for sharding in (Sharding.Pipeline, Sharding.Tensor):
+            for instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl):
+                instance_combinations.extend(
+                    [(sharding, instance_meta, i) for i in range(1, node_count + 1)]
+                )
+        if any(self.state.node_vllm.values()):
             instance_combinations.append((Sharding.Pipeline, InstanceMeta.Vllm, 1))
 
         for sharding, instance_meta, min_nodes in instance_combinations:
@@ -484,6 +471,7 @@ class API:
                     ),
                     node_memory=self.state.node_memory,
                     node_network=self.state.node_network,
+                    node_vllm=self.state.node_vllm,
                     topology=self.state.topology,
                     current_instances=self.state.instances,
                     required_nodes=required_nodes,
@@ -761,20 +749,30 @@ class API:
         return await self._collect_text_generation_with_stats(command.command_id)
 
     async def _resolve_and_validate_text_model(self, model_id: ModelId) -> ModelId:
-        """Validate a text model exists and return the resolved model ID.
+        from exo.shared.models.model_cards import derive_base_model
 
-        Raises HTTPException 404 if no instance is found for the model.
-        """
-        if not any(
+        if any(
             instance.shard_assignments.model_id == model_id
             for instance in self.state.instances.values()
         ):
-            await self._trigger_notify_user_to_download_model(model_id)
-            raise HTTPException(
-                status_code=404,
-                detail=f"No instance found for model {model_id}",
+            return model_id
+
+        request_base = derive_base_model(str(model_id))
+        for instance in self.state.instances.values():
+            first_shard = next(
+                iter(instance.shard_assignments.runner_to_shard.values()), None
             )
-        return model_id
+            if (
+                first_shard is not None
+                and first_shard.model_card.base_model.lower() == request_base.lower()
+            ):
+                return instance.shard_assignments.model_id
+
+        await self._trigger_notify_user_to_download_model(model_id)
+        raise HTTPException(
+            status_code=404,
+            detail=f"No instance found for model {model_id}",
+        )
 
     async def _validate_image_model(self, model: ModelId) -> ModelId:
         """Validate model exists and return resolved model ID.

src/exo/master/main.py

@@ -59,7 +59,8 @@ from exo.shared.types.tasks import (
 from exo.shared.types.tasks import (
     TextGeneration as TextGenerationTask,
 )
-from exo.shared.types.worker.instances import InstanceId
+from exo.shared.types.worker.instances import Instance, InstanceId, VllmInstance
+from exo.shared.types.worker.runners import RunnerReady, RunnerRunning
 from exo.utils.channels import Receiver, Sender
 from exo.utils.event_buffer import MultiSourceBuffer
 from exo.utils.task_group import TaskGroup
@@ -93,8 +94,102 @@ class Master:
         self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {}
         self._expected_ranks: dict[TaskId, set[int]] = {}
 
+    def _find_prefill_endpoints(
+        self, decode_instance: Instance, decode_model_base: str
+    ) -> list[str]:
+        from exo.master.placement_utils import (
+            _find_ip_prioritised as find_ip_prioritised,  # pyright: ignore[reportPrivateUsage]
+        )
+        from exo.shared.models.model_cards import derive_base_model
+
+        endpoints: list[tuple[int, str]] = []
+        vllm_instance_count = 0
+        for instance in self.state.instances.values():
+            if not isinstance(instance, VllmInstance):
+                continue
+            if instance.instance_id == decode_instance.instance_id:
+                continue
+            vllm_instance_count += 1
+            first_shard = next(
+                iter(instance.shard_assignments.runner_to_shard.values()), None
+            )
+            if first_shard is None:
+                logger.info(
+                    f"Prefill routing: VllmInstance {instance.instance_id} has no shards"
+                )
+                continue
+            if (
+                derive_base_model(first_shard.model_card.base_model).lower()
+                != decode_model_base.lower()
+            ):
+                logger.info(
+                    f"Prefill routing: VllmInstance {instance.instance_id} base_model "
+                    f"{first_shard.model_card.base_model!r} != decode {decode_model_base!r}"
+                )
+                continue
+
+            pass
+
+            for node_id, runner_id in instance.shard_assignments.node_to_runner.items():
+                runner_status = self.state.runners.get(runner_id)
+                if not isinstance(runner_status, (RunnerReady, RunnerRunning)):
+                    logger.info(
+                        f"Prefill routing: runner {runner_id} not ready ({type(runner_status).__name__})"
+                    )
+                    continue
+                port = runner_status.prefill_server_port
+                if port is None:
+                    logger.info(
+                        f"Prefill routing: runner {runner_id} has no prefill_server_port"
+                    )
+                    continue
+
+                decode_node = next(
+                    iter(decode_instance.shard_assignments.node_to_runner.keys()), None
+                )
+                if decode_node is None:
+                    continue
+
+                ip = find_ip_prioritised(
+                    decode_node,
+                    node_id,
+                    self.state.topology,
+                    self.state.node_network,
+                    ring=True,
+                )
+                if ip is None:
+                    logger.info(
+                        f"Prefill routing: no IP route from {decode_node} to {node_id}"
+                    )
+                    continue
+
+                ip_type = "unknown"
+                node_net = self.state.node_network.get(node_id)
+                if node_net:
+                    for iface in node_net.interfaces:
+                        if iface.ip_address == ip:
+                            ip_type = iface.interface_type
+                            break
+                priority = {
+                    "thunderbolt": 0,
+                    "maybe_ethernet": 1,
+                    "ethernet": 2,
+                    "wifi": 3,
+                    "unknown": 4,
+                }.get(ip_type, 4)
+                endpoints.append((priority, f"{ip}:{port}"))
+
+        if not endpoints:
+            logger.info(
+                f"Prefill routing: no endpoints found for base_model={decode_model_base!r} "
+                f"(total VllmInstances in cluster: {vllm_instance_count})"
+            )
+
+        endpoints.sort(key=lambda x: x[0])
+        return [ep for _, ep in endpoints]
+
     async def run(self):
-        logger.info("Starting Master")
+        logger.debug("Starting Master")
 
         try:
             async with self._tg as tg:
@@ -108,14 +203,14 @@ class Master:
             self.command_receiver.close()
 
     async def shutdown(self):
-        logger.info("Stopping Master")
+        logger.debug("Stopping Master")
         self._tg.cancel_tasks()
 
     async def _command_processor(self) -> None:
         with self.command_receiver as commands:
             async for forwarder_command in commands:
                 try:
-                    logger.info(f"Executing command: {forwarder_command.command}")
+                    logger.debug(f"Executing command: {forwarder_command.command}")
 
                     generated_events: list[Event] = []
                     command = forwarder_command.command
@@ -124,19 +219,36 @@ class Master:
                         case TestCommand():
                             pass
                         case TextGeneration():
+                            from exo.shared.models.model_cards import derive_base_model
+
+                            request_base = derive_base_model(
+                                str(command.task_params.model)
+                            )
+
                             for instance in self.state.instances.values():
-                                if (
+                                exact_match = (
                                     instance.shard_assignments.model_id
                                     == command.task_params.model
-                                ):
-                                    task_count = sum(
-                                        1
-                                        for task in self.state.tasks.values()
-                                        if task.instance_id == instance.instance_id
-                                    )
-                                    instance_task_counts[instance.instance_id] = (
-                                        task_count
-                                    )
+                                )
+                                first_shard = next(
+                                    iter(
+                                        instance.shard_assignments.runner_to_shard.values()
+                                    ),
+                                    None,
+                                )
+                                base_match = (
+                                    first_shard is not None
+                                    and first_shard.model_card.base_model.lower()
+                                    == request_base.lower()
+                                )
+                                if not (exact_match or base_match):
+                                    continue
+                                task_count = sum(
+                                    1
+                                    for task in self.state.tasks.values()
+                                    if task.instance_id == instance.instance_id
+                                )
+                                instance_task_counts[instance.instance_id] = task_count
 
                             if not instance_task_counts:
                                 raise ValueError(
@@ -145,12 +257,38 @@ class Master:
 
                             available_instance_ids = sorted(
                                 instance_task_counts.keys(),
-                                key=lambda instance_id: instance_task_counts[
-                                    instance_id
-                                ],
+                                key=lambda instance_id: (
+                                    0
+                                    if not isinstance(
+                                        self.state.instances[instance_id], VllmInstance
+                                    )
+                                    else 1,
+                                    instance_task_counts[instance_id],
+                                ),
                             )
 
                             task_id = TaskId()
+                            decode_instance = self.state.instances[
+                                available_instance_ids[0]
+                            ]
+                            logger.info(
+                                f"Decode routing: model={command.task_params.model} base={request_base} "
+                                f"instance={available_instance_ids[0]} type={type(decode_instance).__name__} "
+                                f"candidates={len(instance_task_counts)}"
+                            )
+                            task_params = command.task_params
+                            if not task_params.prefill_endpoints:
+                                prefill_eps = self._find_prefill_endpoints(
+                                    decode_instance, request_base
+                                )
+                                logger.info(
+                                    f"Prefill endpoints resolved: {prefill_eps}"
+                                )
+                                if prefill_eps:
+                                    task_params = task_params.model_copy(
+                                        update={"prefill_endpoints": prefill_eps}
+                                    )
+
                             generated_events.append(
                                 TaskCreated(
                                     task_id=task_id,
@@ -159,7 +297,7 @@ class Master:
                                         command_id=command.command_id,
                                         instance_id=available_instance_ids[0],
                                         task_status=TaskStatus.Pending,
-                                        task_params=command.task_params,
+                                        task_params=task_params,
                                     ),
                                 )
                             )
@@ -294,6 +432,7 @@ class Master:
                                 self.state.instances,
                                 self.state.node_memory,
                                 self.state.node_network,
+                                self.state.node_vllm,
                             )
                             transition_events = get_transition_events(
                                 self.state.instances, placement, self.state.tasks
@@ -375,7 +514,7 @@ class Master:
             for node_id, time in self.state.last_seen.items():
                 now = datetime.now(tz=timezone.utc)
                 if now - time > timedelta(seconds=30):
-                    logger.info(f"Manually removing node {node_id} due to inactivity")
+                    logger.debug(f"Manually removing node {node_id} due to inactivity")
                     await self.event_sender.send(NodeTimedOut(node_id=node_id))
 
             await anyio.sleep(10)

src/exo/master/placement.py

@@ -67,11 +67,44 @@ def place_instance(
     current_instances: Mapping[InstanceId, Instance],
     node_memory: Mapping[NodeId, MemoryUsage],
     node_network: Mapping[NodeId, NodeNetworkInfo],
+    node_vllm: Mapping[NodeId, bool],
     required_nodes: set[NodeId] | None = None,
 ) -> dict[InstanceId, Instance]:
     cycles = topology.get_cycles()
     candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
 
+    # vLLM instances can only be placed on nodes that have vLLM available.
+    # vLLM does not support quantized mlx-community models (only bf16 or unquantized).
+    if command.instance_meta == InstanceMeta.Vllm:
+        is_mlx_community = str(command.model_card.model_id).startswith("mlx-community/")
+        if is_mlx_community and command.model_card.quantization not in ("", "bf16"):
+            raise ValueError("vLLM does not support quantized mlx-community models")
+        candidate_cycles = [
+            cycle
+            for cycle in candidate_cycles
+            if all(node_vllm.get(nid, False) for nid in cycle.node_ids)
+        ]
+
+    # QMM/quantized ops are not available on MLX CUDA — exclude CUDA nodes for quantized MLX models.
+    if command.instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl):
+        if command.model_card.quantization not in ("", "bf16"):
+            candidate_cycles = [
+                cycle
+                for cycle in candidate_cycles
+                if not any(node_vllm.get(nid, False) for nid in cycle.node_ids)
+            ]
+
+    # mlx-community models should prefer Apple Silicon nodes over CUDA nodes.
+    if command.instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl):
+        if str(command.model_card.model_id).startswith("mlx-community/"):
+            apple_silicon_cycles = [
+                cycle
+                for cycle in candidate_cycles
+                if not any(node_vllm.get(nid, False) for nid in cycle.node_ids)
+            ]
+            if apple_silicon_cycles:
+                candidate_cycles = apple_silicon_cycles
+
     # Filter to cycles containing all required nodes (subset matching)
     if required_nodes:
         candidate_cycles = [

src/exo/master/tests/test_placement.py

@@ -137,7 +137,7 @@ def test_get_instance_placements_create_instance(
     topology.add_connection(conn_b_a)
 
     # act
-    placements = place_instance(cic, topology, {}, node_memory, node_network)
+    placements = place_instance(cic, topology, {}, node_memory, node_network, {})
 
     # assert
     assert len(placements) == 1
@@ -179,7 +179,7 @@ def test_get_instance_placements_one_node_exact_fit() -> None:
             tasks=[ModelTask.TextGeneration],
         ),
     )
-    placements = place_instance(cic, topology, {}, node_memory, node_network)
+    placements = place_instance(cic, topology, {}, node_memory, node_network, {})
 
     assert len(placements) == 1
     instance_id = list(placements.keys())[0]
@@ -206,7 +206,7 @@ def test_get_instance_placements_one_node_fits_with_extra_memory() -> None:
             tasks=[ModelTask.TextGeneration],
         ),
     )
-    placements = place_instance(cic, topology, {}, node_memory, node_network)
+    placements = place_instance(cic, topology, {}, node_memory, node_network, {})
 
     assert len(placements) == 1
     instance_id = list(placements.keys())[0]
@@ -235,7 +235,7 @@ def test_get_instance_placements_one_node_not_fit() -> None:
     )
 
     with pytest.raises(ValueError, match="No cycles found with sufficient memory"):
-        place_instance(cic, topology, {}, node_memory, node_network)
+        place_instance(cic, topology, {}, node_memory, node_network, {})
 
 
 def test_get_transition_events_no_change(instance: Instance):
@@ -334,7 +334,7 @@ def test_placement_selects_leaf_nodes(
     cic = place_instance_command(model_card=model_card)
 
     # act
-    placements = place_instance(cic, topology, {}, node_memory, node_network)
+    placements = place_instance(cic, topology, {}, node_memory, node_network, {})
 
     # assert
     assert len(placements) == 1
@@ -422,7 +422,7 @@ def test_tensor_rdma_backend_connectivity_matrix(
     )
 
     # act
-    placements = place_instance(cic, topology, {}, node_memory, node_network)
+    placements = place_instance(cic, topology, {}, node_memory, node_network, {})
 
     # assert
     assert len(placements) == 1

src/exo/shared/models/model_cards.py

@@ -38,6 +38,35 @@ CARD_SEARCH_PATH = [
 
 _card_cache: dict[ModelId, "ModelCard"] = {}
 
+import re
+
+_QUANT_SUFFIXES = re.compile(
+    r"[-_ ](?:MLX|MXFP[0-9]+|NVFP[0-9]+|GPTQ|AWQ|GGUF|fp16|bf16|fp8|int[0-9]+|[0-9]+(?:\.[0-9]+)?bit|Q[0-9]+(?:_[A-Z0-9]+)?|gs[0-9]+)(?:[-_ ](?:MLX|Q[0-9]+|Int[0-9]+|[A-Z0-9]+|gs[0-9]+))*$",
+    re.IGNORECASE,
+)
+
+
+def _normalize_base_model(s: str) -> str:
+    return s.replace("-", " ").replace("_", " ").replace("  ", " ").strip()
+
+
+def derive_base_model(model_id: str) -> str:
+    short = model_id.split("/")[-1] if "/" in model_id else model_id
+    base = _QUANT_SUFFIXES.sub("", short)
+    return _normalize_base_model(base)
+
+
+def derive_family(model_id: str) -> str:
+    short = model_id.split("/")[-1] if "/" in model_id else model_id
+    short = _QUANT_SUFFIXES.sub("", short).lower().replace("_", "-")
+    parts = re.split(r"[-.]", short)
+    family_parts: list[str] = []
+    for p in parts:
+        if p.isdigit() or re.match(r"^\d+[bm]?$", p, re.IGNORECASE):
+            break
+        family_parts.append(p)
+    return "-".join(family_parts) if family_parts else short
+
 
 async def _refresh_card_cache():
     for path in CARD_SEARCH_PATH:
@@ -93,6 +122,15 @@ class ModelCard(CamelCaseModel):
     uses_cfg: bool = False
     trust_remote_code: bool = True
 
+    @model_validator(mode="after")
+    def _ensure_derived_fields(self) -> "ModelCard":
+        if not self.base_model:
+            self.base_model = derive_base_model(self.model_id)
+        else:
+            stripped = _QUANT_SUFFIXES.sub("", self.base_model)
+            self.base_model = _normalize_base_model(stripped)
+        return self
+
     @field_validator("tasks", mode="before")
     @classmethod
     def _validate_tasks(cls, v: list[str | ModelTask]) -> list[ModelTask]:
@@ -132,6 +170,9 @@ class ModelCard(CamelCaseModel):
         num_layers = config_data.layer_count
         mem_size_bytes = await fetch_safetensors_size(model_id)
 
+        base_model = derive_base_model(model_id)
+        family = (config_data.model_type or "").replace("_", "-")
+
         mc = ModelCard(
             model_id=ModelId(model_id),
             storage_size=mem_size_bytes,
@@ -141,6 +182,8 @@ class ModelCard(CamelCaseModel):
             num_key_value_heads=config_data.num_key_value_heads,
             tasks=[ModelTask.TextGeneration],
             trust_remote_code=False,
+            base_model=base_model,
+            family=family,
         )
         await mc.save_to_custom_dir()
         _card_cache[model_id] = mc
@@ -170,6 +213,7 @@ def is_custom_card(model_id: ModelId) -> bool:
 class ConfigData(BaseModel):
     model_config = {"extra": "ignore"}  # Allow unknown fields
 
+    model_type: str | None = None
     architectures: list[str] | None = None
     hidden_size: Annotated[int, Field(ge=0)] | None = None
     num_key_value_heads: PositiveInt | None = None

src/exo/shared/tests/test_apply/test_apply_runner_deleted.py

@@ -8,7 +8,7 @@ def test_apply_runner_shutdown_removes_runner():
     runner_id = RunnerId()
     state = State(runners={runner_id: RunnerIdle()})
 
-    new_state = apply_runner_status_updated(
+    new_state = appprefilly_runner_status_updated(
         RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerShutdown()), state
     )
 

src/exo/shared/types/api.py

@@ -221,6 +221,7 @@ class ChatCompletionRequest(BaseModel):
     tool_choice: str | dict[str, Any] | None = None
     parallel_tool_calls: bool | None = None
     user: str | None = None
+    prefill_endpoints: list[str] | None = None
 
 
 class BenchChatCompletionRequest(ChatCompletionRequest):

src/exo/shared/types/text_generation.py

@@ -70,3 +70,4 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     min_p: float | None = None
     repetition_penalty: float | None = None
     repetition_context_size: int | None = None
+    prefill_endpoints: list[str] | None = None

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

@@ -47,11 +47,11 @@ class RunnerWarmingUp(BaseRunnerStatus):
 
 
 class RunnerReady(BaseRunnerStatus):
-    pass
+    prefill_server_port: int | None = None
 
 
 class RunnerRunning(BaseRunnerStatus):
-    pass
+    prefill_server_port: int | None = None
 
 
 class RunnerShuttingDown(BaseRunnerStatus):

src/exo/utils/info_gatherer/system_info.py

@@ -91,31 +91,86 @@ async def _get_interface_types_from_networksetup() -> dict[str, InterfaceType]:
     return types
 
 
+def _classify_unknown_darwin_interface(name: str) -> InterfaceType:
+    if name.lower().startswith("anpi"):
+        return "thunderbolt"
+    return "unknown"
+
+
+async def _get_linux_network_interfaces() -> list[NetworkInterfaceInfo]:
+    import json as _json
+
+    try:
+        result = await run_process(["ip", "-j", "addr", "show"])
+    except (CalledProcessError, FileNotFoundError):
+        return []
+
+    data: list[dict[str, object]] = _json.loads(result.stdout)  # pyright: ignore[reportAny]
+    interfaces: list[NetworkInterfaceInfo] = []
+    for iface in data:
+        name: str = iface.get("ifname", "")  # pyright: ignore[reportAssignmentType, reportAny]
+        link_type: str = iface.get("link_type", "")  # pyright: ignore[reportAssignmentType, reportAny]
+
+        iface_type: InterfaceType
+        if link_type == "loopback":
+            continue
+        elif link_type == "ether":
+            if name.startswith(("wl", "wlan")):
+                iface_type = "wifi"
+            elif name.startswith(("docker", "br-", "veth")):
+                iface_type = "unknown"
+            elif name.startswith(("thunderbolt", "tb", "enx")):
+                iface_type = "thunderbolt"
+            else:
+                iface_type = "ethernet"
+        elif link_type in ("none", "tun"):
+            iface_type = "unknown"
+        else:
+            iface_type = "unknown"
+
+        for addr_info in iface.get("addr_info", []):  # pyright: ignore[reportAny]
+            family: str = addr_info.get("family", "")  # pyright: ignore[reportAny]
+            ip: str = addr_info.get("local", "")  # pyright: ignore[reportAny]
+            if family in ("inet", "inet6") and ip:
+                interfaces.append(
+                    NetworkInterfaceInfo(
+                        name=name, ip_address=ip, interface_type=iface_type
+                    )
+                )
+
+    return interfaces
+
+
 async def get_network_interfaces() -> list[NetworkInterfaceInfo]:
     """
-    Retrieves detailed network interface information on macOS.
-    Parses output from 'networksetup -listallhardwareports' and 'ifconfig'
+    Retrieves detailed network interface information on macOS or Linux.
+    On MacOS: parses output from 'networksetup -listallhardwareports' and 'ifconfig'
     to determine interface names, IP addresses, and types (ethernet, wifi, vpn, other).
+    Falls back to using ip -j addr show on other platforms.
     Returns a list of NetworkInterfaceInfo objects.
     """
-    interfaces_info: list[NetworkInterfaceInfo] = []
-    interface_types = await _get_interface_types_from_networksetup()
-
-    for iface, services in psutil.net_if_addrs().items():
-        for service in services:
-            match service.family:
-                case socket.AF_INET | socket.AF_INET6:
-                    interfaces_info.append(
-                        NetworkInterfaceInfo(
-                            name=iface,
-                            ip_address=service.address,
-                            interface_type=interface_types.get(iface, "unknown"),
+    if sys.platform == "darwin":
+        interfaces_info: list[NetworkInterfaceInfo] = []
+        interface_types = await _get_interface_types_from_networksetup()
+        for iface, services in psutil.net_if_addrs().items():
+            for service in services:
+                match service.family:
+                    case socket.AF_INET | socket.AF_INET6:
+                        iface_type = interface_types.get(iface, "unknown")
+                        if iface_type == "unknown":
+                            iface_type = _classify_unknown_darwin_interface(iface)
+                        interfaces_info.append(
+                            NetworkInterfaceInfo(
+                                name=iface,
+                                ip_address=service.address,
+                                interface_type=iface_type,
+                            )
                         )
-                    )
-                case _:
-                    pass
+                    case _:
+                        pass
+        return interfaces_info
 
-    return interfaces_info
+    return await _get_linux_network_interfaces()
 
 
 def _read_dmi_field(name: str) -> str | None:

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

@@ -194,10 +194,18 @@ class KVPrefixCache:
         # This ensures stream_generate always has at least one token to start with
         mlx_cache = self._get_mlx_cache(best_index)
         has_ssm = has_non_kv_caches(mlx_cache)
+        snapshots_available = self._snapshots[best_index] is not None
+
+        if is_exact and has_ssm and not snapshots_available:
+            prompt_cache = deepcopy(mlx_cache)
+            self._access_counter += 1
+            self._last_used[best_index] = self._access_counter
+            remaining = prompt_tokens[best_length:]
+            return prompt_cache, remaining, best_index
+
         target = (max_length - 1) if is_exact and not has_ssm else best_length
         restore_pos, restore_snap = self._get_snapshot(best_index, target)
 
-        # No usable snapshot — need fresh cache
         if restore_snap is None and has_ssm:
             return make_kv_cache(model), prompt_tokens, None
 

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

@@ -0,0 +1,58 @@
+"""Patch mlx_lm's GDN gated_delta_update to match vLLM's float32 precision.
+
+vLLM computes both softplus (gating) and sigmoid (beta) in float32.
+mlx_lm computes them in bfloat16. The precision difference compounds
+through the SSM recurrence over thousands of tokens.
+"""
+
+import sys
+from functools import partial
+from typing import Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+@partial(mx.compile, shapeless=True)
+def _compute_g_f32(A_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array:
+    return mx.exp(
+        -mx.exp(A_log.astype(mx.float32))
+        * nn.softplus((a + dt_bias).astype(mx.float32))
+    )
+
+
+def patch_gdn_softplus() -> None:
+    from mlx_lm.models import gated_delta
+
+    orig_update = gated_delta.gated_delta_update
+    orig_ops = gated_delta.gated_delta_ops
+    orig_kernel = gated_delta.gated_delta_kernel
+
+    def patched_gated_delta_update(
+        q: mx.array,
+        k: mx.array,
+        v: mx.array,
+        a: mx.array,
+        b: mx.array,
+        A_log: mx.array,
+        dt_bias: mx.array,
+        state: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+        use_kernel: bool = True,
+    ) -> Tuple[mx.array, mx.array]:
+        beta = mx.sigmoid(b.astype(mx.float32)).astype(b.dtype)
+        g = _compute_g_f32(A_log, a, dt_bias)
+        if state is None:
+            B, _, Hk, Dk = q.shape
+            Hv, Dv = v.shape[-2:]
+            state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
+
+        return orig_ops(q, k, v, g, beta, state, mask)
+
+    gated_delta.gated_delta_update = patched_gated_delta_update
+
+    for mod in list(sys.modules.values()):
+        if mod is None or mod is gated_delta:
+            continue
+        if getattr(mod, "gated_delta_update", None) is orig_update:
+            mod.gated_delta_update = patched_gated_delta_update

src/exo/worker/engines/mlx/generator/batch_generate.py

@@ -6,7 +6,7 @@ import mlx.core as mx
 from mlx_lm.generate import (
     BatchGenerator as MlxBatchGenerator,
 )
-from mlx_lm.models.cache import RotatingKVCache
+from mlx_lm.models.cache import KVCache, RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper
 
@@ -27,6 +27,7 @@ from exo.shared.types.worker.runner_response import GenerationResponse
 from exo.worker.engines.mlx.cache import (
     CacheSnapshot,
     KVPrefixCache,
+    cache_length,
     encode_prompt,
     make_kv_cache,
 )
@@ -149,16 +150,49 @@ class ExoBatchGenerator:
             top_k=task_params.top_k if task_params.top_k is not None else 0,
         )
 
-        _prefill_tps, _prefill_tokens, cache_snapshots = prefill(
-            self.model,
-            self.tokenizer,
-            sampler,
-            prompt_tokens[:-1],
-            cache,
-            self.group,
-            on_prefill_progress,
-            distributed_prompt_progress_callback,
-        )
+        _prefill_tps: float = 0.0
+        cache_snapshots: list[CacheSnapshot] | None = None
+        used_remote_prefill = False
+        uncached_count = len(prompt_tokens)
+        if uncached_count > 1000 and task_params.prefill_endpoints and not is_bench:
+            from exo.disaggregated.prefill_client import remote_prefill
+
+            t0 = time.perf_counter()
+            try:
+                injected_cache, total_tokens = remote_prefill(
+                    endpoint=task_params.prefill_endpoints[0],
+                    token_ids=[int(t) for t in all_prompt_tokens.tolist()],  # type: ignore
+                    model_id=str(task_params.model),
+                    mlx_model=self.model,
+                    on_prefill_progress=on_prefill_progress,
+                    existing_cache=list(cache) if prefix_hit_length > 0 else None,
+                    start_pos=cache_length(cache) if prefix_hit_length > 0 else 0,
+                )
+                cache = injected_cache
+                from exo.worker.engines.mlx.cache import snapshot_ssm_states
+
+                cache_snapshots = [snapshot_ssm_states(cache)]
+                _prefill_tps = total_tokens / max(time.perf_counter() - t0, 0.001)
+                used_remote_prefill = True
+                logger.info(
+                    f"Remote prefill: {total_tokens} tokens at {_prefill_tps:.0f} tok/s"
+                )
+            except Exception:
+                logger.opt(exception=True).warning(
+                    "Remote prefill failed, falling back to local"
+                )
+
+        if not used_remote_prefill:
+            _prefill_tps, _prefill_tokens, cache_snapshots = prefill(
+                self.model,
+                self.tokenizer,
+                sampler,
+                prompt_tokens[:-1],
+                cache,
+                self.group,
+                on_prefill_progress,
+                distributed_prompt_progress_callback,
+            )
 
         # We need to clamp rotating kv caches to max size so that mlx lm's _merge_caches behaves
         for c in cache:
@@ -182,7 +216,11 @@ class ExoBatchGenerator:
                 matched_index,
             )
 
-        last_tokens = prompt_tokens[-2:]
+        last_tokens = (
+            mx.array(all_prompt_tokens[-2:])
+            if used_remote_prefill
+            else prompt_tokens[-2:]
+        )
 
         logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = (
             make_logits_processors(

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

@@ -551,7 +551,7 @@ def apply_chat_template(
         )
         if partial_assistant_content:
             prompt += partial_assistant_content
-        logger.info(prompt)
+        logger.debug(prompt)
         return prompt
 
     extra_kwargs: dict[str, Any] = {}
@@ -588,7 +588,7 @@ def apply_chat_template(
     if partial_assistant_content:
         prompt += partial_assistant_content
 
-    logger.info(prompt)
+    logger.debug(prompt)
 
     return prompt
 

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

@@ -0,0 +1,123 @@
+"""Patch mlx_lm's YarnRoPE to match vLLM's inverse-frequency blending formula.
+
+mlx_lm's YarnRoPE uses a harmonic blend of frequencies. vLLM uses a linear blend
+of inverse frequencies. These produce different rotation angles, causing KV cache
+mismatch in disaggregated prefill. This patch replaces the frequency computation
+to match vLLM exactly, including support for the `truncate` parameter.
+"""
+
+import math
+
+import mlx.core as mx
+from mlx_lm.models import rope_utils
+
+
+_original_YarnRoPE_init = rope_utils.YarnRoPE.__init__
+
+
+def _patched_yarn_init(
+    self,  # type: ignore
+    dims,  # type: ignore
+    traditional=False,
+    max_position_embeddings=2048,
+    base=10000,
+    scaling_factor=1.0,
+    original_max_position_embeddings=4096,
+    beta_fast=32,
+    beta_slow=1,
+    mscale=1,
+    mscale_all_dim=0,
+    truncate=True,
+) -> None:
+    super(rope_utils.YarnRoPE, self).__init__()
+
+    def yarn_find_correction_dim(num_rotations: float) -> float:
+        return (
+            dims
+            * math.log(original_max_position_embeddings / (num_rotations * 2 * math.pi))
+        ) / (2 * math.log(base))
+
+    def yarn_find_correction_range() -> tuple[float, float]:
+        low: float = yarn_find_correction_dim(beta_fast)
+        high: float = yarn_find_correction_dim(beta_slow)
+        if truncate:
+            low = math.floor(low)
+            high = math.ceil(high)
+        return max(low, 0), min(high, dims - 1)
+
+    def yarn_get_mscale(scale: float = 1, ms: float = 1) -> float:
+        if scale <= 1:
+            return 1.0
+        return 0.1 * ms * math.log(scale) + 1.0
+
+    def yarn_linear_ramp_mask(min_val: float, max_val: float, dim: int) -> mx.array:
+        if min_val == max_val:
+            max_val += 0.001
+        linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (max_val - min_val)
+        return mx.clip(linear_func, 0, 1)
+
+    self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
+        scaling_factor, mscale_all_dim
+    )
+    pos_freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
+    inv_freq_extrapolation = 1.0 / pos_freqs
+    inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
+    low, high = yarn_find_correction_range()
+    inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dims // 2)
+    inv_freq = (
+        inv_freq_interpolation * (1 - inv_freq_mask)
+        + inv_freq_extrapolation * inv_freq_mask
+    )
+    self._freqs = 1.0 / inv_freq
+    self.dims = dims
+    self.traditional = traditional
+
+
+def _patched_initialize_rope(
+    dims: int,
+    base: float,
+    traditional: bool,
+    scaling_config: dict | None = None,
+    max_position_embeddings: int | None = None,
+) -> object:  # type: ignore
+    if scaling_config is not None:
+        rope_type = scaling_config.get("type") or scaling_config.get(
+            "rope_type", "default"
+        )
+    else:
+        rope_type = "default"
+
+    if rope_type in ("yarn", "deepseek_yarn", "telechat3-yarn"):
+        scaling_factor = scaling_config["factor"]  # type: ignore
+        rope_kwargs = {
+            key: scaling_config[key]  # type: ignore
+            for key in [
+                "original_max_position_embeddings",
+                "beta_fast",
+                "beta_slow",
+                "mscale",
+                "mscale_all_dim",
+                "truncate",
+            ]
+            if key in scaling_config  # type: ignore
+        }
+        return rope_utils.YarnRoPE(
+            dims=dims,
+            max_position_embeddings=max_position_embeddings,
+            traditional=traditional,
+            scaling_factor=scaling_factor,
+            base=base,
+            **rope_kwargs,
+        )
+
+    return _original_initialize_rope(
+        dims, base, traditional, scaling_config, max_position_embeddings
+    )
+
+
+_original_initialize_rope = rope_utils.initialize_rope
+
+
+def patch_yarn_rope() -> None:
+    rope_utils.YarnRoPE.__init__ = _patched_yarn_init  # type: ignore
+    rope_utils.initialize_rope = _patched_initialize_rope  # type: ignore

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

@@ -56,17 +56,25 @@ def _patch_determine_available_memory() -> None:
 
     @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)
+
+        real_empty_cache = torch.cuda.empty_cache
+        torch.cuda.empty_cache = lambda: None  # type: ignore
         try:
             original(self)
-        except AssertionError:
-            logger.warning(
-                "vLLM memory profiling assertion failed (free memory changed during init, "
-                "likely another process released GPU memory). Continuing with growable cache."
-            )
-        torch.cuda.empty_cache()
+        except (AssertionError, Exception):
+            pass
+        finally:
+            torch.cuda.empty_cache = real_empty_cache  # type: ignore
         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)"
@@ -164,12 +172,10 @@ def _try_grow_cache(kv_cache_manager: "object") -> bool:
     model_runner = kv_cache_manager._growable_model_runner  # type: ignore
 
     if model_runner is None:
-        logger.debug("No model_runner reference — cannot grow cache")
         return False
 
     free_bytes, _ = torch.cuda.mem_get_info()
     if free_bytes < GROWTH_HEADROOM_BYTES:
-        logger.debug(f"Only {free_bytes / (1024**3):.2f} GiB free — not enough to grow")
         return False
 
     kv_cache_config = model_runner._growable_kv_cache_config  # type: ignore
@@ -182,7 +188,6 @@ def _try_grow_cache(kv_cache_manager: "object") -> bool:
     growth_blocks = min(usable_bytes // per_block_bytes, old_num_blocks)
 
     if growth_blocks < MIN_GROWTH_BLOCKS:
-        logger.debug(f"Growth too small ({growth_blocks} blocks)")
         return False
 
     new_num_blocks = old_num_blocks + growth_blocks
@@ -193,11 +198,11 @@ def _try_grow_cache(kv_cache_manager: "object") -> bool:
     )
 
     try:
-        _grow_tensors(model_runner, kv_cache_config, old_num_blocks, new_num_blocks)
-        _grow_block_pool(block_pool, old_num_blocks, new_num_blocks)
         kv_cache_config.num_blocks = new_num_blocks
         for tensor_spec in kv_cache_config.kv_cache_tensors:
             tensor_spec.size = int(tensor_spec.size * new_num_blocks / old_num_blocks)
+        _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:
@@ -243,7 +248,6 @@ def _grow_tensors(
         model_runner.compilation_config.static_forward_context
     )  # type: ignore
     runner_kv_caches: list[torch.Tensor] = model_runner.kv_caches  # type: ignore
-    runner_kv_caches.clear()
 
     from collections import defaultdict
 
@@ -258,12 +262,59 @@ def _grow_tensors(
     for ln in new_kv_caches:
         index2name[extract_layer_index(ln, num_attn_module)].append(ln)
 
+    new_ordered: list[torch.Tensor | list[torch.Tensor]] = []
     for layer_index in sorted(index2name.keys()):
         for ln in index2name[layer_index]:
-            runner_kv_caches.append(new_kv_caches[ln])
+            new_ordered.append(new_kv_caches[ln])
 
-    for layer_name, kv_cache in new_kv_caches.items():
-        forward_context[layer_name].kv_cache = [kv_cache]  # type: ignore
+    for i, new_kv in enumerate(new_ordered):
+        if i < len(runner_kv_caches):
+            old_kv = runner_kv_caches[i]
+            if isinstance(old_kv, list) and isinstance(new_kv, list):
+                for j, (old_t, new_t) in enumerate(zip(old_kv, new_kv)):
+                    old_t.set_(
+                        new_t.storage(),
+                        new_t.storage_offset(),
+                        new_t.shape,
+                        new_t.stride(),
+                    )  # type: ignore
+            elif isinstance(old_kv, torch.Tensor) and isinstance(new_kv, torch.Tensor):
+                old_kv.set_(
+                    new_kv.storage(),
+                    new_kv.storage_offset(),
+                    new_kv.shape,
+                    new_kv.stride(),
+                )  # type: ignore
+            else:
+                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  # type: ignore
+        if old_kv_list and len(old_kv_list) > 0:
+            old_entry = old_kv_list[0]
+            if isinstance(old_entry, list) and isinstance(new_kv, list):
+                for j, (old_t, new_t) in enumerate(zip(old_entry, new_kv)):
+                    old_t.set_(
+                        new_t.storage(),
+                        new_t.storage_offset(),
+                        new_t.shape,
+                        new_t.stride(),
+                    )  # type: ignore
+            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(),
+                )  # type: ignore
+            else:
+                forward_context[layer_name].kv_cache = [new_kv]  # type: ignore
+        else:
+            forward_context[layer_name].kv_cache = [new_kv]  # type: ignore
 
 
 def _grow_block_pool(

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

@@ -228,6 +228,7 @@ class TorchKVCache:
         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
@@ -245,6 +246,18 @@ class TorchKVCache:
                 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()
@@ -265,20 +278,47 @@ class TorchKVCache:
         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)
-            n_blocks = min(len(bt), layer.keys.shape[0])
+
+            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]] = layer.keys[:n_blocks].to(
-                    device, non_blocking=True
-                )
-                v_all[bt[:n_blocks]] = layer.values[:n_blocks].to(
-                    device, non_blocking=True
-                )
+                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]:

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

@@ -3,6 +3,7 @@ import math
 import os
 import re
 import sys
+import tempfile
 import time
 from collections.abc import Callable, Generator
 from dataclasses import dataclass, field
@@ -206,15 +207,19 @@ def vllm_generate(
     on_generation_token: Callable[[], None] | None = None,
 ) -> Generator[GenerationResponse, None, None]:
     token_ids, prompt_text, prompt_token_count = format_vllm_prompt(engine, task)
-    logger.info(prompt_text)
+    logger.debug(prompt_text)
     request_id = f"vllm-seq-{time.monotonic_ns()}"
     sampling_params = make_vllm_sampling_params(engine, task, model_id)
     engine.add_request(request_id, {"prompt_token_ids": token_ids}, sampling_params)
 
     tokenizer = engine.get_tokenizer()
     stop_ids = _stop_token_ids(tokenizer, model_id)
+    DEFAULT_PREFILL_STEP_SIZE = 8192
     max_batch_tokens: int = (
-        getattr(engine.model_config, "max_num_batched_tokens", 2048) or 2048
+        getattr(
+            engine.model_config, "max_num_batched_tokens", DEFAULT_PREFILL_STEP_SIZE
+        )
+        or DEFAULT_PREFILL_STEP_SIZE
     )  # type: ignore[reportUnknownMemberType]
     start_time = time.perf_counter()
     first_token_time: float | None = None
@@ -334,7 +339,7 @@ class VllmBatchEngine:
         token_ids, prompt_text, prompt_token_count = format_vllm_prompt(
             self.engine, task_params
         )
-        logger.info(prompt_text)
+        logger.debug(prompt_text)
         sampling_params = make_vllm_sampling_params(
             self.engine, task_params, self.model_id
         )
@@ -554,30 +559,75 @@ def load_vllm_engine(
     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, KVPrefixCache]:
     patch_vllm()
     _patch_weight_loading_progress()
 
+    if kv_connector_cls is not None:
+        from exo.disaggregated.prefill_server import _patch_vllm_for_connector
+
+        _patch_vllm_for_connector(kv_connector_cls)
+
     os.environ.setdefault("FASTSAFETENSORS_NOGDS", "1")
 
     prefix_cache = KVPrefixCache(group=None)
     set_prefix_cache(prefix_cache)
     set_n_layers(n_layers)
 
-    engine_args = EngineArgs(
-        model=model_path,
-        served_model_name=str(model_id),
-        gpu_memory_utilization=0.05,
-        trust_remote_code=trust_remote_code,
-        load_format="fastsafetensors",
-        enable_prefix_caching=False,
-        attention_backend="TRITON_ATTN",
-        enforce_eager=True,
-        disable_log_stats=True,
-    )
+    kv_transfer_config: dict[str, str] | None = None
+    if kv_connector_cls is not None:
+        kv_transfer_config = {
+            "kv_connector": f"{kv_connector_cls.__module__}:{kv_connector_cls.__name__}",
+            "kv_role": "kv_both",
+        }
 
-    set_weight_loading_callback(on_layer_loaded)
-    engine = LLMEngine.from_engine_args(engine_args)
+    import json
+    from pathlib import Path
+
+    is_nvfp4 = "nvfp4" in model_path.lower() or "nvfp4" in str(model_id).lower()
+    has_mamba = False
+    config_path = Path(model_path) / "config.json"
+    if config_path.exists():
+        with open(config_path) as f:
+            model_config = json.load(f)
+        text_config = model_config.get("text_config", model_config)
+        has_mamba = "mamba_ssm_dtype" in text_config or "linear_attention" in (
+            text_config.get("layer_types") or []
+        )
+    if is_nvfp4 and not has_mamba:
+        backends = ["FLASHINFER", "FLASH_ATTN", "TRITON_ATTN"]
+    else:
+        backends = ["FLASH_ATTN", "TRITON_ATTN"]
+
+    engine: LLMEngine | None = None
+    for backend in backends:
+        try:
+            engine_args = EngineArgs(
+                model=model_path,
+                served_model_name=str(model_id),
+                gpu_memory_utilization=0.05,
+                trust_remote_code=trust_remote_code,
+                load_format="fastsafetensors",
+                enable_prefix_caching=False,
+                attention_backend=backend,
+                compilation_config={"cudagraph_mode": "none"},
+                disable_log_stats=True,
+                max_num_batched_tokens=4096,
+                kv_transfer_config=kv_transfer_config,  # type: ignore
+                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) as e:
+            logger.warning(f"Attention backend {backend} failed: {e}, trying next")
+            continue
+
+    if engine is None:
+        raise RuntimeError(f"No attention backend worked for {model_id}")
 
     tool_parser: ToolParser | None = None
     tokenizer = engine.get_tokenizer()

src/exo/worker/main.py

@@ -1,3 +1,4 @@
+import json
 from collections import defaultdict
 from datetime import datetime, timezone
 
@@ -7,7 +8,7 @@ from loguru import logger
 
 from exo.download.download_utils import resolve_model_in_path
 from exo.shared.apply import apply
-from exo.shared.models.model_cards import ModelId
+from exo.shared.models.model_cards import ModelId, derive_base_model
 from exo.shared.types.api import ImageEditsTaskParams
 from exo.shared.types.commands import (
     ForwarderCommand,
@@ -92,6 +93,7 @@ class Worker:
                 tg.start_soon(self.plan_step)
                 tg.start_soon(self._event_applier)
                 tg.start_soon(self._poll_connection_updates)
+                tg.start_soon(self._update_prefill_endpoints)
         finally:
             # Actual shutdown code - waits for all tasks to complete before executing.
             logger.info("Stopping Worker")
@@ -130,6 +132,67 @@ class Worker:
                         event.chunk.data
                     )
 
+    _IFACE_PRIORITY = {
+        "ethernet": 0,
+        "maybe_ethernet": 1,
+        "wifi": 2,
+        "unknown": 3,
+        "thunderbolt": 4,
+    }
+
+    def _best_ip_for_node(self, node_id: NodeId) -> str | None:
+        net = self.state.node_network.get(node_id)
+        if not net or not net.interfaces:
+            return None
+        candidates = [
+            iface
+            for iface in net.interfaces
+            if iface.ip_address not in ("127.0.0.1", "::1")
+            and not iface.ip_address.startswith("fe80:")
+        ]
+        if not candidates:
+            return None
+        candidates.sort(key=lambda i: self._IFACE_PRIORITY.get(i.interface_type, 3))
+        return candidates[0].ip_address
+
+    async def _update_prefill_endpoints(self) -> None:
+        while True:
+            await anyio.sleep(5)
+            try:
+                for runner_sup in self.runners.values():
+                    instance = runner_sup.bound_instance.instance
+                    my_model_id = instance.shard_assignments.model_id
+                    my_runner_id = runner_sup.bound_instance.bound_runner_id
+
+                    endpoints: list[dict[str, object]] = []
+                    for rid, status in self.state.runners.items():
+                        if rid == my_runner_id:
+                            continue
+                        port = getattr(status, "prefill_server_port", None)
+                        if not port:
+                            continue
+                        for other_inst in self.state.instances.values():
+                            if rid not in other_inst.shard_assignments.runner_to_shard:
+                                continue
+                            other_base = derive_base_model(
+                                other_inst.shard_assignments.model_id
+                            )
+                            my_base = derive_base_model(my_model_id)
+                            if other_base != my_base:
+                                continue
+                            for node_id in other_inst.shard_assignments.node_to_runner:
+                                ip = self._best_ip_for_node(node_id)
+                                if ip:
+                                    endpoints.append({"host": ip, "port": port})
+
+                    safe_model = str(my_model_id).replace("/", "--")
+                    # TODO: Change this to be in the task with a list of optional prefill endpoints.
+                    path = f"/tmp/exo_prefill_endpoints_{safe_model}.json"
+                    with open(path, "w") as f:
+                        json.dump(endpoints, f)
+            except:
+                logger.warning("Updating prefill endpoints failed")
+
     async def plan_step(self):
         while True:
             await anyio.sleep(0.1)

src/exo/worker/runner/bootstrap.py

@@ -2,6 +2,7 @@ import ctypes
 import os
 import resource
 import sys
+import urllib.request
 from pathlib import Path
 
 import loguru
@@ -14,6 +15,9 @@ from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
 
 logger: "loguru.Logger" = loguru.logger
 
+_TIKTOKEN_BASE_URL = "https://openaipublic.blob.core.windows.net/encodings"
+_TIKTOKEN_FILES = ["o200k_base.tiktoken", "cl100k_base.tiktoken"]
+
 _CUDA_HOST_LIBS = ["libcuda.so.1", "libnvidia-ml.so.1", "libnvidia-ptxjitcompiler.so.1"]
 _CUDA_HOST_SEARCH_DIRS = [
     Path("/usr/lib/aarch64-linux-gnu"),
@@ -25,6 +29,28 @@ _CUDA_HOST_SEARCH_DIRS = [
 ]
 
 
+def _ensure_tiktoken_encodings() -> None:
+    if os.environ.get("TIKTOKEN_ENCODINGS_BASE"):
+        return
+    from exo.shared.constants import EXO_CACHE_HOME
+
+    enc_dir = EXO_CACHE_HOME / "encodings"
+    enc_dir.mkdir(parents=True, exist_ok=True)
+    for fname in _TIKTOKEN_FILES:
+        dest = enc_dir / fname
+        if dest.exists():
+            continue
+        url = f"{_TIKTOKEN_BASE_URL}/{fname}"
+        logger.info(f"Downloading {url} -> {dest}")
+        try:
+            urllib.request.urlretrieve(url, dest)
+        except Exception:
+            logger.warning(f"Failed to download {fname}, harmony encoding may fail")
+            return
+    os.environ["TIKTOKEN_ENCODINGS_BASE"] = str(enc_dir)
+    logger.info(f"Set TIKTOKEN_ENCODINGS_BASE={enc_dir}")
+
+
 def _ensure_cuda_libs() -> None:
     if sys.platform != "linux":
         return
@@ -65,13 +91,22 @@ def entrypoint(
 
     logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")
 
+    from exo.worker.engines.mlx.yarn_rope_patch import patch_yarn_rope
+
+    patch_yarn_rope()
+
+    from exo.worker.engines.mlx.gdn_softplus_patch import patch_gdn_softplus
+
+    patch_gdn_softplus()
+
     # Import main after setting global logger - this lets us just import logger from this module
     try:
         if isinstance(bound_instance.instance, VllmInstance):
             os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
             os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
-            os.environ["VLLM_BATCH_INVARIANT"] = "1"
+            # os.environ["VLLM_BATCH_INVARIANT"] = "1"
             _ensure_cuda_libs()
+            _ensure_tiktoken_encodings()
             from exo.shared.constants import EXO_MODELS_DIR
             from exo.worker.runner.llm_inference.runner import Runner, VllmBuilder
 

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

@@ -137,6 +137,7 @@ class Runner:
             TaskId,
             TextGeneration,
         ] = {}
+        self.prefill_server_port: int | None = None
 
         logger.info("runner created")
         self.update_status(RunnerIdle())
@@ -237,7 +238,11 @@ class Runner:
                     on_timeout=on_model_load_timeout,
                     on_layer_loaded=on_layer_loaded,
                 )
+                builder_ref = self.generator
                 self.generator = self.generator.build()
+                self.prefill_server_port = getattr(
+                    builder_ref, "_prefill_server_port", None
+                )
 
                 self.send_task_status(task.task_id, TaskStatus.Complete)
                 self.update_status(RunnerLoaded())
@@ -257,7 +262,9 @@ class Runner:
                 )
 
                 self.send_task_status(task.task_id, TaskStatus.Complete)
-                self.update_status(RunnerReady())
+                self.update_status(
+                    RunnerReady(prefill_server_port=self.prefill_server_port)
+                )
                 logger.info("runner ready")
 
             case TextGeneration() if isinstance(self.current_status, RunnerReady):
@@ -342,7 +349,7 @@ class Runner:
             except WouldBlock:
                 pass
 
-        self.update_status(RunnerReady())
+        self.update_status(RunnerReady(prefill_server_port=self.prefill_server_port))
         logger.info("runner ready")
 
         return ExitCode.AllTasksComplete
@@ -529,12 +536,25 @@ class VllmBuilder(Builder):
     ) -> None:
         from exo.worker.engines.vllm.vllm_generator import load_vllm_engine
 
+        kv_connector_cls: type[object] | None = None
+        overlapping = not os.environ.get("EXO_NO_OVERLAPPING_PREFILL_SENDS")
+        if overlapping:
+            from exo.disaggregated.streaming_connector import StreamingConnector
+
+            kv_connector_cls = StreamingConnector
+        else:
+            from exo.disaggregated.batch_connector import BatchConnector
+
+            kv_connector_cls = BatchConnector
+
+        self._bound_runner_id = bound_instance.bound_runner_id
         self._engine, self._tool_parser, self._prefix_cache = load_vllm_engine(
             model_path=self.model_path,
             model_id=self.model_id,
             trust_remote_code=self.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,
         )
 
     def build(self) -> InferenceGenerator:
@@ -548,6 +568,49 @@ class VllmBuilder(Builder):
         tokenizer = TokenizerWrapper(self._engine.get_tokenizer())
         max_concurrent = 1 if os.environ.get("EXO_NO_BATCH") else 8
 
+        from exo.master.placement import random_ephemeral_port
+
+        prefill_port = random_ephemeral_port()
+        overlapping = not os.environ.get("EXO_NO_OVERLAPPING_PREFILL_SENDS")
+        try:
+            from exo.disaggregated.prefill_server import start_prefill_server
+
+            from exo.shared.types.events import RunnerStatusUpdated
+            from exo.shared.types.worker.runners import RunnerReady, RunnerRunning
+
+            runner_id = self._bound_runner_id
+
+            def _on_prefill_status(running: bool) -> None:
+                port = prefill_port
+                if running:
+                    self.event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id,
+                            runner_status=RunnerRunning(prefill_server_port=port),
+                        )
+                    )
+                else:
+                    self.event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id,
+                            runner_status=RunnerReady(prefill_server_port=port),
+                        )
+                    )
+
+            self._prefill_server = start_prefill_server(
+                engine=self._engine,
+                bind_address="0.0.0.0",
+                port=prefill_port,
+                overlapping=overlapping,
+                prefix_cache=self._prefix_cache,
+                on_status_change=_on_prefill_status,
+            )
+            self._prefill_server_port = prefill_port
+        except Exception:
+            logger.opt(exception=True).warning("Failed to start prefill server")
+            self._prefill_server = None
+            self._prefill_server_port = None
+
         logger.info(f"using BatchGenerator (vLLM, max_concurrent={max_concurrent})")
         return BatchGenerator(
             tokenizer=tokenizer,
@@ -564,4 +627,6 @@ class VllmBuilder(Builder):
 
     def close(self) -> None:
         with contextlib.suppress(NameError, AttributeError):
+            if hasattr(self, "_prefill_server") and self._prefill_server is not None:
+                self._prefill_server.shutdown()
             del self._engine, self._prefix_cache, self._tool_parser

src/exo/worker/runner/runner_supervisor.py

@@ -81,7 +81,18 @@ class RunnerSupervisor:
         task_sender, task_recv = mp_channel[Task]()
         cancel_sender, cancel_recv = mp_channel[TaskId]()
 
-        runner_process = mp.Process(
+        from exo.shared.types.worker.instances import VllmInstance
+
+        # vLLM runners use "spawn" to avoid inheriting the parent's CUDA state.
+        # With "fork", the parent's partial CUDA init (from device detection) is
+        # inherited by the child, which conflicts with torch.compile's inductor
+        # backend (cudagraph_mode=none) and causes CUDA illegal instruction errors.
+        ctx = (
+            mp.get_context("spawn")
+            if isinstance(bound_instance.instance, VllmInstance)
+            else mp
+        )
+        runner_process = ctx.Process(
             target=entrypoint,
             args=(
                 bound_instance,

test_growable_compile.py

@@ -0,0 +1,158 @@
+"""Test hybrid prefix cache: _extract_vllm_cache for attn + captured SSM for mamba."""
+
+import os, time
+
+os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
+from exo.worker.engines.vllm.growable_cache import (
+    patch_vllm,
+    set_prefix_cache,
+    get_model_runner,
+)
+
+patch_vllm()
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+
+KVConnectorFactory.register_connector(
+    "StreamingConnector", "exo.disaggregated.streaming_connector", "StreamingConnector"
+)
+from vllm.engine.arg_utils import EngineArgs
+from vllm.sampling_params import SamplingParams
+from vllm.v1.engine.llm_engine import LLMEngine
+
+MODEL = os.path.expanduser("~/.local/share/exo/models/Sehyo--Qwen3.5-35B-A3B-NVFP4")
+GEN = 600
+ea = EngineArgs(
+    model=MODEL,
+    served_model_name="test",
+    gpu_memory_utilization=0.05,
+    trust_remote_code=False,
+    load_format="fastsafetensors",
+    enable_prefix_caching=True,
+    attention_backend="FLASH_ATTN",
+    compilation_config={"cudagraph_mode": "none"},
+    disable_log_stats=True,
+    max_num_batched_tokens=4096,
+    kv_transfer_config={"kv_connector": "StreamingConnector", "kv_role": "kv_both"},
+    disable_hybrid_kv_cache_manager=False,
+)
+engine = LLMEngine.from_engine_args(ea)
+tok = engine.get_tokenizer()
+from exo.worker.engines.mlx.cache import KVPrefixCache
+
+pc = KVPrefixCache(group=None)
+set_prefix_cache(pc)
+
+from exo.disaggregated.prefill_server import (
+    _patch_gdn_capture,
+    _init_gdn_layer_order,
+    _gdn_states,
+    _gdn_call_idx,
+    _ssm_call_idx,
+    _extract_vllm_cache,
+)
+from exo.disaggregated.streaming_connector import reset_shared_queue
+
+_patch_gdn_capture()
+_init_gdn_layer_order()
+print(f"Engine loaded")
+
+article = (
+    "The European Union announced sweeping new regulations on artificial intelligence. "
+    * 500
+)
+tids = tok.encode(article)[:22000]
+msgs = [
+    {
+        "role": "user",
+        "content": tok.decode(tids) + "\nSummarize the key points of this article.",
+    }
+]
+tids = tok.encode(
+    tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+)
+ptids = tids[:-2]
+print(f"Prompt: {len(ptids)} tokens")
+
+reset_shared_queue()
+_gdn_states.clear()
+_gdn_call_idx[0] = 0
+_ssm_call_idx[0] = 0
+
+engine.add_request(
+    "r1", {"prompt_token_ids": ptids}, SamplingParams(max_tokens=2, temperature=0.7)
+)
+done = False
+tc = None
+while engine.has_unfinished_requests() and not done:
+    for out in engine.step():
+        if out.outputs and out.outputs[0].token_ids:
+            tc = _extract_vllm_cache(engine, "r1", len(ptids))
+            engine.abort_request(["r1"])
+            done = True
+            break
+print(f"Extracted: {tc.num_layers if tc else 'NONE'} layers")
+
+if tc and _gdn_states:
+    from exo.worker.engines.vllm.kv_cache import ArraysLayerState, KVLayerState
+
+    replaced = 0
+    for layer_idx in sorted(_gdn_states.keys()):
+        state = _gdn_states[layer_idx]
+        arrays = []
+        if "conv" in state:
+            arrays.append(state["conv"])
+        if "ssm" in state:
+            arrays.append(state["ssm"])
+        if arrays and layer_idx < len(tc.layers):
+            tc.layers[layer_idx] = ArraysLayerState(arrays=arrays)
+            replaced += 1
+    print(f"Replaced {replaced} GDN layers with clean prefill state")
+    kv_c = sum(
+        1 for l in tc.layers if isinstance(l, KVLayerState) and l.keys.numel() > 0
+    )
+    arr_c = sum(1 for l in tc.layers if isinstance(l, ArraysLayerState))
+    print(f"Final cache: {kv_c} KV layers, {arr_c} Arrays layers")
+
+import mlx.core as mx
+
+pc.add_kv_cache(mx.array(ptids), tc, None)
+print("Stored hybrid cache")
+
+engine.add_request(
+    "r2", {"prompt_token_ids": ptids}, SamplingParams(max_tokens=GEN, temperature=0.7)
+)
+t2 = time.perf_counter()
+prev = 0
+text2 = ""
+done2 = False
+while engine.has_unfinished_requests() and not done2:
+    for out in engine.step():
+        if out.outputs:
+            prev = len(out.outputs[0].token_ids)
+            if out.outputs[0].text:
+                text2 = out.outputs[0].text
+        if out.finished:
+            done2 = True
+            break
+e2 = time.perf_counter() - t2
+print(f"\nRequest 2: {prev} tokens in {e2:.1f}s ({prev / max(e2, 0.01):.1f} tok/s)")
+print(f"Output: {text2[:500]}")
+
+keywords = [
+    "regulation",
+    "AI",
+    "high-risk",
+    "compliance",
+    "transparency",
+    "ban",
+    "EU",
+    "framework",
+]
+hits = sum(1 for kw in keywords if kw.lower() in text2.lower())
+print(f"\nKeyword hits: {hits}/{len(keywords)}")
+if hits >= 2:
+    print("PASS")
+else:
+    print(f"FAIL ({hits} hits)")
+    exit(1)

tmp/bench_overnight.log

@@ -0,0 +1,82 @@
+=== Starting overnight bench runs at Tue Mar 10 22:27:20 GMT 2026 ===
+--- [1/8] Qwen3.5-27B-GPTQ-Int4 ---
+2026-03-10 22:27:22.370 | INFO     | __main__:main:300 - pp/tg mode: combinations (product) - 2 pairs
+You are using a model of type qwen3_5 to instantiate a model of type . This is not supported for all configurations of models and can yield errors.
+2026-03-10 22:27:23.659 | DEBUG    | __main__:main:317 - [exo-bench] loaded tokenizer: mlx-community/Qwen3.5-27B-GPTQ-Int4 for prompt sizer
+2026-03-10 22:27:23.661 | DEBUG    | __main__:main:339 - exo-bench model: short_id=Qwen3.5-27B-GPTQ-Int4 full_id=mlx-community/Qwen3.5-27B-GPTQ-Int4
+2026-03-10 22:27:23.661 | INFO     | __main__:main:340 - placements: 1
+2026-03-10 22:27:23.661 | INFO     | __main__:main:342 -   - Pipeline / MlxRing / nodes=1
+2026-03-10 22:27:23.661 | INFO     | __main__:main:353 - Planning phase: checking downloads...
+2026-03-10 22:27:23.670 | INFO     | harness:run_planning_phase:415 - Started download on 12D3KooWGXXhpS3kzjfDVuBGX8AeARLjVdAFaDouYJtXDVXkyq7f
+2026-03-10 22:27:23.674 | INFO     | __main__:main:365 - Download: model already cached
+2026-03-10 22:27:23.675 | INFO     | __main__:main:377 - ================================================================================
+2026-03-10 22:27:23.675 | INFO     | __main__:main:378 - PLACEMENT: Pipeline / MlxRing / nodes=1 / instance_id=725d8b5f-cc83-4dd9-8a4e-c6e2bc5b0607
+2026-03-10 22:27:31.871 | INFO     | __main__:main:409 - --- pp=700 tg=32067 concurrency=1 ---
+2026-03-10 22:27:34.896 | INFO     | __main__:build:224 - tok=700
+2026-03-10 23:00:45.925 | INFO     | __main__:main:519 - prompt_tps=11.76 gen_tps=16.13    prompt_tokens=700 gen_tokens=32067    peak_memory=31.81GB
+
+2026-03-10 23:00:47.935 | INFO     | __main__:main:409 - --- pp=700 tg=33085 concurrency=1 ---
+2026-03-10 23:00:50.948 | INFO     | __main__:build:224 - tok=700
+2026-03-10 23:35:06.460 | INFO     | __main__:main:519 - prompt_tps=11.87 gen_tps=16.12    prompt_tokens=700 gen_tokens=33085    peak_memory=31.89GB
+
+2026-03-10 23:35:08.978 | DEBUG    | __main__:main:532 - Deleted instance 725d8b5f-cc83-4dd9-8a4e-c6e2bc5b0607
+2026-03-10 23:35:13.985 | DEBUG    | __main__:main:541 - 
+Wrote results JSON: bench/results.json
+--- [2/8] NVIDIA-Nemotron-3-Nano-30B-A3B (1120,1330,23100) ---
+2026-03-10 23:35:17.156 | INFO     | __main__:main:300 - pp/tg mode: combinations (product) - 3 pairs
+2026-03-10 23:35:18.698 | DEBUG    | __main__:main:317 - [exo-bench] loaded tokenizer: mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-BF16 for prompt sizer
+2026-03-10 23:35:18.701 | DEBUG    | __main__:main:339 - exo-bench model: short_id=NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-BF16 full_id=mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-BF16
+2026-03-10 23:35:18.701 | INFO     | __main__:main:340 - placements: 1
+2026-03-10 23:35:18.701 | INFO     | __main__:main:342 -   - Pipeline / MlxRing / nodes=1
+2026-03-10 23:35:18.701 | INFO     | __main__:main:353 - Planning phase: checking downloads...
+2026-03-10 23:35:18.710 | INFO     | harness:run_planning_phase:415 - Started download on 12D3KooWGXXhpS3kzjfDVuBGX8AeARLjVdAFaDouYJtXDVXkyq7f
+2026-03-10 23:35:18.716 | INFO     | __main__:main:365 - Download: model already cached
+2026-03-10 23:35:18.716 | INFO     | __main__:main:377 - ================================================================================
+2026-03-10 23:35:18.716 | INFO     | __main__:main:378 - PLACEMENT: Pipeline / MlxRing / nodes=1 / instance_id=fb4d3883-2570-42e9-b3ee-aa1e4a1f8121
+2026-03-10 23:35:43.422 | INFO     | __main__:main:409 - --- pp=700 tg=1120 concurrency=1 ---
+2026-03-10 23:35:46.449 | INFO     | __main__:build:224 - tok=700
+2026-03-10 23:36:05.631 | INFO     | __main__:main:519 - prompt_tps=42.56 gen_tps=62.91    prompt_tokens=700 gen_tokens=1120    peak_memory=64.47GB
+
+2026-03-10 23:36:07.636 | INFO     | __main__:main:409 - --- pp=700 tg=1330 concurrency=1 ---
+2026-03-10 23:36:10.654 | INFO     | __main__:build:224 - tok=700
+2026-03-10 23:36:33.088 | INFO     | __main__:main:519 - prompt_tps=42.28 gen_tps=62.93    prompt_tokens=700 gen_tokens=1330    peak_memory=64.47GB
+
+2026-03-10 23:36:35.097 | INFO     | __main__:main:409 - --- pp=700 tg=23100 concurrency=1 ---
+2026-03-10 23:36:38.107 | INFO     | __main__:build:224 - tok=700
+2026-03-10 23:43:00.943 | INFO     | __main__:main:519 - prompt_tps=42.50 gen_tps=60.57    prompt_tokens=700 gen_tokens=23100    peak_memory=64.47GB
+
+2026-03-10 23:43:03.952 | DEBUG    | __main__:main:532 - Deleted instance fb4d3883-2570-42e9-b3ee-aa1e4a1f8121
+2026-03-10 23:43:08.954 | DEBUG    | __main__:main:541 - 
+Wrote results JSON: bench/results.json
+--- [3/8] Qwen3.5-35B-A3B-bf16 ---
+2026-03-11 10:33:18.576 | INFO     | __main__:main:300 - pp/tg mode: combinations (product) - 4 pairs
+2026-03-11 10:33:18.581 | INFO     | harness:resolve_model_short_id:187 - Model not in /models, adding from HuggingFace: mlx-community/Qwen3.5-35B-A3B-bf16
+You are using a model of type qwen3_5_moe to instantiate a model of type . This is not supported for all configurations of models and can yield errors.
+2026-03-11 10:33:22.909 | DEBUG    | __main__:main:317 - [exo-bench] loaded tokenizer: mlx-community/Qwen3.5-35B-A3B-bf16 for prompt sizer
+2026-03-11 10:33:22.913 | DEBUG    | __main__:main:339 - exo-bench model: short_id=Qwen3.5-35B-A3B-bf16 full_id=mlx-community/Qwen3.5-35B-A3B-bf16
+2026-03-11 10:33:22.913 | INFO     | __main__:main:340 - placements: 1
+2026-03-11 10:33:22.913 | INFO     | __main__:main:342 -   - Pipeline / MlxRing / nodes=1
+2026-03-11 10:33:22.913 | INFO     | __main__:main:353 - Planning phase: checking downloads...
+2026-03-11 10:33:22.923 | INFO     | harness:run_planning_phase:415 - Started download on 12D3KooWGXXhpS3kzjfDVuBGX8AeARLjVdAFaDouYJtXDVXkyq7f
+2026-03-11 10:44:01.613 | INFO     | __main__:main:363 - Download: 638.7s (freshly downloaded)
+2026-03-11 10:44:01.613 | INFO     | __main__:main:377 - ================================================================================
+2026-03-11 10:44:01.613 | INFO     | __main__:main:378 - PLACEMENT: Pipeline / MlxRing / nodes=1 / instance_id=3a5aa42f-a0e1-4273-b61e-56b4e34e0c1d
+2026-03-11 10:44:27.758 | INFO     | __main__:main:409 - --- pp=700 tg=6200 concurrency=1 ---
+2026-03-11 10:44:30.785 | INFO     | __main__:build:224 - tok=700
+2026-03-11 10:46:16.814 | INFO     | __main__:main:519 - prompt_tps=39.17 gen_tps=59.35    prompt_tokens=700 gen_tokens=6200    peak_memory=70.10GB
+
+2026-03-11 10:46:18.816 | INFO     | __main__:main:409 - --- pp=700 tg=6450 concurrency=1 ---
+2026-03-11 10:46:21.834 | INFO     | __main__:build:224 - tok=700
+2026-03-11 10:48:11.923 | INFO     | __main__:main:519 - prompt_tps=38.95 gen_tps=59.50    prompt_tokens=700 gen_tokens=6450    peak_memory=70.10GB
+
+2026-03-11 10:48:13.927 | INFO     | __main__:main:409 - --- pp=700 tg=25600 concurrency=1 ---
+2026-03-11 10:48:16.940 | INFO     | __main__:build:224 - tok=700
+2026-03-11 10:55:53.583 | INFO     | __main__:main:519 - prompt_tps=39.36 gen_tps=56.27    prompt_tokens=700 gen_tokens=25600    peak_memory=70.34GB
+
+2026-03-11 10:55:55.585 | INFO     | __main__:main:409 - --- pp=700 tg=38000 concurrency=1 ---
+2026-03-11 10:55:58.608 | INFO     | __main__:build:224 - tok=700
+2026-03-11 11:07:35.629 | INFO     | __main__:main:519 - prompt_tps=39.43 gen_tps=54.66    prompt_tokens=700 gen_tokens=38000    peak_memory=70.61GB
+
+2026-03-11 11:07:38.598 | DEBUG    | __main__:main:532 - Deleted instance 3a5aa42f-a0e1-4273-b61e-56b4e34e0c1d
+2026-03-11 11:07:43.607 | DEBUG    | __main__:main:541 -
+Wrote results JSON: bench/results.json