From 670b814c0a518fccf7e640b1b7fc013db68c250d Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Thu, 29 Oct 2026 11:13:02 -0700 Subject: [PATCH] fix: save the prompt cache at the turn boundary so early exits keep it (#245) The save moved to the producer's `.info` event to avoid racing the decode step, but a request that ends on a stop sequence or a client disconnect leaves the stream before `.info`, so its prompt was never saved and the next turn re-prefilled the whole history. For the full-attention path (no hybrid/rotating layers, draft model, MTP or TurboKV), prefill the prompt but its last token synchronously, save, then let generation feed the last token, as the hybrid and rotating paths already do. No decode step has started when the cache is read, so there is no race, and the save no longer depends on how the request ends. The entry is one token short of the prompt; restore replays that token anyway. Contract test 16 covers a stop-sequence turn and a client-disconnect turn; test 15 now allows the one-token difference. Co-Authored-By: Claude Sonnet 5.5 --- Sources/SwiftLM/Server.swift | 35 ++++++++++++++++++--- tests/test-contract.sh | 60 ++++++++++++++++++++++++++++++++++-- 2 files changed, 89 insertions(+), 6 deletions(-) diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 1bba8f5..a4f818f 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -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[.. if let boundary = hybridBoundary { let start = await promptCache.restoreExactPrefix( @@ -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( @@ -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 ) } } @@ -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, @@ -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() } @@ -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 diff --git a/tests/test-contract.sh b/tests/test-contract.sh index 35c1755..f8023a2 100755 --- a/tests/test-contract.sh +++ b/tests/test-contract.sh @@ -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; } @@ -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" @@ -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 "═══════════════════════════════════════"