aclnnMoeDistributeCombineV2

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 perform ReduceScatterV communication, then perform AllToAllV communication, and finally integrate the received data (multiply by weights and sum). When no TP domain communication exists, perform AllToAllV communication, and finally integrate the received data (multiply by weights and sum).
Formula:
- When no TP domain communication exists:
$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 aclnnMoeDistributeDispatchV2, which means data is returned along the same path used by the MoeDistributeDispatchV2 operator to collect data.

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

More detailed token information is input to assist aclnnMoeDistributeCombineV2 in efficiently performing cross-card synchronization. Therefore, the expandIdx input with the shape (BS * K,) in the original API is replaced by the assistInfoForCombine parameter with the shape (A * 128,).
The sharedExpertXOptional input is added, allowing the user to input tokens computed by the shared expert when sharedExpertNum is 0.
The commAlg input is added, replacing the HCCL_INTRA_PCIE_ENABLE and HCCL_INTRA_ROCE_ENABLE environment variables.

For details, see the parameter descriptions below.

Function Prototype

Each operator is divided into a two-phase API. You must first call the aclnnMoeDistributeCombineV2GetWorkspaceSize API to obtain the workspace size required for computation and the executor that contains the operator computation process, and then call the aclnnMoeDistributeCombineV2 API to execute the computation.

aclnnStatus aclnnMoeDistributeCombineV2GetWorkspaceSize(
    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 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,
    aclTensor*       xOut,
    uint64_t*        workspaceSize,
    aclOpExecutor**  executor)

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

aclnnMoeDistributeCombineV2GetWorkspaceSize

Parameters

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
expandX	Input	Token features expanded based on expertIds.	Must be a 2D tensor.	FLOAT16, BFLOAT16	ND	(max(tpWorldSize, 1) * A , H)	√
expertIds	Input	Indices of the topK experts for each token.	Must be a 2D tensor.	INT32	ND	(BS, K)	√
assistInfoForCombine	Input	Corresponds to the assistInfoForCombineOut output in aclnnMoeDistributeDispatchV2.	Must be a 1D tensor.	INT32	N/A	(A * 128, )	√
epSendCounts	Input	Corresponds to the epRecvCounts output in aclnnMoeDistributeDispatchV2.	Requires a 1D tensor.	INT32	ND	-	√
expertScales	Input	Weights of the topK experts for each token.	It must be a 2D tensor.	FLOAT32	ND	(BS, K)	√
tpSendCountsOptional	Input	Corresponds to the tpRecvCounts output in aclnnMoeDistributeDispatchV2.	Pass this parameter when TP domain communication is required; pass a null pointer when no TP domain communication is required.	INT32	ND	-	√
xActiveMaskOptional	Input	Indicates whether a token participates in communication.	A valid data or a null pointer can be passed. By default, all tokens participate in communication. When the BS varies across cards, all tokens must be valid.	BOOL	ND	-	√
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 in aclnnMoeDistributeDispatchV2.	-	FLOAT32	ND	-	√
sharedExpertXOptional	Input	Represents the tokens after shared expert computation.	The data type must be consistent with expandX.	FLOAT16, BFLOAT16	ND	-	√
groupEp	Input	Name of the EP communication domain (expert parallel communication domain).	The string length ranges from [1, 128) and 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.	Satisfies moeExpertNum % (epWorldSize - sharedExpertRankNum) = 0.	INT64	-	-	-
groupTp	Input	Name of the TP communication domain (data parallel communication domain).	Cannot be the same as groupEp.	STRING	-	-	-
tpWorldSize	Input	Size of the TP communication domain.	-	INT64	-	-	-
tpRankId	Input	ID of the local card in the TP domain.	The tpRankId of each card in the same TP communication domain must be unique.	INT64	-	-	-
expertShardType	Input	Indicates the shared expert shard distribution type.	-	INT64	-	-	-
sharedExpertNum	Input	Indicates the number of shared experts (a shared expert can be replicated and deployed on multiple cards).	-	INT64	-	-	-
sharedExpertRankNum	Input	Indicates the 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 (maxBS is the maximum BS value of a single card).	INT64	-	-	-
outDtype	Input	Specifies the data type of the output x. Reserved parameter.	Not supported in the current version. Pass 0.	INT64	-	-	-
commQuantMode	Input	Communication quantization type.	Value range: 0 or 2 (0 indicates no quantization, 2 indicates int8 quantization).	INT64	-	-	-
groupListType	Input.	Group List format. Reserved parameter.	Not supported in the current version. Pass 0.	INT64	-	-	-
commAlg	Input	Indicates the communication affinity memory layout algorithm. Data type: string.	-	STRING	-	-	-
xOut	Output	Represents the processed tokens.	Requires a 2D tensor.	FLOAT16, BFLOAT16	ND	(BS, H)	-
workspaceSize	Output	Size of the workspace to be requested on the Device.	-	-	-	-	-
executor	Output	Returns the operator executor, which contains the computation flow of the operator.	-	-	-	-	-

Atlas A2 training products/Atlas A2 inference products:
- commAlg supports nullptr, "", "fullmesh", and "hierarchy". "hierarchy" is recommended and requires 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" reduces inter-server data volume through two transmissions: intra-server and inter-server.
- 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. The last 2 * globalBS * K * serverNum values are used to store the number of tokens and communication domain offsets for which combine can perform reduce in advance before inter-server/intra-server communication. When globalBS=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 shape (BS, ). true must precede false (for example, {true, false, true} is invalid). "hierarchy" 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 a null string.
- tpWorldSize: Not supported in the current version. Set this parameter to 0.
- tpRankId: Not supported in the current version. Set this parameter to 0.
- expertShardType: Not supported in the current version. Set this parameter to 0.
- sharedExpertNum: Not supported in the current version. Set this parameter to 0.
- sharedExpertRankNum: Not supported in the current version. Set this parameter to 0.
- commQuantMode: The value range is 0 or 2 (0 indicates no quantization, and 2 indicates int8 quantization). The value 2 is supported only when commAlg is "hierarchy" or HCCL_INTRA_PCIE_ENABLE=1, HCCL_INTRA_ROCE_ENABLE=0, and the driver version is ≥ 25.0.RC1.1.
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 exists, tpSendCountsOptional is a 1D tensor with the shape (tpWorldSize, ).
- 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 is a reserved parameter and is 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). It can be passed or not. When passed, sharedExpertRankNum must be 0.
- epWorldSize supports values in the range [2, 768].
- moeExpertNum supports values in the range (0, 1024].
- The string length range of groupTp is [0, 128). It cannot be the same as groupEp. Passing an empty string is supported only when there is no TP domain communication.
- The value range of tpWorldSize is [0, 2]. 0 and 1 indicate no TP domain communication. Only 2 is supported when TP domain communication exists.
- The value range of tpRankId is [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 supports passing 0, which indicates that shared expert cards are placed before MoE expert cards.
- The current value range of sharedExpertNum is [0, 4].
- The value range of sharedExpertRankNum is [0, epWorldSize). When it is 0, sharedExpertNum must be 0 or 1. When it is not 0, sharedExpertRankNum % sharedExpertNum = 0 must hold.
- The value range of commQuantMode is 0 or 2 (0 indicates no quantization, and 2 indicates int8 quantization). The value 2 can be enabled only when tpWorldSize < 2.
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, ).
- tpSendCountsOptional is not supported in the current version. Pass a null pointer.
- xActiveMaskOptional must be a 1D or 2D tensor (shape: (BS, ) for 1D, (BS, K) for 2D). For 1D, true values must precede false values (e.g., {true, false, true} is invalid). For 2D, if all K values corresponding to a token are false, the token does not participate in communication.
- expandScalesOptional is a reserved parameter and is 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). It can be passed or omitted. When passed, sharedExpertRankNum must be 0.
- epWorldSize supports values in the range [2, 768].
- moeExpertNum has a value range of (0, 1024].
- groupTp is 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.
- The value range of sharedExpertNum is [0, 4].
- The value range of sharedExpertRankNum is [0, epWorldSize). When the value is 0, sharedExpertNum must be 0 or 1. When the value is not 0, sharedExpertRankNum % sharedExpertNum = 0 must be satisfied.
- The value range of commQuantMode is 0 or 2 (0 indicates no quantization, and 2 indicates int8 quantization).

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 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	The shape of the input or output is not within the supported range.
ACLNN_ERR_INNER_TILING_ERROR	561002	The parameter value is out of the supported range.

aclnnMoeDistributeCombineV2

Parameters

Parameter	Input/Output	Description
workspace	Input.	Workspace memory address applied for on the device side.
workspaceSize	Input.	Workspace size applied for on the device side, obtained from the first-phase API aclnnMoeDistributeCombineV2GetWorkspaceSize.
executor	Input	Operator executor that contains the operator computation flow.
stream	Input	Specifies the stream for task execution.

Return Value

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

Constraints

Deterministic computation:

- aclnnMoeDistributeCombineV2 defaults to a deterministic implementation.

Driver constraints:

- The driver versions of all nodes in the operator communication domain must be the same.

The aclnnMoeDistributeDispatchV2 API and the aclnnMoeDistributeCombineV2 API must be used together. For details, see Example.
Across different product models, communication algorithms, or versions, the element values in the Tensor outputs assistInfoForCombineOut, epRecvCountsOut, tpRecvCountsOut, and expandScalesOut of aclnnMoeDistributeDispatchV2 may vary. When using them, directly pass the above Tensors to the corresponding parameters of aclnnMoeDistributeCombineV2. Other model business logic must not depend on these values.
The groupEp, epWorldSize, moeExpertNum, groupTp, tpWorldSize, expertShardType, sharedExpertNum, sharedExpertRankNum, globalBs, and commAlg parameters, as well as the HCCL_BUFFSIZE value used during the API call, must be consistent across all cards, consistent across different layers in the network, and consistent with the corresponding parameters of aclnnMoeDistributeDispatchV2.
Atlas A3 training products/Atlas A3 inference products: In this scenario, a single card contains two dies (referred to as "die" or "bare die"). Therefore, "this card" in the parameter description refers to a single die.
Description of the shape format in the parameter description:
- A: Indicates the maximum number of tokens to be distributed by the current card. The value range is as follows:
  - For shared experts, the condition (A = BS * epWorldSize * sharedExpertNum / sharedExpertRankNum) must be satisfied.
  - For MoE experts, when globalBS is 0, the condition (A >= BS * epWorldSize * min(localExpertNum, K)) must be satisfied; when globalBS is non-zero, the condition (A >= globalBS * min(localExpertNum, K)) must be satisfied.
- H: Indicates the hidden size:
  - Atlas A2 training products/Atlas A2 inference products: Depends 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: The value range is [1024, 8192].
  - Ascend 950PR/Ascend 950DT: The value range is [1024, 8192].
- BS: Indicates the batch sequence size (the number of tokens finally output by 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: The value range is (0 < BS ≤ 512).
  - Ascend 950PR/Ascend 950DT: The value range is (0 < BS ≤ 512).
- K: Indicates the number of top K experts selected. The value range is (0 < K ≤ 16) and satisfies (0 < K ≤ moeExpertNum).
- serverNum: Indicates the number of server nodes. Only 2, 4, and 8 are supported.
- localExpertNum: Indicates the number of experts on the current card:
  - For a shared expert card, (localExpertNum = 1).
  - For MoE expert cards, (localExpertNum = moeExpertNum / (epWorldSize - sharedExpertRankNum)); when (localExpertNum > 1), TP domain communication is not supported.
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. It defaults to 200 MB when not configured.
- Atlas A2 training products/Atlas A2 inference products:
  - commAlg is "" or nullptr: The "fullmesh" or "hierarchy" formula is selected based on the HCCL environment variable.
  - commAlg is "fullmesh": The workspace size requirement is (≥ 2 (BS epWorldSize min(localExpertNum, K) H * sizeof(uint16) + 2 MB)).
  - commAlg is "hierarchy": The workspace size requirement is (≥ (moeExpertNum + epWorldSize / 4) Align512(maxBS (H 2 + 16 Align8(K))) 1 B + 8 MB, where Align8(x) = ((x + 8 - 1) / 8) 8, and Align512(x) = ((x + 512 - 1) / 512) * 512).
- Atlas A3 training products/Atlas A3 inference products:
  - Within the ep communication domain: The workspace size requirement is (≥ 2) and satisfies (≥ 2 (localExpertNum maxBS epWorldSize Align512(Align32(2 H) + 44) + (K + sharedExpertNum) maxBS Align512(2 H))) (localExpertNum must use the number of experts on the local MoE expert card; Align512(x) = ((x + 512 - 1) / 512) * 512; Align32(x) = ((x + 32 - 1) / 32) * 32).
  - Within the tp communication domain: The workspace size requirement is ≥ A (H 2 + 128) * 2.
- Ascend 950PR/Ascend 950DT: Requires (≥ aivNum * 512 + 2 * epWorldSize * BS * H * 2 * localExpertNum), where aivNum indicates the number of cores, and localExpertNum must use the number of experts on the local MoE expert card.
HCCL_INTRA_PCIE_ENABLE and HCCL_INTRA_ROCE_ENABLE:

Atlas A2 training products/Atlas A2 inference products: These environment variables are no longer recommended. It is recommended to configure commAlg to "hierarchy".
The "/" in the formulas in this document indicates integer division.
Communication domain usage constraints:
- In a model, aclnnMoeDistributeCombineV2 and aclnnMoeDistributeDispatchV2 support only the same EP communication domain, and no other operators are allowed in this communication domain.
- In a model, aclnnMoeDistributeCombineV2 and aclnnMoeDistributeDispatchV2 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 reside within a single super node. Cross-super-node configuration is not supported.
Networking constraints:
- Atlas A2 training products/Atlas A2 inference products: In multi-machine scenarios, only switch networking is supported. Direct dual-machine networking is not supported.

Example

File preparation:
1. Create the rank_table_m2.json file as instructed below and modify it.
2. Copy the project to both servers, and configure the rank_table_m2.json file content based on the device IP of each machine. Ensure that the rank_table_m2.json files on both machines remain consistent.
3. Install the CANN package, and compile and run the sample by following the instructions in Operator Invocation.
About rankTable:
1. Developers can configure the NPU resource information for collective communication through the ranktable file. For detailed configuration, refer to "Communication Function Development > Cluster Information Configuration > Configuring Resource Information in the ranktable File" in the HCCL User Guide.
2. Use cat /etc/hccn.conf or for i in seq 0 7; do echo "===================> dev$i, NPU$((i+1))"; hccn_tool -i $i -ip -g; done to query the device IP of the machine. Then fill in the JSON file by referring to the HCCL documentation.
NOTE: In a two-machine 16-card scenario, the device_id of both machines is 0–7, the rank_id of one machine is 0–7, and the rank_id of the other machine is 8–15. In a single-machine 16-card scenario, both device_id and rank_id are 0–15.

Environment variable configuration:

# Set three environment variables before running.
## FIRST_RANK_ID description: In a two-machine 16-card scenario, set one machine to 0 and the other to 8.
## For example, export FIRST_RANK_ID=0.
export RANK_TABLE_FILE=/home/path/to/rank_table_m2.json
export FIRST_RANK_ID=<starting rank_id of the device>
## ENV_DEV_NUM description: Set this variable based on the number of cards on the current machine. In a two-machine 16-card scenario, set both machines to 16.
export ENV_DEV_NUM=16

Number of Machines Setting:

In a two-machine 16-card scenario, set the MACHINE_NUM parameter to 2, that is,
```
const uint32_t MACHINE_NUM = 2;
```
No modification is required for a single-machine 16-card scenario.
Atlas A2 training products/Atlas A2 inference products :

No ranktable file or the RANK_TABLE_FILE and FIRST_RANK_ID environment variables need to be configured.

This example supports running the A2 operator in a single-machine environment with [2, 8] cards. Before running, set IS_TEST_A2 to true in the example code to ensure the A2 branch is executed. Meanwhile, you can set EP_WORLD_SIZE_A2 to the number of cards in the example code as needed, and change moeExpertNum in the launchOneThreadDispatchV2AndCombineV2_A2 function so that moeExpertNum is divisible by EP_WORLD_SIZE_A2.

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
```
The example operator execution command is as follows:
```
bash build.sh --run_example --ops=moe_distribute_combine_v2 eager cust
```
Atlas A3 training products/Atlas A3 inference products:

No ranktable file or the RANK_TABLE_FILE and FIRST_RANK_ID environment variables need to be configured.

The following provides example code for reference only. For details about compilation and running, see Compile and Run Samples.

Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products, Ascend 950PR/Ascend 950DT:

#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_v2.h"
#include "aclnnop/aclnn_moe_distribute_combine_v2.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 dispatchV2Stream;
    aclrtStream combineV2Stream;
    aclrtContext context;
};

const uint32_t MACHINE_NUM = 1;
const char* rank_table_file = std::getenv("RANK_TABLE_FILE");
const char* first_rank_id = std::getenv("FIRST_RANK_ID");
const char* env_dev_num = std::getenv("ENV_DEV_NUM");

const uint32_t EP_WORLD_SIZE = (!first_rank_id) ? 2 : 16;
const uint32_t TP_WORLD_SIZE = (!first_rank_id) ? 1 : 0;
const uint32_t DEV_NUM = (!first_rank_id) ? EP_WORLD_SIZE * TP_WORLD_SIZE : EP_WORLD_SIZE;

const bool IS_TEST_A2 = false;
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 launchOneThreadDispatchV2AndCombineV2_A3A5(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};
    if (!rank_table_file && !first_rank_id) {
        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, dispatchV2Stream = %p, combineV2Stream = %p, context = %p\n",
        args.rankId, hcomEpName, hcomTpName, args.dispatchV2Stream, args.combineV2Stream, args.context
    );

    // Set the scenario.
    int64_t BS = 8;
    int64_t H = 7168;
    int64_t K = 1;
    int64_t expertShardType = 0;
    int64_t sharedExpertNum = 0;
    int64_t sharedExpertRankNum = 0;
    if (!rank_table_file && !first_rank_id) {
        sharedExpertNum = 1;
        sharedExpertRankNum = 1;
    }
    int64_t moeExpertNum = EP_WORLD_SIZE - sharedExpertRankNum;
    int64_t quantMode = 0;
    int64_t globalBS = BS * EP_WORLD_SIZE;
    int64_t expertTokenNumsType = 0;
    int64_t outDtype = 0;
    int64_t commQuantMode = 0;
    int64_t groupListType = 0;
    int64_t localExpertNum;
    int64_t A;
    if (args.epRankId < sharedExpertRankNum) {
        // Shared expert card.
        localExpertNum = 1;
        A = globalBS / sharedExpertRankNum;
    } else {
        // MoE expert card
        localExpertNum = moeExpertNum / (EP_WORLD_SIZE - sharedExpertRankNum);
        A = globalBS * (localExpertNum < K ? localExpertNum : K);
    }
    std::string commAlg = "";

    /*Construct the input and output variables on the device side based on the current scenario.*/
    // Declare the input and output variables on the device side.
    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;

    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;

    // 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{(sharedExpertRankNum > 0) ? 1 + moeExpertNum : moeExpertNum, H};
    std::vector<int64_t> expertScalesShape{BS, K};
    std::vector<int64_t> expandXShape{(TP_WORLD_SIZE > 0 ? TP_WORLD_SIZE : 1) * A, H};
    std::vector<int64_t> dynamicScalesShape{(TP_WORLD_SIZE > 0 ? TP_WORLD_SIZE : 1) * A};
    std::vector<int64_t> expandIdxShape{A * 128};
    std::vector<int64_t> expertTokenNumsShape{localExpertNum};
    std::vector<int64_t> epRecvCountsShape{(TP_WORLD_SIZE > 0 ? TP_WORLD_SIZE : 1) * localExpertNum * EP_WORLD_SIZE};
    std::vector<int64_t> tpRecvCountsShape{TP_WORLD_SIZE > 0 ? TP_WORLD_SIZE : 1};
    std::vector<int64_t> expandScalesShape{A};

    long long xShapeSize = GetShapeSize(xShape);
    long long expertIdsShapeSize = GetShapeSize(expertIdsShape);
    long long scalesShapeSize = GetShapeSize(scalesShape);
    long long expertScalesShapeSize = GetShapeSize(expertScalesShape);
    long long expandXShapeSize = GetShapeSize(expandXShape);
    long long dynamicScalesShapeSize = GetShapeSize(dynamicScalesShape);
    long long expandIdxShapeSize = GetShapeSize(expandIdxShape);
    long long expertTokenNumsShapeSize = GetShapeSize(expertTokenNumsShape);
    long long epRecvCountsShapeSize = GetShapeSize(epRecvCountsShape);
    long long tpRecvCountsShapeSize = GetShapeSize(tpRecvCountsShape);
    long long expandScalesShapeSize = GetShapeSize(expandScalesShape);

    // Construct the host-side variables.
    std::vector<op::fp16_t> xHostData(xShapeSize, 1);
    std::vector<int32_t> expertIdsHostData;
    for (int32_t token_id = 0; token_id < expertIdsShape[0]; token_id++) {
        // Send each token to MoE experts {0, 1, ... k - 1}.
        for (int32_t k_id = 0; k_id < expertIdsShape[1]; k_id++) {
            expertIdsHostData.push_back(k_id);
        }
    }
    std::vector<float> scalesHostData(scalesShapeSize, 0);
    std::vector<float> expertScalesHostData(expertScalesShapeSize, 0);
    std::vector<op::fp16_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);

    // Construct device-side variables.
    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);

    /*Declare the variables required for operator execution.*/
    uint64_t dispatchV2WorkspaceSize = 0;
    aclOpExecutor *dispatchV2Executor = nullptr;
    void *dispatchV2WorkspaceAddr = nullptr;

    uint64_t combineV2WorkspaceSize = 0;
    aclOpExecutor *combineV2Executor = nullptr;
    void *combineV2WorkspaceAddr = nullptr;

    /*Execute the dispatchV2 and combineV2 operators in sequence.*/
    // Call the first-phase API of the dispatchV2 operator.
    ret = aclnnMoeDistributeDispatchV2GetWorkspaceSize(
        x, expertIds,
        (quantMode > 0 ? scales : nullptr), nullptr,
        expertScales,
        hcomEpName, EP_WORLD_SIZE, args.epRankId,
        moeExpertNum, hcomTpName, TP_WORLD_SIZE,
        args.tpRankId, expertShardType, sharedExpertNum,
        sharedExpertRankNum, quantMode, globalBS,
        expertTokenNumsType, commAlg.c_str(),
        expandX, dynamicScales,
        expandIdx, expertTokenNums,
        epRecvCounts, tpRecvCounts,
        expandScales, &dispatchV2WorkspaceSize,
        &dispatchV2Executor
    );
    CHECK_RET(
        ret == ACL_SUCCESS,
        LOG_PRINT("[ERROR] aclnnMoeDistributedispatchV2GetWorkspaceSize failed. ret = %d\n", ret); return ret
    );
    // Apply for device memory based on the workspaceSize calculated by the first-phase API of the dispatchV2 operator.
    if (dispatchV2WorkspaceSize > 0) {
        ret = aclrtMalloc(&dispatchV2WorkspaceAddr, dispatchV2WorkspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc failed. ret = %d\n", ret); return ret);
    }
    // Call the second-phase API of the dispatchV2 operator.
    ret = aclnnMoeDistributeDispatchV2(dispatchV2WorkspaceAddr, dispatchV2WorkspaceSize, dispatchV2Executor, args.dispatchV2Stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMoeDistributedispatchV2 failed. ret = %d\n", ret); return ret);
    // (Boilerplate) Synchronously wait for task execution to complete.
    ret = aclrtSynchronizeStreamWithTimeout(args.dispatchV2Stream, 10000);
    CHECK_RET(
        ret == ACL_SUCCESS,
        LOG_PRINT("[ERROR] aclrtSynchronizeStreamWithTimeout failed. ret = %d\n", ret); return ret
    );

    // Call the first-phase API of the combineV2 operator.
    ret = aclnnMoeDistributeCombineV2GetWorkspaceSize(
        expandX, expertIds, expandIdx, epRecvCounts, expertScales, tpRecvCounts,
        nullptr, nullptr, nullptr, nullptr, nullptr, nullptr,
        hcomEpName, EP_WORLD_SIZE, args.epRankId, moeExpertNum, hcomTpName, TP_WORLD_SIZE, args.tpRankId,
        expertShardType, sharedExpertNum, sharedExpertRankNum, globalBS, outDtype, commQuantMode, groupListType,
        commAlg.c_str(), x, &combineV2WorkspaceSize, &combineV2Executor);
    CHECK_RET(
        ret == ACL_SUCCESS,
        LOG_PRINT("[ERROR] aclnnMoeDistributecombineV2GetWorkspaceSize failed. ret = %d\n", ret); return ret
    );
    // Apply for device memory based on the workspaceSize calculated by the first-phase API of the combineV2 operator.
    if (combineV2WorkspaceSize > 0) {
        ret = aclrtMalloc(&combineV2WorkspaceAddr, combineV2WorkspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc failed. ret = %d\n", ret); return ret);
    }
    // Call the second-phase API of the combineV2 operator.
    ret = aclnnMoeDistributeCombineV2(combineV2WorkspaceAddr, combineV2WorkspaceSize, combineV2Executor, args.combineV2Stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMoeDistributecombineV2 failed. ret = %d\n", ret); return ret);
    // (Boilerplate) Synchronously wait for task execution to complete.
    ret = aclrtSynchronizeStreamWithTimeout(args.combineV2Stream, 10000);
    CHECK_RET(
        ret == ACL_SUCCESS,
        LOG_PRINT("[ERROR] aclrtSynchronizeStreamWithTimeout failed. ret = %d\n", ret); return ret
    );

    LOG_PRINT("[INFO] device_%d aclnnMoeDistributedispatchV2 and aclnnMoeDistributecombineV2 execute successfully.\n", args.rankId);

    // Release device resources.
    if (dispatchV2WorkspaceSize > 0) {
        aclrtFree(dispatchV2WorkspaceAddr);
    }
    if (combineV2WorkspaceSize > 0) {
        aclrtFree(combineV2WorkspaceAddr);
    }
    DestroyTensor(x);
    DestroyTensor(expertIds);
    DestroyTensor(scales);
    DestroyTensor(expertScales);
    DestroyTensor(expandX);
    DestroyTensor(dynamicScales);
    DestroyTensor(expandIdx);
    DestroyTensor(expertTokenNums);
    DestroyTensor(epRecvCounts);
    DestroyTensor(tpRecvCounts);
    DestroyTensor(expandScales);

    FreeDeviceAddr(xDeviceAddr);
    FreeDeviceAddr(expertIdsDeviceAddr);
    FreeDeviceAddr(scalesDeviceAddr);
    FreeDeviceAddr(expertScalesDeviceAddr);
    FreeDeviceAddr(expandXDeviceAddr);
    FreeDeviceAddr(dynamicScalesDeviceAddr);
    FreeDeviceAddr(expandIdxDeviceAddr);
    FreeDeviceAddr(expertTokenNumsDeviceAddr);
    FreeDeviceAddr(epRecvCountsDeviceAddr);
    FreeDeviceAddr(expandScalesDeviceAddr);
    FreeDeviceAddr(tpRecvCountsDeviceAddr);

    HcclCommDestroy(args.hcclEpComm);
    HcclCommDestroy(args.hcclTpComm);
    aclrtDestroyStream(args.dispatchV2Stream);
    aclrtDestroyStream(args.combineV2Stream);
    aclrtDestroyContext(args.context);
    aclrtResetDevice(args.rankId);

    return 0;
}

int launchOneThreadDispatchV2AndCombineV2_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, dispatchV2Stream = %p, combineV2Stream = %p, \
                context = %p\n", args.rankId, hcomEpName, args.dispatchV2Stream, args.combineV2Stream,                 \
                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
    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;

    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 *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> 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 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> 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(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 = aclnnMoeDistributeDispatchV2GetWorkspaceSize(x, expertIds, (quantMode > 0 ? scales : nullptr), xActiveMask,
            expertScales, hcomEpName, EP_WORLD_SIZE_A2, args.epRankId, moeExpertNum, "", TP_WORLD_SIZE_A2,
            args.tpRankId, expertShardType, sharedExpertNum,sharedExpertRankNum, quantMode, globalBS,
            expertTokenNumsType, commAlg.c_str(), expandX, dynamicScales, assistInfoForCombine, expertTokenNums, epRecvCounts,
            tpRecvCounts, expandScales, &dispatchWorkspaceSize, &dispatchExecutor);

    CHECK_RET(ret == ACL_SUCCESS,
        LOG_PRINT("[ERROR] aclnnMoeDistributeDispatchV2GetWorkspaceSize 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 = aclnnMoeDistributeDispatchV2(dispatchWorkspaceAddr, dispatchWorkspaceSize,
                                        dispatchExecutor, args.dispatchV2Stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMoeDistributeDispatchV2 failed. ret = %d \n", ret);  \
            return ret);
    ret = aclrtSynchronizeStreamWithTimeout(args.dispatchV2Stream, 10000);
                CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] dispatch aclrtSynchronizeStreamWithTimeout failed. ret = %d \n", ret);  \
            return ret);
    LOG_PRINT("[INFO] device_%d aclnnMoeDistributeDispatchV2 execute successfully.\n", args.rankId);
    /**************************************** Call combine ********************************************/
    // Call the first-phase API.
    ret = aclnnMoeDistributeCombineV2GetWorkspaceSize(expandX, expertIds,
                                                        assistInfoForCombine, epRecvCounts,
                                                        expertScales, tpRecvCounts,
                                                        xActiveMask, activationScale, weightScale,
                                                        groupList, expandScales, sharedExpertX,
                                                        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(), xOut,
                                                        &combineWorkspaceSize, &combineExecutor);
    CHECK_RET(ret == ACL_SUCCESS,
        LOG_PRINT("[ERROR] aclnnMoeDistributeCombineV2GetWorkspaceSize 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 = aclnnMoeDistributeCombineV2(combineWorkspaceAddr, combineWorkspaceSize, combineExecutor, args.combineV2Stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMoeDistributeCombineV2 failed. ret = %d \n", ret);
        return ret);
    // (Boilerplate) Synchronously wait for the task execution to finish.
    ret = aclrtSynchronizeStreamWithTimeout(args.combineV2Stream, 10000);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSynchronizeStreamWithTimeout failed. ret = %d \n", ret);
        return ret);
    LOG_PRINT("[INFO] device_%d aclnnMoeDistributeDispatchV2 and aclnnMoeDistributeCombineV2                      \
                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(xOut);

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

    HcclCommDestroy(args.hcclEpComm);
    aclrtDestroyStream(args.dispatchV2Stream);
    aclrtDestroyStream(args.combineV2Stream);
    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 run_example_on_A2()
{
    int ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclInit failed, ret = %d\n", ret); return ret);

    aclrtStream dispatchV2Stream[DEV_NUM_A2];
    aclrtStream combineV2Stream[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(&dispatchV2Stream[rankId]);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateStream failed, ret = %d\n", ret); return ret);
        ret = aclrtCreateStream(&combineV2Stream[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].dispatchV2Stream = dispatchV2Stream[rankId];
        args[rankId].combineV2Stream = combineV2Stream[rankId];
        args[rankId].context = context[rankId];
        threads[rankId].reset(new(std::nothrow) std::thread(&launchOneThreadDispatchV2AndCombineV2_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;
}

int run_example_on_A3A5()
{
    int ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclInit failed. ret = %d\n", ret); return ret);

    aclrtStream dispatchV2Stream[DEV_NUM];
    aclrtStream combineV2Stream[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(&dispatchV2Stream[rankId]);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateStream failed. ret = %d\n", ret); return ret);
        ret = aclrtCreateStream(&combineV2Stream[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;
        }
    }
    // Initialize the EP communication domain, ep = 8 {0,2,4,6,8,10,12,14} {1,3,5,7,9,11,13,15}.
    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 world %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;
        }
    }
    // Initialize the TP communication domain, tp = 2 {0,1} {2,3} {4,5} {6,7} {8,9} {10,11} {12,13} {14,15}.
    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 world %d failed. ret = %d\n", epId, ret); return ret
        );
    }

    Args args[DEV_NUM];
    // Each thread calls the corresponding device to execute the operator.
    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].dispatchV2Stream = dispatchV2Stream[rankId];
        args[rankId].combineV2Stream = combineV2Stream[rankId];
        args[rankId].context = context[rankId];
        threads[rankId].reset(new(std::nothrow) std::thread(&launchOneThreadDispatchV2AndCombineV2_A3A5, 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;
}


int main(int argc, char *argv[])
{
    const char* env_var_name = "RANK_TABLE_FILE and FIRST_RANK_ID";
    if (IS_TEST_A2) {
        LOG_PRINT("[INFO] %s are not identified and example on <Atlas A2> will be executed!\n", env_var_name);
        int ret = run_example_on_A2();
        return 0;
    }
    if (!env_dev_num) {
        LOG_PRINT("[ERROR] Please check whether environment variable ENV_DEV_NUM is set correctly.\n");
        LOG_PRINT("[WARNING] For details related to ENV_DEV_NUM, see aclnnMoeDistributeCombineV2.md.\n");
        return 0;
    }
    int actual_env_dev_num = std::stoi(std::string(env_dev_num));
    if (actual_env_dev_num < DEV_NUM) {
        LOG_PRINT("[INFO] ENV_DEV_NUM = %d is less than %d, currently not supported\n", actual_env_dev_num, DEV_NUM);
        return 0;
    }
    if (!rank_table_file && !first_rank_id) {
        LOG_PRINT("[INFO] %s are not identified and example on <Atlas A3> will be executed!\n", env_var_name);
        int ret = run_example_on_A3A5();
    }
    else if (rank_table_file && !first_rank_id) {
        LOG_PRINT("[INFO] %s are not identified and example on <Atlas A5> will be executed!\n", env_var_name);
        int ret = run_example_on_A3A5();
    }
    else {
        LOG_PRINT("[WARNING] Please check whether %s are set correctly.\n", env_var_name);
    }

    return 0;
}