server : allow cache reuse for text-only prompts with mmproj loaded

tools/server/server-common.cpp

@@ -643,9 +643,10 @@ llama_tokens server_tokens::get_text_tokens() const {
     return res;
 }
 
-void server_tokens::set_token(llama_pos pos, llama_token id) {
-    GGML_ASSERT(!has_mtmd); // only allow this if mtmd is disabled
-    tokens[pos] = id;
+void server_tokens::set_token(size_t idx, llama_token id) {
+    // a media placeholder must stay in sync with the media map, so never write over one
+    GGML_ASSERT(tokens[idx] != LLAMA_TOKEN_NULL && id != LLAMA_TOKEN_NULL);
+    tokens[idx] = id;
 }
 
 void server_tokens::keep_first(size_t n) {

tools/server/server-common.h

@@ -215,13 +215,17 @@ struct server_tokens {
     std::vector<char> serialize() const;
     static server_tokens deserialize(const llama_tokens & packed, bool has_mtmd);
 
-    // for compatibility with speculative decoding
-    void set_token(llama_pos pos, llama_token id);
+    // overwrite a text token, media placeholders are not writable
+    void set_token(size_t idx, llama_token id);
 
     size_t size() const { return tokens.size(); }
 
     bool empty() const { return tokens.empty(); }
 
+    // true if the token list holds real media chunks
+    // note: this differs from has_mtmd, which only means an mmproj is loaded
+    bool has_media_chunks() const { return !map_idx_to_media.empty(); }
+
     void clear() {
         map_idx_to_media.clear();
         tokens.clear();

tools/server/server-context.cpp

@@ -1211,10 +1211,7 @@ struct server_context_impl {
                 SRV_WRN("%s\n", "ctx_shift is not supported by multimodal, it will be disabled");
             }
 
-            if (params_base.n_cache_reuse) {
-                params_base.n_cache_reuse = 0;
-                SRV_WRN("%s\n", "cache_reuse is not supported by multimodal, it will be disabled");
-            }
+            // keep cache_reuse: it is applied per request and only while the prompt has no media
         }
 
         if (!llama_memory_can_shift(llama_get_memory(ctx_tgt))) {
@@ -3419,25 +3416,34 @@ struct server_context_impl {
 
                                 const auto n_cache_reuse = slot.task->params.n_cache_reuse;
 
-                                const bool can_cache_reuse =
-                                    llama_memory_can_shift(llama_get_memory(ctx_tgt)) &&
-                                    !slot.prompt.tokens.has_mtmd;
+                                // the shift is applied to the draft context too, so both must support it
+                                const bool can_shift = llama_memory_can_shift(llama_get_memory(ctx_tgt)) &&
+                                                       (!ctx_dft || llama_memory_can_shift(llama_get_memory(ctx_dft)));
+
+                                // the loop below uses token indices as positions, which a media chunk breaks.
+                                // an mmproj alone is fine, only a real media chunk in either prompt is not
+                                const bool has_media = slot.prompt.tokens.has_media_chunks() || input_tokens.has_media_chunks();
 
-                                if (!can_cache_reuse && n_cache_reuse > 0) {
-                                    SLT_WRN(slot, "cache reuse is not supported - ignoring n_cache_reuse = %d\n", n_cache_reuse);
+                                // the loop moves n_past past the alora cap applied above
+                                const bool has_alora = slot.alora_invocation_start > 0;
+
+                                const bool can_cache_reuse = n_cache_reuse > 0 && can_shift && !has_media && !has_alora;
+
+                                if (n_cache_reuse > 0 && !can_cache_reuse) {
+                                    if (!can_shift) {
+                                        SLT_WRN(slot, "cache reuse is not supported - ignoring n_cache_reuse = %d\n", n_cache_reuse);
+                                    } else {
+                                        // expected on every request with media or an alora, so keep it out of the log
+                                        SLT_DBG(slot, "cache reuse is disabled for this prompt - ignoring n_cache_reuse = %d\n", n_cache_reuse);
+                                    }
                                 }
 
                                 // reuse chunks from the cached prompt by shifting their KV cache in the new position
-                                if (can_cache_reuse && n_cache_reuse > 0) {
-                                    GGML_ASSERT(!slot.prompt.tokens.has_mtmd);
-
+                                if (can_cache_reuse) {
                                     size_t head_c = n_past; // cache
                                     size_t head_p = n_past; // current prompt
 
-                                    if (mctx) {
-                                        // we should never reach this
-                                        GGML_ABORT("not supported by multimodal");
-                                    }
+                                    bool kv_shifted = false;
 
                                     SLT_DBG(slot, "trying to reuse chunks with size > %d, n_past = %d\n", n_cache_reuse, n_past);
 
@@ -3462,6 +3468,8 @@ struct server_context_impl {
                                             slot.mem.seq_rm (slot.id, head_p, head_c);
                                             slot.mem.seq_add(slot.id, head_c, head_c + n_match, kv_shift);
 
+                                            kv_shifted |= kv_shift != 0;
+
                                             for (size_t i = 0; i < n_match; i++) {
                                                 slot.prompt.tokens.set_token(head_p + i, slot.prompt.tokens[head_c + i]);
                                                 n_past++;
@@ -3474,6 +3482,11 @@ struct server_context_impl {
                                         }
                                     }
 
+                                    if (kv_shifted) {
+                                        // the checkpoints were taken before the shift, they no longer match the cache
+                                        slot.prompt.checkpoints.clear();
+                                    }
+
                                     SLT_DBG(slot, "after context reuse, new n_past = %d\n", n_past);
                                 }
                             } else {

tools/server/tests/unit/test_slot_save.py

@@ -493,6 +493,71 @@ def test_slot_save_restore_image_payload_larger_than_context(mmproj_server):
     assert res.body["timings"]["prompt_n"] == 1
 
 
+#
+# Prompt cache reuse on a multimodal server (mmproj loaded).
+#
+# Cache reuse is gated on real media chunks, not on has_mtmd.
+# swa_full keeps the shifted match valid, cache_ram 0 leaves the KV shift as the only reuse path.
+#
+
+CACHE_REUSE_LEAD = "Throw away this opening line."
+
+# starts on a newline, so the shared chunk tokenizes the same with and without the lead
+CACHE_REUSE_TEXT = (
+    "\nAlpha beta gamma delta epsilon zeta eta theta iota kappa"
+    " lambda mu nu xi omicron pi rho sigma tau upsilon phi chi psi omega."
+)
+
+
+def test_cache_reuse_with_mmproj(mmproj_server):
+    server = mmproj_server
+    server.cache_reuse = 4
+    server.swa_full = True
+    server.cache_ram = 0
+    server.start()
+
+    img = _get_img_base64(IMG_URL_CAT)
+
+    def send_text(prompt, id_slot):
+        res = server.make_request("POST", "/completion", data={
+            "prompt": prompt,
+            "id_slot": id_slot,
+            "cache_prompt": True,
+            "n_predict": 1,
+        })
+        assert res.status_code == 200
+        return res.body["timings"]["cache_n"]
+
+    def send_media(prompt_string, id_slot):
+        res = server.make_request("POST", "/completions", data={
+            "id_slot": id_slot,
+            "cache_prompt": True,
+            "n_predict": 1,
+            "prompt": {
+                "prompt_string": prompt_string,
+                "multimodal_data": [img],
+            },
+        })
+        assert res.status_code == 200
+        return res.body["timings"]["cache_n"]
+
+    # text-only: dropping the lead must shift the shared chunk and reuse it
+    send_text(CACHE_REUSE_LEAD + CACHE_REUSE_TEXT, 0)
+    assert send_text(CACHE_REUSE_TEXT, 0) > 10
+
+    # media in the cached prompt blocks the very same shift
+    send_media(CACHE_REUSE_LEAD + CACHE_REUSE_TEXT + " <__media__>", 1)
+    assert send_text(CACHE_REUSE_TEXT, 1) < 10
+
+    # media in the incoming prompt blocks it too, even though it sits after the shared chunk
+    send_text(CACHE_REUSE_LEAD + CACHE_REUSE_TEXT, 1)
+    assert send_media(CACHE_REUSE_TEXT + " <__media__>", 1) < 10
+
+    # reuse resumes as soon as the slot holds text only again
+    send_text(CACHE_REUSE_LEAD + CACHE_REUSE_TEXT, 1)
+    assert send_text(CACHE_REUSE_TEXT, 1) > 10
+
+
 def test_slot_restore_media_file_without_mmproj(mmproj_server):
     server = mmproj_server
     server.start()

tools/server/tests/utils.py

@@ -113,6 +113,7 @@ class ServerProcess:
     media_path: str | None = None
     sleep_idle_seconds: int | None = None
     cache_ram: int | None = None
+    cache_reuse: int | None = None
     no_cache_idle_slots: bool = False
     log_path: str | None = None
     ui_mcp_proxy: bool = False
@@ -278,6 +279,8 @@ def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None:
             server_args.extend(["--sleep-idle-seconds", self.sleep_idle_seconds])
         if self.cache_ram is not None:
             server_args.extend(["--cache-ram", self.cache_ram])
+        if self.cache_reuse is not None:
+            server_args.extend(["--cache-reuse", self.cache_reuse])
         if self.no_cache_idle_slots:
             server_args.append("--no-cache-idle-slots")
         if self.ui_mcp_proxy: