From 5d9e5ac30e469d44c0a5a52556de0ead03aaa5b0 Mon Sep 17 00:00:00 2001 From: Chipmunk <101038159+CHIPMUNK-T0T@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:20:28 +0900 Subject: [PATCH] server : support slot save/restore with media inputs (#26640) * server : save serialized image chunks at the end of the llama state * server : support multimodal slot state save/restore with packed payload * server : refine image slot state serialization * server : support media slot state and centralize media validation * server : remove unnecessary comment * server : remove defensive media checks and move the chunk type check to validate() --- include/llama.h | 1 + src/llama-context.cpp | 11 + tools/server/server-common.cpp | 185 ++++++++++- tools/server/server-common.h | 10 +- tools/server/server-context.cpp | 81 +++-- tools/server/tests/unit/test_slot_save.py | 377 +++++++++++++++++++--- tools/server/tests/utils.py | 6 + 7 files changed, 593 insertions(+), 78 deletions(-) diff --git a/include/llama.h b/include/llama.h index bfef0e1d11..ef278c9c32 100644 --- a/include/llama.h +++ b/include/llama.h @@ -883,6 +883,7 @@ extern "C" { const llama_token * tokens, size_t n_token_count); + // If tokens_out is NULL, only the token count is reported through n_token_count_out and no state is loaded LLAMA_API size_t llama_state_seq_load_file( struct llama_context * ctx, const char * filepath, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0de3a68d1c..aa9fb2c3b4 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3110,6 +3110,17 @@ size_t llama_context::state_seq_load_file(llama_seq_id seq_id, const char * file { const uint32_t n_token_count = file.read_u32(); + if (tokens_out == nullptr) { + const size_t n_token_max = (file.size() - file.tell()) / sizeof(llama_token); + if (n_token_count > n_token_max) { + LLAMA_LOG_ERROR("%s: token count in sequence state file exceeds the file size! %u > %zu\n", __func__, n_token_count, n_token_max); + return 0; + } + + *n_token_count_out = n_token_count; + return file.tell(); + } + if (n_token_count > n_token_capacity) { LLAMA_LOG_ERROR("%s: token count in sequence state file exceeded capacity! %u > %zu\n", __func__, n_token_count, n_token_capacity); return 0; diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index c9109fc962..148b7ea6d8 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -13,6 +13,8 @@ #include #include #include +#include +#include json format_error_response(const std::string & message, const enum error_type type) { std::string type_str; @@ -235,6 +237,102 @@ static inline raw_buffer base64_decode(const std::string & encoded_string) { // server_tokens implementation // +namespace { + +constexpr uint32_t SERVER_TOKENS_STATE_VERSION = 1; + +uint32_t server_tokens_state_u32(size_t value) { + if (value > std::numeric_limits::max()) { + throw std::runtime_error("Server tokens state is too large"); + } + return value; +} + +class server_tokens_state_writer { +public: + template + void write(T value) { + static_assert(std::is_trivially_copyable::value, "T must be trivially copyable"); + const auto * ptr = reinterpret_cast(&value); + data.insert(data.end(), ptr, ptr + sizeof(value)); + } + + template + void write(const std::vector & values) { + static_assert(std::is_trivially_copyable::value, "T must be trivially copyable"); + write(server_tokens_state_u32(values.size())); + if (values.empty()) { + return; + } + const auto * ptr = reinterpret_cast(values.data()); + data.insert(data.end(), ptr, ptr + values.size() * sizeof(T)); + } + + void write_media_chunk(const mtmd_input_chunk * chunk) { + size_t chunk_size = 0; + if (mtmd_input_chunk_save(chunk, nullptr, 0, &chunk_size) != 0 || chunk_size == 0) { + throw std::runtime_error("Cannot serialize media chunk in server tokens"); + } + std::vector chunk_data(server_tokens_state_u32(chunk_size)); + if (mtmd_input_chunk_save(chunk, chunk_data.data(), chunk_data.size(), nullptr) != 0) { + throw std::runtime_error("Cannot serialize media chunk in server tokens"); + } + write(chunk_data); + } + + std::vector take() { + data.resize((data.size() + sizeof(llama_token) - 1) / sizeof(llama_token) * sizeof(llama_token), 0); + return std::move(data); + } + +private: + std::vector data; +}; + +class server_tokens_state_reader { +public: + server_tokens_state_reader(const char * data, size_t size) : data(data), size(size) {} + + template + T read() { + static_assert(std::is_trivially_copyable::value, "T must be trivially copyable"); + if (size - pos < sizeof(T)) { + throw std::runtime_error("Unexpected end of server tokens state"); + } + T value; + std::memcpy(&value, data + pos, sizeof(value)); + pos += sizeof(value); + return value; + } + + template + std::vector read_vector() { + static_assert(std::is_trivially_copyable::value, "T must be trivially copyable"); + const uint32_t n_values = read(); + // reject before resizing, so that a small corrupted payload cannot request a huge allocation + if (n_values > remaining() / sizeof(T)) { + throw std::runtime_error("Unexpected end of server tokens state"); + } + std::vector values(n_values); + if (n_values > 0) { + std::memcpy(values.data(), data + pos, values.size() * sizeof(T)); + pos += values.size() * sizeof(T); + } + return values; + } + + size_t remaining() const { + return size - pos; + } + +private: + const char * data; + size_t size; + size_t pos = 0; +}; + +} // namespace + server_tokens::server_tokens(mtmd::input_chunks & mtmd_chunks, bool has_mtmd) : has_mtmd(has_mtmd) { for (size_t i = 0; i < mtmd_chunks.size(); ++i) { push_back(mtmd_chunks[i]); @@ -408,6 +506,73 @@ const llama_tokens & server_tokens::get_tokens() const { return tokens; } +std::vector server_tokens::serialize() const { + static_assert(sizeof(llama_token) == sizeof(uint32_t), "unexpected llama_token size"); + + server_tokens_state_writer writer; + writer.write((llama_token) LLAMA_TOKEN_NULL); + writer.write(SERVER_TOKENS_STATE_VERSION); + writer.write(tokens); + + std::vector media_keys; + media_keys.reserve(map_idx_to_media.size()); + for (const auto & item : map_idx_to_media) { + media_keys.push_back(server_tokens_state_u32(item.first)); + } + writer.write(media_keys); + + for (const auto & item : map_idx_to_media) { + writer.write_media_chunk(item.second.get()); + } + + return writer.take(); +} + +server_tokens server_tokens::deserialize(const llama_tokens & packed, bool has_mtmd) { + static_assert(sizeof(llama_token) == sizeof(uint32_t), "unexpected llama_token size"); + + if (packed.empty() || packed[0] != LLAMA_TOKEN_NULL) { + // plain token list, as written by older versions + return server_tokens(packed, has_mtmd); + } + + server_tokens_state_reader reader(reinterpret_cast(packed.data()), packed.size() * sizeof(llama_token)); + reader.read(); // format marker + if (reader.read() != SERVER_TOKENS_STATE_VERSION) { + throw std::runtime_error("Unsupported server tokens state version"); + } + + const llama_tokens tokens = reader.read_vector(); + + // the media start indices, followed by the media chunks in the same order + const std::vector media_keys = reader.read_vector(); + if (!media_keys.empty() && !has_mtmd) { + throw std::runtime_error("Cannot restore media tokens without an mmproj"); + } + + server_tokens result(tokens, has_mtmd); + + for (const uint32_t key : media_keys) { + const size_t start_idx = key; + const std::vector chunk_data = reader.read_vector(); + if (chunk_data.empty()) { + throw std::runtime_error("Cannot load media chunk from server tokens state"); + } + + mtmd::input_chunk_ptr chunk(mtmd_input_chunk_load(chunk_data.data(), chunk_data.size())); + if (!chunk) { + throw std::runtime_error("Cannot load media chunk from server tokens state"); + } + result.map_idx_to_media[start_idx] = std::move(chunk); + } + + if (reader.remaining() >= sizeof(llama_token)) { + throw std::runtime_error("Trailing data in server tokens state"); + } + + return result; +} + llama_tokens server_tokens::get_text_tokens() const { llama_tokens res; res.reserve(tokens.size()); @@ -530,14 +695,28 @@ bool server_tokens::validate(const struct llama_context * ctx) const { const llama_model * model = llama_get_model(ctx); const llama_vocab * vocab = llama_model_get_vocab(model); const int32_t n_vocab = llama_vocab_n_tokens(vocab); + size_t n_media = 0; for (size_t i = 0; i < tokens.size(); ++i) { const auto & t = tokens[i]; if (t == LLAMA_TOKEN_NULL) { try { const auto & chunk = find_chunk(i); - size_t n_tokens = mtmd_input_chunk_get_n_tokens(chunk.get()); - i += n_tokens - 1; // will be +1 by the for loop + if (mtmd_input_chunk_get_type(chunk.get()) == MTMD_INPUT_CHUNK_TYPE_TEXT) { + return false; + } + const size_t n_tokens = mtmd_input_chunk_get_n_tokens(chunk.get()); + const llama_pos n_pos = mtmd_input_chunk_get_n_pos(chunk.get()); + if (n_tokens == 0 || n_pos <= 0 || n_tokens > tokens.size() - i) { + return false; + } + for (size_t j = i; j < i + n_tokens; ++j) { + if (tokens[j] != LLAMA_TOKEN_NULL) { + return false; + } + } + ++n_media; + i += n_tokens - 1; } catch (const std::exception & e) { return false; } @@ -545,7 +724,7 @@ bool server_tokens::validate(const struct llama_context * ctx) const { return false; } } - return true; + return n_media == map_idx_to_media.size(); } server_tokens server_tokens::clone() const { diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 6ef797ebb4..79607e288c 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -201,11 +201,14 @@ public: // for compatibility with context shift and prompt truncation void insert(const llama_tokens & inp_tokens); - // for compatibility with speculative decoding, ctx shift, slot save/load + // for compatibility with speculative decoding, ctx shift const llama_tokens & get_tokens() const; llama_tokens get_text_tokens() const; + 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); @@ -213,9 +216,6 @@ public: bool empty() const { return tokens.empty(); } - // true if the sequence actually contains image/audio chunks. - bool has_media() const { return !map_idx_to_media.empty(); } - void clear() { map_idx_to_media.clear(); tokens.clear(); @@ -230,7 +230,7 @@ public: // split the tokens into message spans, skipping over media chunks common_chat_msg_spans find_message_spans(const common_chat_msg_delimiters & delims) const; - // make sure all text tokens are within the vocab range + // check text token IDs and the mapping between media chunks and token ranges bool validate(const struct llama_context * ctx) const; server_tokens clone() const; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index d75c6856df..8b2ae72a4c 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2072,18 +2072,6 @@ private: queue_results.send(std::move(res)); } - // Gate slot save/restore/erase on slot content (does it hold media), - // not model capability: a multimodal model may hold a pure-text slot. - bool check_slot_no_media(const server_slot & slot, const int id_task) { - if (slot.prompt.tokens.has_media()) { - send_error(id_task, - "This operation is not supported while the slot holds image/audio tokens (a pure-text prefix is supported)", - ERROR_TYPE_NOT_SUPPORTED); - return false; - } - return true; - } - void send_partial_response(server_slot & slot, const completion_token_output & tkn, bool is_progress, bool is_begin = false) { auto res = std::make_unique(); @@ -2577,9 +2565,6 @@ private: send_error(task, "Invalid slot ID", ERROR_TYPE_INVALID_REQUEST); break; } - if (!check_slot_no_media(*slot, task.id)) { - break; - } if (slot->is_processing()) { // if requested slot is unavailable, we defer this task for processing later SRV_DBG("requested slot is unavailable, defer task, id_task = %d\n", task.id); @@ -2592,9 +2577,22 @@ private: std::string filename = task.slot_action.filename; std::string filepath = task.slot_action.filepath; - const llama_tokens tokens = slot->prompt.tokens.get_text_tokens(); - const size_t token_count = tokens.size(); - const size_t nwrite = llama_state_seq_save_file(ctx_tgt, filepath.c_str(), slot->id, tokens.data(), token_count); + std::vector packed; + try { + packed = slot->prompt.tokens.serialize(); + } catch (const std::exception & err) { + send_error(task, err.what(), ERROR_TYPE_NOT_SUPPORTED); + break; + } + + GGML_ASSERT(packed.size() % sizeof(llama_token) == 0); + const size_t nwrite = llama_state_seq_save_file( + ctx_tgt, filepath.c_str(), slot->id, + reinterpret_cast(packed.data()), packed.size() / sizeof(llama_token)); + if (nwrite == 0) { + send_error(task, "Unable to save slot", ERROR_TYPE_SERVER); + break; + } const int64_t t_end = ggml_time_us(); const double t_save_ms = (t_end - t_start) / 1000.0; @@ -2604,7 +2602,7 @@ private: res->id_slot = id_slot; res->filename = filename; res->is_save = true; - res->n_tokens = token_count; + res->n_tokens = slot->prompt.tokens.size(); res->n_bytes = nwrite; res->t_ms = t_save_ms; queue_results.send(std::move(res)); @@ -2629,18 +2627,37 @@ private: std::string filename = task.slot_action.filename; std::string filepath = task.slot_action.filepath; - llama_tokens tokens; - tokens.resize(slot->n_ctx); - size_t token_count = 0; - size_t nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, tokens.data(), tokens.size(), &token_count); - if (nread == 0) { - slot->prompt.clear(); // KV may already been invalidated? - send_error(task, "Unable to restore slot, no available space in KV cache or invalid slot save file", ERROR_TYPE_INVALID_REQUEST); + size_t nread = 0; + try { + size_t n_packed = 0; + llama_tokens packed; + nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, nullptr, 0, &n_packed); + if (nread != 0) { + packed.resize(std::max(1, n_packed)); + nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, packed.data(), packed.size(), &n_packed); + } + if (nread == 0) { + throw std::runtime_error("No available space in KV cache or invalid slot save file"); + } + packed.resize(n_packed); + + server_tokens restored = server_tokens::deserialize(packed, mctx != nullptr); + + if (restored.size() > (size_t) slot->n_ctx) { + throw std::runtime_error("Restored prompt does not fit in the slot context"); + } + + if (!restored.validate(ctx_tgt)) { + throw std::runtime_error("Invalid tokens in slot save file"); + } + + slot->prompt.clear(); + slot->prompt.tokens = std::move(restored); + } catch (const std::exception & err) { + slot->prompt_clear(); + send_error(task, std::string("Unable to restore slot: ") + err.what(), ERROR_TYPE_INVALID_REQUEST); break; } - tokens.resize(token_count); - slot->prompt.clear(); - slot->prompt.tokens.insert(tokens); const int64_t t_end = ggml_time_us(); const double t_restore_ms = (t_end - t_start) / 1000.0; @@ -2650,7 +2667,7 @@ private: res->id_slot = id_slot; res->filename = filename; res->is_save = false; - res->n_tokens = token_count; + res->n_tokens = slot->prompt.tokens.size(); res->n_bytes = nread; res->t_ms = t_restore_ms; queue_results.send(std::move(res)); @@ -2663,10 +2680,6 @@ private: send_error(task, "Invalid slot ID", ERROR_TYPE_INVALID_REQUEST); break; } - // Gate on slot content, consistent with save/restore. - if (!check_slot_no_media(*slot, task.id)) { - break; - } if (slot->is_processing()) { // if requested slot is unavailable, we defer this task for processing later SRV_DBG("requested slot is unavailable, defer task, id_task = %d\n", task.id); diff --git a/tools/server/tests/unit/test_slot_save.py b/tools/server/tests/unit/test_slot_save.py index be22d9859e..05acb1be14 100644 --- a/tools/server/tests/unit/test_slot_save.py +++ b/tools/server/tests/unit/test_slot_save.py @@ -2,6 +2,10 @@ import pytest from utils import * import base64 import requests +import struct + +# sequence state file: magic(4) version(4) payload_size(4), then payload_size llama_token words +STATE_FILE_HEADER_SIZE = 12 server = ServerPreset.tinyllama2() @@ -72,6 +76,60 @@ def test_slot_save_restore(): assert res.body["timings"]["prompt_n"] == 1 +def test_slot_restore_legacy_token_list(): + global server + server.start() + + res = server.make_request("POST", "/completion", data={ + "prompt": "What is the capital of France?", + "id_slot": 1, + "cache_prompt": True, + }) + assert res.status_code == 200 + + res = server.make_request("POST", "/slots/1?action=save", data={ + "filename": "slot_legacy.bin", + }) + assert res.status_code == 200 + assert res.body["n_saved"] == 84 + + # rewrite the token payload into a plain token list, as written by servers that predate the packed server_tokens format + path = os.path.join("tmp", "slot_legacy.bin") + with open(path, "rb") as f: + data = bytearray(f.read()) + + # the payload written by this server starts with a packed header: LLAMA_TOKEN_NULL(4) version(4) n_tokens(4) + packed_header_size = 12 + + payload_size = struct.unpack_from("=I", data, STATE_FILE_HEADER_SIZE - 4)[0] + payload_end = STATE_FILE_HEADER_SIZE + payload_size * 4 + n_tokens = struct.unpack_from("=I", data, STATE_FILE_HEADER_SIZE + 8)[0] + assert n_tokens == 84 + + tokens_start = STATE_FILE_HEADER_SIZE + packed_header_size + data = data[:STATE_FILE_HEADER_SIZE] + data[tokens_start:tokens_start + n_tokens * 4] + data[payload_end:] + struct.pack_into("=I", data, STATE_FILE_HEADER_SIZE - 4, n_tokens) + + with open(path, "wb") as f: + f.write(data) + + # the plain token list must restore, and the restored KV must be reusable + res = server.make_request("POST", "/slots/0?action=restore", data={ + "filename": "slot_legacy.bin", + }) + assert res.status_code == 200 + assert res.body["n_restored"] == 84 + + res = server.make_request("POST", "/completion", data={ + "prompt": "What is the capital of Germany?", + "id_slot": 0, + "cache_prompt": True, + }) + assert res.status_code == 200 + assert res.body["timings"]["prompt_n"] == 6 # only the different part is processed + + + def test_slot_erase(): global server server.start() @@ -103,14 +161,12 @@ def test_slot_erase(): # # Multimodal server (mmproj loaded) slot save/restore. # -# Regression coverage for issue #21133: slot save/restore/erase must be gated on -# the slot's CONTENT (does it actually hold image/audio tokens) rather than the -# model's CAPABILITY (is an mmproj loaded). A pure-text slot on a multimodal -# server must save/restore/erase normally; a slot that actually holds an image -# must be rejected with ERROR_TYPE_NOT_SUPPORTED (HTTP 501). +# A pure-text slot on a multimodal server and a slot containing images must both support save/restore. +# Erase remains gated on the slot's content. # IMG_URL_CAT = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/91_cat.png" +IMG_URL_TRUCK = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/11_truck.png" def _get_img_base64(url: str) -> str: @@ -121,8 +177,7 @@ def _get_img_base64(url: str) -> str: @pytest.fixture def mmproj_server(): - # tinygemma3 is a small multimodal model: the mmproj is provided by the HF - # registry API and auto-downloaded on first run. + # tinygemma3 is a small multimodal model: the mmproj is provided by the HF registry API and auto-downloaded on first run. os.environ['LLAMA_MEDIA_MARKER'] = '<__media__>' mm_server = ServerPreset.tinygemma3() mm_server.slot_save_path = "./tmp" @@ -159,10 +214,7 @@ def test_slot_save_restore_text_only_on_multimodal(mmproj_server): assert res.status_code == 200 assert res.body["n_restored"] == n_saved - # The restored slot is usable for a follow-up completion. We do NOT assert - # prefix reuse here: tinygemma3 is a SWA model, which forces full prompt - # re-processing after a restore (a model property, not the save/restore gate - # under test). + # Prefix reuse is not checked with the default SWA cache. res = server.make_request("POST", "/completion", data={ "prompt": "The quick brown fox jumps over the lazy dog.", "id_slot": 0, @@ -171,54 +223,307 @@ def test_slot_save_restore_text_only_on_multimodal(mmproj_server): assert res.status_code == 200 -def test_slot_save_rejected_when_slot_holds_image(mmproj_server): +def test_slot_save_restore_with_image(mmproj_server): server = mmproj_server + # Use the full SWA cache so the restored image prefix can be reused. + server.swa_full = True server.start() - # Process a prompt that actually contains an image on slot 1. + prompt_cat = { + "prompt_string": "What is this: <__media__>\n", + "multimodal_data": [_get_img_base64(IMG_URL_CAT)], + } res = server.make_request("POST", "/completions", data={ "temperature": 0.0, "top_k": 1, "id_slot": 1, "cache_prompt": True, + "prompt": prompt_cat, + }) + assert res.status_code == 200 + content_cat = res.body["content"] + prompt_n_full = res.body["timings"]["prompt_n"] + assert res.body["timings"]["cache_n"] == 0 + assert prompt_n_full > 32 # text plus image tokens are all processed + + res = server.make_request("POST", "/slots/1?action=save", data={ + "filename": "mm_slot_image.bin", + }) + assert res.status_code == 200 + n_saved = res.body["n_saved"] + n_written = res.body["n_written"] + assert n_saved > 0 + assert n_written > 0 + + res = server.make_request("POST", "/slots/1?action=erase") + assert res.status_code == 200 + + res = server.make_request("POST", "/slots/0?action=restore", data={ + "filename": "mm_slot_image.bin", + }) + assert res.status_code == 200 + assert res.body["n_restored"] == n_saved + assert res.body["n_read"] == n_written + + # a different image must not reuse the restored image tokens; only the text prefix before the image is common + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, "prompt": { "prompt_string": "What is this: <__media__>\n", - "multimodal_data": [ _get_img_base64(IMG_URL_CAT) ], + "multimodal_data": [_get_img_base64(IMG_URL_TRUCK)], + }, + }) + assert res.status_code == 200 + cache_n = res.body["timings"]["cache_n"] + assert cache_n < 16 + assert res.body["timings"]["prompt_n"] == prompt_n_full - cache_n + + # restore again and resend the same image: the image tokens must be reused and greedy sampling must reproduce the original content + res = server.make_request("POST", "/slots/0?action=restore", data={ + "filename": "mm_slot_image.bin", + }) + assert res.status_code == 200 + assert res.body["n_restored"] == n_saved + + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": prompt_cat, + }) + assert res.status_code == 200 + assert res.body["timings"]["cache_n"] == prompt_n_full - 1 + assert res.body["timings"]["prompt_n"] == 1 + assert res.body["content"] == content_cat + + +def test_slot_save_restore_with_two_images(mmproj_server): + server = mmproj_server + server.swa_full = True + server.n_ctx = 2048 # two images need more than the default 512 per slot + server.start() + + prompt = { + "prompt_string": "A: <__media__> B: <__media__>\n", + "multimodal_data": [_get_img_base64(IMG_URL_CAT), _get_img_base64(IMG_URL_TRUCK)], + } + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 1, + "cache_prompt": True, + "prompt": prompt, + }) + assert res.status_code == 200 + content = res.body["content"] + prompt_n_full = res.body["timings"]["prompt_n"] + assert prompt_n_full > 64 + + res = server.make_request("POST", "/slots/1?action=save", data={ + "filename": "mm_slot_two_images.bin", + }) + assert res.status_code == 200 + n_saved = res.body["n_saved"] + + res = server.make_request("POST", "/slots/0?action=restore", data={ + "filename": "mm_slot_two_images.bin", + }) + assert res.status_code == 200 + assert res.body["n_restored"] == n_saved + + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": prompt, + }) + assert res.status_code == 200 + assert res.body["timings"]["cache_n"] == prompt_n_full - 1 + assert res.body["timings"]["prompt_n"] == 1 + assert res.body["content"] == content + + +def test_slot_save_restore_with_image_across_restart(mmproj_server): + server = mmproj_server + server.swa_full = True + server.start() + + prompt_cat = { + "prompt_string": "What is this: <__media__>\n", + "multimodal_data": [_get_img_base64(IMG_URL_CAT)], + } + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": prompt_cat, + }) + assert res.status_code == 200 + content = res.body["content"] + prompt_n_full = res.body["timings"]["prompt_n"] + + res = server.make_request("POST", "/slots/0?action=save", data={ + "filename": "mm_slot_restart.bin", + }) + assert res.status_code == 200 + n_saved = res.body["n_saved"] + + # restart the server with the same model and mmproj: the saved file must restore in the new process and the image KV must be reused + server.stop() + server.start() + + res = server.make_request("POST", "/slots/0?action=restore", data={ + "filename": "mm_slot_restart.bin", + }) + assert res.status_code == 200 + assert res.body["n_restored"] == n_saved + + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": prompt_cat, + }) + assert res.status_code == 200 + assert res.body["timings"]["cache_n"] == prompt_n_full - 1 + assert res.body["timings"]["prompt_n"] == 1 + assert res.body["content"] == content + + +def test_slot_save_restore_image_payload_larger_than_context(mmproj_server): + server = mmproj_server + server.swa_full = True + server.start() + + # the slot context, as the server computed it (n_ctx split across the slots) + res = server.make_request("GET", "/props") + assert res.status_code == 200 + n_ctx_slot = res.body["default_generation_settings"]["n_ctx"] + + # a filler token, used to grow the prompt up to the slot context + res = server.make_request("POST", "/tokenize", data={"content": " hello" * 8}) + assert res.status_code == 200 + assert len(res.body["tokens"]) == 8 + + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": { + "prompt_string": "What is this: <__media__>\n", + "multimodal_data": [_get_img_base64(IMG_URL_CAT)], }, }) assert res.status_code == 200 - # Saving a slot that holds image tokens must be rejected (HTTP 501, - # not_supported_error). - res = server.make_request("POST", "/slots/1?action=save", data={ - "filename": "mm_slot_image.bin", + prompt_cat = { + "prompt_string": "What is this: <__media__>\n" + " hello" * (n_ctx_slot - res.body["timings"]["prompt_n"] - 8), + "multimodal_data": [_get_img_base64(IMG_URL_CAT)], + } + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": prompt_cat, }) - assert res.status_code != 200 - assert res.body["error"]["type"] == "not_supported_error" + assert res.status_code == 200 + prompt_n_full = res.body["timings"]["cache_n"] + res.body["timings"]["prompt_n"] + + res = server.make_request("POST", "/slots/0?action=save", data={ + "filename": "mm_slot_large_payload.bin", + }) + assert res.status_code == 200 + + path = os.path.join("tmp", "mm_slot_large_payload.bin") + with open(path, "rb") as f: + data = bytearray(f.read()) + payload_size = struct.unpack_from("=I", data, STATE_FILE_HEADER_SIZE - 4)[0] + assert payload_size > n_ctx_slot # the scenario under test: the payload does not fit in n_ctx + + # drop the image from the slot, then restore it from the file + res = server.make_request("POST", "/completion", data={ + "prompt": "The quick brown fox", + "id_slot": 0, + "cache_prompt": True, + }) + assert res.status_code == 200 + + res = server.make_request("POST", "/slots/0?action=restore", data={ + "filename": "mm_slot_large_payload.bin", + }) + assert res.status_code == 200 + + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": prompt_cat, + }) + assert res.status_code == 200 + assert res.body["timings"]["cache_n"] == prompt_n_full - 1 + assert res.body["timings"]["prompt_n"] == 1 -def test_slot_erase_text_only_on_multimodal(mmproj_server): +def test_slot_restore_media_file_without_mmproj(mmproj_server): server = mmproj_server server.start() - res = server.make_request("POST", "/completion", data={ - "prompt": "The quick brown fox jumps over the lazy dog.", - "id_slot": 1, + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, "cache_prompt": True, + "prompt": { + "prompt_string": "What is this: <__media__>\n", + "multimodal_data": [_get_img_base64(IMG_URL_CAT)], + }, }) assert res.status_code == 200 - prompt_n = res.body["timings"]["prompt_n"] - assert prompt_n > 0 # all tokens are processed - # Erasing a pure-text slot must succeed even though an mmproj is loaded. - res = server.make_request("POST", "/slots/1?action=erase") - assert res.status_code == 200 - - # Re-running the same prompt should process all tokens again. - res = server.make_request("POST", "/completion", data={ - "prompt": "The quick brown fox jumps over the lazy dog.", - "id_slot": 1, - "cache_prompt": True, + res = server.make_request("POST", "/slots/0?action=save", data={ + "filename": "mm_slot_no_mmproj.bin", }) assert res.status_code == 200 - assert res.body["timings"]["prompt_n"] == prompt_n # all tokens are processed again + + # restart the same model without the mmproj: restoring the media file must fail gracefully and leave the slot usable + server.stop() + server.no_mmproj = True + server.start() + + res = server.make_request("POST", "/slots/0?action=restore", data={ + "filename": "mm_slot_no_mmproj.bin", + }) + assert res.status_code == 400 + assert "Cannot restore media tokens without an mmproj" in res.body["error"]["message"] + + # A failed restore must leave the slot empty and usable. + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 1, + "cache_prompt": True, + "prompt": "The quick brown fox", + }) + assert res.status_code == 200 + content = res.body["content"] + + res = server.make_request("POST", "/completions", data={ + "temperature": 0.0, + "top_k": 1, + "id_slot": 0, + "cache_prompt": True, + "prompt": "The quick brown fox", + }) + assert res.status_code == 200 + assert res.body["timings"]["cache_n"] == 0 + assert res.body["content"] == content diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index fffe07a671..9171dbc029 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -86,6 +86,7 @@ class ServerProcess: server_reranking: bool | None = False server_metrics: bool | None = False kv_unified: bool | None = False + swa_full: bool | None = False server_slots: bool | None = False pooling: str | None = None api_key: str | None = None @@ -106,6 +107,7 @@ class ServerProcess: chat_template_file: str | None = None server_path: str | None = None mmproj_url: str | None = None + no_mmproj: bool | None = None media_path: str | None = None sleep_idle_seconds: int | None = None cache_ram: int | None = None @@ -198,6 +200,8 @@ class ServerProcess: server_args.append("--metrics") if self.kv_unified: server_args.append("--kv-unified") + if self.swa_full: + server_args.append("--swa-full") if self.server_slots: server_args.append("--slots") else: @@ -259,6 +263,8 @@ class ServerProcess: server_args.extend(["--chat-template-file", self.chat_template_file]) if self.mmproj_url: server_args.extend(["--mmproj-url", self.mmproj_url]) + if self.no_mmproj: + server_args.append("--no-mmproj") if self.media_path: server_args.extend(["--media-path", self.media_path]) if self.sleep_idle_seconds is not None: