server : support slot save/restore with media inputs - #26640
Conversation
ngxson
left a comment
There was a problem hiding this comment.
The mtmd change looks quite hacky. I already acknowledge such use case and will push a separate PR for deserialize / serialize mtmd_input_chunk
| std::vector<uint8_t> serialize_media_state() const; | ||
| static server_tokens deserialize_media_state(const llama_tokens & tokens, bool has_mtmd, const uint8_t * data, size_t size); |
There was a problem hiding this comment.
API design: should we simply allow serialize the whole server_tokens object instead of just media state?
There was a problem hiding this comment.
@ngxson
Rather than serializing the whole server_tokens object separately, for now I am planning to serialize only what the llama state does not already store: the media chunks and their positions.
The token sequence is already saved by llama_state_seq_save_file(), and the position and length of each image are already carried there as runs of LLAMA_TOKEN_NULL. Saving the token sequence on the server_tokens side as well would hold the same information twice.
On restore, server_tokens is rebuilt by matching the token sequence recovered from the llama state against the saved chunk positions and the serialized chunks. That keeps the existing checks: that each image range lines up with its LLAMA_TOKEN_NULL run, and that no orphan media token is left without a descriptor.
This is how I plan to proceed.
There was a problem hiding this comment.
If I understand correctly, the llama_state_seq_save_file takes the list of llama_token and directly serialize it as raw bytes, completely unrelated to the saved state. In other words, I see the API like this:
// this is what we currently have:
llama_state_seq_save_file(..., const llama_token * tokens, size_t n_token_count);
// equivalent to:
// data_len % sizeof(llama_token) == 0
llama_state_seq_save_file(..., const char * data, size_t data_len);My idea is that instead of storing the list of plain tokens, we can just have a server-specific serializer and the data will be stored in the place of tokens. Something like this:
std::vector<char> data = slot.serialize();
GGML_ASSERT(data.size() % 4 == 0);
const llama_token * data_ptr = reinterpret_cast<const llama_token *>(data);
llama_state_seq_save_file(..., data_ptr, data.size() / 4);CC @ggerganov not sure if you are agree on this solution, or we should probably have a specific version of llama_state_seq_save_file that stores raw bytes as addition info ?
There was a problem hiding this comment.
Yes, that's a good point. The llama_state_seq_save/load API should be promoted to save any user data - no reason to restrict this to tokens only. We can do the trick that @ngxson proposes for now, and in a follow-up PR we can update the API.
There was a problem hiding this comment.
@ngxson @ggerganov
Understood. One question before I update the serialization.
My preference would be to keep this simple and use a single server_tokens packed representation for both text-only and multimodal slots, with the media state simply empty for text-only slots.
Would you prefer that, or should we preserve the existing text-only on-disk format for compatibility and only use the packed representation for media slots?
The current PR says that below additional infomation.
Text-only slot save / restore behavior is unchanged.
So I’d like your opinion on whether we should keep that compatibility guarantee or simplify the serialization by using a single representation.
I assume the generic raw user-data API itself can remain a follow-up, as discussed.
There was a problem hiding this comment.
A proper serializer is cleaner and will save us some headaches in the long run, I don't see why we don't impl it right now (plus, we can impl versioning along the way, similar to what mtmd save/load already having)
On compatibility: Simply write a LLAMA_TOKEN_NULL as the magic bytes for this new serializing format. Old save files generated prior to this PR never have LLAMA_TOKEN_NULL inside it anyway. If we read a file that doesn't start with null token, fallback to old version
So header for this format will be: 4 bytes LLAMA_TOKEN_NULL, followed by 4 bytes for version number
|
PTAL on the chunk save/load API: #26645 |
|
@ngxson The cache-only |
|
@ggerganov Once #26645 is merged, I will rebase this PR and update the implementation to use its |
0478749 to
badb515
Compare
badb515 to
5bddd7b
Compare
|
Following the suggested direction, I have updated the implementation to serialize / deserialize the whole
The PR description summarizes the changes and the validation results. The implementation is ready for another review. Please take a look. |
| STATE_FILE_HEADER_SIZE = 12 | ||
| # the token payload holds a packed server_tokens object (see server_tokens::serialize()): | ||
| # LLAMA_TOKEN_NULL(4) version(4) n_tokens(4) tokens, media list, zero padding to whole tokens | ||
| PACKED_HEADER_SIZE = 12 # LLAMA_TOKEN_NULL, version, n_tokens | ||
| LLAMA_TOKEN_NULL = 0xFFFFFFFF # -1 read back as an unsigned word | ||
|
|
||
| # media list layout in the packed payload: n_media(4), then per image: start_idx(4) chunk_size(4) chunk blob | ||
| N_MEDIA_FIELD_SIZE = 4 | ||
| START_IDX_FIELD_SIZE = 4 | ||
| CHUNK_SIZE_FIELD_SIZE = 4 |
There was a problem hiding this comment.
We don't need to test these. Those are detailed impl, the test system only need to test the API surface
There was a problem hiding this comment.
Done.
I kept coverage of observable API behavior, and the other tests were removed because they exercised packed-format implementation details outside the API boundary.
| template <typename T> | ||
| void server_tokens_state_write(std::vector<char> & data, T value) { | ||
| const auto * ptr = reinterpret_cast<const char *>(&value); | ||
| data.insert(data.end(), ptr, ptr + sizeof(value)); |
There was a problem hiding this comment.
better to make sure T is copyable trivially (same pattern as mtmd_serialization)
There was a problem hiding this comment.
Added std::is_trivially_copyable<T> checks to the reader / writer templates.
| // the token payload of a sequence state file holds a packed server_tokens object (see server_tokens::serialize()), so it can be longer than the number of tokens the slot holds. | ||
| // read its size from the file header, falling back to n_ctx if the header cannot be trusted - llama_state_seq_load_file() then reports the malformed file. | ||
| // TODO: remove this once llama_state_seq_save_file() can store arbitrary user data | ||
| static size_t state_file_payload_size(const std::string & filepath, size_t fallback) { | ||
| constexpr std::streamoff header_size = 3 * sizeof(uint32_t); // magic, version, payload size |
There was a problem hiding this comment.
this is quite hacky tbh, we can simply modify llama_state_seq_load_file to return wanted size via n_token_count_out if tokens_out == nullptr (actual state load will be skipped in that case). such change will be less than 10 lines of code
There was a problem hiding this comment.
Done. Removed the server-side header parsing and updated llama_state_seq_load_file() to return the required size via n_token_count_out when tokens_out == nullptr.
| static_assert(sizeof(llama_token) == sizeof(uint32_t), "unexpected llama_token size"); | ||
|
|
||
| std::vector<char> data; | ||
| server_tokens_state_write(data, (llama_token) LLAMA_TOKEN_NULL); |
There was a problem hiding this comment.
If we already had the RAII server_tokens_state_reader , why not also having server_tokens_state_writer ?
You can abstract the mtmd_chunk save into server_tokens_state_writer so that server_tokens::serialize() will be minimal and more readable, similar to serialize() functions in mtmd.
Plus, having write<std::vector<T>> will make the code much cleaner:
writer.write(LLAMA_TOKEN_NULL);
writer.write(SERVER_TOKENS_STATE_VERSION);
writer.write(tokens); // overload write<std::vector<T>>
std::vector<size_t> map_keys;
// TODO: copy map_idx_to_media keys to map_keys
writer.write(map_keys);
// custom writer for mtmd_chunks
for (const auto & item : map_idx_to_media) {
const auto * chunk = item.second.get();
std::vector<char> chunk_data;
// TODO: get chunk_size, resize chunk_data, then write to chunk_data
writer.write(chunk_data); // will copy, but it's cleaner than mtmd_input_chunk_save() writing directly to output buf
}There was a problem hiding this comment.
Done. I added server_tokens_state_writer to pair with the reader and moved scalar/vector and image-chunk serialization into it. server_tokens::serialize() now only defines the field order and performs media-specific validation.
|
I've updated the implementation based on the review feedback:
|
| if (n_tokens == 0 || n_pos <= 0) { | ||
| throw std::runtime_error("Invalid image chunk in server tokens"); | ||
| } | ||
| if (id == nullptr || id[0] == '\0') { | ||
| throw std::runtime_error("Slot save with image tokens without an ID is not supported"); |
There was a problem hiding this comment.
these checks are too defensive, they never trigger in practice - should remove them
server_tokens can't handle either of these, something else will break before we even get into this code
There was a problem hiding this comment.
Done. Removed the redundant serialize-side checks as suggested.
| if (mtmd_input_chunk_get_type(chunk.get()) != MTMD_INPUT_CHUNK_TYPE_IMAGE) { | ||
| throw std::runtime_error("Unsupported media type in server tokens state"); |
There was a problem hiding this comment.
There was a problem hiding this comment.
ok for removing this
in general, these checks are too defensive as I pointed out. most of the checks you added are either existing somewhere else in the code base, or can be absorbed into something like server_tokens::validate()
| const char * id = mtmd_input_chunk_get_id(chunk.get()); | ||
| if (id == nullptr || id[0] == '\0') { | ||
| throw std::runtime_error("Image ID is missing in server tokens state"); | ||
| } | ||
| const size_t n_chunk_tokens = mtmd_input_chunk_get_n_tokens(chunk.get()); | ||
| const llama_pos n_pos = mtmd_input_chunk_get_n_pos(chunk.get()); | ||
| if (n_chunk_tokens == 0 || n_pos <= 0 || start_idx < last_end || | ||
| start_idx > tokens.size() || n_chunk_tokens > tokens.size() - start_idx) { | ||
| throw std::runtime_error("Invalid image range in server tokens state"); | ||
| } | ||
| for (size_t j = start_idx; j < start_idx + n_chunk_tokens; ++j) { | ||
| if (tokens[j] != LLAMA_TOKEN_NULL) { | ||
| throw std::runtime_error("Image range does not match server tokens"); | ||
| } | ||
| } |
There was a problem hiding this comment.
all these checks should be absorbed by server_tokens::validate()
There was a problem hiding this comment.
Done. Moved the media/token consistency checks into server_tokens::validate().
There was a problem hiding this comment.
llama.h documentation need to be updated to reflect this new behavior
There was a problem hiding this comment.
Done, add comment
| for (size_t i = 0; i < tokens.size();) { | ||
| if (tokens[i] != LLAMA_TOKEN_NULL) { | ||
| ++i; | ||
| continue; | ||
| } | ||
| const auto it = result.map_idx_to_media.find(i); | ||
| if (it == result.map_idx_to_media.end()) { | ||
| throw std::runtime_error("Image token has no descriptor"); | ||
| } | ||
| i += mtmd_input_chunk_get_n_tokens(it->second.get()); | ||
| } |
There was a problem hiding this comment.
this should already be part of one of existing method inside server_tokens, consider reuse existing functions, not to duplicate the logic
There was a problem hiding this comment.
Done. Removed the duplicate token scan and now rely on server_tokens::validate().
4a67feb to
94ab2a6
Compare
|
I've updated the implementation based on the latest review feedback:
|
| } | ||
| // Gate on slot content, consistent with save/restore. | ||
| if (!check_slot_no_media(*slot, task.id)) { | ||
| if (!check_slot_no_media_for_erase(*slot, task.id)) { |
There was a problem hiding this comment.
Do we actually need to check for "no media" here?
There was a problem hiding this comment.
@ggerganov
I agree. Since this PR changes how media chunks are handled, I don't think this check is necessary anymore.
On the other hand, you've already approved the PR. Should I remove it in this PR?
There was a problem hiding this comment.
Done. Removed the media check for slot erase.
| 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; | ||
| } |
There was a problem hiding this comment.
In a follow-up PR, should be relatively simple to append the prompt.checkpoints so that we can restore them too.
|
|
||
| llama_tokens get_text_tokens() const; | ||
|
|
||
| // packed into the token payload of a sequence state file: [LLAMA_TOKEN_NULL][version][n_tokens][tokens][n_media][start_idx]...([chunk_size][media chunk])...[zero padding] |
There was a problem hiding this comment.
no need to specify this in the comment, the code also reflect the details
There was a problem hiding this comment.
Done. Removed this comment.
| if (reader.read<uint8_t>() != 0) { | ||
| throw std::runtime_error("Invalid padding in server tokens state"); | ||
| } | ||
| } |
There was a problem hiding this comment.
do we really need to be this defensive? what's the problem if padding is not 0 ?
There was a problem hiding this comment.
Done. Removed padding validation.
| if (type != MTMD_INPUT_CHUNK_TYPE_IMAGE && type != MTMD_INPUT_CHUNK_TYPE_AUDIO) { | ||
| throw std::runtime_error("Unsupported media type in server tokens state"); |
There was a problem hiding this comment.
I thought you said to remove this check?
There was a problem hiding this comment.
I apologize. It must have slipped my mind.
Removed this check and moved to validate().
Signed-off-by: Gabe Goodhart <ghart@us.ibm.com> * origin/master: (383 commits) cmake : introduce semantic versioning (ggml-org#26839) gguf : harden loader against malformed tensor dims and metadata types (ggml-org#25596) kleidiai: Add runtime feature detection mechanism for aarch64/kleidiai (ggml-org#26076) model : disallow integer dflash sliding_window_pattern (ggml-org#26900) sync : ggml cmake : add config version support (ggml/1582) server : support slot save/restore with media inputs (ggml-org#26640) ui: add read_media tool (ggml-org#25877) opencl: default FA c8 cluster width to 16 on X1E (ggml-org#26433) tests : update speculative params (ggml-org#26925) vulkan: add TQ2_0 (ternary) support (ggml-org#25850) wavtokenizer-dec : bound posnet/convnext block_count against n_layer_all (ggml-org#26892) convert : handle per_layer_config in Gemma4 (transformers 5.15) (ggml-org#26882) opencl: use flat mv q5_k when weight exceeds image1d_buffer_t limit (ggml-org#26880) chat : fix muse-glimmer detection of tool calls after EOM (ggml-org#26879) ci : add missing release check (ggml-org#26923) CUDA: only disable CUDA graphs when mul_mat_id actually needs a stream sync (ggml-org#26802) cuda : add warp-per-row wkv7 kernel for single-token decode (ggml-org#26111) spec : update speculative-simple (ggml-org#26904) chat : tighten bare function parsing for Qwen models (ggml-org#26793) ...
* 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()
Overview
This PR adds support for saving and restoring slots that contain media input. Upstream stores the KV sequence state and a plain token list, but does not preserve the media positions and
mtmd_input_chunkstate required to reconstruct a media-containing slot, so saving such a slot is currently rejected with HTTP 501.This PR uses
mtmd_input_chunk_save()/mtmd_input_chunk_load()from #26645 and stores the completeserver_tokensstate - tokens, media start positions, and serialized media chunks - as a single versioned packed payload. The same format is used for both text-only and media-containing slots.Because the packed payload can be larger than the slot's
n_ctx, the restore path first queries the saved payload size throughllama_state_seq_load_file()before allocating the restore buffer. After restore, the reconstructed media state can participate in the existing media matching / prefix reuse path.Closes #25854
Serialization format
Newly saved slots use the same packed
server_tokensformat for both text-only and media-containing slots.The current format version is
1.n_tokensandn_mediaare stored as element counts for their respective vectors. All media start indices are written first, followed by the serialized media chunks in the same order.n_media = 0llama_tokenwordsBackward compatibility
Slot files written before this PR contain a plain token list in the token payload.
The restore path distinguishes the formats as follows:
LLAMA_TOKEN_NULL: versioned packedserver_tokensformatTherefore, text-only slot files saved by older servers can still be restored after this PR.
Newly saved text-only slots use the packed format as well, so restoring a new-format slot file on an older server is not supported.
server_tokensserializationThis PR adds
server_tokens::serialize()andserver_tokens::deserialize().Internal reader / writer helpers handle scalar and vector fields. Media chunks are serialized through
mtmd_input_chunk_save()and restored throughmtmd_input_chunk_load().The packed state stores the token sequence, media start positions, and serialized media chunks.
deserialize()validates the format and payload boundaries, reconstructs the media chunks, checks the supported media types (IMAGE and AUDIO), and validates trailing padding. After deserialization,server_tokens::validate()validates text token IDs against the current vocabulary and checks the mapping between media chunks and their token ranges.An
mmprojis required only when the restored state actually contains media.Slot save
llama_state_seq_save_file()returning0is handled as an error. Then_savedAPI field continues to report the logical number of tokens held by the slot, rather than the size of the packed representation.Slot restore
Serialized media chunks can make the packed payload larger than the slot's
n_ctx.This PR therefore allows
llama_state_seq_load_file()to query the saved payload count whentokens_out == nullptr, without restoring the sequence state. The server uses that count to allocate the payload buffer before performing the actual load.After deserialization, the logical token count is checked against the slot's
n_ctxand the restored tokens are validated against the current model. If restore or validation fails, the slot is cleared before returning the error.Slot erase
Erase remains unsupported for slots containing media; its media gate is unchanged.
Testing
Local validation:
test-mtmd-c-apiwith assertions enabled: passedtest_slot_save.py+test_vision_api.py: 28 passedMain API-level cases verified:
n_ctxmmproj, leaving the slot usableBenchmark
Qwen3.5-2B Q8_0 + BF16 projector / RTX 4070 / Flash Attention (
-fa on) /--cache-ram 03316-token multimodal prompt with a 3303-token reusable prefix, median of 3 runs:
With a different image, the restored image prefix was not reused.
Additional information
mmprojconfiguration.--swa-full.llama_token-width packing and padding workaround.Requirements
I have read and agree with the contributing guidelines
AI usage disclosure: YES
mtmd_input_chunkserialization API introduced in mtmd: add chunk save/load function #26645.