← back to Exo
Reduce memory consumption by adding Flash Attention to Qwen3.5 and Gemma 4, and fix RotatingKVCache prefix cache memory leak (#1886)
3f0df404a5e484b2ed062aa21303fa4b0273e964 · 2026-04-13 18:32:17 +0100 · rltakashige
## Motivation
Part 1 of many memory improvements.
## Changes
As written in the title
## Test Plan
### Manual Testing
Gemma 4 26B cache reduced from 54GB -> 10GB per 100k tokens, Qwen3.5 35B
A3B cache reduced from 21GB every 100000 tokens to 7GB.
Files touched
M nix/mlx.nixM src/exo/worker/engines/mlx/cache.pyM uv.lock
Diff
commit 3f0df404a5e484b2ed062aa21303fa4b0273e964
Author: rltakashige <rl.takashige@gmail.com>
Date: Mon Apr 13 18:32:17 2026 +0100
Reduce memory consumption by adding Flash Attention to Qwen3.5 and Gemma 4, and fix RotatingKVCache prefix cache memory leak (#1886)
## Motivation
Part 1 of many memory improvements.
## Changes
As written in the title
## Test Plan
### Manual Testing
Gemma 4 26B cache reduced from 54GB -> 10GB per 100k tokens, Qwen3.5 35B
A3B cache reduced from 21GB every 100000 tokens to 7GB.
---
nix/mlx.nix | 2 +-
src/exo/worker/engines/mlx/cache.py | 36 +++++++++++++++++++++++++++++++++++-
uv.lock | 12 ++++++------
3 files changed, 42 insertions(+), 8 deletions(-)
diff --git a/nix/mlx.nix b/nix/mlx.nix
index 93c4519b..c94d9551 100644
--- a/nix/mlx.nix
+++ b/nix/mlx.nix
@@ -49,7 +49,7 @@ let
owner = "rltakashige";
repo = "mlx-jaccl-fix-small-recv";
rev = uvLockMlxRev;
- hash = "sha256-M9x9QBYxwHv2z47qGZNJ4FgJyqLSIZ/3G1fEFQ421Lo=";
+ hash = "sha256-0bHRXhw+8jkyRVScXZQsuVbLyY521fNFo39cwME/sRw=";
};
patches = [
diff --git a/src/exo/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py
index ad941f90..681abd62 100644
--- a/src/exo/worker/engines/mlx/cache.py
+++ b/src/exo/worker/engines/mlx/cache.py
@@ -3,6 +3,7 @@ from copy import deepcopy
from typing import TYPE_CHECKING
import mlx.core as mx
+import numpy as np
import psutil
from mlx_lm.models.cache import (
ArraysCache,
@@ -50,11 +51,44 @@ class CacheSnapshot:
self.token_count = token_count
+def _detached_copy(a: mx.array) -> mx.array:
+ dtype = a.dtype
+ if dtype == mx.bfloat16:
+ return mx.array(np.array(a.astype(mx.float32))).astype(mx.bfloat16)
+ return mx.array(np.array(a))
+
+
+def copy_rotating_kv_cache(cache: RotatingKVCache) -> RotatingKVCache | None:
+ """
+ Deepcopy copies the metadata associated with an mx array.
+ Specifically, it shares a shared_ptr to the underlying data and
+ the mlx graph inputs of the array. This causes a memory leak for rotating
+ kv cache. By creating an np array, no metadata is stored so the old cache
+ can be cleaned up nicely.
+ """
+ if cache.keys is None or cache.values is None:
+ return None
+ n = min(cache.max_size, cache.keys.shape[2])
+ k_slice = _detached_copy(cache.keys[..., -n:, :])
+ v_slice = _detached_copy(cache.values[..., -n:, :])
+ mx.eval(k_slice, v_slice)
+ snap = RotatingKVCache.__new__(RotatingKVCache)
+ snap.keys = k_slice
+ snap.values = v_slice
+ snap.offset = cache.offset
+ snap._idx = n
+ snap.keep = cache.keep
+ snap.max_size = cache.max_size
+ return snap
+
+
def snapshot_ssm_states(cache: KVCacheType) -> CacheSnapshot:
states: list[ArraysCache | RotatingKVCache | None] = []
for c in cache:
- if isinstance(c, (ArraysCache, RotatingKVCache)):
+ if isinstance(c, ArraysCache):
states.append(deepcopy(c))
+ elif isinstance(c, RotatingKVCache):
+ states.append(copy_rotating_kv_cache(c))
else:
states.append(None)
token_count = cache_length(cache)
diff --git a/uv.lock b/uv.lock
index 850dbe0b..9257f3a0 100644
--- a/uv.lock
+++ b/uv.lock
@@ -558,7 +558,7 @@ dependencies = [
{ name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" },
- { name = "mlx", version = "0.31.2.dev20260406+90dd61a5", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#90dd61a5f0837f9bbbab4fd3fbfedba1ca5d33e7" }, marker = "sys_platform == 'darwin'" },
+ { name = "mlx", version = "0.31.2.dev20260413+fc5d577f", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#fc5d577f4c8172c7855bb3e9f382682bad2244f2" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx-vlm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1436,7 +1436,7 @@ dependencies = [
{ name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" },
- { name = "mlx", version = "0.31.2.dev20260406+90dd61a5", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#90dd61a5f0837f9bbbab4fd3fbfedba1ca5d33e7" }, marker = "sys_platform == 'darwin'" },
+ { name = "mlx", version = "0.31.2.dev20260413+fc5d577f", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#fc5d577f4c8172c7855bb3e9f382682bad2244f2" }, marker = "sys_platform == 'darwin'" },
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1494,8 +1494,8 @@ cuda13 = [
[[package]]
name = "mlx"
-version = "0.31.2.dev20260406+90dd61a5"
-source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#90dd61a5f0837f9bbbab4fd3fbfedba1ca5d33e7" }
+version = "0.31.2.dev20260413+fc5d577f"
+source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#fc5d577f4c8172c7855bb3e9f382682bad2244f2" }
resolution-markers = [
"python_full_version >= '3.14' and sys_platform == 'darwin'",
"python_full_version < '3.14' and sys_platform == 'darwin'",
@@ -1531,7 +1531,7 @@ version = "0.31.2"
source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak#b3540361c2fac915dc0f61ae0ce0de1583bfaa90" }
dependencies = [
{ name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
- { name = "mlx", version = "0.31.2.dev20260406+90dd61a5", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#90dd61a5f0837f9bbbab4fd3fbfedba1ca5d33e7" }, marker = "sys_platform == 'darwin'" },
+ { name = "mlx", version = "0.31.2.dev20260413+fc5d577f", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#fc5d577f4c8172c7855bb3e9f382682bad2244f2" }, marker = "sys_platform == 'darwin'" },
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1548,7 +1548,7 @@ dependencies = [
{ name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "miniaudio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" },
- { name = "mlx", version = "0.31.2.dev20260406+90dd61a5", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#90dd61a5f0837f9bbbab4fd3fbfedba1ca5d33e7" }, marker = "sys_platform == 'darwin'" },
+ { name = "mlx", version = "0.31.2.dev20260413+fc5d577f", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#fc5d577f4c8172c7855bb3e9f382682bad2244f2" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
← 9b381f7b bump and simplify flake (#1866)
·
back to Exo
·
Complete responses api usage response field (#1885) 77ffe039 →