Android use new prefill API

Author: kirklandsign

extension/android/jni/jni_layer_llama.cpp

@@ -123,7 +123,6 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
   std::unique_ptr<llm::IRunner> runner_;
   std::unique_ptr<executorch::extension::llm::MultimodalRunner>
       multi_modal_runner_;
-  std::vector<llm::MultimodalInput> prefill_inputs_;
 
  public:
   constexpr static auto kJavaDescriptor =
@@ -213,8 +212,7 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
       facebook::jni::alias_ref<ExecuTorchLlmCallbackJni> callback,
       jboolean echo) {
     if (model_type_category_ == MODEL_TYPE_CATEGORY_MULTIMODAL) {
-      std::vector<llm::MultimodalInput> inputs = prefill_inputs_;
-      prefill_inputs_.clear();
+      std::vector<llm::MultimodalInput> inputs;
       if (!prompt->toStdString().empty()) {
         inputs.emplace_back(llm::MultimodalInput{prompt->toStdString()});
       }
@@ -247,9 +245,11 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
   // Contract is valid within an AAR (JNI + corresponding Java code)
   jint prefill_text_input(facebook::jni::alias_ref<jstring> prompt) {
     if (model_type_category_ == MODEL_TYPE_CATEGORY_LLM) {
-      prefill_inputs_.emplace_back(llm::MultimodalInput{prompt->toStdString()});
+      runner_->prefill(prompt->toStdString(), {});
+      return 0;
     } else if (model_type_category_ == MODEL_TYPE_CATEGORY_MULTIMODAL) {
-      multi_modal_runner_->prefill(llm::MultimodalInput{prompt->toStdString()});
+      multi_modal_runner_->prefill(
+          {llm::MultimodalInput{prompt->toStdString()}});
       return 0;
     }
   }
@@ -260,10 +260,10 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
       jint height,
       jint channels) {
     if (model_type_category_ != MODEL_TYPE_CATEGORY_MULTIMODAL) {
-      return Error::InvalidArgument;
+      return static_cast<jint>(Error::InvalidArgument);
     }
     if (image == nullptr) {
-      return Error::InvalidArgument;
+      return static_cast<jint>(Error::InvalidArgument);
     }
     std::vector<llm::Image> images;
     if (image == nullptr) {
@@ -278,7 +278,8 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
         image_data[i] = image_data_jint[i];
       }
       llm::Image image_runner{std::move(image_data), width, height, channels};
-      multi_modal_runner_->prefill(llm::MultimodalInput{std::move(image_runner)});
+      multi_modal_runner_->prefill(
+          {llm::MultimodalInput{std::move(image_runner)}});
     }
 
     return 0;
@@ -290,10 +291,10 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
       jint n_channels,
       jint n_samples) {
     if (model_type_category_ != MODEL_TYPE_CATEGORY_MULTIMODAL) {
-      return Error::InvalidArgument;
+      return static_cast<jint>(Error::InvalidArgument);
     }
     if (audio == nullptr) {
-      return Error::InvalidArgument;
+      return static_cast<jint>(Error::InvalidArgument);
     }
     auto audio_size = audio->size();
     std::vector<uint8_t> audio_data(audio_size);
@@ -304,7 +305,8 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
         audio_data[i] = audio_data_jint[i];
       }
       llm::RawAudio audio_input{audio_data, batch_size, n_channels, n_samples};
-      multi_modal_runner_->prefill(llm::MultimodalInput{std::move(audio_input)});
+      multi_modal_runner_->prefill(
+          {llm::MultimodalInput{std::move(audio_input)}});
     }
     return 0;
   }

extension/llm/runner/irunner.h

@@ -125,6 +125,17 @@ class ET_EXPERIMENTAL IRunner {
       std::function<void(const std::string&)> token_callback,
       std::function<void(const Stats&)> stats_callback) = 0;
 
+  /**
+   * Prefill text inputs, for example to reload chat history.
+   * @param prompt Text prompt to prefill.
+   * @param config Configuration parameters (if non-zero num_bos and num_eos
+   * used)
+   * @return The error code. KV cache position is tracked internally in pos_.
+   */
+  virtual ::executorch::runtime::Error prefill(
+      const std::string& prompt,
+      const GenerationConfig& config = {}) = 0;
+
   /**
    * Stop the generation process.
    */

extension/llm/runner/multimodal_runner.cpp

@@ -62,7 +62,7 @@ Error MultimodalRunner::load() {
     ET_LOG(Info, format, __VA_ARGS__);     \
   }
 
-Error MultimodalRunner::prefill(std::vector<MultimodalInput>& inputs) {
+Error MultimodalRunner::prefill(const std::vector<MultimodalInput>& inputs) {
   if (!is_loaded()) {
     ET_CHECK_OK_OR_RETURN_ERROR(load());
   }

extension/llm/runner/multimodal_runner.h

@@ -126,7 +126,7 @@ class ET_EXPERIMENTAL MultimodalRunner {
    * @return The error code. KV cache position is tracked internally in pos_.
    */
   virtual ::executorch::runtime::Error prefill(
-      std::vector<MultimodalInput>& inputs);
+      const std::vector<MultimodalInput>& inputs);
 
   inline void stop() {
     text_token_generator_->stop();

extension/llm/runner/text_llm_runner.h

@@ -104,13 +104,13 @@ class ET_EXPERIMENTAL TextLLMRunner : public IRunner {
   /**
    * Prefill text inputs, for example to reload chat history.
    * @param prompt Text prompt to prefill.
-   * @param config Configuration parameters for text generation (e.g.,
-   * max_new_tokens, temperature)
+   * @param config Configuration parameters (if non-zero num_bos and num_eos
+   * used)
    * @return The error code. KV cache position is tracked internally in pos_.
    */
   ::executorch::runtime::Error prefill(
       const std::string& prompt,
-      const GenerationConfig& config);
+      const GenerationConfig& config = {}) override;
 
   /**
    * @brief Warms up the model with a sample prompt