From 15586e2d7165570fb3aa7c26e0d442e289ef69de Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Thu, 6 Aug 2026 19:46:40 +0200 Subject: [PATCH] mtmd: add chunk save/load function (#26645) * mtmd: add chunk save/load function * nits * add tests * rn _MAX --> _COUNT --- tests/test-mtmd-c-api.c | 68 +++++++++++ tools/mtmd/clip-impl.h | 8 ++ tools/mtmd/mtmd.cpp | 248 ++++++++++++++++++++++++++++++++++++++++ tools/mtmd/mtmd.h | 10 ++ 4 files changed, 334 insertions(+) diff --git a/tests/test-mtmd-c-api.c b/tests/test-mtmd-c-api.c index 46a038f4a..970d8a600 100644 --- a/tests/test-mtmd-c-api.c +++ b/tests/test-mtmd-c-api.c @@ -1,4 +1,6 @@ #include +#include +#include #include #include "mtmd.h" @@ -62,6 +64,72 @@ int main(void) { } } + // test chunk save/load round-trip + for (size_t i = 0; i < n_chunks; i++) { + const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i); + assert(chunk != NULL); + enum mtmd_input_chunk_type type = mtmd_input_chunk_get_type(chunk); + + // query the required buffer size (out_buf == NULL) + size_t expected_len = 0; + int32_t rc = mtmd_input_chunk_save(chunk, NULL, 0, &expected_len); + printf(" Chunk %zu: save query rc = %d, expected_len = %zu\n", i, rc, expected_len); + assert(rc == 0); + assert(expected_len > 0); + + // saving into a too-small buffer must fail, not crash + char tiny_buf[1]; + rc = mtmd_input_chunk_save(chunk, tiny_buf, sizeof(tiny_buf), NULL); + printf(" Chunk %zu: save into too-small buffer rc = %d (expect non-zero)\n", i, rc); + assert(rc != 0); + + // save into a properly-sized buffer + char * buf = (char *) malloc(expected_len); + assert(buf != NULL); + rc = mtmd_input_chunk_save(chunk, buf, expected_len, NULL); + assert(rc == 0); + + // loading from a truncated buffer must fail gracefully, not crash + if (expected_len > 1) { + mtmd_input_chunk * bad = mtmd_input_chunk_load(buf, expected_len - 1); + printf(" Chunk %zu: load from truncated buffer = %p (expect NULL)\n", i, (void *) bad); + assert(bad == NULL); + } + + // load it back + mtmd_input_chunk * loaded = mtmd_input_chunk_load(buf, expected_len); + assert(loaded != NULL); + + // metadata must match the original chunk + assert(mtmd_input_chunk_get_type(loaded) == type); + assert(mtmd_input_chunk_get_n_tokens(loaded) == mtmd_input_chunk_get_n_tokens(chunk)); + assert(mtmd_input_chunk_get_n_pos(loaded) == mtmd_input_chunk_get_n_pos(chunk)); + + if (type == MTMD_INPUT_CHUNK_TYPE_TEXT) { + size_t n_tok_orig, n_tok_loaded; + const llama_token * tok_orig = mtmd_input_chunk_get_tokens_text(chunk, &n_tok_orig); + const llama_token * tok_loaded = mtmd_input_chunk_get_tokens_text(loaded, &n_tok_loaded); + printf(" Chunk %zu: loaded %zu text tokens (orig %zu), first token %d (orig %d)\n", + i, n_tok_loaded, n_tok_orig, + n_tok_loaded > 0 ? tok_loaded[0] : -1, + n_tok_orig > 0 ? tok_orig[0] : -1); + assert(n_tok_orig == n_tok_loaded); + for (size_t j = 0; j < n_tok_orig; j++) { + assert(tok_orig[j] == tok_loaded[j]); + } + } else if (type == MTMD_INPUT_CHUNK_TYPE_IMAGE || type == MTMD_INPUT_CHUNK_TYPE_AUDIO) { + const char * id_orig = mtmd_input_chunk_get_id(chunk); + const char * id_loaded = mtmd_input_chunk_get_id(loaded); + printf(" Chunk %zu: loaded id '%s' (orig '%s')\n", i, id_loaded, id_orig); + assert(id_orig != NULL && id_loaded != NULL); + assert(strcmp(id_orig, id_loaded) == 0); + } + + mtmd_input_chunk_free(loaded); + free(buf); + } + printf("Chunk save/load round-trip OK\n"); + // Free the chunks mtmd_input_chunks_free(chunks); diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index e1567ee5b..acfecdde8 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -591,6 +591,8 @@ struct clip_image_u8 { } }; +struct mtmd_serialization; // forward declaration + // For images, buf.size() == nx*ny*3 // Memory layout: RGBRGBRGB... // For seq, buf.size() == nx*ny*3*nt @@ -671,6 +673,9 @@ struct clip_image_f32 { return buf.empty(); } + void serialize(struct mtmd_serialization & ser) const; + void deserialize(struct mtmd_serialization & ser); + private: std::vector buf; int nx_ = 0; @@ -752,6 +757,9 @@ struct clip_image_f32_batch { } return new_batch; } + + void serialize(struct mtmd_serialization & ser) const; + void deserialize(struct mtmd_serialization & ser); }; // diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index ff90d6818..6ef4a9d3a 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -22,8 +22,123 @@ #include #include #include +#include #include +// remember to bump this if the serialization format changes +#define MTMD_SERIALIZATION_VERSION 1 + +struct mtmd_serialization { + // note: using 64-bit here for future-proofing + uint64_t version = MTMD_SERIALIZATION_VERSION; + std::vector data; + size_t read_pos = 0; // cursor used when reading + + // for writing + mtmd_serialization(uint64_t version) : version(version) { + write(version); + } + + // for reading + mtmd_serialization(uint64_t version, const char * buf, size_t len) { + // copy buf to data + data.assign(buf, buf + len); + uint64_t ver_in = read(); + if (ver_in != version) { + throw std::runtime_error("version mismatch"); + } + this->version = ver_in; + } + + template + void write(T value) { + static_assert(std::is_trivially_copyable::value && !std::is_same::value, + "T must be trivially copyable and not bool"); + const char * p = reinterpret_cast(&value); + data.insert(data.end(), p, p + sizeof(T)); + } + + template + T read() { + static_assert(std::is_trivially_copyable::value && !std::is_same::value, + "T must be trivially copyable and not bool"); + if (read_pos + sizeof(T) > data.size()) { + throw std::runtime_error("read OOB"); + } + T value; + std::memcpy(&value, data.data() + read_pos, sizeof(T)); + read_pos += sizeof(T); + return value; + } + +}; + +template <> +void mtmd_serialization::write(bool value) { + write(value ? 1 : 0); +} +template <> +bool mtmd_serialization::read() { + return read() != 0; +} + +template <> +void mtmd_serialization::write(std::string value) { + write(value.size()); + data.insert(data.end(), value.begin(), value.end()); +} +template <> +std::string mtmd_serialization::read() { + uint64_t len = read(); + if (read_pos + len > data.size()) { + throw std::runtime_error("read_string OOB"); + } + std::string str(data.data() + read_pos, len); + read_pos += len; + return str; +} + +// only mtmd.cpp needs these, so they're implemented here rather than in clip-impl.h +void clip_image_f32::serialize(mtmd_serialization & ser) const { + // remember to bump MTMD_SERIALIZATION_VERSION if this is changed + // note: buf is intentionally NOT serialized; the loaded clip_image_f32 will always be a placeholder + ser.write(add_viewsep); + ser.write(add_newline); + ser.write((int32_t)nx_); + ser.write((int32_t)ny_); +} +void clip_image_f32::deserialize(mtmd_serialization & ser) { + add_viewsep = ser.read(); + add_newline = ser.read(); + nx_ = ser.read(); + ny_ = ser.read(); + buf.clear(); // always a placeholder after loading +} + +void clip_image_f32_batch::serialize(mtmd_serialization & ser) const { + // remember to bump MTMD_SERIALIZATION_VERSION if this is changed + ser.write(is_audio); + ser.write(entries.size()); + for (const auto & entry : entries) { + entry.serialize(ser); + } +} +void clip_image_f32_batch::deserialize(mtmd_serialization & ser) { + is_audio = ser.read(); + uint64_t n = ser.read(); + constexpr size_t min_entry_bytes = sizeof(uint8_t) * 2 + sizeof(int32_t) * 2; + if (n > (ser.data.size() - ser.read_pos) / min_entry_bytes) { + throw std::runtime_error("entries count exceeds buffer size"); + } + entries.clear(); + entries.reserve(n); + for (uint64_t i = 0; i < n; i++) { + clip_image_f32 entry; + entry.deserialize(ser); + entries.push_back(std::move(entry)); + } +} + // for still image data, layout is RGBRGBRGB... // length of data must be nx * ny * 3 bytes // @@ -83,6 +198,7 @@ enum mtmd_pos_type { MTMD_POS_TYPE_NORMAL, // number of positions equals to number of tokens MTMD_POS_TYPE_MROPE, // qwen-vl mrope style, each image takes max(t,h,w) position indexes MTMD_POS_TYPE_HUNYUANVL, // HunyuanVL mrope + BOI/EOI/newline layout with XD-RoPE dim-3 + MTMD_POS_TYPE_COUNT, // for validation }; struct mtmd_image_tokens { @@ -136,6 +252,30 @@ struct mtmd_image_tokens { id }; } + + void serialize(mtmd_serialization & ser) const { + // remember to bump MTMD_SERIALIZATION_VERSION if this is changed + ser.write(nx); + ser.write(ny); + ser.write((uint32_t)pos); + ser.write(image_idx); + ser.write(n_temporal_merge); + ser.write(id); + batch_f32.serialize(ser); + } + void deserialize(mtmd_serialization & ser) { + nx = ser.read(); + ny = ser.read(); + uint32_t pos_raw = ser.read(); + if (pos_raw >= MTMD_POS_TYPE_COUNT) { + throw std::runtime_error("invalid pos type"); + } + pos = (mtmd_pos_type)pos_raw; + image_idx = ser.read(); + n_temporal_merge = ser.read(); + id = ser.read(); + batch_f32.deserialize(ser); + } }; using mtmd_image_tokens_ptr = std::unique_ptr; @@ -161,6 +301,18 @@ struct mtmd_audio_tokens { id }; } + + void serialize(mtmd_serialization & ser) const { + // remember to bump MTMD_SERIALIZATION_VERSION if this is changed + ser.write(n_tokens); + ser.write(id); + batch_f32.serialize(ser); + } + void deserialize(mtmd_serialization & ser) { + n_tokens = ser.read(); + id = ser.read(); + batch_f32.deserialize(ser); + } }; using mtmd_audio_tokens_ptr = std::unique_ptr; @@ -192,6 +344,66 @@ struct mtmd_input_chunk { } return false; } + + void serialize(mtmd_serialization & ser) const { + // remember to bump MTMD_SERIALIZATION_VERSION if this is changed + ser.write((uint32_t)type); + + ser.write(tokens_text.size()); + for (llama_token tok : tokens_text) { + ser.write((int32_t)tok); + } + + ser.write(tokens_image != nullptr); + if (tokens_image) { + tokens_image->serialize(ser); + } + + ser.write(tokens_audio != nullptr); + if (tokens_audio) { + tokens_audio->serialize(ser); + } + } + void deserialize(mtmd_serialization & ser) { + uint32_t type_raw = ser.read(); + if (type_raw >= MTMD_INPUT_CHUNK_TYPE_COUNT) { + throw std::runtime_error("invalid chunk type"); + } + type = (mtmd_input_chunk_type)type_raw; + + uint64_t n_tokens_text = ser.read(); + // reject before resize() so a tiny corrupted/malicious buffer can't force a huge allocation + if (n_tokens_text > (ser.data.size() - ser.read_pos) / sizeof(int32_t)) { + throw std::runtime_error("tokens_text length exceeds buffer size"); + } + tokens_text.resize(n_tokens_text); + for (uint64_t i = 0; i < n_tokens_text; i++) { + tokens_text[i] = (llama_token)ser.read(); + } + + if (ser.read()) { + tokens_image = std::make_unique(); + tokens_image->deserialize(ser); + } else { + tokens_image.reset(); + } + + if (ser.read()) { + tokens_audio = std::make_unique(); + tokens_audio->deserialize(ser); + } else { + tokens_audio.reset(); + } + + // catch buffers where the declared type doesn't match which payload is actually present, + // so a mismatched chunk can't slip through and null-deref/abort later in an accessor + if (type == MTMD_INPUT_CHUNK_TYPE_IMAGE && !tokens_image) { + throw std::runtime_error("type is IMAGE but tokens_image is missing"); + } + if (type == MTMD_INPUT_CHUNK_TYPE_AUDIO && !tokens_audio) { + throw std::runtime_error("type is AUDIO but tokens_audio is missing"); + } + } }; struct mtmd_input_chunks { @@ -2043,6 +2255,42 @@ void mtmd_input_chunk_free(mtmd_input_chunk * chunk) { } } +int32_t mtmd_input_chunk_save(const mtmd_input_chunk * chunk, char * out_buf, size_t out_len, size_t * expected_out_len) { + try { + mtmd_serialization ser(MTMD_SERIALIZATION_VERSION); + chunk->serialize(ser); + + if (expected_out_len) { + *expected_out_len = ser.data.size(); + } + if (!out_buf) { + // caller is only querying the required size + return 0; + } + if (out_len < ser.data.size()) { + LOG_ERR("%s: out_buf is too small, need %zu bytes, got %zu\n", __func__, ser.data.size(), out_len); + return -1; + } + std::memcpy(out_buf, ser.data.data(), ser.data.size()); + return 0; + } catch (const std::exception & e) { + LOG_ERR("%s: %s\n", __func__, e.what()); + return -1; + } +} + +mtmd_input_chunk * mtmd_input_chunk_load(const char * buf, size_t len) { + try { + mtmd_serialization ser(MTMD_SERIALIZATION_VERSION, buf, len); + mtmd::input_chunk_ptr chunk(new mtmd_input_chunk()); + chunk->deserialize(ser); + return chunk.release(); + } catch (const std::exception & e) { + LOG_ERR("%s: %s\n", __func__, e.what()); + return nullptr; + } +} + // mtmd_image_tokens size_t mtmd_image_tokens_get_n_tokens(const mtmd_image_tokens * image_tokens) { diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 84651f8dc..e5063d114 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -55,6 +55,7 @@ enum mtmd_input_chunk_type { MTMD_INPUT_CHUNK_TYPE_TEXT, MTMD_INPUT_CHUNK_TYPE_IMAGE, MTMD_INPUT_CHUNK_TYPE_AUDIO, + MTMD_INPUT_CHUNK_TYPE_COUNT, // for validation }; // opaque types @@ -232,6 +233,15 @@ MTMD_API llama_pos mtmd_input_chunk_get_n_pos (const mtmd MTMD_API mtmd_input_chunk * mtmd_input_chunk_copy(const mtmd_input_chunk * chunk); MTMD_API void mtmd_input_chunk_free(mtmd_input_chunk * chunk); +// save/load an input chunk to/from a buffer (useful for KV save/load) +// important: only chunk's metadata will be saved, the actual image/audio data will not be saved +// the loaded chunk will always be a placeholder, cannot be used for mtmd_encode() or mtmd_batch_encode() +// out_buf can be nullptr (to query expected_out_len) +// returns 0 on success, non-zero on failure +MTMD_API int32_t mtmd_input_chunk_save(const mtmd_input_chunk * chunk, char * out_buf, size_t out_len, size_t * expected_out_len); +// returns nullptr on failure +MTMD_API mtmd_input_chunk * mtmd_input_chunk_load(const char * buf, size_t len); + // mtmd_image_tokens //