[xpu] multi_encoder_xpu support adaptive seqlen fuse v3 (#63825)

NeroLoh · web-flow · commit 3205f865762e · 2024-04-29T11:22:07.000+08:00
diff --git a/paddle/fluid/framework/ir/xpu/multi_encoder_xpu_adaptive_seqlen_fuse_pass.cc b/paddle/fluid/framework/ir/xpu/multi_encoder_xpu_adaptive_seqlen_fuse_pass.cc
@@ -328,6 +328,134 @@ int MultiEncoderXPUAdaptiveSeqlenFusePass::ApplyAdaptiveSeqlenPassV2(
   return found_subgraph_count;
 }
 
+namespace patterns {
+struct AdaptiveSeqlenPatternV3 : public PatternBase {
+  AdaptiveSeqlenPatternV3(PDPattern* pattern,
+                          const std::string& name_scope,
+                          const std::string& matmul_type);
+
+  // declare operator node's name
+  PATTERN_DECL_NODE(multi_encoder_xpu);
+  PATTERN_DECL_NODE(matmul);
+  PATTERN_DECL_NODE(scale);
+  PATTERN_DECL_NODE(stack);
+  // declare variable node's name
+  PATTERN_DECL_NODE(x);
+  PATTERN_DECL_NODE(mask);
+  PATTERN_DECL_NODE(matmul_out);
+  PATTERN_DECL_NODE(scale_out);
+  PATTERN_DECL_NODE(stack_out);
+};
+
+AdaptiveSeqlenPatternV3::AdaptiveSeqlenPatternV3(PDPattern* pattern,
+                                                 const std::string& name_scope,
+                                                 const std::string& matmul_type)
+    : PatternBase(pattern, name_scope, name_scope) {
+  auto* x =
+      pattern->NewNode(x_repr())->assert_is_op_input("multi_encoder_xpu", "x");
+
+  auto* mask = pattern->NewNode(mask_repr())
+                   ->assert_is_op_input(matmul_type, "X")
+                   ->assert_is_op_input(matmul_type, "Y");
+  auto* matmul = pattern->NewNode(matmul_repr())->assert_is_op(matmul_type);
+  auto* matmul_out = pattern->NewNode(matmul_out_repr())
+                         ->assert_is_op_output(matmul_type, "Out")
+                         ->assert_is_op_input("scale", "X");
+  auto* scale = pattern->NewNode(scale_repr())->assert_is_op("scale");
+  auto* scale_out = pattern->NewNode(scale_out_repr())
+                        ->assert_is_op_output("scale", "Out")
+                        ->assert_is_op_input("stack", "X");
+  auto* stack = pattern->NewNode(stack_repr())->assert_is_op("stack");
+  auto* stack_out = pattern->NewNode(stack_out_repr())
+                        ->assert_is_op_output("stack", "Y")
+                        ->assert_is_op_input("multi_encoder_xpu", "mask");
+
+  auto* multi_encoder_xpu = pattern->NewNode(multi_encoder_xpu_repr())
+                                ->assert_is_op("multi_encoder_xpu");
+
+  matmul->LinksFrom({mask}).LinksTo({matmul_out});
+  scale->LinksFrom({matmul_out}).LinksTo({scale_out});
+  stack->LinksFrom({scale_out}).LinksTo({stack_out});
+  multi_encoder_xpu->LinksFrom({x, stack_out});
+}
+
+}  // namespace patterns
+
+int MultiEncoderXPUAdaptiveSeqlenFusePass::ApplyAdaptiveSeqlenPassV3(
+    ir::Graph* graph, const std::string& matmul_type) const {
+  GraphPatternDetector gpd;
+  patterns::AdaptiveSeqlenPatternV3 pattern(
+      gpd.mutable_pattern(), name_scope_, matmul_type);
+
+  int found_subgraph_count = 0;
+  auto handler = [&](const GraphPatternDetector::subgraph_t& subgraph,
+                     Graph* graph) {
+    VLOG(4) << "handle ApplyAdaptiveSeqlenPassV3 fuse";
+    GET_IR_NODE(multi_encoder_xpu);
+    GET_IR_NODE(matmul);
+    GET_IR_NODE(scale);
+    GET_IR_NODE(stack);
+    GET_IR_NODE(x);
+    GET_IR_NODE(mask);
+    GET_IR_NODE(matmul_out);
+    GET_IR_NODE(scale_out);
+    GET_IR_NODE(stack_out);
+
+    std::string mask_name = mask->Name();
+    std::string seq_len_name = mask_name + "_seq_len";
+    VarDesc seq_len_desc(seq_len_name);
+    auto* seq_len = graph->CreateVarNode(&seq_len_desc);
+    std::string seq_lod_name = mask_name + "_seq_lod";
+    VarDesc seq_lod_desc(seq_lod_name);
+    auto* seq_lod = graph->CreateVarNode(&seq_lod_desc);
+    std::string max_seq_len_name = mask_name + "_max_seq_len";
+    VarDesc max_seq_len_desc(max_seq_len_name);
+    auto* max_seq_len = graph->CreateVarNode(&max_seq_len_desc);
+    std::string x_vsl_name = x->Name() + "_vsl_packed";
+    VarDesc x_vsl_desc(x_vsl_name);
+    auto* x_vsl = graph->CreateVarNode(&x_vsl_desc);
+
+    framework::OpDesc op_desc;
+    op_desc.SetType("mask_adaptive_xpu");
+    op_desc.SetInput("mask", {mask_name});
+    op_desc.SetOutput("length", {seq_len_name});
+    op_desc.SetOutput("seq_lod", {seq_lod_name});
+    op_desc.SetOutput("pad_seq_len", {max_seq_len_name});
+    auto* mask_adaptive_xpu = graph->CreateOpNode(&op_desc);
+
+    framework::OpDesc new_op_desc;
+    new_op_desc.SetType("sequence_unpad_xpu");
+    new_op_desc.SetInput("x", {x->Name()});
+    new_op_desc.SetInput("length", {seq_len_name});
+    new_op_desc.SetOutput("out", {x_vsl_name});
+    auto* sequence_unpad = graph->CreateOpNode(&new_op_desc);
+
+    multi_encoder_xpu->Op()->SetInput("x", {x_vsl_name});
+    multi_encoder_xpu->Op()->SetInput("seq_lod", {seq_lod_name});
+    multi_encoder_xpu->Op()->SetInput("max_seq_len", {max_seq_len_name});
+    multi_encoder_xpu->Op()->RemoveInput("mask");
+    IR_NODE_LINK_TO(mask, mask_adaptive_xpu);
+    IR_NODE_LINK_TO(mask_adaptive_xpu, seq_len);
+    IR_NODE_LINK_TO(mask_adaptive_xpu, seq_lod);
+    IR_NODE_LINK_TO(mask_adaptive_xpu, max_seq_len);
+    IR_NODE_LINK_TO(x, sequence_unpad);
+    IR_NODE_LINK_TO(seq_len, sequence_unpad);
+    IR_NODE_LINK_TO(sequence_unpad, x_vsl);
+    IR_NODE_LINK_TO(x_vsl, multi_encoder_xpu);
+    IR_NODE_LINK_TO(seq_lod, multi_encoder_xpu);
+    IR_NODE_LINK_TO(max_seq_len, multi_encoder_xpu);
+
+    // delete useless node
+    std::unordered_set<const Node*> delete_nodes{
+        matmul, scale, stack, matmul_out, scale_out, stack_out};
+    GraphSafeRemoveNodes(graph, delete_nodes);
+    found_subgraph_count++;
+  };
+
+  gpd(graph, handler);
+  return found_subgraph_count;
+}
+
 void MultiEncoderXPUAdaptiveSeqlenFusePass::ApplyImpl(ir::Graph* graph) const {
   PADDLE_ENFORCE_NOT_NULL(
       graph, platform::errors::PreconditionNotMet("graph should not be null."));
@@ -337,6 +465,7 @@ void MultiEncoderXPUAdaptiveSeqlenFusePass::ApplyImpl(ir::Graph* graph) const {
   for (auto& matmul_type : matmul_types) {
     found_subgraph_count += ApplyAdaptiveSeqlenPassV1(graph, matmul_type);
     found_subgraph_count += ApplyAdaptiveSeqlenPassV2(graph, matmul_type);
+    found_subgraph_count += ApplyAdaptiveSeqlenPassV3(graph, matmul_type);
   }
 
   AddStatis(found_subgraph_count);
diff --git a/paddle/fluid/framework/ir/xpu/multi_encoder_xpu_adaptive_seqlen_fuse_pass.h b/paddle/fluid/framework/ir/xpu/multi_encoder_xpu_adaptive_seqlen_fuse_pass.h
@@ -136,6 +136,47 @@ class MultiEncoderXPUAdaptiveSeqlenFusePass : public FusePassBase {
   int ApplyAdaptiveSeqlenPassV2(ir::Graph* graph,
                                 const std::string& matmul_type) const;
 
+  /*
+    adaptive seqlen V3, before:
+                      x            mask
+                      |             |
+                      |             |
+                      |           matmul
+                      |             |
+                      |             |
+                      |           scale
+                      |           /
+                      |        stack
+                      |         |
+                      |        /
+                      |      /
+                  multi_encoder_xpu
+                      |
+                      |
+                     out
+  -------------------------------------------
+   After the pass apply:
+                      x      mask
+                      |        |
+                      |        |
+                      | xpu_adaptive_mask
+                      |        |     |
+            sequence_unpad<--length  |
+                      |              |
+                      |            pad_seq_len
+                      |            seq_lod
+                      |            /
+                      |          /
+                      |        /
+              multi_encoder_xpu
+                      |
+                      |
+                    out
+  -------------------------------------------
+  */
+  int ApplyAdaptiveSeqlenPassV3(ir::Graph* graph,
+                                const std::string& matmul_type) const;
+
  private:
   const std::string name_scope_{"multi_encoder_xpu_adaptive_seqlen_fuse_pass"};
 };
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -651,6 +651,13 @@ void AnalysisPredictor::InitDeviceContexts() {
               config_.xpu_config_.fc_autotune_level,
               config_.xpu_config_.fc_autotune_file_writeback,
               place_);
+          if (config_.xpu_config_.transformer_softmax_optimize_level > 0) {
+            xpu_context->SetContextOption(
+                "XPU_SOFTMAX_OPT",
+                std::to_string(
+                    config_.xpu_config_.transformer_softmax_optimize_level)
+                    .c_str());
+          }
           xpu_context->SetAllocator(instance.GetAllocator(place_).get());
           xpu_context->SetGenerator(
               phi::DefaultXPUGenerator(place_.GetDeviceId()).get());
diff --git a/paddle/fluid/inference/api/infer_context.cc b/paddle/fluid/inference/api/infer_context.cc
@@ -180,6 +180,16 @@ void InferXPUContext::SetConvAutotuneInfo(std::string conv_autotune_file,
                           conv_autotune_file_writeback));
   }
 }
+void InferXPUContext::SetContextOption(const char* name, const char* value) {
+  phi::backends::xpu::XPUDeviceGuard guard(GetPlace().GetDeviceId());
+  VLOG(5) << "XPU Set Option name:" << name << " value:" << value;
+  int ret;
+  ret = x_context()->set_option(name, value);
+  PADDLE_ENFORCE_EQ(
+      ret,
+      0,
+      platform::errors::Unavailable("Failed to set XPU option %s.", name));
+}
 
 void InferXPUContext::SetFcAutotuneInfo(std::string fc_autotune_file,
                                         int fc_autotune_level,
diff --git a/paddle/fluid/inference/api/infer_context.h b/paddle/fluid/inference/api/infer_context.h
@@ -78,6 +78,7 @@ class InferXPUContext : public phi::XPUContext {
                          int fc_autotune_level,
                          bool fc_autotune_file_writeback,
                          const phi::Place& place);
+  void SetContextOption(const char* name, const char* value);
 
  private:
   size_t l3_size_{0};
diff --git a/paddle/phi/api/yaml/fused_ops.yaml b/paddle/phi/api/yaml/fused_ops.yaml
@@ -403,6 +403,15 @@
     func : layer_norm_act_xpu
     data_type : x
 
+- op : mask_adaptive_xpu
+  args : (Tensor mask)
+  output : Tensor (length), Tensor (seq_lod), Tensor (pad_seq_len)
+  infer_meta :
+    func : MaskAdaptiveXPUInferMeta
+  kernel :
+    func : mask_adaptive_xpu
+    data_type : mask
+
 # This op is implemented using CUDNN Frontend API, which serves as a supplement to
 # legacy max pooling implementation. It shows better performance with NHWC layout and
 # half precision.
@@ -484,6 +493,15 @@
     func : self_dp_attention
     data_type : x
 
+- op : sequence_unpad_xpu
+  args : (Tensor x, Tensor length)
+  output : Tensor(out)
+  infer_meta :
+    func : SequenceUnpadXPUInferMeta
+  kernel :
+    func : sequence_unpad_xpu
+    data_type : x
+
 - op : sine_pos_xpu
   args : (Tensor x, Tensor y)
   output : Tensor(out)
diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -1236,6 +1236,10 @@ XPUOpMap& get_kl2_ops() {
        XPUKernelSet({phi::DataType::FLOAT16, phi::DataType::FLOAT32})},
       {"flash_attn_unpadded",
        XPUKernelSet({phi::DataType::FLOAT16, phi::DataType::FLOAT32})},
+      {"mask_adaptive_xpu",
+       XPUKernelSet({phi::DataType::FLOAT16, phi::DataType::FLOAT32})},
+      {"sequence_unpad_xpu",
+       XPUKernelSet({phi::DataType::FLOAT16, phi::DataType::FLOAT32})},
   };
 
   return s_xpu2_kernels;
diff --git a/paddle/phi/infermeta/fusion.cc b/paddle/phi/infermeta/fusion.cc
@@ -3981,6 +3981,53 @@ void CrossAttentionXPUInferMeta(
   qkv_max->set_layout(input_q.layout());
 }
 
+void MaskAdaptiveXPUInferMeta(const MetaTensor& mask,
+                              MetaTensor* length,
+                              MetaTensor* seq_lod,
+                              MetaTensor* pad_seq_len) {
+  auto mask_dims = mask.dims();
+  auto mask_dims_size = mask_dims.size();
+  PADDLE_ENFORCE_EQ(
+      mask_dims_size,
+      3,
+      phi::errors::InvalidArgument(
+          "mask_dims_size should be 3, but received mask_dims_size is %d",
+          mask_dims_size));
+  length->set_dims({mask_dims[0]});
+  seq_lod->set_dims({mask_dims[0] + 1});
+  pad_seq_len->set_dims({1});
+  length->set_dtype(phi::DataType::INT64);
+  seq_lod->set_dtype(phi::DataType::INT32);
+  pad_seq_len->set_dtype(phi::DataType::INT32);
+}
+
+void SequenceUnpadXPUInferMeta(const MetaTensor& x,
+                               const MetaTensor& length,
+                               MetaTensor* out) {
+  auto x_dims = x.dims();
+  auto len_dims = length.dims();
+  PADDLE_ENFORCE_GE(
+      x_dims.size(),
+      2,
+      phi::errors::InvalidArgument(
+          "Rank of X can't be less than 2, but received x_dims.size() is %d",
+          x_dims.size()));
+  PADDLE_ENFORCE_EQ(
+      len_dims.size(),
+      1,
+      phi::errors::InvalidArgument(
+          "Rank of Length should be 1, but received en_dims.size() is %d",
+          len_dims.size()));
+  PADDLE_ENFORCE_EQ(x_dims[0],
+                    len_dims[0],
+                    phi::errors::InvalidArgument(
+                        "X and Length should have the same 1st dim, but "
+                        "received X.dims[0] is %d, Length.dims[0] is %d",
+                        x_dims[0],
+                        len_dims[0]));
+  out->set_dtype(x.dtype());
+}
+
 void MultiGruInferMeta(
     const MetaTensor& x,
     const std::vector<const MetaTensor*>& weight_x,
@@ -4015,7 +4062,8 @@ void MultiGruInferMeta(
         phi::errors::InvalidArgument(
             "The first dimension of flattened WeightX #%d"
             "should equal to last dimension of flattened input X, but "
-            "received fattened WeightX dimension is:%d, flattened X dimension "
+            "received fattened WeightX dimension is:%d, flattened X "
+            "dimension "
             "is:%d",
             i,
             weight_x[i]->dims()[0],
@@ -4279,15 +4327,15 @@ void RoformerRelativePosXPUInferMeta(const MetaTensor& x,
   PADDLE_ENFORCE_EQ(
       sin_emb_dims_size,
       4,
-      phi::errors::InvalidArgument(
-          "sin_emb_dims_size should be 4, but received sin_emb_dims_size is %d",
-          sin_emb_dims_size));
+      phi::errors::InvalidArgument("sin_emb_dims_size should be 4, but "
+                                   "received sin_emb_dims_size is %d",
+                                   sin_emb_dims_size));
   PADDLE_ENFORCE_EQ(
       cos_emb_dims_size,
       4,
-      phi::errors::InvalidArgument(
-          "cos_emb_dims_size should be 4, but received cos_emb_dims_size is %d",
-          cos_emb_dims_size));
+      phi::errors::InvalidArgument("cos_emb_dims_size should be 4, but "
+                                   "received cos_emb_dims_size is %d",
+                                   cos_emb_dims_size));
   for (int i = 0; i < sin_emb_dims_size; i++) {
     PADDLE_ENFORCE_EQ(
         sin_emb_dims[i],
diff --git a/paddle/phi/infermeta/fusion.h b/paddle/phi/infermeta/fusion.h
@@ -941,6 +941,15 @@ void MultiGruInferMeta(
     bool force_fp32_output,
     MetaTensor* hidden);
 
+void MaskAdaptiveXPUInferMeta(const MetaTensor& mask,
+                              MetaTensor* length,
+                              MetaTensor* seq_lod,
+                              MetaTensor* pad_seq_len);
+
+void SequenceUnpadXPUInferMeta(const MetaTensor& x,
+                               const MetaTensor& length,
+                               MetaTensor* out);
+
 void FusionLstmInferMeta(const MetaTensor& x,
                          const MetaTensor& weight_x,
                          const MetaTensor& weight_h,
diff --git a/paddle/phi/kernels/fusion/xpu/mask_adaptive_xpu_kernerl.cc b/paddle/phi/kernels/fusion/xpu/mask_adaptive_xpu_kernerl.cc
diff --git a/paddle/phi/kernels/fusion/xpu/multi_encoder_xpu_kernel.cc b/paddle/phi/kernels/fusion/xpu/multi_encoder_xpu_kernel.cc
diff --git a/paddle/phi/kernels/fusion/xpu/sequance_unpad_xpu_kernerl.cc b/paddle/phi/kernels/fusion/xpu/sequance_unpad_xpu_kernerl.cc
