Qualcomm AI Engine Direct - Refactor llama runner (#10578)

Summary:
- Refactored io_manager into five distinct components:
  - DecoderRunner: Module wrapper class.
- PromptProcessor: Handles prompt processing using the decoder and
key-value manager.
- TokenGenerator: Generates tokens using the decoder and key-value
manager.
- KVManager: Manages key-value cache with kv_updater, including data
buffer allocation, cache updates, and buffer updates in TensorImpl.
- IBufferAlloc: Allocates data buffers from RPC memory or client buffer.
  - Validated story llama with CL=128, prefill_ar_len=16, QNN SDK: 2.32
  - Original :

| CL | prefill_ar_len | eval_mode | kv_updater | Model Load Time
(seconds) | Prompt evaluation (seconds) | Generated token rate
(tokens/seconds) | Time to first generated token (seconds) |
| --- | --- | --- | --- | --- | --- | --- | --- |
| 128 | 16 | KV | shift_pointer | 0.3082 | 0.0105 | 237.5553131 | 0.0152
|
| 128 | 16 | KV | smart_mask | 0.2691 | 0.0501 | 258.9103433 | 0.0544 |
| 128 | 16 | hybrid | shift_pointer | 0.3408 | 0.008 | 232.1754892 |
0.008 |
| 128 | 16 | hybrid | smart_mask | 0.3175 | 0.0447 | 237.7134587 |
0.0447 |
  - Refactor:

| CL | prefill_ar_len | eval_mode | kv_updater | Model Load Time
(seconds) | Prompt evaluation (seconds) | Generated token rate
(tokens/seconds) | Time to first generated token (seconds) |
| --- | --- | --- | --- | --- | --- | --- | --- |
| 128 | 16 | KV | shift_pointer |0.2808 | 0.0124 | 234.835 | 0.0124 |
| 128 | 16 | KV | smart_mask | 0.238 | 0.027 | 251.004016 | 0.027 |
| 128 | 16 | hybrid | shift_pointer | 0.3305 | 0.0082 | 229.1122162 |
0.0082 |
| 128 | 16 | hybrid | smart_mask | 0.258| 0.013 |239.463602 | 0.013 | 
- Support multi-turn use case. 
- Validated on story llama. To simulate the scenario, I forced decode
mode to generate 5 tokens each time. Tokens with random length are
inserted after one round of prefill->decode finished.
- Reproduce command: (Note that some whitespaces are missing due to
decoding. But token is actually the same as golden.)
    ```
python examples/qualcomm/oss_scripts/llama/llama.py -b build-android
--checkpoint stories110M.pt --params params.json --tokenizer_model
tokenizer.model --prompt "Once" "a little girl named Lily." "toys and
her favorite toy was a big, red ball." "s mom asked her to help her with
the laundry." "and she put all the clothes in the washing machine."
--temperature 0 --tokenizer_bin tokenizer.bin --llama_model stories110m
--model_mode hybrid --ptq 16a4w -m SM8650 -H ${HOST} -s ${DEVICE}-a
${ARTIFACTS}--max_seq_len 128 --prefill_ar_len 16
  Result:
Once upon a time, there wasa little girl named Lily. She loved to play
with hertoys and her favorite toy was a big, red ball. One day, Lily's
mom asked her to help her with the laundry. Lily was happy to helpand
she put all the clothes in the washing machine.
  After the clothes were
    ```
- Need to apply the below patch to forced decode mode to generate 5
tokens each time.
```
diff --git a/examples/qualcomm/oss_scripts/llama/runner/token_generator.cpp b/examples/qualcomm/oss_scripts/llama/runner/token_generator.cpp
index 8a81b598d..a8ec53cdb 100644
--- a/examples/qualcomm/oss_scripts/llama/runner/token_generator.cpp
+++ b/examples/qualcomm/oss_scripts/llama/runner/token_generator.cpp
@@ -170,7 +170,10 @@ Result<int64_t> TokenGenerator::generate(
       "Failed to set output tensor for module %s",
       forward_name_.c_str());
   // Generate our tokens
-  while (pos < seq_len - 1) {
+  // force decode to generate 5 runs at most
+  int64_t max_pos = std::min(pos + 5, (int64_t)seq_len - 1);
+//   while (pos < seq_len - 1) {
+  while (pos < max_pos) {
     // Fill in the token and position data
     prepare_io(cur_token, pos);
     // Only update data pointer of the cache to the tensor for SHIFT_POINTER
```

Author: shewu-quic

backends/qualcomm/runtime/SharedBuffer.cpp

@@ -22,7 +22,7 @@ std::size_t std::hash<CustomMemTensorInfo>::operator()(
   hash_val ^= std::hash<size_t>()(info.pos);
   hash_val ^= std::hash<size_t>()(info.tensor_bytes);
   for (int i = 0; i < info.rank; ++i) {
-    hash_val ^= info.shape[i];
+    hash_val ^= std::hash<uint32_t>()(info.shape[i]);
   }
   hash_val ^= std::hash<uint32_t>()(info.rank);
   hash_val ^= std::hash<executorch::aten::ScalarType>()(info.dtype);

backends/qualcomm/runtime/backends/QnnBackendFactory.cpp

@@ -80,7 +80,9 @@ std::unique_ptr<BackendConfigParameters> QnnBackendFactory::Create(
           options->soc_info(),
           htp_options);
       backend_params->qnn_mem_manager_ptr_ = std::make_unique<QnnMemManager>(
-          implementation, backend_params->qnn_context_ptr_.get());
+          implementation,
+          backend_params->qnn_context_ptr_.get(),
+          options->log_level());
       backend_params->backend_init_state_ = BackendInitializeState::INITIALIZED;
     } break;
     case QnnExecuTorchBackendType::kGpuBackend:

backends/qualcomm/runtime/backends/QnnMemManager.cpp

@@ -47,9 +47,12 @@ Error QnnMemManager::RegisterIonMem(
   }
   tensor_wrapper->SetMemHandle(handle);
   registered_map_.insert({handle, mem_ptr});
-  QNN_EXECUTORCH_LOG_INFO(
-      "Tensor %s is successfully registered to ION shared memory.",
-      tensor_wrapper->GetName().c_str());
+  if (log_level_ >= QnnExecuTorchLogLevel::kLogLevelInfo) {
+    QNN_EXECUTORCH_LOG_INFO(
+        "Tensor %s is successfully registered to ION shared memory.",
+        tensor_wrapper->GetName().c_str());
+  }
+
   return Error::Ok;
 }
 
@@ -92,9 +95,11 @@ Error QnnMemManager::RegisterCustomMem(
   }
   tensor_wrapper->SetMemHandle(handle);
   registered_map_.insert({handle, mem_ptr});
-  QNN_EXECUTORCH_LOG_INFO(
-      "Tensor %s is successfully registered to custom shared memory.",
-      tensor_wrapper->GetName().c_str());
+  if (log_level_ >= QnnExecuTorchLogLevel::kLogLevelInfo) {
+    QNN_EXECUTORCH_LOG_INFO(
+        "Tensor %s is successfully registered to custom shared memory.",
+        tensor_wrapper->GetName().c_str());
+  }
   return Error::Ok;
 }
 

backends/qualcomm/runtime/backends/QnnMemManager.h

@@ -21,8 +21,11 @@ class QnnMemManager {
  public:
   explicit QnnMemManager(
       const QnnImplementation& implementation,
-      QnnContext* context)
-      : implementation_(implementation), context_(context) {}
+      QnnContext* context,
+      QnnExecuTorchLogLevel log_level)
+      : implementation_(implementation),
+        context_(context),
+        log_level_(log_level) {}
   ~QnnMemManager() {
     DeRegisterMem();
   }
@@ -63,6 +66,7 @@ class QnnMemManager {
 
   const QnnImplementation& implementation_;
   QnnContext* context_;
+  QnnExecuTorchLogLevel log_level_;
   std::unordered_map<Qnn_MemHandle_t, void*> registered_map_;
   std::unordered_map<CustomMemTensorInfo, void*> pre_registered_handles_;
   std::unordered_map<executorch::aten::ScalarType, Qnn_DataType_t>

backends/qualcomm/tests/test_qnn_delegate.py

@@ -3681,7 +3681,7 @@ def test_llama3_2_1b(self):
         if self.pre_gen_pte:
             cmds.extend(["--pre_gen_pte", self.pre_gen_pte])
 
-        golden_start_with = "<|begin_of_text|><|start_header_id|>user<|end_header_id|>"
+        golden_start_with = "<|start_header_id|>user<|end_header_id|>"
         p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL)
         with Listener((self.ip, self.port)) as listener:
             conn = listener.accept()

examples/qualcomm/oss_scripts/llama/CMakeLists.txt

@@ -28,8 +28,18 @@ list(
   ${CMAKE_CURRENT_LIST_DIR}/qnn_llama_runner.cpp
   ${CMAKE_CURRENT_LIST_DIR}/runner/runner.cpp
   ${CMAKE_CURRENT_LIST_DIR}/runner/runner.h
-  ${CMAKE_CURRENT_LIST_DIR}/runner/io_manager.cpp
-  ${CMAKE_CURRENT_LIST_DIR}/runner/io_manager.h
+  ${CMAKE_CURRENT_LIST_DIR}/runner/decoder_runner.cpp
+  ${CMAKE_CURRENT_LIST_DIR}/runner/decoder_runner.h
+  ${CMAKE_CURRENT_LIST_DIR}/runner/prompt_processor.cpp
+  ${CMAKE_CURRENT_LIST_DIR}/runner/prompt_processor.h
+  ${CMAKE_CURRENT_LIST_DIR}/runner/token_generator.cpp
+  ${CMAKE_CURRENT_LIST_DIR}/runner/token_generator.h
+  ${CMAKE_CURRENT_LIST_DIR}/runner/imem_alloc.h
+  ${CMAKE_CURRENT_LIST_DIR}/runner/client_mem.h
+  ${CMAKE_CURRENT_LIST_DIR}/runner/rpc_mem.cpp
+  ${CMAKE_CURRENT_LIST_DIR}/runner/rpc_mem.h
+  ${CMAKE_CURRENT_LIST_DIR}/runner/kv_manager.cpp
+  ${CMAKE_CURRENT_LIST_DIR}/runner/kv_manager.h
 )
 
 list(
@@ -42,7 +52,7 @@ list(
 # build qnn llama runner
 add_executable(qnn_llama_runner ${_llama_runner__srcs})
 target_include_directories(
-  qnn_llama_runner PUBLIC ${_common_include_directories}
+  qnn_llama_runner PUBLIC ${_common_include_directories} ${CMAKE_CURRENT_SOURCE_DIR}/../../../../extension/llm/tokenizers/include
 )
 
 target_link_options_shared_lib(quantized_ops_lib)

examples/qualcomm/oss_scripts/llama/llama.py

@@ -403,7 +403,7 @@ def quantize(self, quant_dtype, args, tokenizer, custom_annotations=()):
         logging.info("Quantizing the model...")
         calibrate(
             self.get_example_inputs(self.llama_meta["get_use_kv_cache"]),
-            args.prompt,
+            args.prompt[0],
             fx_graph_module,
             tokenizer=tokenizer,
             ar_len=self.llama_meta["get_ar_len"],
@@ -756,7 +756,7 @@ def permute(w, heads):
     return quant_attrs
 
 
-def inference(args, quant_attrs, pte_filename, runtime_tokenizer_path, pre_gen_pte=""):
+def inference(args, pte_filename, runtime_tokenizer_path, pre_gen_pte=""):
     workspace = f"/data/local/tmp/{getpass.getuser()}/executorch/single_llama"
 
     if args.model_mode == "kv":
@@ -782,14 +782,13 @@ def post_process():
             outputs.append(f.read())
 
     seq_len = args.max_seq_len
+    multi_prompts = " ".join([f'--prompt "{prompt}"' for prompt in args.prompt])
     runner_args = " ".join(
         [
-            f'--prompt "{args.prompt}"',
+            multi_prompts,
             f"--eval_mode {eval_mode}",
             f"--temperature {args.temperature}",
             f"--system_prompt '{args.system_prompt}'",
-            f"--logits_scale {quant_attrs['scale']}",
-            f"--logits_offset {quant_attrs['zero_point']}",
         ]
     )
 
@@ -932,9 +931,10 @@ def _build_parser():
 
     parser.add_argument(
         "--prompt",
-        help="User prompts for llama.",
+        help="User prompts for Llama. When multiple prompts are entered, a multi-turn conversation will be initiated. Note that this feature is currently for testing purposes only.",
         required=True,
         type=str,
+        nargs="+",
     )
 
     parser.add_argument(
@@ -1018,7 +1018,7 @@ def _build_parser():
 
 def export_llama(args) -> None:
     if args.compile_only and args.pre_gen_pte:
-        exit("Cannot set both compile_only and pre_gen_pte as true")
+        raise RuntimeError("Cannot set both compile_only and pre_gen_pte as true")
 
     if args.model_mode == "kv":
         pte_filename = "kv_llama_qnn"
@@ -1054,29 +1054,15 @@ def export_llama(args) -> None:
     elif args.kv_updater == "shift_pointer":
         args.kv_updater = shift_pointer_updater
     else:
-        exit(f"Using an unkown kv update {args.kv_updater}")
+        raise RuntimeError(f"Using an unknown kv update {args.kv_updater}")
 
     if args.pre_gen_pte:
-        quant_attrs = json.load(
-            open(f"{args.pre_gen_pte}/{pte_filename}_quant_attrs.txt")
-        )
-        inference(
-            args, quant_attrs, pte_filename, runtime_tokenizer_path, args.pre_gen_pte
-        )
-        exit(f"Finish the running pre_gen_pte from {args.pre_gen_pte}")
+        inference(args, pte_filename, runtime_tokenizer_path, args.pre_gen_pte)
+        print(f"Finish the running pre_gen_pte from {args.pre_gen_pte}")
+        return
 
     if args.compile_only:
-        quant_attrs = compile(args, pte_filename, tokenizer)
-        if quant_attrs:
-            json.dump(
-                {
-                    "scale": quant_attrs["scale"],
-                    "zero_point": quant_attrs["zero_point"],
-                },
-                open(f"{args.artifact}/{pte_filename}_quant_attrs.txt", "w"),
-            )
-        else:
-            logging.warning("Quant attributes of the logit is None.")
+        compile(args, pte_filename, tokenizer)
 
         if args.ip and args.port != -1:
             pte_path = f"{args.artifact}/{pte_filename}.pte"
@@ -1089,24 +1075,18 @@ def export_llama(args) -> None:
                         }
                     )
                 )
-        exit(f"Finish compile_only and save to {args.artifact}")
+        print(f"Finish compile_only and save to {args.artifact}")
+        return
+
+    compile(args, pte_filename, tokenizer)
+    inference(args, pte_filename, runtime_tokenizer_path)
 
+
+def main():
+    parser = _build_parser()
+    args = parser.parse_args()
     try:
-        quant_attrs = compile(args, pte_filename, tokenizer)
-        if quant_attrs:
-            logging.info(
-                f"Logit scale: {quant_attrs['scale']}; Logit offset: {quant_attrs['zero_point']}"
-            )
-            json.dump(
-                {
-                    "scale": quant_attrs["scale"],
-                    "zero_point": quant_attrs["zero_point"],
-                },
-                open(f"{args.artifact}/{pte_filename}_quant_attrs.txt", "w"),
-            )
-        else:
-            logging.warning("Quant attributes of the logit is None.")
-        inference(args, quant_attrs, pte_filename, runtime_tokenizer_path)
+        export_llama(args)
     except Exception as e:
         if args.ip and args.port != -1:
             with Client((args.ip, args.port)) as conn:
@@ -1115,12 +1095,6 @@ def export_llama(args) -> None:
             raise Exception(e)
 
 
-def main():
-    parser = _build_parser()
-    args = parser.parse_args()
-    export_llama(args)
-
-
 # flake8: noqa: C901
 if __name__ == "__main__":
     main()

examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp

@@ -34,7 +34,10 @@ DEFINE_string(
     "inference_speed.txt",
     "Records inference speed. For CI purpose.");
 DEFINE_string(tokenizer_path, "tokenizer.bin", "Tokenizer stuff.");
-DEFINE_string(prompt, "The answer to the ultimate question is", "Prompt.");
+DEFINE_string(
+    prompt,
+    "The answer to the ultimate question is",
+    "User prompts for Llama. When multiple prompts are entered, a multi-turn conversation will be initiated. Note that this feature is currently for testing purposes only.");
 DEFINE_string(
     system_prompt,
     "",
@@ -49,10 +52,8 @@ DEFINE_int32(
     "Total number of tokens to generate (prompt + output).");
 DEFINE_int32(
     eval_mode,
-    1,
+    0,
     "0: TokenGenerator(kv) / 1: HybridMode (prefill+kv)");
-DEFINE_double(logits_scale, 0.0, "Logits scale");
-DEFINE_int32(logits_offset, 0, "Logits offset");
 DEFINE_string(
     kv_updater,
     "How to update kv cache. Choose between SmartMask and ShiftPointer",
@@ -72,20 +73,46 @@ std::vector<std::string> CollectPrompts(int argc, char** argv) {
   return prompts;
 }
 
+std::string get_formatted_prompt(
+    const std::string& prompt,
+    const std::string& system_prompt,
+    example::LlamaVersion llama_version) {
+  std::string formatted_prompt;
+  switch (llama_version) {
+    case example::LlamaVersion::kLlama2:
+      formatted_prompt.append(prompt);
+      break;
+    case example::LlamaVersion::kLlama3:
+      if (!system_prompt.empty()) {
+        formatted_prompt.append(
+            "<|start_header_id|>system<|end_header_id|>\n\n");
+        formatted_prompt.append(system_prompt);
+        formatted_prompt.append("<|eot_id|>");
+      }
+      formatted_prompt.append("<|start_header_id|>user<|end_header_id|>\n\n");
+      formatted_prompt.append(prompt);
+      formatted_prompt.append(
+          "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n");
+      break;
+    default:
+      ET_CHECK_MSG(false, "unsupported llama version");
+      break;
+  }
+  return formatted_prompt;
+}
+
 int main(int argc, char** argv) {
   std::vector<std::string> prompts = CollectPrompts(argc, argv);
   gflags::ParseCommandLineFlags(&argc, &argv, true);
   // create llama runner
   example::Runner runner(
-      {FLAGS_model_path},
+      FLAGS_model_path.c_str(),
       FLAGS_tokenizer_path.c_str(),
       FLAGS_performance_output_path.c_str(),
-      FLAGS_logits_scale,
-      FLAGS_logits_offset,
       FLAGS_temperature,
       FLAGS_eval_mode,
-      FLAGS_kv_updater,
-      FLAGS_num_iters);
+      FLAGS_kv_updater);
+  auto llama_version = runner.get_llama_version();
   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());
@@ -97,8 +124,10 @@ int main(int argc, char** argv) {
   // generate tokens & store inference output
   for (int i = 0; i < FLAGS_num_iters; i++) {
     for (const auto& prompt : prompts) {
-      runner.generate(
-          FLAGS_seq_len, prompt.c_str(), FLAGS_system_prompt.c_str(), callback);
+      std::string formatted_prompt;
+      formatted_prompt = get_formatted_prompt(
+          prompt, FLAGS_system_prompt, llama_version.get());
+      runner.generate(formatted_prompt.c_str(), FLAGS_seq_len, callback);
     }
   }
   fout.write(buf.data(), buf.size());

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

@@ -0,0 +1,48 @@
+/*
+ * Copyright (c) Qualcomm Innovation Center, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <executorch/examples/qualcomm/oss_scripts/llama/runner/imem_alloc.h>
+#include <vector>
+
+namespace example {
+/**
+ * @class ClientMem
+ * @brief Final class for client buffer allocation, implementing IBufferAlloc
+ * interface. Used for SHIFT_POINTER mode.
+ */
+class ClientMem final : public IMemAlloc {
+ public:
+  ClientMem(){};
+  // Disable copy constructors, r-value referencing, etc
+  ClientMem(const ClientMem&) = delete;
+  ClientMem& operator=(const ClientMem&) = delete;
+  ClientMem(ClientMem&&) = delete;
+  ClientMem& operator=(ClientMem&&) = delete;
+  virtual ~ClientMem(){};
+  /**
+   * @brief Allocate buffer of specified size with vector.
+   * @param data_size Size of the data to allocate.
+   * @return Pointer to the allocated buffer.
+   */
+  std::byte* allocate(size_t data_size) override {
+    allocated_buffers_.push_back(std::vector<std::byte>(data_size));
+    return allocated_buffers_.back().data();
+  };
+  // Only used for SMART_MASK mode
+  void add_memory_info(
+      void* data_ptr,
+      size_t data_size,
+      executorch::runtime::TensorInfo tensor_info) override {};
+
+ private:
+  std::vector<std::vector<std::byte>> allocated_buffers_;
+};
+
+} // namespace example