server: save processed mtmd chunks as placeholder (#27278)

tools/mtmd/mtmd.cpp

@@ -2322,23 +2322,12 @@ 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) {
+// returns 0 on success
+static int32_t mtmd_input_chunk_save_impl(const mtmd_input_chunk * chunk, std::vector<char> & out_buf) {
     try {
         mtmd_serialization ser(MTMD_SERIALIZATION_VERSION);
         chunk->serialize(ser);
-
+        out_buf = std::move(ser.data);
-        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());
@@ -2346,6 +2335,35 @@ int32_t mtmd_input_chunk_save(const mtmd_input_chunk * chunk, char * out_buf, si
     }
 }
 
+mtmd_input_chunk * mtmd_input_chunk_get_placeholder(const mtmd_input_chunk * chunk) {
+    // this is hacky, but still faster than copy the whole batch data
+    std::vector<char> buf;
+    if (mtmd_input_chunk_save_impl(chunk, buf) != 0) {
+        return nullptr;
+    }
+    return mtmd_input_chunk_load(buf.data(), buf.size());
+}
+
+int32_t mtmd_input_chunk_save(const mtmd_input_chunk * chunk, char * out_buf, size_t out_len, size_t * expected_out_len) {
+    std::vector<char> buf;
+    if (mtmd_input_chunk_save_impl(chunk, buf) != 0) {
+        return -1;
+    }
+    if (expected_out_len) {
+        *expected_out_len = buf.size();
+    }
+    if (!out_buf) {
+        // caller is only querying the required size
+        return 0;
+    }
+    if (out_len < buf.size()) {
+        LOG_ERR("%s: out_buf is too small, need %zu bytes, got %zu\n", __func__, buf.size(), out_len);
+        return -1;
+    }
+    std::memcpy(out_buf, buf.data(), buf.size());
+    return 0;
+}
+
 mtmd_input_chunk * mtmd_input_chunk_load(const char * buf, size_t len) {
     try {
         mtmd_serialization ser(MTMD_SERIALIZATION_VERSION, buf, len);

tools/mtmd/mtmd.h

@@ -233,6 +233,9 @@ 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);
 
+// similar to mtmd_input_chunk_copy, but returns a placeholder chunk
+MTMD_API mtmd_input_chunk * mtmd_input_chunk_get_placeholder(const 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()

tools/server/server-common.cpp

@@ -507,6 +507,23 @@ void server_tokens::push_back(const mtmd_input_chunk * chunk) {
     }
 }
 
+void server_tokens::push_back_placeholder(const mtmd_input_chunk * chunk) {
+    auto type = mtmd_input_chunk_get_type(chunk);
+    if (type == MTMD_INPUT_CHUNK_TYPE_IMAGE || type == MTMD_INPUT_CHUNK_TYPE_AUDIO) {
+        GGML_ASSERT(has_mtmd);
+        mtmd::input_chunk_ptr new_chunk(mtmd_input_chunk_get_placeholder(chunk));
+        GGML_ASSERT(new_chunk != nullptr && "failed to create placeholder chunk");
+        const size_t n_tokens = mtmd_input_chunk_get_n_tokens(chunk);
+        size_t start_idx = tokens.size();
+        for (size_t i = 0; i < n_tokens; ++i) {
+            tokens.emplace_back(LLAMA_TOKEN_NULL);
+        }
+        map_idx_to_media[start_idx] = std::move(new_chunk);
+    } else {
+        push_back(chunk);
+    }
+}
+
 void server_tokens::push_back(server_tokens & tokens) {
     size_t start_idx = size();
     for (size_t i = 0; i < tokens.size(); i++) {

tools/server/server-common.h

@@ -195,6 +195,10 @@ public:
     // will create a copy of the chunk if it contains non-text data
     void push_back(const mtmd_input_chunk * chunk);
 
+    // same as push_back, but media chunks are stored as placeholders (no image/audio data)
+    // only use this if the chunk will never be encoded again (e.g. it is already in the KV cache)
+    void push_back_placeholder(const mtmd_input_chunk * chunk);
+
     // appends server tokens, updates the media map. copies media chunks.
     void push_back(server_tokens & tokens);
 

tools/server/server-context.cpp

@@ -3416,7 +3416,8 @@ private:
                         // add the mtmd chunk to cache
                         {
                             const auto & chunk = input_tokens.find_chunk(cur_token_idx);
-                            slot.prompt.tokens.push_back(chunk.get()); // copy
+                            // the chunk is already in the KV cache at this point, so we don't need to keep its data around
+                            slot.prompt.tokens.push_back_placeholder(chunk.get());
                         }
 
                         has_mtmd = true;