PaddlePaddle
diff --git a/‎cmake/external/xpu.cmake
Copy file name to clipboard
+9-2Lines changed: 9 additions & 2 deletions b/‎cmake/external/xpu.cmake
Copy file name to clipboard
+9-2Lines changed: 9 additions & 2 deletions
diff --git a/‎paddle/fluid/framework/ir/xpu/decoder_attention_xpu_fuse_pass.cc
Copy file name to clipboard
+53-4Lines changed: 53 additions & 4 deletions b/‎paddle/fluid/framework/ir/xpu/decoder_attention_xpu_fuse_pass.cc
Copy file name to clipboard
+53-4Lines changed: 53 additions & 4 deletions
diff --git a/‎paddle/fluid/framework/ir/xpu/fc_xpu_fuse_pass.cc
Copy file name to clipboard
+29Lines changed: 29 additions & 0 deletions b/‎paddle/fluid/framework/ir/xpu/fc_xpu_fuse_pass.cc
Copy file name to clipboard
+29Lines changed: 29 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/ir/xpu/qk_qkv_attention_xpu_fuse_pass.cc
Copy file name to clipboard
+56-4Lines changed: 56 additions & 4 deletions b/‎paddle/fluid/framework/ir/xpu/qk_qkv_attention_xpu_fuse_pass.cc
Copy file name to clipboard
+56-4Lines changed: 56 additions & 4 deletions
diff --git a/‎paddle/fluid/framework/ir/xpu/xpu_quantize_op_pass.cc
Copy file name to clipboard
+72-3Lines changed: 72 additions & 3 deletions b/‎paddle/fluid/framework/ir/xpu/xpu_quantize_op_pass.cc
Copy file name to clipboard
+72-3Lines changed: 72 additions & 3 deletions
diff --git a/‎paddle/fluid/framework/ir/xpu/xpu_quantize_op_pass.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/framework/ir/xpu/xpu_quantize_op_pass.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/phi/api/yaml/fused_ops.yaml
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions b/‎paddle/phi/api/yaml/fused_ops.yaml
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions
@@ -28,6 +28,9 @@ set(XPU_XFA_LIB_NAME "libxpu_flash_attention.so")
 if(NOT DEFINED XPU_BASE_DATE)
   set(XPU_BASE_DATE "20240104")
 endif()
+if(NOT DEFINED XPU_XDNN_BASE_DATE)
+  set(XPU_XDNN_BASE_DATE "20240327")
+endif()
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
   set(XPU_XHPC_BASE_DATE "20240328")
 endif()
@@ -45,6 +48,10 @@ else()
   set(XPU_BASE_URL "${XPU_BASE_URL}")
 endif()
 
+set(XPU_XDNN_BASE_URL
+    "https://klx-sdk-release-public.su.bcebos.com/xdnn/stable/${XPU_XDNN_BASE_DATE}"
+)
+
 set(XPU_XCCL_BASE_URL
     "https://klx-sdk-release-public.su.bcebos.com/xccl/release/${XPU_XCCL_BASE_VERSION}"
 )
@@ -105,7 +112,7 @@ set(XPU_XRE_URL
     "${XPU_BASE_URL}/${XPU_XRE_DIR_NAME}.tar.gz"
     CACHE STRING "" FORCE)
 set(XPU_XDNN_URL
-    "${XPU_BASE_URL}/${XPU_XDNN_DIR_NAME}.tar.gz"
+    "${XPU_XDNN_BASE_URL}/${XPU_XDNN_DIR_NAME}.tar.gz"
     CACHE STRING "" FORCE)
 set(XPU_XCCL_URL
     "${XPU_XCCL_BASE_URL}/${XPU_XCCL_DIR_NAME}.tar.gz"
@@ -229,7 +236,7 @@ if(WITH_XPTI)
 endif()
 
 if(WITH_XPU_XHPC)
-  target_link_libraries(xpulib ${XPU_API_LIB} ${XPU_RT_LIB} ${XPU_XBLAS_LIB}
+  target_link_libraries(xpulib ${XPU_RT_LIB} ${XPU_XBLAS_LIB} ${XPU_API_LIB}
                         ${XPU_XFA_LIB})
 endif()
 
 
@@ -17,6 +17,7 @@
 #include "glog/logging.h"
 
 #include "paddle/fluid/framework/ir/graph_pattern_detector.h"
+#include "paddle/fluid/framework/ir/quantize_helper.h"
 #include "paddle/fluid/framework/ir/xpu/pass_utils.h"
 #include "paddle/fluid/framework/op_version_registry.h"
 #include "paddle/fluid/platform/enforce.h"
@@ -222,6 +223,7 @@ void DecoderAttentionXPUFusePass::ApplyDecoderAttentionXPUFuse(
     GET_IR_NODE(output);
 
     // Generate fuse op
+    auto* scope = param_scope();
     auto* block = reshape2_1->Op()->Block();
     framework::OpDesc fused_op_desc(block);
     fused_op_desc.SetType("qkv_attention_xpu");
@@ -230,6 +232,54 @@ void DecoderAttentionXPUFusePass::ApplyDecoderAttentionXPUFuse(
     fused_op_desc.SetInput("q", {input_q->Name()});
     fused_op_desc.SetInput("k", {input_k->Name()});
     fused_op_desc.SetInput("v", {input_v->Name()});
+    std::unordered_map<std::string, std::vector<float>> var_quant_scales =
+        GetQuantInfoFromTheGraph(graph, "has_quant_info", "var_quant_scales");
+    // recored q/k/v max, qk_max, and qkv_max
+    std::vector<Node*> input_max_nodes;
+    if (var_quant_scales.find(input_q->Name()) != var_quant_scales.end() &&
+        var_quant_scales.find(input_k->Name()) != var_quant_scales.end() &&
+        var_quant_scales.find(input_v->Name()) != var_quant_scales.end() &&
+        var_quant_scales.find(qk_matmul_out->Name()) !=
+            var_quant_scales.end() &&
+        var_quant_scales.find(qkv_matmul_out->Name()) !=
+            var_quant_scales.end()) {
+      std::vector<float> input_max_vec;
+      input_max_vec.push_back(var_quant_scales.at(input_q->Name())[0]);
+      input_max_vec.push_back(var_quant_scales.at(input_k->Name())[0]);
+      input_max_vec.push_back(var_quant_scales.at(input_v->Name())[0]);
+      input_max_vec.push_back(var_quant_scales.at(qk_matmul_out->Name())[0]);
+      input_max_vec.push_back(var_quant_scales.at(qkv_matmul_out->Name())[0]);
+      std::vector<std::string> quant_max_names = {
+          "q_max", "k_max", "v_max", "qk_max", "qkv_max"};
+      for (size_t i = 0; i < input_max_vec.size(); i++) {
+        std::string input_max_name =
+            input_q->Name() + "_" + std::to_string(i) + "_max_in";
+        int max_ptr_size = phi::backends::xpu::get_xpu_max_ptr_size(-1);
+        VarDesc input_max_desc(input_max_name);
+        input_max_desc.SetPersistable(true);
+        input_max_desc.SetShape({static_cast<int64_t>(max_ptr_size)});
+        input_max_desc.SetDataType(proto::VarType::Type::VarType_Type_FP32);
+        Node* input_max_in = graph->CreateVarNode(&input_max_desc);
+        auto* block_input_max_in_desc = block->Var(input_max_name);
+        block_input_max_in_desc->SetPersistable(input_max_desc.Persistable());
+        block_input_max_in_desc->SetShape(input_max_desc.GetShape());
+        block_input_max_in_desc->SetDataType(input_max_desc.GetDataType());
+        phi::DenseTensor input_max_in_cpu_tensor;
+        auto* cpu_ctx = static_cast<phi::CPUContext*>(
+            platform::DeviceContextPool::Instance().Get(phi::CPUPlace()));
+        input_max_in_cpu_tensor.set_type(phi::DataType::FLOAT32);
+        input_max_in_cpu_tensor.Resize({max_ptr_size});
+        std::vector<float> input_max(max_ptr_size, input_max_vec[i]);
+        memcpy(cpu_ctx->Alloc<float>(&input_max_in_cpu_tensor),
+               input_max.data(),
+               max_ptr_size * sizeof(float));
+        Assign(input_max_in_cpu_tensor,
+               scope->Var(input_max_name)->GetMutable<phi::DenseTensor>());
+        fused_op_desc.SetInput(quant_max_names[i], {input_max_name});
+
+        input_max_nodes.push_back(input_max_in);
+      }
+    }
 
     // set attributes of fuse_op
     float scale_val = PADDLE_GET_CONST(float, scale->Op()->GetAttr("scale"));
@@ -245,9 +295,6 @@ void DecoderAttentionXPUFusePass::ApplyDecoderAttentionXPUFuse(
     fused_op_desc.SetAttr("out_dtype", input_q->Var()->GetDataType());
 
     // set output of fuse_op
-    VarDesc fused_op_out_max_desc("qkv_max");
-    Node* fused_op_out_max = graph->CreateVarNode(&fused_op_out_max_desc);
-    fused_op_desc.SetOutput("qkv_max", {"qkv_max"});
     fused_op_desc.SetOutput("qkv", {output->Name()});
 
     auto* fused_op = graph->CreateOpNode(&fused_op_desc);
@@ -256,7 +303,9 @@ void DecoderAttentionXPUFusePass::ApplyDecoderAttentionXPUFuse(
     IR_NODE_LINK_TO(input_k, fused_op);
     IR_NODE_LINK_TO(input_v, fused_op);
     IR_NODE_LINK_TO(fused_op, output);
-    IR_NODE_LINK_TO(fused_op, fused_op_out_max);
+    for (size_t i = 0; i < input_max_nodes.size(); i++) {
+      IR_NODE_LINK_TO(input_max_nodes[i], fused_op);
+    }
 
     // delete useless node
     std::unordered_set<const Node*> del_node_set;
 
@@ -841,6 +841,35 @@ int FcXPUFusePass::ApplyImpl(ir::Graph* graph,
     } else if (filter_data_type == phi::DataType::FLOAT16) {
       op_weights_precision = "float16";
     }
+    if (op_weights_precision == "float32" &&
+        AreScalesPresentForNodes(&var_quant_scales, {mul_w})) {
+      // convert weight to int8
+      auto* var = scope->FindVar(mul_w_name);
+      PADDLE_ENFORCE_NOT_NULL(
+          var,
+          platform::errors::NotFound(
+              "The input persistable [%s] var of [%s] op is not found.",
+              mul_w_name));
+      auto* weight_tensor = var->GetMutable<phi::DenseTensor>();
+      float* fp32_weight_data = weight_tensor->data<float>();
+      std::vector<int8_t> weight_data;
+      weight_data.resize(weight_tensor->numel());
+      for (int i = 0; i < weight_tensor->numel(); i++) {
+        weight_data[i] = static_cast<int8_t>(fp32_weight_data[i]);
+      }
+      const auto weight_dims = weight_tensor->dims();
+      weight_tensor->clear();  // clear int weight
+      weight_tensor->set_type(phi::DataType::INT8);
+      weight_tensor->Resize(common::make_ddim(common::vectorize(weight_dims)));
+      auto* cpu_ctx = static_cast<phi::CPUContext*>(
+          platform::DeviceContextPool::Instance().Get(phi::CPUPlace()));
+      auto* new_weight_data = cpu_ctx->Alloc<int8_t>(weight_tensor);
+      memcpy(new_weight_data,
+             weight_data.data(),
+             weight_tensor->numel() * sizeof(int8_t));
+      op_weights_precision = "int8";
+    }
+
     VLOG(4) << "FC fusion fuse pass is running on " << op_weights_precision
             << " precision!";
     auto* block = mul->Op()->Block();
 
@@ -17,6 +17,7 @@
 #include "glog/logging.h"
 
 #include "paddle/fluid/framework/ir/graph_pattern_detector.h"
+#include "paddle/fluid/framework/ir/quantize_helper.h"
 #include "paddle/fluid/framework/ir/xpu/pass_utils.h"
 #include "paddle/fluid/framework/op_version_registry.h"
 #include "paddle/fluid/platform/enforce.h"
@@ -212,13 +213,65 @@ void QkQkvAttentionXPUFusePass::ApplyQkQkvAttentionXPUFuse(
     GET_IR_NODE(output);
 
     // Generate fuse op
+    auto* scope = param_scope();
     auto* block = reshape_1->Op()->Block();
     framework::OpDesc fused_op_desc(block);
     fused_op_desc.SetType("qkv_attention_xpu");
     // set input of fuse_op
     fused_op_desc.SetInput("q", {input->Name()});
     fused_op_desc.SetInput("k", {input->Name()});
     fused_op_desc.SetInput("v", {input->Name()});
+    std::unordered_map<std::string, std::vector<float>> var_quant_scales =
+        GetQuantInfoFromTheGraph(graph, "has_quant_info", "var_quant_scales");
+    // recored q/k/v max, qk_max, and qkv_max
+    std::vector<Node*> input_max_nodes;
+    if (var_quant_scales.find(input->Name()) != var_quant_scales.end() &&
+        var_quant_scales.find(qk_matmul_out->Name()) !=
+            var_quant_scales.end() &&
+        var_quant_scales.find(qkv_matmul_out->Name()) !=
+            var_quant_scales.end()) {
+      std::vector<float> input_max_vec;
+      input_max_vec.push_back(var_quant_scales.at(input->Name())[0]);
+      input_max_vec.push_back(var_quant_scales.at(qk_matmul_out->Name())[0]);
+      input_max_vec.push_back(var_quant_scales.at(qkv_matmul_out->Name())[0]);
+      std::vector<std::string> quant_max_names = {
+          "input_max", "qk_max", "qkv_max"};
+      for (size_t i = 0; i < input_max_vec.size(); i++) {
+        std::string input_max_name =
+            input->Name() + "_" + std::to_string(i) + "_max_in";
+        int max_ptr_size = phi::backends::xpu::get_xpu_max_ptr_size(-1);
+        VarDesc input_max_desc(input_max_name);
+        input_max_desc.SetPersistable(true);
+        input_max_desc.SetShape({static_cast<int64_t>(max_ptr_size)});
+        input_max_desc.SetDataType(proto::VarType::Type::VarType_Type_FP32);
+        Node* input_max_in = graph->CreateVarNode(&input_max_desc);
+        auto* block_input_max_in_desc = block->Var(input_max_name);
+        block_input_max_in_desc->SetPersistable(input_max_desc.Persistable());
+        block_input_max_in_desc->SetShape(input_max_desc.GetShape());
+        block_input_max_in_desc->SetDataType(input_max_desc.GetDataType());
+
+        phi::DenseTensor input_max_in_cpu_tensor;
+        auto* cpu_ctx = static_cast<phi::CPUContext*>(
+            platform::DeviceContextPool::Instance().Get(phi::CPUPlace()));
+        input_max_in_cpu_tensor.set_type(phi::DataType::FLOAT32);
+        input_max_in_cpu_tensor.Resize({max_ptr_size});
+        std::vector<float> input_max(max_ptr_size, input_max_vec[i]);
+        memcpy(cpu_ctx->Alloc<float>(&input_max_in_cpu_tensor),
+               input_max.data(),
+               max_ptr_size * sizeof(float));
+        Assign(input_max_in_cpu_tensor,
+               scope->Var(input_max_name)->GetMutable<phi::DenseTensor>());
+        if (i == 0) {
+          fused_op_desc.SetInput("q_max", {input_max_name});
+          fused_op_desc.SetInput("k_max", {input_max_name});
+          fused_op_desc.SetInput("v_max", {input_max_name});
+        } else {
+          fused_op_desc.SetInput(quant_max_names[i], {input_max_name});
+        }
+        input_max_nodes.push_back(input_max_in);
+      }
+    }
+
     // set attributes of fuse_op
     if (with_q_scale) {
       float scale_val = PADDLE_GET_CONST(float, scale->Op()->GetAttr("scale"));
@@ -239,16 +292,15 @@ void QkQkvAttentionXPUFusePass::ApplyQkQkvAttentionXPUFuse(
     fused_op_desc.SetAttr("out_dtype", input->Var()->GetDataType());
 
     // set output of fuse_op
-    VarDesc fused_op_out_max_desc("qkv_max");
-    Node* fused_op_out_max = graph->CreateVarNode(&fused_op_out_max_desc);
-    fused_op_desc.SetOutput("qkv_max", {"qkv_max"});
     fused_op_desc.SetOutput("qkv", {output->Name()});
 
     auto* fused_op = graph->CreateOpNode(&fused_op_desc);
 
     IR_NODE_LINK_TO(input, fused_op);
     IR_NODE_LINK_TO(fused_op, output);
-    IR_NODE_LINK_TO(fused_op, fused_op_out_max);
+    for (size_t i = 0; i < input_max_nodes.size(); i++) {
+      IR_NODE_LINK_TO(input_max_nodes[i], fused_op);
+    }
 
     // delete useless node
     std::unordered_set<const Node*> del_node_set;
 
@@ -164,7 +164,9 @@ void XPUQuantizeOpPass::QuantizeConv(ir::Graph* graph) const {
           out_var_node = output_node;
         }
       }
-      if (!AreScalesPresentForNodes(&var_quant_scales_, {x_var_node})) {
+      if (!AreScalesPresentForNodes(&var_quant_scales_, {x_var_node}) ||
+          w_var_node->Var()->GetDataType() !=
+              proto::VarType::Type::VarType_Type_INT8) {
         VLOG(4) << "Skip quantize op: " << n->Name()
                 << "x_var_node_name:" << x_var_node->Name()
                 << " w_var_node_name:" << w_var_node->Name();
@@ -239,8 +241,9 @@ void XPUQuantizeOpPass::QuantizeFC(ir::Graph* graph) const {
           out_var_node = output_node;
         }
       }
-      if (!AreScalesPresentForNodes(&var_quant_scales_,
-                                    {x_var_node, w_var_node})) {
+      if (!AreScalesPresentForNodes(&var_quant_scales_, {x_var_node}) ||
+          w_var_node->Var()->GetDataType() !=
+              proto::VarType::Type::VarType_Type_INT8) {
         MarkAndLogCannotQuantizeOp(n, "No scale available for the operator");
         continue;
       }
@@ -261,6 +264,71 @@ void XPUQuantizeOpPass::QuantizeFC(ir::Graph* graph) const {
   }
 }
 
+void XPUQuantizeOpPass::QuantizeQkvAttention(ir::Graph* graph) const {
+  for (auto* n : graph->Nodes()) {
+    if (n->IsOp()) {
+      auto* op = n->Op();
+      if (op->Type() != "qkv_attention_xpu") {
+        continue;
+      }
+      std::vector<std::string> max_node_names = {
+          "q_max", "k_max", "v_max", "qk_max"};
+      std::unordered_map<std::string, Node*> input_node_map;
+      for (auto* input_node : n->inputs) {
+        if (!input_node->IsVar()) {
+          continue;
+        }
+        for (auto input_name : op->InputNames()) {
+          if (op->Input(input_name)[0] == input_node->Var()->Name()) {
+            input_node_map[input_name] = input_node;
+          }
+        }
+      }
+      bool continue_flag = false;
+      for (auto max_name : max_node_names) {
+        if (input_node_map.find(max_name) == input_node_map.end()) {
+          continue_flag = true;
+          break;
+        }
+      }
+      if (continue_flag) {
+        continue;
+      }
+      Node* out_var_node = nullptr;
+      for (auto* output_node : n->outputs) {
+        if (!output_node->IsVar()) {
+          continue;
+        }
+        if (output_node->Var()->Name() == op->Output("qkv")[0]) {
+          out_var_node = output_node;
+        }
+      }
+      if (input_node_map["q"]->Name() == input_node_map["k"]->Name() &&
+          input_node_map["q"]->Name() == input_node_map["v"]->Name()) {
+        QuantizeInput(graph, n, input_node_map["q"], "q");
+        op->SetInput("k", op->Input("q"));
+        op->SetInput("v", op->Input("q"));
+        UnlinkNodes(input_node_map["k"], n);
+        UnlinkNodes(input_node_map["v"], n);
+      } else {
+        QuantizeInput(graph, n, input_node_map["q"], "q");
+        QuantizeInput(graph, n, input_node_map["k"], "k");
+        QuantizeInput(graph, n, input_node_map["v"], "v");
+      }
+      auto has_output_scale =
+          AreScalesPresentForNodes(&var_quant_scales_, {out_var_node});
+      if (has_output_scale) {
+        DequantizeOutput(graph, n, out_var_node, "qkv");
+        n->Op()->SetAttr(
+            "out_dtype",
+            static_cast<int>(proto::VarType::Type::VarType_Type_INT8));
+      } else {
+        n->Op()->SetAttr("out_dtype",
+                         input_node_map["q"]->Var()->GetDataType());
+      }
+    }
+  }
+}
 void XPUQuantizeOpPass::ApplyImpl(ir::Graph* graph) const {
   VLOG(3) << "Insert quantize/dequantize op to the graph.";
   PADDLE_ENFORCE_NOT_NULL(
@@ -273,6 +341,7 @@ void XPUQuantizeOpPass::ApplyImpl(ir::Graph* graph) const {
   GetQuantInfo(graph);
   QuantizeConv(graph);
   QuantizeFC(graph);
+  QuantizeQkvAttention(graph);
 }
 
 }  // namespace ir
 
@@ -38,6 +38,7 @@ class XPUQuantizeOpPass : public FusePassBase {
  protected:
   void ApplyImpl(Graph* graph) const override;
   void QuantizeConv(Graph* graph) const;
+  void QuantizeQkvAttention(Graph* graph) const;
   void QuantizeFC(Graph* graph) const;
 
  private:
 
@@ -430,14 +430,14 @@
   optional : bias_qk
 
 - op : qkv_attention_xpu
-  args : (Tensor q, Tensor k, Tensor v, Tensor q_max, Tensor k_max, Tensor v_max, float alpha, int head_num, int head_dim, bool qkv_fc_fusion, DataType out_dtype)
-  output : Tensor(qkv), Tensor(qkv_max)
+  args : (Tensor q, Tensor k, Tensor v, Tensor q_max, Tensor k_max, Tensor v_max, Tensor qk_max, Tensor qkv_max, float alpha, int head_num, int head_dim, bool qkv_fc_fusion, DataType out_dtype)
+  output : Tensor(qkv)
   infer_meta :
     func : QKVAttentionXPUInferMeta
   kernel :
     func : qkv_attention_xpu
     data_type : q
-  optional : q_max, k_max, v_max
+  optional : q_max, k_max, v_max, qk_max, qkv_max
 
 - op : quantize_xpu
   args : (Tensor x, DataType out_dtype, float scale = 1.0f)