[MLX] Add off-graph KV-cache flat runtime (#21501)
Summary Adds the MLX runtime for the off-graph KV cache: a Pool-backed sequence cache plus the update_and_attend op handler, wiring the neutral SequenceCache/FlatPolicy bookkeeping to real MLX tensors. The graph carries a cache-agnostic UpdateAndAttendNode — no cache operand, since the cache is off-graph runtime state keyed by layer_id. At runtime the handler asks the cache for the K/V window and mask kind (AttendSpec) and calls SDPA: the cache owns KV byte movement and attention semantics, the handler owns q/scale and the SDPA call. Flat (full-history) layers only for now; ring layers are rejected at construction rather than silently mis-served. The per-query mask stays in the backend as MLX's fused "causal"/none, with no materialized tensor. Nothing assigns ExecutionState::cache yet — the registry wiring and emitter land in a follow-up, so the op is inert in this PR. Files - backends/mlx/serialization/schema.fbs — UpdateAndAttendNode appended to the OpNode union. layer_id/scale are null-defaulted so an omitted field fails loudly instead of defaulting to 0; out_dtype is optional. - backends/mlx/runtime/MLXCache.h — op-facing MLXCache interface + AttendSpec (K/V window + mask kind None/Causal/Explicit; SWA window and score/softcap deferred). - backends/mlx/runtime/MLXSequenceCache.h — Pool + MLXSequenceCache, one pool per layer sized from that layer's policy. - backends/mlx/runtime/MLXExecutor.h — ExecutionState gains a per-session MLXCache* (survives reset(), mirrors mutable_buffers). - backends/mlx/runtime/MLXInterpreter.h — exec_update_and_attend handler + dispatch case. - extension/llm/cache/cache.h — kv_dtype (required storage precision) and a valid(cfg) predicate for the config contract. - backends/mlx/test/mlx_sequence_cache_test.cpp, backends/mlx/test/CMakeLists.txt, .github/workflows/mlx.yml — GTest coverage and CI wiring. Testing C++ op-level GTest: drives MLXSequenceCache::update_and_fetch directly and compares the returned AttendSpec against the K/V history the cache should have assembled — prefill (Causal), decode (None), storage-dtype cast, run placement, and the capacity/config reject paths. Runs on Apple Silicon (MLX needs Metal). Added to the test-mlx CI job. 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
bc13a7b5337b41d02c7007e00ee4f069616339ef
Parent: e0c610a
Committed by GitHub <noreply@github.com>
on 7/31/2026, 8:29:50 PM