aclnnMoeFinalizeRoutingV2

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: In MoE computation, it finalizes and merges the FFN output results of MoE.
Formula:
$expertid=expertIdx[i,k]$ $out(i,j)=x1_{i,j}+x2_{i,j}+\sum_{k=0}^{K}(scales_{i,k}*(expandedX_{expandedRowIdx_{i+k*num\_rows},j}+bias_{expertid,j}))$

Function Prototype

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

aclnnStatus aclnnMoeFinalizeRoutingV2GetWorkspaceSize(
    const aclTensor *expandedX,
    const aclTensor *expandedRowIdx,
    const aclTensor *x1Optional,
    const aclTensor *x2Optional,
    const aclTensor *biasOptional,
    const aclTensor *scalesOptional,
    const aclTensor *expertIdxOptional,
    int64_t          dropPadMode,
    const aclTensor *out,
    uint64_t        *workspaceSize,
    aclOpExecutor  **executor)

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

aclnnMoeFinalizeRoutingV2GetWorkspaceSize

Parameters:

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
expandedX	Input	expandedX in the formula, the FFN output of MoE.	-	FLOAT16, BFLOAT16, FLOAT32	ND	drop less scenario: (NUM_ROWS * K, H), drop pad scenario: (E, C, H).	√
expandedRowIdx	Input	expandedRowIdx in the formula.	-	INT32	ND	(NUM_ROWS * K)	√
x1Optional	Input	x1 in the formula, representing the first shared expert.	-	Same as expandedX.	ND	Same as out.	√
x2Optional	Input	x2 in the formula represents the second shared expert.	-	Same as expandedX.	ND	Same as out.	√
biasOptional	Input	bias in the formula, representing the bias value.	-	Same as expandedX.	ND	(E, H)	√
scalesOptional	Input	scales in the formula.	-	FLOAT16, BFLOAT16, FLOAT32	ND	(NUM_ROWS, K)	√
expertIdxOptional	Input	expertIdx in the formula.	The value range in the tensor is [0, E-1].	INT32	ND	(NUM_ROWS, K)	√
dropPadMode	Input	Indicates whether the drop mode is supported, and the permutation of expandedRowIdx.	The value range is [0, 3].	-	-	-	-
out	Output	Output in the formula.	-	Same as expandedX.	ND	(NUM_ROWS, H)	×
workspaceSize	Output	Returns the workspace size to be allocated on the Device side.	-	-	-	-	-
executor	Output	Returns the operator executor, which contains the operator computation flow.	-	-	-	-	-

Atlas A2 training products/Atlas A2 inference products, Atlas A3 training products/Atlas A3 inference products:
- expandedX must be a 2D/3D tensor. Supported data types are FLOAT16, BFLOAT16, and FLOAT32. Both drop-less and drop-pad scenarios are supported.
- scalesOptional: In mixed precision mode, when expandedX is BFLOAT16, scalesOptional must be FLOAT32. In non-mixed precision mode, the data type must be the same as expandedX.
Ascend 950PR/Ascend 950DT:
- expandedX must be a 2D or 3D tensor. Supported data types are FLOAT16, BFLOAT16, and FLOAT32. Both drop-less and drop-pad scenarios are supported.
- The data type of scalesOptional can differ from that of expandedX.
Atlas inference products:
- expandedX must be a 2D tensor. Supported data types are FLOAT16 and FLOAT32. The last axis H must be 32-aligned.
- x1Optional, x2Optional, biasOptional, and expertIdxOptional only support passing nullptr.
- Only dropPadMode set to 2 is supported.
- The data type of scalesOptional supports FLOAT16 and FLOAT32, and must be consistent with expandedX.

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, required output, or required attribute passed in is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data types and data formats of the inputs and outputs are not within the supported range.
ACLNN_ERR_INNER_TILING_ERROR	561002	The shape information of multiple input tensors does not match.
ACLNN_ERR_INNER_TILING_ERROR	561002	The shape information between the input attributes and the input tensors does not match.

aclnnMoeFinalizeRoutingV2

Parameters:

Parameter	Input/Output	Description
workspace	Input	Address of the workspace memory applied for on the Device side.
workspaceSize	Input	Size of the workspace applied for on the Device side, obtained from the first-phase API `aclnnMoeFinalizeRoutingV2GetWorkspaceSize`.
executor	Input	Operator executor, which contains the operator computation flow.
stream	Input	Specifies the stream on which the task is executed.

Return Value

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

Constraints

Deterministic computation:
- aclnnMoeFinalizeRoutingV2 defaults to a deterministic implementation.
NUM_ROWS: Indicates the number of rows;
- K: Indicates that K experts are selected from the total E experts;
- H: Indicates the hidden size, i.e., the sequence length of each token, which is the number of columns.
- E: Indicates the expert num, i.e., the number of experts. E must be greater than or equal to K.
- C: Indicates the expert capacity, i.e., the capacity threshold for the number of tokens an expert can process.
expandedRowIdx: When the dropPadMode parameter value is 0 or 2, the value range of the tensor is [0, NUM_ROWS * K - 1]; when the dropPadMode parameter value is 1 or 3, the value range of the tensor is [-1, E * C - 1].
When the x1Optional parameter is not input, the x2Optional parameter cannot be input either.
When scalesOptional does not exist, K is 1.
When biasOptional exists, expertIdxOptional must also exist.
The values and meanings of dropPadMode are as follows:
- 0: drop less scenario, where expandedRowIdx is arranged by column (corresponding to the output format of aclnnMoeInitRouting).
- 1: drop pad scenario, where expandedRowIdx is arranged by column (corresponding to the output format of aclnnMoeInitRouting).
- 2: drop less scenario, where expandedRowIdx is arranged by row (corresponding to the output format of aclnnMoeInitRoutingV2).
- 3: drop pad scenario, where expandedRowIdx is arranged by row (corresponding to the output format of aclnnMoeInitRoutingV2).

Example

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

#include "acl/acl.h"
#include "aclnnop/aclnn_moe_finalize_routing_v2.h"
#include <iostream>
#include <vector>

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

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

int64_t GetShapeSize(const std::vector<int64_t>& shape) {
    int64_t shape_size = 1;
    for (auto i : shape) {
        shape_size *= i;
    }
    return shape_size;
}
int Init(int32_t deviceId, aclrtStream* stream) {
    // Boilerplate for resource initialization.
    auto  ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
    ret = aclrtSetDevice(deviceId);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
    ret = aclrtCreateStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
    return 0;
}

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);
    // Call aclrtMalloc to allocate memory on the device side.
    auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
    // Call aclrtMemcpy to copy data from the host side to the device memory.
    ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
    // Calculate the strides of the contiguous tensor.
    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];
    }
    // Call the aclCreateTensor API to create an aclTensor.
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                              shape.data(), shape.size(), *deviceAddr);
    return 0;
}

int main() {
  // 1. (Boilerplate) Initialize the device/stream. See the ACL API manual.
  // Fill in the deviceId based on the actual device.
  int32_t deviceId = 0;
  aclrtStream stream;
  auto ret = Init(deviceId, &stream);
  // Handle the check based on your requirements.
  CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

  // 2. Construct inputs and outputs. Customize the construction based on the API.
  std::vector<int64_t> expandedXShape = {3 * 2, 4};
  std::vector<int64_t> x1Shape = {3, 4};
  std::vector<int64_t> x2OptionalShape = {3, 4};
  std::vector<int64_t> biasShape = {2, 4};
  std::vector<int64_t> scalesShape = {3, 2};
  std::vector<int64_t> expandedExpertIdxShape = {3, 2};
  std::vector<int64_t> expandedRowIdxShape = {3 * 2};
  std::vector<int64_t> outShape = {3, 4};
  void* expandedXAddr = nullptr;
  void* x1Addr = nullptr;
  void* x2OptionalAddr = nullptr;
  void* biasAddr = nullptr;
  void* scalesDeviceAddr = nullptr;
  void* expandedExpertIdxAddr = nullptr;
  void* expandedRowIdxAddr = nullptr;
  void* outDeviceAddr = nullptr;

  aclTensor* expandedX = nullptr;
  aclTensor* x1 = nullptr;
  aclTensor* x2Optional = nullptr;
  aclTensor* bias = nullptr;
  aclTensor* scales = nullptr;
  aclTensor* expandedExpertIdx = nullptr;
  aclTensor* expandedRowIdx = nullptr;
  aclTensor* out = nullptr;
  std::vector<float> expandedXHostData = {0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1, 8.1, 9.1, 10.1, 11.1,
                                                     0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1, 8.1, 9.1, 10.1, 11.1};
  std::vector<float> x1HostData = {0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2, 8.2, 9.2, 10.2, 11.2};
  std::vector<float> x2OptionalHostData = {0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2, 8.2, 9.2, 10.2, 11.2};
  std::vector<float> biasHostData = {0.2, 0.4, 0.2, 0.4, 0.2, 0.4, 0.2, 0.4};
  std::vector<float> scalesHostData = {1.3, 1.6, 1.2, 1.8, 1.2, 2.3};
  std::vector<int32_t> expandedExpertIdxHostData = {0, 1, 0, 1, 0, 1};
  std::vector<int32_t> expandedRowIdxHostData = {2, 1, 4, 3, 0, 5};
  std::vector<float> outHostData(12, 0.0f);
  int64_t dropPadMode = 0;
  // Create the expandedX aclTensor.
  ret = CreateAclTensor(expandedXHostData, expandedXShape, &expandedXAddr,
                        aclDataType::ACL_FLOAT, &expandedX);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the x1 aclTensor.
  ret = CreateAclTensor(x1HostData, x1Shape, &x1Addr, aclDataType::ACL_FLOAT, &x1);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the x2Optional aclTensor.
  ret = CreateAclTensor(x2OptionalHostData, x2OptionalShape, &x2OptionalAddr, aclDataType::ACL_FLOAT, &x2Optional);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the bias aclTensor.
  ret = CreateAclTensor(biasHostData, biasShape, &biasAddr, aclDataType::ACL_FLOAT, &bias);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the totalWeightOut aclTensor.
  ret = CreateAclTensor(scalesHostData, scalesShape, &scalesDeviceAddr, aclDataType::ACL_FLOAT, &scales);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  // Create the expandedExpertIdx aclTensor.
  ret = CreateAclTensor(expandedExpertIdxHostData, expandedExpertIdxShape, &expandedExpertIdxAddr,
                        aclDataType::ACL_INT32, &expandedExpertIdx);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the expandedRowIdx aclTensor.
  ret = CreateAclTensor(expandedRowIdxHostData, expandedRowIdxShape, &expandedRowIdxAddr,
                        aclDataType::ACL_INT32, &expandedRowIdx);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the Out aclTensor.
  ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  // 3. Call the CANN operator library API. Replace this with the specific operator API.
  uint64_t workspaceSize = 0;
  aclOpExecutor* executor;

  // Call the first-phase API of aclnnMoeFinalizeRoutingV2.
  ret = aclnnMoeFinalizeRoutingV2GetWorkspaceSize(expandedX, expandedRowIdx, x1, x2Optional, bias, scales,
                                                  expandedExpertIdx, dropPadMode, out, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMoeFinalizeRoutingV2GetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  // Apply for device memory based on the workspaceSize calculated by the first-phase API.
  void* workspaceAddr = nullptr;
  if (workspaceSize > 0) {
      ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
  }
  // Call the second-phase API of aclnnMoeFinalizeRoutingV2.
  ret = aclnnMoeFinalizeRoutingV2(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMoeFinalizeRoutingV2 failed. ERROR: %d\n", ret); return ret);

  // 4. (Boilerplate) Synchronously wait for task execution to complete.
  ret = aclrtSynchronizeStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

  // 5. Obtain the output value and copy the result from the device memory to the host side. Modify based on the specific API definition.
  auto size = GetShapeSize(outShape);
  std::vector<float> resultData(size, 0.0f);
  ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]),
                    outDeviceAddr, size * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
  for (int64_t i = 0; i < size; i++) {
      LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]);
  }

  // 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
  aclDestroyTensor(expandedX);
  aclDestroyTensor(x1);
  aclDestroyTensor(x2Optional);
  aclDestroyTensor(bias);
  aclDestroyTensor(scales);
  aclDestroyTensor(expandedExpertIdx);
  aclDestroyTensor(expandedRowIdx);
  aclDestroyTensor(out);

  // 7. Release device resources. Modify based on the specific API definition.
  aclrtFree(expandedXAddr);
  aclrtFree(x1Addr);
  aclrtFree(x2OptionalAddr);
  aclrtFree(biasAddr);
  aclrtFree(scalesDeviceAddr);
  aclrtFree(expandedExpertIdxAddr);
  aclrtFree(expandedRowIdxAddr);
  aclrtFree(outDeviceAddr);
  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();
  return 0;
}