Skip to content

fix: save the prompt cache at the turn boundary so early exits keep it (#245) #253

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
solderzzc merged 1 commit into from
Oct 29, 2026
Merged

fix: save the prompt cache at the turn boundary so early exits keep it (#245) #253

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 31 additions & 4 deletions Sources/SwiftLM/Server.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2799,6 +2799,26 @@ func handleChatCompletion(
context.tokenizer.convertTokenToId("<|turn>"),
context.tokenizer.convertTokenToId("<|im_start|>"),
])
// ── Full-attention prompt cache ──
// Saving once the producer has emitted `.info` misses every request that ends early:
// a stop sequence or a client disconnect leaves the stream before `.info`, so the
// next turn re-prefills the whole history (#245). Snapshot at the turn boundary
// instead, as the paths above do: prefill the prompt but its last token
// synchronously, save, and let generation feed the last token. No decode step has
// started, so nothing is writing the cache while it is read. The saved prefix is one
// token short of the prompt; restore replays that token anyway.
let saveAtBoundary = !skipPromptCache && hybridBoundary == nil && rotatingBoundary == nil
&& draftModelRef == nil && !config.mtp && !config.turboKV && !hasRotatingLayer
&& promptLength > 1
func snapshotPrompt(from start: Int) async throws -> Int {
let boundary = promptLength - 1
guard saveAtBoundary, start < boundary else { return start }
_ = try TokenIterator(
input: LMInput(text: sliceText(lmInput.text, from: start, to: boundary)),
model: context.model, cache: cache, parameters: params)
await promptCache.save(tokens: Array(promptTokens[..<boundary]), cache: cache)
return boundary
}
var stream: AsyncStream<Generation>
if let boundary = hybridBoundary {
let start = await promptCache.restoreExactPrefix(
Expand Down Expand Up @@ -2877,6 +2897,7 @@ func handleChatCompletion(
// Trim the KV cache back by 1 to avoid double-counting the replayed token.
for layer in cache { layer.trim(1) }
}
startIndex = try await snapshotPrompt(from: startIndex)
let trimmedInput = LMInput(text: sliceText(lmInput.text, from: startIndex))
if config.mtp, let mtpCtx = mtpContext(main: context, assistant: mtpAssistant) {
stream = try MLXLMCommon.generateMTP(
Expand All @@ -2895,9 +2916,10 @@ func handleChatCompletion(
input: lmInput, cache: cache, parameters: params, context: mtpCtx, numMTPTokens: config.numMtpTokens
)
} else {
let start = try await snapshotPrompt(from: 0)
stream = try MLXLMCommon.generate(
input: lmInput, cache: cache, parameters: params, context: context,
components: genComponents
input: start == 0 ? lmInput : LMInput(text: sliceText(lmInput.text, from: start)),
cache: cache, parameters: params, context: context, components: genComponents
)
}
}
Expand All @@ -2922,6 +2944,8 @@ func handleChatCompletion(
// The hybrid path already saved at its boundary; a save here would capture
// recurrent state one decode token too late.
guard hybridBoundary == nil, rotatingBoundary == nil else { return }
// Saved at the boundary above, before generation started.
guard !saveAtBoundary else { return }
// Nothing would ever restore it (see skipPromptCache above).
guard !skipPromptCache else { return }
// A ring buffer cannot be snapshotted after the first decode token (the MTP,
Expand Down Expand Up @@ -3433,7 +3457,9 @@ func handleChatStreaming(
// `.info` is the producer's last yield, so the cache is no longer written once
// it arrives. The prompt is saved here, not on the first chunk: that raced the
// decode step the producer had already started on the same KVCacheSimple. An
// early exit (stop sequence, disconnect) never reaches `.info` and skips the save.
// early exit (stop sequence, disconnect) never reaches `.info`, so the
// full-attention path saves at the prompt boundary before generating instead
// (`saveAtBoundary`); only the paths that still save here skip on early exit.
// The save does not wait for a chunk: a turn that is only a tool call never
// yields one, and its prompt is exactly what the next turn of the loop needs.
if let onPrefillDone { await onPrefillDone() }
Expand Down Expand Up @@ -3615,7 +3641,8 @@ func handleChatNonStreaming(
print("[SwiftLM] Rejected tool call: reason=\(rejection.reason) tool=\(rejection.toolName ?? "?") detail=\(rejection.detail ?? "n/a")")
case .info(let info):
// The producer's last yield; see the matching comment in `handleChatStreaming`.
// The stop-sequence `break` above exits before `.info`, so it skips the save.
// The stop-sequence `break` above exits before `.info`; the paths that save at the
// prompt boundary are unaffected, the others skip the save.
if let onPrefillDone { await onPrefillDone() }
generationStopReason = info.stopReason
completionTokenCount = info.generationTokenCount // real count; see handleChatStreaming
Expand Down
60 changes: 58 additions & 2 deletions tests/test-contract.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -439,7 +439,7 @@ fi
# #245: the save at `.info` waited for a text chunk, and a turn the model answers with
# nothing but a tool call never produces one. The next turn of an agent loop (tool
# result, follow-up) then re-prefilled the whole conversation. The server log shows the
# hit: the follow-up must reuse at least the whole prompt of the tool-call turn.
# hit: the follow-up must reuse the prompt of the tool-call turn (less its last token).
log "Test 15: a tool-call-only turn saves the prompt cache"
SERVER_LOG=/tmp/SwiftLM-test-contract.log
hit_count() { grep -c 'Prompt cache HIT' "$SERVER_LOG" || true; }
Expand Down Expand Up @@ -488,7 +488,9 @@ tool_turn_then_follow_up() {
return
fi
REUSED=$(last_hit_reused)
if [ "${REUSED:-0}" -ge "$PROMPT_TOKENS" ]; then
# The prompt is saved at the turn boundary without its last token (that token is fed to
# start generation), so the entry is one token short of the prompt.
if [ "${REUSED:-0}" -ge $((PROMPT_TOKENS - 1)) ]; then
pass "$LABEL: the follow-up reused $REUSED cached tokens, the tool-call turn's prompt was $PROMPT_TOKENS"
else
fail "$LABEL: the follow-up reused only $REUSED cached tokens; the tool-call turn's prompt was $PROMPT_TOKENS"
Expand All @@ -497,6 +499,60 @@ tool_turn_then_follow_up() {
tool_turn_then_follow_up "non-streaming" "You are a weather desk. Call get_weather for any weather question and say nothing else." ""
tool_turn_then_follow_up "streaming" "You are a weather kiosk. Call get_weather for any weather question and say nothing else." ',"stream":true,"stream_options":{"include_usage":true}'

# ── 16. A turn that ends early still saves the prompt cache ──────────────────
# #245: the save moved to the producer's `.info`, which a stop sequence or a client
# disconnect never reaches, so the next turn re-prefilled the whole history. The prompt is
# now saved at the turn boundary, before generation starts. The follow-up repeats the
# first turn and adds a short exchange, so it must reuse nearly all of its own prompt.
log "Test 16: a turn that ends on a stop sequence or a disconnect saves the prompt cache"
# $1: label, $2: a tag no other test uses, $3: stop sequence ("" for none), $4: "leave" to
# disconnect after the first streamed token.
early_exit_then_follow_up() {
local LABEL="$1" TAG="$2" STOP="$3" MODE="$4"
local SYS HISTORY HITS_BEFORE REUSED FOLLOW PROMPT_TOKENS
SYS="Scenario $TAG. You are a terse assistant. $(python3 -c 'print("The committee reviews the quarterly infrastructure budget and the migration schedule for the storage cluster. " * 30)')"
HISTORY="{\"role\":\"system\",\"content\":\"$SYS\"},{\"role\":\"user\",\"content\":\"Name three rivers, one sentence.\"}"
if [ "$MODE" = "leave" ]; then
HOST="$HOST" PORT="$PORT" HISTORY="$HISTORY" python3 - <<'PY'
import http.client, json, os
body = '{"model":"x","stream":true,"max_tokens":200,"temperature":0,"messages":[' + os.environ["HISTORY"] + ']}'
conn = http.client.HTTPConnection(os.environ["HOST"], int(os.environ["PORT"]), timeout=120)
conn.request("POST", "/v1/chat/completions", body, {"Content-Type": "application/json"})
resp = conn.getresponse()
for line in resp:
line = line.decode().strip()
if line.startswith("data: {"):
ch = json.loads(line[6:]).get("choices") or []
if ch and (ch[0].get("delta") or {}).get("content"):
break
conn.close()
PY
else
chat "{\"model\":\"x\",\"messages\":[$HISTORY],\"max_tokens\":60,\"temperature\":0,\"stop\":[\"$STOP\"]}" >/dev/null
fi
sleep 1
HITS_BEFORE=$(hit_count)
FOLLOW=$(chat "{\"model\":\"x\",\"messages\":[$HISTORY,{\"role\":\"assistant\",\"content\":\"Three rivers are\"},{\"role\":\"user\",\"content\":\"Now name three lakes.\"}],\"max_tokens\":4,\"temperature\":0}")
PROMPT_TOKENS=$(echo "$FOLLOW" | python3 -c 'import json,sys; print(json.load(sys.stdin)["usage"]["prompt_tokens"])')
for _ in $(seq 1 20); do
[ "$(hit_count)" -gt "$HITS_BEFORE" ] && break
sleep 0.25
done
if [ "$(hit_count)" -le "$HITS_BEFORE" ]; then
fail "$LABEL: the follow-up got no prompt cache hit at all"
return
fi
REUSED=$(last_hit_reused)
# The follow-up adds an assistant reply and a user message on top of the first turn.
if [ "${REUSED:-0}" -ge $((PROMPT_TOKENS - 60)) ]; then
pass "$LABEL: the follow-up reused $REUSED of its $PROMPT_TOKENS prompt tokens"
else
fail "$LABEL: the follow-up reused only $REUSED of its $PROMPT_TOKENS prompt tokens"
fi
}
early_exit_then_follow_up "stop sequence" "stop-seq" "the" ""
early_exit_then_follow_up "client disconnect" "disconnect" "" "leave"

log "═══════════════════════════════════════"
log "Results: $PASS passed, $FAIL failed, $SKIP skipped, $TOTAL total"
log "═══════════════════════════════════════"
Expand Down
Loading