mindspore 2.1.0__cp39-none-any.whl → 2.2.11__cp39-none-any.whl

mindspore/ops/operations/_grad_ops.py

@@ -390,7 +390,7 @@ class Conv2DBackpropFilter(Primitive):
         stride (tuple): The stride to be applied to the convolution filter. Default: (1, 1).
         dilation (tuple): Specifies the dilation rate to be used for the dilated convolution. Default: (1, 1, 1, 1).
         group (int): Splits input into groups. Default: 1.
-        data_format (str) - The format of input and output data. It should be 'NHWC' or 'NCHW'，\
+        data_format (str) - The format of input and output data. It should be 'NHWC' or 'NCHW', \
             default is 'NCHW'.
 
     Returns:
@@ -636,7 +636,7 @@ class EinsumGrad(PrimitiveWithInfer):
 
     @prim_attr_register
     def __init__(self, equation):
-        self.add_prim_attr('equation', equation)
+        pass
 
     def infer_shape(self, x_shapes, dout_shape):
         out_shape = ()
@@ -1521,9 +1521,11 @@ class LSTMGrad(Primitive):
     """Computes the data and weight gradients of LSTM."""
 
     @prim_attr_register
-    def __init__(self, input_size, hidden_size, num_layers, has_bias, bidirectional, dropout):
+    def __init__(self, input_size, hidden_size, num_layers, has_bias, bidirectional, dropout, proj_size=0):
         self.input_size = validator.check_positive_int(input_size, 'input_size', self.name)
         self.hidden_size = validator.check_positive_int(hidden_size, 'hidden_size', self.name)
+        self.proj_size = validator.check_int_range(proj_size, 0, hidden_size, validator.INC_LEFT,
+                                                   'proj_size', self.name)
         self.num_layers = validator.check_positive_int(num_layers, 'num_layers', self.name)
         self.has_bias = validator.check_value_type('has_bias', has_bias, (bool,), self.name)
         self.bidirectional = validator.check_value_type('bidirectional', bidirectional, (bool,), self.name)
@@ -2573,7 +2575,12 @@ class MultilabelMarginLossGrad(Primitive):
     Compute the gradients of MultilabelMarginLoss operation.
 
     Args:
-        reduction (str): Apply specific reduction method to the output: 'none', 'mean', 'sum'. Default: "mean".
+        reduction (str, optional): Apply specific reduction method to the output: ``'none'`` , ``'mean'`` ,
+            ``'sum'`` . Default: ``'mean'`` .
+
+            - ``'none'``: no reduction will be applied.
+            - ``'mean'``: compute and return the mean of elements in the output.
+            - ``'sum'``: the output elements will be summed.
 
     Inputs:
         - **y_grad** (Tensor) - The gradients of loss to output of MultilabelMarginLoss function, with
@@ -2595,7 +2602,7 @@ class MultilabelMarginLossGrad(Primitive):
         TypeError: If dtype of `y_grad` is not the same as `x`.
         ValueError: If length of shape of `x` is neither 1 nor 2.
         ValueError: If shape of `x` is not the same as `target`.
-        ValueError: If `reduction` is not one of 'none', 'mean', 'sum'.
+        ValueError: If `reduction` is not one of ``'none'``, ``'mean'``, ``'sum'``.
         ValueError: If shape of `y_grad` is not the same as forward output `y`.
 
     Supported Platforms:
@@ -2862,7 +2869,9 @@ class Dilation2DBackpropFilter(Primitive):
         self.pad_mode = validator.check_string(self.pad_mode, ["SAME", "VALID", 'same', "valid"], "pad_mode", self.name)
         self.add_prim_attr("pad_mode", self.pad_mode.upper())
         self.stride = _check_format_stride_or_dilation("stride", stride, self.name, self.data_format)
-        if self.stride[2] < 1 or self.stride[2] > 255 or self.stride[3] < 1 or self.stride[3] > 255:
+        def is_in_range(x):
+            return 1 <= x <= 255
+        if not is_in_range(self.stride[2]) or not is_in_range(self.stride[3]):
             raise ValueError(f"For '{self.name}', size of stride is not supported, "
                              f'stride should be in the range of [1, 255], '
                              f'but got stride_h: `{self.stride[2]}`, stride_w: `{self.stride[3]}`.')
@@ -2917,7 +2926,12 @@ class MultiMarginLossGrad(Primitive):
     Args:
         p (int): Optional. The norm degree for pairwise distance.Should be 1 or 2. Default: 1.
         margin (float): Optional. A parameter to change pairwise distance. Default: 1.0.
-        reduction (str): Apply specific reduction method to the output: 'none', 'mean', 'sum'. Default: "mean".
+        reduction (str, optional): Apply specific reduction method to the output: ``'none'`` , ``'mean'`` ,
+            ``'sum'`` . Default: ``'mean'`` .
+
+            - ``'none'``: no reduction will be applied.
+            - ``'mean'``: compute and return the weighted mean of elements in the output.
+            - ``'sum'``: the output elements will be summed.
 
     Inputs:
         - **y_grad** (Tensor) - If it's not a scalar, the shape of 'y_grad' :math:`(N, C)`.
@@ -3818,3 +3832,53 @@ class WKVGrad(Primitive):
         """Initialize WKVGrad."""
         self.init_prim_io_names(inputs=["time_first", "time_decay", "key", "value", "gy"],
                                 outputs=["gw", "gu", "gk", "gv"])
+
+
+class FlashAttentionScoreGrad(Primitive):
+    r"""
+    Calculates the gradient of FlashAttentionScore operation.
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Supported Platforms:
+        ``Ascend``
+    """
+    @prim_attr_register
+    def __init__(self, head_num, keep_prob=1.0, scale_value=1.0, pre_tokens=65536, next_tokens=65536, inner_precise=1,
+                 input_layout='BSH', sparse_mode=0):
+        """Initialize FlashAttentionScoreGrad."""
+        validator.check_value_type('head_num', head_num, [int], self.name)
+        validator.check_value_type('keep_prob', keep_prob, [int, float], self.name)
+        validator.check_float(keep_prob, 0.0, validator.GE, "keep_prob", self.name)
+        validator.check_float(keep_prob, 1.0, validator.LE, "keep_prob", self.name)
+        validator.check_value_type('scale_value', scale_value, [float], self.name)
+        validator.check_value_type('pre_tokens', pre_tokens, [int], self.name)
+        validator.check_value_type('next_tokens', next_tokens, [int], self.name)
+        validator.check_value_type('inner_precise', inner_precise, [int], self.name)
+        validator.check_value_type('sparse_mode', sparse_mode, [int], self.name)
+        if inner_precise not in [0, 1]:
+            raise ValueError(f"Attribute 'inner_precise' must be either 0 or 1, but got {inner_precise}")
+        validator.check_value_type('input_layout', input_layout, [str], self.name)
+        if input_layout not in ["BSH", "BNSD"]:
+            raise ValueError(f"Attribute 'input_layout' must be either 'BSH' or 'BNSD', but got {input_layout}")
+        self.init_prim_io_names(inputs=['query', 'key', 'value', 'dy', 'pse_shift', 'drop_mask', "padding_mask",
+                                        'attn_mask', 'softmax_max', 'softmax_sum', 'softmax_out', 'attention_in',
+                                        'prefix'],
+                                outputs=['dq', 'dk', 'dv', 'dpse'])
+
+
+class RmsNormGrad(Primitive):
+    r"""
+    Calculates the gradient of RmsNorm operation.
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Supported Platforms:
+        ``Ascend``
+    """
+
+    @prim_attr_register
+    def __init__(self):
+        """Initialize RmsNormGrad."""
+        self.init_prim_io_names(inputs=["dy", "x", "rstd", "gamma"],
+                                outputs=["dx", "dgamma"])

mindspore/ops/operations/_inner_ops.py

@@ -23,16 +23,17 @@ from mindspore.common._stub_tensor import StubTensor
 from mindspore.ops import composite as C
 from mindspore.ops.operations.array_ops import Cast
 from mindspore.ops.operations._scalar_ops import bit_or, bit_and
+from mindspore.ops.operations.comm_ops import ReduceOp
 from mindspore.ops import signature as sig
 from mindspore.ops.operations.math_ops import _infer_shape_reduce
-from mindspore.ops.primitive import PrimitiveWithCheck, PrimitiveWithInfer, prim_attr_register, Primitive, _run_op
-from mindspore import context
+from mindspore.ops.primitive import PrimitiveWithCheck, PrimitiveWithInfer, prim_attr_register, Primitive, \
+    _run_op, _check_contains_variable
 from mindspore._c_expression import Tensor as Tensor_
 from mindspore._c_expression import typing
 from mindspore import _checkparam as validator
 from mindspore.common import dtype as mstype
 from mindspore.common.parameter import Parameter
-from mindspore.communication.management import GlobalComm
+from mindspore.communication.management import GlobalComm, get_rank
 from mindspore.common.api import _pynative_executor
 from mindspore.common._register_for_adapter import ms_adapter_registry
 from mindspore import ops
@@ -74,11 +75,11 @@ class ExtractImagePatches(Primitive):
             - valid: Means that the taken patch area must be completely covered in the original image.
 
     Inputs:
-        - **input_x** (Tensor) - A 4-D tensor whose shape is :math:`(in_batch, in_depth, in_row, in_col)`.
+        - **input_x** (Tensor) - A 4-D tensor whose shape is :math:`(in\_batch, in\_depth, in\_row, in\_col)`.
 
     Outputs:
         Tensor, a 4-D tensor whose data type is same as 'input_x', and the shape
-        is :math:`(out_batch, out_depth, out_row, out_col)`,where the out_batch is the same as the in_batch
+        is :math:`(out\_batch, out\_depth, out\_row, out\_col)`,where the out_batch is the same as the in_batch
         and
 
         .. math::
@@ -121,7 +122,6 @@ class ExtractImagePatches(Primitive):
         validator.check_value_type('padding', padding, [str], self.name)
         self.padding = validator.check_string(padding.upper(), ['VALID', 'SAME'], 'padding', self.name)
         self.add_prim_attr("padding", self.padding)
-        self.is_ge = context.get_context("enable_ge")
 
 
 class Quant(PrimitiveWithInfer):
@@ -167,6 +167,7 @@ class Quant(PrimitiveWithInfer):
         self.sqrt_mode = validator.check_value_type("sqrt_mode", sqrt_mode, [bool], self.name)
         self.round_mode = validator.check_string(round_mode, ["Round", "Floor", "Ceil", "Trunc"],
                                                  "round_mode", self.name)
+        self.add_prim_attr("dst_type", mstype.int8)
 
     def infer_shape(self, x_shape):
         return x_shape
@@ -174,7 +175,7 @@ class Quant(PrimitiveWithInfer):
     def infer_dtype(self, x_type):
         validator.check_subclass("input_x", x_type, mstype.tensor_type, self.name)
         validator.check_type_name("input_x", x_type, [mstype.float16, mstype.float32], self.name)
-        return mstype.int8
+        return self.get_attr_dict()['dst_type']
 
 
 class Lamb(PrimitiveWithInfer):
@@ -491,7 +492,7 @@ class Receive(PrimitiveWithInfer):
         self.dtype = dtype
         self.group = group
         self.add_prim_attr("no_eliminate", True)
-        valid_type = [mstype.float16, mstype.float32, mstype.int32, mstype.int8, mstype.uint8]
+        valid_type = [mstype.float16, mstype.bfloat16, mstype.float32, mstype.int32, mstype.int8, mstype.uint8]
         args = {"dtype": dtype}
         validator.check_scalar_or_tensor_types_same(args, valid_type, self.name)
 
@@ -502,6 +503,109 @@ class Receive(PrimitiveWithInfer):
         return self.get_attr_dict()['dtype']
 
 
+class Reduce(PrimitiveWithInfer):
+    """
+    Reduces tensor across the processes in the specified communication group.
+
+    Note:
+        Only process with destination rank receives the reduced output.
+        Other processes only get a tensor with shape [1], which has no mathematical meaning.
+
+    Args:
+        dest_rank (int): Specifies the rank of the process that receives the reduced output.
+        op (str, optional): Specifies an operation used for element-wise reductions, like sum, prod, max, and min.
+                On the CPU, only 'sum' is supported. Default: ``ReduceOp.SUM`` .
+        group (str, optional): The communication group to work on.
+            Default: "hccl_world_group" on Ascend, "nccl_world_group" on GPU.
+
+    Inputs:
+        - **input_x** (Tensor) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
+
+    Examples:
+        >>> import mindspore.ops as ops
+        >>> import mindspore.nn as nn
+        >>> from mindspore.communication import init
+        >>> from mindspore import Tensor
+        >>> import numpy as np
+        >>> # Launch 4 processes.
+        >>> init()
+        >>> class ReduceNet(nn.Cell):
+        >>>     def __init__(self):
+        >>>         super(Net, self).__init__()
+        >>>         self.reduce = ops.Reduce(dest_rank=1)
+        >>>
+        >>>     def construct(self, x):
+        >>>         out = self.reduce(x)
+        >>>         return out
+        >>> input = Tensor(np.ones([2, 8]).astype(np.float32))
+        >>> net = ReduceNet()
+        >>> output = net(input)
+        >>> print(output)
+        Process with rank 1: [[4. 4. 4. 4. 4. 4. 4. 4.]
+                             [4. 4. 4. 4. 4. 4. 4. 4.]],
+        Other proesses: [0.].
+    """
+
+    @prim_attr_register
+    def __init__(self, dest_rank, op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_GROUP):
+        self.dest_rank = dest_rank
+        self.op = op
+        self.group = group
+
+    def infer_shape(self, x_shape):
+        # The process with dest_rank returns the reduced output.
+        # Other processes only gets a tensor with shape [1], which has no mathematical meaning.
+        if self.dest_rank == get_rank():
+            return x_shape
+        return [1]
+
+    def infer_dtype(self, x_dtype):
+        return x_dtype
+
+
+class Barrier(PrimitiveWithInfer):
+    """
+    Synchronizes all processes in the specified group.
+
+    Note:
+        After calling this collective operator,
+        this process will be blocked until all other processes in the group call this operator.
+
+    Args:
+        group (str, optional): The communication group to work on.
+            Default: "hccl_world_group" on Ascend, "nccl_world_group" on GPU.
+
+    Examples:
+        >>> import mindspore.ops as ops
+        >>> import mindspore.nn as nn
+        >>> from mindspore.communication import init
+        >>> from mindspore import Tensor
+        >>> import numpy as np
+        >>> # Launch 4 processes.
+        >>> init()
+        >>> class BarrierNet(nn.Cell):
+        >>>     def __init__(self):
+        >>>         super(Net, self).__init__()
+        >>>         self.barrier = ops.Barrier()
+        >>>
+        >>>     def construct(self):
+        >>>         self.barrier()
+        >>> net = BarrierNet()
+        >>> net()
+    """
+
+    @prim_attr_register
+    def __init__(self, group=GlobalComm.WORLD_COMM_GROUP):
+        self.group = group
+        self.add_prim_attr("side_effect_mem", True)
+
+    def infer_shape(self):
+        return [1]
+
+    def infer_dtype(self):
+        return mstype.float32
+
+
 class MatrixSetDiag(PrimitiveWithInfer):
     r"""
     Modifies the batched diagonal part of a batched tensor.
@@ -1843,16 +1947,32 @@ class Format(PrimitiveWithInfer):
     def __init__(self):
         self.init_prim_io_names(inputs=['string', 'args'], outputs=['string'])
 
+
     def __infer__(self, str_, *var):
-        str_value = str_["value"]
+        def check_variable(str_, var):
+            if _check_contains_variable(str_['dtype'], str_['value']):
+                return True
+
+            for item in var:
+                if _check_contains_variable(item['dtype'], item['value']):
+                    return True
+            return False
+
+
+        if check_variable(str_, var):
+            return {'dtype': mstype.string, 'shape': [], 'value': None}
+
+
+        str_value = str_['value']
+        kwargs = dict()
         var_value = list()
-        if str_value is None and str_["dtype"] is not None:
-            raise ValueError("str.format not support to input a variable.")
+
         for item in var:
-            if item["value"] is None and item["dtype"] is not None:
-                raise ValueError("str.format not support to input a variable.")
+            if isinstance(item["dtype"], typing.Keyword):
+                kwargs.update(item["value"])
             var_value.append(item["value"])
-        value = str_value.format(*var_value)
+
+        value = str_value.format(*var_value, **kwargs)
         return {'dtype': mstype.string, 'shape': [], 'value': value}
 
 
@@ -2027,13 +2147,14 @@ class ClipByNorm(PrimitiveWithInfer):
     @prim_attr_register
     def __init__(self, axis=None):
         """Initialize ClipByNorm"""
+        self.axis_str = 'axis'
         self.axis = () if axis is None else axis
-        validator.check_value_type('axis', self.axis, [int, tuple, list], self.name)
+        validator.check_value_type(self.axis_str, self.axis, [int, tuple, list], self.name)
         axis_check = self.axis if isinstance(self.axis, Iterable) else (self.axis,)
         for i, value in enumerate(axis_check):
             validator.check_value_type('axis[%d]' % i, value, [int], self.name)
-        self.init_attrs['axis'] = self.axis
-        self.add_prim_attr('axis', self.axis)
+        self.init_attrs[self.axis_str] = self.axis
+        self.add_prim_attr(self.axis_str, self.axis)
         self.init_prim_io_names(inputs=['x', 'clip_norm'], outputs=['output'])
 
     def infer_shape(self, x_shape, clip_norm_shape):
@@ -2588,3 +2709,215 @@ class IsConstant(Primitive):
 
     def __call__(self, x):
         return True
+
+
+class SelectView(Primitive):
+    r"""
+        Select tensor of view
+    """
+
+    @prim_attr_register
+    def __init__(self):
+        self.init_prim_io_names(inputs=['input_tensor', 'input_indices', 'axis'], outputs=['output'])
+
+
+class CopyWithSlice(Primitive):
+    r"""
+        Copy data to discontinuous tensor
+    """
+    @prim_attr_register
+    def __init__(self):
+        self.add_prim_attr('side_effect_mem', True)
+        self.init_prim_io_names(inputs=['x', 'y'], outputs=['x'])
+
+
+class FFN(Primitive):
+    r"""
+    The FFN computation is similar to Feed-Forward Network, it contains matmul + gelu + matmul.
+
+    Args:
+        activation (string): The activation type, set to 'fastgelu' or 'gelu'.
+            Only support 'fastgelu' for now. Default: "fastgelu".
+        inner_precise (int): The precise mode, set to 0 for high precision or 1 for high performance.
+            Only support 1 for now. Default: 0.
+
+    Inputs:
+        - **x** (Tensor) - The input tensor with data type of int8, float16.
+          Input tensor of shape :math:`(batch\_size * seq\_length, hidden\_size)`.
+        - **weight1** (Tensor) - The weight1 tensor with data type of float16.
+          Weight1 tensor of shape :math:`(expert\_num, hidden\_size, ffn\_hidden\_size)`.
+        - **weight2** (Tensor) - The weight2 tensor with data type of float16.
+          Weight2 tensor of shape :math:`(expert\_num, ffn\_hidden\_size, hidden\_size)`.
+        - **expert_tokens** (Tensor]) - The expert tokens tensor with data type of int64.
+          Expert tokens tensor of shape :math:`(16,)`. For example, `(2, 1, 0, .., 9)`
+          indicate that the 0th expert deals with 2 tokens, the 1th expert deals with 1 tokens,
+          the 2th expert do noting and so on.
+        - **bias1** (Tensor) - The bias1 tensor with data type of float16.
+          Bias1 tensor of shape :math:`(expert\_num, ffn\_hidden\_size)`.
+        - **bias2** (Tensor) - The bias2 tensor with data type of float16.
+          Bias2 tensor of shape :math:`(expert\_num, hidden\_size)`.
+        - **scale** (Tensor) - The scale tensor with data type of float16. Not enable now.
+        - **offset** (Tensor) - The offset tensor with data type of float16. Not enable now.
+        - **deq_scale1** (Tensor) - The deq_scale1 tensor with data type of float16. Not enable now.
+        - **deq_scale2** (Tensor) - The deq_scale2 tensor with data type of float16. Not enable now.
+
+    Outputs:
+        Tensor of shape :math:`(batch\_size * seq\_length, hidden\_size)`. With data type of float16.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> from mindspore.ops.operations import _inner_ops
+        >>> b = 4
+        >>> s = 128
+        >>> h = 1024
+        >>> h_f = 4 * h
+        >>> e = 16
+        >>> x = Tensor(np.random.randn(b * s, h).astype(np.float16))
+        >>> w1 = Tensor(np.random.randn(e, h, h_f).astype(np.float16))
+        >>> w2 = Tensor(np.random.randn(e, h_f, h).astype(np.float16))
+        >>> expert_tokens = Tensor(np.random.randn(e).astype(np.int64))
+        >>> bias1 = Tensor(np.random.randn(e, h_f).astype(np.float16))
+        >>> bias2 = Tensor(np.random.randn(e, h).astype(np.float16))
+        >>> ffn = _inner_ops.FFN("fastgelu", 1)
+        >>> output = ffn(x, w1, w2, expert_tokens, bias1, bias2)
+        >>> print(output)
+    """
+
+    @prim_attr_register
+    def __init__(self, activation, inner_precise):
+        """Initialize FFN."""
+        self.init_prim_io_names(inputs=["x", "weight1", "weight2", "expert_tokens", "bias1",
+                                        "bias2", "scale", "offset", "deq_scale1", "deq_scale2"],
+                                outputs=["y"])
+        cls_name = self.name
+        validator.check_value_type("activation", activation, [str], cls_name)
+        validator.check_value_type("inner_precise", inner_precise, [int], cls_name)
+
+
+class DecoderKVCache(Primitive):
+    r"""
+    The DecoderKVCache is used for decoding the KVCache of transformer network.
+
+    Args:
+        cache (Tensor): The cahe tensor with data type of int8, uint8, int16, uint16, float16, float32 and int32.
+            When seq_len_axis is 2, cache tensor of shape
+            :math:`(batch\_size, num_head, max\_seq\_length, hidden\_size)`.
+            When seq_len_axis is 1, cache tensor of shape
+            :math:`(batch\_size, max\_seq\_length, num_head, hidden\_size)`.
+        update (Tensor]): The tensor which is used to update the cache tensor. Same data type as cache tensor.
+            When seq_len_axis is 2, update tensor of shape
+            :math:`(batch\_size, num_head, update\_seq\_length, hidden\_size)`.
+            When seq_len_axis is 1, update tensor of shape
+            :math:`(batch\_size, update\_seq\_length, num_head, hidden\_size)`.
+        valid_seq_len (Tensor): The valid_seq_len tensor with data type of int64.
+            Valid_seq_len tensor of shape :math:`(batch\_size)`.
+        batch_index (Tensor): The batch_index tensor with data type of int64.
+            Batch_index tensor of shape :math:`(1)`. Indicate that which batch of cache tensor is going to be update.
+        seq_len_axis (int64): The seq_len_axis indicate which axis is seq_eln, set to '1' or '2'. Default: "2".
+        new_max_seq_len (Tensor): The new_max_seq_len tensor with data type of int64.
+            New_max_seq_len tensor of shape :math:`(1)`.
+            Indicate that user want to change the shape of cache tensor from
+            :math:`(batch\_size, num_head, max\_seq\_length, hidden\_size)` to
+            :math:
+            `(batch\_size * max\_seq\_length / new\_max\_seq\_length, num_head, new\_max\_seq\_length, hidden\_size)`
+            to update the cache tensor. This will not real change the shape of `cache` tensor. Not able for now.
+        cur_max_seq_len (Tensor): The new_max_seq_len tensor with data type of int64.
+            Cur_max_seq_len tensor of shape :math:`(1)`. Keep the current seq_len of cache tensor. Not abel for now.
+
+    Outputs:
+        With same data type and same shape as `cache` tensor.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> from mindspore.ops.operations import _inner_ops
+        >>> b = 4
+        >>> h = 40
+        >>> max_s = 1024
+        >>> s = 1
+        >>> d = 128
+        >>> cache = Tensor(np.random.randn(b, h, max_s, d).astype(np.float16))
+        >>> update = Tensor(np.random.randn(b, h, s, d).astype(np.float16))
+        >>> valid_seq_len = Tensor(np.random.randn(b).astype(np.int64))
+        >>> batch_index = Tensor(np.random.randn(1).astype(np.int64))
+        >>> new_max_seq_len = Tensor(np.random.randn(1).astype(np.int64))
+        >>> cur_max_seq_len = Tensor(np.random.randn(1).astype(np.int64))
+        >>> decoder_kv_cache = _inner_ops.DecoderKVCache()
+        >>> output = decoder_kv_cache(cache, update, valid_seq_len, batch_index, 2, new_max_seq_len, cur_max_seq_len)
+        >>> print(cache)
+    """
+    @prim_attr_register
+    def __init__(self):
+        """Initialize DecoderKVCache."""
+        self.init_prim_io_names(inputs=["cache", "update", "valid_seq_len", "batch_index", "seq_len_axis",
+                                        "new_max_seq_len", "cur_max_seq_len"],
+                                outputs=["out"])
+        self.add_prim_attr('side_effect_mem', True)
+
+
+class PromptKVCache(Primitive):
+    r"""
+    The PromptKVCache is used for prefill the KVCache of transformer network.
+
+    Args:
+        cache (Tensor): The cahe tensor with data type of int8, uint8, int16, uint16, float16, float32 and int32.
+            When seq_len_axis is 2, cache tensor of shape
+            :math:`(batch\_size, num_head, max\_seq\_length, hidden\_size)`.
+            When seq_len_axis is 1, cache tensor of shape
+            :math:`(batch\_size, max\_seq\_length, num_head, hidden\_size)`.
+        update (Tensor]): The tensor which is used to update the cache tensor. Same data type as cache tensor.
+            When seq_len_axis is 2, update tensor of shape
+            :math:`(batch\_size, num_head, update\_seq\_length, hidden\_size)`.
+            When seq_len_axis is 1, update tensor of shape
+            :math:`(batch\_size, update\_seq\_length, num_head, hidden\_size)`.
+        valid_seq_len (Tensor): The valid_seq_len tensor with data type of int64.
+            Valid_seq_len tensor of shape :math:`(batch\_size)`.
+        batch_index (Tensor): The batch_index tensor with data type of int64.
+            Batch_index tensor of shape :math:`(1)`. Indicate that which batch of cache tensor is going to be update.
+        seq_len_axis (int64): The seq_len_axis indicate which axis is seq_eln, set to '1' or '2'. Default: "2".
+        new_max_seq_len (Tensor): The new_max_seq_len tensor with data type of int64.
+            New_max_seq_len tensor of shape :math:`(1)`.
+            Indicate that user want to change the shape of cache tensor from
+            :math:`(batch\_size, num_head, max\_seq\_length, hidden\_size)` to
+            :math:
+            `(batch\_size * max\_seq\_length / new\_max\_seq\_length, num_head, new\_max\_seq\_length, hidden\_size)`
+            to update the cache tensor. This will not real change the shape of `cache` tensor. Not able for now.
+        cur_max_seq_len (Tensor): The new_max_seq_len tensor with data type of int64.
+            Cur_max_seq_len tensor of shape :math:`(1)`. Keep the current seq_len of cache tensor. Not abel for now.
+        align_mode (int64): indicate which axis is seq_eln, 0 is 'right', 1 is 'left'. Default: 0.
+
+    Outputs:
+        With same data type and same shape as `cache` tensor.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> from mindspore import Tensor
+        >>> from mindspore.ops.operations import _inner_ops
+        >>> b = 4
+        >>> h = 40
+        >>> max_s = 1024
+        >>> s = 256
+        >>> d = 128
+        >>> cache = Tensor(np.random.randn(b, h, max_s, d).astype(np.float16))
+        >>> update = Tensor(np.random.randn(b, h, s, d).astype(np.float16))
+        >>> valid_seq_len = Tensor(np.random.randn(b).astype(np.int64))
+        >>> batch_index = Tensor(np.random.randn(1).astype(np.int64))
+        >>> new_max_seq_len = Tensor(np.random.randn(1).astype(np.int64))
+        >>> cur_max_seq_len = Tensor(np.random.randn(1).astype(np.int64))
+        >>> prompt_kv_cache = _inner_ops.PromptKVCache(0)
+        >>> output = prompt_kv_cache(cache, update, valid_seq_len, batch_index, 2, new_max_seq_len, cur_max_seq_len)
+        >>> print(cache)
+    """
+    @prim_attr_register
+    def __init__(self, padding_mode="right"):
+        """Initialize PromptKVCache."""
+        self.init_prim_io_names(inputs=["cache", "update", "valid_seq_len", "batch_index", "seq_len_axis",
+                                        "new_max_seq_len", "cur_max_seq_len"],
+                                outputs=["out"])
+        self.add_prim_attr('side_effect_mem', True)
+        self.padding_mode = padding_mode

mindspore/ops/operations/_quant_ops.py

@@ -269,7 +269,7 @@ class FakeLearnedScaleQuantPerLayer(PrimitiveWithInfer):
         - **quant_max** (Tensor) : Value of the quantization range.
 
     Outputs:
-        - Tensor: Simulates quantize tensor of `input_x`，with the same type and shape as the `input_x`.
+        - Tensor: Simulates quantize tensor of `input_x`, with the same type and shape as the `input_x`.
 
     Examples:
         >>> input_tensor = Tensor(np.random.rand(3, 16, 5, 5), mstype.float32)
@@ -419,7 +419,7 @@ class FakeLearnedScaleQuantPerChannel(PrimitiveWithInfer):
         - **quant_max** (Tensor) : Value of the quantization range.
 
     Outputs:
-        - Tensor: Simulates quantize tensor of `input_x`，with the same type and shape as the `input_x`.
+        - Tensor: Simulates quantize tensor of `input_x`, with the same type and shape as the `input_x`.
 
     Examples:
         >>> input_tensor = Tensor(np.random.rand(3, 16, 5, 5), mstype.float32)
@@ -975,7 +975,7 @@ class FakeQuantPerChannel(PrimitiveWithInfer):
         >>> result = fake_quant(input_x, _min, _max)
     """
     support_quant_bit = [4, 7, 8]
-    ascend_support_x_rank = [2, 4]
+    ascend_support_x_rank = [2, 3, 4]
 
     @prim_attr_register
     def __init__(self,
@@ -1008,11 +1008,7 @@ class FakeQuantPerChannel(PrimitiveWithInfer):
         self.ema_decay = validator.check_float_range(ema_decay, 0, 1, validator.INC_BOTH, 'ema_decay', self.name)
         self.num_bits = validator.check_positive_int(num_bits, 'num_bits', self.name)
         self.quant_delay = validator.check_non_negative_int(quant_delay, 'quant_delay', self.name)
-        if self.is_ascend:
-            self.channel_axis = validator.check_int_range(channel_axis, 0, 1, validator.INC_BOTH,
-                                                          'channel_axis', self.name)
-        else:
-            self.channel_axis = validator.check_non_negative_int(channel_axis, 'channel_axis', self.name)
+        self.channel_axis = validator.check_non_negative_int(channel_axis, 'channel_axis', self.name)
         self.init_prim_io_names(inputs=['x', 'min', 'max'], outputs=['out'])
 
     def infer_shape(self, x_shape, min_shape, max_shape):