[CINN]add convert to unpack mha pass (#65142)

phlrain · 6clc · Vvsmile · web-flow · commit 2d113eb7a10c · 2024-06-27T09:00:12.000+08:00
* cinn(feat): run on X86

* cinn(x86): add conditions to run x86

* fix SASS opt

* update

* add convert fa to qkvmha pass

* revert code

---------

Co-authored-by: 6clc &lt;chaoliu.lc@foxmail.com&gt;
Co-authored-by: Vvsmile &lt;17864154871@163.com&gt;
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc
@@ -31,6 +31,7 @@
 #include "paddle/cinn/hlir/dialect/operator/transforms/add_store_in_group_op_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/cinn_group_cluster_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/conv2d_transpose_filter_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/convert_fa_to_qkvmha_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/convert_memory_effec_attn_to_flash_attn_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/dynamic_reshape_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/fold_manipulation_ops_pass.h"
@@ -98,6 +99,7 @@ void ApplyPdToCinnPass(
   pass_manager->AddPass(cinn::dialect::ir::CreateRemoveAssignOutPass());
   pass_manager->AddPass(cinn::dialect::ir::CreateConv2dTransposeFilterPass());
   pass_manager->AddPass(cinn::dialect::ir::CreateConvertMEA2FAPass());
+  pass_manager->AddPass(cinn::dialect::ir::CreateConvertFA2QKVMHAPass());
   pass_manager->AddPass(cinn::dialect::ir::CreatePdOpToCinnOpPass());
 
   pass_manager->AddPass(pir::CreateDeadCodeEliminationPass());
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/convert_fa_to_qkvmha_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/convert_fa_to_qkvmha_pass.cc
@@ -0,0 +1,95 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/cinn/hlir/dialect/operator/transforms/convert_fa_to_qkvmha_pass.h"
+
+#include "paddle/cinn/hlir/dialect/operator/ir/cinn_op.h"
+#include "paddle/cinn/hlir/dialect/operator/ir/manual_op.h"
+#include "paddle/cinn/hlir/framework/pir/utils.h"
+#include "paddle/common/ddim.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/pir/include/core/builtin_dialect.h"
+#include "paddle/pir/include/pass/pass.h"
+#include "paddle/pir/include/pass/pass_registry.h"
+#include "paddle/pir/include/pattern_rewrite/frozen_rewrite_pattern_set.h"
+#include "paddle/pir/include/pattern_rewrite/pattern_applicator.h"
+#include "paddle/pir/include/pattern_rewrite/pattern_match.h"
+#include "paddle/pir/include/pattern_rewrite/pattern_rewrite_driver.h"
+
+namespace cinn {
+namespace dialect {
+namespace ir {
+
+class ConvertFA2QKVMHAPattern
+    : public pir::OpRewritePattern<paddle::dialect::FlashAttnOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::FlashAttnOp>::OpRewritePattern;
+
+  bool Match(paddle::dialect::FlashAttnOp op) const override {
+    bool is_test =
+        op->attribute("is_test").dyn_cast<pir::BoolAttribute>().data();
+    if (!is_test) {
+      return false;
+    }
+    auto& shape_analysis =
+        pir::ShapeAnalysisManager::Instance().Get(op->GetParentProgram());
+
+    auto in_shape =
+        shape_analysis.GetShapeOrDataForValue(op->operand_source(0));
+
+    auto dim1 = in_shape.shape()[1];
+
+    if (dim1.isa<int64_t>() &&
+        (dim1.dyn_cast<int64_t>() == static_cast<int64_t>(1))) {
+      return true;
+    }
+
+    return false;
+  }
+
+  void Rewrite(paddle::dialect::FlashAttnOp op,
+               pir::PatternRewriter& rewriter) const override {
+    auto q = op->operand_source(0);
+    auto k = op->operand_source(1);
+    auto v = op->operand_source(2);
+
+    pir::Value attn_mask;
+    auto fa =
+        rewriter.Build<paddle::dialect::QkvUnpackMhaOp>(q, k, v, attn_mask);
+
+    rewriter.ReplaceAllUsesWith(op->result(0), fa.result(0));
+    rewriter.EraseOp(op);
+  }
+};
+
+class ConvertFA2QKVMHAPass : public pir::PatternRewritePass {
+ public:
+  ConvertFA2QKVMHAPass() : pir::PatternRewritePass("convert_FA_to_QKVMHA", 1) {}
+
+  pir::RewritePatternSet InitializePatterns(pir::IrContext* context) override {
+    pir::RewritePatternSet ps(context);
+    ps.Add<ConvertFA2QKVMHAPattern>(context);
+    return ps;
+  }
+};
+
+std::unique_ptr<pir::Pass> CreateConvertFA2QKVMHAPass() {
+  return std::make_unique<ConvertFA2QKVMHAPass>();
+}
+
+}  // namespace ir
+}  // namespace dialect
+}  // namespace cinn
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/convert_fa_to_qkvmha_pass.h b/paddle/cinn/hlir/dialect/operator/transforms/convert_fa_to_qkvmha_pass.h
@@ -0,0 +1,28 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <memory>
+#include "paddle/pir/include/pass/pass.h"
+
+namespace cinn {
+namespace dialect {
+namespace ir {
+
+std::unique_ptr<pir::Pass> CreateConvertFA2QKVMHAPass();
+
+}  // namespace ir
+}  // namespace dialect
+}  // namespace cinn
diff --git a/paddle/phi/kernels/fusion/gpu/qkv_unpack_mha_kernel.cu b/paddle/phi/kernels/fusion/gpu/qkv_unpack_mha_kernel.cu
@@ -168,7 +168,8 @@ __global__ void qkv_attention_kernel(QkvUnpackMhaParams<T> params,
     for (int ii = 0; ii < K_VECS_PER_THREAD; ++ii) {
       if (ti < act_time_step) {
         k[ii] = *reinterpret_cast<const K_vec *>(
-            &params.k[ti * params.num_head * Dh + ki +
+            &params.k[bi * params.timestep * params.num_head * Dh +
+                      ti * params.num_head * Dh + ki +
                       ii * THREADS_PER_KEY * K_VEC_SIZE + hi * Dh]);
       }
     }
@@ -246,7 +247,8 @@ __global__ void qkv_attention_kernel(QkvUnpackMhaParams<T> params,
 
       // update here
       v = *reinterpret_cast<const V_vec *>(
-          &params.v[ti * params.num_head * Dh + vi + hi * Dh]);
+          &params.v[bi * params.timestep * params.num_head * Dh +
+                    ti * params.num_head * Dh + vi + hi * Dh]);
 
 #if defined(MMHA_USE_FP32_ACUM_FOR_LOGITS)
       float logit = logits_smem[ti];
@@ -291,7 +293,8 @@ __global__ void qkv_attention_kernel(QkvUnpackMhaParams<T> params,
 #ifdef MMHA_USE_FP32_ACUM_FOR_OUT
     V_vec tmp_out;
     convert_from_float(tmp_out, out);
-    store_func.template store<V_vec>(tmp_out, vi + hi * Dh);
+    store_func.template store<V_vec>(
+        tmp_out, bi * (params.num_head) * Dh + vi + hi * Dh);
 #else
 
     store_func.template store<V_vec>(out, vi + hi * Dh);
