[PIR]Migrate scaled_dot_product_attention to pir (#61097)

0x45f · web-flow · commit 606b5538e332 · 2024-01-26T10:54:47.000+08:00
diff --git a/python/paddle/incubate/nn/memory_efficient_attention.py b/python/paddle/incubate/nn/memory_efficient_attention.py
@@ -22,7 +22,7 @@
 import paddle
 from paddle import _C_ops
 from paddle.base.layer_helper import LayerHelper
-from paddle.framework import in_dynamic_mode
+from paddle.framework import in_dynamic_or_pir_mode
 
 from .attn_bias import (
     BlockDiagonalCausalMask,
@@ -99,7 +99,7 @@ def memory_efficient_attention(
     bias = _get_tensor_bias(attn_bias)
     is_test = not training
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         output, logsumexp, seed_and_offset = _C_ops.memory_efficient_attention(
             query,
             key,
diff --git a/python/paddle/nn/functional/flash_attention.py b/python/paddle/nn/functional/flash_attention.py
@@ -15,6 +15,7 @@
 import paddle
 import paddle.nn.functional as F
 from paddle import _C_ops, in_dynamic_mode
+from paddle.base.framework import in_dynamic_or_pir_mode
 from paddle.base.layer_helper import LayerHelper
 from paddle.base.wrapped_decorator import signature_safe_contextmanager
 
@@ -225,7 +226,7 @@ def flash_attention(
     sdp_func_name = _select_sdp(head_dim)
 
     if sdp_func_name == "flash_attn":
-        if in_dynamic_mode():
+        if in_dynamic_or_pir_mode():
             (result_attention, result_softmax, _, _) = _C_ops.flash_attn(
                 query,
                 key,
@@ -505,7 +506,7 @@ def scaled_dot_product_attention(
         out, _ = flash_attention(query, key, value, dropout_p, is_causal)
         return out
     else:
-        if in_dynamic_mode():
+        if in_dynamic_or_pir_mode():
             fixed_seed_offset = (None,)
             return_softmax = False
             rng_name = ""
diff --git a/test/legacy_test/test_flash_attention.py b/test/legacy_test/test_flash_attention.py
@@ -28,6 +28,7 @@
     flash_attn_unpadded,
     scaled_dot_product_attention,
 )
+from paddle.pir_utils import test_with_pir_api
 
 logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s")
 
@@ -204,6 +205,7 @@ def test_unpadded(self):
 
         paddle.disable_static()
 
+    @test_with_pir_api
     def test_all(self):
         print(
             f"Test case shape {self.shape} dtype {self.dtype} causal {self.causal}"
@@ -383,7 +385,7 @@ class TestFlashAttentionAPITest1(TestFlashAttentionAPI):
     def setUp(self):
         self.place = paddle.CUDAPlace(0)
         self.shape = (2, 128, 8, 16)
-        self.dtype = paddle.float16
+        self.dtype = 'float16'
         self.dropout = 0.0
         self.causal = False
         self.return_softmax = False
@@ -394,7 +396,7 @@ class TestFlashAttentionAPITest2(TestFlashAttentionAPI):
     def setUp(self):
         self.place = paddle.CUDAPlace(0)
         self.shape = (2, 256, 8, 16)
-        self.dtype = paddle.float16
+        self.dtype = 'float16'
         self.dropout = 0.0
         self.causal = False
         self.return_softmax = False
@@ -405,7 +407,7 @@ class TestFlashAttentionAPITest3(TestFlashAttentionAPI):
     def setUp(self):
         self.place = paddle.CUDAPlace(0)
         self.shape = (2, 512, 8, 16)
-        self.dtype = paddle.float16
+        self.dtype = 'float16'
         self.dropout = 0.0
         self.causal = True
         self.return_softmax = False
@@ -416,7 +418,7 @@ class TestFlashAttentionAPITest4(TestFlashAttentionAPI):
     def setUp(self):
         self.place = paddle.CUDAPlace(0)
         self.shape = (8, 1024, 16, 128)
-        self.dtype = paddle.float16
+        self.dtype = 'float16'
         self.dropout = 0.0
         self.causal = False
         self.return_softmax = False
@@ -427,7 +429,7 @@ class TestFlashAttentionAPITest5(TestFlashAttentionAPI):
     def setUp(self):
         self.place = paddle.CUDAPlace(0)
         self.shape = (8, 1024, 16, 256)
-        self.dtype = paddle.float16
+        self.dtype = 'float16'
         self.dropout = 0.0
         self.causal = False
         self.return_softmax = False
@@ -438,7 +440,7 @@ class TestMathAttentionAPITest(TestFlashAttentionAPI):
     def setUp(self):
         self.place = paddle.CUDAPlace(0)
         self.shape = (8, 1024, 16, 128)
-        self.dtype = paddle.float16
+        self.dtype = 'float16'
         self.dropout = 0.0
         self.causal = False
         self.return_softmax = False
@@ -453,7 +455,7 @@ class TestSDPAttentionAPITest(TestFlashAttentionAPI):
     def setUp(self):
         self.place = paddle.CUDAPlace(0)
         self.shape = (8, 1024, 16, 128)
-        self.dtype = paddle.float16
+        self.dtype = 'float16'
         self.dropout = 0.0
         self.causal = False
         self.return_softmax = False
@@ -468,7 +470,7 @@ class TestFlashAttenionWithMaskAPITest(TestFlashAttentionWithMaskAPI):
     def setUp(self):
         self.place = paddle.CUDAPlace(0)
         self.shape = (8, 1024, 16, 128)
-        self.dtype = paddle.float16
+        self.dtype = 'float16'
         self.dropout = 0.0
         self.causal = False
 
@@ -552,7 +554,7 @@ def setUp(self):
         self.seq_len = 8192
         self.head_dim = 128
         self.num_group = 2
-        self.dtype = paddle.bfloat16
+        self.dtype = 'bfloat16'
 
     def gen_unpadded_data(self, dtype):
         seq_len_q = np.random.randint(
