Enable workspace optimization for cudnn flash attention (#60295)

Wong4j · web-flow · commit dd4ac87358c9 · 2024-02-03T18:57:46.000+08:00
* set workspace_opt to true

* test=cuda117

* update workspace opt

* test=cuda117
diff --git a/paddle/phi/kernels/fusion/gpu/fused_dot_product_attention_op.cu b/paddle/phi/kernels/fusion/gpu/fused_dot_product_attention_op.cu
@@ -161,7 +161,8 @@ void FusedDotProductAttentionGradKernel(const Context &dev_ctx,
                                         DenseTensor *q_grad,
                                         DenseTensor *k_grad,
                                         DenseTensor *v_grad) {
-  PADDLE_ENFORCE_GE(dev_ctx.GetComputeCapability(),
+  auto sm_arch = dev_ctx.GetComputeCapability();
+  PADDLE_ENFORCE_GE(sm_arch,
                     80,
                     phi::errors::PreconditionNotMet(
                         "This op only supports Ampere and later devices, "
@@ -225,6 +226,28 @@ void FusedDotProductAttentionGradKernel(const Context &dev_ctx,
   void *offset_dev_ptr = reinterpret_cast<void *>(
       const_cast<int64_t *>(rng_state.data<int64_t>()) + 1);
 
+  bool use_workspace_opt = false;
+  if (sm_arch >= 90) {
+    // quick estimate of dp workspace size
+    size_t max_seqlen_div_up_q = ((q_seq_len + 64 - 1) / 64) * 64;
+    size_t max_seqlen_div_up_kv = ((kv_seq_len + 64 - 1) / 64) * 64;
+    size_t required_dp_workspace =
+        (batch_size * num_heads * max_seqlen_div_up_q * max_seqlen_div_up_kv *
+             2 +
+         1048576 - 1) /
+        1048576;
+    // default upper limit for dp workspace 256MB
+    size_t max_allowed_dp_workspace = 256;
+    if (required_dp_workspace <= max_allowed_dp_workspace) {
+      use_workspace_opt = true;
+    }
+    auto use_workspace_opt_str =
+        std::getenv("CUDNN_FUSE_ATTN_USE_WORKSPACE_OPT");
+    if (use_workspace_opt_str != nullptr) {
+      use_workspace_opt = static_cast<bool>(std::stoi(use_workspace_opt_str));
+    }
+  }
+
   fused_attn_arbitrary_seqlen_bwd(batch_size,
                                   num_heads,
                                   q_seq_len,
@@ -249,7 +272,7 @@ void FusedDotProductAttentionGradKernel(const Context &dev_ctx,
                                   tensor_dtype,
                                   dev_ctx.stream(),
                                   handle,
-                                  false);
+                                  use_workspace_opt);
 }
 
 }  // namespace fusion
diff --git a/python/paddle/incubate/nn/functional/fused_dot_product_attention.py b/python/paddle/incubate/nn/functional/fused_dot_product_attention.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+
 import paddle
 from paddle import _C_ops
 from paddle.framework import LayerHelper, in_dynamic_mode
@@ -26,6 +28,7 @@ def fused_dot_product_attention(
     dropout_prob,
     is_training,
     is_causal_masking,
+    use_workspace_opt=None,
     return_softmax=False,
 ):
     r"""
@@ -41,6 +44,9 @@ def fused_dot_product_attention(
         dropout_prob (float): The dropout probability.
         is_training (bool): A flag indicating whether it is in train phrase or not.
         is_causal_masking (bool): A flag indicating whether it is causal masking or not. If True, the mask will be ignored.
+        use_workspace_opt (bool, optional): A flag indicating whether to use workspace optimization or not.
+            Default: None. When set to None, the code path will be decided based on its internal logic.
+            Currently, it's only supported on Hopper GPU and will be ignored on other devices.
         return_softmax (bool, optional): A flag indicating whether to return softmax_output or not. Default: False.
 
 
@@ -69,6 +75,11 @@ def fused_dot_product_attention(
         ), "mask shape must be [batch_size, 1, q_seqlen, k_seqlen]"
         mask = mask.astype('int32')
 
+    if use_workspace_opt is False:
+        os.environ['CUDNN_FUSE_ATTN_USE_WORKSPACE_OPT'] = '0'
+    elif use_workspace_opt is True:
+        os.environ['CUDNN_FUSE_ATTN_USE_WORKSPACE_OPT'] = '1'
+
     if in_dynamic_mode():
         out, softmax, _ = _C_ops.fused_dot_product_attention(
             q,
diff --git a/test/legacy_test/test_fuse_dot_product_attention_pass.py b/test/legacy_test/test_fuse_dot_product_attention_pass.py
@@ -25,7 +25,7 @@
 def skip_unit_test():
     return (
         not paddle.is_compiled_with_cuda()
-        or paddle.device.cuda.get_device_capability() != (8, 0)
+        or paddle.device.cuda.get_device_capability()[0] < 8
         or paddle.get_cudnn_version() < 8906
     )
 
diff --git a/test/legacy_test/test_fused_dot_product_attention_op.py b/test/legacy_test/test_fused_dot_product_attention_op.py
@@ -31,7 +31,7 @@
 def skip_unit_test():
     return (
         not paddle.is_compiled_with_cuda()
-        or paddle.device.cuda.get_device_capability() != (8, 0)
+        or paddle.device.cuda.get_device_capability()[0] < 8
         or paddle.get_cudnn_version() < 8906
     )
 
@@ -247,6 +247,7 @@ def _get_fused_attn_out(self):
             self.dropout_prob,
             True,
             self.is_causal_masking,
+            None,
             True,
         )
 

Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,7 @@`
`25`	`25`	`def skip_unit_test():`
`26`	`26`	`return (`
`27`	`27`	`not paddle.is_compiled_with_cuda()`
`28`		`- or paddle.device.cuda.get_device_capability() != (8, 0)`
	`28`	`+ or paddle.device.cuda.get_device_capability()[0] < 8`
`29`	`29`	`or paddle.get_cudnn_version() < 8906`
`30`	`30`	`)`
`31`	`31`
Original file line number	Diff line number	Diff line change
`@@ -31,7 +31,7 @@`
`31`	`31`	`def skip_unit_test():`
`32`	`32`	`return (`
`33`	`33`	`not paddle.is_compiled_with_cuda()`
`34`		`- or paddle.device.cuda.get_device_capability() != (8, 0)`
	`34`	`+ or paddle.device.cuda.get_device_capability()[0] < 8`
`35`	`35`	`or paddle.get_cudnn_version() < 8906`
`36`	`36`	`)`
`37`	`37`
`@@ -247,6 +247,7 @@ def _get_fused_attn_out(self):`
`247`	`247`	`self.dropout_prob,`
`248`	`248`	`True,`
`249`	`249`	`self.is_causal_masking,`
	`250`	`+ None,`
`250`	`251`	`True,`
`251`	`252`	`)`
`252`	`253`