Skip to content

Commit b3b8c9a

Browse files
Add a batched-sequence KV cache to the eager reference (pytorch#22646)
### Summary Adds a third reference layout: a private history per sequence instead of one shared cell table. Attention splits the flat token axis into its declared spans and runs independently over each sequence's own history, so no sequence attends another and no cross-sequence mask is built. That means a step can be answered by more than one attention, so update_and_fetch returns a list of AttendSpec and the op attends each over the query tokens it covers. A spec now carries its own K/V and how many queries it answers; the offsets follow from the running total, and the op checks the specs cover the step. The single-history layouts return a one-element list and are behaviourally unchanged; most of the test diff is call sites reshaping around that. Read reference_cache.py first (AttendSpec, then BatchedSequenceReferenceCache), then update_and_attend.py for the loop, then the tests.
1 parent db3390d commit b3b8c9a

4 files changed

Lines changed: 687 additions & 133 deletions

File tree

‎backends/mlx/test/test_ops.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8598,14 +8598,14 @@ def compute_expected_outputs(self, model, test_inputs):
85988598
# an oracle cache installed for its duration.
85998599
from executorch.extension.llm.cache.reference_cache import (
86008600
CacheConfig,
8601-
ContiguousReferenceCache,
8601+
SequenceReferenceCache,
86028602
)
86038603
from executorch.extension.llm.cache.update_and_attend import REGISTRY
86048604

86058605
key = f"{self.name}-oracle"
86068606
REGISTRY.install(
86078607
key,
8608-
ContiguousReferenceCache(
8608+
SequenceReferenceCache(
86098609
CacheConfig(
86108610
n_layers=self.n_layers,
86118611
n_kv_heads=self.n_kv_heads,

0 commit comments

Comments
 (0)