mindspore 2.0.0rc1__cp38-none-any.whl → 2.2.0__cp38-none-any.whl

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

@@ -0,0 +1,193 @@
+# 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 impl of flash attention"""
+from __future__ import absolute_import
+import mindspore.ops as ops
+import mindspore.common.dtype as mstype
+from mindspore.ops import Custom
+from mindspore.ops import DataType
+from mindspore.ops import TBERegOp
+from mindspore.ops._op_impl._custom_op.flash_attention.flash_attention_bwd import flash_attention_grad
+from mindspore.ops._op_impl._custom_op.flash_attention.flash_attention_fwd import flash_attention
+from mindspore.ops.composite.multitype_ops.zeros_like_impl import zeros_like
+
+KERNEL_NAME = "flash_attention"
+
+cus_flash_atten_op_info = TBERegOp("FlashAttentionPrimitive") \
+    .fusion_type("OPAQUE") \
+    .partial_flag(True) \
+    .async_flag(False) \
+    .binfile_name("flash_attention.so") \
+    .compute_cost(10) \
+    .kernel_name(KERNEL_NAME) \
+    .attr("prev_block_num", "required", "int", "all", "65536") \
+    .attr("next_block_num", "required", "int", "all", "65536") \
+    .attr("high_precision", "required", "bool", "all", "false") \
+    .attr("tiling_stgy_name", "required", "str", "all", "sparse") \
+    .input(0, "query", False, "required", "all") \
+    .input(1, "key", False, "required", "all") \
+    .input(2, "value", False, "required", "all") \
+    .input(3, "attn_mask", False, "optional", "all") \
+    .input(4, "dropout_mask", False, "optional", "all") \
+    .input(5, "alibi_mask", False, "optional", "all") \
+    .output(0, "output", False, "required", "all") \
+    .output(1, "rowsum", False, "required", "all") \
+    .output(2, "rowmax", False, "required", "all") \
+    .dtype_format(DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_Default,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_Default,
+                  DataType.F16_Default) \
+    .dtype_format(DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_Default,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F32_Default,
+                  DataType.F16_Default) \
+    .get_op_info()
+
+GRAD_KERNEL_NAME = "flash_attention_grad"
+
+cus_flash_atten_grad_op_info = TBERegOp("FlashAttentionGradPrimitive") \
+    .fusion_type("OPAQUE") \
+    .partial_flag(True) \
+    .async_flag(False) \
+    .binfile_name("flash_attention_grad.so") \
+    .compute_cost(10) \
+    .kernel_name(GRAD_KERNEL_NAME) \
+    .attr("prev_block_num", "required", "int", "all", "65536") \
+    .attr("next_block_num", "required", "int", "all", "65536") \
+    .attr("high_precision", "required", "bool", "all", "false") \
+    .attr("tiling_stgy_name", "required", "str", "all", "sparse") \
+    .input(0, "query", False, "required", "all") \
+    .input(1, "key", False, "required", "all") \
+    .input(2, "value", False, "required", "all") \
+    .input(3, "output", False, "required", "all") \
+    .input(4, "do", False, "required", "all") \
+    .input(5, "rowsum", False, "required", "all") \
+    .input(6, "rowmax", False, "required", "all") \
+    .input(7, "attn_mask", False, "optional", "all") \
+    .input(8, "dropout_mask", False, "optional", "all") \
+    .input(9, "alibi_mask", False, "optional", "all") \
+    .output(0, "dq", False, "required", "all") \
+    .output(1, "dk", False, "required", "all") \
+    .output(2, "dv", False, "required", "all") \
+    .dtype_format(DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_Default,
+                  DataType.F16_Default,
+                  DataType.F16_FracNZ,
+                  DataType.F16_Default,
+                  DataType.F16_FracNZ,
+                  DataType.F32_FracNZ,
+                  DataType.F32_FracNZ,
+                  DataType.F32_FracNZ) \
+    .dtype_format(DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F32_Default,
+                  DataType.F16_Default,
+                  DataType.F16_FracNZ,
+                  DataType.F16_Default,
+                  DataType.F16_FracNZ,
+                  DataType.F32_FracNZ,
+                  DataType.F32_FracNZ,
+                  DataType.F32_FracNZ) \
+    .get_op_info()
+
+
+def get_flash_attention_grad(prev_block_num=65536, next_block_num=65536,
+                             tiling_stgy_name='sparse', high_precision=False):
+    """get flash attention grad"""
+
+    def infer_shape(q_shape, k_shape, v_shape, o_shape, do_shape, l_shape, m_shape,
+                    att_mask_shape, dropout_mask_shape, alibi_mask_shape):
+        return q_shape, k_shape, v_shape
+
+    def infer_dtype(q_dtype, k_dtype, v_dtype, o_dytpe, do_dtype, l_dtype, m_dtype,
+                    attn_mask_dtype, dropout_mask_dtype, alibi_mask_type):
+        return mstype.float32, mstype.float32, mstype.float32
+
+    fa_grad = Custom(flash_attention_grad, out_shape=infer_shape,
+                     out_dtype=infer_dtype, func_type="tbe", reg_info=cus_flash_atten_grad_op_info)
+    fa_grad.add_prim_attr("prev_block_num", prev_block_num)
+    fa_grad.add_prim_attr("next_block_num", next_block_num)
+    fa_grad.add_prim_attr("high_precision", high_precision)
+    fa_grad.add_prim_attr("tiling_stgy_name", tiling_stgy_name)
+    fa_grad.init_prim_io_names(
+        inputs=["query", "key", "value", "output", "do", "rowsum", "rowmax", "attn_mask", "dropout_mask",
+                "alibi_mask"],
+        outputs=["dq", "dk", "dv"]
+    )
+
+    def bprop(query, key, value, attn_mask, dropout_mask, alibi_mask, out, douts):
+        output, rowsum, rowmax = out
+        dout, _, _ = douts
+        dq, dk, dv = fa_grad(query, key, value, output, dout, rowsum, rowmax, attn_mask, dropout_mask,
+                             alibi_mask)
+        dq = ops.cast(dq, mstype.float16)
+        dk = ops.cast(dk, mstype.float16)
+        dv = ops.cast(dv, mstype.float16)
+        return dq, dk, dv, zeros_like(attn_mask), \
+            zeros_like(dropout_mask), zeros_like(alibi_mask)
+
+    return bprop
+
+
+def get_flash_attention(prev_block_num=65536, next_block_num=65536, tiling_stgy_name='sparse', high_precision=False):
+    """get_flash_attention"""
+
+    def infer_shape(q_shape, k_shape, v_shape, attn_mask_shape=None,
+                    dropout_mask_shape=None, alibi_mask_shape=None):
+        """infer shape"""
+        batch, hidden_size, seq_len, _ = q_shape
+        l_shape = (batch, hidden_size, seq_len)
+        m_shape = (batch, hidden_size, seq_len)
+        return q_shape, l_shape, m_shape
+
+    def infer_dtype(q_dtype, k_dtype, v_dtype, attn_mask_dtype=None,
+                    dropout_mask_dtype=None, alibi_mask_type=None):
+        """infer type"""
+        l_dtype = mstype.float16
+        if high_precision:
+            l_dtype = mstype.float32
+        return q_dtype, l_dtype, q_dtype
+
+    fa_grad = get_flash_attention_grad(prev_block_num, next_block_num, tiling_stgy_name, high_precision)
+    fa_forward = Custom(flash_attention, out_shape=infer_shape,
+                        out_dtype=infer_dtype, func_type="tbe", bprop=fa_grad,
+                        reg_info=cus_flash_atten_op_info)
+    fa_forward.add_prim_attr("prev_block_num", prev_block_num)
+    fa_forward.add_prim_attr("next_block_num", next_block_num)
+    fa_forward.add_prim_attr("high_precision", high_precision)
+    fa_forward.add_prim_attr("tiling_stgy_name", tiling_stgy_name)
+    fa_forward.init_prim_io_names(
+        inputs=["query", "key", "value", "attn_mask", "dropout_mask", "alibi_mask"],
+        outputs=["output", "rowsum", "rowmax"]
+    )
+
+    return fa_forward

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

@@ -0,0 +1,435 @@
+# 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 common about tik ops"""
+from functools import partial
+
+from mindspore.ops._op_impl._custom_op.flash_attention.constants import DTYPE_SIZE
+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 L0C
+from mindspore.ops._op_impl._custom_op.flash_attention.constants import UB
+
+
+class TikOpsUtils:
+    """Utils function class about tik ops"""
+
+    def __init__(self, tik_instance):
+        self.tik_instance = tik_instance
+        self.dtype = "float16"
+        self.cont_data_mv_1_bust = partial(self.tik_instance.data_move, sid=0, nburst=1,
+                                           src_stride=0,
+                                           dst_stride=0)
+
+    def MK_TO_K1MK0(self, mk_input_tensor, workspace_tensor=None):
+        """change data shape from (M, K) to (K1, M, K0), K1 = K // K0, the effect is equant to:
+        new_tensor =  np.stack(np.hsplit(mk_input_tensor, K1), axis=0)
+
+        :param mk_input_tensor: input tensor in GM with shape: (M, K)
+        :param workspace_tensor: workspace tensor with shape: (K1, M, K0)
+        tensor will be changed, otherwise the new data will be copied to the workspace tensor,
+        and input tensor will stay unchanged.
+        :return: Tensor with shape (K1,M, K0)
+        """
+        dtype = mk_input_tensor.dtype
+        m, k = mk_input_tensor.shape
+        K0 = 16
+        K1 = k // K0
+        M = self.up_align_to_K0(m)
+        try:
+            dtype_size = DTYPE_SIZE[dtype]
+        except KeyError:
+            raise ValueError("The argument 'dtype' is not valid.")
+        if workspace_tensor is not None:
+            with self.tik_instance.for_range(0, K1) as i:
+                self.tik_instance.data_move(
+                    workspace_tensor[i * M * K0:],
+                    mk_input_tensor[i * K0:],
+                    0,
+                    M,
+                    K0 * dtype_size // 32,
+                    (K1 - 1) * K0 * dtype_size // 32,
+                    0,
+                )
+            return workspace_tensor.reshape((K1, M, K0))
+
+        with self.tik_instance.new_stmt_scope(disable_sync=False):
+            tmp_ub = self.tik_instance.Tensor(dtype, (K1, M, K0), name="tmp_ub", scope=UB)
+            # data_move(m,k) --> (k1,m,K0)
+            with self.tik_instance.for_range(0, K1) as i:
+                self.tik_instance.data_move(
+                    tmp_ub[i * M * K0:],
+                    mk_input_tensor[i * K0:],
+                    0,
+                    M,
+                    K0 * dtype_size // 32,
+                    (K1 - 1) * K0 * dtype_size // 32,
+                    0,
+                )
+            self.cont_data_mv_1_bust(
+                dst=mk_input_tensor, src=tmp_ub, burst=K1 * M * K0 * dtype_size // 32)
+            return mk_input_tensor.reshape((K1, M, K0))
+
+    def transpose_matrix(self, src_ub, dst_ub, N, nk0=False):
+        """ transpose matrix, default support shape: (16, n) -> (n, 16)
+        if nk0 is true, support shape: (n, 16) -> (16, n)
+        """
+        K0 = 16
+        rep_times = N // K0
+        if nk0:
+            src_list = [src_ub[16 * i] for i in range(16)]
+            dst_list = [dst_ub[N * i] for i in range(16)]
+        else:
+            src_list = [src_ub[N * i] for i in range(16)]
+            dst_list = [dst_ub[16 * i] for i in range(16)]
+
+        dst_rep_stride = K0
+        src_rep_stride = 1
+        if rep_times == 1:
+            dst_rep_stride = 0
+            src_rep_stride = 0
+
+        if nk0:
+            src_rep_stride, dst_rep_stride = dst_rep_stride, src_rep_stride
+
+        self.tik_instance.vec_trans_scatter(
+            False, False, dst_list, src_list, rep_times, dst_rep_stride, src_rep_stride
+        )
+        return dst_ub
+
+    def KN_TO_K1NK0(self, kn_input_tensor, workspace_tensor=None):
+        """change data shape from (K,N) to (K1, N, K0), K1 = K // K0, the effect is equvilent to:
+        new_tensor =  np.reshape(kn_input_tensor, newshape=(K1, K0, N)).swapaxes(1, 2)
+
+        :param kn_input_tensor: input tensor with shape: (K, N)
+        :param workspace_tensor: workspace tensor with shape: (K1, N, K0)
+        tensor will be changed, otherwise the new data will be copied to the workspace tensor,
+        and input tensor will stay unchanged.
+        :return: Tensor with shape: (K1, N, K0)
+        """
+        dtype = kn_input_tensor.dtype
+        k, n = kn_input_tensor.shape
+        K0 = 16
+        K1 = k // K0
+        N = n
+        try:
+            dtype_size = DTYPE_SIZE[dtype]
+        except KeyError:
+            raise ValueError("The argument 'dtype' is not valid.")
+        with self.tik_instance.for_range(0, K1) as index:
+            k1nk0_ub = self.tik_instance.Tensor(dtype, (N, K0), UB, "k1nk0_ub")
+            src_ub = self.tik_instance.Tensor(dtype, (K0, N), UB, "src_ub")
+            burst_len = K0 * N * dtype_size // 32
+            self.cont_data_mv_1_bust(dst=src_ub, src=kn_input_tensor[index * K0 * N],
+                                     burst=burst_len)
+            k1nk0_ub = self.transpose_matrix(src_ub, k1nk0_ub, N)
+            if workspace_tensor is None:
+                self.cont_data_mv_1_bust(dst=kn_input_tensor[index * K0 * N], src=k1nk0_ub,
+                                         burst=burst_len)
+            else:
+                self.cont_data_mv_1_bust(dst=workspace_tensor[index * K0 * N], src=k1nk0_ub,
+                                         burst=burst_len)
+        if workspace_tensor is None:
+            return kn_input_tensor.reshape((K1, N, K0))
+
+        return workspace_tensor.reshape((K1, N, K0))
+
+    def N1MN0_TO_MN(self, N1MN0_input):
+        """change data shape from (N1, M, N0) to (M, N), N0=16, N = N1 * K0, the effect is equant to:
+        N1MN0_input = np.concatenate(list(map(np.squeeze, np.split(N1MN0_input, N1))), axis=1)
+
+        :param N1MN0_input: input tensor with shape (N, M, N0) in GM or L1.
+        :return:
+        """
+        dtype = N1MN0_input.dtype
+        N1, M, N0 = N1MN0_input.shape
+        try:
+            dtype_size = DTYPE_SIZE[dtype]
+        except KeyError:
+            raise ValueError("The argument 'dtype' is not valid.")
+        with self.tik_instance.new_stmt_scope(disable_sync=False):
+            tmp_ub = self.tik_instance.Tensor(dtype, (M, N1 * N0), name="tmp_ub", scope=UB)
+            # data_move (n1,m,n0) --> (m,n)
+            with self.tik_instance.for_range(0, N1) as i:
+                self.tik_instance.data_move(
+                    tmp_ub[i * N0:],
+                    N1MN0_input[i * M * N0:],
+                    0,
+                    M,
+                    N0 * dtype_size // 32,
+                    0,
+                    (N1 - 1) * N0 * dtype_size // 32,
+                )
+            # data_move out
+            self.cont_data_mv_1_bust(dst=N1MN0_input, src=tmp_ub, burst=M * N1 * N0 * dtype_size // 32)
+        return N1MN0_input.reshape((M, N1 * N0))
+
+    def broadcast(self, vec_ub, shape):
+        """ broadcast a vector to a matrix
+        :param vec_ub: a tensor in UB with shape of (M,), and dtype is float16
+        :param shape: the target shape, a tuple with value (M, N), M and N are integer multiples of 16
+        :return: a tensor in UB with shape of (M, N)
+        """
+        M, N = shape
+        dst_ub = self.tik_instance.Tensor(FP16, shape, name="dst_ub", scope=UB)
+
+        with self.tik_instance.new_stmt_scope(disable_sync=False):
+            # (M,) -> (2, M) -> (4, M) -> (8, M) -> (16, M)
+            tmp_ub1 = self.tik_instance.Tensor(FP16, (16, M), name="tmp_ub1", scope=UB)
+            self.tik_instance.data_move(tmp_ub1, vec_ub, 0, 1, M // 16, 0, 0)
+            times = self.tik_instance.Scalar("int32", name="times", init_value=1)
+            with self.tik_instance.for_range(begint=0, endt=16):
+                with self.tik_instance.if_scope(times <= 8):
+                    offset = times * M
+                    burst = times * M // 16
+                    self.cont_data_mv_1_bust(dst=tmp_ub1[offset], src=tmp_ub1, burst=burst)
+                with self.tik_instance.else_scope():
+                    self.tik_instance.tik_break()
+                times.set_as(times * 2)
+
+            # (16, M) -> (M, 16)
+            tmp_ub2 = self.tik_instance.Tensor(FP16, (M, 16), name="tmp_ub2", scope=UB)
+            tmp_ub2_transposed = self.transpose_matrix(tmp_ub1, tmp_ub2, M)
+
+            # (M, 16) -> (M, 32) -> (M, 64) -> ... -> (M, N)
+            self.tik_instance.data_move(dst_ub, tmp_ub2_transposed, 0, M, 1, 0, N // 16 - 1)
+            times.set_as(1)
+            with self.tik_instance.for_range(begint=0, endt=N):
+                offset = times * 16
+                with self.tik_instance.if_scope(offset * 2 <= N):
+                    burst = offset // 16
+                    src_stride = N // 16 - burst
+                    dst_stride = N // 16 - burst
+                    self.tik_instance.data_move(dst_ub[offset], dst_ub, 0, M, burst, src_stride,
+                                                dst_stride)
+                with self.tik_instance.else_scope():
+                    burst = (N - offset) // 16
+                    src_stride = N // 16 - burst
+                    dst_stride = N // 16 - burst
+                    with self.tik_instance.if_scope(burst > 0):
+                        self.tik_instance.data_move(dst_ub[offset], dst_ub, 0, M, burst, src_stride,
+                                                    dst_stride)
+                    self.tik_instance.tik_break()
+                times.set_as(times * 2)
+        return dst_ub
+
+    def broadcast_row(self, vec_ub, shape):
+        """broadcast row"""
+        M, N = shape
+        dst_ub = self.tik_instance.Tensor(FP16, shape, name="dst_ub", scope=UB)
+        self.tik_instance.data_move(dst_ub, vec_ub, 0, 1, N // 16, 0, 0)
+        times = self.tik_instance.Scalar("int32", name="times", init_value=1)
+        # (1, N) -> (2, M) -> (4, N) -> ... -> (M, N)
+        with self.tik_instance.for_range(begint=0, endt=M):
+            with self.tik_instance.if_scope(times * 2 <= M):
+                burst = times * N // 16
+                offset = times * N
+                self.tik_instance.data_move(dst_ub[offset], dst_ub, 0, 1, burst, 0, 0)
+            with self.tik_instance.else_scope():
+                burst = (M - times) * N // 16
+                offset = times * N
+                with self.tik_instance.if_scope(burst > 0):
+                    self.tik_instance.data_move(dst_ub[offset], dst_ub, 0, 1, burst, 0, 0)
+                self.tik_instance.tik_break()
+            times.set_as(times * 2)
+        return dst_ub
+
+    def get_K0(self, dtype=None):
+        """get K0"""
+        if dtype is None:
+            dtype = self.dtype
+        try:
+            dtype_size = DTYPE_SIZE[dtype]
+        except KeyError:
+            raise ValueError("The argument 'dtype' is not valid.")
+        return 32 // dtype_size
+
+    def up_align_to_K0(self, n, dtype=None):
+        """byte alignment by dtype"""
+        if dtype is None:
+            dtype = self.dtype
+        try:
+            dtype_size = DTYPE_SIZE[dtype]
+        except KeyError:
+            raise ValueError("The argument 'dtype' is not valid.")
+        K0 = 32 // dtype_size
+        return (n + K0 - 1) // K0 * K0
+
+    def calc_vec_rec(self, vec_ub, vec_len):
+        """cal the reciprocal of a vector"""
+        dtype = vec_ub.dtype
+        vec_len_aligned = self.up_align_to_K0(vec_len)
+        vec_rec_ub = self.tik_instance.Tensor(dtype, (vec_len_aligned,), scope=UB, name="li_new_rec_ub")
+        try:
+            dtype_size = DTYPE_SIZE[dtype]
+        except KeyError:
+            raise ValueError("The argument 'dtype' is not valid.")
+        mask_len = 256 // dtype_size
+        block_len = 32 // dtype_size
+        work_size = 8 // dtype_size
+
+        with self.tik_instance.new_stmt_scope(disable_sync=False):
+            repeat_times = vec_len // mask_len
+            if repeat_times > 0:
+                dst_rep_stride = 8
+                src_rep_stride = 8
+
+                src_extent_size = (repeat_times - 1) * src_rep_stride * block_len + mask_len
+                wk_size_unit = ((src_extent_size + block_len - 1) // block_len) * block_len
+                wk_size = work_size * wk_size_unit
+                # 定义work_tensor
+                work_tensor_ub = self.tik_instance.Tensor(
+                    "float32", (wk_size,), name="work_tensor_ub", scope=UB
+                )
+                # 如果work_tensor有索引，需要写成 work_tensor[index:]
+                self.tik_instance.vec_rec_high_preci(
+                    mask_len,
+                    vec_rec_ub[0:],
+                    vec_ub[0:],
+                    work_tensor_ub[0:],
+                    repeat_times,
+                    dst_rep_stride,
+                    src_rep_stride,
+                )
+
+            mask_len = vec_len - repeat_times * mask_len
+            if mask_len > 0:
+                wk_size = work_size * ((mask_len + block_len - 1) // block_len) * block_len
+                work_tensor_ub2 = self.tik_instance.Tensor(
+                    "float32", (wk_size,), name="work_tensor_ub2", scope=UB
+                )
+                self.tik_instance.vec_rec_high_preci(
+                    mask_len,
+                    vec_rec_ub[repeat_times * 128:],
+                    vec_ub[repeat_times * 128:],
+                    work_tensor_ub2[0:],
+                    1,
+                    0,
+                    0,
+                )
+        return vec_rec_ub
+
+    def row_sum_cube_impl(self, matrix_l1_K1MK0_ed, right_all_one_matrix_l1, rowsum_ub, m, k, precision_type):
+        """用cube实现矩阵行和：右乘一个shape=(n,1)全一矩阵
+        :param matrix_l1_K1MK0_ed: input tensor with shape (K1, M, K0)
+        :param right_all_one_matrix_l1: input tensor with shape (K, 16)
+        :param rowsum_ub: output tensor stores the row sum of input tensor
+        :param m: actual tensor height
+        :param k: actual tensor width
+        :return: row sum of the output tensor
+        """
+        K1, M, K0 = matrix_l1_K1MK0_ed.shape
+        # 调用matmul实现rowsum，结果shape=(m, 16)，取每行的第一个数
+        with self.tik_instance.new_stmt_scope(disable_sync=False):
+            row_sum_ub_N1MN0 = self.matmul_compute(matrix_l1_K1MK0_ed, right_all_one_matrix_l1, m, k, 16,
+                                                   N1MN0_to_MN=False, precision_type=precision_type)
+            row_sum_ub_MN_ed = row_sum_ub_N1MN0.reshape((M, 16))
+            if precision_type == FP32:
+                for idx in range(0, m):
+                    cur_row_sum = self.tik_instance.Scalar(FP32, init_value=row_sum_ub_MN_ed[idx, 0])
+                    rowsum_ub[idx].set_as(cur_row_sum)
+            else:
+                # row_sum_ub_MN_ed 先转置，然后取一行, 替换原来按行操作: lij_ub[i].set_as(row_sum_ub_MN_ed[i, 0])
+                row_sum_ub_trans = self.tik_instance.Tensor(FP16, (16, M), name="row_sum_ub_trans", scope=UB)
+                row_sum_ub_trans = self.transpose_matrix(row_sum_ub_MN_ed, row_sum_ub_trans, M, True)
+                self.cont_data_mv_1_bust(dst=rowsum_ub, src=row_sum_ub_trans, burst=M // 16)
+
+        return rowsum_ub
+
+    def matmul_compute(self, A_l1, B_l1, m, k, n, N1MN0_to_MN=True, precision_type=FP16):
+        """calculate matrix multiplication A_l1 * B_l1, and move the result to C_ub,
+        then rearrange C_ub
+        :param A_l1: input tensor in L1 with shape of (K1, M, K0)
+        :param B_l1: input tensor in L1 with shape of (K1, N, K0)
+        :param m: the actual number of rows of A_l1
+        :param k: the actual number of cols of A_l1
+        :param n: the actual number of cols of B_l1
+        :param N1MN0_to_MN: Whether reorder the result tensor.
+        :return: C_ub with tensor with shape of (M, N) if N1MN0_to_MN else (N1, M, N0)
+        """
+        M = self.up_align_to_K0(m)
+        N = self.up_align_to_K0(n)
+        C_ub = self.tik_instance.Tensor(precision_type, (N // 16, M, 16), name="C_ub", scope=UB)
+        try:
+            dtype_size = DTYPE_SIZE[FP32]
+        except KeyError:
+            raise ValueError("The argument 'dtype' is not valid.")
+        with self.tik_instance.new_stmt_scope(disable_sync=False):
+            # matmul
+            C_l0c = self.tik_instance.Tensor(
+                FP32, (N // 16, M, 16), scope=L0C, name="C_l0c"
+            )  # n1mn0 (n0=16)
+            self.tik_instance.matmul(C_l0c, A_l1, B_l1, m, k, n)
+            # L0C -> ub, fp32 -> fp16 (tensor_mov可做随路转换)
+            self.tik_instance.tensor_mov(C_ub, C_l0c, "m", 1, M * N * dtype_size // 1024, 0, 0)
+        if N1MN0_to_MN:
+            return self.N1MN0_TO_MN(C_ub)
+        return C_ub
+
+    def move_vector_from_gm_to_ub(self, dst_tensor, src_tensor, gm_offset, vec_len):
+        """load the vector from gm to ub
+        :param dst_tensor:
+        :param src_tensor:
+        :param gm_offset:
+        :return:
+        """
+        try:
+            dtype_size = DTYPE_SIZE[src_tensor.dtype]
+        except KeyError:
+            raise ValueError("The argument 'src_tensor dtype' is not valid.")
+        a_burst_num = 32 // dtype_size
+        full_tik_blk_num, tail_num = divmod(vec_len, a_burst_num)
+        with self.tik_instance.if_scope(full_tik_blk_num > 0):
+            self.cont_data_mv_1_bust(dst=dst_tensor, src=src_tensor[gm_offset],
+                                     burst=full_tik_blk_num)
+        # 地址回退处理尾部数据
+        with self.tik_instance.if_scope(tail_num > 0):
+            offset = vec_len - a_burst_num
+            last_blk_ub = self.tik_instance.Tensor(FP16, (a_burst_num,), name="last_blk_ub", scope=UB)
+            self.cont_data_mv_1_bust(dst=last_blk_ub, src=src_tensor[gm_offset + offset], burst=1)
+            with self.tik_instance.for_range(0, a_burst_num) as idx:  # offset非32bytes对齐, 无法用datamove
+                dst_tensor[offset + idx].set_as(last_blk_ub[idx])
+
+    def move_vector_from_ub_to_gm(self, dst_tensor, src_tensor, gm_offset, block_h):
+        """write the vector back to gm
+        :param dst_tensor:
+        :param src_tensor:
+        :param gm_offset:
+        :param block_h:
+        :return:
+        """
+        try:
+            dtype_size = DTYPE_SIZE[src_tensor.dtype]
+        except KeyError:
+            raise ValueError("The argument 'src_tensor dtype' is not valid.")
+        a_burst_num = 32 // dtype_size
+        full_tik_blk_num = block_h // a_burst_num
+        with self.tik_instance.if_scope(full_tik_blk_num > 0):
+            self.cont_data_mv_1_bust(dst=dst_tensor[gm_offset], src=src_tensor,
+                                     burst=full_tik_blk_num)
+        tail_num = block_h % a_burst_num
+        with self.tik_instance.if_scope(tail_num > 0):
+            offset = block_h - a_burst_num
+            tmp_ub = self.tik_instance.Tensor(FP16, (a_burst_num,), name="tmp_ub", scope=UB)
+            with self.tik_instance.for_range(0, a_burst_num) as idx:
+                tmp_ub[idx].set_as(src_tensor[offset + idx])
+            self.cont_data_mv_1_bust(dst=dst_tensor[gm_offset + offset], src=tmp_ub, burst=1)
+
+    def scale_compute_vector(self, Sij_ub, dim):
+        """scale compute vector"""
+        scale_value = dim ** -0.5
+        scale = self.tik_instance.Scalar(dtype=FP16)
+        scale.set_as(scale_value)
+        self.tik_instance.h_mul(Sij_ub, Sij_ub, scale)
+        return Sij_ub

mindspore/ops/_op_impl/_custom_op/flash_attention/tiling_strategy/sparse_tiling.py

@@ -0,0 +1,45 @@
+# 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.
+# ============================================================================
+"""tiling for sparse """
+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
+
+
+class SparseTiling(TilingStrategy):
+    """A tiling strategy implementation for sparse shape"""
+
+    @classmethod
+    def strategy_name(cls):
+        return "sparse"
+
+    def tiling(self) -> TilingPara:
+        self.Br = min(128, self.Nq)
+        self.Bc = min(128, self.N)
+
+        self.Tr = self.Nq // self.Br
+        self.Tc = self.N // self.Bc
+
+        if self.Nq % self.Br != 0:
+            self.last_Br = self.Nq - self.Tr * self.Br
+            self.Tr += 1
+        else:
+            self.last_Br = self.Br
+        if self.N % self.Bc != 0:
+            self.last_Bc = self.N - self.Tc * self.Bc
+            self.Tc += 1
+        else:
+            self.last_Bc = self.Bc
+
+        return self.gen_tiling_para()