Repository navigation
Prompt cache is not saved when a request ends on a stop sequence or disconnect (regression from b620ae8 bump) #245
Description
Summary
After the mlx-swift-lm bump to b620ae8 (#243, merged as 57ea94b), the prompt KV cache is saved only on the producer's .info event. A request that ends early, either on a text-level stop sequence or on a client disconnect, never reaches .info, so its prompt is not saved. The next request that shares the prefix re-prefills the whole prompt.
Before the bump, the save ran on the first decode token, so an early stop still kept the prompt.
Where (at 57ea94b)
Sources/SwiftLM/Server.swift:3418: the streaming path saves incase .info, and the comment there says an early exit "never reaches.infoand skips the save".Sources/SwiftLM/Server.swift:3594: the non-streaming path has the same behavior at the stop-sequencebreak.
Impact
- Any request that stops on a stop sequence or a disconnect leaves no prompt cache entry. The next turn of that conversation re-prefills the full history, so the cost grows with conversation length.
- Not verified: whether the built-in stop strings (
<|im_end|>,<end_of_turn>, and others) fire in practice before the EOS token id. If they do, most ordinary turns are affected, not only aborted ones. - Not measured: the time-to-first-token cost.
Suggested direction
Save the prompt on the first decode token again, or add a save on the stop and disconnect paths. The .info move exists to avoid a race with the decode step that is still writing the same KVCacheSimple, so any fix has to keep that race closed.
References
- chore: bump mlx-swift-lm to b620ae8 #243
chore: bump mlx-swift-lm to b620ae8, merged as57ea94b.
Activity
HamedShams commentedon Oct 28, 2026
One more early-exit case for this list, found while working on #249, on main at
ca8f3a2: a turn the model answers with a tool call and nothing else is never saved either.Why it happens: both handlers save at
.infoonly oncefirstTokenhas been cleared (Server.swift:3419streaming,:3595non-streaming), andfirstTokenis cleared only incase .chunk(:3305,:3575). A reply that is only a tool call never produces a chunk, because the parser holds the whole<tool_call>…</tool_call>body and emits a single.toolCall. So when.infoarrives the guard is still set andonPrefillDoneis skipped, even though generation finished normally. This is older than the bump: before2d8a787the save sat inside thefirstTokenblock, which such a turn never entered either.What I saw with the CI model (Qwen2.5-0.5B-Instruct-4bit) and a
get_weathertool:prefill doneand saves nothing;In an agent loop most turns end in a tool call, so this is the common case rather than an edge one.
Possible fix: by the time
.infoarrives, prefill has always completed, so the!firstTokenguard looks unnecessary there and could be dropped. AlternativelyfirstTokencould be cleared incase .toolCallas well. Either one covers this case on its own, independently of the stop-sequence and disconnect cases.fix: save the prompt cache after a tool-call-only turn
Merge pull request #253 from SharpAI/fix/prompt-cache-early-exit