Add off-graph KV-cache registry

[kiymetakdemir]

Summary

Adds a shared, backend-agnostic registry for the off-graph KV cache: a small set of C++ pieces that let a runtime own a cache and look it up by key.

Files

• extension/llm/cache/cache.h — the opaque cache handle (CacheBase) and the CacheConfig struct.
• extension/llm/cache/cache_registry.h / cache_registry.cpp — a process-global registry that stores caches by key, a
  per-(backend_id, kind) builder registry, unique-key generation, and a RAII session that installs on construction
  and erases on destruction.
• extension/llm/cache/CMakeLists.txt — builds the library and pulls in the test folder (cmake).
• extension/llm/cache/test/cache_test.cpp — unit tests for the above.
• extension/llm/cache/test/CMakeLists.txt — builds the test (cmake).
• CMakeLists.txt — includes the new folder in the build.

Testing

Added C++ unit tests (test/cache_test.cpp) covering the registry (install/get/erase, unique keys), the builder
registry (build + missing-kind NotFound error), and the session's install-on-create / erase-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.

⏳ No Failures, 8 Pending

As of commit 40c1666 with merge base 55d693b ():
💚 Looks good so far! There are no failures yet. 💚

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?

[kiymetakdemir]

It is the initial capacity added a comment for this.

[metascroy]

What does this mean?

[kiymetakdemir]

We don't use it actually, I removed.

[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?

[kiymetakdemir]

Applied this together with claude's suggestion on explicit as_planner()/as_control().

[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

[metascroy]

@claude review this code

[claude]

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

---

Code Review

• Read the cache core, registry, and tests
• Review build wiring (cmake)
• Assess against reviewer feedback & follow-up changes
• Post review

The revision addresses the major points from the earlier round cleanly. Quick status on prior feedback, then a few remaining nits.

Resolved since last review ✅

• Exceptions → ET_CHECK_MSG (Add support for quantized LeakyReLU #1): No more throw std::runtime_error. rewind/plan/build now use ET_CHECK_MSG, so the core is -fno-exceptions-safe. 👍
• Interface split (Re-sync with internal repository #2):CacheBase + SequenceControl + SequencePlanner with as_control()/as_planner(). The three bases are independent (not virtual-inheriting from CacheBase), so recovery is a static upcast returning this — no diamond, no RTTI. This is a cleaner solution than the dynamic_cast path floated earlier.
• Naming (Add unlifting pass under private config #4):ContiguousBookkeeping → FlatCache, paired conceptually with a future RingCache. Reads as a policy, not a layout claim.
• min_chunk / contiguous (Rename _pt2e to pt2e #3):min_chunk now documented as initial pool size; the ambiguous contiguous field is gone from SeqStepPlan.
• buck (Re-sync with internal repository #6): Went CMake-only. The new-core diff is 7 files, no .bzl/BUCK — matches your "buckify later" call.

Remaining items

1. Stale doc comment in cache_registry.h:63

// Throws if no builder is registered for (backend_id, kind).
std::shared_ptr<CacheBase> build(...) const;

It no longer throws — build() uses ET_CHECK_MSG (aborts). Update the comment to match, otherwise a reader will expect a catchable exception. Fix this →

2. plan() mutates state (cache.h:122)
plan() advances length_ as a side effect, so it both computes and commits the step. Worth a one-line comment on the method that "plan commits the length" — a backend that expects plan() to be a pure query (e.g. to peek a layout before deciding) would be surprised. If a peek-vs-commit split is ever needed, this is the seam.

3. valid_len semantics on re-plan / overwrite (cache.h:132-135)
valid_len = position + T, and length_ only grows (if (end > length_) length_ = end). If a caller re-plans at a position < length_ (e.g. overwrite after a rewind), valid_len reflects position + T, not the previously-tracked length_. That's a reasonable "read bound = what you just wrote" contract, but it's implicit — a comment stating that plan returns the read bound for this step (not necessarily the full tracked length) would remove ambiguity for the ring/tree planners that follow.

Minor

• Default arg on a virtual (cache.h:44):can_extend(int n = 1) — default args on virtuals bind to the static type, a known footgun if a future override changes the default. Harmless with one impl today; just flagging before more caches inherit the face.
• FlatCache is not internally thread-safe. The registry is mutex-guarded, but the cache object's length_ isn't. Fine for the single-sequence generation-loop model; worth an explicit "single-threaded per sequence" note if that assumption isn't documented elsewhere.
• Builder-registry test isolation:register_builder("TestBackend", ...) still mutates the process-global and persists across cases in the binary. Serviceable at this size (noted last round); keep an eye as more kinds register.
• int for capacity/length/position throughout — fine for sequence lengths; be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions from the last review are settled well, and the as_control()/as_planner() recovery is a nice RTTI-free resolution. Nothing blocking — items above are a stale comment plus a couple of clarifying comments on plan()/valid_len semantics.
· kv-cache-shared-registry

[metascroy]

I think initial_capacity is a better name?

[metascroy]

What does ET_CHECK_MSG do on failure? Call std::terminate?

[kiymetakdemir]

It calls std::abort()

[metascroy]

Are you planning to use this for Ring as well?

[kiymetakdemir]

Yes, I guess in that case it needs one more field for the window's start position.

[metascroy]

Looks great! Main concern is calling std::abort. Instead wrap in Result instead.

See how ET runtime / delegates avoids exception without calling abort.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 31s —— View job

---

Code Review

• Read the cache core, registry, and tests
• Verify the abort→Result change requested last round
• Review build wiring (CMake)
• Post review

The last blocker — std::abort via ET_CHECK_MSG — is fully resolved. This round is clean; only minor hardening notes remain.

Resolved since last review ✅

• abort → Result/Error (main concern):plan() and build() now return Result<T> and rewind() returns Error, all via ET_CHECK_OR_RETURN_ERROR, which logs and returns rather than aborting (runtime/core/error.h:182). No std::abort/std::terminate on any error path — this matches how the ET runtime/delegates avoid exceptions. 👍
• Death tests → Result checks:cache_test.cpp now asserts .error() == Error::InvalidArgument and .ok() instead of EXPECT_DEATH. Matches the requested pattern.
• min_chunk → initial_capacity: renamed and documented as the lazy-doubling pool's starting size (cache.h:84-87).
• Interface split + RTTI-free recovery:CacheBase / SequenceControl / SequencePlanner with as_control()/as_planner() returning this — static upcast, no dynamic_cast.
• buck: CMake-only; no .bzl/BUCK in the diff, matching the "buckify later" call.

Minor / non-blocking

1. SeqStepPlan fields have no default initializers (cache.h:57-61)

structSeqStepPlan {
int valid_len;
int write_start;
int write_len;
};

FlatCache::plan aggregate-inits all three, so it's safe today. But once ring/tree planners return partially-filled plans, an uninitialized field reads as garbage. The proposal in the review thread had = 0 defaults — worth keeping them here to stay safe as planners multiply. Fix this →

2. plan() / rewind() don't guard negative inputs (cache.h:117,130)
plan(position, T) checks end <= capacity_ but not position >= 0 || T >= 0; a negative position yields a negative write_start. rewind(new_len) checks new_len <= length_ but not new_len >= 0, so rewind(-1) sets length_ = -1 and silently corrupts can_extend. Both are internal-API misuse rather than user input, but a lower-bound ET_CHECK_OR_RETURN_ERROR is cheap insurance.

3. plan() commits length as a side effect (cache.h:138-140)
plan() both computes the step and advances length_. A one-line comment noting "plan commits the length" would help a future planner that expects a pure peek before deciding a layout. (Raised last round; still worth the comment.)

4. Carried-forward nits (all still fine at this size):

• can_extend(int n = 1) — default arg on a virtual binds to the static type; harmless with one impl.
• FlatCache isn't internally thread-safe (length_ unguarded); fine under the single-sequence generation-loop model — worth an explicit note.
• Builder-registry test mutates the process-global (register_builder("TestBackend", …)) and persists across cases.
• int for capacity/length/position — be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions are settled and the error-handling convention is now consistent with the rest of the runtime. Nothing blocking — the only one I'd fold in before backends land is the SeqStepPlan default initializers (#1), since those plans get consumed by byte-layer code shortly.
· kv-cache-shared-registry

[metascroy]

Let's keep gemma4 in mind for the first model we enable here, which has a mix of flat/ring, which I think must be captured in config. I think both must be driven from the same controller.

What if we do:

struct LayerPolicy {
enum class Kind : int { Flat = 0, Ring = 1 }; // serialized values: append-only
Kind kind = Kind::Flat;
int window = 0; // Ring: sliding-window size; must be 0 when Flat
};
struct LayerConfig {
LayerPolicy policy; // default Flat
int n_kv_heads;
int head_dim;
// dtype / scale / sink size land here later (not this PR)
};
struct CacheConfig {
int capacity; // logical cap
int n_layers;
std::vector<LayerConfig> layers; // size 1 == uniform, else == n_layers
int initial_capacity = 512;
};

[digantdesai]

Also see CacheLayerMixIn from hf/transformers

[metascroy]

Thinking ahead to mixed gemma4 model, maybe we have one SequenceCache, with flat/ring being policies. We need one control driver for both caches across all layers I think. Something like this:

class SequenceCache : public CacheBase,
public SequenceControl,
public SequencePlanner {
public:
// CacheBase: static upcasts, no RTTI (each returns `this`).
SequenceControl* as_control() override { return this; }
SequencePlanner* as_planner() override { return this; }
// SequenceControl (runner face): length / admission / rewind / clear.
bool can_extend(int n = 1) const override; // min over policies' retention
int capacity() const override { return capacity_; }
Error rewind(int new_len) override; // bounded by smallest window
void clear() override { length_ = 0; }
// SequencePlanner (backend face): dispatch to the layout policy.
Result<SeqStepPlan> plan(int policy_id, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(
policy_id >= 0 && policy_id < (int)policies_.size(),
InvalidArgument, "bad policy_id %d", policy_id);
return policies_[policy_id]->plan(position, T);
}
private:
int capacity_;
int length_ = 0;
std::vector<std::unique_ptr<LayoutPolicy>> policies_; // FlatPolicy / RingPolicy
std::vector<int> layer_to_policy_; // built from CacheConfig
};

Here LayoutPolicy is something like:

class LayoutPolicy {
public:
virtual ~LayoutPolicy() = default;
virtual SeqStepPlan plan(int position, int T) const = 0; // const == no length side effect
virtual int depth() const = 0; // physical rows to allocate
};

with FlatPolicy having the logic you have for FlatCache here, e.g.,

class FlatPolicy final : public LayoutPolicy {
public:
explicit FlatPolicy(int capacity) : capacity_(capacity) {}
int depth() const override { return capacity_; }
SeqStepPlan plan(int position, int T) const override {
const int end = position + T;
SeqStepPlan p{};
p.n_write = 1; p.write[0] = {position, T}; // contiguous, never wraps
p.n_read = 1; p.read[0] = {0, end}; // attend over all history
return p;
}
private:
int capacity_;
};

Although this diff is flat cache, we need SeqStepPlan to work for both. A ring can have 1 or 2 contiguous runs, a flat cache will always have 1:

struct Run { int start; int len; }; // a contiguous span of physical rows
struct SeqStepPlan {
Run write[2]; int n_write; // where to place this step's new KV
Run read[2]; int n_read; // valid history to attend over, oldest→newest
};

[metascroy]

@claude review again in light of latest comments on policy

[claude]

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

---

Review — the gemma4 mixed flat/ring policy proposal

• Read cache core, registry, tests
• Assess current code against the mixed flat/ring policy design
• Post review

I read the current cache.h/cache_registry.* against your two latest comments (the CacheConfig/LayerPolicy shape at cache.h:79, and the single-SequenceCache-with-LayoutPolicy shape at cache.h:92). I think the direction is right — below is where I'd push on the details before it lands, since these are the surfaces backends will bind to.

The core call: one controller, per-layer policies ✅

Agreed, and I'd go further — for a mixed model this isn't just nicer, the separate-class design is actually wrong. Gemma-style alternating local/global attention means flat and ring layers coexist in one model, but there is exactly one logical sequence length across all layers. Two independent FlatCache/RingCache objects would each track their own length_, and can_extend/rewind would no longer be coherent for the sequence as a whole. A single SequenceCache that owns length_ and dispatches layout to per-layer policies is the only thing that keeps admission/rewind consistent. So the SequenceCache + vector<LayoutPolicy> + layer_to_policy_ shape is the one to build.

1. Dispatch key: prefer layer over policy_id

plan(int policy_id, int position, int T) leaks an internal id to the backend. The byte layer knows which layer it's materializing, not which policy bucket that maps to. Keeping the mapping private is cleaner:

Result<SeqStepPlan> plan(int layer, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(layer >= 0 && layer < (int)layer_to_policy_.size(),
InvalidArgument, "bad layer %d", layer);
return policies_[layer_to_policy_[layer]]->plan(position, T);
}

layer_to_policy_ is already in your sketch — this just makes it the public parameter and keeps policy_id an implementation detail.

2. Length commit must be idempotent (plan runs N_layers times per step)

Making LayoutPolicy::planconst (pure layout, no side effect) is the right resolution to the "plan commits length" concern from earlier rounds — the controller owns length now. But watch the seam: the backend calls plan()once per layer per step, so SequenceCache must not advance length_ additively inside plan. Idempotent length_ = max(length_, position + T) is safe to call every layer; an additive length_ += T would multiply by n_layers. Either document that plan's length commit is idempotent, or split it out into a single control()->advance(T) the runner calls once per step and keep plan fully read-only. I'd lean to the explicit advance — it makes "one commit per step" impossible to get wrong.

3. can_extend / rewind semantics across mixed policies

Your one-liners ("min over policies' retention", "bounded by smallest window") are the correct contract, but they carry real logic worth stating in code:

• can_extend: a ring layer never refuses (it evicts), so admission is gated only by the flat/full layers → effectively bounded by capacity. Ring layers contribute no bound. Fine, but that means can_extend is really "any non-evicting policy near its cap," so the controller needs each policy to expose whether it evicts.
• rewind(new_len): this is the subtle one. A ring layer with window W has physically evicted everything older than length_ - W, so rewind to a new_len < length_ - min_window is unrecoverable — that history is gone from the ring layers even though the flat layers still have it. The bound isn't just new_len <= length_; it's new_len >= length_ - smallest_window too. That lower bound needs the window sizes in the controller, so LayoutPolicy should expose retained_from(length) (or similar) alongside depth().

4. SeqStepPlan: read[] runs are needed, but so is the read base position

Reintroducing Run read[2]; int n_read is correct — for a ring cache the valid history is not[0, length), it wraps, so a single valid_len can't describe it (that's why the earlier flat-only valid_len field doesn't generalize). Two points:

• Run write[2]/read[2] caps at 2 runs, which is exactly right for flat (1) and ring (≤2). Make it explicit in a comment that SeqStepPlan is deliberately the flat∪ring shape and that paged/tree planners return their own plan type (you noted this earlier) — otherwise [2] reads as an arbitrary limit.
• Your first proposal (the 18:10 sketch) had read_base_pos — the logical position of read[0].start. The latest slimmed SeqStepPlan dropped it. For ring, the backend needs the logical positions of the read runs to apply RoPE and the attention mask correctly (the oldest retained row is not at logical position 0 after eviction). Unless positions are reconstructed elsewhere, I'd bring read_base_pos back (or a per-run base) — dropping it looks like it'll break RoPE on the ring path.

5. Config shape

The LayerPolicy/LayerConfig/CacheConfig nesting is good and folds the current parallel n_kv_heads/head_dim vectors into an array-of-structs, which is cleaner. Keep the "size 1 == uniform, else == n_layers" convention you already have for the head vectors — most models are uniform and shouldn't pay for a full per-layer list. Two small things:

• Keep the enum class Kind : int with the "serialized values: append-only" comment — good instinct if this ever hits .pte metadata.
• State the window == 0 iff Flat invariant as an ET_CHECK_OR_RETURN_ERROR in whatever validates the config, not just a comment, since it's now load-bearing for depth()/rewind.

Scope: adopt at least the shapes in this PR

You've already approved, so this is a "now vs. next PR" judgment. My take: the internal SequenceCache/LayoutPolicy refactor is cheapest to do now, while FlatPolicy is the only implementation to migrate — the same argument that justified doing the control/planner split in this PR rather than deferring it. At minimum, land the CacheConfig/LayerConfig and SeqStepPlan (write[2]/read[2] + read_base_pos) shapes now, since those are the surfaces the byte layer and any serialized metadata will bind to and are the expensive ones to change later. The FlatCache→SequenceCache+FlatPolicy mechanical move can follow if you'd rather keep this PR small, but the config/plan structs I'd fix here.

Nothing above contradicts the approval — it's design hardening for the mixed-model case you're aiming at. The two I'd treat as near-blocking for the refactor are #2 (idempotent length commit) and #4 (read_base_pos for ring RoPE), since both are silent-correctness traps rather than cosmetic.
· kv-cache-shared-registry

[digantdesai]

can we wrap the base class in ET runtime friendly cache class? This is to avoid dependence on ET from runner side. I.e. I want this to be useful even when using outside ET runner. Or we can split this file into two, runner facing which can use ET stuff because for non-ET runner we can have another class. And the backend cache interface, which shouldn't use ET stuff.

Code location is also a problem, but we can discuss that later, and I don't want to block on that one.

[kiymetakdemir]

Sure! After looking into it, I'm leaning toward implementing an ET wrapper so the whole core stays a single, ET-free implementation.

[digantdesai]

Ok sorry I wasn't very clear, given I didn't see this PR.

To rephrase, let's keep the ET interfaces separate (and can be used elsewhere outside of 'et::'). The rationale is we might want to use this in a new backend layer which is not an ET delegate or linked against ET runtime libs.

For now the ET specific registry design seems OK.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 25s —— View job

---

Code Review

• Read the current cache core, registry, and tests
• Review build wiring (CMake) and test gating
• Assess against the latest ET-independence feedback (digantdesai)
• Post review

The PR has been trimmed substantially since the last round. cache.h is now just the ET-free CacheBase anchor + CacheConfig — the SequenceControl/SequencePlanner/FlatCache/SeqStepPlan machinery from earlier rounds is gone, deferred to the follow-up that lands the concrete cache. That's a clean response to @digantdesai's request to keep the neutral core usable outside an ET runner.

Resolved / consistent with feedback ✅

• ET-free core (digantdesai):cache.h has zero ExecuTorch includes and no et:: dependency — it's a plain CacheBase + CacheConfig. The ET-specific Error/Result usage is confined to cache_registry.*, and the header comments state the split explicitly. This matches "keep the ET interfaces separate so this can be used in a non-delegate backend layer."
• abort → Result/Error:build() returns Result<...> via ET_CHECK_OR_RETURN_ERROR (logs + returns, no std::abort), and the test asserts .error() == Error::NotFound rather than a death test. Consistent with the resolution from the earlier rounds.
• CMake-only, no buck: matches the "buckify later" call. Test gating uses if(BUILD_TESTING), consistent with the other extension/*/CMakeLists.txt.

Observations on the current (slimmed) surface

1. CacheConfig fields are uninitialized (cache.h:31-35)

structCacheConfig {
int capacity;
int n_layers;
int initial_capacity = 512;
};

capacity and n_layers have no default initializer. The test aggregate-inits both (CacheConfig cfg{32, 1}), so it's safe there, but a caller doing CacheConfig cfg; gets garbage in the two most important fields. Given initial_capacity already has a default, adding = 0 to the other two keeps the struct uniformly safe to default-construct. Fix this →

2. The gemma4 mixed flat/ring design (from the last round) isn't in this PR
The LayerPolicy/LayerConfig per-layer config and the SeqStepPlan (write[2]/read[2] + read_base_pos) shapes discussed at length aren't here — CacheConfig is still the flat capacity/n_layers/initial_capacity triple. That's fine as a scoping decision (this PR is now just the registry + neutral anchor), but worth confirming that's intentional: my earlier note stands that the config and step-plan structs are the expensive surfaces to change later because the byte layer and any serialized metadata bind to them. If gemma4 is the first target, folding at least the LayerConfig array-of-structs into CacheConfig now is cheaper than after consumers exist. Not blocking for a registry-only PR, but flag it so it's a deliberate "next PR" rather than an oversight.

3. Builder-registry key concatenation (cache_registry.cpp:53,63)
The map key is backend_id + ":" + kind. If either token can contain :, "a:b" + ":" + "c" and "a" + ":" + "b:c" collide. Backend ids/kinds are internal constants today so it's theoretical, but a one-line note (or a std::pair key) removes the ambiguity if these ever become user-facing strings.

Minor / carried forward (all still fine at this size)

• Process-global singletons + test isolation:register_builder("TestBackend", "stub", ...) mutates the process-global CacheBuilderRegistry and persists for the test binary's lifetime. Serviceable at one test; watch it as more kinds register.
• int for capacity/layers: fine for these counts; be deliberate about widening if any of these ever feed byte-offset math in the pool layer.
• Registry comment (cache_registry.h:69): "Returns Error::NotFound" now correctly matches the Result return — good, the stale "throws" comment from earlier rounds is fixed.

Overall: the scope reduction is the right move and directly addresses the ET-independence concern. Nothing blocking. The one I'd fold in here is the CacheConfig default initializers (#1); the mixed-policy config/plan shapes (#2) are a scoping judgment call for the follow-up.
· kv-cache-shared-registry

[metascroy]

Use std::pair, other than that, LGTM! Let's land, we can always fix later

[kiymetakdemir]

Summary

Adds a shared, backend-agnostic registry for the off-graph KV cache: a small set of C++ pieces that let a runtime own a cache and look it up by key.

Files

• extension/llm/cache/cache.h — the opaque cache handle (CacheBase) and the CacheConfig struct.
• extension/llm/cache/cache_registry.h / cache_registry.cpp — a process-global registry that stores caches by key, a
  per-(backend_id, kind) builder registry, unique-key generation, and a RAII session that installs on construction
  and erases on destruction.
• extension/llm/cache/CMakeLists.txt — builds the library and pulls in the test folder (cmake).
• extension/llm/cache/test/cache_test.cpp — unit tests for the above.
• extension/llm/cache/test/CMakeLists.txt — builds the test (cmake).
• CMakeLists.txt — includes the new folder in the build.

Testing

Added C++ unit tests (test/cache_test.cpp) covering the registry (install/get/erase, unique keys), the builder
registry (build + missing-kind NotFound error), and the session's install-on-create / erase-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.

⏳ No Failures, 8 Pending

As of commit 40c1666 with merge base 55d693b ():
💚 Looks good so far! There are no failures yet. 💚

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?

[kiymetakdemir]

It is the initial capacity added a comment for this.

[metascroy]

What does this mean?

[kiymetakdemir]

We don't use it actually, I removed.

[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?

[kiymetakdemir]

Applied this together with claude's suggestion on explicit as_planner()/as_control().

[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

[metascroy]

@claude review this code

[claude]

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

---

Code Review

• Read the cache core, registry, and tests
• Review build wiring (cmake)
• Assess against reviewer feedback & follow-up changes
• Post review

The revision addresses the major points from the earlier round cleanly. Quick status on prior feedback, then a few remaining nits.

Resolved since last review ✅

• Exceptions → ET_CHECK_MSG (Add support for quantized LeakyReLU #1): No more throw std::runtime_error. rewind/plan/build now use ET_CHECK_MSG, so the core is -fno-exceptions-safe. 👍
• Interface split (Re-sync with internal repository #2):CacheBase + SequenceControl + SequencePlanner with as_control()/as_planner(). The three bases are independent (not virtual-inheriting from CacheBase), so recovery is a static upcast returning this — no diamond, no RTTI. This is a cleaner solution than the dynamic_cast path floated earlier.
• Naming (Add unlifting pass under private config #4):ContiguousBookkeeping → FlatCache, paired conceptually with a future RingCache. Reads as a policy, not a layout claim.
• min_chunk / contiguous (Rename _pt2e to pt2e #3):min_chunk now documented as initial pool size; the ambiguous contiguous field is gone from SeqStepPlan.
• buck (Re-sync with internal repository #6): Went CMake-only. The new-core diff is 7 files, no .bzl/BUCK — matches your "buckify later" call.

Remaining items

1. Stale doc comment in cache_registry.h:63

// Throws if no builder is registered for (backend_id, kind).
std::shared_ptr<CacheBase> build(...) const;

It no longer throws — build() uses ET_CHECK_MSG (aborts). Update the comment to match, otherwise a reader will expect a catchable exception. Fix this →

2. plan() mutates state (cache.h:122)
plan() advances length_ as a side effect, so it both computes and commits the step. Worth a one-line comment on the method that "plan commits the length" — a backend that expects plan() to be a pure query (e.g. to peek a layout before deciding) would be surprised. If a peek-vs-commit split is ever needed, this is the seam.

3. valid_len semantics on re-plan / overwrite (cache.h:132-135)
valid_len = position + T, and length_ only grows (if (end > length_) length_ = end). If a caller re-plans at a position < length_ (e.g. overwrite after a rewind), valid_len reflects position + T, not the previously-tracked length_. That's a reasonable "read bound = what you just wrote" contract, but it's implicit — a comment stating that plan returns the read bound for this step (not necessarily the full tracked length) would remove ambiguity for the ring/tree planners that follow.

Minor

• Default arg on a virtual (cache.h:44):can_extend(int n = 1) — default args on virtuals bind to the static type, a known footgun if a future override changes the default. Harmless with one impl today; just flagging before more caches inherit the face.
• FlatCache is not internally thread-safe. The registry is mutex-guarded, but the cache object's length_ isn't. Fine for the single-sequence generation-loop model; worth an explicit "single-threaded per sequence" note if that assumption isn't documented elsewhere.
• Builder-registry test isolation:register_builder("TestBackend", ...) still mutates the process-global and persists across cases in the binary. Serviceable at this size (noted last round); keep an eye as more kinds register.
• int for capacity/length/position throughout — fine for sequence lengths; be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions from the last review are settled well, and the as_control()/as_planner() recovery is a nice RTTI-free resolution. Nothing blocking — items above are a stale comment plus a couple of clarifying comments on plan()/valid_len semantics.
· kv-cache-shared-registry

[metascroy]

I think initial_capacity is a better name?

[metascroy]

What does ET_CHECK_MSG do on failure? Call std::terminate?

[kiymetakdemir]

It calls std::abort()

[metascroy]

Are you planning to use this for Ring as well?

[kiymetakdemir]

Yes, I guess in that case it needs one more field for the window's start position.

[metascroy]

Looks great! Main concern is calling std::abort. Instead wrap in Result instead.

See how ET runtime / delegates avoids exception without calling abort.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 31s —— View job

---

Code Review

• Read the cache core, registry, and tests
• Verify the abort→Result change requested last round
• Review build wiring (CMake)
• Post review

The last blocker — std::abort via ET_CHECK_MSG — is fully resolved. This round is clean; only minor hardening notes remain.

Resolved since last review ✅

• abort → Result/Error (main concern):plan() and build() now return Result<T> and rewind() returns Error, all via ET_CHECK_OR_RETURN_ERROR, which logs and returns rather than aborting (runtime/core/error.h:182). No std::abort/std::terminate on any error path — this matches how the ET runtime/delegates avoid exceptions. 👍
• Death tests → Result checks:cache_test.cpp now asserts .error() == Error::InvalidArgument and .ok() instead of EXPECT_DEATH. Matches the requested pattern.
• min_chunk → initial_capacity: renamed and documented as the lazy-doubling pool's starting size (cache.h:84-87).
• Interface split + RTTI-free recovery:CacheBase / SequenceControl / SequencePlanner with as_control()/as_planner() returning this — static upcast, no dynamic_cast.
• buck: CMake-only; no .bzl/BUCK in the diff, matching the "buckify later" call.

Minor / non-blocking

1. SeqStepPlan fields have no default initializers (cache.h:57-61)

structSeqStepPlan {
int valid_len;
int write_start;
int write_len;
};

FlatCache::plan aggregate-inits all three, so it's safe today. But once ring/tree planners return partially-filled plans, an uninitialized field reads as garbage. The proposal in the review thread had = 0 defaults — worth keeping them here to stay safe as planners multiply. Fix this →

2. plan() / rewind() don't guard negative inputs (cache.h:117,130)
plan(position, T) checks end <= capacity_ but not position >= 0 || T >= 0; a negative position yields a negative write_start. rewind(new_len) checks new_len <= length_ but not new_len >= 0, so rewind(-1) sets length_ = -1 and silently corrupts can_extend. Both are internal-API misuse rather than user input, but a lower-bound ET_CHECK_OR_RETURN_ERROR is cheap insurance.

3. plan() commits length as a side effect (cache.h:138-140)
plan() both computes the step and advances length_. A one-line comment noting "plan commits the length" would help a future planner that expects a pure peek before deciding a layout. (Raised last round; still worth the comment.)

4. Carried-forward nits (all still fine at this size):

• can_extend(int n = 1) — default arg on a virtual binds to the static type; harmless with one impl.
• FlatCache isn't internally thread-safe (length_ unguarded); fine under the single-sequence generation-loop model — worth an explicit note.
• Builder-registry test mutates the process-global (register_builder("TestBackend", …)) and persists across cases.
• int for capacity/length/position — be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions are settled and the error-handling convention is now consistent with the rest of the runtime. Nothing blocking — the only one I'd fold in before backends land is the SeqStepPlan default initializers (#1), since those plans get consumed by byte-layer code shortly.
· kv-cache-shared-registry

[metascroy]

Let's keep gemma4 in mind for the first model we enable here, which has a mix of flat/ring, which I think must be captured in config. I think both must be driven from the same controller.

What if we do:

struct LayerPolicy {
enum class Kind : int { Flat = 0, Ring = 1 }; // serialized values: append-only
Kind kind = Kind::Flat;
int window = 0; // Ring: sliding-window size; must be 0 when Flat
};
struct LayerConfig {
LayerPolicy policy; // default Flat
int n_kv_heads;
int head_dim;
// dtype / scale / sink size land here later (not this PR)
};
struct CacheConfig {
int capacity; // logical cap
int n_layers;
std::vector<LayerConfig> layers; // size 1 == uniform, else == n_layers
int initial_capacity = 512;
};

[digantdesai]

Also see CacheLayerMixIn from hf/transformers

[metascroy]

Thinking ahead to mixed gemma4 model, maybe we have one SequenceCache, with flat/ring being policies. We need one control driver for both caches across all layers I think. Something like this:

class SequenceCache : public CacheBase,
public SequenceControl,
public SequencePlanner {
public:
// CacheBase: static upcasts, no RTTI (each returns `this`).
SequenceControl* as_control() override { return this; }
SequencePlanner* as_planner() override { return this; }
// SequenceControl (runner face): length / admission / rewind / clear.
bool can_extend(int n = 1) const override; // min over policies' retention
int capacity() const override { return capacity_; }
Error rewind(int new_len) override; // bounded by smallest window
void clear() override { length_ = 0; }
// SequencePlanner (backend face): dispatch to the layout policy.
Result<SeqStepPlan> plan(int policy_id, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(
policy_id >= 0 && policy_id < (int)policies_.size(),
InvalidArgument, "bad policy_id %d", policy_id);
return policies_[policy_id]->plan(position, T);
}
private:
int capacity_;
int length_ = 0;
std::vector<std::unique_ptr<LayoutPolicy>> policies_; // FlatPolicy / RingPolicy
std::vector<int> layer_to_policy_; // built from CacheConfig
};

Here LayoutPolicy is something like:

class LayoutPolicy {
public:
virtual ~LayoutPolicy() = default;
virtual SeqStepPlan plan(int position, int T) const = 0; // const == no length side effect
virtual int depth() const = 0; // physical rows to allocate
};

with FlatPolicy having the logic you have for FlatCache here, e.g.,

class FlatPolicy final : public LayoutPolicy {
public:
explicit FlatPolicy(int capacity) : capacity_(capacity) {}
int depth() const override { return capacity_; }
SeqStepPlan plan(int position, int T) const override {
const int end = position + T;
SeqStepPlan p{};
p.n_write = 1; p.write[0] = {position, T}; // contiguous, never wraps
p.n_read = 1; p.read[0] = {0, end}; // attend over all history
return p;
}
private:
int capacity_;
};

Although this diff is flat cache, we need SeqStepPlan to work for both. A ring can have 1 or 2 contiguous runs, a flat cache will always have 1:

struct Run { int start; int len; }; // a contiguous span of physical rows
struct SeqStepPlan {
Run write[2]; int n_write; // where to place this step's new KV
Run read[2]; int n_read; // valid history to attend over, oldest→newest
};

[metascroy]

@claude review again in light of latest comments on policy

[claude]

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

---

Review — the gemma4 mixed flat/ring policy proposal

• Read cache core, registry, tests
• Assess current code against the mixed flat/ring policy design
• Post review

I read the current cache.h/cache_registry.* against your two latest comments (the CacheConfig/LayerPolicy shape at cache.h:79, and the single-SequenceCache-with-LayoutPolicy shape at cache.h:92). I think the direction is right — below is where I'd push on the details before it lands, since these are the surfaces backends will bind to.

The core call: one controller, per-layer policies ✅

Agreed, and I'd go further — for a mixed model this isn't just nicer, the separate-class design is actually wrong. Gemma-style alternating local/global attention means flat and ring layers coexist in one model, but there is exactly one logical sequence length across all layers. Two independent FlatCache/RingCache objects would each track their own length_, and can_extend/rewind would no longer be coherent for the sequence as a whole. A single SequenceCache that owns length_ and dispatches layout to per-layer policies is the only thing that keeps admission/rewind consistent. So the SequenceCache + vector<LayoutPolicy> + layer_to_policy_ shape is the one to build.

1. Dispatch key: prefer layer over policy_id

plan(int policy_id, int position, int T) leaks an internal id to the backend. The byte layer knows which layer it's materializing, not which policy bucket that maps to. Keeping the mapping private is cleaner:

Result<SeqStepPlan> plan(int layer, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(layer >= 0 && layer < (int)layer_to_policy_.size(),
InvalidArgument, "bad layer %d", layer);
return policies_[layer_to_policy_[layer]]->plan(position, T);
}

layer_to_policy_ is already in your sketch — this just makes it the public parameter and keeps policy_id an implementation detail.

2. Length commit must be idempotent (plan runs N_layers times per step)

Making LayoutPolicy::planconst (pure layout, no side effect) is the right resolution to the "plan commits length" concern from earlier rounds — the controller owns length now. But watch the seam: the backend calls plan()once per layer per step, so SequenceCache must not advance length_ additively inside plan. Idempotent length_ = max(length_, position + T) is safe to call every layer; an additive length_ += T would multiply by n_layers. Either document that plan's length commit is idempotent, or split it out into a single control()->advance(T) the runner calls once per step and keep plan fully read-only. I'd lean to the explicit advance — it makes "one commit per step" impossible to get wrong.

3. can_extend / rewind semantics across mixed policies

Your one-liners ("min over policies' retention", "bounded by smallest window") are the correct contract, but they carry real logic worth stating in code:

• can_extend: a ring layer never refuses (it evicts), so admission is gated only by the flat/full layers → effectively bounded by capacity. Ring layers contribute no bound. Fine, but that means can_extend is really "any non-evicting policy near its cap," so the controller needs each policy to expose whether it evicts.
• rewind(new_len): this is the subtle one. A ring layer with window W has physically evicted everything older than length_ - W, so rewind to a new_len < length_ - min_window is unrecoverable — that history is gone from the ring layers even though the flat layers still have it. The bound isn't just new_len <= length_; it's new_len >= length_ - smallest_window too. That lower bound needs the window sizes in the controller, so LayoutPolicy should expose retained_from(length) (or similar) alongside depth().

4. SeqStepPlan: read[] runs are needed, but so is the read base position

Reintroducing Run read[2]; int n_read is correct — for a ring cache the valid history is not[0, length), it wraps, so a single valid_len can't describe it (that's why the earlier flat-only valid_len field doesn't generalize). Two points:

• Run write[2]/read[2] caps at 2 runs, which is exactly right for flat (1) and ring (≤2). Make it explicit in a comment that SeqStepPlan is deliberately the flat∪ring shape and that paged/tree planners return their own plan type (you noted this earlier) — otherwise [2] reads as an arbitrary limit.
• Your first proposal (the 18:10 sketch) had read_base_pos — the logical position of read[0].start. The latest slimmed SeqStepPlan dropped it. For ring, the backend needs the logical positions of the read runs to apply RoPE and the attention mask correctly (the oldest retained row is not at logical position 0 after eviction). Unless positions are reconstructed elsewhere, I'd bring read_base_pos back (or a per-run base) — dropping it looks like it'll break RoPE on the ring path.

5. Config shape

The LayerPolicy/LayerConfig/CacheConfig nesting is good and folds the current parallel n_kv_heads/head_dim vectors into an array-of-structs, which is cleaner. Keep the "size 1 == uniform, else == n_layers" convention you already have for the head vectors — most models are uniform and shouldn't pay for a full per-layer list. Two small things:

• Keep the enum class Kind : int with the "serialized values: append-only" comment — good instinct if this ever hits .pte metadata.
• State the window == 0 iff Flat invariant as an ET_CHECK_OR_RETURN_ERROR in whatever validates the config, not just a comment, since it's now load-bearing for depth()/rewind.

Scope: adopt at least the shapes in this PR

You've already approved, so this is a "now vs. next PR" judgment. My take: the internal SequenceCache/LayoutPolicy refactor is cheapest to do now, while FlatPolicy is the only implementation to migrate — the same argument that justified doing the control/planner split in this PR rather than deferring it. At minimum, land the CacheConfig/LayerConfig and SeqStepPlan (write[2]/read[2] + read_base_pos) shapes now, since those are the surfaces the byte layer and any serialized metadata will bind to and are the expensive ones to change later. The FlatCache→SequenceCache+FlatPolicy mechanical move can follow if you'd rather keep this PR small, but the config/plan structs I'd fix here.

Nothing above contradicts the approval — it's design hardening for the mixed-model case you're aiming at. The two I'd treat as near-blocking for the refactor are #2 (idempotent length commit) and #4 (read_base_pos for ring RoPE), since both are silent-correctness traps rather than cosmetic.
· kv-cache-shared-registry

[digantdesai]

can we wrap the base class in ET runtime friendly cache class? This is to avoid dependence on ET from runner side. I.e. I want this to be useful even when using outside ET runner. Or we can split this file into two, runner facing which can use ET stuff because for non-ET runner we can have another class. And the backend cache interface, which shouldn't use ET stuff.

Code location is also a problem, but we can discuss that later, and I don't want to block on that one.

[kiymetakdemir]

Sure! After looking into it, I'm leaning toward implementing an ET wrapper so the whole core stays a single, ET-free implementation.

[digantdesai]

Ok sorry I wasn't very clear, given I didn't see this PR.

To rephrase, let's keep the ET interfaces separate (and can be used elsewhere outside of 'et::'). The rationale is we might want to use this in a new backend layer which is not an ET delegate or linked against ET runtime libs.

For now the ET specific registry design seems OK.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 25s —— View job

---

Code Review

• Read the current cache core, registry, and tests
• Review build wiring (CMake) and test gating
• Assess against the latest ET-independence feedback (digantdesai)
• Post review

The PR has been trimmed substantially since the last round. cache.h is now just the ET-free CacheBase anchor + CacheConfig — the SequenceControl/SequencePlanner/FlatCache/SeqStepPlan machinery from earlier rounds is gone, deferred to the follow-up that lands the concrete cache. That's a clean response to @digantdesai's request to keep the neutral core usable outside an ET runner.

Resolved / consistent with feedback ✅

• ET-free core (digantdesai):cache.h has zero ExecuTorch includes and no et:: dependency — it's a plain CacheBase + CacheConfig. The ET-specific Error/Result usage is confined to cache_registry.*, and the header comments state the split explicitly. This matches "keep the ET interfaces separate so this can be used in a non-delegate backend layer."
• abort → Result/Error:build() returns Result<...> via ET_CHECK_OR_RETURN_ERROR (logs + returns, no std::abort), and the test asserts .error() == Error::NotFound rather than a death test. Consistent with the resolution from the earlier rounds.
• CMake-only, no buck: matches the "buckify later" call. Test gating uses if(BUILD_TESTING), consistent with the other extension/*/CMakeLists.txt.

Observations on the current (slimmed) surface

1. CacheConfig fields are uninitialized (cache.h:31-35)

structCacheConfig {
int capacity;
int n_layers;
int initial_capacity = 512;
};

capacity and n_layers have no default initializer. The test aggregate-inits both (CacheConfig cfg{32, 1}), so it's safe there, but a caller doing CacheConfig cfg; gets garbage in the two most important fields. Given initial_capacity already has a default, adding = 0 to the other two keeps the struct uniformly safe to default-construct. Fix this →

2. The gemma4 mixed flat/ring design (from the last round) isn't in this PR
The LayerPolicy/LayerConfig per-layer config and the SeqStepPlan (write[2]/read[2] + read_base_pos) shapes discussed at length aren't here — CacheConfig is still the flat capacity/n_layers/initial_capacity triple. That's fine as a scoping decision (this PR is now just the registry + neutral anchor), but worth confirming that's intentional: my earlier note stands that the config and step-plan structs are the expensive surfaces to change later because the byte layer and any serialized metadata bind to them. If gemma4 is the first target, folding at least the LayerConfig array-of-structs into CacheConfig now is cheaper than after consumers exist. Not blocking for a registry-only PR, but flag it so it's a deliberate "next PR" rather than an oversight.

3. Builder-registry key concatenation (cache_registry.cpp:53,63)
The map key is backend_id + ":" + kind. If either token can contain :, "a:b" + ":" + "c" and "a" + ":" + "b:c" collide. Backend ids/kinds are internal constants today so it's theoretical, but a one-line note (or a std::pair key) removes the ambiguity if these ever become user-facing strings.

Minor / carried forward (all still fine at this size)

• Process-global singletons + test isolation:register_builder("TestBackend", "stub", ...) mutates the process-global CacheBuilderRegistry and persists for the test binary's lifetime. Serviceable at one test; watch it as more kinds register.
• int for capacity/layers: fine for these counts; be deliberate about widening if any of these ever feed byte-offset math in the pool layer.
• Registry comment (cache_registry.h:69): "Returns Error::NotFound" now correctly matches the Result return — good, the stale "throws" comment from earlier rounds is fixed.

Overall: the scope reduction is the right move and directly addresses the ET-independence concern. Nothing blocking. The one I'd fold in here is the CacheConfig default initializers (#1); the mixed-policy config/plan shapes (#2) are a scoping judgment call for the follow-up.
· kv-cache-shared-registry

[metascroy]

Use std::pair, other than that, LGTM! Let's land, we can always fix later

[kiymetakdemir]

Summary

Adds a shared, backend-agnostic registry for the off-graph KV cache: a small set of C++ pieces that let a runtime own a cache and look it up by key.

Files

• extension/llm/cache/cache.h — the opaque cache handle (CacheBase) and the CacheConfig struct.
• extension/llm/cache/cache_registry.h / cache_registry.cpp — a process-global registry that stores caches by key, a
  per-(backend_id, kind) builder registry, unique-key generation, and a RAII session that installs on construction
  and erases on destruction.
• extension/llm/cache/CMakeLists.txt — builds the library and pulls in the test folder (cmake).
• extension/llm/cache/test/cache_test.cpp — unit tests for the above.
• extension/llm/cache/test/CMakeLists.txt — builds the test (cmake).
• CMakeLists.txt — includes the new folder in the build.

Testing

Added C++ unit tests (test/cache_test.cpp) covering the registry (install/get/erase, unique keys), the builder
registry (build + missing-kind NotFound error), and the session's install-on-create / erase-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.

⏳ No Failures, 8 Pending

As of commit 40c1666 with merge base 55d693b ():
💚 Looks good so far! There are no failures yet. 💚

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?

[kiymetakdemir]

It is the initial capacity added a comment for this.

[metascroy]

What does this mean?

[kiymetakdemir]

We don't use it actually, I removed.

[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?

[kiymetakdemir]

Applied this together with claude's suggestion on explicit as_planner()/as_control().

[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

[metascroy]

@claude review this code

[claude]

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

---

Code Review

• Read the cache core, registry, and tests
• Review build wiring (cmake)
• Assess against reviewer feedback & follow-up changes
• Post review

The revision addresses the major points from the earlier round cleanly. Quick status on prior feedback, then a few remaining nits.

Resolved since last review ✅

• Exceptions → ET_CHECK_MSG (Add support for quantized LeakyReLU #1): No more throw std::runtime_error. rewind/plan/build now use ET_CHECK_MSG, so the core is -fno-exceptions-safe. 👍
• Interface split (Re-sync with internal repository #2):CacheBase + SequenceControl + SequencePlanner with as_control()/as_planner(). The three bases are independent (not virtual-inheriting from CacheBase), so recovery is a static upcast returning this — no diamond, no RTTI. This is a cleaner solution than the dynamic_cast path floated earlier.
• Naming (Add unlifting pass under private config #4):ContiguousBookkeeping → FlatCache, paired conceptually with a future RingCache. Reads as a policy, not a layout claim.
• min_chunk / contiguous (Rename _pt2e to pt2e #3):min_chunk now documented as initial pool size; the ambiguous contiguous field is gone from SeqStepPlan.
• buck (Re-sync with internal repository #6): Went CMake-only. The new-core diff is 7 files, no .bzl/BUCK — matches your "buckify later" call.

Remaining items

1. Stale doc comment in cache_registry.h:63

// Throws if no builder is registered for (backend_id, kind).
std::shared_ptr<CacheBase> build(...) const;

It no longer throws — build() uses ET_CHECK_MSG (aborts). Update the comment to match, otherwise a reader will expect a catchable exception. Fix this →

2. plan() mutates state (cache.h:122)
plan() advances length_ as a side effect, so it both computes and commits the step. Worth a one-line comment on the method that "plan commits the length" — a backend that expects plan() to be a pure query (e.g. to peek a layout before deciding) would be surprised. If a peek-vs-commit split is ever needed, this is the seam.

3. valid_len semantics on re-plan / overwrite (cache.h:132-135)
valid_len = position + T, and length_ only grows (if (end > length_) length_ = end). If a caller re-plans at a position < length_ (e.g. overwrite after a rewind), valid_len reflects position + T, not the previously-tracked length_. That's a reasonable "read bound = what you just wrote" contract, but it's implicit — a comment stating that plan returns the read bound for this step (not necessarily the full tracked length) would remove ambiguity for the ring/tree planners that follow.

Minor

• Default arg on a virtual (cache.h:44):can_extend(int n = 1) — default args on virtuals bind to the static type, a known footgun if a future override changes the default. Harmless with one impl today; just flagging before more caches inherit the face.
• FlatCache is not internally thread-safe. The registry is mutex-guarded, but the cache object's length_ isn't. Fine for the single-sequence generation-loop model; worth an explicit "single-threaded per sequence" note if that assumption isn't documented elsewhere.
• Builder-registry test isolation:register_builder("TestBackend", ...) still mutates the process-global and persists across cases in the binary. Serviceable at this size (noted last round); keep an eye as more kinds register.
• int for capacity/length/position throughout — fine for sequence lengths; be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions from the last review are settled well, and the as_control()/as_planner() recovery is a nice RTTI-free resolution. Nothing blocking — items above are a stale comment plus a couple of clarifying comments on plan()/valid_len semantics.
· kv-cache-shared-registry

[metascroy]

I think initial_capacity is a better name?

[metascroy]

What does ET_CHECK_MSG do on failure? Call std::terminate?

[kiymetakdemir]

It calls std::abort()

[metascroy]

Are you planning to use this for Ring as well?

[kiymetakdemir]

Yes, I guess in that case it needs one more field for the window's start position.

[metascroy]

Looks great! Main concern is calling std::abort. Instead wrap in Result instead.

See how ET runtime / delegates avoids exception without calling abort.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 31s —— View job

---

Code Review

• Read the cache core, registry, and tests
• Verify the abort→Result change requested last round
• Review build wiring (CMake)
• Post review

The last blocker — std::abort via ET_CHECK_MSG — is fully resolved. This round is clean; only minor hardening notes remain.

Resolved since last review ✅

• abort → Result/Error (main concern):plan() and build() now return Result<T> and rewind() returns Error, all via ET_CHECK_OR_RETURN_ERROR, which logs and returns rather than aborting (runtime/core/error.h:182). No std::abort/std::terminate on any error path — this matches how the ET runtime/delegates avoid exceptions. 👍
• Death tests → Result checks:cache_test.cpp now asserts .error() == Error::InvalidArgument and .ok() instead of EXPECT_DEATH. Matches the requested pattern.
• min_chunk → initial_capacity: renamed and documented as the lazy-doubling pool's starting size (cache.h:84-87).
• Interface split + RTTI-free recovery:CacheBase / SequenceControl / SequencePlanner with as_control()/as_planner() returning this — static upcast, no dynamic_cast.
• buck: CMake-only; no .bzl/BUCK in the diff, matching the "buckify later" call.

Minor / non-blocking

1. SeqStepPlan fields have no default initializers (cache.h:57-61)

structSeqStepPlan {
int valid_len;
int write_start;
int write_len;
};

FlatCache::plan aggregate-inits all three, so it's safe today. But once ring/tree planners return partially-filled plans, an uninitialized field reads as garbage. The proposal in the review thread had = 0 defaults — worth keeping them here to stay safe as planners multiply. Fix this →

2. plan() / rewind() don't guard negative inputs (cache.h:117,130)
plan(position, T) checks end <= capacity_ but not position >= 0 || T >= 0; a negative position yields a negative write_start. rewind(new_len) checks new_len <= length_ but not new_len >= 0, so rewind(-1) sets length_ = -1 and silently corrupts can_extend. Both are internal-API misuse rather than user input, but a lower-bound ET_CHECK_OR_RETURN_ERROR is cheap insurance.

3. plan() commits length as a side effect (cache.h:138-140)
plan() both computes the step and advances length_. A one-line comment noting "plan commits the length" would help a future planner that expects a pure peek before deciding a layout. (Raised last round; still worth the comment.)

4. Carried-forward nits (all still fine at this size):

• can_extend(int n = 1) — default arg on a virtual binds to the static type; harmless with one impl.
• FlatCache isn't internally thread-safe (length_ unguarded); fine under the single-sequence generation-loop model — worth an explicit note.
• Builder-registry test mutates the process-global (register_builder("TestBackend", …)) and persists across cases.
• int for capacity/length/position — be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions are settled and the error-handling convention is now consistent with the rest of the runtime. Nothing blocking — the only one I'd fold in before backends land is the SeqStepPlan default initializers (#1), since those plans get consumed by byte-layer code shortly.
· kv-cache-shared-registry

[metascroy]

Let's keep gemma4 in mind for the first model we enable here, which has a mix of flat/ring, which I think must be captured in config. I think both must be driven from the same controller.

What if we do:

struct LayerPolicy {
enum class Kind : int { Flat = 0, Ring = 1 }; // serialized values: append-only
Kind kind = Kind::Flat;
int window = 0; // Ring: sliding-window size; must be 0 when Flat
};
struct LayerConfig {
LayerPolicy policy; // default Flat
int n_kv_heads;
int head_dim;
// dtype / scale / sink size land here later (not this PR)
};
struct CacheConfig {
int capacity; // logical cap
int n_layers;
std::vector<LayerConfig> layers; // size 1 == uniform, else == n_layers
int initial_capacity = 512;
};

[digantdesai]

Also see CacheLayerMixIn from hf/transformers

[metascroy]

Thinking ahead to mixed gemma4 model, maybe we have one SequenceCache, with flat/ring being policies. We need one control driver for both caches across all layers I think. Something like this:

class SequenceCache : public CacheBase,
public SequenceControl,
public SequencePlanner {
public:
// CacheBase: static upcasts, no RTTI (each returns `this`).
SequenceControl* as_control() override { return this; }
SequencePlanner* as_planner() override { return this; }
// SequenceControl (runner face): length / admission / rewind / clear.
bool can_extend(int n = 1) const override; // min over policies' retention
int capacity() const override { return capacity_; }
Error rewind(int new_len) override; // bounded by smallest window
void clear() override { length_ = 0; }
// SequencePlanner (backend face): dispatch to the layout policy.
Result<SeqStepPlan> plan(int policy_id, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(
policy_id >= 0 && policy_id < (int)policies_.size(),
InvalidArgument, "bad policy_id %d", policy_id);
return policies_[policy_id]->plan(position, T);
}
private:
int capacity_;
int length_ = 0;
std::vector<std::unique_ptr<LayoutPolicy>> policies_; // FlatPolicy / RingPolicy
std::vector<int> layer_to_policy_; // built from CacheConfig
};

Here LayoutPolicy is something like:

class LayoutPolicy {
public:
virtual ~LayoutPolicy() = default;
virtual SeqStepPlan plan(int position, int T) const = 0; // const == no length side effect
virtual int depth() const = 0; // physical rows to allocate
};

with FlatPolicy having the logic you have for FlatCache here, e.g.,

class FlatPolicy final : public LayoutPolicy {
public:
explicit FlatPolicy(int capacity) : capacity_(capacity) {}
int depth() const override { return capacity_; }
SeqStepPlan plan(int position, int T) const override {
const int end = position + T;
SeqStepPlan p{};
p.n_write = 1; p.write[0] = {position, T}; // contiguous, never wraps
p.n_read = 1; p.read[0] = {0, end}; // attend over all history
return p;
}
private:
int capacity_;
};

Although this diff is flat cache, we need SeqStepPlan to work for both. A ring can have 1 or 2 contiguous runs, a flat cache will always have 1:

struct Run { int start; int len; }; // a contiguous span of physical rows
struct SeqStepPlan {
Run write[2]; int n_write; // where to place this step's new KV
Run read[2]; int n_read; // valid history to attend over, oldest→newest
};

[metascroy]

@claude review again in light of latest comments on policy

[claude]

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

---

Review — the gemma4 mixed flat/ring policy proposal

• Read cache core, registry, tests
• Assess current code against the mixed flat/ring policy design
• Post review

I read the current cache.h/cache_registry.* against your two latest comments (the CacheConfig/LayerPolicy shape at cache.h:79, and the single-SequenceCache-with-LayoutPolicy shape at cache.h:92). I think the direction is right — below is where I'd push on the details before it lands, since these are the surfaces backends will bind to.

The core call: one controller, per-layer policies ✅

Agreed, and I'd go further — for a mixed model this isn't just nicer, the separate-class design is actually wrong. Gemma-style alternating local/global attention means flat and ring layers coexist in one model, but there is exactly one logical sequence length across all layers. Two independent FlatCache/RingCache objects would each track their own length_, and can_extend/rewind would no longer be coherent for the sequence as a whole. A single SequenceCache that owns length_ and dispatches layout to per-layer policies is the only thing that keeps admission/rewind consistent. So the SequenceCache + vector<LayoutPolicy> + layer_to_policy_ shape is the one to build.

1. Dispatch key: prefer layer over policy_id

plan(int policy_id, int position, int T) leaks an internal id to the backend. The byte layer knows which layer it's materializing, not which policy bucket that maps to. Keeping the mapping private is cleaner:

Result<SeqStepPlan> plan(int layer, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(layer >= 0 && layer < (int)layer_to_policy_.size(),
InvalidArgument, "bad layer %d", layer);
return policies_[layer_to_policy_[layer]]->plan(position, T);
}

layer_to_policy_ is already in your sketch — this just makes it the public parameter and keeps policy_id an implementation detail.

2. Length commit must be idempotent (plan runs N_layers times per step)

Making LayoutPolicy::planconst (pure layout, no side effect) is the right resolution to the "plan commits length" concern from earlier rounds — the controller owns length now. But watch the seam: the backend calls plan()once per layer per step, so SequenceCache must not advance length_ additively inside plan. Idempotent length_ = max(length_, position + T) is safe to call every layer; an additive length_ += T would multiply by n_layers. Either document that plan's length commit is idempotent, or split it out into a single control()->advance(T) the runner calls once per step and keep plan fully read-only. I'd lean to the explicit advance — it makes "one commit per step" impossible to get wrong.

3. can_extend / rewind semantics across mixed policies

Your one-liners ("min over policies' retention", "bounded by smallest window") are the correct contract, but they carry real logic worth stating in code:

• can_extend: a ring layer never refuses (it evicts), so admission is gated only by the flat/full layers → effectively bounded by capacity. Ring layers contribute no bound. Fine, but that means can_extend is really "any non-evicting policy near its cap," so the controller needs each policy to expose whether it evicts.
• rewind(new_len): this is the subtle one. A ring layer with window W has physically evicted everything older than length_ - W, so rewind to a new_len < length_ - min_window is unrecoverable — that history is gone from the ring layers even though the flat layers still have it. The bound isn't just new_len <= length_; it's new_len >= length_ - smallest_window too. That lower bound needs the window sizes in the controller, so LayoutPolicy should expose retained_from(length) (or similar) alongside depth().

4. SeqStepPlan: read[] runs are needed, but so is the read base position

Reintroducing Run read[2]; int n_read is correct — for a ring cache the valid history is not[0, length), it wraps, so a single valid_len can't describe it (that's why the earlier flat-only valid_len field doesn't generalize). Two points:

• Run write[2]/read[2] caps at 2 runs, which is exactly right for flat (1) and ring (≤2). Make it explicit in a comment that SeqStepPlan is deliberately the flat∪ring shape and that paged/tree planners return their own plan type (you noted this earlier) — otherwise [2] reads as an arbitrary limit.
• Your first proposal (the 18:10 sketch) had read_base_pos — the logical position of read[0].start. The latest slimmed SeqStepPlan dropped it. For ring, the backend needs the logical positions of the read runs to apply RoPE and the attention mask correctly (the oldest retained row is not at logical position 0 after eviction). Unless positions are reconstructed elsewhere, I'd bring read_base_pos back (or a per-run base) — dropping it looks like it'll break RoPE on the ring path.

5. Config shape

The LayerPolicy/LayerConfig/CacheConfig nesting is good and folds the current parallel n_kv_heads/head_dim vectors into an array-of-structs, which is cleaner. Keep the "size 1 == uniform, else == n_layers" convention you already have for the head vectors — most models are uniform and shouldn't pay for a full per-layer list. Two small things:

• Keep the enum class Kind : int with the "serialized values: append-only" comment — good instinct if this ever hits .pte metadata.
• State the window == 0 iff Flat invariant as an ET_CHECK_OR_RETURN_ERROR in whatever validates the config, not just a comment, since it's now load-bearing for depth()/rewind.

Scope: adopt at least the shapes in this PR

You've already approved, so this is a "now vs. next PR" judgment. My take: the internal SequenceCache/LayoutPolicy refactor is cheapest to do now, while FlatPolicy is the only implementation to migrate — the same argument that justified doing the control/planner split in this PR rather than deferring it. At minimum, land the CacheConfig/LayerConfig and SeqStepPlan (write[2]/read[2] + read_base_pos) shapes now, since those are the surfaces the byte layer and any serialized metadata will bind to and are the expensive ones to change later. The FlatCache→SequenceCache+FlatPolicy mechanical move can follow if you'd rather keep this PR small, but the config/plan structs I'd fix here.

Nothing above contradicts the approval — it's design hardening for the mixed-model case you're aiming at. The two I'd treat as near-blocking for the refactor are #2 (idempotent length commit) and #4 (read_base_pos for ring RoPE), since both are silent-correctness traps rather than cosmetic.
· kv-cache-shared-registry

[digantdesai]

can we wrap the base class in ET runtime friendly cache class? This is to avoid dependence on ET from runner side. I.e. I want this to be useful even when using outside ET runner. Or we can split this file into two, runner facing which can use ET stuff because for non-ET runner we can have another class. And the backend cache interface, which shouldn't use ET stuff.

Code location is also a problem, but we can discuss that later, and I don't want to block on that one.

[kiymetakdemir]

Sure! After looking into it, I'm leaning toward implementing an ET wrapper so the whole core stays a single, ET-free implementation.

[digantdesai]

Ok sorry I wasn't very clear, given I didn't see this PR.

To rephrase, let's keep the ET interfaces separate (and can be used elsewhere outside of 'et::'). The rationale is we might want to use this in a new backend layer which is not an ET delegate or linked against ET runtime libs.

For now the ET specific registry design seems OK.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 25s —— View job

---

Code Review

• Read the current cache core, registry, and tests
• Review build wiring (CMake) and test gating
• Assess against the latest ET-independence feedback (digantdesai)
• Post review

The PR has been trimmed substantially since the last round. cache.h is now just the ET-free CacheBase anchor + CacheConfig — the SequenceControl/SequencePlanner/FlatCache/SeqStepPlan machinery from earlier rounds is gone, deferred to the follow-up that lands the concrete cache. That's a clean response to @digantdesai's request to keep the neutral core usable outside an ET runner.

Resolved / consistent with feedback ✅

• ET-free core (digantdesai):cache.h has zero ExecuTorch includes and no et:: dependency — it's a plain CacheBase + CacheConfig. The ET-specific Error/Result usage is confined to cache_registry.*, and the header comments state the split explicitly. This matches "keep the ET interfaces separate so this can be used in a non-delegate backend layer."
• abort → Result/Error:build() returns Result<...> via ET_CHECK_OR_RETURN_ERROR (logs + returns, no std::abort), and the test asserts .error() == Error::NotFound rather than a death test. Consistent with the resolution from the earlier rounds.
• CMake-only, no buck: matches the "buckify later" call. Test gating uses if(BUILD_TESTING), consistent with the other extension/*/CMakeLists.txt.

Observations on the current (slimmed) surface

1. CacheConfig fields are uninitialized (cache.h:31-35)

structCacheConfig {
int capacity;
int n_layers;
int initial_capacity = 512;
};

capacity and n_layers have no default initializer. The test aggregate-inits both (CacheConfig cfg{32, 1}), so it's safe there, but a caller doing CacheConfig cfg; gets garbage in the two most important fields. Given initial_capacity already has a default, adding = 0 to the other two keeps the struct uniformly safe to default-construct. Fix this →

2. The gemma4 mixed flat/ring design (from the last round) isn't in this PR
The LayerPolicy/LayerConfig per-layer config and the SeqStepPlan (write[2]/read[2] + read_base_pos) shapes discussed at length aren't here — CacheConfig is still the flat capacity/n_layers/initial_capacity triple. That's fine as a scoping decision (this PR is now just the registry + neutral anchor), but worth confirming that's intentional: my earlier note stands that the config and step-plan structs are the expensive surfaces to change later because the byte layer and any serialized metadata bind to them. If gemma4 is the first target, folding at least the LayerConfig array-of-structs into CacheConfig now is cheaper than after consumers exist. Not blocking for a registry-only PR, but flag it so it's a deliberate "next PR" rather than an oversight.

3. Builder-registry key concatenation (cache_registry.cpp:53,63)
The map key is backend_id + ":" + kind. If either token can contain :, "a:b" + ":" + "c" and "a" + ":" + "b:c" collide. Backend ids/kinds are internal constants today so it's theoretical, but a one-line note (or a std::pair key) removes the ambiguity if these ever become user-facing strings.

Minor / carried forward (all still fine at this size)

• Process-global singletons + test isolation:register_builder("TestBackend", "stub", ...) mutates the process-global CacheBuilderRegistry and persists for the test binary's lifetime. Serviceable at one test; watch it as more kinds register.
• int for capacity/layers: fine for these counts; be deliberate about widening if any of these ever feed byte-offset math in the pool layer.
• Registry comment (cache_registry.h:69): "Returns Error::NotFound" now correctly matches the Result return — good, the stale "throws" comment from earlier rounds is fixed.

Overall: the scope reduction is the right move and directly addresses the ET-independence concern. Nothing blocking. The one I'd fold in here is the CacheConfig default initializers (#1); the mixed-policy config/plan shapes (#2) are a scoping judgment call for the follow-up.
· kv-cache-shared-registry

[metascroy]

Use std::pair, other than that, LGTM! Let's land, we can always fix later

[kiymetakdemir]

Summary

Adds a shared, backend-agnostic registry for the off-graph KV cache: a small set of C++ pieces that let a runtime own a cache and look it up by key.

Files

• extension/llm/cache/cache.h — the opaque cache handle (CacheBase) and the CacheConfig struct.
• extension/llm/cache/cache_registry.h / cache_registry.cpp — a process-global registry that stores caches by key, a
  per-(backend_id, kind) builder registry, unique-key generation, and a RAII session that installs on construction
  and erases on destruction.
• extension/llm/cache/CMakeLists.txt — builds the library and pulls in the test folder (cmake).
• extension/llm/cache/test/cache_test.cpp — unit tests for the above.
• extension/llm/cache/test/CMakeLists.txt — builds the test (cmake).
• CMakeLists.txt — includes the new folder in the build.

Testing

Added C++ unit tests (test/cache_test.cpp) covering the registry (install/get/erase, unique keys), the builder
registry (build + missing-kind NotFound error), and the session's install-on-create / erase-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.

⏳ No Failures, 8 Pending

As of commit 40c1666 with merge base 55d693b ():
💚 Looks good so far! There are no failures yet. 💚

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?

[kiymetakdemir]

It is the initial capacity added a comment for this.

[metascroy]

What does this mean?

[kiymetakdemir]

We don't use it actually, I removed.

[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?

[kiymetakdemir]

Applied this together with claude's suggestion on explicit as_planner()/as_control().

[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

[metascroy]

@claude review this code

[claude]

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

---

Code Review

• Read the cache core, registry, and tests
• Review build wiring (cmake)
• Assess against reviewer feedback & follow-up changes
• Post review

The revision addresses the major points from the earlier round cleanly. Quick status on prior feedback, then a few remaining nits.

Resolved since last review ✅

• Exceptions → ET_CHECK_MSG (Add support for quantized LeakyReLU #1): No more throw std::runtime_error. rewind/plan/build now use ET_CHECK_MSG, so the core is -fno-exceptions-safe. 👍
• Interface split (Re-sync with internal repository #2):CacheBase + SequenceControl + SequencePlanner with as_control()/as_planner(). The three bases are independent (not virtual-inheriting from CacheBase), so recovery is a static upcast returning this — no diamond, no RTTI. This is a cleaner solution than the dynamic_cast path floated earlier.
• Naming (Add unlifting pass under private config #4):ContiguousBookkeeping → FlatCache, paired conceptually with a future RingCache. Reads as a policy, not a layout claim.
• min_chunk / contiguous (Rename _pt2e to pt2e #3):min_chunk now documented as initial pool size; the ambiguous contiguous field is gone from SeqStepPlan.
• buck (Re-sync with internal repository #6): Went CMake-only. The new-core diff is 7 files, no .bzl/BUCK — matches your "buckify later" call.

Remaining items

1. Stale doc comment in cache_registry.h:63

// Throws if no builder is registered for (backend_id, kind).
std::shared_ptr<CacheBase> build(...) const;

It no longer throws — build() uses ET_CHECK_MSG (aborts). Update the comment to match, otherwise a reader will expect a catchable exception. Fix this →

2. plan() mutates state (cache.h:122)
plan() advances length_ as a side effect, so it both computes and commits the step. Worth a one-line comment on the method that "plan commits the length" — a backend that expects plan() to be a pure query (e.g. to peek a layout before deciding) would be surprised. If a peek-vs-commit split is ever needed, this is the seam.

3. valid_len semantics on re-plan / overwrite (cache.h:132-135)
valid_len = position + T, and length_ only grows (if (end > length_) length_ = end). If a caller re-plans at a position < length_ (e.g. overwrite after a rewind), valid_len reflects position + T, not the previously-tracked length_. That's a reasonable "read bound = what you just wrote" contract, but it's implicit — a comment stating that plan returns the read bound for this step (not necessarily the full tracked length) would remove ambiguity for the ring/tree planners that follow.

Minor

• Default arg on a virtual (cache.h:44):can_extend(int n = 1) — default args on virtuals bind to the static type, a known footgun if a future override changes the default. Harmless with one impl today; just flagging before more caches inherit the face.
• FlatCache is not internally thread-safe. The registry is mutex-guarded, but the cache object's length_ isn't. Fine for the single-sequence generation-loop model; worth an explicit "single-threaded per sequence" note if that assumption isn't documented elsewhere.
• Builder-registry test isolation:register_builder("TestBackend", ...) still mutates the process-global and persists across cases in the binary. Serviceable at this size (noted last round); keep an eye as more kinds register.
• int for capacity/length/position throughout — fine for sequence lengths; be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions from the last review are settled well, and the as_control()/as_planner() recovery is a nice RTTI-free resolution. Nothing blocking — items above are a stale comment plus a couple of clarifying comments on plan()/valid_len semantics.
· kv-cache-shared-registry

[metascroy]

I think initial_capacity is a better name?

[metascroy]

What does ET_CHECK_MSG do on failure? Call std::terminate?

[kiymetakdemir]

It calls std::abort()

[metascroy]

Are you planning to use this for Ring as well?

[kiymetakdemir]

Yes, I guess in that case it needs one more field for the window's start position.

[metascroy]

Looks great! Main concern is calling std::abort. Instead wrap in Result instead.

See how ET runtime / delegates avoids exception without calling abort.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 31s —— View job

---

Code Review

• Read the cache core, registry, and tests
• Verify the abort→Result change requested last round
• Review build wiring (CMake)
• Post review

The last blocker — std::abort via ET_CHECK_MSG — is fully resolved. This round is clean; only minor hardening notes remain.

Resolved since last review ✅

• abort → Result/Error (main concern):plan() and build() now return Result<T> and rewind() returns Error, all via ET_CHECK_OR_RETURN_ERROR, which logs and returns rather than aborting (runtime/core/error.h:182). No std::abort/std::terminate on any error path — this matches how the ET runtime/delegates avoid exceptions. 👍
• Death tests → Result checks:cache_test.cpp now asserts .error() == Error::InvalidArgument and .ok() instead of EXPECT_DEATH. Matches the requested pattern.
• min_chunk → initial_capacity: renamed and documented as the lazy-doubling pool's starting size (cache.h:84-87).
• Interface split + RTTI-free recovery:CacheBase / SequenceControl / SequencePlanner with as_control()/as_planner() returning this — static upcast, no dynamic_cast.
• buck: CMake-only; no .bzl/BUCK in the diff, matching the "buckify later" call.

Minor / non-blocking

1. SeqStepPlan fields have no default initializers (cache.h:57-61)

structSeqStepPlan {
int valid_len;
int write_start;
int write_len;
};

FlatCache::plan aggregate-inits all three, so it's safe today. But once ring/tree planners return partially-filled plans, an uninitialized field reads as garbage. The proposal in the review thread had = 0 defaults — worth keeping them here to stay safe as planners multiply. Fix this →

2. plan() / rewind() don't guard negative inputs (cache.h:117,130)
plan(position, T) checks end <= capacity_ but not position >= 0 || T >= 0; a negative position yields a negative write_start. rewind(new_len) checks new_len <= length_ but not new_len >= 0, so rewind(-1) sets length_ = -1 and silently corrupts can_extend. Both are internal-API misuse rather than user input, but a lower-bound ET_CHECK_OR_RETURN_ERROR is cheap insurance.

3. plan() commits length as a side effect (cache.h:138-140)
plan() both computes the step and advances length_. A one-line comment noting "plan commits the length" would help a future planner that expects a pure peek before deciding a layout. (Raised last round; still worth the comment.)

4. Carried-forward nits (all still fine at this size):

• can_extend(int n = 1) — default arg on a virtual binds to the static type; harmless with one impl.
• FlatCache isn't internally thread-safe (length_ unguarded); fine under the single-sequence generation-loop model — worth an explicit note.
• Builder-registry test mutates the process-global (register_builder("TestBackend", …)) and persists across cases.
• int for capacity/length/position — be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions are settled and the error-handling convention is now consistent with the rest of the runtime. Nothing blocking — the only one I'd fold in before backends land is the SeqStepPlan default initializers (#1), since those plans get consumed by byte-layer code shortly.
· kv-cache-shared-registry

[metascroy]

Let's keep gemma4 in mind for the first model we enable here, which has a mix of flat/ring, which I think must be captured in config. I think both must be driven from the same controller.

What if we do:

struct LayerPolicy {
enum class Kind : int { Flat = 0, Ring = 1 }; // serialized values: append-only
Kind kind = Kind::Flat;
int window = 0; // Ring: sliding-window size; must be 0 when Flat
};
struct LayerConfig {
LayerPolicy policy; // default Flat
int n_kv_heads;
int head_dim;
// dtype / scale / sink size land here later (not this PR)
};
struct CacheConfig {
int capacity; // logical cap
int n_layers;
std::vector<LayerConfig> layers; // size 1 == uniform, else == n_layers
int initial_capacity = 512;
};

[digantdesai]

Also see CacheLayerMixIn from hf/transformers

[metascroy]

Thinking ahead to mixed gemma4 model, maybe we have one SequenceCache, with flat/ring being policies. We need one control driver for both caches across all layers I think. Something like this:

class SequenceCache : public CacheBase,
public SequenceControl,
public SequencePlanner {
public:
// CacheBase: static upcasts, no RTTI (each returns `this`).
SequenceControl* as_control() override { return this; }
SequencePlanner* as_planner() override { return this; }
// SequenceControl (runner face): length / admission / rewind / clear.
bool can_extend(int n = 1) const override; // min over policies' retention
int capacity() const override { return capacity_; }
Error rewind(int new_len) override; // bounded by smallest window
void clear() override { length_ = 0; }
// SequencePlanner (backend face): dispatch to the layout policy.
Result<SeqStepPlan> plan(int policy_id, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(
policy_id >= 0 && policy_id < (int)policies_.size(),
InvalidArgument, "bad policy_id %d", policy_id);
return policies_[policy_id]->plan(position, T);
}
private:
int capacity_;
int length_ = 0;
std::vector<std::unique_ptr<LayoutPolicy>> policies_; // FlatPolicy / RingPolicy
std::vector<int> layer_to_policy_; // built from CacheConfig
};

Here LayoutPolicy is something like:

class LayoutPolicy {
public:
virtual ~LayoutPolicy() = default;
virtual SeqStepPlan plan(int position, int T) const = 0; // const == no length side effect
virtual int depth() const = 0; // physical rows to allocate
};

with FlatPolicy having the logic you have for FlatCache here, e.g.,

class FlatPolicy final : public LayoutPolicy {
public:
explicit FlatPolicy(int capacity) : capacity_(capacity) {}
int depth() const override { return capacity_; }
SeqStepPlan plan(int position, int T) const override {
const int end = position + T;
SeqStepPlan p{};
p.n_write = 1; p.write[0] = {position, T}; // contiguous, never wraps
p.n_read = 1; p.read[0] = {0, end}; // attend over all history
return p;
}
private:
int capacity_;
};

Although this diff is flat cache, we need SeqStepPlan to work for both. A ring can have 1 or 2 contiguous runs, a flat cache will always have 1:

struct Run { int start; int len; }; // a contiguous span of physical rows
struct SeqStepPlan {
Run write[2]; int n_write; // where to place this step's new KV
Run read[2]; int n_read; // valid history to attend over, oldest→newest
};

[metascroy]

@claude review again in light of latest comments on policy

[claude]

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

---

Review — the gemma4 mixed flat/ring policy proposal

• Read cache core, registry, tests
• Assess current code against the mixed flat/ring policy design
• Post review

I read the current cache.h/cache_registry.* against your two latest comments (the CacheConfig/LayerPolicy shape at cache.h:79, and the single-SequenceCache-with-LayoutPolicy shape at cache.h:92). I think the direction is right — below is where I'd push on the details before it lands, since these are the surfaces backends will bind to.

The core call: one controller, per-layer policies ✅

Agreed, and I'd go further — for a mixed model this isn't just nicer, the separate-class design is actually wrong. Gemma-style alternating local/global attention means flat and ring layers coexist in one model, but there is exactly one logical sequence length across all layers. Two independent FlatCache/RingCache objects would each track their own length_, and can_extend/rewind would no longer be coherent for the sequence as a whole. A single SequenceCache that owns length_ and dispatches layout to per-layer policies is the only thing that keeps admission/rewind consistent. So the SequenceCache + vector<LayoutPolicy> + layer_to_policy_ shape is the one to build.

1. Dispatch key: prefer layer over policy_id

plan(int policy_id, int position, int T) leaks an internal id to the backend. The byte layer knows which layer it's materializing, not which policy bucket that maps to. Keeping the mapping private is cleaner:

Result<SeqStepPlan> plan(int layer, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(layer >= 0 && layer < (int)layer_to_policy_.size(),
InvalidArgument, "bad layer %d", layer);
return policies_[layer_to_policy_[layer]]->plan(position, T);
}

layer_to_policy_ is already in your sketch — this just makes it the public parameter and keeps policy_id an implementation detail.

2. Length commit must be idempotent (plan runs N_layers times per step)

Making LayoutPolicy::planconst (pure layout, no side effect) is the right resolution to the "plan commits length" concern from earlier rounds — the controller owns length now. But watch the seam: the backend calls plan()once per layer per step, so SequenceCache must not advance length_ additively inside plan. Idempotent length_ = max(length_, position + T) is safe to call every layer; an additive length_ += T would multiply by n_layers. Either document that plan's length commit is idempotent, or split it out into a single control()->advance(T) the runner calls once per step and keep plan fully read-only. I'd lean to the explicit advance — it makes "one commit per step" impossible to get wrong.

3. can_extend / rewind semantics across mixed policies

Your one-liners ("min over policies' retention", "bounded by smallest window") are the correct contract, but they carry real logic worth stating in code:

• can_extend: a ring layer never refuses (it evicts), so admission is gated only by the flat/full layers → effectively bounded by capacity. Ring layers contribute no bound. Fine, but that means can_extend is really "any non-evicting policy near its cap," so the controller needs each policy to expose whether it evicts.
• rewind(new_len): this is the subtle one. A ring layer with window W has physically evicted everything older than length_ - W, so rewind to a new_len < length_ - min_window is unrecoverable — that history is gone from the ring layers even though the flat layers still have it. The bound isn't just new_len <= length_; it's new_len >= length_ - smallest_window too. That lower bound needs the window sizes in the controller, so LayoutPolicy should expose retained_from(length) (or similar) alongside depth().

4. SeqStepPlan: read[] runs are needed, but so is the read base position

Reintroducing Run read[2]; int n_read is correct — for a ring cache the valid history is not[0, length), it wraps, so a single valid_len can't describe it (that's why the earlier flat-only valid_len field doesn't generalize). Two points:

• Run write[2]/read[2] caps at 2 runs, which is exactly right for flat (1) and ring (≤2). Make it explicit in a comment that SeqStepPlan is deliberately the flat∪ring shape and that paged/tree planners return their own plan type (you noted this earlier) — otherwise [2] reads as an arbitrary limit.
• Your first proposal (the 18:10 sketch) had read_base_pos — the logical position of read[0].start. The latest slimmed SeqStepPlan dropped it. For ring, the backend needs the logical positions of the read runs to apply RoPE and the attention mask correctly (the oldest retained row is not at logical position 0 after eviction). Unless positions are reconstructed elsewhere, I'd bring read_base_pos back (or a per-run base) — dropping it looks like it'll break RoPE on the ring path.

5. Config shape

The LayerPolicy/LayerConfig/CacheConfig nesting is good and folds the current parallel n_kv_heads/head_dim vectors into an array-of-structs, which is cleaner. Keep the "size 1 == uniform, else == n_layers" convention you already have for the head vectors — most models are uniform and shouldn't pay for a full per-layer list. Two small things:

• Keep the enum class Kind : int with the "serialized values: append-only" comment — good instinct if this ever hits .pte metadata.
• State the window == 0 iff Flat invariant as an ET_CHECK_OR_RETURN_ERROR in whatever validates the config, not just a comment, since it's now load-bearing for depth()/rewind.

Scope: adopt at least the shapes in this PR

You've already approved, so this is a "now vs. next PR" judgment. My take: the internal SequenceCache/LayoutPolicy refactor is cheapest to do now, while FlatPolicy is the only implementation to migrate — the same argument that justified doing the control/planner split in this PR rather than deferring it. At minimum, land the CacheConfig/LayerConfig and SeqStepPlan (write[2]/read[2] + read_base_pos) shapes now, since those are the surfaces the byte layer and any serialized metadata will bind to and are the expensive ones to change later. The FlatCache→SequenceCache+FlatPolicy mechanical move can follow if you'd rather keep this PR small, but the config/plan structs I'd fix here.

Nothing above contradicts the approval — it's design hardening for the mixed-model case you're aiming at. The two I'd treat as near-blocking for the refactor are #2 (idempotent length commit) and #4 (read_base_pos for ring RoPE), since both are silent-correctness traps rather than cosmetic.
· kv-cache-shared-registry

[digantdesai]

can we wrap the base class in ET runtime friendly cache class? This is to avoid dependence on ET from runner side. I.e. I want this to be useful even when using outside ET runner. Or we can split this file into two, runner facing which can use ET stuff because for non-ET runner we can have another class. And the backend cache interface, which shouldn't use ET stuff.

Code location is also a problem, but we can discuss that later, and I don't want to block on that one.

[kiymetakdemir]

Sure! After looking into it, I'm leaning toward implementing an ET wrapper so the whole core stays a single, ET-free implementation.

[digantdesai]

Ok sorry I wasn't very clear, given I didn't see this PR.

To rephrase, let's keep the ET interfaces separate (and can be used elsewhere outside of 'et::'). The rationale is we might want to use this in a new backend layer which is not an ET delegate or linked against ET runtime libs.

For now the ET specific registry design seems OK.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 25s —— View job

---

Code Review

• Read the current cache core, registry, and tests
• Review build wiring (CMake) and test gating
• Assess against the latest ET-independence feedback (digantdesai)
• Post review

The PR has been trimmed substantially since the last round. cache.h is now just the ET-free CacheBase anchor + CacheConfig — the SequenceControl/SequencePlanner/FlatCache/SeqStepPlan machinery from earlier rounds is gone, deferred to the follow-up that lands the concrete cache. That's a clean response to @digantdesai's request to keep the neutral core usable outside an ET runner.

Resolved / consistent with feedback ✅

• ET-free core (digantdesai):cache.h has zero ExecuTorch includes and no et:: dependency — it's a plain CacheBase + CacheConfig. The ET-specific Error/Result usage is confined to cache_registry.*, and the header comments state the split explicitly. This matches "keep the ET interfaces separate so this can be used in a non-delegate backend layer."
• abort → Result/Error:build() returns Result<...> via ET_CHECK_OR_RETURN_ERROR (logs + returns, no std::abort), and the test asserts .error() == Error::NotFound rather than a death test. Consistent with the resolution from the earlier rounds.
• CMake-only, no buck: matches the "buckify later" call. Test gating uses if(BUILD_TESTING), consistent with the other extension/*/CMakeLists.txt.

Observations on the current (slimmed) surface

1. CacheConfig fields are uninitialized (cache.h:31-35)

structCacheConfig {
int capacity;
int n_layers;
int initial_capacity = 512;
};

capacity and n_layers have no default initializer. The test aggregate-inits both (CacheConfig cfg{32, 1}), so it's safe there, but a caller doing CacheConfig cfg; gets garbage in the two most important fields. Given initial_capacity already has a default, adding = 0 to the other two keeps the struct uniformly safe to default-construct. Fix this →

2. The gemma4 mixed flat/ring design (from the last round) isn't in this PR
The LayerPolicy/LayerConfig per-layer config and the SeqStepPlan (write[2]/read[2] + read_base_pos) shapes discussed at length aren't here — CacheConfig is still the flat capacity/n_layers/initial_capacity triple. That's fine as a scoping decision (this PR is now just the registry + neutral anchor), but worth confirming that's intentional: my earlier note stands that the config and step-plan structs are the expensive surfaces to change later because the byte layer and any serialized metadata bind to them. If gemma4 is the first target, folding at least the LayerConfig array-of-structs into CacheConfig now is cheaper than after consumers exist. Not blocking for a registry-only PR, but flag it so it's a deliberate "next PR" rather than an oversight.

3. Builder-registry key concatenation (cache_registry.cpp:53,63)
The map key is backend_id + ":" + kind. If either token can contain :, "a:b" + ":" + "c" and "a" + ":" + "b:c" collide. Backend ids/kinds are internal constants today so it's theoretical, but a one-line note (or a std::pair key) removes the ambiguity if these ever become user-facing strings.

Minor / carried forward (all still fine at this size)

• Process-global singletons + test isolation:register_builder("TestBackend", "stub", ...) mutates the process-global CacheBuilderRegistry and persists for the test binary's lifetime. Serviceable at one test; watch it as more kinds register.
• int for capacity/layers: fine for these counts; be deliberate about widening if any of these ever feed byte-offset math in the pool layer.
• Registry comment (cache_registry.h:69): "Returns Error::NotFound" now correctly matches the Result return — good, the stale "throws" comment from earlier rounds is fixed.

Overall: the scope reduction is the right move and directly addresses the ET-independence concern. Nothing blocking. The one I'd fold in here is the CacheConfig default initializers (#1); the mixed-policy config/plan shapes (#2) are a scoping judgment call for the follow-up.
· kv-cache-shared-registry

[metascroy]

Use std::pair, other than that, LGTM! Let's land, we can always fix later

[kiymetakdemir]

Summary

Adds a shared, backend-agnostic registry for the off-graph KV cache: a small set of C++ pieces that let a runtime own a cache and look it up by key.

Files

• extension/llm/cache/cache.h — the opaque cache handle (CacheBase) and the CacheConfig struct.
• extension/llm/cache/cache_registry.h / cache_registry.cpp — a process-global registry that stores caches by key, a
  per-(backend_id, kind) builder registry, unique-key generation, and a RAII session that installs on construction
  and erases on destruction.
• extension/llm/cache/CMakeLists.txt — builds the library and pulls in the test folder (cmake).
• extension/llm/cache/test/cache_test.cpp — unit tests for the above.
• extension/llm/cache/test/CMakeLists.txt — builds the test (cmake).
• CMakeLists.txt — includes the new folder in the build.

Testing

Added C++ unit tests (test/cache_test.cpp) covering the registry (install/get/erase, unique keys), the builder
registry (build + missing-kind NotFound error), and the session's install-on-create / erase-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.

⏳ No Failures, 8 Pending

As of commit 40c1666 with merge base 55d693b ():
💚 Looks good so far! There are no failures yet. 💚

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?

[kiymetakdemir]

It is the initial capacity added a comment for this.

[metascroy]

What does this mean?

[kiymetakdemir]

We don't use it actually, I removed.

[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?

[kiymetakdemir]

Applied this together with claude's suggestion on explicit as_planner()/as_control().

[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

[metascroy]

@claude review this code

[claude]

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

---

Code Review

• Read the cache core, registry, and tests
• Review build wiring (cmake)
• Assess against reviewer feedback & follow-up changes
• Post review

The revision addresses the major points from the earlier round cleanly. Quick status on prior feedback, then a few remaining nits.

Resolved since last review ✅

• Exceptions → ET_CHECK_MSG (Add support for quantized LeakyReLU #1): No more throw std::runtime_error. rewind/plan/build now use ET_CHECK_MSG, so the core is -fno-exceptions-safe. 👍
• Interface split (Re-sync with internal repository #2):CacheBase + SequenceControl + SequencePlanner with as_control()/as_planner(). The three bases are independent (not virtual-inheriting from CacheBase), so recovery is a static upcast returning this — no diamond, no RTTI. This is a cleaner solution than the dynamic_cast path floated earlier.
• Naming (Add unlifting pass under private config #4):ContiguousBookkeeping → FlatCache, paired conceptually with a future RingCache. Reads as a policy, not a layout claim.
• min_chunk / contiguous (Rename _pt2e to pt2e #3):min_chunk now documented as initial pool size; the ambiguous contiguous field is gone from SeqStepPlan.
• buck (Re-sync with internal repository #6): Went CMake-only. The new-core diff is 7 files, no .bzl/BUCK — matches your "buckify later" call.

Remaining items

1. Stale doc comment in cache_registry.h:63

// Throws if no builder is registered for (backend_id, kind).
std::shared_ptr<CacheBase> build(...) const;

It no longer throws — build() uses ET_CHECK_MSG (aborts). Update the comment to match, otherwise a reader will expect a catchable exception. Fix this →

2. plan() mutates state (cache.h:122)
plan() advances length_ as a side effect, so it both computes and commits the step. Worth a one-line comment on the method that "plan commits the length" — a backend that expects plan() to be a pure query (e.g. to peek a layout before deciding) would be surprised. If a peek-vs-commit split is ever needed, this is the seam.

3. valid_len semantics on re-plan / overwrite (cache.h:132-135)
valid_len = position + T, and length_ only grows (if (end > length_) length_ = end). If a caller re-plans at a position < length_ (e.g. overwrite after a rewind), valid_len reflects position + T, not the previously-tracked length_. That's a reasonable "read bound = what you just wrote" contract, but it's implicit — a comment stating that plan returns the read bound for this step (not necessarily the full tracked length) would remove ambiguity for the ring/tree planners that follow.

Minor

• Default arg on a virtual (cache.h:44):can_extend(int n = 1) — default args on virtuals bind to the static type, a known footgun if a future override changes the default. Harmless with one impl today; just flagging before more caches inherit the face.
• FlatCache is not internally thread-safe. The registry is mutex-guarded, but the cache object's length_ isn't. Fine for the single-sequence generation-loop model; worth an explicit "single-threaded per sequence" note if that assumption isn't documented elsewhere.
• Builder-registry test isolation:register_builder("TestBackend", ...) still mutates the process-global and persists across cases in the binary. Serviceable at this size (noted last round); keep an eye as more kinds register.
• int for capacity/length/position throughout — fine for sequence lengths; be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions from the last review are settled well, and the as_control()/as_planner() recovery is a nice RTTI-free resolution. Nothing blocking — items above are a stale comment plus a couple of clarifying comments on plan()/valid_len semantics.
· kv-cache-shared-registry

[metascroy]

I think initial_capacity is a better name?

[metascroy]

What does ET_CHECK_MSG do on failure? Call std::terminate?

[kiymetakdemir]

It calls std::abort()

[metascroy]

Are you planning to use this for Ring as well?

[kiymetakdemir]

Yes, I guess in that case it needs one more field for the window's start position.

[metascroy]

Looks great! Main concern is calling std::abort. Instead wrap in Result instead.

See how ET runtime / delegates avoids exception without calling abort.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 31s —— View job

---

Code Review

• Read the cache core, registry, and tests
• Verify the abort→Result change requested last round
• Review build wiring (CMake)
• Post review

The last blocker — std::abort via ET_CHECK_MSG — is fully resolved. This round is clean; only minor hardening notes remain.

Resolved since last review ✅

• abort → Result/Error (main concern):plan() and build() now return Result<T> and rewind() returns Error, all via ET_CHECK_OR_RETURN_ERROR, which logs and returns rather than aborting (runtime/core/error.h:182). No std::abort/std::terminate on any error path — this matches how the ET runtime/delegates avoid exceptions. 👍
• Death tests → Result checks:cache_test.cpp now asserts .error() == Error::InvalidArgument and .ok() instead of EXPECT_DEATH. Matches the requested pattern.
• min_chunk → initial_capacity: renamed and documented as the lazy-doubling pool's starting size (cache.h:84-87).
• Interface split + RTTI-free recovery:CacheBase / SequenceControl / SequencePlanner with as_control()/as_planner() returning this — static upcast, no dynamic_cast.
• buck: CMake-only; no .bzl/BUCK in the diff, matching the "buckify later" call.

Minor / non-blocking

1. SeqStepPlan fields have no default initializers (cache.h:57-61)

structSeqStepPlan {
int valid_len;
int write_start;
int write_len;
};

FlatCache::plan aggregate-inits all three, so it's safe today. But once ring/tree planners return partially-filled plans, an uninitialized field reads as garbage. The proposal in the review thread had = 0 defaults — worth keeping them here to stay safe as planners multiply. Fix this →

2. plan() / rewind() don't guard negative inputs (cache.h:117,130)
plan(position, T) checks end <= capacity_ but not position >= 0 || T >= 0; a negative position yields a negative write_start. rewind(new_len) checks new_len <= length_ but not new_len >= 0, so rewind(-1) sets length_ = -1 and silently corrupts can_extend. Both are internal-API misuse rather than user input, but a lower-bound ET_CHECK_OR_RETURN_ERROR is cheap insurance.

3. plan() commits length as a side effect (cache.h:138-140)
plan() both computes the step and advances length_. A one-line comment noting "plan commits the length" would help a future planner that expects a pure peek before deciding a layout. (Raised last round; still worth the comment.)

4. Carried-forward nits (all still fine at this size):

• can_extend(int n = 1) — default arg on a virtual binds to the static type; harmless with one impl.
• FlatCache isn't internally thread-safe (length_ unguarded); fine under the single-sequence generation-loop model — worth an explicit note.
• Builder-registry test mutates the process-global (register_builder("TestBackend", …)) and persists across cases.
• int for capacity/length/position — be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions are settled and the error-handling convention is now consistent with the rest of the runtime. Nothing blocking — the only one I'd fold in before backends land is the SeqStepPlan default initializers (#1), since those plans get consumed by byte-layer code shortly.
· kv-cache-shared-registry

[metascroy]

Let's keep gemma4 in mind for the first model we enable here, which has a mix of flat/ring, which I think must be captured in config. I think both must be driven from the same controller.

What if we do:

struct LayerPolicy {
enum class Kind : int { Flat = 0, Ring = 1 }; // serialized values: append-only
Kind kind = Kind::Flat;
int window = 0; // Ring: sliding-window size; must be 0 when Flat
};
struct LayerConfig {
LayerPolicy policy; // default Flat
int n_kv_heads;
int head_dim;
// dtype / scale / sink size land here later (not this PR)
};
struct CacheConfig {
int capacity; // logical cap
int n_layers;
std::vector<LayerConfig> layers; // size 1 == uniform, else == n_layers
int initial_capacity = 512;
};

[digantdesai]

Also see CacheLayerMixIn from hf/transformers

[metascroy]

Thinking ahead to mixed gemma4 model, maybe we have one SequenceCache, with flat/ring being policies. We need one control driver for both caches across all layers I think. Something like this:

class SequenceCache : public CacheBase,
public SequenceControl,
public SequencePlanner {
public:
// CacheBase: static upcasts, no RTTI (each returns `this`).
SequenceControl* as_control() override { return this; }
SequencePlanner* as_planner() override { return this; }
// SequenceControl (runner face): length / admission / rewind / clear.
bool can_extend(int n = 1) const override; // min over policies' retention
int capacity() const override { return capacity_; }
Error rewind(int new_len) override; // bounded by smallest window
void clear() override { length_ = 0; }
// SequencePlanner (backend face): dispatch to the layout policy.
Result<SeqStepPlan> plan(int policy_id, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(
policy_id >= 0 && policy_id < (int)policies_.size(),
InvalidArgument, "bad policy_id %d", policy_id);
return policies_[policy_id]->plan(position, T);
}
private:
int capacity_;
int length_ = 0;
std::vector<std::unique_ptr<LayoutPolicy>> policies_; // FlatPolicy / RingPolicy
std::vector<int> layer_to_policy_; // built from CacheConfig
};

Here LayoutPolicy is something like:

class LayoutPolicy {
public:
virtual ~LayoutPolicy() = default;
virtual SeqStepPlan plan(int position, int T) const = 0; // const == no length side effect
virtual int depth() const = 0; // physical rows to allocate
};

with FlatPolicy having the logic you have for FlatCache here, e.g.,

class FlatPolicy final : public LayoutPolicy {
public:
explicit FlatPolicy(int capacity) : capacity_(capacity) {}
int depth() const override { return capacity_; }
SeqStepPlan plan(int position, int T) const override {
const int end = position + T;
SeqStepPlan p{};
p.n_write = 1; p.write[0] = {position, T}; // contiguous, never wraps
p.n_read = 1; p.read[0] = {0, end}; // attend over all history
return p;
}
private:
int capacity_;
};

Although this diff is flat cache, we need SeqStepPlan to work for both. A ring can have 1 or 2 contiguous runs, a flat cache will always have 1:

struct Run { int start; int len; }; // a contiguous span of physical rows
struct SeqStepPlan {
Run write[2]; int n_write; // where to place this step's new KV
Run read[2]; int n_read; // valid history to attend over, oldest→newest
};

[metascroy]

@claude review again in light of latest comments on policy

[claude]

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

---

Review — the gemma4 mixed flat/ring policy proposal

• Read cache core, registry, tests
• Assess current code against the mixed flat/ring policy design
• Post review

I read the current cache.h/cache_registry.* against your two latest comments (the CacheConfig/LayerPolicy shape at cache.h:79, and the single-SequenceCache-with-LayoutPolicy shape at cache.h:92). I think the direction is right — below is where I'd push on the details before it lands, since these are the surfaces backends will bind to.

The core call: one controller, per-layer policies ✅

Agreed, and I'd go further — for a mixed model this isn't just nicer, the separate-class design is actually wrong. Gemma-style alternating local/global attention means flat and ring layers coexist in one model, but there is exactly one logical sequence length across all layers. Two independent FlatCache/RingCache objects would each track their own length_, and can_extend/rewind would no longer be coherent for the sequence as a whole. A single SequenceCache that owns length_ and dispatches layout to per-layer policies is the only thing that keeps admission/rewind consistent. So the SequenceCache + vector<LayoutPolicy> + layer_to_policy_ shape is the one to build.

1. Dispatch key: prefer layer over policy_id

plan(int policy_id, int position, int T) leaks an internal id to the backend. The byte layer knows which layer it's materializing, not which policy bucket that maps to. Keeping the mapping private is cleaner:

Result<SeqStepPlan> plan(int layer, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(layer >= 0 && layer < (int)layer_to_policy_.size(),
InvalidArgument, "bad layer %d", layer);
return policies_[layer_to_policy_[layer]]->plan(position, T);
}

layer_to_policy_ is already in your sketch — this just makes it the public parameter and keeps policy_id an implementation detail.

2. Length commit must be idempotent (plan runs N_layers times per step)

Making LayoutPolicy::planconst (pure layout, no side effect) is the right resolution to the "plan commits length" concern from earlier rounds — the controller owns length now. But watch the seam: the backend calls plan()once per layer per step, so SequenceCache must not advance length_ additively inside plan. Idempotent length_ = max(length_, position + T) is safe to call every layer; an additive length_ += T would multiply by n_layers. Either document that plan's length commit is idempotent, or split it out into a single control()->advance(T) the runner calls once per step and keep plan fully read-only. I'd lean to the explicit advance — it makes "one commit per step" impossible to get wrong.

3. can_extend / rewind semantics across mixed policies

Your one-liners ("min over policies' retention", "bounded by smallest window") are the correct contract, but they carry real logic worth stating in code:

• can_extend: a ring layer never refuses (it evicts), so admission is gated only by the flat/full layers → effectively bounded by capacity. Ring layers contribute no bound. Fine, but that means can_extend is really "any non-evicting policy near its cap," so the controller needs each policy to expose whether it evicts.
• rewind(new_len): this is the subtle one. A ring layer with window W has physically evicted everything older than length_ - W, so rewind to a new_len < length_ - min_window is unrecoverable — that history is gone from the ring layers even though the flat layers still have it. The bound isn't just new_len <= length_; it's new_len >= length_ - smallest_window too. That lower bound needs the window sizes in the controller, so LayoutPolicy should expose retained_from(length) (or similar) alongside depth().

4. SeqStepPlan: read[] runs are needed, but so is the read base position

Reintroducing Run read[2]; int n_read is correct — for a ring cache the valid history is not[0, length), it wraps, so a single valid_len can't describe it (that's why the earlier flat-only valid_len field doesn't generalize). Two points:

• Run write[2]/read[2] caps at 2 runs, which is exactly right for flat (1) and ring (≤2). Make it explicit in a comment that SeqStepPlan is deliberately the flat∪ring shape and that paged/tree planners return their own plan type (you noted this earlier) — otherwise [2] reads as an arbitrary limit.
• Your first proposal (the 18:10 sketch) had read_base_pos — the logical position of read[0].start. The latest slimmed SeqStepPlan dropped it. For ring, the backend needs the logical positions of the read runs to apply RoPE and the attention mask correctly (the oldest retained row is not at logical position 0 after eviction). Unless positions are reconstructed elsewhere, I'd bring read_base_pos back (or a per-run base) — dropping it looks like it'll break RoPE on the ring path.

5. Config shape

The LayerPolicy/LayerConfig/CacheConfig nesting is good and folds the current parallel n_kv_heads/head_dim vectors into an array-of-structs, which is cleaner. Keep the "size 1 == uniform, else == n_layers" convention you already have for the head vectors — most models are uniform and shouldn't pay for a full per-layer list. Two small things:

• Keep the enum class Kind : int with the "serialized values: append-only" comment — good instinct if this ever hits .pte metadata.
• State the window == 0 iff Flat invariant as an ET_CHECK_OR_RETURN_ERROR in whatever validates the config, not just a comment, since it's now load-bearing for depth()/rewind.

Scope: adopt at least the shapes in this PR

You've already approved, so this is a "now vs. next PR" judgment. My take: the internal SequenceCache/LayoutPolicy refactor is cheapest to do now, while FlatPolicy is the only implementation to migrate — the same argument that justified doing the control/planner split in this PR rather than deferring it. At minimum, land the CacheConfig/LayerConfig and SeqStepPlan (write[2]/read[2] + read_base_pos) shapes now, since those are the surfaces the byte layer and any serialized metadata will bind to and are the expensive ones to change later. The FlatCache→SequenceCache+FlatPolicy mechanical move can follow if you'd rather keep this PR small, but the config/plan structs I'd fix here.

Nothing above contradicts the approval — it's design hardening for the mixed-model case you're aiming at. The two I'd treat as near-blocking for the refactor are #2 (idempotent length commit) and #4 (read_base_pos for ring RoPE), since both are silent-correctness traps rather than cosmetic.
· kv-cache-shared-registry

[digantdesai]

can we wrap the base class in ET runtime friendly cache class? This is to avoid dependence on ET from runner side. I.e. I want this to be useful even when using outside ET runner. Or we can split this file into two, runner facing which can use ET stuff because for non-ET runner we can have another class. And the backend cache interface, which shouldn't use ET stuff.

Code location is also a problem, but we can discuss that later, and I don't want to block on that one.

[kiymetakdemir]

Sure! After looking into it, I'm leaning toward implementing an ET wrapper so the whole core stays a single, ET-free implementation.

[digantdesai]

Ok sorry I wasn't very clear, given I didn't see this PR.

To rephrase, let's keep the ET interfaces separate (and can be used elsewhere outside of 'et::'). The rationale is we might want to use this in a new backend layer which is not an ET delegate or linked against ET runtime libs.

For now the ET specific registry design seems OK.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 25s —— View job

---

Code Review

• Read the current cache core, registry, and tests
• Review build wiring (CMake) and test gating
• Assess against the latest ET-independence feedback (digantdesai)
• Post review

The PR has been trimmed substantially since the last round. cache.h is now just the ET-free CacheBase anchor + CacheConfig — the SequenceControl/SequencePlanner/FlatCache/SeqStepPlan machinery from earlier rounds is gone, deferred to the follow-up that lands the concrete cache. That's a clean response to @digantdesai's request to keep the neutral core usable outside an ET runner.

Resolved / consistent with feedback ✅

• ET-free core (digantdesai):cache.h has zero ExecuTorch includes and no et:: dependency — it's a plain CacheBase + CacheConfig. The ET-specific Error/Result usage is confined to cache_registry.*, and the header comments state the split explicitly. This matches "keep the ET interfaces separate so this can be used in a non-delegate backend layer."
• abort → Result/Error:build() returns Result<...> via ET_CHECK_OR_RETURN_ERROR (logs + returns, no std::abort), and the test asserts .error() == Error::NotFound rather than a death test. Consistent with the resolution from the earlier rounds.
• CMake-only, no buck: matches the "buckify later" call. Test gating uses if(BUILD_TESTING), consistent with the other extension/*/CMakeLists.txt.

Observations on the current (slimmed) surface

1. CacheConfig fields are uninitialized (cache.h:31-35)

structCacheConfig {
int capacity;
int n_layers;
int initial_capacity = 512;
};

capacity and n_layers have no default initializer. The test aggregate-inits both (CacheConfig cfg{32, 1}), so it's safe there, but a caller doing CacheConfig cfg; gets garbage in the two most important fields. Given initial_capacity already has a default, adding = 0 to the other two keeps the struct uniformly safe to default-construct. Fix this →

2. The gemma4 mixed flat/ring design (from the last round) isn't in this PR
The LayerPolicy/LayerConfig per-layer config and the SeqStepPlan (write[2]/read[2] + read_base_pos) shapes discussed at length aren't here — CacheConfig is still the flat capacity/n_layers/initial_capacity triple. That's fine as a scoping decision (this PR is now just the registry + neutral anchor), but worth confirming that's intentional: my earlier note stands that the config and step-plan structs are the expensive surfaces to change later because the byte layer and any serialized metadata bind to them. If gemma4 is the first target, folding at least the LayerConfig array-of-structs into CacheConfig now is cheaper than after consumers exist. Not blocking for a registry-only PR, but flag it so it's a deliberate "next PR" rather than an oversight.

3. Builder-registry key concatenation (cache_registry.cpp:53,63)
The map key is backend_id + ":" + kind. If either token can contain :, "a:b" + ":" + "c" and "a" + ":" + "b:c" collide. Backend ids/kinds are internal constants today so it's theoretical, but a one-line note (or a std::pair key) removes the ambiguity if these ever become user-facing strings.

Minor / carried forward (all still fine at this size)

• Process-global singletons + test isolation:register_builder("TestBackend", "stub", ...) mutates the process-global CacheBuilderRegistry and persists for the test binary's lifetime. Serviceable at one test; watch it as more kinds register.
• int for capacity/layers: fine for these counts; be deliberate about widening if any of these ever feed byte-offset math in the pool layer.
• Registry comment (cache_registry.h:69): "Returns Error::NotFound" now correctly matches the Result return — good, the stale "throws" comment from earlier rounds is fixed.

Overall: the scope reduction is the right move and directly addresses the ET-independence concern. Nothing blocking. The one I'd fold in here is the CacheConfig default initializers (#1); the mixed-policy config/plan shapes (#2) are a scoping judgment call for the follow-up.
· kv-cache-shared-registry

[metascroy]

Use std::pair, other than that, LGTM! Let's land, we can always fix later

[kiymetakdemir]

Summary

Adds a shared, backend-agnostic registry for the off-graph KV cache: a small set of C++ pieces that let a runtime own a cache and look it up by key.

Files

• extension/llm/cache/cache.h — the opaque cache handle (CacheBase) and the CacheConfig struct.
• extension/llm/cache/cache_registry.h / cache_registry.cpp — a process-global registry that stores caches by key, a
  per-(backend_id, kind) builder registry, unique-key generation, and a RAII session that installs on construction
  and erases on destruction.
• extension/llm/cache/CMakeLists.txt — builds the library and pulls in the test folder (cmake).
• extension/llm/cache/test/cache_test.cpp — unit tests for the above.
• extension/llm/cache/test/CMakeLists.txt — builds the test (cmake).
• CMakeLists.txt — includes the new folder in the build.

Testing

Added C++ unit tests (test/cache_test.cpp) covering the registry (install/get/erase, unique keys), the builder
registry (build + missing-kind NotFound error), and the session's install-on-create / erase-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.

⏳ No Failures, 8 Pending

As of commit 40c1666 with merge base 55d693b ():
💚 Looks good so far! There are no failures yet. 💚

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?

[kiymetakdemir]

It is the initial capacity added a comment for this.

[metascroy]

What does this mean?

[kiymetakdemir]

We don't use it actually, I removed.

[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?

[kiymetakdemir]

Applied this together with claude's suggestion on explicit as_planner()/as_control().

[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

[metascroy]

@claude review this code

[claude]

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

---

Code Review

• Read the cache core, registry, and tests
• Review build wiring (cmake)
• Assess against reviewer feedback & follow-up changes
• Post review

The revision addresses the major points from the earlier round cleanly. Quick status on prior feedback, then a few remaining nits.

Resolved since last review ✅

• Exceptions → ET_CHECK_MSG (Add support for quantized LeakyReLU #1): No more throw std::runtime_error. rewind/plan/build now use ET_CHECK_MSG, so the core is -fno-exceptions-safe. 👍
• Interface split (Re-sync with internal repository #2):CacheBase + SequenceControl + SequencePlanner with as_control()/as_planner(). The three bases are independent (not virtual-inheriting from CacheBase), so recovery is a static upcast returning this — no diamond, no RTTI. This is a cleaner solution than the dynamic_cast path floated earlier.
• Naming (Add unlifting pass under private config #4):ContiguousBookkeeping → FlatCache, paired conceptually with a future RingCache. Reads as a policy, not a layout claim.
• min_chunk / contiguous (Rename _pt2e to pt2e #3):min_chunk now documented as initial pool size; the ambiguous contiguous field is gone from SeqStepPlan.
• buck (Re-sync with internal repository #6): Went CMake-only. The new-core diff is 7 files, no .bzl/BUCK — matches your "buckify later" call.

Remaining items

1. Stale doc comment in cache_registry.h:63

// Throws if no builder is registered for (backend_id, kind).
std::shared_ptr<CacheBase> build(...) const;

It no longer throws — build() uses ET_CHECK_MSG (aborts). Update the comment to match, otherwise a reader will expect a catchable exception. Fix this →

2. plan() mutates state (cache.h:122)
plan() advances length_ as a side effect, so it both computes and commits the step. Worth a one-line comment on the method that "plan commits the length" — a backend that expects plan() to be a pure query (e.g. to peek a layout before deciding) would be surprised. If a peek-vs-commit split is ever needed, this is the seam.

3. valid_len semantics on re-plan / overwrite (cache.h:132-135)
valid_len = position + T, and length_ only grows (if (end > length_) length_ = end). If a caller re-plans at a position < length_ (e.g. overwrite after a rewind), valid_len reflects position + T, not the previously-tracked length_. That's a reasonable "read bound = what you just wrote" contract, but it's implicit — a comment stating that plan returns the read bound for this step (not necessarily the full tracked length) would remove ambiguity for the ring/tree planners that follow.

Minor

• Default arg on a virtual (cache.h:44):can_extend(int n = 1) — default args on virtuals bind to the static type, a known footgun if a future override changes the default. Harmless with one impl today; just flagging before more caches inherit the face.
• FlatCache is not internally thread-safe. The registry is mutex-guarded, but the cache object's length_ isn't. Fine for the single-sequence generation-loop model; worth an explicit "single-threaded per sequence" note if that assumption isn't documented elsewhere.
• Builder-registry test isolation:register_builder("TestBackend", ...) still mutates the process-global and persists across cases in the binary. Serviceable at this size (noted last round); keep an eye as more kinds register.
• int for capacity/length/position throughout — fine for sequence lengths; be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions from the last review are settled well, and the as_control()/as_planner() recovery is a nice RTTI-free resolution. Nothing blocking — items above are a stale comment plus a couple of clarifying comments on plan()/valid_len semantics.
· kv-cache-shared-registry

[metascroy]

I think initial_capacity is a better name?

[metascroy]

What does ET_CHECK_MSG do on failure? Call std::terminate?

[kiymetakdemir]

It calls std::abort()

[metascroy]

Are you planning to use this for Ring as well?

[kiymetakdemir]

Yes, I guess in that case it needs one more field for the window's start position.

[metascroy]

Looks great! Main concern is calling std::abort. Instead wrap in Result instead.

See how ET runtime / delegates avoids exception without calling abort.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 31s —— View job

---

Code Review

• Read the cache core, registry, and tests
• Verify the abort→Result change requested last round
• Review build wiring (CMake)
• Post review

The last blocker — std::abort via ET_CHECK_MSG — is fully resolved. This round is clean; only minor hardening notes remain.

Resolved since last review ✅

• abort → Result/Error (main concern):plan() and build() now return Result<T> and rewind() returns Error, all via ET_CHECK_OR_RETURN_ERROR, which logs and returns rather than aborting (runtime/core/error.h:182). No std::abort/std::terminate on any error path — this matches how the ET runtime/delegates avoid exceptions. 👍
• Death tests → Result checks:cache_test.cpp now asserts .error() == Error::InvalidArgument and .ok() instead of EXPECT_DEATH. Matches the requested pattern.
• min_chunk → initial_capacity: renamed and documented as the lazy-doubling pool's starting size (cache.h:84-87).
• Interface split + RTTI-free recovery:CacheBase / SequenceControl / SequencePlanner with as_control()/as_planner() returning this — static upcast, no dynamic_cast.
• buck: CMake-only; no .bzl/BUCK in the diff, matching the "buckify later" call.

Minor / non-blocking

1. SeqStepPlan fields have no default initializers (cache.h:57-61)

structSeqStepPlan {
int valid_len;
int write_start;
int write_len;
};

FlatCache::plan aggregate-inits all three, so it's safe today. But once ring/tree planners return partially-filled plans, an uninitialized field reads as garbage. The proposal in the review thread had = 0 defaults — worth keeping them here to stay safe as planners multiply. Fix this →

2. plan() / rewind() don't guard negative inputs (cache.h:117,130)
plan(position, T) checks end <= capacity_ but not position >= 0 || T >= 0; a negative position yields a negative write_start. rewind(new_len) checks new_len <= length_ but not new_len >= 0, so rewind(-1) sets length_ = -1 and silently corrupts can_extend. Both are internal-API misuse rather than user input, but a lower-bound ET_CHECK_OR_RETURN_ERROR is cheap insurance.

3. plan() commits length as a side effect (cache.h:138-140)
plan() both computes the step and advances length_. A one-line comment noting "plan commits the length" would help a future planner that expects a pure peek before deciding a layout. (Raised last round; still worth the comment.)

4. Carried-forward nits (all still fine at this size):

• can_extend(int n = 1) — default arg on a virtual binds to the static type; harmless with one impl.
• FlatCache isn't internally thread-safe (length_ unguarded); fine under the single-sequence generation-loop model — worth an explicit note.
• Builder-registry test mutates the process-global (register_builder("TestBackend", …)) and persists across cases.
• int for capacity/length/position — be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions are settled and the error-handling convention is now consistent with the rest of the runtime. Nothing blocking — the only one I'd fold in before backends land is the SeqStepPlan default initializers (#1), since those plans get consumed by byte-layer code shortly.
· kv-cache-shared-registry

[metascroy]

Let's keep gemma4 in mind for the first model we enable here, which has a mix of flat/ring, which I think must be captured in config. I think both must be driven from the same controller.

What if we do:

struct LayerPolicy {
enum class Kind : int { Flat = 0, Ring = 1 }; // serialized values: append-only
Kind kind = Kind::Flat;
int window = 0; // Ring: sliding-window size; must be 0 when Flat
};
struct LayerConfig {
LayerPolicy policy; // default Flat
int n_kv_heads;
int head_dim;
// dtype / scale / sink size land here later (not this PR)
};
struct CacheConfig {
int capacity; // logical cap
int n_layers;
std::vector<LayerConfig> layers; // size 1 == uniform, else == n_layers
int initial_capacity = 512;
};

[digantdesai]

Also see CacheLayerMixIn from hf/transformers

[metascroy]

Thinking ahead to mixed gemma4 model, maybe we have one SequenceCache, with flat/ring being policies. We need one control driver for both caches across all layers I think. Something like this:

class SequenceCache : public CacheBase,
public SequenceControl,
public SequencePlanner {
public:
// CacheBase: static upcasts, no RTTI (each returns `this`).
SequenceControl* as_control() override { return this; }
SequencePlanner* as_planner() override { return this; }
// SequenceControl (runner face): length / admission / rewind / clear.
bool can_extend(int n = 1) const override; // min over policies' retention
int capacity() const override { return capacity_; }
Error rewind(int new_len) override; // bounded by smallest window
void clear() override { length_ = 0; }
// SequencePlanner (backend face): dispatch to the layout policy.
Result<SeqStepPlan> plan(int policy_id, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(
policy_id >= 0 && policy_id < (int)policies_.size(),
InvalidArgument, "bad policy_id %d", policy_id);
return policies_[policy_id]->plan(position, T);
}
private:
int capacity_;
int length_ = 0;
std::vector<std::unique_ptr<LayoutPolicy>> policies_; // FlatPolicy / RingPolicy
std::vector<int> layer_to_policy_; // built from CacheConfig
};

Here LayoutPolicy is something like:

class LayoutPolicy {
public:
virtual ~LayoutPolicy() = default;
virtual SeqStepPlan plan(int position, int T) const = 0; // const == no length side effect
virtual int depth() const = 0; // physical rows to allocate
};

with FlatPolicy having the logic you have for FlatCache here, e.g.,

class FlatPolicy final : public LayoutPolicy {
public:
explicit FlatPolicy(int capacity) : capacity_(capacity) {}
int depth() const override { return capacity_; }
SeqStepPlan plan(int position, int T) const override {
const int end = position + T;
SeqStepPlan p{};
p.n_write = 1; p.write[0] = {position, T}; // contiguous, never wraps
p.n_read = 1; p.read[0] = {0, end}; // attend over all history
return p;
}
private:
int capacity_;
};

Although this diff is flat cache, we need SeqStepPlan to work for both. A ring can have 1 or 2 contiguous runs, a flat cache will always have 1:

struct Run { int start; int len; }; // a contiguous span of physical rows
struct SeqStepPlan {
Run write[2]; int n_write; // where to place this step's new KV
Run read[2]; int n_read; // valid history to attend over, oldest→newest
};

[metascroy]

@claude review again in light of latest comments on policy

[claude]

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

---

Review — the gemma4 mixed flat/ring policy proposal

• Read cache core, registry, tests
• Assess current code against the mixed flat/ring policy design
• Post review

I read the current cache.h/cache_registry.* against your two latest comments (the CacheConfig/LayerPolicy shape at cache.h:79, and the single-SequenceCache-with-LayoutPolicy shape at cache.h:92). I think the direction is right — below is where I'd push on the details before it lands, since these are the surfaces backends will bind to.

The core call: one controller, per-layer policies ✅

Agreed, and I'd go further — for a mixed model this isn't just nicer, the separate-class design is actually wrong. Gemma-style alternating local/global attention means flat and ring layers coexist in one model, but there is exactly one logical sequence length across all layers. Two independent FlatCache/RingCache objects would each track their own length_, and can_extend/rewind would no longer be coherent for the sequence as a whole. A single SequenceCache that owns length_ and dispatches layout to per-layer policies is the only thing that keeps admission/rewind consistent. So the SequenceCache + vector<LayoutPolicy> + layer_to_policy_ shape is the one to build.

1. Dispatch key: prefer layer over policy_id

plan(int policy_id, int position, int T) leaks an internal id to the backend. The byte layer knows which layer it's materializing, not which policy bucket that maps to. Keeping the mapping private is cleaner:

Result<SeqStepPlan> plan(int layer, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(layer >= 0 && layer < (int)layer_to_policy_.size(),
InvalidArgument, "bad layer %d", layer);
return policies_[layer_to_policy_[layer]]->plan(position, T);
}

layer_to_policy_ is already in your sketch — this just makes it the public parameter and keeps policy_id an implementation detail.

2. Length commit must be idempotent (plan runs N_layers times per step)

Making LayoutPolicy::planconst (pure layout, no side effect) is the right resolution to the "plan commits length" concern from earlier rounds — the controller owns length now. But watch the seam: the backend calls plan()once per layer per step, so SequenceCache must not advance length_ additively inside plan. Idempotent length_ = max(length_, position + T) is safe to call every layer; an additive length_ += T would multiply by n_layers. Either document that plan's length commit is idempotent, or split it out into a single control()->advance(T) the runner calls once per step and keep plan fully read-only. I'd lean to the explicit advance — it makes "one commit per step" impossible to get wrong.

3. can_extend / rewind semantics across mixed policies

Your one-liners ("min over policies' retention", "bounded by smallest window") are the correct contract, but they carry real logic worth stating in code:

• can_extend: a ring layer never refuses (it evicts), so admission is gated only by the flat/full layers → effectively bounded by capacity. Ring layers contribute no bound. Fine, but that means can_extend is really "any non-evicting policy near its cap," so the controller needs each policy to expose whether it evicts.
• rewind(new_len): this is the subtle one. A ring layer with window W has physically evicted everything older than length_ - W, so rewind to a new_len < length_ - min_window is unrecoverable — that history is gone from the ring layers even though the flat layers still have it. The bound isn't just new_len <= length_; it's new_len >= length_ - smallest_window too. That lower bound needs the window sizes in the controller, so LayoutPolicy should expose retained_from(length) (or similar) alongside depth().

4. SeqStepPlan: read[] runs are needed, but so is the read base position

Reintroducing Run read[2]; int n_read is correct — for a ring cache the valid history is not[0, length), it wraps, so a single valid_len can't describe it (that's why the earlier flat-only valid_len field doesn't generalize). Two points:

• Run write[2]/read[2] caps at 2 runs, which is exactly right for flat (1) and ring (≤2). Make it explicit in a comment that SeqStepPlan is deliberately the flat∪ring shape and that paged/tree planners return their own plan type (you noted this earlier) — otherwise [2] reads as an arbitrary limit.
• Your first proposal (the 18:10 sketch) had read_base_pos — the logical position of read[0].start. The latest slimmed SeqStepPlan dropped it. For ring, the backend needs the logical positions of the read runs to apply RoPE and the attention mask correctly (the oldest retained row is not at logical position 0 after eviction). Unless positions are reconstructed elsewhere, I'd bring read_base_pos back (or a per-run base) — dropping it looks like it'll break RoPE on the ring path.

5. Config shape

The LayerPolicy/LayerConfig/CacheConfig nesting is good and folds the current parallel n_kv_heads/head_dim vectors into an array-of-structs, which is cleaner. Keep the "size 1 == uniform, else == n_layers" convention you already have for the head vectors — most models are uniform and shouldn't pay for a full per-layer list. Two small things:

• Keep the enum class Kind : int with the "serialized values: append-only" comment — good instinct if this ever hits .pte metadata.
• State the window == 0 iff Flat invariant as an ET_CHECK_OR_RETURN_ERROR in whatever validates the config, not just a comment, since it's now load-bearing for depth()/rewind.

Scope: adopt at least the shapes in this PR

You've already approved, so this is a "now vs. next PR" judgment. My take: the internal SequenceCache/LayoutPolicy refactor is cheapest to do now, while FlatPolicy is the only implementation to migrate — the same argument that justified doing the control/planner split in this PR rather than deferring it. At minimum, land the CacheConfig/LayerConfig and SeqStepPlan (write[2]/read[2] + read_base_pos) shapes now, since those are the surfaces the byte layer and any serialized metadata will bind to and are the expensive ones to change later. The FlatCache→SequenceCache+FlatPolicy mechanical move can follow if you'd rather keep this PR small, but the config/plan structs I'd fix here.

Nothing above contradicts the approval — it's design hardening for the mixed-model case you're aiming at. The two I'd treat as near-blocking for the refactor are #2 (idempotent length commit) and #4 (read_base_pos for ring RoPE), since both are silent-correctness traps rather than cosmetic.
· kv-cache-shared-registry

[digantdesai]

can we wrap the base class in ET runtime friendly cache class? This is to avoid dependence on ET from runner side. I.e. I want this to be useful even when using outside ET runner. Or we can split this file into two, runner facing which can use ET stuff because for non-ET runner we can have another class. And the backend cache interface, which shouldn't use ET stuff.

Code location is also a problem, but we can discuss that later, and I don't want to block on that one.

[kiymetakdemir]

Sure! After looking into it, I'm leaning toward implementing an ET wrapper so the whole core stays a single, ET-free implementation.

[digantdesai]

Ok sorry I wasn't very clear, given I didn't see this PR.

To rephrase, let's keep the ET interfaces separate (and can be used elsewhere outside of 'et::'). The rationale is we might want to use this in a new backend layer which is not an ET delegate or linked against ET runtime libs.

For now the ET specific registry design seems OK.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 25s —— View job

---

Code Review

• Read the current cache core, registry, and tests
• Review build wiring (CMake) and test gating
• Assess against the latest ET-independence feedback (digantdesai)
• Post review

The PR has been trimmed substantially since the last round. cache.h is now just the ET-free CacheBase anchor + CacheConfig — the SequenceControl/SequencePlanner/FlatCache/SeqStepPlan machinery from earlier rounds is gone, deferred to the follow-up that lands the concrete cache. That's a clean response to @digantdesai's request to keep the neutral core usable outside an ET runner.

Resolved / consistent with feedback ✅

• ET-free core (digantdesai):cache.h has zero ExecuTorch includes and no et:: dependency — it's a plain CacheBase + CacheConfig. The ET-specific Error/Result usage is confined to cache_registry.*, and the header comments state the split explicitly. This matches "keep the ET interfaces separate so this can be used in a non-delegate backend layer."
• abort → Result/Error:build() returns Result<...> via ET_CHECK_OR_RETURN_ERROR (logs + returns, no std::abort), and the test asserts .error() == Error::NotFound rather than a death test. Consistent with the resolution from the earlier rounds.
• CMake-only, no buck: matches the "buckify later" call. Test gating uses if(BUILD_TESTING), consistent with the other extension/*/CMakeLists.txt.

Observations on the current (slimmed) surface

1. CacheConfig fields are uninitialized (cache.h:31-35)

structCacheConfig {
int capacity;
int n_layers;
int initial_capacity = 512;
};

capacity and n_layers have no default initializer. The test aggregate-inits both (CacheConfig cfg{32, 1}), so it's safe there, but a caller doing CacheConfig cfg; gets garbage in the two most important fields. Given initial_capacity already has a default, adding = 0 to the other two keeps the struct uniformly safe to default-construct. Fix this →

2. The gemma4 mixed flat/ring design (from the last round) isn't in this PR
The LayerPolicy/LayerConfig per-layer config and the SeqStepPlan (write[2]/read[2] + read_base_pos) shapes discussed at length aren't here — CacheConfig is still the flat capacity/n_layers/initial_capacity triple. That's fine as a scoping decision (this PR is now just the registry + neutral anchor), but worth confirming that's intentional: my earlier note stands that the config and step-plan structs are the expensive surfaces to change later because the byte layer and any serialized metadata bind to them. If gemma4 is the first target, folding at least the LayerConfig array-of-structs into CacheConfig now is cheaper than after consumers exist. Not blocking for a registry-only PR, but flag it so it's a deliberate "next PR" rather than an oversight.

3. Builder-registry key concatenation (cache_registry.cpp:53,63)
The map key is backend_id + ":" + kind. If either token can contain :, "a:b" + ":" + "c" and "a" + ":" + "b:c" collide. Backend ids/kinds are internal constants today so it's theoretical, but a one-line note (or a std::pair key) removes the ambiguity if these ever become user-facing strings.

Minor / carried forward (all still fine at this size)

• Process-global singletons + test isolation:register_builder("TestBackend", "stub", ...) mutates the process-global CacheBuilderRegistry and persists for the test binary's lifetime. Serviceable at one test; watch it as more kinds register.
• int for capacity/layers: fine for these counts; be deliberate about widening if any of these ever feed byte-offset math in the pool layer.
• Registry comment (cache_registry.h:69): "Returns Error::NotFound" now correctly matches the Result return — good, the stale "throws" comment from earlier rounds is fixed.

Overall: the scope reduction is the right move and directly addresses the ET-independence concern. Nothing blocking. The one I'd fold in here is the CacheConfig default initializers (#1); the mixed-policy config/plan shapes (#2) are a scoping judgment call for the follow-up.
· kv-cache-shared-registry

[metascroy]

Use std::pair, other than that, LGTM! Let's land, we can always fix later

[kiymetakdemir]

Summary

Adds a shared, backend-agnostic registry for the off-graph KV cache: a small set of C++ pieces that let a runtime own a cache and look it up by key.

Files

• extension/llm/cache/cache.h — the opaque cache handle (CacheBase) and the CacheConfig struct.
• extension/llm/cache/cache_registry.h / cache_registry.cpp — a process-global registry that stores caches by key, a
  per-(backend_id, kind) builder registry, unique-key generation, and a RAII session that installs on construction
  and erases on destruction.
• extension/llm/cache/CMakeLists.txt — builds the library and pulls in the test folder (cmake).
• extension/llm/cache/test/cache_test.cpp — unit tests for the above.
• extension/llm/cache/test/CMakeLists.txt — builds the test (cmake).
• CMakeLists.txt — includes the new folder in the build.

Testing

Added C++ unit tests (test/cache_test.cpp) covering the registry (install/get/erase, unique keys), the builder
registry (build + missing-kind NotFound error), and the session's install-on-create / erase-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.

⏳ No Failures, 8 Pending

As of commit 40c1666 with merge base 55d693b ():
💚 Looks good so far! There are no failures yet. 💚

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?

[kiymetakdemir]

It is the initial capacity added a comment for this.

[metascroy]

What does this mean?

[kiymetakdemir]

We don't use it actually, I removed.

[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?

[kiymetakdemir]

Applied this together with claude's suggestion on explicit as_planner()/as_control().

[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

[metascroy]

@claude review this code

[claude]

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

---

Code Review

• Read the cache core, registry, and tests
• Review build wiring (cmake)
• Assess against reviewer feedback & follow-up changes
• Post review

The revision addresses the major points from the earlier round cleanly. Quick status on prior feedback, then a few remaining nits.

Resolved since last review ✅

• Exceptions → ET_CHECK_MSG (Add support for quantized LeakyReLU #1): No more throw std::runtime_error. rewind/plan/build now use ET_CHECK_MSG, so the core is -fno-exceptions-safe. 👍
• Interface split (Re-sync with internal repository #2):CacheBase + SequenceControl + SequencePlanner with as_control()/as_planner(). The three bases are independent (not virtual-inheriting from CacheBase), so recovery is a static upcast returning this — no diamond, no RTTI. This is a cleaner solution than the dynamic_cast path floated earlier.
• Naming (Add unlifting pass under private config #4):ContiguousBookkeeping → FlatCache, paired conceptually with a future RingCache. Reads as a policy, not a layout claim.
• min_chunk / contiguous (Rename _pt2e to pt2e #3):min_chunk now documented as initial pool size; the ambiguous contiguous field is gone from SeqStepPlan.
• buck (Re-sync with internal repository #6): Went CMake-only. The new-core diff is 7 files, no .bzl/BUCK — matches your "buckify later" call.

Remaining items

1. Stale doc comment in cache_registry.h:63

// Throws if no builder is registered for (backend_id, kind).
std::shared_ptr<CacheBase> build(...) const;

It no longer throws — build() uses ET_CHECK_MSG (aborts). Update the comment to match, otherwise a reader will expect a catchable exception. Fix this →

2. plan() mutates state (cache.h:122)
plan() advances length_ as a side effect, so it both computes and commits the step. Worth a one-line comment on the method that "plan commits the length" — a backend that expects plan() to be a pure query (e.g. to peek a layout before deciding) would be surprised. If a peek-vs-commit split is ever needed, this is the seam.

3. valid_len semantics on re-plan / overwrite (cache.h:132-135)
valid_len = position + T, and length_ only grows (if (end > length_) length_ = end). If a caller re-plans at a position < length_ (e.g. overwrite after a rewind), valid_len reflects position + T, not the previously-tracked length_. That's a reasonable "read bound = what you just wrote" contract, but it's implicit — a comment stating that plan returns the read bound for this step (not necessarily the full tracked length) would remove ambiguity for the ring/tree planners that follow.

Minor

• Default arg on a virtual (cache.h:44):can_extend(int n = 1) — default args on virtuals bind to the static type, a known footgun if a future override changes the default. Harmless with one impl today; just flagging before more caches inherit the face.
• FlatCache is not internally thread-safe. The registry is mutex-guarded, but the cache object's length_ isn't. Fine for the single-sequence generation-loop model; worth an explicit "single-threaded per sequence" note if that assumption isn't documented elsewhere.
• Builder-registry test isolation:register_builder("TestBackend", ...) still mutates the process-global and persists across cases in the binary. Serviceable at this size (noted last round); keep an eye as more kinds register.
• int for capacity/length/position throughout — fine for sequence lengths; be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions from the last review are settled well, and the as_control()/as_planner() recovery is a nice RTTI-free resolution. Nothing blocking — items above are a stale comment plus a couple of clarifying comments on plan()/valid_len semantics.
· kv-cache-shared-registry

[metascroy]

I think initial_capacity is a better name?

[metascroy]

What does ET_CHECK_MSG do on failure? Call std::terminate?

[kiymetakdemir]

It calls std::abort()

[metascroy]

Are you planning to use this for Ring as well?

[kiymetakdemir]

Yes, I guess in that case it needs one more field for the window's start position.

[metascroy]

Looks great! Main concern is calling std::abort. Instead wrap in Result instead.

See how ET runtime / delegates avoids exception without calling abort.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 31s —— View job

---

Code Review

• Read the cache core, registry, and tests
• Verify the abort→Result change requested last round
• Review build wiring (CMake)
• Post review

The last blocker — std::abort via ET_CHECK_MSG — is fully resolved. This round is clean; only minor hardening notes remain.

Resolved since last review ✅

• abort → Result/Error (main concern):plan() and build() now return Result<T> and rewind() returns Error, all via ET_CHECK_OR_RETURN_ERROR, which logs and returns rather than aborting (runtime/core/error.h:182). No std::abort/std::terminate on any error path — this matches how the ET runtime/delegates avoid exceptions. 👍
• Death tests → Result checks:cache_test.cpp now asserts .error() == Error::InvalidArgument and .ok() instead of EXPECT_DEATH. Matches the requested pattern.
• min_chunk → initial_capacity: renamed and documented as the lazy-doubling pool's starting size (cache.h:84-87).
• Interface split + RTTI-free recovery:CacheBase / SequenceControl / SequencePlanner with as_control()/as_planner() returning this — static upcast, no dynamic_cast.
• buck: CMake-only; no .bzl/BUCK in the diff, matching the "buckify later" call.

Minor / non-blocking

1. SeqStepPlan fields have no default initializers (cache.h:57-61)

structSeqStepPlan {
int valid_len;
int write_start;
int write_len;
};

FlatCache::plan aggregate-inits all three, so it's safe today. But once ring/tree planners return partially-filled plans, an uninitialized field reads as garbage. The proposal in the review thread had = 0 defaults — worth keeping them here to stay safe as planners multiply. Fix this →

2. plan() / rewind() don't guard negative inputs (cache.h:117,130)
plan(position, T) checks end <= capacity_ but not position >= 0 || T >= 0; a negative position yields a negative write_start. rewind(new_len) checks new_len <= length_ but not new_len >= 0, so rewind(-1) sets length_ = -1 and silently corrupts can_extend. Both are internal-API misuse rather than user input, but a lower-bound ET_CHECK_OR_RETURN_ERROR is cheap insurance.

3. plan() commits length as a side effect (cache.h:138-140)
plan() both computes the step and advances length_. A one-line comment noting "plan commits the length" would help a future planner that expects a pure peek before deciding a layout. (Raised last round; still worth the comment.)

4. Carried-forward nits (all still fine at this size):

• can_extend(int n = 1) — default arg on a virtual binds to the static type; harmless with one impl.
• FlatCache isn't internally thread-safe (length_ unguarded); fine under the single-sequence generation-loop model — worth an explicit note.
• Builder-registry test mutates the process-global (register_builder("TestBackend", …)) and persists across cases.
• int for capacity/length/position — be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions are settled and the error-handling convention is now consistent with the rest of the runtime. Nothing blocking — the only one I'd fold in before backends land is the SeqStepPlan default initializers (#1), since those plans get consumed by byte-layer code shortly.
· kv-cache-shared-registry

[metascroy]

Let's keep gemma4 in mind for the first model we enable here, which has a mix of flat/ring, which I think must be captured in config. I think both must be driven from the same controller.

What if we do:

struct LayerPolicy {
enum class Kind : int { Flat = 0, Ring = 1 }; // serialized values: append-only
Kind kind = Kind::Flat;
int window = 0; // Ring: sliding-window size; must be 0 when Flat
};
struct LayerConfig {
LayerPolicy policy; // default Flat
int n_kv_heads;
int head_dim;
// dtype / scale / sink size land here later (not this PR)
};
struct CacheConfig {
int capacity; // logical cap
int n_layers;
std::vector<LayerConfig> layers; // size 1 == uniform, else == n_layers
int initial_capacity = 512;
};

[digantdesai]

Also see CacheLayerMixIn from hf/transformers

[metascroy]

Thinking ahead to mixed gemma4 model, maybe we have one SequenceCache, with flat/ring being policies. We need one control driver for both caches across all layers I think. Something like this:

class SequenceCache : public CacheBase,
public SequenceControl,
public SequencePlanner {
public:
// CacheBase: static upcasts, no RTTI (each returns `this`).
SequenceControl* as_control() override { return this; }
SequencePlanner* as_planner() override { return this; }
// SequenceControl (runner face): length / admission / rewind / clear.
bool can_extend(int n = 1) const override; // min over policies' retention
int capacity() const override { return capacity_; }
Error rewind(int new_len) override; // bounded by smallest window
void clear() override { length_ = 0; }
// SequencePlanner (backend face): dispatch to the layout policy.
Result<SeqStepPlan> plan(int policy_id, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(
policy_id >= 0 && policy_id < (int)policies_.size(),
InvalidArgument, "bad policy_id %d", policy_id);
return policies_[policy_id]->plan(position, T);
}
private:
int capacity_;
int length_ = 0;
std::vector<std::unique_ptr<LayoutPolicy>> policies_; // FlatPolicy / RingPolicy
std::vector<int> layer_to_policy_; // built from CacheConfig
};

Here LayoutPolicy is something like:

class LayoutPolicy {
public:
virtual ~LayoutPolicy() = default;
virtual SeqStepPlan plan(int position, int T) const = 0; // const == no length side effect
virtual int depth() const = 0; // physical rows to allocate
};

with FlatPolicy having the logic you have for FlatCache here, e.g.,

class FlatPolicy final : public LayoutPolicy {
public:
explicit FlatPolicy(int capacity) : capacity_(capacity) {}
int depth() const override { return capacity_; }
SeqStepPlan plan(int position, int T) const override {
const int end = position + T;
SeqStepPlan p{};
p.n_write = 1; p.write[0] = {position, T}; // contiguous, never wraps
p.n_read = 1; p.read[0] = {0, end}; // attend over all history
return p;
}
private:
int capacity_;
};

Although this diff is flat cache, we need SeqStepPlan to work for both. A ring can have 1 or 2 contiguous runs, a flat cache will always have 1:

struct Run { int start; int len; }; // a contiguous span of physical rows
struct SeqStepPlan {
Run write[2]; int n_write; // where to place this step's new KV
Run read[2]; int n_read; // valid history to attend over, oldest→newest
};

[metascroy]

@claude review again in light of latest comments on policy

[claude]

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

---

Review — the gemma4 mixed flat/ring policy proposal

• Read cache core, registry, tests
• Assess current code against the mixed flat/ring policy design
• Post review

I read the current cache.h/cache_registry.* against your two latest comments (the CacheConfig/LayerPolicy shape at cache.h:79, and the single-SequenceCache-with-LayoutPolicy shape at cache.h:92). I think the direction is right — below is where I'd push on the details before it lands, since these are the surfaces backends will bind to.

The core call: one controller, per-layer policies ✅

Agreed, and I'd go further — for a mixed model this isn't just nicer, the separate-class design is actually wrong. Gemma-style alternating local/global attention means flat and ring layers coexist in one model, but there is exactly one logical sequence length across all layers. Two independent FlatCache/RingCache objects would each track their own length_, and can_extend/rewind would no longer be coherent for the sequence as a whole. A single SequenceCache that owns length_ and dispatches layout to per-layer policies is the only thing that keeps admission/rewind consistent. So the SequenceCache + vector<LayoutPolicy> + layer_to_policy_ shape is the one to build.

1. Dispatch key: prefer layer over policy_id

plan(int policy_id, int position, int T) leaks an internal id to the backend. The byte layer knows which layer it's materializing, not which policy bucket that maps to. Keeping the mapping private is cleaner:

Result<SeqStepPlan> plan(int layer, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(layer >= 0 && layer < (int)layer_to_policy_.size(),
InvalidArgument, "bad layer %d", layer);
return policies_[layer_to_policy_[layer]]->plan(position, T);
}

layer_to_policy_ is already in your sketch — this just makes it the public parameter and keeps policy_id an implementation detail.

2. Length commit must be idempotent (plan runs N_layers times per step)

Making LayoutPolicy::planconst (pure layout, no side effect) is the right resolution to the "plan commits length" concern from earlier rounds — the controller owns length now. But watch the seam: the backend calls plan()once per layer per step, so SequenceCache must not advance length_ additively inside plan. Idempotent length_ = max(length_, position + T) is safe to call every layer; an additive length_ += T would multiply by n_layers. Either document that plan's length commit is idempotent, or split it out into a single control()->advance(T) the runner calls once per step and keep plan fully read-only. I'd lean to the explicit advance — it makes "one commit per step" impossible to get wrong.

3. can_extend / rewind semantics across mixed policies

Your one-liners ("min over policies' retention", "bounded by smallest window") are the correct contract, but they carry real logic worth stating in code:

• can_extend: a ring layer never refuses (it evicts), so admission is gated only by the flat/full layers → effectively bounded by capacity. Ring layers contribute no bound. Fine, but that means can_extend is really "any non-evicting policy near its cap," so the controller needs each policy to expose whether it evicts.
• rewind(new_len): this is the subtle one. A ring layer with window W has physically evicted everything older than length_ - W, so rewind to a new_len < length_ - min_window is unrecoverable — that history is gone from the ring layers even though the flat layers still have it. The bound isn't just new_len <= length_; it's new_len >= length_ - smallest_window too. That lower bound needs the window sizes in the controller, so LayoutPolicy should expose retained_from(length) (or similar) alongside depth().

4. SeqStepPlan: read[] runs are needed, but so is the read base position

Reintroducing Run read[2]; int n_read is correct — for a ring cache the valid history is not[0, length), it wraps, so a single valid_len can't describe it (that's why the earlier flat-only valid_len field doesn't generalize). Two points:

• Run write[2]/read[2] caps at 2 runs, which is exactly right for flat (1) and ring (≤2). Make it explicit in a comment that SeqStepPlan is deliberately the flat∪ring shape and that paged/tree planners return their own plan type (you noted this earlier) — otherwise [2] reads as an arbitrary limit.
• Your first proposal (the 18:10 sketch) had read_base_pos — the logical position of read[0].start. The latest slimmed SeqStepPlan dropped it. For ring, the backend needs the logical positions of the read runs to apply RoPE and the attention mask correctly (the oldest retained row is not at logical position 0 after eviction). Unless positions are reconstructed elsewhere, I'd bring read_base_pos back (or a per-run base) — dropping it looks like it'll break RoPE on the ring path.

5. Config shape

The LayerPolicy/LayerConfig/CacheConfig nesting is good and folds the current parallel n_kv_heads/head_dim vectors into an array-of-structs, which is cleaner. Keep the "size 1 == uniform, else == n_layers" convention you already have for the head vectors — most models are uniform and shouldn't pay for a full per-layer list. Two small things:

• Keep the enum class Kind : int with the "serialized values: append-only" comment — good instinct if this ever hits .pte metadata.
• State the window == 0 iff Flat invariant as an ET_CHECK_OR_RETURN_ERROR in whatever validates the config, not just a comment, since it's now load-bearing for depth()/rewind.

Scope: adopt at least the shapes in this PR

You've already approved, so this is a "now vs. next PR" judgment. My take: the internal SequenceCache/LayoutPolicy refactor is cheapest to do now, while FlatPolicy is the only implementation to migrate — the same argument that justified doing the control/planner split in this PR rather than deferring it. At minimum, land the CacheConfig/LayerConfig and SeqStepPlan (write[2]/read[2] + read_base_pos) shapes now, since those are the surfaces the byte layer and any serialized metadata will bind to and are the expensive ones to change later. The FlatCache→SequenceCache+FlatPolicy mechanical move can follow if you'd rather keep this PR small, but the config/plan structs I'd fix here.

Nothing above contradicts the approval — it's design hardening for the mixed-model case you're aiming at. The two I'd treat as near-blocking for the refactor are #2 (idempotent length commit) and #4 (read_base_pos for ring RoPE), since both are silent-correctness traps rather than cosmetic.
· kv-cache-shared-registry

[digantdesai]

can we wrap the base class in ET runtime friendly cache class? This is to avoid dependence on ET from runner side. I.e. I want this to be useful even when using outside ET runner. Or we can split this file into two, runner facing which can use ET stuff because for non-ET runner we can have another class. And the backend cache interface, which shouldn't use ET stuff.

Code location is also a problem, but we can discuss that later, and I don't want to block on that one.

[kiymetakdemir]

Sure! After looking into it, I'm leaning toward implementing an ET wrapper so the whole core stays a single, ET-free implementation.

[digantdesai]

Ok sorry I wasn't very clear, given I didn't see this PR.

To rephrase, let's keep the ET interfaces separate (and can be used elsewhere outside of 'et::'). The rationale is we might want to use this in a new backend layer which is not an ET delegate or linked against ET runtime libs.

For now the ET specific registry design seems OK.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 25s —— View job

---

Code Review

• Read the current cache core, registry, and tests
• Review build wiring (CMake) and test gating
• Assess against the latest ET-independence feedback (digantdesai)
• Post review

The PR has been trimmed substantially since the last round. cache.h is now just the ET-free CacheBase anchor + CacheConfig — the SequenceControl/SequencePlanner/FlatCache/SeqStepPlan machinery from earlier rounds is gone, deferred to the follow-up that lands the concrete cache. That's a clean response to @digantdesai's request to keep the neutral core usable outside an ET runner.

Resolved / consistent with feedback ✅

• ET-free core (digantdesai):cache.h has zero ExecuTorch includes and no et:: dependency — it's a plain CacheBase + CacheConfig. The ET-specific Error/Result usage is confined to cache_registry.*, and the header comments state the split explicitly. This matches "keep the ET interfaces separate so this can be used in a non-delegate backend layer."
• abort → Result/Error:build() returns Result<...> via ET_CHECK_OR_RETURN_ERROR (logs + returns, no std::abort), and the test asserts .error() == Error::NotFound rather than a death test. Consistent with the resolution from the earlier rounds.
• CMake-only, no buck: matches the "buckify later" call. Test gating uses if(BUILD_TESTING), consistent with the other extension/*/CMakeLists.txt.

Observations on the current (slimmed) surface

1. CacheConfig fields are uninitialized (cache.h:31-35)

structCacheConfig {
int capacity;
int n_layers;
int initial_capacity = 512;
};

capacity and n_layers have no default initializer. The test aggregate-inits both (CacheConfig cfg{32, 1}), so it's safe there, but a caller doing CacheConfig cfg; gets garbage in the two most important fields. Given initial_capacity already has a default, adding = 0 to the other two keeps the struct uniformly safe to default-construct. Fix this →

2. The gemma4 mixed flat/ring design (from the last round) isn't in this PR
The LayerPolicy/LayerConfig per-layer config and the SeqStepPlan (write[2]/read[2] + read_base_pos) shapes discussed at length aren't here — CacheConfig is still the flat capacity/n_layers/initial_capacity triple. That's fine as a scoping decision (this PR is now just the registry + neutral anchor), but worth confirming that's intentional: my earlier note stands that the config and step-plan structs are the expensive surfaces to change later because the byte layer and any serialized metadata bind to them. If gemma4 is the first target, folding at least the LayerConfig array-of-structs into CacheConfig now is cheaper than after consumers exist. Not blocking for a registry-only PR, but flag it so it's a deliberate "next PR" rather than an oversight.

3. Builder-registry key concatenation (cache_registry.cpp:53,63)
The map key is backend_id + ":" + kind. If either token can contain :, "a:b" + ":" + "c" and "a" + ":" + "b:c" collide. Backend ids/kinds are internal constants today so it's theoretical, but a one-line note (or a std::pair key) removes the ambiguity if these ever become user-facing strings.

Minor / carried forward (all still fine at this size)

• Process-global singletons + test isolation:register_builder("TestBackend", "stub", ...) mutates the process-global CacheBuilderRegistry and persists for the test binary's lifetime. Serviceable at one test; watch it as more kinds register.
• int for capacity/layers: fine for these counts; be deliberate about widening if any of these ever feed byte-offset math in the pool layer.
• Registry comment (cache_registry.h:69): "Returns Error::NotFound" now correctly matches the Result return — good, the stale "throws" comment from earlier rounds is fixed.

Overall: the scope reduction is the right move and directly addresses the ET-independence concern. Nothing blocking. The one I'd fold in here is the CacheConfig default initializers (#1); the mixed-policy config/plan shapes (#2) are a scoping judgment call for the follow-up.
· kv-cache-shared-registry

[metascroy]

Use std::pair, other than that, LGTM! Let's land, we can always fix later

[kiymetakdemir]

Summary

Adds a shared, backend-agnostic registry for the off-graph KV cache: a small set of C++ pieces that let a runtime own a cache and look it up by key.

Files

• extension/llm/cache/cache.h — the opaque cache handle (CacheBase) and the CacheConfig struct.
• extension/llm/cache/cache_registry.h / cache_registry.cpp — a process-global registry that stores caches by key, a
  per-(backend_id, kind) builder registry, unique-key generation, and a RAII session that installs on construction
  and erases on destruction.
• extension/llm/cache/CMakeLists.txt — builds the library and pulls in the test folder (cmake).
• extension/llm/cache/test/cache_test.cpp — unit tests for the above.
• extension/llm/cache/test/CMakeLists.txt — builds the test (cmake).
• CMakeLists.txt — includes the new folder in the build.

Testing

Added C++ unit tests (test/cache_test.cpp) covering the registry (install/get/erase, unique keys), the builder
registry (build + missing-kind NotFound error), and the session's install-on-create / erase-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.

⏳ No Failures, 8 Pending

As of commit 40c1666 with merge base 55d693b ():
💚 Looks good so far! There are no failures yet. 💚

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?

[kiymetakdemir]

It is the initial capacity added a comment for this.

[metascroy]

What does this mean?

[kiymetakdemir]

We don't use it actually, I removed.

[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?

[kiymetakdemir]

Applied this together with claude's suggestion on explicit as_planner()/as_control().

[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

[metascroy]

@claude review this code

[claude]

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

---

Code Review

• Read the cache core, registry, and tests
• Review build wiring (cmake)
• Assess against reviewer feedback & follow-up changes
• Post review

The revision addresses the major points from the earlier round cleanly. Quick status on prior feedback, then a few remaining nits.

Resolved since last review ✅

• Exceptions → ET_CHECK_MSG (Add support for quantized LeakyReLU #1): No more throw std::runtime_error. rewind/plan/build now use ET_CHECK_MSG, so the core is -fno-exceptions-safe. 👍
• Interface split (Re-sync with internal repository #2):CacheBase + SequenceControl + SequencePlanner with as_control()/as_planner(). The three bases are independent (not virtual-inheriting from CacheBase), so recovery is a static upcast returning this — no diamond, no RTTI. This is a cleaner solution than the dynamic_cast path floated earlier.
• Naming (Add unlifting pass under private config #4):ContiguousBookkeeping → FlatCache, paired conceptually with a future RingCache. Reads as a policy, not a layout claim.
• min_chunk / contiguous (Rename _pt2e to pt2e #3):min_chunk now documented as initial pool size; the ambiguous contiguous field is gone from SeqStepPlan.
• buck (Re-sync with internal repository #6): Went CMake-only. The new-core diff is 7 files, no .bzl/BUCK — matches your "buckify later" call.

Remaining items

1. Stale doc comment in cache_registry.h:63

// Throws if no builder is registered for (backend_id, kind).
std::shared_ptr<CacheBase> build(...) const;

It no longer throws — build() uses ET_CHECK_MSG (aborts). Update the comment to match, otherwise a reader will expect a catchable exception. Fix this →

2. plan() mutates state (cache.h:122)
plan() advances length_ as a side effect, so it both computes and commits the step. Worth a one-line comment on the method that "plan commits the length" — a backend that expects plan() to be a pure query (e.g. to peek a layout before deciding) would be surprised. If a peek-vs-commit split is ever needed, this is the seam.

3. valid_len semantics on re-plan / overwrite (cache.h:132-135)
valid_len = position + T, and length_ only grows (if (end > length_) length_ = end). If a caller re-plans at a position < length_ (e.g. overwrite after a rewind), valid_len reflects position + T, not the previously-tracked length_. That's a reasonable "read bound = what you just wrote" contract, but it's implicit — a comment stating that plan returns the read bound for this step (not necessarily the full tracked length) would remove ambiguity for the ring/tree planners that follow.

Minor

• Default arg on a virtual (cache.h:44):can_extend(int n = 1) — default args on virtuals bind to the static type, a known footgun if a future override changes the default. Harmless with one impl today; just flagging before more caches inherit the face.
• FlatCache is not internally thread-safe. The registry is mutex-guarded, but the cache object's length_ isn't. Fine for the single-sequence generation-loop model; worth an explicit "single-threaded per sequence" note if that assumption isn't documented elsewhere.
• Builder-registry test isolation:register_builder("TestBackend", ...) still mutates the process-global and persists across cases in the binary. Serviceable at this size (noted last round); keep an eye as more kinds register.
• int for capacity/length/position throughout — fine for sequence lengths; be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions from the last review are settled well, and the as_control()/as_planner() recovery is a nice RTTI-free resolution. Nothing blocking — items above are a stale comment plus a couple of clarifying comments on plan()/valid_len semantics.
· kv-cache-shared-registry

[metascroy]

I think initial_capacity is a better name?

[metascroy]

What does ET_CHECK_MSG do on failure? Call std::terminate?

[kiymetakdemir]

It calls std::abort()

[metascroy]

Are you planning to use this for Ring as well?

[kiymetakdemir]

Yes, I guess in that case it needs one more field for the window's start position.

[metascroy]

Looks great! Main concern is calling std::abort. Instead wrap in Result instead.

See how ET runtime / delegates avoids exception without calling abort.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 31s —— View job

---

Code Review

• Read the cache core, registry, and tests
• Verify the abort→Result change requested last round
• Review build wiring (CMake)
• Post review

The last blocker — std::abort via ET_CHECK_MSG — is fully resolved. This round is clean; only minor hardening notes remain.

Resolved since last review ✅

• abort → Result/Error (main concern):plan() and build() now return Result<T> and rewind() returns Error, all via ET_CHECK_OR_RETURN_ERROR, which logs and returns rather than aborting (runtime/core/error.h:182). No std::abort/std::terminate on any error path — this matches how the ET runtime/delegates avoid exceptions. 👍
• Death tests → Result checks:cache_test.cpp now asserts .error() == Error::InvalidArgument and .ok() instead of EXPECT_DEATH. Matches the requested pattern.
• min_chunk → initial_capacity: renamed and documented as the lazy-doubling pool's starting size (cache.h:84-87).
• Interface split + RTTI-free recovery:CacheBase / SequenceControl / SequencePlanner with as_control()/as_planner() returning this — static upcast, no dynamic_cast.
• buck: CMake-only; no .bzl/BUCK in the diff, matching the "buckify later" call.

Minor / non-blocking

1. SeqStepPlan fields have no default initializers (cache.h:57-61)

structSeqStepPlan {
int valid_len;
int write_start;
int write_len;
};

FlatCache::plan aggregate-inits all three, so it's safe today. But once ring/tree planners return partially-filled plans, an uninitialized field reads as garbage. The proposal in the review thread had = 0 defaults — worth keeping them here to stay safe as planners multiply. Fix this →

2. plan() / rewind() don't guard negative inputs (cache.h:117,130)
plan(position, T) checks end <= capacity_ but not position >= 0 || T >= 0; a negative position yields a negative write_start. rewind(new_len) checks new_len <= length_ but not new_len >= 0, so rewind(-1) sets length_ = -1 and silently corrupts can_extend. Both are internal-API misuse rather than user input, but a lower-bound ET_CHECK_OR_RETURN_ERROR is cheap insurance.

3. plan() commits length as a side effect (cache.h:138-140)
plan() both computes the step and advances length_. A one-line comment noting "plan commits the length" would help a future planner that expects a pure peek before deciding a layout. (Raised last round; still worth the comment.)

4. Carried-forward nits (all still fine at this size):

• can_extend(int n = 1) — default arg on a virtual binds to the static type; harmless with one impl.
• FlatCache isn't internally thread-safe (length_ unguarded); fine under the single-sequence generation-loop model — worth an explicit note.
• Builder-registry test mutates the process-global (register_builder("TestBackend", …)) and persists across cases.
• int for capacity/length/position — be deliberate about widening if these ever feed byte-offset math in the pool layer.

Overall: the design questions are settled and the error-handling convention is now consistent with the rest of the runtime. Nothing blocking — the only one I'd fold in before backends land is the SeqStepPlan default initializers (#1), since those plans get consumed by byte-layer code shortly.
· kv-cache-shared-registry

[metascroy]

Let's keep gemma4 in mind for the first model we enable here, which has a mix of flat/ring, which I think must be captured in config. I think both must be driven from the same controller.

What if we do:

struct LayerPolicy {
enum class Kind : int { Flat = 0, Ring = 1 }; // serialized values: append-only
Kind kind = Kind::Flat;
int window = 0; // Ring: sliding-window size; must be 0 when Flat
};
struct LayerConfig {
LayerPolicy policy; // default Flat
int n_kv_heads;
int head_dim;
// dtype / scale / sink size land here later (not this PR)
};
struct CacheConfig {
int capacity; // logical cap
int n_layers;
std::vector<LayerConfig> layers; // size 1 == uniform, else == n_layers
int initial_capacity = 512;
};

[digantdesai]

Also see CacheLayerMixIn from hf/transformers

[metascroy]

Thinking ahead to mixed gemma4 model, maybe we have one SequenceCache, with flat/ring being policies. We need one control driver for both caches across all layers I think. Something like this:

class SequenceCache : public CacheBase,
public SequenceControl,
public SequencePlanner {
public:
// CacheBase: static upcasts, no RTTI (each returns `this`).
SequenceControl* as_control() override { return this; }
SequencePlanner* as_planner() override { return this; }
// SequenceControl (runner face): length / admission / rewind / clear.
bool can_extend(int n = 1) const override; // min over policies' retention
int capacity() const override { return capacity_; }
Error rewind(int new_len) override; // bounded by smallest window
void clear() override { length_ = 0; }
// SequencePlanner (backend face): dispatch to the layout policy.
Result<SeqStepPlan> plan(int policy_id, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(
policy_id >= 0 && policy_id < (int)policies_.size(),
InvalidArgument, "bad policy_id %d", policy_id);
return policies_[policy_id]->plan(position, T);
}
private:
int capacity_;
int length_ = 0;
std::vector<std::unique_ptr<LayoutPolicy>> policies_; // FlatPolicy / RingPolicy
std::vector<int> layer_to_policy_; // built from CacheConfig
};

Here LayoutPolicy is something like:

class LayoutPolicy {
public:
virtual ~LayoutPolicy() = default;
virtual SeqStepPlan plan(int position, int T) const = 0; // const == no length side effect
virtual int depth() const = 0; // physical rows to allocate
};

with FlatPolicy having the logic you have for FlatCache here, e.g.,

class FlatPolicy final : public LayoutPolicy {
public:
explicit FlatPolicy(int capacity) : capacity_(capacity) {}
int depth() const override { return capacity_; }
SeqStepPlan plan(int position, int T) const override {
const int end = position + T;
SeqStepPlan p{};
p.n_write = 1; p.write[0] = {position, T}; // contiguous, never wraps
p.n_read = 1; p.read[0] = {0, end}; // attend over all history
return p;
}
private:
int capacity_;
};

Although this diff is flat cache, we need SeqStepPlan to work for both. A ring can have 1 or 2 contiguous runs, a flat cache will always have 1:

struct Run { int start; int len; }; // a contiguous span of physical rows
struct SeqStepPlan {
Run write[2]; int n_write; // where to place this step's new KV
Run read[2]; int n_read; // valid history to attend over, oldest→newest
};

[metascroy]

@claude review again in light of latest comments on policy

[claude]

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

---

Review — the gemma4 mixed flat/ring policy proposal

• Read cache core, registry, tests
• Assess current code against the mixed flat/ring policy design
• Post review

I read the current cache.h/cache_registry.* against your two latest comments (the CacheConfig/LayerPolicy shape at cache.h:79, and the single-SequenceCache-with-LayoutPolicy shape at cache.h:92). I think the direction is right — below is where I'd push on the details before it lands, since these are the surfaces backends will bind to.

The core call: one controller, per-layer policies ✅

Agreed, and I'd go further — for a mixed model this isn't just nicer, the separate-class design is actually wrong. Gemma-style alternating local/global attention means flat and ring layers coexist in one model, but there is exactly one logical sequence length across all layers. Two independent FlatCache/RingCache objects would each track their own length_, and can_extend/rewind would no longer be coherent for the sequence as a whole. A single SequenceCache that owns length_ and dispatches layout to per-layer policies is the only thing that keeps admission/rewind consistent. So the SequenceCache + vector<LayoutPolicy> + layer_to_policy_ shape is the one to build.

1. Dispatch key: prefer layer over policy_id

plan(int policy_id, int position, int T) leaks an internal id to the backend. The byte layer knows which layer it's materializing, not which policy bucket that maps to. Keeping the mapping private is cleaner:

Result<SeqStepPlan> plan(int layer, int position, int T) override {
ET_CHECK_OR_RETURN_ERROR(layer >= 0 && layer < (int)layer_to_policy_.size(),
InvalidArgument, "bad layer %d", layer);
return policies_[layer_to_policy_[layer]]->plan(position, T);
}

layer_to_policy_ is already in your sketch — this just makes it the public parameter and keeps policy_id an implementation detail.

2. Length commit must be idempotent (plan runs N_layers times per step)

Making LayoutPolicy::planconst (pure layout, no side effect) is the right resolution to the "plan commits length" concern from earlier rounds — the controller owns length now. But watch the seam: the backend calls plan()once per layer per step, so SequenceCache must not advance length_ additively inside plan. Idempotent length_ = max(length_, position + T) is safe to call every layer; an additive length_ += T would multiply by n_layers. Either document that plan's length commit is idempotent, or split it out into a single control()->advance(T) the runner calls once per step and keep plan fully read-only. I'd lean to the explicit advance — it makes "one commit per step" impossible to get wrong.

3. can_extend / rewind semantics across mixed policies

Your one-liners ("min over policies' retention", "bounded by smallest window") are the correct contract, but they carry real logic worth stating in code:

• can_extend: a ring layer never refuses (it evicts), so admission is gated only by the flat/full layers → effectively bounded by capacity. Ring layers contribute no bound. Fine, but that means can_extend is really "any non-evicting policy near its cap," so the controller needs each policy to expose whether it evicts.
• rewind(new_len): this is the subtle one. A ring layer with window W has physically evicted everything older than length_ - W, so rewind to a new_len < length_ - min_window is unrecoverable — that history is gone from the ring layers even though the flat layers still have it. The bound isn't just new_len <= length_; it's new_len >= length_ - smallest_window too. That lower bound needs the window sizes in the controller, so LayoutPolicy should expose retained_from(length) (or similar) alongside depth().

4. SeqStepPlan: read[] runs are needed, but so is the read base position

Reintroducing Run read[2]; int n_read is correct — for a ring cache the valid history is not[0, length), it wraps, so a single valid_len can't describe it (that's why the earlier flat-only valid_len field doesn't generalize). Two points:

• Run write[2]/read[2] caps at 2 runs, which is exactly right for flat (1) and ring (≤2). Make it explicit in a comment that SeqStepPlan is deliberately the flat∪ring shape and that paged/tree planners return their own plan type (you noted this earlier) — otherwise [2] reads as an arbitrary limit.
• Your first proposal (the 18:10 sketch) had read_base_pos — the logical position of read[0].start. The latest slimmed SeqStepPlan dropped it. For ring, the backend needs the logical positions of the read runs to apply RoPE and the attention mask correctly (the oldest retained row is not at logical position 0 after eviction). Unless positions are reconstructed elsewhere, I'd bring read_base_pos back (or a per-run base) — dropping it looks like it'll break RoPE on the ring path.

5. Config shape

The LayerPolicy/LayerConfig/CacheConfig nesting is good and folds the current parallel n_kv_heads/head_dim vectors into an array-of-structs, which is cleaner. Keep the "size 1 == uniform, else == n_layers" convention you already have for the head vectors — most models are uniform and shouldn't pay for a full per-layer list. Two small things:

• Keep the enum class Kind : int with the "serialized values: append-only" comment — good instinct if this ever hits .pte metadata.
• State the window == 0 iff Flat invariant as an ET_CHECK_OR_RETURN_ERROR in whatever validates the config, not just a comment, since it's now load-bearing for depth()/rewind.

Scope: adopt at least the shapes in this PR

You've already approved, so this is a "now vs. next PR" judgment. My take: the internal SequenceCache/LayoutPolicy refactor is cheapest to do now, while FlatPolicy is the only implementation to migrate — the same argument that justified doing the control/planner split in this PR rather than deferring it. At minimum, land the CacheConfig/LayerConfig and SeqStepPlan (write[2]/read[2] + read_base_pos) shapes now, since those are the surfaces the byte layer and any serialized metadata will bind to and are the expensive ones to change later. The FlatCache→SequenceCache+FlatPolicy mechanical move can follow if you'd rather keep this PR small, but the config/plan structs I'd fix here.

Nothing above contradicts the approval — it's design hardening for the mixed-model case you're aiming at. The two I'd treat as near-blocking for the refactor are #2 (idempotent length commit) and #4 (read_base_pos for ring RoPE), since both are silent-correctness traps rather than cosmetic.
· kv-cache-shared-registry

[digantdesai]

can we wrap the base class in ET runtime friendly cache class? This is to avoid dependence on ET from runner side. I.e. I want this to be useful even when using outside ET runner. Or we can split this file into two, runner facing which can use ET stuff because for non-ET runner we can have another class. And the backend cache interface, which shouldn't use ET stuff.

Code location is also a problem, but we can discuss that later, and I don't want to block on that one.

[kiymetakdemir]

Sure! After looking into it, I'm leaning toward implementing an ET wrapper so the whole core stays a single, ET-free implementation.

[digantdesai]

Ok sorry I wasn't very clear, given I didn't see this PR.

To rephrase, let's keep the ET interfaces separate (and can be used elsewhere outside of 'et::'). The rationale is we might want to use this in a new backend layer which is not an ET delegate or linked against ET runtime libs.

For now the ET specific registry design seems OK.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 25s —— View job

---

Code Review

• Read the current cache core, registry, and tests
• Review build wiring (CMake) and test gating
• Assess against the latest ET-independence feedback (digantdesai)
• Post review

The PR has been trimmed substantially since the last round. cache.h is now just the ET-free CacheBase anchor + CacheConfig — the SequenceControl/SequencePlanner/FlatCache/SeqStepPlan machinery from earlier rounds is gone, deferred to the follow-up that lands the concrete cache. That's a clean response to @digantdesai's request to keep the neutral core usable outside an ET runner.

Resolved / consistent with feedback ✅

• ET-free core (digantdesai):cache.h has zero ExecuTorch includes and no et:: dependency — it's a plain CacheBase + CacheConfig. The ET-specific Error/Result usage is confined to cache_registry.*, and the header comments state the split explicitly. This matches "keep the ET interfaces separate so this can be used in a non-delegate backend layer."
• abort → Result/Error:build() returns Result<...> via ET_CHECK_OR_RETURN_ERROR (logs + returns, no std::abort), and the test asserts .error() == Error::NotFound rather than a death test. Consistent with the resolution from the earlier rounds.
• CMake-only, no buck: matches the "buckify later" call. Test gating uses if(BUILD_TESTING), consistent with the other extension/*/CMakeLists.txt.

Observations on the current (slimmed) surface

1. CacheConfig fields are uninitialized (cache.h:31-35)

structCacheConfig {
int capacity;
int n_layers;
int initial_capacity = 512;
};

capacity and n_layers have no default initializer. The test aggregate-inits both (CacheConfig cfg{32, 1}), so it's safe there, but a caller doing CacheConfig cfg; gets garbage in the two most important fields. Given initial_capacity already has a default, adding = 0 to the other two keeps the struct uniformly safe to default-construct. Fix this →

2. The gemma4 mixed flat/ring design (from the last round) isn't in this PR
The LayerPolicy/LayerConfig per-layer config and the SeqStepPlan (write[2]/read[2] + read_base_pos) shapes discussed at length aren't here — CacheConfig is still the flat capacity/n_layers/initial_capacity triple. That's fine as a scoping decision (this PR is now just the registry + neutral anchor), but worth confirming that's intentional: my earlier note stands that the config and step-plan structs are the expensive surfaces to change later because the byte layer and any serialized metadata bind to them. If gemma4 is the first target, folding at least the LayerConfig array-of-structs into CacheConfig now is cheaper than after consumers exist. Not blocking for a registry-only PR, but flag it so it's a deliberate "next PR" rather than an oversight.

3. Builder-registry key concatenation (cache_registry.cpp:53,63)
The map key is backend_id + ":" + kind. If either token can contain :, "a:b" + ":" + "c" and "a" + ":" + "b:c" collide. Backend ids/kinds are internal constants today so it's theoretical, but a one-line note (or a std::pair key) removes the ambiguity if these ever become user-facing strings.

Minor / carried forward (all still fine at this size)

• Process-global singletons + test isolation:register_builder("TestBackend", "stub", ...) mutates the process-global CacheBuilderRegistry and persists for the test binary's lifetime. Serviceable at one test; watch it as more kinds register.
• int for capacity/layers: fine for these counts; be deliberate about widening if any of these ever feed byte-offset math in the pool layer.
• Registry comment (cache_registry.h:69): "Returns Error::NotFound" now correctly matches the Result return — good, the stale "throws" comment from earlier rounds is fixed.

Overall: the scope reduction is the right move and directly addresses the ET-independence concern. Nothing blocking. The one I'd fold in here is the CacheConfig default initializers (#1); the mixed-policy config/plan shapes (#2) are a scoping judgment call for the follow-up.
· kv-cache-shared-registry

[metascroy]

Use std::pair, other than that, LGTM! Let's land, we can always fix later