cache

package
v0.32.1 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Jul 15, 2026 License: MIT Imports: 6 Imported by: 0

Documentation

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

This section is empty.

Types

type Attention added in v0.22.1 

type Attention interface {
	Cache

	// 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
}

Attention is the contract for caches that back attention layers (KVCache, RotatingKVCache).

type Cache 

type Cache interface {
	// State returns the cache-owned state roots that should be kept/evaluated.
	State() []*mlx.Array
	Free()
	Offset() int

	// Snapshot copies cache state from fromOffset to current offset into
	// 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).
	Restore(snapshot Snapshot, target int) bool

	// Merge combines two sequential snapshots [a,b) and [b,c) into [a,c).
	// Takes ownership of both inputs.
	Merge(parent, child Snapshot) Snapshot

	// Split divides a snapshot [a,c) at offset b into [a,b) and [b,c).
	// Takes ownership of the input. Cache types that cannot split
	// (e.g. recurrent) return (nil, snapshot).
	Split(snapshot Snapshot, at int) (parent, child Snapshot)
}

Cache is common state management shared by every cache kind. Writers live on the specific caches

type KVCache 

type KVCache struct {
	// contains filtered or unexported fields
}

func NewKVCache 

func NewKVCache() *KVCache

func (*KVCache) Free 

func (c *KVCache) Free()

func (*KVCache) Merge added in v0.18.3 

func (c *KVCache) Merge(parent, child Snapshot) Snapshot

func (*KVCache) Offset 

func (c *KVCache) Offset() int

func (*KVCache) PrepareSnapshots added in v0.30.8 

func (c *KVCache) PrepareSnapshots(offsets []int)

func (*KVCache) Restore added in v0.18.3 

func (c *KVCache) Restore(snapshot Snapshot, target int) bool

func (*KVCache) Snapshot added in v0.18.3 

func (c *KVCache) Snapshot(fromOffset int) Snapshot

func (*KVCache) Split added in v0.18.3 

func (c *KVCache) Split(snapshot Snapshot, at int) (Snapshot, Snapshot)

func (*KVCache) State 

func (c *KVCache) State() []*mlx.Array

func (*KVCache) TakeSnapshots added in v0.30.8 

func (c *KVCache) TakeSnapshots() []Snapshot

func (*KVCache) Update 

func (c *KVCache) Update(_ *batch.Batch, keys, values *mlx.Array) *nn.KVHistory

Assumes B = 1; heterogeneous batches are not supported.

func (*KVCache) View added in v0.23.1 

func (c *KVCache) View(_ *batch.Batch) *nn.KVHistory

View returns the current cache contents as attention history without writing.

type RecurrentCache 

type RecurrentCache struct {
	// contains filtered or unexported fields
}

RecurrentCache stores state for linear-recurrent layers.

Conv state takes its dtype from the first Get call (the activation dtype). Delta state is always float32: the gated-delta recurrent accumulator runs for the full sequence length and needs the extra precision regardless of activation dtype.

Conv state shape: [B, convTail, convDim] Delta state shape: [B, numVHeads, headVDim, headKDim]

func NewRecurrentCache 

func NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim int32) *RecurrentCache

func (*RecurrentCache) Free 

func (c *RecurrentCache) Free()

func (*RecurrentCache) Get added in v0.22.1 

func (c *RecurrentCache) Get(b *batch.Batch, dtype mlx.DType) *nn.RecurrentHistory

Get returns the current conv/delta state for the SSM layer's read phase. On first call it lazy-initializes zero-filled state tensors sized from b.InputIDs and dtyped from the caller's activation dtype. On subsequent calls it returns the existing state; batch size and dtype must match the first call, since recurrent state is cumulative and cannot be reshaped without losing history.

func (*RecurrentCache) Merge added in v0.18.3 

func (c *RecurrentCache) Merge(parent, child Snapshot) Snapshot

func (*RecurrentCache) Offset 

func (c *RecurrentCache) Offset() int

func (*RecurrentCache) PrepareSnapshots added in v0.30.8 

func (c *RecurrentCache) PrepareSnapshots(offsets []int)

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 (*RecurrentCache) Put added in v0.22.1 

func (c *RecurrentCache) Put(b *batch.Batch, convStates, deltaStates []*mlx.Array)

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 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 (*RecurrentCache) Restore added in v0.18.3 

func (c *RecurrentCache) Restore(snapshot Snapshot, target int) bool

func (*RecurrentCache) Snapshot added in v0.18.3 

func (c *RecurrentCache) Snapshot(fromOffset int) Snapshot

func (*RecurrentCache) SnapshotSplits added in v0.30.8 

func (c *RecurrentCache) SnapshotSplits(forwardLen int) []int

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 (*RecurrentCache) Split added in v0.18.3 

func (c *RecurrentCache) Split(snapshot Snapshot, at int) (Snapshot, Snapshot)

func (*RecurrentCache) State 

func (c *RecurrentCache) State() []*mlx.Array

func (*RecurrentCache) TakeSnapshots added in v0.30.8 

func (c *RecurrentCache) TakeSnapshots() []Snapshot

type RotatingKVCache 

type RotatingKVCache struct {
	// contains filtered or unexported fields
}

RotatingKVCache implements sliding window attention with bounded memory.

func NewRotatingKVCache 

func NewRotatingKVCache(maxSize int) *RotatingKVCache

func (*RotatingKVCache) Free added in v0.18.3 

func (c *RotatingKVCache) Free()

func (*RotatingKVCache) Merge added in v0.18.3 

func (c *RotatingKVCache) Merge(parent, child Snapshot) Snapshot

func (*RotatingKVCache) Offset added in v0.30.8 

func (c *RotatingKVCache) Offset() int

func (*RotatingKVCache) PrepareSnapshots added in v0.30.8 

func (c *RotatingKVCache) PrepareSnapshots(offsets []int)

func (*RotatingKVCache) Restore added in v0.18.3 

func (c *RotatingKVCache) Restore(snapshot Snapshot, target int) bool

func (*RotatingKVCache) Snapshot added in v0.18.3 

func (c *RotatingKVCache) Snapshot(fromOffset int) Snapshot

func (*RotatingKVCache) Split added in v0.18.3 

func (c *RotatingKVCache) Split(snapshot Snapshot, at int) (Snapshot, Snapshot)

func (*RotatingKVCache) State 

func (c *RotatingKVCache) State() []*mlx.Array

func (*RotatingKVCache) TakeSnapshots added in v0.30.8 

func (c *RotatingKVCache) TakeSnapshots() []Snapshot

func (*RotatingKVCache) Update 

func (c *RotatingKVCache) Update(b *batch.Batch, keys, values *mlx.Array) *nn.KVHistory

Assumes B = 1; heterogeneous batches are not supported.

func (*RotatingKVCache) View added in v0.23.1 

func (c *RotatingKVCache) View(b *batch.Batch) *nn.KVHistory

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.

type Snapshot added in v0.18.3 

type Snapshot interface {
	// 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()
}

Snapshot is paged-out cache state that can be restored later.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.