SIGN IN SIGN UP

[MLX] Grow the KV-cache pool on demand (#21525)

Summary

Makes the MLX KV-cache pool grow on demand instead of allocating the
full cap at construction. A pool now starts at
CacheConfig::initial_capacity and doubles as writes need room, up to the
layer's cap; so a short session doesn't pay for a long one's worth of
memory.

Files

- backends/mlx/runtime/MLXSequenceCache.h — Pool takes both bounds
(initial_slots, max_slots) and allocates min of
them. A write bounds its run against the cap, then grow_for doubles from
the current size and clamps to the cap,
since the last doubling can overshoot. MLXSequenceCache passes
cfg.initial_capacity alongside the per-layer cap.
- extension/llm/cache/cache.h — valid(cfg) now rejects a negative
initial_capacity; zero is allowed and means
  "allocate nothing up front."
  - backends/mlx/test/mlx_sequence_cache_test.cpp — four growth cases.

  Testing

  Four cases added to the existing GTest suite (11 total), all passing:

- a first write larger than the initial allocation grows instead of
failing
- a decode crossing the allocated boundary still reads back the full
history (growth preserves written cells)
- doubling stops once the run fits, a 16→32 doubling clamps to a cap of
20, and an initial above the cap clamps at
  construction
- a zero initial capacity grows on first write; a negative one is
rejected

  cmake --preset mlx-release -DEXECUTORCH_BUILD_TESTS=ON
  cmake --build cmake-out --target mlx_sequence_cache_test
  ctest --test-dir cmake-out -R mlx_sequence_cache --output-on-failure
K
Kiymet Akdemir committed
849f9d87de694a42414af12a0a2de830381754db
Parent: e4f6748
Committed by GitHub <noreply@github.com> on 7/31/2026, 11:10:08 PM