aclnnMoeDistributeCombineV3

Applicable Products

Product	Supported
Ascend 950PR/Ascend 950DT	√
Atlas A3 training products/Atlas A3 inference products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas 200I/500 A2 inference products	×
Atlas inference products	×
Atlas training products	×

Function

Description: When TP domain communication exists, first performs ReduceScatterV communication, then AllToAllV communication, and finally integrates the received data (multiplies by weights and sums). When TP domain communication does not exist, performs AllToAllV communication, and finally integrates the received data (multiplies by weights and sums).
Formula:
- When TP domain communication does not exist:
$ataOut = AllToAllV(expandX)\\ xOut = Sum(expertScales * ataOut + expertScales * sharedExpertX)$
- When TP domain communication exists:
$rsOut = ReduceScatterV(expandX)\\ ataOut = AllToAllV(rsOut)\\ xOut = Sum(expertScales * ataOut + expertScales * sharedExpertX)$

Note: This API must be used together with aclnnMoeDistributeDispatchV3. It is equivalent to returning data along the same path as the data collection in the aclnnMoeDistributeDispatchV3 API.

Compared with the aclnnMoeDistributeCombineV2 API, this API has the following changes:

Added support for special expert scenarios:
- zeroExpertNum≠0: Enabled by passing a zeroExpertNum parameter greater than 0.
  
  $Moe(oriXOptional) = 0$
- copyExpertNum≠0: Enabled by passing a copyExpertNum parameter greater than 0, and requires a valid oriXOptional parameter.
  
  $Moe(oriXOptional) = oriXOptional$
- constExpertNum≠0: Enabled by passing a constExpertNum parameter greater than 0, and valid oriXOptional, constExpertAlpha1Optional, constExpertAlpha2Optional, and constExpertVOptional parameters must be passed.
  
  $Moe(oriXOptional) = constExpertAlpha1Optional * oriXOptional + constExpertAlpha2Optional * constExpertVOptional$

Function Prototype

Each operator uses a two-phase API. You must first call the "aclnnMoeDistributeCombineV3GetWorkspaceSize" API to obtain the workspace size required for computation and the executor that contains the operator computation process, and then call the "aclnnMoeDistributeCombineV3" API to perform the computation.

aclnnStatus aclnnMoeDistributeCombineV3GetWorkspaceSize(
    const aclTensor* expandX,
    const aclTensor* expertIds,
    const aclTensor* assistInfoForCombine,
    const aclTensor* epSendCounts,
    const aclTensor* expertScales,
    const aclTensor* tpSendCountsOptional,
    const aclTensor* xActiveMaskOptional,
    const aclTensor* activationScaleOptional,
    const aclTensor* weightScaleOptional,
    const aclTensor* groupListOptional,
    const aclTensor* expandScalesOptional,
    const aclTensor* sharedExpertXOptional,
    const aclTensor* elasticInfoOptional,
    const aclTensor* oriXOptional,
    const aclTensor* constExpertAlpha1Optional,
    const aclTensor* constExpertAlpha2Optional,
    const aclTensor* constExpertVOptional,
    const char*      groupEp,
    int64_t          epWorldSize,
    int64_t          epRankId,
    int64_t          moeExpertNum,
    const char*      groupTp,
    int64_t          tpWorldSize,
    int64_t          tpRankId,
    int64_t          expertShardType,
    int64_t          sharedExpertNum,
    int64_t          sharedExpertRankNum,
    int64_t          globalBS,
    int64_t          outDtype,
    int64_t          commQuantMode,
    int64_t          groupListType,
    const char*      commAlg,
    int64_t          zeroExpertNum,
    int64_t          copyExpertNum,
    int64_t          constExpertNum,
    aclTensor*       xOut,
    uint64_t*        workspaceSize,
    aclOpExecutor**  executor)

aclnnStatus aclnnMoeDistributeCombineV3(
    void*           workspace,
    uint64_t        workspaceSize,
    aclOpExecutor*  executor,
    aclrtStream     stream)

aclnnMoeDistributeCombineV3GetWorkspaceSize

Parameters

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Shape	Non-contiguous Tensor
expandX	Input	Token features expanded based on expertIds.	Requires a 2D tensor.	FLOAT16, BFLOAT16	ND	(max(tpWorldSize, 1) * A , H)	√
expertIds	Input	TopK expert indices for each token.	Must be a 2D tensor.	INT32	ND	(BS, K)	√
assistInfoForCombine	Input	Corresponds to the `assistInfoForCombineOut` output of `aclnnMoeDistributeDispatchV3`.	Must be a 1D tensor.	INT32	ND	(A * 128, )	√
epSendCounts	Input	Corresponds to the `epRecvCounts` output of `aclnnMoeDistributeDispatchV3`.	Requires a 1D tensor.	INT32	ND	-	√
expertScales	Input	Weights of the topK experts for each token.	Requires a 2D tensor.	FLOAT32	ND	`(BS, K)`	√
tpSendCountsOptional	Input	Corresponds to the `tpRecvCounts` output of `aclnnMoeDistributeDispatchV3`.	Pass this parameter when TP domain communication is involved; otherwise, pass a null pointer.	INT32	ND	When TP communication is involved, the shape is `(tpWorldSize, )`.	√
xActiveMaskOptional	Input	Identifies whether a token participates in communication.	Specifies a 1D or 2D tensor. Either valid data or a null pointer can be passed in.	BOOL	ND	When the input is 1D, the shape is `(BS, )`; when the input is 2D, the shape is `(BS, K)`.	√
activationScaleOptional	Input	Reserved parameter.	Not supported in the current version. Pass a null pointer.	-	ND	-	-
weightScaleOptional	Input	Reserved parameter.	Not supported in the current version. Pass a null pointer.	-	ND	-	-
groupListOptional	Input	Reserved parameter.	Not supported in the current version. Pass a null pointer.	-	ND	-	-
expandScalesOptional	Input	Corresponds to the `expandScales` output of `aclnnMoeDistributeDispatchV3`.	-	FLOAT32	ND	`(A, )`	√
sharedExpertXOptional	Input	Indicates the tokens after shared expert computation.	-	FLOAT16, BFLOAT16	ND	`(BS, H)`	√
elasticInfoOptional	Input	Dynamic scale-in information of the EP communication domain.	-	INT32	ND	-	√
oriXOptional	Input	Represents token data that has not passed through the FFN (Feed-Forward Neural network).	This input data is required when copyExpert or constExpert is enabled. Either valid data or a null pointer can be passed. When `copyExpertNum` is not 0 or `constExpertNum` is not 0, valid input must be passed. When valid data is passed, it must be a 2D tensor whose data type is consistent with expandX.	FLOAT16, BFLOAT16	ND	`(BS, H)`	√
constExpertAlpha1Optional	Input	Specifies the computation coefficient required when the constExpert scenario is enabled.	-	FLOAT16, BFLOAT16	ND	-	√
constExpertAlpha2Optional	Input	Specifies the computation coefficient to be input when the constExpert scenario is enabled.	-	FLOAT16, BFLOAT16	ND	-	√
constExpertVOptional	Input	Device-side aclTensor, representing the computation coefficient required when the constExpert scenario is enabled.	-	FLOAT16, BFLOAT16	ND	-	√
groupEp	Input	Name of the EP communication domain for expert parallelism.	The string length ranges from [1, 128). It cannot be the same as groupTp.	STRING	-	-	-
epWorldSize	Input	Size of the EP communication domain.	-	INT64	-	-	-
epRankId	Input	ID of the local card in the EP domain.	Value range: [0, epWorldSize). The epRankId of each card in the same EP communication domain must be unique.	INT64	-	-	-
moeExpertNum	Input	Number of MoE experts.	-	INT64	-	-	-
groupTp	Input	Name of the TP communication domain (data parallelism).	-	STRING	-	-	-
tpWorldSize	Input	TP communication domain size.	-	INT64	-	-	-
tpRankId	Input	ID of the local card in the TP domain.	-	INT64	-	-	-
expertShardType	Input	Specifies the distribution type of shared expert cards.	-	INT64	-	-	-
sharedExpertNum	Input	Number of shared experts.	-	INT64	-	-	-
sharedExpertRankNum	Input	Number of shared expert cards.	-	INT64	-	-	-
globalBS	Input	Global batch size in the EP domain.	When the BS of each rank is consistent, `globalBS = BS * epWorldSize` or 0; When the BS of each rank is inconsistent, `globalBS = maxBS * epWorldSize`, where maxBS indicates the maximum BS of a single card.	INT64	-	-	-
outDtype	Input	Reserved parameter that specifies the output data type.	Not supported in the current version. Pass 0.	INT64	-	-	-
commQuantMode	Input	Communication quantization type.	Value range: 0 or 2. 0 indicates no quantization during communication, and 2 indicates int8 quantization during communication.	INT64	-	-	-
Group list type.	Input	Reserved parameter, in group list format.	Not supported in the current version. Pass 0.	INT64	-	-	-
commAlg	Input	Communication affinity memory layout algorithm.	-	STRING	-	-	-
zeroExpertNum	Input	Number of zero experts.	-	INT64	-	-	-
copyExpertNum	Input	Number of copy experts.	-	INT64	-	-	-
constExpertNum	Input	Number of constant experts.	-	INT64	-	-	-
xOut	Output	Processed tokens.	2D tensor whose data type and format are the same as those of expandX.	FLOAT16, BFLOAT16	ND	`(BS, H)`	-
workspaceSize	Output	Returns the workspace size to be requested on the device side.	-	UINT64	-	-	-
executor	Output	Returns the op executor that contains the operator computation process.	-	aclOpExecutor*	-	-	-

Atlas A2 training products/Atlas A2 inference products:

commAlg supports nullptr, "", "fullmesh", and "hierarchy". It is recommended to configure "hierarchy" with driver version ≥ 25.0.RC1.1. nullptr and "" select the algorithm based on HCCL environment variables (not recommended). "fullmesh" directly transmits tokens via RDMA. "hierarchy" sends data in two stages (intra-server and inter-server) to reduce cross-server data volume.
The shared expert scenario is not supported.
The shape of epSendCounts is (moeExpertNum + 2 * globalBS * K * serverNum, ), where K indicates the number of top-K experts. The first moeExpertNum values indicate the number of tokens received from each card in the EP communication domain, and the last 2 * globalBS * K * serverNum values are used to store the number of tokens and communication domain offsets that can be reduced in advance by combine before inter-server/intra-server communication. When globalBS is 0, it is calculated as BS * epWorldSize.
TP domain communication is not supported currently.
xActiveMaskOptional depends on the value of commAlg. For "fullmesh", it must be a 1D tensor with the shape (BS, ), and true must precede false (for example, {true, false, true} is invalid). For "hierarchy", it is not supported in the current version; pass a null pointer.
expandScalesOptional must be a 1D tensor with the shape (A, ).
sharedExpertXOptional is a reserved parameter. It is not supported in the current version. Pass a null pointer.
epWorldSize depends on the value of commAlg. "fullmesh" supports 2, 3, 4, 5, 6, 7, 8, 16, 32, 64, 128, 192, 256, and 384. "hierarchy" supports 16, 32, and 64.
moeExpertNum has a value range of (0, 512].
groupTp is not supported in the current version. Pass an empty string.
tpWorldSize is not supported in the current version. Pass 0.
tpRankId: Not supported in the current version. Pass 0.
expertShardType: Not supported in the current version. Pass 0.
sharedExpertNum: Not supported in the current version. Pass 0.
sharedExpertRankNum: Not supported in the current version. Pass 0.
commQuantMode: The value ranges from 0 or 2 (0 indicates no quantization, and 2 indicates int8 quantization). The value 2 is supported only when commAlg is "hierarchy" or when HCCL_INTRA_PCIE_ENABLE=1, HCCL_INTRA_ROCE_ENABLE=0, and the driver version ≥ 25.0.RC1.1.
expandScalesOptional must be a 1D tensor.
elasticInfoOptional is not supported in the current version. Pass a null pointer.
When commAlg is set to "hierarchy", oriXOptional is not supported in the current version. Pass a null pointer.
constExpertAlpha1Optional is a reserved parameter and is not supported in the current version. Pass a null pointer.
constExpertAlpha2Optional is a reserved parameter and is not supported in the current version. Pass a null pointer.
constExpertVOptional: Reserved parameter. Not supported in the current version. Pass a null pointer.
zeroExpertNum: When commAlg="fullmesh", the value range is [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. Valid zero expert IDs are in the range [moeExpertNum, moeExpertNum + zeroExpertNum).
copyExpertNum: When commAlg="fullmesh", the value range is [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. Valid copy expert IDs are in the range [moeExpertNum + zeroExpertNum, moeExpertNum + zeroExpertNum + copyExpertNum).
constExpertNum: Not supported in the current version. Pass 0.

Atlas A3 training products/Atlas A3 inference products:

commAlg is not supported in the current version. Pass a null pointer.
The shape of epSendCounts is (epWorldSize * max(tpWorldSize, 1) * localExpertNum, ).
When TP domain communication is present, tpSendCountsOptional is a 1D tensor with the shape (tpWorldSize, ).
xActiveMaskOptional must be a 1D or 2D tensor (shape (BS, ) for 1D, shape (BS, K) for 2D). For 1D, true values must precede false values. For 2D, if all K values corresponding to a token are false, the token does not participate in communication.
expandScalesOptional: Reserved parameter. Not supported in the current version. Pass a null pointer.
sharedExpertXOptional: Must be a 2D or 3D tensor (shape is (BS, H) for 2D; for 3D, the product of the first two dimensions equals BS and the third dimension equals H). This parameter is optional. When passed, sharedExpertRankNum must be 0.
epWorldSize: Value range: [2, 768].
moeExpertNum: Value range: (0, 1024].
groupTp: String length range: [0, 128). Cannot be the same as groupEp. Passing an empty string is supported only when there is no TP domain communication.
tpWorldSize: Value range: [0, 2]. 0 and 1 indicate no TP domain communication. When TP domain communication is enabled, only 2 is supported.
tpRankId: Value range: [0, 1]. The tpRankId of each card in the same TP communication domain must be unique. Pass 0 when there is no TP domain communication.
expertShardType: Currently only 0 is supported, indicating that shared expert cards are placed before MoE expert cards.
sharedExpertNum: Current value range: [0, 4].
sharedExpertRankNum: Value range: [0, epWorldSize). When set to 0, sharedExpertNum must be 0 or 1. When not 0, sharedExpertRankNum % sharedExpertNum = 0 must hold.
commQuantMode: The value range is 0 or 2 (0 indicates no quantization, and 2 indicates int8 quantization). The value 2 can be enabled only when tpWorldSize < 2.
expandScalesOptional: Reserved parameter. It is not supported in the current version. Pass a null pointer.
elasticInfoOptional: It is not supported in the current version. Pass a null pointer.
constExpertAlpha1Optional: You can choose to pass valid data or a null pointer. When constExpertNum is not 0, valid input must be passed. When valid data is passed, it must be a 2D tensor with the shape (constExpertNum, H), and the data type must be consistent with expandX.
constExpertAlpha2Optional: You can choose to pass valid data or a null pointer. When constExpertNum is not 0, valid input must be passed. When valid data is passed, it must be a 2D tensor with the shape (constExpertNum, H), and the data type must be consistent with expandX.
constExpertVOptional: You can choose to pass valid data or a null pointer. When constExpertNum is not 0, you must pass valid input. When valid data is passed, it must be a 2D tensor with shape (constExpertNum, H), and its data type must be consistent with expandX.
zeroExpertNum: Value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. Valid zero expert IDs are in the range [moeExpertNum, moeExpertNum + zeroExpertNum).
copyExpertNum: Value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. Valid copy expert IDs are in the range [moeExpertNum + zeroExpertNum, moeExpertNum + zeroExpertNum + copyExpertNum).
constExpertNum: Value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. Valid constant expert IDs are in the range [moeExpertNum + zeroExpertNum + copyExpertNum, moeExpertNum + zeroExpertNum + copyExpertNum + constExpertNum).

Ascend 950PR/Ascend 950DT:

commAlg is not supported in the current version. Pass a null pointer.
The shape of epSendCounts is (epWorldSize max(tpWorldSize, 1) localExpertNum, ).
TP communication domain is not supported.
xActiveMaskOptional must be a 1D or 2D tensor (shape: (BS, ) for 1D, (BS, K) for 2D). For 1D, true values must precede false values. For 2D, if all K values corresponding to a token are false, the token does not participate in communication.
expandScalesOptional: Reserved parameter. Not supported in the current version. Pass a null pointer.
sharedExpertXOptional: Must be a 2D or 3D tensor (shape (BS, H) for 2D; for 3D, the product of the first two dimensions equals BS and the third dimension equals H). This parameter is optional. When passed, sharedExpertRankNum must be 0.
epWorldSize: Value range: [2, 768].
moeExpertNum: Value range: (0, 1024].
groupTp: Not supported in the current version. Pass an empty string.
tpWorldSize is not supported in the current version. Pass 0.
tpRankId is not supported in the current version. Pass 0.
expertShardType currently supports only 0, indicating that shared expert cards are placed before MoE expert cards.
sharedExpertNum currently ranges from 0 to 4.
sharedExpertRankNum ranges from 0 to epWorldSize). When set to 0, sharedExpertNum must be 0 or 1. When not 0, sharedExpertRankNum % sharedExpertNum must equal 0.
commQuantMode has a value range of 0 or 2 (0 indicates no quantization, 2 indicates int8 quantization). The value 2 can be enabled only when tpWorldSize < 2.
expandScalesOptional is a reserved parameter. It is not supported in the current version. Pass a null pointer.
elasticInfoOptional is a reserved parameter. It is not supported in the current version. Pass a null pointer.
constExpertAlpha1Optional can be set to valid data or a null pointer. When constExpertNum is not 0, a valid input must be passed. When valid data is passed, it must be a 2D tensor with the shape (constExpertNum, H), and its data type must be consistent with expandX.
constExpertAlpha2Optional can be set to valid data or a null pointer. When constExpertNum is not 0, a valid input must be passed. When valid data is passed, it must be a 2D tensor with the shape (constExpertNum, H), and its data type must be consistent with expandX.
constExpertVOptional: Optionally accepts valid data or a null pointer. When constExpertNum is not 0, a valid input must be provided. When valid data is provided, it must be a 2D tensor with shape (constExpertNum, H), and the data type must be consistent with expandX.
zeroExpertNum: Value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. Valid zero-expert IDs are in the range [moeExpertNum, moeExpertNum + zeroExpertNum).
copyExpertNum: Value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. Valid copy-expert IDs are in the range [moeExpertNum + zeroExpertNum, moeExpertNum + zeroExpertNum + copyExpertNum).
constExpertNum: Value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. Valid constant-expert IDs are in the range [moeExpertNum + zeroExpertNum + copyExpertNum, moeExpertNum + zeroExpertNum + copyExpertNum + constExpertNum).

Return Value

aclnnStatus: return code. For details, see aclnn Return Code.

The first-phase API performs input parameter validation and returns an error in the following scenarios:

Return Value	Error Code	Description
ACLNN_ERR_PARAM_NULLPTR	161001	The required input or output parameter Tensor is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data type of the input or output is not within the supported range.
ACLNN_ERR_INNER_TILING_ERROR	561002	Input or output shapes outside the supported range.
ACLNN_ERR_INNER_TILING_ERROR	561002	Parameter values outside the supported range.

aclnnMoeDistributeCombineV3

Parameters

Parameter	Input/Output	Description
workspace	Input	Workspace memory address applied on the Device side.
workspaceSize	Input	Workspace size applied on the Device side, obtained from the first-phase API aclnnMoeDistributeCombineV3GetWorkspaceSize.
executor	Input	Operator executor that contains the operator computation flow.
stream	Input	Specifies the stream on which the task is executed.

Return Value

aclnnStatus: return code. For details, see aclnn Return Code.

Constraints

Deterministic computation:
- aclnnMoeDistributeCombineV3 defaults to a deterministic implementation.
Driver Constraints:
- The driver versions of all nodes in the operator communication domain must be the same.
API Matching Constraints:
- aclnnMoeDistributeDispatchV3 and aclnnMoeDistributeCombineV3 must be used together. The assistInfoForCombineOut, epRecvCountsOut, tpRecvCountsOut, and expandScalesOut output by the former must be directly passed to the corresponding parameters of the latter. The service logic must not rely on the specific values of these tensors.
Parameter Consistency Constraints:
- The groupEp, epWorldSize, moeExpertNum, groupTp, tpWorldSize, expertShardType, sharedExpertNum, sharedExpertRankNum, globalBS, and commAlg parameters, as well as the HCCL_BUFFSIZE value, of all cards must be consistent, and consistent with the corresponding parameters of aclnnMoeDistributeDispatchV3.
Product-Specific Constraints:
- Atlas A3 training products/Atlas A3 inference products: A single card contains two DIE (dies). "This card" in the parameter description refers to a single DIE.

Shape Variable Constraints:

Variable	Definition and Value Range
A	Maximum number of tokens to be distributed on the current card. The value range is as follows: For shared experts, A = BS * epWorldSize * sharedExpertNum / sharedExpertRankNum must be satisfied. For MoE experts, when globalBS is 0, A >= BS * epWorldSize * min(localExpertNum, K) must be satisfied; when globalBS is non-zero, A >= globalBS * min(localExpertNum, K) must be satisfied.
H	Indicates the hidden size: Atlas A2 training products/Atlas A2 inference products: Depending on the commAlg value, "fullmesh" supports (0, 7168] and must be an integer multiple of 32; "hierarchy" with driver version ≥ 25.0.RC1.1 supports (0, 10*1024] and must be an integer multiple of 32. Atlas A3 training products/Atlas A3 inference products: [1024, 8192].
BS	Number of final output tokens on the current card: Atlas A2 training products/Atlas A2 inference products: Depending on the commAlg value, "fullmesh" supports (0, 256]; "hierarchy" with driver version ≥ 25.0.RC1.1 supports (0, 512]. Atlas A3 training products/Atlas A3 inference products: 0 < BS ≤ 512.
K	Indicates the selection of top-K experts: 0 < K ≤ 16, and 0 < K ≤ moeExpertNum + zeroExpertNum + copyExpertNum + constExpertNum.
serverNum	Number of server nodes: Atlas A2 training products/Atlas A2 inference products: This variable is used only in the shape of this scenario. Only 2, 4, and 8 are supported.
localExpertNum	Number of local experts: For a shared expert card, localExpertNum = 1. For an MoE expert card, localExpertNum = moeExpertNum/(epWorldSize-sharedExpertRankNum). TP communication is not supported when localExpertNum > 1. Atlas A3 training products/Atlas A3 inference products: 0 < localExpertNum * epWorldSize ≤ 2048 must be satisfied.

Environment variable constraints:
- HCCL_BUFFSIZE: Before calling this API, check whether the value of the HCCL_BUFFSIZE environment variable is reasonable. This environment variable indicates the memory size occupied by a single communication domain, in MB. The default value is 200 MB when not configured.
  - Atlas A2 training products/Atlas A2 inference products:
    - commAlg configured as "" or nullptr: Selects the "fullmesh" or "hierarchy" formula based on the HCCL_INTRA_PCIE_ENABLE and HCCL_INTRA_ROCE_ENABLE environment variable configurations.
    - commAlg set to "fullmesh": Set the size requirement to = 2 * (BS * epWorldSize * min(localExpertNum, K) * H * sizeof(uint16) + 2MB).
    - commAlg set to "hierarchy": Set the size requirement to (≥ (moeExpertNum + epWorldSize / 4) * Align512(maxBS * (H * 2 + 16 * Align8(K))) * 1B + 8MB, where Align8(x) = ((x + 8 - 1) / 8) * 8, Align512(x) = ((x + 512 - 1) / 512) * 512).
  - Atlas A3 training products/Atlas A3 inference products:
    - Within the EP communication domain: Set the size requirement to >= 2 and >= 2 * (localExpertNum * maxBS * epWorldSize * Align512(Align32(2 * H) + 44) + (K + sharedExpertNum) * maxBS * Align512(2 * H)), where localExpertNum must use the number of experts on the local MoE expert card, and Align512(x) = ((x + 512 - 1) / 512) * 512, Align32(x) = ((x + 32 - 1) / 32) * 32.
    - Within the TP communication domain: Set the size requirement to >= A * (H * 2 + 128) * 2.
- HCCL_INTRA_PCIE_ENABLE/HCCL_INTRA_ROCE_ENABLE:
  - Atlas A2 training products/Atlas A2 inference products: This environment variable is no longer recommended. Configure commAlg to "hierarchy" instead.
Communication domain usage constraints:
- The aclnnMoeDistributeCombineV3 series operators and aclnnMoeDistributeDispatchV3 in a model support only the same EP communication domain, and no other operators are allowed in this communication domain.
- The aclnnMoeDistributeCombineV3 series operators and aclnnMoeDistributeDispatchV3 in a model support only the same TP communication domain or neither supports a TP communication domain. When a TP communication domain exists, no other operators are allowed in this communication domain.
- Atlas A3 training products/Atlas A3 inference products: Nodes within a communication domain must be within a supernode. Cross-supernode communication is not supported.
Networking Constraints:
- Atlas A2 training products/Atlas A2 inference products: In multi-server scenarios, only switch networking is supported. Direct dual-server networking is not supported.
Other Constraints:
- The "/" in the formula indicates exact division.
- moeExpertNum + zeroExpertNum + copyExpertNum + constExpertNum < MAX_INT32.

Example

Atlas A2 training products/Atlas A2 inference products:

This sample supports running the A2 operator in a single-node environment with a card count of [2, 8]. You can set EP_WORLD_SIZE_A2 to the card count in the sample code as needed, and modify moeExpertNum so that moeExpertNum is divisible by EP_WORLD_SIZE_A2.

Compile the operator: The operator compilation command is as follows. Both the moe_distribute_dispatch_v2 and moe_distribute_combine_v2 operators need to be compiled, and these two operators must be executed in pairs.
```
bash build.sh --pkg --soc=ascend910b --ops=moe_distribute_dispatch_v2,moe_distribute_combine_v2
```
Create the sample code for Atlas A2 training products/Atlas A2 inference products: After compilation, refer to the existing test_aclnn_moe_distribute_combine_v2.cpp file in the operator examples directory, and create a new test file test_aclnn_moe_distribute_combine_v3.cpp using the A2 sample code.
Operator execution sample: The following command executes the operator sample, which runs all sample code files in the examples directory.
```
bash build.sh --run_example --ops=moe_distribute_combine_v2 eager cust
```

Sample code for Atlas A2 training products/Atlas A2 inference products:

#include <thread>
#include <iostream>
#include <string>
#include <cstring>
#include <vector>
#include <memory>
#include <cstdio>
#include "acl/acl.h"
#include "hccl/hccl.h"
#include "aclnn/opdev/fp16_t.h"
#include "aclnnop/aclnn_moe_distribute_dispatch_v3.h"
#include "aclnnop/aclnn_moe_distribute_combine_v3.h"

#define CHECK_RET(cond, return_expr) \
    do {                             \
        if (!(cond)) {               \
            return_expr;             \
        }                            \
    } while (0)

#define LOG_PRINT(message, ...)         \
    do {                                \
        printf(message, ##__VA_ARGS__); \
    } while(0)

struct Args {
    uint32_t rankId;
    uint32_t epRankId;
    uint32_t tpRankId;
    HcclComm hcclEpComm;
    HcclComm hcclTpComm;
    aclrtStream dispatchV3Stream;
    aclrtStream combineV3Stream;
    aclrtContext context;
};

const uint32_t EP_WORLD_SIZE_A2 = 8;
const uint32_t TP_WORLD_SIZE_A2 = 1;
const uint32_t DEV_NUM_A2 = EP_WORLD_SIZE_A2 * TP_WORLD_SIZE_A2;

int64_t GetShapeSize(const std::vector<int64_t> &shape)
{
    int64_t shape_size = 1;
    for (auto i : shape) {
        shape_size *= i;
    }
    return shape_size;
}

template<typename T>
int CreateAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr,
    aclDataType dataType, aclTensor **tensor)
{
    auto size = GetShapeSize(shape) * sizeof(T);
    auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc failed. ret: %d\n", ret); return ret);
    ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMemcpy failed. ret: %d\n", ret); return ret);
    std::vector<int64_t> strides(shape.size(), 1);
    for (int64_t i = shape.size() - 2; i >= 0; i--) {
        strides[i] = shape[i +1] * strides[i + 1];
    }
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
        shape.data(), shape.size(), *deviceAddr);
    return 0;
}

void DestroyTensor(aclTensor *tensor) {
    if (tensor != nullptr) {
        aclDestroyTensor(tensor);
    }
}

void FreeDeviceAddr(void *deviceAddr) {
    if (deviceAddr != nullptr) {
        aclrtFree(deviceAddr);
    }
}

int launchOneThreadDispatchV3AndCombineV3_A2(Args &args)
{
    int ret = aclrtSetCurrentContext(args.context);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSetCurrentContext failed, ret %d\n", ret); return ret);
    char hcomEpName[128] = {0};
    ret = HcclGetCommName(args.hcclEpComm, hcomEpName);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetEpCommName failed, ret %d\n", ret); return -1);
    LOG_PRINT("[INFO] rank = %d, hcomEpName = %s, dispatchV3Stream = %p, combineV3Stream = %p, \
                context = %p\n", args.rankId, hcomEpName, args.dispatchV3Stream, args.combineV3Stream,                 \
                args.context);

    int64_t BS = 32;
    int64_t H = 7168;
    int64_t K = 8;
    int64_t expertShardType = 0;
    int64_t sharedExpertNum = 0;
    int64_t sharedExpertRankNum = 0;
    int64_t moeExpertNum = 256;
    int64_t quantMode = 0;
    int64_t globalBS = BS * EP_WORLD_SIZE_A2;
    int64_t expertTokenNumsType = 1;
    int64_t outDtype = 0;
    int64_t commQuantMode = 0;
    int64_t groupList_type = 1;
    int64_t localExpertNum;
    int64_t A;
    int64_t zeroExpertNum = 0;
    int64_t copyExpertNum = 0;
    int64_t constExpertNum = 0; // A3 only.
    std::string commAlg = "fullmesh";
    if (args.epRankId < sharedExpertRankNum) {
        localExpertNum = 1;
        A = globalBS / sharedExpertRankNum;
    } else {
        localExpertNum = moeExpertNum / (EP_WORLD_SIZE_A2 - sharedExpertRankNum);
        A = globalBS * (localExpertNum < K ? localExpertNum : K);
    }

    void *xDeviceAddr = nullptr;
    void *expertIdsDeviceAddr = nullptr;
    void *scalesDeviceAddr = nullptr;
    void *expertScalesDeviceAddr = nullptr;

    void *expandXDeviceAddr = nullptr;
    void *dynamicScalesDeviceAddr = nullptr;
    void *assistInfoForCombineDeviceAddr = nullptr;
    void *expertTokenNumsDeviceAddr = nullptr;
    void *epRecvCountsDeviceAddr = nullptr;
    void *tpRecvCountsDeviceAddr = nullptr;
    void *expandScalesDeviceAddr = nullptr;

    // Zero-expert scenario input.
    void *oriXDeviceAddr = nullptr;

    void *xOutDeviceAddr = nullptr;

    aclTensor *x = nullptr;
    aclTensor *expertIds = nullptr;
    aclTensor *scales = nullptr;
    aclTensor *xActiveMask = nullptr;
    aclTensor *expertScales = nullptr;

    aclTensor *elasticInfo = nullptr; // A3
    aclTensor *expandX = nullptr;
    aclTensor *dynamicScales = nullptr;
    aclTensor *assistInfoForCombine = nullptr; // expandIdx
    aclTensor *expertTokenNums = nullptr;
    aclTensor *epRecvCounts = nullptr;
    aclTensor *tpRecvCounts = nullptr;
    aclTensor *expandScales = nullptr;

    aclTensor *activationScale = nullptr; // Reserved parameter.
    aclTensor *weightScale = nullptr; // Reserved parameter.
    aclTensor *groupList = nullptr; // Reserved parameter.

    aclTensor *sharedExpertX = nullptr; // A3

    aclTensor *oriX = nullptr;
    aclTensor *constExpertAlpha1 = nullptr; // A3
    aclTensor *constExpertAlpha2 = nullptr; // A3
    aclTensor *constExpertV = nullptr; // A3

    aclTensor *xOut = nullptr;

    //Define the dimensions of each variable in the current scenario.
    std::vector<int64_t> xShape{BS, H};
    std::vector<int64_t> expertIdsShape{BS, K};
    std::vector<int64_t> scalesShape{moeExpertNum + 1, H};
    std::vector<int64_t> expertScalesShape{BS, K};

    std::vector<int64_t> expandXShape{TP_WORLD_SIZE_A2 * A, H};
    std::vector<int64_t> dynamicScalesShape{TP_WORLD_SIZE_A2 * A};
    std::vector<int64_t> assistInfoForCombineShape{A * 128};
    std::vector<int64_t> expertTokenNumsShape{localExpertNum};
    std::vector<int64_t> epRecvCountsShape{TP_WORLD_SIZE_A2 * localExpertNum * EP_WORLD_SIZE_A2}; // Non-hierarchical.
    std::vector<int64_t> tpRecvCountsShape{TP_WORLD_SIZE_A2};
    std::vector<int64_t> expandScalesShape{A};

    std::vector<int64_t> oriXShape{BS, H};
    std::vector<int64_t> xOutShape{BS, H};

    int64_t xShapeSize = GetShapeSize(xShape);
    int64_t expertIdsShapeSize = GetShapeSize(expertIdsShape);
    int64_t scalesShapeSize = GetShapeSize(scalesShape);
    int64_t expertScalesShapeSize = GetShapeSize(expertScalesShape);

    int64_t expandXShapeSize = GetShapeSize(expandXShape);
    int64_t dynamicScalesShapeSize = GetShapeSize(dynamicScalesShape);
    int64_t assistInfoForCombineShapeSize = GetShapeSize(assistInfoForCombineShape);
    int64_t expertTokenNumsShapeSize = GetShapeSize(expertTokenNumsShape);
    int64_t epRecvCountsShapeSize = GetShapeSize(epRecvCountsShape);
    int64_t tpRecvCountsShapeSize = GetShapeSize(tpRecvCountsShape);
    int64_t expandScalesShapeSize = GetShapeSize(expandScalesShape);

    int64_t oriXSize = GetShapeSize(oriXShape);

    int64_t xOutShapeSize = GetShapeSize(xOutShape);

    std::vector<int16_t> xHostData(xShapeSize, 1);
    std::vector<int32_t> expertIdsHostData;
    for (int32_t token_id = 0; token_id < expertIdsShape[0]; token_id++) {
        for (int32_t k_id = 0; k_id < expertIdsShape[1]; k_id++) {
            expertIdsHostData.push_back(k_id);
        }
    }

    std::vector<float> scalesHostData(scalesShapeSize, 0.1);
    std::vector<float> expertScalesHostData(expertScalesShapeSize, 0.1);

    std::vector<int16_t> expandXHostData(expandXShapeSize, 0);
    std::vector<float> dynamicScalesHostData(dynamicScalesShapeSize, 0);
    std::vector<int32_t> assistInfoForCombineHostData(assistInfoForCombineShapeSize, 0);
    std::vector<int64_t> expertTokenNumsHostData(expertTokenNumsShapeSize, 0);
    std::vector<int32_t> epRecvCountsHostData(epRecvCountsShapeSize, 0);
    std::vector<int32_t> tpRecvCountsHostData(tpRecvCountsShapeSize, 0);
    std::vector<float> expandScalesHostData(expandScalesShapeSize, 0);

    std::vector<int16_t> oriXHostData(oriXSize, 1);
    std::vector<int16_t> xOutHostData(xOutShapeSize, 0);


    ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_BF16, &x);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(expertIdsHostData, expertIdsShape, &expertIdsDeviceAddr, aclDataType::ACL_INT32, &expertIds);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(scalesHostData, scalesShape, &scalesDeviceAddr, aclDataType::ACL_FLOAT, &scales);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(expertScalesHostData, expertScalesShape, &expertScalesDeviceAddr, aclDataType::ACL_FLOAT, &expertScales);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    ret = CreateAclTensor(expandXHostData, expandXShape, &expandXDeviceAddr, (quantMode > 0) ? aclDataType::ACL_INT8 : aclDataType::ACL_BF16, &expandX);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(dynamicScalesHostData, dynamicScalesShape, &dynamicScalesDeviceAddr, aclDataType::ACL_FLOAT, &dynamicScales);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(assistInfoForCombineHostData, assistInfoForCombineShape, &assistInfoForCombineDeviceAddr, aclDataType::ACL_INT32, &assistInfoForCombine);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(expertTokenNumsHostData, expertTokenNumsShape, &expertTokenNumsDeviceAddr, aclDataType::ACL_INT64, &expertTokenNums);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(epRecvCountsHostData, epRecvCountsShape, &epRecvCountsDeviceAddr, aclDataType::ACL_INT32, &epRecvCounts);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(tpRecvCountsHostData, tpRecvCountsShape, &tpRecvCountsDeviceAddr, aclDataType::ACL_INT32, &tpRecvCounts);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(expandScalesHostData, expandScalesShape, &expandScalesDeviceAddr, aclDataType::ACL_FLOAT, &expandScales);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    ret = CreateAclTensor(oriXHostData, oriXShape, &oriXDeviceAddr, aclDataType::ACL_BF16, &oriX);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    ret = CreateAclTensor(xOutHostData, xOutShape, &xOutDeviceAddr, aclDataType::ACL_BF16, &xOut);
    CHECK_RET(ret == ACL_SUCCESS, return ret);


    uint64_t dispatchWorkspaceSize = 0;
    aclOpExecutor *dispatchExecutor = nullptr;
    void *dispatchWorkspaceAddr = nullptr;

    uint64_t combineWorkspaceSize = 0;
    aclOpExecutor *combineExecutor = nullptr;
    void *combineWorkspaceAddr = nullptr;

    /**************************************** Call dispatch ********************************************/
    // Call the first-phase API.
    ret = aclnnMoeDistributeDispatchV3GetWorkspaceSize(x, expertIds, (quantMode > 0 ? scales : nullptr), xActiveMask,
            expertScales, elasticInfo, hcomEpName, EP_WORLD_SIZE_A2, args.epRankId, moeExpertNum, "", TP_WORLD_SIZE_A2,
            args.tpRankId, expertShardType, sharedExpertNum,sharedExpertRankNum, quantMode, globalBS,
            expertTokenNumsType, commAlg.c_str(), zeroExpertNum, copyExpertNum, constExpertNum, expandX, dynamicScales, assistInfoForCombine, expertTokenNums, epRecvCounts,
            tpRecvCounts, expandScales, &dispatchWorkspaceSize, &dispatchExecutor);

    CHECK_RET(ret == ACL_SUCCESS,
        LOG_PRINT("[ERROR] aclnnMoeDistributeDispatchV3GetWorkspaceSize failed. ret = %d \n", ret); return ret);

    if (dispatchWorkspaceSize > 0) {
        ret = aclrtMalloc(&dispatchWorkspaceAddr, dispatchWorkspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc workspace failed. ret = %d \n", ret); return ret);
    }
    // Call the second-phase API.
    ret = aclnnMoeDistributeDispatchV3(dispatchWorkspaceAddr, dispatchWorkspaceSize,
                                        dispatchExecutor, args.dispatchV3Stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMoeDistributeDispatchV3 failed. ret = %d \n", ret);  \
            return ret);
    ret = aclrtSynchronizeStreamWithTimeout(args.dispatchV3Stream, 10000);
                CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] dispatch aclrtSynchronizeStreamWithTimeout failed. ret = %d \n", ret);  \
            return ret);
    LOG_PRINT("[INFO] device_%d aclnnMoeDistributeDispatchV3 execute successfully.\n", args.rankId);
    /**************************************** Call combine ********************************************/
    // Call the first-phase API.
    ret = aclnnMoeDistributeCombineV3GetWorkspaceSize(expandX, expertIds,
                                                        assistInfoForCombine, epRecvCounts,
                                                        expertScales, tpRecvCounts,
                                                        xActiveMask, activationScale, weightScale,
                                                        groupList, expandScales, sharedExpertX,
                                                        elasticInfo, oriX, constExpertAlpha1, constExpertAlpha2, constExpertV,
                                                        hcomEpName, EP_WORLD_SIZE_A2, args.epRankId, moeExpertNum,
                                                        "", TP_WORLD_SIZE_A2, args.tpRankId, expertShardType,
                                                        sharedExpertNum, sharedExpertRankNum, globalBS, outDtype,
                                                        commQuantMode, groupList_type, commAlg.c_str(), zeroExpertNum, copyExpertNum, constExpertNum, xOut,
                                                        &combineWorkspaceSize, &combineExecutor);
    CHECK_RET(ret == ACL_SUCCESS,
        LOG_PRINT("[ERROR] aclnnMoeDistributeCombineV3GetWorkspaceSize failed. ret = %d \n", ret); return ret);
    // Apply for device memory based on the workspaceSize calculated by the first-phase API.
    if (combineWorkspaceSize > 0) {
        ret = aclrtMalloc(&combineWorkspaceAddr, combineWorkspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc workspace failed. ret = %d \n", ret); return ret);
    }

    // Call the second-phase API.
    ret = aclnnMoeDistributeCombineV3(combineWorkspaceAddr, combineWorkspaceSize, combineExecutor, args.combineV3Stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMoeDistributeCombineV3 failed. ret = %d \n", ret);
        return ret);
    // (Boilerplate) Synchronously wait for task execution to complete.
    ret = aclrtSynchronizeStreamWithTimeout(args.combineV3Stream, 10000);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSynchronizeStreamWithTimeout failed. ret = %d \n", ret);
        return ret);
    LOG_PRINT("[INFO] device_%d aclnnMoeDistributeDispatchV3 and aclnnMoeDistributeCombineV3                      \
                execute successfully.\n", args.rankId);
    // Release device resources.
    if (dispatchWorkspaceSize > 0) {
        aclrtFree(dispatchWorkspaceAddr);
    }
    if (combineWorkspaceSize > 0) {
        aclrtFree(combineWorkspaceAddr);
    }
    DestroyTensor(x);
    DestroyTensor(expertIds);
    DestroyTensor(scales);
    DestroyTensor(xActiveMask);
    DestroyTensor(expertScales);
    DestroyTensor(elasticInfo);
    DestroyTensor(expandX);
    DestroyTensor(dynamicScales);
    DestroyTensor(assistInfoForCombine);
    DestroyTensor(expertTokenNums);
    DestroyTensor(epRecvCounts);
    DestroyTensor(tpRecvCounts);
    DestroyTensor(expandScales);
    DestroyTensor(activationScale);
    DestroyTensor(weightScale);
    DestroyTensor(groupList);
    DestroyTensor(sharedExpertX);
    DestroyTensor(oriX);
    DestroyTensor(constExpertAlpha1);
    DestroyTensor(constExpertAlpha2);
    DestroyTensor(constExpertV);
    DestroyTensor(xOut);

    FreeDeviceAddr(xDeviceAddr);
    FreeDeviceAddr(expertIdsDeviceAddr);
    FreeDeviceAddr(scalesDeviceAddr);
    FreeDeviceAddr(expertScalesDeviceAddr);
    FreeDeviceAddr(expandXDeviceAddr);
    FreeDeviceAddr(dynamicScalesDeviceAddr);
    FreeDeviceAddr(assistInfoForCombineDeviceAddr);
    FreeDeviceAddr(expertTokenNumsDeviceAddr);
    FreeDeviceAddr(epRecvCountsDeviceAddr);
    FreeDeviceAddr(tpRecvCountsDeviceAddr);
    FreeDeviceAddr(expandScalesDeviceAddr);
    FreeDeviceAddr(oriXDeviceAddr);
    FreeDeviceAddr(xOutDeviceAddr);

    HcclCommDestroy(args.hcclEpComm);
    aclrtDestroyStream(args.dispatchV3Stream);
    aclrtDestroyStream(args.combineV3Stream);
    aclrtDestroyContext(args.context);
    LOG_PRINT("[INFO] device_%d DeStroy.\n", args.rankId);
    aclrtResetDevice(args.rankId);
    LOG_PRINT("[INFO] device_%d Reset.\n", args.rankId);
    return 0;
}
int main(int argc, char *argv[])
{
    LOG_PRINT("[INFO] run_example_on_A2.\n");
    int ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclInit failed, ret = %d\n", ret); return ret);
    aclrtStream dispatchV3Stream[DEV_NUM_A2];
    aclrtStream combineV3Stream[DEV_NUM_A2];
    aclrtContext context[DEV_NUM_A2];
    for (uint32_t rankId = 0; rankId < DEV_NUM_A2; rankId++) {
        ret = aclrtSetDevice(rankId);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSetDevice failed, ret = %d\n", ret); return ret);
        ret = aclrtCreateContext(&context[rankId], rankId);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateContext failed, ret = %d\n", ret); return ret);
        ret = aclrtCreateStream(&dispatchV3Stream[rankId]);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateStream failed, ret = %d\n", ret); return ret);
        ret = aclrtCreateStream(&combineV3Stream[rankId]);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateStream failed, ret = %d\n", ret); return ret);
    }

    int32_t devicesEp[EP_WORLD_SIZE_A2];
    for (int32_t epId = 0; epId < EP_WORLD_SIZE_A2; epId++) {
        devicesEp[epId] = epId;
    }

    HcclComm commsEp[EP_WORLD_SIZE_A2];
    ret = HcclCommInitAll(EP_WORLD_SIZE_A2, devicesEp, commsEp);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclCommInitAll ep failed, ret %d\n", ret); return ret);

    Args args[DEV_NUM_A2];
    std::vector<std::unique_ptr<std::thread>> threads(DEV_NUM_A2);
    for (uint32_t rankId = 0; rankId < DEV_NUM_A2; rankId++) {
        uint32_t epRankId = rankId / TP_WORLD_SIZE_A2;
        uint32_t tpRankId = rankId % TP_WORLD_SIZE_A2;

        args[rankId].rankId = rankId;
        args[rankId].epRankId = epRankId;
        args[rankId].tpRankId = tpRankId;
        args[rankId].hcclEpComm = commsEp[epRankId];
        args[rankId].dispatchV3Stream = dispatchV3Stream[rankId];
        args[rankId].combineV3Stream = combineV3Stream[rankId];
        args[rankId].context = context[rankId];
        threads[rankId].reset(new(std::nothrow) std::thread(&launchOneThreadDispatchV3AndCombineV3_A2, std::ref(args[rankId])));
    }

    for(uint32_t rankId = 0; rankId < DEV_NUM_A2; rankId++) {
        threads[rankId]->join();
    }

    aclFinalize();
    LOG_PRINT("[INFO] aclFinalize success\n");
    return 0;
}

Ascend 950PR/Ascend 950DT: Refer to the preparation section and sample code in the example of aclnnMoeDistributeCombineV2. Reset the involved variables according to the constraints described above. For the scenario parameters newly added in the V4 API compared to the V3 API, pass values as described in the parameter description above.
Atlas A3 training products/Atlas A3 inference products:

For the detailed compilation and execution process, see Compile and Run Samples.

The sample code is as follows, for reference only

#include <thread>
#include <iostream>
#include <string>
#include <vector>
#include <unordered_set>
#include "acl/acl.h"
#include "hccl/hccl.h"
#include "aclnnop/aclnn_moe_distribute_dispatch_v3.h"
#include "aclnnop/aclnn_moe_distribute_combine_v3.h"

#define CHECK_RET(cond, return_expr) \
    do {                             \
        if (!(cond)) {               \
            return_expr;             \
        }                            \
    } while (0)

#define LOG_PRINT(message, ...)         \
    do {                                \
        printf(message, ##__VA_ARGS__); \
    } while(0)

struct Args {
    uint32_t rankId;
    uint32_t epRankId;
    uint32_t tpRankId;
    HcclComm hcclEpComm;
    HcclComm hcclTpComm;
    aclrtStream dispatchStream;
    aclrtStream combineStream;
    aclrtContext context;
};

constexpr uint32_t EP_WORLD_SIZE = 8;
constexpr uint32_t TP_WORLD_SIZE = 1;
constexpr uint32_t DEV_NUM = EP_WORLD_SIZE * TP_WORLD_SIZE;

int64_t GetShapeSize(const std::vector<int64_t> &shape)
{
    int64_t shape_size = 1;
    for (auto i : shape) {
        shape_size *= i;
    }
    return shape_size;
}

template<typename T>
int CreateAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr,
    aclDataType dataType, aclTensor **tensor)
{
    auto size = GetShapeSize(shape) * sizeof(T);
    auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc failed. ret: %d\n", ret); return ret);
    ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMemcpy failed. ret: %d\n", ret); return ret);
    std::vector<int64_t> strides(shape.size(), 1);
    for (int64_t i = shape.size() - 2; i >= 0; i--) {
        strides[i] = shape[i +1] * strides[i + 1];
    }
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
        shape.data(), shape.size(), *deviceAddr);
    return 0;
}

int LaunchOneProcessDispatchAndCombine(Args &args)
{
    int ret = aclrtSetCurrentContext(args.context);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSetCurrentContext failed, ret %d\n", ret); return ret);

    char hcomEpName[128] = {0};
    ret = HcclGetCommName(args.hcclEpComm, hcomEpName);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetEpCommName failed, ret %d\n", ret); return -1);
    char hcomTpName[128] = {0};
    ret = HcclGetCommName(args.hcclTpComm, hcomTpName);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetTpCommName failed, ret %d\n", ret); return -1);
    LOG_PRINT("[INFO] rank = %d, hcomEpName = %s, hcomTpName = %s, dispatchStream = %p, combineStream = %p, \
                context = %p\n", args.rankId, hcomEpName, hcomTpName, args.dispatchStream, args.combineStream,                 \
                args.context);

    int64_t BS = 8;
    int64_t H = 7168;
    int64_t K = 3;
    int64_t expertShardType = 0;
    int64_t sharedExpertNum = 1;
    int64_t sharedExpertRankNum = 1;
    int64_t moeExpertNum = 7;
    int64_t quantMode = 0;
    int64_t globalBS = BS * EP_WORLD_SIZE;
    int64_t expertTokenNumsType = 1;
    int64_t outDtype = 0;
    int64_t commQuantMode = 0;
    int64_t groupList_type = 1;
    int64_t localExpertNum;
    int64_t A;
    int64_t zeroExpertNum = 1;
    int64_t copyExpertNum = 1;
    int64_t constExpertNum = 1;
    if (args.epRankId < sharedExpertRankNum) {
        localExpertNum = 1;
        A = globalBS / sharedExpertRankNum;
    } else {
        localExpertNum = moeExpertNum / (EP_WORLD_SIZE - sharedExpertRankNum);
        A = globalBS * (localExpertNum < K ? localExpertNum : K);
    }

    void *xDeviceAddr = nullptr;
    void *expertIdsDeviceAddr = nullptr;
    void *scalesDeviceAddr = nullptr;
    void *expertScalesDeviceAddr = nullptr;

    void *expandXDeviceAddr = nullptr;
    void *dynamicScalesDeviceAddr = nullptr;
    void *expandIdxDeviceAddr = nullptr;
    void *expertTokenNumsDeviceAddr = nullptr;
    void *epRecvCountsDeviceAddr = nullptr;
    void *tpRecvCountsDeviceAddr = nullptr;
    void *expandScalesDeviceAddr = nullptr;
    void *residualXDeviceAddr = nullptr;
    void *sharedExpertXDeviceAddr = nullptr;

    void *elasticInfoDeviceAddr = nullptr;
    void *oriXDeviceAddr = nullptr;
    void *constExpertAlpha1DeviceAddr = nullptr;
    void *constExpertAlpha2DeviceAddr = nullptr;
    void *constExpertVDeviceAddr = nullptr;

    void *xOutDeviceAddr = nullptr;

    aclTensor *x = nullptr;
    aclTensor *expertIds = nullptr;
    aclTensor *scales = nullptr;
    aclTensor *expertScales = nullptr;

    aclTensor *expandX = nullptr;
    aclTensor *dynamicScales = nullptr;
    aclTensor *expandIdx = nullptr;
    aclTensor *expertTokenNums = nullptr;
    aclTensor *epRecvCounts = nullptr;
    aclTensor *tpRecvCounts = nullptr;
    aclTensor *expandScales = nullptr;
    aclTensor *residualX = nullptr;
    aclTensor *sharedExpertX = nullptr;


    aclTensor *elasticInfo = nullptr;
    aclTensor *oriX = nullptr;
    aclTensor *constExpertAlpha1 = nullptr;
    aclTensor *constExpertAlpha2 = nullptr;
    aclTensor *constExpertV = nullptr;

    aclTensor *xOut = nullptr;

    //Define the dimensions of each variable in the current scenario.
    std::vector<int64_t> xShape{BS, H};
    std::vector<int64_t> expertIdsShape{BS, K};
    std::vector<int64_t> scalesShape{moeExpertNum + 1, H};
    std::vector<int64_t> expertScalesShape{BS, K};

    std::vector<int64_t> expandXShape{TP_WORLD_SIZE * A, H};
    std::vector<int64_t> dynamicScalesShape{TP_WORLD_SIZE * A};
    std::vector<int64_t> expandIdxShape{A * 128};
    std::vector<int64_t> expertTokenNumsShape{localExpertNum};
    std::vector<int64_t> epRecvCountsShape{TP_WORLD_SIZE * localExpertNum * EP_WORLD_SIZE};
    std::vector<int64_t> tpRecvCountsShape{TP_WORLD_SIZE};
    std::vector<int64_t> expandScalesShape{A};
    std::vector<int64_t> sharedExpertXShape{BS, 1, H};


    std::vector<int64_t> elasticInfoShape{4 + EP_WORLD_SIZE * 2};
    std::vector<int64_t> oriXShape{BS, H};
    std::vector<int64_t> constExpertAlpha1Shape{constExpertNum, H};
    std::vector<int64_t> constExpertAlpha2Shape{constExpertNum, H};
    std::vector<int64_t> constExpertVShape{constExpertNum, H};

    std::vector<int64_t> xOutShape{BS, H};

    int64_t xShapeSize = GetShapeSize(xShape);
    int64_t expertIdsShapeSize = GetShapeSize(expertIdsShape);
    int64_t scalesShapeSize = GetShapeSize(scalesShape);
    int64_t expertScalesShapeSize = GetShapeSize(expertScalesShape);

    int64_t expandXShapeSize = GetShapeSize(expandXShape);
    int64_t dynamicScalesShapeSize = GetShapeSize(dynamicScalesShape);
    int64_t expandIdxShapeSize = GetShapeSize(expandIdxShape);
    int64_t expertTokenNumsShapeSize = GetShapeSize(expertTokenNumsShape);
    int64_t epRecvCountsShapeSize = GetShapeSize(epRecvCountsShape);
    int64_t tpRecvCountsShapeSize = GetShapeSize(tpRecvCountsShape);
    int64_t expandScalesShapeSize = GetShapeSize(expandScalesShape);
    int64_t sharedExpertXShapeSize = GetShapeSize(sharedExpertXShape);

    int64_t elasticInfoSize = GetShapeSize(elasticInfoShape);
    int64_t oriXSize = GetShapeSize(oriXShape);
    int64_t constExpertAlpha1Size = GetShapeSize(constExpertAlpha1Shape);
    int64_t constExpertAlpha2Size = GetShapeSize(constExpertAlpha2Shape);
    int64_t constExpertVSize = GetShapeSize(constExpertVShape);

    int64_t xOutShapeSize = GetShapeSize(xOutShape);

    std::vector<int16_t> xHostData(xShapeSize, 1);
    std::vector<int32_t> expertIdsHostData;
    for (int32_t token_id = 0; token_id < expertIdsShape[0]; token_id++) {
        for (int32_t k_id = 0; k_id < expertIdsShape[1]; k_id++) {
            expertIdsHostData.push_back(k_id);
        }
    }

    std::vector<float> scalesHostData(scalesShapeSize, 0.1);
    std::vector<float> expertScalesHostData(expertScalesShapeSize, 0.1);

    std::vector<int16_t> expandXHostData(expandXShapeSize, 0);
    std::vector<float> dynamicScalesHostData(dynamicScalesShapeSize, 0);
    std::vector<int32_t> expandIdxHostData(expandIdxShapeSize, 0);
    std::vector<int64_t> expertTokenNumsHostData(expertTokenNumsShapeSize, 0);
    std::vector<int32_t> epRecvCountsHostData(epRecvCountsShapeSize, 0);
    std::vector<int32_t> tpRecvCountsHostData(tpRecvCountsShapeSize, 0);
    std::vector<float> expandScalesHostData(expandScalesShapeSize, 0);
    std::vector<int16_t> sharedExpertXHostData(sharedExpertXShapeSize, 1);

    int32_t isElastic = 1;
    int32_t rankNumAfterElastic = 4;
    int32_t sharedExpertRankNumAfterElastic = sharedExpertRankNum;
    int32_t moeExpertNumAfterElastic = rankNumAfterElastic - sharedExpertRankNumAfterElastic;
    std::unordered_set<int16_t> availableRank{
        0, 1, /*2, 3, 4, 5,*/ 6, 7
    };
    std::vector<int32_t> elasticInfoHostData{
        isElastic, rankNumAfterElastic, sharedExpertRankNumAfterElastic, moeExpertNumAfterElastic,
        0, 1, -1, -1, -1, -1, 2, 3,
        0, 1, 6, 7, -1, -1, -1, -1
    };
    std::vector<int16_t> oriXHostData(oriXSize, 1);
    std::vector<int16_t> constExpertAlpha1HostData(constExpertAlpha1Size, 0);
    std::vector<int16_t> constExpertAlpha2HostData(constExpertAlpha2Size, 0);
    std::vector<int16_t> constExpertVHostData(constExpertVSize, 0);

    std::vector<int16_t> xOutHostData(xOutShapeSize, 0);


    ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_BF16, &x);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(expertIdsHostData, expertIdsShape, &expertIdsDeviceAddr, aclDataType::ACL_INT32, &expertIds);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(scalesHostData, scalesShape, &scalesDeviceAddr, aclDataType::ACL_FLOAT, &scales);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(expertScalesHostData, expertScalesShape, &expertScalesDeviceAddr, aclDataType::ACL_FLOAT, &expertScales);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    ret = CreateAclTensor(expandXHostData, expandXShape, &expandXDeviceAddr, (quantMode > 0) ? aclDataType::ACL_INT8 : aclDataType::ACL_BF16, &expandX);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(dynamicScalesHostData, dynamicScalesShape, &dynamicScalesDeviceAddr, aclDataType::ACL_FLOAT, &dynamicScales);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
        ret = CreateAclTensor(expandIdxHostData, expandIdxShape, &expandIdxDeviceAddr, aclDataType::ACL_INT32, &expandIdx);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(expertTokenNumsHostData, expertTokenNumsShape, &expertTokenNumsDeviceAddr, aclDataType::ACL_INT64, &expertTokenNums);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(epRecvCountsHostData, epRecvCountsShape, &epRecvCountsDeviceAddr, aclDataType::ACL_INT32, &epRecvCounts);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(tpRecvCountsHostData, tpRecvCountsShape, &tpRecvCountsDeviceAddr, aclDataType::ACL_INT32, &tpRecvCounts);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(expandScalesHostData, expandScalesShape, &expandScalesDeviceAddr, aclDataType::ACL_FLOAT, &expandScales);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(sharedExpertXHostData, sharedExpertXShape, &sharedExpertXDeviceAddr, aclDataType::ACL_BF16, &sharedExpertX);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    ret = CreateAclTensor(elasticInfoHostData, elasticInfoShape, &elasticInfoDeviceAddr, aclDataType::ACL_INT32, &elasticInfo);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(oriXHostData, oriXShape, &oriXDeviceAddr, aclDataType::ACL_BF16, &oriX);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(constExpertAlpha1HostData, constExpertAlpha1Shape, &constExpertAlpha1DeviceAddr, aclDataType::ACL_BF16, &constExpertAlpha1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(constExpertAlpha2HostData, constExpertAlpha2Shape, &constExpertAlpha2DeviceAddr, aclDataType::ACL_BF16, &constExpertAlpha2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(constExpertVHostData, constExpertVShape, &constExpertVDeviceAddr, aclDataType::ACL_BF16, &constExpertV);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(xOutHostData, xOutShape, &xOutDeviceAddr, aclDataType::ACL_BF16, &xOut);
    CHECK_RET(ret == ACL_SUCCESS, return ret);



    uint64_t dispatchWorkspaceSize = 0;
    aclOpExecutor *dispatchExecutor = nullptr;
    void *dispatchWorkspaceAddr = nullptr;

    uint64_t combineWorkspaceSize = 0;
    aclOpExecutor *combineExecutor = nullptr;
    void *combineWorkspaceAddr = nullptr;
    /**************************************** Call dispatch warm up. ********************************************/
    ret = aclnnMoeDistributeDispatchV3GetWorkspaceSize(x, expertIds, (quantMode > 0 ? scales : nullptr), nullptr,
            expertScales, nullptr, hcomEpName, EP_WORLD_SIZE, args.epRankId, moeExpertNum, hcomTpName, TP_WORLD_SIZE,
            args.tpRankId, expertShardType, sharedExpertNum,sharedExpertRankNum, quantMode, globalBS,
            expertTokenNumsType, nullptr, zeroExpertNum, copyExpertNum, constExpertNum, expandX, dynamicScales, expandIdx, expertTokenNums, epRecvCounts,
            tpRecvCounts, expandScales, &dispatchWorkspaceSize, &dispatchExecutor);

    CHECK_RET(ret == ACL_SUCCESS,
        LOG_PRINT("[ERROR] warm up aclnnMoeDistributeDispatchV3GetWorkspaceSize failed. ret = %d \n", ret); return ret);

    if (dispatchWorkspaceSize > 0) {
        ret = aclrtMalloc(&dispatchWorkspaceAddr, dispatchWorkspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] warm up aclrtMalloc workspace failed. ret = %d \n", ret); return ret);
    }
    // Call the second-phase API.
    ret = aclnnMoeDistributeDispatchV3(dispatchWorkspaceAddr, dispatchWorkspaceSize,
                                        dispatchExecutor, args.dispatchStream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] warm up aclnnMoeDistributeDispatchV3 failed. ret = %d \n", ret);  \
            return ret);
    ret = aclrtSynchronizeStreamWithTimeout(args.dispatchStream, 10000);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] warm up aclrtSynchronizeStreamWithTimeout failed. ret = %d \n", ret);  \
        return ret);

    /**************************************** Call dispatch. ********************************************/
    if (availableRank.find(args.rankId) != availableRank.end()) {
        // Call the first-phase API.
    ret = aclnnMoeDistributeDispatchV3GetWorkspaceSize(x, expertIds, (quantMode > 0 ? scales : nullptr), nullptr,
            expertScales, elasticInfo, hcomEpName, EP_WORLD_SIZE, args.epRankId, moeExpertNum, hcomTpName, TP_WORLD_SIZE,
            args.tpRankId, expertShardType, sharedExpertNum,sharedExpertRankNum, quantMode, globalBS,
            expertTokenNumsType, nullptr, zeroExpertNum, copyExpertNum, constExpertNum, expandX, dynamicScales, expandIdx, expertTokenNums, epRecvCounts,
            tpRecvCounts, expandScales, &dispatchWorkspaceSize, &dispatchExecutor);

    CHECK_RET(ret == ACL_SUCCESS,
        LOG_PRINT("[ERROR] aclnnMoeDistributeDispatchV3GetWorkspaceSize failed. ret = %d \n", ret); return ret);

    if (dispatchWorkspaceSize > 0) {
        ret = aclrtMalloc(&dispatchWorkspaceAddr, dispatchWorkspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc workspace failed. ret = %d \n", ret); return ret);
    }
    // Call the second-phase API.
    ret = aclnnMoeDistributeDispatchV3(dispatchWorkspaceAddr, dispatchWorkspaceSize,
                                        dispatchExecutor, args.dispatchStream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMoeDistributeDispatchV3 failed. ret = %d \n", ret);  \
            return ret);
    ret = aclrtSynchronizeStreamWithTimeout(args.dispatchStream, 10000);
                CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] dispatch aclrtSynchronizeStreamWithTimeout failed. ret = %d \n", ret);  \
            return ret);
    }
    /**************************************** Call combine. ********************************************/
    // Call the first-phase API.
    if (availableRank.find(args.rankId) != availableRank.end()) {
    ret = aclnnMoeDistributeCombineV3GetWorkspaceSize(expandX, expertIds,
                                                        expandIdx, epRecvCounts,
                                                        expertScales, tpRecvCounts,
                                                        nullptr, nullptr, nullptr,
                                                        nullptr, nullptr, nullptr,
                                                        elasticInfo, oriX, constExpertAlpha1, constExpertAlpha2, constExpertV,
                                                        hcomEpName, EP_WORLD_SIZE, args.epRankId, moeExpertNum,
                                                        hcomTpName, TP_WORLD_SIZE, args.tpRankId, expertShardType,
                                                        sharedExpertNum, sharedExpertRankNum, globalBS, outDtype,
                                                        commQuantMode, groupList_type, nullptr, zeroExpertNum, copyExpertNum, constExpertNum, xOut,
                                                        &combineWorkspaceSize, &combineExecutor);
    CHECK_RET(ret == ACL_SUCCESS,
        LOG_PRINT("[ERROR] aclnnMoeDistributeCombineV3GetWorkspaceSize failed. ret = %d \n", ret); return ret);
    // Apply for device memory based on the workspaceSize calculated by the first-phase API.
    if (combineWorkspaceSize > 0) {
        ret = aclrtMalloc(&combineWorkspaceAddr, combineWorkspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc workspace failed. ret = %d \n", ret); return ret);
    }

    // Call the second-phase API.
    ret = aclnnMoeDistributeCombineV3(combineWorkspaceAddr, combineWorkspaceSize, combineExecutor, args.combineStream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMoeDistributeCombineV3 failed. ret = %d \n", ret);
        return ret);
    // (Boilerplate) Synchronously wait for task execution to complete.
    ret = aclrtSynchronizeStreamWithTimeout(args.combineStream, 10000);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSynchronizeStreamWithTimeout failed. ret = %d \n", ret);
        return ret);
    LOG_PRINT("[INFO] device_%d aclnnMoeDistributeDispatchV3 and aclnnMoeDistributeCombineV3                      \
                execute successfully.\n", args.rankId);
    }
    // Release device resources.
    if (dispatchWorkspaceSize > 0) {
        aclrtFree(dispatchWorkspaceAddr);
    }
    if (combineWorkspaceSize > 0) {
        aclrtFree(combineWorkspaceAddr);
    }
    if (x != nullptr) {
        aclDestroyTensor(x);
    }
    if (expertIds != nullptr) {
        aclDestroyTensor(expertIds);
    }
    if (scales != nullptr) {
        aclDestroyTensor(scales);
    }
    if (expertScales != nullptr) {
        aclDestroyTensor(expertScales);
    }

    if (expandX != nullptr) {
        aclDestroyTensor(expandX);
    }
    if (dynamicScales != nullptr) {
        aclDestroyTensor(dynamicScales);
    }
    if (expandIdx != nullptr) {
        aclDestroyTensor(expandIdx);
    }
    if (expertTokenNums != nullptr) {
        aclDestroyTensor(expertTokenNums);
    }
    if (epRecvCounts != nullptr) {
        aclDestroyTensor(epRecvCounts);
    }
    if (tpRecvCounts != nullptr) {
        aclDestroyTensor(tpRecvCounts);
    }
    if (expandScales != nullptr) {
        aclDestroyTensor(expandScales);
    }
    if (residualX != nullptr) {
        aclDestroyTensor(residualX);
    }
    if (sharedExpertX != nullptr) {
        aclDestroyTensor(sharedExpertX);
    }
    if (elasticInfo != nullptr) {
        aclDestroyTensor(elasticInfo);
    }
    if (oriX != nullptr) {
        aclDestroyTensor(oriX);
    }
    if (constExpertAlpha1 != nullptr) {
        aclDestroyTensor(constExpertAlpha1);
    }
    if (constExpertAlpha2 != nullptr) {
        aclDestroyTensor(constExpertAlpha2);
    }
    if (constExpertV != nullptr) {
        aclDestroyTensor(constExpertV);
    }

    if (xOut != nullptr) {
        aclDestroyTensor(xOut);
    }
    if (xDeviceAddr != nullptr) {
        aclrtFree(xDeviceAddr);
    }
    if (expertIdsDeviceAddr != nullptr) {
        aclrtFree(expertIdsDeviceAddr);
    }
    if (scalesDeviceAddr != nullptr) {
        aclrtFree(scalesDeviceAddr);
    }
    if (expertScalesDeviceAddr != nullptr) {
        aclrtFree(expertScalesDeviceAddr);
    }
    if (expandXDeviceAddr != nullptr) {
        aclrtFree(expandXDeviceAddr);
    }
    if (dynamicScalesDeviceAddr != nullptr) {
        aclrtFree(dynamicScalesDeviceAddr);
    }
    if (expandIdxDeviceAddr != nullptr) {
        aclrtFree(expandIdxDeviceAddr);
    }
    if (expertTokenNumsDeviceAddr != nullptr) {
        aclrtFree(expertTokenNumsDeviceAddr);
    }
    if (epRecvCountsDeviceAddr != nullptr) {
        aclrtFree(epRecvCountsDeviceAddr);
    }
    if (expandScalesDeviceAddr != nullptr) {
        aclrtFree(expandScalesDeviceAddr);
    }
    if (tpRecvCountsDeviceAddr != nullptr) {
        aclrtFree(tpRecvCountsDeviceAddr);
    }
    if (sharedExpertXDeviceAddr != nullptr) {
        aclrtFree(sharedExpertXDeviceAddr);
    }

    if (elasticInfoDeviceAddr != nullptr) {
        aclrtFree(elasticInfoDeviceAddr);
    }
    if (oriXDeviceAddr != nullptr) {
        aclrtFree(oriXDeviceAddr);
    }
    if (constExpertAlpha1DeviceAddr != nullptr) {
        aclrtFree(constExpertAlpha1DeviceAddr);
    }
    if (constExpertAlpha2DeviceAddr != nullptr) {
        aclrtFree(constExpertAlpha2DeviceAddr);
    }
    if (constExpertVDeviceAddr != nullptr) {
        aclrtFree(constExpertVDeviceAddr);
    }

    if (xOutDeviceAddr != nullptr) {
        aclrtFree(xOutDeviceAddr);
    }

    HcclCommDestroy(args.hcclEpComm);
    HcclCommDestroy(args.hcclTpComm);
    aclrtDestroyStream(args.dispatchStream);
    aclrtDestroyStream(args.combineStream);
    aclrtDestroyContext(args.context);
    aclrtResetDevice(args.rankId);

    return 0;
}

int main(int argc, char *argv[])
{
    int ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclInit failed, ret = %d\n", ret); return ret);

    aclrtStream dispatchStream[DEV_NUM];
    aclrtStream combineStream[DEV_NUM];
    aclrtContext context[DEV_NUM];
    for (uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
        ret = aclrtSetDevice(rankId);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSetDevice failed, ret = %d\n", ret); return ret);
        ret = aclrtCreateContext(&context[rankId], rankId);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateContext failed, ret = %d\n", ret); return ret);
        ret = aclrtCreateStream(&dispatchStream[rankId]);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateStream failed, ret = %d\n", ret); return ret);
        ret = aclrtCreateStream(&combineStream[rankId]);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateStream failed, ret = %d\n", ret); return ret);
    }

    int32_t devicesEp[TP_WORLD_SIZE][EP_WORLD_SIZE];
    for (int32_t tpId = 0; tpId < TP_WORLD_SIZE; tpId++) {
        for (int32_t epId = 0; epId < EP_WORLD_SIZE; epId++) {
            devicesEp[tpId][epId] = epId * TP_WORLD_SIZE + tpId;
        }
    }

    HcclComm commsEp[TP_WORLD_SIZE][EP_WORLD_SIZE];
    for (int32_t tpId = 0; tpId < TP_WORLD_SIZE; tpId++) {
        ret = HcclCommInitAll(EP_WORLD_SIZE, devicesEp[tpId], commsEp[tpId]);
        CHECK_RET(ret == ACL_SUCCESS,
                    LOG_PRINT("[ERROR] HcclCommInitAll ep %d failed, ret %d\n", tpId, ret); return ret);
    }

    int32_t devicesTp[EP_WORLD_SIZE][TP_WORLD_SIZE];
    for (int32_t epId = 0; epId < EP_WORLD_SIZE; epId++) {
        for (int32_t tpId = 0; tpId < TP_WORLD_SIZE; tpId++) {
            devicesTp[epId][tpId] = epId * TP_WORLD_SIZE + tpId;
        }
    }

    HcclComm commsTp[EP_WORLD_SIZE][TP_WORLD_SIZE];
    for (int32_t epId = 0; epId < EP_WORLD_SIZE; epId++) {
        ret = HcclCommInitAll(TP_WORLD_SIZE, devicesTp[epId], commsTp[epId]);
        CHECK_RET(ret == ACL_SUCCESS,
                    LOG_PRINT("[ERROR] HcclCommInitAll tp %d failed, ret %d\n", epId, ret); return ret);
    }

    Args args[DEV_NUM];
    std::vector<std::unique_ptr<std::thread>> threads(DEV_NUM);
    for (uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
        uint32_t epRankId = rankId / TP_WORLD_SIZE;
        uint32_t tpRankId = rankId % TP_WORLD_SIZE;

        args[rankId].rankId = rankId;
        args[rankId].epRankId = epRankId;
        args[rankId].tpRankId = tpRankId;
        args[rankId].hcclEpComm = commsEp[tpRankId][epRankId];
        args[rankId].hcclTpComm = commsTp[epRankId][tpRankId];
        args[rankId].dispatchStream = dispatchStream[rankId];
        args[rankId].combineStream = combineStream[rankId];
        args[rankId].context = context[rankId];
        threads[rankId].reset(new(std::nothrow) std::thread(&LaunchOneProcessDispatchAndCombine, std::ref(args[rankId])));
    }

    for(uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
        threads[rankId]->join();
    }

    aclFinalize();
    LOG_PRINT("[INFO] aclFinalize success\n");

    return 0;
}