QNN Module API(new frontend) Preliminary Support

CMakeLists.txt

@@ -68,6 +68,13 @@ endif()
 
 # backend options
 option(QNN "Enable QNN" OFF)
+option(QNN_OLD_FRONTEND "Enable Old QNN" OFF)
+if(QNN)
+    add_definitions(-DUSE_QNN) # the USE_QNN should come before cpu subdirectory
+endif()
+if(QNN_OLD_FRONTEND)
+    add_definitions(-DOLD_QNN)
+endif()
 
 if (CMAKE_VERSION VERSION_GREATER_EQUAL "3.24.0")
     cmake_policy(SET CMP0135 NEW)
@@ -116,8 +123,7 @@ include_directories(${PROJECT_SOURCE_DIR}/third_party/pybind11/include)
 
 add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/src/backends/cpu)
 
-if(QNN)
-    add_definitions(-DUSE_QNN)
+if(QNN) # QNN lib
     add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/src/backends/qnn)
 endif()
 

examples/CMakeLists.txt

@@ -53,7 +53,6 @@ endmacro()
 
 
 ## new demos
-
 func_llm_add_executable(benchmark)
 func_llm_add_executable(demo_llama)
 func_llm_add_executable(demo_tinyllama)
@@ -82,6 +81,7 @@ func_vlm_add_executable(demo_imagebind_1mod)
 
 # QNN demo
 if(QNN)
+    func_llm_add_executable(demo_qnn)
     func_llm_add_executable(main_qwen_npu)
 endif()
 

examples/demo_qnn.cpp

@@ -1,6 +1,7 @@
 #include "cmdline.h"
 #include "models/qwen/configuration_qwen.hpp"
 #include "models/qwen/modeling_qwen_npu.hpp"
+#include "models/qwen/modeling_qwen.hpp"
 #include "models/qwen/tokenization_qwen.hpp"
 
 using namespace mllm;
@@ -9,7 +10,7 @@ int main(int argc, char **argv) {
     cmdline::parser cmdParser;
     cmdParser.add<string>("vocab", 'v', "specify mllm tokenizer model path", false, "../vocab/qwen_vocab.mllm");
     cmdParser.add<string>("merge", 'e', "specify mllm merge file path", false, "../vocab/qwen_merges.txt");
-    cmdParser.add<string>("model", 'm', "specify mllm model path", false, "../models/qwen-1.5-1.8b-q8_0.mllm");
+    cmdParser.add<string>("model", 'm', "specify mllm model path", false, "../models/qwen-1.5-1.8b-chat-int8.mllm");
     cmdParser.add<string>("billion", 'b', "[0.5B | 1.8B]", false, "1.8B");
     cmdParser.add<int>("limits", 'l', "max KV cache size", false, 400);
     cmdParser.add<int>("thread", 't', "num of threads", false, 4);
@@ -24,25 +25,29 @@ int main(int argc, char **argv) {
 
     auto tokenizer = QWenTokenizer(vocab_path, merge_path);
     QWenConfig config(tokens_limit, model_billion, RoPEType::HFHUBROPE);
-    auto model = QWenForCausalLM(config);
+    auto model = QWenForCausalLM_NPU(config);
     model.load(model_path);
+    // auto decoding_model = QWenForCausalLM(config);
+    // decoding_model.load("../models/qwen-1.5-1.8b-chat-q4k.mllm");
 
     vector<string> in_strs = {
         " Give me a short introduction to large language model.",
     };
 
     for (int i = 0; i < in_strs.size(); ++i) {
         auto input_str = tokenizer.apply_chat_template(in_strs[i]);
-        auto input_tensor = tokenizer.tokenize(input_str);
+        auto [real_seq_length, input_tensor] = tokenizer.tokenizeWithPadding(input_str, 64, config.vocab_size);
         std::cout << "[Q] " << in_strs[i] << std::endl;
         std::cout << "[A] " << std::flush;
 
         LlmTextGeneratorOpts opt{
             .max_new_tokens = 100,
-            .do_sample = true,
+            .do_sample = false,
             .temperature = 0.3f,
             .top_k = 50,
             .top_p = 0.f,
+            .is_padding = true,
+            .seq_before_padding = real_seq_length,
         };
         model.generate(input_tensor, opt, [&](unsigned int out_token) -> bool {
             auto out_string = tokenizer.detokenize({out_token});
@@ -51,6 +56,24 @@ int main(int argc, char **argv) {
             std::cout << output_string << std::flush;
             return true;
         });
-        std::cout << "FINISH\n";
+
+        LlmTextGeneratorOpts decoding_opt{
+            .max_new_tokens = 100,
+            .do_sample = false,
+            .temperature = 0.3f,
+            .top_k = 50,
+            .top_p = 0.f,
+            .is_padding = false,
+        };
+        // decoding_model.generate(input_tensor, decoding_opt, [&](unsigned int out_token) -> bool {
+        //     auto out_string = tokenizer.detokenize({out_token});
+        //     auto [isOk, print_string] = processOutput(out_string);
+        //     if (isOk) {
+        //         std::cout << print_string << std::flush;
+        //     } else {
+        //         return false;
+        //     }
+        //     return true;
+        // });
     }
 }

examples/main_qwen_npu.cpp

@@ -138,17 +138,12 @@ int main(int argc, char **argv) {
         std::cout << "[Q] " << in_str << std::endl;
         std::cout << "[A] " << std::flush;
 
-        // cpuExe.run(&cpuNet, {input});
-        // auto result = cpuExe.result();
-        // auto token_idx = postProcessing(result[0], input);
-
-        // auto out_token = tokenizer.detokenize({token_idx});
-        // std::cout << out_token << std::flush;
-        // exit(0);
-
         do {
             // 1: Prefill stage using NPU chunk execute
-            npuExe.run(npu_ctx, &npuNet, {input});
+            if (chunk == 1)
+                npuExe.run(npu_ctx, &npuNet, {input});
+            else
+                npuExe.runExp(npu_ctx, &npuNet, {input});
             auto result = npuExe.result();
 
             // inter model for prefill-decode

examples/main_qwen_npu.hpp

@@ -10,11 +10,12 @@ namespace modeling {
 NetTensor *Qwen_FFN_NPU(Context *c, NetTensor *i, int hidden_dim, int ffn_hidden_dim, string name) {
     auto *x = _LinearINT8({i}, hidden_dim, ffn_hidden_dim, false, name + ".gate_proj");
     auto *y = _LinearINT8({i}, hidden_dim, ffn_hidden_dim, false, name + ".up_proj");
-    x = _Dequantize({x}, true, (string)name + ".gate_proj.dequantize", true);
-    y = _Dequantize({y}, true, (string)name + ".up_proj.dequantize", true);
-    x = _SiLU({x}, name + ".silu");
-    x = *x * y;
-    x = _Quantize({x}, true, (string)name + ".down_proj.quantize");
+    x = _SuperSiLU({x,y}, name + ".supersilu");
+    // x = _Dequantize({x}, true, (string)name + ".gate_proj.dequantize", false);
+    // y = _Dequantize({y}, true, (string)name + ".up_proj.dequantize", false);
+    // x = _SiLU({x}, name + ".silu");
+    // x = *x * y;
+    // x = _Quantize({x}, true, (string)name + ".down_proj.quantize");
     x = _LinearINT8({x}, ffn_hidden_dim, hidden_dim, false, name + ".down_proj");
     x = _Dequantize({x}, true, (string)name + ".down_proj.dequantize");
     return x;
@@ -29,9 +30,9 @@ std::vector<NetTensor *> Qwen_CPUNPUAttention(Context *c, NetTensor *x, NetTenso
     k = k->view(1, head_size, seq / chunk, hidden_size);
     v = v->view(1, head_size, seq / chunk, hidden_size);
 
-    q = _Dequantize({q}, true, (string)name + ".q_proj.dequantize");
-    k = _Dequantize({k}, true, (string)name + ".k_proj.dequantize");
-    v = _Dequantize({v}, true, (string)name + ".v_proj.dequantize");
+    q = _Dequantize({q}, true, (string)name + ".q_proj.dequantize", true);
+    k = _Dequantize({k}, true, (string)name + ".k_proj.dequantize", false);
+    v = _Dequantize({v}, true, (string)name + ".v_proj.dequantize", false);
 
     v = _Transpose({v}, {0, 2, 3, 1}, (string)name + ".v_proj.transpose");
 
@@ -153,41 +154,6 @@ NetTensor *Qwen_FFN_CPU_q4k(Context *c, NetTensor *i, int hidden_dim, int ffn_hi
     return x;
 }
 
-void qwen_cpu(Context *c, int vocab_size = 32000, int hidden_dim = 4096, int ffn_hidden_dim = 11008, int mutil_head_size = 32, int cache_max = 200, int seq = 256, int chunk = 2) {
-    auto *i = _Input(c);
-    i = _Embedding({i}, vocab_size, hidden_dim, (string) "model.embed_tokens");
-
-    for (int layer = 0; layer < 24; ++layer) {
-        auto res = _RMSNorm({i}, hidden_dim, 1e-6, (string) "model.layers." + std::to_string(layer) + ".input_layernorm");
-
-        i = *Qwen_CPUAttention(c, res, hidden_dim, hidden_dim / mutil_head_size, mutil_head_size, cache_max, (string) "model.layers." + std::to_string(layer) + ".self_attn", seq, chunk) + i;
-
-        res = _RMSNorm({i}, hidden_dim, 1e-6, (string) "model.layers." + std::to_string(layer) + ".post_attention_layernorm");
-
-        if (layer != 6 && layer != 1 && layer != 2) {
-            i = *Qwen_FFN_CPU(c, res, hidden_dim, ffn_hidden_dim, (string) "model.layers." + std::to_string(layer) + ".mlp") + i;
-        } else {
-            auto name = (string) "model.layers." + std::to_string(layer) + ".mlp";
-
-            auto *x = _LinearINT8({res}, hidden_dim, ffn_hidden_dim, false, name + ".gate_proj");
-            x = _SiLU({x}, name + ".silu");
-            auto *y = _LinearINT8({res}, hidden_dim, ffn_hidden_dim, false, name + ".up_proj");
-            x = *x * y; // x = _Mul( {x, y}, name+".dot");
-
-            auto *i1 = x;
-            x = _LinearINT8({x}, ffn_hidden_dim, hidden_dim, false, name + ".down_proj");
-
-            auto *i2 = x;
-
-            i = *x + i;
-
-            i = _LinearINT8Shadow({i1, i2, i}, ffn_hidden_dim, hidden_dim, false, name + ".down_proj.shadow");
-        }
-    }
-    i = _RMSNorm({i}, hidden_dim, 1e-6, (string) "model.norm");
-    i = _Linear({i}, hidden_dim, vocab_size, false, "lm_head");
-}
-
 void qwen_cpu_q4k(Context *c, int vocab_size = 32000, int hidden_dim = 4096, int ffn_hidden_dim = 11008, int mutil_head_size = 32, int cache_max = 200, int seq = 256, int chunk = 2) {
     auto *i = _Input(c);
     i = _Embedding({i}, vocab_size, hidden_dim, (string) "model.embed_tokens");
@@ -242,9 +208,9 @@ void qwen_npu(Context *c, int vocab_size = 32000, int hidden_dim = 4096, int ffn
 
         res = i;
 
-        i = _RMSNorm({i}, hidden_dim, 1e-6, (string) "model.layers." + std::to_string(layer) + ".post_attention_layernorm");
+        i = _RMSNorm({i}, hidden_dim, 1e-6, (string) "model.layers." + std::to_string(layer) + ".post_attention_layernorm", false);
 
-        i = _Quantize({i}, true, (string) "model.layers." + std::to_string(layer) + ".mlp.up_proj.quantize");
+        // i = _Quantize({i}, true, (string) "model.layers." + std::to_string(layer) + ".mlp.up_proj.quantize");
 
         i = i->view(1, static_cast<int>(seq / chunk / 32), static_cast<int>(32), hidden_dim);
 

include/OpDefined.hpp

@@ -61,6 +61,7 @@ enum OpType {
     MERGEOUTPUT,
     SPLITINPUT,
     IROPE,
+    SUPERSILU,
     OP_NUM
 };
 
@@ -107,6 +108,7 @@ static const vector<string> OpNames = {
     "Range",
     "Where",
     "Replace",
+    "Predictor",
     "SparseLinear",
     "SparseIdLinear",
     "ElasticLinear",
@@ -117,6 +119,7 @@ static const vector<string> OpNames = {
     "MergeOutput",
     "SplitInput",
     "IRoPE",
+    "SuperSiLU",
     "OP_NUM"};
 
 enum TensorFuncType {

include/Types.hpp

@@ -93,9 +93,9 @@ inline std::map<std::vector<int>, ChlType> Chls2Type = {
     {{0, 3, 4, 1, 2}, BWCTH}};
 
 enum TensorType {
-    INPUT_TENSOR = 0,
+    INPUT_TENSOR = 0, // used for input of the model
     NORMAL_TENSOR,
-    OUTPUT_TENSOR,
+    GRAPH_OUTPUT, // used for output of a graph
 };
 
 enum Chl {

scripts/build_qnn_android.sh

@@ -12,6 +12,7 @@ cmake .. \
 -DQNN=ON \
 -DDEBUG=OFF \
 -DTEST=OFF \
--DQUANT=OFF
+-DQUANT=OFF \
+-DQNN_OLD_FRONTEND=ON
 
 make -j4

src/Backend.cpp

@@ -3,6 +3,7 @@
 #include <memory>
 #include <unordered_map>
 #include <mutex>
+#include "Layer.hpp"
 
 namespace mllm {
 extern void registerCPUBackendCreator();
@@ -29,6 +30,9 @@ static std::unordered_map<BackendType, std::shared_ptr<BackendCreator>> &GetBack
 }
 
 const std::shared_ptr<BackendCreator> GetBackendCreator(BackendType type) {
+    if (type == MLLM_QNN) {
+        Layer::use_layername_2_tensorname = false;
+    }
     registerBackend();
 
     auto &gExtraCreator = GetBackendCreatorMap();

src/Backend.hpp

@@ -13,6 +13,11 @@ class Op;
 class Tensor;
 class Backend;
 
+// KVCache map for QNN-CPU KVCache sharing
+#ifdef USE_QNN
+static std::unordered_map<string, Op *> kv_cache_map;
+#endif
+
 class TensorFunction {
 public:
     virtual void setup(vector<Tensor *> outputs, vector<Tensor *> inputs, vector<float> args) = 0;

src/Generate.hpp

@@ -27,6 +27,8 @@ struct LlmTextGeneratorOpts {
     float temperature = 0.7;
     int top_k = 5;
     float top_p = 0.92;
+    bool is_padding = false;
+    int seq_before_padding = 0;
 };
 
 template <typename T>
@@ -47,14 +49,24 @@ enum class LLmTextGeneratorType : int32_t {
 };
 
 class _LlmTextGenerateMethod {
+    bool is_padding = false;
+    int seq_before_padding = 0;
 public:
     virtual ~_LlmTextGenerateMethod() = default;
     virtual unsigned int generate(Tensor &t) = 0;
+    inline void setPadding(bool is_padding, int seq_before_padding) {
+        this->is_padding = is_padding;
+        this->seq_before_padding = seq_before_padding;
+    }
     inline void _tensor_to_vec(Tensor &t, std::vector<float> &scores) {
         assert(t.batch() == 1 && "Batch size of result is not 1. Which is not supported for now.");
         assert(t.head() == 1 && "The 3rd dim of result should be one. e.g.:[1, 1, seq, hidden]");
         int _dims = t.dimension();
         int _seq = t.sequence() - 1;
+        // padding prefill for QNN
+        if (is_padding) {
+            _seq = seq_before_padding - 1;
+        }
         for (int i = 0; i < _dims; ++i) {
             auto value = t.dataAt<float>(0, 0, _seq, i);
             scores.push_back(value);
@@ -144,6 +156,11 @@ class LlmTextGenerator {
             assert(false && "NIY");
             break;
         }
+
+        // padding prefill for QNN
+        if (opt.is_padding) {
+            m_method_class->setPadding(opt.is_padding, opt.seq_before_padding);
+        }
     }
 
     inline unsigned int generate(Tensor &t) {

src/Graph.cpp

@@ -18,9 +18,6 @@ std::string intToStringWithLeadingZero(int num) {
 
 namespace mllm {
 
-#ifdef USE_QNN
-static unordered_map<string, Op *> kv_cache_map;
-#endif
 
 Graph::Graph(const NetParameter &param, Backend *bn,
              unordered_map<string, shared_ptr<Tensor>> &external_tensors,
@@ -149,7 +146,7 @@ void Graph::setUpTensors() {
     // set graph out tensor TensorType
     auto &graph_out_tensors = ops_output_tensors_[op_names_[op_names_.size() - 1]];
     for (auto &t : graph_out_tensors) {
-        t->setTtype(OUTPUT_TENSOR);
+        t->setTtype(GRAPH_OUTPUT);
     }
 
     this->backend_->onSetUpStart(graph_in_tensors, graph_out_tensors);