[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