Add generate_from_pos and prefill_prompt APIs in llama runner (#3)

* Add generate_from_pos and prefill_prompt APIs in llama runner

* Update runner with correct implementation and update tokenizer commit

* Update prefill_prompt function with proper error checking and remove extra logs from generate_from_pos

* Address review comments

Author: Arpit2601

examples/models/llama/runner/runner.cpp

@@ -356,6 +356,121 @@ Error Runner::generate(
   return Error::Ok;
 }
 
+Result<uint64_t> Runner::prefill_prompt(
+    const std::string& prompt,
+    int64_t& start_pos,
+    int8_t bos,
+    int8_t eos) {
+  std::vector<uint64_t> prompt_tokens =
+      ET_UNWRAP_TOKENIZER(tokenizer_->encode(prompt, bos, eos));
+  int64_t total_context_len = start_pos + prompt_tokens.size();
+  if (total_context_len >= metadata_.at(kMaxContextLen)) {
+    ET_LOG(
+        Error,
+        "num_prompt_tokens + start_pos %ld >= max_seq_len_ %ld, Max seq length exceeded - please increase max seq len value in your export script ",
+        total_context_len,
+        metadata_.at(kMaxContextLen));
+    return Result<uint64_t>(Error::InvalidState);
+  }
+  return text_prefiller_->prefill(prompt_tokens, start_pos);
+}
+
+Error Runner::generate_from_pos(
+    const std::string& prompt,
+    const ::executorch::extension::llm::GenerationConfig& config,
+    int64_t start_pos,
+    std::function<void(const std::string&)> token_callback,
+    std::function<void(const ::executorch::extension::llm::Stats&)>
+        stats_callback) {
+  ET_CHECK_MSG(!prompt.empty(), "Prompt cannot be null");
+  if (!is_loaded()) {
+    stats_->model_load_start_ms = llm::time_in_ms();
+    ET_CHECK_OK_OR_RETURN_ERROR(load());
+    stats_->model_load_end_ms = llm::time_in_ms();
+  }
+  if (config.warming) {
+    ET_LOG(Info, "Doing a warmup run...");
+  }
+
+  RUNNER_ET_LOG(
+      config.warming,
+      "RSS after loading model: %f MiB (0 if unsupported)",
+      llm::get_rss_bytes() / 1024.0 / 1024.0);
+
+  // Wrap the token_callback with print function
+  std::function<void(const std::string&)> wrapped_callback =
+      [token_callback, config](const std::string& piece) {
+        if (token_callback) {
+          token_callback(piece);
+        }
+      };
+  // First token time only measures the time it takes to encode the prompt and
+  // return a response token.
+  stats_->inference_start_ms = llm::time_in_ms();
+
+  // print prompts
+  if (config.echo) {
+    wrapped_callback(prompt);
+  }
+
+  uint64_t prefill_next_token =
+      ET_UNWRAP(prefill_prompt(prompt, start_pos, /*bos=*/0, /*eos=*/0));
+  stats_->first_token_ms = llm::time_in_ms();
+  stats_->prompt_eval_end_ms = llm::time_in_ms();
+  stats_->num_prompt_tokens = start_pos;
+
+  ET_CHECK_MSG(
+      start_pos < metadata_.at(kMaxContextLen),
+      "start_pos %ld >= max_context_len_ %" PRId64
+      ", Max context length exceeded - please increase max context len value in your export script",
+      start_pos,
+      metadata_.at(kMaxContextLen));
+
+  int max_new_tokens =
+      config.resolve_max_new_tokens(metadata_.at(kMaxContextLen), start_pos);
+
+  ET_LOG(Info, "Max new tokens resolved: %d", max_new_tokens);
+
+  // print the first token from prefill. No prev_token so use prefill_next_token
+  wrapped_callback(ET_UNWRAP_TOKENIZER(
+      tokenizer_->decode(prefill_next_token, prefill_next_token)));
+  RUNNER_ET_LOG(
+      config.warming,
+      "RSS after prompt prefill: %f MiB (0 if unsupported)",
+      llm::get_rss_bytes() / 1024.0 / 1024.0);
+
+  // Generate tokens
+  int64_t num_generated_tokens = ET_UNWRAP(text_token_generator_->generate(
+      /*tokens=*/{prefill_next_token},
+      /*start_pos=*/start_pos,
+      /*max_new_tokens=*/max_new_tokens + start_pos - 1,
+      /*temperature=*/temperature_ == -1.0f ? config.temperature : temperature_,
+      /*token_callback=*/wrapped_callback));
+
+  stats_->inference_end_ms = llm::time_in_ms();
+  if (!config.warming) {
+    printf("\n");
+  }
+
+  RUNNER_ET_LOG(
+      config.warming,
+      "RSS after finishing text generation: %f MiB (0 if unsupported)",
+      llm::get_rss_bytes() / 1024.0 / 1024.0);
+
+  stats_->num_generated_tokens = num_generated_tokens;
+
+  if (config.warming) {
+    ET_LOG(Info, "Warmup run finished!");
+  } else {
+    // Do not print report during warmup
+    ::executorch::llm::print_report(*stats_);
+  }
+  if (stats_callback) {
+    stats_callback(*stats_);
+  }
+  return Error::Ok;
+}
+
 Error Runner::warmup(const std::string& prompt, int32_t max_new_tokens) {
   // Create a GenerationConfig for warmup
   llm::GenerationConfig config{

examples/models/llama/runner/runner.h

@@ -64,6 +64,20 @@ class ET_EXPERIMENTAL Runner : public executorch::extension::llm::IRunner {
       int32_t max_new_tokens);
   void stop() override;
 
+  ::executorch::runtime::Error generate_from_pos(
+      const std::string& prompt,
+      const ::executorch::extension::llm::GenerationConfig& config,
+      int64_t start_pos,
+      std::function<void(const std::string&)> token_callback = {},
+      std::function<void(const ::executorch::extension::llm::Stats&)>
+          stats_callback = {}) override;
+
+  ::executorch::runtime::Result<uint64_t> prefill_prompt(
+      const std::string& prompt,
+      int64_t& start_pos,
+      int8_t bos = 0,
+      int8_t eos = 0) override;
+
  private:
   bool shouldStop_{false};
 

extension/android/jni/jni_layer_llama.cpp

@@ -17,6 +17,7 @@
 #include <executorch/examples/models/llava/runner/llava_runner.h>
 #include <executorch/extension/llm/runner/image.h>
 #include <executorch/extension/llm/runner/irunner.h>
+#include <executorch/runtime/core/result.h>
 #include <executorch/runtime/platform/log.h>
 #include <executorch/runtime/platform/platform.h>
 #include <executorch/runtime/platform/runtime.h>
@@ -237,13 +238,17 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
       jint eos) {
     facebook::jni::local_ref<jlongArray> tuple_result =
         facebook::jni::make_long_array(2);
-    if (model_type_category_ != MODEL_TYPE_CATEGORY_MULTIMODAL) {
+    if (model_type_category_ != MODEL_TYPE_CATEGORY_MULTIMODAL &&
+        model_type_category_ != MODEL_TYPE_CATEGORY_LLM) {
       tuple_result->pin()[0] = static_cast<jint>(Error::NotSupported);
       return tuple_result;
     }
 
-    auto&& result = multi_modal_runner_->prefill_prompt(
-        prompt->toStdString(), start_pos, bos, eos);
+    auto&& result = model_type_category_ == MODEL_TYPE_CATEGORY_MULTIMODAL
+        ? multi_modal_runner_->prefill_prompt(
+              prompt->toStdString(), start_pos, bos, eos)
+        : runner_->prefill_prompt(prompt->toStdString(), start_pos, bos, eos);
+
     tuple_result->pin()[0] = static_cast<jint>(Error::Ok);
     if (result.ok()) {
       tuple_result->pin()[1] = static_cast<jlong>(start_pos);
@@ -295,16 +300,28 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
       jlong start_pos,
       facebook::jni::alias_ref<ExecuTorchLlmCallbackJni> callback,
       jboolean echo) {
-    if (model_type_category_ != MODEL_TYPE_CATEGORY_MULTIMODAL) {
-      return static_cast<jint>(Error::NotSupported);
+    if (model_type_category_ == MODEL_TYPE_CATEGORY_MULTIMODAL) {
+      return static_cast<jint>(multi_modal_runner_->generate_from_pos(
+          prompt->toStdString(),
+          seq_len,
+          start_pos,
+          [callback](const std::string& result) { callback->onResult(result); },
+          [callback](const llm::Stats& stats) { callback->onStats(stats); },
+          echo));
+    } else if (model_type_category_ == MODEL_TYPE_CATEGORY_LLM) {
+      executorch::extension::llm::GenerationConfig config{
+          .echo = static_cast<bool>(echo),
+          .seq_len = seq_len,
+          .temperature = temperature_,
+      };
+      return static_cast<jint>(runner_->generate_from_pos(
+          prompt->toStdString(),
+          config,
+          start_pos,
+          [callback](const std::string& result) { callback->onResult(result); },
+          [callback](const llm::Stats& stats) { callback->onStats(stats); }));
     }
-    return static_cast<jint>(multi_modal_runner_->generate_from_pos(
-        prompt->toStdString(),
-        seq_len,
-        start_pos,
-        [callback](const std::string& result) { callback->onResult(result); },
-        [callback](const llm::Stats& stats) { callback->onStats(stats); },
-        echo));
+    return static_cast<jint>(Error::NotSupported);
   }
 
   void stop() {

extension/llm/runner/irunner.h

@@ -10,14 +10,15 @@
 
 #pragma once
 
+#include <executorch/extension/llm/runner/stats.h>
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/result.h>
+
 #include <cstdint>
 #include <functional>
 #include <memory>
 #include <string>
 
-#include <executorch/extension/llm/runner/stats.h>
-#include <executorch/runtime/core/error.h>
-
 namespace executorch {
 namespace extension {
 namespace llm {
@@ -125,6 +126,38 @@ class ET_EXPERIMENTAL IRunner {
    * Stop the generation process.
    */
   virtual void stop() = 0;
+
+  /**
+   * Generate tokens from the given prompt, starting from the given position.
+   * @param prompt The text prompt to LLaMa.
+   * @param start_pos The starting position in KV cache of the input in the LLM.
+   * @param config Generation configuration parameters
+   * @param token_callback What to do after a token is generated.
+   * @param stats_callback What to do with Stats.
+   * @return Error::Ok if successful, an error otherwise
+   */
+  virtual runtime::Error generate_from_pos(
+      const std::string& prompt,
+      const GenerationConfig& config,
+      int64_t start_pos = 0,
+      std::function<void(const std::string&)> token_callback = {},
+      std::function<void(const ::executorch::extension::llm::Stats&)>
+          stats_callback = {}) = 0;
+
+  /**
+   * Prefill an LLama Module with the given text input.
+   * @param prompt The text prompt to LLama.
+   * @param start_pos The starting position in KV cache of the input in the LLM.
+   * It's passed as reference and will be updated inside this function.
+   * @param bos The number of BOS (begin of sequence) token.
+   * @param eos The number of EOS (end of sequence) token.
+   * @return The generated token of the LLama Module after prefill prompt.
+   */
+  virtual ::executorch::runtime::Result<uint64_t> prefill_prompt(
+      const std::string& prompt,
+      int64_t& start_pos,
+      int8_t bos = 0,
+      int8_t eos = 0) = 0;
 };
 
 } // namespace llm

extension/llm/runner/text_token_generator.h

@@ -151,7 +151,7 @@ class ET_EXPERIMENTAL TextTokenGenerator {
    */
   bool inline is_loaded() const {
     // Implementation to check if resources are loaded
-    return tokenizer_->is_loaded() && text_decoder_runner_->is_method_loaded();
+    return text_decoder_runner_->is_method_loaded();
   }
 
  private: