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

mlxrunner: rename kvCache to prefixCache

Browse Source 
The type coordinates every cache kind — KV, sliding-window, recurrent —
around prefix matching over the trie, so kv was a misnomer.
This commit is contained in:
Jesse Gross
2026-07-06 13:17:59 -07:00
parent 4e96f4dbf2
commit dd49563d55
4 changed files with 134 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
x/mlxrunner/cache_trie.go
+1 -1
View File
@@ -8,7 +8,7 @@ import (
"github.com/ollama/ollama/x/mlxrunner/cache" "github.com/ollama/ollama/x/mlxrunner/cache"
) )
// trieNode represents a node in the compressed prefix trie for KV cache branching. // trieNode represents a node in the compressed prefix trie for cache branching.
// Each node stores a compressed edge (multiple tokens) and optional paged-out // Each node stores a compressed edge (multiple tokens) and optional paged-out
// snapshot data per cache layer. // snapshot data per cache layer.
type trieNode struct { type trieNode struct {
x/mlxrunner/cache.go → x/mlxrunner/prefix_cache.go
+17 -16
View File
@@ -1,6 +1,7 @@
// cache.go manages a shared KV cache across conversations using a compressed // prefix_cache.go manages cache state shared across conversations using a
// prefix trie. Each trie node stores a token sequence (edge) and optional // compressed prefix trie. Each trie node stores a token sequence (edge) and
// per-layer snapshots that can be paged in/out of the live MLX cache arrays. // optional per-layer snapshots that can be paged in/out of the live MLX cache
// arrays.
// //
// Key properties: // Key properties:
// - Only one path through the trie is "active" (backed by live MLX arrays) // - Only one path through the trie is "active" (backed by live MLX arrays)
@@ -31,7 +32,7 @@ import (
const maxPagedOutBytes int64 = 8 << 30 // 8 GiB eviction threshold for paged-out snapshot memory const maxPagedOutBytes int64 = 8 << 30 // 8 GiB eviction threshold for paged-out snapshot memory
type kvCache struct { type prefixCache struct {
root *trieNode // root of the prefix trie root *trieNode // root of the prefix trie
activePath []*trieNode // current root→leaf path with live MLX arrays activePath []*trieNode // current root→leaf path with live MLX arrays
caches []cache.Cache caches []cache.Cache
@@ -48,7 +49,7 @@ type pendingSnapshot struct {
// Callers should append generated tokens to outputs and // Callers should append generated tokens to outputs and
// defer close to save the cache state. // defer close to save the cache state.
type cacheSession struct { type cacheSession struct {
cache *kvCache cache *prefixCache
inputs []int32 inputs []int32
outputs []int32 outputs []int32
@@ -61,7 +62,7 @@ type cacheSession struct {
pendingSnapshots []pendingSnapshot pendingSnapshots []pendingSnapshot
} }
func (c *kvCache) ensureCaches(m base.Model) { func (c *prefixCache) ensureCaches(m base.Model) {
if len(c.caches) != 0 { if len(c.caches) != 0 {
return return
} }
@@ -75,7 +76,7 @@ func (c *kvCache) ensureCaches(m base.Model) {
} }
} }
func (c *kvCache) ensureRoot() { func (c *prefixCache) ensureRoot() {
if c.root == nil { if c.root == nil {
c.root = &trieNode{ c.root = &trieNode{
lastUsed: time.Now(), lastUsed: time.Now(),
@@ -86,7 +87,7 @@ func (c *kvCache) ensureRoot() {
// begin prepares caches for a new request. It finds the nearest // begin prepares caches for a new request. It finds the nearest
// matching cache or creates new caches if none match. // matching cache or creates new caches if none match.
func (c *kvCache) begin(m base.Model, inputs []int32) *cacheSession { func (c *prefixCache) begin(m base.Model, inputs []int32) *cacheSession {
c.ensureCaches(m) c.ensureCaches(m)
c.ensureRoot() c.ensureRoot()
@@ -130,7 +131,7 @@ func (c *kvCache) begin(m base.Model, inputs []int32) *cacheSession {
// switchToPath transitions from the current active path to a new path, // switchToPath transitions from the current active path to a new path,
// paging out diverging segments and paging in the new path. // paging out diverging segments and paging in the new path.
func (c *kvCache) switchToPath(newPath []*trieNode, matched int) { func (c *prefixCache) switchToPath(newPath []*trieNode, matched int) {
defer c.enforceEvictionPolicy() defer c.enforceEvictionPolicy()
// Find common ancestor index. // Find common ancestor index.
@@ -401,7 +402,7 @@ func (s *cacheSession) attachCapturedSnapshots(node *trieNode, snaps []cache.Sna
// advancePath advances the active path from the current frontier by matching // advancePath advances the active path from the current frontier by matching
// tokens against existing trie children, splitting partial matches, and // tokens against existing trie children, splitting partial matches, and
// appending any remaining tokens as new nodes. Returns the new frontier. // appending any remaining tokens as new nodes. Returns the new frontier.
func (c *kvCache) advancePath(frontier *trieNode, tokens []int32, endOffset int) *trieNode { func (c *prefixCache) advancePath(frontier *trieNode, tokens []int32, endOffset int) *trieNode {
// Check if existing children already cover some or all of tokens. // Check if existing children already cover some or all of tokens.
// tokens may span multiple trie nodes when extending a previous run's // tokens may span multiple trie nodes when extending a previous run's
// leaf and this snapshot now overlaps that same range. // leaf and this snapshot now overlaps that same range.
@@ -438,7 +439,7 @@ func (c *kvCache) advancePath(frontier *trieNode, tokens []int32, endOffset int)
} }
// freeAll releases all cache layers. // freeAll releases all cache layers.
func (c *kvCache) freeAll() { func (c *prefixCache) freeAll() {
for _, kv := range c.caches { for _, kv := range c.caches {
if kv != nil { if kv != nil {
kv.Free() kv.Free()
@@ -446,7 +447,7 @@ func (c *kvCache) freeAll() {
} }
} }
func (c *kvCache) minCacheOffset() int { func (c *prefixCache) minCacheOffset() int {
offset := 0 offset := 0
found := false found := false
for _, kv := range c.caches { for _, kv := range c.caches {
@@ -503,7 +504,7 @@ func (s *cacheSession) close() {
} }
// enforceEvictionPolicy evicts eligible nodes until paged-out memory is within limits. // enforceEvictionPolicy evicts eligible nodes until paged-out memory is within limits.
func (c *kvCache) enforceEvictionPolicy() { func (c *prefixCache) enforceEvictionPolicy() {
if c.pagedOutBytes <= maxPagedOutBytes { if c.pagedOutBytes <= maxPagedOutBytes {
return return
} }
@@ -537,7 +538,7 @@ func (c *kvCache) enforceEvictionPolicy() {
} }
// evictNode evicts a single node from the trie, freeing its snapshot memory. // evictNode evicts a single node from the trie, freeing its snapshot memory.
func (c *kvCache) evictNode(node *trieNode) { func (c *prefixCache) evictNode(node *trieNode) {
if len(node.children) == 0 { if len(node.children) == 0 {
// Leaf: remove entirely. // Leaf: remove entirely.
slog.Debug("evicting leaf", "offset", node.startOffset(), "tokens", len(node.tokens), "freed", mlx.PrettyBytes(int(node.snapshotBytes()))) slog.Debug("evicting leaf", "offset", node.startOffset(), "tokens", len(node.tokens), "freed", mlx.PrettyBytes(int(node.snapshotBytes())))
@@ -553,7 +554,7 @@ func (c *kvCache) evictNode(node *trieNode) {
} }
} }
func (c *kvCache) dumpTree() { func (c *prefixCache) dumpTree() {
// Summary stats // Summary stats
var cacheBytes int var cacheBytes int
for _, kv := range c.caches { for _, kv := range c.caches {
@@ -640,7 +641,7 @@ func (c *kvCache) dumpTree() {
dump(c.root, "", true) dump(c.root, "", true)
offset := c.minCacheOffset() offset := c.minCacheOffset()
logutil.Trace(fmt.Sprintf("kv cache active_tokens: %d, active_size: %s, paged_out: %s, trie: nodes=%d, snapshots=%d", logutil.Trace(fmt.Sprintf("prefix cache active_tokens: %d, active_size: %s, paged_out: %s, trie: nodes=%d, snapshots=%d",
offset, mlx.PrettyBytes(cacheBytes), mlx.PrettyBytes(int(pagedBytes)), nodeCount, snapshotCount)) offset, mlx.PrettyBytes(cacheBytes), mlx.PrettyBytes(int(pagedBytes)), nodeCount, snapshotCount))
for i, l := range lines { for i, l := range lines {
if i == 0 { if i == 0 {
x/mlxrunner/cache_test.go → x/mlxrunner/prefix_cache_test.go
+115 -115
View File
@@ -398,9 +398,9 @@ type feedableCache interface {
feed(tokens []int32) feed(tokens []int32)
} }
// testEnv encapsulates a kvCache and its fake caches for a test scenario. // testEnv encapsulates a prefixCache and its fake caches for a test scenario.
type testEnv struct { type testEnv struct {
kvc *kvCache pc *prefixCache
caches []cache.Cache // typed references for assertions caches []cache.Cache // typed references for assertions
tracker *snapshotTracker tracker *snapshotTracker
rewindable bool // true when all caches support arbitrary Restore(nil, target) rewindable bool // true when all caches support arbitrary Restore(nil, target)
@@ -412,7 +412,7 @@ func newTransformerEnv() *testEnv {
tracker := &snapshotTracker{} tracker := &snapshotTracker{}
caches := []cache.Cache{&fakeRewindableCache{tracker: tracker}} caches := []cache.Cache{&fakeRewindableCache{tracker: tracker}}
return &testEnv{ return &testEnv{
kvc: &kvCache{caches: caches}, pc: &prefixCache{caches: caches},
caches: caches, caches: caches,
tracker: tracker, tracker: tracker,
rewindable: true, rewindable: true,
@@ -430,7 +430,7 @@ func newSlidingWindowEnv() *testEnv {
sw := &fakeSlidingWindowCache{maxSize: 4, tracker: tr} sw := &fakeSlidingWindowCache{maxSize: 4, tracker: tr}
caches := []cache.Cache{rc, sw} caches := []cache.Cache{rc, sw}
return &testEnv{ return &testEnv{
kvc: &kvCache{caches: caches}, pc: &prefixCache{caches: caches},
caches: caches, caches: caches,
tracker: tr, tracker: tr,
rewindable: false, rewindable: false,
@@ -445,7 +445,7 @@ func newRecurrentEnv() *testEnv {
nrc := &fakeRecurrentCache{tracker: tr} nrc := &fakeRecurrentCache{tracker: tr}
caches := []cache.Cache{rc, nrc} caches := []cache.Cache{rc, nrc}
return &testEnv{ return &testEnv{
kvc: &kvCache{caches: caches}, pc: &prefixCache{caches: caches},
caches: caches, caches: caches,
tracker: tr, tracker: tr,
rewindable: false, rewindable: false,
@@ -476,10 +476,10 @@ type requestResult struct {
// simulateRequest runs a request through the harness. If userSnapshotAt > 0, // simulateRequest runs a request through the harness. If userSnapshotAt > 0,
// a user snapshot is requested at that offset during prefill. // a user snapshot is requested at that offset during prefill.
func simulateRequest(t *testing.T, kvc *kvCache, inputs, generated []int32, userSnapshotAt ...int) requestResult { func simulateRequest(t *testing.T, pc *prefixCache, inputs, generated []int32, userSnapshotAt ...int) requestResult {
t.Helper() t.Helper()
session := kvc.begin(nil, inputs) session := pc.begin(nil, inputs)
var snapshotOffsets []int var snapshotOffsets []int
for _, at := range userSnapshotAt { for _, at := range userSnapshotAt {
if at > 0 { if at > 0 {
@@ -492,9 +492,9 @@ func simulateRequest(t *testing.T, kvc *kvCache, inputs, generated []int32, user
pendingSnapshots: len(session.pendingSnapshots), pendingSnapshots: len(session.pendingSnapshots),
} }
assertCacheOffsetAlignment(t, kvc, "after begin") assertCacheOffsetAlignment(t, pc, "after begin")
baseOffset := kvc.minCacheOffset() baseOffset := pc.minCacheOffset()
remaining := inputs[baseOffset:] remaining := inputs[baseOffset:]
// Prefill: schedule the pending snapshots, feed the whole prompt in one pass // Prefill: schedule the pending snapshots, feed the whole prompt in one pass
@@ -502,19 +502,19 @@ func simulateRequest(t *testing.T, kvc *kvCache, inputs, generated []int32, user
// captures to the trie. // captures to the trie.
session.schedulePrefillSnapshots(snapshotOffsets) session.schedulePrefillSnapshots(snapshotOffsets)
if len(remaining) > 0 { if len(remaining) > 0 {
feedAll(kvc.caches, remaining) feedAll(pc.caches, remaining)
} }
session.attachPrefillSnapshots() session.attachPrefillSnapshots()
assertCacheOffsetAlignment(t, kvc, "after prefill") assertCacheOffsetAlignment(t, pc, "after prefill")
// Generate tokens. // Generate tokens.
if len(generated) > 0 { if len(generated) > 0 {
session.outputs = generated session.outputs = generated
feedAll(kvc.caches, generated) feedAll(pc.caches, generated)
} }
assertCacheOffsetAlignment(t, kvc, "before close") assertCacheOffsetAlignment(t, pc, "before close")
session.close() session.close()
return result return result
} }
@@ -528,14 +528,14 @@ func feedAll(caches []cache.Cache, tokens []int32) {
} }
// assertCacheOffsetAlignment verifies all caches report the same offset. // assertCacheOffsetAlignment verifies all caches report the same offset.
func assertCacheOffsetAlignment(t *testing.T, kvc *kvCache, label string) { func assertCacheOffsetAlignment(t *testing.T, pc *prefixCache, label string) {
t.Helper() t.Helper()
if len(kvc.caches) < 2 { if len(pc.caches) < 2 {
return return
} }
expected := kvc.caches[0].Offset() expected := pc.caches[0].Offset()
for i := 1; i < len(kvc.caches); i++ { for i := 1; i < len(pc.caches); i++ {
if got := kvc.caches[i].Offset(); got != expected { if got := pc.caches[i].Offset(); got != expected {
t.Errorf("%s: cache %d offset=%d != cache 0 offset=%d", label, i, got, expected) t.Errorf("%s: cache %d offset=%d != cache 0 offset=%d", label, i, got, expected)
} }
} }
@@ -649,7 +649,7 @@ func forEachEnv(t *testing.T, fn func(t *testing.T, env *testEnv)) {
t.Helper() t.Helper()
run := func(t *testing.T, env *testEnv) { run := func(t *testing.T, env *testEnv) {
t.Cleanup(func() { t.Cleanup(func() {
checkSnapshotLeaks(t, env.tracker, env.kvc.root) checkSnapshotLeaks(t, env.tracker, env.pc.root)
}) })
fn(t, env) fn(t, env)
} }
@@ -664,17 +664,17 @@ func forEachEnv(t *testing.T, fn func(t *testing.T, env *testEnv)) {
// hit lengths, and that semantic caches contain the correct token sequences. // hit lengths, and that semantic caches contain the correct token sequences.
func TestBranchCreationAndReuse(t *testing.T) { func TestBranchCreationAndReuse(t *testing.T) {
forEachEnv(t, func(t *testing.T, env *testEnv) { forEachEnv(t, func(t *testing.T, env *testEnv) {
kvc := env.kvc pc := env.pc
// Request A: [1,2,3,4,5,6,7,8] + generate [20,21] — full miss. // Request A: [1,2,3,4,5,6,7,8] + generate [20,21] — full miss.
resA := simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5, 6, 7, 8}, []int32{20, 21}) resA := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 6, 7, 8}, []int32{20, 21})
if len(resA.remaining) != 8 { if len(resA.remaining) != 8 {
t.Fatalf("A: remaining = %d, want 8 (full miss)", len(resA.remaining)) 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, 21})
// Verify trie was populated by close(). // Verify trie was populated by close().
_, mA := findBestMatch(kvc.root, []int32{1, 2, 3, 4, 5, 6, 7, 8, 20, 21}) _, mA := findBestMatch(pc.root, []int32{1, 2, 3, 4, 5, 6, 7, 8, 20, 21})
if mA != 10 { if mA != 10 {
t.Fatalf("A findable: expected 10 matched, got %d", mA) t.Fatalf("A findable: expected 10 matched, got %d", mA)
} }
@@ -683,7 +683,7 @@ func TestBranchCreationAndReuse(t *testing.T) {
// For rewindable caches, switchToPath rewinds to the match point // For rewindable caches, switchToPath rewinds to the match point
// so only the non-matching suffix needs evaluation. For non-rewindable // so only the non-matching suffix needs evaluation. For non-rewindable
// caches (RecurrentCache), the rewind fails and freeAll fires. // caches (RecurrentCache), the rewind fails and freeAll fires.
resB := simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5, 10, 11, 12}, []int32{30, 31}) resB := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 12}, []int32{30, 31})
if env.rewindable { if env.rewindable {
if resB.pendingSnapshots != 0 { if resB.pendingSnapshots != 0 {
t.Fatalf("B: pendingSnapshots = %d, want 0 (rewind succeeded)", resB.pendingSnapshots) t.Fatalf("B: pendingSnapshots = %d, want 0 (rewind succeeded)", resB.pendingSnapshots)
@@ -702,24 +702,24 @@ func TestBranchCreationAndReuse(t *testing.T) {
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, 31})
// Both A and B should be findable in the trie. // Both A and B should be findable in the trie.
_, mA2 := findBestMatch(kvc.root, []int32{1, 2, 3, 4, 5, 6, 7, 8, 20, 21}) _, mA2 := findBestMatch(pc.root, []int32{1, 2, 3, 4, 5, 6, 7, 8, 20, 21})
if mA2 < 5 { if mA2 < 5 {
t.Fatalf("A still findable: expected >= 5 matched, got %d", mA2) t.Fatalf("A still findable: expected >= 5 matched, got %d", mA2)
} }
_, mB := findBestMatch(kvc.root, []int32{1, 2, 3, 4, 5, 10, 11, 12, 30, 31}) _, mB := findBestMatch(pc.root, []int32{1, 2, 3, 4, 5, 10, 11, 12, 30, 31})
if mB < 5 { if mB < 5 {
t.Fatalf("B findable: expected >= 5 matched, got %d", mB) t.Fatalf("B findable: expected >= 5 matched, got %d", mB)
} }
// Request C: [1,2,3,4,5,6,7,8,40,41] — extends A's prefix. // Request C: [1,2,3,4,5,6,7,8,40,41] — extends A's prefix.
// Should get a cache hit for the shared prefix. // Should get a cache hit for the shared prefix.
resC := simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5, 6, 7, 8, 40, 41}, nil) resC := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 6, 7, 8, 40, 41}, nil)
if len(resC.remaining) >= 10 { if len(resC.remaining) >= 10 {
t.Fatalf("C: remaining = %d, want < 10 (should get cache hit)", len(resC.remaining)) 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, 41})
checkTrieInvariants(t, kvc.root) checkTrieInvariants(t, pc.root)
}) })
} }
@@ -728,16 +728,16 @@ func TestBranchCreationAndReuse(t *testing.T) {
// The last token must be re-evaluated to seed generation. // The last token must be re-evaluated to seed generation.
func TestExactMatchSeedBehavior(t *testing.T) { func TestExactMatchSeedBehavior(t *testing.T) {
forEachEnv(t, func(t *testing.T, env *testEnv) { forEachEnv(t, func(t *testing.T, env *testEnv) {
kvc := env.kvc pc := env.pc
// Request A: first time. // Request A: first time.
simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5}, []int32{10, 11}) simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10, 11})
// Request B: identical prompt. Holdback means matched=4, partial in // Request B: identical prompt. Holdback means matched=4, partial in
// the 5-token edge. For rewindable caches, switchToPath rewinds to // the 5-token edge. For rewindable caches, switchToPath rewinds to
// offset 4, so only the held-back token needs re-evaluation. For // offset 4, so only the held-back token needs re-evaluation. For
// non-rewindable caches, the rewind fails and freeAll fires. // non-rewindable caches, the rewind fails and freeAll fires.
resB := simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5}, []int32{20, 21}) resB := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{20, 21})
if env.rewindable { if env.rewindable {
if len(resB.remaining) != 1 { if len(resB.remaining) != 1 {
t.Fatalf("B: remaining = %d, want 1 (rewind to holdback point)", len(resB.remaining)) t.Fatalf("B: remaining = %d, want 1 (rewind to holdback point)", len(resB.remaining))
@@ -755,7 +755,7 @@ func TestExactMatchSeedBehavior(t *testing.T) {
} }
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, 21})
checkTrieInvariants(t, kvc.root) checkTrieInvariants(t, pc.root)
}) })
} }
@@ -764,28 +764,28 @@ func TestExactMatchSeedBehavior(t *testing.T) {
// prefix (system prompt + first turn + assistant response). // prefix (system prompt + first turn + assistant response).
func TestConversationResumption(t *testing.T) { func TestConversationResumption(t *testing.T) {
forEachEnv(t, func(t *testing.T, env *testEnv) { forEachEnv(t, func(t *testing.T, env *testEnv) {
kvc := env.kvc pc := env.pc
// Turn 1: system prompt + user message, assistant generates response. // Turn 1: system prompt + user message, assistant generates response.
simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5}, []int32{10, 11, 12}) 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, 12})
// Turn 2: full history + new user message. Should get a cache hit on // 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]. // the prefix [1,2,3,4,5,10,11,12] and only need to evaluate [20,21].
resB := simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21}, []int32{30}) resB := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21}, []int32{30})
if len(resB.remaining) > 5 { if len(resB.remaining) > 5 {
t.Fatalf("turn 2: remaining = %d, want <= 5 (should reuse most of history)", len(resB.remaining)) 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, 30})
// Turn 3: even longer history. // Turn 3: even longer history.
resC := simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21, 30, 40, 41}, nil) resC := simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 10, 11, 12, 20, 21, 30, 40, 41}, nil)
if len(resC.remaining) > 5 { if len(resC.remaining) > 5 {
t.Fatalf("turn 3: remaining = %d, want <= 5", len(resC.remaining)) 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, 41})
checkTrieInvariants(t, kvc.root) checkTrieInvariants(t, pc.root)
}) })
} }
@@ -794,18 +794,18 @@ func TestConversationResumption(t *testing.T) {
// active path and shared prefix survive while memory stays bounded. // active path and shared prefix survive while memory stays bounded.
func TestEvictionPreservesActiveConversations(t *testing.T) { func TestEvictionPreservesActiveConversations(t *testing.T) {
forEachEnv(t, func(t *testing.T, env *testEnv) { forEachEnv(t, func(t *testing.T, env *testEnv) {
kvc := env.kvc pc := env.pc
systemPrompt := []int32{1, 2, 3, 4, 5} systemPrompt := []int32{1, 2, 3, 4, 5}
// Create 5 conversations with unique suffixes. // Create 5 conversations with unique suffixes.
for i := range 5 { for i := range 5 {
suffix := []int32{int32(100 + i*10), int32(101 + i*10), int32(102 + i*10)} suffix := []int32{int32(100 + i*10), int32(101 + i*10), int32(102 + i*10)}
inputs := append(slices.Clone(systemPrompt), suffix...) inputs := append(slices.Clone(systemPrompt), suffix...)
simulateRequest(t, kvc, inputs, []int32{int32(200 + i)}) simulateRequest(t, pc, inputs, []int32{int32(200 + i)})
} }
// Inflate snapshot sizes to trigger eviction. // Inflate snapshot sizes to trigger eviction.
walkNodes(kvc.root, func(n *trieNode) bool { walkNodes(pc.root, func(n *trieNode) bool {
if !n.hasSnapshots() { if !n.hasSnapshots() {
return true return true
} }
@@ -815,31 +815,31 @@ func TestEvictionPreservesActiveConversations(t *testing.T) {
snaps[i] = &fakeSnapshot{byteSize: 2 * 1024 * 1024 * 1024} // 2 GiB per snapshot snaps[i] = &fakeSnapshot{byteSize: 2 * 1024 * 1024 * 1024} // 2 GiB per snapshot
} }
} }
n.setSnapshots(snaps, &kvc.pagedOutBytes) n.setSnapshots(snaps, &pc.pagedOutBytes)
return true return true
}) })
// Run eviction. // Run eviction.
kvc.enforceEvictionPolicy() pc.enforceEvictionPolicy()
// Memory should be within limits. // Memory should be within limits.
if kvc.pagedOutBytes > maxPagedOutBytes { if pc.pagedOutBytes > maxPagedOutBytes {
t.Fatalf("pagedOutBytes = %d, want <= %d", kvc.pagedOutBytes, maxPagedOutBytes) t.Fatalf("pagedOutBytes = %d, want <= %d", pc.pagedOutBytes, maxPagedOutBytes)
} }
// Active path should be untouched. // Active path should be untouched.
if len(kvc.activePath) < 2 { if len(pc.activePath) < 2 {
t.Fatalf("activePath should have >= 2 nodes, got %d", len(kvc.activePath)) t.Fatalf("activePath should have >= 2 nodes, got %d", len(pc.activePath))
} }
// System prompt prefix should still be findable (multi-child // System prompt prefix should still be findable (multi-child
// branch points are protected from eviction entirely). // branch points are protected from eviction entirely).
_, matched := findBestMatch(kvc.root, systemPrompt) _, matched := findBestMatch(pc.root, systemPrompt)
if matched < len(systemPrompt) { if matched < len(systemPrompt) {
t.Fatalf("system prompt match = %d, want %d", matched, len(systemPrompt)) t.Fatalf("system prompt match = %d, want %d", matched, len(systemPrompt))
} }
checkTrieInvariants(t, kvc.root) checkTrieInvariants(t, pc.root)
}) })
} }
@@ -849,27 +849,27 @@ func TestEvictionPreservesActiveConversations(t *testing.T) {
// - The snapshot remains restorable after other branches are added // - The snapshot remains restorable after other branches are added
func TestUserSnapshotPreservesRestorePoint(t *testing.T) { func TestUserSnapshotPreservesRestorePoint(t *testing.T) {
forEachEnv(t, func(t *testing.T, env *testEnv) { forEachEnv(t, func(t *testing.T, env *testEnv) {
kvc := env.kvc pc := env.pc
// Request A: user snapshot at offset 5, then generate. // Request A: user snapshot at offset 5, then generate.
simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5}, []int32{10, 11}, 5) simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10, 11}, 5)
assertUserNodeExists(t, kvc, "after A") assertUserNodeExists(t, pc, "after A")
// Request B: extends A's prefix. The user node at offset 5 should // Request B: extends A's prefix. The user node at offset 5 should
// force tokens into a child rather than extending in-place. // force tokens into a child rather than extending in-place.
simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5, 10, 11, 20, 21}, nil) 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, 21})
assertUserNodeExists(t, kvc, "after B") assertUserNodeExists(t, pc, "after B")
// Request C: diverge from the user node. // Request C: diverge from the user node.
simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5, 30, 31}, []int32{40}) simulateRequest(t, pc, []int32{1, 2, 3, 4, 5, 30, 31}, []int32{40})
// Request D: switch back to A's branch — user snapshot still restorable. // Request D: switch back to A's branch — user snapshot still restorable.
simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5, 10, 11, 20, 21, 50}, nil) 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, 50})
checkTrieInvariants(t, kvc.root) checkTrieInvariants(t, pc.root)
}) })
} }
@@ -877,22 +877,22 @@ func TestUserSnapshotPreservesRestorePoint(t *testing.T) {
// a user-marked parent node is not auto-merged with its remaining single child. // a user-marked parent node is not auto-merged with its remaining single child.
func TestUserSnapshotResistsAutoMerge(t *testing.T) { func TestUserSnapshotResistsAutoMerge(t *testing.T) {
forEachEnv(t, func(t *testing.T, env *testEnv) { forEachEnv(t, func(t *testing.T, env *testEnv) {
kvc := env.kvc pc := env.pc
// Request A: user snapshot at offset 3, then continue to offset 5. // Request A: user snapshot at offset 3, then continue to offset 5.
simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5}, []int32{10}, 3) simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10}, 3)
// Request B: diverges at the user node, creating a second child. // Request B: diverges at the user node, creating a second child.
simulateRequest(t, kvc, []int32{1, 2, 3, 6, 7}, []int32{20}) simulateRequest(t, pc, []int32{1, 2, 3, 6, 7}, []int32{20})
userNode := findUserNode(t, kvc) userNode := findUserNode(t, pc)
if len(userNode.children) != 2 { if len(userNode.children) != 2 {
t.Fatalf("user node children = %d, want 2", len(userNode.children)) t.Fatalf("user node children = %d, want 2", len(userNode.children))
} }
// Inflate snapshot sizes and evict. The non-active branch should be // Inflate snapshot sizes and evict. The non-active branch should be
// evicted, leaving the user node with one child. // evicted, leaving the user node with one child.
walkNodes(kvc.root, func(n *trieNode) bool { walkNodes(pc.root, func(n *trieNode) bool {
if !n.hasSnapshots() { if !n.hasSnapshots() {
return true return true
} }
@@ -902,15 +902,15 @@ func TestUserSnapshotResistsAutoMerge(t *testing.T) {
snaps[i] = &fakeSnapshot{byteSize: 5 * 1024 * 1024 * 1024} snaps[i] = &fakeSnapshot{byteSize: 5 * 1024 * 1024 * 1024}
} }
} }
n.setSnapshots(snaps, &kvc.pagedOutBytes) n.setSnapshots(snaps, &pc.pagedOutBytes)
return true return true
}) })
kvc.enforceEvictionPolicy() pc.enforceEvictionPolicy()
// The user node should still exist (not auto-merged) even with one child. // The user node should still exist (not auto-merged) even with one child.
assertUserNodeExists(t, kvc, "after eviction") assertUserNodeExists(t, pc, "after eviction")
checkTrieInvariants(t, kvc.root) checkTrieInvariants(t, pc.root)
}) })
} }
@@ -920,26 +920,26 @@ func TestUserSnapshotResistsAutoMerge(t *testing.T) {
// than materialized as a trie node claiming tokens the cache never wrote. // than materialized as a trie node claiming tokens the cache never wrote.
func TestSnapshotBeyondPrefillSkipped(t *testing.T) { func TestSnapshotBeyondPrefillSkipped(t *testing.T) {
forEachEnv(t, func(t *testing.T, env *testEnv) { forEachEnv(t, func(t *testing.T, env *testEnv) {
kvc := env.kvc pc := env.pc
inputs := []int32{1, 2, 3, 4, 5} inputs := []int32{1, 2, 3, 4, 5}
session := kvc.begin(nil, inputs) session := pc.begin(nil, inputs)
// Request a reachable snapshot at 3 and one at len(inputs), which a // Request a reachable snapshot at 3 and one at len(inputs), which a
// prefill that stops one token short never crosses. // prefill that stops one token short never crosses.
session.schedulePrefillSnapshots([]int{3, len(inputs)}) session.schedulePrefillSnapshots([]int{3, len(inputs)})
// Prefill writes all but the final token (mirrors total-processed > 1). // Prefill writes all but the final token (mirrors total-processed > 1).
feedAll(kvc.caches, inputs[kvc.minCacheOffset():len(inputs)-1]) feedAll(pc.caches, inputs[pc.minCacheOffset():len(inputs)-1])
session.attachPrefillSnapshots() session.attachPrefillSnapshots()
// The reachable offset became a node; the unreached one did not. // The reachable offset became a node; the unreached one did not.
if !nodeExistsAtOffset(kvc.root, 3) { if !nodeExistsAtOffset(pc.root, 3) {
t.Errorf("no trie node at reached offset 3") t.Errorf("no trie node at reached offset 3")
} }
if nodeExistsAtOffset(kvc.root, len(inputs)) { if nodeExistsAtOffset(pc.root, len(inputs)) {
t.Errorf("trie node materialized at unreached offset %d", len(inputs)) t.Errorf("trie node materialized at unreached offset %d", len(inputs))
} }
checkTrieInvariants(t, kvc.root) checkTrieInvariants(t, pc.root)
}) })
} }
@@ -950,20 +950,20 @@ func TestSnapshotBeyondPrefillSkipped(t *testing.T) {
// leaking the snapshots (caught by checkSnapshotLeaks in the env cleanup). // leaking the snapshots (caught by checkSnapshotLeaks in the env cleanup).
func TestPrefillSnapshotsDiscardedOnCancel(t *testing.T) { func TestPrefillSnapshotsDiscardedOnCancel(t *testing.T) {
forEachEnv(t, func(t *testing.T, env *testEnv) { forEachEnv(t, func(t *testing.T, env *testEnv) {
kvc := env.kvc pc := env.pc
inputs := []int32{1, 2, 3, 4, 5} inputs := []int32{1, 2, 3, 4, 5}
session := kvc.begin(nil, inputs) session := pc.begin(nil, inputs)
session.schedulePrefillSnapshots([]int{3}) session.schedulePrefillSnapshots([]int{3})
// Cross offset 3 so the caches capture it, then close the session as a // Cross offset 3 so the caches capture it, then close the session as a
// canceled prefill would, before the captures are attached to the trie. // canceled prefill would, before the captures are attached to the trie.
feedAll(kvc.caches, inputs[kvc.minCacheOffset():3]) feedAll(pc.caches, inputs[pc.minCacheOffset():3])
session.close() session.close()
// close advances the trie over the committed tokens, but the abandoned // close advances the trie over the committed tokens, but the abandoned
// captures must not be attached as snapshots to any node. // captures must not be attached as snapshots to any node.
walkNodes(kvc.root, func(n *trieNode) bool { walkNodes(pc.root, func(n *trieNode) bool {
if n != kvc.root && n.hasSnapshots() { if n != pc.root && n.hasSnapshots() {
t.Errorf("abandoned capture attached as snapshot at offset %d", n.endOffset) t.Errorf("abandoned capture attached as snapshot at offset %d", n.endOffset)
} }
return true return true
@@ -971,9 +971,9 @@ func TestPrefillSnapshotsDiscardedOnCancel(t *testing.T) {
// A second request re-prepares snapshots on the same caches: if the // A second request re-prepares snapshots on the same caches: if the
// discarded ones were not closed, prepare() orphans them here. // discarded ones were not closed, prepare() orphans them here.
simulateRequest(t, kvc, inputs, nil, 4) simulateRequest(t, pc, inputs, nil, 4)
checkTrieInvariants(t, kvc.root) checkTrieInvariants(t, pc.root)
}) })
} }
@@ -988,10 +988,10 @@ func nodeExistsAtOffset(root *trieNode, offset int) bool {
return found return found
} }
func findUserNode(t *testing.T, kvc *kvCache) *trieNode { func findUserNode(t *testing.T, pc *prefixCache) *trieNode {
t.Helper() t.Helper()
var found *trieNode var found *trieNode
walkNodes(kvc.root, func(n *trieNode) bool { walkNodes(pc.root, func(n *trieNode) bool {
if n.user { if n.user {
found = n found = n
} }
@@ -1003,10 +1003,10 @@ func findUserNode(t *testing.T, kvc *kvCache) *trieNode {
return found return found
} }
func assertUserNodeExists(t *testing.T, kvc *kvCache, label string) { func assertUserNodeExists(t *testing.T, pc *prefixCache, label string) {
t.Helper() t.Helper()
var exists bool var exists bool
walkNodes(kvc.root, func(n *trieNode) bool { walkNodes(pc.root, func(n *trieNode) bool {
if n.user { if n.user {
exists = true exists = true
} }
@@ -1023,21 +1023,21 @@ func assertUserNodeExists(t *testing.T, kvc *kvCache, label string) {
// non-rewindable caches. // non-rewindable caches.
func TestBranchSwitchRestoresCorrectState(t *testing.T) { func TestBranchSwitchRestoresCorrectState(t *testing.T) {
forEachEnv(t, func(t *testing.T, env *testEnv) { forEachEnv(t, func(t *testing.T, env *testEnv) {
kvc := env.kvc pc := env.pc
// Request A: [1,2,3,4,5] + generate [10,11] // Request A: [1,2,3,4,5] + generate [10,11]
simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5}, []int32{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, 11})
// Request B: [1,2,3,6,7] — diverges at token 4 // Request B: [1,2,3,6,7] — diverges at token 4
simulateRequest(t, kvc, []int32{1, 2, 3, 6, 7}, []int32{12, 13}) 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, 13})
// Request C: switch back to A's branch [1,2,3,4,5,10,11,20] // Request C: switch back to A's branch [1,2,3,4,5,10,11,20]
simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5, 10, 11, 20}, nil) 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, 20})
checkTrieInvariants(t, kvc.root) checkTrieInvariants(t, pc.root)
}) })
} }
@@ -1046,18 +1046,18 @@ func TestBranchSwitchRestoresCorrectState(t *testing.T) {
// refreshed, allowing them to age out and collapse. // refreshed, allowing them to age out and collapse.
func TestLRUOnlyUpdatesUsedNodes(t *testing.T) { func TestLRUOnlyUpdatesUsedNodes(t *testing.T) {
forEachEnv(t, func(t *testing.T, env *testEnv) { forEachEnv(t, func(t *testing.T, env *testEnv) {
kvc := env.kvc pc := env.pc
// Request A: creates path [1,2,3,4,5] + generate [10,11] // Request A: creates path [1,2,3,4,5] + generate [10,11]
simulateRequest(t, kvc, []int32{1, 2, 3, 4, 5}, []int32{10, 11}) simulateRequest(t, pc, []int32{1, 2, 3, 4, 5}, []int32{10, 11})
// Request B: diverges at token 4, creating a branch point at offset 3 // Request B: diverges at token 4, creating a branch point at offset 3
// with a split snapshot. // with a split snapshot.
simulateRequest(t, kvc, []int32{1, 2, 3, 6, 7}, []int32{20, 21}) simulateRequest(t, pc, []int32{1, 2, 3, 6, 7}, []int32{20, 21})
// Set all lastUsed to a known old time. // Set all lastUsed to a known old time.
oldTime := time.Now().Add(-1 * time.Hour) oldTime := time.Now().Add(-1 * time.Hour)
walkNodes(kvc.root, func(n *trieNode) bool { walkNodes(pc.root, func(n *trieNode) bool {
n.lastUsed = oldTime n.lastUsed = oldTime
return true return true
}) })
@@ -1066,15 +1066,15 @@ func TestLRUOnlyUpdatesUsedNodes(t *testing.T) {
// and extend it. The branch point's snapshot may be paged in // and extend it. The branch point's snapshot may be paged in
// for some cache types but not others. // for some cache types but not others.
beforeRequest := time.Now() beforeRequest := time.Now()
simulateRequest(t, kvc, []int32{1, 2, 3, 6, 7, 20, 21, 30}, nil) simulateRequest(t, pc, []int32{1, 2, 3, 6, 7, 20, 21, 30}, nil)
// The path must have enough depth to exercise intermediate nodes. // The path must have enough depth to exercise intermediate nodes.
if len(kvc.activePath) < 3 { if len(pc.activePath) < 3 {
t.Fatalf("activePath too short to test intermediate nodes: got %d nodes", len(kvc.activePath)) t.Fatalf("activePath too short to test intermediate nodes: got %d nodes", len(pc.activePath))
} }
// The frontier (deepest node on the active path) must be updated. // The frontier (deepest node on the active path) must be updated.
frontier := kvc.activePath[len(kvc.activePath)-1] frontier := pc.activePath[len(pc.activePath)-1]
if frontier.lastUsed.Before(beforeRequest) { if frontier.lastUsed.Before(beforeRequest) {
t.Errorf("frontier lastUsed was not updated: got %v, want >= %v", t.Errorf("frontier lastUsed was not updated: got %v, want >= %v",
frontier.lastUsed, beforeRequest) frontier.lastUsed, beforeRequest)
@@ -1082,14 +1082,14 @@ func TestLRUOnlyUpdatesUsedNodes(t *testing.T) {
// Every non-frontier node on the active path (including root) // Every non-frontier node on the active path (including root)
// should retain its old lastUsed — only the frontier gets refreshed. // should retain its old lastUsed — only the frontier gets refreshed.
for i, node := range kvc.activePath[:len(kvc.activePath)-1] { for i, node := range pc.activePath[:len(pc.activePath)-1] {
if !node.lastUsed.Before(beforeRequest) { if !node.lastUsed.Before(beforeRequest) {
t.Errorf("activePath[%d] (endOffset=%d) lastUsed was refreshed: got %v, want < %v", t.Errorf("activePath[%d] (endOffset=%d) lastUsed was refreshed: got %v, want < %v",
i, node.endOffset, node.lastUsed, beforeRequest) i, node.endOffset, node.lastUsed, beforeRequest)
} }
} }
checkTrieInvariants(t, kvc.root) checkTrieInvariants(t, pc.root)
}) })
} }
@@ -1098,24 +1098,24 @@ func TestLRUOnlyUpdatesUsedNodes(t *testing.T) {
// state), the trie's pagedOutBytes counter picks up the delta via the // state), the trie's pagedOutBytes counter picks up the delta via the
// installed materialize hook. // installed materialize hook.
func TestPagedOutBytesUpdatesOnMaterialize(t *testing.T) { func TestPagedOutBytesUpdatesOnMaterialize(t *testing.T) {
kvc := &kvCache{} pc := &prefixCache{}
kvc.ensureRoot() pc.ensureRoot()
node := &trieNode{parent: kvc.root, tokens: []int32{1, 2, 3}, endOffset: 3} node := &trieNode{parent: pc.root, tokens: []int32{1, 2, 3}, endOffset: 3}
kvc.root.children = append(kvc.root.children, node) pc.root.children = append(pc.root.children, node)
snap := &fakeSnapshot{from: 0, to: 3, byteSize: 0} snap := &fakeSnapshot{from: 0, to: 3, byteSize: 0}
node.setSnapshots([]cache.Snapshot{snap}, &kvc.pagedOutBytes) node.setSnapshots([]cache.Snapshot{snap}, &pc.pagedOutBytes)
if kvc.pagedOutBytes != 0 { if pc.pagedOutBytes != 0 {
t.Fatalf("pagedOutBytes after install = %d, want 0 (lazy snapshot)", kvc.pagedOutBytes) t.Fatalf("pagedOutBytes after install = %d, want 0 (lazy snapshot)", pc.pagedOutBytes)
} }
const materialized = 1 << 20 const materialized = 1 << 20
snap.materialize(materialized) snap.materialize(materialized)
if kvc.pagedOutBytes != materialized { if pc.pagedOutBytes != materialized {
t.Fatalf("pagedOutBytes after materialize = %d, want %d", kvc.pagedOutBytes, materialized) t.Fatalf("pagedOutBytes after materialize = %d, want %d", pc.pagedOutBytes, materialized)
} }
} }
@@ -1124,28 +1124,28 @@ func TestPagedOutBytesUpdatesOnMaterialize(t *testing.T) {
// materialize. Without detach, a Split/Merge that folds an old snapshot // materialize. Without detach, a Split/Merge that folds an old snapshot
// elsewhere would double-count its bytes if it copied out afterward. // elsewhere would double-count its bytes if it copied out afterward.
func TestSwapSnapshotsDetachesHook(t *testing.T) { func TestSwapSnapshotsDetachesHook(t *testing.T) {
kvc := &kvCache{} pc := &prefixCache{}
kvc.ensureRoot() pc.ensureRoot()
node := &trieNode{parent: kvc.root, tokens: []int32{1, 2, 3}, endOffset: 3} node := &trieNode{parent: pc.root, tokens: []int32{1, 2, 3}, endOffset: 3}
kvc.root.children = append(kvc.root.children, node) pc.root.children = append(pc.root.children, node)
snap := &fakeSnapshot{from: 0, to: 3, byteSize: 0} snap := &fakeSnapshot{from: 0, to: 3, byteSize: 0}
node.setSnapshots([]cache.Snapshot{snap}, &kvc.pagedOutBytes) node.setSnapshots([]cache.Snapshot{snap}, &pc.pagedOutBytes)
replacement := &fakeSnapshot{from: 0, to: 3, byteSize: 0} replacement := &fakeSnapshot{from: 0, to: 3, byteSize: 0}
old := node.swapSnapshots([]cache.Snapshot{replacement}, &kvc.pagedOutBytes) old := node.swapSnapshots([]cache.Snapshot{replacement}, &pc.pagedOutBytes)
if len(old) != 1 || old[0] != snap { if len(old) != 1 || old[0] != snap {
t.Fatalf("swapSnapshots returned %v, want the original snap", old) t.Fatalf("swapSnapshots returned %v, want the original snap", old)
} }
snap.materialize(1 << 20) snap.materialize(1 << 20)
if kvc.pagedOutBytes != 0 { if pc.pagedOutBytes != 0 {
t.Fatalf("pagedOutBytes after detached materialize = %d, want 0", kvc.pagedOutBytes) t.Fatalf("pagedOutBytes after detached materialize = %d, want 0", pc.pagedOutBytes)
} }
replacement.materialize(2 << 20) replacement.materialize(2 << 20)
if kvc.pagedOutBytes != 2<<20 { if pc.pagedOutBytes != 2<<20 {
t.Fatalf("pagedOutBytes after replacement materialize = %d, want %d", kvc.pagedOutBytes, 2<<20) t.Fatalf("pagedOutBytes after replacement materialize = %d, want %d", pc.pagedOutBytes, 2<<20)
} }
} }
x/mlxrunner/runner.go
+1 -1
View File
@@ -38,7 +38,7 @@ type Runner struct {
Tokenizer *tokenizer.Tokenizer Tokenizer *tokenizer.Tokenizer
Requests chan Request Requests chan Request
Sampler *sample.Sampler Sampler *sample.Sampler
cache kvCache cache prefixCache
contextLength int contextLength int
mlxThread *mlxthread.Thread mlxThread *mlxthread.Thread
// spec is the speculative-decoding subsystem. Nil when the model ships no // spec is the speculative-decoding subsystem. Nil when the model ships no