diff --git a/x/mlxrunner/prefix_cache.go b/x/mlxrunner/prefix_cache.go index f48f3e2f5..1adbb2376 100644 --- a/x/mlxrunner/prefix_cache.go +++ b/x/mlxrunner/prefix_cache.go @@ -63,7 +63,7 @@ type cacheSession struct { // pendingSnapshots lists offsets where snapshots should be captured // during prefill, sorted by offset. Entries are scheduled on the caches - // before prefill and drained or discarded after. + // before prefill and drained when the captures are attached. pendingSnapshots []pendingSnapshot } @@ -326,29 +326,6 @@ func (s *cacheSession) schedulePrefillSnapshots(offsets []int) { } } -// 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 @@ -538,13 +515,10 @@ func (c *prefixCache) 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() + // A cancelled prefill never reaches the success-path attach; attaching + // here keeps its crossed captures for the retry and drains the schedule + // PrepareSnapshots would otherwise overwrite, leaking them. + s.attachPrefillSnapshots() offset := s.cache.minCacheOffset() if offset <= 0 { diff --git a/x/mlxrunner/prefix_cache_scenario_test.go b/x/mlxrunner/prefix_cache_scenario_test.go index d0ee003b7..4f560c6db 100644 --- a/x/mlxrunner/prefix_cache_scenario_test.go +++ b/x/mlxrunner/prefix_cache_scenario_test.go @@ -265,3 +265,126 @@ func TestScenarioConversationTurns(t *testing.T) { } checkSnapshotCoverage(t, e.pc, e.kvLayers()) } + +// runCancelledPrefill mirrors a prefill cancelled by a client timeout at +// (roughly) cancelAt prompt tokens: begin, schedule, feed whole chunks until +// cancelAt with the drafter's pairs lagging one token, settle the drafter with +// the next prompt token, then close, which attaches the crossed captures. +func (e *hybridEnv) runCancelledPrefill(t *testing.T, inputs []int32, chunk, cancelAt int) { + t.Helper() + pc := e.pc + session := pc.begin(inputs, nil) + session.schedulePrefillSnapshots(periodic(len(inputs))) + + pos := pc.minCacheOffset() + for pos < cancelAt && pos < len(inputs)-1 { + n := min(chunk, len(inputs)-1-pos) + for _, c := range pc.caches { + if c == cache.Cache(e.draft) { + continue + } + if fc, ok := c.(feedableCache); ok { + fc.feed(inputs[pos : pos+n]) + } + } + if d := e.draft.Offset(); d < pos+n-1 { + e.draft.feed(inputs[d : pos+n-1]) + } + pos += n + } + + if d := e.draft.Offset(); d < pos { + e.draft.feed(inputs[d:pos]) + } + session.close() +} + +// TestScenarioCancelledPrefills covers the cancellation invariants: a retry +// of a cancelled prefill resumes exactly where the previous attempt stopped, +// and the captures the cancelled attempt crossed become restore points. +func TestScenarioCancelledPrefills(t *testing.T) { + logs := captureWarns(t) + e := newHybridEnv() + ts := &tokenStream{} + begins := &beginLog{t: t, logs: logs} + + prompt := ts.fresh(47) + + // Cancelled after two chunks of 9. The prefill crossed the captures + // scheduled at interval and 2*interval — key offsets 7 and 15 after the + // draft look-ahead shift — and close attached them. + e.runCancelledPrefill(t, prompt, 9, 2*9) + begins.next() + + // A session diverging inside the cancelled span restores to the deepest + // crossed capture: 17 keys match, and the entry at 15 is the closest + // restore point below. + probe := slices.Concat(slices.Clone(prompt[:18]), ts.fresh(10)) + pr := e.pc.begin(probe, nil) + if _, m, c, _ := begins.next(); m != 17 || c != 15 { + t.Errorf("probe into cancelled span: matched=%d cached=%d, want 17/15", m, c) + } + pr.close() + + // The retry resumes exactly at the 18 tokens the first attempt recorded, + // and is cancelled again two chunks deeper. + e.runCancelledPrefill(t, prompt, 9, 2*9+2*9) + if _, m, c, _ := begins.next(); m != 18 || c != 18 { + t.Errorf("first retry: matched=%d cached=%d, want 18/18", m, c) + } + + // The final retry resumes at 36 and completes. + e.runRequest(t, prompt, ts.fresh(6)) + if _, m, c, _ := begins.next(); m != 36 || c != 36 { + t.Errorf("second retry: matched=%d cached=%d, want 36/36", m, c) + } + + if out := logs(); strings.Contains(out, "failed to restore cache") { + t.Errorf("freeAll warn fired:\n%s", out) + } + checkSnapshotCoverage(t, e.pc, e.kvLayers()) +} + +// TestScenarioDivergentCancels interleaves cancelled prefills that diverge +// from a conversation mid-history with completed requests, growing branch +// points above and below the conversation's capture offsets, then re-requests +// the full conversation, which must restore to within the capture cadence. +func TestScenarioDivergentCancels(t *testing.T) { + logs := captureWarns(t) + e := newHybridEnv() + ts := &tokenStream{} + begins := &beginLog{t: t, logs: logs} + + prompt := ts.fresh(24) + gen := ts.fresh(8) + e.runRequest(t, prompt, gen) + stream := slices.Concat(prompt, gen) + begins.next() + + // A cancelled prefill diverging between the second and third captures. + e.runCancelledPrefill(t, slices.Concat(slices.Clone(stream[:2*interval+2]), ts.fresh(20)), 5, 12) + begins.next() + + // A completed turn extends the conversation. + p2 := slices.Concat(stream, ts.fresh(10)) + g2 := ts.fresh(4) + e.runRequest(t, p2, g2) + stream = slices.Concat(p2, g2) + begins.next() + + // A cancelled prefill diverging below the first capture. + e.runCancelledPrefill(t, slices.Concat(slices.Clone(stream[:interval-2]), ts.fresh(16)), 5, 14) + begins.next() + + // Re-requesting the full conversation restores to within the cadence of + // the last completed turn. + e.runRequest(t, slices.Concat(stream, ts.fresh(12)), ts.fresh(5)) + if _, m, c, _ := begins.next(); c < m-e.restoreBound(len(g2), 0) { + t.Errorf("re-request: cached=%d fell more than %d below matched=%d", c, e.restoreBound(len(g2), 0), m) + } + + if out := logs(); strings.Contains(out, "failed to restore cache") { + t.Errorf("freeAll warn fired:\n%s", out) + } + checkSnapshotCoverage(t, e.pc, e.kvLayers()) +} diff --git a/x/mlxrunner/prefix_cache_test.go b/x/mlxrunner/prefix_cache_test.go index e00d46cbc..cb7165b66 100644 --- a/x/mlxrunner/prefix_cache_test.go +++ b/x/mlxrunner/prefix_cache_test.go @@ -56,12 +56,18 @@ func (p *fakePending) take() []cache.Snapshot { func (p *fakePending) feedCapturing(start int, tokens []int32, capture func(from, reached int) cache.Snapshot, advance func([]int32)) { end := start + len(tokens) captureAt := func(reached int) { + fired := false for i, o := range p.offsets { if p.captured[i] == nil && o == reached { p.captured[i] = capture(p.base, reached) + fired = true } } - p.base = reached + // The base advances only when a capture fires, mirroring + // pendingSnapshots.captureReached. + if fired { + p.base = reached + } } if len(p.offsets) == 0 { @@ -1005,12 +1011,13 @@ func TestSnapshotBeyondPrefillSkipped(t *testing.T) { }) } -// 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) { +// TestPrefillSnapshotsKeptOnCancel mirrors a prefill canceled after the caches +// captured interior snapshots but before the success-path attach ran. Closing +// the session attaches the crossed captures so a retry can resume from them, +// and drains the capture schedule; otherwise the next request's +// PrepareSnapshots would overwrite it without closing the captures, leaking +// them (caught by checkSnapshotLeaks in the env cleanup). +func TestPrefillSnapshotsKeptOnCancel(t *testing.T) { forEachEnv(t, func(t *testing.T, env *testEnv) { pc := env.pc inputs := []int32{1, 2, 3, 4, 5} @@ -1018,12 +1025,17 @@ func TestPrefillSnapshotsDiscardedOnCancel(t *testing.T) { session := pc.begin(inputs, nil) 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. + // canceled prefill would, before the success-path attach. feedAll(pc.caches, inputs[pc.minCacheOffset():3]) session.close() + // The crossed capture becomes a restore point for the retry. + if at := 3 - pc.draftLookahead; !nodeExistsAtOffset(pc.root, at) { + t.Errorf("no trie node at capture point %d after cancel", at) + } + // A second request re-prepares snapshots on the same caches: if the - // discarded ones were not closed, prepare() orphans them here. + // pending ones were not drained, prepare() orphans them here. simulateRequest(t, pc, inputs, nil, 5) checkTrieInvariants(t, pc.root)