Add the off-graph KV-cache cell layout to the neutral C++ layer

[kiymetakdemir]

Summary

Adds the cell layout to the neutral C++ cache. Cells are addressed absolutely: each holds a position and a bitset of the sequences that own it, so a sequence need not be contiguous, a fork sets a second bit instead of copying K/V, and a cell frees when its bitset empties.

plan() places a step's tokens once and every later layer of that forward reuses the placement. It classifies: a single sequence appending a contiguous run at the tail of a window it owns outright gets a fused kind and a run write; anything else
— multi-sequence, forks, holes from removals — gets each token's cell plus a dense mask. Layers may window differently, so the mask is memoized per policy.

A step is declared before the forward through begin_step, which is also the admission gate: it reports whether the tokens fit before any compute is spent. A step may only extend its sequences, and one that would rewrite a position a sequence still holds is refused rather than stored twice.

Nothing constructs a CellCache yet — this is the neutral layer only. The existing SequenceCache is untouched behaviourally.

Files

• extension/llm/cache/cache.h — CacheControl factored out of SequenceControl to hold can_extend/capacity/clear; new BatchControl (the sequence verbs) and CellPlanner faces; MaskKind; the four as_* accessors now default to nullptr so a cache exposes only the faces it has.
• extension/llm/cache/cell_cache.h — new. CellStepPlan, CellPlanner, CellCache: placement, the seq verbs, classification, and the per-policy mask.
• extension/llm/cache/test/cache_test.cpp — 9 cell tests.

Testing

cmake --build cmake-out --target extension_llm_cache_test && ctest --test-dir cmake-out -R extension_llm_cache --output-on-failure

Covers the fused path, a second sequence forcing an explicit mask, one plan shared
across layers, a fork sharing cells and the refcounted free, ranged removal, holes
giving up the fused path, a rejected rewrite, layers sharing a window sharing a
plan, and the step protocol.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 Cancelled Job, 2 Unclassified Failures

As of commit 9369873 with merge base c461421 ():

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

• Test ARM Backend / test-arm-vgf / test-backend-linux (arm_vgf_fp, models) / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
• Test ARM Backend / test-arm-vgf / test-backend-linux (arm_vgf_fp, operators) / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)

CANCELLED JOB - The following job was cancelled. Please retry:

• pull / test-arm-backend-public-api-backward-compatibility / linux-job (gh)
  ##[error]The operation was canceled.

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]

Why raw pointer and not something like vector? What does n_tok mean?

Is begin_step the best name here for what this function does?

[kiymetakdemir]

True we don't need a pointer here, changing begin_step to take a std::vector<int32_t>.

[metascroy]

Why raw pointer?

[kiymetakdemir]

The positions live in the graph's position tensor, and the backend calls plan once per layer with that tensor's buffer. A vector would build a copy on every layer and 31 of them would go unread, since only the step's first layer looks at the positions and the rest reuse the placement.

[metascroy]

These comments are overly explanatory I think

[metascroy]

Can we move the implementation to a cpp file

[metascroy]

In the same way we had helper for python one, can we have for C++

[metascroy]

@claude review this code

[metascroy]

Design is sound and the invariant commentary is unusually careful — the cell/owner-bitset model and the "placement happens once per forward, mask memoized per policy" split both read well. Comments below are mostly about invariants the code claims but doesn't enforce, plus the new face pair being less disciplined than the SequenceControl/SequencePlanner pair it sits beside.

Blocking

1. fused() compares a position-window against a cell count (cell_cache.h:334) — produces a wrong mask when a sequence's positions have gaps. Reproducer inline.
2. assert(cell >= 0) compiles out over an OOB write (cell_cache.h:283).
3. Constructor indexes cfg.layers unchecked (cell_cache.h:76) — valid(cfg) is never called, and cache.h:185-188 explicitly warns this reads past the end.

Should fix before landing

4. seq_cp can leave one sequence owning two cells for one position (cell_cache.h:128) — the same invariant extends() exists to protect, unenforced on the other path.
5. begin_step's stated admission guarantee doesn't hold (cell_cache.h:189).
6. CellPlanner::plan — naming, return type, and pointer lifetime (cell_cache.h:54).
7. lowest_free() is O(capacity) per token (cell_cache.h:240).
8. Test coverage gaps, and the Cells::step helper hides which gate rejected a step (cache_test.cpp:261).

Smaller items (not filed inline)

API clarity

• SequencePlanner::plan returns std::optional, CellPlanner::plan returns a raw pointer. cache_et.h adapts the first to Error/Result and has nothing for BatchControl/CellPlanner, so the new face pair has no ET-facing error mapping at all.
• CacheBase::as_control() went from pure virtual to defaulting nullptr (cache.h:44-49). A cache that forgets a face now fails at runtime instead of compile time, and CacheSession::control() in cache_registry.h forwards it out with no null check — a CellCache in a session yields a null SequenceControl* silently. Not live yet, but worth the null check now. cache_registry.h's header comment is also stale (still says faces are recovered "via as_control()/as_planner()").
• CacheControl::capacity() is documented as "logical cap" (cache.h:64), but for CellCache it's a cell-pool size, and per-sequence lengths can legitimately sum past it once a prefix is shared via seq_cp. Worth calling out on the face — a runner will get this wrong when sizing.
• CellStepPlan (cell_cache.h:38-46) is a de-facto union: write_start is meaningful only when fused, cells/mask_bits only when Explicit, and write_start = -1 is a sentinel. Two structs or a variant would make the invalid states unrepresentable.
• free_cells()/used_end() (cell_cache.h:169-174) are public non-virtual accessors on the concrete class, used only by tests, on no face. Either they're contract or they're inspection — worth saying which.

Organization / DRY

• SeqInfo::min_cell and max_cell (cell_cache.h:203-204) are written in both claim() and rescan() and read nowhere. Dead — delete.
• SeqInfo is derived state maintained two ways: incrementally in claim() (cell_cache.h:288-292), wholesale in rescan() (cell_cache.h:260-277). The fast path buys little — rescan() is O(used_end_) and already runs on every verb, so one call after place() would collapse both. Same argument for used_count_, which could be derived.
• Three overlapping reset helpers (clear / invalidate_step / invalidate_plan), and begin_step hand-rolls two of invalidate_step's four actions inline (cell_cache.h:122-123) instead of calling it. The subtle part — that invalidate_plan() deliberately spares declared_/served_ so a mid-step verb can't disguise a forward that skipped begin_step — isn't written down. That's the one comment I'd actually add to this file.
• Agreeing with my earlier note on the cache.h:11-21 header: it's now a paragraph covering two cache kinds and the RTTI strategy. Two or three lines is plenty; the per-declaration comments are the ones carrying weight.
• build_mask (cell_cache.h:375-379) implicitly converts a bool expression to uint8_t — add the static_cast.
• int vs int32_t is inconsistent: pos_/cells_/step_pos_ are int32_t; CellStepPlan::write_start/read_len and seq_rm's p0/p1 are int. Pick one type for positions and one for cell indices.
• Per-step allocation churn: plan_for copies cells_ (cell_cache.h:322), build_mask allocates n_tok × used_end_ bytes (2 MB for a 512-token prefill into 4096 cells), and plans_ is a std::map cleared and rebuilt every step.
• targets.bzl is still def define_common_targets(): pass — none of this C++ is built or tested under Buck. Worth fixing here or filing.

On moving the implementation to a .cpp (my earlier comment): CMakeLists.txt already has add_library(extension_llm_cache cache_registry.cpp) and the install(DIRECTORY ... PATTERN "*.h") glob already picks up the header, so it's about a one-line change.

[metascroy]

Blocking — this produces a wrong mask.

window is measured in positions; used_end_ is a cell extent. Those only coincide when the sequence's positions are dense and ascending in cell order, and nothing enforces that — extends() below requires positions to be strictly increasing, not contiguous.

Cells c(16, {ring_layer(2)});
c.step({0}, {0});           // cell 0, pos 0
auto* p = c.step({0}, {5}); // cell 1, pos 5 -- legal, 5 > 0
// used_end_ == 2, window == 2 -> !(2 < 2)  -> passes
// info.count == 2 == used_end_             -> passes
// cells_ == {1} == tail run                -> passes
// => kind == MaskKind::None, read_len == 2

MaskKind::None tells the backend the query may attend the entire window, so the query at position 5 attends cell 0 (position 0) — five positions outside a window of 2.

The tail-run check at L352-357 does catch the hole-refill cases (I worked through several and couldn't break it that way), but it doesn't catch this one. The guard needs to be in position space, e.g. reject unless every occupied cell's pos_ is >= min(step positions) - window + 1.

Worth a test at each of window == used_end_, window == used_end_ - 1, window > used_end_, plus the gap case above.

[kiymetakdemir]

Fused now checks the span instead of the cell count: it scans the read window for the oldest and newest positions and refuses unless newest - oldest < window. Added CellFusesOnlyWhileTheWindowCoversTheSpan test cases.

[metascroy]

Blocking. Under NDEBUG this assert vanishes and the next line is pos_[-1] = pos — a heap write before the buffer.

The comment says only a used_count_ slip reaches here, but used_count_ is duplicated derived state (it's maintained incrementally in claim/seq_rm rather than derived), which is exactly the kind of thing that slips. place() is the only caller and plan() already has a nullptr failure path, so this can propagate instead of asserting.

ET runtime code also generally avoids <cassert> — worth dropping the include.

[metascroy]

Blocking. This indexes cfg.layers directly without ever calling valid(cfg). cache.h:185-188 warns about exactly this case:

Callers must check this before constructing a cache: the layers broadcast rule is indexed directly, so a list that is neither size 1 nor n_layers reads past the end.

An empty layers also makes .front() UB, since size() == 1 is false and it falls through to cfg.layers[l]. The Cells test helper doesn't call valid() either, so nothing in the PR exercises the guard. Since the config is caller-supplied, I'd guard in the ctor rather than rely on the precondition.

[kiymetakdemir]

valid(cfg) now runs in CacheBuilderRegistry::build. Both CellCache and SequenceCache constructors additionally assert(valid(cfg)).

[metascroy]

Two things here.

1. No guard that dst is position-disjoint from src. If dst already owns positions 0–3 and src owns 0–3, after seq_cp(src, dst, nullopt) dst owns 8 cells covering each position twice, and build_mask will let every dst query attend both copies — duplicated K/V in attention. seq_len(dst) then reports 8 while next_pos(dst) reports 4.

What makes this worth fixing rather than documenting away: extends() exists specifically to stop "a sequence owning two cells for one token" on the begin_step path (see its comment at L215-218). The identical invariant is simply unenforced here. Enforce it in both places, or state it as a caller precondition on BatchControl in cache.h:140-143.

2. Silent no-op on a bad seq id. Within this one class a bad seq id is reported three ways: begin_step returns false, plan() returns nullptr, and seq_cp/seq_rm return void and silently do nothing. A caller that typos a seq id gets no signal and finds out later as a wrong mask. These two should return bool.

[kiymetakdemir]

seq_cp now refuses when dst already holds slots. Since upto only bounds a prefix, a non-empty dst always overlaps its own positions, so a fork is only ever onto a fresh sequence. seq_cp and seq_rm return bool now.

[metascroy]

The admission guarantee in begin_step's doc doesn't hold. From cache.h:136-139:

the admission gate: false = rejected, nothing changed. Deciding it here means a step that passes cannot later fail to place.

But extends() runs here, not in begin_step, because positions don't exist until the forward. So begin_step can return true and plan() still return nullptr. Worse: declared_ is consumed on the line above, so the step isn't retryable — the caller has to go all the way back to begin_step, and the only signal it gets is a nullptr that's indistinguishable from "layer out of range" or "layer served twice".

Either take positions in begin_step so admission really is decided in one place, or soften the doc and give the caller a way to tell the failures apart.

[kiymetakdemir]

I softened the comment instead of supplying positions twice.

[metascroy]

Three things about this face, expanding on my raw-pointer question below.

Name. SequencePlanner::plan (cache.h:99-107) is const, pure, and paired with an explicit commit(). This plan() is non-const, mutates the cell table, and is the commit — as the comment above says, "placing the cells is what commits them." Same verb, same slot in the design, opposite contract. A reader will assume it's safe to call speculatively. place_step() or similar.

Return type. The sibling face returns std::optional<SeqStepPlan>; this returns a raw pointer, and nullptr collapses four distinct failures (no declaration, token-count mismatch, layer out of range, layer served twice) into one value. cache_et.h also has an adapter for the optional and nothing for this.

Lifetime. The returned pointer points into plans_, which invalidate_plan() clears on every begin_step/seq_cp/seq_rm/clear. So the backend's pointer dangles the moment the runner touches a verb, and that rule isn't documented anywhere. Related: std::map is load-bearing here for reference stability — if someone later swaps it for a vector this silently becomes a use-after-free. Worth a comment saying so.

[metascroy]

This restarts the scan from index 0 for every token, so place() is O(capacity × n_tok) — a 512-token prefill into a 4096-cell pool is ~2M iterations per forward, and it runs on the first layer of every step.

The comment above notes the placement choice is semantically free ("the mask keys off pos/owners, never the index"), so this is pure implementation — a free list, or even just a rolling cursor that survives across claim() calls within a step, fixes it without changing any result.

[kiymetakdemir]

Added a rolling cursor.

[metascroy]

This helper collapses two distinct failures into one nullptr: begin_step rejecting the step, and plan() rejecting it. That weakens the tests that depend on the distinction — in CellRejectsAPositionASequenceStillHolds, EXPECT_EQ(c.step({0}, {0}), nullptr) passes even if the admission gate had wrongly rejected it, which is the opposite of what that test means to assert.

ASSERT_TRUE the begin_step inside the helper and return only the plan. (This is also roughly what I was after with my other comment about a C++ helper.)

While here — row() just below indexes plan.mask_bits[...] unchecked, and on a fused plan mask_bits is empty, so that's UB the moment someone calls it on the wrong plan. An ASSERT_EQ(plan.kind, MaskKind::Explicit) would make it safe.

Coverage gaps, mapping the 9 tests against the branches in cell_cache.h — several of the riskiest paths have nothing on them:

• seq_cp with upto — the prefix-fork predicate at cell_cache.h:133 is never exercised; CellForkSharesCellsAndEvictionRefcounts only passes std::nullopt. A whole documented feature with zero coverage.
• fused() boundaries — six early-outs, one test touching them. See my comment on cell_cache.h:334.
• Windowed + multi-sequence together — every windowed test uses one sequence, every multi-sequence test is flat, so the oldest bound in build_mask is never exercised alongside the owner-bit check.
• clear() — never called.
• can_extend() / capacity() — never called directly, only reached through begin_step.
• Out-of-range layer in plan() — neither negative nor >= n_layers.
• A verb invalidating a memoized plan — invalidate_plan() from seq_cp/seq_rm is untested, and that's the subtle step-protocol behaviour worth pinning.
• seq_cp/seq_rm with a bad seq id — the silent no-op path.
• mask_bits.size() == n_tok * read_len is never asserted anywhere.

[kiymetakdemir]

• Fused path removed. It only worked when one sequence owned the whole read window.
• Mask is uint8_t, std::vector is bit-packed with no data(), so it can't hand a buffer to the backend. The byte layer converts once when it builds the array.
• step_for returns a reference, place_step a pointer. step_for is private and always produces a step, place_step has to express refusal, which is what the nullptr carries.
• seq_cp is now seq_clone(src, upto), returning a new id, and seq_new() allocates one. Ids now come from the cache, and declare_step refuses any it didn't hand out.

[metascroy]

Do users call these? Or are they private methods?

[kiymetakdemir]

next_pos is call ed to get the position to feed the model next, instead of tracking a counter.
seq_len doesn't have a caller yet, could be useful if scheduler has an eviction/trimming policy for long sequences. free_cells and used_end are only used by tests.

[metascroy]

@claude review this code

[claude]

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

---

Code review — cell layout (updated pass)

• Read the changed files (cache.h, cell_cache.h/.cpp, tests, registry)
• Verify the fixes made in response to the earlier review
• Check remaining open items
• Post review

This is a re-review against the current tip (2728718). The five commits close nearly all of my earlier blocking/should-fix items — I'll confirm those first, then flag what's left.

Resolved since the last round

• Fused mask bug (was blocking Add support for quantized LeakyReLU #1). The whole fused path is gone; every step now emits an explicit position-space mask via build_mask, and the window bound (oldest = tok_pos - window + 1) is computed in position space, so the window-vs-used_end_ category error can't recur. CellWindowBoundsEachQueryByPosition pins the <, ==, and sparse cases. Good.
• assert over an OOB write (was blocking Re-sync with internal repository #2). place() now checks cell < 0 and returns false before any claim(), so there's no write past the buffer to guard. ✓
• Unchecked cfg.layers (was blocking Rename _pt2e to pt2e #3). valid(cfg) now runs in CacheBuilderRegistry::build (cache_registry.cpp:74). ✓ (caveat below)
• seq_cp disjointness (was Add unlifting pass under private config #4). seq_clone always allocates a fresh dst via seq_new(), so the destination is provably empty and the double-owned-position case is unrepresentable. Cleaner than enforcing the invariant. ✓
• Admission/retryability ([build] Add BUCK.oss #5). The !extends() path returns nullptr before declared_ = false, so a rejected placement leaves the declaration standing and is retryable — CellRejectsAPositionASequenceStillHolds exercises exactly this. ✓
• plan face (Re-sync with internal repository #6). Renamed to place_step, and the header now documents lifetime ("owned by the cache and valid until the next verb") and the std::map reference-stability requirement. ✓
• lowest_free O(capacity·n_tok) (change internal code to only import from _quantize_pt2e #7). The rolling from cursor in place() makes it O(capacity + n_tok) per step. ✓
• Test helper (Executorch OSS CI tracking issue (2023) #8). step() now EXPECT_TRUEs declare_step and returns only the placement, so admission and placement failures no longer collapse. Coverage is much broader — seq_clone with upto, clear(), windowed+multi-seq, a verb invalidating a memoized step, out-of-range layers, and bad-seq-id bool returns are all covered now.

Remaining items

1. Dead SeqInfo::min_cell / max_cell. Still written in both claim() (cell_cache.cpp:244-245) and rescan() (cpp:222,227-228) and read nowhere — I grepped the whole tree, only writes. This was called out last round; it's leftover. Delete both fields and their assignments. Fix this →

2. build_mask implicit bool→uint8_t (cell_cache.cpp:294). row[j] = (pos_[j] >= 0 && ...) narrows a bool expression into uint8_t. It's correct but reads as accidental; a static_cast<uint8_t>(...) makes the intent explicit. (Also flagged last round.)

3. CacheSession::control() has no null check (cache_registry.h:105-107). as_control() now defaults to nullptr, so a CellCache bound into a session yields a null SequenceControl* that's forwarded silently. Not live yet, but a null check (or at least a comment that a batch cache must go through a different accessor) would turn a future segfault into a diagnosable error. The header comment at cache_registry.h:16 is also stale — it still says faces are recovered "via as_control()/as_planner()", but there are four faces now including as_batch_control()/as_cell_stepper().

4. Direct-construction validation is assert-only. valid(cfg) is enforced in the builder registry, but direct construction (the Cells test helper, and any non-registry caller) only has assert(valid(cfg)), which compiles out under NDEBUG — the exact scenario blocking #3 warned about, just moved to the direct path. Since the ctor indexes cfg.layers directly, consider either clamping defensively in the loop or keeping the precondition but dropping <cassert> (ET runtime code generally avoids it) in favor of the existing ET_CHECK idiom.

5. Minor / non-blocking

• int vs int32_t is still mixed: CellStep::length/read_len are int, cells is int32_t, seq_rm's p0/p1 are int. Picking one type for positions and one for indices would help.
• capacity() is documented as "logical cap" (cache.h:54), but for CellCache it's the cell-pool size and per-sequence lengths can legitimately sum past it once a prefix is shared via seq_clone. Worth a one-line note on BatchControl so a runner doesn't size against it.
• Per-step allocation churn is unchanged (full n_tok × used_end_ mask, steps_ map rebuilt each step). Fine for landing the neutral layer; worth a follow-up if prefill masks get large.

None of the remaining items are correctness-blocking — the design and invariants are sound and the test suite now covers the paths that worried me. Item #1 (dead fields) is the only thing I'd insist on before landing since it was already raised.
· branch kvcache-cell-faces

[metascroy]

Address the nits like "Dead SeqInfo::min_cell / max_cell" before landing