PaddlePaddle
diff --git a/‎paddle/phi/api/yaml/backward.yaml
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/phi/api/yaml/backward.yaml
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/phi/api/yaml/generator/dist_api_gen.py
Copy file name to clipboard
+12-1Lines changed: 12 additions & 1 deletion b/‎paddle/phi/api/yaml/generator/dist_api_gen.py
Copy file name to clipboard
+12-1Lines changed: 12 additions & 1 deletion
diff --git a/‎paddle/phi/api/yaml/ops.yaml
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/phi/api/yaml/ops.yaml
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
@@ -830,6 +830,7 @@
   infer_meta :
     func : FlashAttnGradInferMeta
     param : [q, k, v]
+    spmd_rule : FlashAttGradInferSpmd
   kernel :
     func : flash_attn_grad
     data_type: q
 
@@ -769,6 +769,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>&"
@@ -901,7 +909,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
 
@@ -930,6 +930,7 @@
   infer_meta :
     func : FlashAttnInferMeta
     param : [q, k, v]
+    spmd_rule : FlashAttInferSpmd
   kernel :
     func : flash_attn
     data_type : q