PaddlePaddle
diff --git a/‎paddle/fluid/framework/details/CMakeLists.txt
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎paddle/fluid/framework/details/CMakeLists.txt
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/details/build_strategy.cc
Copy file name to clipboard
+7Lines changed: 7 additions & 0 deletions b/‎paddle/fluid/framework/details/build_strategy.cc
Copy file name to clipboard
+7Lines changed: 7 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/details/build_strategy.h
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion b/‎paddle/fluid/framework/details/build_strategy.h
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion
diff --git a/‎paddle/fluid/framework/distributed_strategy.proto
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/framework/distributed_strategy.proto
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/fluid/framework/ir/CMakeLists.txt
Copy file name to clipboard
+7Lines changed: 7 additions & 0 deletions b/‎paddle/fluid/framework/ir/CMakeLists.txt
Copy file name to clipboard
+7Lines changed: 7 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/ir/fuse_dot_product_attention_pass.cc
Copy file name to clipboard
+493Lines changed: 493 additions & 0 deletions b/‎paddle/fluid/framework/ir/fuse_dot_product_attention_pass.cc
Copy file name to clipboard
+493Lines changed: 493 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/ir/fuse_dot_product_attention_pass.h
Copy file name to clipboard
+120Lines changed: 120 additions & 0 deletions b/‎paddle/fluid/framework/ir/fuse_dot_product_attention_pass.h
Copy file name to clipboard
+120Lines changed: 120 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/ir/fuse_gemm_epilogue_pass.h
Copy file name to clipboard
+7-4Lines changed: 7 additions & 4 deletions b/‎paddle/fluid/framework/ir/fuse_gemm_epilogue_pass.h
Copy file name to clipboard
+7-4Lines changed: 7 additions & 4 deletions
@@ -414,6 +414,10 @@ set(IR_PASS_DEPS
     fused_feedforward_pass
     delete_dropout_op_pass)
 
+if(WITH_CUDNN_FRONTEND)
+  set(IR_PASS_DEPS ${IR_PASS_DEPS} fuse_dot_product_attention_pass)
+endif()
+
 if(WITH_CINN)
   set(IR_PASS_DEPS ${IR_PASS_DEPS} paddle2cinn)
 endif()
 
@@ -182,6 +182,10 @@ class ParallelExecutorPassBuilder : public ir::PassBuilder {
         "mkldnn_placement_pass");
 
     // 2. trainning pass
+#ifdef PADDLE_WITH_CUDNN_FRONTEND
+    AppendPassWithCheck(strategy_.fuse_dot_product_attention_,
+                        "fuse_dot_product_attention_pass");
+#endif
     AppendPassWithCheck(strategy_.fuse_relu_depthwise_conv_,
                         "fuse_relu_depthwise_conv_pass");
     AppendPassWithCheck(strategy_.fuse_bn_act_ops_, "fuse_bn_act_pass");
@@ -552,3 +556,6 @@ USE_PASS(fusion_group_pass);
 #if (defined(PADDLE_WITH_CUDA) && CUDA_VERSION >= 11060)
 USE_PASS(fuse_gemm_epilogue_pass);
 #endif
+#ifdef PADDLE_WITH_CUDNN_FRONTEND
+USE_PASS(fuse_dot_product_attention_pass);
+#endif
@@ -136,7 +136,8 @@ struct BuildStrategy {
   // Fused feed forward
   bool fused_feedforward_{false};
   bool sequential_run_{false};
-
+  // Fuse dot product attention
+  bool fuse_dot_product_attention_{false};
   // mkldnn_enabled_op_types specify the operator type list to
   // use MKLDNN acceleration. It is null in default, means
   // that all the operators supported by MKLDNN will be
 
@@ -172,6 +172,7 @@ message BuildStrategy {
   optional string debug_graphviz_path = 17;
   optional bool fused_attention = 18 [ default = false];
   optional bool fused_feedforward = 19 [ default = false];
+  optional bool fuse_dot_product_attention = 20 [ default = false ];
 }
 
 message ExecutionStrategy {
 
@@ -351,6 +351,13 @@ cc_library(
   SRCS fused_feedforward_pass.cc
   DEPS pass graph_pattern_detector)
 
+if(WITH_CUDNN_FRONTEND)
+  cc_library(
+    fuse_dot_product_attention_pass
+    SRCS fuse_dot_product_attention_pass.cc
+    DEPS pass graph_pattern_detector)
+endif()
+
 set(GLOB_PASS_LIB
     ${INFER_IR_PASSES}
     CACHE INTERNAL "Global PASS library")
 
@@ -0,0 +1,120 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// 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
+//
+//     https://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.
+
+#pragma once
+
+#include "paddle/fluid/framework/ir/fuse_pass_base.h"
+#include "paddle/fluid/framework/ir/graph.h"
+#include "paddle/fluid/framework/ir/graph_pattern_detector.h"
+#include "paddle/fluid/framework/ir/pass.h"
+
+namespace paddle {
+namespace framework {
+namespace ir {
+
+/*
+ * Fuse the ElewiseAdd and activation
+ */
+class Graph;
+class Node;
+
+enum class AttentionType { kSelfAttention, kCrossAttention };
+
+class MaskMetaData {
+ public:
+  ir::Node *mask_node = nullptr;
+};
+class QKVMetaData {
+ public:
+  ir::Node *q_node = nullptr;
+  ir::Node *k_node = nullptr;
+  ir::Node *v_node = nullptr;
+};
+class OutputMetaData {
+ public:
+  ir::Node *output_node = nullptr;
+  ir::Node *softmax_output_node = nullptr;
+  ir::Node *rng_state_node = nullptr;
+};
+
+template <typename T, char const *NAME>
+class OpCache {
+ public:
+  OpCache() {}
+  OpCache(const OpCache &) = delete;
+  void operator=(const OpCache &) = delete;
+
+  bool Exist(const std::string &key) const {
+    std::lock_guard<std::mutex> lock(mtx_);
+    return map_.count(key);
+  }
+
+  T Get(const std::string &key) const {
+    std::lock_guard<std::mutex> lock(mtx_);
+    return map_.find(key)->second;
+  }
+
+  void Insert(const std::string &key, const T &value) {
+    std::lock_guard<std::mutex> lock(mtx_);
+    map_[key] = value;
+  }
+
+  void Erase(const std::string &key) {
+    std::lock_guard<std::mutex> lock(mtx_);
+    map_.erase(key);
+  }
+
+ private:
+  std::unordered_map<std::string, T> map_;
+  mutable std::mutex mtx_;
+};
+
+const char NAME1[] = "MaskMetaData";
+const char NAME2[] = "QKVMetaData";
+const char NAME3[] = "OutputMetaData";
+using MaskCache = OpCache<MaskMetaData, NAME1>;
+using QKVCache = OpCache<QKVMetaData, NAME2>;
+using OutputCache = OpCache<OutputMetaData, NAME3>;
+
+class FuseDotProductAttentionPass : public FusePassBase {
+ public:
+  virtual ~FuseDotProductAttentionPass() {}
+
+ protected:
+  void ApplyImpl(ir::Graph *graph) const override;
+
+  ir::Graph *FuseDotProductAttentionFwd(
+      ir::Graph *graph,
+      bool with_dropout,
+      QKVCache *qkv_cache,
+      MaskCache *mask_cache,
+      OutputCache *output_cache,
+      std::unordered_set<const Node *> *nodes_to_remove) const;
+  ir::Graph *FuseDotProductAttentionBwd(
+      ir::Graph *graph,
+      bool with_dropout,
+      QKVCache *qkv_cache,
+      MaskCache *mask_cache,
+      OutputCache *output_cache,
+      std::unordered_set<const Node *> *nodes_to_remove) const;
+
+ private:
+  std::string GenerateMetaKey_(const std::string &q_name,
+                               const std::string &k_name,
+                               const std::string &v_name) const;
+};
+
+}  // namespace ir
+}  // namespace framework
+}  // namespace paddle
@@ -41,7 +41,10 @@ class EpiloguePassActivationCache {
   void operator=(const EpiloguePassActivationCache &) = delete;
 
   bool HasFusedActivation(const std::string &key) const {
-    return fused_activation_space_map_.count(key);
+    mtx_.lock();
+    bool has = fused_activation_space_map_.count(key);
+    mtx_.unlock();
+    return has;
   }
 
   ir::Node *GetFusedActivationSpace(const std::string &key) {
@@ -54,9 +57,9 @@ class EpiloguePassActivationCache {
 
   void InsertFusedActivation(const std::string &key, ir::Node *const value) {
     if (!HasFusedActivation(key)) {
-      mtx.lock();
+      mtx_.lock();
       fused_activation_space_map_.insert({key, value});
-      mtx.unlock();
+      mtx_.unlock();
     } else {
       PADDLE_THROW(platform::errors::AlreadyExists(
           "The key (%d) of EpiloguePassActivationCache already exist.", key));
@@ -65,7 +68,7 @@ class EpiloguePassActivationCache {
 
  private:
   std::unordered_map<std::string, ir::Node *> fused_activation_space_map_;
-  std::mutex mtx;
+  mutable std::mutex mtx_;
 };
 
 class FuseGemmEpiloguePass : public FusePassBase {
Original file line number	Diff line number	Diff line change
`@@ -172,6 +172,7 @@ message BuildStrategy {`
`172`	`172`	`optional string debug_graphviz_path = 17;`
`173`	`173`	`optional bool fused_attention = 18 [ default = false];`
`174`	`174`	`optional bool fused_feedforward = 19 [ default = false];`
	`175`	`+ optional bool fuse_dot_product_attention = 20 [ default = false ];`
`175`	`176`	`}`
`176`	`177`
`177`	`178`	`message ExecutionStrategy {`