[Kernels] Reference impl and UT for Dense MHA with dynamic quantizati…

…on (#612)



---------

Co-authored-by: Wang,Zhe <zhe1.wang@intel.com>

intel_extension_for_transformers/backends/neural_engine/kernels/include/interface.hpp

@@ -88,6 +88,7 @@ class SPARSE_API_ kernel_proxy : public proxy_base<kernel_t, std::shared_ptr<con
  public:
   inline const jd::kernel_kind& kernel_kind() const { return get_sp()->kd()->kernel_kind(); }
   void execute(const std::vector<const void*>& rt_data) const;
+  size_t get_workspace_size() const;
 };
 
 //// The following paragraphs are the various derived kernels and its descriptors.
@@ -164,6 +165,13 @@ class SPARSE_API_ transpose_mha_desc : public kernel_desc_proxy {
   virtual ~transpose_mha_desc() {}
 };
 
+class SPARSE_API_ dyn_quantize_mha_desc : public kernel_desc_proxy {
+ public:
+  dyn_quantize_mha_desc() {}
+  explicit dyn_quantize_mha_desc(const operator_desc& op_desc) : kernel_desc_proxy(op_desc) {}
+  virtual ~dyn_quantize_mha_desc() {}
+};
+
 /**
  * @brief Derived proxy class, interfacing to the real/cached sparse_matmul_t.
  */
@@ -237,5 +245,12 @@ class SPARSE_API_ transpose_mha : public kernel_proxy {
   virtual ~transpose_mha() {}
 };
 
+class SPARSE_API_ dyn_quantize_mha : public kernel_proxy {
+ public:
+  dyn_quantize_mha() {}
+  explicit dyn_quantize_mha(const kernel_desc_proxy& kdp) : kernel_proxy(kdp) {}
+  virtual ~dyn_quantize_mha() {}
+};
+
 }  // namespace jd
 #endif  // ENGINE_SPARSELIB_INCLUDE_INTERFACE_HPP_

intel_extension_for_transformers/backends/neural_engine/kernels/include/kernel.hpp

@@ -55,6 +55,7 @@ class kernel_t {
   // init kernel_t
   virtual bool init() = 0;
   virtual bool execute(const std::vector<const void*>& rt_data) const = 0;
+  virtual size_t get_workspace_size() const { return 0; }
 
  public:
   const std::shared_ptr<const kernel_desc_t>& kd() const { return kd_; }

intel_extension_for_transformers/backends/neural_engine/kernels/include/kernels/dyn_quantize_mha_ref.hpp

@@ -0,0 +1,100 @@
+//  Copyright (c) 2021 Intel Corporation
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+
+#ifndef ENGINE_SPARSELIB_INCLUDE_KERNELS_DYN_QUANTIZE_MHA_REF_HPP_
+#define ENGINE_SPARSELIB_INCLUDE_KERNELS_DYN_QUANTIZE_MHA_REF_HPP_
+
+#include <memory>
+#include <vector>
+
+#include "amx_utils.hpp"
+#include "cpu_isa.hpp"
+#include "dyn_quantize_mha_types.hpp"
+#include "kernel.hpp"
+#include "kernel_desc.hpp"
+#include "operator_desc.hpp"
+#include "utils.hpp"
+
+namespace jd {
+
+/**
+ * @brief
+ *       Q       K       V
+ *       |       |       |
+ *       |       |       |
+ *       |    Reorder    |
+ *        \     /        |
+ *         \   /      Reorder
+ *        Matmul        /
+ *           |         /
+ *           |        /
+ *         Softmax   /
+ *            \     /
+ *             \   /
+ *             Matmul
+ *               |
+ *               |
+ *             Output
+ */
+class dyn_quantize_mha_ref_k_t;
+
+class SPARSE_API_ dyn_quantize_mha_ref_kd_t : public kernel_desc_t {
+ public:
+  using io = ssd::dyn_quantize_mha_io::io;
+  explicit dyn_quantize_mha_ref_kd_t(const jd::operator_desc& op_desc)
+      : kernel_desc_t(kernel_kind::dyn_quantize_mha), op_desc_(op_desc) {}
+  virtual ~dyn_quantize_mha_ref_kd_t() {}
+
+  bool init() override;
+  DECLARE_COMMON_PD_T(dyn_quantize_mha_ref_k_t, dyn_quantize_mha_ref_kd_t);
+
+  const jd::operator_desc& get_operator_desc() const override { return op_desc_; }
+  inline std::vector<dim_t> shape() const override {
+    return {
+        op_desc_.tensor_descs()[io::Q].shape()[0],  // batch_size
+        op_desc_.tensor_descs()[io::Q].shape()[2],  // head_num
+        op_desc_.tensor_descs()[io::Q].shape()[1],  // M
+        op_desc_.tensor_descs()[io::Q].shape()[3],  // head_size
+        op_desc_.tensor_descs()[io::K].shape()[1],  // N
+    };
+  }
+
+ private:
+  jd::operator_desc op_desc_;
+};
+
+class SPARSE_API_ dyn_quantize_mha_ref_k_t : public kernel_t {
+ public:
+  using io = ssd::dyn_quantize_mha_io::io;
+  using kd_t = dyn_quantize_mha_ref_kd_t;
+  explicit dyn_quantize_mha_ref_k_t(const std::shared_ptr<const kernel_desc_t>& kd);
+  virtual ~dyn_quantize_mha_ref_k_t() {}
+  // Delete move constructor and move operator
+  dyn_quantize_mha_ref_k_t(dyn_quantize_mha_ref_k_t&&) = delete;
+  dyn_quantize_mha_ref_k_t& operator=(dyn_quantize_mha_ref_k_t&&) = delete;
+  // Delete copy constructor and copy operator
+  dyn_quantize_mha_ref_k_t(const dyn_quantize_mha_ref_k_t&) = delete;
+  dyn_quantize_mha_ref_k_t& operator=(const dyn_quantize_mha_ref_k_t&) = delete;
+
+  bool init() override;
+  bool execute(const std::vector<const void*>& rt_data) const override;
+  const std::shared_ptr<const kd_t> derived_kd() const { return std::static_pointer_cast<const kd_t>(kd_); }
+
+ private:
+  std::vector<std::vector<dim_t>> t_shapes_;
+  int32_t batch_size_, head_num_, M_, head_size_, N_;
+};
+
+}  // namespace jd
+#endif  // ENGINE_SPARSELIB_INCLUDE_KERNELS_DYN_QUANTIZE_MHA_REF_HPP_

intel_extension_for_transformers/backends/neural_engine/kernels/include/kernels/dyn_quantize_mha_types.hpp

@@ -0,0 +1,54 @@
+//  Copyright (c) 2022 Intel Corporation
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+
+#ifndef ENGINE_SPARSELIB_INCLUDE_KERNELS_DYN_QUANTIZE_MHA_TYPES_HPP_
+#define ENGINE_SPARSELIB_INCLUDE_KERNELS_DYN_QUANTIZE_MHA_TYPES_HPP_
+
+#include <vector>
+
+#include "amx_utils.hpp"
+#include "param_types.hpp"
+
+namespace jd {
+namespace ssd {
+namespace dyn_quantize_mha_io {
+enum io {
+  Q,
+  K,
+  MASK,
+  V,
+  DST,
+  TMP,  // size of K + size of V + ~1M
+
+  Q_SCALE,
+  Q_ZP,
+  K_SCALE,
+  K_ZP,
+  V_SCALE,
+  V_ZP,
+  DST_SCALE,
+  DST_ZP,
+
+  BATCH_SIZE,
+  HEAD_NUM,
+  HEAD_SIZE,
+  M,  // "seq_len" for Q & DST
+  N,  // "seq_len" for K & V
+  dyn_quantize_mha_io_MAX = N,
+};
+}  // namespace dyn_quantize_mha_io
+
+}  // namespace ssd
+}  // namespace jd
+#endif  // ENGINE_SPARSELIB_INCLUDE_KERNELS_DYN_QUANTIZE_MHA_TYPES_HPP_

intel_extension_for_transformers/backends/neural_engine/kernels/include/kernels/matmul_types.hpp

@@ -23,16 +23,18 @@
 
 namespace jd {
 namespace ssd {
-/**
- * @brief tensors index configuration of this kernel.
- * TODO(Yi): potential confliction with indices of other op types
- */
-static constexpr int SRC0 = 0;
-static constexpr int SRC1 = 1;
-static constexpr int DST0 = 2;
-static constexpr int SRC2 = 3;  // for binary add
-static constexpr int SCALE0 = 4;
-static constexpr int ZP0 = 5;
+
+namespace matmul_io {
+enum io {
+  SRC0,
+  SRC1,
+  DST0,
+  SRC2,
+  SCALE0,
+  ZP0,
+  matmul_io_MAX = ZP0,
+};
+}  // namespace matmul_io
 
 struct matmul_param_t {
   dim_t M;

intel_extension_for_transformers/backends/neural_engine/kernels/include/param_types.hpp

@@ -31,7 +31,8 @@ enum class kernel_kind : uint8_t {
   logsoftmax,
   gather,
   attention,
-  transpose_mha
+  transpose_mha,
+  dyn_quantize_mha,
 };
 
 enum class postop_alg : uint8_t { undef, exp, tanh, gelu, relu, quantize, dequantize, linear, eltop_int_lut };
@@ -82,6 +83,7 @@ enum class format_type : uint8_t {
   ab,  // shape permutation = {0, 1}
   ba,  // shape permutation = {1, 0}
   abc,
+  abcd,
 
   // encoding format of sparse matrix
   uncoded,

intel_extension_for_transformers/backends/neural_engine/kernels/src/cpu_engine.cpp

@@ -29,6 +29,7 @@ DECLARE_IMPL_LIST(softmax);
 DECLARE_IMPL_LIST(gather);
 DECLARE_IMPL_LIST(attention);
 DECLARE_IMPL_LIST(transpose_mha);
+DECLARE_IMPL_LIST(dyn_quantize_mha);
 
 #undef DECLARE_IMPL_LIST
 
@@ -48,6 +49,7 @@ const std::vector<impl_list_item_t>* cpu_engine::get_implementation_list(const o
     CASE(softmax);
     CASE(attention);
     CASE(transpose_mha);
+    CASE(dyn_quantize_mha);
     default:
       return &cpu_engine::empty_list;
   }

intel_extension_for_transformers/backends/neural_engine/kernels/src/interface.cpp

@@ -76,6 +76,8 @@ bool kernel_proxy::create_proxy_object(std::shared_ptr<const kernel_t>& result_r
   return true;
 }
 
+size_t kernel_proxy::get_workspace_size() const { return get_sp()->get_workspace_size(); }
+
 void kernel_proxy::execute(const std::vector<const void*>& rt_data) const {
   bool status = false;
 #ifdef SPARSE_LIB_USE_VTUNE

intel_extension_for_transformers/backends/neural_engine/kernels/src/kernels/attention.cpp

@@ -63,6 +63,8 @@ enum SubKernel {
 }  // namespace
 
 namespace jd {
+using matmul_io = ssd::matmul_io::io;
+
 template <typename T_kd>
 inline bool attention_kd_t::add_kernel_desc(const operator_desc& op_desc, const char* name) {
   std::shared_ptr<const kernel_desc_t> kd;
@@ -283,7 +285,7 @@ void attention_k_t::setup_memory() {
   const auto tensor_bytes = [](const jd::tensor_desc& d) { return d.size() * type2bytes[d.dtype()]; };
 
   offset.push_back(tensor_bytes(ker_opdesc(SubKernel::QK_SPMM).tensor_descs()[ssd::DST]));
-  offset.push_back(tensor_bytes(ker_opdesc(SubKernel::Q_K_GEMM).tensor_descs()[ssd::DST0]));
+  offset.push_back(tensor_bytes(ker_opdesc(SubKernel::Q_K_GEMM).tensor_descs()[matmul_io::DST0]));
   offset.push_back(tensor_bytes(ker_opdesc(SubKernel::SOFTMAX).tensor_descs()[1]));
   offset.push_back(tensor_bytes(ker_opdesc(SubKernel::V_SPMM).tensor_descs()[ssd::DST]));
   // the last kernel(QK(softmax) * V) don't need alloc memory
@@ -301,15 +303,16 @@ void attention_k_t::setup_memory() {
 
   // part1 Q X K
   mem_[SubKernel::Q_K_GEMM].resize(4);
-  mem_[SubKernel::Q_K_GEMM][ssd::SRC0] = mem_[SubKernel::QK_SPMM][ssd::DST];
-  mem_[SubKernel::Q_K_GEMM][ssd::SRC1] = mem_[SubKernel::QK_SPMM][ssd::DST] + offset[0] / 2;  // split qk out to q and k
-  mem_[SubKernel::Q_K_GEMM][ssd::DST0] = mem_[SubKernel::QK_SPMM][ssd::DST] + offset[0];      // dst
-  mem_[SubKernel::Q_K_GEMM][ssd::SRC2] = nullptr;
+  mem_[SubKernel::Q_K_GEMM][matmul_io::SRC0] = mem_[SubKernel::QK_SPMM][ssd::DST];
+  mem_[SubKernel::Q_K_GEMM][matmul_io::SRC1] =
+      mem_[SubKernel::QK_SPMM][ssd::DST] + offset[0] / 2;  // split qk out to q and k
+  mem_[SubKernel::Q_K_GEMM][matmul_io::DST0] = mem_[SubKernel::QK_SPMM][ssd::DST] + offset[0];  // dst
+  mem_[SubKernel::Q_K_GEMM][matmul_io::SRC2] = nullptr;
 
   // part2 Softmax
   mem_[SubKernel::SOFTMAX].resize(2);
-  mem_[SubKernel::SOFTMAX][0] = mem_[SubKernel::Q_K_GEMM][ssd::DST0];
-  mem_[SubKernel::SOFTMAX][1] = mem_[SubKernel::Q_K_GEMM][ssd::DST0] + offset[1];
+  mem_[SubKernel::SOFTMAX][0] = mem_[SubKernel::Q_K_GEMM][matmul_io::DST0];
+  mem_[SubKernel::SOFTMAX][1] = mem_[SubKernel::Q_K_GEMM][matmul_io::DST0] + offset[1];
 
   // part5 spmm for V
   mem_[SubKernel::V_SPMM].resize(ssd::SCALES + 1);
@@ -321,13 +324,13 @@ void attention_k_t::setup_memory() {
   mem_[SubKernel::V_SPMM][ssd::DST] = mem_[SubKernel::SOFTMAX][1] + offset[2];
 
   // part6  V X QK(softmax out)
-  mem_[SubKernel::QK_V_MATMUL].resize(ssd::ZP0 + 1);
-  mem_[SubKernel::QK_V_MATMUL][ssd::SRC0] = mem_[SubKernel::SOFTMAX][1];
-  mem_[SubKernel::QK_V_MATMUL][ssd::SRC1] = mem_[SubKernel::V_SPMM][ssd::DST];
-  mem_[SubKernel::QK_V_MATMUL][ssd::DST0] = nullptr;
-  mem_[SubKernel::QK_V_MATMUL][ssd::SRC2] = nullptr;
-  mem_[SubKernel::QK_V_MATMUL][ssd::SCALE0] = nullptr;
-  mem_[SubKernel::QK_V_MATMUL][ssd::ZP0] = nullptr;
+  mem_[SubKernel::QK_V_MATMUL].resize(matmul_io::ZP0 + 1);
+  mem_[SubKernel::QK_V_MATMUL][matmul_io::SRC0] = mem_[SubKernel::SOFTMAX][1];
+  mem_[SubKernel::QK_V_MATMUL][matmul_io::SRC1] = mem_[SubKernel::V_SPMM][ssd::DST];
+  mem_[SubKernel::QK_V_MATMUL][matmul_io::DST0] = nullptr;
+  mem_[SubKernel::QK_V_MATMUL][matmul_io::SRC2] = nullptr;
+  mem_[SubKernel::QK_V_MATMUL][matmul_io::SCALE0] = nullptr;
+  mem_[SubKernel::QK_V_MATMUL][matmul_io::ZP0] = nullptr;
 }
 bool attention_k_t::init() {
   // Create kernel
@@ -357,12 +360,12 @@ std::vector<const void*> attention_k_t::set_input_output(int index, const std::v
     // part0 QK spmm_vnni and part5 V spmm_vnni
     data[ssd::SRC] = rt_data[attention_io::MERGE_SRC];
   } else if (index == SubKernel::Q_K_GEMM) {
-    data[ssd::SRC2] = rt_data[attention_io::Q_K_SRC2];
+    data[matmul_io::SRC2] = rt_data[attention_io::Q_K_SRC2];
   } else if (index == SubKernel::QK_V_MATMUL) {
     // part4 transpose matmul for QK x V
-    data[ssd::DST0] = rt_data[attention_io::MERGE_DST];
-    data[ssd::SCALE0] = rt_data[attention_io::QK_V_OUTPUT_SCALES];
-    data[ssd::ZP0] = rt_data[attention_io::QK_V_OUTPUT_ZERO_POINT];
+    data[matmul_io::DST0] = rt_data[attention_io::MERGE_DST];
+    data[matmul_io::SCALE0] = rt_data[attention_io::QK_V_OUTPUT_SCALES];
+    data[matmul_io::ZP0] = rt_data[attention_io::QK_V_OUTPUT_ZERO_POINT];
   }
   return data;
 }