diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index eade7db21256..710632f76f8f 100644 --- a/tools/server/server-common.cpp +++ b/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) { diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 9894f5f06fb0..bd429e995057 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -215,13 +215,17 @@ struct server_tokens { std::vector 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(); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index bc5fbf937fef..4006098e82cc 100644 --- a/tools/server/server-context.cpp +++ b/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 { diff --git a/tools/server/tests/unit/test_slot_save.py b/tools/server/tests/unit/test_slot_save.py index 5eca46cb292d..e5885dfcec09 100644 --- a/tools/server/tests/unit/test_slot_save.py +++ b/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() diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index 826aef2d5bcb..d245f3bd6437 100644 --- a/tools/server/tests/utils.py +++ b/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: