Mirror_Repos/ollama_ollama
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

mlxrunner: key the cache trie by token pairs for draft caches

Browse Source 
A draft cache pairs each slot with the token that follows it, so the
deepest stored pair always names one token past what a prefix match
can verify - at generation end, the sampled-but-never-committed
final token. Restoring at the match point reuses that pair blind: a
stop token stripped from the next prompt, or any divergence at the
boundary, leaves it stale, and pairing never rewrites below the
resume position, quietly lowering draft acceptance.

Key the trie by token pairs instead: the key for offset i packs
(token i, token i+1), so matching k keys verifies k+1 tokens and
every match is a valid restore point. A pair is reused only if the
token it names matched, and prefill re-evaluates the boundary token,
rebuilding its pair with the token that actually follows. A token
gets a key only once its successor is recorded, so endings record the
final sampled token - never forwarded - and the trie stays level with
the caches. Without a look-ahead the keys are the tokens and behavior
is unchanged.

The recorded tokens' slice bounds used to reject state past them for
free; close now checks the invariant against the stored keys
directly. The test harness rests requests the way the pipeline does -
the deepest recorded token never enters the caches.
This commit is contained in:
Jesse Gross
2026-07-06 13:21:36 -07:00
parent c963822dca
commit 556245843a
8 changed files with 355 additions and 223 deletions
Download Patch File Download Diff File Expand all files Collapse all files
x/mlxrunner/cache_trie.go
+9 -6
View File
@@ -8,12 +8,15 @@ import (
"github.com/ollama/ollama/x/mlxrunner/cache"
)
// trieKey encodes a token for trie matching (see prefixCache.key).
type trieKey int64
// trieNode represents a node in the compressed prefix trie for cache branching.
// Each node stores a compressed edge (multiple tokens) and optional paged-out
// snapshot data per cache layer.
type trieNode struct {
tokens []int32 // compressed edge — multiple tokens per node
endOffset int // cumulative tokens from root to end of this node
tokens []trieKey // compressed edge — multiple tokens per node
endOffset int // cumulative tokens from root to end of this node
parent *trieNode
children []*trieNode
lastUsed time.Time // for LRU eviction
@@ -88,7 +91,7 @@ func (n *trieNode) hasAllSnapshots() bool {
// findBestMatch walks the trie matching input tokens, returning the path of
// nodes traversed and the total number of tokens matched.
func findBestMatch(root *trieNode, tokens []int32) (path []*trieNode, matched int) {
func findBestMatch(root *trieNode, tokens []trieKey) (path []*trieNode, matched int) {
if root == nil {
return nil, 0
}
@@ -145,10 +148,10 @@ func findBestMatch(root *trieNode, tokens []int32) (path []*trieNode, matched in
// appendTokens either creates a new child node or extends the leaf in place,
// returning the node that now holds the tokens.
func (n *trieNode) appendTokens(root *trieNode, tokens []int32, endOffset int) *trieNode {
func (n *trieNode) appendTokens(root *trieNode, tokens []trieKey, endOffset int) *trieNode {
if n == root || len(n.children) > 0 || n.hasSnapshots() {
child := &trieNode{
tokens: make([]int32, len(tokens)),
tokens: make([]trieKey, len(tokens)),
endOffset: endOffset,
parent: n,
lastUsed: n.lastUsed,
@@ -193,7 +196,7 @@ func splitNode(node *trieNode, at int, caches []cache.Cache, counter *int64) *tr
// Create new parent with the prefix of the edge.
newParent := &trieNode{
tokens: make([]int32, at),
tokens: make([]trieKey, at),
endOffset: node.startOffset() + at,
parent: node.parent,
children: []*trieNode{node},
x/mlxrunner/cache_trie_test.go
+48 -48
View File
@@ -8,7 +8,7 @@ import (
"github.com/ollama/ollama/x/mlxrunner/cache"
)
func newTestTrie(tokens []int32) *trieNode {
func newTestTrie(tokens []trieKey) *trieNode {
root := &trieNode{lastUsed: time.Now()}
if len(tokens) > 0 {
child := &trieNode{
@@ -26,13 +26,13 @@ func TestFindBestMatchMultipleBranches(t *testing.T) {
root := &trieNode{lastUsed: time.Now()}
branch1 := &trieNode{
tokens: []int32{1, 2, 3},
tokens: []trieKey{1, 2, 3},
endOffset: 3,
parent: root,
lastUsed: time.Now(),
}
branch2 := &trieNode{
tokens: []int32{4, 5, 6},
tokens: []trieKey{4, 5, 6},
endOffset: 3,
parent: root,
lastUsed: time.Now(),
@@ -40,7 +40,7 @@ func TestFindBestMatchMultipleBranches(t *testing.T) {
root.children = []*trieNode{branch1, branch2}
// Match branch 1.
path, matched := findBestMatch(root, []int32{1, 2, 3, 7})
path, matched := findBestMatch(root, []trieKey{1, 2, 3, 7})
if matched != 3 {
t.Fatalf("expected 3 matched, got %d", matched)
}
@@ -49,7 +49,7 @@ func TestFindBestMatchMultipleBranches(t *testing.T) {
}
// Match branch 2.
path, matched = findBestMatch(root, []int32{4, 5, 6, 8})
path, matched = findBestMatch(root, []trieKey{4, 5, 6, 8})
if matched != 3 {
t.Fatalf("expected 3 matched, got %d", matched)
}
@@ -58,7 +58,7 @@ func TestFindBestMatchMultipleBranches(t *testing.T) {
}
// Match neither.
_, matched = findBestMatch(root, []int32{7, 8, 9})
_, matched = findBestMatch(root, []trieKey{7, 8, 9})
if matched != 0 {
t.Fatalf("expected 0 matched, got %d", matched)
}
@@ -68,7 +68,7 @@ func TestFindBestMatchPrefersFullEdge(t *testing.T) {
root := &trieNode{lastUsed: time.Now()}
shared := &trieNode{
tokens: []int32{1, 2, 3},
tokens: []trieKey{1, 2, 3},
endOffset: 3,
parent: root,
lastUsed: time.Now(),
@@ -76,13 +76,13 @@ func TestFindBestMatchPrefersFullEdge(t *testing.T) {
root.children = []*trieNode{shared}
longer := &trieNode{
tokens: []int32{10, 11, 12, 13, 14},
tokens: []trieKey{10, 11, 12, 13, 14},
endOffset: 8,
parent: shared,
lastUsed: time.Now(),
}
shorter := &trieNode{
tokens: []int32{10, 11, 12},
tokens: []trieKey{10, 11, 12},
endOffset: 6,
parent: shared,
lastUsed: time.Now(),
@@ -90,7 +90,7 @@ func TestFindBestMatchPrefersFullEdge(t *testing.T) {
// Put longer first so naive first-match would pick it.
shared.children = []*trieNode{longer, shorter}
input := []int32{1, 2, 3, 10, 11, 12, 99, 100}
input := []trieKey{1, 2, 3, 10, 11, 12, 99, 100}
path, matched := findBestMatch(root, input)
if matched != 6 {
@@ -108,20 +108,20 @@ func TestFindBestMatchPrefersLongerPartial(t *testing.T) {
root := &trieNode{lastUsed: time.Now()}
child1 := &trieNode{
tokens: []int32{1, 2, 3, 4, 5},
tokens: []trieKey{1, 2, 3, 4, 5},
endOffset: 5,
parent: root,
lastUsed: time.Now(),
}
child2 := &trieNode{
tokens: []int32{1, 2, 9},
tokens: []trieKey{1, 2, 9},
endOffset: 3,
parent: root,
lastUsed: time.Now(),
}
root.children = []*trieNode{child2, child1}
input := []int32{1, 2, 3, 7, 8}
input := []trieKey{1, 2, 3, 7, 8}
path, matched := findBestMatch(root, input)
if matched != 3 {
@@ -133,7 +133,7 @@ func TestFindBestMatchPrefersLongerPartial(t *testing.T) {
}
func TestSplitNodeWithSnapshots(t *testing.T) {
root := newTestTrie([]int32{1, 2, 3, 4, 5})
root := newTestTrie([]trieKey{1, 2, 3, 4, 5})
child := root.children[0]
rc := &fakeRewindableCache{tracker: &snapshotTracker{}, tokens: []int32{1, 2, 3, 4, 5}}
@@ -159,9 +159,9 @@ func TestSplitNodeWithSnapshots(t *testing.T) {
}
func TestFindSplitAppendSequence(t *testing.T) {
root := newTestTrie([]int32{1, 2, 3, 4, 5})
root := newTestTrie([]trieKey{1, 2, 3, 4, 5})
path, matched := findBestMatch(root, []int32{1, 2, 3, 6, 7})
path, matched := findBestMatch(root, []trieKey{1, 2, 3, 6, 7})
if matched != 3 {
t.Fatalf("expected 3 matched, got %d", matched)
}
@@ -170,28 +170,28 @@ func TestFindSplitAppendSequence(t *testing.T) {
matchedInEdge := matched - lastNode.startOffset()
split := splitNode(lastNode, matchedInEdge, nil, nil)
split.appendTokens(root, []int32{6, 7}, 5)
split.appendTokens(root, []trieKey{6, 7}, 5)
if len(root.children) != 1 {
t.Fatalf("root should have 1 child, got %d", len(root.children))
}
shared := root.children[0]
if !slices.Equal(shared.tokens, []int32{1, 2, 3}) {
if !slices.Equal(shared.tokens, []trieKey{1, 2, 3}) {
t.Fatalf("shared tokens = %v, want [1,2,3]", shared.tokens)
}
if len(shared.children) != 2 {
t.Fatalf("shared should have 2 children, got %d", len(shared.children))
}
_, m1 := findBestMatch(root, []int32{1, 2, 3, 4, 5})
_, m1 := findBestMatch(root, []trieKey{1, 2, 3, 4, 5})
if m1 != 5 {
t.Fatalf("original branch: expected 5 matched, got %d", m1)
}
_, m2 := findBestMatch(root, []int32{1, 2, 3, 6, 7})
_, m2 := findBestMatch(root, []trieKey{1, 2, 3, 6, 7})
if m2 != 5 {
t.Fatalf("new branch: expected 5 matched, got %d", m2)
}
_, m3 := findBestMatch(root, []int32{1, 2, 3, 9, 9})
_, m3 := findBestMatch(root, []trieKey{1, 2, 3, 9, 9})
if m3 != 3 {
t.Fatalf("unrelated input: expected 3 matched, got %d", m3)
}
@@ -200,32 +200,32 @@ func TestFindSplitAppendSequence(t *testing.T) {
func TestRepeatedBranching(t *testing.T) {
root := &trieNode{lastUsed: time.Now()}
root.appendTokens(root, []int32{1, 2, 3, 4, 5}, 5)
root.appendTokens(root, []trieKey{1, 2, 3, 4, 5}, 5)
_, matchedB := findBestMatch(root, []int32{1, 2, 3, 6, 7})
_, matchedB := findBestMatch(root, []trieKey{1, 2, 3, 6, 7})
if matchedB != 3 {
t.Fatalf("B: expected 3 matched, got %d", matchedB)
}
nodeA := root.children[0]
split1 := splitNode(nodeA, 3, nil, nil)
split1.appendTokens(root, []int32{6, 7}, 5)
split1.appendTokens(root, []trieKey{6, 7}, 5)
_, matchedC := findBestMatch(root, []int32{1, 2, 8, 9})
_, matchedC := findBestMatch(root, []trieKey{1, 2, 8, 9})
if matchedC != 2 {
t.Fatalf("C: expected 2 matched, got %d", matchedC)
}
split2 := splitNode(split1, 2, nil, nil)
split2.appendTokens(root, []int32{8, 9}, 4)
split2.appendTokens(root, []trieKey{8, 9}, 4)
_, mA := findBestMatch(root, []int32{1, 2, 3, 4, 5})
_, mA := findBestMatch(root, []trieKey{1, 2, 3, 4, 5})
if mA != 5 {
t.Fatalf("A: expected 5 matched, got %d", mA)
}
_, mB := findBestMatch(root, []int32{1, 2, 3, 6, 7})
_, mB := findBestMatch(root, []trieKey{1, 2, 3, 6, 7})
if mB != 5 {
t.Fatalf("B: expected 5 matched, got %d", mB)
}
_, mC := findBestMatch(root, []int32{1, 2, 8, 9})
_, mC := findBestMatch(root, []trieKey{1, 2, 8, 9})
if mC != 4 {
t.Fatalf("C: expected 4 matched, got %d", mC)
}
@@ -239,21 +239,21 @@ func TestMergeWithChild(t *testing.T) {
now := time.Now()
root := &trieNode{lastUsed: now}
a := &trieNode{
tokens: []int32{1, 2, 3},
tokens: []trieKey{1, 2, 3},
endOffset: 3,
parent: root,
lastUsed: now,
snapshots: []cache.Snapshot{&fakeSnapshot{tokens: []int32{1, 2, 3}, from: 0, to: 3}},
}
b := &trieNode{
tokens: []int32{4, 5},
tokens: []trieKey{4, 5},
endOffset: 5,
parent: a,
lastUsed: now,
snapshots: []cache.Snapshot{&fakeSnapshot{tokens: []int32{4, 5}, from: 3, to: 5}},
}
c := &trieNode{tokens: []int32{6}, endOffset: 6, parent: b, lastUsed: now}
d := &trieNode{tokens: []int32{7}, endOffset: 6, parent: b, lastUsed: now}
c := &trieNode{tokens: []trieKey{6}, endOffset: 6, parent: b, lastUsed: now}
d := &trieNode{tokens: []trieKey{7}, endOffset: 6, parent: b, lastUsed: now}
root.children = []*trieNode{a}
a.children = []*trieNode{b}
b.children = []*trieNode{c, d}
@@ -262,7 +262,7 @@ func TestMergeWithChild(t *testing.T) {
mergeWithChild(a, []cache.Cache{mc}, nil)
// Tokens concatenated.
if !slices.Equal(a.tokens, []int32{1, 2, 3, 4, 5}) {
if !slices.Equal(a.tokens, []trieKey{1, 2, 3, 4, 5}) {
t.Fatalf("merged tokens = %v, want [1,2,3,4,5]", a.tokens)
}
if a.endOffset != 5 {
@@ -294,11 +294,11 @@ func TestMergeWithChild(t *testing.T) {
t.Run("UserFlag", func(t *testing.T) {
root := &trieNode{lastUsed: time.Now()}
parent := &trieNode{
tokens: []int32{1, 2}, endOffset: 2, parent: root,
tokens: []trieKey{1, 2}, endOffset: 2, parent: root,
lastUsed: time.Now(), user: false,
}
child := &trieNode{
tokens: []int32{3, 4}, endOffset: 4, parent: parent,
tokens: []trieKey{3, 4}, endOffset: 4, parent: parent,
lastUsed: time.Now(), user: true,
}
root.children = []*trieNode{parent}
@@ -315,11 +315,11 @@ func TestMergeWithChild(t *testing.T) {
now := time.Now()
root := &trieNode{lastUsed: now}
parent := &trieNode{
tokens: []int32{1}, endOffset: 1, parent: root,
tokens: []trieKey{1}, endOffset: 1, parent: root,
lastUsed: now.Add(-1 * time.Hour),
}
child := &trieNode{
tokens: []int32{2}, endOffset: 2, parent: parent,
tokens: []trieKey{2}, endOffset: 2, parent: parent,
lastUsed: now.Add(1 * time.Hour),
}
root.children = []*trieNode{parent}
@@ -340,10 +340,10 @@ func TestMergeWithChild(t *testing.T) {
}()
root := &trieNode{lastUsed: time.Now()}
node := &trieNode{
tokens: []int32{1}, endOffset: 1, parent: root, lastUsed: time.Now(),
tokens: []trieKey{1}, endOffset: 1, parent: root, lastUsed: time.Now(),
children: []*trieNode{
{tokens: []int32{2}, endOffset: 2, lastUsed: time.Now()},
{tokens: []int32{3}, endOffset: 2, lastUsed: time.Now()},
{tokens: []trieKey{2}, endOffset: 2, lastUsed: time.Now()},
{tokens: []trieKey{3}, endOffset: 2, lastUsed: time.Now()},
},
}
root.children = []*trieNode{node}
@@ -354,7 +354,7 @@ func TestMergeWithChild(t *testing.T) {
func TestSplitMergeRoundTrip(t *testing.T) {
root := &trieNode{lastUsed: time.Now()}
leaf := &trieNode{
tokens: []int32{1, 2, 3, 4, 5},
tokens: []trieKey{1, 2, 3, 4, 5},
endOffset: 5,
parent: root,
lastUsed: time.Now(),
@@ -367,17 +367,17 @@ func TestSplitMergeRoundTrip(t *testing.T) {
// Split at 3: [1,2,3] -> [4,5]
newParent := splitNode(leaf, 3, caches, nil)
if !slices.Equal(newParent.tokens, []int32{1, 2, 3}) {
if !slices.Equal(newParent.tokens, []trieKey{1, 2, 3}) {
t.Fatalf("after split: parent tokens = %v, want [1,2,3]", newParent.tokens)
}
if !slices.Equal(leaf.tokens, []int32{4, 5}) {
if !slices.Equal(leaf.tokens, []trieKey{4, 5}) {
t.Fatalf("after split: child tokens = %v, want [4,5]", leaf.tokens)
}
checkTrieInvariants(t, root)
// Merge back: should restore [1,2,3,4,5]
mergeWithChild(newParent, caches, nil)
if !slices.Equal(newParent.tokens, []int32{1, 2, 3, 4, 5}) {
if !slices.Equal(newParent.tokens, []trieKey{1, 2, 3, 4, 5}) {
t.Fatalf("after merge: tokens = %v, want [1,2,3,4,5]", newParent.tokens)
}
if newParent.endOffset != 5 {
@@ -402,14 +402,14 @@ func TestRemoveNode(t *testing.T) {
t.Run("Leaf", func(t *testing.T) {
root := &trieNode{lastUsed: time.Now()}
shared := &trieNode{
tokens: []int32{1, 2, 3}, endOffset: 3, parent: root, lastUsed: time.Now(),
tokens: []trieKey{1, 2, 3}, endOffset: 3, parent: root, lastUsed: time.Now(),
}
leafA := &trieNode{
tokens: []int32{4, 5}, endOffset: 5, parent: shared, lastUsed: time.Now(),
tokens: []trieKey{4, 5}, endOffset: 5, parent: shared, lastUsed: time.Now(),
snapshots: []cache.Snapshot{&fakeSnapshot{from: 3, to: 5}},
}
leafB := &trieNode{
tokens: []int32{6, 7}, endOffset: 5, parent: shared, lastUsed: time.Now(),
tokens: []trieKey{6, 7}, endOffset: 5, parent: shared, lastUsed: time.Now(),
snapshots: []cache.Snapshot{&fakeSnapshot{from: 3, to: 5}},
}
root.children = []*trieNode{shared}
x/mlxrunner/mtp.go
+12 -20
View File
@@ -14,12 +14,18 @@ import (
const mtpPendingFlushTokens = 32
// mtpDrafter drafts with a model's multi-token-prediction head. Constructed
// at load, it opens each request's drafting session.
// at load, it fixes the trie keys' draft look-ahead for the model's lifetime
// and opens each request's drafting session.
type mtpDrafter struct {
spec *speculation
}
func newMTPDrafter(s *speculation) *mtpDrafter {
if len(s.draftKV) > 0 {
// The pairing references one token past each slot; trie keys carry
// that look-ahead so a match verifies it (see prefixCache.draftLookahead).
s.r.cache.draftLookahead = 1
}
return &mtpDrafter{spec: s}
}
@@ -94,29 +100,15 @@ func (d *mtpDraftSession) committed(tokens, hiddens *mlx.Array, position int) {
d.setFrontierHidden(lastHiddenRow(hiddens))
}
// finish settles the drafter when generation ends: current completes the
// frontier pair, leveling the draft caches with the target's resting offset.
//
// TODO: leveling the draft to the target writes a boundary entry whose
// look-ahead token is outside the stored prefix (here, the never-committed
// current). When a later request restores this prefix and diverges at the
// boundary, that entry is stale and lowers draft acceptance. EAGLE keeps the
// draft one slot behind the target so the unconfirmed boundary entry is never
// written (its bigram partner token[S+1] does not exist yet); we should do the
// same rather than level here. Regenerating hidden[S] to re-pair the boundary
// on the next request is a separate, re-prefill-bound concern for recurrent
// targets.
func (d *mtpDraftSession) finish(current *mlx.Array) {
// settle completes any open frontier pair with next — the token after the
// last committed slot — and flushes, leveling the draft caches with the
// target.
func (d *mtpDraftSession) settle(next *mlx.Array) {
if len(d.drafter.spec.draftKV) == 0 {
return
}
d.settle(current)
}
// settle completes any open frontier pair with current, then flushes.
func (d *mtpDraftSession) settle(current *mlx.Array) {
if d.frontierHidden != nil && d.frontier-1 == d.committedDraftOffset+d.pendingCount {
d.queueCacheWrites(current.ExpandDims(-1), d.frontierHidden)
d.queueCacheWrites(next.ExpandDims(-1), d.frontierHidden)
}
d.flush()
}
x/mlxrunner/mtp_test.go
+77 -23
View File
@@ -456,7 +456,7 @@ func TestRunMTPDecodeSampled(t *testing.T) {
t.Fatalf("open rejected a sampled request")
}
pinDraftLimit(spec, 4)
d := spec.decoder([]int32{1}, position)
d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
if err := r.decode(context.Background(), req, session, d, 0); err != nil {
t.Fatalf("decode: %v", err)
}
@@ -501,7 +501,7 @@ func TestRunMTPDecodeWarmDrafter(t *testing.T) {
// hidden row, leaving the drafter ready to propose from slot 1.
spec.committed(mlx.FromValues([]int32{0}, 1, 1), oneHotLogits([]int32{1}), 0)
d := spec.decoder([]int32{1}, position)
d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
if err := r.decode(context.Background(), req, session, d, 0); err != nil {
t.Fatalf("decode: %v", err)
}
@@ -565,7 +565,7 @@ func TestRunMTPDecodeEOSCutLeavesPositionsUnjudged(t *testing.T) {
spec.limit = 4
spec.committed(mlx.FromValues([]int32{0}, 1, 1), oneHotLogits([]int32{1}), 0)
d := spec.decoder([]int32{1}, position)
d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
if err := r.decode(context.Background(), req, session, d, 0); err != nil {
t.Fatalf("decode: %v", err)
}
@@ -625,14 +625,14 @@ func TestDecodePlain(t *testing.T) {
if final.DoneReason != 0 || final.EvalCount != 3 {
t.Fatalf("final DoneReason = %d, EvalCount = %d, want 0 (EOS) and 3", final.DoneReason, final.EvalCount)
}
if got := []int32{2, 3, 4, eos}; !slices.Equal(session.outputs, got) {
if got := []int32{2, 3, 4, eos, 0}; !slices.Equal(session.outputs, got) {
t.Fatalf("session outputs = %v, want %v", session.outputs, got)
}
// One forward per token: the seed at offset 1, then each sampled token
// in turn — including the ending EOS, whose forward is already in
// flight when the token is checked. Every forwarded token is recorded,
// so the caches rest exactly at the recorded path.
// and the EOS's sampled successor is recorded without being forwarded.
wantForwards := []forwardCall{{1, 1}, {2, 1}, {3, 1}, {4, 1}, {5, 1}}
model := r.Model.(*fakeMTPModel)
if !slices.Equal(model.forwards, wantForwards) {
@@ -710,7 +710,7 @@ func testDecoder(r *Runner, req Request, caches []cache.Cache, seed []int32, pos
if spec.enabled {
pinDraftLimit(spec, 4)
}
return spec.decoder(seed, position)
return spec.decoder(mlx.FromValues(seed, len(seed)), position)
}
return r.pipelinedDecoder(nil, caches, mlx.FromValues(seed, 1, len(seed)), position)
}
@@ -746,7 +746,7 @@ func TestDecodeKVDraft(t *testing.T) {
}
pinDraftLimit(spec, 4)
defer spec.close()
d := spec.decoder([]int32{1}, position)
d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
if err := r.decode(context.Background(), req, session, d, 0); err != nil {
t.Fatalf("decode: %v", err)
}
@@ -827,7 +827,7 @@ func TestDecodeKVDraftRejectionRebuildsFromTarget(t *testing.T) {
spec := r.spec.open(req, caches)
pinDraftLimit(spec, 4)
defer spec.close()
d := spec.decoder([]int32{1}, position)
d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
if err := r.decode(context.Background(), req, session, d, 0); err != nil {
t.Fatalf("decode: %v", err)
}
@@ -897,7 +897,7 @@ func TestDecodeMaintainsDraftCacheWithoutDrafting(t *testing.T) {
if spec == nil || spec.enabled {
t.Fatalf("want a permanent-park speculationSession, got %+v", spec)
}
d := spec.decoder([]int32{1}, position)
d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
if err := r.decode(context.Background(), req, session, d, 0); err != nil {
t.Fatalf("decode: %v", err)
}
@@ -933,13 +933,12 @@ func TestDecodeMaintainsDraftCacheWithoutDrafting(t *testing.T) {
}
}
func TestFlushLevelsDraftCacheWithPrefill(t *testing.T) {
func TestSettleLevelsDraftCacheWithPrefill(t *testing.T) {
skipIfNoMLX(t)
// Prefill attaches its scheduled snapshots only at offsets every cache
// has crossed, so the pipeline flushes the drafter's buffered pairs
// first: after a prefill-sized committed report and a flush, the
// draft cache covers every completed pair, one slot behind the target
// (the frontier pair still awaits its look-ahead token).
// has crossed, so the pipeline settles the drafter with the seed first:
// the completed frontier pair brings the draft cache level with the
// target, in one batched extend with the buffered pairs.
const eos int32 = 7
predict := map[int32]int32{2: 3, 3: 4, 4: 5}
r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
@@ -954,15 +953,16 @@ func TestFlushLevelsDraftCacheWithPrefill(t *testing.T) {
spec := r.spec.open(req, caches)
defer spec.close()
// The prompt's only chunk: tokens 1..4 at slots 0..3 with their hiddens.
// The prompt's only chunk: tokens 1..4 at slots 0..3 with their hiddens;
// token 5 is the seed.
spec.committed(mlx.FromValues([]int32{1, 2, 3, 4}, 1, 4), oneHotLogits([]int32{1, 2, 3, 4}), 0)
spec.flush()
spec.settle(mlx.FromValues([]int32{5}, 1))
if got := caches[1].Offset(); got != 3 {
t.Fatalf("draft cache offset after flush = %d, want 3 (all completed pairs)", got)
if got := caches[1].Offset(); got != 4 {
t.Fatalf("draft cache offset after settle = %d, want 4 (level with target)", got)
}
wantExtends := []extendCall{
{offset: 0, ids: []int32{2, 3, 4}, hiddens: []int32{1, 2, 3}},
{offset: 0, ids: []int32{2, 3, 4, 5}, hiddens: []int32{1, 2, 3, 4}},
}
if !reflect.DeepEqual(draft.extends, wantExtends) {
t.Fatalf("draft extends = %+v, want %+v", draft.extends, wantExtends)
@@ -995,8 +995,8 @@ func TestCommittedRunBatchesPastFlushCap(t *testing.T) {
defer spec.close()
// One prefill-sized chunk: n tokens at slots 0..n-1 with their hiddens.
// The run crosses the flush cap, so the write happens inside committed.
spec.committed(mlx.FromValues(tokens, 1, n), oneHotLogits(tokens), 0)
spec.flush()
if got := len(draft.extends); got != 1 {
t.Fatalf("draft extends = %d calls, want 1 batched extend", got)
@@ -1012,6 +1012,61 @@ func TestCommittedRunBatchesPastFlushCap(t *testing.T) {
}
}
func TestRestoredPrefixRewritesBoundaryPair(t *testing.T) {
skipIfNoMLX(t)
// A finished generation levels the draft with the target, its boundary
// pair naming the never-committed EOS. The next request restores one
// token below the match (the draft look-ahead), so the re-evaluated
// boundary token's report rewrites that pair with the token that follows.
const eos int32 = 7
predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: eos, eos: 0}
r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
draft := &fakeKVDraft{predict: predict}
caches, _ := newMTPTestCaches(2)
r.cache.caches = caches
r.spec = newSpeculation(r, draft)
session, ch := newMTPTestSession(caches)
req := Request{
Responses: ch,
Tokens: []int32{1},
CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
SamplerOpts: sampler.Options{},
}
spec := r.spec.open(req, caches)
pinDraftLimit(spec, 4)
d := spec.decoder(mlx.FromValues([]int32{1}, 1), 0)
if err := r.decode(context.Background(), req, session, d, 0); err != nil {
t.Fatalf("decode: %v", err)
}
d.close()
spec.close()
// The leveled rest: the draft's last pair names the EOS.
if got, want := caches[1].(*fakeRewindableCache).tokens, []int32{2, 3, 4, 5, 6, eos}; !slices.Equal(got, want) {
t.Fatalf("draft cache = %v, want %v", got, want)
}
// The next request matches the six stored tokens and restores at 5, as
// begin does for the draft look-ahead; its prefill re-evaluates token 6,
// and the token that follows is 1, not the EOS.
for _, c := range caches {
if !c.Restore(nil, 5) {
t.Fatal("restore to 5 failed")
}
}
spec = r.spec.open(req, caches)
spec.committed(mlx.FromValues([]int32{6, 1}, 1, 2), oneHotLogits([]int32{eos, 2}), 5)
spec.close()
last := draft.extends[len(draft.extends)-1]
if want := (extendCall{offset: 5, ids: []int32{1}, hiddens: []int32{eos}}); !reflect.DeepEqual(last, want) {
t.Fatalf("boundary rewrite = %+v, want %+v", last, want)
}
if got, want := caches[1].(*fakeRewindableCache).tokens, []int32{2, 3, 4, 5, 6, 1}; !slices.Equal(got, want) {
t.Fatalf("draft cache = %v, want %v (stale EOS pair rewritten)", got, want)
}
}
func TestDecodeParkedDraftResume(t *testing.T) {
skipIfNoMLX(t)
// Leaving a parked stretch: the inner decoder's in-flight sample is
@@ -1035,7 +1090,7 @@ func TestDecodeParkedDraftResume(t *testing.T) {
t.Fatalf("want a drafting speculationSession, got %+v", spec)
}
pinDraftLimit(spec, 0)
d := spec.decoder([]int32{1}, 0).(*speculativeDecoder)
d := spec.decoder(mlx.FromValues([]int32{1}, 1), 0).(*speculativeDecoder)
// Two parked calls arrive pipelined, one token each.
for _, want := range []int{2, 3} {
@@ -1136,8 +1191,7 @@ type nopDrafter struct{}
func (nopDrafter) propose(*mlx.Array, int) *draftCandidates { return nil }
func (nopDrafter) committed(_, _ *mlx.Array, _ int) {}
func (nopDrafter) finish(*mlx.Array) {}
func (nopDrafter) flush() {}
func (nopDrafter) settle(*mlx.Array) {}
func (nopDrafter) close() {}
// scriptedCandidates builds draft candidates by running the real drafter
x/mlxrunner/pipeline.go
+33 -21
View File
@@ -92,7 +92,7 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er
if spec != nil {
d = spec.decoder(seed, position)
} else {
d = r.pipelinedDecoder(nil, caches, mlx.FromValues(seed, 1, len(seed)), position)
d = r.pipelinedDecoder(nil, caches, seed.ExpandDims(-1), position)
}
defer d.close()
return r.decode(ctx, request, session, d, promptEval)
@@ -100,8 +100,8 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er
// prefill evaluates the prompt in chunks, leaving one token for decode to
// seed from, and schedules the prompt's periodic snapshots. It returns the
// seed tokens, the resume position, and the prompt-evaluation duration.
func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *speculationSession) ([]int32, int, time.Duration, error) {
// seed token, the resume position, and the prompt-evaluation duration.
func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *speculationSession) (*mlx.Array, int, time.Duration, error) {
start := time.Now()
inputs := session.inputs
tokens := session.remaining
@@ -161,13 +161,14 @@ func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *specu
mlx.ClearCache()
}
// Flush before attaching: snapshots attach only at offsets every cache
// has crossed, and a drafter with draft caches keeps buffered pairs that
// would otherwise hold those caches short of the scheduled offsets.
spec.flush()
// Settle before attaching: snapshots attach only at offsets every cache
// has crossed, and the draft caches stay a pair short of the target
// until the seed completes the frontier pair.
seed := mlx.FromValues(tokens[processed:], 1)
spec.settle(seed)
session.attachPrefillSnapshots()
return tokens[processed:], position, time.Since(start), nil
return seed, position, time.Since(start), nil
}
// A decoder produces each run of tokens to emit, owning its own dispatch and
@@ -175,6 +176,12 @@ func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *specu
// cancellation. next may return none while its first tokens are in flight.
type decoder interface {
next(remaining int) ([]sampler.Result, error)
// drain ends production, returning any results sampled but never
// delivered through next and the position the next forward would have
// taken; the decoder remains closeable.
drain() ([]sampler.Result, int)
close()
}
@@ -183,6 +190,14 @@ type decoder interface {
// never rest ahead of session.outputs; tokens past the stop are recorded but
// not streamed or counted.
func (r *Runner) decode(ctx context.Context, request Request, session *cacheSession, d decoder, promptEval time.Duration) error {
// A sampled-but-undelivered result is still a produced token; record it.
defer func() {
results, _ := d.drain()
for _, res := range results {
session.outputs = append(session.outputs, int32(res.Token.Int()))
}
}()
detok := detokenizer{
tokenizer: r.Tokenizer,
wantLogprobs: request.SamplerOpts.Logprobs,
@@ -276,7 +291,6 @@ type pipelinedDecoder struct {
caches []cache.Cache
position int
sample sampler.Result // in flight: sampled, not yet forwarded
emitted sampler.Result // last call's result, pinned until the next call
}
func (r *Runner) pipelinedDecoder(spec *speculationSession, caches []cache.Cache, seed *mlx.Array, position int) *pipelinedDecoder {
@@ -305,25 +319,23 @@ func (t *pipelinedDecoder) dispatch(token *mlx.Array) sampler.Result {
}
func (t *pipelinedDecoder) next(int) ([]sampler.Result, error) {
mlx.Unpin(t.emitted.Arrays()...)
t.emitted, t.sample = t.sample, t.dispatch(t.sample.Token.ExpandDims(-1))
return []sampler.Result{t.emitted}, nil
out := t.sample
t.sample = t.dispatch(out.Token.ExpandDims(-1))
mlx.Unpin(out.Arrays()...)
return []sampler.Result{out}, nil
}
// detach ends a parked stretch: it hands the in-flight sample (sampled but
// never forwarded) and the resume position to the caller, releasing the
// decoder's emitted pin but not settling the drafter. The caller drafts from
// the sample next, and that round completes the still-open frontier pair.
func (t *pipelinedDecoder) detach() (sampler.Result, int) {
mlx.Unpin(t.emitted.Arrays()...)
return t.sample, t.position
// drain ends production: it returns the in-flight sample (sampled but never
// forwarded) and the position its forward would have taken. The decoder
// keeps the sample for close.
func (t *pipelinedDecoder) drain() ([]sampler.Result, int) {
return []sampler.Result{t.sample}, t.position
}
func (t *pipelinedDecoder) close() {
// The in-flight sample's forward was never dispatched; its report settles
// the drafter level with the caches' resting offset.
t.spec.finish(t.sample.Token)
mlx.Unpin(t.emitted.Arrays()...)
t.spec.settle(t.sample.Token)
mlx.Unpin(t.sample.Arrays()...)
}
x/mlxrunner/prefix_cache.go
+45 -10
View File
@@ -37,6 +37,10 @@ type prefixCache struct {
activePath []*trieNode // current root→leaf path with live MLX arrays
caches []cache.Cache
pagedOutBytes int64 // total bytes in paged-out snapshots across the trie
// draftLookahead is how far the draft caches' entries reference past
// their own slot; trie keys pack each token with its look-ahead (see key).
draftLookahead int
}
// pendingSnapshot is a snapshot scheduled to be taken during prefill.
@@ -89,13 +93,14 @@ func (c *prefixCache) ensureRoot() {
func (c *prefixCache) begin(inputs []int32) *cacheSession {
c.ensureRoot()
matchPath, matched := findBestMatch(c.root, inputs)
keys := c.key(inputs)
matchPath, matched := findBestMatch(c.root, keys)
originalMatched := matched
// Always keep at least one token to re-evaluate so the
// pipeline can seed token generation from it.
if matched == len(inputs) && matched > 0 {
matchPath, matched = findBestMatch(c.root, inputs[:len(inputs)-1])
matchPath, matched = findBestMatch(c.root, keys[:matched-1])
}
// Switch to the matched path, paging in/out as needed.
@@ -127,6 +132,28 @@ func (c *prefixCache) begin(inputs []int32) *cacheSession {
return session
}
// key converts tokens to trie keys, one per restorable cache offset. A model
// that drafts through MTP-style draft caches pairs each cache slot with the
// token after it, so slot i is reusable only if token i+1 also matched. The
// key for offset i then packs (token i, token i+1): matching k keys verifies
// k+1 tokens, making every match a valid restore point.
func (c *prefixCache) key(tokens []int32) []trieKey {
keys := make([]trieKey, max(len(tokens)-c.draftLookahead, 0))
switch c.draftLookahead {
case 0:
for i, t := range tokens {
keys[i] = trieKey(t)
}
case 1:
for i := range keys {
keys[i] = trieKey(uint32(tokens[i]))<<32 | trieKey(uint32(tokens[i+1]))
}
default:
panic(fmt.Sprintf("prefixCache: unsupported draft look-ahead %d", c.draftLookahead))
}
return keys
}
// switchToPath transitions from the current active path to a new path,
// paging out diverging segments and paging in the new path.
func (c *prefixCache) switchToPath(newPath []*trieNode, matched int) {
@@ -252,9 +279,11 @@ pageIn:
// schedulePrefillSnapshots schedules every cache to capture snapshots as the
// forward pass crosses the given absolute token offsets, so a single full-size
// prefill records interior states without the caller breaking the batch. The
// passed offsets are user-requested restore points; they are merged with any
// snapshots begin already scheduled (e.g. a branch point), with coinciding
// prefill records interior states without the caller breaking the batch. A
// passed offset names a token prefix; the capture lands at the deepest
// state that prefix alone determines (offset - draftLookahead), which is where
// a prompt sharing exactly that prefix restores. The offsets are merged with
// any snapshots begin already scheduled (e.g. a branch point), with coinciding
// offsets upgraded to user so eviction preserves them.
//
// Offsets at or before the current cache position, or past the end of the
@@ -264,6 +293,7 @@ func (s *cacheSession) schedulePrefillSnapshots(offsets []int) {
c := s.cache
base := c.minCacheOffset()
for _, offset := range offsets {
offset -= c.draftLookahead
if offset <= base || offset > len(s.inputs) {
continue
}
@@ -360,7 +390,7 @@ func (s *cacheSession) attachPrefillSnapshots() {
// no captured state. Skip it rather than materialize a node whose edge
// claims tokens the cache never wrote. Closing its (nil) row is a no-op.
reached := c.minCacheOffset()
stored := append(s.inputs, s.outputs...)
stored := c.key(append(s.inputs, s.outputs...))
for i, p := range pending {
if p.offset > reached {
// Never crossed by a write, so the row is nil; close any entry
@@ -400,7 +430,7 @@ func (s *cacheSession) attachCapturedSnapshots(node *trieNode, snaps []cache.Sna
// advancePath advances the active path from the current frontier by matching
// tokens against existing trie children, splitting partial matches, and
// appending any remaining tokens as new nodes. Returns the new frontier.
func (c *prefixCache) advancePath(frontier *trieNode, tokens []int32, endOffset int) *trieNode {
func (c *prefixCache) advancePath(frontier *trieNode, tokens []trieKey, endOffset int) *trieNode {
// Check if existing children already cover some or all of tokens.
// tokens may span multiple trie nodes when extending a previous run's
// leaf and this snapshot now overlaps that same range.
@@ -487,12 +517,17 @@ func (s *cacheSession) close() {
// that we also actually have the data.
mlx.AsyncEval(arrays...)
// Advance the trie frontier with any newly generated tokens.
// The caches never advance past the stored keys; anything more
// means positions desynced.
c := s.cache
stored := c.key(append(s.inputs, s.outputs...))
if offset > len(stored) {
panic(fmt.Sprintf("cache: offset %d exceeds %d stored keys", offset, len(stored)))
}
// Advance the trie frontier with any newly generated tokens.
if len(c.activePath) > 0 {
frontier := c.activePath[len(c.activePath)-1]
stored := append(s.inputs, s.outputs...)
if offset > frontier.endOffset {
newTokens := stored[frontier.endOffset:offset]
c.advancePath(frontier, newTokens, offset)
x/mlxrunner/prefix_cache_test.go
+108 -69
View File
@@ -495,23 +495,24 @@ func simulateRequest(t *testing.T, pc *prefixCache, inputs, generated []int32, u
assertCacheOffsetAlignment(t, pc, "after begin")
baseOffset := pc.minCacheOffset()
remaining := inputs[baseOffset:]
seed := len(inputs) - 1
// Prefill: schedule the pending snapshots, feed the whole prompt in one pass
// (the caches self-segment at the scheduled offsets), then attach the
// captures to the trie.
// Prefill: schedule the pending snapshots, feed the prompt up to the
// seed token in one pass (the caches self-segment at the scheduled
// offsets), then attach the captures to the trie.
session.schedulePrefillSnapshots(snapshotOffsets)
if len(remaining) > 0 {
feedAll(pc.caches, remaining)
if baseOffset < seed {
feedAll(pc.caches, inputs[baseOffset:seed])
}
session.attachPrefillSnapshots()
assertCacheOffsetAlignment(t, pc, "after prefill")
// Generate tokens.
// Decode: feed the seed and all but the last generated token.
if len(generated) > 0 {
session.outputs = generated
feedAll(pc.caches, generated)
feedAll(pc.caches, inputs[seed:])
feedAll(pc.caches, generated[:len(generated)-1])
}
assertCacheOffsetAlignment(t, pc, "before close")
@@ -586,7 +587,7 @@ func checkTrieInvariants(t *testing.T, root *trieNode) {
}
}
// No two siblings should start with the same token.
seen := make(map[int32]bool)
seen := make(map[trieKey]bool)
for _, c := range n.children {
if len(c.tokens) > 0 {
first := c.tokens[0]
@@ -641,21 +642,36 @@ func checkSnapshotLeaks(t *testing.T, tracker *snapshotTracker, root *trieNode)
}
}
// forEachEnv runs fn as subtests for three realistic model configurations:
// forEachEnv runs fn as subtests for three realistic model configurations —
// pure transformer, transformer + sliding window (Mistral-style), and
// transformer + recurrent (Jamba-style). Leak checking runs automatically
// at the end of each subtest.
// transformer + recurrent (Jamba-style) — each with and without a draft
// look-ahead. Leak checking runs automatically at the end of each subtest.
func forEachEnv(t *testing.T, fn func(t *testing.T, env *testEnv)) {
t.Helper()
run := func(t *testing.T, env *testEnv) {
t.Cleanup(func() {
checkSnapshotLeaks(t, env.tracker, env.pc.root)
})
fn(t, env)
envs := []struct {
name string
make func() *testEnv
}{
{"Transformer", newTransformerEnv},
{"SlidingWindow", newSlidingWindowEnv},
{"Recurrent", newRecurrentEnv},
}
for _, e := range envs {
for _, lookahead := range []int{0, 1} {
name := e.name
if lookahead > 0 {
name += "+Lookahead"
}
t.Run(name, func(t *testing.T) {
env := e.make()
env.pc.draftLookahead = lookahead
t.Cleanup(func() {
checkSnapshotLeaks(t, env.tracker, env.pc.root)
})
fn(t, env)
})
}
}
t.Run("Transformer", func(t *testing.T) { run(t, newTransformerEnv()) })
t.Run("SlidingWindow", func(t *testing.T) { run(t, newSlidingWindowEnv()) })
t.Run("Recurrent", func(t *testing.T) { run(t, newRecurrentEnv()) })
}
// TestBranchCreationAndReuse exercises the core multi-conversation lifecycle:
@@ -671,25 +687,28 @@ func TestBranchCreationAndReuse(t *testing.T) {
if len(resA.remaining) != 8 {
t.Fatalf("A: remaining = %d, want 8 (full miss)", len(resA.remaining))
}
env.assertAllTokens(t, "after A", []int32{1, 2, 3, 4, 5, 6, 7, 8, 20, 21})
env.assertAllTokens(t, "after A", []int32{1, 2, 3, 4, 5, 6, 7, 8, 20})
// Verify trie was populated by close().
_, mA := findBestMatch(pc.root, []int32{1, 2, 3, 4, 5, 6, 7, 8, 20, 21})
if mA != 10 {
t.Fatalf("A findable: expected 10 matched, got %d", mA)
// Verify trie was populated by close(): everything in the caches
// is findable; the last generated token is not.
seqA := []int32{1, 2, 3, 4, 5, 6, 7, 8, 20, 21}
_, mA := findBestMatch(pc.root, pc.key(seqA))
if want := len(seqA) - 1; mA != want {
t.Fatalf("A findable: expected %d matched, got %d", want, mA)
}
// Request B: [1,2,3,4,5,10,11,12] — shares 5-token prefix with A.
// For rewindable caches, switchToPath rewinds to the match point
// so only the non-matching suffix needs evaluation. For non-rewindable
// caches (RecurrentCache), the rewind fails and freeAll fires.
// (one below the shared tokens with a look-ahead) so only the suffix
// needs evaluation. For non-rewindable caches (RecurrentCache), the
// rewind fails and freeAll fires.
resB := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 12}, []int32{30, 31})
if env.rewindable {
if resB.pendingSnapshots != 0 {
t.Fatalf("B: pendingSnapshots = %d, want 0 (rewind succeeded)", resB.pendingSnapshots)
}
if len(resB.remaining) != 3 {
t.Fatalf("B: remaining = %d, want 3 (rewind to match point)", len(resB.remaining))
if want := 3 + pc.draftLookahead; len(resB.remaining) != want {
t.Fatalf("B: remaining = %d, want %d (rewind to match point)", len(resB.remaining), want)
}
} else {
if resB.pendingSnapshots != 1 {
@@ -699,14 +718,14 @@ func TestBranchCreationAndReuse(t *testing.T) {
t.Fatalf("B: remaining = %d, want 8 (freeAll fallback)", len(resB.remaining))
}
}
env.assertAllTokens(t, "after B", []int32{1, 2, 3, 4, 5, 10, 11, 12, 30, 31})
env.assertAllTokens(t, "after B", []int32{1, 2, 3, 4, 5, 10, 11, 12, 30})
// Both A and B should be findable in the trie.
_, mA2 := findBestMatch(pc.root, []int32{1, 2, 3, 4, 5, 6, 7, 8, 20, 21})
_, mA2 := findBestMatch(pc.root, pc.key(seqA))
if mA2 < 5 {
t.Fatalf("A still findable: expected >= 5 matched, got %d", mA2)
}
_, mB := findBestMatch(pc.root, []int32{1, 2, 3, 4, 5, 10, 11, 12, 30, 31})
_, mB := findBestMatch(pc.root, pc.key([]int32{1, 2, 3, 4, 5, 10, 11, 12, 30, 31}))
if mB < 5 {
t.Fatalf("B findable: expected >= 5 matched, got %d", mB)
}
@@ -717,7 +736,7 @@ func TestBranchCreationAndReuse(t *testing.T) {
if len(resC.remaining) >= 10 {
t.Fatalf("C: remaining = %d, want < 10 (should get cache hit)", len(resC.remaining))
}
env.assertAllTokens(t, "after C", []int32{1, 2, 3, 4, 5, 6, 7, 8, 40, 41})
env.assertAllTokens(t, "after C", []int32{1, 2, 3, 4, 5, 6, 7, 8, 40})
checkTrieInvariants(t, pc.root)
})
@@ -753,7 +772,7 @@ func TestExactMatchSeedBehavior(t *testing.T) {
t.Fatalf("B: pendingSnapshots = %d, want 1", resB.pendingSnapshots)
}
}
env.assertAllTokens(t, "after B", []int32{1, 2, 3, 4, 5, 20, 21})
env.assertAllTokens(t, "after B", []int32{1, 2, 3, 4, 5, 20})
checkTrieInvariants(t, pc.root)
})
@@ -768,22 +787,22 @@ func TestConversationResumption(t *testing.T) {
// Turn 1: system prompt + user message, assistant generates response.
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10, 11, 12})
env.assertAllTokens(t, "turn 1", []int32{1, 2, 3, 4, 5, 10, 11, 12})
env.assertAllTokens(t, "turn 1", []int32{1, 2, 3, 4, 5, 10, 11})
// Turn 2: full history + new user message. Should get a cache hit on
// the prefix [1,2,3,4,5,10,11,12] and only need to evaluate [20,21].
// Turn 2: full history + new user message; the match lands exactly
// at the caches' offset, so even exact-offset state is reused.
resB := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21}, []int32{30})
if len(resB.remaining) > 5 {
t.Fatalf("turn 2: remaining = %d, want <= 5 (should reuse most of history)", len(resB.remaining))
}
env.assertAllTokens(t, "turn 2", []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21, 30})
env.assertAllTokens(t, "turn 2", []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21})
// Turn 3: even longer history.
resC := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21, 30, 40, 41}, nil)
if len(resC.remaining) > 5 {
t.Fatalf("turn 3: remaining = %d, want <= 5", len(resC.remaining))
}
env.assertAllTokens(t, "turn 3", []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21, 30, 40, 41})
env.assertAllTokens(t, "turn 3", []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21, 30, 40})
checkTrieInvariants(t, pc.root)
})
@@ -834,9 +853,9 @@ func TestEvictionPreservesActiveConversations(t *testing.T) {
// System prompt prefix should still be findable (multi-child
// branch points are protected from eviction entirely).
_, matched := findBestMatch(pc.root, systemPrompt)
if matched < len(systemPrompt) {
t.Fatalf("system prompt match = %d, want %d", matched, len(systemPrompt))
_, matched := findBestMatch(pc.root, pc.key(systemPrompt))
if want := len(pc.key(systemPrompt)); matched < want {
t.Fatalf("system prompt match = %d, want %d", matched, want)
}
checkTrieInvariants(t, pc.root)
@@ -844,30 +863,39 @@ func TestEvictionPreservesActiveConversations(t *testing.T) {
}
// TestUserSnapshotPreservesRestorePoint verifies that user-created snapshots
// (snapshot(true)) resist structural changes that would destroy them:
// (snapshot(true)) are exact restore points that resist structural changes:
// - A user node forces new tokens into a child instead of extending in-place
// - A prompt diverging at the snapshot resumes there, even for caches that
// cannot rewind
// - The snapshot remains restorable after other branches are added
func TestUserSnapshotPreservesRestorePoint(t *testing.T) {
forEachEnv(t, func(t *testing.T, env *testEnv) {
pc := env.pc
inputs := []int32{1, 2, 3, 4, 5}
// Request A: user snapshot at offset 5, then generate.
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10, 11}, 5)
// Request A: user snapshot at offset 4, then generate.
simulateRequest(t, pc, inputs, []int32{10, 11}, 4)
assertUserNodeExists(t, pc, "after A")
// Request B: extends A's prefix. The user node at offset 5 should
// force tokens into a child rather than extending in-place.
// Request B: extends A's prefix. The user node should force tokens
// into a child rather than extending in-place.
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 20, 21}, nil)
env.assertAllTokens(t, "after B", []int32{1, 2, 3, 4, 5, 10, 11, 20, 21})
env.assertAllTokens(t, "after B", []int32{1, 2, 3, 4, 5, 10, 11, 20})
assertUserNodeExists(t, pc, "after B")
// Request C: diverge from the user node.
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 30, 31}, []int32{40})
// Request C: diverge at the snapshot — prefill resumes at its capture
// point, one token lower with a look-ahead (the boundary token
// re-evaluates to rebuild its draft pair).
divergeC := append(slices.Clone(inputs[:4]), 30, 31)
resC := simulateRequest(t, pc, divergeC, []int32{40})
if want := divergeC[4-pc.draftLookahead:]; !slices.Equal(resC.remaining, want) {
t.Fatalf("C: remaining = %v, want %v", resC.remaining, want)
}
// Request D: switch back to A's branch — user snapshot still restorable.
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 20, 21, 50}, nil)
env.assertAllTokens(t, "back to A", []int32{1, 2, 3, 4, 5, 10, 11, 20, 21, 50})
env.assertAllTokens(t, "back to A", []int32{1, 2, 3, 4, 5, 10, 11, 20, 21})
checkTrieInvariants(t, pc.root)
})
@@ -878,9 +906,10 @@ func TestUserSnapshotPreservesRestorePoint(t *testing.T) {
func TestUserSnapshotResistsAutoMerge(t *testing.T) {
forEachEnv(t, func(t *testing.T, env *testEnv) {
pc := env.pc
inputs := []int32{1, 2, 3, 4, 5}
// Request A: user snapshot at offset 3, then continue to offset 5.
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10}, 3)
simulateRequest(t, pc, inputs, []int32{10}, 3)
// Request B: diverges at the user node, creating a second child.
simulateRequest(t, pc, []int32{1, 2, 3, 6, 7}, []int32{20})
@@ -924,20 +953,25 @@ func TestSnapshotBeyondPrefillSkipped(t *testing.T) {
inputs := []int32{1, 2, 3, 4, 5}
session := pc.begin(inputs)
// Request a reachable snapshot at 3 and one at len(inputs), which a
// prefill that stops one token short never crosses.
// Request a snapshot at 3 and one at len(inputs); captures land at
// the requested prefix minus the look-ahead.
session.schedulePrefillSnapshots([]int{3, len(inputs)})
// Prefill writes all but the final token (mirrors total-processed > 1).
feedAll(pc.caches, inputs[pc.minCacheOffset():len(inputs)-1])
session.attachPrefillSnapshots()
// The reachable offset became a node; the unreached one did not.
if !nodeExistsAtOffset(pc.root, 3) {
t.Errorf("no trie node at reached offset 3")
}
if nodeExistsAtOffset(pc.root, len(inputs)) {
t.Errorf("trie node materialized at unreached offset %d", len(inputs))
// The first request became a node at its capture point; nothing may
// claim offsets the prefill never wrote.
if at := 3 - pc.draftLookahead; !nodeExistsAtOffset(pc.root, at) {
t.Errorf("no trie node at capture point %d", at)
}
reached := pc.minCacheOffset()
walkNodes(pc.root, func(n *trieNode) bool {
if n.endOffset > reached {
t.Errorf("trie node materialized at unwritten offset %d", n.endOffset)
}
return true
})
checkTrieInvariants(t, pc.root)
})
@@ -971,7 +1005,7 @@ func TestPrefillSnapshotsDiscardedOnCancel(t *testing.T) {
// A second request re-prepares snapshots on the same caches: if the
// discarded ones were not closed, prepare() orphans them here.
simulateRequest(t, pc, inputs, nil, 4)
simulateRequest(t, pc, inputs, nil, 5)
checkTrieInvariants(t, pc.root)
})
@@ -1027,15 +1061,15 @@ func TestBranchSwitchRestoresCorrectState(t *testing.T) {
// Request A: [1,2,3,4,5] + generate [10,11]
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10, 11})
env.assertAllTokens(t, "after A", []int32{1, 2, 3, 4, 5, 10, 11})
env.assertAllTokens(t, "after A", []int32{1, 2, 3, 4, 5, 10})
// Request B: [1,2,3,6,7] — diverges at token 4
simulateRequest(t, pc, []int32{1, 2, 3, 6, 7}, []int32{12, 13})
env.assertAllTokens(t, "after B", []int32{1, 2, 3, 6, 7, 12, 13})
env.assertAllTokens(t, "after B", []int32{1, 2, 3, 6, 7, 12})
// Request C: switch back to A's branch [1,2,3,4,5,10,11,20]
simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 20}, nil)
env.assertAllTokens(t, "after C (back to A)", []int32{1, 2, 3, 4, 5, 10, 11, 20})
env.assertAllTokens(t, "after C (back to A)", []int32{1, 2, 3, 4, 5, 10, 11})
checkTrieInvariants(t, pc.root)
})
@@ -1066,7 +1100,9 @@ func TestLRUOnlyUpdatesUsedNodes(t *testing.T) {
// and extend it. The branch point's snapshot may be paged in
// for some cache types but not others.
beforeRequest := time.Now()
simulateRequest(t, pc, []int32{1, 2, 3, 6, 7, 20, 21, 30}, nil)
inputsC := []int32{1, 2, 3, 6, 7, 20, 21, 30}
resC := simulateRequest(t, pc, inputsC, nil)
landing := len(inputsC) - len(resC.remaining)
// The path must have enough depth to exercise intermediate nodes.
if len(pc.activePath) < 3 {
@@ -1080,9 +1116,12 @@ func TestLRUOnlyUpdatesUsedNodes(t *testing.T) {
frontier.lastUsed, beforeRequest)
}
// Every non-frontier node on the active path (including root)
// should retain its old lastUsed — only the frontier gets refreshed.
// Only used nodes refresh — the frontier and the restore landing;
// merely traversed nodes keep their age so they can still evict.
for i, node := range pc.activePath[:len(pc.activePath)-1] {
if node.endOffset == landing {
continue
}
if !node.lastUsed.Before(beforeRequest) {
t.Errorf("activePath[%d] (endOffset=%d) lastUsed was refreshed: got %v, want < %v",
i, node.endOffset, node.lastUsed, beforeRequest)
@@ -1101,7 +1140,7 @@ func TestPagedOutBytesUpdatesOnMaterialize(t *testing.T) {
pc := &prefixCache{}
pc.ensureRoot()
node := &trieNode{parent: pc.root, tokens: []int32{1, 2, 3}, endOffset: 3}
node := &trieNode{parent: pc.root, tokens: []trieKey{1, 2, 3}, endOffset: 3}
pc.root.children = append(pc.root.children, node)
snap := &fakeSnapshot{from: 0, to: 3, byteSize: 0}
@@ -1127,7 +1166,7 @@ func TestSwapSnapshotsDetachesHook(t *testing.T) {
pc := &prefixCache{}
pc.ensureRoot()
node := &trieNode{parent: pc.root, tokens: []int32{1, 2, 3}, endOffset: 3}
node := &trieNode{parent: pc.root, tokens: []trieKey{1, 2, 3}, endOffset: 3}
pc.root.children = append(pc.root.children, node)
snap := &fakeSnapshot{from: 0, to: 3, byteSize: 0}
x/mlxrunner/speculate.go
+23 -26
View File
@@ -25,13 +25,10 @@ type drafter interface {
// chunks, the decode seed, then each round's validated tokens.
committed(tokens, hiddens *mlx.Array, position int)
// finish reports generation ended with current sampled but never
// committed, so the drafter can settle state tracking the target caches.
finish(current *mlx.Array)
// flush writes any buffered committed reports through to the draft
// caches. A drafter without draft caches has nothing to write.
flush()
// settle completes any open frontier pair with next — the token after
// the last committed slot — and writes buffered reports through,
// leveling the draft caches with the target caches.
settle(next *mlx.Array)
close()
}
@@ -184,21 +181,13 @@ func (s *speculationSession) committed(tokens, hiddens *mlx.Array, position int)
s.drafter.committed(tokens, hiddens, position)
}
// finish reports the end of generation to the drafter: current was sampled
// after the last committed slot and will never be committed.
func (s *speculationSession) finish(current *mlx.Array) {
// settle completes the drafter's open frontier pair with next and writes
// buffered reports through to the draft caches.
func (s *speculationSession) settle(next *mlx.Array) {
if s == nil {
return
}
s.drafter.finish(current)
}
// flush writes the drafter's buffered committed reports to the draft caches.
func (s *speculationSession) flush() {
if s == nil {
return
}
s.drafter.flush()
s.drafter.settle(next)
}
func (s *speculationSession) close() {
@@ -225,8 +214,8 @@ type speculativeDecoder struct {
// decoder returns the decoder for this engine's session. A speculationSession that
// cannot draft (logprobs) has no depth controller and permanently parks,
// running the inner pipelined decoder whose reports keep the draft KV level.
func (s *speculationSession) decoder(seed []int32, position int) decoder {
current := sampler.Result{Token: mlx.FromValues(seed, len(seed))}
func (s *speculationSession) decoder(seed *mlx.Array, position int) decoder {
current := sampler.Result{Token: seed}
mlx.Pin(current.Arrays()...)
return &speculativeDecoder{s: s, position: position, current: current}
}
@@ -284,14 +273,13 @@ func (st *speculativeDecoder) advance(next sampler.Result) {
// but never forwarded) is exactly the current token a drafting round expects,
// so emit it and let the next call draft from it.
func (st *speculativeDecoder) resume() []sampler.Result {
next, position := st.inner.detach()
next, position := st.inner.drain()
st.position = position
st.inner.close()
st.inner = nil
// No round spans this call, so the next beginRound attributes no cost.
st.s.roundDrafts = -1
// detach handed over the pin; advance re-pins, no sweep in between.
mlx.Unpin(next.Arrays()...)
return []sampler.Result{next}
return next
}
// park decodes one pipelined plain token while the engine cannot draft. Each
@@ -305,6 +293,15 @@ func (st *speculativeDecoder) park(remaining int) ([]sampler.Result, error) {
return st.inner.next(remaining)
}
// drain surrenders the inner decoder's undelivered sample while parked; a
// drafting decoder has already delivered everything it sampled.
func (st *speculativeDecoder) drain() ([]sampler.Result, int) {
if st.inner != nil {
return st.inner.drain()
}
return nil, st.position
}
func (st *speculativeDecoder) close() {
if st.inner != nil {
// Ended while parked: the inner decoder's close settles the drafter
@@ -313,7 +310,7 @@ func (st *speculativeDecoder) close() {
} else {
// The final token was emitted but never forwarded; its report settles
// the drafter level with the caches' resting offset.
st.s.finish(st.current.Token)
st.s.settle(st.current.Token)
}
mlx.Unpin(st.current.Arrays()...)
st.s.logStats()