[Infer Symbolic Shape No.199] Add symbol_infer_interface for flash_attn_qkvpacked (#68984)

ooooo-create · web-flow · commit 70ae127f17ca · 2024-11-05T17:13:30.000+08:00
* Add symbol_infer_interface for flash_attn_qkvpacked

* Refine logic

* Fix ci
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
@@ -1723,6 +1723,11 @@ bool FlashAttnVarlenQkvpackedOpInferSymbolicShape(
   return true;
 }
 
+bool FlashAttnQkvpackedOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return FlashAttnVarlenQkvpackedOpInferSymbolicShape(op, infer_context);
+}
+
 // bool FlashAttnUnpaddedOpInferSymbolicShape(pir::Operation *op,
 //                                            pir::InferSymbolicShapeContext
 //                                            *infer_context) {
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h
@@ -55,6 +55,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(FullWithTensor)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FlashAttn)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FusedFeedforward)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(FusedAttention)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FlashAttnQkvpacked)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FlashAttnVarlenQkvpacked)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(FlashAttnUnpadded)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FlashmaskAttention)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1958,7 +1958,7 @@
     func : flash_attn_qkvpacked
     data_type : qkv
   backward : flash_attn_qkvpacked_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : flash_attn_unpadded
   args : (Tensor q, Tensor k, Tensor v, Tensor cu_seqlens_q,  Tensor cu_seqlens_k, Tensor fixed_seed_offset, Tensor attn_mask, int64_t max_seqlen_q, int64_t max_seqlen_k, float scale, float dropout = 0.0, bool causal = false, bool return_softmax = false, bool is_test = false, str rng_name = "")
