PaddlePaddle
diff --git a/‎paddle/fluid/inference/api/helper.cc
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions b/‎paddle/fluid/inference/api/helper.cc
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions
diff --git a/‎paddle/phi/infermeta/fusion.cc
Copy file name to clipboard
+11-5Lines changed: 11 additions & 5 deletions b/‎paddle/phi/infermeta/fusion.cc
Copy file name to clipboard
+11-5Lines changed: 11 additions & 5 deletions
@@ -421,9 +421,9 @@ void RegisterAllCustomOperator(bool use_pir) {
     if (all_op_kernels.find(pair.first) == all_op_kernels.end()) {
       framework::RegisterOperatorWithMetaInfo(pair.second);
     } else {
-      LOG(INFO) << "The operator `" << pair.first
-                << "` has been registered. "
-                   "Therefore, we will not repeat the registration here.";
+      VLOG(3) << "The operator `" << pair.first
+              << "` has been registered. Therefore, we will not repeat the "
+                 "registration here.";
     }
   }
 }
 
@@ -290,9 +290,15 @@ void BlockMultiheadAttentionInferMeta(const MetaTensor& qkv,
                                       MetaTensor* value_cache_out) {
   auto input_dims = qkv.dims();
   auto key_cache_dims = key_cache.dims();
-  const int num_head = key_cache_dims[1];
+  const int kv_num_head = key_cache_dims[1];
   const int dim_head = key_cache_dims[3];
+  const int total_num_head = qkv.dims()[qkv.dims().size() - 1] / dim_head;
+  const int q_num_head = total_num_head - 2 * kv_num_head;
 
+  PADDLE_ENFORCE_EQ(q_num_head % kv_num_head,
+                    0,
+                    errors::InvalidArgument(
+                        "The q num_head must be divisible by kv_num_head"));
   PADDLE_ENFORCE_EQ(
       input_dims.size(),
       2UL,
@@ -302,12 +308,12 @@ void BlockMultiheadAttentionInferMeta(const MetaTensor& qkv,
       4UL,
       errors::InvalidArgument("The input(key_cache) must be a 4D Tensor."));
   PADDLE_ENFORCE_EQ(
-      3 * num_head * dim_head,
+      (2 * kv_num_head + q_num_head) * dim_head,
       input_dims[1],
-      errors::InvalidArgument(
-          "The input_dims[1] must be equal to 3 * num_head * dim_head"));
+      errors::InvalidArgument("The input_dims[1] must be equal to (2 * "
+                              "kv_num_head + q_num_head) * dim_head"));
 
-  fmha_out->set_dims({input_dims[0], num_head * dim_head});
+  fmha_out->set_dims({input_dims[0], q_num_head * dim_head});
   qkv_out->set_dims(qkv.dims());
   key_cache_out->set_dims(key_cache_dims);
   key_cache_out->set_dtype(key_cache.dtype());
Original file line number	Diff line number	Diff line change
`@@ -421,9 +421,9 @@ void RegisterAllCustomOperator(bool use_pir) {`
`421`	`421`	`if (all_op_kernels.find(pair.first) == all_op_kernels.end()) {`
`422`	`422`	`framework::RegisterOperatorWithMetaInfo(pair.second);`
`423`	`423`	`} else {`
`424`		- LOG(INFO) << "The operator `" << pair.first
`425`		- << "` has been registered. "
`426`		`- "Therefore, we will not repeat the registration here.";`
	`424`	+ VLOG(3) << "The operator `" << pair.first
	`425`	+ << "` has been registered. Therefore, we will not repeat the "
	`426`	`+ "registration here.";`
`427`	`427`	`}`
`428`	`428`	`}`
`429`	`429`	`}`