fix(server): skip the prompt cache for VLM-loaded models (b769 text chat failures)

Sources/SwiftLM/Server.swift

@@ -973,6 +973,9 @@ struct MLXServer: AsyncParsableCommand {
 
         // Final after the VLM→LLM fallback above; later closures capture this `let`.
         let loadedAsVision = isVision
+        if loadedAsVision || isAudio {
+            print("[SwiftLM] Note: the prompt cache is not used for VLM/Omni loads; each text request re-prefills its full prompt.")
+        }
 
         print("[SwiftLM] Loaded model configuration. Inferred tool call format: \(String(describing: await container.configuration.toolCallFormat))")
 
@@ -2132,7 +2135,15 @@ func handleChatCompletion(
         // Skip prompt cache for quantized-KV requests: the prompt cache stores KV state
         // produced with KVCacheSimple; restoring it into a QuantizedKVCache (or vice-versa)
         // is unsafe and produces incorrect results or runtime failures.
-        let skipPromptCache = isMultimodalRequest || params.kvBits != nil
+        //
+        // Skip it for VLM/Omni-loaded models too. Their processors mostly return [1, T]
+        // tokens, which the axis-0 slices below mishandle (a generic hit re-feeds the whole
+        // prompt on top of the restored KV), and models such as Qwen3.5/Qwen3-VL need the
+        // LMOutput.State (ropeDeltas) of the cached prefix, which the prompt cache does not
+        // store. Checked on the model type rather than config.isVision so --audio
+        // (OmniModelFactory) loads are covered as well.
+        let isVLM = context.model is any VLMModel
+        let skipPromptCache = isMultimodalRequest || params.kvBits != nil || isVLM
 
         // ── Hybrid (recurrent + attention) prompt cache ──
         // Qwen3.5/3.6-style models pair MambaCache (linear attention) with KVCacheSimple
@@ -2227,6 +2238,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 else { return }
+            // Nothing would ever restore it (see skipPromptCache above).
+            guard !isVLM else { return }
             if turboHasCompressed {
                 print("[SwiftLM] 🧠 Skipping prompt cache save — TurboQuant has compressed \(cache.compactMap { ($0 as? KVCacheSimple)?.compressedOffset }.max() ?? 0) tokens. Saving would decode ~37 GB back to fp16.")
             } else if params.kvBits != nil {

tests/test-speculative.sh

@@ -213,6 +213,50 @@ else
     fail "Server became unresponsive"
 fi
 
+# ── Test 7: same checkpoint, no draft -> auto-detected qwen3_5 VLM ──
+# Every other Qwen3.5 launch in CI passes --draft-model/--dflash, which forces the LLM
+# factory and bypasses the hybrid prompt-cache split. b769 shipped with every text-only
+# chat to an auto-detected Qwen3.5 VLM failing with HTTP 500
+# (unsupportedBatchContinuation), and CI never ran that path.
+log "Test 7: $MAIN_MODEL without a draft (auto-detected VLM, text-only chat)"
+cleanup
+"$BINARY" --model "$MAIN_MODEL" --port "$PORT" --host "$HOST" >> "$LOG_FILE" 2>&1 &
+SERVER_PID=$!
+for _ in $(seq 1 300); do
+    kill -0 "$SERVER_PID" 2>/dev/null || break
+    curl -sf "$URL/health" >/dev/null 2>&1 && break
+    sleep 1
+done
+T1='{"role":"user","content":"Say hi."}'
+T2="$T1"',{"role":"assistant","content":"Hi."},{"role":"user","content":"Say bye."}'
+for MSGS in "$T1" "$T2"; do
+    N=$(echo "[$MSGS]" | jq length)
+    R=$(curl -sf --max-time 120 -X POST "$URL/v1/chat/completions" \
+        -H "Content-Type: application/json" \
+        -d "{\"model\":\"$MAIN_MODEL\",\"max_tokens\":8,\"temperature\":0,\"messages\":[$MSGS]}" 2>/dev/null || echo "")
+    if echo "$R" | jq -e '.choices[0].finish_reason and .usage.prompt_tokens > 0' >/dev/null 2>&1; then
+        pass "No-draft VLM text chat, $N message(s): completed"
+    else
+        fail "No-draft VLM text chat, $N message(s): no completion (${R:0:200})"
+    fi
+done
+# Checked after the requests, which fflush stdout. Guards against this test silently
+# covering the LLM path (auto-detect off, or the VLM->LLM fallback, which still prints
+# the auto-detect line).
+if grep -q "Loading VLM (vision-language model)" "$LOG_FILE" \
+    && ! grep -q "Auto-detected VLM failed to load" "$LOG_FILE"; then
+    pass "No-draft: $MAIN_MODEL was loaded as a VLM"
+else
+    fail "No-draft: $MAIN_MODEL was not loaded as a VLM, so Test 7 did not cover the VLM path"
+fi
+# The earlier --draft-model runs never restore from the prompt cache, so any hit here
+# comes from the VLM server, which must bypass the cache.
+if grep -q "Prompt cache HIT" "$LOG_FILE"; then
+    fail "No-draft: the VLM server restored from the prompt cache"
+else
+    pass "No-draft: the VLM server bypassed the prompt cache"
+fi
+
 # ── Results ──────────────────────────────────────────────────────────
 echo ""
 log "═══════════════════════════════════════"