aclnnMoeDistributeCombineV4

📄 View Source Code

Applicable Products

Function

• Description: When TP domain communication exists, first perform ReduceScatterV communication, then perform AllToAllV communication, and finally combine the received data (multiply by weight and sum); when TP domain communication does not exist, perform AllToAllV communication, and finally combine the received data (multiply by weight and sum).

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

  Added the communication latency collection feature, which records the communication time of each card. This feature is enabled by passing the performanceInfoOptional parameter. It is recommended to use this feature together with the DeepXTrace tool. The communication latency of each card for a single operator call is accumulated onto this tensor, and users should clear it to zero as needed before use.

• 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 The API must be used together with aclnnMoeDistributeDispatchV4, which means data is returned along the same path used by aclnnMoeDistributeDispatchV4 to collect data.

Function Prototype

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

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

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

aclnnMoeDistributeCombineV4GetWorkspaceSize

• Parameters

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

  Atlas A2 training products/Atlas A2 inference products:

  • commAlg supports nullptr, "", "fullmesh", and "hierarchy". It is recommended to configure "hierarchy" with driver version ≥ 25.0.RC1.1. nullptr and "" select the algorithm based on the HCCL environment variable (not recommended). "fullmesh" directly transmits tokens via RDMA. "hierarchy" sends data in two stages (intra-node and inter-node) to reduce the amount of cross-node data.

  • The shared expert scenario is not supported.

  • The shape of epSendCounts is (moeExpertNum + 2 * globalBS * K * serverNum, ), where K refers to the number of topK experts. The first moeExpertNum values indicate the number of tokens received from each card in the EP communication domain, and the last 2 * globalBS * K * serverNum values are used to store the number of tokens and communication area offsets for which combine can perform reduce in advance before inter-node/intra-node communication. When globalBS=0, it is calculated as BS * epWorldSize.

  • TP domain communication is not supported currently.

  • xActiveMaskOptional depends on commAlg. For "fullmesh", it is required to be a 1D tensor with shape (BS, ). true must be placed before false (e.g., {true, false, true} is invalid). "hierarchy" is not supported in the current version; pass a null pointer.

  • expandScalesOptional is required to be a 1D tensor with shape (A, ).

  • sharedExpertXOptional is a reserved parameter, not supported in the current version; pass a null pointer.

  • epWorldSize depends on commAlg. For "fullmesh", supported values are 2, 3, 4, 5, 6, 7, 8, 16, 32, 64, 128, 192, 256, and 384. For "hierarchy", supported values are 16, 32, and 64.

  • moeExpertNum: value range (0, 512].

  • groupTp: not supported in the current version. Pass an empty string.

  • tpWorldSize: not supported in the current version. Pass 0.

  • tpRankId: not supported in the current version. Pass 0.

  • expertShardType: not supported in the current version. Pass 0.

  • sharedExpertNum is not supported in the current version. Pass 0.

  • sharedExpertRankNum is not supported in the current version. Pass 0.

  • commQuantMode has a value range of 0 or 2 (0 indicates non-quantized, 2 indicates int8 quantization). The value 2 is supported only when commAlg is "hierarchy" or when HCCL_INTRA_PCIE_ENABLE=1, HCCL_INTRA_ROCE_ENABLE=0, and the driver version is ≥ 25.0.RC1.1.

  • expandScalesOptional is required to be a 1D Tensor.

  • elasticInfoOptional is not supported in the current version. Pass a null pointer.

  • oriXOptional: When commAlg="hierarchy", this is not supported in the current version. Pass a null pointer.

  • constExpertAlpha1Optional: Reserved parameter. Not supported in the current version. Pass a null pointer.

  • constExpertAlpha2Optional: Reserved parameter. Not supported in the current version. Pass a null pointer.

  • constExpertVOptional: Reserved parameter. Not supported in the current version. Pass a null pointer.

  • zeroExpertNum: When commAlg="fullmesh", value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. The valid zero expert ID value is [moeExpertNum, moeExpertNum + zeroExpertNum).

  • copyExpertNum: When commAlg="fullmesh", value range: [0, MAX_INT32), MAX_INT32 = 2^31 - 1. The valid copy expert ID value is [moeExpertNum + zeroExpertNum, moeExpertNum + zeroExpertNum + copyExpertNum).

  • constExpertNum: Not supported in the current version. Pass 0.

  • performanceInfoOptional: Optionally pass valid data or pass null pointer. Passing a null pointer indicates that the communication latency recording function is not enabled. When passing valid data, it is required to be a 1D tensor with shape (ep_world_size,) and data type int64. The data format is required to be ND.

  Atlas A3 training products/Atlas A3 inference products:

  • commAlg supports two input modes: "" and "hierarchy". "": default value, hierarchy across super templates is not enabled; "hierarchy": enables hierarchy across super templates. This template only supports scenarios where tpWorldSize is 1 and the shared expert is 0, and does not support variable BS, 2D mask, special expert, or performanceInfo scenarios.

  • 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 is required to be a 1D or 2D tensor (1D with shape (BS, ), 2D with shape (BS, K)); in 1D, true must be placed before false; in 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: required to be a 2D or 3D Tensor (when 2D, shape is (BS, H); when 3D, the product of the first two dimensions equals BS and the third dimension equals H); can be passed or not, and when passed, sharedExpertRankNum is required to be 0.

  • epWorldSize: value range [2, 768]; in the scenario where commAlg="hierarchy", the value range is [16, 256] and must be an integer multiple of 16.

  • moeExpertNum: value range (0, 1024]; in the scenario where commAlg="hierarchy", the value range is (0, 512].

  • groupTp: string length range is [0, 128), cannot be the same as groupEp, and passing empty is supported only when there is no TP domain communication.

  • tpWorldSize: value range [0, 2], where 0 and 1 indicate no TP domain communication, and when TP domain communication exists, only 2 is supported.

  • tpRankId: value range [0, 1]. The tpRankId of each card in the same TP communication domain must be unique. Pass 0 when there is no TP domain communication.

  • expertShardType: currently only supports passing 0, which indicates that shared expert cards are placed before MoE expert cards.

  • sharedExpertNum: current value range [0, 4].

  • sharedExpertRankNum: value range [0, epWorldSize). When it is 0, sharedExpertNum must be 0 or 1. When it is not 0, sharedExpertRankNum % sharedExpertNum = 0 must be satisfied.

  • commQuantMode: value range 0 or 2 (0 indicates non-quantized, 2 indicates int8 quantization). The value 2 can be enabled only when tpWorldSize < 2.

  • expandScalesOptional: reserved parameter, not supported in the current version. Pass a null pointer.

  • elasticInfoOptional: reserved parameter, not supported in the current version. Pass a null pointer.

  • constExpertAlpha1Optional: optionally pass valid data or a null pointer. When constExpertNum is not 0, valid input must be passed. When passing valid data, it is required to be a 2D tensor with shape (constExpertNum, H), and the data type must keep consistent with expandX.

  • constExpertAlpha2Optional: optionally pass valid data or a null pointer. When constExpertNum is not 0, valid input must be passed. When passing valid data, it is required to be a 2D tensor with shape (constExpertNum, H), and the data type must keep consistent with expandX.

  • constExpertVOptional: optionally pass valid data or a null pointer. When constExpertNum is not 0, valid input must be passed. When passing valid data, it is required to be a 2D tensor with shape (constExpertNum, H), and the data type must keep consistent with expandX.

  • zeroExpertNum value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. The valid zero expert ID value is [moeExpertNum, moeExpertNum + zeroExpertNum).

  • copyExpertNum value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. The valid copy expert ID value is [moeExpertNum + zeroExpertNum, moeExpertNum + zeroExpertNum + copyExpertNum).

  • constExpertNum value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. The valid constant expert ID value is [moeExpertNum + zeroExpertNum + copyExpertNum, moeExpertNum + zeroExpertNum + copyExpertNum + constExpertNum).

  • performanceInfoOptional: optionally pass valid data or a null pointer. Passing a null pointer indicates that the communication latency recording function is not enabled. When passing valid data, it is required to be a 1D tensor with shape (ep_world_size,), supporting int64 data type. The data format is required to be ND.

  Ascend 950PR/Ascend 950DT：

  • commAlg is not supported in the current version. Pass a null pointer.

  • The shape of epSendCounts is (epWorldSize * max(tpWorldSize, 1) * localExpertNum, ).

  • TP communication domain is not supported.

  • xActiveMaskOptional is required to be a 1D or 2D tensor (shape is (BS, ) for 1D, and (BS, K) for 2D). For 1D, true values must be placed before false values. For 2D, if all K values corresponding to a token are false, it does not participate in communication.

  • expandScalesOptional: reserved parameter, not supported in the current version. Pass null pointer.

  • sharedExpertXOptional: required to be a 2D or 3D tensor (shape is (BS, H) for 2D; for 3D, the product of the first two dimensions equals BS and the third dimension equals H). Optionally pass or not; when passed, sharedExpertRankNum must be 0.

  • epWorldSize: value range [2, 768].

  • moeExpertNum: value range (0, 1024].

  • groupTp: not supported in the current version. Pass empty characters.

  • tpWorldSize is not supported in the current version. Pass 0.

  • tpRankId is not supported in the current version. Pass 0.

  • expertShardType currently only supports passing 0, which indicates that shared expert cards are placed before MoE expert cards.

  • sharedExpertNum currently has a value range of [0, 4].

  • sharedExpertRankNum has a value range of [0, epWorldSize). When it is 0, sharedExpertNum must be 0 or 1. When it is not 0, sharedExpertRankNum % sharedExpertNum = 0 must be satisfied.

  • commQuantMode value range: 0 or 2 (0 indicates non-quantized, 2 indicates int8 quantization). The value is 2 only when tpWorldSize < 2 can be enabled.

  • expandScalesOptional is a reserved parameter, not supported in the current version. Pass a null pointer.

  • elasticInfoOptional is a reserved parameter, not supported in the current version. Pass a null pointer.

  • constExpertAlpha1Optional optionally pass valid data or empty. When constExpertNum is not 0, valid input must be passed. When passing valid data, it is required to be a 2D tensor with shape (constExpertNum, H), and the data type must keep consistent with expandX.

  • constExpertAlpha2Optional optionally pass valid data or empty. When constExpertNum is not 0, valid input must be passed. When passing valid data, it is required to be a 2D tensor with shape (constExpertNum, H), and the data type must keep consistent with expandX.

  • constExpertVOptional: optionally pass valid data or a null pointer. When constExpertNum is not 0, valid input must be passed. When passing valid data, it is required to be a 2D tensor with shape (constExpertNum, H), and the data type must be consistent with expandX.

  • zeroExpertNum: value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. Valid zero expert IDs are in the range [moeExpertNum, moeExpertNum + zeroExpertNum).

  • copyExpertNum: value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. Valid copy expert IDs are in the range [moeExpertNum + zeroExpertNum, moeExpertNum + zeroExpertNum + copyExpertNum).

  • constExpertNum: value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. Valid constant expert IDs are in the range [moeExpertNum + zeroExpertNum + copyExpertNum, moeExpertNum + zeroExpertNum + copyExpertNum + constExpertNum).

  • performanceInfoOptional: reserved parameter, not supported in the current version. Pass a null pointer.

• 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	Return Code	Description
  ACLNN_ERR_PARAM_NULLPTR	161001	The required input and output parameter Tensor is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The input and output data types are not within the supported range.
  ACLNN_ERR_INNER_TILING_ERROR	561002	The shapes of the input and output are not within the supported range.
  		The parameter value is not within the supported range.

aclnnMoeDistributeCombineV4

• Parameters

  Parameter	Input/Output	Description
  workspace	Input	Address of the workspace memory applied on the device side.
  workspaceSize	Input	Size of the workspace applied on the device side, obtained from the first-phase API aclnnMoeDistributeCombineV4GetWorkspaceSize.
  executor	Input	Operator executor, which contains the operator computation flow.
  stream	Input	Specifies the stream that executes the task.

• Return Value

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

Constraints

• Deterministic Computation:

  • aclnnMoeDistributeCombineV4 defaults to a deterministic implementation.

• Driver Constraints:

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

• API Matching Constraints:

  • aclnnMoeDistributeDispatchV4 and aclnnMoeDistributeCombineV4 must be used together. The assistInfoForCombineOut, epRecvCountsOut, tpRecvCountsOut, and expandScalesOut output by the former must be directly passed to the corresponding parameters of the latter. The business logic must not rely on the specific values of these tensors.

• Parameter Consistency Constraints:

  • The groupEp, epWorldSize, moeExpertNum, groupTp, tpWorldSize, expertShardType, sharedExpertNum, sharedExpertRankNum, globalBS, commAlg parameters and the HCCL_BUFFSIZE value must be consistent across all cards, and must be consistent with the corresponding parameters of aclnnMoeDistributeDispatchV4.

• Product-Specific Constraints:

  • Atlas A3 training products/Atlas A3 inference products : A single card contains dual DIE (dies). "This card" in the parameter description refers to a single DIE.

• Shape variable constraints:

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

• Environment Variable Constraints:

  • HCCL_BUFFSIZE: Before calling this API, check whether the value of the HCCL_BUFFSIZE environment variable is reasonable. This environment variable indicates the memory size occupied by a single communication domain, in MB. The default value is 200 MB when not configured.

    • Atlas A2 training products/Atlas A2 inference products :

      • commAlg configured as "" or nullptr: Select the "fullmesh" or "hierarchy" formula based on the HCCL_INTRA_PCIE_ENABLE and HCCL_INTRA_ROCE_ENABLE environment variable configurations.

      • commAlg configured as "fullmesh": required to be = 2 * (BS * epWorldSize * min(localExpertNum, K) * H * sizeof(uint16) + 2MB).

      • commAlg configured as "hierarchy": required to be >= (moeExpertNum + epWorldSize / 4) * Align512(maxBS * (H * 2 + 16 * Align8(K))) * 1B + 8MB, where Align8(x) = ((x + 8 - 1) / 8) * 8, Align512(x) = ((x + 512 - 1) / 512) * 512.

    • Atlas A3 training products/Atlas A3 inference products :

      • When commAlg is configured as "" or not configured: required to be >= 2 and satisfy >= 2 * (localExpertNum * maxBS * epWorldSize * Align512(Align32(2 * H) + 44) + (K + sharedExpertNum) * maxBS * Align512(2 * H)), where localExpertNum must use the number of experts on the local MoE expert card, and Align512(x) = ((x + 512 - 1) / 512) * 512, Align32(x) = ((x + 32 - 1) / 32) * 32.

      • When commAlg is configured as "hierarchy": the value must satisfy (moeExpertNum * maxBS * (H * 2 + (3 * (K + 7) / 8 * 8)) * 4 + 64) + 404 * 1024 * 1024.

  • HCCL_INTRA_PCIE_ENABLE/HCCL_INTRA_ROCE_ENABLE:

    • Atlas A2 training products/Atlas A2 inference products: This environment variable is no longer recommended. It is recommended to configure commAlg as "hierarchy".

  • HCCL_LOGIC_SUPERPOD_ID

    • Atlas A3 training products/Atlas A3 inference products: When commAlg is configured as "hierarchy", this environment variable must be configured based on different super nodes. For example, set them as export HCCL_LOGIC_SUPERPOD_ID=0 and export HCCL_LOGIC_SUPERPOD_ID=1 for two machines respectively.

• Communication Domain Usage Constraints:

  • The aclnnMoeDistributeCombineV4 series operators and aclnnMoeDistributeDispatchV4 in a model only support the same EP communication domain, and no other operators are allowed in this communication domain.

  • The aclnnMoeDistributeCombineV4 series operators and aclnnMoeDistributeDispatchV4 in a model 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 : The communication domain supports cross-supernode only when commAlg is configured as "hierarchy". In other scenarios, nodes within a communication domain must be within a single supernode, and cross-supernode is not supported.

• Networking constraints:

  • Atlas A2 training products/Atlas A2 inference products : In multi-machine scenarios, only switch networking is supported, and dual-machine direct-connect networking is not supported.

• Other Constraints:

  • The "/" in the formula indicates integer division.

  • moeExpertNum + zeroExpertNum + copyExpertNum + constExpertNum < MAX_INT32.

Example

• Atlas A2 training products/Atlas A2 inference products:

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

  • Compile operators: The operator compilation command is as follows. Both moe_distribute_dispatch_v2 and moe_distribute_combine_v2 operators need to be compiled, and these two operators must be executed in pairs.

    bash build.sh --pkg --soc=ascend910b --ops=moe_distribute_dispatch_v2,moe_distribute_combine_v2

  • Create sample code for Atlas A2 training products/Atlas A2 inference products: After compilation, refer to the existing test_aclnn_moe_distribute_combine_v2.cpp file in the operator examples directory, and create a new test file test_aclnn_moe_distribute_combine_v4.cpp using the A2 sample code.

  • Execute operator samples: The sample operator execution command is as follows. This command executes all sample code files in the operator examples directory.

    bash build.sh --run_example --ops=moe_distribute_combine_v2 eager cust

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

    #include <thread>
    #include <iostream>
    #include <string>
    #include <cstring>
    #include <vector>
    #include <memory>
    #include <cstdio>
    #include "acl/acl.h"
    #include "hccl/hccl.h"
    #include "aclnn/opdev/fp16_t.h"
    #include "aclnnop/aclnn_moe_distribute_dispatch_v4.h"
    #include "aclnnop/aclnn_moe_distribute_combine_v4.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 dispatchV4Stream;
        aclrtStream combineV4Stream;
        aclrtContext context;
    };
    
    const uint32_t EP_WORLD_SIZE_A2 = 8;
    const uint32_t TP_WORLD_SIZE_A2 = 1;
    const uint32_t DEV_NUM_A2 = EP_WORLD_SIZE_A2 * TP_WORLD_SIZE_A2;
    
    int64_t GetShapeSize(const std::vector<int64_t> &shape)
    {
        int64_t shape_size = 1;
        for (auto i : shape) {
            shape_size *= i;
        }
        return shape_size;
    }
    
    template<typename T>
    int CreateAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr,
        aclDataType dataType, aclTensor **tensor)
    {
        auto size = GetShapeSize(shape) * sizeof(T);
        auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc failed. ret: %d\n", ret); return ret);
        ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMemcpy failed. ret: %d\n", ret); return ret);
        std::vector<int64_t> strides(shape.size(), 1);
        for (int64_t i = shape.size() - 2; i >= 0; i--) {
            strides[i] = shape[i +1] * strides[i + 1];
        }
        *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
            shape.data(), shape.size(), *deviceAddr);
        return 0;
    }
    
    void DestroyTensor(aclTensor *tensor) {
        if (tensor != nullptr) {
            aclDestroyTensor(tensor);
        }
    }
    
    void FreeDeviceAddr(void *deviceAddr) {
        if (deviceAddr != nullptr) {
            aclrtFree(deviceAddr);
        }
    }
    
    int launchOneThreadDispatchV4AndCombineV4_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, dispatchV4Stream = %p, combineV4Stream = %p, \
                    context = %p\n", args.rankId, hcomEpName, args.dispatchV4Stream, args.combineV4Stream,                 \
                    args.context);
    
        int64_t BS = 32;
        int64_t H = 7168;
        int64_t K = 8;
        int64_t expertShardType = 0;
        int64_t sharedExpertNum = 0;
        int64_t sharedExpertRankNum = 0;
        int64_t moeExpertNum = 256;
        int64_t quantMode = 0;
        int64_t globalBS = BS * EP_WORLD_SIZE_A2;
        int64_t expertTokenNumsType = 1;
        int64_t outDtype = 0;
        int64_t commQuantMode = 0;
        int64_t groupList_type = 1;
        int64_t localExpertNum;
        int64_t A;
        int64_t zeroExpertNum = 0;
        int64_t copyExpertNum = 0;
        int64_t constExpertNum = 0; // A3 only
        std::string commAlg = "fullmesh";
        if (args.epRankId < sharedExpertRankNum) {
            localExpertNum = 1;
            A = globalBS / sharedExpertRankNum;
        } else {
            localExpertNum = moeExpertNum / (EP_WORLD_SIZE_A2 - sharedExpertRankNum);
            A = globalBS * (localExpertNum < K ? localExpertNum : K);
        }
    
        void *xDeviceAddr = nullptr;
        void *expertIdsDeviceAddr = nullptr;
        void *scalesDeviceAddr = nullptr;
        void *expertScalesDeviceAddr = nullptr;
    
        void *expandXDeviceAddr = nullptr;
        void *dynamicScalesDeviceAddr = nullptr;
        void *assistInfoForCombineDeviceAddr = nullptr;
        void *expertTokenNumsDeviceAddr = nullptr;
        void *epRecvCountsDeviceAddr = nullptr;
        void *tpRecvCountsDeviceAddr = nullptr;
        void *expandScalesDeviceAddr = nullptr;
    
        // Zero-expert scenario input
        void *oriXDeviceAddr = nullptr;
    
        // Performance timing
        void *performanceInfoDeviceAddr = nullptr;
    
        void *xOutDeviceAddr = nullptr;
    
        aclTensor *x = nullptr;
        aclTensor *expertIds = nullptr;
        aclTensor *scales = nullptr;
        aclTensor *xActiveMask = nullptr;
        aclTensor *expertScales = nullptr;
    
        aclTensor *elasticInfo = nullptr; // A3
        aclTensor *performanceInfo = nullptr;
        aclTensor *expandX = nullptr;
        aclTensor *dynamicScales = nullptr;
        aclTensor *assistInfoForCombine = nullptr; // expandIdx
        aclTensor *expertTokenNums = nullptr;
        aclTensor *epRecvCounts = nullptr;
        aclTensor *tpRecvCounts = nullptr;
        aclTensor *expandScales = nullptr;
    
        aclTensor *activationScale = nullptr; // Reserved parameter
        aclTensor *weightScale = nullptr; // reserved parameter
        aclTensor *groupList = nullptr; // reserved parameter
    
        aclTensor *sharedExpertX = nullptr; // A3
    
        aclTensor *oriX = nullptr;
        aclTensor *constExpertAlpha1 = nullptr; // A3
        aclTensor *constExpertAlpha2 = nullptr; // A3
        aclTensor *constExpertV = nullptr; // A3
    
        aclTensor *xOut = nullptr;
    
        //Define the dimensions of each variable in the current scenario.
        std::vector<int64_t> xShape{BS, H};
        std::vector<int64_t> expertIdsShape{BS, K};
        std::vector<int64_t> scalesShape{moeExpertNum + 1, H};
        std::vector<int64_t> expertScalesShape{BS, K};
    
        std::vector<int64_t> expandXShape{TP_WORLD_SIZE_A2 * A, H};
        std::vector<int64_t> dynamicScalesShape{TP_WORLD_SIZE_A2 * A};
        std::vector<int64_t> assistInfoForCombineShape{A * 128};
        std::vector<int64_t> expertTokenNumsShape{localExpertNum};
        std::vector<int64_t> epRecvCountsShape{TP_WORLD_SIZE_A2 * localExpertNum * EP_WORLD_SIZE_A2}; // Non-hierarchical
        std::vector<int64_t> tpRecvCountsShape{TP_WORLD_SIZE_A2};
        std::vector<int64_t> expandScalesShape{A};
    
        std::vector<int64_t> oriXShape{BS, H};
        std::vector<int64_t> xOutShape{BS, H};
        std::vector<int64_t> performanceInfoShape{EP_WORLD_SIZE_A2};
    
        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 performanceInfoShapeSize = GetShapeSize(performanceInfoShape);
        int64_t oriXSize = GetShapeSize(oriXShape);
    
        int64_t xOutShapeSize = GetShapeSize(xOutShape);
    
        std::vector<int16_t> xHostData(xShapeSize, 1);
        std::vector<int32_t> expertIdsHostData;
        for (int32_t token_id = 0; token_id < expertIdsShape[0]; token_id++) {
            for (int32_t k_id = 0; k_id < expertIdsShape[1]; k_id++) {
                expertIdsHostData.push_back(k_id);
            }
        }
    
        std::vector<float> scalesHostData(scalesShapeSize, 0.1);
        std::vector<float> expertScalesHostData(expertScalesShapeSize, 0.1);
    
        std::vector<int16_t> expandXHostData(expandXShapeSize, 0);
        std::vector<float> dynamicScalesHostData(dynamicScalesShapeSize, 0);
        std::vector<int32_t> assistInfoForCombineHostData(assistInfoForCombineShapeSize, 0);
        std::vector<int64_t> expertTokenNumsHostData(expertTokenNumsShapeSize, 0);
        std::vector<int32_t> epRecvCountsHostData(epRecvCountsShapeSize, 0);
        std::vector<int32_t> tpRecvCountsHostData(tpRecvCountsShapeSize, 0);
        std::vector<float> expandScalesHostData(expandScalesShapeSize, 0);
    
        std::vector<int16_t> oriXHostData(oriXSize, 1);
        std::vector<int16_t> performanceInfoHostData(performanceInfoShapeSize, 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(oriXHostData, oriXShape, &oriXDeviceAddr, aclDataType::ACL_BF16, &oriX);
        CHECK_RET(ret == ACL_SUCCESS, return ret);
        ret = CreateAclTensor(performanceInfoHostData, performanceInfoShape, &performanceInfoDeviceAddr, aclDataType::ACL_BF16, &performanceInfo);
        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 = aclnnMoeDistributeDispatchV4GetWorkspaceSize(x, expertIds, (quantMode > 0 ? scales : nullptr), xActiveMask,
                expertScales, elasticInfo, performanceInfo, hcomEpName, EP_WORLD_SIZE_A2, args.epRankId, moeExpertNum, "", TP_WORLD_SIZE_A2,
                args.tpRankId, expertShardType, sharedExpertNum,sharedExpertRankNum, quantMode, globalBS,
                expertTokenNumsType, commAlg.c_str(), zeroExpertNum, copyExpertNum, constExpertNum, expandX, dynamicScales, assistInfoForCombine, expertTokenNums, epRecvCounts,
                tpRecvCounts, expandScales, &dispatchWorkspaceSize, &dispatchExecutor);
    
        CHECK_RET(ret == ACL_SUCCESS,
            LOG_PRINT("[ERROR] aclnnMoeDistributeDispatchV4GetWorkspaceSize 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 = aclnnMoeDistributeDispatchV4(dispatchWorkspaceAddr, dispatchWorkspaceSize,
                                            dispatchExecutor, args.dispatchV4Stream);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMoeDistributeDispatchV4 failed. ret = %d \n", ret);  \
                return ret);
        ret = aclrtSynchronizeStreamWithTimeout(args.dispatchV4Stream, 10000);
                    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] dispatch aclrtSynchronizeStreamWithTimeout failed. ret = %d \n", ret);  \
                return ret);
        LOG_PRINT("[INFO] device_%d aclnnMoeDistributeDispatchV4 execute successfully.\n", args.rankId);
        /**************************************** Call combine ********************************************/
        // Call the first-phase API.
        ret = aclnnMoeDistributeCombineV4GetWorkspaceSize(expandX, expertIds,
                                                            assistInfoForCombine, epRecvCounts,
                                                            expertScales, tpRecvCounts,
                                                            xActiveMask, activationScale, weightScale,
                                                            groupList, expandScales, sharedExpertX,
                                                            elasticInfo, oriX, constExpertAlpha1, constExpertAlpha2, constExpertV, performanceInfo,
                                                            hcomEpName, EP_WORLD_SIZE_A2, args.epRankId, moeExpertNum,
                                                            "", TP_WORLD_SIZE_A2, args.tpRankId, expertShardType,
                                                            sharedExpertNum, sharedExpertRankNum, globalBS, outDtype,
                                                            commQuantMode, groupList_type, commAlg.c_str(), zeroExpertNum, copyExpertNum, constExpertNum, xOut,
                                                            &combineWorkspaceSize, &combineExecutor);
        CHECK_RET(ret == ACL_SUCCESS,
            LOG_PRINT("[ERROR] aclnnMoeDistributeCombineV4GetWorkspaceSize 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 = aclnnMoeDistributeCombineV4(combineWorkspaceAddr, combineWorkspaceSize, combineExecutor, args.combineV4Stream);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMoeDistributeCombineV4 failed. ret = %d \n", ret);
            return ret);
        // (Boilerplate) Synchronously wait for task execution to complete.
        ret = aclrtSynchronizeStreamWithTimeout(args.combineV4Stream, 10000);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSynchronizeStreamWithTimeout failed. ret = %d \n", ret);
            return ret);
        LOG_PRINT("[INFO] device_%d aclnnMoeDistributeDispatchV4 and aclnnMoeDistributeCombineV4                      \
                    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(performanceInfo);
        DestroyTensor(expandX);
        DestroyTensor(dynamicScales);
        DestroyTensor(assistInfoForCombine);
        DestroyTensor(expertTokenNums);
        DestroyTensor(epRecvCounts);
        DestroyTensor(tpRecvCounts);
        DestroyTensor(expandScales);
        DestroyTensor(activationScale);
        DestroyTensor(weightScale);
        DestroyTensor(groupList);
        DestroyTensor(sharedExpertX);
        DestroyTensor(oriX);
        DestroyTensor(constExpertAlpha1);
        DestroyTensor(constExpertAlpha2);
        DestroyTensor(constExpertV);
        DestroyTensor(xOut);
    
        FreeDeviceAddr(xDeviceAddr);
        FreeDeviceAddr(expertIdsDeviceAddr);
        FreeDeviceAddr(scalesDeviceAddr);
        FreeDeviceAddr(expertScalesDeviceAddr);
        FreeDeviceAddr(expandXDeviceAddr);
        FreeDeviceAddr(dynamicScalesDeviceAddr);
        FreeDeviceAddr(assistInfoForCombineDeviceAddr);
        FreeDeviceAddr(expertTokenNumsDeviceAddr);
        FreeDeviceAddr(epRecvCountsDeviceAddr);
        FreeDeviceAddr(tpRecvCountsDeviceAddr);
        FreeDeviceAddr(expandScalesDeviceAddr);
        FreeDeviceAddr(oriXDeviceAddr);
        FreeDeviceAddr(performanceInfoDeviceAddr);
        FreeDeviceAddr(xOutDeviceAddr);
    
        HcclCommDestroy(args.hcclEpComm);
        aclrtDestroyStream(args.dispatchV4Stream);
        aclrtDestroyStream(args.combineV4Stream);
        aclrtDestroyContext(args.context);
        LOG_PRINT("[INFO] device_%d DeStroy.\n", args.rankId);
        aclrtResetDevice(args.rankId);
        LOG_PRINT("[INFO] device_%d Reset.\n", args.rankId);
        return 0;
    }
    int main(int argc, char *argv[])
    {
        LOG_PRINT("[INFO] run_example_on_A2.\n");
        int ret = aclInit(nullptr);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclInit failed, ret = %d\n", ret); return ret);
        aclrtStream dispatchV4Stream[DEV_NUM_A2];
        aclrtStream combineV4Stream[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(&dispatchV4Stream[rankId]);
            CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateStream failed, ret = %d\n", ret); return ret);
            ret = aclrtCreateStream(&combineV4Stream[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].dispatchV4Stream = dispatchV4Stream[rankId];
            args[rankId].combineV4Stream = combineV4Stream[rankId];
            args[rankId].context = context[rankId];
            threads[rankId].reset(new(std::nothrow) std::thread(&launchOneThreadDispatchV4AndCombineV4_A2, std::ref(args[rankId])));
        }
    
        for(uint32_t rankId = 0; rankId < DEV_NUM_A2; rankId++) {
            threads[rankId]->join();
        }
    
        aclFinalize();
        LOG_PRINT("[INFO] aclFinalize success\n");
        return 0;
    }

• Ascend 950PR/Ascend 950DT: Refer to the preparation section and sample code in the example of aclnnMoeDistributeCombineV2, and reset the involved variables according to the constraints above. For the scenario parameters newly added in V4 compared to V3, pass values as described in the Parameters section above.

• Atlas A3 training products/Atlas A3 inference products:

  For specific compilation and execution, see Compile and Run Samples.

• The sample code is as follows, for reference only

  #include <thread>
  #include <iostream>
  #include <string>
  #include <vector>
  #include <unordered_set>
  #include "acl/acl.h"
  #include "hccl/hccl.h"
  #include "aclnnop/aclnn_moe_distribute_dispatch_v4.h"
  #include "aclnnop/aclnn_moe_distribute_combine_v4.h"
  
  #define CHECK_RET(cond, return_expr) \
      do {                             \
          if (!(cond)) {               \
              return_expr;             \
          }                            \
      } while (0)
  
  #define LOG_PRINT(message, ...)         \
      do {                                \
          printf(message, ##__VA_ARGS__); \
      } while(0)
  
  struct Args {
      uint32_t rankId;
      uint32_t epRankId;
      uint32_t tpRankId;
      HcclComm hcclEpComm;
      HcclComm hcclTpComm;
      aclrtStream dispatchStream;
      aclrtStream combineStream;
      aclrtContext context;
  };
  
  constexpr uint32_t EP_WORLD_SIZE = 8;
  constexpr uint32_t TP_WORLD_SIZE = 1;
  constexpr uint32_t DEV_NUM = EP_WORLD_SIZE * TP_WORLD_SIZE;
  
  int64_t GetShapeSize(const std::vector<int64_t> &shape)
  {
      int64_t shape_size = 1;
      for (auto i : shape) {
          shape_size *= i;
      }
      return shape_size;
  }
  
  template<typename T>
  int CreateAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr,
      aclDataType dataType, aclTensor **tensor)
  {
      auto size = GetShapeSize(shape) * sizeof(T);
      auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc failed. ret: %d\n", ret); return ret);
      ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMemcpy failed. ret: %d\n", ret); return ret);
      std::vector<int64_t> strides(shape.size(), 1);
      for (int64_t i = shape.size() - 2; i >= 0; i--) {
          strides[i] = shape[i +1] * strides[i + 1];
      }
      *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
          shape.data(), shape.size(), *deviceAddr);
      return 0;
  }
  
  int LaunchOneProcessDispatchAndCombine(Args &args)
  {
      int ret = aclrtSetCurrentContext(args.context);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSetCurrentContext failed, ret %d\n", ret); return ret);
  
      char hcomEpName[128] = {0};
      ret = HcclGetCommName(args.hcclEpComm, hcomEpName);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetEpCommName failed, ret %d\n", ret); return -1);
      char hcomTpName[128] = {0};
      ret = HcclGetCommName(args.hcclTpComm, hcomTpName);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetTpCommName failed, ret %d\n", ret); return -1);
      LOG_PRINT("[INFO] rank = %d, hcomEpName = %s, hcomTpName = %s, dispatchStream = %p, combineStream = %p, \
                  context = %p\n", args.rankId, hcomEpName, hcomTpName, args.dispatchStream, args.combineStream,                 \
                  args.context);
  
      int64_t BS = 8;
      int64_t H = 7168;
      int64_t K = 3;
      int64_t expertShardType = 0;
      int64_t sharedExpertNum = 1;
      int64_t sharedExpertRankNum = 1;
      int64_t moeExpertNum = 7;
      int64_t quantMode = 0;
      int64_t globalBS = BS * EP_WORLD_SIZE;
      int64_t expertTokenNumsType = 1;
      int64_t outDtype = 0;
      int64_t commQuantMode = 0;
      int64_t groupList_type = 1;
      int64_t localExpertNum;
      int64_t A;
      int64_t zeroExpertNum = 1;
      int64_t copyExpertNum = 1;
      int64_t constExpertNum = 1;
      if (args.epRankId < sharedExpertRankNum) {
          localExpertNum = 1;
          A = globalBS / sharedExpertRankNum;
      } else {
          localExpertNum = moeExpertNum / (EP_WORLD_SIZE - sharedExpertRankNum);
          A = globalBS * (localExpertNum < K ? localExpertNum : K);
      }
  
      void *xDeviceAddr = nullptr;
      void *expertIdsDeviceAddr = nullptr;
      void *scalesDeviceAddr = nullptr;
      void *expertScalesDeviceAddr = nullptr;
  
      void *expandXDeviceAddr = nullptr;
      void *dynamicScalesDeviceAddr = nullptr;
      void *expandIdxDeviceAddr = nullptr;
      void *expertTokenNumsDeviceAddr = nullptr;
      void *epRecvCountsDeviceAddr = nullptr;
      void *tpRecvCountsDeviceAddr = nullptr;
      void *expandScalesDeviceAddr = nullptr;
      void *residualXDeviceAddr = nullptr;
      void *sharedExpertXDeviceAddr = nullptr;
  
      void *elasticInfoDeviceAddr = nullptr;
      void *oriXDeviceAddr = nullptr;
      void *constExpertAlpha1DeviceAddr = nullptr;
      void *constExpertAlpha2DeviceAddr = nullptr;
      void *constExpertVDeviceAddr = nullptr;
  
      void *xOutDeviceAddr = nullptr;
  
      aclTensor *x = nullptr;
      aclTensor *expertIds = nullptr;
      aclTensor *scales = nullptr;
      aclTensor *expertScales = nullptr;
  
      aclTensor *expandX = nullptr;
      aclTensor *dynamicScales = nullptr;
      aclTensor *expandIdx = nullptr;
      aclTensor *expertTokenNums = nullptr;
      aclTensor *epRecvCounts = nullptr;
      aclTensor *tpRecvCounts = nullptr;
      aclTensor *expandScales = nullptr;
      aclTensor *residualX = nullptr;
      aclTensor *sharedExpertX = nullptr;
  
  
      aclTensor *elasticInfo = nullptr;
      aclTensor *oriX = nullptr;
      aclTensor *constExpertAlpha1 = nullptr;
      aclTensor *constExpertAlpha2 = nullptr;
      aclTensor *constExpertV = nullptr;
  
      aclTensor *xOut = nullptr;
  
      //Define the dimensions of variables in the current scenario.
      std::vector<int64_t> xShape{BS, H};
      std::vector<int64_t> expertIdsShape{BS, K};
      std::vector<int64_t> scalesShape{moeExpertNum + 1, H};
      std::vector<int64_t> expertScalesShape{BS, K};
  
      std::vector<int64_t> expandXShape{TP_WORLD_SIZE * A, H};
      std::vector<int64_t> dynamicScalesShape{TP_WORLD_SIZE * A};
      std::vector<int64_t> expandIdxShape{A * 128};
      std::vector<int64_t> expertTokenNumsShape{localExpertNum};
      std::vector<int64_t> epRecvCountsShape{TP_WORLD_SIZE * localExpertNum * EP_WORLD_SIZE};
      std::vector<int64_t> tpRecvCountsShape{TP_WORLD_SIZE};
      std::vector<int64_t> expandScalesShape{A};
      std::vector<int64_t> sharedExpertXShape{BS, 1, H};
  
  
      std::vector<int64_t> elasticInfoShape{4 + EP_WORLD_SIZE * 2};
      std::vector<int64_t> oriXShape{BS, H};
      std::vector<int64_t> constExpertAlpha1Shape{constExpertNum, H};
      std::vector<int64_t> constExpertAlpha2Shape{constExpertNum, H};
      std::vector<int64_t> constExpertVShape{constExpertNum, H};
  
      std::vector<int64_t> xOutShape{BS, H};
  
      int64_t xShapeSize = GetShapeSize(xShape);
      int64_t expertIdsShapeSize = GetShapeSize(expertIdsShape);
      int64_t scalesShapeSize = GetShapeSize(scalesShape);
      int64_t expertScalesShapeSize = GetShapeSize(expertScalesShape);
  
      int64_t expandXShapeSize = GetShapeSize(expandXShape);
      int64_t dynamicScalesShapeSize = GetShapeSize(dynamicScalesShape);
      int64_t expandIdxShapeSize = GetShapeSize(expandIdxShape);
      int64_t expertTokenNumsShapeSize = GetShapeSize(expertTokenNumsShape);
      int64_t epRecvCountsShapeSize = GetShapeSize(epRecvCountsShape);
      int64_t tpRecvCountsShapeSize = GetShapeSize(tpRecvCountsShape);
      int64_t expandScalesShapeSize = GetShapeSize(expandScalesShape);
      int64_t sharedExpertXShapeSize = GetShapeSize(sharedExpertXShape);
  
      int64_t elasticInfoSize = GetShapeSize(elasticInfoShape);
      int64_t oriXSize = GetShapeSize(oriXShape);
      int64_t constExpertAlpha1Size = GetShapeSize(constExpertAlpha1Shape);
      int64_t constExpertAlpha2Size = GetShapeSize(constExpertAlpha2Shape);
      int64_t constExpertVSize = GetShapeSize(constExpertVShape);
  
      int64_t xOutShapeSize = GetShapeSize(xOutShape);
  
      std::vector<int16_t> xHostData(xShapeSize, 1);
      std::vector<int32_t> expertIdsHostData;
      for (int32_t token_id = 0; token_id < expertIdsShape[0]; token_id++) {
          for (int32_t k_id = 0; k_id < expertIdsShape[1]; k_id++) {
              expertIdsHostData.push_back(k_id);
          }
      }
  
      std::vector<float> scalesHostData(scalesShapeSize, 0.1);
      std::vector<float> expertScalesHostData(expertScalesShapeSize, 0.1);
  
      std::vector<int16_t> expandXHostData(expandXShapeSize, 0);
      std::vector<float> dynamicScalesHostData(dynamicScalesShapeSize, 0);
      std::vector<int32_t> expandIdxHostData(expandIdxShapeSize, 0);
      std::vector<int64_t> expertTokenNumsHostData(expertTokenNumsShapeSize, 0);
      std::vector<int32_t> epRecvCountsHostData(epRecvCountsShapeSize, 0);
      std::vector<int32_t> tpRecvCountsHostData(tpRecvCountsShapeSize, 0);
      std::vector<float> expandScalesHostData(expandScalesShapeSize, 0);
      std::vector<int16_t> sharedExpertXHostData(sharedExpertXShapeSize, 1);
  
      int32_t isElastic = 1;
      int32_t rankNumAfterElastic = 4;
      int32_t sharedExpertRankNumAfterElastic = sharedExpertRankNum;
      int32_t moeExpertNumAfterElastic = rankNumAfterElastic - sharedExpertRankNumAfterElastic;
      std::unordered_set<int16_t> availableRank{
          0, 1, /*2, 3, 4, 5,*/ 6, 7
      };
      std::vector<int32_t> elasticInfoHostData{
          isElastic, rankNumAfterElastic, sharedExpertRankNumAfterElastic, moeExpertNumAfterElastic,
          0, 1, -1, -1, -1, -1, 2, 3,
          0, 1, 6, 7, -1, -1, -1, -1
      };
      std::vector<int16_t> oriXHostData(oriXSize, 1);
      std::vector<int16_t> constExpertAlpha1HostData(constExpertAlpha1Size, 0);
      std::vector<int16_t> constExpertAlpha2HostData(constExpertAlpha2Size, 0);
      std::vector<int16_t> constExpertVHostData(constExpertVSize, 0);
  
      std::vector<int16_t> xOutHostData(xOutShapeSize, 0);
  
  
      ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_BF16, &x);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(expertIdsHostData, expertIdsShape, &expertIdsDeviceAddr, aclDataType::ACL_INT32, &expertIds);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(scalesHostData, scalesShape, &scalesDeviceAddr, aclDataType::ACL_FLOAT, &scales);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(expertScalesHostData, expertScalesShape, &expertScalesDeviceAddr, aclDataType::ACL_FLOAT, &expertScales);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
  
      ret = CreateAclTensor(expandXHostData, expandXShape, &expandXDeviceAddr, (quantMode > 0) ? aclDataType::ACL_INT8 : aclDataType::ACL_BF16, &expandX);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(dynamicScalesHostData, dynamicScalesShape, &dynamicScalesDeviceAddr, aclDataType::ACL_FLOAT, &dynamicScales);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
          ret = CreateAclTensor(expandIdxHostData, expandIdxShape, &expandIdxDeviceAddr, aclDataType::ACL_INT32, &expandIdx);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(expertTokenNumsHostData, expertTokenNumsShape, &expertTokenNumsDeviceAddr, aclDataType::ACL_INT64, &expertTokenNums);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(epRecvCountsHostData, epRecvCountsShape, &epRecvCountsDeviceAddr, aclDataType::ACL_INT32, &epRecvCounts);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(tpRecvCountsHostData, tpRecvCountsShape, &tpRecvCountsDeviceAddr, aclDataType::ACL_INT32, &tpRecvCounts);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(expandScalesHostData, expandScalesShape, &expandScalesDeviceAddr, aclDataType::ACL_FLOAT, &expandScales);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(sharedExpertXHostData, sharedExpertXShape, &sharedExpertXDeviceAddr, aclDataType::ACL_BF16, &sharedExpertX);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
  
      ret = CreateAclTensor(elasticInfoHostData, elasticInfoShape, &elasticInfoDeviceAddr, aclDataType::ACL_INT32, &elasticInfo);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(oriXHostData, oriXShape, &oriXDeviceAddr, aclDataType::ACL_BF16, &oriX);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(constExpertAlpha1HostData, constExpertAlpha1Shape, &constExpertAlpha1DeviceAddr, aclDataType::ACL_BF16, &constExpertAlpha1);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(constExpertAlpha2HostData, constExpertAlpha2Shape, &constExpertAlpha2DeviceAddr, aclDataType::ACL_BF16, &constExpertAlpha2);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(constExpertVHostData, constExpertVShape, &constExpertVDeviceAddr, aclDataType::ACL_BF16, &constExpertV);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
  
      ret = CreateAclTensor(xOutHostData, xOutShape, &xOutDeviceAddr, aclDataType::ACL_BF16, &xOut);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
  
  
  
      uint64_t dispatchWorkspaceSize = 0;
      aclOpExecutor *dispatchExecutor = nullptr;
      void *dispatchWorkspaceAddr = nullptr;
  
      uint64_t combineWorkspaceSize = 0;
      aclOpExecutor *combineExecutor = nullptr;
      void *combineWorkspaceAddr = nullptr;
      /**************************************** Call dispatch warm up ********************************************/
      ret = aclnnMoeDistributeDispatchV4GetWorkspaceSize(x, expertIds, (quantMode > 0 ? scales : nullptr), nullptr,
              expertScales, nullptr, nullptr, // performanceInfoOptional
              hcomEpName, EP_WORLD_SIZE, args.epRankId, moeExpertNum, hcomTpName, TP_WORLD_SIZE,
              args.tpRankId, expertShardType, sharedExpertNum,sharedExpertRankNum, quantMode, globalBS,
              expertTokenNumsType, nullptr, zeroExpertNum, copyExpertNum, constExpertNum, expandX, dynamicScales, expandIdx, expertTokenNums, epRecvCounts,
              tpRecvCounts, expandScales, &dispatchWorkspaceSize, &dispatchExecutor);
  
      CHECK_RET(ret == ACL_SUCCESS,
          LOG_PRINT("[ERROR] warm up aclnnMoeDistributeDispatchV4GetWorkspaceSize failed. ret = %d \n", ret); return ret);
  
      if (dispatchWorkspaceSize > 0) {
          ret = aclrtMalloc(&dispatchWorkspaceAddr, dispatchWorkspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
          CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] warm up aclrtMalloc workspace failed. ret = %d \n", ret); return ret);
      }
      // Call the second-phase API.
      ret = aclnnMoeDistributeDispatchV4(dispatchWorkspaceAddr, dispatchWorkspaceSize,
                                          dispatchExecutor, args.dispatchStream);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] warm up aclnnMoeDistributeDispatchV4 failed. ret = %d \n", ret);  \
              return ret);
      ret = aclrtSynchronizeStreamWithTimeout(args.dispatchStream, 10000);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] warm up aclrtSynchronizeStreamWithTimeout failed. ret = %d \n", ret);  \
          return ret);
  
      /**************************************** Call dispatch ********************************************/
      if (availableRank.find(args.rankId) != availableRank.end()) {
          // Call the first-phase API.
      ret = aclnnMoeDistributeDispatchV4GetWorkspaceSize(x, expertIds, (quantMode > 0 ? scales : nullptr), nullptr,
              expertScales, elasticInfo, nullptr, // performanceInfoOptional
              hcomEpName, EP_WORLD_SIZE, args.epRankId, moeExpertNum, hcomTpName, TP_WORLD_SIZE,
              args.tpRankId, expertShardType, sharedExpertNum,sharedExpertRankNum, quantMode, globalBS,
              expertTokenNumsType, nullptr, zeroExpertNum, copyExpertNum, constExpertNum, expandX, dynamicScales, expandIdx, expertTokenNums, epRecvCounts,
              tpRecvCounts, expandScales, &dispatchWorkspaceSize, &dispatchExecutor);
  
      CHECK_RET(ret == ACL_SUCCESS,
          LOG_PRINT("[ERROR] aclnnMoeDistributeDispatchV4GetWorkspaceSize 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 = aclnnMoeDistributeDispatchV4(dispatchWorkspaceAddr, dispatchWorkspaceSize,
                                          dispatchExecutor, args.dispatchStream);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMoeDistributeDispatchV4 failed. ret = %d \n", ret);  \
              return ret);
      ret = aclrtSynchronizeStreamWithTimeout(args.dispatchStream, 10000);
                  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] dispatch aclrtSynchronizeStreamWithTimeout failed. ret = %d \n", ret);  \
              return ret);
      }
      /**************************************** Call combine ********************************************/
      // Call the first-phase API.
      if (availableRank.find(args.rankId) != availableRank.end()) {
      ret = aclnnMoeDistributeCombineV4GetWorkspaceSize(expandX, expertIds,
                                                          expandIdx, epRecvCounts,
                                                          expertScales, tpRecvCounts,
                                                          nullptr, nullptr, nullptr,
                                                          nullptr, nullptr, nullptr,
                                                          elasticInfo, oriX, constExpertAlpha1, constExpertAlpha2, constExpertV,
                                                          nullptr, // performanceInfoOptional
                                                          hcomEpName, EP_WORLD_SIZE, args.epRankId, moeExpertNum,
                                                          hcomTpName, TP_WORLD_SIZE, args.tpRankId, expertShardType,
                                                          sharedExpertNum, sharedExpertRankNum, globalBS, outDtype,
                                                          commQuantMode, groupList_type, nullptr, zeroExpertNum, copyExpertNum, constExpertNum, xOut,
                                                          &combineWorkspaceSize, &combineExecutor);
      CHECK_RET(ret == ACL_SUCCESS,
          LOG_PRINT("[ERROR] aclnnMoeDistributeCombineV4GetWorkspaceSize 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 = aclnnMoeDistributeCombineV4(combineWorkspaceAddr, combineWorkspaceSize, combineExecutor, args.combineStream);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnMoeDistributeCombineV4 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 aclnnMoeDistributeDispatchV4 and aclnnMoeDistributeCombineV4                      \
                  execute successfully.\n", args.rankId);
      }
      // Release device resources.
      if (dispatchWorkspaceSize > 0) {
          aclrtFree(dispatchWorkspaceAddr);
      }
      if (combineWorkspaceSize > 0) {
          aclrtFree(combineWorkspaceAddr);
      }
      if (x != nullptr) {
          aclDestroyTensor(x);
      }
      if (expertIds != nullptr) {
          aclDestroyTensor(expertIds);
      }
      if (scales != nullptr) {
          aclDestroyTensor(scales);
      }
      if (expertScales != nullptr) {
          aclDestroyTensor(expertScales);
      }
  
      if (expandX != nullptr) {
          aclDestroyTensor(expandX);
      }
      if (dynamicScales != nullptr) {
          aclDestroyTensor(dynamicScales);
      }
      if (expandIdx != nullptr) {
          aclDestroyTensor(expandIdx);
      }
      if (expertTokenNums != nullptr) {
          aclDestroyTensor(expertTokenNums);
      }
      if (epRecvCounts != nullptr) {
          aclDestroyTensor(epRecvCounts);
      }
      if (tpRecvCounts != nullptr) {
          aclDestroyTensor(tpRecvCounts);
      }
      if (expandScales != nullptr) {
          aclDestroyTensor(expandScales);
      }
      if (residualX != nullptr) {
          aclDestroyTensor(residualX);
      }
      if (sharedExpertX != nullptr) {
          aclDestroyTensor(sharedExpertX);
      }
      if (elasticInfo != nullptr) {
          aclDestroyTensor(elasticInfo);
      }
      if (oriX != nullptr) {
          aclDestroyTensor(oriX);
      }
      if (constExpertAlpha1 != nullptr) {
          aclDestroyTensor(constExpertAlpha1);
      }
      if (constExpertAlpha2 != nullptr) {
          aclDestroyTensor(constExpertAlpha2);
      }
      if (constExpertV != nullptr) {
          aclDestroyTensor(constExpertV);
      }
  
      if (xOut != nullptr) {
          aclDestroyTensor(xOut);
      }
      if (xDeviceAddr != nullptr) {
          aclrtFree(xDeviceAddr);
      }
      if (expertIdsDeviceAddr != nullptr) {
          aclrtFree(expertIdsDeviceAddr);
      }
      if (scalesDeviceAddr != nullptr) {
          aclrtFree(scalesDeviceAddr);
      }
      if (expertScalesDeviceAddr != nullptr) {
          aclrtFree(expertScalesDeviceAddr);
      }
      if (expandXDeviceAddr != nullptr) {
          aclrtFree(expandXDeviceAddr);
      }
      if (dynamicScalesDeviceAddr != nullptr) {
          aclrtFree(dynamicScalesDeviceAddr);
      }
      if (expandIdxDeviceAddr != nullptr) {
          aclrtFree(expandIdxDeviceAddr);
      }
      if (expertTokenNumsDeviceAddr != nullptr) {
          aclrtFree(expertTokenNumsDeviceAddr);
      }
      if (epRecvCountsDeviceAddr != nullptr) {
          aclrtFree(epRecvCountsDeviceAddr);
      }
      if (expandScalesDeviceAddr != nullptr) {
          aclrtFree(expandScalesDeviceAddr);
      }
      if (tpRecvCountsDeviceAddr != nullptr) {
          aclrtFree(tpRecvCountsDeviceAddr);
      }
      if (sharedExpertXDeviceAddr != nullptr) {
          aclrtFree(sharedExpertXDeviceAddr);
      }
  
      if (elasticInfoDeviceAddr != nullptr) {
          aclrtFree(elasticInfoDeviceAddr);
      }
      if (oriXDeviceAddr != nullptr) {
          aclrtFree(oriXDeviceAddr);
      }
      if (constExpertAlpha1DeviceAddr != nullptr) {
          aclrtFree(constExpertAlpha1DeviceAddr);
      }
      if (constExpertAlpha2DeviceAddr != nullptr) {
          aclrtFree(constExpertAlpha2DeviceAddr);
      }
      if (constExpertVDeviceAddr != nullptr) {
          aclrtFree(constExpertVDeviceAddr);
      }
  
      if (xOutDeviceAddr != nullptr) {
          aclrtFree(xOutDeviceAddr);
      }
  
      HcclCommDestroy(args.hcclEpComm);
      HcclCommDestroy(args.hcclTpComm);
      aclrtDestroyStream(args.dispatchStream);
      aclrtDestroyStream(args.combineStream);
      aclrtDestroyContext(args.context);
      aclrtResetDevice(args.rankId);
  
      return 0;
  }
  
  int main(int argc, char *argv[])
  {
      int ret = aclInit(nullptr);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclInit failed, ret = %d\n", ret); return ret);
  
      aclrtStream dispatchStream[DEV_NUM];
      aclrtStream combineStream[DEV_NUM];
      aclrtContext context[DEV_NUM];
      for (uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
          ret = aclrtSetDevice(rankId);
          CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSetDevice failed, ret = %d\n", ret); return ret);
          ret = aclrtCreateContext(&context[rankId], rankId);
          CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateContext failed, ret = %d\n", ret); return ret);
          ret = aclrtCreateStream(&dispatchStream[rankId]);
          CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateStream failed, ret = %d\n", ret); return ret);
          ret = aclrtCreateStream(&combineStream[rankId]);
          CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateStream failed, ret = %d\n", ret); return ret);
      }
  
      int32_t devicesEp[TP_WORLD_SIZE][EP_WORLD_SIZE];
      for (int32_t tpId = 0; tpId < TP_WORLD_SIZE; tpId++) {
          for (int32_t epId = 0; epId < EP_WORLD_SIZE; epId++) {
              devicesEp[tpId][epId] = epId * TP_WORLD_SIZE + tpId;
          }
      }
  
      HcclComm commsEp[TP_WORLD_SIZE][EP_WORLD_SIZE];
      for (int32_t tpId = 0; tpId < TP_WORLD_SIZE; tpId++) {
          ret = HcclCommInitAll(EP_WORLD_SIZE, devicesEp[tpId], commsEp[tpId]);
          CHECK_RET(ret == ACL_SUCCESS,
                      LOG_PRINT("[ERROR] HcclCommInitAll ep %d failed, ret %d\n", tpId, ret); return ret);
      }
  
      int32_t devicesTp[EP_WORLD_SIZE][TP_WORLD_SIZE];
      for (int32_t epId = 0; epId < EP_WORLD_SIZE; epId++) {
          for (int32_t tpId = 0; tpId < TP_WORLD_SIZE; tpId++) {
              devicesTp[epId][tpId] = epId * TP_WORLD_SIZE + tpId;
          }
      }
  
      HcclComm commsTp[EP_WORLD_SIZE][TP_WORLD_SIZE];
      for (int32_t epId = 0; epId < EP_WORLD_SIZE; epId++) {
          ret = HcclCommInitAll(TP_WORLD_SIZE, devicesTp[epId], commsTp[epId]);
          CHECK_RET(ret == ACL_SUCCESS,
                      LOG_PRINT("[ERROR] HcclCommInitAll tp %d failed, ret %d\n", epId, ret); return ret);
      }
  
      Args args[DEV_NUM];
      std::vector<std::unique_ptr<std::thread>> threads(DEV_NUM);
      for (uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
          uint32_t epRankId = rankId / TP_WORLD_SIZE;
          uint32_t tpRankId = rankId % TP_WORLD_SIZE;
  
          args[rankId].rankId = rankId;
          args[rankId].epRankId = epRankId;
          args[rankId].tpRankId = tpRankId;
          args[rankId].hcclEpComm = commsEp[tpRankId][epRankId];
          args[rankId].hcclTpComm = commsTp[epRankId][tpRankId];
          args[rankId].dispatchStream = dispatchStream[rankId];
          args[rankId].combineStream = combineStream[rankId];
          args[rankId].context = context[rankId];
          threads[rankId].reset(new(std::nothrow) std::thread(&LaunchOneProcessDispatchAndCombine, std::ref(args[rankId])));
      }
  
      for(uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
          threads[rankId]->join();
      }
  
      aclFinalize();
      LOG_PRINT("[INFO] aclFinalize success\n");
  
      return 0;
  }