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

@@ -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,

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;

tools/server/server-common.cpp

@@ -13,6 +13,8 @@
 #include <sstream>
 #include <fstream>
 #include <limits>
+#include <cstring>
+#include <type_traits>
 
 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<uint32_t>::max()) {
+        throw std::runtime_error("Server tokens state is too large");
+    }
+    return value;
+}
+
+class server_tokens_state_writer {
+public:
+    template <typename T>
+    void write(T value) {
+        static_assert(std::is_trivially_copyable<T>::value, "T must be trivially copyable");
+        const auto * ptr = reinterpret_cast<const char *>(&value);
+        data.insert(data.end(), ptr, ptr + sizeof(value));
+    }
+
+    template <typename T>
+    void write(const std::vector<T> & values) {
+        static_assert(std::is_trivially_copyable<T>::value, "T must be trivially copyable");
+        write(server_tokens_state_u32(values.size()));
+        if (values.empty()) {
+            return;
+        }
+        const auto * ptr = reinterpret_cast<const char *>(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<char> 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<char> take() {
+        data.resize((data.size() + sizeof(llama_token) - 1) / sizeof(llama_token) * sizeof(llama_token), 0);
+        return std::move(data);
+    }
+
+private:
+    std::vector<char> data;
+};
+
+class server_tokens_state_reader {
+public:
+    server_tokens_state_reader(const char * data, size_t size) : data(data), size(size) {}
+
+    template <typename T>
+    T read() {
+        static_assert(std::is_trivially_copyable<T>::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 <typename T>
+    std::vector<T> read_vector() {
+        static_assert(std::is_trivially_copyable<T>::value, "T must be trivially copyable");
+        const uint32_t n_values = read<uint32_t>();
+        // 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<T> 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<char> 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<uint32_t> 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<const char *>(packed.data()), packed.size() * sizeof(llama_token));
+    reader.read<llama_token>(); // format marker
+    if (reader.read<uint32_t>() != SERVER_TOKENS_STATE_VERSION) {
+        throw std::runtime_error("Unsupported server tokens state version");
+    }
+
+    const llama_tokens tokens = reader.read_vector<llama_token>();
+
+    // the media start indices, followed by the media chunks in the same order
+    const std::vector<uint32_t> media_keys = reader.read_vector<uint32_t>();
+    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<char> chunk_data = reader.read_vector<char>();
+        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 {

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<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);
 
@@ -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;

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<server_task_result_cmpl_partial>();
 
@@ -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<char> 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<const llama_token *>(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<size_t>(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);

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

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: