mindspore 2.1.0__cp38-cp38-manylinux1_x86_64.whl → 2.2.10__cp38-cp38-manylinux1_x86_64.whl

mindspore/nn/layer/dense.py

@@ -77,6 +77,7 @@ class BiDense(Cell):
         bias_init (Union[Tensor, str, Initializer, numbers.Number]): The trainable bias_init parameter.
             The values of str refer to the function `initializer`. Default: ``None`` .
         has_bias (bool): Specifies whether the layer uses :math:`\text{bias}` vector. Default: ``True`` .
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Shape:
         - **input1** - :math:`(*, H_{in1})` where :math:`H_{in1}=\text{in1_channels}` and
@@ -90,8 +91,8 @@ class BiDense(Cell):
           are the same shape as the inputs.
 
     Dtype:
-        - **input1** (Tensor) - The dtype must be float16 or float32 and be same as **input2**.
-        - **input1** (Tensor) - The dtype must be float16 or float32 and be same as **input1**.
+        - **input1** (Tensor) - The dtype must be float16 or float32 and be same as **input2** .
+        - **input2** (Tensor) - The dtype must be float16 or float32 and be same as **input1** .
         - **output** (Tensor) - With the same dtype as the inputs.
 
     Weights:
@@ -133,7 +134,8 @@ class BiDense(Cell):
                  out_channels,
                  weight_init=None,
                  bias_init=None,
-                 has_bias=True):
+                 has_bias=True,
+                 dtype=mstype.float32):
         super().__init__()
         self.in_channels = Validator.check_positive_int(in1_channels, "in1_channels", self.cls_name)
         self.in_channels = Validator.check_positive_int(in2_channels, "in2_channels", self.cls_name)
@@ -156,7 +158,8 @@ class BiDense(Cell):
                                  f"equal to 'in2_channels'. But got 'weight_init': {weight_init}, "
                                  f"'out_channels': {out_channels}, 'in_channels': {in1_channels}, "
                                  f"'in2_channels': {in2_channels}")
-        self.weight = Parameter(initializer(weight_init, (out_channels, in1_channels, in2_channels)), 'weight')
+        self.weight = Parameter(initializer(weight_init, (out_channels, in1_channels, in2_channels), dtype=dtype),
+                                'weight')
 
         if self.has_bias:
             if bias_init is None:
@@ -166,7 +169,7 @@ class BiDense(Cell):
                     raise ValueError(f"For '{self.cls_name}', bias init shape error. The ndim of 'bias_init' should "
                                      f"be equal to 1, and the first dim must be equal to 'out_channels'. But got "
                                      f"'bias_init': {bias_init}, 'out_channels': {out_channels}.")
-            self.bias = Parameter(initializer(bias_init, [out_channels]), name="bias")
+            self.bias = Parameter(initializer(bias_init, [out_channels], dtype=dtype), name="bias")
             self.bias_add = P.BiasAdd()
         self.matmul = P.MatMul()
 

mindspore/nn/layer/embedding.py

@@ -64,11 +64,13 @@ class Embedding(Cell):
         embedding_size (int): The size of each embedding vector.
         use_one_hot (bool): Specifies whether to apply one_hot encoding form. Default: ``False`` .
         embedding_table (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the embedding_table.
-            Refer to class `initializer` for the values of string when a string
-            is specified. Default: ``'normal'`` .
+            Refer to class `mindspore.common.initializer
+            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            for the values of string when a string is specified. Default: ``'normal'`` .
         dtype (:class:`mindspore.dtype`): Data type of `x`. Default: ``mstype.float32`` .
         padding_idx (int, None): When the padding_idx encounters index, the output embedding vector of this index
                                  will be initialized to zero. Default: ``None`` . The feature is inactivated.
+
     Inputs:
         - **x** (Tensor) - Tensor of shape :math:`(\text{batch_size}, \text{x_length})`. The elements of
           the Tensor must be integer and not larger than vocab_size. Otherwise the corresponding embedding vector will
@@ -145,9 +147,8 @@ class Embedding(Cell):
         return output
 
     def extend_repr(self):
-        s = 'vocab_size={}, embedding_size={}, use_one_hot={}, embedding_table={}, dtype={}, padding_idx={}'.format(
-            self.vocab_size, self.embedding_size, self.use_one_hot, self.embedding_table, self.dtype, self.padding_idx)
-        return s
+        return f'vocab_size={self.vocab_size}, embedding_size={self.embedding_size}, use_one_hot={self.use_one_hot}, ' \
+            f'embedding_table={self.embedding_table}, dtype={self.dtype}, padding_idx={self.padding_idx}'
 
 
 @_primexpr
@@ -190,6 +191,7 @@ class EmbeddingLookup(Cell):
             parameter server trainning mode and 'DEVICE' target. And the moment parameter of corresponding
             optimizer will also be set to the cache size. In addition, it should be noted that it will cost the 'DEVICE'
             memory, so suggests setting a reasonable value to avoid insufficient memory.
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Inputs:
         - **input_indices** (Tensor) - The shape of tensor is :math:`(y_1, y_2, ..., y_S)`.
@@ -231,7 +233,7 @@ class EmbeddingLookup(Cell):
 
     def __init__(self, vocab_size, embedding_size, param_init='normal',
                  target='CPU', slice_mode='batch_slice', manual_shapes=None,
-                 max_norm=None, sparse=True, vocab_cache_size=0):
+                 max_norm=None, sparse=True, vocab_cache_size=0, dtype=mstype.float32):
         """Initialize EmbeddingLookup."""
         super(EmbeddingLookup, self).__init__()
         Validator.check_value_type('sparse', sparse, [bool], self.cls_name)
@@ -255,8 +257,8 @@ class EmbeddingLookup(Cell):
         if enable_ps:
             self._process_vocab_cache(slice_mode)
         self.embedding_size = Validator.check_positive_int(embedding_size, 'embedding_size', self.cls_name)
-        self.embedding_table = Parameter(initializer(param_init, [self.vocab_size, self.embedding_size]),
-                                         name='embedding_table')
+        self.embedding_table = Parameter(initializer(param_init, [self.vocab_size, self.embedding_size],
+                                                     dtype=dtype), name='embedding_table')
         parallel_mode = _get_parallel_mode()
         is_auto_parallel = parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
         self.gather_revert = P.Gather()
@@ -267,7 +269,7 @@ class EmbeddingLookup(Cell):
         if is_auto_parallel:
             self.unique = P.Unique().shard(((1,),))
         if self.cache_enable and enable_ps:
-            self._set_voacb_cache_enable_for_ps(vocab_cache_size, embedding_size, vocab_size, param_init)
+            self._set_voacb_cache_enable_for_ps(vocab_cache_size, embedding_size, vocab_size, param_init, dtype=dtype)
             if is_auto_parallel:
                 self.unique.add_prim_attr('cache_enable', True)
         indices_shape_size = 2
@@ -310,8 +312,8 @@ class EmbeddingLookup(Cell):
         else:
             if is_auto_parallel:
                 support_mode = ["field_slice", "table_row_slice", "table_column_slice", "batch_slice"]
-                raise ValueError("For '{}', the 'slice_mode' must be in {}, "
-                                 "but got \"{}\".".format(self.cls_name, support_mode, slice_mode))
+                raise ValueError(f"For '{self.cls_name}', the 'slice_mode' must be in {support_mode}, "
+                                 f"but got \"{slice_mode}\".")
         if self.cache_enable and not enable_ps:
             raise ValueError(f"For '{self.cls_name}', haven't supported cache enable for not ps mode.")
         self.embedding_table.unique = self.forward_unique
@@ -354,7 +356,8 @@ class EmbeddingLookup(Cell):
             if _is_role_worker():
                 self.vocab_size = self.vocab_cache_size
 
-    def _set_voacb_cache_enable_for_ps(self, vocab_cache_size, embedding_size, vocab_size, param_init):
+    def _set_voacb_cache_enable_for_ps(self, vocab_cache_size, embedding_size, vocab_size, param_init,
+                                       dtype=mstype.float32):
         """PS embeddingLookup cache enable set."""
         if self.sparse:
             self.forward_unique = True
@@ -368,10 +371,10 @@ class EmbeddingLookup(Cell):
         if _enable_distributed_mindrt():
             self.rank_id = get_rank()
             if self.is_ps_server:
-                self._slice_pserver_embeddings("zeros")
+                self._slice_pserver_embeddings("zeros", dtype=dtype)
                 self._set_cache_enable_and_key_for_pserver(param_key)
 
-    def _slice_pserver_embeddings(self, param_init):
+    def _slice_pserver_embeddings(self, param_init, dtype=mstype.float32):
         '''
         Method to slice embedding tables on Parameter Servers.
         It helps to train with a large scale embedding table and is used only in Parameter Server training mode.
@@ -399,7 +402,7 @@ class EmbeddingLookup(Cell):
         for i in range(server_num):
             self.embedding_table_list.append(Parameter(initializer(param_init,
                                                                    [self.embedding_table_vocab_dim_list[i],
-                                                                    self.embedding_size]),
+                                                                    self.embedding_size], dtype=dtype),
                                                        name="embedding_table_server_" + str(i)))
 
             self.embedding_offset.append(offset)
@@ -505,12 +508,13 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
             :class:`mindspore.nn.EmbeddingLookup`. Default: ``'batch_slice'``.
         feature_num_list (tuple): The accompaniment array in field slice mode. This is unused currently.
             Default:  ``None`` .
-        max_norm (Union[float, None]): A maximum clipping value. The data type must be float16, float32
-                                       or None. Default: ``None`` .
+        max_norm (Union[float, None]): A maximum clipping value. The data type must be float16, float32.
+            Default: ``None`` .
         sparse (bool): Using sparse mode. When 'target' is set to ``'CPU'`` , 'sparse' has to be true.
             Default: ``True`` .
         operator (str): The pooling method for the features in one field. Support ``'SUM'`` , ``'MEAN'`` and
             ``'MAX'`` . Default: ``'SUM'`` .
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Inputs:
         - **input_indices** (Tensor) - The shape of tensor is :math:`(batch\_size, seq\_length)`.
@@ -529,12 +533,12 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
         TypeError: If `vocab_size` or `embedding_size` or `field_size` is not an int.
         TypeError: If `sparse` is not a bool or `feature_num_list` is not a tuple.
         ValueError: If `vocab_size` or `embedding_size` or `field_size` is less than 1.
-        ValueError: If `target` is neither 'CPU' nor 'DEVICE'.
-        ValueError: If `slice_mode` is not one of 'batch_slice', 'field_slice', 'table_row_slice',
-                    'table_column_slice'.
-        ValueError: If `sparse` is False and `target` is 'CPU'.
-        ValueError: If `slice_mode` is 'field_slice' and `feature_num_list` is None.
-        ValueError: If `operator` is not one of 'SUM', 'MAX', 'MEAN'.
+        ValueError: If `target` is neither ``'CPU'`` nor ``'DEVICE'``.
+        ValueError: If `slice_mode` is not one of ``'batch_slice'``, ``'field_slice'``, ``'table_row_slice'``,
+                    ``'table_column_slice'`` .
+        ValueError: If `sparse` is False and `target` is ``'CPU'`` .
+        ValueError: If `slice_mode` is ``'field_slice'`` and `feature_num_list` is None.
+        ValueError: If `operator` is not one of ``'SUM'``, ``'MAX'``, ``'MEAN'`` .
 
     Supported Platforms:
         ``Ascend`` ``GPU``
@@ -555,10 +559,11 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
     OPERATOR_MAX = 'MAX'
 
     def __init__(self, vocab_size, embedding_size, field_size, param_init='normal', target='CPU',
-                 slice_mode='batch_slice', feature_num_list=None, max_norm=None, sparse=True, operator='SUM'):
+                 slice_mode='batch_slice', feature_num_list=None, max_norm=None, sparse=True, operator='SUM',
+                 dtype=mstype.float32):
         """Initialize MultiFieldEmbeddingLookup."""
         super(MultiFieldEmbeddingLookup, self).__init__(vocab_size, embedding_size, param_init, target,
-                                                        slice_mode, feature_num_list, max_norm, sparse)
+                                                        slice_mode, feature_num_list, max_norm, sparse, dtype=dtype)
         self.field_size = Validator.check_positive_int(field_size, 'field_size', self.cls_name)
         self.operator = operator
 
@@ -622,8 +627,9 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
                 self.inf_add.shard(((1, 1, get_group_size()), (1, 1, 1)))
         else:
             if is_auto_parallel:
-                raise ValueError("For '{}', the 'slice_mode' must be in ['table_row_slice', 'batch_slice' and \
-                                       'table_column_slice'], but got {}".format(self.cls_name, str(slice_mode)))
+                raise ValueError(
+                    f"For '{self.cls_name}', the 'slice_mode' must be in ['table_row_slice', 'batch_slice' "
+                    f"and 'table_column_slice'], but got {str(slice_mode)}.")
 
         # Min value for fp32
         self.negative_inf_value = -3.402823466E+38

mindspore/nn/layer/flash_attention.py

@@ -17,12 +17,13 @@ A FlashAttention Layer.
 """
 import math
 
-import mindspore.numpy as mnp
-from mindspore import ops
-from mindspore.common import dtype as mstype
+import mindspore.common.dtype as mstype
 from mindspore.common.tensor import Tensor
+from mindspore import ops
 from mindspore.nn.cell import Cell
 from mindspore.ops._op_impl._custom_op.flash_attention.flash_attention_impl import get_flash_attention
+from mindspore.ops.operations.nn_ops import FlashAttentionScore
+from mindspore._c_expression import MSContext
 
 __all__ = ['FlashAttention']
 
@@ -56,14 +57,15 @@ class FlashAttention(Cell):
             Default True
         alibi(bool): This parameter indicates whether the flashattention supports the Alibi.
             Default: False
+        use_mqa(bool): Using MHA if True, only take effect under 910B. Default: False.
 
 
     Inputs:
       - **query** (Tensor) - Tensor query (:class:`mstype.fp16` [batch_size, head_num, seq_length, head_dim])
       - **key** (Tensor) - Tensor key (:class:`mstype.fp16` [batch_size, head_num, seq_length, head_dim])
       - **value** (Tensor) - Tensor value (:class:`mstype.fp16` [batch_size, head_num, seq_length, head_dim])
-      - **attention_mask** (Tensor) - Float Tensor the mask of (:class:`mstype.fp16` [batch_size, seq_length,
-          seq_length]): A matrix to pass masked information.
+      - **attention_mask** (Tensor) - Float Tensor the mask of (:class:`mstype.fp16` `mstype.uint8`
+        [batch_size, seq_length, seq_length]): A matrix to pass masked information.
 
     Outputs:
         A Tensor. The output of the attention with shape [batch_size, head_num, seq_length, head_dim]
@@ -92,6 +94,7 @@ class FlashAttention(Cell):
 
     def __init__(self,
                  head_dim,
+                 head_num,
                  dropout_rate=0.0,
                  prev_block_num=65536,
                  next_block_num=65536,
@@ -100,27 +103,63 @@ class FlashAttention(Cell):
                  mp=1,
                  high_precision=False,
                  have_attention_mask_batch=True,
-                 alibi=False
+                 alibi=False,
+                 use_mqa=False
                  ):
         super(FlashAttention, self).__init__()
 
-        self.flash_attention = get_flash_attention(
-            prev_block_num=prev_block_num,
-            next_block_num=next_block_num,
-            tiling_stgy_name=tiling_stgy_name,
-            high_precision=high_precision
-        )
-        self.flash_attention.add_prim_attr("primitive_target", "Ascend")
         scaling_constant = math.sqrt(head_dim)
-        if scaling_constant != 0:
-            self.scale_factor = Tensor([1. / scaling_constant], dtype=mstype.float16)
-        else:
+        if scaling_constant == 0:
             raise ValueError("the scaling constant must not be 0.")
-        self.dim_mask = Tensor([1 for _ in range(head_dim)], dtype=mstype.int8)
-        self.scale_mul = ops.Mul().shard(((dp, mp, 1, 1), (1,)))
         self.dropout_rate = dropout_rate
-        self.have_attention_mask_batch = have_attention_mask_batch
-        self.alibi = alibi
+        self.is_910A = MSContext.get_instance().get_ascend_soc_version() == "ascend910"
+        if self.is_910A:
+            self.scale_factor = Tensor([1. / math.sqrt(scaling_constant)], dtype=mstype.float16)
+            self.scale_mul = ops.Mul().shard(((dp, mp, 1, 1), (1,)))
+            self.ones = ops.Ones()
+            self.dim_mask = Tensor([1 for _ in range(head_dim)], dtype=mstype.int8)
+            self.have_attention_mask_batch = have_attention_mask_batch
+            self.alibi = alibi
+            self.flash_attention = get_flash_attention(
+                prev_block_num=prev_block_num,
+                next_block_num=next_block_num,
+                tiling_stgy_name=tiling_stgy_name,
+                high_precision=high_precision
+            )
+            self.flash_attention.add_prim_attr("primitive_target", "Ascend")
+            fa_strategies = ((dp, mp, 1, 1),
+                             (dp, mp, 1, 1),
+                             (dp, mp, 1, 1))
+            self.shard(fa_strategies)
+        else:
+            if alibi:
+                raise ValueError(f"When soc_version is not Ascend910A, alibi must be False")
+            self.transpose_4d_pre = ops.Transpose().shard(((dp, mp, 1, 1),))
+            self.transpose_4d_post = ops.Transpose().shard(((dp, 1, mp, 1),))
+            self.reshape = ops.Reshape()
+            self.zeros_like = ops.ZerosLike().shard(((dp, mp, 1, 1),))
+            self.zeros = ops.Zeros()
+            self.attn_cast = ops.Cast()
+            if use_mqa:
+                fa_strategies = ((dp, mp, 1, 1),
+                                 (dp, 1, 1, 1),
+                                 (dp, 1, 1, 1),
+                                 (dp, 1, 1, 1))
+            else:
+                fa_strategies = ((dp, mp, 1, 1),
+                                 (dp, mp, 1, 1),
+                                 (dp, mp, 1, 1),
+                                 (dp, 1, 1, 1))
+            if dropout_rate > 1e-5:
+                fa_strategies += ((dp, mp, 1, 1),)
+            self.flash_attention = FlashAttentionScore(head_num=head_num, pre_tokens=prev_block_num,
+                                                       next_tokens=next_block_num,
+                                                       keep_prob=1 - dropout_rate,
+                                                       scale_value=1. / scaling_constant,
+                                                       inner_precise=0 if high_precision else 1,
+                                                       input_layout="BNSD").shard(fa_strategies)
+
+        self.dropout_rate = dropout_rate
         if self.dropout_rate > 1e-5:
             self.keep_prob = Tensor(1 - self.dropout_rate, dtype=mstype.float16)
             self.fill_v2 = ops.FillV2().shard(((dp, mp, 1, 1), ()))
@@ -136,49 +175,49 @@ class FlashAttention(Cell):
                                   such as MatMul. Default: None.
         :return:
         """
-        if in_strategy is not None:
-            shard_stgy = list(in_strategy)
-            shard_stgy.insert(3, (1,))  # dim_mask
-            shard_stgy = tuple(shard_stgy)
-        else:
-            # default: dp=1, mp=1, construct inputs only contain query, key, value
-            shard_stgy = (
-                (1, 1, 1, 1),
-                (1, 1, 1, 1),
-                (1, 1, 1, 1),
-                (1,),  # dim_mask
-            )
-        self.flash_attention.shard(shard_stgy)
-        dp = shard_stgy[0][0]
-        mp = shard_stgy[0][1]
-        self.flash_attention.add_prim_attr("dev_matrix_shape", [dp, mp, 1, 1])
-        inputs_tensor_map = [
-            [3, 2, 1, 0],
-            [3, 2, 1, 0],
-            [3, 2, 1, 0],
-            [-1]
-        ]
-        if self.have_attention_mask_batch:
-            inputs_tensor_map.append([3, 1, 0])
-        else:
-            inputs_tensor_map.append([-1, 1, 0])
+        if self.is_910A:
+            if in_strategy is None:
+                # default: dp=1, mp=1, construct inputs only contain query, key, value
+                in_strategy = (
+                    (1, 1, 1, 1),
+                    (1, 1, 1, 1),
+                    (1, 1, 1, 1),
+                )
+            self.flash_attention.shard(in_strategy)
+            dp = in_strategy[0][0]
+            mp = in_strategy[0][1]
+            self.flash_attention.add_prim_attr("dev_matrix_shape", [dp, mp, 1, 1])
+            inputs_tensor_map = [
+                [3, 2, 1, 0],
+                [3, 2, 1, 0],
+                [3, 2, 1, 0],
+            ]
+            if self.have_attention_mask_batch:
+                inputs_tensor_map.append([3, 1, 0])
+            else:
+                inputs_tensor_map.append([-1, 1, 0])
 
-        # dropout_mask
-        if self.dropout_rate > 1e-5:
-            inputs_tensor_map.append([3, 2, 1, 0])
+            input_empty_args_num = 2
+            # dropout_mask
+            if self.dropout_rate > 1e-5:
+                input_empty_args_num -= 1
+                inputs_tensor_map.append([3, 2, 1, 0])
 
-        if self.alibi:
-            inputs_tensor_map.append([3, 2, 1, 0])
+            if self.alibi:
+                input_empty_args_num -= 1
+                inputs_tensor_map.append([3, 2, 1, 0])
 
-        self.flash_attention.add_prim_attr("inputs_tensor_map", inputs_tensor_map)
+            self.flash_attention.add_prim_attr("inputs_tensor_map", inputs_tensor_map)
 
-        self.flash_attention.add_prim_attr("outputs_tensor_map", [
-            [3, 2, 1, 0],  # O
-            [3, 2, 1],  # L
-            [3, 2, 1]  # M
-        ])
-        self.flash_attention.add_prim_attr("as_loss_divisor", 0)
-        self.flash_attention.add_prim_attr("empty_mirror_ops", 1)
+            self.flash_attention.add_prim_attr("outputs_tensor_map", [
+                [3, 2, 1, 0],  # O
+                [3, 2, 1],  # L
+                [3, 2, 1]  # M
+            ])
+            self.flash_attention.add_prim_attr("as_loss_divisor", 0)
+            self.flash_attention.add_prim_attr("empty_mirror_ops", input_empty_args_num)
+        else:
+            self.flash_attention.shard(in_strategy)
 
     def construct(self, query, key, value, attn_mask=None, alibi_mask=None):
         """FlashAttention forward
@@ -189,35 +228,49 @@ class FlashAttention(Cell):
         :param alibi_mask: [bsz, head_num, 1, seq_len], if not None
         :return: output          [bsz, head_num, seq_len, head_dim]
         """
-        query = self.scale_mul(query, self.scale_factor)
         bsz, head_num, seq_len, head_dim = query.shape
-        _, k_head_num, k_seq_len, _ = key.shape
-        _, v_head_num, v_seq_len, _ = value.shape
-        if head_num != k_head_num or head_num != v_head_num:
-            raise ValueError(
-                "the head_num of query, key and value must be the same, "
-                "If different head_num are used, users need to change themselves to be same by tile.")
-        if seq_len % 16 != 0 or k_seq_len % 16 != 0 or k_seq_len != v_seq_len:
-            raise ValueError(
-                "query, key, value seq_len must be a multiple of 16, and key seq_len, value seq_len must be the same.")
-        if self.dropout_rate > 1e-5:
-            drop_mask_bits = self.drop_gen_mask((bsz, head_num, seq_len, seq_len), self.keep_prob)
-            tensor_shape = Tensor((bsz, head_num, seq_len, seq_len), mstype.int32)
-            ones = self.fill_v2(tensor_shape, self.tensor_one)
-            ones = self.depend(ones, query)
-            drop_mask = self.do_dropout(ones, drop_mask_bits, self.keep_prob)
-        else:
-            drop_mask = None
-        if head_dim > 304:
-            raise ValueError(
-                "the head_dim must be less than 304, otherwise the ub would be OOM.")
-        if head_dim % 16 != 0:
-            padding_size = 16 - head_dim % 16
-            query = mnp.pad(query, ((0, 0), (0, 0), (0, 0), (0, padding_size)), constant_values=0)
-            key = mnp.pad(key, ((0, 0), (0, 0), (0, 0), (0, padding_size)), constant_values=0)
-            value = mnp.pad(value, ((0, 0), (0, 0), (0, 0), (0, padding_size)), constant_values=0)
-            output, _, _ = self.flash_attention(query, key, value, self.dim_mask, attn_mask, drop_mask, alibi_mask)
-            output = ops.slice(output, [0, 0, 0, 0], [bsz, head_num, seq_len, head_dim])
+        if self.is_910A:
+            _, k_head_num, k_seq_len, _ = key.shape
+            _, v_head_num, v_seq_len, _ = value.shape
+            if head_num != k_head_num or head_num != v_head_num:
+                raise ValueError(
+                    "the head_num of query, key and value must be the same, "
+                    "If different head_num are used, users need to change themselves to be same by tile.")
+            if seq_len % 16 != 0 or k_seq_len % 16 != 0 or k_seq_len != v_seq_len:
+                raise ValueError(
+                    "query, key, value seq_len must be a multiple of 16, "
+                    "and the seq_len between key and value must be equal.")
+            # 910A -- FlashAttentionPrimtive
+            if head_dim > 304:
+                raise ValueError(
+                    "the head_dim must be less than 304, otherwise the ub would be OOM.")
+            if self.dropout_rate > 1e-5:
+                drop_mask_bits = self.drop_gen_mask((bsz, head_num, seq_len, seq_len), self.keep_prob)
+                tensor_shape = Tensor((bsz, head_num, seq_len, seq_len), mstype.int32)
+                ones = self.fill_v2(tensor_shape, self.tensor_one)
+                ones = self.depend(ones, query)
+                drop_mask = self.do_dropout(ones, drop_mask_bits, self.keep_prob)
+            else:
+                drop_mask = None
+            query = self.scale_mul(query, self.scale_factor)
+            key = self.scale_mul(key, self.scale_factor)
+            attn_mask = self.cast(attn_mask, mstype.float16)
+            output, _, _ = self.flash_attention(query, key, value, attn_mask, drop_mask, alibi_mask)
         else:
-            output, _, _ = self.flash_attention(query, key, value, self.dim_mask, attn_mask, drop_mask, alibi_mask)
+            # 910B -- FlashAttentionScore
+            if self.dropout_rate > 1e-5:
+                drop_mask_bits = self.reshape(self.drop_gen_mask((bsz, head_num, seq_len, seq_len), self.keep_prob),
+                                              (bsz, head_num, seq_len, seq_len // 8))
+            else:
+                drop_mask_bits = None
+            # (B, S, S) -> (B, 1, S, S)
+            attn_mask = self.cast(self.reshape(attn_mask, (bsz, 1, seq_len, seq_len)), mstype.uint8)
+            output, _, _ = self.flash_attention(query,
+                                                key,
+                                                value,
+                                                attn_mask,
+                                                drop_mask_bits,
+                                                None,
+                                                None,
+                                                None)
         return output

mindspore/nn/layer/image.py

@@ -83,17 +83,17 @@ class ImageGradients(Cell):
         _check_input_4d(F.shape(images), "images", self.cls_name)
         batch_size, depth, height, width = P.Shape()(images)
         if height == 1:
-            dy = P.Fill()(P.DType()(images), (batch_size, depth, 1, width), 0)
+            dy = F.fill(P.DType()(images), (batch_size, depth, 1, width), 0)
         else:
             dy = images[:, :, 1:, :] - images[:, :, :height - 1, :]
-            dy_last = P.Fill()(P.DType()(images), (batch_size, depth, 1, width), 0)
+            dy_last = F.fill(P.DType()(images), (batch_size, depth, 1, width), 0)
             dy = P.Concat(2)((dy, dy_last))
 
         if width == 1:
-            dx = P.Fill()(P.DType()(images), (batch_size, depth, height, 1), 0)
+            dx = F.fill(P.DType()(images), (batch_size, depth, height, 1), 0)
         else:
             dx = images[:, :, :, 1:] - images[:, :, :, :width - 1]
-            dx_last = P.Fill()(P.DType()(images), (batch_size, depth, height, 1), 0)
+            dx_last = F.fill(P.DType()(images), (batch_size, depth, height, 1), 0)
             dx = P.Concat(3)((dx, dx_last))
         return dy, dx
 
@@ -571,7 +571,8 @@ class PixelShuffle(Cell):
     <https://arxiv.org/abs/1609.05158>`_ .
 
     Typically, the input is of shape :math:`(*, C \times r^2, H, W)` , and the output is of shape
-    :math:`(*, C, H \times r, W \times r)`, where r is an upscale factor and * is zero or more batch dimensions.
+    :math:`(*, C, H \times r, W \times r)`,
+    where :math:`r` is an upscale factor and :math:`*` is zero or more batch dimensions.
 
     Note:
         The dimension of input Tensor on Ascend should be less than 7.
@@ -621,7 +622,8 @@ class PixelUnshuffle(Cell):
     <https://arxiv.org/abs/1609.05158>`_ .
 
     Typically, the input is of shape :math:`(*, C, H \times r, W \times r)` , and the output is of shape
-    :math:`(*, C \times r^2, H, W)` , where r is a downscale factor and * is zero or more batch dimensions.
+    :math:`(*, C \times r^2, H, W)` ,
+    where :math:`r` is a downscale factor and :math:`*` is zero or more batch dimensions.
 
     Args:
         downscale_factor (int): factor to unshuffle the input, and is a positive integer.