aclnnMoeDistributeCombineAddRmsNorm

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 combines the received data (multiplies by weights and sums them). When TP domain communication does not exist, performs AlltoAllV communication and then combines the received data (multiplies by weights and sums them), followed by fused Add and RmsNorm operations.
Formula:

rsOut = ReduceScatterV(expandX)\\ ataOut = AllToAllV(rsOut)\\ combineOut = Sum(expertScales * ataOut + expertScales * sharedExpertX)\\ x = combineOut + residualX\\ y = \frac{x}{RMS(x)} * gamma,\quad\text{where}\ RMS(x) = \sqrt{\frac{1}{H}\sum_{i=1}^{H}x_{i}^{2}+normEps}

Note: This API must be used together with aclnnMoeDistributeDispatchV2, which is equivalent to returning data along the original path used by the MoeDistributeDispatchV2 operator for data collection.

Function Prototype

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

aclnnStatus aclnnMoeDistributeCombineAddRmsNormGetWorkspaceSize(
    const aclTensor* expandX,
    const aclTensor* expertIds,
    const aclTensor* assistInfoForCombine,
    const aclTensor* epSendCounts,
    const aclTensor* expertScales,
    const aclTensor* residualX,
    const aclTensor* gamma,
    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,
    float            normEps,
    aclTensor*       yOut,
    aclTensor*       rstdOut,
    aclTensor*       xOut,
    uint64_t*        workspaceSize,
    aclOpExecutor**  executor)

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

aclnnMoeDistributeCombineAddRmsNormGetWorkspaceSize

Parameters

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension	Non-contiguous
expandX	Input	Token features expanded based on expertIds.	Must be a 2D tensor with shape (max(tpWorldSize, 1) * A, H).	BFLOAT16	ND	2	√
expertIds	Input	Indices of the topK experts for each token.	Must be a 2D tensor with shape (Bs, K).	INT32	ND	2	√
assistInfoForCombine	Input	Corresponds to the assistInfoForCombineOut output in aclnnMoeDistributeDispatchV2.	Must be a 1D tensor with shape (A * 128, ).	INT32	ND	1	√
epSendCounts	Input	Corresponds to the epRecvCounts output in aclnnMoeDistributeDispatchV2.	Must be a 1D tensor with shape (epWorldSize * max(tpWorldSize, 1) * localExpertNum, ).	INT32	ND	1	√
expertScales	Input	Weights of the topK experts for each token.	Must be a 2D tensor with shape (Bs, K).	FLOAT32	ND	2	✓
residualX	Input	Right matrix of the Add operation in AddRmsNorm.	Must be a 3D tensor with shape (Bs, 1, H).	BFLOAT16	ND	3	√
gamma	Input	Gamma input in RmsNorm.	Must be a 1D tensor with shape (H, ).	BFLOAT16	ND	1	√
tpSendCountsOptional	Input	Corresponds to the tpRecvCounts output in aclnnMoeDistributeDispatchV2.	A parameter must be passed when TP communication is involved, and a null pointer must be passed when no TP communication is involved. When TP communication is involved, it is a 1D tensor with the shape (tpWorldSize, ).	INT32	ND	-	√
xActiveMaskOptional	Input	Indicates whether tokens participate in communication.	Valid data or a null pointer can be passed. By default, all tokens participate in communication. The shape is (BS, ) for 1D and (BS, K) for 2D. When BS is inconsistent 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; reserved parameter.	Not supported in the current version. Pass a null pointer.	-	ND	-	-
sharedExpertXOptional	Input	Token after shared expert computation.	Valid data or a null pointer can be passed. The shape is (Bs, H) in 2D and (Bs, 1, H) in 3D.	BFLOAT16	ND	-	√
groupEp	Input	Name of the EP communication domain (expert parallel communication domain).	The string length ranges from [1, 128). It cannot be the same as groupTp.	STRING	ND	-	-
epWorldSize	Input	EP communication domain size.	-	INT64	ND	-	-
epRankId	Input	ID of the current card in the EP domain.	Value range: [0, epWorldSize). The epRankId of each card in the same EP communication domain must be unique.	INT64	ND	-	-
moeExpertNum	Input	Number of MoE experts.	Satisfies moeExpertNum % (epWorldSize - sharedExpertRankNum) = 0.	INT64	ND	-	-
groupTp	Input	Name of the TP communication domain (data parallel communication domain).	Must not be the same as groupEp.	STRING	ND	-	-
tpWorldSize	Input	Size of the TP communication domain.	Value range: [0, 2]. 0 and 1 indicate no TP communication, and only 2 is supported when TP communication is enabled.	INT64	ND	-	-
tpRankId	Input	ID of the local card in the TP communication domain.	Value range: [0, 1]. The tpRankId of each card in the same TP communication domain must be unique. Pass 0 when no TP communication domain is involved.	INT64	ND	-	-
expertShardType	Input	Indicates the distribution type of shared expert cards.	Currently, only 0 is supported, indicating that shared expert cards are placed before MoE expert cards.	INT64	ND	-	-
sharedExpertNum	Input	Indicates the number of shared experts.	Not supported in the current version. Pass 0.	INT64	ND	-	-
sharedExpertRankNum	Input	Indicates the number of shared expert cards.	Not supported in the current version. Pass 0.	INT64	ND	-	-
globalBS	Input	Global batch size in the EP domain.	When Bs is consistent across ranks, globalBS = Bs * epWorldSize or 0. When Bs is inconsistent across ranks, globalBS = maxBs * epWorldSize (maxBs is the maximum Bs value on a single card).	INT64	ND	-	-
outDtype	Input	Specifies the data type of the output x. Reserved parameter.	Not supported in the current version. Pass 0.	INT64	ND	-	-
commQuantMode	Input	Communication quantization type.	Not supported in the current version. Pass 0.	INT64	ND	-	-
groupListType	Input	Group list format. Reserved parameter.	Not supported in the current version. Pass 0.	INT64	ND	-	-
commAlg	Input	Indicates the communication-affinity memory layout algorithm.	Reserved field. Not supported in the current version. Pass a null pointer.	STRING	ND	-	-
normEps	Input	Prevents division-by-zero errors in AddRmsNorm.	Can be set to 1e-6.	FLOAT	ND	-	-
yOut	Output	Output result after RmsNorm.	Must be a 3D tensor.	BFLOAT16	ND	3	-
rstdOut	Output	Output result after RmsNorm.	Must be a 3D tensor with shape (Bs, 1, 1).	FLOAT32	ND	3	√
xOut	Output	Output result after the Add operation.	Must be a 3D tensor with shape (Bs, 1, H).	BFLOAT16	ND	3	-
workspaceSize	Output	Returns the workspace size to be applied for on the Device side.	-	-	-	-	-
executor	Output	Returns the op executor, which contains the computation flow of the operator.	-	-	-	-	-

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	1. The input and output shapes are not within the supported range. 2. The parameter values are not within the supported range.

aclnnMoeDistributeCombineAddRmsNorm

Parameters

Parameter	Input/Output	Description
workspace	Input	Specifies the workspace memory address applied for on the device.
workspaceSize	Input	Specifies the workspace size applied for on the device, which is obtained through the first-phase API `aclnnMoeDistributeCombineAddRmsNormGetWorkspaceSize`.
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:

- aclnnMoeDistributeCombineAddRmsNorm defaults to a deterministic implementation.

The aclnnMoeDistributeDispatchV2 API and the aclnnMoeDistributeCombineAddRmsNorm API must be used together. For details, see the example.
The values of the groupEp, epWorldSize, moeExpertNum, groupTp, tpWorldSize, expertShardType, sharedExpertNum, sharedExpertRankNum, and globalBS parameters used during the API call must be consistent across all cards, must also be consistent across different layers in the network, and must be consistent with the corresponding parameters of aclnnMoeDistributeDispatchV2.
Atlas A3 training products/Atlas A3 inference products: In this scenario, a single card contains two DIEs (short for "die"), so "this card" in the parameter description refers to a single DIE.
Shape format description in the parameter description:
- A: Indicates the maximum number of tokens to be distributed by this card. The value range is as follows:
  - For shared experts, the condition (A = BsepWorldSizesharedExpertNum / sharedExpertRankNum) must be met.
  - For MoE experts, when globalBS is 0, the condition (A>= BsepWorldSizemin(localExpertNum, K)) must be met; when globalBS is non-zero, the condition (A >= globalBS min(localExpertNum, K)) must be met.
- H: Indicates the hidden size, with a value range of [1024, 8192].
- Bs: Indicates the batch sequence size (the number of tokens finally output by the current card), with a value range of (0 < Bs ≤ 512).
- K: Indicates the top-K experts to be selected, with a value range of (0 < K ≤ 16) and satisfying (0 < K ≤ moeExpertNum).
- 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 specifies the memory size occupied by a single communication domain, in MB. It defaults to 200 MB when not configured.
- In the ep communication domain: The configured size must be (≥ 2) and satisfy (1024^2*(HCCL_BUFFSIZE - 2)/2 ≥ Bs2(H + 128)(epWorldSizelocalExpertNum+K+1)), where localExpertNum uses the local expert count on the MoE expert card.
- In the tp communication domain: The configured size must be >= (A * Align512(Align32(h * 2) + 44) + A * Align512(h * 2)) * 2.
Communication domain usage constraints:
- In a model, aclnnMoeDistributeCombineAddRmsNorm and aclnnMoeDistributeDispatchV2 only support the same EP communication domain, and no other operators are allowed in this communication domain.
- In a model, aclnnMoeDistributeCombineAddRmsNorm and aclnnMoeDistributeDispatchV2 only support the same TP communication domain or neither supports a TP communication domain. When a TP communication domain is present, 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 single super node. Cross-super-node configuration is not supported.

Example

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

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

#include <thread>
#include <iostream>
#include <string>
#include <cstring>
#include <vector>
#include "acl/acl.h"
#include "hccl/hccl.h"
#include "aclnnop/aclnn_moe_distribute_dispatch_v2.h"
#include "aclnnop/aclnn_moe_distribute_combine_add_rms_norm.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 = 2;
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};
    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 = 0;
    int64_t sharedExpertRankNum = 0;
    int64_t moeExpertNum = 8;
    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;
    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);
    }

    /*Construct the device-side input and output variables 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;
    void *residualXDeviceAddr = nullptr;
    void *sharedExpertXDeviceAddr = nullptr;
    void *gammaDeviceAddr = nullptr;
    void *yOutDeviceAddr = nullptr;
    void *rstdOutDeviceAddr = 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 *gamma = nullptr;
    aclTensor *yOut = nullptr;
    aclTensor *rstdOut = 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{(sharedExpertRankNum > 0) ? moeExpertNum + 1 : moeExpertNum, 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> residualXShape{BS, 1, H};
    std::vector<int64_t> sharedExpertXShape{BS, 1, H};
    std::vector<int64_t> gammaShape{H};
    std::vector<int64_t> yOutShape{BS, 1, H};
    std::vector<int64_t> rstdOutShape{BS, 1, 1};
    std::vector<int64_t> xOutShape{BS, 1, 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 residualXShapeSize = GetShapeSize(residualXShape);
    int64_t sharedExpertXShapeSize = GetShapeSize(sharedExpertXShape);
    int64_t gammaShapeSize = GetShapeSize(gammaShape);
    int64_t yOutShapeSize = GetShapeSize(yOutShape);
    int64_t rstdOutShapeSize = GetShapeSize(rstdOutShape);
    int64_t xOutShapeSize = GetShapeSize(xOutShape);

    // Construct the host-side variables.
    std::vector<int16_t> xHostData(xShapeSize, 1);
    std::vector<int32_t> expertIdsHostData;
    for (int32_t token_id = 0; token_id < expertIdsShape[0]; token_id++) {
        // Each token is sent 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.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> residualXHostData(residualXShapeSize, 1);
    std::vector<int16_t> sharedExpertXHostData(sharedExpertXShapeSize, 1);
    std::vector<int16_t> gammaHostData(gammaShapeSize, 1);
    std::vector<int16_t> yOutHostData(yOutShapeSize, 0);
    std::vector<float> rstdOutHostData(rstdOutShapeSize, 0);
    std::vector<int16_t> xOutHostData(xOutShapeSize, 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);
    ret = CreateAclTensor(residualXHostData, residualXShape, &residualXDeviceAddr, aclDataType::ACL_BF16, &residualX);
    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(gammaHostData, gammaShape, &gammaDeviceAddr, aclDataType::ACL_BF16, &gamma);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(yOutHostData, yOutShape, &yOutDeviceAddr, aclDataType::ACL_BF16, &yOut);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(rstdOutHostData, rstdOutShape, &rstdOutDeviceAddr, aclDataType::ACL_FLOAT, &rstdOut);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(xOutHostData, xOutShape, &xOutDeviceAddr, aclDataType::ACL_BF16, &xOut);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    /*Declare the variables required for operator execution.*/
    uint64_t dispatchWorkspaceSize = 0;
    aclOpExecutor *dispatchExecutor = nullptr;
    void *dispatchWorkspaceAddr = nullptr;

    uint64_t combineAddRmsNormWorkspaceSize = 0;
    aclOpExecutor *combineAddRmsNormExecutor = nullptr;
    void *combineWorkspaceAddr = nullptr;

    /**************************************** Call dispatch ********************************************/

    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, nullptr, expandX, dynamicScales, expandIdx, 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.dispatchStream);
    ret = aclrtSynchronizeStreamWithTimeout(args.dispatchStream, 10000);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMoeDistributeDispatchV2 failed. ret = %d \n", ret);  \
        return ret);

    /**************************************** Call combineAddRmsNorm ********************************************/
    // Call the first-phase API.
    ret = aclnnMoeDistributeCombineAddRmsNormGetWorkspaceSize(
        expandX, expertIds, expandIdx, epRecvCounts, expertScales, residualX, gamma, tpRecvCounts, nullptr, nullptr,
        nullptr, nullptr, nullptr, sharedExpertX, hcomEpName, EP_WORLD_SIZE, args.epRankId, moeExpertNum, hcomTpName, TP_WORLD_SIZE,
        args.tpRankId, expertShardType, sharedExpertNum, sharedExpertRankNum, globalBS, outDtype, commQuantMode,
        groupList_type, nullptr, 1e-6, yOut, rstdOut, xOut, &combineAddRmsNormWorkspaceSize, &combineAddRmsNormExecutor);
    CHECK_RET(ret == ACL_SUCCESS,
        LOG_PRINT("[ERROR] aclnnMoeDistributeCombineAddRmsNormGetWorkspaceSize failed. ret = %d \n", ret); return ret);
    // Apply for device memory based on the workspaceSize calculated by the first-phase API.
    if (combineAddRmsNormWorkspaceSize > 0) {
        ret = aclrtMalloc(&combineWorkspaceAddr, combineAddRmsNormWorkspaceSize, 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 = aclnnMoeDistributeCombineAddRmsNorm(combineWorkspaceAddr, combineAddRmsNormWorkspaceSize, combineAddRmsNormExecutor, args.combineStream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMoeDistributeCombineAddRmsNorm 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 aclnnMoeDistributeDispatchV2 and aclnnMoeDistributeCombineAddRmsNorm                      \
                execute successfully.\n", args.rankId);

    // Release device resources.
    if (dispatchWorkspaceSize > 0) {
        aclrtFree(dispatchWorkspaceAddr);
    }
    if (combineAddRmsNormWorkspaceSize > 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 (gamma != nullptr) {
        aclDestroyTensor(gamma);
    }
    if (yOut != nullptr) {
        aclDestroyTensor(yOut);
    }
    if (rstdOut != nullptr) {
        aclDestroyTensor(rstdOut);
    }
    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 (residualXDeviceAddr != nullptr) {
        aclrtFree(residualXDeviceAddr);
    }
    if (sharedExpertXDeviceAddr != nullptr) {
        aclrtFree(sharedExpertXDeviceAddr);
    }
    if (gammaDeviceAddr != nullptr) {
        aclrtFree(gammaDeviceAddr);
    }
    if (yOutDeviceAddr != nullptr) {
        aclrtFree(yOutDeviceAddr);
    }
    if (rstdOutDeviceAddr != nullptr) {
        aclrtFree(rstdOutDeviceAddr);
    }
    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];
    // Each thread calls its respective card 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].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;
}