Add neutral KV-cache core: registry and bookkeeping

[kiymetakdemir]

Summary

Adds a shared, backend-agnostic core for the off-graph KV cache: a set of C++ pieces that manage cache length/limits and let a runtime look up the right cache by a key. It has no backend or tensor code.

Files

• extension/llm/cache/cache.h — the cache interface, the contiguous length-tracking logic, and the config struct.
• extension/llm/cache/cache_registry.h / cache_registry.cpp — the registry that stores caches by key, a per-kind builder, and an auto-install/remove helper.
• extension/llm/cache/CMakeLists.txt — builds the library and pulls in the test folder (cmake).
• extension/llm/cache/targets.bzl — builds the library (buck).
• extension/llm/cache/test/cache_test.cpp — unit tests for the above.
• extension/llm/cache/test/CMakeLists.txt — builds the test (cmake).
• extension/llm/cache/test/targets.bzl / test/BUCK — builds the test (buck).
• shim_et/xplat/executorch/build/build_variables.bzl — registers the new source file.
• tools/cmake/Codegen.cmake — wires that source list into cmake.
• CMakeLists.txt — includes the new folder in the build.

Testing

Added C++ unit tests (test/cache_test.cpp) covering the bookkeeping (plan/can_extend/rewind/clear), the registry (install/get/erase, unique keys), the builder registry (build + missing-kind error), and the session's install-on-create / remove-on-destroy. Built and ran through the standard cmake/ctest flow, all pass:

cmake -B cmake-out -DEXECUTORCH_BUILD_EXTENSION_LLM=ON -DEXECUTORCH_BUILD_TESTS=ON
cmake --build cmake-out --target extension_llm_cache_test -j
ctest --test-dir cmake-out -R extension_llm_cache --output-on-failure

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21383

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 5 New Failures, 4 Pending, 2 Unrelated Failures, 1 Unclassified Failure

As of commit 94744ea with merge base 1fb8b57 ():

NEW FAILURES - The following jobs have failed:

• Build Presets / apple (ios-simulator) / build (gh)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• Build Presets / apple (ios) / build (gh)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• Build Presets / apple (llm) / build (gh)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• Build Presets / linux (linux, linux.arm64.2xlarge, executorch-ubuntu-22.04-gcc11-aarch64) / build (gh)
  RuntimeError: Command docker exec -t 8321c5624ed3dce3fe2b5b1d278bb20b8c0cd1d2f99c1de9a5fe46ccc0d05e00 /exec failed with exit code 1
• pull / android / build-android (gh)
  Process completed with exit code 1.

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• Check Labels / Check labels (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  Process completed with exit code 1.

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• pull / test-arm-backend-public-api-backward-compatibility / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• pull / unittest-nxp-neutron / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[metascroy]

What is min_chunk mean?

[metascroy]

What does this mean?

[metascroy]

Either stick with CMake only (we can buckify later), or if doing buck, let's import to internal to make sure it works

[metascroy]

Two things about this I'm hesitant about:

1. I don't like the name Contiguous b/c it implies something about memory
2. The interface here has two audiences that want separate APIs: 1) the backends which want "plan", and 2) the application/user which wants "can_extend" / "clear" / "rewind".

What do you think of splitting the "Bookkeeping" class to serve these two audiences. Something like "Control" class for application, and "Planner" for backend. A cache inherits both "FlatCache : public SequenceControl, public SequencePlanner"

class CacheBase {                          // registry ownership / erasure anchor
 public: virtual ~CacheBase() = default;
};

class SequenceControl : public virtual CacheBase {   // application face
 public:
  virtual bool can_extend(int n = 1) const = 0;
  virtual void rewind(int new_len) = 0;
  virtual void clear() = 0;
};

class SequencePlanner : public virtual CacheBase {   // backend face
 public:
  virtual SeqStepPlan plan(int position, int T) = 0;
};

// ---- Flat: non-wrapping, full history [0,length) --------------------------
class FlatCache : public SequenceControl, public SequencePlanner

// ---- Ring: sliding window, wraps, evicts ---------------------------------
class RingCache : public SequenceControl, public SequencePlanner

The planner and step plan are tied. SequencePlanner (single sequence) returns a step plan like this:

struct Run { int start = 0; int len = 0; };     // contiguous slot range in the pool

struct SeqStepPlan {
  bool contiguous = true;   // hint: single write run AND single read run (no wrap)
  int  valid_len  = 0;      // # valid context rows (= sum of read lens)
  Run  write[2];            // where the T new rows land (1 run, or 2 on ring wrap)
  Run  read[2];             // pool rows holding valid context, logical order
  int  read_base_pos = 0;   // logical pos of read[0].start (RoPE/mask). Flat: 0.
};

A TreePlanner/ BatchedPaged one would return a different kind of step plan.

Thoughts?

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 43s —— View job

---

Code Review

• Read the cache core, registry, and tests
• Review build wiring (cmake/buck)
• Assess design against reviewer's feedback
• Post review

Nicely scoped, well-commented, and the tests cover the bookkeeping/registry/session paths cleanly. A few substantive points below.

---

1. Exceptions vs. ExecuTorch error handling (most important)

cache.h and cache_registry.cpp use throw std::runtime_error in plan(), rewind(), and build(). ExecuTorch's runtime deliberately avoids C++ exceptions and compiles size-optimized builds with -fno-exceptions (see CMakeLists.txt:301-316), using Result<T> / Error / ET_CHECK_MSG instead. Since this is billed as a neutral core shared across backends — including embedded targets — throw is a portability hazard: any consumer built with -fno-exceptions will std::terminate on these paths, and the library won't compile under some stricter configs.

Consider mirroring the rest of the runtime: return Error/Result<StepPlan> or use ET_CHECK_MSG. Worth confirming the intended error convention with maintainers before this pattern spreads to the backend layers.

Fix this →

2. On @metascroy's audience-split proposal (cache.h:46)

I agree with the direction. Today plan()/rewind() live only on the concrete ContiguousBookkeeping, while the Cache base exposes can_extend/capacity/clear. That means a runner holding a Cache* (or a CacheSession<Cache>) can't rewind or plan without downcasting — the interface already leaks the "two audiences" problem. Splitting into a SequenceControl (application) + SequencePlanner (backend) face, both anchored on a minimal CacheBase for registry ownership/erasure, resolves that and keeps StepPlan polymorphic per planner kind (Flat/Ring/Tree). The CacheSession<Face> template is already written to carry a derived face, so it fits.

If you go this way, the registry's value type stays shared_ptr<CacheBase> and the two interfaces are recovered via dynamic_cast — which loops back to point 1, since dynamic_cast needs RTTI (also disabled in size builds). Worth designing the recovery mechanism (typed builders returning the right face, or an explicit as_planner()/as_control()) alongside the split.

3. StepPlan::contiguous (cache.h:37) & min_chunk (cache.h:105)

Answering the two inline questions in-code would help. contiguous is currently hardcoded true for the only cache kind, so its meaning (a hint that write+read are each a single non-wrapping run — only meaningful once a ring/wrapping cache exists) isn't discoverable from this PR alone. Same for min_chunk = 512: it's declared but unused here, so the "512" reads as a magic default. Either a one-line comment on each field, or deferring min_chunk until the pool-sizing code that consumes it lands, would remove the ambiguity.

4. Naming — ContiguousBookkeeping

Agree with the reviewer: "Contiguous" reads as a memory-layout claim, but this class tracks a logical length counter and does no byte/memory work. Something like FlatSequence / LinearSequence (paired with a future RingSequence) describes the eviction/history policy rather than implying storage contiguity.

5. Registry as process-global mutable singletons

CacheRegistry/CacheBuilderRegistry::global() are process-wide singletons the tests mutate directly (register_builder("TestBackend", ...) persists across the test binary's lifetime). It's serviceable and the atomic make_unique_key() avoids key collisions, but two notes:

• Test isolation: builder registrations leak between test cases. Fine at this size; worth an eye as more kinds register.
• The design comment says the registry, session guard, and delegate handle all co-own the cache so mid-method erase is safe — good, and the shared_ptr value type backs that up.

6. targets.bzl / buck (targets.bzl:5)

@metascroy's call stands: either CMake-only for now (buckify later) or import to internal to confirm the buck targets actually build. The current targets.bzl/BUCK aren't exercised by the public CI, so they're untested as written.

Minor

• int is used throughout for capacity/length/position. Fine for sequence lengths, but if these ever feed byte-offset math in the pool layer, be deliberate about widening then.
• CI failures on this PR (linux aarch64, android, apple) are all HTTP 504s on third-party downloads (KleidiAI zip, Gradle distribution) — infra flakes, not caused by this change. The cache lib builds and installs fine in the Android job.

Overall: solid foundation. The exception-handling convention (#1) is the one I'd resolve before layering backends on top, and the interface split (#2) is worth doing now while there's a single implementation to migrate.
· kv-cache-shared-registry