fix fa4 test (PaddlePaddle#6408)

Author: ckl117

fastdeploy/model_executor/layers/attention/flash_attn_backend.py

@@ -118,9 +118,9 @@ def flash_attn_func(
     head_dim: int = 128,
     version: Optional[int] = None,
 ):
+    if FLASH_ATTN_VERSION is None:
+        init_flash_attn_version()
     if version is None:
-        if FLASH_ATTN_VERSION is None:
-            init_flash_attn_version()
         version = FLASH_ATTN_VERSION
     if version == 4:
         assert (

tests/layers/test_flash_attn_func.py

@@ -198,6 +198,9 @@ def test_fa4(self):
             k,
             v,
             attn_mask_q=attn_mask_q,
+            num_heads=num_heads,
+            kv_num_heads=kv_num_heads,
+            head_dim=head_dim,
             version=4,
         )
 

tests/operators/test_flash_mask_attn.py

@@ -91,6 +91,8 @@ def paddle_flash_attn_mask(self, q_input, k_input, v_input, attn_out, mask):
         )
 
     def test_flash_mask_attention(self):
+        if self.sm_version < 89 or self.sm_version >= 100:
+            self.skipTest("flash_mask_attention V3 requires SM89+ but less than SM100.")
         q_input = paddle.randn([self.q_len, self.num_head * self.head_dim], dtype="bfloat16")
         k_input = paddle.randn([self.q_len + self.k_len, self.num_kv_head, self.head_dim], dtype="bfloat16")
         v_input = paddle.randn(k_input.shape, dtype="bfloat16")