Repository navigation
Commit b3b8c9a
authored
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
- extension/llm/cache
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
8598 | 8598 | | |
8599 | 8599 | | |
8600 | 8600 | | |
8601 | | - | |
| 8601 | + | |
8602 | 8602 | | |
8603 | 8603 | | |
8604 | 8604 | | |
8605 | 8605 | | |
8606 | 8606 | | |
8607 | 8607 | | |
8608 | | - | |
| 8608 | + | |
8609 | 8609 | | |
8610 | 8610 | | |
8611 | 8611 | | |
| |||
0 commit comments