mindspore 2.0.0rc1__cp38-cp38-manylinux1_x86_64.whl → 2.2.0__cp38-cp38-manylinux1_x86_64.whl

mindspore/ops/_op_impl/_custom_op/flash_attention/attention.py

@@ -0,0 +1,406 @@
+# Copyright 2023 Huawei Technologies Co., Ltd
+#
+# 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
+#
+# http://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.
+# ============================================================================
+"""the base class of flash attention"""
+from abc import ABCMeta
+from abc import abstractmethod
+from functools import partial
+from collections import defaultdict
+
+import te.platform as tbe_platform
+from tbe import tik
+from tbe.common.platform import get_soc_spec
+
+from mindspore.ops._op_impl._custom_op.flash_attention.constants import FP16
+from mindspore.ops._op_impl._custom_op.flash_attention.constants import FP32
+from mindspore.ops._op_impl._custom_op.flash_attention.constants import GM
+from mindspore.ops._op_impl._custom_op.flash_attention.constants import MASK_FILL_VALUE
+from mindspore.ops._op_impl._custom_op.flash_attention.constants import UB
+from mindspore.ops._op_impl._custom_op.flash_attention.tik_ops_utils import TikOpsUtils
+from mindspore.ops._op_impl._custom_op.flash_attention.tiling_strategy.strategy import TilingPara
+from mindspore.ops._op_impl._custom_op.flash_attention.tiling_strategy.strategy import TilingStrategy
+from mindspore.ops._op_impl._custom_op.flash_attention.tiling_strategy.sparse_tiling import SparseTiling
+
+
+class FlashAttention(metaclass=ABCMeta):
+    """The base class of FlashAttention"""
+
+    def __init__(self, q, k, v, attn_mask, dropout_mask, alibi_mask, kernel_name,
+                 tiling_stgy_cls,
+                 prev_block_num=65536,
+                 next_block_num=65536,
+                 high_precision=False,
+                 disable_debug=True):
+        """
+        Init parameter shape
+        :param q: with shape: (B, h, N, d)
+        :param k: with shape: (B, h, N, d)
+        :param v: with shape: (B, h, N, d)
+        :param attn_mask: with shape: (1, N, N) or (B, N, N)
+        :param dropout_mask: with shape: (B, h, N, N)
+        :param alibi_mask: with shape: (B, h, 1, N)
+        :param kernel_name:
+        :param tiling_stgy_cls:
+        :param prev_block_num:
+        :param next_block_num:
+        :param disable_debug:
+        """
+        self.tik_instance = tik.Tik(disable_debug=disable_debug)
+        self.core_num = get_soc_spec(tbe_platform.CORE_NUM)
+        self.M = tbe_platform.get_soc_spec(tbe_platform.L1_SIZE)
+        self.kernel_name = kernel_name
+        self.cont_data_mv_1_bust = partial(self.tik_instance.data_move, sid=0, nburst=1,
+                                           src_stride=0,
+                                           dst_stride=0)
+        self.tik_ops_utils = TikOpsUtils(self.tik_instance)
+        self.parse_input_shape(alibi_mask, attn_mask, dropout_mask, k, q, v)
+        # NZ
+        _, _, N1, M1, M0, N0 = self.q_shape
+        self.M1 = M1
+        self.N1 = N1
+        self.M0 = M0
+        self.N0 = N0
+        self.d = N1 * N0
+        # ND
+        batch_size, h, Nq, actual_d = self.q_ori_shape
+        self.head_num = h
+        self.B, self.Nq = batch_size * h, Nq
+        self.N = self.k_ori_shape[2]
+        self.actual_d = actual_d
+
+        self.l_shape = [batch_size, h, self.Nq]
+        self.m_shape = [batch_size, h, self.Nq]
+        self.O_shape = self.q_shape
+
+        self.prev_block_num = prev_block_num
+        self.next_block_num = next_block_num
+        self.high_precision = high_precision
+        if self.high_precision:
+            self.precision_type = FP32
+        else:
+            self.precision_type = FP16
+        if tiling_stgy_cls is None:
+            self.tiling_stgy = SparseTiling(self.Nq, self.N, self.d)
+        else:
+            self.tiling_stgy: TilingStrategy = tiling_stgy_cls(self.Nq, self.N, self.d)
+        self.Br = None
+        self.last_Br = None
+        self.Bc = None
+        self.last_Bc = None
+        self.Tr = None
+        self.Tc = None
+        self.Q_gm = None
+        self.K_gm = None
+        self.V_gm = None
+        self.dim_mask_gm = None
+        self.att_mask_gm = None
+        self.drop_mask_gm = None
+        self.alibi_mask_gm = None
+
+    @staticmethod
+    def get_l_m_gm_offset(batch_start, batch_idx, h, block_h, block_idx):
+        """get l m gm offset"""
+        gm_offset = (batch_start + batch_idx) * h + block_idx * block_h
+        return gm_offset
+
+    @staticmethod
+    def get_alibi_gm_offset(batch_start, batch_idx, w, block_w, block_idx):
+        """get alibi gm offset"""
+        gm_offset = (batch_start + batch_idx) * w + block_idx * block_w
+        return gm_offset
+
+    @staticmethod
+    def get_drop_mask_gm_offset(batch_start, batch_idx, h, w, block_h, block_h_idx, block_w, block_w_idx):
+        """get drop mask gm offset"""
+        gm_offset = (batch_start + batch_idx) * h * w + block_h_idx * (w * block_h) + block_w_idx * block_w
+        return gm_offset
+
+    @abstractmethod
+    def define_custom_inputs(self):
+        """define custom inputs"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def define_outputs(self):
+        """define outputs"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def collect_inputs(self):
+        """collect inputs"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def collect_outputs(self):
+        """collect outputs"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def compute_one_core(self, batch_start_s, batch_num_s, core_idx_to_tr_info, core_idx):
+        """compute one core"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def prepare_global_ones(self):
+        """prepare global ones"""
+        raise NotImplementedError
+
+    def get_gm_offset(self, batch_start, batch_idx, h, w, block_h, block_idx):
+        """get gm offset"""
+        gm_offset = (batch_start + batch_idx) * h * w + block_idx * block_h * self.N0
+        return gm_offset
+
+    def get_cur_tr_block_num(self, tr_idx):
+        """get cur tr block_num"""
+        cur_prev_block_num = min(tr_idx, self.prev_block_num)
+        cur_next_block_num = min(self.next_block_num, self.Tc - tr_idx - 1)
+        block_num = cur_prev_block_num + 1 + cur_next_block_num
+        return block_num
+
+    def get_total_block_num(self):
+        """get total block num"""
+        block_num = 0
+        for b_idx in range(self.B):
+            for tr_idx in range(self.Tr):
+                block_num += self.get_cur_tr_block_num(tr_idx)
+        return block_num
+
+    def update_core_task_map(self,
+                             core_b_map,
+                             core_b_tr_map,
+                             core_idx,
+                             b_start,
+                             b_end,
+                             tr_start,
+                             tr_end):
+        """update core task map"""
+        core_b_map[core_idx][0] = min(core_b_map[core_idx][0], b_start)
+        if tr_end == 0:  # 跨head，但跨过的head不会被当前的core处理
+            core_b_map[core_idx][1] = max(core_b_map[core_idx][1], b_end - 1)
+        else:
+            core_b_map[core_idx][1] = max(core_b_map[core_idx][1], b_end)
+        for b_idx in range(b_start, b_end + 1):
+            if b_idx == b_end and tr_end == 0:  # 跨head，但跨过的head不会被当前的core处理
+                break
+            elif b_idx == b_start and b_idx == b_end:  # 没跨head
+                core_b_tr_map[core_idx][b_idx] = (tr_start, tr_end)
+            elif b_idx == b_start:  # 跨head，第一个head
+                core_b_tr_map[core_idx][b_idx] = (tr_start, self.Tr)
+            elif b_idx == b_end:  # 跨head，最后一个head
+                core_b_tr_map[core_idx][b_idx] = (0, tr_end)
+            else:  # 跨head，中间的head
+                core_b_tr_map[core_idx][b_idx] = (0, self.Tr)
+
+    def convert_py_dict_to_tik_tensor(self, core_b_map, core_b_tr_map):
+        """convert py dict to tik tensor"""
+        # python dict -> tik tensor
+        # [batch_start, batch_idx_end] -> [batch_start, batch_num]
+        # [tr_start, tr_idx_end] -> [tr_start, tr_idx_end)
+        core_idx_to_batch_info = self.tik_instance.Tensor(
+            "int32", (self.core_num, 2), name="core_idx_to_batch_info", scope=UB
+        )
+        core_idx_to_tr_info = self.tik_instance.Tensor(
+            "int32", (self.core_num, self.B, 2), name="core_idx_to_tr_info", scope=UB
+        )
+        for core_idx in core_b_map.keys():
+            batch_start, batch_end = core_b_map[core_idx]
+            core_idx_to_batch_info[core_idx, 0] = batch_start
+            core_idx_to_batch_info[core_idx, 1] = batch_end - batch_start + 1
+            for batch_idx in core_b_tr_map[core_idx].keys():
+                tr_start, tr_end = core_b_tr_map[core_idx][batch_idx]
+                core_idx_to_tr_info[core_idx, batch_idx, 0] = tr_start
+                core_idx_to_tr_info[core_idx, batch_idx, 1] = tr_end
+
+        return core_idx_to_batch_info, core_idx_to_tr_info
+
+    def get_core_task_info(self):
+        """
+        Get batch start and batch number of each NPU core.
+        :return: Tensor([[core_1_batch_start, core_1_batch_num],...,[core_m_batch_start,
+        core_m_batch_num]]), Tensor([[[core_1_batch_1_Tr_start, core_1_batch_1_Tr_end],...[core_1_batch_n_Tr_start,
+        core_1_batch_n_Tr_end]],...,[[core_m_batch_1_Tr_start, core_m_batch_1_Tr_end],...[core_m_batch_n_Tr_start,
+        core_m_batch_n_Tr_end]]
+        """
+        if self.core_num > self.B * self.Tr:
+            self.core_num = self.B * self.Tr
+
+        total_blk_num = self.get_total_block_num()
+        b_start = 0
+        tr_start = 0
+        remain_blk_num = total_blk_num
+        core_b_map = defaultdict(lambda: [100000, -1])
+        core_b_tr_map = defaultdict(lambda: defaultdict(list))
+        for core_idx in range(self.core_num):
+            cur_core_blk_num = 0
+            cur_each_core_blk_num = remain_blk_num // (self.core_num - core_idx)
+            cur_core_finished = False
+            b_end = b_start
+            tr_end = tr_start
+            while b_end < self.B:
+                while tr_end < self.Tr:
+                    cur_tr_blk_num = self.get_cur_tr_block_num(tr_end)
+                    if abs(cur_core_blk_num - cur_each_core_blk_num) <= \
+                            (cur_core_blk_num + cur_tr_blk_num - cur_each_core_blk_num):
+                        self.update_core_task_map(core_b_map, core_b_tr_map, core_idx, b_start, b_end, tr_start, tr_end)
+                        remain_blk_num -= cur_core_blk_num
+                        cur_core_finished = True
+                        break
+                    else:
+                        cur_core_blk_num += cur_tr_blk_num
+                        tr_end += 1
+                        if tr_end == self.Tr:
+                            tr_end = 0
+                            b_end += 1
+                if cur_core_finished:
+                    b_start = b_end
+                    tr_start = tr_end
+                    break
+        core_idx_to_batch_info, core_idx_to_tr_info = self.convert_py_dict_to_tik_tensor(core_b_map, core_b_tr_map)
+        return core_idx_to_batch_info, core_idx_to_tr_info
+
+    def get_attn_mask_gm_offset(self, batch_start, batch_idx, h, w, block_h, block_h_idx, block_w, block_w_idx):
+        """get attn mask gm offset"""
+        if self.att_mask_shape[0] == 1:
+            gm_offset = block_w_idx * (h * block_w) + block_h_idx * block_h * self.N0
+        else:
+            gm_offset = ((batch_start + batch_idx) // self.head_num) * h * w \
+                        + block_w_idx * (h * block_w) + block_h_idx * block_h * self.N0
+        return gm_offset
+
+    def parse_input_shape(self, alibi_mask, attn_mask, dropout_mask, k, q, v):
+        """parser input shape"""
+        self.has_attn_mask = False
+        self.has_drop_mask = False
+        self.has_alibi_mask = False
+        # NZ
+        self.q_shape = q["shape"]
+        self.k_shape = k["shape"]
+        self.v_shape = v["shape"]
+        # ND
+        self.q_ori_shape = q["ori_shape"]
+        self.k_ori_shape = k["ori_shape"]
+        if attn_mask is not None:
+            self.has_attn_mask = True
+            self.att_mask_shape = attn_mask["shape"]
+        if dropout_mask is not None:
+            self.has_drop_mask = True
+            self.drop_mask_shape = dropout_mask["shape"]
+        if alibi_mask is not None:
+            self.has_alibi_mask = True
+            self.alibi_mask_shape = alibi_mask["shape"]
+
+    def define_inputs_outputs(self):
+        """define inputs outputs"""
+        self.define_common_inputs()
+
+        self.define_custom_inputs()
+
+        self.define_outputs()
+
+    def init(self):
+        """init parameters"""
+        tiling_para: TilingPara = self.tiling_stgy.tiling()
+
+        self.Br = tiling_para.Br
+        self.last_Br = tiling_para.last_Br
+        self.Bc = tiling_para.Bc
+        self.last_Bc = tiling_para.last_Bc
+        self.Tr = tiling_para.Tr
+        self.Tc = tiling_para.Tc
+
+        self.define_inputs_outputs()
+
+    def define_common_inputs(self):
+        """define common input gm tensors"""
+        self.Q_gm = self.tik_instance.Tensor(FP16, self.q_shape, name="Q_gm", scope=GM)
+        self.K_gm = self.tik_instance.Tensor(FP16, self.k_shape, name="K_gm", scope=GM)
+        self.V_gm = self.tik_instance.Tensor(FP16, self.v_shape, name="V_gm", scope=GM)
+        if self.has_attn_mask:
+            self.att_mask_gm = self.tik_instance.Tensor(FP16, self.att_mask_shape,
+                                                        name="att_mask_gm", scope=GM)
+        if self.has_drop_mask:
+            self.drop_mask_gm = self.tik_instance.Tensor(FP16, self.drop_mask_shape,
+                                                         name="drop_mask_gm", scope=GM)
+        if self.has_alibi_mask:
+            self.alibi_mask_gm = self.tik_instance.Tensor(FP16, self.alibi_mask_shape,
+                                                          name="alibi_mask_gm", scope=GM)
+
+    def do_alibi_mask(self, Sij_ub, alibi_mask_gm_offset, m_aligned, n_aligned):
+        """load alibi mask from gm to ub, then add Sij"""
+        with self.tik_instance.new_stmt_scope(disable_sync=False):
+            alibi_mask_ub = self.tik_instance.Tensor(FP16, (1, n_aligned),
+                                                     scope=UB, name="alibi_mask_ub")
+            self.tik_instance.data_move(alibi_mask_ub, self.alibi_mask_gm[alibi_mask_gm_offset], 0, 1,
+                                        n_aligned // 16, 0, 0)
+            alibi_mask_ub_broadcast = self.tik_ops_utils.broadcast_row(alibi_mask_ub, (m_aligned, n_aligned))
+            self.tik_instance.h_add(Sij_ub, Sij_ub, alibi_mask_ub_broadcast)
+
+    def do_att_mask(self, Sij_ub_N1MN0, attn_mask_gm_offset, q_blk_height, kv_blk_height,
+                    q_blk_h_aligned, kv_blk_h_aligned):
+        """load attn mask from gm to ub, then mul it by MASK_FILL_VALUE and add Sij"""
+        with self.tik_instance.new_stmt_scope(disable_sync=False):
+            att_mask_ub = self.tik_instance.Tensor(FP16, (kv_blk_h_aligned // self.N0, q_blk_h_aligned, self.N0),
+                                                   scope=UB, name="att_mask_ub")
+            self.tik_instance.data_move(att_mask_ub, self.att_mask_gm[attn_mask_gm_offset], 0,
+                                        kv_blk_height // self.N0, q_blk_height * self.N0 // 16,
+                                        (self.Nq - q_blk_height) * self.N0 // 16, 0)
+            self.tik_instance.h_mul(att_mask_ub, att_mask_ub, MASK_FILL_VALUE)
+            self.tik_instance.h_add(Sij_ub_N1MN0, Sij_ub_N1MN0, att_mask_ub)
+
+    def do_dropout_mask(self, Pij_ub, dropout_mask_gm_offset, kv_blk_h_aligned, kv_blk_height,
+                        q_blk_h_aligned, q_blk_height, precision_type=FP16, workspace=None):
+        """load drop mask from gm to ub, then mul it by Pij"""
+        with self.tik_instance.new_stmt_scope(disable_sync=False):
+            dropout_mask_ub = self.tik_instance.Tensor(FP16, (q_blk_h_aligned, kv_blk_h_aligned),
+                                                       scope=UB, name="drop_mask_ub")
+            self.tik_instance.data_move(dropout_mask_ub, self.drop_mask_gm[dropout_mask_gm_offset], 0,
+                                        q_blk_height, kv_blk_height // 16, (self.N - kv_blk_height) // 16, 0)
+            dropout_mask_ub = dropout_mask_ub.reshape((kv_blk_height // self.N0, q_blk_height, self.N0))
+            if precision_type == FP32:
+                dropout_mask_ub_fp32 = self.tik_instance.Tensor(FP32,
+                                                                (kv_blk_h_aligned // self.N0, q_blk_h_aligned, self.N0),
+                                                                scope=UB, name="dropout_mask_ub_fp32")
+                self.tik_instance.h_cast(dropout_mask_ub_fp32, dropout_mask_ub, "none")
+                if workspace is None:
+                    self.tik_instance.h_mul(Pij_ub, Pij_ub, dropout_mask_ub_fp32)
+                else:
+                    self.tik_instance.h_mul(workspace, Pij_ub, dropout_mask_ub_fp32)
+            else:
+                if workspace is None:
+                    self.tik_instance.h_mul(Pij_ub, Pij_ub, dropout_mask_ub)
+                else:
+                    self.tik_instance.h_mul(workspace, Pij_ub, dropout_mask_ub)
+
+    def compute_process(self):
+        """The compute process of FlashAttention"""
+        self.init()
+        self.prepare_global_ones()
+        core_idx_to_batch_info, core_idx_to_tr_info = self.get_core_task_info()
+        with self.tik_instance.for_range(begint=0, endt=self.core_num, name="core_index",
+                                         block_num=self.core_num) as core_idx:
+            batch_start_s = self.tik_instance.Scalar("int32", name="batch_start_s")
+            batch_num_s = self.tik_instance.Scalar("int32", name="batch_num_s")
+
+            batch_start_s.set_as(core_idx_to_batch_info[core_idx, 0])
+            batch_num_s.set_as(core_idx_to_batch_info[core_idx, 1])
+
+            self.compute_one_core(batch_start_s, batch_num_s, core_idx_to_tr_info, core_idx)
+
+        self.tik_instance.BuildCCE(
+            kernel_name=self.kernel_name,
+            inputs=self.collect_inputs(),
+            outputs=self.collect_outputs(),
+            config={"dump_cce_code": False, "save_temp_cce_file": True, "enable_const_fold": True},
+            enable_l2=True
+        )

mindspore/{_extends/graph_kernel/expanders/complex/__init__.py → ops/_op_impl/_custom_op/flash_attention/constants.py}

@@ -1,4 +1,4 @@
-# Copyright 2021-2022 Huawei Technologies Co., Ltd
+# Copyright 2023 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,11 +12,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""complex expanders init"""
+"""define constants"""
+from tbe import tik
 
-from .abs import CAbs
-from .add import CAdd
-from .div import CDiv
-from .mul import CMul
-from .sub import CSub
-from .real_div import CRealDiv
+BLOCK_NUM = 16
+FP16 = "float16"
+INT8 = "int8"
+INT32 = "int32"
+FP32 = "float32"
+REPEAT_SZ = 128
+BLK_STRIDE = 1
+REPEAT_STRIDE = 8
+TRANS_CUBE_TGT = 8
+FP16_MIN_VAL = -65504.0
+MASK_FILL_VALUE = -10000.0
+GM = tik.scope_gm
+L1 = tik.scope_cbuf
+L1OUT = tik.scope_cbuf_out
+UB = tik.scope_ubuf
+L0A = tik.scope_ca
+L0B = tik.scope_cb
+L0C = tik.scope_cc
+DTYPE_SIZE = {
+    "int8": 1,
+    "float16": 2,
+    "int16": 2,
+    "float32": 4,
+}