Flash att infer

paddle/phi/api/yaml/backward.yaml

@@ -830,6 +830,7 @@
   infer_meta :
     func : FlashAttnGradInferMeta
     param : [q, k, v]
+    spmd_rule : FlashAttGradInferSpmd
   kernel :
     func : flash_attn_grad
     data_type: q

paddle/phi/api/yaml/generator/dist_api_gen.py

@@ -768,6 +768,14 @@ def generate_specialized_infer_spmd_code(self) -> str:
                         name=param
                     )
                     input_args_code += "meta_dist_input_" + param + ", "
+                elif (
+                    self.inputs['input_info'][param]
+                    == "const paddle::optional<Tensor>&"
+                ):
+                    input_decl_code += (
+                        OPTIONAL_SINGLE_DIST_META_IN_TEMPLATE.format(name=param)
+                    )
+                    input_args_code += "meta_dist_input_" + param + ", "
                 elif (
                     self.inputs['input_info'][param]
                     == "const std::vector<Tensor>&"
@@ -900,7 +908,10 @@ def generate_output_creation_code(self) -> str:
             # kernel output generate
             self.dist_output_args.append('dist_out')
             self.dense_output_args.append('dense_out')
-            if self.outputs['types'][0] == 'Tensor':
+            if (
+                self.outputs['types'][0] == 'Tensor'
+                or self.outputs['types'][0] == 'const paddle::optional<Tensor>'
+            ):
                 if (
                     self.need_to_generate_code_for_inplace_impl(0)
                     and self.generate_general_infer_spmd

paddle/phi/api/yaml/ops.yaml

@@ -929,6 +929,7 @@
   infer_meta :
     func : FlashAttnInferMeta
     param : [q, k, v]
+    spmd_rule : FlashAttInferSpmd
   kernel :
     func : flash_attn
     data_type : q