PaddlePaddle
diff --git a/‎paddle/fluid/operators/fused/fused_gate_attention_op.cu
Copy file name to clipboard
-652Lines changed: 0 additions & 652 deletions b/‎paddle/fluid/operators/fused/fused_gate_attention_op.cu
Copy file name to clipboard
-652Lines changed: 0 additions & 652 deletions
diff --git a/‎paddle/fluid/operators/ops_signature/fused_gate_attention_sig.cc
Copy file name to clipboard
+102Lines changed: 102 additions & 0 deletions b/‎paddle/fluid/operators/ops_signature/fused_gate_attention_sig.cc
Copy file name to clipboard
+102Lines changed: 102 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/phi/infermeta/fusion.cc
Copy file name to clipboard
+149Lines changed: 149 additions & 0 deletions b/‎paddle/phi/infermeta/fusion.cc
Copy file name to clipboard
+149Lines changed: 149 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/fusion.h
Copy file name to clipboard
+63Lines changed: 63 additions & 0 deletions b/‎paddle/phi/infermeta/fusion.h
Copy file name to clipboard
+63Lines changed: 63 additions & 0 deletions
@@ -0,0 +1,102 @@
+/* Copyright (c) 2024 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. */
+
+#include "paddle/phi/core/compat/op_utils.h"
+
+namespace phi {
+
+KernelSignature FusedGateAttentionOpArgumentMapping(
+    const ArgumentMappingContext& ctx) {
+  paddle::small_vector<const char*> inputs{"Query",
+                                           "Key",
+                                           "QueryWeight",
+                                           "KeyWeight",
+                                           "ValueWeight",
+                                           "QKVWeight",
+                                           "NonbatchedBias",
+                                           "SrcMask",
+                                           "GateWeight",
+                                           "GateBias",
+                                           "OutLinearWeight",
+                                           "OutLinearBias"};
+  paddle::small_vector<const char*> attrs;
+  attrs.emplace_back("has_gating");
+  attrs.emplace_back("merge_qkv");
+  attrs.emplace_back("use_flash_attn");
+  paddle::small_vector<const char*> outputs{"QueryTransposeOut",
+                                            "KeyTransposeOut",
+                                            "ValueTransposeOut",
+                                            "QKVTransposeOut",
+                                            "SoftmaxOut",
+                                            "SoftmaxLse",
+                                            "FMHAOut",
+                                            "GateOut",
+                                            "Out"};
+  return KernelSignature("fused_gate_attention",
+                         std::move(inputs),
+                         std::move(attrs),
+                         std::move(outputs));
+}
+
+KernelSignature FusedGateAttentionGradOpArgumentMapping(
+    const ArgumentMappingContext& ctx) {
+  paddle::small_vector<const char*> inputs{"Query",
+                                           "Key",
+                                           "QueryWeight",
+                                           "KeyWeight",
+                                           "ValueWeight",
+                                           "QKVWeight",
+                                           "NonbatchedBias",
+                                           "SrcMask",
+                                           "GateWeight",
+                                           "GateBias",
+                                           "OutLinearWeight",
+                                           "OutLinearBias",
+                                           "QueryTransposeOut",
+                                           "KeyTransposeOut",
+                                           "ValueTransposeOut",
+                                           "QKVTransposeOut",
+                                           "SoftmaxOut",
+                                           "SoftmaxLse",
+                                           "FMHAOut",
+                                           "GateOut",
+                                           "Out@GRAD"};
+  paddle::small_vector<const char*> attrs;
+  attrs.emplace_back("has_gating");
+  attrs.emplace_back("merge_qkv");
+  attrs.emplace_back("use_flash_attn");
+  paddle::small_vector<const char*> outputs{"Query@GRAD",
+                                            "Key@GRAD",
+                                            "QueryWeight@GRAD",
+                                            "KeyWeight@GRAD",
+                                            "ValueWeight@GRAD",
+                                            "QKVWeight@GRAD",
+                                            "NonbatchedBias@GRAD",
+                                            "GateWeight@GRAD",
+                                            "GateBias@GRAD",
+                                            "OutLinearWeight@GRAD",
+                                            "OutLinearBias@GRAD"};
+  return KernelSignature("fused_gate_attention_grad",
+                         std::move(inputs),
+                         std::move(attrs),
+                         std::move(outputs));
+}
+
+}  // namespace phi
+
+PD_REGISTER_ARG_MAPPING_FN(fused_gate_attention,
+                           phi::FusedGateAttentionOpArgumentMapping);
+
+PD_REGISTER_ARG_MAPPING_FN(fused_gate_attention_grad,
+                           phi::FusedGateAttentionGradOpArgumentMapping);
@@ -193,6 +193,7 @@
     'fused_elementwise_mul',
     'fused_elementwise_sub',
     'fused_embedding_fc_lstm',
+    'fused_gate_attention',
     'fused_multi_transformer_int8',
     'fusion_group',
     'fusion_lstm',
 
@@ -5321,6 +5321,155 @@ void ResnetUnitGradInferMeta(const MetaTensor& x,
   }
 }
 
+void FusedGateAttentionInferMeta(const MetaTensor& query,
+                                 const MetaTensor& key,
+                                 const MetaTensor& query_weight,
+                                 const MetaTensor& key_weight,
+                                 const MetaTensor& value_weight,
+                                 const MetaTensor& qkv_weight,
+                                 const MetaTensor& nonbatched_bias,
+                                 const MetaTensor& src_mask,
+                                 const MetaTensor& gate_weight,
+                                 const MetaTensor& gate_bias,
+                                 const MetaTensor& out_linear_weight,
+                                 const MetaTensor& out_linear_bias,
+                                 bool has_gating,
+                                 bool merge_qkv,
+                                 bool use_flash_attn,
+                                 MetaTensor* query_transpose_out,
+                                 MetaTensor* key_transpose_out,
+                                 MetaTensor* value_transpose_out,
+                                 MetaTensor* qkv_transpose_out,
+                                 MetaTensor* softmax_out,
+                                 MetaTensor* softmax_lse,
+                                 MetaTensor* fmha_out,
+                                 MetaTensor* gate_out,
+                                 MetaTensor* out,
+                                 MetaConfig config) {
+  const auto& input_q_dims = query.dims();
+  int batch_size = input_q_dims[0];
+  int seq_len_m = input_q_dims[1];
+  int seq_len_r = input_q_dims[2];
+
+  int num_head, m_size, head_dim;
+  if (merge_qkv) {
+    // QKV's input: [batch_size, seq_len_m, seq_len_r, qkv_dim]
+    // QKV's weight: [3, num_head, head_dim, qkv_dim]
+    const auto& qkv_w_dims = qkv_weight.dims();
+
+    num_head = qkv_w_dims[1];
+    head_dim = qkv_w_dims[2];
+    m_size = seq_len_r;
+
+    qkv_transpose_out->set_dims(
+        {3, batch_size, seq_len_m, num_head, seq_len_r, head_dim});
+    qkv_transpose_out->set_dtype(query.dtype());
+  } else {
+    const auto& input_k_dims = key.dims();
+    const auto& q_w_dims = query_weight.dims();
+
+    num_head = q_w_dims[1];
+    head_dim = q_w_dims[2];
+    m_size = input_k_dims[2];
+
+    query_transpose_out->set_dims(
+        {batch_size, seq_len_m, num_head, seq_len_r, head_dim});
+    key_transpose_out->set_dims(
+        {batch_size, seq_len_m, num_head, m_size, head_dim});
+    value_transpose_out->set_dims(
+        {batch_size, seq_len_m, num_head, m_size, head_dim});
+    query_transpose_out->set_dtype(query.dtype());
+    key_transpose_out->set_dtype(query.dtype());
+    value_transpose_out->set_dtype(query.dtype());
+  }
+
+  softmax_out->set_dims({batch_size, seq_len_m, num_head, seq_len_r, m_size});
+  fmha_out->set_dims({batch_size, seq_len_m, seq_len_r, num_head, head_dim});
+  softmax_out->set_dtype(query.dtype());
+  fmha_out->set_dtype(query.dtype());
+
+  if (has_gating) {
+    gate_out->set_dims({batch_size, seq_len_m, seq_len_r, num_head, head_dim});
+    gate_out->set_dtype(query.dtype());
+  }
+  out->set_dims(query.dims());
+  out->set_dtype(query.dtype());
+}
+
+void FusedGateAttentionGradInferMeta(const MetaTensor& query,
+                                     const MetaTensor& key,
+                                     const MetaTensor& query_weight,
+                                     const MetaTensor& key_weight,
+                                     const MetaTensor& value_weight,
+                                     const MetaTensor& qkv_weight,
+                                     const MetaTensor& nonbatched_bias,
+                                     const MetaTensor& src_mask,
+                                     const MetaTensor& gate_weight,
+                                     const MetaTensor& gate_bias,
+                                     const MetaTensor& out_linear_weight,
+                                     const MetaTensor& out_linear_bias,
+                                     const MetaTensor& query_transpose_out,
+                                     const MetaTensor& key_transpose_out,
+                                     const MetaTensor& value_transpose_out,
+                                     const MetaTensor& qkv_transpose_out,
+                                     const MetaTensor& softmax_out,
+                                     const MetaTensor& softmax_lse,
+                                     const MetaTensor& fmha_out,
+                                     const MetaTensor& gate_out,
+                                     const MetaTensor& out_grad,
+                                     bool has_gating,
+                                     bool merge_qkv,
+                                     bool use_flash_attn,
+                                     MetaTensor* query_grad,
+                                     MetaTensor* key_grad,
+                                     MetaTensor* query_weight_grad,
+                                     MetaTensor* key_weight_grad,
+                                     MetaTensor* value_weight_grad,
+                                     MetaTensor* qkv_weight_grad,
+                                     MetaTensor* nonbatched_bias_grad,
+                                     MetaTensor* gate_weight_grad,
+                                     MetaTensor* gate_bias_grad,
+                                     MetaTensor* out_linear_weight_grad,
+                                     MetaTensor* out_linear_bias_grad,
+                                     MetaConfig config) {
+  if (query_grad != nullptr) {
+    query_grad->set_dims(query.dims());
+    query_grad->set_dtype(query.dtype());
+  }
+  if (key_grad != nullptr) {
+    key_grad->set_dims(key.dims());
+    key_grad->set_dtype(key.dtype());
+  }
+
+  if (merge_qkv) {
+    qkv_weight_grad->set_dims(qkv_weight.dims());
+    qkv_weight_grad->set_dtype(qkv_weight.dtype());
+  } else {
+    query_weight_grad->set_dims(query_weight.dims());
+    key_weight_grad->set_dims(key_weight.dims());
+    value_weight_grad->set_dims(value_weight.dims());
+    query_weight_grad->set_dtype(query_weight.dtype());
+    key_weight_grad->set_dtype(key_weight.dtype());
+    value_weight_grad->set_dtype(value_weight.dtype());
+  }
+
+  if (has_gating) {
+    gate_weight_grad->set_dims(gate_weight.dims());
+    gate_bias_grad->set_dims(gate_bias.dims());
+    gate_weight_grad->set_dtype(gate_weight.dtype());
+    gate_bias_grad->set_dtype(gate_bias.dtype());
+  }
+
+  if (nonbatched_bias_grad != nullptr) {
+    nonbatched_bias_grad->set_dims(nonbatched_bias.dims());
+    nonbatched_bias_grad->set_dtype(nonbatched_bias.dtype());
+  }
+  out_linear_weight_grad->set_dims(out_linear_weight.dims());
+  out_linear_bias_grad->set_dims(out_linear_bias.dims());
+  out_linear_weight_grad->set_dtype(out_linear_weight.dtype());
+  out_linear_bias_grad->set_dtype(out_linear_bias.dtype());
+}
+
 void ResnetBasicBlockInferMeta(const MetaTensor& x,
                                const MetaTensor& filter1,
                                const MetaTensor& scale1,
 
@@ -1237,6 +1237,69 @@ void ResnetUnitGradInferMeta(const MetaTensor& x,
                              MetaTensor* scale_z_grad,
                              MetaTensor* bias_z_grad);
 
+void FusedGateAttentionInferMeta(const MetaTensor& query,
+                                 const MetaTensor& key,
+                                 const MetaTensor& query_weight,
+                                 const MetaTensor& key_weight,
+                                 const MetaTensor& value_weight,
+                                 const MetaTensor& qkv_weight,
+                                 const MetaTensor& nonbatched_bias,
+                                 const MetaTensor& src_mask,
+                                 const MetaTensor& gate_weight,
+                                 const MetaTensor& gate_bias,
+                                 const MetaTensor& out_linear_weight,
+                                 const MetaTensor& out_linear_bias,
+                                 bool has_gating,
+                                 bool merge_qkv,
+                                 bool use_flash_attn,
+                                 MetaTensor* query_transpose_out,
+                                 MetaTensor* key_transpose_out,
+                                 MetaTensor* value_transpose_out,
+                                 MetaTensor* qkv_transpose_out,
+                                 MetaTensor* softmax_out,
+                                 MetaTensor* softmax_lse,
+                                 MetaTensor* fmha_out,
+                                 MetaTensor* gate_out,
+                                 MetaTensor* out,
+                                 MetaConfig config = MetaConfig());
+
+void FusedGateAttentionGradInferMeta(const MetaTensor& query,
+                                     const MetaTensor& key,
+                                     const MetaTensor& query_weight,
+                                     const MetaTensor& key_weight,
+                                     const MetaTensor& value_weight,
+                                     const MetaTensor& qkv_weight,
+                                     const MetaTensor& nonbatched_bias,
+                                     const MetaTensor& src_mask,
+                                     const MetaTensor& gate_weight,
+                                     const MetaTensor& gate_bias,
+                                     const MetaTensor& out_linear_weight,
+                                     const MetaTensor& out_linear_bias,
+                                     const MetaTensor& query_transpose_out,
+                                     const MetaTensor& key_transpose_out,
+                                     const MetaTensor& value_transpose_out,
+                                     const MetaTensor& qkv_transpose_out,
+                                     const MetaTensor& softmax_out,
+                                     const MetaTensor& softmax_lse,
+                                     const MetaTensor& fmha_out,
+                                     const MetaTensor& gate_out,
+                                     const MetaTensor& out_grad,
+                                     bool has_gating,
+                                     bool merge_qkv,
+                                     bool use_flash_attn,
+                                     MetaTensor* query_grad,
+                                     MetaTensor* key_grad,
+                                     MetaTensor* query_weight_grad,
+                                     MetaTensor* key_weight_grad,
+                                     MetaTensor* value_weight_grad,
+                                     MetaTensor* qkv_weight_grad,
+                                     MetaTensor* nonbatched_bias_grad,
+                                     MetaTensor* gate_weight_grad,
+                                     MetaTensor* gate_bias_grad,
+                                     MetaTensor* out_linear_weight_grad,
+                                     MetaTensor* out_linear_bias_grad,
+                                     MetaConfig config = MetaConfig());
+
 void ResnetBasicBlockInferMeta(const MetaTensor& x,
                                const MetaTensor& filter1,
                                const MetaTensor& scale1,