mirror of
https://github.com/ollama/ollama.git
synced 2026-09-21 05:28:00 -05:00
mlxrunner: evict the active conversation's own checkpoints under the budget
Eviction skipped every node on the active path, so a conversation's own
turn checkpoints were never reclaimed no matter how far over budget the
trie was. On models with sliding-window or recurrent layers each turn's
checkpoint is a full copy of that state, 800 MiB per turn on
gemma4:31b-mlx, and a long chat grows without bound. The scheduler then
counts that memory as in use and evicts the model to load anything
else.
Only the frontier and branch points are protected now. Any other node,
active or not, is evicted least recently used first. On the active path
that merges a turn into the next one: the merged node keeps the newer
whole-state, and the KV snapshots there are lazy views of the live
buffer, so nothing is copied. Rewinding to an evicted turn resumes at
the newest surviving checkpoint before it.
qwen3.8:27b-mlx on an M5 Max, the same short question every turn with
24 tokens generated per reply, 8 GiB budget, 17.2 GiB of weights:
turn | before: paged out nodes reported | after: paged out nodes reported
11 | 4.61 GiB 33 21.6 GiB | 4.61 GiB 33 21.6 GiB
21 | 7.91 GiB 56 24.9 GiB | 7.92 GiB 56 24.9 GiB
31 | 8.46 GiB 60 25.5 GiB | 7.94 GiB 56 25.0 GiB
41 | 9.90 GiB 70 26.9 GiB | 7.96 GiB 56 25.0 GiB
50 | 11.19 GiB 79 28.2 GiB | 7.98 GiB 56 25.0 GiB
Fixes #17783
This commit is contained in:
@@ -606,15 +606,13 @@ func (c *prefixCache) enforceEvictionPolicy() {
|
||||
return
|
||||
}
|
||||
|
||||
activeSet := make(map[*trieNode]bool, len(c.activePath))
|
||||
for _, n := range c.activePath {
|
||||
activeSet[n] = true
|
||||
}
|
||||
|
||||
for c.pagedOutBytes > maxPagedOutBytes {
|
||||
// Evicting the frontier's parent merges the frontier into it, so
|
||||
// resolve the frontier again after every eviction.
|
||||
frontier := c.activePath[len(c.activePath)-1]
|
||||
var best *trieNode
|
||||
walkNodes(c.root, func(n *trieNode) bool {
|
||||
if n == c.root || activeSet[n] || len(n.children) > 1 {
|
||||
if n == c.root || n == frontier || len(n.children) > 1 {
|
||||
return true
|
||||
}
|
||||
// Evict: oldest, then deepest, then largest.
|
||||
@@ -644,7 +642,11 @@ func (c *prefixCache) evictNode(node *trieNode) {
|
||||
// Interior node with one child: merge with child.
|
||||
before := c.pagedOutBytes
|
||||
tokens := len(node.tokens)
|
||||
child := node.children[0]
|
||||
mergeWithChild(node, c.caches, &c.pagedOutBytes)
|
||||
if i := slices.Index(c.activePath, child); i >= 0 {
|
||||
c.activePath = slices.Delete(c.activePath, i, i+1)
|
||||
}
|
||||
slog.Debug("evicting interior node", "offset", node.startOffset(), "tokens", tokens, "freed", mlx.PrettyBytes(int(before-c.pagedOutBytes)))
|
||||
} else {
|
||||
panic("evictNode called on multi-child branch point")
|
||||
|
||||
@@ -881,7 +881,7 @@ func TestEvictionPreservesActiveConversations(t *testing.T) {
|
||||
t.Fatalf("pagedOutBytes = %d, want <= %d", pc.pagedOutBytes, maxPagedOutBytes)
|
||||
}
|
||||
|
||||
// Active path should be untouched.
|
||||
// The branch point and the frontier survive.
|
||||
if len(pc.activePath) < 2 {
|
||||
t.Fatalf("activePath should have >= 2 nodes, got %d", len(pc.activePath))
|
||||
}
|
||||
@@ -954,8 +954,26 @@ func TestUserSnapshotResistsAutoMerge(t *testing.T) {
|
||||
t.Fatalf("user node children = %d, want 2", len(userNode.children))
|
||||
}
|
||||
|
||||
// Inflate snapshot sizes and evict. The non-active branch should be
|
||||
// evicted, leaving the user node with one child.
|
||||
// Inflate snapshot sizes so that evicting the non-active branch alone
|
||||
// brings the trie under budget, leaving the user node with one child.
|
||||
var kept, evicted int
|
||||
walkNodes(pc.root, func(n *trieNode) bool {
|
||||
for _, s := range n.snapshots {
|
||||
if s == nil {
|
||||
continue
|
||||
}
|
||||
if n.parent == userNode && !slices.Contains(pc.activePath, n) {
|
||||
evicted++
|
||||
} else {
|
||||
kept++
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
if evicted == 0 {
|
||||
t.Fatal("no snapshots on the non-active branch")
|
||||
}
|
||||
size := int(maxPagedOutBytes) / kept
|
||||
walkNodes(pc.root, func(n *trieNode) bool {
|
||||
if !n.hasSnapshots() {
|
||||
return true
|
||||
@@ -963,7 +981,7 @@ func TestUserSnapshotResistsAutoMerge(t *testing.T) {
|
||||
snaps := make([]cache.Snapshot, len(n.snapshots))
|
||||
for i, s := range n.snapshots {
|
||||
if s != nil {
|
||||
snaps[i] = &fakeSnapshot{byteSize: 5 * 1024 * 1024 * 1024}
|
||||
snaps[i] = &fakeSnapshot{byteSize: size}
|
||||
}
|
||||
}
|
||||
n.setSnapshots(snaps, &pc.pagedOutBytes)
|
||||
|
||||
Reference in New Issue
Block a user