aclnnMoeGatingTopKSoftmax

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, performs Softmax calculation on the output of x and then takes the topK operation. yOut is the topK result of softmax; expertIdxOut is the index result of the topK values, i.e., the corresponding expert number; rowIdxOut is the column-wise result with the same shape as expertIdxOut, indicating the original row position corresponding to each position. If the finished flag of the corresponding row in expertIdxOut is True, the value is directly filled with the num_expert value (i.e., the size of the last axis of x).
Formula:
$softmaxOut=softmax(x,axis=-1)$ $yOut,expertIdxOut=topK(softmaxOut,k=k)$ $rowIdxRange=arange(expertIdxOut.shape[0]*expertIdxOut.shape[1])$ $rowIdxOut=rowIdxRange.reshape([expertIdxOut.shape[1],expertIdxOut.shape[0]]).transpose(1,0)$

Function Prototype

Each operator consists of a two-phase API. You must first call the "aclnnMoeGatingTopKSoftmaxGetWorkspaceSize" API to obtain the workspace size required for computation and the executor that encapsulates the operator computation process, and then call the "aclnnMoeGatingTopKSoftmax" API to perform the computation.

aclnnStatus aclnnMoeGatingTopKSoftmaxGetWorkspaceSize(
    const aclTensor *x,
    const aclTensor *finishedOptional,
    int64_t          k,
    const aclTensor *yOut,
    const aclTensor *expertIdxOut,
    const aclTensor *rowIdxOut,
    uint64_t        *workspaceSize,
    aclOpExecutor  **executor)

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

aclnnMoeGatingTopKSoftmaxGetWorkspaceSize

Parameters:

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
x	Input	x in the formula, the input to be computed.	Must be a 2D or 3D tensor, with each dimension size not exceeding the maximum int32 value of 2147483647.	FLOAT16, BFLOAT16, FLOAT32	ND	2-3	√
finishedOptional	Input	Indicates whether the row participates in the computation, corresponding to "finished" in the formula.	shape is x_shape[:-1].	BOOL	ND	1-2	Yes
k	Input	The k value for topK, indicating the number of experts.	0 <= k <= size of the -1 axis of x, and k is not greater than 1024.	INT64	-	-	-
yOut	Output	TopK values obtained after applying softmax to x.	The non-(-1) axes of its shape must match the corresponding axes of x, and the size of its -1 axis must equal the value of k.	Must be the same as the input x.	ND	2-3	x
expertIdxOut	Output	Indices of the top-K values after applying softmax to x, i.e., the expert indices.	The shape must be the same as that of yOut.	INT32	ND	2-3	x
rowIdxOut	Output	rowIdxOut in the formula.	The shape must be the same as yOut.	INT32	ND	2-3	x
workspaceSize	Output	Returns the workspace size to be requested on the Device side.	-	-	-	-	-
executor	Output	Returns the op executor, which contains the operator computation process.	-	-	-	-	-

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	A required input, required output, or required attribute passed in is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data type or data format of the input or output is not within the supported range.
ACLNN_ERR_INNER_TILING_ERROR	561002	The shape information of multiple input tensors does not match.
		The shape information between the input attributes and the input tensor does not match.
		The value of k is less than 0 or greater than the size of the axis of x-1.
		The value of k is greater than 1024.

aclnnMoeGatingTopKSoftmax

Parameters:

Parameter	Input/Output	Description
workspace	Input	Memory address of the workspace applied on the device side.
workspaceSize	Input	Size of the workspace applied on the device side, obtained by the first-phase API aclnnMoeGatingTopKSoftmaxGetWorkspaceSize.
executor	Input	Operator executor, which contains the operator computation flow.
stream	Input	Specifies the stream for task execution.

Return Value

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

Constraints

Deterministic computation:
- aclnnMoeGatingTopKSoftmax defaults to a deterministic implementation.
Ascend 950PR/Ascend 950DT:
- The size of the last dimension of x (i.e., the number of experts) ranges from [1, 2048].

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_gating_top_k_softmax.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: 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 device memory.
  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 host data 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);
  // Compute 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 aclCreateTensor 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. Refer to the ACL API manual.
  int32_t deviceId = 0;
  aclrtStream stream;
  auto ret = Init(deviceId, &stream);
  // Handle the check as needed.
  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> inputShape = {3, 4};
  std::vector<int64_t> outShape = {3, 2};
  std::vector<int64_t> expertIdOutShape = {3, 2};
  std::vector<int64_t> rowOutShape = {3, 2};

  void* inputAddr = nullptr;
  void* outAddr = nullptr;
  void* expertIdOutAddr = nullptr;
  void* rowOutAddr = nullptr;

  aclTensor* input = nullptr;
  aclTensor* out = nullptr;
  aclTensor* expertIdOut = nullptr;
  aclTensor* rowOut = nullptr;

  std::vector<float> inputHostData = {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> outHostData = {0.1, 1.1, 2.1, 3.1, 4.1, 5.1};
  std::vector<int32_t> expertIdOutHostData = {1, 1, 1, 1, 1, 1};
  std::vector<int32_t> rowOutHostData = {1, 1, 1, 1, 1, 1};

  // Create the expandedPermutedRows aclTensor.
  ret = CreateAclTensor(inputHostData, inputShape, &inputAddr, aclDataType::ACL_FLOAT, &input);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the expertForSourceRow aclTensor.
  ret = CreateAclTensor(outHostData, outShape, &outAddr, aclDataType::ACL_FLOAT, &out);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the expandedSrcToDstRow aclTensor.
  ret = CreateAclTensor(expertIdOutHostData, expertIdOutShape, &expertIdOutAddr, aclDataType::ACL_INT32, &expertIdOut);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the Out aclTensor.
  ret = CreateAclTensor(rowOutHostData, rowOutShape, &rowOutAddr, aclDataType::ACL_INT32, &rowOut);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

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

  // Call the first-phase API of aclnnMoeGatingTopKSoftmax.
  ret = aclnnMoeGatingTopKSoftmaxGetWorkspaceSize(input, nullptr, 2, out, expertIdOut, rowOut, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMoeGatingTopKSoftmaxGetWorkspaceSize 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 aclnnMoeGatingTopKSoftmax.
  ret = aclnnMoeGatingTopKSoftmax(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMoeGatingTopKSoftmax 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. 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]),
                    outAddr, 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(input);
  aclDestroyTensor(out);
  aclDestroyTensor(expertIdOut);
  aclDestroyTensor(rowOut);

  // 7. Release device resources. Modify based on the specific API definition.
  aclrtFree(inputAddr);
  aclrtFree(outAddr);
  aclrtFree(expertIdOutAddr);
  aclrtFree(rowOutAddr);
  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();
  return 0;
}