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()
This commit is contained in:
Chipmunk
2026-08-12 12:20:28 +02:00
committed by GitHub
parent 4dd127584b
commit 5d9e5ac30e
7 changed files with 593 additions and 78 deletions
+47 -34
View File
@@ -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);