Add off-graph KV-cache SequenceCache and flat policy

[kiymetakdemir]

Summary

Adds the neutral single-sequence cache: a SequenceCache controller that owns one logical length for the whole model and produces per-layer, integer-only layout plans, plus its full-history FlatPolicy. Backend and tensor-free the cache logic stays ET-independent, with a thin adapter for ExecuTorch callers. The SeqStepPlan and config shapes are already general for sliding-window policies, but only full-history flat is implemented here.

Files

• extension/llm/cache/cache.h — adds the two neutral faces recovered from CacheBase without RTTI (SequenceControl runner face; SequencePlanner backend face), the SeqStepPlan planner→backend handoff, the LayoutPolicy interface, and per-layer config (LayerPolicy, LayerConfig, CacheConfig.layers).
• extension/llm/cache/sequence_cache.h — the SequenceCache controller (single logical length, admission, rewind,
  per-layer plan dispatch with policy dedup) and FlatPolicy (full-history layout). Header-only, ET-independent.
• extension/llm/cache/cache_et.h — thin ExecuTorch adapter mapping the core's bool/std::optional results to
  Error/Result (et::plan → OutOfResources, et::rewind → InvalidArgument).
• extension/llm/cache/cache_registry.h / cache_registry.cpp — CacheSession gains control() (the typed runner face);
  builder registry now keyed by std::pair<backend_id, kind> instead of a concatenated string.
• extension/llm/cache/test/cache_test.cpp — tests for flat plan (append/read-all-history), admission/rewind, face
  recovery, the registry with a real SequenceCache, and the ET adapter.

Testing

Added C++ unit tests covering the flat bookkeeping (plan/can_extend/rewind/clear), face recovery, the registry with
a real cache, and the ET adapter's result/code mapping. 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/21412

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 139 Pending

As of commit a7438ce with merge base 739fe55 ():
💚 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]

@claude review this code

[claude]

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

---

Code Review — Off-graph KV-cache SequenceCache and flat policy

Read all six changed files plus the surrounding build wiring. Overall this is clean, well-documented, and the layering (neutral core → ET adapter → registry) is nicely separated. The tensor-free integer-only planner face is a good abstraction and the tests cover the happy paths well. A few correctness gaps and smells below.

Correctness

1. depth(int layer) is unchecked while plan(int layer, …) is guarded — sequence_cache.h:121-123
plan() validates layer against layer_to_policy_.size() and returns nullopt on out-of-range, but depth() indexes layer_to_policy_[layer] (and then policies_[…]) with no bounds check. A backend sizing pools by calling depth(layer) with a bad index gets silent out-of-bounds vector access (UB), and the SequencePlanner interface gives no way to signal failure from depth(). At minimum this asymmetry is surprising given how carefully plan() guards. Consider an assert/clamp, or documenting layer as a hard precondition for depth().
Fix this →

2. Constructor doesn't validate layers.size() — sequence_cache.h:61-69
The per-layer lookup is cfg.layers.size() == 1 ? front() : cfg.layers[l]. If layers is empty, or has size ≠ 1 and ≠ n_layers (e.g. 3 configs but n_layers == 8), the else branch reads cfg.layers[l] out of bounds. The header comment documents the "size 1 == uniform, else == n_layers" contract, but nothing enforces it. An assert(cfg.layers.size() == 1 || cfg.layers.size() == static_cast<size_t>(cfg.n_layers)) would turn a UB read into a clear failure.
Fix this →

3. plan() accepts non-contiguous / overlapping writes silently — sequence_cache.h:108-120
length_ = std::max(length_, position + T) commits the new length regardless of whether position == length_. A caller passing position ahead of the current length (leaving a gap of uninitialized rows) or behind it still succeeds, and for flat the read run [0, end) will read that gap. This is presumably a caller precondition, but unlike capacity and layer-range there's no guard and no comment stating "position must equal length for an append." Worth a one-line precondition note at least.

Concerns

4. SequenceCache instance is not thread-safe. The registry (cache_registry.cpp) locks its maps, and the header notes the cache is co-owned by the registry entry, the session, and the delegate handle. But plan()/rewind()/clear() all mutate length_ with no synchronization. If that's fine because generation is single-threaded, a sentence in the class comment saying so would prevent a future footgun.

5. Default argument on a virtual function — cache.h:46, sequence_cache.h:80
bool can_extend(int n = 1) declares the default in both the base and the override. Default args on virtuals bind statically to the declared type at the call site — harmless here since both say = 1, but it's a classic trap if they ever drift. Declaring the default only on the interface (and omitting it on the override) is the usual convention.

Minor

• sequence_cache.h:127 uses std::size_t and sequence_cache.h:98 uses std::max; <algorithm> is included but consider <cstddef> for std::size_t rather than relying on transitive inclusion.
• LayerConfig::n_kv_heads / head_dim (cache.h:109-110) are unused in this PR (integer-only core) and have no default initializers, so a default-constructed LayerConfig leaves them indeterminate. Fine if always aggregate-initialized as in the tests, but a = 0 default is cheap insurance.
• Tests are solid but don't cover multi-layer configs (n_layers > 1 with mixed/uniform policies) or the policy-dedup path (policy_index returning an existing index) — both are core logic in this PR. A two-layer test would exercise layer_to_policy_ broadcasting and dedup.

Nothing here blocks the direction; items 1–2 are the ones I'd fix before merge since they're latent UB on misuse.
· branch kv-cache-sequence-cache

[metascroy]

Where is the cpp impl file?

[metascroy]

Is this always read[0].start? If so, does it need to be another field?

[kiymetakdemir]

read[0].start is a physical index, read_base_pos is the logical sequence position. In flat it equals read[0].start, but in ring it will be the logical window start (length - window)

[metascroy]

What is depth?

[kiymetakdemir]

It is how many physical rows the backend should allocate for that layer's KV pool, it is capacity for flat but window for ring.

[metascroy]

Why does this take length?

[kiymetakdemir]

retained_from returns the oldest logical position the cache still holds at the current length, it is 0 for full-history flat, but will be length - window for a windowing policy

[metascroy]

I still don't follow what depth is. If it's related to specific memory layout (and not logical bookkeeping), let's leave it out here

[kiymetakdemir]

Yes it is, okay, I will remove it.