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:
Jesse Gross
2026-09-10 15:20:11 -07:00
parent b859a94509
commit 6137793ac4
2 changed files with 30 additions and 10 deletions
+8 -6
View File
@@ -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")
+22 -4
View File
@@ -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)