support qnn runner multi iter run

billmguo · facebook-github-bot · commit 09f1baf3b96c · 2025-03-08T10:16:10.000-08:00
Summary: support qnn runner multi iter run

Differential Revision: D70842764
diff --git a/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp b/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp
@@ -25,7 +25,6 @@ DEFINE_string(
     model_path,
     "kv_llama_qnn.pte",
     "Model serialized in flatbuffer format.");
-
 DEFINE_string(
     output_path,
     "outputs.txt",
@@ -48,7 +47,6 @@ DEFINE_int32(
     seq_len,
     128,
     "Total number of tokens to generate (prompt + output).");
-
 DEFINE_int32(
     eval_mode,
     1,
@@ -59,6 +57,7 @@ DEFINE_string(
     kv_updater,
     "How to update kv cache. Choose between SmartMask and ShiftPointer",
     "SmartMask");
+DEFINE_int32(num_iters, 1, "total num of iterations to run.");
 
 int main(int argc, char** argv) {
   gflags::ParseCommandLineFlags(&argc, &argv, true);
@@ -72,7 +71,8 @@ int main(int argc, char** argv) {
       FLAGS_logits_offset,
       FLAGS_temperature,
       FLAGS_eval_mode,
-      FLAGS_kv_updater);
+      FLAGS_kv_updater,
+      FLAGS_num_iters);
   std::vector<char> buf;
   buf.reserve(5 * FLAGS_seq_len); // assume each token is around 5 char
   std::ofstream fout(FLAGS_output_path.c_str());
@@ -82,11 +82,13 @@ int main(int argc, char** argv) {
     }
   };
   // generate tokens & store inference output
-  runner.generate(
-      FLAGS_seq_len,
-      FLAGS_prompt.c_str(),
-      FLAGS_system_prompt.c_str(),
-      callback);
+  for (int i = 0; i < FLAGS_num_iters; i++) {
+    runner.generate(
+        FLAGS_seq_len,
+        FLAGS_prompt.c_str(),
+        FLAGS_system_prompt.c_str(),
+        callback);
+  }
   fout.write(buf.data(), buf.size());
   fout.close();
   return 0;
diff --git a/examples/qualcomm/oss_scripts/llama/runner/io_manager.cpp b/examples/qualcomm/oss_scripts/llama/runner/io_manager.cpp
@@ -167,7 +167,101 @@ void ShiftPointerIoMgr::init_io() {
       break;
   }
 }
+void ShiftPointerIoMgr::reset_io(
+    const std::vector<Result<MethodMeta>>& prefill_methods_meta,
+    const std::vector<Result<MethodMeta>>& kv_methods_meta) {
+  IO* ptr = static_cast<IO*>(data_ptr_.get());
+  std::memset(ptr, 0, sizeof(IO));
+  int32_t k_in_size = (head_dim_ + 1) * kv_cache_len_;
+  int32_t max_ar_len = std::max(kv_ar_len_, prefill_ar_len_);
+
+  int32_t v_cache_size = (num_heads_ + 1) * context_len_ * head_dim_;
+  int32_t k_cache_out_size = num_heads_ * max_ar_len * head_dim_;
+
+  ptr->k_cache_out.clear();
+  ptr->v_cache.clear();
+  // Optionally, reserve space again if you plan to refill them
+  ptr->k_cache_out.reserve(num_layers_);
+  ptr->v_cache.reserve(num_layers_);
+  // Refill the vectors if needed
+  for (int layer = 0; layer < num_layers_; layer++) {
+    ptr->k_cache_out.emplace_back(std::vector<uint8_t>(k_cache_out_size));
+    ptr->v_cache.emplace_back(std::vector<uint8_t>(v_cache_size));
+  }
+
+  auto reset_kv = [&]() {
+    ptr->kv_logits.clear();
+    ptr->kv_logits.resize(kv_ar_len_ * vocab_size_);
+
+    ptr->kv_attention_mask.clear();
+    ptr->kv_attention_mask.resize((kv_ar_len_ * context_len_), 0);
+
+    ptr->k_cache.clear();
+    ptr->k_cache.reserve(num_layers_);
+    for (int layer = 0; layer < num_layers_; layer++) {
+      ptr->k_cache.emplace_back();
+      ptr->k_cache[layer].reserve(num_heads_);
+      for (int head = 0; head < num_heads_; head++) {
+        ptr->k_cache[layer].emplace_back(std::vector<uint8_t>(k_in_size));
+      }
+    }
+  };
+
+  auto reset_prefill = [&]() {
+    ptr->prefill_input_toks.clear();
+    ptr->prefill_input_toks.resize(prefill_ar_len_, 0);
+
+    ptr->prefill_input_pos.clear();
+    ptr->prefill_input_pos.resize(prefill_ar_len_, 0);
+
+    ptr->prefill_attention_mask.clear();
+    ptr->prefill_attention_mask.resize((prefill_ar_len_ * context_len_), 0);
 
+    ptr->prefill_logits.clear();
+    ptr->prefill_logits.resize(prefill_ar_len_ * vocab_size_);
+  };
+  switch (eval_mode_) {
+    case EvalMode::kKVCached:
+      reset_kv();
+      break;
+    case EvalMode::kHybrid:
+      reset_prefill();
+      reset_kv();
+      break;
+    default:
+      break;
+  }
+
+  input_tensors_[kv_forward_name_].clear();
+  input_tensors_[kv_forward_name_].resize(modules_.size());
+  output_tensors_[kv_forward_name_].clear();
+  output_tensors_[kv_forward_name_].resize(modules_.size());
+  k_cache_in_[kv_forward_name_].clear();
+  v_cache_in_[kv_forward_name_].clear();
+  k_cache_out_[kv_forward_name_].clear();
+  v_cache_out_[kv_forward_name_].clear();
+  input_tensors_[prefill_forward_name_].clear();
+  input_tensors_[prefill_forward_name_].resize(modules_.size());
+  output_tensors_[prefill_forward_name_].clear();
+  output_tensors_[prefill_forward_name_].resize(modules_.size());
+  k_cache_in_[prefill_forward_name_].clear();
+  v_cache_in_[prefill_forward_name_].clear();
+  k_cache_out_[prefill_forward_name_].clear();
+  v_cache_out_[prefill_forward_name_].clear();
+
+  switch (eval_mode_) {
+    case EvalMode::kKVCached:
+      prepare_kv_io(kv_methods_meta);
+      break;
+    case EvalMode::kHybrid:
+      prepare_prefill_io(prefill_methods_meta);
+      prepare_kv_io(kv_methods_meta);
+      break;
+    default:
+      ET_CHECK_MSG(false, "unsupported mode");
+      break;
+  }
+}
 void ShiftPointerIoMgr::prepare_kv_io(
     const std::vector<Result<MethodMeta>>& methods_meta) {
   for (int i = 0; i < modules_.size(); ++i) {
@@ -179,7 +273,6 @@ void ShiftPointerIoMgr::prepare_kv_io(
 
   ET_CHECK_MSG(!(kv_forward_name_.empty()), "kv forward name is empty");
   IO* ptr = static_cast<IO*>(data_ptr_.get());
-
   // [I]: input_tokens
   Result<TensorInfo> kv_input_toks = methods_meta[0]->input_tensor_meta(0);
   kv_input_toks_ = std::make_unique<TensorImpl>(
@@ -406,7 +499,6 @@ void ShiftPointerIoMgr::prepare_prefill_io(
       const_cast<TensorImpl::DimOrderType*>(logits->dim_order().data()));
   output_tensors_[prefill_forward_name_][modules_.size() - 1].push_back(
       prefill_logits_.get());
-
   // [O] kv_cache
   int index = 1;
   // In hybrid mode, we use kv mode cache len for v stride since we want to
@@ -885,6 +977,44 @@ void SmartMaskIoMgr::init_io() {
   ptr->init_io_ptrs(shared_ptr, io_bytes_map);
 }
 
+void SmartMaskIoMgr::reset_io(
+    const std::vector<Result<MethodMeta>>& prefill_methods_meta,
+    const std::vector<Result<MethodMeta>>& kv_methods_meta) {
+  init_io();
+  input_tensors_[kv_forward_name_].clear();
+  input_tensors_[kv_forward_name_].resize(modules_.size());
+  output_tensors_[kv_forward_name_].clear();
+  output_tensors_[kv_forward_name_].resize(modules_.size());
+
+  k_cache_in_[kv_forward_name_].clear();
+  v_cache_in_[kv_forward_name_].clear();
+  k_cache_out_[kv_forward_name_].clear();
+  v_cache_out_[kv_forward_name_].clear();
+
+  input_tensors_[prefill_forward_name_].clear();
+  input_tensors_[prefill_forward_name_].resize(modules_.size());
+  output_tensors_[prefill_forward_name_].clear();
+  output_tensors_[prefill_forward_name_].resize(modules_.size());
+
+  k_cache_in_[prefill_forward_name_].clear();
+  v_cache_in_[prefill_forward_name_].clear();
+  k_cache_out_[prefill_forward_name_].clear();
+  v_cache_out_[prefill_forward_name_].clear();
+
+  switch (eval_mode_) {
+    case EvalMode::kKVCached:
+      prepare_kv_io(prefill_methods_meta);
+      break;
+    case EvalMode::kHybrid:
+      prepare_prefill_io(prefill_methods_meta);
+      prepare_kv_io(kv_methods_meta);
+      break;
+    default:
+      ET_CHECK_MSG(false, "unsupported mode");
+      break;
+  }
+}
+
 void SmartMaskIoMgr::prepare_kv_io(
     const std::vector<Result<MethodMeta>>& methods_meta) {
   for (int i = 0; i < modules_.size(); ++i) {
diff --git a/examples/qualcomm/oss_scripts/llama/runner/io_manager.h b/examples/qualcomm/oss_scripts/llama/runner/io_manager.h
@@ -33,6 +33,12 @@ class IoMgrBase {
       std::vector<std::shared_ptr<executorch::extension::Module>>& modules);
   virtual ~IoMgrBase();
   virtual void init_io() = 0;
+  virtual void reset_io(
+      const std::vector<executorch::runtime::Result<
+          executorch::runtime::MethodMeta>>& prefill_methods_meta,
+      const std::vector<
+          executorch::runtime::Result<executorch::runtime::MethodMeta>>&
+          kv_methods_meta) = 0;
   virtual void prepare_prefill_io(
       const std::vector<
           executorch::runtime::Result<executorch::runtime::MethodMeta>>&
@@ -97,6 +103,12 @@ class ShiftPointerIoMgr : public IoMgrBase {
       const bool use_int64_token);
 
   void init_io() override;
+  void reset_io(
+      const std::vector<executorch::runtime::Result<
+          executorch::runtime::MethodMeta>>& prefill_methods_meta,
+      const std::vector<
+          executorch::runtime::Result<executorch::runtime::MethodMeta>>&
+          kv_methods_meta) override;
   void prepare_prefill_io(
       const std::vector<
           executorch::runtime::Result<executorch::runtime::MethodMeta>>&
@@ -199,6 +211,12 @@ class SmartMaskIoMgr : public IoMgrBase {
       const bool use_int64_token);
 
   void init_io() override;
+  void reset_io(
+      const std::vector<executorch::runtime::Result<
+          executorch::runtime::MethodMeta>>& prefill_methods_meta,
+      const std::vector<
+          executorch::runtime::Result<executorch::runtime::MethodMeta>>&
+          kv_methods_meta) override;
   void prepare_prefill_io(
       const std::vector<
           executorch::runtime::Result<executorch::runtime::MethodMeta>>&
diff --git a/examples/qualcomm/oss_scripts/llama/runner/runner.cpp b/examples/qualcomm/oss_scripts/llama/runner/runner.cpp
@@ -48,7 +48,8 @@ Runner::Runner(
     const int32_t logits_offset,
     const float temperature,
     const int eval_mode,
-    const std::string& kv_updater)
+    const std::string& kv_updater,
+    const int num_iters)
     : n_bos_(1),
       n_eos_(1),
       tokenizer_path_(tokenizer_path),
@@ -57,7 +58,8 @@ Runner::Runner(
       logits_offset_(logits_offset),
       temperature_(temperature),
       eval_mode_(static_cast<EvalMode>(eval_mode)),
-      kv_updater_(kv_updater) {
+      kv_updater_(kv_updater),
+      num_iters_(num_iters) {
   for (size_t i = 0; i < models_path.size(); ++i) {
     modules_.push_back(std::make_shared<Module>(
         models_path[i], Module::LoadMode::MmapUseMlockIgnoreErrors));
@@ -280,7 +282,7 @@ Error Runner::generate(
   std::unordered_map<std::string, std::vector<std::vector<Tensor>>>
       input_tensors, output_tensors;
   std::unordered_map<std::string, std::vector<std::vector<EValue>>> inputs;
-  if (!is_loaded()) {
+  if (!is_loaded() || num_iters_ > 1) {
     stats_.model_load_start_ms = time_in_ms();
     ET_CHECK_OK_OR_RETURN_ERROR(load());
     for (auto method_name : method_names_) {
@@ -445,7 +447,11 @@ Error Runner::generate(
   if (stats_callback) {
     stats_callback(stats_);
   }
+  io_mgr_->reset_io(
+      get_methods_meta(prefill_forward_name_),
+      get_methods_meta(kv_forward_name_));
 
+  prompt_.clear();
   return Error::Ok;
 }
 
diff --git a/examples/qualcomm/oss_scripts/llama/runner/runner.h b/examples/qualcomm/oss_scripts/llama/runner/runner.h
@@ -34,7 +34,8 @@ class Runner {
       const int32_t logits_offset,
       const float temperature,
       const int eval_mode,
-      const std::string& kv_updater);
+      const std::string& kv_updater,
+      const int num_iters);
 
   struct Stats {
     // Scaling factor for timestamps - in this case, we use ms.
@@ -117,6 +118,7 @@ class Runner {
   std::vector<std::string> method_names_;
   LlamaVersion llama_version_;
   std::string kv_updater_;
+  int num_iters_;
 };
 
 } // namespace example
