SIGN IN SIGN UP

[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