Use sliding window

## Motivation

(I think it) Makes Evan's massive PR easier to merge later on

## Changes

- Renamed exo_pyo3_bindings to exo_rs
- Upgraded versions of pyo3-based dependencies
- Renamed PyFromSwarm to just FromSwarm, and PyNetworkingHandle to just
NetworkingHandle

.envrc

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

Cargo.toml

@@ -1,6 +1,6 @@
 [workspace]
 resolver = "3"
-members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
+members = ["rust/networking", "rust/exo_rs", "rust/util"]
 
 [workspace.package]
 version = "0.0.1"

flake.nix

@@ -110,7 +110,7 @@
               nixpkgs-fmt.enable = true;
               ruff-format = {
                 enable = true;
-                excludes = [ "rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi" ];
+                excludes = [ "rust/exo_rs/exo_rs.pyi" ];
               };
               rustfmt = {
                 enable = true;

justfile

@@ -23,7 +23,7 @@ sync-clean:
 
 rust-rebuild:
     PYO3_PYTHON="$(uv run python -c 'import sys; print(sys.executable)')" cargo run --bin stub_gen
-    uv sync --reinstall-package exo_pyo3_bindings
+    uv sync --reinstall-package exo_rs
 
 build-dashboard:
     #!/usr/bin/env bash

pyproject.toml

@@ -15,7 +15,7 @@ dependencies = [
   "huggingface-hub>=1.8.0",
   "psutil>=7.0.0",
   "loguru>=0.7.3",
-  "exo-pyo3-bindings",               # rust bindings
+  "exo-rs",                          # rust bindings
   "anyio==4.11.0",
   "tiktoken>=0.12.0",                # required for kimi k2 tokenizer
   "hypercorn>=0.18.0",
@@ -76,15 +76,14 @@ mlx-cuda13 = [
 ###
 
 [tool.uv.workspace]
-members = ["rust/exo_pyo3_bindings", "bench", "tools"]
+members = ["rust/exo_rs", "bench", "tools"]
 
 [tool.uv.sources]
-exo-pyo3-bindings = { workspace = true }
+exo-rs = { workspace = true }
 mlx = [
   { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" },
   { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
   { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine != 'aarch64'" },
-
 ]
 mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
 mflux = { git = "https://github.com/evanev7/mflux", branch = "exo2" }
@@ -241,7 +240,7 @@ torchaudio = ["torch"]
 ###
 
 [tool.ruff]
-extend-exclude = [".typings/**", "rust/exo_pyo3_bindings/**", "bench/vendor/**"]
+extend-exclude = [".typings/**", "rust/exo_rs/**", "bench/vendor/**"]
 
 [tool.ruff.lint]
 extend-select = ["I", "N", "B", "A", "PIE", "SIM"]

python/parts.nix

@@ -44,20 +44,21 @@ let
         paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (cudaLibs ++ [ cudaPackages.cuda_nvcc cuda_cccl_compat ]);
       };
       exoOverlay = final: prev: {
-        # Replace workspace exo_pyo3_bindings with Nix-built wheel.
+        # Replace workspace exo_rs 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";
+        exo-rs = pkgs.stdenv.mkDerivation {
+          pname = "exo-rs";
           version = "0.1.0";
-          src = self'.packages.exo_pyo3_bindings;
+          src = self'.packages.exo-rs;
           # Install from pre-built wheel
           nativeBuildInputs = [ final.pyprojectWheelHook ];
           dontStrip = true;
-          passthru = prev.exo-pyo3-bindings.passthru or { };
+          passthru = prev.exo-rs.passthru or { };
+
           postInstall = ''
-            local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings
-            cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/
+            local siteDir=$out/${final.python.sitePackages}/exo_rs
+            cp ${inputs.self}/rust/exo_rs/exo_rs.pyi $siteDir/
             touch $siteDir/py.typed
           '';
         };
@@ -223,7 +224,7 @@ let
           ++ lib.optionals isDarwin [ pkgs.macmon ];
           passthru = {
             venv = venv name;
-            evenv = ((pythonSet.overrideScope editableOverlay).mkVirtualEnv "${name}-evenv" (members // { exo = (members.exo or [ ]) ++ [ "dev" ]; })).overrideAttrs (_: {
+            evenv = ((pythonSet.overrideScope editableOverlay).mkVirtualEnv "${name}-evenv" (members // { exo = (members.exo or [ ]) ++ [ "dev" ]; exo-rs = [ ]; })).overrideAttrs (_: {
               venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ];
             });
           } // lib.optionalAttrs cudaSupport {

rust/exo_rs/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "exo_pyo3_bindings"
+name = "exo_rs"
 version = { workspace = true }
 edition = { workspace = true }
 publish = false
@@ -7,7 +7,7 @@ publish = false
 [lib]
 doctest = false
 path = "src/lib.rs"
-name = "exo_pyo3_bindings"
+name = "exo_rs"
 
 # "cdylib" needed to produce shared library for Python to import
 # "rlib" needed for stub-gen to run
@@ -25,7 +25,7 @@ workspace = true
 networking = { workspace = true }
 
 # interop
-pyo3 = { version = "0.27.2", features = [
+pyo3 = { version = "0.28.3", features = [
   # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
   # "nightly", # enables better-supported GIL integration
   "experimental-async", # async support in #[pyfunction] & #[pymethods]
@@ -38,13 +38,13 @@ pyo3 = { version = "0.27.2", features = [
   # "ordered-float", "rust_decimal", "smallvec",
   # "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time",  "serde",
 ] }
-pyo3-stub-gen = { version = "0.17.2" }
-pyo3-async-runtimes = { version = "0.27.0", features = [
+pyo3-stub-gen = { version = "0.22.3" }
+pyo3-async-runtimes = { version = "0.28.0", features = [
   "attributes",
   "tokio-runtime",
   "testing",
 ] }
-pyo3-log = "0.13.2"
+pyo3-log = "0.13.3"
 
 pidfile-rs = { git = "https://github.com/AndreiCravtov/pidfile-rs" }
 

rust/exo_rs/exo_rs.pyi

@@ -1,10 +1,20 @@
 # This file is automatically generated by pyo3_stub_gen
-# ruff: noqa: E501, F401
+# ruff: noqa: E501, F401, F403, F405
 
 import builtins
 import os
 import pathlib
 import typing
+__all__ = [
+    "AllQueuesFullError",
+    "FromSwarm",
+    "Keypair",
+    "MessageTooLargeError",
+    "NetworkingHandle",
+    "NoPeersSubscribedToTopicError",
+    "Pidfile",
+    "PidfileError",
+]
 
 @typing.final
 class AllQueuesFullError(builtins.Exception):
@@ -12,6 +22,29 @@ class AllQueuesFullError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
+class FromSwarm:
+    @typing.final
+    class Connection(FromSwarm):
+        __match_args__ = ("peer_id", "connected",)
+        @property
+        def peer_id(self) -> builtins.str: ...
+        @property
+        def connected(self) -> builtins.bool: ...
+        def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> FromSwarm.Connection: ...
+    
+    @typing.final
+    class Message(FromSwarm):
+        __match_args__ = ("origin", "topic", "data",)
+        @property
+        def origin(self) -> builtins.str: ...
+        @property
+        def topic(self) -> builtins.str: ...
+        @property
+        def data(self) -> bytes: ...
+        def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> FromSwarm.Message: ...
+    
+    ...
+
 @typing.final
 class Keypair:
     r"""
@@ -45,6 +78,7 @@ class MessageTooLargeError(builtins.Exception):
 @typing.final
 class NetworkingHandle:
     def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
+    def recv(self) -> typing.Awaitable[FromSwarm]: ...
     async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Subscribe to a `GossipSub` topic.
@@ -63,7 +97,6 @@ class NetworkingHandle:
         
         If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
         """
-    async def recv(self) -> PyFromSwarm: ...
 
 @typing.final
 class NoPeersSubscribedToTopicError(builtins.Exception):
@@ -77,7 +110,7 @@ class Pidfile:
     A PID file protected with a lock.
     
     An instance of `Pidfile` can be used to manage a PID file: create it,
-    lock it, detect already running daemons. It is backed by [`pidfile`][]
+    lock it, detect already running daemons. It is backed by [`pidfile`]
     functions of `libbsd`/`libutil` which use `flopen` to lock the PID
     file.
     
@@ -127,26 +160,3 @@ class PidfileError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
-class PyFromSwarm:
-    @typing.final
-    class Connection(PyFromSwarm):
-        __match_args__ = ("peer_id", "connected",)
-        @property
-        def peer_id(self) -> builtins.str: ...
-        @property
-        def connected(self) -> builtins.bool: ...
-        def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> PyFromSwarm.Connection: ...
-    
-    @typing.final
-    class Message(PyFromSwarm):
-        __match_args__ = ("origin", "topic", "data",)
-        @property
-        def origin(self) -> builtins.str: ...
-        @property
-        def topic(self) -> builtins.str: ...
-        @property
-        def data(self) -> bytes: ...
-        def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
-    
-    ...
-

rust/exo_rs/pyproject.toml

@@ -3,8 +3,8 @@ requires = ["maturin>=1.0,<2.0"]
 build-backend = "maturin"
 
 [project]
-name = "exo_pyo3_bindings"
-version = "0.2.10"
+name = "exo_rs"
+version = "0.2.16"
 description = "Add your description here"
 readme = "README.md"
 authors = [
@@ -15,14 +15,17 @@ requires-python = ">=3.13"
 dependencies = []
 
 [dependency-groups]
-dev = ["exo_pyo3_bindings", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
+dev = ["exo_rs", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
 
 [tool.maturin]
 #purelib = true
 #python-source = "python"
-module-name = "exo_pyo3_bindings"
+module-name = "exo_rs"
 features = ["pyo3/extension-module", "pyo3/experimental-async"]
 
+[tool.pyo3-stub-gen]
+generate-init-py = true
+
 [tool.pytest.ini_options]
 log_cli = true
 log_cli_level = "INFO"

rust/exo_rs/src/bin/stub_gen.rs

@@ -2,7 +2,7 @@ use pyo3_stub_gen::Result;
 
 fn main() -> Result<()> {
     env_logger::Builder::from_env(env_logger::Env::default().filter_or("RUST_LOG", "info")).init();
-    let stub = exo_pyo3_bindings::stub_info()?;
+    let stub = exo_rs::stub_info()?;
     stub.generate()?;
     Ok(())
 }

rust/exo_rs/src/lib.rs

@@ -153,7 +153,7 @@ pub(crate) mod ext {
 /// A Python module implemented in Rust. The name of this function must match
 /// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
 /// import the module.
-#[pymodule(name = "exo_pyo3_bindings")]
+#[pymodule(name = "exo_rs", gil_used = true)]
 fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // install logger
     pyo3_log::init();

rust/exo_rs/src/networking.rs

@@ -16,9 +16,7 @@ use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::{PyModule, PyModuleMethods as _};
 use pyo3::types::PyBytes;
 use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
-use pyo3_stub_gen::derive::{
-    gen_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods,
-};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods};
 use tokio::sync::{Mutex, mpsc, oneshot};
 
 mod exception {
@@ -138,7 +136,7 @@ struct PyNetworkingHandle {
 }
 
 #[gen_stub_pyclass_complex_enum]
-#[pyclass]
+#[pyclass(name = "FromSwarm")]
 enum PyFromSwarm {
     Connection {
         peer_id: String,
@@ -204,9 +202,11 @@ impl PyNetworkingHandle {
         })
     }
 
-    #[gen_stub(skip)]
+    #[gen_stub(override_return_type(
+        type_repr="typing.Awaitable[FromSwarm]", imports=("typing")
+    ))]
     fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        let swarm = Arc::clone(&self.swarm);
+        let swarm = self.swarm.clone();
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
             swarm
                 .try_lock()
@@ -297,15 +297,6 @@ impl PyNetworkingHandle {
     }
 }
 
-pyo3_stub_gen::inventory::submit! {
-    gen_methods_from_python! {
-        r#"
-            class PyNetworkingHandle:
-                async def recv() -> PyFromSwarm: ...
-        "#
-    }
-}
-
 pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<exception::PyNoPeersSubscribedToTopicError>()?;
     m.add_class::<exception::PyAllQueuesFullError>()?;

rust/exo_rs/src/pidfile.rs

@@ -38,7 +38,7 @@ impl PyPidfileError {
 /// A PID file protected with a lock.
 ///
 /// An instance of `Pidfile` can be used to manage a PID file: create it,
-/// lock it, detect already running daemons. It is backed by [`pidfile`][]
+/// lock it, detect already running daemons. It is backed by [`pidfile`]
 /// functions of `libbsd`/`libutil` which use `flopen` to lock the PID
 /// file.
 ///

rust/exo_rs/tests/test_python.py

@@ -2,12 +2,12 @@ import asyncio
 
 import pytest
 from _pytest.capture import CaptureFixture
-from exo_pyo3_bindings import (
+from exo_rs import (
     Keypair,
     NetworkingHandle,
     NoPeersSubscribedToTopicError,
     Pidfile,
-    PyFromSwarm,
+    FromSwarm,
 )
 
 
@@ -39,9 +39,9 @@ async def _await_recv(h: NetworkingHandle):
     while True:
         event = await h.recv()
         match event:
-            case PyFromSwarm.Connection() as c:
+            case FromSwarm.Connection() as c:
                 print(f"PYTHON: connection update: {c}")
-            case PyFromSwarm.Message() as m:
+            case FromSwarm.Message() as m:
                 print(f"PYTHON: message: {m}")
 
 

rust/parts.nix

@@ -55,6 +55,7 @@
         ];
 
         OPENSSL_NO_VENDOR = "1";
+        MATURIN_NO_INSTALL_RUST = "1";
 
         # Required for pyo3 tests to find libpython
         LD_LIBRARY_PATH = lib.makeLibraryPath [ pkgs.python313 ];
@@ -81,11 +82,11 @@
       config = {
         packages = {
           # Python bindings wheel via maturin
-          exo_pyo3_bindings = craneLib.buildPackage (
+          exo-rs = craneLib.buildPackage (
             commonArgs
             // {
               inherit cargoArtifacts;
-              pname = "exo_pyo3_bindings";
+              pname = "exo-rs";
 
               nativeBuildInputs = commonArgs.nativeBuildInputs ++ [
                 pkgs.maturin
@@ -95,7 +96,7 @@
                 maturin build \
                   --release \
                   --manylinux off \
-                  --manifest-path rust/exo_pyo3_bindings/Cargo.toml \
+                  --manifest-path rust/exo_rs/Cargo.toml \
                   --features "pyo3/extension-module,pyo3/experimental-async" \
                   --interpreter ${pkgs.python313}/bin/python \
                   --out dist

src/exo/main.py

@@ -10,7 +10,7 @@ from typing import Self
 import anyio
 from anyio.lowlevel import checkpoint as anyio_checkpoint
 from daemon import DaemonContext  # pyright: ignore[reportMissingTypeStubs]
-from exo_pyo3_bindings import Pidfile, PidfileError
+from exo_rs import Pidfile, PidfileError
 from loguru import logger
 from pydantic import PositiveInt
 

src/exo/routing/connection_message.py

@@ -1,4 +1,4 @@
-from exo_pyo3_bindings import PyFromSwarm
+from exo_rs import FromSwarm
 
 from exo.shared.types.common import NodeId
 from exo.utils.pydantic_ext import FrozenModel
@@ -11,5 +11,5 @@ class ConnectionMessage(FrozenModel):
     connected: bool
 
     @classmethod
-    def from_update(cls, update: PyFromSwarm.Connection) -> "ConnectionMessage":
+    def from_update(cls, update: FromSwarm.Connection) -> "ConnectionMessage":
         return cls(node_id=NodeId(update.peer_id), connected=update.connected)

src/exo/routing/router.py

@@ -12,13 +12,13 @@ from anyio import (
     move_on_after,
     sleep_forever,
 )
-from exo_pyo3_bindings import (
+from exo_rs import (
     AllQueuesFullError,
+    FromSwarm,
     Keypair,
     MessageTooLargeError,
     NetworkingHandle,
     NoPeersSubscribedToTopicError,
-    PyFromSwarm,
 )
 from filelock import FileLock
 from loguru import logger
@@ -191,7 +191,7 @@ class Router:
                 from_swarm = await self._net.recv()
                 logger.debug(from_swarm)
                 match from_swarm:
-                    case PyFromSwarm.Message(origin, topic, data):
+                    case FromSwarm.Message(origin, topic, data):
                         logger.trace(
                             f"Received message on {topic} from {origin} with payload {data}"
                         )
@@ -202,7 +202,7 @@ class Router:
                             continue
                         router = self.topic_routers[topic]
                         await router.publish_bytes(data)
-                    case PyFromSwarm.Connection():
+                    case FromSwarm.Connection():
                         message = ConnectionMessage.from_update(from_swarm)
                         logger.trace(
                             f"Received message on connection_messages with payload {message}"

src/exo/shared/logging.py

@@ -46,7 +46,7 @@ class _InterceptHandler(logging.Handler):
 def logger_setup(log_file: Path | None, verbosity: int = 0):
     """Set up logging for this process - formatting, file handles, verbosity and output"""
 
-    logging.getLogger("exo_pyo3_bindings").setLevel(logging.WARNING)
+    logging.getLogger("exo_rs").setLevel(logging.WARNING)
     logging.getLogger("httpx").setLevel(logging.WARNING)
     logging.getLogger("httpcore").setLevel(logging.WARNING)
 

src/exo/utils/tests/test_pidfile.py

@@ -8,14 +8,14 @@ import textwrap
 from pathlib import Path
 from typing import Final
 
-from exo_pyo3_bindings.exo_pyo3_bindings import Pidfile
+from exo_rs import Pidfile
 
 _CHILD_ACQUIRE_PIDFILE_SCRIPT: Final = textwrap.dedent(
     """
     import sys
     from pathlib import Path
 
-    from exo_pyo3_bindings.exo_pyo3_bindings import Pidfile, PidfileError
+    from exo_rs import Pidfile, PidfileError
 
     path = Path(sys.argv[1])
     try:

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

@@ -229,6 +229,47 @@ def has_non_kv_caches(cache: KVCacheType) -> bool:
     return any(is_non_trimmable_cache_entry(c) for c in cache)
 
 
+# Max snapshots retained per cache entry. Each CacheSnapshot pins detached GPU
+# copies of every non-trimmable (SSM/ArraysCache, RotatingKVCache) layer, so
+# retaining one per ~4096-token prefill chunk makes snapshot memory grow linearly
+# with context — the dominant residual cost when a single entry is grown to long
+# contexts on hybrid models (~56 MB/snapshot on Qwen3.5-122B, so a full 256K
+# context = 64 snapshots ≈ 3.6 GB). A sliding window of the most-recent N caps
+# this at N×per-snapshot (~0.9 GB here) while preserving the restore points
+# in-place grows actually use (they always extend from the tip).
+_MAX_RETAINED_SNAPSHOTS = 16
+
+
+def _bounded_snapshots(snapshots: list[CacheSnapshot]) -> list[CacheSnapshot]:
+    """Deduplicate snapshots by token position and bound the retained count.
+
+    Returned list is sorted ascending by ``token_count``.
+    """
+    # Deduplicate by position, keeping the most-recently-appended snapshot per
+    # position. Repeated in-place grows re-snapshot positions the kept old
+    # snapshots already cover, which would otherwise grow `_snapshots`
+    # unbounded even at constant context.
+    # TODO: keying on token_count alone is safe only while a position uniquely
+    # identifies the prefix within an entry (grows are strict prefix-extensions).
+    # If edit-and-regenerate, sliding-window/prefix trimming, cross-entry
+    # snapshot sharing, per-request adapter/LoRA swap, or branchy decoding
+    # (beam/parallel/speculative) is added, enrich the key to
+    # (token_count, prefix_hash[, media/adapter id]) — else a stale snapshot
+    # could be restored for a different prefix (silent wrong output).
+    by_position: dict[int, CacheSnapshot] = {}
+    for snapshot in snapshots:
+        by_position[snapshot.token_count] = snapshot
+    deduped = [by_position[pos] for pos in sorted(by_position)]
+
+    # Sliding window: keep only the most-recent N positions. In-place grows
+    # always extend from the tip, so the newest snapshots are the ones future
+    # grows restore from — dropping the oldest is never incorrect: a later hit on
+    # a prefix older than the window finds no snapshot <= target, so get_kv_cache
+    # returns a fresh cache (matched_index=None) and the request takes a full cold
+    # prefill — correct, just slower than a partial-hit reuse for that one request.
+    return deduped[-_MAX_RETAINED_SNAPSHOTS:]
+
+
 class KVPrefixCache:
     def __init__(self, group: mx.distributed.Group | None):
         self.prompts: list[mx.array] = []  # mx array of tokens (ints)
@@ -261,7 +302,9 @@ class KVPrefixCache:
         self._evict_if_needed()
         self.prompts.append(prompt_tokens)
         self.caches.append(deepcopy(cache))
-        self._snapshots.append(ssm_snapshots)
+        self._snapshots.append(
+            _bounded_snapshots(ssm_snapshots) if ssm_snapshots else None
+        )
         self._media_regions.append(media_regions or [])
         self.prefill_tps.append(prefill_tps)
         self._access_counter += 1
@@ -288,7 +331,7 @@ class KVPrefixCache:
 
         self.prompts[index] = prompt_tokens
         self.caches[index] = deepcopy(cache)
-        self._snapshots[index] = merged or None
+        self._snapshots[index] = _bounded_snapshots(merged) or None
         self._media_regions[index] = media_regions or []
         self.prefill_tps[index] = prefill_tps
         self._access_counter += 1

src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py

@@ -11,6 +11,7 @@ from mlx_lm.sample_utils import make_sampler
 from exo.shared.types.common import ModelId
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.worker.engines.mlx.cache import (
+    CacheSnapshot,
     KVPrefixCache,
     cache_length,
     encode_prompt,
@@ -77,6 +78,74 @@ class TestGetPrefixLength:
         assert get_prefix_length(a, b) == 0
 
 
+class TestSnapshotAccumulation:
+    """Locks in the fix for the actual per-grow Metal leak on hybrid (SSM)
+    models: `update_kv_cache` must not let `_snapshots` grow without bound when
+    the same entry is grown in place many times."""
+
+    def test_repeated_update_does_not_accumulate_snapshots(self):
+        with patch(
+            "exo.worker.engines.mlx.cache.get_memory_used_percentage",
+            return_value=0.0,
+        ):
+            kv_prefix_cache = KVPrefixCache(None)
+            initial = [
+                CacheSnapshot(states=[None], token_count=4096),
+                CacheSnapshot(states=[None], token_count=8192),
+            ]
+            kv_prefix_cache.add_kv_cache(
+                mx.arange(10000), [KVCache()], ssm_snapshots=initial
+            )
+
+            # Each in-place grow re-prefills from restore_pos and produces a
+            # fresh snapshot at a position the retained old snapshots already
+            # cover. Pre-fix this appended one snapshot per grow forever.
+            for _ in range(50):
+                fresh = [CacheSnapshot(states=[None], token_count=8192)]
+                kv_prefix_cache.update_kv_cache(
+                    0, mx.arange(10000), [KVCache()], fresh, restore_pos=8192
+                )
+
+            stored = kv_prefix_cache._snapshots[0]
+            assert stored is not None
+            # Bounded by the number of distinct snapshot positions (here 2),
+            # not by the 50 grows.
+            assert len(stored) == 2
+            assert sorted(s.token_count for s in stored) == [4096, 8192]
+            # The kept 8192 snapshot must be the most recently supplied one.
+            assert stored[1] is fresh[0]
+
+    def test_extension_caps_snapshots_to_sliding_window(self):
+        """Extending a single entry to a long context (one snapshot per ~4096
+        tokens) must cap retained snapshots to a sliding window of the most-recent
+        N, not keep all of them — that linear-in-context retention was the
+        residual OOM cause."""
+        from exo.worker.engines.mlx.cache import _MAX_RETAINED_SNAPSHOTS
+
+        with patch(
+            "exo.worker.engines.mlx.cache.get_memory_used_percentage",
+            return_value=0.0,
+        ):
+            kv_prefix_cache = KVPrefixCache(None)
+            # 64 distinct positions = a 262144-token context at 4096/chunk.
+            num_positions = 64
+            snaps = [
+                CacheSnapshot(states=[None], token_count=4096 * (i + 1))
+                for i in range(num_positions)
+            ]
+            kv_prefix_cache.add_kv_cache(
+                mx.arange(10), [KVCache()], ssm_snapshots=snaps
+            )
+
+            stored = kv_prefix_cache._snapshots[0]
+            assert stored is not None
+            # Capped at the window; the most-recent N positions are retained
+            # (in-place grows extend from the tip, so these are what get used).
+            assert len(stored) == _MAX_RETAINED_SNAPSHOTS
+            assert stored == snaps[-_MAX_RETAINED_SNAPSHOTS:]
+            assert stored[-1] is snaps[-1]  # tip always kept
+
+
 class TestKVPrefix:
     @pytest.fixture
     def mock_tokenizer(self):

uv.lock

@@ -24,7 +24,7 @@ prerelease-mode = "allow"
 members = [
     "exo",
     "exo-bench",
-    "exo-pyo3-bindings",
+    "exo-rs",
     "exo-tools",
 ]
 overrides = [{ name = "opencv-python", marker = "python_full_version < '0'" }]
@@ -423,7 +423,7 @@ dependencies = [
     { name = "aiofiles", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "aiohttp", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "anyio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
-    { name = "exo-pyo3-bindings", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
+    { name = "exo-rs", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "httpx", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
@@ -523,7 +523,7 @@ requires-dist = [
     { name = "exo", extras = ["mlx"], marker = "extra == 'mlx-cpu'" },
     { name = "exo", extras = ["mlx"], marker = "extra == 'mlx-cuda12'" },
     { name = "exo", extras = ["mlx"], marker = "extra == 'mlx-cuda13'" },
-    { name = "exo-pyo3-bindings", editable = "rust/exo_pyo3_bindings" },
+    { name = "exo-rs", editable = "rust/exo_rs" },
     { name = "fastapi", specifier = ">=0.116.1" },
     { name = "filelock", specifier = ">=3.18.0" },
     { name = "httpx", specifier = ">=0.28.1" },
@@ -622,13 +622,13 @@ requires-dist = [
 ]
 
 [[package]]
-name = "exo-pyo3-bindings"
-version = "0.2.10"
-source = { editable = "rust/exo_pyo3_bindings" }
+name = "exo-rs"
+version = "0.2.16"
+source = { editable = "rust/exo_rs" }
 
 [package.dev-dependencies]
 dev = [
-    { name = "exo-pyo3-bindings", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
+    { name = "exo-rs", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "pytest", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
     { name = "pytest-asyncio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda12') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cpu' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-cuda13') or (extra == 'extra-3-exo-mlx-cuda12' and extra == 'extra-3-exo-mlx-none') or (extra == 'extra-3-exo-mlx-cuda13' and extra == 'extra-3-exo-mlx-none')" },
 ]
@@ -637,7 +637,7 @@ dev = [
 
 [package.metadata.requires-dev]
 dev = [
-    { name = "exo-pyo3-bindings", editable = "rust/exo_pyo3_bindings" },
+    { name = "exo-rs", editable = "rust/exo_rs" },
     { name = "pytest", specifier = ">=8.4.0" },
     { name = "pytest-asyncio", specifier = ">=1.0.0" },
 ]