Cache cleanup

backends/mlx/runtime/MLXBackend.cpp

@@ -25,6 +25,7 @@
 #include <executorch/runtime/core/evalue.h>
 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
 #include <executorch/runtime/core/named_data_map.h>
+#include <executorch/runtime/platform/assert.h>
 
 #include <mlx/mlx.h>
 
@@ -190,8 +191,8 @@ struct MLXHandle {
 
   // Keep-alive for the off-graph KV cache bound in init(). state.cache is a
   // non-owning view of the same object, so the cache must outlive the handle
-  // even if the runner's session is torn down first.
-  std::shared_ptr<::executorch::extension::llm::cache::CacheBase> cache_shared;
+  // even if the runner drops its InstallGuard first.
+  std::shared_ptr<::executorch::extension::llm::cache::Cache> cache_shared;
 
   // Keep the constant buffers alive for zero-copy constants
   // Each FreeableBuffer must outlive the MLX arrays that reference it
@@ -371,7 +372,8 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
       // Bind the off-graph KV cache, if the runner installed one under a key it
       // passed as a runtime spec. Bound before the init chain runs so an
       // update_and_attend node there sees the same cache execute() will.
-      if (auto spec = context.get_runtime_spec<const char*>(kCacheKeyKey);
+      if (auto spec =
+              context.get_runtime_spec<const char*>(cache::kCacheKeyOption);
           spec.ok() && spec.get() != nullptr && *spec.get() != '\0') {
         const char* cache_key = spec.get();
         handle->cache_shared =
@@ -381,12 +383,10 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
               std::string("init: cache_key '") + cache_key +
               "' is not installed in the CacheRegistry");
         }
-        // Cross-cast from the neutral ownership anchor to this backend's
-        // tensor-typed op face; the two are deliberately unrelated bases (see
-        // MLXCache.h), so nullptr here means the key names another backend's
-        // cache.
-        handle->state.cache =
-            dynamic_cast<MLXCache*>(handle->cache_shared.get());
+        // Ask the neutral ownership anchor for this backend's tensor-typed op
+        // face. It is named by MLXCache itself rather than by cache.h, so
+        // nullptr here means the key names another backend's cache.
+        handle->state.cache = handle->cache_shared->as<MLXCache>();
         if (handle->state.cache == nullptr) {
           throw std::runtime_error(
               std::string("init: cache under key '") + cache_key +
@@ -603,18 +603,30 @@ static auto success_with_compiler = register_backend(backend);
 
 // Cache kind is named by the builder tag rather than an enum on the config: a
 // runner asks the registry for (backend_id, kind) and gets back a neutral
-// CacheBase it installs under a cache_key. Adding a kind is a new builder here.
+// Cache it installs under a cache_key. Adding a kind is a new builder here.
 const int cache_builders_registered = [] {
-  cache::CacheBuilderRegistry::global().register_builder(
-      kMLXBackendId, "seq", [](const cache::CacheConfig& cfg) {
-        return std::shared_ptr<cache::CacheBase>(
+  const Error single = cache::CacheFactory::global().register_builder(
+      kMLXBackendId, cache::kind::kSingle, [](const cache::CacheConfig& cfg) {
+        return std::shared_ptr<cache::Cache>(
             std::make_shared<MLXSequenceCache>(cfg));
       });
-  cache::CacheBuilderRegistry::global().register_builder(
-      kMLXBackendId, "cell", [](const cache::CacheConfig& cfg) {
-        return std::shared_ptr<cache::CacheBase>(
+  ET_CHECK_MSG(
+      single == Error::Ok,
+      "Failed to register cache builder for %s:%s",
+      kMLXBackendId,
+      cache::kind::kSingle);
+  const Error batched_cell = cache::CacheFactory::global().register_builder(
+      kMLXBackendId,
+      cache::kind::kBatchedCell,
+      [](const cache::CacheConfig& cfg) {
+        return std::shared_ptr<cache::Cache>(
             std::make_shared<MLXCellCache>(cfg));
       });
+  ET_CHECK_MSG(
+      batched_cell == Error::Ok,
+      "Failed to register cache builder for %s:%s",
+      kMLXBackendId,
+      cache::kind::kBatchedCell);
   return 0;
 }();
 } // namespace

backends/mlx/runtime/MLXCache.h

@@ -30,12 +30,16 @@ struct AttendSpec {
 };
 
 // Tensor-typed op face of the off-graph KV cache, kept separate from the
-// neutral CacheBase (which is tensor-free) so a cache can expose both without a
+// neutral Cache (which is tensor-free) so a cache can expose both without a
 // diamond. ExecutionState holds one; nothing assigns it yet -- the registry
 // that owns the cache and hands this pointer to the executor lands in a
 // follow-up, until which exec_update_and_attend is unreachable.
 class MLXCache {
  public:
+  // Named here, not in cache.h: a backend face is tensor-typed and the
+  // neutral header cannot know about it.
+  static constexpr const char* kFaceName = "mlx.MLXCache";
+
   virtual ~MLXCache() = default;
 
   // Write this step's K/V for `layer` at `positions`, one host int per query

backends/mlx/runtime/MLXCellCache.h

@@ -84,6 +84,14 @@ class MLXCellCache : public cache::CellCache, public MLXCache {
         mask(*step)};
   }
 
+ protected:
+  void* face(cache::FaceId id) override {
+    if (void* p = cache::CellCache::face(id)) {
+      return p;
+    }
+    return cache::expose<MLXCache>(this, id);
+  }
+
  private:
   // The step's bits as SDPA wants them: [1, 1, length, read_len], one row per
   // query token.

backends/mlx/runtime/MLXSequenceCache.h

@@ -125,6 +125,14 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache {
     return AttendSpec{K, V, AttendSpec::Mask::Causal, std::nullopt};
   }
 
+ protected:
+  void* face(cache::FaceId id) override {
+    if (void* p = cache::SequenceCache::face(id)) {
+      return p;
+    }
+    return cache::expose<MLXCache>(this, id);
+  }
+
  private:
   // A sequence cache holds one run of one sequence, so the step is described by
   // where it starts; the remaining positions carry no information beyond

backends/mlx/runtime/backend_options.h

@@ -42,14 +42,6 @@ inline constexpr char kClearCacheIntervalKey[] = "clear_cache_interval";
 // errors otherwise). Saves one full mutable-buffer (KV-cache) copy per handle.
 inline constexpr char kSkipMutableBufferInitKey[] = "skip_mutable_buffer_init";
 
-// Per-model runtime-spec key (string). Names the off-graph KV cache this handle
-// binds to: the runner creates the cache, installs it in the process-global
-// CacheRegistry under this key, and the delegate looks it up in init(). The
-// DelegateHandle is opaque to the host, so the key is the only rendezvous
-// channel. Unset means no cache, and any update_and_attend node then fails at
-// execute() rather than silently attending nothing.
-inline constexpr char kCacheKeyKey[] = "cache_key";
-
 } // namespace mlx
 } // namespace backends
 } // namespace executorch

backends/mlx/test/mlx_cell_cache_test.cpp

@@ -268,12 +268,16 @@ TEST_F(MLXCellCacheTest, InvalidConfigThrows) {
 // A runner reaches a layout by (backend_id, kind), so the builder registration
 // is as much a part of the layout as the class.
 TEST_F(MLXCellCacheTest, RegistryBuildsCellLayout) {
-  auto built = cache::CacheBuilderRegistry::global().build(
-      kMLXBackendId, "cell", flat_config(32, 1, H, D, kHalf));
+  auto built = cache::CacheFactory::global().build(
+      kMLXBackendId,
+      cache::kind::kBatchedCell,
+      flat_config(32, 1, H, D, kHalf));
   ASSERT_TRUE(built.ok());
-  const std::shared_ptr<cache::CacheBase>& c = *built;
-  EXPECT_NE(c->as_batch_control(), nullptr);
-  EXPECT_EQ(c->as_control(), nullptr);
+  const std::shared_ptr<cache::Cache>& c = *built;
+  EXPECT_NE(c->as<cache::BatchControl>(), nullptr);
+  EXPECT_NE(c->as<MLXCache>(), nullptr) << "the backend face comes back too";
+  // A cell layout is multi-sequence, so it offers no single-sequence face.
+  EXPECT_EQ(c->as<cache::SequenceControl>(), nullptr);
 }
 
 } // namespace

backends/mlx/test/op_test_runner.cpp

@@ -300,38 +300,33 @@ int main(int argc, char* argv[]) {
 
     namespace cache = ::executorch::extension::llm::cache;
 
-    // Build and install the off-graph KV cache before the Module, so the
-    // registry entry exists by the time the delegate's init() looks it up.
-    // Declared here so the session outlives the module.
-    std::optional<cache::CacheSession> cache_session;
+    // Publish the off-graph KV cache until the delegate resolves its key while
+    // loading the method.
+    std::optional<cache::InstallGuard> cache_install_guard;
     if (!kv_cache_spec.empty()) {
       cache::CacheConfig cfg{};
       if (!parse_kv_cache_spec(kv_cache_spec, cfg)) {
         std::cerr << "Invalid --kv-cache spec: " << kv_cache_spec << std::endl;
         return 1;
       }
-      auto built = cache::CacheBuilderRegistry::global().build(
-          ::executorch::backends::mlx::kMLXBackendId, "seq", cfg);
+      auto built = cache::CacheFactory::global().build(
+          ::executorch::backends::mlx::kMLXBackendId,
+          cache::kind::kSingle,
+          cfg);
       if (!built.ok()) {
         std::cerr << "Failed to build KV cache: "
                   << static_cast<int>(built.error()) << std::endl;
         return 1;
       }
-      cache_session.emplace(cache::make_unique_key(), built.get());
-      if (verbose) {
-        std::cout << "Installed KV cache under key " << cache_session->key()
-                  << std::endl;
-      }
+      cache_install_guard.emplace(built.get());
     }
 
     Module module(pte_path);
     Error load_error = Error::Ok;
-    if (cache_session) {
+    if (cache_install_guard) {
       ::executorch::runtime::BackendOptions<1> mlx_opts;
       ::executorch::runtime::LoadBackendOptionsMap options_map;
-      if (mlx_opts.set_option(
-              ::executorch::backends::mlx::kCacheKeyKey,
-              cache_session->key().c_str()) != Error::Ok ||
+      if (cache_install_guard->set_option(mlx_opts) != Error::Ok ||
           options_map.set_options(
               ::executorch::backends::mlx::kMLXBackendId, mlx_opts.view()) !=
               Error::Ok) {
@@ -358,6 +353,7 @@ int main(int argc, char* argv[]) {
                 << static_cast<int>(load_method_error) << std::endl;
       return 1;
     }
+    cache_install_guard.reset();
 
     if (verbose) {
       std::cout << "Reading inputs from: " << input_path << std::endl;