mindspore 2.6.0__cp310-cp310-manylinux1_x86_64.whl → 2.7.0rc1__cp310-cp310-manylinux1_x86_64.whl

mindspore/ops/auto_generate/gen_ops_def.py

@@ -79,7 +79,7 @@ def cummin_ext(input, dim):
         \end{array}
 
     .. note::
-        O2 mode is not supported in Ascend.
+        GE backend is not supported in Ascend.
 
     Args:
         input (Tensor): The input Tensor, The dimension must be greater than 0.
@@ -456,7 +456,7 @@ def rfft2(input, s=None, dim=(-2, -1), norm=None):
 
     Note:
         - `rfft2` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `rfft2` is not supported on Windows platform yet.
 
     Args:
@@ -1169,7 +1169,7 @@ def ihfft2(input, s=None, dim=(-2, -1), norm=None):
 
     Note:
         - `ihfft2` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `ihfft2` is not supported on Windows platform yet.
 
     Args:
@@ -1874,7 +1874,7 @@ def broadcast_to(input, shape):
 
     Args:
         input (Tensor): The input tensor.
-        shape (tuple): The target shape.
+        shape (tuple[int]): The target shape.
 
     Returns:
         Tensor
@@ -2169,65 +2169,6 @@ def diagonal(input, offset=0, dim1=0, dim2=1):
     return diagonal_op(input)
 
 
-def prompt_k_v_cache(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, align_mode='LEFT'):
-    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 format is BHSD, cache tensor of shape
-          :math:`(cache\_batch\_size, num\_head, max\_seq\_length, size\_pre\_head)`.
-          When format is BSD, cache tensor of shape
-          :math:`(cache\_batch\_size, max\_seq\_length, hidden\_size)`.
-        update (Tensor]): The tensor which is used to update the cache tensor. Same data type as cache tensor.
-          When format is BHSD, cache tensor of shape
-          :math:`(update\_batch\_size, num\_head, max\_seq\_length, size\_pre\_head)`.
-          When format is BSD, cache tensor of shape
-          :math:`(update\_batch\_size, max\_seq\_length, hidden\_size)`.
-        valid_seq_len (Tensor): The valid_seq_len tensor with data type of int64.
-          Valid_seq_len tensor of shape :math:`(update\_batch\_size)`.
-        batch_index (Tensor): The batch_index tensor with data type of int64.
-          Batch_index tensor of shape :math:`(update\_batch\_size)`. Indicate that which batch of cache tensor is going to be update.
-        seq_len_axis (Tensor): The seq_len_axis indicate which axis is seq_eln, set to '1' or '2'. Not able for now.
-        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 (str): indicate which axis is seq_len. Default: left.
-
-
-    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.randint(-1, s, size=ub).astype(np.int64))
-        >>> batch_index = Tensor(np.random.choice(np.arange(-1, b), size=ub, replace=False).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, Tensor(2), new_max_seq_len, cur_max_seq_len)
-        >>> print(cache)
-    """
-    prompt_k_v_cache_op = _get_cache_prim(PromptKVCache)(align_mode)
-    return prompt_k_v_cache_op(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len)
-
-
 def mse_loss_ext(input, target, reduction='mean'):
     r"""
     Calculates the mean squared error between the predicted value and the label value.
@@ -2403,9 +2344,6 @@ def inplace_threshold(input, threshold, value):
         \text{value}, &\text{ otherwise }
         \end{cases}
 
-    .. warning::
-        This is an experimental API that is subject to change or deletion.
-
     Args:
         input (Tensor): The input Tensor.
         threshold (Union[int, float]): The value of the threshold.
@@ -2433,11 +2371,11 @@ def inplace_threshold(input, threshold, value):
     return inplace_threshold_op(input, threshold, value)
 
 
-def fused_add_topk_div(x, add_num, group_num, group_topk, n, k, activate_type=0, is_norm=True, scale=2.5):
+def fused_add_topk_div(x, add_num, group_num, group_topk, n, k, activate_type=0, is_norm=True, scale=2.5, mapping_num=None, mapping_table=None, enable_expert_mapping=False):
     r"""
     
     """
-    return fused_add_topk_div_op(x, add_num, group_num, group_topk, n, k, activate_type, is_norm, scale)
+    return fused_add_topk_div_op(x, add_num, group_num, group_topk, n, k, activate_type, is_norm, scale, mapping_num, mapping_table, enable_expert_mapping)
 
 
 def batch_norm_elemt(input, weight=None, bias=None, mean=None, invstd=None, eps=1e-5):
@@ -2694,6 +2632,129 @@ def embedding_feature_mapping_import(file_path, teble_name, feature_size, global
     return embedding_feature_mapping_import_op(file_path, teble_name, feature_size, global_step, embedding_dim, only_offset_flag, num)
 
 
+def moe_distribute_combine(expand_x, expert_ids, expand_idx, ep_send_counts, expert_scales, ep_world_size, ep_rank_id, moe_expert_num, tp_send_counts=None, x_active_mask=None, activate_scale=None, weight_scale=None, group_list=None, expand_scales=None, group_ep=None, group_tp=None, tp_world_size=0, tp_rank_id=0, expert_shard_type=0, shared_expert_num=0, shared_export_rank_num=0, global_bs=0, out_dtype=0, common_quant_mode=0, group_list_type=0):
+    r"""
+    Parallel communication for Mixture of Experts (MoE). When Tensor Parallelism (TP) communication exists,
+    it first ReduceScatter performs communication followed by Expert Parallelism (EP) AllToAllV communication.
+    Otherwise, only EP AllToAllV communication is performed. Finally multiply the received data by weight and
+    add them up.
+
+    Notes:
+        This function must be used in conjunction with function `moe_distribute_dispatch`.
+        - A: Maximum tokens to dispatch per rank:
+            - For shared experts: A = BS * ep_world_size * shared_expert_num / shared_expert_rank_num
+            - For MoE experts:
+                - When global_bs = 0: A >= BS * ep_world_size * min(local_expert_num, K)
+                - When global_bs != 0: A >= global_bs * min(local_expert_num, K)
+        - H (hidden size): Dimension of each token's hidden state
+            - Ascend 910B: 0 < H <= 7168, must be multiple of 32
+            - Ascend 910_93: H = 7168
+        - BS (batch sequence size): Number of tokens processed per rank
+            - Ascend 910B: 0 < BS <= 256
+            - Ascend 910_93: 0 < BS <= 512
+        - K: Number of experts selected per token (0 < K <= 8 and K <= moe_expert_num)
+        - server_num: Number of server nodes (supports 2, 4, 8)
+        - local_expert_num: Number of experts per rank:
+            - Shared expert ranks: local_expert_num = 1
+            - MoE expert ranks: local_expert_num = moe_expert_num / (ep_world_size - shared_expert_rank_num)
+            (TP communication not supported when localExpertNum > 1)
+
+    Inputs:
+        - **expand_x** (Tensor) - Expanded token features. 2D tensor [A, H] with dtype matching input.
+            Supported dtypes: float16, bfloat16, int8. Format: ND, non-contiguous allowed.
+        - **expert_ids** (Tensor) - Top-K expert indices for each token. 2D int32 tensor with shape [BS, K].
+            Format: ND, non-contiguous allowed.
+        - **expert_idx** (Tensor) - Token counts per expert, it's the output of dispatch operation. 
+            1D int32 tensor [BS*K]. Format: ND, non-contiguous allowed.
+        - **ep_send_counts** (Tensor) - Tokens that each EP rank needs to send, it's the output of dispatch operation.
+            - Ascend 910B: 1D int32 tensor [moe_expert_num + 2 * global_bs * K * server_num]
+            - Ascend 910_93: 1D int32 tensor [ep_world_size * max(tp_world_size,1) * local_expert_num]
+            Format: ND, non-contiguous allowed.
+        - **expert_scales** (Tensor) - Top-K expert weights per token.
+        - **ep_world_size** (int) - EP domain size.
+            - Ascend 910B: Supports 16, 32, 64.
+            - Ascend 910_93: Supports 8, 16, 32, 64, 128, 144, 256, 288.
+        - **ep_rank_id** (int) - Local rank ID in EP domain [0, ep_world_size), must be unique per domain.
+        - **moe_expert_num** (int) - Number of MoE experts (0 < moe_expert_num <= 256), 
+            must satisfy moe_expert_num % (ep_world_size-shared_expert_rank_num) = 0.
+        - **tp_send_counts** (Tensor) - Tokens that each TP rank needs to send (when TP exists). It's the output of dispatch operation. Default: ``None``.
+            - Ascend 910B: Not supported.
+            - Ascend 910_93: 1D int32 tensor [tp_world_size] when TP exists. Format: ND, non-contiguous allowed.
+        - **x_active_mask** (Tensor) - Reserved parameter. Default: ``None``.
+        - **activate_scale** (Tensor) - Reserved parameter. Default: ``None``.
+        - **weight_scale** (Tensor) - Reserved parameter. Default: ``None``.
+        - **group_list** (Tensor) - Reserved parameter. Default: ``None``.
+        - **expand_scales** (Tensor) - Output of dispatch operation. Default: ``None``.
+            - Ascend 910B: 1D float32 tensor [A]. Format: ND, non-contiguous allowed.
+            - Ascend 910_93: Unsupported.
+        - **group_ep** (str) - EP communication domain name (string length 1-127), must differ from group_tp. Default: ``None``.
+        - **group_tp** (str) - TP communication domain name. Default: ``None``.
+            - Ascend 910B: Unsupported (pass empty string).
+            - Ascend 910_93: When TP communication exists, string length 1-127, must differ from group_ep.
+        - **tp_world_size** (int) - TP domain size. Default: ``0``.
+            - Ascend 910B: Unsupported (pass 0).
+            - Ascend 910_93: 0/1 means no TP communication; only 2 supported when TP exists.
+        - **tp_rank_id** (int) - Local rank ID in TP domain. Default: ``0``.
+            - Ascend 910B: Unsupported (pass 0).
+            - Ascend 910_93: [0,1], unique per domain; pass 0 when no TP communication.
+        - **expert_shard_type** (int) - Shared expert distribution type. Default: ``0``.
+            - Ascend 910B: Unsupported (pass 0).
+            - Ascend 910_93: Currently only 0 (shared experts precede MoE experts).
+        - **shared_expert_num** (int) - Number of shared experts. Default: ``0``.
+            - Ascend 910B: Unsupported (pass 0).
+            - Ascend 910_93: Currently 0 (none) or 1 (one shared expert).
+        - **shared_expert_rank_num** (int) - Number of ranks hosting shared experts. Default: ``0``.
+            - Ascend 910B: Unsupported (pass 0).
+            - Ascend 910_93: [0, ep_world_size-1), must satisfy ep_world_size % shared_expert_rank_num = 0 when non-zero.
+        - **global_bs** (int) - Global batch size across EP domain. Default: ``0``.
+            - Ascend 910B: 256*ep_world_size when BS varies per rank; 0 or BS*ep_world_size when uniform.
+            - Ascend 910_93: 0 or BS*ep_world_size.
+        - **out_dtype** (int) - Specify the type of output x. Reserved parameter (pass 0 in current version). Default: ``0``.
+        - **common_quant_mode** (int) - Communication quantification type. Reserved parameter (pass 0 in current version). Default: ``0``.
+        - **group_list_type** (int) - The format of group_list. Reserved parameter (pass 0 in current version). Default: ``0``.
+
+    Outputs:
+        - **x** (Tensor) - Processed tokens. 2D tensor [BS, H] with dtype matching input `expand_x`.
+
+    Raises:
+        TypeError: If input dtypes don't match specifications.
+        ValueError: If input values violate constraints (e.g., invalid expert indices).
+        RuntimeError: If communication domain configuration is invalid.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+    >>> # EP-only communication example (Ascend 910B)
+    >>> import mindspore as ms
+    >>> from mindspore import Tensor
+    >>> from mindspore import ops
+    >>> from mindspore.communication import init, get_rank, GlobalComm
+    >>> from mindspore.ops.auto_generate import moe_distribute_dispatch, moe_distribute_combine
+    >>> import numpy as np
+    >>> bs = 8
+    >>> h = 7168
+    >>> k = 8
+    >>> ep_world_size = 16
+    >>> moe_expert_num = 16
+    >>> global_bs = bs * ep_world_size
+    >>> x = Tensor(np.random.randn(bs, h), ms.float16)
+    >>> expert_ids = Tensor(np.random.randint(0, moe_expert_num, (bs, k)), ms.int32)
+    >>> expert_scales = Tensor(np.random.randn(bs, k), ms.float32)
+    >>> init()
+    >>> rank_id = get_rank()
+    >>> expand_x, _, expand_idx, _, ep_recv_count, _, expand_scale = moe_distribute_dispatch(
+    ...     x, expert_ids, expert_scales, ep_world_size, rank_id, moe_expert_num,
+    ...     group_ep=GlobalComm.WORLD_COMM_GROUP)
+    >>> out_x = moe_distribute_combine(
+    ...     expand_x, expert_ids, expand_idx, ep_recv_count, expert_scales, ep_world_size, rank_id,
+    ...     moe_expert_num, group_ep=GlobalComm.WORLD_COMM_GROUP)
+    >>> print(out_x.shape)
+    (8, 7168)
+    """
+    return moe_distribute_combine_op(expand_x, expert_ids, expand_idx, ep_send_counts, expert_scales, ep_world_size, ep_rank_id, moe_expert_num, tp_send_counts, x_active_mask, activate_scale, weight_scale, group_list, expand_scales, group_ep, group_tp, tp_world_size, tp_rank_id, expert_shard_type, shared_expert_num, shared_export_rank_num, global_bs, out_dtype, common_quant_mode, group_list_type)
+
+
 def inplace_mul(input, other):
     r"""
     
@@ -3088,7 +3149,7 @@ def fft2(input, s=None, dim=(-2, -1), norm=None):
 
     Note:
         - `fft2` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `fft2` is not supported on Windows platform yet.
 
     Args:
@@ -3348,7 +3409,7 @@ def ifftn(input, s=None, dim=None, norm=None):
 
     Note:
         - `ifftn` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `ifftn` is not supported on Windows platform yet.
 
     Args:
@@ -3452,7 +3513,7 @@ def hfft2(input, s=None, dim=(-2, -1), norm=None):
 
     Note:
         - `hfft2` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `hfft2` is not supported on Windows platform yet.
 
     Args:
@@ -3524,7 +3585,7 @@ def rfftfreq(n, d=1.0, dtype=None):
 
     Note:
         - `rfftfreq` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `rfftfreq` is not supported on Windows platform yet.
 
     Args:
@@ -3729,60 +3790,6 @@ def less(input, other):
     return less_op(input, other)
 
 
-def moe_token_unpermute(permuted_tokens, sorted_indices, probs=None, padded_mode=False, restore_shape=None):
-    r"""
-    Unpermute a tensor of permuted tokens based on sorted indices, and optionally merge the tokens with their corresponding probabilities.
-
-    .. warning::
-        - It is only supported on Atlas A2 Training Series Products.
-        - The inputs `permuted_tokens` and `probs` only support the bfloat16 data type in the current version.
-        - `sorted_indices` must not have duplicate values, otherwise the result is undefined.
-
-    Args:
-        permuted_tokens (Tensor): The tensor of permuted tokens to be unpermuted.
-            The shape is :math:`[num\_tokens * topk, hidden\_size]` , where `num_tokens`, `topk` and `hidden_size` are positive integers.
-        sorted_indices (Tensor): The tensor of sorted indices used to unpermute the tokens.
-            The shape is :math:`[num\_tokens * topk,]`, where `num_tokens` and `topk` are positive integers.
-            It only supports the int32 data type.
-        probs (Tensor, optional): The tensor of probabilities corresponding to the permuted tokens.
-            If provided, the unpermuted tokens will be merged with their respective probabilities.
-            The shape is :math:`[num\_tokens, topk]`, where `num_tokens` and `topk` are positive integers. Default: ``None`` .
-        padded_mode (bool, optional): If ``True``, indicating the indices are padded to denote selected tokens per expert. Default: ``False`` .
-        restore_shape (Union[tuple[int], list[int]], optional): The input shape before permutation, only used in padding mode. Default: ``None`` .
-
-    Returns:
-        Tensor, with the same dtype as `permuted_tokens`. If `padded_mode` is ``False``, the shape will be [`num_tokens`, `hidden_size`].
-        If `padded_mode` is ``True``, the shape will be specified by `restore_shape`.
-
-    Raises:
-        TypeError: If `permuted_tokens` is not a Tensor.
-        ValueError: Only supported when `padded_mode` is ``False``.
-
-    Supported Platforms:
-        ``Ascend``
-
-    Examples:
-        >>> import mindspore
-        >>> from mindspore import Tensor, ops
-        >>> permuted_token = Tensor([
-        ...                          [1, 1, 1],
-        ...                          [0, 0, 0],
-        ...                          [0, 0, 0],
-        ...                          [3, 3, 3],
-        ...                          [2, 2, 2],
-        ...                          [1, 1, 1],
-        ...                          [2, 2, 2],
-        ...                          [3, 3, 3]], dtype=mindspore.bfloat16)
-        >>> sorted_indices = Tensor([0, 6, 7, 5, 3, 1, 2, 4], dtype=mindspore.int32)
-        >>> out = ops.moe_token_unpermute(permuted_token, sorted_indices)
-        >>> out.shape
-        (8, 3)
-        
-        
-    """
-    return moe_token_unpermute_op(permuted_tokens, sorted_indices, probs, padded_mode, restore_shape)
-
-
 def sign(input):
     r"""
     Return an element-wise indication of the sign of a number.
@@ -4033,7 +4040,7 @@ def hfftn(input, s=None, dim=None, norm=None):
 
     Note:
         - `hfftn` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `hfftn` is not supported on Windows platform yet.
 
     Args:
@@ -4121,7 +4128,7 @@ def fftshift(input, dim=None):
 
     Note:
         - `fftshift` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `fftshift` is not supported on Windows platform yet.
 
     Args:
@@ -4161,7 +4168,7 @@ def fftfreq(n, d=1.0, dtype=None):
 
     Note:
         - `fftfreq` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `fftfreq` is not supported on Windows platform yet.
 
     Args:
@@ -4251,9 +4258,6 @@ def swiglu(input, dim=-1):
     Computes SwiGLU (Swish-Gated Linear Unit activation function) of input tensor.
     SwiGLU is a variant of the :class:`mindspore.ops.GLU` activation function, it is defined as:
 
-    .. warning::
-        This is an experimental API that is subject to change or deletion.
-
     .. math::
         {SwiGLU}(a, b)= Swish(a) \otimes b
 
@@ -4261,6 +4265,9 @@ def swiglu(input, dim=-1):
     Swish(a)=a :math:`\sigma` (a), :math:`\sigma` is the :func:`mindspore.ops.sigmoid` activation function
     and :math:`\otimes` is the Hadamard product.
 
+    .. warning::
+        Only support on Atlas A2 training series.
+
     Args:
         input (Tensor): Tensor to be split. It has shape :math:`(\ast_1, N, \ast_2)`
             where `*` means, any number of additional dimensions. :math:`N` must be divisible by 2.
@@ -4615,56 +4622,6 @@ def t_ext(input):
     return t_ext_op(input)
 
 
-def index_add_ext(input, dim, index, source, alpha=1):
-    r"""
-    Accumulate the elements of `alpha` times `source` into the `input` by adding to the index in the order given in `index`. For example, if ``dim == 0`` , ``index[i] == j`` , and ``alpha = -1`` , then the `i` th row of `source` is subtracted from the `j` th row of `input` . The `dim` th dimension of `source` must have the same size as the length of `index` , and all other dimensions must match `input`, or an error will be raised. For a 3-D tensor, the output is defined as follows:
-
-    .. math::
-        \begin{array}{ll}
-        input[index[i],\ :,\ :]\ +=\ alpha * source[i,\ :,\ :]  \qquad \#if\ dim == 0 \\
-        input[:,\ \ index[i],\ :]\ +=\ alpha * source[:,\ \ i,\ :]  \qquad \#if\ dim == 1 \\
-        input[:,\ :,\ \ index[i]]\ +=\ alpha * source[:,\ :,\ \ i]  \qquad\#if\ dim == 2 \\
-        \end{array}
-        
-    .. warning::
-        This is an experimental API that is subject to change or deletion.
-        
-    Args:
-        input (Tensor): The input Tensor.
-        dim (int): The dimension along which to index.
-        index (Tensor): Add the value of "input Tensor" and `source` along the dimension of the `dim` according to the specified index value, with data type int32. The `index` must be 1D with the same size as the size of `source` in the `dim` dimension. The values of `index` should be in [0, b), where the b is the size of "input Tensor" in the `dim` dimension.
-        source (Tensor): The input tensor with the value to add. Must have same data type as "input Tensor". The shape must be the same as "input Tensor" except the `dim` th dimension.
-        alpha (number, optional): The scalar multiplier for source. Default: ``1``.
-        
-    Returns:
-        Tensor, has the same shape and dtype as `input`.
-
-    Raises:
-        TypeError: If neither `index` nor `source` is a Tensor.
-        ValueError: If the value of `dim` is out of the dimension range of `source` shape.
-        ValueError: If `index` rank is not the same as `source` rank.
-        ValueError: If shape of `index` is not 1D or size of `index` is not equal to dimension of source[dim].
-        ValueError: If the shape of `source` is not the same as that of `input` except the `dim` axis.
-
-    Supported Platforms:
-        ``Ascend``
-
-    Examples:
-        >>> import numpy as np
-        >>> import mindspore
-        >>> from mindspore import Tensor, ops
-        >>> x = Tensor(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]), mindspore.float32)
-        >>> index = Tensor(np.array([0, 2]), mindspore.int32)
-        >>> y = Tensor(np.array([[0.5, 1.0], [1.0, 1.5], [2.0, 2.5]]), mindspore.float32)
-        >>> output = ops.auto_generate.index_add_ext(x, 1, index, y, alpha=1)
-        >>> print(output)
-        [[ 1.5  2.   4. ]
-         [ 5.   5.   7.5]
-         [ 9.   8.  11.5]]
-    """
-    return index_add_ext_op(input, dim, index, source, alpha)
-
-
 def assign_add(variable, value):
     r"""
     Updates a parameter or tensor by adding a value to it.
@@ -4696,8 +4653,7 @@ def apply_rotary_pos_emb_(query, key, cos, sin, position_ids, cos_format=0):
     r"""
     
     """
-    apply_rotary_pos_emb_op = _get_cache_prim(ApplyRotaryPosEmb)(cos_format)
-    return apply_rotary_pos_emb_op(query, key, cos, sin, position_ids)
+    return apply_rotary_pos_emb_impl(query, key, cos, sin, position_ids, cos_format)
 
 
 def soft_margin_loss(input, target, reduction='mean'):
@@ -4873,7 +4829,7 @@ def rfft(input, n=None, dim=-1, norm=None):
 
     Note:
         - `rfft` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `rfft` is not supported on Windows platform yet.
 
     Args:
@@ -5092,7 +5048,7 @@ def ifft2(input, s=None, dim=(-2, -1), norm=None):
 
     Note:
         - `ifft2` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `ifft2` is not supported on Windows platform yet.
 
     Args:
@@ -5393,7 +5349,7 @@ def ihfftn(input, s=None, dim=None, norm=None):
 
     Note:
         - `ihfftn` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `ihfftn` is not supported on Windows platform yet.
 
     Args:
@@ -5539,61 +5495,6 @@ def shard_identity(input):
     return shard_identity_op(input)
 
 
-def decoder_k_v_cache(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len):
-    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 format is BHSD, cache tensor of shape
-          :math:`(batch\_size, num\_head, max\_seq\_length, size\_pre\_head)`.
-          When format is BSD, cache tensor of shape
-          :math:`(batch\_size, max\_seq\_length, hidden\_size)`.
-        update (Tensor]): The tensor which is used to update the cache tensor. Same data type as cache tensor.
-          When format is BHSD, update tensor of shape
-          :math:`(batch\_size, num\_head, update\_seq\_length, size\_pre\_head)`.
-          When format is BSD, update tensor of shape
-          :math:`(batch\_size, update\_seq\_length, 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:`(batch\_size)`. Indicate that which batch of cache tensor is going to be update. Not abel for now.
-        seq_len_axis (Tensor): The seq_len_axis indicate which axis is seq_eln, set to '1' or '2'. Not able for now.
-        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.
-        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.randint(-1, s, size=b).astype(np.int64))
-        >>> batch_index = Tensor(np.random.choice(np.arange(-1, b), size=b, replace=False).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, Tensor(2), new_max_seq_len, cur_max_seq_len)
-        >>> print(cache)
-    """
-    return decoder_k_v_cache_op(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len)
-
-
 def embedding_table_evict(var_handle, global_step, steps_to_live=0):
     r"""
     
@@ -5947,7 +5848,7 @@ def hfft(input, n=None, dim=-1, norm=None):
 
     Note:
         - `hfft` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `hfft` is not supported on Windows platform yet.
 
     Args:
@@ -6009,7 +5910,7 @@ def irfft(input, n=None, dim=-1, norm=None):
 
     Note:
         - `irfft` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `irfft` is not supported on Windows platform yet.
 
     Args:
@@ -6094,7 +5995,7 @@ def fftn(input, s=None, dim=None, norm=None):
 
     Note:
         - `fftn` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `fftn` is not supported on Windows platform yet.
 
     Args:
@@ -6238,6 +6139,52 @@ def zero_(input):
     return inplace_zero_op(input)
 
 
+def inplace_silu(input):
+    r"""
+    Computes Sigmoid Linear Unit of input element-wise. The SiLU function is defined as:
+
+    .. math::
+
+        \text{SiLU}(x) = x * \sigma(x),
+
+    where :math:`x` is an element of the input, :math:`\sigma(x)` is Sigmoid function.
+
+    .. math::
+
+        \text{sigma}(x_i) = \frac{1}{1 + \exp(-x_i)},
+
+    SiLU Function Graph:
+
+    .. image:: ../images/SiLU.png
+        :align: center
+
+    Args:
+        input (Tensor): `input` is :math:`x` in the preceding formula. Input with the data type
+            float16 or float32.
+        inplace (bool, optional): If it is ``True``, enable the in place update function.
+            Default value: ``False``.
+
+    Returns:
+        Tensor, with the same type and shape as the `input`.
+
+    Raises:
+        TypeError: If dtype of `input` is neither float16 nor float32.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, mint
+        >>> import numpy as np
+        >>> input = Tensor(np.array([-1, 2, -3, 2, -1]), mindspore.float16)
+        >>> output = mint.nn.functional.silu(input, inplace=True)
+        >>> print(output)
+        [-0.269  1.762  -0.1423  1.762  -0.269]
+    """
+    return inplace_silu_op(input)
+
+
 def rotary_position_embedding(x, cos, sin, mode=0):
     r"""
     Implements the Rotary Position Embedding algorithm.
@@ -6395,6 +6342,30 @@ def as_strided(input, size, stride, storage_offset=0):
     return as_strided_op(input, size, stride, storage_offset)
 
 
+def topprouter(input, capacity, expert_num, drop_type=0, threshold=0.0, router_prob=0.0):
+    r"""
+    TopPRouter implementation in MOE.
+
+    Inputs:
+        - **x** (Tensor) - Input Tensor of 3D, supporting types:[int32, int64]
+        - **capacity** (Int64) - The maximum number of tokens each expert can handle.
+        - **expert_num** (Int64) - The number of expert.
+        - **drop_type** (Int64) - S-Drop/K-Drop, 0 means S-Drop, 1 means K-Drop, default 0.
+        - **threshold** (float32) - Expert threshold, default 0.
+        - **router_prob** (Tensor) - Topk prob Tensor of 2D, supporting types:[float32], default 0.
+
+    Outputs:
+        tuple(Tensor), tuple of 2 tensors, `dispatch_index` and `combine_inex`.
+
+        - dispatch_index (Tensor) - Token ID processed by each expert.
+        - combine_index (Tensor) - The combine index of each token.
+
+    Supported Platforms:
+        ``Ascend``
+    """
+    return topprouter_op(input, capacity, expert_num, drop_type, threshold, router_prob)
+
+
 def asin_ext(input):
     r"""
     Computes arcsine of input tensors element-wise.
@@ -6475,7 +6446,7 @@ def ihfft(input, n=None, dim=-1, norm=None):
 
     Note:
         - `ihfft` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `ihfft` is not supported on Windows platform yet.
 
     Args:
@@ -6571,11 +6542,10 @@ def moe_token_permute(tokens, indices, num_out_tokens=None, padded_mode=False):
 
     .. warning::
         - It is only supported on Atlas A2 Training Series Products.
-        - The input `tokens` only supports the bfloat16 data type in the current version.
         - When `indices` is 2-D, the size of the second dim must be less than or equal to 512.
 
     Args:
-        tokens (Tensor): The input token tensor to be permuted. The dtype is bfloat16.
+        tokens (Tensor): The input token tensor to be permuted. The dtype is bfloat16, float16 or float32.
             The shape is :math:`(num\_tokens, hidden\_size)` , where `num_tokens` and `hidden_size` are positive integers.
         indices (Tensor): The tensor specifies indices used to permute the tokens. The dtype is int32 or int64.
             The shape is :math:`(num\_tokens, topk)` or :math:`(num\_tokens,)`, where `num_tokens` and `topk` are positive integers.
@@ -6591,7 +6561,6 @@ def moe_token_permute(tokens, indices, num_out_tokens=None, padded_mode=False):
 
     Raises:
         TypeError: If `tokens` or `indices` is not a Tensor.
-        TypeError: If dtype of `tokens` is not bfloat16.
         TypeError: If dtype of `indices` is not int32 or int64.
         TypeError: If specified `num_out_tokens` is not an integer.
         TypeError: If specified `padded_mode` is not a bool.
@@ -7838,6 +7807,13 @@ def nextafter(input, other):
     return next_after_op(input, other)
 
 
+def inner_moe_token_unpermute(permuted_tokens, sorted_indices, probs=None, padded_mode=False, restore_shape=None):
+    r"""
+    
+    """
+    return inner_moe_token_unpermute_op(permuted_tokens, sorted_indices, probs, padded_mode, restore_shape)
+
+
 def fill_diagonal_(input, fill_value, wrap=False):
     r"""
     
@@ -8683,7 +8659,7 @@ def irfftn(input, s=None, dim=None, norm=None):
 
     Note:
         - `irfftn` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `irfftn` is not supported on Windows platform yet.
 
     Args:
@@ -8906,7 +8882,7 @@ def fft(input, n=None, dim=-1, norm=None):
 
     Note:
         - `fft` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `fft` is not supported on Windows platform yet.
 
     Args:
@@ -8962,7 +8938,7 @@ def solve_triangular(a, b, trans=0, lower=False, unit_diagonal=False):
 
     Note:
         - `solve_triangular` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `solve_triangular` is not supported on Windows platform yet.
 
     Args:
@@ -9067,7 +9043,7 @@ def ifft(input, n=None, dim=-1, norm=None):
 
     Note:
         - `ifft` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `ifft` is not supported on Windows platform yet.
 
     Args:
@@ -9123,7 +9099,7 @@ def ifftshift(input, dim=None):
 
     Note:
         - `ifftshift` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `ifftshift` is not supported on Windows platform yet.
 
     Args:
@@ -9363,7 +9339,7 @@ def irfft2(input, s=None, dim=(-2, -1), norm=None):
 
     Note:
         - `irfft2` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `irfft2` is not supported on Windows platform yet.
 
     Args:
@@ -9539,7 +9515,7 @@ def correlate(a, v, pad_mode='valid'):
 
     Note:
         - `correlate` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `correlate` is not supported on Windows platform yet.
 
     Args:
@@ -9593,7 +9569,7 @@ def rfftn(input, s=None, dim=None, norm=None):
 
     Note:
         - `rfftn` is currently only used in `mindscience` scientific computing scenarios and
-          dose not support other usage scenarios.
+          does not support other usage scenarios.
         - `rfftn` is not supported on Windows platform yet.
 
     Args:
@@ -9719,6 +9695,51 @@ def neg(input):
     return neg_op(input)
 
 
+def matmul_add_(x, weight, C):
+    r"""
+    Fusion Operator of Transpose, Matmul, and InplaceAdd.
+
+    .. warning::
+        - This is an experimental API that is subject to change or deletion.
+        - This API is only supported in Atlas A2 training series for now.
+        - This API is only supported on GRAPH mode.
+
+    Args:
+        x (Tensor): Matrix A in matrix multiplication, with shape :math:`(k, m)` or :math:`(batch, k, m)`,
+            whose type should be float16 or bfloat16.
+        weight (Tensor): Matrix B in matrix multiplication, with shape :math:`(k, n)` or :math:`(batch, k, n)`,
+            whose type should be float16 or bfloat16.
+        C (Tensor): A Tensor acting as both input and output, with type of float32.
+            It's shape should be :math:`(m, n)` or :math:`(batch, m, n)`.
+
+    Returns:
+        Tensor, has the same shape and data type as `C`.
+
+    Raises:
+        TypeError: If the dtype of `weight` is not the same as `x`.
+        ValueError: If the ranks of `x` , `weight` and `C` are not the same.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> import numpy as np
+        >>> from mindspore import Tensor, ops, nn, context
+        >>> context.set_context(mode=context.GRAPH_MODE, jit_config={"jit_level": "O0"})
+        >>> class Net(nn.Cell):
+        ...     def construct(self, x, weight, C):
+        ...        return ops.auto_generate.inplace_matmul_add_op(x, weight, C)
+        >>> x = Tensor(np.random.randn(10, 20), mindspore.float16)
+        >>> weight = Tensor(np.random.randn(10, 8), mindspore.float16)
+        >>> C = Tensor(np.random.randn(20, 8), mindspore.float32)
+        >>> output = Net()(x, weight, C)
+        >>> print(output.shape)
+        (20, 8)
+    """
+    return inplace_matmul_add_op(x, weight, C)
+
+
 def gather(input_params, input_indices, axis, batch_dims=0):
     r"""
     Returns the slice of the input tensor corresponding to the elements of `input_indices` on the specified `axis`.
@@ -9812,6 +9833,131 @@ def deepcopy(input_x):
     return identity_op(input_x)
 
 
+def moe_distribute_dispatch(x, expert_ids, ep_world_size, ep_rank_id, moe_expert_num, expert_scales=None, scales=None, x_active_mask=None, group_ep=None, group_tp=None, tp_world_size=0, tp_rank_id=0, expert_shard_type=0, shared_expert_num=0, shared_expert_rank_num=0, quant_mode=0, global_bs=0, expert_token_nums_type=0):
+    r"""
+    Performs token data quantization (optional) and parallel communication for Mixture of Experts (MoE).
+    When Tensor Parallelism (TP) communication exists, it first performs Expert Parallelism (EP) AllToAllV
+    communication followed by TP AllGatherV communication. Otherwise, only EP AllToAllV communication is performed.
+
+    Notes:
+        - A: Maximum tokens to dispatch per rank:
+            - For shared experts: A = BS * ep_world_size * shared_expert_num / shared_expert_rank_num
+            - For MoE experts:
+                - When global_bs = 0: A >= BS * ep_world_size * min(local_expert_num, K)
+                - When global_bs != 0: A >= global_bs * min(local_expert_num, K)
+        - H (hidden size): Dimension of each token's hidden state
+            - Ascend 910B: 0 < H <= 7168, must be multiple of 32
+            - Ascend 910_93: H = 7168
+        - BS (batch sequence size): Number of tokens processed per rank
+            - Ascend 910B: 0 < BS <= 256
+            - Ascend 910_93: 0 < BS <= 512
+        - K: Number of experts selected per token (0 < K <= 8 and K <= moe_expert_num)
+        - server_num: Number of server nodes (supports 2, 4, 8)
+        - local_expert_num: Number of experts per rank:
+            - Shared expert ranks: local_expert_num = 1
+            - MoE expert ranks: local_expert_num = moe_expert_num / (ep_world_size - shared_expert_rank_num)
+            (TP communication not supported when localExpertNum > 1)
+
+    Inputs:
+        - **x** (Tensor) - Input token data to be sent. 2D tensor with shape [BS, H].
+            Supported dtypes: float16, bfloat16. Format: ND, non-contiguous allowed.
+        - **expert_ids** (Tensor) - Top-K expert indices for each token. 2D int32 tensor with shape [BS, K].
+            Format: ND, non-contiguous allowed.
+        - **ep_world_size** (int64) - EP domain size.
+            - Ascend 910B: Supports 16, 32, 64.
+            - Ascend 910_93: Supports 8, 16, 32, 64, 128, 144, 256, 288.
+        - **ep_rank_id** (int64) - Local rank ID in EP domain [0, ep_world_size), must be unique per domain.
+        - **moe_expert_num** (int64) - Number of MoE experts (0 < moe_expert_num <= 256),
+            must satisfy moe_expert_num % (ep_world_size-shared_expert_rank_num) = 0.
+        - **expert_scales** (Tensor) - Top-K expert weights per token.
+            - Ascend 910B: 2D float32 tensor [BS, K], ND format, non-contiguous allowed.
+            - Ascend 910_93: Unsupported (pass nullptr).
+        - **scales** (Tensor) - Expert weights. 2D float32 tensor with shape [shared_expert_num + moe_expert_num, H].
+            Pass nullptr for non-quantized scenarios. Format: ND, non-contiguous allowed.
+            Note: On Ascend 910B, must be nullptr when HCCL_INTRA_PCIE_ENABLE=1 and HCCL_INTRA_ROCE_ENABLE=0.
+        - **x_active_mask** (Tensor) - Reserved parameter (pass nullptr in current version).
+        - **group_ep** (str) - EP communication domain name (string length 1-127), must differ from group_tp.
+        - **group_tp** (str) - TP communication domain name.
+            - Ascend 910B: Unsupported (pass empty string).
+            - Ascend 910_93: When TP communication exists, string length 1-127, must differ from group_ep.
+        - **tp_world_size** (int64) - TP domain size.
+            - Ascend 910B: Unsupported (pass 0).
+            - Ascend 910_93: 0/1 means no TP communication; only 2 supported when TP exists.
+        - **tp_rank_id** (int64) - Local rank ID in TP domain.
+            - Ascend 910B: Unsupported (pass 0).
+            - Ascend 910_93: [0,1], unique per domain; pass 0 when no TP communication.
+        - **expert_shard_type** (int64) - Shared expert distribution type.
+            - Ascend 910B: Unsupported (pass 0).
+            - Ascend 910_93: Currently only 0 (shared experts precede MoE experts).
+        - **shared_expert_num** (int64) - Number of shared experts.
+            - Ascend 910B: Unsupported (pass 0).
+            - Ascend 910_93: Currently 0 (none) or 1 (one shared expert).
+        - **shared_expert_rank_num** (int64) - Number of ranks hosting shared experts.
+            - Ascend 910B: Unsupported (pass 0).
+            - Ascend 910_93: [0, ep_world_size-1), must satisfy ep_world_size % shared_expert_rank_num = 0 when non-zero.
+        - **quant_mode** (int64) - Quantization mode: 0 (none), 2 (dynamic quantization).
+        - **global_bs** (int64) - Global batch size across EP domain.
+            - Ascend 910B: 256*ep_world_size when BS varies per rank; 0 or BS*ep_world_size when uniform.
+            - Ascend 910_93: 0 or BS*ep_world_size.
+        - **expert_token_nums_type** (int64) - Semantic meaning of expert_token_nums output:
+            0 (prefix sums), 1 (raw counts).
+
+    Outputs:
+        - **expand_x** (Tensor) - Expanded token features. 2D tensor [A, H] with dtype matching input.
+            Supported dtypes: float16, bfloat16, int8. Format: ND, non-contiguous allowed.
+        - **dynamic_scales** (Tensor) - Dynamic quantization scales (when quant_mode=2).
+            1D float32 tensor [A]. Format: ND, non-contiguous allowed.
+        - **expand_idx** (Tensor) - Token counts per expert for combine operation.
+            1D int32 tensor [BS*K]. Format: ND, non-contiguous allowed.
+        - **expert_token_nums** (Tensor) - Tokens received per expert.
+            1D int64 tensor [local_expert_num]. Format: ND, non-contiguous allowed.
+        - **ep_recv_counts** (Tensor) - Tokens received from each EP rank.
+            - Ascend 910B: 1D int32 tensor [moe_expert_num + 2 * global_bs * K * server_num]
+            - Ascend 910_93: 1D int32 tensor [ep_world_size * max(tp_world_size,1) * local_expert_num]
+            Format: ND, non-contiguous allowed.
+        - **tp_recv_counts** (Tensor) - Tokens received from each TP rank (when TP exists).
+            - Ascend 910B: Not supported.
+            - Ascend 910_93: 1D int32 tensor [tp_world_size] when TP exists. Format: ND, non-contiguous allowed.
+        - **expand_scales** (Tensor) - Output token weights for combine operation.
+            - Ascend 910B: 1D float32 tensor [A]. Format: ND, non-contiguous allowed.
+            - Ascend 910_93: Unsupported.
+
+    Raises:
+        TypeError: If input dtypes don't match specifications.
+        ValueError: If input values violate constraints (e.g., invalid expert indices).
+        RuntimeError: If communication domain configuration is invalid.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+    >>> # EP-only communication example (Ascend 910B)
+    >>> import mindspore as ms
+    >>> from mindspore import Tensor
+    >>> from mindspore import ops
+    >>> from mindspore.communication import init, get_rank, GlobalComm
+    >>> from mindspore.ops.auto_generate import moe_distribute_dispatch
+    >>> import numpy as np
+    >>> bs = 8
+    >>> h = 7168
+    >>> k = 8
+    >>> ep_world_size = 16
+    >>> moe_expert_num = 16
+    >>> global_bs = bs * ep_world_size
+    >>> x = Tensor(np.random.randn(bs, h), ms.float16)
+    >>> expert_ids = Tensor(np.random.randint(0, moe_expert_num, (bs, k)), ms.int32)
+    >>> expert_scales = Tensor(np.random.randn(bs, k), ms.float32)
+    >>> init()
+    >>> rank_id = get_rank()
+    >>> out = moe_distribute_dispatch(
+    ...     x, expert_ids, ep_world_size, rank_id, moe_expert_num, expert_scales=expert_scales,
+    ...     group_ep=GlobalComm.WORLD_COMM_GROUP)
+    >>> print(out[0].shape)  # expand_x
+    (128, 7168)
+    """
+    return moe_distribute_dispatch_op(x, expert_ids, ep_world_size, ep_rank_id, moe_expert_num, expert_scales, scales, x_active_mask, group_ep, group_tp, tp_world_size, tp_rank_id, expert_shard_type, shared_expert_num, shared_expert_rank_num, quant_mode, global_bs, expert_token_nums_type)
+
+
 def pow(input, exponent):
     r"""
     Calculate the `exponent` power of each element in `input`.
@@ -10006,6 +10152,13 @@ def embedding_feature_mapping_v2(table_name, feature_id, table_total_size, table
     return embedding_feature_mapping_v2_op(table_name, feature_id, table_total_size, table_actual_size)
 
 
+def quant_matmul(x1, x2, scale, offset=None, pertoken_scale=None, bias=None, output_dtype=None, x1_dtype=None, x2_dtype=None, pertoken_scale_dtype=None, scale_dtype=None, group_sizes=None):
+    r"""
+    
+    """
+    return quant_matmul_op(x1, x2, scale, offset, pertoken_scale, bias, output_dtype, x1_dtype, x2_dtype, pertoken_scale_dtype, scale_dtype, group_sizes)
+
+
 def kv_cache_scatter_update(var, indices, updates, axis, reduce='none'):
     r"""
     Update var with updates and indices along sequence axis.
@@ -10483,3 +10636,56 @@ def grouped_matmul_v2(x, weight, bias=None, scale=None, offset=None, antiquant_s
     
     """
     return grouped_matmul_v2_op(x, weight, bias, scale, offset, antiquant_scale, antiquant_offset, group_list, split_item, group_type)
+
+
+def moe_token_unpermute(permuted_tokens, sorted_indices, probs=None, padded_mode=False, restore_shape=None):
+    r"""
+    Unpermute a tensor of permuted tokens based on sorted indices, and optionally merge the tokens with their corresponding probabilities.
+
+    .. warning::
+        - It is only supported on Atlas A2 Training Series Products.
+        - `sorted_indices` must not have duplicate values, otherwise the result is undefined.
+
+    Args:
+        permuted_tokens (Tensor): The tensor of permuted tokens to be unpermuted.
+            The shape is :math:`[num\_tokens * topk, hidden\_size]` , where `num_tokens`, `topk` and `hidden_size` are positive integers.
+        sorted_indices (Tensor): The tensor of sorted indices used to unpermute the tokens.
+            The shape is :math:`[num\_tokens * topk,]`, where `num_tokens` and `topk` are positive integers.
+            It only supports the int32 data type.
+        probs (Tensor, optional): The tensor of probabilities corresponding to the permuted tokens.
+            If provided, the unpermuted tokens will be merged with their respective probabilities.
+            The shape is :math:`[num\_tokens, topk]`, where `num_tokens` and `topk` are positive integers. Default: ``None`` .
+        padded_mode (bool, optional): If ``True``, indicating the indices are padded to denote selected tokens per expert. Default: ``False`` .
+        restore_shape (Union[tuple[int], list[int]], optional): The input shape before permutation, only used in padding mode. Default: ``None`` .
+
+    Returns:
+        Tensor, with the same dtype as `permuted_tokens`. If `padded_mode` is ``False``, the shape will be [`num_tokens`, `hidden_size`].
+        If `padded_mode` is ``True``, the shape will be specified by `restore_shape`.
+
+    Raises:
+        TypeError: If `permuted_tokens` is not a Tensor.
+        ValueError: Only supported when `padded_mode` is ``False``.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, ops
+        >>> permuted_token = Tensor([
+        ...                          [1, 1, 1],
+        ...                          [0, 0, 0],
+        ...                          [0, 0, 0],
+        ...                          [3, 3, 3],
+        ...                          [2, 2, 2],
+        ...                          [1, 1, 1],
+        ...                          [2, 2, 2],
+        ...                          [3, 3, 3]], dtype=mindspore.bfloat16)
+        >>> sorted_indices = Tensor([0, 6, 7, 5, 3, 1, 2, 4], dtype=mindspore.int32)
+        >>> out = ops.moe_token_unpermute(permuted_token, sorted_indices)
+        >>> out.shape
+        (8, 3)
+        
+        
+    """
+    return moe_token_unpermute_op(permuted_tokens, sorted_indices, probs, padded_mode, restore_shape)