【Hackathon 6th Fundable Projects 3 No.360】fluid operator sparse_attention(#64175)

co63oc · web-flow · commit 5859bb9b4dc0 · 2024-05-15T10:14:48.000+08:00
diff --git a/paddle/fluid/operators/CMakeLists.txt b/paddle/fluid/operators/CMakeLists.txt
@@ -83,7 +83,7 @@ endif()
 set(OP_HEADER_DEPS ${OP_HEADER_DEPS} phi common phi_utils static_prim_api get_expected_kernel_func)
 
 register_operators(EXCLUDES py_func_op dgc_op generated_op1 generated_op2 generated_op3 generated_op4 load_combine_op lstm_op run_program_op quantize_linear_op
-        recurrent_op save_combine_op sparse_attention_op sync_batch_norm_op activation_op ${OP_MKL_DEPS} DEPS ${OP_HEADER_DEPS} processgroup_comm_utils)
+        recurrent_op save_combine_op sync_batch_norm_op activation_op ${OP_MKL_DEPS} DEPS ${OP_HEADER_DEPS} processgroup_comm_utils)
 
 op_library(generated_op UNITY SRCS generated_op1.cc generated_op2.cc generated_op3.cc generated_op4.cc DEPS ${OP_HEADER_DEPS})
 op_library(run_program_op DEPS executor_cache ${OP_HEADER_DEPS})
@@ -106,9 +106,6 @@ endif()
 
 if (WITH_GPU OR WITH_ROCM)
     op_library(sync_batch_norm_op DEPS processgroup_comm_utils)
-    if ((NOT WIN32) AND (NOT WITH_ROCM) AND (NOT PADDLE_WITH_ARM) AND (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_LESS 11.3) )
-        op_library(sparse_attention_op DEPS processgroup_comm_utils)
-    endif()
 endif()
 
 op_library(lstm_op DEPS ${OP_HEADER_DEPS})
diff --git a/paddle/fluid/operators/sparse_attention_op.cc b/paddle/fluid/operators/sparse_attention_op.cc
diff --git a/paddle/phi/api/yaml/backward.yaml b/paddle/phi/api/yaml/backward.yaml
@@ -2879,6 +2879,19 @@
     param : [x, x]
   composite: silu_double_grad(x, out, grad_out, grad_x_grad, x_grad, grad_out_grad)
 
+- backward_op: sparse_attention_grad
+  forward: sparse_attention(Tensor q, Tensor k, Tensor v, Tensor offset, Tensor columns, Tensor key_padding_mask,
+    Tensor attn_mask) -> Tensor (out), Tensor (sparse_dot_sdd), Tensor (softmax)
+  args: (Tensor q, Tensor k, Tensor v, Tensor offset, Tensor columns, Tensor sparse_dot_sdd,
+    Tensor softmax, Tensor out_grad)
+  output: Tensor (q_grad), Tensor (k_grad), Tensor (v_grad)
+  infer_meta:
+    func: GeneralTernaryGradInferMeta
+    param: [q, k, v]
+  kernel:
+    func: sparse_attention_grad
+    data_type: out_grad
+
 - backward_op: stft_grad
   forward: stft (Tensor x, Tensor window, int n_fft, int hop_length, bool normalized, bool onesided) -> Tensor (out)
   args: (Tensor x, Tensor window, Tensor out_grad, int n_fft, int hop_length, bool normalized, bool onesided)
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -4198,6 +4198,13 @@
   outputs :
     out : Out
 
+- op: sparse_attention
+  backward: sparse_attention_grad
+  inputs:
+    {q : Q, k : K, v : V, offset : Offset, columns : Columns, key_padding_mask : KeyPaddingMask, attn_mask : AttnMask}
+  outputs:
+    {out : Out, sparse_dot_sdd : SparseDotSdd, softmax : Softmax}
+
 - op: sparse_momentum
   inputs :
     {param: Param, grad: Grad, velocity: Velocity, index: Index, axis: Axis, learning_rate: LearningRate,master_param: MasterParam}
diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml
@@ -2888,6 +2888,19 @@
     data_type : x
   backward : solve_grad
 
+- op : sparse_attention
+  args: (Tensor q, Tensor k, Tensor v, Tensor offset, Tensor columns, Tensor key_padding_mask,
+    Tensor attn_mask)
+  output: Tensor (out), Tensor (sparse_dot_sdd), Tensor (softmax)
+  infer_meta:
+    func: SparseAttentionInferMeta
+  kernel:
+    func: sparse_attention
+    data_type: q
+  optional: key_padding_mask, attn_mask
+  intermediate: sparse_dot_sdd, softmax
+  backward: sparse_attention_grad
+
 - op : spectral_norm
   args : (Tensor weight, Tensor u, Tensor v, int dim = 0, int power_iters = 1, float eps = 1e-12f)
   output : Tensor
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -4314,6 +4314,46 @@ void SendUVInferMeta(const MetaTensor& x,
   out->set_dims(common::make_ddim(out_dims_array));
 }
 
+void SparseAttentionInferMeta(const MetaTensor& q,
+                              const MetaTensor& k,
+                              const MetaTensor& v,
+                              const MetaTensor& offset,
+                              const MetaTensor& columns,
+                              const MetaTensor& key_padding_mask,
+                              const MetaTensor& attn_mask,
+                              MetaTensor* out,
+                              MetaTensor* sparse_dot_sdd,
+                              MetaTensor* softmax) {
+  const auto& dims_q = q.dims();
+  const auto& dims_k = k.dims();
+  const auto& dims_v = v.dims();
+  const auto& dims_columns = columns.dims();
+
+  PADDLE_ENFORCE_EQ(
+      dims_q.size(),
+      static_cast<size_t>(4),
+      phi::errors::InvalidArgument("Dimension in query' shapes should be 4."));
+  PADDLE_ENFORCE_EQ(
+      dims_k.size(),
+      static_cast<size_t>(4),
+      phi::errors::InvalidArgument("Dimension in key' shapes should be 4."));
+  PADDLE_ENFORCE_EQ(
+      dims_v.size(),
+      static_cast<size_t>(4),
+      phi::errors::InvalidArgument("Dimension in value' shapes should be 4."));
+
+  auto batch_size = dims_q[0];
+  auto num_heads = dims_q[1];
+  auto M = dims_q[2];
+  auto N = dims_q[3];
+  auto sparse_nnz = dims_columns[2];
+  out->set_dims({batch_size, num_heads, M, N});
+  sparse_dot_sdd->set_dims({batch_size, num_heads, sparse_nnz});
+  softmax->set_dims({batch_size, num_heads, sparse_nnz});
+  out->share_lod(q);
+  out->set_dtype(q.dtype());
+}
+
 void StackInferMeta(const std::vector<const MetaTensor*>& x,
                     int axis,
                     MetaTensor* out,
diff --git a/paddle/phi/infermeta/multiary.h b/paddle/phi/infermeta/multiary.h
@@ -798,6 +798,17 @@ void SigmoidCrossEntropyWithLogitsInferMeta(const MetaTensor& x,
                                             MetaTensor* out,
                                             MetaConfig config = MetaConfig());
 
+void SparseAttentionInferMeta(const MetaTensor& q,
+                              const MetaTensor& k,
+                              const MetaTensor& v,
+                              const MetaTensor& offset,
+                              const MetaTensor& columns,
+                              const MetaTensor& key_padding_mask,
+                              const MetaTensor& attn_mask,
+                              MetaTensor* out,
+                              MetaTensor* sparse_dot_sdd,
+                              MetaTensor* softmax);
+
 void StackInferMeta(const std::vector<const MetaTensor*>& x,
                     int axis,
                     MetaTensor* out,
diff --git a/paddle/phi/kernels/sparse/gpu/sparse_attention_kernel.cu b/paddle/phi/kernels/sparse/gpu/sparse_attention_kernel.cu
diff --git a/test/legacy_test/test_sparse_attention_op.py b/test/legacy_test/test_sparse_attention_op.py
