server : allow text-only slot save/restore with mtmd (#25076)

This commit is contained in:
Chipmunk
2026-07-16 21:26:44 +09:00
committed by GitHub
parent a55a8c5266
commit a8dc0e3269
3 changed files with 145 additions and 14 deletions

View File

@@ -207,6 +207,9 @@ 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();

View File

@@ -2011,10 +2011,13 @@ private:
queue_results.send(std::move(res));
}
// if multimodal is enabled, send an error and return false
bool check_no_mtmd(const int id_task) {
if (mctx) {
send_error(id_task, "This feature is not supported by multimodal", ERROR_TYPE_NOT_SUPPORTED);
// 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;
@@ -2502,16 +2505,15 @@ private:
} break;
case SERVER_TASK_TYPE_SLOT_SAVE:
{
if (!check_no_mtmd(task.id)) {
break;
}
const int id_slot = task.slot_action.id_slot;
server_slot * slot = get_slot_by_id(id_slot);
if (slot == nullptr) {
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);
@@ -2519,13 +2521,13 @@ private:
break;
}
const size_t token_count = slot->prompt.tokens.size();
const int64_t t_start = ggml_time_us();
std::string filename = task.slot_action.filename;
std::string filepath = task.slot_action.filepath;
const llama_tokens & tokens = slot->prompt.tokens.get_tokens();
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);
const int64_t t_end = ggml_time_us();
@@ -2543,7 +2545,6 @@ private:
} break;
case SERVER_TASK_TYPE_SLOT_RESTORE:
{
if (!check_no_mtmd(task.id)) break;
const int id_slot = task.slot_action.id_slot;
server_slot * slot = get_slot_by_id(id_slot);
if (slot == nullptr) {
@@ -2590,15 +2591,16 @@ private:
} break;
case SERVER_TASK_TYPE_SLOT_ERASE:
{
if (!check_no_mtmd(task.id)) {
break;
}
const int id_slot = task.slot_action.id_slot;
server_slot * slot = get_slot_by_id(id_slot);
if (slot == nullptr) {
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);

View File

@@ -1,5 +1,7 @@
import pytest
from utils import *
import base64
import requests
server = ServerPreset.tinyllama2()
@@ -96,3 +98,127 @@ def test_slot_erase():
assert res.status_code == 200
assert match_regex("(Whiskers|Flana)+", res.body["content"])
assert res.body["timings"]["prompt_n"] == 21 # all tokens are processed
#
# 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).
#
IMG_URL_CAT = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/91_cat.png"
def _get_img_base64(url: str) -> str:
response = requests.get(url)
response.raise_for_status() # Raise an exception for bad status codes
return base64.b64encode(response.content).decode("utf-8")
@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.
os.environ['LLAMA_MEDIA_MARKER'] = '<__media__>'
mm_server = ServerPreset.tinygemma3()
mm_server.slot_save_path = "./tmp"
mm_server.temperature = 0.0
return mm_server
def test_slot_save_restore_text_only_on_multimodal(mmproj_server):
server = mmproj_server
server.start()
# A pure-text prompt processed on slot 1 of a multimodal server.
res = server.make_request("POST", "/completion", data={
"prompt": "The quick brown fox jumps over the lazy dog.",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
prompt_n = res.body["timings"]["prompt_n"]
assert prompt_n > 0 # all tokens are processed
# Saving a pure-text slot must succeed even though an mmproj is loaded.
res = server.make_request("POST", "/slots/1?action=save", data={
"filename": "mm_slot1.bin",
})
assert res.status_code == 200
n_saved = res.body["n_saved"]
assert n_saved > 0 # the slot KV (prompt + generated tokens) was written
# Restore the saved state into slot 0; it must round-trip exactly.
res = server.make_request("POST", "/slots/0?action=restore", data={
"filename": "mm_slot1.bin",
})
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).
res = server.make_request("POST", "/completion", data={
"prompt": "The quick brown fox jumps over the lazy dog.",
"id_slot": 0,
"cache_prompt": True,
})
assert res.status_code == 200
def test_slot_save_rejected_when_slot_holds_image(mmproj_server):
server = mmproj_server
server.start()
# Process a prompt that actually contains an image on slot 1.
res = server.make_request("POST", "/completions", data={
"temperature": 0.0,
"top_k": 1,
"id_slot": 1,
"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",
})
assert res.status_code != 200
assert res.body["error"]["type"] == "not_supported_error"
def test_slot_erase_text_only_on_multimodal(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,
"cache_prompt": True,
})
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,
})
assert res.status_code == 200
assert res.body["timings"]["prompt_n"] == prompt_n # all tokens are processed again