Add neutral KV-cache core: registry and bookkeeping

CMakeLists.txt

@@ -903,6 +903,8 @@ if(EXECUTORCH_BUILD_EXTENSION_LLM)
     )
   endif()
   list(APPEND _executorch_extensions tokenizers)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/cache)
+  list(APPEND _executorch_extensions extension_llm_cache)
 endif()
 
 if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL)

extension/llm/cache/CMakeLists.txt

@@ -0,0 +1,38 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Neutral, tensor-free KV-cache core shared across backends (registry,
+# bookkeeping, config). No backend or tensor dependency. CMake-only for now;
+# buck targets are added later (see the buckify follow-up).
+
+if(NOT EXECUTORCH_ROOT)
+  set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..)
+endif()
+
+add_library(extension_llm_cache cache_registry.cpp)
+target_link_libraries(extension_llm_cache executorch_core)
+target_include_directories(
+  extension_llm_cache PUBLIC ${_common_include_directories}
+)
+target_compile_options(extension_llm_cache PUBLIC ${_common_compile_options})
+
+install(
+  TARGETS extension_llm_cache
+  EXPORT ExecuTorchTargets
+  DESTINATION ${CMAKE_INSTALL_LIBDIR}
+  INCLUDES
+  DESTINATION ${_common_include_directories}
+)
+install(
+  DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/
+  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/extension/llm/cache
+  FILES_MATCHING
+  PATTERN "*.h"
+)
+
+if(BUILD_TESTING)
+  add_subdirectory(test)
+endif()

extension/llm/cache/cache.h

@@ -0,0 +1,155 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+// Neutral, tensor-free KV-cache core shared across backends. A cache exposes
+// two faces: a runner-facing control face (SequenceControl) and a
+// backend-facing planner face (SequencePlanner). Both are recovered from the
+// owning CacheBase* via as_control()/as_planner() -- static upcasts, so no
+// dynamic_cast/RTTI and no diamond inheritance. Errors propagate as
+// Error/Result<T> (no exceptions, no abort). See extension/llm/cache for the
+// full design.
+
+#include <vector>
+
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/result.h>
+
+namespace executorch {
+namespace extension {
+namespace llm {
+namespace cache {
+
+using ::executorch::runtime::Error;
+using ::executorch::runtime::Result;
+
+class SequenceControl;
+class SequencePlanner;
+
+// Registry ownership / erasure anchor. The registry owns a cache as a
+// CacheBase*; the runner recovers the control face and the backend recovers the
+// planner face through these accessors (each concrete cache returns `this`).
+class CacheBase {
+ public:
+  virtual ~CacheBase() = default;
+  virtual SequenceControl* as_control() = 0;
+  virtual SequencePlanner* as_planner() = 0;
+};
+
+// Application (runner) face: lifecycle + admission, tensor-free.
+class SequenceControl {
+ public:
+  virtual ~SequenceControl() = default;
+  virtual bool can_extend(int n = 1) const = 0; // admission / hard-stop
+  virtual int capacity() const = 0; // logical cap
+  virtual Error rewind(int new_len) = 0; // truncate (agent backtracking)
+  virtual void clear() = 0; // reset for reuse
+};
+
+// Integer-only handoff from the planner to the backend byte layer: where to
+// write this step and how much history to read.
+struct SeqStepPlan {
+  int valid_len; // contiguous read bound
+  int write_start; // contiguous write run start
+  int write_len; // contiguous write run length
+};
+
+// Backend face: turns (position, T) into a write/read plan over the pool.
+class SequencePlanner {
+ public:
+  virtual ~SequencePlanner() = default;
+  virtual Result<SeqStepPlan> plan(int position, int T) = 0;
+};
+
+// Model facts + runtime policy the byte layer sizes its pools from.
+// Architecture facts (n_layers/n_kv_heads/head_dim) come from .pte metadata;
+// capacity and initial_capacity are runtime policy. KV storage dtype joins this
+// with the quantized pool; kept tensor-free here so the neutral core carries no
+// backend types.
+//
+// n_kv_heads/head_dim are per-layer so hybrid / mixed-head / MLA models (KV
+// dims vary by layer) need no config change; a single-element vector means
+// uniform across all layers (the common case).
+struct CacheConfig {
+  int capacity;
+  int n_layers;
+  std::vector<int> n_kv_heads;
+  std::vector<int> head_dim;
+  // Initial pool size in cells: the byte layer's lazy-doubling pool starts at
+  // initial_capacity and doubles up to `capacity`, so short sequences avoid
+  // reallocation.
+  int initial_capacity = 512;
+};
+
+// Flat single-sequence cache: non-wrapping history [0, length) tracked by a
+// length counter. Implements both faces; a backend byte layer builds on it.
+class FlatCache : public CacheBase,
+                  public SequenceControl,
+                  public SequencePlanner {
+ public:
+  explicit FlatCache(int capacity) : capacity_(capacity) {}
+
+  // CacheBase: face recovery without RTTI.
+  SequenceControl* as_control() override {
+    return this;
+  }
+  SequencePlanner* as_planner() override {
+    return this;
+  }
+
+  // SequenceControl.
+  bool can_extend(int n = 1) const override {
+    return length_ + n <= capacity_;
+  }
+  int capacity() const override {
+    return capacity_;
+  }
+  // Rows are overwritten on the next write, so reuse/truncation is byte-free.
+  void clear() override {
+    length_ = 0;
+  }
+  Error rewind(int new_len) override {
+    ET_CHECK_OR_RETURN_ERROR(
+        new_len <= length_,
+        InvalidArgument,
+        "rewind: cannot grow (new_len=%d > length=%d)",
+        new_len,
+        length_);
+    length_ = new_len;
+    return Error::Ok;
+  }
+
+  // SequencePlanner: place T tokens at `position`, return the write run + read
+  // bound.
+  Result<SeqStepPlan> plan(int position, int T) override {
+    const int end = position + T;
+    ET_CHECK_OR_RETURN_ERROR(
+        end <= capacity_,
+        InvalidArgument,
+        "cache: %d cells exceeds capacity %d",
+        end,
+        capacity_);
+    if (end > length_) {
+      length_ = end;
+    }
+    return SeqStepPlan{
+        /*valid_len=*/end,
+        /*write_start=*/position,
+        /*write_len=*/T};
+  }
+
+ private:
+  int capacity_;
+  int length_ = 0;
+};
+
+} // namespace cache
+} // namespace llm
+} // namespace extension
+} // namespace executorch

extension/llm/cache/cache_registry.cpp

@@ -0,0 +1,83 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/extension/llm/cache/cache_registry.h>
+
+#include <atomic>
+
+#include <executorch/runtime/core/error.h>
+
+namespace executorch {
+namespace extension {
+namespace llm {
+namespace cache {
+
+CacheRegistry& CacheRegistry::global() {
+  static CacheRegistry registry;
+  return registry;
+}
+
+void CacheRegistry::install(
+    const std::string& key,
+    std::shared_ptr<CacheBase> cache) {
+  std::lock_guard<std::mutex> lock(mu_);
+  caches_[key] = std::move(cache);
+}
+
+std::shared_ptr<CacheBase> CacheRegistry::get(const std::string& key) const {
+  std::lock_guard<std::mutex> lock(mu_);
+  const auto it = caches_.find(key);
+  return it == caches_.end() ? nullptr : it->second;
+}
+
+void CacheRegistry::erase(const std::string& key) {
+  std::lock_guard<std::mutex> lock(mu_);
+  caches_.erase(key);
+}
+
+CacheBuilderRegistry& CacheBuilderRegistry::global() {
+  static CacheBuilderRegistry registry;
+  return registry;
+}
+
+void CacheBuilderRegistry::register_builder(
+    const std::string& backend_id,
+    const std::string& kind,
+    CacheBuilder builder) {
+  std::lock_guard<std::mutex> lock(mu_);
+  builders_[backend_id + ":" + kind] = std::move(builder);
+}
+
+Result<std::shared_ptr<CacheBase>> CacheBuilderRegistry::build(
+    const std::string& backend_id,
+    const std::string& kind,
+    const CacheConfig& cfg) const {
+  CacheBuilder builder;
+  {
+    std::lock_guard<std::mutex> lock(mu_);
+    const auto it = builders_.find(backend_id + ":" + kind);
+    ET_CHECK_OR_RETURN_ERROR(
+        it != builders_.end(),
+        InvalidArgument,
+        "no cache builder registered for %s:%s",
+        backend_id.c_str(),
+        kind.c_str());
+    builder = it->second;
+  }
+  return builder(cfg);
+}
+
+std::string make_unique_key() {
+  static std::atomic<uint64_t> counter{0};
+  return "cache-" + std::to_string(counter.fetch_add(1));
+}
+
+} // namespace cache
+} // namespace llm
+} // namespace extension
+} // namespace executorch