From a0c7e3d12e8324463a888b85566c5fc0861ef34b Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Mon, 10 Aug 2026 16:10:33 -0700 Subject: [PATCH 1/3] Add off-graph KV-cache cell reference: multi-sequence batching --- extension/llm/cache/reference_cache.py | 264 +++++++++++++++++- extension/llm/cache/test_update_and_attend.py | 196 +++++++++++++ 2 files changed, 458 insertions(+), 2 deletions(-) diff --git a/extension/llm/cache/reference_cache.py b/extension/llm/cache/reference_cache.py index f633fbcd144..ee0f2015705 100644 --- a/extension/llm/cache/reference_cache.py +++ b/extension/llm/cache/reference_cache.py @@ -20,14 +20,16 @@ The cache places K/V and returns the history plus an ``AttendSpec`` (a mask *semantic*). The attend mechanism (``attend`` below) is applied by the op/backend from that spec. -Scope for this initial slice: single sequence, contiguous placement, float KV. +Two caches share the op: ``ContiguousReferenceCache`` (one sequence appended in +place) and ``CellReferenceCache`` (many sequences over a pool of per-token cells, +with sharing and eviction). Both store float KV. """ from __future__ import annotations from dataclasses import dataclass from enum import Enum -from typing import List, Optional, Tuple +from typing import List, Optional, Sequence, Set, Tuple import torch import torch.nn.functional as F @@ -171,6 +173,264 @@ def _spec(q_len: int, total: int, device: torch.device) -> AttendSpec: return AttendSpec(kind=MaskKind.EXPLICIT, mask=offsets <= total - q_len) +# A cell's owners are a bitset in a torch int64, so bit 63 (the sign bit) is out. +MAX_SEQS = 63 + + +@dataclass +class _CellStepPlan: + """One step's allocation, shared by every layer of that forward.""" + + cells: torch.Tensor # [n_tok] long -- the cell each query token was given + mask: torch.Tensor # [n_tok, read_len] bool -- true = attend + + +@experimental( + "update_and_attend KV cache is experimental and may change without notice." +) +class CellReferenceCache: + """Per-cell KV history for several sequences sharing one pool. + + Each cell holds one token's K/V plus that token's position and the set of + sequences owning it, so a sequence need not be contiguous and two may share + cells -- a fork sets a second bit instead of copying K/V. Visibility is then + a property of the cell rather than of the layout: query i attends cell j iff + j is occupied, shares a sequence with i, and is no newer than i. No causal + alignment can express that, so the spec is always EXPLICIT. + + The batch is flat: tokens from every sequence sit on one axis with B = 1, + and sequence identity is supplied out-of-band. ``begin_step`` declares which + sequence each of the next forward's tokens belongs to; the positions arrive + with the forward itself, in the op's ``position`` tensor, so cells are + allocated on the first layer of the step and memoized for the rest of it. + + DYNAMIC sizing grows the pool to the occupied extent, so a short session + reserves a short pool rather than the whole context. Growth must keep every + cell's index and its bytes -- a cell's index is its name, held by the plan + and by ``_pos``/``_seq`` -- so it appends rows and never renumbers. + """ + + def __init__(self, config: CacheConfig): + if config.batch_size != 1: + raise ValueError( + "cell cache is flat on the token axis: batch_size must be 1" + ) + self.config = config + cap = config.capacity + self._pos: List[int] = [-1] * cap # per cell; -1 = free + self._seq: List[int] = [0] * cap # per cell; owning-sequence bitset + self._used_end = 0 # every occupied cell is in [0, used_end): the read window + h, d = config.n_kv_heads, config.head_dim + rows = cap if config.sizing == CacheSizing.STATIC else 0 + self._k = [ + torch.zeros(1, h, rows, d, dtype=config.dtype) + for _ in range(config.n_layers) + ] + self._v = [ + torch.zeros(1, h, rows, d, dtype=config.dtype) + for _ in range(config.n_layers) + ] + self._step_seq: List[int] = [] + self._declared = False # a declaration is spent by the step it allocates + self._plan: Optional[_CellStepPlan] = None + self._served: Set[int] = set() + + # -- runner face: admission, lifecycle, sequence verbs ------------------ + + def free_cells(self) -> int: + return self._pos.count(-1) + + def can_extend(self, n: int = 1) -> bool: + return self.free_cells() >= n + + def seq_len(self, seq: int) -> int: + bit = 1 << seq + return sum(1 for owners in self._seq if owners & bit) + + def begin_step(self, seq_ids: Sequence[int]) -> None: + """Declare the sequence each of the next forward's tokens belongs to. + + Admission is decided here, before the forward: the token count is known + without the positions, and cells are interchangeable, so a step that + passes this check cannot then fail to allocate. + """ + for seq in seq_ids: + if not 0 <= seq < MAX_SEQS: + raise ValueError(f"seq_id {seq} outside [0, {MAX_SEQS})") + if not self.can_extend(len(seq_ids)): + raise RuntimeError( + f"KV cache full: {len(seq_ids)} tokens need as many cells, " + f"{self.free_cells()} free" + ) + self._step_seq = list(seq_ids) + self._declared = True + self._plan = None + self._served.clear() + + def seq_cp(self, src: int, dst: int, p0: int = 0, p1: Optional[int] = None) -> None: + """Give dst a claim on src's cells -- a fork that copies no K/V.""" + src_bit, dst_bit = 1 << src, 1 << dst + for i in range(self._used_end): + if self._seq[i] & src_bit and self._in_range(self._pos[i], p0, p1): + self._seq[i] |= dst_bit + self._invalidate_plan() + + def seq_rm(self, seq: int, p0: int = 0, p1: Optional[int] = None) -> None: + """Drop seq's claim; a cell frees only once no sequence owns it.""" + bit = 1 << seq + for i in range(self._used_end): + if self._seq[i] & bit and self._in_range(self._pos[i], p0, p1): + self._seq[i] &= ~bit + if self._seq[i] == 0: + self._pos[i] = -1 + self._shrink() + self._invalidate_plan() + + def seq_keep(self, keep: int) -> None: + """Commit one sequence, freeing every cell no longer owned.""" + bit = 1 << keep + for i in range(self._used_end): + self._seq[i] &= bit + if self._seq[i] == 0: + self._pos[i] = -1 + self._shrink() + self._invalidate_plan() + + def reset(self): + self._pos = [-1] * self.config.capacity + self._seq = [0] * self.config.capacity + self._used_end = 0 + self._step_seq = [] + self._declared = False + self._plan = None + self._served.clear() + if self.config.sizing == CacheSizing.DYNAMIC: + h, d = self.config.n_kv_heads, self.config.head_dim + for i in range(self.config.n_layers): + self._k[i] = torch.zeros(1, h, 0, d, dtype=self.config.dtype) + self._v[i] = torch.zeros(1, h, 0, d, dtype=self.config.dtype) + + # -- op face ------------------------------------------------------------ + + def update_and_fetch( + self, + layer_id: int, + k: torch.Tensor, + v: torch.Tensor, + position: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, AttendSpec]: + """Scatter this step's K/V into its cells and return the read window. + + The first layer of a step allocates; the rest reuse that allocation, so + the cells and the mask are computed once per forward, not once per + layer. Args are as ``ContiguousReferenceCache.update_and_fetch``. + """ + if layer_id in self._served: + raise RuntimeError( + f"layer {layer_id} served twice for one step: " + "begin_step must precede every forward" + ) + if self._plan is None: + self._plan = self._allocate(position, k.device) + self._served.add(layer_id) + + read_len = self._plan.mask.shape[-1] + self._ensure(layer_id, read_len) + cells = self._plan.cells + self._k[layer_id][:, :, cells, :] = k.to(self.config.dtype) + self._v[layer_id][:, :, cells, :] = v.to(self.config.dtype) + return ( + self._k[layer_id][:, :, :read_len, :], + self._v[layer_id][:, :, :read_len, :], + AttendSpec(kind=MaskKind.EXPLICIT, mask=self._plan.mask), + ) + + # -- internals ---------------------------------------------------------- + + def _allocate(self, position: torch.Tensor, device: torch.device) -> _CellStepPlan: + if not self._declared: + raise RuntimeError( + "no step declared: begin_step must precede every forward" + ) + if position.shape[-1] != 1: + raise NotImplementedError( + "cell placement needs one position per token, got " + f"{position.shape[-1]}" + ) + positions = position.reshape(-1).tolist() + if len(positions) != len(self._step_seq): + raise ValueError( + f"begin_step declared {len(self._step_seq)} tokens, " + f"the forward carries {len(positions)}" + ) + cells = [ + self._claim(pos, 1 << seq) for pos, seq in zip(positions, self._step_seq) + ] + # Occupied, sharing a sequence, and no newer than the query. The step's + # own cells are already placed, so a query sees itself and any earlier + # token of its sequence in the same batch. + n = self._used_end + cell_pos = torch.tensor(self._pos[:n], device=device) + cell_seq = torch.tensor(self._seq[:n], device=device) + tok_pos = torch.tensor(positions, device=device).unsqueeze(-1) + tok_seq = torch.tensor( + [1 << seq for seq in self._step_seq], device=device + ).unsqueeze(-1) + mask = (cell_pos >= 0) & ((cell_seq & tok_seq) != 0) & (cell_pos <= tok_pos) + self._declared = False # one declaration, one allocation + return _CellStepPlan( + cells=torch.tensor(cells, dtype=torch.long, device=device), mask=mask + ) + + def _ensure(self, layer_id: int, rows: int) -> None: + """Make room for `rows` cells, doubling as the byte layer's pool does. + + Rows are appended, so a cell keeps the index it was claimed under and + the K/V already stored there stays where the plan expects it. + """ + have = self._k[layer_id].shape[2] + if rows <= have: + return + grown = max(have, 1) + while grown < rows: + grown *= 2 + grown = min(grown, self.config.capacity) + pad = torch.zeros( + 1, + self.config.n_kv_heads, + grown - have, + self.config.head_dim, + dtype=self.config.dtype, + ) + self._k[layer_id] = torch.cat([self._k[layer_id], pad], dim=2) + self._v[layer_id] = torch.cat([self._v[layer_id], pad.clone()], dim=2) + + def _claim(self, pos: int, owners: int) -> int: + # Lowest free cell, which keeps the read window tight. The byte layer + # keeps a free list rather than scanning. + for i in range(self.config.capacity): + if self._pos[i] < 0: + self._pos[i] = pos + self._seq[i] = owners + self._used_end = max(self._used_end, i + 1) + return i + raise RuntimeError("no free cell") # begin_step admitted the step + + def _shrink(self): + while self._used_end > 0 and self._pos[self._used_end - 1] < 0: + self._used_end -= 1 + + def _invalidate_plan(self): + # A mutated cell table leaves a built plan's cells and mask stale. The + # step protocol state is deliberately left alone: a mutation must not + # disguise a forward that skipped begin_step. + self._plan = None + + @staticmethod + def _in_range(pos: int, p0: int, p1: Optional[int]) -> bool: + return pos >= p0 and (p1 is None or pos < p1) + + def attend( q: torch.Tensor, k: torch.Tensor, diff --git a/extension/llm/cache/test_update_and_attend.py b/extension/llm/cache/test_update_and_attend.py index e1721891020..8bfa41e38d9 100644 --- a/extension/llm/cache/test_update_and_attend.py +++ b/extension/llm/cache/test_update_and_attend.py @@ -13,6 +13,7 @@ AttendSpec, CacheConfig, CacheSizing, + CellReferenceCache, ContiguousReferenceCache, MaskKind, ) @@ -258,6 +259,201 @@ def forward(self, q, k, v, position): self.assertEqual(tuple(node.meta["val"].shape), (1, 4, 3, 5)) +class CellCacheTest(unittest.TestCase): + # Many sequences over one pool of per-token cells, flat on the token axis. + # The baseline throughout is the cacheless model: whatever a sequence would + # have computed alone, it must still compute when batched beside others. + + CAPACITY = 32 + + def setUp(self): + torch.manual_seed(0) + self.n_layers, self.hidden = 2, 16 + self.n_heads, self.n_kv_heads, self.head_dim = 4, 2, 8 + self.model = TinyAttentionModel( + self.n_layers, + self.hidden, + self.n_heads, + self.n_kv_heads, + self.head_dim, + 40, + ).eval() + self.cache_key = "cells" + + def tearDown(self): + REGISTRY.uninstall(self.cache_key) + + def _cache(self, capacity=CAPACITY, sizing=CacheSizing.DYNAMIC): + cache = CellReferenceCache( + CacheConfig( + n_layers=self.n_layers, + n_kv_heads=self.n_kv_heads, + head_dim=self.head_dim, + capacity=capacity, + sizing=sizing, + ) + ) + REGISTRY.install(self.cache_key, cache) + return cache + + def _step(self, cache, x, positions, seqs): + """One forward carrying `x`, whose tokens have these positions/seqs.""" + cache.begin_step(seqs) + pos = torch.tensor(positions, dtype=torch.long).unsqueeze(-1) + with REGISTRY.active(self.cache_key): + return self.model(x, pos, torch.arange(x.shape[1])) + + def test_single_sequence_matches_baseline(self): + x = torch.randn(1, 5, self.hidden) + out = self._step(self._cache(), x, list(range(5)), [0] * 5) + torch.testing.assert_close( + out, self.model.reference_forward(x, torch.arange(5)), atol=1e-4, rtol=1e-4 + ) + + def test_batched_sequences_match_separate_runs(self): + # Two prefills in ONE forward. Each must equal what it computes alone, + # which is exactly the isolation the per-cell seq bitset buys. + a, b = torch.randn(1, 4, self.hidden), torch.randn(1, 3, self.hidden) + out = self._step( + self._cache(), + torch.cat([a, b], dim=1), + [0, 1, 2, 3, 0, 1, 2], + [0, 0, 0, 0, 1, 1, 1], + ) + torch.testing.assert_close( + out[:, :4, :], + self.model.reference_forward(a, torch.arange(4)), + atol=1e-4, + rtol=1e-4, + ) + torch.testing.assert_close( + out[:, 4:, :], + self.model.reference_forward(b, torch.arange(3)), + atol=1e-4, + rtol=1e-4, + ) + + def test_batched_decode_continues_each_sequence(self): + # Prefill both, then one forward carrying a new token for each. + a, b = torch.randn(1, 4, self.hidden), torch.randn(1, 3, self.hidden) + cache = self._cache() + self._step( + cache, + torch.cat([a[:, :3], b[:, :2]], dim=1), + [0, 1, 2, 0, 1], + [0, 0, 0, 1, 1], + ) + out = self._step(cache, torch.cat([a[:, 3:], b[:, 2:]], dim=1), [3, 2], [0, 1]) + torch.testing.assert_close( + out[:, 0, :], + self.model.reference_forward(a, torch.arange(4))[:, -1, :], + atol=1e-4, + rtol=1e-4, + ) + torch.testing.assert_close( + out[:, 1, :], + self.model.reference_forward(b, torch.arange(3))[:, -1, :], + atol=1e-4, + rtol=1e-4, + ) + + def test_fork_shares_cells_and_history(self): + trunk, tail = torch.randn(1, 4, self.hidden), torch.randn(1, 1, self.hidden) + cache = self._cache() + self._step(cache, trunk, [0, 1, 2, 3], [0] * 4) + + free_before = cache.free_cells() + cache.seq_cp(0, 1) + self.assertEqual(cache.free_cells(), free_before) # no cell, no byte copied + self.assertEqual(cache.seq_len(1), 4) + + out = self._step(cache, tail, [4], [1]) # the branch continues the trunk + torch.testing.assert_close( + out[:, 0, :], + self.model.reference_forward( + torch.cat([trunk, tail], dim=1), torch.arange(5) + )[:, -1, :], + atol=1e-4, + rtol=1e-4, + ) + + def test_seq_rm_frees_only_unowned_cells(self): + cache = self._cache() + self._step(cache, torch.randn(1, 3, self.hidden), [0, 1, 2], [0] * 3) + cache.seq_cp(0, 1) + + cache.seq_rm(0) + self.assertEqual(cache.seq_len(0), 0) + self.assertEqual(cache.seq_len(1), 3) # the fork still owns them + self.assertEqual(cache.free_cells(), self.CAPACITY - 3) + + cache.seq_rm(1) + self.assertEqual(cache.free_cells(), self.CAPACITY) + + def test_admission_fails_before_the_forward(self): + cache = self._cache(capacity=4) + self.assertFalse(cache.can_extend(5)) + with self.assertRaises(RuntimeError): + cache.begin_step([0] * 5) + + def test_step_protocol_is_enforced(self): + cache = self._cache() + kv = torch.randn(1, self.n_kv_heads, 1, self.head_dim) + pos = torch.tensor([[0]]) + + cache.begin_step([0, 0]) # declares two tokens, forward carries one + with self.assertRaises(ValueError): + cache.update_and_fetch(0, kv, kv, pos) + + cache.begin_step([0]) + cache.update_and_fetch(0, kv, kv, pos) + with self.assertRaises(RuntimeError): # a second step, no begin_step + cache.update_and_fetch(0, kv, kv, pos) + + def test_growth_keeps_cell_indices_and_bytes(self): + # A grown pool must append rows only: a cell's index is its name, held + # by the plan and by _pos/_seq, so renumbering or dropping rows would + # move history without anything noticing. + cache = self._cache() + first = torch.randn(1, self.n_kv_heads, 2, self.head_dim) + cache.begin_step([0, 0]) + k, _, _ = cache.update_and_fetch(0, first, first, torch.tensor([[0], [1]])) + self.assertEqual(k.shape[2], 2) # a short session reserves a short pool + + rest = torch.randn(1, self.n_kv_heads, 6, self.head_dim) + cache.begin_step([0] * 6) + k, v, _ = cache.update_and_fetch( + 0, rest, rest, torch.tensor([[p] for p in range(2, 8)]) + ) + self.assertEqual(k.shape[2], 8) + torch.testing.assert_close(k[:, :, :2, :], first) # cells 0,1 unmoved + torch.testing.assert_close(v[:, :, 2:, :], rest) + + def test_sizings_agree(self): + x = torch.randn(1, 5, self.hidden) + out = [ + self._step(self._cache(sizing=s), x, list(range(5)), [0] * 5) + for s in (CacheSizing.DYNAMIC, CacheSizing.STATIC) + ] + torch.testing.assert_close(out[0], out[1]) + + def test_a_verb_does_not_hide_a_missing_begin_step(self): + # A sequence verb drops the memoized plan, which must not be mistaken + # for the start of a step -- that would silently reuse the previous + # step's sequence assignment for the new tokens. + cache = self._cache() + kv = torch.randn(1, self.n_kv_heads, 2, self.head_dim) + pos = torch.tensor([[0], [0]]) + cache.begin_step([0, 1]) + cache.update_and_fetch(0, kv, kv, pos) + + cache.seq_rm(2) # any verb; a no-op here beyond dropping the plan + with self.assertRaises(RuntimeError): # layer 0 was already served + cache.update_and_fetch(0, kv, kv, pos) + with self.assertRaises(RuntimeError): # and the declaration is spent + cache.update_and_fetch(1, kv, kv, pos) + + class ContiguousSpecTest(unittest.TestCase): # Which mask semantic the cache declares for each shape of step. From 91c390119cb622df29a1a2cddca5b08d84dddf4e Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Wed, 12 Aug 2026 11:15:54 -0700 Subject: [PATCH 2/3] Address review: flatten_step helper, seq_id guards, tighter sequence verbs --- extension/llm/cache/reference_cache.py | 151 ++++++++++++------ extension/llm/cache/test_update_and_attend.py | 93 +++++++++-- 2 files changed, 185 insertions(+), 59 deletions(-) diff --git a/extension/llm/cache/reference_cache.py b/extension/llm/cache/reference_cache.py index ee0f2015705..7038243ad4c 100644 --- a/extension/llm/cache/reference_cache.py +++ b/extension/llm/cache/reference_cache.py @@ -29,7 +29,7 @@ from dataclasses import dataclass from enum import Enum -from typing import List, Optional, Sequence, Set, Tuple +from typing import List, Mapping, Optional, Sequence, Set, Tuple import torch import torch.nn.functional as F @@ -177,6 +177,44 @@ def _spec(q_len: int, total: int, device: torch.device) -> AttendSpec: MAX_SEQS = 63 +def flatten_step( + sequences: Mapping[int, Tuple[torch.Tensor, int]], +) -> Tuple[torch.Tensor, torch.Tensor, List[int], torch.Tensor]: + """Lay out one step's sequences on a single token axis. + + A step is flat: every sequence's tokens share one axis with B = 1, and the + per-token arrays must stay aligned. Building them together is what keeps + them so. + + It is a host helper, not part of the cache: the cache is handed only the + sequence ids, and never sees the tokens themselves. + + Args: + sequences: ``{seq_id: (tokens, start_pos)}`` -- each sequence's tokens + with the token axis second (``[1, n]`` ids, or ``[1, n, hidden]`` + where the model takes embeddings), and the position its first + token takes. + + Returns: + ``(tokens, positions, seq_ids, logits_indices)`` -- tokens concatenated + on the token axis and ``positions`` (``[n_tok, 1]``) as model inputs, + ``seq_ids`` for ``begin_step``, and ``logits_indices`` selecting each + sequence's last token, the rows worth running the LM head on. + """ + tokens, positions, seq_ids, logits_indices = [], [], [], [] + for seq_id, (toks, start_pos) in sequences.items(): + tokens.append(toks) + positions.extend(range(start_pos, start_pos + toks.shape[1])) + seq_ids.extend([seq_id] * toks.shape[1]) + logits_indices.append(len(seq_ids) - 1) + return ( + torch.cat(tokens, dim=1), + torch.tensor(positions, dtype=torch.long).unsqueeze(-1), + seq_ids, + torch.tensor(logits_indices, dtype=torch.long), + ) + + @dataclass class _CellStepPlan: """One step's allocation, shared by every layer of that forward.""" @@ -207,7 +245,7 @@ class CellReferenceCache: DYNAMIC sizing grows the pool to the occupied extent, so a short session reserves a short pool rather than the whole context. Growth must keep every cell's index and its bytes -- a cell's index is its name, held by the plan - and by ``_pos``/``_seq`` -- so it appends rows and never renumbers. + and by ``_pos``/``_owners`` -- so it appends rows and never renumbers. """ def __init__(self, config: CacheConfig): @@ -218,7 +256,7 @@ def __init__(self, config: CacheConfig): self.config = config cap = config.capacity self._pos: List[int] = [-1] * cap # per cell; -1 = free - self._seq: List[int] = [0] * cap # per cell; owning-sequence bitset + self._owners: List[int] = [0] * cap # per cell; owning-sequence bitset self._used_end = 0 # every occupied cell is in [0, used_end): the read window h, d = config.n_kv_heads, config.head_dim rows = cap if config.sizing == CacheSizing.STATIC else 0 @@ -230,8 +268,8 @@ def __init__(self, config: CacheConfig): torch.zeros(1, h, rows, d, dtype=config.dtype) for _ in range(config.n_layers) ] - self._step_seq: List[int] = [] - self._declared = False # a declaration is spent by the step it allocates + self._step_seq_ids: List[int] = [] + self._declared = False # set by begin_step, cleared by the step it authorizes self._plan: Optional[_CellStepPlan] = None self._served: Set[int] = set() @@ -241,11 +279,17 @@ def free_cells(self) -> int: return self._pos.count(-1) def can_extend(self, n: int = 1) -> bool: + """Whether `n` more tokens fit: cache-wide, one cell per token. + + The bound is on cells, so a prefix shared by several sequences counts + once and their lengths can sum past `capacity` while a step still fits. + """ return self.free_cells() >= n - def seq_len(self, seq: int) -> int: - bit = 1 << seq - return sum(1 for owners in self._seq if owners & bit) + def seq_len(self, seq_id: int) -> int: + self._check_seq_id(seq_id) + bit = 1 << seq_id + return sum(1 for owners in self._owners if owners & bit) def begin_step(self, seq_ids: Sequence[int]) -> None: """Declare the sequence each of the next forward's tokens belongs to. @@ -254,53 +298,58 @@ def begin_step(self, seq_ids: Sequence[int]) -> None: without the positions, and cells are interchangeable, so a step that passes this check cannot then fail to allocate. """ - for seq in seq_ids: - if not 0 <= seq < MAX_SEQS: - raise ValueError(f"seq_id {seq} outside [0, {MAX_SEQS})") + if not seq_ids: + raise ValueError("a step carries at least one token") + for seq_id in seq_ids: + self._check_seq_id(seq_id) if not self.can_extend(len(seq_ids)): raise RuntimeError( f"KV cache full: {len(seq_ids)} tokens need as many cells, " f"{self.free_cells()} free" ) - self._step_seq = list(seq_ids) + self._step_seq_ids = list(seq_ids) self._declared = True self._plan = None self._served.clear() - def seq_cp(self, src: int, dst: int, p0: int = 0, p1: Optional[int] = None) -> None: - """Give dst a claim on src's cells -- a fork that copies no K/V.""" - src_bit, dst_bit = 1 << src, 1 << dst - for i in range(self._used_end): - if self._seq[i] & src_bit and self._in_range(self._pos[i], p0, p1): - self._seq[i] |= dst_bit - self._invalidate_plan() + def seq_cp(self, src_id: int, dst_id: int, upto: Optional[int] = None) -> None: + """Give dst_id a claim on src_id's cells -- a fork that copies no K/V. - def seq_rm(self, seq: int, p0: int = 0, p1: Optional[int] = None) -> None: - """Drop seq's claim; a cell frees only once no sequence owns it.""" - bit = 1 << seq + Shares src_id's cells at positions below `upto`; None shares all of them, + forking at src_id's end. There is no lower bound: a shared cell keeps one + position, so what can be shared is a prefix, not an arbitrary range. + """ + self._check_seq_id(src_id) + self._check_seq_id(dst_id) + src_bit, dst_bit = 1 << src_id, 1 << dst_id for i in range(self._used_end): - if self._seq[i] & bit and self._in_range(self._pos[i], p0, p1): - self._seq[i] &= ~bit - if self._seq[i] == 0: - self._pos[i] = -1 - self._shrink() + if self._owners[i] & src_bit and (upto is None or self._pos[i] < upto): + self._owners[i] |= dst_bit self._invalidate_plan() - def seq_keep(self, keep: int) -> None: - """Commit one sequence, freeing every cell no longer owned.""" - bit = 1 << keep + def seq_rm(self, seq_id: int, p0: int = 0, p1: Optional[int] = None) -> None: + """Drop seq_id's claim on positions [p0, p1); p1 = None runs to the end. + + A cell frees only once no sequence owns it, so removing a shared range + reclaims nothing until the last owner lets go. seq_rm(s) drops the whole + sequence, seq_rm(s, 0, k) evicts its oldest k positions, and seq_rm(s, k) + truncates it at position k. + """ + self._check_seq_id(seq_id) + bit = 1 << seq_id for i in range(self._used_end): - self._seq[i] &= bit - if self._seq[i] == 0: - self._pos[i] = -1 + if self._owners[i] & bit and self._in_range(self._pos[i], p0, p1): + self._owners[i] &= ~bit + if self._owners[i] == 0: + self._pos[i] = -1 self._shrink() self._invalidate_plan() def reset(self): self._pos = [-1] * self.config.capacity - self._seq = [0] * self.config.capacity + self._owners = [0] * self.config.capacity self._used_end = 0 - self._step_seq = [] + self._step_seq_ids = [] self._declared = False self._plan = None self._served.clear() @@ -331,7 +380,7 @@ def update_and_fetch( "begin_step must precede every forward" ) if self._plan is None: - self._plan = self._allocate(position, k.device) + self._plan = self._allocate(position) self._served.add(layer_id) read_len = self._plan.mask.shape[-1] @@ -347,37 +396,40 @@ def update_and_fetch( # -- internals ---------------------------------------------------------- - def _allocate(self, position: torch.Tensor, device: torch.device) -> _CellStepPlan: + def _allocate(self, position: torch.Tensor) -> _CellStepPlan: + # The plan indexes and masks the pools, so it is built where they live. + device = self._k[0].device if not self._declared: raise RuntimeError( "no step declared: begin_step must precede every forward" ) + self._declared = False # one declaration, one attempt at allocating it if position.shape[-1] != 1: raise NotImplementedError( "cell placement needs one position per token, got " f"{position.shape[-1]}" ) positions = position.reshape(-1).tolist() - if len(positions) != len(self._step_seq): + if len(positions) != len(self._step_seq_ids): raise ValueError( - f"begin_step declared {len(self._step_seq)} tokens, " + f"begin_step declared {len(self._step_seq_ids)} tokens, " f"the forward carries {len(positions)}" ) cells = [ - self._claim(pos, 1 << seq) for pos, seq in zip(positions, self._step_seq) + self._claim(pos, 1 << seq_id) + for pos, seq_id in zip(positions, self._step_seq_ids) ] # Occupied, sharing a sequence, and no newer than the query. The step's # own cells are already placed, so a query sees itself and any earlier # token of its sequence in the same batch. n = self._used_end cell_pos = torch.tensor(self._pos[:n], device=device) - cell_seq = torch.tensor(self._seq[:n], device=device) + cell_owners = torch.tensor(self._owners[:n], device=device) tok_pos = torch.tensor(positions, device=device).unsqueeze(-1) - tok_seq = torch.tensor( - [1 << seq for seq in self._step_seq], device=device + tok_bit = torch.tensor( + [1 << seq_id for seq_id in self._step_seq_ids], device=device ).unsqueeze(-1) - mask = (cell_pos >= 0) & ((cell_seq & tok_seq) != 0) & (cell_pos <= tok_pos) - self._declared = False # one declaration, one allocation + mask = (cell_pos >= 0) & ((cell_owners & tok_bit) != 0) & (cell_pos <= tok_pos) return _CellStepPlan( cells=torch.tensor(cells, dtype=torch.long, device=device), mask=mask ) @@ -411,7 +463,7 @@ def _claim(self, pos: int, owners: int) -> int: for i in range(self.config.capacity): if self._pos[i] < 0: self._pos[i] = pos - self._seq[i] = owners + self._owners[i] = owners self._used_end = max(self._used_end, i + 1) return i raise RuntimeError("no free cell") # begin_step admitted the step @@ -426,6 +478,13 @@ def _invalidate_plan(self): # disguise a forward that skipped begin_step. self._plan = None + @staticmethod + def _check_seq_id(seq_id: int) -> None: + # An id past the bitset silently makes owners a Python big-int, which + # only surfaces much later as an int64 overflow building the mask. + if not 0 <= seq_id < MAX_SEQS: + raise ValueError(f"seq_id {seq_id} outside [0, {MAX_SEQS})") + @staticmethod def _in_range(pos: int, p0: int, p1: Optional[int]) -> bool: return pos >= p0 and (p1 is None or pos < p1) diff --git a/extension/llm/cache/test_update_and_attend.py b/extension/llm/cache/test_update_and_attend.py index 8bfa41e38d9..a574ede87d7 100644 --- a/extension/llm/cache/test_update_and_attend.py +++ b/extension/llm/cache/test_update_and_attend.py @@ -15,7 +15,9 @@ CacheSizing, CellReferenceCache, ContiguousReferenceCache, + flatten_step, MaskKind, + MAX_SEQS, ) from executorch.extension.llm.cache.update_and_attend import REGISTRY, update_and_attend @@ -314,12 +316,15 @@ def test_batched_sequences_match_separate_runs(self): # Two prefills in ONE forward. Each must equal what it computes alone, # which is exactly the isolation the per-cell seq bitset buys. a, b = torch.randn(1, 4, self.hidden), torch.randn(1, 3, self.hidden) - out = self._step( - self._cache(), - torch.cat([a, b], dim=1), - [0, 1, 2, 3, 0, 1, 2], - [0, 0, 0, 0, 1, 1, 1], - ) + cache = self._cache() + + # {seq_id: (tokens, start_pos)} -> the step's parallel arrays + tokens, positions, seq_ids, _ = flatten_step({0: (a, 0), 1: (b, 0)}) + cache.begin_step(seq_ids) + with REGISTRY.active(self.cache_key): + # every row, not one per sequence: each token is compared below + out = self.model(tokens, positions, torch.arange(tokens.shape[1])) + torch.testing.assert_close( out[:, :4, :], self.model.reference_forward(a, torch.arange(4)), @@ -334,16 +339,25 @@ def test_batched_sequences_match_separate_runs(self): ) def test_batched_decode_continues_each_sequence(self): - # Prefill both, then one forward carrying a new token for each. + # Prefill both, then one forward carrying a new token for each, laid + # out by flatten_step -- one logits row per sequence, not per token. a, b = torch.randn(1, 4, self.hidden), torch.randn(1, 3, self.hidden) cache = self._cache() - self._step( - cache, - torch.cat([a[:, :3], b[:, :2]], dim=1), - [0, 1, 2, 0, 1], - [0, 0, 0, 1, 1], + + tokens, positions, seq_ids, logits_indices = flatten_step( + {0: (a[:, :3], 0), 1: (b[:, :2], 0)} ) - out = self._step(cache, torch.cat([a[:, 3:], b[:, 2:]], dim=1), [3, 2], [0, 1]) + cache.begin_step(seq_ids) + with REGISTRY.active(self.cache_key): + self.model(tokens, positions, logits_indices) + + tokens, positions, seq_ids, logits_indices = flatten_step( + {0: (a[:, 3:], 3), 1: (b[:, 2:], 2)} + ) + cache.begin_step(seq_ids) + with REGISTRY.active(self.cache_key): + out = self.model(tokens, positions, logits_indices) + torch.testing.assert_close( out[:, 0, :], self.model.reference_forward(a, torch.arange(4))[:, -1, :], @@ -390,6 +404,54 @@ def test_seq_rm_frees_only_unowned_cells(self): cache.seq_rm(1) self.assertEqual(cache.free_cells(), self.CAPACITY) + def test_flatten_step_lays_out_the_parallel_arrays(self): + tokens, positions, seq_ids, logits_indices = flatten_step( + { + 0: (torch.zeros(1, 3, self.hidden), 5), + 1: (torch.ones(1, 2, self.hidden), 0), + } + ) + self.assertEqual(tokens.shape[1], 5) # one axis, both sequences + self.assertEqual(positions.squeeze(-1).tolist(), [5, 6, 7, 0, 1]) + self.assertEqual(seq_ids, [0, 0, 0, 1, 1]) + self.assertEqual(logits_indices.tolist(), [2, 4]) # each sequence's last + + def test_fork_at_a_position_shares_only_the_prefix(self): + cache = self._cache() + self._step(cache, torch.randn(1, 4, self.hidden), [0, 1, 2, 3], [0] * 4) + + cache.seq_cp(0, 1, upto=2) + self.assertEqual(cache.seq_len(0), 4) + self.assertEqual(cache.seq_len(1), 2) # only positions 0 and 1 + self.assertEqual(cache.free_cells(), self.CAPACITY - 4) # still no copy + + def test_seq_rm_over_a_range_frees_only_that_window(self): + cache = self._cache() + self._step(cache, torch.randn(1, 5, self.hidden), [0, 1, 2, 3, 4], [0] * 5) + + cache.seq_rm(0, 0, 2) # sliding window: drop the oldest two + self.assertEqual(cache.seq_len(0), 3) + self.assertEqual(cache.free_cells(), self.CAPACITY - 3) + + cache.seq_rm(0, 4) # backtrack: drop position 4 onwards + self.assertEqual(cache.seq_len(0), 2) + self.assertEqual(cache.free_cells(), self.CAPACITY - 2) + + def test_every_verb_range_checks_the_seq_id(self): + # An id past the bitset would set a bit no int64 can hold, surfacing + # much later as an overflow while building the mask. + cache = self._cache() + for call in ( + lambda: cache.begin_step([MAX_SEQS]), + lambda: cache.seq_cp(0, MAX_SEQS), + lambda: cache.seq_cp(MAX_SEQS, 0), + lambda: cache.seq_rm(MAX_SEQS), + lambda: cache.seq_len(MAX_SEQS), + lambda: cache.seq_len(-1), + ): + with self.assertRaises(ValueError): + call() + def test_admission_fails_before_the_forward(self): cache = self._cache(capacity=4) self.assertFalse(cache.can_extend(5)) @@ -401,9 +463,14 @@ def test_step_protocol_is_enforced(self): kv = torch.randn(1, self.n_kv_heads, 1, self.head_dim) pos = torch.tensor([[0]]) + with self.assertRaises(ValueError): # a step with no tokens + cache.begin_step([]) + cache.begin_step([0, 0]) # declares two tokens, forward carries one with self.assertRaises(ValueError): cache.update_and_fetch(0, kv, kv, pos) + with self.assertRaises(RuntimeError): # the failed attempt still cleared it + cache.update_and_fetch(0, kv, kv, torch.tensor([[0], [1]])) cache.begin_step([0]) cache.update_and_fetch(0, kv, kv, pos) From 490bfea25d34cb07bd4abc8f8e714ebe2cc2d089 Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Wed, 12 Aug 2026 14:37:12 -0700 Subject: [PATCH 3/3] Address review: test that freeing the tail shrinks the read window --- extension/llm/cache/test_update_and_attend.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/extension/llm/cache/test_update_and_attend.py b/extension/llm/cache/test_update_and_attend.py index a574ede87d7..dec27a8bcba 100644 --- a/extension/llm/cache/test_update_and_attend.py +++ b/extension/llm/cache/test_update_and_attend.py @@ -425,6 +425,23 @@ def test_fork_at_a_position_shares_only_the_prefix(self): self.assertEqual(cache.seq_len(1), 2) # only positions 0 and 1 self.assertEqual(cache.free_cells(), self.CAPACITY - 4) # still no copy + def test_freeing_the_tail_shrinks_the_read_window(self): + cache = self._cache() + kv = torch.randn(1, self.n_kv_heads, 4, self.head_dim) + cache.begin_step([0] * 4) + k, _, _ = cache.update_and_fetch(0, kv, kv, _positions(0, 4)) + self.assertEqual(k.shape[2], 4) # four cells held, so a window of four + + cache.seq_rm(0) # frees all four, so used_end walks back to 0 + self.assertEqual(cache.free_cells(), self.CAPACITY) + + # one token reclaims cell 0, so the window is its own single cell + kv = torch.randn(1, self.n_kv_heads, 1, self.head_dim) + cache.begin_step([1]) + k, _, spec = cache.update_and_fetch(0, kv, kv, torch.tensor([[0]])) + self.assertEqual(k.shape[2], 1) # the window length is 1, not the old 4 + self.assertEqual(spec.mask.shape[-1], 1) + def test_seq_rm_over_a_range_frees_only_that_window(self): cache = self._cache() self._step(cache, torch.randn(1, 5, self.hidden), [0, 1, 2, 3, 4], [0] * 5)