286 lines
8.2 KiB
Go
286 lines
8.2 KiB
Go
package cache
|
|
|
|
import (
|
|
"github.com/ollama/ollama/x/mlxrunner/batch"
|
|
"github.com/ollama/ollama/x/mlxrunner/mlx"
|
|
"github.com/ollama/ollama/x/models/nn"
|
|
)
|
|
|
|
// Recurrent is the contract for caches that back recurrent linear-attention layers.
|
|
type Recurrent interface {
|
|
Cache
|
|
Get(b *batch.Batch, dtype mlx.DType) *nn.RecurrentHistory
|
|
Put(b *batch.Batch, newConv, newDelta *mlx.Array)
|
|
}
|
|
|
|
// RecurrentRecorder records the per-token scan inputs needed to commit an
|
|
// accepted prefix after a speculative recurrent forward.
|
|
type RecurrentRecorder interface {
|
|
Record(qkv, q, k, v, gDecay, beta *mlx.Array)
|
|
}
|
|
|
|
// RecurrentCache stores state for linear-recurrent layers.
|
|
//
|
|
// Conv state shape: [B, convTail, convDim]
|
|
// Delta state shape: [B, numVHeads, headVDim, headKDim]
|
|
type RecurrentCache struct {
|
|
convState *mlx.Array
|
|
deltaState *mlx.Array
|
|
offset int
|
|
|
|
convTail int
|
|
convDim int
|
|
numVHeads int
|
|
headVDim int
|
|
headKDim int
|
|
}
|
|
|
|
func (c *RecurrentCache) setState(old, v *mlx.Array, contiguous bool) *mlx.Array {
|
|
if v == nil || !v.Valid() {
|
|
return old
|
|
}
|
|
|
|
if contiguous {
|
|
v = mlx.Contiguous(v, false)
|
|
}
|
|
v = v.Clone()
|
|
|
|
mlx.Pin(v)
|
|
mlx.Unpin(old)
|
|
|
|
return v
|
|
}
|
|
|
|
func NewRecurrentCache(convTail, convDim, numVHeads, headVDim, headKDim int32) *RecurrentCache {
|
|
return &RecurrentCache{
|
|
convTail: int(convTail),
|
|
convDim: int(convDim),
|
|
numVHeads: int(numVHeads),
|
|
headVDim: int(headVDim),
|
|
headKDim: int(headKDim),
|
|
}
|
|
}
|
|
|
|
func (c *RecurrentCache) ensure(batch int, dtype mlx.DType) {
|
|
if batch <= 0 {
|
|
batch = 1
|
|
}
|
|
|
|
// Keep the gated-delta recurrent state in float32 even when activations are
|
|
// bf16/fp16. The convolution tail stays in the activation dtype.
|
|
deltaDType := mlx.DTypeFloat32
|
|
needConv := c.convState == nil || !c.convState.Valid() || c.convState.DType() != dtype ||
|
|
c.convState.Dim(0) != batch || c.convState.Dim(1) != c.convTail || c.convState.Dim(2) != c.convDim
|
|
needDelta := c.deltaState == nil || !c.deltaState.Valid() || c.deltaState.DType() != deltaDType ||
|
|
c.deltaState.Dim(0) != batch || c.deltaState.Dim(1) != c.numVHeads || c.deltaState.Dim(2) != c.headVDim || c.deltaState.Dim(3) != c.headKDim
|
|
if !needConv && !needDelta {
|
|
return
|
|
}
|
|
|
|
if needConv {
|
|
c.convState = c.setState(c.convState, mlx.Zeros(dtype, batch, c.convTail, c.convDim), false)
|
|
}
|
|
if needDelta {
|
|
c.deltaState = c.setState(c.deltaState, mlx.Zeros(deltaDType, batch, c.numVHeads, c.headVDim, c.headKDim), false)
|
|
}
|
|
}
|
|
|
|
// Get returns the current conv/delta state for the SSM layer's read
|
|
// phase. Lazy-initializes zero-filled state tensors using b.InputIDs
|
|
// for the batch size; reallocates if the existing state's batch size
|
|
// or dtype no longer matches.
|
|
func (c *RecurrentCache) Get(b *batch.Batch, dtype mlx.DType) *nn.RecurrentHistory {
|
|
c.ensure(b.InputIDs.Dim(0), dtype)
|
|
return nn.NewRecurrentHistory(c.convState, c.deltaState)
|
|
}
|
|
|
|
// Put stores the post-computation conv/delta states for the SSM
|
|
// layer's write phase and advances the cache offset by the current
|
|
// forward's real token count.
|
|
//
|
|
// Assumes B = 1; heterogeneous batches are not supported.
|
|
func (c *RecurrentCache) Put(b *batch.Batch, newConv, newDelta *mlx.Array) {
|
|
c.convState = c.setState(c.convState, newConv, true)
|
|
c.deltaState = c.setState(c.deltaState, newDelta, false)
|
|
c.offset += int(b.SeqQueryLens[0])
|
|
}
|
|
|
|
func (c *RecurrentCache) State() []*mlx.Array {
|
|
return []*mlx.Array{c.convState, c.deltaState}
|
|
}
|
|
|
|
// recurrentSnapshot holds paged-out recurrent state. Self-contained —
|
|
// does not depend on any parent state.
|
|
type recurrentSnapshot struct {
|
|
convState, deltaState *mlx.Array
|
|
offset int
|
|
}
|
|
|
|
func (s *recurrentSnapshot) Size() int { return s.convState.NumBytes() + s.deltaState.NumBytes() }
|
|
func (s *recurrentSnapshot) Close() { mlx.Unpin(s.convState, s.deltaState) }
|
|
|
|
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
|
|
}
|
|
|
|
func (c *RecurrentCache) Restore(snapshot Snapshot, target int) bool {
|
|
if snapshot == nil {
|
|
// Recurrent state is cumulative and can't rewind. Only succeed
|
|
// if we're already at the target (no-op).
|
|
return target == c.offset
|
|
}
|
|
|
|
snap := snapshot.(*recurrentSnapshot)
|
|
|
|
// Recurrent snapshots encode cumulative state up to exactly
|
|
// snap.offset. Target must match — rewinding would leave stale
|
|
// state, and advancing isn't possible without feeding tokens.
|
|
if target != snap.offset {
|
|
return false
|
|
}
|
|
|
|
c.convState = c.setState(c.convState, snap.convState, false)
|
|
c.deltaState = c.setState(c.deltaState, snap.deltaState, false)
|
|
c.offset = snap.offset
|
|
|
|
return true
|
|
}
|
|
|
|
func (c *RecurrentCache) Merge(parent, child Snapshot) Snapshot {
|
|
// Recurrent snapshots are self-contained — child supersedes parent.
|
|
if parent != nil {
|
|
parent.Close()
|
|
}
|
|
return child
|
|
}
|
|
|
|
func (c *RecurrentCache) Split(snapshot Snapshot, at int) (Snapshot, Snapshot) {
|
|
// Recurrent state is cumulative and not position-sliceable.
|
|
// Cannot recover intermediate state at the split point.
|
|
return nil, snapshot
|
|
}
|
|
|
|
func (c *RecurrentCache) Free() {
|
|
mlx.Unpin(c.convState, c.deltaState)
|
|
c.convState, c.deltaState = nil, nil
|
|
c.offset = 0
|
|
}
|
|
|
|
func (c *RecurrentCache) Offset() int { return c.offset }
|
|
|
|
type speculativeRecurrentCache struct {
|
|
speculativeBase
|
|
target *RecurrentCache
|
|
|
|
start int
|
|
|
|
initialConv *mlx.Array
|
|
initialDelta *mlx.Array
|
|
|
|
qkv, q, k, v, gDecay, beta *mlx.Array
|
|
fullConv, fullDelta *mlx.Array
|
|
length int
|
|
}
|
|
|
|
func newSpeculativeRecurrentCache(target *RecurrentCache) *speculativeRecurrentCache {
|
|
return &speculativeRecurrentCache{
|
|
speculativeBase: speculativeBase{offset: target.Offset()},
|
|
target: target,
|
|
start: target.Offset(),
|
|
}
|
|
}
|
|
|
|
func (c *speculativeRecurrentCache) Get(b *batch.Batch, dtype mlx.DType) *nn.RecurrentHistory {
|
|
if c.fullConv != nil && c.fullDelta != nil {
|
|
return nn.NewRecurrentHistory(c.fullConv, c.fullDelta)
|
|
}
|
|
|
|
history := c.target.Get(b, dtype)
|
|
if c.initialConv == nil {
|
|
c.initialConv = history.ConvState()
|
|
}
|
|
if c.initialDelta == nil {
|
|
c.initialDelta = history.DeltaState()
|
|
}
|
|
return history
|
|
}
|
|
|
|
func (c *speculativeRecurrentCache) Record(qkv, q, k, v, gDecay, beta *mlx.Array) {
|
|
c.qkv, c.q, c.k, c.v, c.gDecay, c.beta = qkv, q, k, v, gDecay, beta
|
|
if qkv != nil {
|
|
c.length = qkv.Dim(1)
|
|
}
|
|
}
|
|
|
|
func (c *speculativeRecurrentCache) Put(b *batch.Batch, newConv, newDelta *mlx.Array) {
|
|
c.fullConv, c.fullDelta = newConv, newDelta
|
|
c.offset += int(b.SeqQueryLens[0])
|
|
}
|
|
|
|
func (c *speculativeRecurrentCache) State() []*mlx.Array {
|
|
if c.fullConv != nil && c.fullDelta != nil {
|
|
return []*mlx.Array{c.fullConv, c.fullDelta}
|
|
}
|
|
return c.target.State()
|
|
}
|
|
|
|
func (c *speculativeRecurrentCache) commit(n int) {
|
|
if n <= 0 {
|
|
return
|
|
}
|
|
if c.length > 0 && n > c.length {
|
|
n = c.length
|
|
}
|
|
|
|
if c.length > 0 && n == c.length && c.fullConv != nil && c.fullDelta != nil {
|
|
c.target.convState = c.target.setState(c.target.convState, c.fullConv, true)
|
|
c.target.deltaState = c.target.setState(c.target.deltaState, c.fullDelta, false)
|
|
c.target.offset = c.start + n
|
|
return
|
|
}
|
|
|
|
if c.initialConv == nil || c.initialDelta == nil || c.qkv == nil || c.q == nil || c.k == nil || c.v == nil || c.gDecay == nil || c.beta == nil {
|
|
return
|
|
}
|
|
|
|
qkv := sliceSeq(c.qkv, n)
|
|
convConcat := mlx.Concatenate([]*mlx.Array{c.initialConv, qkv}, 1)
|
|
total := convConcat.Dim(1)
|
|
nextConv := convConcat.Slice(mlx.Slice(), mlx.Slice(total-c.target.convTail, total), mlx.Slice())
|
|
|
|
_, delta := mlx.FastGatedDelta(
|
|
sliceSeq(c.q, n),
|
|
sliceSeq(c.k, n),
|
|
sliceSeq(c.v, n),
|
|
sliceSeq(c.gDecay, n),
|
|
sliceSeq(c.beta, n),
|
|
c.initialDelta,
|
|
nil,
|
|
)
|
|
|
|
c.target.convState = c.target.setState(c.target.convState, nextConv, true)
|
|
c.target.deltaState = c.target.setState(c.target.deltaState, delta, false)
|
|
c.target.offset = c.start + n
|
|
}
|
|
|
|
func sliceSeq(a *mlx.Array, n int) *mlx.Array {
|
|
switch a.NumDims() {
|
|
case 3:
|
|
return a.Slice(mlx.Slice(), mlx.Slice(0, n), mlx.Slice())
|
|
case 4:
|
|
return a.Slice(mlx.Slice(), mlx.Slice(0, n), mlx.Slice(), mlx.Slice())
|
|
default:
|
|
panic("recurrent speculative sequence tensor must be rank 3 or 4")
|
|
}
|
|
}
|