aclnnMoeGatingTopK

📄 View Source Code

Applicable Products

Function

• Description: In MoE computation, performs Sigmoid or Softmax computation on the input x, performs grouped sorting on the computation results, and finally selects the top k experts based on the grouped sorting results.

• Formula:

  Step 1: Normalization

  Normalizes the input x based on normType:

  $$normOut = \begin{cases} SoftMax(x), & normType = 0 \\ Sigmoid(x), & normType = 1 \end{cases}$$

  Step 2: Add bias

  If bias is not empty, adds the bias to obtain the value used for selection:

  $$normValue = normOut + bias$$

  Otherwise, $normValue = normOut$.

  Step 3: Grouped filtering (executed only when groupCount > 1)

  Group normValue by groupCount, and compute the score of each group based on groupSelectMode:

  $$groupedValue = Reshape(normValue,\ [batch,\ groupCount,\ -1])$$ $$groupScore = \begin{cases} ReduceMax(groupedValue,\ dim=-1), & groupSelectMode = 0 \\ ReduceSum(TopK(groupedValue,\ k=2,\ dim=-1),\ dim=-1), & groupSelectMode = 1 \end{cases}$$

  Select the kGroup groups with the highest scores, and set the corresponding positions of unselected groups to $-\infty$:

  $$groupIdx = TopK(groupScore,\ k=kGroup).indices$$ $$normValue = Mask(groupedValue,\ groupIdx,\ fillValue=-\infty)$$

  Step 4: Top-K expert selection

  Apply Top-K to normValue to obtain the expert indices. Only expertIdxOut is needed here:

  $$y, expertIdxOut = TopK(normValue[groupIdx, :],\ k=k)$$

  Step 5: Renorm and Scaling

  When normType=1, normalization is performed; when normType=0, the renorm parameter takes effect, and renorm is performed when renorm=1:

  $$if\ (normType = 1)\ or\ (normType = 0\ and\ renorm = 1):$$ $$\quad yOut = \frac{normOut}{ReduceSum(normOut,\ dim=-1) + eps}$$

  Final output:

  $$yOut = yOut \times routedScalingFactor$$

  Step 6: Optional output

  If outFlag is True, the third output is normOut; otherwise, it is empty.

Function Prototype

Each operator is divided into a two-phase API. The "aclnnMoeGatingTopKGetWorkspaceSize" API must be called first to obtain the workspace size required for computation and the executor that encapsulates the operator computation flow. Then, the "aclnnMoeGatingTopK" API is called to perform the computation.

aclnnStatus aclnnMoeGatingTopKGetWorkspaceSize(
  const aclTensor *x,
  const aclTensor *biasOptional,
  int64_t          k,
  int64_t          kGroup,
  int64_t          groupCount,
  int64_t          groupSelectMode,
  int64_t          renorm,
  int64_t          normType,
  bool             outFlag,
  double           routedScalingFactor,
  double           eps,
  const aclTensor *yOut,
  const aclTensor *expertIdxOut,
  const aclTensor *outOut,
  uint64_t        *workspaceSize,
  aclOpExecutor  **executor)

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

aclnnMoeGatingTopKGetWorkspaceSize

• Parameters:

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-Contiguous Tensor
  x	Input	Input parameter to be computed, corresponding to x in the formula.	None	FLOAT16, BFLOAT16, FLOAT32	ND	0-8	√
  biasOptional	Input	Bias value computed with the input x, corresponding to `bias` in the formula.	The shape value equals the last dimension of x.	FLOAT16, BFLOAT16, FLOAT32	ND	0-8	√
  k	Input	Specifies the k value of TopK, corresponding to `k` in the formula.	None	INT64	-	-	-
  kGroup	Input	Number of groups selected after grouped sorting, corresponding to `kGroup` in the formula.	None	INT64	-	-	-
  groupCount	Input	Total number of groups, corresponding to `groupCount` in the formula.	None	INT64	-	-	-
  groupSelectMode	Input	Grouped sorting mode.	None	INT64	-	-	-
  renorm	Input.	renorm flag.	None.	INT64	-	-	-
  normType	Input	Type of the norm function.	None	INT64	-	-	-
  outFlag	Input	Indicates whether to output the norm operation result.	None	BOOL	-	-	-
  routedScalingFactor	Input	`routedScalingFactor` coefficient used to compute yOut, corresponding to `routedScalingFactor` in the formula.	None	DOUBLE	-	-	-
  eps	Input.	Eps coefficient used to compute yOut, corresponding to `eps` in the formula.	None.	DOUBLE	-	-	-
  yOut	Output	Result computed after performing norm, grouped sorting, and topk on x, corresponding to `yOut` in the formula.	The data type must be consistent with x.	FLOAT16, BFLOAT16, FLOAT32	ND	0-8	-
  expertIdxOut	Output	Index after performing norm and grouped sorting TopK on x, corresponding to `expertIdxOut` in the formula.	The shape must be consistent with yOut.	INT32	ND	0-8	-
  outOut	Output	Output of the norm computation, corresponding to `normOut` in the formula.	The shape must be consistent with x.	FLOAT32	ND	0-8	-
  workspaceSize	Output	Returns the workspace size to be allocated on the Device side.	-	-	-	-	-
  executor	Output	Returns the op executor, which contains the operator computation flow.	-	-	-	-	-

• 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 input and output are null pointers.
  ACLNN_ERR_PARAM_NULLPTR	161002	The data type of the input or output is not within the supported range.
  ACLNN_ERR_INNER_TILING_ERROR	561002	The shape of x does not meet the requirement. The shapes of x and biasOptional do not match. The value of k is not within the range from 1 to x_shape[-1] / groupCount * kGroup. The value of kGroup is not within the range from 1 to groupCount. After aligning the number of experts in each group to 32, the computed input parameter value does not meet the requirement.

aclnnMoeGatingTopK

• Parameters:

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

• Return value

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

Constraints

• Deterministic computation:

  • aclnnMoeGatingTopK defaults to a deterministic implementation.

• Input shape constraints:

  • The last dimension of x (i.e., the number of experts) must not exceed 2048.

• Input value range restrictions:

  • Requires 1 <= k <= x_shape[-1] / groupCount * kGroup.

  • Requires 1 <= kGroup <= groupCount, and the value of kGroup * x_shape[-1] / groupCount must be greater than or equal to k.

  • Requires groupCount > 0, x_shape[-1] must be divisible by groupCount and the quotient must be greater than groupSelectMode, and the result of the quotient aligned to 32 elements multiplied by groupCount must not exceed 2048.

  • renorm supports only 0 and 1, indicating that the norm operation is performed first, followed by topk computation.

• Other constraints:

  • groupSelectMode takes values 0 and 1. 0 indicates sorting groups by the maximum value, and 1 indicates sorting groups by the sum of topk2 values.

  • normType takes values 0 and 1, where 0 indicates using the Softmax function and 1 indicates using the Sigmoid function.

  • outFlag takes values true and false, where true indicates output and false indicates no output.

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.h"
#include <iostream>
#include <vector>
#include <random>

#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;
}

std::vector<float> GenerateRandomFloats(int64_t count) {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<float> dist(0.0f, 10.0f);

    std::vector<float> result(count);
    for (auto& num : result) {
        num = dist(gen);
    }
    return result;
}
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 device-side 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-side data to the device-side 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 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. For details, see the AscendCL API list.
  // Set 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 interface.
  std::vector<int64_t> inputShape = {3, 256};
  std::vector<int64_t> biasShape = {256};
  std::vector<int64_t> outShape = {3, 8};
  std::vector<int64_t> expertIdOutShape = {3, 8};
  std::vector<int64_t> normOutShape = {3, 256};

  void* inputAddr = nullptr;
  void* biasAddr = nullptr;
  void* outAddr = nullptr;
  void* expertIdOutAddr = nullptr;
  void* normOutAddr = nullptr;

  aclTensor* input = nullptr;
  aclTensor* bias = nullptr;
  aclTensor* out = nullptr;
  aclTensor* expertIdOut = nullptr;
  aclTensor* normOut = nullptr;

  std::vector<float> inputHostData = GenerateRandomFloats(GetShapeSize(inputShape));
  std::vector<float> biasHostData = GenerateRandomFloats(GetShapeSize(biasShape));
  std::vector<float> outHostData(GetShapeSize(outShape));
  std::vector<int32_t> expertIdOutHostData(GetShapeSize(expertIdOutShape));
  std::vector<float> normOutHostData(GetShapeSize(normOutShape));

  // Create the expandedPermutedRows aclTensor.
  ret = CreateAclTensor(inputHostData, inputShape, &inputAddr, aclDataType::ACL_FLOAT, &input);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the expandedPermutedRows aclTensor.
  ret = CreateAclTensor(biasHostData, biasShape, &biasAddr, aclDataType::ACL_FLOAT, &bias);
  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 expandedSrcToDstRow aclTensor.
  ret = CreateAclTensor(normOutHostData, normOutShape, &normOutAddr, aclDataType::ACL_FLOAT, &normOut);
  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 aclnnMoeGatingTopK.
  ret = aclnnMoeGatingTopKGetWorkspaceSize(input, bias, 8, 4, 8, 1, 0, 1, false, 1, 1, out, expertIdOut, normOut, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMoeGatingTopKGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  // Allocate device memory based on the workspaceSize computed 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 aclnnMoeGatingTopK.
  ret = aclnnMoeGatingTopK(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMoeGatingTopK 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(bias);
  aclDestroyTensor(out);
  aclDestroyTensor(expertIdOut);
  aclDestroyTensor(normOut);

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