mlxrunner: drive MTP speculation through cache snapshots
Speculation used a parallel hierarchy of wrapper cache types that shadowed the live caches and reconciled against them on commit. Replace it with snapshot/restore on the live caches themselves: a cache snapshots itself as a write crosses each offset, and the runner commits a batched draft by restoring to the accepted count. The wrappers and the comparison plumbing around them are gone. Snapshots are lazy. A KV or rotating capture indexes into the live buffer and owns no memory until a destructive write forces a copy-out, so rejecting a draft is free. Recurrent layers now validate in the same batched pass rather than falling back to serial. A gated-delta layer reports its interior split offsets and hands back the recurrent state at each one, which the cache records as a snapshot.
This commit is contained in:
Jesse Gross
170
x/mlxrunner/cache/cache.go
vendored
170
x/mlxrunner/cache/cache.go
vendored
@@ -1,9 +1,9 @@
|
||||
package cache
|
||||
|
||||
import (
|
||||
"github.com/ollama/ollama/x/mlxrunner/batch"
|
||||
"fmt"
|
||||
|
||||
"github.com/ollama/ollama/x/mlxrunner/mlx"
|
||||
"github.com/ollama/ollama/x/models/nn"
|
||||
)
|
||||
|
||||
// Cache is common state management shared by every cache kind. Writers
|
||||
@@ -18,6 +18,26 @@ type Cache interface {
|
||||
// pinned VRAM arrays. The active cache is unchanged.
|
||||
Snapshot(fromOffset int) Snapshot
|
||||
|
||||
// PrepareSnapshots schedules the cache to capture a snapshot as its
|
||||
// storage offset reaches each listed offset during subsequent writes.
|
||||
// Offsets are storage offsets, must not already be passed (current
|
||||
// offset <= offset), and must be sorted ascending and unique. Scheduled
|
||||
// offsets persist across multiple writes until TakeSnapshots is called;
|
||||
// captures accumulate. The previous schedule must be drained first.
|
||||
//
|
||||
// Unlike Snapshot, capture happens at the moment a write crosses the
|
||||
// scheduled offset, so it can record states interior to a batched write
|
||||
// without the caller breaking the write into pieces.
|
||||
PrepareSnapshots(offsets []int)
|
||||
|
||||
// TakeSnapshots returns the snapshots captured since PrepareSnapshots,
|
||||
// one per scheduled offset in the caller's order, and clears the
|
||||
// schedule. An entry is nil when its scheduled offset captured a
|
||||
// zero-width range (the offset equalled the previous boundary, e.g. an
|
||||
// offset scheduled at the current position): rolling back there needs
|
||||
// only a live rewind, so there is nothing to page out.
|
||||
TakeSnapshots() []Snapshot
|
||||
|
||||
// Restore brings the cache to target. If snapshot is nil, rewinds
|
||||
// using the cache's own live state. Returns false if the target is
|
||||
// unreachable (e.g. target > current offset, or negative).
|
||||
@@ -35,94 +55,96 @@ type Cache interface {
|
||||
|
||||
// Snapshot is paged-out cache state that can be restored later.
|
||||
type Snapshot interface {
|
||||
// Size returns the byte size of the paged-out data (in VRAM).
|
||||
// Size returns the byte size of the paged-out data (in VRAM). A lazy
|
||||
// snapshot that still indexes a live cache buffer returns 0 — it owns
|
||||
// no extra memory yet. Once materialized (the cache copies the range
|
||||
// out before overwriting its slots), Size returns the owned bytes.
|
||||
Size() int
|
||||
|
||||
// SetMaterializeHook installs a callback fired once when a lazy
|
||||
// snapshot materializes (allocates its owned arrays). delta is the
|
||||
// newly-allocated byte count. Pass nil to detach. Snapshots that are
|
||||
// never lazy may treat this as a no-op.
|
||||
SetMaterializeHook(func(delta int))
|
||||
|
||||
// Close unpins the snapshot's arrays so they can be freed by Sweep.
|
||||
Close()
|
||||
}
|
||||
|
||||
// Viewer exposes a read-only attention history for a cache.
|
||||
type Viewer interface {
|
||||
View(b *batch.Batch) *nn.KVHistory
|
||||
// pendingSnapshots holds the per-token snapshot capture state shared by all
|
||||
// cache kinds. The owning cache calls capture(offset) from its write path
|
||||
// after each token's storage is in place; capture materializes a snapshot
|
||||
// for every scheduled offset that the write has now reached.
|
||||
type pendingSnapshots struct {
|
||||
offsets []int // scheduled storage offsets, in caller order
|
||||
captured []Snapshot // captured[i] corresponds to offsets[i]; nil until reached
|
||||
base int // running capture cursor: the from-offset of the next capture
|
||||
}
|
||||
|
||||
type speculativeCommitter interface {
|
||||
Cache
|
||||
commit(n int)
|
||||
}
|
||||
|
||||
// Speculation is an isolated cache transaction for speculative target
|
||||
// validation. Updates record generated K/V without mutating the live caches;
|
||||
// Commit appends only the accepted prefix to the live caches.
|
||||
type Speculation struct {
|
||||
layers []speculativeCommitter
|
||||
}
|
||||
|
||||
// BeginSpeculation returns cache wrappers suitable for a speculative target
|
||||
// forward. The returned caches must only be used for that forward.
|
||||
func BeginSpeculation(caches []Cache) ([]Cache, *Speculation, bool) {
|
||||
specCaches := make([]Cache, len(caches))
|
||||
layers := make([]speculativeCommitter, len(caches))
|
||||
|
||||
for i, c := range caches {
|
||||
switch c := c.(type) {
|
||||
case nil:
|
||||
case *RotatingKVCache:
|
||||
sc := newSpeculativeRotatingKVCache(c)
|
||||
specCaches[i] = sc
|
||||
layers[i] = sc
|
||||
case *KVCache:
|
||||
sc := newSpeculativeKVCache(c)
|
||||
specCaches[i] = sc
|
||||
layers[i] = sc
|
||||
default:
|
||||
return nil, nil, false
|
||||
// prepare schedules the listed storage offsets. See Cache.PrepareSnapshots.
|
||||
func (p *pendingSnapshots) prepare(currentOffset int, offsets []int) {
|
||||
if p.captured != nil {
|
||||
panic("PrepareSnapshots: previous schedule not drained; TakeSnapshots first or its captures leak")
|
||||
}
|
||||
for i, o := range offsets {
|
||||
if o < currentOffset {
|
||||
panic(fmt.Sprintf("PrepareSnapshots: offset %d already passed (current %d)", o, currentOffset))
|
||||
}
|
||||
// captureReached and scheduledIn walk offsets in storage order and rely
|
||||
// on it being ascending and unique.
|
||||
if i > 0 && o <= offsets[i-1] {
|
||||
panic(fmt.Sprintf("PrepareSnapshots: offsets must be sorted and unique, got %v", offsets))
|
||||
}
|
||||
}
|
||||
|
||||
return specCaches, &Speculation{layers: layers}, true
|
||||
p.offsets = append([]int(nil), offsets...)
|
||||
p.captured = make([]Snapshot, len(offsets))
|
||||
p.base = currentOffset
|
||||
}
|
||||
|
||||
// BeginIsolatedSpeculation returns cache wrappers that never mutate live cache
|
||||
// state. It is intended for correctness instrumentation, not the hot path.
|
||||
func BeginIsolatedSpeculation(caches []Cache) ([]Cache, bool) {
|
||||
specCaches := make([]Cache, len(caches))
|
||||
// take returns the captured snapshots and clears the schedule. See
|
||||
// Cache.TakeSnapshots. Captures fire in ascending offset order, so by take time
|
||||
// every scheduled offset the writes crossed has been visited; a nil entry is a
|
||||
// zero-width capture, not a missed one.
|
||||
func (p *pendingSnapshots) take() []Snapshot {
|
||||
out := p.captured
|
||||
p.offsets, p.captured = nil, nil
|
||||
return out
|
||||
}
|
||||
|
||||
for i, c := range caches {
|
||||
switch c := c.(type) {
|
||||
case nil:
|
||||
case *RotatingKVCache:
|
||||
specCaches[i] = newSpeculativeRotatingKVCache(c)
|
||||
case *KVCache:
|
||||
specCaches[i] = newIsolatedKVCache(c)
|
||||
default:
|
||||
return nil, false
|
||||
// captureReached materializes a snapshot for every scheduled offset that equals
|
||||
// reached and hasn't been captured yet, using snap to produce the rollback
|
||||
// state. The capture cursor base holds the previous scheduled boundary: snap
|
||||
// reads it (yielding a [base, reached) range for position-sliceable caches),
|
||||
// then base advances to reached so the next capture starts there. base only
|
||||
// advances when a capture actually fires — write boundaries between scheduled
|
||||
// offsets must not move it, or a later capture would record a fromOffset past
|
||||
// the previous scheduled offset and break the alignment between each capture's
|
||||
// range and the trie edge the caller attaches it to.
|
||||
func (p *pendingSnapshots) captureReached(reached int, snap func(offset int) Snapshot) {
|
||||
captured := false
|
||||
for i, o := range p.offsets {
|
||||
if p.captured[i] == nil && o == reached {
|
||||
p.captured[i] = snap(o)
|
||||
captured = true
|
||||
}
|
||||
}
|
||||
|
||||
return specCaches, true
|
||||
if captured {
|
||||
p.base = reached
|
||||
}
|
||||
}
|
||||
|
||||
// Commit appends the accepted prefix from the speculative forward to the live
|
||||
// caches. The target bonus token is intentionally not committed.
|
||||
func (s *Speculation) Commit(n int) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
for _, layer := range s.layers {
|
||||
if layer != nil {
|
||||
layer.commit(n)
|
||||
// scheduledIn returns the scheduled offsets in [start, end], ascending. A write
|
||||
// spanning [start, end) drives captureReached at each of these so any capture
|
||||
// due there fires; for position-sliceable caches, the resulting snapshots cover
|
||||
// [previous-scheduled-offset, this-offset) via captureReached's running cursor.
|
||||
// p.offsets is ascending and unique (prepare enforces it), so the range filter
|
||||
// preserves both.
|
||||
func (p *pendingSnapshots) scheduledIn(start, end int) []int {
|
||||
var boundaries []int
|
||||
for _, o := range p.offsets {
|
||||
if o >= start && o <= end {
|
||||
boundaries = append(boundaries, o)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func concatKV(prev, next *mlx.Array) *mlx.Array {
|
||||
if prev == nil {
|
||||
return next
|
||||
}
|
||||
return prev.Concatenate(2, next)
|
||||
}
|
||||
|
||||
func prefixKV(a *mlx.Array, n int) *mlx.Array {
|
||||
return a.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, n), mlx.Slice())
|
||||
return boundaries
|
||||
}
|
||||
|
||||
296
x/mlxrunner/cache/kvcache.go
vendored
296
x/mlxrunner/cache/kvcache.go
vendored
@@ -1,6 +1,8 @@
|
||||
package cache
|
||||
|
||||
import (
|
||||
"slices"
|
||||
|
||||
"github.com/ollama/ollama/x/mlxrunner/batch"
|
||||
"github.com/ollama/ollama/x/mlxrunner/mlx"
|
||||
"github.com/ollama/ollama/x/models/nn"
|
||||
@@ -14,12 +16,27 @@ type Attention interface {
|
||||
// Update appends (k, v) and returns an opaque nn.KVHistory for
|
||||
// this layer's SDPA.
|
||||
Update(b *batch.Batch, keys, values *mlx.Array) *nn.KVHistory
|
||||
|
||||
// View returns the current attention history without writing.
|
||||
View(b *batch.Batch) *nn.KVHistory
|
||||
}
|
||||
|
||||
type KVCache struct {
|
||||
keys, values *mlx.Array
|
||||
offset int
|
||||
step int
|
||||
|
||||
snapshots pendingSnapshots
|
||||
|
||||
// lazySnapshots index into the live keys/values buffer rather than owning a
|
||||
// copy (see kvSnapshot); the cache copies them out before overwriting or
|
||||
// freeing the slots they name.
|
||||
lazySnapshots []*kvSnapshot
|
||||
|
||||
// rewound is set when a restore moves offset backward. The buffer is
|
||||
// append-only, so only an append after a rewind can clobber a still-lazy
|
||||
// snapshot.
|
||||
rewound bool
|
||||
}
|
||||
|
||||
func NewKVCache() *KVCache {
|
||||
@@ -28,25 +45,32 @@ func NewKVCache() *KVCache {
|
||||
|
||||
// Assumes B = 1; heterogeneous batches are not supported.
|
||||
func (c *KVCache) Update(_ *batch.Batch, keys, values *mlx.Array) *nn.KVHistory {
|
||||
start := c.offset
|
||||
newK, newV := c.appendKV(keys, values)
|
||||
c.captureLazySnapshots(start, c.offset)
|
||||
return nn.NewKVHistory(newK, newV, nil)
|
||||
}
|
||||
|
||||
// View returns the current cache contents as attention history without writing.
|
||||
func (c *KVCache) View(_ *batch.Batch) *nn.KVHistory {
|
||||
state := c.State()
|
||||
if len(state) < 2 {
|
||||
return nil
|
||||
}
|
||||
return nn.NewKVHistory(state[0], state[1], nil)
|
||||
}
|
||||
|
||||
// appendKV is the raw write path shared by Update and Restore.
|
||||
func (c *KVCache) appendKV(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) {
|
||||
B, H, L, Dk, Dv := keys.Dim(0), keys.Dim(1), keys.Dim(2), keys.Dim(3), values.Dim(3)
|
||||
|
||||
prev := c.offset
|
||||
|
||||
// This write fills slots [prev, prev+L). Only an append after a rewind can
|
||||
// land on slots a still-lazy snapshot names and overwrite its data, so copy
|
||||
// the overlapping snapshots out first. copyOut removes the snapshot from
|
||||
// c.lazySnapshots, so range over a clone to avoid skipping entries as the
|
||||
// slice shrinks.
|
||||
if c.rewound {
|
||||
for _, s := range slices.Clone(c.lazySnapshots) {
|
||||
if s.fromOffset < prev+L && s.toOffset > prev {
|
||||
s.copyOut()
|
||||
}
|
||||
}
|
||||
c.rewound = false
|
||||
}
|
||||
|
||||
// Grow buffer if needed
|
||||
if c.keys == nil || (prev+L) > c.keys.Dim(2) {
|
||||
steps := (c.step + L - 1) / c.step
|
||||
@@ -74,6 +98,12 @@ func (c *KVCache) appendKV(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) {
|
||||
c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, c.offset), mlx.Slice())
|
||||
}
|
||||
|
||||
// View returns the current cache contents as attention history without writing.
|
||||
func (c *KVCache) View(_ *batch.Batch) *nn.KVHistory {
|
||||
state := c.State()
|
||||
return nn.NewKVHistory(state[0], state[1], nil)
|
||||
}
|
||||
|
||||
func (c *KVCache) State() []*mlx.Array {
|
||||
if c.keys == nil || c.values == nil {
|
||||
return nil
|
||||
@@ -84,37 +114,115 @@ func (c *KVCache) State() []*mlx.Array {
|
||||
}
|
||||
}
|
||||
|
||||
func (c *KVCache) PrepareSnapshots(offsets []int) { c.snapshots.prepare(c.offset, offsets) }
|
||||
func (c *KVCache) TakeSnapshots() []Snapshot { return c.snapshots.take() }
|
||||
|
||||
// captureLazySnapshots records edge-local snapshots for the scheduled offsets the
|
||||
// write [start, end) reached. A KVCache snapshot is a pure index into the
|
||||
// contiguous append-only buffer, so the write needs no segmenting: one appendKV
|
||||
// lays down [start, end), then each scheduled offset o captures the edge
|
||||
// [base, o) by arithmetic, where base is the previous boundary (running cursor).
|
||||
// Offsets are ascending, so the edges match what segmentation would produce. An
|
||||
// offset scheduled at start captures a zero-width range and stays nil; rolling
|
||||
// back there is a live rewind.
|
||||
func (c *KVCache) captureLazySnapshots(start, end int) {
|
||||
for _, o := range c.snapshots.scheduledIn(start, end) {
|
||||
c.snapshots.captureReached(o, func(int) Snapshot { return c.lazySnapshot(c.snapshots.base, o) })
|
||||
}
|
||||
}
|
||||
|
||||
// kvSnapshot holds paged-out KV data for a range [fromOffset, toOffset).
|
||||
//
|
||||
// A snapshot is initially lazy: keys/values are nil and the data lives in the
|
||||
// issuing cache's buffer at [fromOffset, toOffset). It costs nothing to capture
|
||||
// and, holding no MLX handle on that buffer, never blocks the in-place append
|
||||
// donation. The cache copies the range into owned keys/values (copyOut) before
|
||||
// it overwrites or frees those slots, after which the snapshot is independent
|
||||
// and cache is nil.
|
||||
type kvSnapshot struct {
|
||||
keys, values *mlx.Array
|
||||
fromOffset, toOffset int
|
||||
cache *KVCache // issuer while lazy; nil once copied out
|
||||
|
||||
// onMaterialize, if set, is fired once from copyOut with the newly-owned
|
||||
// byte count so an owner (e.g. the trie's pagedOutBytes counter) can pick
|
||||
// up bytes that were free while the snapshot was lazy.
|
||||
onMaterialize func(delta int)
|
||||
}
|
||||
|
||||
func (s *kvSnapshot) Size() int { return s.keys.NumBytes() + s.values.NumBytes() }
|
||||
func (s *kvSnapshot) Close() { mlx.Unpin(s.keys, s.values) }
|
||||
|
||||
func (c *KVCache) Snapshot(fromOffset int) Snapshot {
|
||||
if c.keys == nil || c.offset <= fromOffset {
|
||||
return nil
|
||||
func (s *kvSnapshot) Size() int {
|
||||
if s.keys != nil {
|
||||
return s.keys.NumBytes() + s.values.NumBytes()
|
||||
}
|
||||
from := max(0, fromOffset)
|
||||
to := c.offset
|
||||
// Lazy snapshots own no extra memory: the range still lives in the
|
||||
// issuing cache's buffer.
|
||||
return 0
|
||||
}
|
||||
|
||||
kSlice := c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(from, to), mlx.Slice())
|
||||
vSlice := c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(from, to), mlx.Slice())
|
||||
func (s *kvSnapshot) SetMaterializeHook(fn func(delta int)) { s.onMaterialize = fn }
|
||||
|
||||
func (s *kvSnapshot) Close() {
|
||||
mlx.Unpin(s.keys, s.values)
|
||||
if s.cache != nil {
|
||||
s.cache.dropLazySnapshot(s)
|
||||
s.cache = nil
|
||||
}
|
||||
}
|
||||
|
||||
// copyOut converts a lazy snapshot into an owned [fromOffset, toOffset) copy. It
|
||||
// is a no-op once the snapshot already owns its data. The copy is an independent
|
||||
// MLX handle on its own bytes, so a following in-place write to the live buffer
|
||||
// reallocates rather than mutating data the snapshot still names.
|
||||
func (s *kvSnapshot) copyOut() {
|
||||
if s.keys != nil {
|
||||
return
|
||||
}
|
||||
c := s.cache
|
||||
kSlice := c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.fromOffset, s.toOffset), mlx.Slice())
|
||||
vSlice := c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.fromOffset, s.toOffset), mlx.Slice())
|
||||
kCopy := mlx.Contiguous(kSlice, false)
|
||||
vCopy := mlx.Contiguous(vSlice, false)
|
||||
mlx.Pin(kCopy, vCopy)
|
||||
mlx.AsyncEval(kCopy, vCopy)
|
||||
|
||||
return &kvSnapshot{
|
||||
keys: kCopy,
|
||||
values: vCopy,
|
||||
fromOffset: from,
|
||||
toOffset: to,
|
||||
s.keys, s.values = kCopy, vCopy
|
||||
c.dropLazySnapshot(s)
|
||||
s.cache = nil
|
||||
|
||||
if s.onMaterialize != nil {
|
||||
s.onMaterialize(s.keys.NumBytes() + s.values.NumBytes())
|
||||
s.onMaterialize = nil
|
||||
}
|
||||
}
|
||||
|
||||
func (c *KVCache) addLazySnapshot(s *kvSnapshot) { c.lazySnapshots = append(c.lazySnapshots, s) }
|
||||
|
||||
func (c *KVCache) dropLazySnapshot(s *kvSnapshot) {
|
||||
if i := slices.Index(c.lazySnapshots, s); i >= 0 {
|
||||
c.lazySnapshots = slices.Delete(c.lazySnapshots, i, i+1)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *KVCache) Snapshot(fromOffset int) Snapshot {
|
||||
return c.lazySnapshot(fromOffset, c.offset)
|
||||
}
|
||||
|
||||
// lazySnapshot records a lazy [fromOffset, toOffset) snapshot indexing into the
|
||||
// live buffer. It returns nil for an empty range (the zero-width edge of an
|
||||
// offset scheduled at the current position).
|
||||
func (c *KVCache) lazySnapshot(fromOffset, toOffset int) Snapshot {
|
||||
if c.keys == nil || toOffset <= fromOffset {
|
||||
return nil
|
||||
}
|
||||
s := &kvSnapshot{
|
||||
fromOffset: fromOffset,
|
||||
toOffset: toOffset,
|
||||
cache: c,
|
||||
}
|
||||
c.addLazySnapshot(s)
|
||||
return s
|
||||
}
|
||||
|
||||
func (c *KVCache) Restore(snapshot Snapshot, target int) bool {
|
||||
if target < 0 {
|
||||
return false
|
||||
@@ -125,6 +233,7 @@ func (c *KVCache) Restore(snapshot Snapshot, target int) bool {
|
||||
return false
|
||||
}
|
||||
c.offset = target
|
||||
c.rewound = true
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -134,8 +243,24 @@ func (c *KVCache) Restore(snapshot Snapshot, target int) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// Rewind to snapshot start, then feed snapshot.
|
||||
// A lazy snapshot still in our own set indexes data that is, by construction,
|
||||
// still live in our buffer at [fromOffset, toOffset): appendKV copies out any
|
||||
// lazy snapshot before overwriting its slots, so one that has stayed lazy was
|
||||
// never clobbered.
|
||||
if snap.cache == c && snap.keys == nil {
|
||||
c.offset = min(target, snap.toOffset)
|
||||
c.rewound = true
|
||||
return true
|
||||
}
|
||||
|
||||
// Own the data before feeding it: appendKV mutates the buffer a lazy snapshot
|
||||
// may still index into, so copy out first (no-op if already owned).
|
||||
snap.copyOut()
|
||||
|
||||
// Rewind to snapshot start, then feed snapshot. The rewind may expose other
|
||||
// outstanding lazy snapshots to the appendKV write, so flag it for the scan.
|
||||
c.offset = snap.fromOffset
|
||||
c.rewound = true
|
||||
c.appendKV(snap.keys, snap.values)
|
||||
|
||||
// Clamp to target if needed (target may be less than full snapshot).
|
||||
@@ -159,6 +284,21 @@ func (c *KVCache) Merge(parent, child Snapshot) Snapshot {
|
||||
p := parent.(*kvSnapshot)
|
||||
ch := child.(*kvSnapshot)
|
||||
|
||||
// Two adjacent lazy snapshots into the same live buffer merge by arithmetic:
|
||||
// the combined range [p.from, ch.to) is a single contiguous snapshot, no copy.
|
||||
if p.keys == nil && ch.keys == nil && p.cache == ch.cache && p.toOffset == ch.fromOffset {
|
||||
merged := &kvSnapshot{fromOffset: p.fromOffset, toOffset: ch.toOffset, cache: p.cache}
|
||||
p.cache.addLazySnapshot(merged)
|
||||
p.Close()
|
||||
ch.Close()
|
||||
return merged
|
||||
}
|
||||
|
||||
// At least one is an owned copy in its own buffer: concatenate so Restore's
|
||||
// single-array appendKV sees one buffer. Own both first.
|
||||
p.copyOut()
|
||||
ch.copyOut()
|
||||
|
||||
mk := p.keys.Concatenate(2, ch.keys)
|
||||
mv := p.values.Concatenate(2, ch.values)
|
||||
mlx.Pin(mk, mv)
|
||||
@@ -189,6 +329,16 @@ func (c *KVCache) Split(snapshot Snapshot, at int) (Snapshot, Snapshot) {
|
||||
return snapshot, nil
|
||||
}
|
||||
|
||||
// Lazy: split is pure arithmetic into two adjacent lazy snapshots.
|
||||
if snap.keys == nil {
|
||||
p := &kvSnapshot{fromOffset: snap.fromOffset, toOffset: at, cache: snap.cache}
|
||||
ch := &kvSnapshot{fromOffset: at, toOffset: snap.toOffset, cache: snap.cache}
|
||||
snap.cache.addLazySnapshot(p)
|
||||
snap.cache.addLazySnapshot(ch)
|
||||
snap.Close()
|
||||
return p, ch
|
||||
}
|
||||
|
||||
pk := mlx.Contiguous(snap.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, splitIdx), mlx.Slice()), false)
|
||||
pv := mlx.Contiguous(snap.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, splitIdx), mlx.Slice()), false)
|
||||
ck := mlx.Contiguous(snap.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(splitIdx, seqLen), mlx.Slice()), false)
|
||||
@@ -214,97 +364,17 @@ func (c *KVCache) Split(snapshot Snapshot, at int) (Snapshot, Snapshot) {
|
||||
}
|
||||
|
||||
func (c *KVCache) Free() {
|
||||
// Freeing drops the buffer every lazy snapshot indexes into; own their data
|
||||
// first so they survive independently. copyOut drops the snapshot from
|
||||
// c.lazySnapshots, so iterate over a clone to avoid skipping.
|
||||
for _, s := range slices.Clone(c.lazySnapshots) {
|
||||
s.copyOut()
|
||||
}
|
||||
mlx.Unpin(c.keys, c.values)
|
||||
c.keys, c.values = nil, nil
|
||||
c.offset = 0
|
||||
c.rewound = false
|
||||
c.snapshots = pendingSnapshots{}
|
||||
}
|
||||
|
||||
func (c *KVCache) Offset() int { return c.offset }
|
||||
|
||||
type speculativeBase struct {
|
||||
offset int
|
||||
}
|
||||
|
||||
func (s *speculativeBase) Free() {}
|
||||
func (s *speculativeBase) Offset() int { return s.offset }
|
||||
func (s *speculativeBase) Snapshot(int) Snapshot { return nil }
|
||||
func (s *speculativeBase) Restore(Snapshot, int) bool { return false }
|
||||
func (s *speculativeBase) Merge(parent, child Snapshot) Snapshot { return nil }
|
||||
func (s *speculativeBase) Split(snapshot Snapshot, at int) (Snapshot, Snapshot) {
|
||||
return nil, snapshot
|
||||
}
|
||||
|
||||
type speculativeKVCache struct {
|
||||
speculativeBase
|
||||
target *KVCache
|
||||
start int
|
||||
end int
|
||||
}
|
||||
|
||||
func newSpeculativeKVCache(target *KVCache) *speculativeKVCache {
|
||||
return &speculativeKVCache{
|
||||
speculativeBase: speculativeBase{offset: target.Offset()},
|
||||
target: target,
|
||||
start: target.Offset(),
|
||||
end: target.Offset(),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *speculativeKVCache) Update(b *batch.Batch, keys, values *mlx.Array) *nn.KVHistory {
|
||||
history := c.target.Update(b, keys, values)
|
||||
c.offset = c.target.Offset()
|
||||
c.end = c.target.Offset()
|
||||
return history
|
||||
}
|
||||
|
||||
func (c *speculativeKVCache) State() []*mlx.Array {
|
||||
return c.target.State()
|
||||
}
|
||||
|
||||
func (c *speculativeKVCache) commit(n int) {
|
||||
target := max(c.start, c.start+n)
|
||||
if target > c.end {
|
||||
target = c.end
|
||||
}
|
||||
c.target.offset = target
|
||||
c.offset = target
|
||||
}
|
||||
|
||||
type isolatedKVCache struct {
|
||||
speculativeBase
|
||||
target *KVCache
|
||||
keys, values *mlx.Array
|
||||
}
|
||||
|
||||
func newIsolatedKVCache(target *KVCache) *isolatedKVCache {
|
||||
return &isolatedKVCache{
|
||||
speculativeBase: speculativeBase{offset: target.Offset()},
|
||||
target: target,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *isolatedKVCache) Update(_ *batch.Batch, keys, values *mlx.Array) *nn.KVHistory {
|
||||
c.keys = concatKV(c.keys, keys)
|
||||
c.values = concatKV(c.values, values)
|
||||
c.offset += keys.Dim(2)
|
||||
|
||||
state := c.target.State()
|
||||
if len(state) < 2 {
|
||||
return nn.NewKVHistory(c.keys, c.values, nil)
|
||||
}
|
||||
return nn.NewKVHistory(state[0].Concatenate(2, c.keys), state[1].Concatenate(2, c.values), nil)
|
||||
}
|
||||
|
||||
func (c *isolatedKVCache) State() []*mlx.Array {
|
||||
if c.keys == nil || c.values == nil {
|
||||
return c.target.State()
|
||||
}
|
||||
state := c.target.State()
|
||||
if len(state) < 2 {
|
||||
return []*mlx.Array{c.keys, c.values}
|
||||
}
|
||||
return []*mlx.Array{
|
||||
state[0].Concatenate(2, c.keys),
|
||||
state[1].Concatenate(2, c.values),
|
||||
}
|
||||
}
|
||||
|
||||
339
x/mlxrunner/cache/lazy_test.go
vendored
Normal file
339
x/mlxrunner/cache/lazy_test.go
vendored
Normal file
@@ -0,0 +1,339 @@
|
||||
package cache
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/ollama/ollama/x/mlxrunner/mlx"
|
||||
)
|
||||
|
||||
// distinctKV builds a [1, H, L, D] keys/values pair whose values encode the
|
||||
// absolute token position, so a restored range can be checked against the
|
||||
// positions it claims. Keys at position p are filled with float(p); values
|
||||
// with float(-p).
|
||||
func distinctKV(start, L, H, D int) (*mlx.Array, *mlx.Array) {
|
||||
ks := make([]float32, H*L*D)
|
||||
vs := make([]float32, H*L*D)
|
||||
for h := range H {
|
||||
for l := range L {
|
||||
for d := range D {
|
||||
i := (h*L+l)*D + d
|
||||
ks[i] = float32(start + l)
|
||||
vs[i] = -float32(start + l)
|
||||
}
|
||||
}
|
||||
}
|
||||
return mlx.FromValues(ks, 1, H, L, D), mlx.FromValues(vs, 1, H, L, D)
|
||||
}
|
||||
|
||||
// firstKeyAt returns the keys value stored at sequence position p (channel 0).
|
||||
func firstKeyAt(arr *mlx.Array, p, D int) float32 {
|
||||
return arr.Floats()[p*D]
|
||||
}
|
||||
|
||||
// settledActiveMemory drains unpinned arrays and the allocator cache, then
|
||||
// reports active (allocated, in-use) bytes.
|
||||
func settledActiveMemory() int {
|
||||
mlx.Sweep()
|
||||
mlx.ClearCache()
|
||||
return mlx.ActiveMemory()
|
||||
}
|
||||
|
||||
// TestKVSpeculationCaptureAllocatesNothing verifies that an MTP-style capture —
|
||||
// schedule per-token offsets, run one batched write, take the snapshots, rewind
|
||||
// via Restore(nil), and Close — allocates no snapshot buffers, because every
|
||||
// snapshot stays lazy and is discarded before any overwrite. Compare against the
|
||||
// bytes an eager per-token copy would cost.
|
||||
func TestKVSpeculationCaptureAllocatesNothing(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
const before, draft, H, D = 16, 8, 4, 8
|
||||
|
||||
c := NewKVCache()
|
||||
fillKV(c, before)
|
||||
|
||||
offsets := make([]int, draft)
|
||||
for i := range offsets {
|
||||
offsets[i] = before + i
|
||||
}
|
||||
c.PrepareSnapshots(offsets)
|
||||
|
||||
k, v := batchKV(draft)
|
||||
c.Update(newKVBatch(before, draft), k, v)
|
||||
|
||||
baseline := settledActiveMemory()
|
||||
|
||||
snaps := c.TakeSnapshots()
|
||||
// Every captured snapshot is a lazy snapshot (no owned buffer).
|
||||
for i, s := range snaps {
|
||||
if s == nil {
|
||||
continue
|
||||
}
|
||||
if ks := s.(*kvSnapshot); ks.keys != nil {
|
||||
t.Fatalf("snaps[%d] owns a buffer at capture; want a lazy snapshot", i)
|
||||
}
|
||||
}
|
||||
|
||||
// MTP commit: rewind to a partial accept, then discard all snapshots.
|
||||
if !c.Restore(nil, before+draft/2) {
|
||||
t.Fatal("live rewind failed")
|
||||
}
|
||||
for _, s := range snaps {
|
||||
if s != nil {
|
||||
s.Close()
|
||||
}
|
||||
}
|
||||
|
||||
after := settledActiveMemory()
|
||||
// Lazy snapshots allocate nothing; allow a tiny slack for allocator noise but
|
||||
// well under one per-token copy (draft tokens * keys+values).
|
||||
perToken := (c.keys.NumBytes() + c.values.NumBytes()) / c.keys.Dim(2)
|
||||
if after-baseline > perToken {
|
||||
t.Fatalf("capture allocated %d bytes (> one token %d); lazy snapshots should allocate nothing",
|
||||
after-baseline, perToken)
|
||||
}
|
||||
}
|
||||
|
||||
// TestKVLazySnapshotSizeZeroUntilMaterialized verifies the accounting contract:
|
||||
// a lazy snapshot reports Size() == 0 (it owns no extra memory), and once a
|
||||
// destructive write triggers copyOut the materialize hook fires with the
|
||||
// newly-allocated bytes and Size() reports the owned arrays.
|
||||
func TestKVLazySnapshotSizeZeroUntilMaterialized(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
const H, D = 4, 8
|
||||
|
||||
c := NewKVCache()
|
||||
k, v := distinctKV(0, 10, H, D)
|
||||
c.Update(newKVBatch(0, 10), k, v)
|
||||
|
||||
snap := c.Snapshot(5).(*kvSnapshot)
|
||||
defer snap.Close()
|
||||
if snap.Size() != 0 {
|
||||
t.Fatalf("lazy snapshot Size = %d, want 0", snap.Size())
|
||||
}
|
||||
|
||||
var hookDelta int
|
||||
snap.SetMaterializeHook(func(delta int) { hookDelta = delta })
|
||||
|
||||
// Rewind and overwrite to force copyOut.
|
||||
if !c.Restore(nil, 5) {
|
||||
t.Fatal("rewind failed")
|
||||
}
|
||||
nk, nv := distinctKV(100, 3, H, D)
|
||||
c.Update(newKVBatch(5, 3), nk, nv)
|
||||
|
||||
if snap.keys == nil {
|
||||
t.Fatal("snapshot was not materialized by overwriting write")
|
||||
}
|
||||
want := snap.keys.NumBytes() + snap.values.NumBytes()
|
||||
if hookDelta != want {
|
||||
t.Fatalf("hook fired with delta %d, want %d (owned bytes)", hookDelta, want)
|
||||
}
|
||||
if snap.Size() != want {
|
||||
t.Fatalf("materialized Size = %d, want %d", snap.Size(), want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestKVLazySnapshotCopiedOutOnOverwrite verifies that after a rewind, a write
|
||||
// that overwrites a lazy snapshot's slots copies the snapshot out first, so it
|
||||
// still reads the pre-overwrite data — both keys and values, across the whole
|
||||
// captured range (guarding the Slice+Contiguous copy-out representation).
|
||||
func TestKVLazySnapshotCopiedOutOnOverwrite(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
const H, D = 4, 8
|
||||
|
||||
c := NewKVCache()
|
||||
// Fill [0,10) with position-encoded values (keys p, values -p).
|
||||
k, v := distinctKV(0, 10, H, D)
|
||||
c.Update(newKVBatch(0, 10), k, v)
|
||||
|
||||
// Lazy snapshot [5,10).
|
||||
snap := c.Snapshot(5).(*kvSnapshot)
|
||||
if snap.keys != nil {
|
||||
t.Fatal("Snapshot should return a lazy snapshot")
|
||||
}
|
||||
|
||||
// Rewind to 5 and overwrite [5,8) with different data.
|
||||
if !c.Restore(nil, 5) {
|
||||
t.Fatal("rewind failed")
|
||||
}
|
||||
nk, nv := distinctKV(100, 3, H, D) // positions 100,101,102
|
||||
c.Update(newKVBatch(5, 3), nk, nv)
|
||||
|
||||
// The overwrite must have copied the lazy snapshot out beforehand, preserving
|
||||
// the pre-overwrite keys and values across the whole [5,10) range.
|
||||
if snap.keys == nil {
|
||||
t.Fatal("lazy snapshot was not copied out before the overwriting write")
|
||||
}
|
||||
mlx.Eval(snap.keys, snap.values)
|
||||
keys := snap.keys.Floats()
|
||||
vals := snap.values.Floats()
|
||||
for l := range snap.toOffset - snap.fromOffset {
|
||||
wantK := float32(snap.fromOffset + l)
|
||||
if got := keys[l*D]; got != wantK {
|
||||
t.Fatalf("snapshot keys[%d] = %v, want %v (pre-overwrite data)", l, got, wantK)
|
||||
}
|
||||
if got := vals[l*D]; got != -wantK {
|
||||
t.Fatalf("snapshot values[%d] = %v, want %v (pre-overwrite data)", l, got, -wantK)
|
||||
}
|
||||
}
|
||||
snap.Close()
|
||||
}
|
||||
|
||||
// TestKVLazySnapshotCopiedOutOnFree verifies Free copies out outstanding lazy snapshots
|
||||
// so they survive after the cache buffer is gone.
|
||||
func TestKVLazySnapshotCopiedOutOnFree(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
const H, D = 4, 8
|
||||
|
||||
c := NewKVCache()
|
||||
k, v := distinctKV(0, 10, H, D)
|
||||
c.Update(newKVBatch(0, 10), k, v)
|
||||
|
||||
snap := c.Snapshot(5).(*kvSnapshot)
|
||||
c.Free()
|
||||
|
||||
if snap.keys == nil {
|
||||
t.Fatal("lazy snapshot was not copied out on Free")
|
||||
}
|
||||
mlx.Eval(snap.keys)
|
||||
if got := firstKeyAt(snap.keys, 0, D); got != 5 {
|
||||
t.Fatalf("snapshot[0] key = %v, want 5 (data preserved through Free)", got)
|
||||
}
|
||||
snap.Close()
|
||||
}
|
||||
|
||||
// TestKVLazySnapshotSplitMergeNoCopy verifies Split of a lazy snapshot and Merge of two
|
||||
// adjacent lazy snapshots are pure arithmetic — they produce lazy snapshots and allocate
|
||||
// nothing — while still tracking the correct offsets and data.
|
||||
func TestKVLazySnapshotSplitMergeNoCopy(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
const H, D = 4, 8
|
||||
|
||||
c := NewKVCache()
|
||||
k, v := distinctKV(0, 10, H, D)
|
||||
c.Update(newKVBatch(0, 10), k, v)
|
||||
|
||||
base := settledActiveMemory()
|
||||
|
||||
// Lazy snapshot [2,10), split at 5.
|
||||
snap := c.Snapshot(2)
|
||||
p, ch := c.Split(snap, 5)
|
||||
ps, cs := p.(*kvSnapshot), ch.(*kvSnapshot)
|
||||
if ps.keys != nil || cs.keys != nil {
|
||||
t.Fatal("Split of a lazy snapshot should yield lazy snapshots (no copy)")
|
||||
}
|
||||
if ps.fromOffset != 2 || ps.toOffset != 5 || cs.fromOffset != 5 || cs.toOffset != 10 {
|
||||
t.Fatalf("split ranges = [%d,%d)/[%d,%d), want [2,5)/[5,10)", ps.fromOffset, ps.toOffset, cs.fromOffset, cs.toOffset)
|
||||
}
|
||||
|
||||
// Merge them back into [2,10).
|
||||
merged := c.Merge(p, ch).(*kvSnapshot)
|
||||
if merged.keys != nil {
|
||||
t.Fatal("Merge of adjacent lazy snapshots should yield a lazy snapshot (no Concatenate)")
|
||||
}
|
||||
if merged.fromOffset != 2 || merged.toOffset != 10 {
|
||||
t.Fatalf("merged range = [%d,%d), want [2,10)", merged.fromOffset, merged.toOffset)
|
||||
}
|
||||
|
||||
if after := settledActiveMemory(); after > base {
|
||||
t.Fatalf("Split/Merge of lazy snapshots allocated %d bytes; want 0", after-base)
|
||||
}
|
||||
merged.Close()
|
||||
}
|
||||
|
||||
// TestKVLazySnapshotSurvivesPathSwitch reproduces the switchToPath sequence that
|
||||
// would otherwise corrupt a trie-held lazy snapshot: page out the diverging leaf
|
||||
// (Snapshot), rewind to the common ancestor (Restore(nil)), then page in a
|
||||
// different path (Restore feeds new data via appendKV, overwriting the old
|
||||
// leaf's slots). The paged-out snapshot must still hold the original tokens.
|
||||
func TestKVLazySnapshotSurvivesPathSwitch(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
const H, D = 4, 8
|
||||
|
||||
c := NewKVCache()
|
||||
// Active path tokens [0,10): positions 0..9.
|
||||
k, v := distinctKV(0, 10, H, D)
|
||||
c.Update(newKVBatch(0, 10), k, v)
|
||||
|
||||
// Page out the diverging leaf [4,10) as a lazy snapshot (what switchToPath does
|
||||
// before rewinding).
|
||||
leaf := c.Snapshot(4).(*kvSnapshot)
|
||||
|
||||
// Rewind to the common ancestor at 4 (Restore(nil): offset move only).
|
||||
if !c.Restore(nil, 4) {
|
||||
t.Fatal("rewind to ancestor failed")
|
||||
}
|
||||
|
||||
// Page in the new path [4,9): positions 200..204, overwriting the old leaf's
|
||||
// slots. appendKV must copy the leaf lazy snapshot out first.
|
||||
nk, nv := distinctKV(200, 5, H, D)
|
||||
c.Update(newKVBatch(4, 5), nk, nv)
|
||||
|
||||
if leaf.keys == nil {
|
||||
t.Fatal("leaf lazy snapshot was not copied out during page-in")
|
||||
}
|
||||
mlx.Eval(leaf.keys)
|
||||
// leaf covers [4,10): its keys are the original positions 4..9.
|
||||
keys := leaf.keys.Floats()
|
||||
for l := range leaf.toOffset - leaf.fromOffset {
|
||||
if got, want := keys[l*D], float32(leaf.fromOffset+l); got != want {
|
||||
t.Fatalf("paged-out leaf key[%d] = %v, want %v (original path data)", l, got, want)
|
||||
}
|
||||
}
|
||||
leaf.Close()
|
||||
}
|
||||
|
||||
// TestKVRestoreLiveLazySnapshotIsOffsetMove verifies the same-path rewind/rematch
|
||||
// fast path: restoring a snapshot that is still a lazy index into this cache's
|
||||
// own buffer (its slots never overwritten, so the data is already live) advances
|
||||
// the offset without cloning or replaying. This is the switchToPath sequence
|
||||
// where a paged-out leaf is restored before any write displaced it.
|
||||
func TestKVRestoreLiveLazySnapshotIsOffsetMove(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
const H, D = 4, 8
|
||||
|
||||
c := NewKVCache()
|
||||
k, v := distinctKV(0, 10, H, D)
|
||||
c.Update(newKVBatch(0, 10), k, v)
|
||||
|
||||
// Page out the leaf [5,10) as a lazy snapshot, then rewind (offset move only).
|
||||
snap := c.Snapshot(5).(*kvSnapshot)
|
||||
if !c.Restore(nil, 5) {
|
||||
t.Fatal("rewind failed")
|
||||
}
|
||||
|
||||
base := settledActiveMemory()
|
||||
|
||||
// Restore the snapshot back to 10. Its slots [5,10) were never overwritten,
|
||||
// so it is still lazy and the data is already in the buffer — a pure offset
|
||||
// move, no allocation.
|
||||
if !c.Restore(snap, 10) {
|
||||
t.Fatal("restore failed")
|
||||
}
|
||||
if snap.keys != nil {
|
||||
t.Fatal("snapshot was copied out; expected the offset-move fast path")
|
||||
}
|
||||
if c.Offset() != 10 {
|
||||
t.Fatalf("offset after restore = %d, want 10", c.Offset())
|
||||
}
|
||||
if after := settledActiveMemory(); after > base {
|
||||
t.Fatalf("restore of a live lazy snapshot allocated %d bytes; want 0 (offset move)", after-base)
|
||||
}
|
||||
|
||||
// The buffer still holds the original positions 0..9.
|
||||
st := c.State()
|
||||
mlx.Eval(st[0])
|
||||
keys := st[0].Floats()
|
||||
for l := range 10 {
|
||||
if got := keys[l*D]; got != float32(l) {
|
||||
t.Fatalf("restored key[%d] = %v, want %v", l, got, float32(l))
|
||||
}
|
||||
}
|
||||
snap.Close()
|
||||
}
|
||||
97
x/mlxrunner/cache/recurrent.go
vendored
97
x/mlxrunner/cache/recurrent.go
vendored
@@ -27,6 +27,52 @@ type RecurrentCache struct {
|
||||
numVHeads int
|
||||
headVDim int
|
||||
headKDim int
|
||||
|
||||
snapshots pendingSnapshots
|
||||
}
|
||||
|
||||
// PrepareSnapshots schedules snapshot capture. Recurrent state is cumulative;
|
||||
// an interior offset within a forward has no state unless the recurrent kernel
|
||||
// is run in segments cut at that offset (see SnapshotSplits + Put). The
|
||||
// current offset is a boundary now (the pre-forward state) and is captured
|
||||
// immediately. Interior offsets are captured when Put receives the matching
|
||||
// per-boundary state; the end offset is captured by Put's final state.
|
||||
func (c *RecurrentCache) PrepareSnapshots(offsets []int) {
|
||||
c.snapshots.prepare(c.offset, offsets)
|
||||
// The current offset is a valid boundary right now, so capture it.
|
||||
c.captureBoundary(c.offset)
|
||||
}
|
||||
|
||||
func (c *RecurrentCache) TakeSnapshots() []Snapshot { return c.snapshots.take() }
|
||||
|
||||
// SnapshotSplits returns the scheduled offsets strictly interior to the upcoming
|
||||
// forward [offset, offset+forwardLen), expressed relative to the forward
|
||||
// start — the points at which the caller must segment the recurrent kernel so
|
||||
// each interior state can be captured. Empty when nothing is scheduled or no
|
||||
// interior offsets fall in range.
|
||||
func (c *RecurrentCache) SnapshotSplits(forwardLen int) []int {
|
||||
start := c.offset
|
||||
end := start + forwardLen
|
||||
var splits []int
|
||||
for _, o := range c.snapshots.offsets {
|
||||
if o > start && o < end {
|
||||
splits = append(splits, o-start)
|
||||
}
|
||||
}
|
||||
return splits
|
||||
}
|
||||
|
||||
func (c *RecurrentCache) captureBoundary(reached int) {
|
||||
c.snapshots.captureReached(reached, func(int) Snapshot { return c.Snapshot(reached) })
|
||||
}
|
||||
|
||||
// captureBoundaryState captures a scheduled interior offset from the
|
||||
// per-boundary conv/delta states the kernel wrappers produced while segmenting,
|
||||
// rather than from the cache's current (end-of-forward) state.
|
||||
func (c *RecurrentCache) captureBoundaryState(reached int, conv, delta *mlx.Array) {
|
||||
c.snapshots.captureReached(reached, func(int) Snapshot {
|
||||
return newRecurrentSnapshot(conv, delta, reached)
|
||||
})
|
||||
}
|
||||
|
||||
func (c *RecurrentCache) setState(old, v *mlx.Array, contiguous bool) *mlx.Array {
|
||||
@@ -84,19 +130,39 @@ func (c *RecurrentCache) Get(b *batch.Batch, dtype mlx.DType) *nn.RecurrentHisto
|
||||
|
||||
// Put stores the conv/delta states produced by the SSM layer's write phase.
|
||||
// convStates/deltaStates are the per-boundary recurrent states, one per
|
||||
// boundary ending with the forward-end state. The final entry becomes the
|
||||
// committed live state, advancing the cache offset by the forward's real
|
||||
// token count.
|
||||
// boundary ending with the forward-end state. The boundaries align with this
|
||||
// forward's snapshot splits plus the end: the leading entries are captured as
|
||||
// snapshots at the scheduled interior offsets, and the final entry becomes the
|
||||
// committed live state, advancing the cache offset by the forward's real token
|
||||
// count.
|
||||
//
|
||||
// In the common (unsegmented) case both slices have length 1 — just the
|
||||
// forward-end state.
|
||||
//
|
||||
// Assumes B = 1; heterogeneous batches are not supported.
|
||||
func (c *RecurrentCache) Put(b *batch.Batch, convStates, deltaStates []*mlx.Array) {
|
||||
if len(convStates) != len(deltaStates) || len(convStates) == 0 {
|
||||
panic(fmt.Sprintf("recurrent cache: %d conv / %d delta boundary states", len(convStates), len(deltaStates)))
|
||||
}
|
||||
|
||||
start := c.offset
|
||||
splits := c.SnapshotSplits(int(b.SeqQueryLens[0]))
|
||||
if len(splits) != len(convStates)-1 {
|
||||
panic(fmt.Sprintf("recurrent cache: %d interior splits but %d boundary states", len(splits), len(convStates)))
|
||||
}
|
||||
|
||||
// Leading entries are the interior split boundaries; capture each as a
|
||||
// snapshot at its scheduled offset.
|
||||
for i, s := range splits {
|
||||
c.captureBoundaryState(start+s, convStates[i], deltaStates[i])
|
||||
}
|
||||
|
||||
// The final entry is the forward-end state — the committed live state.
|
||||
last := len(convStates) - 1
|
||||
c.convState = c.setState(c.convState, convStates[last], true)
|
||||
c.deltaState = c.setState(c.deltaState, deltaStates[last], false)
|
||||
c.offset += int(b.SeqQueryLens[0])
|
||||
c.captureBoundary(c.offset)
|
||||
}
|
||||
|
||||
func (c *RecurrentCache) State() []*mlx.Array {
|
||||
@@ -113,18 +179,30 @@ type recurrentSnapshot struct {
|
||||
func (s *recurrentSnapshot) Size() int { return s.convState.NumBytes() + s.deltaState.NumBytes() }
|
||||
func (s *recurrentSnapshot) Close() { mlx.Unpin(s.convState, s.deltaState) }
|
||||
|
||||
// SetMaterializeHook is a no-op: recurrent snapshots are always materialized
|
||||
// at construction.
|
||||
func (s *recurrentSnapshot) SetMaterializeHook(func(int)) {}
|
||||
|
||||
// newRecurrentSnapshot clones and pins conv/delta into an owned snapshot at
|
||||
// offset. Recurrent state is not position-sliceable, so a snapshot always owns
|
||||
// a full copy.
|
||||
func newRecurrentSnapshot(conv, delta *mlx.Array, offset int) *recurrentSnapshot {
|
||||
snap := &recurrentSnapshot{
|
||||
convState: conv.Clone(),
|
||||
deltaState: delta.Clone(),
|
||||
offset: offset,
|
||||
}
|
||||
mlx.Pin(snap.convState, snap.deltaState)
|
||||
return snap
|
||||
}
|
||||
|
||||
func (c *RecurrentCache) Snapshot(fromOffset int) Snapshot {
|
||||
// Recurrent state is not position-sliceable — always snapshot the full state.
|
||||
if c.convState == nil && c.deltaState == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
snap := &recurrentSnapshot{offset: c.offset}
|
||||
snap.convState = c.convState.Clone()
|
||||
snap.deltaState = c.deltaState.Clone()
|
||||
mlx.Pin(snap.convState, snap.deltaState)
|
||||
|
||||
return snap
|
||||
return newRecurrentSnapshot(c.convState, c.deltaState, c.offset)
|
||||
}
|
||||
|
||||
func (c *RecurrentCache) Restore(snapshot Snapshot, target int) bool {
|
||||
@@ -168,6 +246,7 @@ func (c *RecurrentCache) Free() {
|
||||
mlx.Unpin(c.convState, c.deltaState)
|
||||
c.convState, c.deltaState = nil, nil
|
||||
c.offset = 0
|
||||
c.snapshots = pendingSnapshots{}
|
||||
}
|
||||
|
||||
func (c *RecurrentCache) Offset() int { return c.offset }
|
||||
|
||||
462
x/mlxrunner/cache/rotating.go
vendored
462
x/mlxrunner/cache/rotating.go
vendored
@@ -2,6 +2,7 @@ package cache
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
|
||||
"github.com/ollama/ollama/logutil"
|
||||
"github.com/ollama/ollama/x/mlxrunner/batch"
|
||||
@@ -9,21 +10,42 @@ import (
|
||||
"github.com/ollama/ollama/x/models/nn"
|
||||
)
|
||||
|
||||
// RotatingKVCache implements sliding window attention with bounded memory
|
||||
// RotatingKVCache implements sliding window attention with bounded memory.
|
||||
type RotatingKVCache struct {
|
||||
maxSize int
|
||||
idx int
|
||||
keys, values *mlx.Array
|
||||
offset int
|
||||
step int
|
||||
maxSize int
|
||||
idx int
|
||||
|
||||
*KVCache
|
||||
snapshots pendingSnapshots
|
||||
|
||||
// lazySnapshots are outstanding snapshots still in their lazy state: they
|
||||
// index into the live keys/values buffer by slot rather than owning a copy
|
||||
// (see rotatingSnapshot). A write that trims, linearizes, or overwrites the
|
||||
// buffer copies them out (copyOutLazySnapshots) before destroying the slots
|
||||
// they name.
|
||||
lazySnapshots []*rotatingSnapshot
|
||||
}
|
||||
|
||||
func NewRotatingKVCache(maxSize int) *RotatingKVCache {
|
||||
return &RotatingKVCache{maxSize: maxSize, KVCache: NewKVCache()}
|
||||
return &RotatingKVCache{maxSize: maxSize, step: 256}
|
||||
}
|
||||
|
||||
// Assumes B = 1; heterogeneous batches are not supported.
|
||||
func (c *RotatingKVCache) Update(b *batch.Batch, keys, values *mlx.Array) *nn.KVHistory {
|
||||
newK, newV := c.appendKV(keys, values)
|
||||
start := c.offset
|
||||
c.captureStartBoundary(start)
|
||||
|
||||
batched := keys.Dim(2) > 1
|
||||
var newK, newV *mlx.Array
|
||||
if batched {
|
||||
newK, newV = c.concat(keys, values)
|
||||
} else {
|
||||
newK, newV = c.update(keys, values)
|
||||
}
|
||||
|
||||
c.captureLazySnapshots(start, c.offset, batched)
|
||||
return nn.NewKVHistory(newK, newV, rotatingApplier{
|
||||
b: b,
|
||||
K: newK.Dim(2),
|
||||
@@ -34,127 +56,13 @@ func (c *RotatingKVCache) Update(b *batch.Batch, keys, values *mlx.Array) *nn.KV
|
||||
})
|
||||
}
|
||||
|
||||
// View returns the current rotating cache contents in logical order for
|
||||
// assistant KV sharing.
|
||||
func (c *RotatingKVCache) View(_ *batch.Batch) *nn.KVHistory {
|
||||
k, v := c.logicalTail(c.maxSize - 1)
|
||||
if k == nil || v == nil {
|
||||
return nil
|
||||
}
|
||||
return nn.NewKVHistory(k, v, nil)
|
||||
}
|
||||
|
||||
func (c *RotatingKVCache) logicalTail(keep int) (*mlx.Array, *mlx.Array) {
|
||||
state := c.State()
|
||||
if len(state) < 2 || keep <= 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
keys, values := state[0], state[1]
|
||||
K := keys.Dim(2)
|
||||
if K == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
keep = min(keep, K)
|
||||
if K > c.maxSize || c.offset < c.maxSize {
|
||||
start := K - keep
|
||||
return keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(start, K), mlx.Slice()),
|
||||
values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(start, K), mlx.Slice())
|
||||
}
|
||||
|
||||
oldest := c.idx % K
|
||||
var logicalK, logicalV *mlx.Array
|
||||
if oldest == 0 {
|
||||
logicalK, logicalV = keys, values
|
||||
} else {
|
||||
tailK := keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(oldest, K), mlx.Slice())
|
||||
tailV := values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(oldest, K), mlx.Slice())
|
||||
headK := keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, oldest), mlx.Slice())
|
||||
headV := values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, oldest), mlx.Slice())
|
||||
logicalK = tailK.Concatenate(2, headK)
|
||||
logicalV = tailV.Concatenate(2, headV)
|
||||
}
|
||||
|
||||
start := K - keep
|
||||
return logicalK.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(start, K), mlx.Slice()),
|
||||
logicalV.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(start, K), mlx.Slice())
|
||||
}
|
||||
|
||||
type speculativeRotatingKVCache struct {
|
||||
speculativeBase
|
||||
target *RotatingKVCache
|
||||
keys, values *mlx.Array
|
||||
}
|
||||
|
||||
func newSpeculativeRotatingKVCache(target *RotatingKVCache) *speculativeRotatingKVCache {
|
||||
return &speculativeRotatingKVCache{
|
||||
speculativeBase: speculativeBase{offset: target.Offset()},
|
||||
target: target,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *speculativeRotatingKVCache) Update(b *batch.Batch, keys, values *mlx.Array) *nn.KVHistory {
|
||||
c.keys = concatKV(c.keys, keys)
|
||||
c.values = concatKV(c.values, values)
|
||||
c.offset += keys.Dim(2)
|
||||
|
||||
oldK, oldV := c.target.logicalTail(c.target.maxSize - 1)
|
||||
histK, histV := c.keys, c.values
|
||||
if oldK != nil && oldV != nil {
|
||||
histK = oldK.Concatenate(2, c.keys)
|
||||
histV = oldV.Concatenate(2, c.values)
|
||||
}
|
||||
|
||||
return nn.NewKVHistory(histK, histV, logicalSlidingApplier{
|
||||
b: b,
|
||||
K: histK.Dim(2),
|
||||
window: c.target.maxSize,
|
||||
dtype: keys.DType(),
|
||||
})
|
||||
}
|
||||
|
||||
func (c *speculativeRotatingKVCache) State() []*mlx.Array {
|
||||
if c.keys == nil || c.values == nil {
|
||||
return c.target.State()
|
||||
}
|
||||
oldK, oldV := c.target.logicalTail(c.target.maxSize - 1)
|
||||
if oldK == nil || oldV == nil {
|
||||
return []*mlx.Array{c.keys, c.values}
|
||||
}
|
||||
return []*mlx.Array{oldK.Concatenate(2, c.keys), oldV.Concatenate(2, c.values)}
|
||||
}
|
||||
|
||||
func (c *speculativeRotatingKVCache) commit(n int) {
|
||||
if c.keys == nil || c.values == nil || n <= 0 {
|
||||
return
|
||||
}
|
||||
n = min(n, c.keys.Dim(2))
|
||||
c.target.appendKV(prefixKV(c.keys, n), prefixKV(c.values, n))
|
||||
}
|
||||
|
||||
type logicalSlidingApplier struct {
|
||||
b *batch.Batch
|
||||
K int
|
||||
window int
|
||||
dtype mlx.DType
|
||||
}
|
||||
|
||||
func (a logicalSlidingApplier) ApplyMask(logical nn.AttentionMask) nn.AttentionMask {
|
||||
return logical.Intersect(nn.SlidingWindowMask(a.b, a.K, a.window, a.dtype))
|
||||
}
|
||||
|
||||
// appendKV is the raw write path shared by Update and Restore —
|
||||
// routes to concat for prefill (L > 1) and update for decode.
|
||||
func (c *RotatingKVCache) appendKV(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) {
|
||||
if keys.Dim(2) > 1 {
|
||||
return c.concat(keys, values)
|
||||
}
|
||||
return c.update(keys, values)
|
||||
}
|
||||
|
||||
func (c *RotatingKVCache) concat(keys, values *mlx.Array) (newK *mlx.Array, newV *mlx.Array) {
|
||||
logutil.Trace("(*RotatingKVCache).concat", "keys_dim", keys.Dims(), "values_dim", values.Dims(), "offset", c.offset, "idx", c.idx, "max_size", c.maxSize)
|
||||
|
||||
// Freeze outstanding lazy snapshots: the linearize/trim/concat below
|
||||
// reorders and drops the slots they name.
|
||||
c.copyOutLazySnapshots()
|
||||
|
||||
if c.keys == nil {
|
||||
c.keys, c.values = keys.Clone(), values.Clone()
|
||||
mlx.Pin(c.keys, c.values)
|
||||
@@ -187,7 +95,6 @@ func (c *RotatingKVCache) concat(keys, values *mlx.Array) (newK *mlx.Array, newV
|
||||
|
||||
c.keys.Set(c.keys.Concatenate(2, keys))
|
||||
c.values.Set(c.values.Concatenate(2, values))
|
||||
c.idx = c.keys.Dim(2)
|
||||
}
|
||||
|
||||
c.offset += keys.Dim(2)
|
||||
@@ -197,6 +104,11 @@ func (c *RotatingKVCache) concat(keys, values *mlx.Array) (newK *mlx.Array, newV
|
||||
|
||||
func (c *RotatingKVCache) update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) {
|
||||
logutil.Trace("(*RotatingKVCache).update", "keys_dim", keys.Dims(), "values_dim", values.Dims(), "offset", c.offset, "idx", c.idx, "max_size", c.maxSize)
|
||||
|
||||
// Freeze outstanding lazy snapshots: the trim/rotate/SliceUpdate below
|
||||
// overwrites the slots they name.
|
||||
c.copyOutLazySnapshots()
|
||||
|
||||
B, H, L, Dk, Dv := keys.Dim(0), keys.Dim(1), keys.Dim(2), keys.Dim(3), values.Dim(3)
|
||||
|
||||
prev := c.offset
|
||||
@@ -239,6 +151,38 @@ func (c *RotatingKVCache) update(keys, values *mlx.Array) (*mlx.Array, *mlx.Arra
|
||||
c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, validLen), mlx.Slice())
|
||||
}
|
||||
|
||||
// View returns the current cache contents as a read-only KV history, used by an
|
||||
// assistant model that shares this cache. It sets L=1 so rotatingApplier treats
|
||||
// the buffer as ring-ordered (its stored layout); L=1 is a layout selector, not
|
||||
// a query length. A post-concat oversize buffer (K > maxSize) is already in
|
||||
// logical order, so View trims to the trailing maxSize tokens and resets ringIdx
|
||||
// to 0, collapsing the applier's gather to identity.
|
||||
func (c *RotatingKVCache) View(b *batch.Batch) *nn.KVHistory {
|
||||
state := c.State()
|
||||
k, v := state[0], state[1]
|
||||
K := k.Dim(2)
|
||||
ringIdx := c.idx
|
||||
if K > c.maxSize {
|
||||
// Post-concat oversize buffer: storage is in logical (oldest-first)
|
||||
// order, so slice the trailing maxSize tokens. The slice is already
|
||||
// in logical layout, so reset ringIdx to make the applier's gather
|
||||
// collapse to identity.
|
||||
start := K - c.maxSize
|
||||
k = k.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(start, K), mlx.Slice())
|
||||
v = v.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(start, K), mlx.Slice())
|
||||
K = c.maxSize
|
||||
ringIdx = 0
|
||||
}
|
||||
return nn.NewKVHistory(k, v, rotatingApplier{
|
||||
b: b,
|
||||
K: K,
|
||||
L: 1,
|
||||
window: c.maxSize,
|
||||
ringIdx: ringIdx,
|
||||
dtype: k.DType(),
|
||||
})
|
||||
}
|
||||
|
||||
func (c *RotatingKVCache) State() []*mlx.Array {
|
||||
if c.keys == nil || c.values == nil {
|
||||
return nil
|
||||
@@ -250,14 +194,166 @@ func (c *RotatingKVCache) State() []*mlx.Array {
|
||||
}
|
||||
}
|
||||
|
||||
// rotatingSnapshot holds paged-out data for a RotatingKVCache.
|
||||
type rotatingSnapshot struct {
|
||||
kvSnapshot // embedded KV data
|
||||
idx int // buffer write position at snapshot time
|
||||
// replaceBuffer swaps in newK/newV as the cache's keys/values, unpinning the old
|
||||
// buffer and pinning the new one.
|
||||
func (c *RotatingKVCache) replaceBuffer(newK, newV *mlx.Array) {
|
||||
mlx.Unpin(c.keys, c.values)
|
||||
c.keys, c.values = newK, newV
|
||||
mlx.Pin(c.keys, c.values)
|
||||
}
|
||||
|
||||
func (s *rotatingSnapshot) Size() int { return s.kvSnapshot.Size() }
|
||||
func (s *rotatingSnapshot) Close() { s.kvSnapshot.Close() }
|
||||
func (c *RotatingKVCache) Free() {
|
||||
// Freeing drops the buffer lazy snapshots index into; copy them out first.
|
||||
c.copyOutLazySnapshots()
|
||||
mlx.Unpin(c.keys, c.values)
|
||||
c.keys, c.values = nil, nil
|
||||
c.offset = 0
|
||||
c.idx = 0
|
||||
c.snapshots = pendingSnapshots{}
|
||||
}
|
||||
|
||||
func (c *RotatingKVCache) Offset() int { return c.offset }
|
||||
|
||||
func (c *RotatingKVCache) PrepareSnapshots(offsets []int) { c.snapshots.prepare(c.offset, offsets) }
|
||||
func (c *RotatingKVCache) TakeSnapshots() []Snapshot { return c.snapshots.take() }
|
||||
|
||||
// captureStartBoundary captures a scheduled offset at the pre-write position via
|
||||
// the clone path (Snapshot), so the rollback point holds the full pre-write
|
||||
// window before concat/update reorders or drops it.
|
||||
func (c *RotatingKVCache) captureStartBoundary(start int) {
|
||||
if len(c.snapshots.offsets) == 0 {
|
||||
return
|
||||
}
|
||||
c.snapshots.captureReached(start, func(int) Snapshot { return c.Snapshot(0) })
|
||||
}
|
||||
|
||||
// captureLazySnapshots records a snapshot for each scheduled offset the write
|
||||
// reached after the start boundary. A batched write (concat) linearizes the
|
||||
// buffer into logical order, so each window is a contiguous slot slice captured
|
||||
// lazily (lazyRotatingSnapshot). A single-token write (update) leaves the buffer
|
||||
// ring-ordered or grown, breaking that slot math, so capture a clone (Snapshot)
|
||||
// instead; update only ever reaches the end boundary (o == c.offset).
|
||||
func (c *RotatingKVCache) captureLazySnapshots(start, end int, batched bool) {
|
||||
if len(c.snapshots.offsets) == 0 {
|
||||
return
|
||||
}
|
||||
for _, o := range c.snapshots.scheduledIn(start, end) {
|
||||
if o == start {
|
||||
continue // captured pre-write by captureStartBoundary
|
||||
}
|
||||
c.snapshots.captureReached(o, func(int) Snapshot {
|
||||
if batched {
|
||||
return c.lazyRotatingSnapshot(o)
|
||||
}
|
||||
return c.Snapshot(o - min(o, c.maxSize))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// lazyRotatingSnapshot records the window ending at offset o as a lazy snapshot
|
||||
// into the current (logically ordered) buffer: window [o-liveLen, o) maps to
|
||||
// slots [sliceStart, sliceEnd), and restoring sets idx == liveLen so the buffer
|
||||
// reads back in logical order. Returns nil for a zero-width range.
|
||||
func (c *RotatingKVCache) lazyRotatingSnapshot(o int) Snapshot {
|
||||
if c.keys == nil {
|
||||
return nil
|
||||
}
|
||||
bufBase := c.offset - c.keys.Dim(2)
|
||||
liveLen := min(o, c.maxSize)
|
||||
sliceStart := o - liveLen - bufBase
|
||||
sliceEnd := o - bufBase
|
||||
if sliceEnd <= sliceStart {
|
||||
return nil
|
||||
}
|
||||
s := &rotatingSnapshot{
|
||||
fromOffset: o - liveLen,
|
||||
toOffset: o,
|
||||
idx: liveLen,
|
||||
cache: c,
|
||||
sliceStart: sliceStart,
|
||||
sliceEnd: sliceEnd,
|
||||
}
|
||||
c.lazySnapshots = append(c.lazySnapshots, s)
|
||||
return s
|
||||
}
|
||||
|
||||
// rotatingSnapshot holds paged-out data for a RotatingKVCache. Initially lazy:
|
||||
// the window lives in the issuing cache's buffer at slots [sliceStart, sliceEnd)
|
||||
// in logical order, and copyOut clones it into owned keys/values before a write
|
||||
// reorders or drops those slots.
|
||||
type rotatingSnapshot struct {
|
||||
keys, values *mlx.Array // owned window once copied out; nil while lazy
|
||||
fromOffset, toOffset int // absolute offset range the window covers
|
||||
idx int // buffer write position a restore installs
|
||||
|
||||
cache *RotatingKVCache // issuer while lazy; nil once copied out
|
||||
sliceStart, sliceEnd int // buffer slot range of the window while lazy
|
||||
|
||||
// onMaterialize, if set, is fired once from copyOut with the newly-owned
|
||||
// byte count so an owner (e.g. the trie's pagedOutBytes counter) can pick
|
||||
// up bytes that were free while the snapshot was lazy.
|
||||
onMaterialize func(delta int)
|
||||
}
|
||||
|
||||
func (s *rotatingSnapshot) Size() int {
|
||||
if s.keys != nil {
|
||||
return s.keys.NumBytes() + s.values.NumBytes()
|
||||
}
|
||||
// Lazy snapshots own no extra memory: the window still lives in the
|
||||
// issuing cache's buffer.
|
||||
return 0
|
||||
}
|
||||
|
||||
func (s *rotatingSnapshot) SetMaterializeHook(fn func(delta int)) { s.onMaterialize = fn }
|
||||
|
||||
func (s *rotatingSnapshot) Close() {
|
||||
mlx.Unpin(s.keys, s.values)
|
||||
if s.cache != nil {
|
||||
s.cache.dropLazySnapshot(s)
|
||||
s.cache = nil
|
||||
}
|
||||
}
|
||||
|
||||
// copyOut converts a lazy snapshot into an owned clone of its window slots. It
|
||||
// is a no-op once the snapshot already owns its data. The slots hold the window
|
||||
// in logical order, so the clone needs no reordering.
|
||||
func (s *rotatingSnapshot) copyOut() {
|
||||
if s.keys != nil {
|
||||
return
|
||||
}
|
||||
c := s.cache
|
||||
kSlice := c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.sliceStart, s.sliceEnd), mlx.Slice())
|
||||
vSlice := c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(s.sliceStart, s.sliceEnd), mlx.Slice())
|
||||
k := mlx.Contiguous(kSlice, false)
|
||||
v := mlx.Contiguous(vSlice, false)
|
||||
mlx.Pin(k, v)
|
||||
mlx.AsyncEval(k, v)
|
||||
|
||||
s.keys, s.values = k, v
|
||||
c.dropLazySnapshot(s)
|
||||
s.cache = nil
|
||||
|
||||
if s.onMaterialize != nil {
|
||||
s.onMaterialize(s.keys.NumBytes() + s.values.NumBytes())
|
||||
s.onMaterialize = nil
|
||||
}
|
||||
}
|
||||
|
||||
func (c *RotatingKVCache) dropLazySnapshot(s *rotatingSnapshot) {
|
||||
if i := slices.Index(c.lazySnapshots, s); i >= 0 {
|
||||
c.lazySnapshots = slices.Delete(c.lazySnapshots, i, i+1)
|
||||
}
|
||||
}
|
||||
|
||||
// copyOutLazySnapshots clones every outstanding lazy snapshot into owned data.
|
||||
// The destructive write paths (concat, update) call this before they trim,
|
||||
// linearize, or overwrite the slots a snapshot names. copyOut removes the
|
||||
// snapshot from the set, so iterate over a clone.
|
||||
func (c *RotatingKVCache) copyOutLazySnapshots() {
|
||||
for _, s := range slices.Clone(c.lazySnapshots) {
|
||||
s.copyOut()
|
||||
}
|
||||
}
|
||||
|
||||
func (c *RotatingKVCache) Snapshot(fromOffset int) Snapshot {
|
||||
if c.keys == nil || c.offset <= fromOffset {
|
||||
@@ -270,13 +366,11 @@ func (c *RotatingKVCache) Snapshot(fromOffset int) Snapshot {
|
||||
mlx.Pin(k, v)
|
||||
|
||||
return &rotatingSnapshot{
|
||||
kvSnapshot: kvSnapshot{
|
||||
keys: k,
|
||||
values: v,
|
||||
fromOffset: fromOffset,
|
||||
toOffset: c.offset,
|
||||
},
|
||||
idx: c.idx,
|
||||
keys: k,
|
||||
values: v,
|
||||
fromOffset: fromOffset,
|
||||
toOffset: c.offset,
|
||||
idx: c.idx,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -289,10 +383,9 @@ func (c *RotatingKVCache) Restore(snapshot Snapshot, target int) bool {
|
||||
if target >= c.offset {
|
||||
return target == c.offset
|
||||
}
|
||||
// Live rewind is only safe when the buffer hasn't filled yet
|
||||
// (offset <= maxSize). Once the window has shifted, rewinding
|
||||
// leaves fewer than maxSize trailing tokens to attend to —
|
||||
// a snapshot is required to restore the full window.
|
||||
// Live rewind is only safe before the buffer fills (offset <= maxSize);
|
||||
// once wrapped, rewinding leaves an incomplete window, so a snapshot is
|
||||
// required.
|
||||
if c.offset > c.maxSize {
|
||||
return false
|
||||
}
|
||||
@@ -312,14 +405,37 @@ func (c *RotatingKVCache) Restore(snapshot Snapshot, target int) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// Restore from snapshot: rebuild buffer state.
|
||||
// Free existing state first.
|
||||
if c.keys != nil {
|
||||
mlx.Unpin(c.keys, c.values)
|
||||
// Fast path: this cache's own still-lazy snapshot names the restored window
|
||||
// in the live buffer, so slice it in as the new buffer (no copy) and re-point
|
||||
// the snapshot to slots [0, liveLen), kept lazy.
|
||||
if snap.cache == c && snap.keys == nil {
|
||||
// Drop snap (re-pointed, not copied), then copy out the siblings: their
|
||||
// slots fall outside the restored window.
|
||||
c.dropLazySnapshot(snap)
|
||||
c.copyOutLazySnapshots()
|
||||
liveLen := snap.sliceEnd - snap.sliceStart
|
||||
c.replaceBuffer(
|
||||
c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(snap.sliceStart, snap.sliceEnd), mlx.Slice()),
|
||||
c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(snap.sliceStart, snap.sliceEnd), mlx.Slice()),
|
||||
)
|
||||
snap.sliceStart, snap.sliceEnd = 0, liveLen
|
||||
c.lazySnapshots = append(c.lazySnapshots, snap)
|
||||
c.offset = snap.toOffset
|
||||
c.idx = snap.idx
|
||||
if target < c.offset {
|
||||
c.offset = target
|
||||
c.idx = target
|
||||
}
|
||||
return true
|
||||
}
|
||||
c.keys = snap.keys.Clone()
|
||||
c.values = snap.values.Clone()
|
||||
mlx.Pin(c.keys, c.values)
|
||||
|
||||
// snap is owned (cross-cache or already materialized; the lazy-own case is
|
||||
// above), so copyOut is a no-op here. Freeze this cache's lazy snapshots off
|
||||
// the buffer the assignment below replaces.
|
||||
snap.copyOut()
|
||||
c.copyOutLazySnapshots()
|
||||
|
||||
c.replaceBuffer(snap.keys.Clone(), snap.values.Clone())
|
||||
c.offset = snap.toOffset
|
||||
c.idx = snap.idx
|
||||
|
||||
@@ -346,20 +462,11 @@ func (c *RotatingKVCache) Split(snapshot Snapshot, at int) (Snapshot, Snapshot)
|
||||
return nil, snapshot
|
||||
}
|
||||
|
||||
func (c *RotatingKVCache) Free() {
|
||||
c.KVCache.Free()
|
||||
c.idx = 0
|
||||
}
|
||||
|
||||
// rotatingApplier composes the sliding-window storage restriction
|
||||
// onto the caller's logical mask.
|
||||
//
|
||||
// ringIdx is the cache's write cursor at Update time. At L=1 decode
|
||||
// the ring buffer is not position-ordered — logical col j lives at
|
||||
// storage slot (ringIdx+j) mod K — so tensor masks built in
|
||||
// logical space must be gathered into this layout before the kernel
|
||||
// sees them. At L>1 prefill the concat path has already linearised
|
||||
// storage, so the gather is identity and ringIdx is unused.
|
||||
// rotatingApplier composes the sliding-window storage restriction onto the
|
||||
// caller's logical mask. ringIdx is the write cursor at Update time: at L=1
|
||||
// decode the ring is not position-ordered (logical col j lives at slot
|
||||
// (ringIdx+j) mod K), so tensor masks must be gathered into ring layout. At L>1
|
||||
// prefill concat has linearized storage, so the gather is identity.
|
||||
type rotatingApplier struct {
|
||||
b *batch.Batch
|
||||
K int
|
||||
@@ -371,20 +478,15 @@ type rotatingApplier struct {
|
||||
|
||||
func (r rotatingApplier) ApplyMask(logical nn.AttentionMask) nn.AttentionMask {
|
||||
if r.L == 1 {
|
||||
// Single-query decode: storage already enforces the window
|
||||
// (Update keeps the last maxSize tokens, all within
|
||||
// [absQ-window+1, absQ]), and every stored key's absolute
|
||||
// position <= absQ. For a zero or plain-causal logical mask
|
||||
// both constraints reduce to "no mask", so return the zero
|
||||
// mask and let SDPA dispatch to mode="".
|
||||
// Single-query decode: storage already enforces the window and every
|
||||
// stored key's position <= absQ, so a zero or causal logical mask
|
||||
// reduces to no mask — let SDPA dispatch to mode="".
|
||||
if logical.IsZero() || logical.IsCausal() {
|
||||
return nn.AttentionMask{}
|
||||
}
|
||||
|
||||
// Tensor-backed mask (user ArrayMask, causal+Relax, causal
|
||||
// with accumulated array): materialize in logical-position
|
||||
// order then gather K cols into ring-slot order so they
|
||||
// align with the cache output the kernel will index.
|
||||
// Tensor-backed mask: materialize in logical order, then gather K cols
|
||||
// into ring-slot order to align with the cache output.
|
||||
arr := logical.AsArray(r.b, r.K, r.dtype)
|
||||
arr = gatherRingCols(arr, r.ringIdx, r.K)
|
||||
return nn.ArrayMask(arr)
|
||||
@@ -393,14 +495,10 @@ func (r rotatingApplier) ApplyMask(logical nn.AttentionMask) nn.AttentionMask {
|
||||
return logical.Intersect(nn.SlidingWindowMask(r.b, r.K, r.window, r.dtype))
|
||||
}
|
||||
|
||||
// gatherRingCols reorders a [B, 1, L, K] mask's K axis from
|
||||
// logical-position order (col 0 = oldest stored position) into the
|
||||
// cache's ring-slot order (col 0 = buffer slot 0). Logical col j
|
||||
// lives at slot (ringIdx+j) mod K, so storage slot s reads from
|
||||
// logical col (s-ringIdx+K) mod K. Returns arr unchanged when the
|
||||
// permutation is a no-op: ringIdx % K == 0 (layouts coincide), or
|
||||
// the K axis broadcasts (dim 3 == 1, i.e. Q-padding-shaped masks
|
||||
// where every key shares the same value).
|
||||
// gatherRingCols reorders a [B, 1, L, K] mask's K axis from logical order
|
||||
// (col 0 = oldest) into ring-slot order (col 0 = slot 0): logical col j lives at
|
||||
// slot (ringIdx+j) mod K. A no-op when ringIdx % K == 0 or the K axis broadcasts
|
||||
// (dim 3 == 1, Q-padding-shaped masks where every key shares one value).
|
||||
func gatherRingCols(arr *mlx.Array, ringIdx, K int) *mlx.Array {
|
||||
if w := arr.Dim(3); w != 1 && w != K {
|
||||
panic(fmt.Sprintf("gatherRingCols: K-axis width %d must be 1 or %d", w, K))
|
||||
|
||||
83
x/mlxrunner/cache/rotating_attention_test.go
vendored
83
x/mlxrunner/cache/rotating_attention_test.go
vendored
@@ -143,14 +143,13 @@ func TestAssistantSharedHistoryL1MasksMatchNoMask(t *testing.T) {
|
||||
}
|
||||
|
||||
b := newKVBatch(total-1, 1)
|
||||
slidingHistory := sliding.View(b)
|
||||
cases := []struct {
|
||||
name string
|
||||
h *nn.KVHistory
|
||||
mask nn.AttentionMask
|
||||
}{
|
||||
{name: "full", h: full.View(b), mask: nn.CausalMask()},
|
||||
{name: "sliding", h: slidingHistory, mask: nn.CausalMask().Intersect(nn.SlidingWindowMask(b, slidingHistory.K().Dim(2), window, q.DType()))},
|
||||
{name: "sliding", h: sliding.View(b), mask: nn.CausalMask()},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
@@ -252,6 +251,86 @@ func TestRotatingKVCachePrefillParity(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestRotatingKVCacheScheduledSnapshotParity verifies that scheduling per-token
|
||||
// snapshots over a batched write past the wrap point does not change the
|
||||
// attention the write returns: the history must produce the same SDPA output as
|
||||
// an identical write with no snapshots scheduled. Capture happens after the
|
||||
// write via lazy snapshots, so it must not perturb the write itself.
|
||||
func TestRotatingKVCacheScheduledSnapshotParity(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
const H, D = 1, 4
|
||||
const window = 4
|
||||
const before = 5 // past wrap before the batched write
|
||||
const draft = 4
|
||||
const scale = 1.0
|
||||
|
||||
perPosKV := func(pos int) (k, v *mlx.Array) {
|
||||
kVals := make([]float32, H*D)
|
||||
vVals := make([]float32, H*D)
|
||||
for i := range kVals {
|
||||
kVals[i] = 0.1*float32(pos+1) + 0.01*float32(i)
|
||||
vVals[i] = -0.1*float32(pos+1) + 0.01*float32(i)
|
||||
}
|
||||
return mlx.FromValues(kVals, 1, H, 1, D), mlx.FromValues(vVals, 1, H, 1, D)
|
||||
}
|
||||
|
||||
// Build the batched K/V for offsets [before, before+draft).
|
||||
batchK := make([]*mlx.Array, draft)
|
||||
batchV := make([]*mlx.Array, draft)
|
||||
for i := range draft {
|
||||
batchK[i], batchV[i] = perPosKV(before + i)
|
||||
}
|
||||
kBatch := mlx.Concatenate(batchK, 2)
|
||||
vBatch := mlx.Concatenate(batchV, 2)
|
||||
|
||||
qVals := make([]float32, H*draft*D)
|
||||
for i := range qVals {
|
||||
qVals[i] = 0.5 + 0.05*float32(i)
|
||||
}
|
||||
q := mlx.FromValues(qVals, 1, H, draft, D)
|
||||
b := newKVBatch(before, draft)
|
||||
|
||||
// Run the same write twice: once with snapshots scheduled, once without.
|
||||
run := func(schedule bool) []float32 {
|
||||
c := NewRotatingKVCache(window)
|
||||
for pos := range before {
|
||||
k, v := perPosKV(pos)
|
||||
c.Update(newKVBatch(c.Offset(), 1), k, v)
|
||||
}
|
||||
if schedule {
|
||||
offsets := make([]int, draft)
|
||||
for i := range offsets {
|
||||
offsets[i] = before + i
|
||||
}
|
||||
c.PrepareSnapshots(offsets)
|
||||
}
|
||||
history := c.Update(b, kBatch, vBatch)
|
||||
out := nn.ScaledDotProductAttention(b, q, scale,
|
||||
nn.WithKVHistory(history),
|
||||
nn.WithMask(nn.CausalMask()))
|
||||
if schedule {
|
||||
for _, s := range c.TakeSnapshots() {
|
||||
if s != nil {
|
||||
s.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
mlx.Eval(out)
|
||||
return out.Floats()
|
||||
}
|
||||
|
||||
withSnap := run(true)
|
||||
noSnap := run(false)
|
||||
if len(withSnap) != len(noSnap) {
|
||||
t.Fatalf("output length %d vs %d", len(withSnap), len(noSnap))
|
||||
}
|
||||
for i := range noSnap {
|
||||
if math.Abs(float64(withSnap[i]-noSnap[i])) > 1e-5 {
|
||||
t.Fatalf("index %d: scheduled=%v, unscheduled=%v", i, withSnap[i], noSnap[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRotatingKVCacheMLAParity drives a rotating cache with the MLA
|
||||
// shape — K = [kvLatent, kPE] concatenated, V = zero-width — then
|
||||
// uses WithMLAHistory to slice V from K and compares output against
|
||||
|
||||
754
x/mlxrunner/cache/snapshot_capture_test.go
vendored
Normal file
754
x/mlxrunner/cache/snapshot_capture_test.go
vendored
Normal file
@@ -0,0 +1,754 @@
|
||||
package cache
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"testing"
|
||||
|
||||
"github.com/ollama/ollama/x/mlxrunner/mlx"
|
||||
)
|
||||
|
||||
// fillKV writes n single-token steps into c, one Update per token, so the
|
||||
// cache reaches offset n the way decode would.
|
||||
func fillKV(c Attention, n int) {
|
||||
for range n {
|
||||
k := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
|
||||
v := mlx.Zeros(mlx.DTypeFloat16, 1, 4, 1, 8)
|
||||
c.Update(newKVBatch(c.Offset(), k.Dim(2)), k, v)
|
||||
}
|
||||
}
|
||||
|
||||
// batchKV builds K/V for an L-token batched write.
|
||||
func batchKV(L int) (*mlx.Array, *mlx.Array) {
|
||||
return mlx.Zeros(mlx.DTypeFloat16, 1, 4, L, 8), mlx.Zeros(mlx.DTypeFloat16, 1, 4, L, 8)
|
||||
}
|
||||
|
||||
// taggedKV builds an L-token K/V batch where every element of the token at
|
||||
// absolute offset p carries the value p+1 (the +1 keeps tags distinct from the
|
||||
// zero grow-padding the cache writes). The tag survives slicing/rotation, so a
|
||||
// restored window's logical order can be read back as the sequence of absolute
|
||||
// positions it holds. Shape is the standard [B=1, H=4, L, D=8].
|
||||
func taggedKV(startOffset, L int) (*mlx.Array, *mlx.Array) {
|
||||
const H, D = 4, 8
|
||||
vals := make([]float32, H*L*D)
|
||||
for l := range L {
|
||||
tag := float32(startOffset + l + 1)
|
||||
for h := range H {
|
||||
for d := range D {
|
||||
vals[(h*L+l)*D+d] = tag
|
||||
}
|
||||
}
|
||||
}
|
||||
k := mlx.FromValues(vals, 1, H, L, D)
|
||||
v := mlx.FromValues(vals, 1, H, L, D)
|
||||
return k, v
|
||||
}
|
||||
|
||||
// fillTagged advances c from offset 0 to n with single-token tagged writes (the
|
||||
// decode update path), so each stored position carries its absolute-offset tag.
|
||||
func fillTagged(c Attention, n int) {
|
||||
for p := range n {
|
||||
k, v := taggedKV(p, 1)
|
||||
c.Update(newKVBatch(p, 1), k, v)
|
||||
}
|
||||
}
|
||||
|
||||
// windowTags reads c's logical window and returns the per-position tag (the
|
||||
// absolute offset each slot holds, recovered as value-1). It uses element 0 of
|
||||
// each token, which taggedKV set uniformly. Returns nil if the window is empty.
|
||||
func windowTags(t *testing.T, c *RotatingKVCache) []int {
|
||||
t.Helper()
|
||||
state := c.State()
|
||||
if len(state) == 0 {
|
||||
return nil
|
||||
}
|
||||
k := state[0]
|
||||
K := k.Dim(2)
|
||||
if K == 0 {
|
||||
return nil
|
||||
}
|
||||
// Linearize ring storage into logical (oldest-first) order: slots
|
||||
// [oldest, K) ++ [0, oldest). After concat the buffer is already
|
||||
// in logical order (oldest == 0), so the concat below is skipped.
|
||||
if oldest := c.idx % K; oldest != 0 {
|
||||
tail := k.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(oldest, K), mlx.Slice())
|
||||
head := k.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, oldest), mlx.Slice())
|
||||
k = tail.Concatenate(2, head)
|
||||
}
|
||||
L, D := k.Dim(2), k.Dim(3)
|
||||
mlx.Eval(k)
|
||||
f := k.Floats()
|
||||
tags := make([]int, L)
|
||||
for l := range L {
|
||||
tags[l] = int(f[l*D]) - 1 // element 0 of token l; undo the +1 tag offset
|
||||
}
|
||||
return tags
|
||||
}
|
||||
|
||||
// wantWindowTags returns the absolute positions a full window at offset would
|
||||
// hold in logical (oldest-first) order: the trailing min(offset, window) tokens.
|
||||
func wantWindowTags(offset, window int) []int {
|
||||
n := min(offset, window)
|
||||
tags := make([]int, n)
|
||||
for i := range tags {
|
||||
tags[i] = offset - n + i
|
||||
}
|
||||
return tags
|
||||
}
|
||||
|
||||
// TestKVCachePerTokenSnapshotRestore schedules per-token offsets across a single
|
||||
// batched write and verifies the captures are edge-local and that the
|
||||
// speculation commit path — a live rewind, since KV is append-only — restores
|
||||
// the cache to each accepted offset with the prefix intact.
|
||||
func TestKVCachePerTokenSnapshotRestore(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
const before = 6
|
||||
const draft = 4
|
||||
|
||||
for accepted := 0; accepted <= draft; accepted++ {
|
||||
c := NewKVCache()
|
||||
fillKV(c, before)
|
||||
|
||||
offsets := make([]int, draft)
|
||||
for i := range offsets {
|
||||
offsets[i] = before + i
|
||||
}
|
||||
c.PrepareSnapshots(offsets)
|
||||
|
||||
k, v := batchKV(draft)
|
||||
c.Update(newKVBatch(before, draft), k, v)
|
||||
if c.Offset() != before+draft {
|
||||
t.Fatalf("accepted=%d: offset after write = %d, want %d", accepted, c.Offset(), before+draft)
|
||||
}
|
||||
|
||||
snaps := c.TakeSnapshots()
|
||||
if len(snaps) != draft {
|
||||
t.Fatalf("accepted=%d: got %d snapshots, want %d", accepted, len(snaps), draft)
|
||||
}
|
||||
|
||||
// Captures are edge-local: offset before is zero-width (nil), and each
|
||||
// later offset holds exactly the single token [before+i-1, before+i).
|
||||
if snaps[0] != nil {
|
||||
t.Fatalf("accepted=%d: snaps[0] = %v, want nil (zero-width base)", accepted, snaps[0])
|
||||
}
|
||||
for i := 1; i < draft; i++ {
|
||||
ks := snaps[i].(*kvSnapshot)
|
||||
if ks.fromOffset != before+i-1 || ks.toOffset != before+i {
|
||||
t.Fatalf("accepted=%d: snaps[%d] = [%d,%d), want [%d,%d)", accepted, i, ks.fromOffset, ks.toOffset, before+i-1, before+i)
|
||||
}
|
||||
}
|
||||
|
||||
// Commit rolls back via a live rewind (Restore(nil)) — the append-only
|
||||
// buffer still holds [0, before+draft), so the edge captures go unused.
|
||||
if accepted < draft {
|
||||
if !c.Restore(nil, before+accepted) {
|
||||
t.Fatalf("accepted=%d: live rewind failed", accepted)
|
||||
}
|
||||
}
|
||||
for _, s := range snaps {
|
||||
if s != nil {
|
||||
s.Close()
|
||||
}
|
||||
}
|
||||
|
||||
want := before + draft
|
||||
if accepted < draft {
|
||||
want = before + accepted
|
||||
}
|
||||
if c.Offset() != want {
|
||||
t.Fatalf("accepted=%d: offset after commit = %d, want %d", accepted, c.Offset(), want)
|
||||
}
|
||||
if st := c.State(); len(st) == 2 && st[0].Dim(2) != want {
|
||||
t.Fatalf("accepted=%d: state seq dim = %d, want %d", accepted, st[0].Dim(2), want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestKVCaptureMergeSplit verifies that two adjacent edge-local captures merge
|
||||
// into the combined edge and split back into the halves — the operations the
|
||||
// trie performs on stored snapshots (mergeWithChild / splitNode) — proving they
|
||||
// work on captured snapshots, not just freshly-taken ones.
|
||||
func TestKVCaptureMergeSplit(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
const before = 6
|
||||
c := NewKVCache()
|
||||
fillKV(c, before)
|
||||
|
||||
// Schedule two interior offsets so the write captures [before, before+1) and
|
||||
// [before+1, before+2).
|
||||
c.PrepareSnapshots([]int{before + 1, before + 2})
|
||||
k, v := batchKV(3)
|
||||
c.Update(newKVBatch(before, 3), k, v)
|
||||
|
||||
snaps := c.TakeSnapshots()
|
||||
a := snaps[0].(*kvSnapshot)
|
||||
b := snaps[1].(*kvSnapshot)
|
||||
if a.fromOffset != before || a.toOffset != before+1 {
|
||||
t.Fatalf("snaps[0] = [%d,%d), want [%d,%d)", a.fromOffset, a.toOffset, before, before+1)
|
||||
}
|
||||
if b.fromOffset != before+1 || b.toOffset != before+2 {
|
||||
t.Fatalf("snaps[1] = [%d,%d), want [%d,%d)", b.fromOffset, b.toOffset, before+1, before+2)
|
||||
}
|
||||
|
||||
// Merge the adjacent edges into [before, before+2).
|
||||
merged := c.Merge(snaps[0], snaps[1]).(*kvSnapshot)
|
||||
if merged.fromOffset != before || merged.toOffset != before+2 {
|
||||
t.Fatalf("merged = [%d,%d), want [%d,%d)", merged.fromOffset, merged.toOffset, before, before+2)
|
||||
}
|
||||
|
||||
// Split back at before+1 and confirm the halves match the originals.
|
||||
p, ch := c.Split(merged, before+1)
|
||||
ps := p.(*kvSnapshot)
|
||||
cs := ch.(*kvSnapshot)
|
||||
if ps.fromOffset != before || ps.toOffset != before+1 {
|
||||
t.Fatalf("split parent = [%d,%d), want [%d,%d)", ps.fromOffset, ps.toOffset, before, before+1)
|
||||
}
|
||||
if cs.fromOffset != before+1 || cs.toOffset != before+2 {
|
||||
t.Fatalf("split child = [%d,%d), want [%d,%d)", cs.fromOffset, cs.toOffset, before+1, before+2)
|
||||
}
|
||||
p.Close()
|
||||
ch.Close()
|
||||
}
|
||||
|
||||
// TestRotatingPerTokenSnapshotRestore exercises per-token capture on a
|
||||
// rotating cache across regimes: ring not yet full, exactly full, and wrapped.
|
||||
// Every restore is exact-match against its own snapshot, so it must succeed
|
||||
// regardless of wrap state, and the restored logical window must hold exactly
|
||||
// the trailing absolute positions it should. Tagged K/V make slot-math errors
|
||||
// observable as wrong positions, not just wrong shapes; the wrapped regime
|
||||
// forces concat's linearize branch to run on entry.
|
||||
func TestRotatingPerTokenSnapshotRestore(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
window int
|
||||
before int
|
||||
}{
|
||||
{"ring-not-full", 32, 4},
|
||||
{"ring-exactly-full", 8, 4},
|
||||
{"ring-wrapped", 4, 10},
|
||||
}
|
||||
|
||||
const draft = 4
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
for accepted := 0; accepted <= draft; accepted++ {
|
||||
c := NewRotatingKVCache(tc.window)
|
||||
fillTagged(c, tc.before)
|
||||
|
||||
offsets := make([]int, draft)
|
||||
for i := range offsets {
|
||||
offsets[i] = tc.before + i
|
||||
}
|
||||
c.PrepareSnapshots(offsets)
|
||||
|
||||
k, v := taggedKV(tc.before, draft)
|
||||
c.Update(newKVBatch(tc.before, draft), k, v)
|
||||
|
||||
snaps := c.TakeSnapshots()
|
||||
if len(snaps) != draft {
|
||||
t.Fatalf("accepted=%d: got %d snapshots, want %d", accepted, len(snaps), draft)
|
||||
}
|
||||
|
||||
if accepted < draft {
|
||||
if !c.Restore(snaps[accepted], tc.before+accepted) {
|
||||
t.Fatalf("accepted=%d: restore failed", accepted)
|
||||
}
|
||||
}
|
||||
for _, s := range snaps {
|
||||
s.Close()
|
||||
}
|
||||
|
||||
want := tc.before + draft
|
||||
if accepted < draft {
|
||||
want = tc.before + accepted
|
||||
}
|
||||
if c.Offset() != want {
|
||||
t.Fatalf("accepted=%d: offset after commit = %d, want %d", accepted, c.Offset(), want)
|
||||
}
|
||||
// The logical window holds the trailing min(offset, window)
|
||||
// absolute positions in order — a slot-math error would keep the
|
||||
// right count but the wrong positions, which a dimension check
|
||||
// would miss. A batched write through concat can leave the raw
|
||||
// buffer larger than the window (it retains maxSize-1+L slots);
|
||||
// View trims to the trailing window at SDPA time, so compare the
|
||||
// trailing window of the linearized buffer.
|
||||
got := windowTags(t, c)
|
||||
if len(got) > tc.window {
|
||||
got = got[len(got)-tc.window:]
|
||||
}
|
||||
if wantTags := wantWindowTags(want, tc.window); !slices.Equal(got, wantTags) {
|
||||
t.Fatalf("accepted=%d: window tags = %v, want %v", accepted, got, wantTags)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestRotatingRestoreLazyOwnSnapshotSlices verifies the restore fast path:
|
||||
// restoring from this cache's own still-lazy, non-trie-owned snapshot slices the
|
||||
// live buffer rather than copying the window out, and does not consume the
|
||||
// snapshot — it stays lazy, re-pointed at the new buffer, so it still names the
|
||||
// same window data. Covers the restored content, a following decode write, and
|
||||
// the snapshot copying out correctly afterward, across wrap regimes.
|
||||
func TestRotatingRestoreLazyOwnSnapshotSlices(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
window int
|
||||
before int
|
||||
}{
|
||||
{"ring-not-full", 32, 4},
|
||||
{"ring-exactly-full", 8, 4},
|
||||
{"ring-wrapped", 4, 10},
|
||||
}
|
||||
|
||||
const draft, accepted = 4, 2
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c := NewRotatingKVCache(tc.window)
|
||||
fillTagged(c, tc.before)
|
||||
|
||||
offsets := make([]int, draft)
|
||||
for i := range offsets {
|
||||
offsets[i] = tc.before + i
|
||||
}
|
||||
c.PrepareSnapshots(offsets)
|
||||
|
||||
k, v := taggedKV(tc.before, draft)
|
||||
c.Update(newKVBatch(tc.before, draft), k, v)
|
||||
|
||||
snaps := c.TakeSnapshots()
|
||||
snap := snaps[accepted].(*rotatingSnapshot)
|
||||
if snap.keys != nil {
|
||||
t.Fatal("snapshot materialized before restore; expected lazy")
|
||||
}
|
||||
|
||||
if !c.Restore(snap, tc.before+accepted) {
|
||||
t.Fatal("restore failed")
|
||||
}
|
||||
|
||||
// Fast path: the snapshot was sliced, not copied out, and stays a
|
||||
// valid lazy snapshot re-pointed at the new buffer.
|
||||
if snap.keys != nil {
|
||||
t.Fatal("snapshot was copied out; expected the slice fast path")
|
||||
}
|
||||
if snap.cache != c {
|
||||
t.Fatal("snapshot lost its lazy cache reference")
|
||||
}
|
||||
if !slices.Contains(c.lazySnapshots, snap) {
|
||||
t.Fatal("snapshot not re-added to the lazy set")
|
||||
}
|
||||
if snap.sliceStart != 0 || snap.sliceEnd != min(tc.before+accepted, tc.window) {
|
||||
t.Fatalf("snapshot slots = [%d,%d), want [0,%d)", snap.sliceStart, snap.sliceEnd, min(tc.before+accepted, tc.window))
|
||||
}
|
||||
|
||||
got := windowTags(t, c)
|
||||
want := wantWindowTags(tc.before+accepted, tc.window)
|
||||
if !slices.Equal(got, want) {
|
||||
t.Fatalf("window tags = %v, want %v", got, want)
|
||||
}
|
||||
|
||||
// A following decode write must produce correct content (the sliced
|
||||
// buffer feeds update's ring math unchanged). The write copies the
|
||||
// re-pointed snapshot out first; it must capture the same window.
|
||||
wk, wv := taggedKV(c.Offset(), 1)
|
||||
c.Update(newKVBatch(c.Offset(), 1), wk, wv)
|
||||
if snap.keys == nil {
|
||||
t.Fatal("following write did not copy out the re-pointed snapshot")
|
||||
}
|
||||
mlx.Eval(snap.keys)
|
||||
if head := int(snap.keys.Floats()[0]) - 1; head != tc.before+accepted-min(tc.before+accepted, tc.window) {
|
||||
t.Fatalf("re-pointed snapshot head tag = %d, want %d", head, tc.before+accepted-min(tc.before+accepted, tc.window))
|
||||
}
|
||||
got = windowTags(t, c)
|
||||
want = wantWindowTags(tc.before+accepted+1, tc.window)
|
||||
if !slices.Equal(got, want) {
|
||||
t.Fatalf("after follow-up write: window tags = %v, want %v", got, want)
|
||||
}
|
||||
|
||||
for _, s := range snaps {
|
||||
if s != nil {
|
||||
s.Close()
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestRotatingRestoreHookedSnapshotStaysLazy verifies a trie-owned snapshot (one
|
||||
// with a materialize hook) takes the same re-point fast path: restore does not
|
||||
// copy it out, so its hook does not fire and pagedOutBytes is not charged while
|
||||
// the window still rides the live buffer. The hook fires exactly once, later,
|
||||
// when a following write would destroy the window and copies the snapshot out —
|
||||
// the lazy mechanism paying for itself only when the data is about to be lost.
|
||||
func TestRotatingRestoreHookedSnapshotStaysLazy(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
const window, before, draft, accepted = 4, 10, 4, 2
|
||||
|
||||
c := NewRotatingKVCache(window)
|
||||
fillTagged(c, before)
|
||||
|
||||
offsets := make([]int, draft)
|
||||
for i := range offsets {
|
||||
offsets[i] = before + i
|
||||
}
|
||||
c.PrepareSnapshots(offsets)
|
||||
|
||||
k, v := taggedKV(before, draft)
|
||||
c.Update(newKVBatch(before, draft), k, v)
|
||||
|
||||
snaps := c.TakeSnapshots()
|
||||
snap := snaps[accepted].(*rotatingSnapshot)
|
||||
// Simulate trie ownership: a node sets a materialize hook on attach.
|
||||
fired := 0
|
||||
snap.SetMaterializeHook(func(int) { fired++ })
|
||||
|
||||
if !c.Restore(snap, before+accepted) {
|
||||
t.Fatal("restore failed")
|
||||
}
|
||||
// Re-pointed, not copied out: still lazy, hook unfired.
|
||||
if snap.keys != nil {
|
||||
t.Fatal("hooked snapshot was copied out; expected the re-point fast path")
|
||||
}
|
||||
if fired != 0 {
|
||||
t.Fatalf("materialize hook fired %d times on restore, want 0 (still lazy)", fired)
|
||||
}
|
||||
|
||||
got := windowTags(t, c)
|
||||
want := wantWindowTags(before+accepted, window)
|
||||
if !slices.Equal(got, want) {
|
||||
t.Fatalf("window tags = %v, want %v", got, want)
|
||||
}
|
||||
|
||||
// A following decode write destroys the window's slots, so it copies the
|
||||
// snapshot out — firing the hook exactly once.
|
||||
wk, wv := taggedKV(c.Offset(), 1)
|
||||
c.Update(newKVBatch(c.Offset(), 1), wk, wv)
|
||||
if snap.keys == nil {
|
||||
t.Fatal("following write did not copy out the snapshot")
|
||||
}
|
||||
if fired != 1 {
|
||||
t.Fatalf("materialize hook fired %d times, want 1", fired)
|
||||
}
|
||||
|
||||
for _, s := range snaps {
|
||||
if s != nil {
|
||||
s.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRotatingSnapshotSingleTokenWrite mirrors the tail of chunked prefill: the
|
||||
// loop leaves the last token for decode seeding, so when two tokens remain it
|
||||
// writes a single-token chunk (L=1) through update rather than concat. A snapshot
|
||||
// scheduled at that write's end offset is captured against a buffer that may be in
|
||||
// ring order, where a lazy slot slice would name the wrong slots. Covers the
|
||||
// not-yet-wrapped and wrapped regimes; the wrapped one exercises the ring-clone
|
||||
// fallback.
|
||||
func TestRotatingSnapshotSingleTokenWrite(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
window int
|
||||
before int // tokens written before the final single-token write
|
||||
}{
|
||||
{"not-wrapped", 32, 5},
|
||||
{"wrapped", 4, 10},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c := NewRotatingKVCache(tc.window)
|
||||
fillTagged(c, tc.before)
|
||||
|
||||
// Schedule the offset the single-token write reaches as its end
|
||||
// boundary, then perform that write.
|
||||
c.PrepareSnapshots([]int{tc.before + 1})
|
||||
k, v := taggedKV(tc.before, 1)
|
||||
c.Update(newKVBatch(tc.before, 1), k, v)
|
||||
|
||||
snaps := c.TakeSnapshots()
|
||||
if len(snaps) != 1 || snaps[0] == nil {
|
||||
t.Fatalf("snapshot not captured: %v", snaps)
|
||||
}
|
||||
|
||||
if !c.Restore(snaps[0], tc.before+1) {
|
||||
t.Fatal("restore failed")
|
||||
}
|
||||
snaps[0].Close()
|
||||
|
||||
if got, want := windowTags(t, c), wantWindowTags(tc.before+1, tc.window); !slices.Equal(got, want) {
|
||||
t.Fatalf("window tags = %v, want %v", got, want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestRotatingSnapshotSurvivesLaterChunk mirrors chunked prefill: a snapshot
|
||||
// captured during one batched write must still restore correctly after a second
|
||||
// batched write trims/rewrites the buffer the snapshot's slots lived in. This is
|
||||
// the case that forces a lazy snapshot to copy out before the later write destroys it.
|
||||
func TestRotatingSnapshotSurvivesLaterChunk(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
const window = 6
|
||||
c := NewRotatingKVCache(window)
|
||||
|
||||
// Schedule an offset in the first chunk and one in the second, then write
|
||||
// both chunks before taking — the second chunk's concat trims past the first
|
||||
// snapshot's window.
|
||||
c.PrepareSnapshots([]int{4, 12})
|
||||
|
||||
k1, v1 := taggedKV(0, 8) // chunk 1: [0, 8)
|
||||
c.Update(newKVBatch(0, 8), k1, v1)
|
||||
k2, v2 := taggedKV(8, 8) // chunk 2: [8, 16)
|
||||
c.Update(newKVBatch(8, 8), k2, v2)
|
||||
|
||||
snaps := c.TakeSnapshots()
|
||||
if len(snaps) != 2 || snaps[0] == nil || snaps[1] == nil {
|
||||
t.Fatalf("snapshots not captured: %v", snaps)
|
||||
}
|
||||
|
||||
// Restore the first-chunk snapshot (offset 4): its window predates the
|
||||
// second chunk entirely, so the data must have survived the chunk-2 write.
|
||||
if !c.Restore(snaps[0], 4) {
|
||||
t.Fatal("restore to offset 4 failed")
|
||||
}
|
||||
if got, want := windowTags(t, c), wantWindowTags(4, window); !slices.Equal(got, want) {
|
||||
t.Fatalf("offset 4 window tags = %v, want %v", got, want)
|
||||
}
|
||||
|
||||
// Restore the second-chunk snapshot (offset 12) on a fresh cache.
|
||||
c2 := NewRotatingKVCache(window)
|
||||
if !c2.Restore(snaps[1], 12) {
|
||||
t.Fatal("restore to offset 12 failed")
|
||||
}
|
||||
if got, want := windowTags(t, c2), wantWindowTags(12, window); !slices.Equal(got, want) {
|
||||
t.Fatalf("offset 12 window tags = %v, want %v", got, want)
|
||||
}
|
||||
for _, s := range snaps {
|
||||
s.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// TestRotatingLazySnapshotSizeZeroUntilMaterialized verifies the speculation
|
||||
// shape — a single batched write with per-token snapshots and no later write —
|
||||
// leaves every interior capture lazy (no copy-out, so rejected drafts cost only
|
||||
// the lazy arithmetic; the start boundary is the one eager clone), and that a
|
||||
// lazy snapshot reports Size() == 0 until a destructive write copies it out, at
|
||||
// which point the materialize hook fires with the newly-allocated bytes.
|
||||
func TestRotatingLazySnapshotSizeZeroUntilMaterialized(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
const window = 4
|
||||
const before = 10 // wrapped
|
||||
const draft = 4
|
||||
c := NewRotatingKVCache(window)
|
||||
fillTagged(c, before)
|
||||
|
||||
offsets := make([]int, draft)
|
||||
for i := range offsets {
|
||||
offsets[i] = before + i
|
||||
}
|
||||
c.PrepareSnapshots(offsets)
|
||||
|
||||
k, v := taggedKV(before, draft)
|
||||
c.Update(newKVBatch(before, draft), k, v)
|
||||
|
||||
snaps := c.TakeSnapshots()
|
||||
defer func() {
|
||||
for _, s := range snaps {
|
||||
if s != nil {
|
||||
s.Close()
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// Interior captures (offsets after the start boundary) stay lazy: no write
|
||||
// destroyed their slots, so keys is still nil and the issuing cache is live.
|
||||
for i := 1; i < draft; i++ {
|
||||
rs := snaps[i].(*rotatingSnapshot)
|
||||
if rs.keys != nil {
|
||||
t.Fatalf("snaps[%d] copied out (keys != nil); expected a live lazy snapshot", i)
|
||||
}
|
||||
if rs.cache == nil {
|
||||
t.Fatalf("snaps[%d] has no issuing cache; expected a live lazy snapshot", i)
|
||||
}
|
||||
}
|
||||
|
||||
lazy := snaps[1].(*rotatingSnapshot)
|
||||
if lazy.Size() != 0 {
|
||||
t.Fatalf("lazy rotating snapshot Size = %d, want 0", lazy.Size())
|
||||
}
|
||||
|
||||
var hookDelta int
|
||||
lazy.SetMaterializeHook(func(delta int) { hookDelta = delta })
|
||||
|
||||
// Free the cache to force every outstanding lazy snapshot to copy out.
|
||||
c.Free()
|
||||
|
||||
if lazy.keys == nil {
|
||||
t.Fatal("Free did not materialize the lazy snapshot")
|
||||
}
|
||||
want := lazy.keys.NumBytes() + lazy.values.NumBytes()
|
||||
if hookDelta != want {
|
||||
t.Fatalf("hook fired with delta %d, want %d", hookDelta, want)
|
||||
}
|
||||
if lazy.Size() != want {
|
||||
t.Fatalf("materialized Size = %d, want %d", lazy.Size(), want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPerTokenSnapshotPersistsAcrossWrites verifies that scheduled offsets
|
||||
// survive multiple writes until TakeSnapshots — the property prefill would rely
|
||||
// on to snapshot interior offsets without splitting its forward.
|
||||
func TestPerTokenSnapshotPersistsAcrossWrites(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
c := NewKVCache()
|
||||
fillKV(c, 2)
|
||||
|
||||
// Schedule offsets that span two separate writes. Offset 2 equals the
|
||||
// schedule-time position, so it captures a zero-width range (nil); the rest
|
||||
// are edge-local.
|
||||
c.PrepareSnapshots([]int{2, 3, 5})
|
||||
|
||||
k1, v1 := batchKV(2) // reaches offsets 2,3
|
||||
c.Update(newKVBatch(2, 2), k1, v1)
|
||||
k2, v2 := batchKV(2) // reaches offsets 4,5
|
||||
c.Update(newKVBatch(4, 2), k2, v2)
|
||||
|
||||
snaps := c.TakeSnapshots()
|
||||
if len(snaps) != 3 {
|
||||
t.Fatalf("got %d snapshots, want 3", len(snaps))
|
||||
}
|
||||
if snaps[0] != nil {
|
||||
t.Fatalf("snaps[0] = %v, want nil (zero-width base)", snaps[0])
|
||||
}
|
||||
if s := snaps[1].(*kvSnapshot); s.fromOffset != 2 || s.toOffset != 3 {
|
||||
t.Fatalf("snaps[1] = [%d,%d), want [2,3)", s.fromOffset, s.toOffset)
|
||||
}
|
||||
// Offset 5 was scheduled across the second write; its edge starts at the
|
||||
// previous scheduled offset (3), confirming the base cursor only advances
|
||||
// on capture so the snapshot range matches the trie edge between scheduled
|
||||
// offsets — write boundaries between captures must not move it.
|
||||
if s := snaps[2].(*kvSnapshot); s.fromOffset != 3 || s.toOffset != 5 {
|
||||
t.Fatalf("snaps[2] = [%d,%d), want [3,5)", s.fromOffset, s.toOffset)
|
||||
}
|
||||
|
||||
// Restore from the [2,3) edge snapshot to offset 3.
|
||||
if !c.Restore(snaps[1], 3) {
|
||||
t.Fatal("restore to offset 3 failed")
|
||||
}
|
||||
if c.Offset() != 3 {
|
||||
t.Fatalf("offset = %d, want 3", c.Offset())
|
||||
}
|
||||
for _, s := range snaps {
|
||||
if s != nil {
|
||||
s.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRecurrentSnapshotSplitsAndSegmentedCapture verifies that SnapshotSplits
|
||||
// reports the interior scheduled offsets and PutSegmented captures them from the
|
||||
// per-boundary states, so each accepted count restores to a distinct state.
|
||||
func TestRecurrentSnapshotSplitsAndSegmentedCapture(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
const convTail, convDim, nv, vd, kd = 3, 8, 2, 4, 4
|
||||
c := NewRecurrentCache(convTail, convDim, nv, vd, kd)
|
||||
c.Get(newKVBatch(0, 1), mlx.DTypeFloat16)
|
||||
// Advance to offset 5 so the speculative forward starts there.
|
||||
c.Put(newKVBatch(0, 5),
|
||||
[]*mlx.Array{mlx.Zeros(mlx.DTypeFloat16, 1, convTail, convDim)},
|
||||
[]*mlx.Array{mlx.Zeros(mlx.DTypeFloat32, 1, nv, vd, kd)})
|
||||
|
||||
const before, draft = 5, 4
|
||||
offsets := []int{before, before + 1, before + 2, before + 3}
|
||||
c.PrepareSnapshots(offsets)
|
||||
|
||||
splits := c.SnapshotSplits(draft)
|
||||
want := []int{1, 2, 3}
|
||||
if len(splits) != len(want) {
|
||||
t.Fatalf("SnapshotSplits = %v, want %v", splits, want)
|
||||
}
|
||||
for i := range want {
|
||||
if splits[i] != want[i] {
|
||||
t.Fatalf("SnapshotSplits = %v, want %v", splits, want)
|
||||
}
|
||||
}
|
||||
|
||||
// Distinct per-boundary states (3 interior splits + the end) so restore
|
||||
// targets are distinguishable.
|
||||
mkConv := func(s float32) *mlx.Array { return mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat16, 1, convTail, convDim), s) }
|
||||
mkDelta := func(s float32) *mlx.Array { return mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, nv, vd, kd), s) }
|
||||
convStates := []*mlx.Array{mkConv(1), mkConv(2), mkConv(3), mkConv(4)}
|
||||
deltaStates := []*mlx.Array{mkDelta(1), mkDelta(2), mkDelta(3), mkDelta(4)}
|
||||
|
||||
c.Put(newKVBatch(before, draft), convStates, deltaStates)
|
||||
if c.Offset() != before+draft {
|
||||
t.Fatalf("offset after segmented put = %d, want %d", c.Offset(), before+draft)
|
||||
}
|
||||
|
||||
snaps := c.TakeSnapshots()
|
||||
if len(snaps) != draft {
|
||||
t.Fatalf("got %d snapshots, want %d", len(snaps), draft)
|
||||
}
|
||||
for i, s := range snaps {
|
||||
if s == nil {
|
||||
t.Fatalf("snapshot %d not captured", i)
|
||||
}
|
||||
}
|
||||
|
||||
// Full accept (no restore): the live state is the committed end boundary
|
||||
// (value 4) at offset before+draft.
|
||||
st := c.State()
|
||||
mlx.Eval(st[1])
|
||||
if got := st[1].Floats()[0]; got != 4 {
|
||||
t.Fatalf("full-accept delta state = %v, want end boundary value 4", got)
|
||||
}
|
||||
|
||||
// Each partial accept restores to offset before+accepted and must recover the
|
||||
// distinct boundary state captured there: snaps[0] is the pre-forward state
|
||||
// (value 0); snaps[i>=1] is the interior split boundary (value i). Recurrent
|
||||
// snapshots are self-contained, so restores need not run in order.
|
||||
for accepted := range draft {
|
||||
if !c.Restore(snaps[accepted], before+accepted) {
|
||||
t.Fatalf("accepted=%d: restore to before+%d failed", accepted, accepted)
|
||||
}
|
||||
if c.Offset() != before+accepted {
|
||||
t.Fatalf("accepted=%d: offset after restore = %d, want %d", accepted, c.Offset(), before+accepted)
|
||||
}
|
||||
st := c.State()
|
||||
mlx.Eval(st[1])
|
||||
if got := st[1].Floats()[0]; got != float32(accepted) {
|
||||
t.Fatalf("accepted=%d: restored delta state = %v, want boundary value %d", accepted, got, accepted)
|
||||
}
|
||||
}
|
||||
for _, s := range snaps {
|
||||
s.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// TestPrepareSnapshotsPastOffsetPanics verifies scheduling an already-passed offset
|
||||
// is rejected.
|
||||
func TestPrepareSnapshotsPastOffsetPanics(t *testing.T) {
|
||||
skipIfNoMLX(t)
|
||||
|
||||
c := NewKVCache()
|
||||
fillKV(c, 5)
|
||||
|
||||
defer func() {
|
||||
if recover() == nil {
|
||||
t.Fatal("expected panic for already-passed offset")
|
||||
}
|
||||
}()
|
||||
c.PrepareSnapshots([]int{3})
|
||||
}
|
||||
@@ -26,6 +26,63 @@ func (tr *snapshotTracker) track(s *fakeSnapshot) {
|
||||
// Fake caches that store actual token sequences so tests can verify the right
|
||||
// data was restored, not just the right offset.
|
||||
|
||||
// fakePending mirrors the production pendingSnapshots capture machinery for the
|
||||
// fakes: it schedules offsets, captures as feed crosses each one (edge-local,
|
||||
// via a running base cursor), and returns the captures in scheduled order.
|
||||
// capture is the owning fake's snapshot-at-an-offset function.
|
||||
type fakePending struct {
|
||||
offsets []int
|
||||
captured []cache.Snapshot
|
||||
base int
|
||||
}
|
||||
|
||||
func (p *fakePending) prepare(currentOffset int, offsets []int) {
|
||||
p.offsets = slices.Clone(offsets)
|
||||
p.captured = make([]cache.Snapshot, len(offsets))
|
||||
p.base = currentOffset
|
||||
}
|
||||
|
||||
func (p *fakePending) take() []cache.Snapshot {
|
||||
out := p.captured
|
||||
p.offsets, p.captured = nil, nil
|
||||
return out
|
||||
}
|
||||
|
||||
// feedCapturing advances the cache from start over the fed tokens, capturing at
|
||||
// every scheduled offset crossed (including start and end). capture(from,
|
||||
// reached) produces the snapshot for the edge ending at reached;
|
||||
// advance(tokens) appends a token segment to the live state. Segment boundaries
|
||||
// fall at scheduled offsets.
|
||||
func (p *fakePending) feedCapturing(start int, tokens []int32, capture func(from, reached int) cache.Snapshot, advance func([]int32)) {
|
||||
end := start + len(tokens)
|
||||
captureAt := func(reached int) {
|
||||
for i, o := range p.offsets {
|
||||
if p.captured[i] == nil && o == reached {
|
||||
p.captured[i] = capture(p.base, reached)
|
||||
}
|
||||
}
|
||||
p.base = reached
|
||||
}
|
||||
|
||||
if len(p.offsets) == 0 {
|
||||
advance(tokens)
|
||||
return
|
||||
}
|
||||
|
||||
captureAt(start)
|
||||
prev := 0
|
||||
for cut := start + 1; cut < end; cut++ {
|
||||
if !slices.Contains(p.offsets, cut) {
|
||||
continue
|
||||
}
|
||||
advance(tokens[prev : cut-start])
|
||||
captureAt(cut)
|
||||
prev = cut - start
|
||||
}
|
||||
advance(tokens[prev:])
|
||||
captureAt(end)
|
||||
}
|
||||
|
||||
// fakeSnapshot stores a copy of the token sub-sequence it covers.
|
||||
type fakeSnapshot struct {
|
||||
tokens []int32
|
||||
@@ -34,11 +91,23 @@ type fakeSnapshot struct {
|
||||
|
||||
tracker *snapshotTracker
|
||||
closeCount int
|
||||
|
||||
onMaterialize func(delta int)
|
||||
}
|
||||
|
||||
func (s *fakeSnapshot) Size() int { return s.byteSize }
|
||||
func (s *fakeSnapshot) Close() {
|
||||
s.closeCount++
|
||||
func (s *fakeSnapshot) Size() int { return s.byteSize }
|
||||
func (s *fakeSnapshot) SetMaterializeHook(fn func(delta int)) { s.onMaterialize = fn }
|
||||
func (s *fakeSnapshot) Close() { s.closeCount++ }
|
||||
|
||||
// materialize simulates a lazy snapshot copying out: grow byteSize by delta
|
||||
// and fire the trie's hook if one is attached. Tests use this to verify the
|
||||
// trie's counter responds to materialization events.
|
||||
func (s *fakeSnapshot) materialize(delta int) {
|
||||
s.byteSize += delta
|
||||
if s.onMaterialize != nil {
|
||||
s.onMaterialize(delta)
|
||||
s.onMaterialize = nil
|
||||
}
|
||||
}
|
||||
|
||||
// fakeRewindableCache tracks the full token sequence and supports
|
||||
@@ -46,10 +115,13 @@ func (s *fakeSnapshot) Close() {
|
||||
type fakeRewindableCache struct {
|
||||
tokens []int32
|
||||
tracker *snapshotTracker
|
||||
pending fakePending
|
||||
}
|
||||
|
||||
func (c *fakeRewindableCache) feed(tokens []int32) {
|
||||
c.tokens = append(c.tokens, tokens...)
|
||||
c.pending.feedCapturing(len(c.tokens), tokens,
|
||||
func(from, reached int) cache.Snapshot { return c.Snapshot(from) },
|
||||
func(seg []int32) { c.tokens = append(c.tokens, seg...) })
|
||||
}
|
||||
|
||||
func (c *fakeRewindableCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) {
|
||||
@@ -159,6 +231,11 @@ func (c *fakeRewindableCache) Split(snapshot cache.Snapshot, at int) (cache.Snap
|
||||
return p, ch
|
||||
}
|
||||
|
||||
func (c *fakeRewindableCache) PrepareSnapshots(offsets []int) {
|
||||
c.pending.prepare(len(c.tokens), offsets)
|
||||
}
|
||||
func (c *fakeRewindableCache) TakeSnapshots() []cache.Snapshot { return c.pending.take() }
|
||||
|
||||
// fakeSlidingWindowCache models RotatingKVCache semantics: stores the full
|
||||
// token sequence but only the trailing maxSize tokens are "live" in the window.
|
||||
// Once the window fills, live rewind is impossible without a snapshot.
|
||||
@@ -166,10 +243,13 @@ type fakeSlidingWindowCache struct {
|
||||
tokens []int32
|
||||
maxSize int
|
||||
tracker *snapshotTracker
|
||||
pending fakePending
|
||||
}
|
||||
|
||||
func (c *fakeSlidingWindowCache) feed(tokens []int32) {
|
||||
c.tokens = append(c.tokens, tokens...)
|
||||
c.pending.feedCapturing(len(c.tokens), tokens,
|
||||
func(from, reached int) cache.Snapshot { return c.Snapshot(0) },
|
||||
func(seg []int32) { c.tokens = append(c.tokens, seg...) })
|
||||
}
|
||||
|
||||
func (c *fakeSlidingWindowCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) {
|
||||
@@ -241,15 +321,23 @@ func (c *fakeSlidingWindowCache) Split(snapshot cache.Snapshot, at int) (cache.S
|
||||
return nil, snapshot
|
||||
}
|
||||
|
||||
func (c *fakeSlidingWindowCache) PrepareSnapshots(offsets []int) {
|
||||
c.pending.prepare(len(c.tokens), offsets)
|
||||
}
|
||||
func (c *fakeSlidingWindowCache) TakeSnapshots() []cache.Snapshot { return c.pending.take() }
|
||||
|
||||
// fakeRecurrentCache models RecurrentCache semantics: stores tokens
|
||||
// but cannot rewind without a snapshot.
|
||||
type fakeRecurrentCache struct {
|
||||
tokens []int32
|
||||
tracker *snapshotTracker
|
||||
pending fakePending
|
||||
}
|
||||
|
||||
func (c *fakeRecurrentCache) feed(tokens []int32) {
|
||||
c.tokens = append(c.tokens, tokens...)
|
||||
c.pending.feedCapturing(len(c.tokens), tokens,
|
||||
func(from, reached int) cache.Snapshot { return c.Snapshot(0) },
|
||||
func(seg []int32) { c.tokens = append(c.tokens, seg...) })
|
||||
}
|
||||
|
||||
func (c *fakeRecurrentCache) Update(keys, values *mlx.Array) (*mlx.Array, *mlx.Array) {
|
||||
@@ -300,6 +388,11 @@ func (c *fakeRecurrentCache) Split(snapshot cache.Snapshot, at int) (cache.Snaps
|
||||
return nil, snapshot // can't split cumulative state
|
||||
}
|
||||
|
||||
func (c *fakeRecurrentCache) PrepareSnapshots(offsets []int) {
|
||||
c.pending.prepare(len(c.tokens), offsets)
|
||||
}
|
||||
func (c *fakeRecurrentCache) TakeSnapshots() []cache.Snapshot { return c.pending.take() }
|
||||
|
||||
type feedableCache interface {
|
||||
cache.Cache
|
||||
feed(tokens []int32)
|
||||
@@ -936,3 +1029,60 @@ func TestLRUOnlyUpdatesUsedNodes(t *testing.T) {
|
||||
checkTrieInvariants(t, kvc.root)
|
||||
})
|
||||
}
|
||||
|
||||
// TestPagedOutBytesUpdatesOnMaterialize verifies that when a snapshot owned
|
||||
// by a trie node materializes (allocates owned bytes from a previously lazy
|
||||
// state), the trie's pagedOutBytes counter picks up the delta via the
|
||||
// installed materialize hook.
|
||||
func TestPagedOutBytesUpdatesOnMaterialize(t *testing.T) {
|
||||
kvc := &kvCache{}
|
||||
kvc.ensureRoot()
|
||||
|
||||
node := &trieNode{parent: kvc.root, tokens: []int32{1, 2, 3}, endOffset: 3}
|
||||
kvc.root.children = append(kvc.root.children, node)
|
||||
|
||||
snap := &fakeSnapshot{from: 0, to: 3, byteSize: 0}
|
||||
node.setSnapshots([]cache.Snapshot{snap}, &kvc.pagedOutBytes)
|
||||
|
||||
if kvc.pagedOutBytes != 0 {
|
||||
t.Fatalf("pagedOutBytes after install = %d, want 0 (lazy snapshot)", kvc.pagedOutBytes)
|
||||
}
|
||||
|
||||
const materialized = 1 << 20
|
||||
snap.materialize(materialized)
|
||||
|
||||
if kvc.pagedOutBytes != materialized {
|
||||
t.Fatalf("pagedOutBytes after materialize = %d, want %d", kvc.pagedOutBytes, materialized)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSwapSnapshotsDetachesHook verifies that snapshots removed from a trie
|
||||
// node via swapSnapshots no longer feed the trie's counter when they later
|
||||
// materialize. Without detach, a Split/Merge that folds an old snapshot
|
||||
// elsewhere would double-count its bytes if it copied out afterward.
|
||||
func TestSwapSnapshotsDetachesHook(t *testing.T) {
|
||||
kvc := &kvCache{}
|
||||
kvc.ensureRoot()
|
||||
|
||||
node := &trieNode{parent: kvc.root, tokens: []int32{1, 2, 3}, endOffset: 3}
|
||||
kvc.root.children = append(kvc.root.children, node)
|
||||
|
||||
snap := &fakeSnapshot{from: 0, to: 3, byteSize: 0}
|
||||
node.setSnapshots([]cache.Snapshot{snap}, &kvc.pagedOutBytes)
|
||||
|
||||
replacement := &fakeSnapshot{from: 0, to: 3, byteSize: 0}
|
||||
old := node.swapSnapshots([]cache.Snapshot{replacement}, &kvc.pagedOutBytes)
|
||||
if len(old) != 1 || old[0] != snap {
|
||||
t.Fatalf("swapSnapshots returned %v, want the original snap", old)
|
||||
}
|
||||
|
||||
snap.materialize(1 << 20)
|
||||
if kvc.pagedOutBytes != 0 {
|
||||
t.Fatalf("pagedOutBytes after detached materialize = %d, want 0", kvc.pagedOutBytes)
|
||||
}
|
||||
|
||||
replacement.materialize(2 << 20)
|
||||
if kvc.pagedOutBytes != 2<<20 {
|
||||
t.Fatalf("pagedOutBytes after replacement materialize = %d, want %d", kvc.pagedOutBytes, 2<<20)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -51,14 +51,27 @@ func (n *trieNode) setSnapshots(snaps []cache.Snapshot, counter *int64) {
|
||||
// swapSnapshots is like setSnapshots but returns the previous snapshots
|
||||
// without closing them. Use this when the old snapshots will be consumed
|
||||
// (e.g. by Split/Merge).
|
||||
//
|
||||
// Snapshots that are lazy when installed contribute 0 to the counter, but
|
||||
// may later materialize via copyOut and grow the counter through a hook.
|
||||
func (n *trieNode) swapSnapshots(snaps []cache.Snapshot, counter *int64) []cache.Snapshot {
|
||||
old := n.snapshots
|
||||
for _, s := range old {
|
||||
if s != nil {
|
||||
s.SetMaterializeHook(nil)
|
||||
}
|
||||
}
|
||||
if counter != nil {
|
||||
*counter -= n.snapshotBytes()
|
||||
}
|
||||
n.snapshots = snaps
|
||||
if counter != nil {
|
||||
*counter += n.snapshotBytes()
|
||||
for _, s := range snaps {
|
||||
if s != nil {
|
||||
s.SetMaterializeHook(func(delta int) { *counter += int64(delta) })
|
||||
}
|
||||
}
|
||||
}
|
||||
return old
|
||||
}
|
||||
|
||||
@@ -29,27 +29,24 @@ const (
|
||||
)
|
||||
|
||||
type mtpStats struct {
|
||||
iterations int
|
||||
drafted int
|
||||
accepted int
|
||||
mismatches int
|
||||
allAccepted int
|
||||
batched int
|
||||
serial int
|
||||
compared int
|
||||
batchSerialMismatches int
|
||||
maxDraft int
|
||||
targetDuration time.Duration
|
||||
draftDuration time.Duration
|
||||
validateDuration time.Duration
|
||||
iterations int
|
||||
drafted int
|
||||
accepted int
|
||||
mismatches int
|
||||
allAccepted int
|
||||
batched int
|
||||
serial int
|
||||
maxDraft int
|
||||
targetDuration time.Duration
|
||||
draftDuration time.Duration
|
||||
validateDuration time.Duration
|
||||
}
|
||||
|
||||
type mtpOptions struct {
|
||||
initialDraftTokens int
|
||||
maxDraftTokens int
|
||||
draftSchedule mtpDraftSchedule
|
||||
serialValidate bool
|
||||
compareSerialValidate bool
|
||||
initialDraftTokens int
|
||||
maxDraftTokens int
|
||||
draftSchedule mtpDraftSchedule
|
||||
serialValidate bool
|
||||
}
|
||||
|
||||
func (r *Runner) mtpDefaults(sample bool) base.MTPDefaults {
|
||||
@@ -90,9 +87,6 @@ func (r *Runner) loadMTPOptions(sample bool) mtpOptions {
|
||||
if b, err := strconv.ParseBool(os.Getenv("OLLAMA_MLX_MTP_SERIAL_VALIDATE")); err == nil {
|
||||
opts.serialValidate = b
|
||||
}
|
||||
if b, err := strconv.ParseBool(os.Getenv("OLLAMA_MLX_MTP_COMPARE_SERIAL_VALIDATE")); err == nil {
|
||||
opts.compareSerialValidate = b
|
||||
}
|
||||
switch schedule := strings.ToLower(strings.TrimSpace(os.Getenv("OLLAMA_MLX_MTP_DRAFT_SCHEDULE"))); schedule {
|
||||
case "", string(mtpDraftScheduleConstant):
|
||||
opts.draftSchedule = mtpDraftScheduleConstant
|
||||
@@ -146,9 +140,6 @@ func (r *Runner) useSampleMTP(opts sampler.Options) bool {
|
||||
if serial, err := strconv.ParseBool(os.Getenv("OLLAMA_MLX_MTP_SERIAL_VALIDATE")); err == nil && serial {
|
||||
return false
|
||||
}
|
||||
if compare, err := strconv.ParseBool(os.Getenv("OLLAMA_MLX_MTP_COMPARE_SERIAL_VALIDATE")); err == nil && compare {
|
||||
return false
|
||||
}
|
||||
if r.Draft == nil {
|
||||
return false
|
||||
}
|
||||
@@ -173,7 +164,7 @@ func (r *Runner) runGreedyMTPDecode(ctx context.Context, request Request, sessio
|
||||
mtpOpts := r.loadMTPOptions(false)
|
||||
stats := mtpStats{maxDraft: mtpOpts.initialDraftTokens}
|
||||
draftLimit := mtpOpts.initialDraftTokens
|
||||
slog.Info("MTP greedy decode enabled", "initial_draft_tokens", mtpOpts.initialDraftTokens, "max_draft_tokens", mtpOpts.maxDraftTokens, "draft_schedule", mtpOpts.draftSchedule, "serial_validate", mtpOpts.serialValidate, "compare_serial_validate", mtpOpts.compareSerialValidate)
|
||||
slog.Info("MTP greedy decode enabled", "initial_draft_tokens", mtpOpts.initialDraftTokens, "max_draft_tokens", mtpOpts.maxDraftTokens, "draft_schedule", mtpOpts.draftSchedule, "serial_validate", mtpOpts.serialValidate)
|
||||
|
||||
targetForward := func(token *mlx.Array) *mlx.Array {
|
||||
fwd := r.Model.Forward(&batch.Batch{
|
||||
@@ -301,7 +292,7 @@ func (r *Runner) runGreedyMTPDecode(ctx context.Context, request Request, sessio
|
||||
avgDraft = float64(stats.drafted) / float64(stats.iterations)
|
||||
avgAccepted = float64(stats.accepted) / float64(stats.iterations)
|
||||
}
|
||||
slog.Info("MTP decode stats", "generated", generated, "drafted", stats.drafted, "accepted", stats.accepted, "acceptance", acceptance, "iterations", stats.iterations, "avg_draft", avgDraft, "avg_accepted", avgAccepted, "batched", stats.batched, "serial", stats.serial, "compared", stats.compared, "batch_serial_mismatches", stats.batchSerialMismatches, "mismatches", stats.mismatches, "all_accepted", stats.allAccepted, "max_draft", stats.maxDraft, "draft_schedule", mtpOpts.draftSchedule, "target_duration", stats.targetDuration, "draft_duration", stats.draftDuration, "validate_duration", stats.validateDuration)
|
||||
slog.Info("MTP decode stats", "generated", generated, "drafted", stats.drafted, "accepted", stats.accepted, "acceptance", acceptance, "iterations", stats.iterations, "avg_draft", avgDraft, "avg_accepted", avgAccepted, "batched", stats.batched, "serial", stats.serial, "mismatches", stats.mismatches, "all_accepted", stats.allAccepted, "max_draft", stats.maxDraft, "draft_schedule", mtpOpts.draftSchedule, "target_duration", stats.targetDuration, "draft_duration", stats.draftDuration, "validate_duration", stats.validateDuration)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
@@ -544,19 +535,68 @@ func (r *Runner) acceptMTPDrafts(ctx context.Context, request Request, session *
|
||||
return r.acceptMTPDraftsSerial(ctx, request, session, dec, caches, position, baseLogits, draftTokens, final, generated)
|
||||
}
|
||||
|
||||
specCaches, spec, ok := cache.BeginSpeculation(caches)
|
||||
if ok {
|
||||
stats.batched++
|
||||
return r.acceptMTPDraftsBatched(ctx, request, session, dec, caches, specCaches, spec, position, baseLogits, draftTokens, final, generated, stats, opts)
|
||||
}
|
||||
|
||||
stats.serial++
|
||||
return r.acceptMTPDraftsSerial(ctx, request, session, dec, caches, position, baseLogits, draftTokens, final, generated)
|
||||
stats.batched++
|
||||
return r.acceptMTPDraftsBatched(ctx, request, session, dec, caches, position, baseLogits, draftTokens, final, generated)
|
||||
}
|
||||
|
||||
func (r *Runner) acceptMTPDraftsBatched(ctx context.Context, request Request, session *cacheSession, dec *decoder, liveCaches []cache.Cache, caches []cache.Cache, spec *cache.Speculation, position *int, baseLogits *mlx.Array, draftTokens *mlx.Array, final *CompletionResponse, generated *int, stats *mtpStats, opts mtpOptions) (sampler.Result, int, bool, error) {
|
||||
// scheduleSpeculation schedules per-token snapshots at offsets
|
||||
// [before, before+draftCount) on every cache, so the speculative forward
|
||||
// captures a rollback point before each draft token's write.
|
||||
func scheduleSpeculation(caches []cache.Cache, before, draftCount int) {
|
||||
offsets := make([]int, draftCount)
|
||||
for i := range offsets {
|
||||
offsets[i] = before + i
|
||||
}
|
||||
for _, c := range caches {
|
||||
if c != nil {
|
||||
c.PrepareSnapshots(offsets)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// commitSpeculation collects the captured snapshots and rolls every cache
|
||||
// back to before+accepted, keeping only the accepted prefix of the
|
||||
// speculative forward. The bonus token (at before+draftCount) is never
|
||||
// committed. On full acceptance no restore is needed.
|
||||
//
|
||||
// Rollback tries a live rewind first (Restore(nil)) and falls back to the
|
||||
// captured snapshot when the live state can't be rewound in place.
|
||||
func commitSpeculation(caches []cache.Cache, accepted, draftCount, before int) {
|
||||
target := before + accepted
|
||||
for _, c := range caches {
|
||||
if c == nil {
|
||||
continue
|
||||
}
|
||||
snaps := c.TakeSnapshots()
|
||||
if accepted < draftCount {
|
||||
// Close the snapshots we won't restore from before restoring: a
|
||||
// snapshot restore on a wrapped RotatingKVCache copies out every
|
||||
// outstanding lazy snapshot before it rebuilds the buffer, so
|
||||
// dropping the unused ones first stops that copy-out from
|
||||
// materializing snapshots we are about to discard anyway.
|
||||
for i, s := range snaps {
|
||||
if s != nil && i != accepted {
|
||||
s.Close()
|
||||
snaps[i] = nil
|
||||
}
|
||||
}
|
||||
if !c.Restore(nil, target) && !c.Restore(snaps[accepted], target) {
|
||||
panic(fmt.Sprintf("mtp: cache restore to %d failed", target))
|
||||
}
|
||||
}
|
||||
for _, s := range snaps {
|
||||
if s != nil {
|
||||
s.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Runner) acceptMTPDraftsBatched(ctx context.Context, request Request, session *cacheSession, dec *decoder, caches []cache.Cache, position *int, baseLogits *mlx.Array, draftTokens *mlx.Array, final *CompletionResponse, generated *int) (sampler.Result, int, bool, error) {
|
||||
before := *position
|
||||
draftCount := draftTokens.Dim(1)
|
||||
|
||||
scheduleSpeculation(caches, before, draftCount)
|
||||
hiddenSeq := r.Model.Forward(&batch.Batch{
|
||||
InputIDs: draftTokens,
|
||||
SeqOffsets: []int32{int32(before)},
|
||||
@@ -572,6 +612,10 @@ func (r *Runner) acceptMTPDraftsBatched(ctx context.Context, request Request, se
|
||||
draftIDs := draftTokens.Ints()
|
||||
selectedIDs := selectedTokens.Ints()
|
||||
if len(selectedIDs) < draftCount+1 {
|
||||
// Drain the scheduled snapshots and roll the speculative forward back
|
||||
// out of the live caches before bailing, so the abandoned drafts don't
|
||||
// reach the trie via session.close().
|
||||
commitSpeculation(caches, 0, draftCount, before)
|
||||
return sampler.Result{}, accepted, false, fmt.Errorf("mtp validation produced %d tokens for %d draft tokens", len(selectedIDs), draftCount)
|
||||
}
|
||||
|
||||
@@ -587,11 +631,7 @@ func (r *Runner) acceptMTPDraftsBatched(ctx context.Context, request Request, se
|
||||
}
|
||||
}
|
||||
|
||||
if opts.compareSerialValidate {
|
||||
spec.Commit(0)
|
||||
r.compareMTPBatchedWithSerial(ctx, liveCaches, before, baseLogits, hiddenSeq, draftIDs, selectedIDs, accepted, draftCount, stats)
|
||||
}
|
||||
spec.Commit(accepted)
|
||||
commitSpeculation(caches, accepted, draftCount, before)
|
||||
*position = before + accepted
|
||||
|
||||
for _, id := range draftIDs[:accepted] {
|
||||
@@ -623,20 +663,16 @@ func (r *Runner) acceptMTPDraftsBatched(ctx context.Context, request Request, se
|
||||
}
|
||||
|
||||
func (r *Runner) acceptSampleMTPDrafts(ctx context.Context, request Request, session *cacheSession, dec *decoder, caches []cache.Cache, position *int, baseLogits *mlx.Array, candidates *mtpDraftCandidates, final *CompletionResponse, generated *int, stats *mtpStats) (sampler.Result, int, bool, error) {
|
||||
specCaches, spec, ok := cache.BeginSpeculation(caches)
|
||||
if !ok {
|
||||
stats.serial++
|
||||
return r.Sampler.Sample([]int{pipelineSlot}, baseLogits), 0, false, nil
|
||||
}
|
||||
stats.batched++
|
||||
|
||||
before := *position
|
||||
draftCount := candidates.tokens.Dim(1)
|
||||
scheduleSpeculation(caches, before, draftCount)
|
||||
hiddenSeq := r.Model.Forward(&batch.Batch{
|
||||
InputIDs: candidates.tokens,
|
||||
SeqOffsets: []int32{int32(before)},
|
||||
SeqQueryLens: []int32{int32(draftCount)},
|
||||
}, specCaches)
|
||||
}, caches)
|
||||
|
||||
targetDist := r.Sampler.Distribution(pipelineSlot, r.mtpValidationLogits(baseLogits, hiddenSeq), candidates.tokens)
|
||||
draftDist := candidates.dist
|
||||
@@ -653,6 +689,10 @@ func (r *Runner) acceptSampleMTPDrafts(ctx context.Context, request Request, ses
|
||||
accepted++
|
||||
}
|
||||
if accepted > draftCount {
|
||||
// Drain the scheduled snapshots and roll the speculative forward back
|
||||
// out of the live caches before bailing, so the abandoned drafts don't
|
||||
// reach the trie via session.close().
|
||||
commitSpeculation(caches, 0, draftCount, before)
|
||||
return sampler.Result{}, 0, false, fmt.Errorf("mtp sample validation accepted %d tokens for %d draft tokens", accepted, draftCount)
|
||||
}
|
||||
|
||||
@@ -668,7 +708,7 @@ func (r *Runner) acceptSampleMTPDrafts(ctx context.Context, request Request, ses
|
||||
}
|
||||
}
|
||||
|
||||
spec.Commit(accepted)
|
||||
commitSpeculation(caches, accepted, draftCount, before)
|
||||
*position = before + accepted
|
||||
|
||||
for _, id := range draftIDs[:accepted] {
|
||||
@@ -749,106 +789,6 @@ func mtpTokenVector(token *mlx.Array) *mlx.Array {
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Runner) compareMTPBatchedWithSerial(ctx context.Context, caches []cache.Cache, before int, baseLogits, hiddenSeq *mlx.Array, draftIDs, selectedIDs []int, accepted, draftCount int, stats *mtpStats) {
|
||||
serialCaches, ok := cache.BeginIsolatedSpeculation(caches)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
compareCount := accepted + 1
|
||||
if accepted == draftCount {
|
||||
// Include the target bonus token when every draft was accepted.
|
||||
compareCount = draftCount + 1
|
||||
}
|
||||
|
||||
serialLogits := baseLogits
|
||||
for i := range compareCount {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return
|
||||
}
|
||||
if i >= len(selectedIDs) {
|
||||
return
|
||||
}
|
||||
|
||||
batchedLogits := baseLogits
|
||||
if i > 0 {
|
||||
batchedLogits = r.targetLogitsAt(hiddenSeq, i-1)
|
||||
}
|
||||
|
||||
batchedToken := greedyTokenFromLogits(batchedLogits)
|
||||
serialToken := greedyTokenFromLogits(serialLogits)
|
||||
mlx.Eval(batchedToken, serialToken)
|
||||
|
||||
batchedID := tokenID(batchedToken)
|
||||
vectorizedID := selectedIDs[i]
|
||||
serialID := tokenID(serialToken)
|
||||
stats.compared++
|
||||
if vectorizedID != serialID {
|
||||
firstMismatch := stats.batchSerialMismatches == 0
|
||||
stats.batchSerialMismatches++
|
||||
if !firstMismatch {
|
||||
return
|
||||
}
|
||||
|
||||
draftID := -1
|
||||
if i < draftCount {
|
||||
draftID = draftIDs[i]
|
||||
}
|
||||
batchedTop := top2FromLogits(batchedLogits)
|
||||
serialTop := top2FromLogits(serialLogits)
|
||||
slog.Warn("MTP batched validation differs from serial validation",
|
||||
"position", before+i,
|
||||
"draft", draftID,
|
||||
"batched", vectorizedID,
|
||||
"batched_slice", batchedID,
|
||||
"serial", serialID,
|
||||
"batched_slice_top1", batchedTop.firstToken,
|
||||
"batched_slice_top2", batchedTop.secondToken,
|
||||
"batched_slice_margin", batchedTop.margin,
|
||||
"serial_top1", serialTop.firstToken,
|
||||
"serial_top2", serialTop.secondToken,
|
||||
"serial_margin", serialTop.margin,
|
||||
)
|
||||
return
|
||||
}
|
||||
|
||||
if i >= draftCount || i >= accepted {
|
||||
return
|
||||
}
|
||||
|
||||
hidden := r.Model.Forward(&batch.Batch{
|
||||
InputIDs: mlx.FromValues([]int32{int32(draftIDs[i])}, 1, 1),
|
||||
SeqOffsets: []int32{int32(before + i)},
|
||||
SeqQueryLens: []int32{1},
|
||||
}, serialCaches)
|
||||
serialLogits = r.lastLogits(hidden)
|
||||
}
|
||||
}
|
||||
|
||||
type mtpTop2 struct {
|
||||
firstToken int
|
||||
secondToken int
|
||||
margin float64
|
||||
}
|
||||
|
||||
func top2FromLogits(logits *mlx.Array) mtpTop2 {
|
||||
indices := logits.Negative().ArgsortAxis(-1).Slice(mlx.Slice(), mlx.Slice(0, 2))
|
||||
indices32 := indices.AsType(mlx.DTypeInt32)
|
||||
values := logits.TakeAlongAxis(indices, -1).AsType(mlx.DTypeFloat32)
|
||||
mlx.Eval(indices32, values)
|
||||
|
||||
tokenIDs := indices32.Ints()
|
||||
logitValues := values.Floats()
|
||||
if len(tokenIDs) < 2 || len(logitValues) < 2 {
|
||||
return mtpTop2{}
|
||||
}
|
||||
return mtpTop2{
|
||||
firstToken: tokenIDs[0],
|
||||
secondToken: tokenIDs[1],
|
||||
margin: float64(logitValues[0] - logitValues[1]),
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Runner) acceptMTPDraftsSerial(ctx context.Context, request Request, session *cacheSession, dec *decoder, caches []cache.Cache, position *int, baseLogits *mlx.Array, draftTokens *mlx.Array, final *CompletionResponse, generated *int) (sampler.Result, int, bool, error) {
|
||||
logits := baseLogits
|
||||
accepted := 0
|
||||
@@ -912,11 +852,6 @@ func (r *Runner) lastLogits(hidden *mlx.Array) *mlx.Array {
|
||||
return r.lastLogitsFromLogits(logits)
|
||||
}
|
||||
|
||||
func (r *Runner) targetLogitsAt(hiddenSeq *mlx.Array, index int) *mlx.Array {
|
||||
hidden := hiddenSeq.Slice(mlx.Slice(), mlx.Slice(index), mlx.Slice())
|
||||
return r.lastLogits(hidden)
|
||||
}
|
||||
|
||||
func (r *Runner) mtpValidationTokens(baseLogits, hiddenSeq *mlx.Array) *mlx.Array {
|
||||
return greedyTokenFromLogits(r.mtpValidationLogits(baseLogits, hiddenSeq))
|
||||
}
|
||||
|
||||
@@ -293,10 +293,10 @@ func (m *AssistantModel) sharedHistories(b *batch.Batch, caches []cache.Cache) (
|
||||
if len(caches) < 2 {
|
||||
return nil, nil
|
||||
}
|
||||
if v, ok := caches[len(caches)-2].(cache.Viewer); ok {
|
||||
if v, ok := caches[len(caches)-2].(cache.Attention); ok {
|
||||
sliding = v.View(b)
|
||||
}
|
||||
if v, ok := caches[len(caches)-1].(cache.Viewer); ok {
|
||||
if v, ok := caches[len(caches)-1].(cache.Attention); ok {
|
||||
full = v.View(b)
|
||||
}
|
||||
return sliding, full
|
||||
@@ -376,12 +376,7 @@ func (a *AssistantAttention) Forward(x *mlx.Array, b *batch.Batch, positions *ml
|
||||
}
|
||||
q = mlx.RoPEWithFreqs(q, ropeDims, false, ropeBase, 1.0, positions, ropeFreqs)
|
||||
|
||||
mask := nn.CausalMask()
|
||||
if isSliding && cfg.SlidingWindow > 0 {
|
||||
mask = mask.Intersect(nn.SlidingWindowMask(b, history.K().Dim(2), int(cfg.SlidingWindow), q.DType()))
|
||||
}
|
||||
|
||||
out := nn.ScaledDotProductAttention(b, q, scale, nn.WithKVHistory(history), nn.WithMask(mask))
|
||||
out := nn.ScaledDotProductAttention(b, q, scale, nn.WithKVHistory(history), nn.WithMask(nn.CausalMask()))
|
||||
out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), B, L, cfg.NumAttentionHeads*headDim)
|
||||
if !mlx.MetalIsAvailable() {
|
||||
out = mlx.Contiguous(out, false)
|
||||
|
||||
@@ -1165,18 +1165,24 @@ func (g *GatedDeltaNet) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, B,
|
||||
}
|
||||
convTail := cfg.LinearConvKernelDim - 1
|
||||
var rc *cache.RecurrentCache
|
||||
var rec nn.RecurrentOption
|
||||
opts := make([]nn.RecurrentOption, 0, 2)
|
||||
if typed, ok := c.(*cache.RecurrentCache); ok {
|
||||
rc = typed
|
||||
rec = nn.WithRecurrentHistory(rc.Get(b, x.DType()))
|
||||
opts = append(opts, nn.WithRecurrentHistory(rc.Get(b, x.DType())))
|
||||
// When the cache has scheduled per-token snapshots, segment the
|
||||
// recurrent kernels at the interior offsets so each boundary state
|
||||
// can be captured.
|
||||
if splits := rc.SnapshotSplits(int(L)); len(splits) > 0 {
|
||||
opts = append(opts, nn.WithSnapshotSplits(splits))
|
||||
}
|
||||
} else {
|
||||
rec = nn.WithRecurrentState(
|
||||
opts = append(opts, nn.WithRecurrentState(
|
||||
mlx.Zeros(x.DType(), int(B), int(convTail), qkv.Dim(2)),
|
||||
mlx.Zeros(mlx.DTypeFloat32, int(B), int(cfg.LinearNumValueHeads), int(cfg.LinearValueHeadDim), int(cfg.LinearKeyHeadDim)),
|
||||
)
|
||||
))
|
||||
}
|
||||
|
||||
convOut, convStates := nn.CausalConv1D(b, qkv, g.Conv1D, g.ConvWeight, int(convTail), rec)
|
||||
convOut, convStates := nn.CausalConv1D(b, qkv, g.Conv1D, g.ConvWeight, int(convTail), opts...)
|
||||
convOut = mlx.SiLU(convOut)
|
||||
|
||||
keyDim := cfg.LinearNumKeyHeads * cfg.LinearKeyHeadDim
|
||||
@@ -1198,7 +1204,7 @@ func (g *GatedDeltaNet) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, B,
|
||||
|
||||
betaGate := mlx.Sigmoid(beta)
|
||||
|
||||
out, deltaStates := nn.GatedDelta(b, q, k, v, gDecay, betaGate, rec)
|
||||
out, deltaStates := nn.GatedDelta(b, q, k, v, gDecay, betaGate, opts...)
|
||||
outDType := out.DType()
|
||||
out = mlx.RMSNormFn(out, g.NormWeight, cfg.RMSNormEps)
|
||||
out = mlx.Mul(out.AsType(mlx.DTypeFloat32), mlx.SiLU(z.AsType(mlx.DTypeFloat32))).AsType(outDType)
|
||||
|
||||
Reference in New Issue
Block a user