[layer] Added GPU implementation of custom MHA attention layer

GPU implementation of custom MHA layer Signed-off-by: Debadri Samaddar <[email protected]>
nnstreamer · Jul 9, 2024 · e52ce37 · e52ce37
1 parent 9c8decb
commit e52ce37
Show file tree

Hide file tree

Showing 8 changed files with 1,238 additions and 0 deletions.
diff --git a/api/ccapi/include/layer.h b/api/ccapi/include/layer.h
@@ -77,6 +77,9 @@ enum LayerType {
     ML_TRAIN_LAYER_TYPE_POSITIONAL_ENCODING, /**< Positional Encoding Layer type
                                               */
   LAYER_IDENTITY = ML_TRAIN_LAYER_TYPE_IDENTITY, /**< Identity Layer type */
+  LAYER_CUSTOM_MULTI_HEAD_ATTENTION =
+    ML_TRAIN_LAYER_TYPE_CUSTOM_MULTI_HEAD_ATTENTION, /**< Multi Head Attention
+                                                 Layer type */
   LAYER_PREPROCESS_FLIP =
     ML_TRAIN_LAYER_TYPE_PREPROCESS_FLIP, /**< Preprocess flip Layer type */
   LAYER_PREPROCESS_TRANSLATE =
@@ -517,6 +520,16 @@ MultiHeadAttention(const std::vector<std::string> &properties = {}) {
   return createLayer(LayerType::LAYER_MULTI_HEAD_ATTENTION, properties);
 }
 
+/**
+ * @brief Helper function to create Custom Multi Head Attention Layer
+ */
+inline std::unique_ptr<Layer> CustomMultiHeadAttention(
+  const std::vector<std::string> &properties = {},
+  const LayerComputeEngine &compute_engine = LayerComputeEngine::CPU) {
+  return createLayer(LayerType::LAYER_CUSTOM_MULTI_HEAD_ATTENTION, properties,
+                     compute_engine);
+}
+
 /**
  * @brief Helper function to create Positional Encoding Layer
  */

diff --git a/api/nntrainer-api-common.h b/api/nntrainer-api-common.h
@@ -64,6 +64,8 @@ typedef enum {
     28, /**< Positional Encoding Layer type (Since 7.0) */
   ML_TRAIN_LAYER_TYPE_IDENTITY = 29, /**< Identity Layer type (Since 8.0) */
   ML_TRAIN_LAYER_TYPE_SWIGLU = 30,   /**< Swiglu Layer type */
+  ML_TRAIN_LAYER_TYPE_CUSTOM_MULTI_HEAD_ATTENTION =
+    31, /**< Custom Multi Head Attention Layer type (Since 8.0) */
   ML_TRAIN_LAYER_TYPE_PREPROCESS_FLIP =
     300, /**< Preprocess flip Layer (Since 6.5) */
   ML_TRAIN_LAYER_TYPE_PREPROCESS_TRANSLATE =

diff --git a/nntrainer/cl_context.cpp b/nntrainer/cl_context.cpp
@@ -15,6 +15,7 @@
 
 #include <addition_layer_cl.h>
 #include <cl_context.h>
+#include <custom_multi_head_attention_layer_cl.h>
 #include <fc_layer_cl.h>
 #include <reshape_cl.h>
 #include <swiglu_cl.h>
@@ -40,6 +41,10 @@ static void add_default_object(ClContext &cc) {
 
   cc.registerFactory(nntrainer::createLayer<ReshapeLayerCl>,
                      ReshapeLayerCl::type, ml::train::LayerType::LAYER_RESHAPE);
+
+  cc.registerFactory(nntrainer::createLayer<CustomMultiHeadAttentionLayerCl>,
+                     CustomMultiHeadAttentionLayerCl::type,
+                     ml::train::LayerType::LAYER_CUSTOM_MULTI_HEAD_ATTENTION);
 }
 
 static void registerer(ClContext &cc) noexcept {