Add neutral KV-cache core: registry and bookkeeping#21383
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21383
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 94aa45e with merge base 30bf6af ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
e16a69c to
116bf40
Compare
This PR needs a
|
116bf40 to
94744ea
Compare
| int n_layers; | ||
| std::vector<int> n_kv_heads; | ||
| std::vector<int> head_dim; | ||
| int min_chunk = 512; |
There was a problem hiding this comment.
It is the initial capacity added a comment for this.
| // Integer-only handoff from bookkeeping to the backend byte layer: where to | ||
| // write this step and how much history to read. | ||
| struct StepPlan { | ||
| bool contiguous; |
There was a problem hiding this comment.
We don't use it actually, I removed.
|
|
||
| def define_common_targets(): | ||
| pass | ||
| """Defines targets that should be shared between fbcode and xplat. |
There was a problem hiding this comment.
Either stick with CMake only (we can buckify later), or if doing buck, let's import to internal to make sure it works
| // Contiguous single-sequence bookkeeping: history is a length counter; writes | ||
| // append at the current position and reads cover [0, length). Backends inherit | ||
| // this for the runner-facing methods and the per-step plan. | ||
| class ContiguousBookkeeping : public Cache { |
There was a problem hiding this comment.
Two things about this I'm hesitant about:
- I don't like the name Contiguous b/c it implies something about memory
- 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?
There was a problem hiding this comment.
Applied this together with claude's suggestion on explicit as_planner()/as_control().
|
@claude review this code |
|
Claude finished @metascroy's task in 2m 43s —— View job Code 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)
Consider mirroring the rest of the runtime: return 2. On @metascroy's audience-split proposal (cache.h:46)I agree with the direction. Today If you go this way, the registry's value type stays 3.
|
94744ea to
5ca9993
Compare
|
@claude review this code |
|
Claude finished @metascroy's task in 2m 3s —— View job Code 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 ✅
Remaining items1. Stale doc comment in // Throws if no builder is registered for (backend_id, kind).
std::shared_ptr<CacheBase> build(...) const;It no longer throws — 2. 3. Minor
Overall: the design questions from the last review are settled well, and the |
| std::vector<int> head_dim; | ||
| // Initial pool size in cells: the byte layer's lazy-doubling pool starts here | ||
| // and doubles up to `capacity`, so short sequences avoid reallocation. | ||
| int min_chunk = 512; |
There was a problem hiding this comment.
I think initial_capacity is a better name?
| length_ = 0; | ||
| } | ||
| void rewind(int new_len) override { | ||
| ET_CHECK_MSG( |
There was a problem hiding this comment.
What does ET_CHECK_MSG do on failure? Call std::terminate?
There was a problem hiding this comment.
It calls std::abort()
|
|
||
| // Integer-only handoff from the planner to the backend byte layer: where to | ||
| // write this step and how much history to read. | ||
| struct SeqStepPlan { |
There was a problem hiding this comment.
Are you planning to use this for Ring as well?
There was a problem hiding this comment.
Yes, I guess in that case it needs one more field for the window's start position.
|
Looks great! Main concern is calling std::abort. Instead wrap in Result instead. See how ET runtime / delegates avoids exception without calling abort. |
5ca9993 to
94aa45e
Compare
|
@claude review this code |
|
Claude finished @metascroy's task in 1m 31s —— View job Code Review
The last blocker — Resolved since last review ✅
Minor / non-blocking1. struct SeqStepPlan {
int valid_len;
int write_start;
int write_len;
};
2. 3. 4. Carried-forward nits (all still fine at this size):
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 |
| // n_kv_heads/head_dim are per-layer so hybrid / mixed-head / MLA models (KV | ||
| // dims vary by layer) need no config change; a single-element vector means | ||
| // uniform across all layers (the common case). | ||
| struct CacheConfig { |
There was a problem hiding this comment.
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;
};
|
|
||
| // Flat single-sequence cache: non-wrapping history [0, length) tracked by a | ||
| // length counter. Implements both faces; a backend byte layer builds on it. | ||
| class FlatCache : public CacheBase, |
There was a problem hiding this comment.
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
};
|
@claude review again in light of latest comments on policy |
|
Claude finished @metascroy's task in 2m 12s —— View job Review — the gemma4 mixed flat/ring policy proposal
I read the current 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 1. Dispatch key: prefer
|
2 participants
Summary
Adds a shared, backend-agnostic core for the off-graph KV cache: a set of C++ pieces that manage cache length/limits and let a runtime look up the right cache by a key. It has no backend or tensor code.
Files
Testing
Added C++ unit tests (test/cache_test.cpp) covering the bookkeeping (plan/can_extend/rewind/clear), the registry (install/get/erase, unique keys), the builder registry (build + missing-kind error), and the session's install-on-create / remove-on-destroy. Built and ran through the standard cmake/ctest flow, all pass:
cmake -B cmake-out -DEXECUTORCH_BUILD_EXTENSION_LLM=ON -DEXECUTORCH_BUILD_TESTS=ON
cmake --build cmake-out --target extension_llm_cache_test -j
ctest --test-dir cmake-out -R extension_llm_cache --output-on-failure