【Infer Symbolic Shape No.208】[BUAA] Add masked_multihead_attention_ op (#67861)

Luohongzhige · web-flow · commit 918ee36da5f1 · 2024-09-06T10:07:15.000+08:00
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
@@ -2118,12 +2118,58 @@ bool NearestInterpOpInferSymbolicShape(
   return BicubicInterpOpInferSymbolicShape(op, infer_context);
 }
 
-// bool MaskedMultiheadAttention_OpInferSymbolicShape(pir::Operation *op,
-//                                                    pir::InferSymbolicShapeContext
-//                                                    *infer_context) {
-//   // pass
-//   return true;
-// }
+bool MaskedMultiheadAttention_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const symbol::ShapeOrDataDimExprs &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const symbol::ShapeOrDataDimExprs &cache_kv_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+
+  const std::vector<symbol::DimExpr> &x_shape = x_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &cache_kv_shape =
+      cache_kv_shape_or_data.shape();
+
+  std::string compute_dtype =
+      op->attribute<pir::StrAttribute>("compute_dtype").AsString();
+
+  PADDLE_ENFORCE_EQ(
+      cache_kv_shape.size(),
+      5,
+      phi::errors::InvalidArgument("The cache_kv must be 5 dims."));
+  infer_context->AddEqualCstr(cache_kv_shape[0], symbol::DimExpr(2));
+  // TODO(Luohongzhige, Buaa): add constrain for the num_head and k_num_head
+
+  symbol::DimExpr bsz = x_shape[0];
+  symbol::DimExpr dim_head = cache_kv_shape[4];
+  symbol::DimExpr k_num_head = cache_kv_shape[2];
+  symbol::DimExpr v_num_head = k_num_head;
+  symbol::DimExpr num_head =
+      (x_shape[x_shape.size() - 1] / dim_head - k_num_head - v_num_head);
+  std::vector<symbol::DimExpr> out_shape = {bsz, num_head * dim_head};
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(cache_kv_shape)});
+
+  if (op->operand_source(7) != nullptr) {
+    const symbol::ShapeOrDataDimExprs &beam_cache_offset_shape_or_data =
+        infer_context->GetShapeOrDataForValue(op->operand_source(7));
+    const std::vector<symbol::DimExpr> &beam_cache_offset_shape =
+        beam_cache_offset_shape_or_data.shape();
+    infer_context->SetShapeOrDataForValue(
+        op->result(2),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(beam_cache_offset_shape)});
+  }
+
+  return true;
+}
 
 bool MemoryEfficientAttentionOpInferSymbolicShape(
     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
@@ -77,7 +77,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Linspace)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(LinearInterp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logspace)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lstm)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(MaskedMultiheadAttention_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(MaskedMultiheadAttention_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedAdam)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedAdam_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedMomentum)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3097,7 +3097,7 @@
     data_type : x
   optional : bias, src_mask, cum_offsets, sequence_lengths, rotary_tensor, beam_cache_offset, qkv_out_scale, out_shift, out_smooth
   inplace : (cache_kv -> cache_kv_out), (beam_cache_offset -> beam_cache_offset_out)
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : masked_select
   args : (Tensor x, Tensor mask)
