mindspore 2.2.14__cp39-cp39-win_amd64.whl → 2.4.0__cp39-cp39-win_amd64.whl

mindspore/ops/operations/_inner_ops.py

@@ -17,6 +17,7 @@
 from types import FunctionType, MethodType
 from collections.abc import Iterable
 import os
+import weakref
 import numpy as np
 
 from mindspore.common import Tensor
@@ -24,20 +25,21 @@ 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, _check_contains_variable
 from mindspore._c_expression import Tensor as Tensor_
-from mindspore._c_expression import typing
+from mindspore._c_expression import typing, HookType
 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, get_rank
+from mindspore.communication.management import GlobalComm, get_rank, _get_group, get_group_size
 from mindspore.common.api import _pynative_executor
 from mindspore.common._register_for_adapter import ms_adapter_registry
 from mindspore import ops
+from ..auto_generate import TensorCopySlices, SiLU, Cummin, TopKRouter, ExtractImagePatches, DecoderKVCache, \
+    PromptKVCache, ApplyCamePart1, ApplyCamePart2, ApplyCamePart3, ApplyCamePart4
 
 # Bit operation
 bit_and = bit_and()
@@ -56,73 +58,28 @@ string_mul = Primitive("string_mul")
 string_getitem = Primitive("string_getitem")
 
 
-class ExtractImagePatches(Primitive):
+class Generator(Primitive):
     r"""
-    Extracts patches from images.
-    The input tensor must be a 4-D tensor and the data format is NCHW.
-
-    Args:
-        ksizes (Union[tuple[int], list[int]]): The size of sliding window, must be a tuple or a list of integers,
-            and the format is [1, 1, ksize_row, ksize_col].
-        strides (Union[tuple[int], list[int]]): Distance between the centers of the two consecutive patches,
-            must be a tuple or list of int, and the format is [1, 1, stride_row, stride_col].
-        rates (Union[tuple[int], list[int]]): In each extracted patch, the gap between the corresponding dimension
-            pixel positions, must be a tuple or a list of integers, and the format is [1, 1, rate_row, rate_col].
-        padding (str): The type of padding algorithm, is a string whose value is "same" or "valid",
-            not case sensitive. Default: "valid".
-
-            - same: Means that the patch can take the part beyond the original image, and this part is filled with 0.
-
-            - valid: Means that the taken patch area must be completely covered in the original image.
+    Manage the state of random number generation.
 
     Inputs:
-        - **input_x** (Tensor) - A 4-D tensor whose shape is :math:`(in\_batch, in\_depth, in\_row, in\_col)`.
+        - **cmd** (int) : operation to be executed.
+        - **inputs** (tuple[tensor]) : inputs for the operation.
 
     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
-        and
-
-        .. math::
-            out_depth=ksize\_row * ksize\_col * in\_depth
-
-        and
-        if 'padding' is "valid":
-
-        .. math::
-            out\_row=floor((in\_row - (ksize\_row + (ksize\_row - 1) * (rate\_row - 1))) / stride\_row) + 1
-            out\_col=floor((in\_col - (ksize\_col + (ksize\_col - 1) * (rate\_col - 1))) / stride\_col) + 1
-
-        if 'padding' is "same":
-
-        .. math::
-            out\_row=floor((in\_row - 1) / stride\_row) + 1
-            out\_col=floor((in\_col - 1) / stride\_col) + 1
-
-    Supported Platforms:
-        ``Ascend`` ``GPU``
+        - **seed** (Tensor): Seed for the random number generation algorithm.
+        - **offset** (Tensor): Offset of the random number sequence.
+        - **state** (Tensor): State tensor, can be used to restore current state.
     """
 
     @prim_attr_register
-    def __init__(self, ksizes, strides, rates, padding="valid"):
-        """init"""
-
-        def _check_tuple_or_list(arg_name, arg_val, prim_name):
-            validator.check_value_type(f"{arg_name}s", arg_val, [tuple, list], self.name)
-            if len(arg_val) != 4 or arg_val[0] != 1 or arg_val[1] != 1:
-                raise ValueError(f"For \'{prim_name}\' the format of {arg_name}s must be [1, {arg_name}_row, "
-                                 f"{arg_name}_col, 1], but got {arg_val}.")
-            if not isinstance(arg_val[2], int) or not isinstance(arg_val[3], int) or arg_val[2] < 1 or arg_val[3] < 1:
-                raise ValueError(f"For '{prim_name}' the {arg_name}_row and {arg_name}_col in {arg_name}s must be "
-                                 f"an positive integer number, but got {arg_name}_row is {arg_val[2]}, "
-                                 f"{arg_name}_col is {arg_val[3]}")
-
-        _check_tuple_or_list("ksize", ksizes, self.name)
-        _check_tuple_or_list("stride", strides, self.name)
-        _check_tuple_or_list("rate", rates, self.name)
-        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)
+    def __init__(self):
+        self.add_prim_attr("side_effect_mem", True)
+
+    def __call__(self, cmd, inputs):
+        if cmd == 0:  # step cmd
+            return inputs[0], inputs[1]
+        return super().__call__(cmd, inputs)
 
 
 class Quant(PrimitiveWithInfer):
@@ -140,7 +97,7 @@ class Quant(PrimitiveWithInfer):
         y = round(scale * x * scale + offset)
 
     Note:
-        This operation only support Ascend 310 inference environment.
+        This operation only support Atlas 200/300/500 inference product.
 
     Args:
         scale (float) : Specifies the scaling ratio.
@@ -253,7 +210,7 @@ class Dequant(PrimitiveWithInfer):
         y = x * deq\_scale * deq\_scale
 
     Note:
-        This operation only support Ascend 310 inference environment.
+        This operation only support Atlas 200/300/500 inference product.
 
     Args:
         sqrt_mode (bool) : Specifies whether to perform square root on `scale`. Default: ``False``.
@@ -274,10 +231,10 @@ class Dequant(PrimitiveWithInfer):
     """
 
     @prim_attr_register
-    def __init__(self, sqrt_mode=False, relu_flag=False):
+    def __init__(self, sqrt_mode=False, relu_flag=False, dtype=mstype.float16):
         self.sqrt_mode = validator.check_value_type("sqrt_mode", sqrt_mode, [bool], self.name)
         self.relu_flag = validator.check_value_type("relu_flag", relu_flag, [bool], self.name)
-        self.add_prim_attr("dtype", mstype.float16)
+        self.dtype = dtype
 
     def infer_shape(self, x_shape, deq_scale_shape):
         return x_shape
@@ -289,6 +246,53 @@ class Dequant(PrimitiveWithInfer):
         return mstype.float16
 
 
+class AntiQuant(Primitive):
+    r"""
+    Returns the antiquantized value of input_x.
+
+    If `sqrt_mode` is False:
+
+    .. math::
+        y = scale * (x + offset)
+
+    If `sqrt_mode` is True:
+
+    .. math::
+        y = scale * scale * (x + offset)
+
+    Note:
+        This operation only support Atlas 200/300/500 inference product.
+
+    Args:
+        scale (float) : Specifies the scaling ratio.
+        offset (float): Specifies the offset.
+        sqrt_mode (bool) : Specifies whether to perform square root on `scale`. Default: ``False``.
+
+    Inputs:
+        - **input_x** (Tensor) : Input tensor. Must be mindspore.int8.
+
+    Outputs:
+        - Tensor: The antiquantized output tensor of type mindspore.float32.
+
+    Examples:
+        >>> from mindspore.ops.operations._inner_ops import AntiQuant
+        >>> input_x = Tensor([50.0, 20.0], mstype.int8)
+        >>> antiquant = AntiQuant(2.0, 1.0, False)
+        >>> y = antiquant(input_x)
+        >>> print(y)
+        [102. 42.]
+    """
+
+    @prim_attr_register
+    def __init__(self, sqrt_mode=False, dtype=mstype.float16):
+        super().__init__("AntiQuant")
+        self.sqrt_mode = validator.check_value_type("sqrt_mode", sqrt_mode, [bool], self.name)
+        self.dtype = dtype
+
+        self.init_prim_io_names(inputs=['x', 'scale', 'offset'],
+                                outputs=['y'])
+
+
 class MatrixDiag(PrimitiveWithInfer):
     """
     Returns a batched diagonal tensor with a given batched diagonal values.
@@ -386,227 +390,6 @@ class MatrixDiagPart(PrimitiveWithInfer):
         return out_shape
 
 
-class Send(PrimitiveWithInfer):
-    """
-    Send tensors from src_rank to the specified dest_rank.
-
-    Note:
-        Send and Receive must be used in combination and have same sr_tag.
-        Send must be used between servers.
-
-    Args:
-        sr_tag (int): A required integer identifying the send/recv message tag. The message will
-                      will be received by the Receive op with the same "sr_tag".
-        dest_rank (int): A required integer identifying the destination rank.
-        group (str): The communication group to work on. Default: "hccl_world_group/nccl_world_group".
-
-    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
-        >>>
-        >>> init()
-        >>> class Net(nn.Cell):
-        >>>     def __init__(self):
-        >>>         super(Net, self).__init__()
-        >>>         self.depend = ops.Depend()
-        >>>         self.send = ops.Send(st_tag=0, dest_rank=8, group="hccl_world_group")
-        >>>
-        >>>     def construct(self, x):
-        >>>         out = self.depend(x, self.send(x))
-        >>>         return out
-        >>>
-        >>> input_ = Tensor(np.ones([2, 8]).astype(np.float32))
-        >>> net = Net()
-        >>> output = net(input_)
-    """
-
-    @prim_attr_register
-    def __init__(self, sr_tag, dest_rank, group=GlobalComm.WORLD_COMM_GROUP, group_back=GlobalComm.WORLD_COMM_GROUP):
-        self.rank = dest_rank
-        self.sr_tag = sr_tag
-        self.group = group
-        self.add_prim_attr("no_eliminate", True)
-
-    def infer_shape(self, x_shape):
-        self.add_prim_attr("shape", x_shape)
-        return x_shape
-
-    def infer_dtype(self, x_dtype):
-        return x_dtype
-
-
-class Receive(PrimitiveWithInfer):
-    """
-    Receive tensors from src_rank.
-
-    Note:
-        Send and Receive must be used in combination and have same sr_tag.
-        Receive must be used between servers.
-
-    Args:
-        sr_tag (int): A required integer identifying the send/recv message tag. The message will
-                      will be send by the Send op with the same "sr_tag".
-        src_rank (int): A required integer identifying the source rank.
-        shape (list[int]): A required list identifying the shape of the tensor to be received.
-        dtype (Type): A required Type identifying the type of the tensor to be received. The supported types:
-                       int8, int16, int32, float16, float32.
-        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
-        >>>
-        >>> init()
-        >>> class Net(nn.Cell):
-        >>>     def __init__(self):
-        >>>         super(Net, self).__init__()
-        >>>         self.recv = ops.Receive(st_tag=0, src_rank=0, shape=[2, 8], dtype=np.float32,
-        >>>                               group="hccl_world_group")
-        >>>
-        >>>     def construct(self):
-        >>>         out = self.recv()
-        >>>         return out
-        >>>
-        >>> net = Net()
-        >>> output = net()
-    """
-
-    @prim_attr_register
-    def __init__(self, sr_tag, src_rank, shape, dtype, group=GlobalComm.WORLD_COMM_GROUP,
-                 group_back=GlobalComm.WORLD_COMM_GROUP):
-        self.rank = src_rank
-        self.tag = sr_tag
-        self.shape = shape
-        self.dtype = dtype
-        self.group = group
-        self.add_prim_attr("no_eliminate", True)
-        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)
-
-    def infer_shape(self, x_shape=None):
-        return self.get_attr_dict()['shape']
-
-    def infer_dtype(self, x_dtype=None):
-        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.
@@ -1329,45 +1112,6 @@ class DynamicBroadcastGradientArgs(Primitive):
         """Init BroadcastGradientArgs"""
 
 
-class TensorCopySlices(Primitive):
-    """
-    Copy continues memory.
-
-    Inputs:
-        - **x** (Tensor) - The target Tensor.
-        - **value** (Tensor) - The tensor to update x.
-        - **begin** (tuple[int]) - A tuple which represents the location where to start. Only
-          constant value is allowed.
-        - **end** (tuple[int]) - A tuple or which represents the maximum location where to end.
-          Only constant value is allowed.
-        - **strides** (tuple[int]) - A tuple which represents the stride is continuously added
-          before reaching the maximum location. Only constant value is allowed.
-
-    Outputs:
-        - **y** (Tensor), has the same shape and data type of x.
-
-    Examples:
-        >>> import numpy as np
-        >>> from mindspore.ops.operations import _inner_ops
-        >>> copy_slices = _inner_ops.TensorCopySlices()
-        >>> out = copy_slices(Tensor(np.zeros((5, 5))), Tensor(np.ones((2, 5))), (3, 0), (5, 5), (1, 1))
-        >>> print(out)
-            [[1., 1., 1., 1., 1.],
-             [1., 1., 1., 1., 1.],
-             [1., 1., 1., 1., 1.],
-             [0., 0., 0., 0., 0.],
-             [0., 0., 0., 0., 0.]]
-
-    Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
-    """
-
-    @prim_attr_register
-    def __init__(self):
-        """Initialize TensorScatterUpdate"""
-        self.init_prim_io_names(inputs=['x', 'value', 'begin', 'end', 'strides'], outputs=['y'])
-
-
 class DSDMatmul(PrimitiveWithInfer):
     """
     The definition of the CusSquare primitive.
@@ -1592,46 +1336,6 @@ class DynamicBroadcastTo(Primitive):
         self.init_prim_io_names(inputs=['x', 'shape'], outputs=['y'])
 
 
-class Cummin(Primitive):
-    r"""
-    Returns the cumulative minimum of elements and the index.
-
-    .. warning::
-        This is an experimental API that is subject to change or deletion.
-
-    Refer to :func:`mindspore.ops.cummin` for more detail.
-
-    Args:
-        axis (int): The axis to accumulate the tensor's value. Must be in the range [-rank(input), rank(input)).
-
-    Inputs:
-        - **input** (Tensor) - The input tensor.
-
-    Outputs:
-        A tuple of 2 Tensors(values, indices), containing the cumulative minimum of elements and the index,
-        The shape of each output tensor is the same as input `input`.
-
-    Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
-
-    Examples:
-        >>> from mindspore import Tensor, ops
-        >>> import mindspore
-        >>> a = Tensor([-0.2284, -0.6628,  0.0975,  0.2680, -1.3298, -0.4220], mindspore.float32)
-        >>> func = ops.Cummin(axis=0)
-        >>> output = func(a)
-        >>> print(output[0])
-        [-0.2284 -0.6628 -0.6628 -0.6628 -1.3298 -1.3298]
-        >>> print(output[1])
-        [0 1 1 1 4 4]
-    """
-
-    @prim_attr_register
-    def __init__(self, axis):
-        """Initialize Cummin"""
-        validator.check_value_type('axis', axis, [int], self.name)
-
-
 class DynamicResizeNearestNeighbor(Primitive):
     r"""
     Resizes the input tensor by using the nearest neighbor algorithm.
@@ -1832,7 +1536,7 @@ class CellBackwardHook(PrimitiveWithInfer):
         ...     print(grad)
         ...
         >>> hook = inner.CellBackwardHook()
-        >>> hook_fn_key = hook.register_backward_hook(hook_fn)
+        >>> hook_fn_key = hook.register_backward_hook()
         >>> def hook_test(x, y):
         ...     z = x * y
         ...     z = hook(z)
@@ -1853,16 +1557,19 @@ class CellBackwardHook(PrimitiveWithInfer):
         (Tensor(shape=[], dtype=Float32, value= 4), Tensor(shape=[], dtype=Float32, value= 4))
     """
 
-    def __init__(self, cell_id=""):
+    def __init__(self, cell_id="", cell=None, hook_dict=None):
         """Initialize CellBackwardHook"""
         super(CellBackwardHook, self).__init__(self.__class__.__name__)
         self.cell_id = cell_id
+        self.cell = cell
+        self.hook_dict = weakref.ref(hook_dict)
         self.add_prim_attr("cell_id", cell_id)
-        self.init_attrs["cell_id"] = cell_id
+        self.grad_output = None
 
-    def __call__(self, args):
-        if not isinstance(args, tuple):
-            args = (args,)
+    def __call__(self, *args):
+        # If args is empty, just return.
+        if not args:
+            return args
         return _run_op(self, self.name, args)
 
     def infer_shape(self, *inputs_shape):
@@ -1875,51 +1582,76 @@ class CellBackwardHook(PrimitiveWithInfer):
             return inputs_type[0]
         return inputs_type
 
-    def register_backward_hook(self, hook_fn):
-        r"""
-        This function is used to register backward hook function. Note that this function is only supported in pynative
-        mode.
-
-        Note:
-            The 'hook_fn' must be defined as the following code.
-            `cell_id` is the information of registered cell. `grad_input` is the gradient passed to the cell.
-            `grad_output` is the gradient computed and passed to the next cell or primitive, which may be modified by
-            returning a new output gradient.
-            The 'hook_fn' should have the following signature:
-            hook_fn(cell_id, grad_input, grad_output) -> New output gradient or none.
-            The 'hook_fn' is executed in the python environment.
+    def register_backward_hook(self):
+        """
+        Register the backward hook function.
 
         Args:
-            hook_fn (Function): Python function. Backward hook function.
+            None
 
         Returns:
-            - **key** (int) - The key of 'hook_fn'.
+            None
 
-        Raises:
-            TypeError: If the `hook_fn` is not a function of python.
+        Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
         """
-        if not isinstance(hook_fn, (FunctionType, MethodType)):
-            raise TypeError(f"When using 'register_backward_hook(hook_fn)', the type of 'hook_fn' must be python "
-                            f"function, but got {type(hook_fn)}.")
-        key = self.add_backward_hook_fn(hook_fn)
-        return key
 
-    def remove_backward_hook(self, key):
-        r"""
-        This function is used to remove backward hook function. Note that this operation is only supported in pynative
-        mode.
-
-        Note:
-            The 'key' is the object returned by 'register_backward_hook' function of the same CellBackwardHook
-            operator.
+        def hook_backward_grad(grad):
+            if self.grad_output is None:
+                self.grad_output = grad
+                # Indicates the first time of call backward hook, and need to wait for the second time call
+                return self.cell_id
+            backward_hook_grad_input = grad
+            if self.hook_dict():
+                backward_hooks = self.hook_dict().values()
+                for hook in backward_hooks:
+                    res = hook(self.cell, backward_hook_grad_input, self.grad_output)
+                    if res is None:
+                        continue
+                    if not isinstance(res, tuple):
+                        res = (res,)
+                    if len(res) != len(grad):
+                        raise TypeError(
+                            "The backward hook return value size is {} not equal to expect grad input size {}".format(
+                                len(res), len(grad)))
+                    backward_hook_grad_input = res
+            self.grad_output = None
+            return backward_hook_grad_input
+
+        self.set_hook_fn(hook_backward_grad, HookType.BackwardHook)
+
+    def register_backward_pre_hook(self):
+        """
+        Register the backward pre hook function.
 
         Args:
-            key (int): The key corresponding to the 'hook_fn'.
+            None
 
         Returns:
-            None.
+            None
+
+        Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
         """
-        self.remove_backward_hook_fn(key)
+
+        def hook_backward_pre_grad(grad):
+            backward_pre_hook_grad = grad
+            if self.hook_dict():
+                backward_pre_hooks = self.hook_dict().values()
+                for hook in backward_pre_hooks:
+                    res = hook(self.cell, backward_pre_hook_grad)
+                    if res is None:
+                        continue
+                    if not isinstance(res, tuple):
+                        res = (res,)
+                    if len(res) != len(grad):
+                        raise TypeError(
+                            "The backward pre hook return value size is {} not equal to expect output size {}".format(
+                                len(res), len(grad)))
+                    backward_pre_hook_grad = res
+            return backward_pre_hook_grad
+
+        self.set_hook_fn(hook_backward_pre_grad, HookType.BackwardPreHook)
 
 
 class Format(PrimitiveWithInfer):
@@ -1948,7 +1680,6 @@ class Format(PrimitiveWithInfer):
     def __init__(self):
         self.init_prim_io_names(inputs=['string', 'args'], outputs=['string'])
 
-
     def __infer__(self, str_, *var):
         def check_variable(str_, var):
             if _check_contains_variable(str_['dtype'], str_['value']):
@@ -1959,11 +1690,9 @@ class Format(PrimitiveWithInfer):
                     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()
@@ -2148,14 +1877,13 @@ 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(self.axis_str, self.axis, [int, tuple, list], self.name)
+        validator.check_value_type('axis', 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[self.axis_str] = self.axis
-        self.add_prim_attr(self.axis_str, self.axis)
+        self.init_attrs['axis'] = self.axis
+        self.add_prim_attr('axis', self.axis)
         self.init_prim_io_names(inputs=['x', 'clip_norm'], outputs=['output'])
 
     def infer_shape(self, x_shape, clip_norm_shape):
@@ -2248,7 +1976,8 @@ class MixedPrecisionCast(Primitive):
 
     def __call__(self, dst_dtype, x):
         def cast_inner(data):
-            if isinstance(data, Tensor) and data.dtype in (mstype.float16, mstype.float32, mstype.float64):
+            if isinstance(data, Tensor) and data.dtype in (mstype.float16, mstype.float32,
+                                                           mstype.float64, mstype.bfloat16):
                 return self.cast(data, dst_dtype)
             return data
 
@@ -2559,7 +2288,7 @@ class ConvertToMsTensor(Primitive):
         """Run in PyNative mode"""
         if isinstance(x, StubTensor):
             return StubTensor(stub=x.stub, tensor=x.tensor)
-        return ops.deepcopy(x)
+        return ops.auto_generate.deepcopy(x)
 
 
 convert_to_ms_tensor = ConvertToMsTensor()
@@ -2621,28 +2350,6 @@ class IsParameter(PrimitiveWithInfer):
                 'value': isinstance(x['dtype'], mstype.RefType)}
 
 
-class SiLU(Primitive):
-    r"""
-    Computes SiLU (Sigmoid Linear Unit activation function) of input tensors element-wise.
-
-    Refer to :func:`mindspore.ops.silu` for more details.
-
-    Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
-
-    Examples:
-        >>> x = Tensor(np.array([-1, 2, -3, 2, -1]), mindspore.float16)
-        >>> output = ops.SiLU(x)
-        >>> print(output)
-        [-0.269  1.762  -0.1423  1.762  -0.269]
-    """
-
-    @prim_attr_register
-    def __init__(self):
-        """Initialize SiLU"""
-        self.init_prim_io_names(inputs=['x'], outputs=['output'])
-
-
 class TileSize(Primitive):
     r"""
         Tile size for matmul
@@ -2726,6 +2433,7 @@ class CopyWithSlice(Primitive):
     r"""
         Copy data to discontinuous tensor
     """
+
     @prim_attr_register
     def __init__(self):
         self.add_prim_attr('side_effect_mem', True)
@@ -2775,10 +2483,10 @@ class FFN(Primitive):
         >>> h = 1024
         >>> h_f = 4 * h
         >>> e = 16
-        >>> x = Tensor(np.random.randn(b * s, h).astype(np.float16))
+        >>> x = Tensor(np.random.randn(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))
+        >>> expert_tokens = Tensor(np.full(e, 8))
         >>> 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)
@@ -2790,189 +2498,47 @@ class FFN(Primitive):
     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"],
+                                        "bias2", "scale", "offset", "deq_scale1", "deq_scale2",
+                                        "antiquant_scale1", "antiquant_scale2",
+                                        "antiquant_offset1", "antiquant_offset2"],
                                 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)
+class _VirtualConverterEnd(PrimitiveWithInfer):
     """
-    @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 _MirrorSilentCheck(PrimitiveWithInfer):
+    Auto parallel virtual operator.
     """
-    The operator _MirrorSilentCheck implements accuracy-sensitive detection on the tensor input in backpropagator.
-    Call _MirrorSilentCheck in method __call__ of derived class to implement accuracy-sensitive detection.
-
-    Inputs:
-        - **input** (Tensor) : The tensor used for detection.
-          Its data type must be mindspore.float16, mindspore.float32 or mindspore.bfloat16.
-        - **pre_val** (Parameter(Tensor)) : Support parameter in accuracy-sensitive detection.
-          Please only generated by method generate_params() of ASDBase.
-        - **min_val** (Parameter(Tensor)) : Support parameter in accuracy-sensitive detection.
-          Please only generated by method generate_params() of ASDBase.
-        - **max_val** (Parameter(Tensor)) : Support parameter in accuracy-sensitive detection.
-          Please only generated by method generate_params() of ASDBase.
-        - **cnt** (Parameter(Tensor)) : Support parameter in accuracy-sensitive detection.
-          Please only generated by method generate_params() of ASDBase.
-          After each invocation of _MirrorSilentCheck, increment the value of cnt by one.
 
-    Outputs:
-        - **output** (Tensor) - Same shape, type and value as `input`.
-    """
     @prim_attr_register
-    def __init__(self, min_steps=8):
-        upper_thresh, sigma_thresh = self.get_thresh()
-        self.min_steps = min_steps
-        self.thresh_l1 = upper_thresh[0]
-        self.coeff_l1 = sigma_thresh[0]
-        self.thresh_l2 = upper_thresh[1]
-        self.coeff_l2 = sigma_thresh[1]
-        self.add_prim_attr('side_effect_mem', True)
-
-    def parse_thresh(self, env_var_name, default_value, min_value):
-        env_var = os.environ.get(env_var_name, default=default_value)
-        thresh = [value.strip() for value in env_var.split(",")]
-        if len(thresh) != 2 or not all(value.isdigit() for value in thresh):
-            thresh = default_value.split(",")
-        thresh = [float(max(int(value), min_value)) for value in thresh]
-        if thresh[0] <= thresh[1]:
-            thresh = [float(value) for value in default_value.split(",")]
-
-        return thresh
-
-    def get_thresh(self):
-        upper_thresh = self.parse_thresh("NPU_ASD_UPPER_THRESH", "1000000,10000", 3)
-        sigma_thresh = self.parse_thresh("NPU_ASD_SIGMA_THRESH", "100000,5000", 3)
-        return upper_thresh, sigma_thresh
+    def __init__(self, input_nums):
+        """Initialize _VirtualConverterEnd."""
+        self.input_nums = input_nums
 
-    def infer_shape(self, x_shape, pre_shape, min_shape, max_shape, n_step, loss_scale_shape):
-        return x_shape
-
-    def infer_dtype(self, x_dtype, pre_dtype, min_dtype, max_dtype, n_dtype, loss_scale_dtype):
-        return x_dtype
+    def infer_shape(self, *args):
+        return (args[0][0] * self.input_nums,) + tuple(args[0][1:])
 
+    def infer_dtype(self, *args):
+        return args[0]
 
-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.
+class _VirtualConverterBegin(PrimitiveWithInfer):
+    """
+    Auto parallel virtual operator.
+    """
 
-    Outputs:
-        With same data type and same shape as `cache` tensor.
+    @prim_attr_register
+    def __init__(self, output_nums):
+        """Initialize _VirtualConverterBegin."""
+        self.output_nums = output_nums
 
-    Supported Platforms:
-        ``Ascend``
+    def infer_shape(self, arg):
+        if self.output_nums == 0:
+            return ValueError("output_nums can\'t be zero.")
+        new_arg = (arg[0] / self.output_nums,) + tuple(arg[1:])
+        return (new_arg,) * self.output_nums
 
-    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
+    def infer_dtype(self, arg):
+        return (arg,) * self.output_nums