diff --git a/x/mlxrunner/cache.go b/x/mlxrunner/cache.go index c1c53b668..4fff61581 100644 --- a/x/mlxrunner/cache.go +++ b/x/mlxrunner/cache.go @@ -56,8 +56,8 @@ type cacheSession struct { remaining []int32 // pendingSnapshots lists offsets where snapshots should be captured - // during prefill, sorted by offset. Entries are consumed as the - // cache advances past them. + // during prefill, sorted by offset. Entries are scheduled on the caches + // before prefill and drained or discarded after. pendingSnapshots []pendingSnapshot } @@ -251,89 +251,151 @@ pageIn: } } -// requestSnapshot schedules a user snapshot at the given absolute token -// offset. The snapshot will be captured during prefill when the cache -// reaches this offset. -func (s *cacheSession) requestSnapshot(offset int) { - baseOffset := len(s.inputs) - len(s.remaining) - if offset <= baseOffset || offset > len(s.inputs) { - return - } - // Deduplicate: if this offset already exists, upgrade to user. - for i := range s.pendingSnapshots { - if s.pendingSnapshots[i].offset == offset { - s.pendingSnapshots[i].user = true - return +// 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 +// offsets upgraded to user so eviction preserves them. +// +// Offsets at or before the current cache position, or past the end of the +// prompt, are dropped: callers only request offsets ahead of the prefill base, +// so this is a defensive guard. +func (s *cacheSession) schedulePrefillSnapshots(offsets []int) { + c := s.cache + base := c.minCacheOffset() + for _, offset := range offsets { + if offset <= base || offset > len(s.inputs) { + continue + } + // Deduplicate: if this offset already exists, upgrade to user. + found := false + for i := range s.pendingSnapshots { + if s.pendingSnapshots[i].offset == offset { + s.pendingSnapshots[i].user = true + found = true + break + } + } + if !found { + s.pendingSnapshots = append(s.pendingSnapshots, pendingSnapshot{offset: offset, user: true}) } } - s.pendingSnapshots = append(s.pendingSnapshots, pendingSnapshot{offset: offset, user: true}) slices.SortFunc(s.pendingSnapshots, func(a, b pendingSnapshot) int { return a.offset - b.offset }) -} -// nextPendingSnapshot returns the offset of the next pending snapshot, -// or 0 if there are none. -func (s *cacheSession) nextPendingSnapshot() int { if len(s.pendingSnapshots) == 0 { - return 0 - } - return s.pendingSnapshots[0].offset -} - -// snapshot creates a snapshot at the current cache position. It determines -// whether this is a user snapshot by consuming pending entries whose offset -// has been reached. -func (s *cacheSession) snapshot() { - c := s.cache - cacheOffset := c.minCacheOffset() - if cacheOffset <= 0 { return } - // Consume pending snapshots up to the current offset and derive - // the user flag from them. - user := false - for len(s.pendingSnapshots) > 0 && cacheOffset >= s.pendingSnapshots[0].offset { - if s.pendingSnapshots[0].user { - user = true + prepared := make([]int, len(s.pendingSnapshots)) + for i, p := range s.pendingSnapshots { + prepared[i] = p.offset + } + for _, kv := range c.caches { + if kv != nil { + kv.PrepareSnapshots(prepared) } - s.pendingSnapshots = s.pendingSnapshots[1:] + } +} + +// discardPrefillSnapshots drains and closes the snapshots scheduled by +// schedulePrefillSnapshots without attaching them to the trie, releasing their +// pinned/lazy state. It is a no-op once attachPrefillSnapshots has drained the +// schedule, so close can call it unconditionally to clean up an abandoned +// prefill. +func (s *cacheSession) discardPrefillSnapshots() { + if len(s.pendingSnapshots) == 0 { + return + } + s.pendingSnapshots = nil + + for _, kv := range s.cache.caches { + if kv == nil { + continue + } + for _, snap := range kv.TakeSnapshots() { + if snap != nil { + snap.Close() + } + } + } +} + +// attachPrefillSnapshots collects the snapshots captured during prefill and +// attaches them to the trie, materializing a node at each requested offset. +// Pending offsets are ascending and were scheduled in the same order, so the +// snapshots each cache returns line up with them. The trie frontier is +// advanced to each offset in turn, so its node edges [prev, offset) match the +// edge-local ranges the caches captured. +func (s *cacheSession) attachPrefillSnapshots() { + if len(s.pendingSnapshots) == 0 { + return } - // The last node in activePath is the frontier where caches are advancing. - // cacheOffset is always >= its endOffset: begin() restores caches to this - // boundary and prefill advances monotonically forward. - frontier := c.activePath[len(c.activePath)-1] + c := s.cache + pending := s.pendingSnapshots + s.pendingSnapshots = nil - // If the frontier already ends at cacheOffset, just ensure it has snapshots. - if frontier.endOffset == cacheOffset { - if user { + // Drain each cache's captures (one per pending offset, in order) into + // per-offset rows. + rows := make([][]cache.Snapshot, len(pending)) + for i := range rows { + rows[i] = make([]cache.Snapshot, len(c.caches)) + } + for j, kv := range c.caches { + if kv == nil { + continue + } + taken := kv.TakeSnapshots() + for i := range pending { + if i < len(taken) { + rows[i][j] = taken[i] + } + } + } + + // Prefill leaves one token unprocessed for decode seeding, so an offset + // at or past the live cache position was never crossed by a write and has + // 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...) + for i, p := range pending { + if p.offset > reached { + // Never crossed by a write, so the row is nil; close any entry + // defensively in case a cache captured one anyway. + for _, snap := range rows[i] { + if snap != nil { + snap.Close() + } + } + continue + } + frontier := c.activePath[len(c.activePath)-1] + if frontier.endOffset < p.offset { + edgeTokens := stored[frontier.endOffset:p.offset] + frontier = c.advancePath(frontier, edgeTokens, p.offset) + } + if p.user { frontier.user = true } - if !frontier.hasAllSnapshots() { - s.attachSnapshots(frontier, cacheOffset) - } - return + s.attachCapturedSnapshots(frontier, rows[i]) } +} - if frontier.endOffset > cacheOffset { - slog.Warn("snapshot skipped: cacheOffset is behind frontier", "cacheOffset", cacheOffset, "frontierEndOffset", frontier.endOffset) - return - } - - // Advance the trie to cacheOffset — find or create a node there. - edgeTokens := append(s.inputs, s.outputs...)[frontier.endOffset:cacheOffset] - frontier = c.advancePath(frontier, edgeTokens, cacheOffset) - - // Attach fresh snapshots from the live caches. Always use fresh - // snapshots even if the node already has some (e.g. from splitNode's - // Cache.Split which may be incomplete for non-splittable caches - // like RecurrentCache). - if user { - frontier.user = true - } - s.attachSnapshots(frontier, cacheOffset) +// attachCapturedSnapshots stores pre-captured snapshots on a trie node. Unlike +// taking a fresh Snapshot from the live cache, this works for an interior node +// whose offset the live cache has already advanced past: the snapshots come +// from the capture scheduled earlier, not from the cache's current state. The +// node takes ownership of the snapshots (TakeSnapshots already transferred it). +func (s *cacheSession) attachCapturedSnapshots(node *trieNode, snaps []cache.Snapshot) { + c := s.cache + node.setSnapshots(snaps, &c.pagedOutBytes) + node.lastUsed = time.Now() + slog.Debug("created snapshot", "offset", node.endOffset) + c.enforceEvictionPolicy() } // advancePath advances the active path from the current frontier by matching @@ -375,33 +437,6 @@ func (c *kvCache) advancePath(frontier *trieNode, tokens []int32, endOffset int) return dest } -// attachSnapshots attaches cache snapshots to a trie node at the given offset. -// The node must be on the active path (and thus protected from eviction; -// lastUsed is updated in close()). All non-nil caches must be at the same -// offset (cacheOffset); a mismatch indicates a bug in the caller. -func (s *cacheSession) attachSnapshots(node *trieNode, cacheOffset int) { - c := s.cache - - if c.activePath[len(c.activePath)-1] != node { - slog.Warn("attachSnapshots skipped: node is not the active frontier", "nodeEndOffset", node.endOffset) - return - } - - snaps := make([]cache.Snapshot, len(c.caches)) - for i, kv := range c.caches { - if kv != nil { - if kv.Offset() != cacheOffset { - panic(fmt.Sprintf("attachSnapshots: cache offset mismatch layer %d: expected %d, got %d", i, cacheOffset, kv.Offset())) - } - snaps[i] = kv.Snapshot(node.startOffset()) - } - } - node.setSnapshots(snaps, &c.pagedOutBytes) - node.lastUsed = time.Now() - slog.Debug("created snapshot", "offset", cacheOffset) - c.enforceEvictionPolicy() -} - // freeAll releases all cache layers. func (c *kvCache) freeAll() { for _, kv := range c.caches { @@ -428,6 +463,14 @@ func (c *kvCache) minCacheOffset() int { // close saves the token state if the forward pass ran. func (s *cacheSession) close() { + // Release any prefill snapshots the session scheduled but never attached to + // the trie. A successful prefill drains them in attachPrefillSnapshots (so + // this is a no-op then); an abandoned one (e.g. cancellation between + // schedule and attach) leaves them in the caches, where the next request's + // PrepareSnapshots would overwrite the schedule without closing them, + // leaking the pinned/lazy snapshots and their VRAM. + s.discardPrefillSnapshots() + offset := s.cache.minCacheOffset() if offset <= 0 { return diff --git a/x/mlxrunner/cache_test.go b/x/mlxrunner/cache_test.go index acd9dab4d..a1a1346bf 100644 --- a/x/mlxrunner/cache_test.go +++ b/x/mlxrunner/cache_test.go @@ -467,7 +467,7 @@ func (e *testEnv) assertAllTokens(t *testing.T, label string, expected []int32) } // simulateRequest mirrors the production pipeline lifecycle: -// begin -> prefill with snapshot(false) at branch points -> generate -> close +// begin -> schedule snapshots -> prefill in one pass -> attach snapshots -> generate -> close type requestResult struct { remaining []int32 @@ -480,9 +480,10 @@ func simulateRequest(t *testing.T, kvc *kvCache, inputs, generated []int32, user t.Helper() session := kvc.begin(nil, inputs) + var snapshotOffsets []int for _, at := range userSnapshotAt { if at > 0 { - session.requestSnapshot(at) + snapshotOffsets = append(snapshotOffsets, at) } } @@ -496,26 +497,14 @@ func simulateRequest(t *testing.T, kvc *kvCache, inputs, generated []int32, user baseOffset := kvc.minCacheOffset() remaining := inputs[baseOffset:] - // Prefill: feed tokens, pausing at each pending snapshot. - for len(session.pendingSnapshots) > 0 { - sp := session.pendingSnapshots[0] - count := sp.offset - baseOffset - if count > len(remaining) { - break - } - if count > 0 { - feedAll(kvc.caches, remaining[:count]) - remaining = remaining[count:] - baseOffset = sp.offset - } - assertCacheOffsetAlignment(t, kvc, "at snapshot point") - session.snapshot() - } - - // Feed rest of input tokens. + // 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. + session.schedulePrefillSnapshots(snapshotOffsets) if len(remaining) > 0 { feedAll(kvc.caches, remaining) } + session.attachPrefillSnapshots() assertCacheOffsetAlignment(t, kvc, "after prefill") @@ -925,6 +914,80 @@ func TestUserSnapshotResistsAutoMerge(t *testing.T) { }) } +// TestSnapshotBeyondPrefillSkipped verifies that a snapshot scheduled at an +// offset the prefill never reaches (prefill leaves one token for decode +// seeding, so the last token is never written during prefill) is dropped rather +// than materialized as a trie node claiming tokens the cache never wrote. +func TestSnapshotBeyondPrefillSkipped(t *testing.T) { + forEachEnv(t, func(t *testing.T, env *testEnv) { + kvc := env.kvc + inputs := []int32{1, 2, 3, 4, 5} + + session := kvc.begin(nil, inputs) + // Request a reachable snapshot at 3 and one at len(inputs), which a + // prefill that stops one token short never crosses. + session.schedulePrefillSnapshots([]int{3, len(inputs)}) + // Prefill writes all but the final token (mirrors total-processed > 1). + feedAll(kvc.caches, inputs[kvc.minCacheOffset():len(inputs)-1]) + session.attachPrefillSnapshots() + + // The reachable offset became a node; the unreached one did not. + if !nodeExistsAtOffset(kvc.root, 3) { + t.Errorf("no trie node at reached offset 3") + } + if nodeExistsAtOffset(kvc.root, len(inputs)) { + t.Errorf("trie node materialized at unreached offset %d", len(inputs)) + } + + checkTrieInvariants(t, kvc.root) + }) +} + +// TestPrefillSnapshotsDiscardedOnCancel mirrors a prefill canceled after the +// caches captured interior snapshots but before attachPrefillSnapshots ran. The +// abandoned captures must be released when the session closes; otherwise the +// next request's PrepareSnapshots overwrites the schedule without closing them, +// leaking the snapshots (caught by checkSnapshotLeaks in the env cleanup). +func TestPrefillSnapshotsDiscardedOnCancel(t *testing.T) { + forEachEnv(t, func(t *testing.T, env *testEnv) { + kvc := env.kvc + inputs := []int32{1, 2, 3, 4, 5} + + session := kvc.begin(nil, inputs) + session.schedulePrefillSnapshots([]int{3}) + // 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. + feedAll(kvc.caches, inputs[kvc.minCacheOffset():3]) + session.close() + + // close advances the trie over the committed tokens, but the abandoned + // captures must not be attached as snapshots to any node. + walkNodes(kvc.root, func(n *trieNode) bool { + if n != kvc.root && n.hasSnapshots() { + t.Errorf("abandoned capture attached as snapshot at offset %d", n.endOffset) + } + return true + }) + + // A second request re-prepares snapshots on the same caches: if the + // discarded ones were not closed, prepare() orphans them here. + simulateRequest(t, kvc, inputs, nil, 4) + + checkTrieInvariants(t, kvc.root) + }) +} + +func nodeExistsAtOffset(root *trieNode, offset int) bool { + var found bool + walkNodes(root, func(n *trieNode) bool { + if n.endOffset == offset && n != root { + found = true + } + return true + }) + return found +} + func findUserNode(t *testing.T, kvc *kvCache) *trieNode { t.Helper() var found *trieNode diff --git a/x/mlxrunner/pipeline.go b/x/mlxrunner/pipeline.go index 777653c04..0f69374e6 100644 --- a/x/mlxrunner/pipeline.go +++ b/x/mlxrunner/pipeline.go @@ -84,13 +84,14 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er // prompt so that long prompts can be partially restored and // thinking/generation can be retried without full reprocessing. const snapshotInterval = 8192 + var snapshotOffsets []int for offset := snapshotInterval; offset < len(inputs); offset += snapshotInterval { - session.requestSnapshot(offset) + snapshotOffsets = append(snapshotOffsets, offset) } const preThinking = 4 if end := len(inputs) - preThinking; end > 0 { - session.requestSnapshot(end) + snapshotOffsets = append(snapshotOffsets, end) } materializeCaches := func() { @@ -104,6 +105,8 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er mlx.Eval(state...) } + session.schedulePrefillSnapshots(snapshotOffsets) + now := time.Now() total, processed := len(tokens), 0 position := len(inputs) - len(tokens) @@ -114,15 +117,6 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er n := min(prefillChunk, total-processed-1) - // If there's a pending snapshot, split the batch so we can - // capture it at the exact offset. - if snapOffset := session.nextPendingSnapshot(); snapOffset > 0 { - tokensUntilSnapshot := snapOffset - position - if tokensUntilSnapshot > 0 && tokensUntilSnapshot < n { - n = tokensUntilSnapshot - } - } - r.Model.Forward(&batch.Batch{ InputIDs: mlx.FromValues(tokens[processed:processed+n], 1, n), SeqOffsets: []int32{int32(position)}, @@ -135,16 +129,12 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er slog.Info("Prompt processing progress", "processed", processed, "total", total) logutil.TraceContext(ctx, "mlx prompt forward", "processed", processed, "total", total, "tokens", n, "memory", mlx.Memory{}) - // Create snapshot if we've reached a pending offset. - if snapOffset := session.nextPendingSnapshot(); snapOffset > 0 { - if position >= snapOffset { - session.snapshot() - } - } - mlx.ClearCache() } + // Attach the snapshots captured during prefill to the trie. + session.attachPrefillSnapshots() + // Register the sampler after prefill completes. r.Sampler.Add(pipelineSlot, request.SamplerOpts, inputs) if r.useGreedyMTP(request.SamplerOpts) {