[CINN] 【Infer Symbolic Shape BUAA 】Add flash_attn_varlen_qkvpacked op (#68318)

crazyxiaoxi · web-flow · commit c50b16e1aa74 · 2024-10-12T14:23:48.000+08:00
* try

* try

* fix bug

* build fix

* fix

* comment
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
@@ -1671,12 +1671,75 @@ bool FusedFeedforwardOpInferSymbolicShape(
 //   return true;
 // }
 
-// bool FlashAttnQkvpackedOpInferSymbolicShape(pir::Operation *op,
-//                                             pir::InferSymbolicShapeContext
-//                                             *infer_context) {
-//   // pass
-//   return true;
-// }
+bool FlashAttnVarlenQkvpackedOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &qkv_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const std::vector<symbol::DimExpr> &qkv_shape = qkv_shape_or_data.shape();
+
+  auto round_multiple = [](symbol::DimExpr x) {
+    auto m = symbol::DimExpr{128};
+    auto m_minus_one = symbol::DimExpr{127};
+    return (x + m_minus_one) / m * m;
+  };
+
+  size_t rank = qkv_shape.size();
+  PADDLE_ENFORCE_EQ((rank == 4UL || rank == 5UL),
+                    true,
+                    common::errors::InvalidArgument(
+                        "qkv dims must be 4(unpadded) or 5(padded batch)"));
+  std::vector<symbol::DimExpr> out_dims;
+  std::vector<symbol::DimExpr> softmax_lse_shape;
+  std::vector<symbol::DimExpr> softmax_shape;
+  if (rank == 4UL) {
+    // qkv [total_*,nheads/nheads_k+2,nheads_k,headdim]
+    out_dims = {qkv_shape[0], (qkv_shape[1] - 2) * qkv_shape[2], qkv_shape[3]};
+    softmax_shape = {qkv_shape[0],
+                     (qkv_shape[1] - 2) * qkv_shape[2],
+                     infer_context->GetNextSymName(),
+                     infer_context->GetNextSymName()};
+    softmax_lse_shape = {qkv_shape[0],
+                         (qkv_shape[1] - 2) * qkv_shape[2],
+                         infer_context->GetNextSymName()};
+  } else if (rank == 5UL) {
+    // qkv [batchsize,seqlen,nheads/nheads_k+2,nheads_k,headdim]
+    out_dims = {qkv_shape[0],
+                qkv_shape[1],
+                (qkv_shape[2] - 2) * qkv_shape[3],
+                qkv_shape[4]};
+    softmax_shape = {qkv_shape[0],
+                     (qkv_shape[2] - 2) * qkv_shape[3],
+                     round_multiple(qkv_shape[1]),
+                     infer_context->GetNextSymName()};
+    softmax_lse_shape = {qkv_shape[0],
+                         (qkv_shape[2] - 2) * qkv_shape[3],
+                         round_multiple(qkv_shape[1])};
+  }
+
+  // Set output tensor shapes
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(out_dims)});
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs({softmax_shape})});
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(2),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(softmax_lse_shape)});
+
+  if (!paddle::dialect::details::IsFakeValue(op->result(3))) {
+    std::vector<symbol::DimExpr> seed_offset_dims = {symbol::DimExpr(2)};
+    infer_context->SetShapeOrDataForValue(
+        op->result(3),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(seed_offset_dims)});
+  }
+  return true;
+}
 
 // bool FlashAttnUnpaddedOpInferSymbolicShape(pir::Operation *op,
 //                                            pir::InferSymbolicShapeContext
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,7 +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(FusedBatchNormAct)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FusedBatchNormAct_)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1925,7 +1925,7 @@
     data_type : qkv
   intermediate : softmax_lse, seed_offset
   backward : flash_attn_varlen_qkvpacked_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : flashmask_attention
   args : (Tensor q, Tensor k, Tensor v, Tensor startend_row_indices,  Tensor fixed_seed_offset, float dropout = 0.0, bool causal = false, bool return_softmax = false, bool is_test = false, str rng_name = "")
