Add prefill API to MultimodalRunner (#14489)

Co-authored-by: Rohan Joshi <[email protected]>

Author: pytorchbot

extension/llm/runner/multimodal_runner.cpp

@@ -62,6 +62,16 @@ Error MultimodalRunner::load() {
     ET_LOG(Info, format, __VA_ARGS__);     \
   }
 
+Error MultimodalRunner::prefill(std::vector<MultimodalInput>& inputs) {
+  if (!is_loaded()) {
+    ET_CHECK_OK_OR_RETURN_ERROR(load());
+  }
+  for (auto& input : inputs) {
+    ET_UNWRAP(multimodal_prefiller_->prefill(input, pos_));
+  }
+  return Error::Ok;
+}
+
 Error MultimodalRunner::generate(
     const std::vector<MultimodalInput>& inputs,
     const GenerationConfig& config,

extension/llm/runner/multimodal_runner.h

@@ -119,6 +119,15 @@ class ET_EXPERIMENTAL MultimodalRunner {
       std::function<void(const std::string&)> token_callback = {},
       std::function<void(const Stats&)> stats_callback = {});
 
+  /**
+   * Prefill multimodal inputs, for example to reload chat history.
+   * @param inputs A vector of MultimodalInput objects containing images and
+   * text.
+   * @return The error code. KV cache position is tracked internally in pos_.
+   */
+  virtual ::executorch::runtime::Error prefill(
+      std::vector<MultimodalInput>& inputs);
+
   inline void stop() {
     text_token_generator_->stop();
   }