Description: This operator corresponds to the routing computation in the MoE model. It uses the computation result of the operator as the input and outputs the routing matrix,
[object Object], and other results for subsequent computation. This API has the following function changes based on . Select a proper API based on your actual requirements.- Added the drop mode. In this mode, the output is based on the number of tokens processed by each expert and equal to the value of·
[object Object]. If the output exceeds the value, exceeded tokens are dropped. If the output falls short of the value, 0s are padded. - Added the optional output
[object Object]in dropless mode. The output is based on the number of cumulative tokens (Cumsum) to be processed by each expert or the number of tokens (Count) to be processed by each expert. - Added the optional output
[object Object]in drop mode. The output is based on the number of tokens to be processed by each expert before the drop operation. - Deleted the input
[object Object].
Note: Routing computation is a phase in an MoE model. The MoE model consists of a group of expert models and a gating model. During computation, the input data is first used to calculate the k experts with the highest weights corresponding to each data element based on the gating network (including the MoeGatingTopKSoftmax operator). Then, the result is input to the MoeInitRouting operator to generate the routing matrix. In subsequent operations, each expert in the model processes the data that it should process based on the routing matrix and generates the corresponding output. The outputs of all experts are weighted and summed up to form the final prediction result.
- Added the drop mode. In this mode, the output is based on the number of tokens processed by each expert and equal to the value of·
Formula:
- Flatten the input
[object Object]with shape [numRows, k] or [numRows] into a row for sorting, and obtain[object Object]and the corresponding[object Object], where[object Object]indicates the number of tokens, and k indicates the number of experts. When[object Object]is 1D, k is[object Object].
- Use
[object Object]for location mapping to obtain[object Object].
- Sort tokens by expert in the order of
[object Object]. When[object Object]is set to[object Object], the number of tokens to be processed by each expert is the same as the value of[object Object]. The tokens that exceed the value of[object Object]are dropped, and the tokens that fall short of the value are padded with 0s. The[object Object]is obtained as follows:
- Collect statistics on the histogram of each expert in
[object Object]and perform Cumsum to obtain[object Object].
- Collect statistics on the histogram of each expert in
[object Object]to obtain[object Object].
- Flatten the input
Each operator has calls. First, [object Object] is called to obtain the workspace size required for computation and the executor that contains the operator computation process. Then, [object Object] is called to perform computation.
Parameters:
[object Object]- [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and [object Object]Atlas A3 training products/Atlas A3 inference products[object Object]: The data type of input
[object Object]supports INT32, and the value is a 2D shape [numRows, k]. - Ascend 950PR/Ascend 950DT: The expertIdx input supports the INT32 and INT64 data types. The shape must be 2D [numRows, k] or 1D [numRows]. When the shape is 1D, k is 1.
- [object Object]Atlas inference products[object Object]: The data type of input
[object Object]can be INT32. The value must be a 2D shape with size [numRows, k]. The[object Object]parameter supports only[object Object].
- [object Object]Atlas A2 training products/Atlas A2 inference products[object Object] and [object Object]Atlas A3 training products/Atlas A3 inference products[object Object]: The data type of input
Returns:
aclnnStatus: status code. For details, see [object Object]aclnn Return Codes[object Object].
The first-phase API implements input parameter verification. The following errors may be thrown.
[object Object]
- Deterministic computation:
[object Object]defaults to deterministic implementation.
The following example is for reference only. For details, see .