QNN Llama Runner implement IRunner (pytorch#13171)

Summary:
Pull Request resolved: pytorch#13171

This PR makes the Runner for running Qualcomm LlamaModels implement the IRunner interface

Using this, enable running static Llama models with QNN backend inside LlamaDemo Android app

Switched default eval mode to hybrid everywhere

Reviewed By: cccclai

Differential Revision: D79759817

Author: rohansjoshi

examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ModelUtils.java

@@ -21,6 +21,9 @@ public class ModelUtils {
   // MediaTek
   static final int MEDIATEK_TEXT_MODEL = 3;
 
+  // QNN static llama
+  static final int QNN_TEXT_MODEL = 4;
+
   public static int getModelCategory(ModelType modelType, BackendType backendType) {
     if (backendType.equals(BackendType.XNNPACK)) {
       switch (modelType) {
@@ -35,6 +38,8 @@ public static int getModelCategory(ModelType modelType, BackendType backendType)
       }
     } else if (backendType.equals(BackendType.MEDIATEK)) {
       return MEDIATEK_TEXT_MODEL;
+    } else if (backendType.equals(BackendType.QUALCOMM)) {
+      return QNN_TEXT_MODEL;
     }
 
     return TEXT_MODEL; // default

examples/models/llama/runner/runner.h

@@ -18,6 +18,7 @@
 #include <string>
 #include <unordered_map>
 
+#include <executorch/examples/models/llama/tokenizer/llama_tiktoken.h>
 #include <executorch/extension/llm/runner/irunner.h>
 #include <executorch/extension/llm/runner/text_llm_runner.h>
 #include <pytorch/tokenizers/tokenizer.h>
@@ -33,6 +34,7 @@ std::unique_ptr<llm::TextLLMRunner> create_llama_runner(
     float temperature = -1.0f);
 
 std::unique_ptr<tokenizers::Tokenizer> load_llama_tokenizer(
-    const std::string& tokenizer_path);
+    const std::string& tokenizer_path,
+    Version version = Version::Default);
 
 } // namespace example

examples/qualcomm/oss_scripts/llama/CMakeLists.txt

@@ -42,6 +42,8 @@ list(
   ${CMAKE_CURRENT_LIST_DIR}/runner/rpc_mem.h
   ${CMAKE_CURRENT_LIST_DIR}/runner/kv_manager.cpp
   ${CMAKE_CURRENT_LIST_DIR}/runner/kv_manager.h
+  ${EXECUTORCH_SOURCE_DIR}/examples/models/llama/runner/runner.cpp
+  ${EXECUTORCH_SOURCE_DIR}/examples/models/llama/runner/runner.h
 )
 
 list(APPEND _llama_runner__srcs)

examples/qualcomm/oss_scripts/llama/llama.py

@@ -1038,7 +1038,7 @@ def _build_parser():
     parser.add_argument(
         "--model_mode",
         help="Export and inference kv mode, hybrid mode, or lookahead decoding mode",
-        default="kv",
+        default="hybrid",
         choices=["kv", "hybrid", "lookahead"],
         type=str,
     )

examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp

@@ -16,6 +16,7 @@
 
 #include <executorch/backends/qualcomm/runtime/QnnExecuTorch.h>
 #include <executorch/examples/qualcomm/oss_scripts/llama/runner/runner.h>
+#include <executorch/extension/llm/runner/irunner.h>
 #include <executorch/runtime/platform/log.h>
 #include <gflags/gflags.h>
 #include <fstream>
@@ -61,7 +62,7 @@ DEFINE_int32(
     "Total number of tokens to generate (prompt + output).");
 DEFINE_int32(
     eval_mode,
-    0,
+    1,
     "0: TokenGenerator(kv) / 1: HybridMode (prefill+kv) / 2: Lookahead Decoding");
 DEFINE_string(
     kv_updater,
@@ -171,25 +172,26 @@ int main(int argc, char** argv) {
       buf.push_back(c);
     }
   };
-
+  executorch::extension::llm::GenerationConfig config{
+      true,
+      -1,
+      false,
+      FLAGS_seq_len,
+      static_cast<float>(FLAGS_temperature),
+      0,
+      0};
   if (use_tokenized_prompt) {
-    runner.generate(
-        FLAGS_tokenized_prompt.c_str(),
-        use_tokenized_prompt,
-        FLAGS_seq_len,
-        callback);
+    runner.generate_from_prompt_or_file(
+        FLAGS_tokenized_prompt.c_str(), use_tokenized_prompt, config, callback);
   } else {
     // generate tokens & store inference output
     for (int i = 0; i < FLAGS_num_iters; i++) {
       for (const auto& prompt : prompts) {
         std::string formatted_prompt;
         formatted_prompt = get_formatted_prompt(
             prompt, FLAGS_system_prompt, decoder_model_version.get());
-        runner.generate(
-            formatted_prompt.c_str(),
-            use_tokenized_prompt,
-            FLAGS_seq_len,
-            callback);
+        runner.generate_from_prompt_or_file(
+            formatted_prompt.c_str(), use_tokenized_prompt, config, callback);
       }
     }
   }

examples/qualcomm/oss_scripts/llama/runner/runner.cpp

@@ -9,6 +9,7 @@
 // A llama 3.2 runner that includes preprocessing and post processing
 // logic. The module takes in a string as input and emits a string as output.
 
+#include <executorch/examples/models/llama/runner/runner.h>
 #include <executorch/examples/models/llama/tokenizer/llama_tiktoken.h>
 #include <executorch/examples/qualcomm/oss_scripts/llama/runner/client_mem.h>
 #include <executorch/examples/qualcomm/oss_scripts/llama/runner/lhd_token_generator.h>
@@ -59,7 +60,7 @@ void print_performance_report(
     outfile << num_tok;
     outfile.close();
   } else {
-    ET_CHECK_MSG(false, "Error saving the inference speed file");
+    ET_LOG(Error, "Error saving the inference speed file");
   }
 }
 
@@ -84,13 +85,6 @@ void save_logits(
 
 } // namespace
 
-std::unique_ptr<::tokenizers::Tokenizer> load_llama_tokenizer(
-    const std::string& tokenizer_path,
-    Version version) {
-  auto special_tokens = get_special_tokens(version);
-  return llm::load_tokenizer(tokenizer_path, std::move(special_tokens));
-}
-
 Runner::Runner(
     const std::string& decoder_model_version,
     const std::string& model_path,
@@ -177,7 +171,8 @@ Error Runner::load() {
     eos_ids->insert(tokenizer_->encode("<|eot|>", 0, 0).get()[0]);
     eos_ids->insert(tokenizer_->encode("<|end_of_text|>", 0, 0).get()[0]);
   } else {
-    tokenizer_ = load_llama_tokenizer(tokenizer_path_, Version::Default);
+    tokenizer_ =
+        example::load_llama_tokenizer(tokenizer_path_, Version::Default);
     if (tokenizer_ == nullptr) {
       ET_LOG(
           Error, "Failed to load tokenizer with %s", tokenizer_path_.c_str());
@@ -317,13 +312,30 @@ Error Runner::load() {
 }
 
 Error Runner::generate(
+    const std::string& prompt,
+    const llm::GenerationConfig& config,
+    std::function<void(const std::string&)> token_callback,
+    std::function<void(const Stats&)> stats_callback) {
+  return generate_from_pos(prompt, 0, config, token_callback, stats_callback);
+}
+
+Error Runner::generate_from_pos(
+    const std::string& prompt,
+    int64_t start_pos,
+    const llm::GenerationConfig& config,
+    std::function<void(const std::string&)> token_callback,
+    std::function<void(const Stats&)> stats_callback) {
+  // TODO: currently only support start_pos == 0
+  return generate_from_prompt_or_file(
+      prompt, false, config, token_callback, stats_callback);
+}
+
+Error Runner::generate_from_prompt_or_file(
     const std::string& prompt,
     bool tokenized_prompt,
-    int32_t seq_len,
+    const llm::GenerationConfig& config,
     std::function<void(const std::string&)> token_callback,
-    std::function<void(const Stats&)> stats_callback,
-    bool echo,
-    bool warming) {
+    std::function<void(const Stats&)> stats_callback) {
   ET_CHECK_MSG(!prompt.empty(), "prompt cannot be null");
   if (!is_loaded()) {
     stats_.model_load_start_ms = time_in_ms();
@@ -332,6 +344,7 @@ Error Runner::generate(
   }
   stats_.inference_start_ms = time_in_ms();
 
+  int32_t seq_len = config.seq_len;
   seq_len = (seq_len > 0 && seq_len <= context_len_) ? seq_len : context_len_;
   int32_t n_bos = (cur_pos_ == 0) ? 1 : 0;
 
@@ -370,7 +383,7 @@ Error Runner::generate(
       "sequence length exceeded - please increase the seq_len value");
 
   // Prompt Processor first
-  if (token_callback) {
+  if (token_callback && config.echo) {
     token_callback(prompt);
   }
   bool dump_logits = dump_logits_path_.empty() ? false : true;

examples/qualcomm/oss_scripts/llama/runner/runner.h

@@ -21,6 +21,7 @@
 #include <executorch/examples/qualcomm/oss_scripts/llama/runner/kv_manager.h>
 #include <executorch/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.h>
 #include <executorch/examples/qualcomm/oss_scripts/llama/runner/token_generator.h>
+#include <executorch/extension/llm/runner/irunner.h>
 #include <executorch/extension/llm/runner/stats.h>
 #include <executorch/extension/module/module.h>
 #include <pytorch/tokenizers/tokenizer.h>
@@ -33,7 +34,7 @@ enum DecoderModelVersion {
   kQwen2_5,
   kPhi4,
 };
-class Runner {
+class Runner : public executorch::extension::llm::IRunner {
  public:
   explicit Runner(
       const std::string& decoder_model,
@@ -42,25 +43,36 @@ class Runner {
       const std::string& performance_output_path,
       const std::string& dump_logits_path,
       const float temperature = 0.8f,
-      const int eval_mode = EvalMode::kKVCached,
+      const int eval_mode = EvalMode::kHybrid,
       const std::string& kv_updater = "SmartMask",
       const int ngram = 0,
       const int window = 0,
       const int gcap = 0,
       std::unique_ptr<tokenizers::Tokenizer> tokenizer = nullptr);
 
-  bool is_loaded() const;
-  executorch::runtime::Error load();
+  bool is_loaded() const override;
+  executorch::runtime::Error load() override;
   // TODO: Support echo and warming
   executorch::runtime::Error generate(
+      const std::string& prompt,
+      const executorch::extension::llm::GenerationConfig& config,
+      std::function<void(const std::string&)> token_callback = {},
+      std::function<void(const executorch::llm::Stats&)> stats_callback = {})
+      override;
+  executorch::runtime::Error generate_from_pos(
+      const std::string& prompt,
+      int64_t start_pos,
+      const executorch::extension::llm::GenerationConfig& config,
+      std::function<void(const std::string&)> token_callback = {},
+      std::function<void(const executorch::llm::Stats&)> stats_callback = {})
+      override;
+  executorch::runtime::Error generate_from_prompt_or_file(
       const std::string& prompt,
       bool tokenized_prompt,
-      int32_t seq_len,
+      const executorch::extension::llm::GenerationConfig& config,
       std::function<void(const std::string&)> token_callback = {},
-      std::function<void(const executorch::llm::Stats&)> stats_callback = {},
-      bool echo = true,
-      bool warming = false);
-  void stop() {};
+      std::function<void(const executorch::llm::Stats&)> stats_callback = {});
+  void stop() override {};
   executorch::runtime::Result<DecoderModelVersion> get_decoder_model_version();
 
  private:

examples/qualcomm/oss_scripts/llama/targets.bzl

@@ -29,6 +29,7 @@ def define_common_targets():
         exported_deps = [
             "//executorch/extension/module:module",
             "//executorch/extension/llm/sampler:sampler",
+            "//executorch/examples/models/llama/runner:runner",
             "//executorch/examples/models/llama/tokenizer:tiktoken",
             "//executorch/extension/evalue_util:print_evalue",
             "//executorch/backends/qualcomm/runtime:runtime",

extension/android/CMakeLists.txt

@@ -179,6 +179,35 @@ if(EXECUTORCH_BUILD_LLAMA_JNI)
     ${CMAKE_CURRENT_BINARY_DIR}/../../examples/models/llama/runner
   )
 
+  target_sources(
+    executorch_jni PRIVATE
+    ${EXECUTORCH_ROOT}/extension/llm/runner/llm_runner_helper.cpp
+  )
+
+  target_include_directories(
+    executorch_jni PRIVATE
+    ${EXECUTORCH_ROOT}/extension/llm/runner
+  )
+
+  if(QNN_SDK_ROOT)
+    target_sources(
+      executorch_jni PRIVATE
+      ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner/runner.cpp
+      ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner/decoder_runner.cpp
+      ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.cpp
+      ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner/token_generator.cpp
+      ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner/lhd_token_generator.cpp
+      ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner/rpc_mem.cpp
+      ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp
+    )
+
+    target_include_directories(
+      executorch_jni PRIVATE
+      ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner
+    )
+    target_compile_definitions(executorch_jni PRIVATE EXECUTORCH_BUILD_QNN=1)
+  endif()
+
   if(NEURON_BUFFER_ALLOCATOR_LIB)
     target_sources(
       executorch_jni

extension/android/jni/jni_layer_llama.cpp

@@ -29,6 +29,10 @@
 #include <fbjni/ByteBuffer.h>
 #include <fbjni/fbjni.h>
 
+#if defined(EXECUTORCH_BUILD_QNN)
+#include <executorch/examples/qualcomm/oss_scripts/llama/runner/runner.h>
+#endif
+
 #if defined(EXECUTORCH_BUILD_MEDIATEK)
 #include <executorch/examples/mediatek/executor_runner/mtk_llama_runner.h>
 #endif
@@ -124,6 +128,7 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
   constexpr static int MODEL_TYPE_CATEGORY_LLM = 1;
   constexpr static int MODEL_TYPE_CATEGORY_MULTIMODAL = 2;
   constexpr static int MODEL_TYPE_MEDIATEK_LLAMA = 3;
+  constexpr static int MODEL_TYPE_QNN_LLAMA = 4;
 
   static facebook::jni::local_ref<jhybriddata> initHybrid(
       facebook::jni::alias_ref<jclass>,
@@ -174,6 +179,17 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
           model_path->toStdString(),
           tokenizer_path->toStdString(),
           data_path_str);
+#if defined(EXECUTORCH_BUILD_QNN)
+    } else if (model_type_category == MODEL_TYPE_QNN_LLAMA) {
+      std::string decoder_model = "llama3"; // use llama3 for now
+      runner_ = std::make_unique<example::Runner>( // QNN runner
+          decoder_model.c_str(),
+          model_path->toStdString().c_str(),
+          tokenizer_path->toStdString().c_str(),
+          data_path->toStdString().c_str(),
+          "");
+      model_type_category_ = MODEL_TYPE_CATEGORY_LLM;
+#endif
 #if defined(EXECUTORCH_BUILD_MEDIATEK)
     } else if (model_type_category == MODEL_TYPE_MEDIATEK_LLAMA) {
       runner_ = std::make_unique<MTKLlamaRunner>(
@@ -318,6 +334,7 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
           [callback](std::string result) { callback->onResult(result); },
           [callback](const llm::Stats& stats) { callback->onStats(stats); }));
     }
+    return static_cast<jint>(executorch::runtime::Error::InvalidArgument);
   }
 
   void stop() {