Documentation
Download
文档中心
Ln

aclnnMoeFusedTopk

Applicable Products

ProductSupported
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 Sigmoid computation on the input x, sorts the computation results by group, and finally selects the top k experts based on the grouped sorting results.

  • Formula:

    Apply sigmoid to the input:

    sigmoidRes=sigmoid(x)sigmoidRes=sigmoid(x)

    Add addNum:

    normOut=sigmoidRes+addNumnormOut = sigmoidRes + addNum

    Group the computation results by groupNum, sort the groups within each group by the sum of topN values, and take the top groupTopk groups:

    groupOut,groupId=TopK(ReduceSum(TopK(Split(normOut,groupCount),k=2,dim=−1),dim=−1),k=kGroup)groupOut, groupId = TopK(ReduceSum(TopK(Split(normOut, groupCount), k=2, dim=-1), dim=-1),k=kGroup)

    Based on the groupId from the previous step, retrieve the corresponding elements from normOut, then perform TopK on the data again to obtain the indices result:

    normY,indices=TopK(normOut[groupId,:],k=k)normY,indices=TopK(normOut[groupId, :],k=k)

    Select y from sigmoidRes based on indices:

    y=gather(sigmoidRes,indices)y = gather(sigmoidRes, indices)

    If isNorm is true, compute y based on the input scale parameter to obtain the result of y:

    y=y/(ReduceSum(y,dim=−1))∗scaley = y / (ReduceSum(y, dim=-1))*scale

    If enableExpertMapping is true, the physical experts in indices are then mapped to logical experts based on the input mappingNum and mappingTable to obtain the output indices.

Function Prototype

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

aclnnMoeFusedTopkGetWorkspaceSize(
  const aclTensor* x,
  const aclTensor* addNum,
  const aclTensor* mappingNum,
  const aclTensor* mappingTable,
  uint32_t         groupNum,
  uint32_t         groupTopk,
  uint32_t         topN,
  uint32_t         topK,
  uint32_t         activateType,
  bool             isNorm,
  float            scale,
  bool             enableExpertMapping,
  aclTensor*       y,
  aclTensor*       indices,
  uint64_t*        workspaceSize,
  aclOpExecutor**  executor)
aclnnStatus aclnnMoeFusedTopk(
  void*          workspace,
  uint64_t       workspaceSize,
  aclOpExecutor* executor,
  aclrtStream    stream)

aclnnMoeFusedTopkGetWorkspaceSize

  • Parameters:

    Parameter Input/Output Description Instruction Data Type Data Format Shape Non-contiguous Tensor
    x Input Score of each token corresponding to each expert. - FLOAT16, BFLOAT16, FLOAT32 ND (numToken, expertNum) √
    addNum Input Bias value used for computation with the input x. - Same as x. ND (expertNum) √
    mappingNum Input The number of logical experts to which each physical expert is actually mapped. Not used when enableExpertMapping is false. INT32 ND (expertNum) √
    mappingTable Input Mapping table for each physical expert/logical expert. Not enabled when enableExpertMapping is set to false.
    Less than or equal to 128.    		 INT32 ND (expertNum, maxMappingNum) √
    groupNum Input Number of groups. The value must be greater than 0. - - - - -
    groupTopk Input Number of selected groups. The value must be greater than 0. - - - - -
    topN Input Number of experts selected within each group for summation. Must be greater than 0. - - - - -
    topK Input Number of experts finally selected. The value must be greater than 0. - - - - -
    activateType Input Activation type. Currently only 0 (ACTIVATION_SIGMOID) is supported. - - - - -
    isNorm Input Whether to normalize the output. - - - - -
    scale Input Normalized coefficient multiplication. - - - - -
    enableExpertMapping Input Whether to enable mapping from physical experts to logical experts. - - - - -
    y Output The aclTensor on the device side. - FLOAT32 ND (numToken, topK) √
    indices Output The aclTensor on the device side. - INT32 ND (numToken, topK) √
    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 x or addNum is a null pointer.
    The output y or indices is a null pointer.
    When enableExpertMapping is true, the input mappingNum or mappingTable 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.
    The input parameters do not satisfy the constraints.
    The shapes of the input and output do not satisfy the constraints.

aclnnMoeFusedTopk

  • Parameters:

    Parameter Input/Output Description
    workspace Input The workspace memory address applied for on the Device side.
    workspaceSize Input The workspace size applied for on the Device side, obtained from the first-phase API aclnnMoeFusedTopkGetWorkspaceSize.
    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:

    • aclnnMoeFusedTopk defaults to a deterministic implementation.

  • expertNum must be an integer multiple of groupNum.

  • groupTopk is less than or equal to groupNum.

  • maxMappingNum must be less than or equal to 128.

  • TopK must be less than or equal to expertNum.

  • TopN must be less than or equal to expertNum / groupNum.

  • expertNum must be less than or equal to 1024.

  • groupNum must be less than or equal to 256.

Example

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

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

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

#define CHECK_FREE_RET(cond, return_expr) \
  do {                                     \
      if (!(cond)) {                       \
          Finalize(deviceId, stream);      \
          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 shapeSize = 1;
  for (auto i : shape) {
    shapeSize *= i;
  }
  return shapeSize;
}

int Init(int32_t deviceId, aclrtStream* stream) {
  // Boilerplate. Initialize.
  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 apply for 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 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 stride of the contiguous tensor.
  std::vector<int64_t> stride(shape.size(), 1);
  for (int64_t i = shape.size() - 2; i >= 0; i--) {
    stride[i] = shape[i + 1] * stride[i + 1];
  }

  // Call the aclCreateTensor API to create an aclTensor.
  *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, stride.data(), 0, aclFormat::ACL_FORMAT_ND,
                            shape.data(), shape.size(), *deviceAddr);
  return 0;
}

void Finalize(int32_t deviceId, aclrtStream stream)
{
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();
}

int aclnnMoeFusedTopkTest(int32_t deviceId, aclrtStream& stream) {
  auto ret = Init(deviceId, &stream);
  CHECK_FREE_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

  // 2. Construct the input and output. Customize the construction based on the API.

  int64_t num_token = 16;
  int64_t expert_num = 32;
  int64_t max_mapping_num = 16;

  uint32_t groupNum = 2;
  uint32_t groupTopk = 2;
  uint32_t topN = 2;
  uint32_t topK = 4;
  uint32_t activateType = 0;
  bool isNorm = false;
  float scale = 1.0;
  bool enableExpertMapping = true;

  std::vector<int64_t> xShape = {num_token, expert_num};
  std::vector<int64_t> addNumShape = {expert_num};
  std::vector<int64_t> mappingNumShape = {expert_num};
  std::vector<int64_t> mappingTableShape = {expert_num, max_mapping_num};

  std::vector<int64_t> yShape = {num_token, topK};
  std::vector<int64_t> indicesShape = {num_token, topK};

  void* xDeviceAddr = nullptr;
  void* addNumDeviceAddr = nullptr;
  void* mappingNumDeviceAddr = nullptr;
  void* mappingTableDeviceAddr = nullptr;

  void* yDeviceAddr = nullptr;
  void* indicesDeviceAddr = nullptr;

  aclTensor* x = nullptr;
  aclTensor* addNum = nullptr;
  aclTensor* mappingNum = nullptr;
  aclTensor* mappingTable = nullptr;

  aclTensor* y = nullptr;
  aclTensor* indices = nullptr;

  std::vector<float> xHostData(GetShapeSize(xShape), 1);
  std::vector<float> addNumHostData(GetShapeSize(addNumShape), 1);
  std::vector<int32_t> mappingNumHostData(GetShapeSize(mappingNumShape), 1);
  std::vector<int32_t> mappingTableHostData(GetShapeSize(mappingTableShape), 1);

  std::vector<float> yHostData(GetShapeSize(yShape), 0);
  std::vector<int32_t> indicesHostData(GetShapeSize(indicesShape), 0);

  // Create the x aclTensor.
  ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT, &x);
  CHECK_FREE_RET(ret == ACL_SUCCESS, return ret);
  std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> xTensorPtr(x, aclDestroyTensor);

  // Create the addNum aclTensor.
  ret = CreateAclTensor(addNumHostData, addNumShape, &addNumDeviceAddr, aclDataType::ACL_FLOAT, &addNum);
  CHECK_FREE_RET(ret == ACL_SUCCESS, return ret);
  std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> addNumTensorPtr(addNum, aclDestroyTensor);

  // Create the mappingNum aclTensor.
  ret = CreateAclTensor(mappingNumHostData, mappingNumShape, &mappingNumDeviceAddr, aclDataType::ACL_INT32, &mappingNum);
  CHECK_FREE_RET(ret == ACL_SUCCESS, return ret);
  std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> mappingNumTensorPtr(mappingNum, aclDestroyTensor);

  // Create the mappingTable aclTensor.
  ret = CreateAclTensor(mappingTableHostData, mappingTableShape, &mappingTableDeviceAddr, aclDataType::ACL_INT32, &mappingTable);
  CHECK_FREE_RET(ret == ACL_SUCCESS, return ret);
  std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> mappingTableTensorPtr(mappingTable, aclDestroyTensor);

  // Create the y aclTensor.
  ret = CreateAclTensor(yHostData, yShape, &yDeviceAddr, aclDataType::ACL_FLOAT, &y);
  CHECK_FREE_RET(ret == ACL_SUCCESS, return ret);
  std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> yTensorPtr(y, aclDestroyTensor);

  // Create the indices aclTensor.
  ret = CreateAclTensor(indicesHostData, indicesShape, &indicesDeviceAddr, aclDataType::ACL_INT32, &indices);
  CHECK_FREE_RET(ret == ACL_SUCCESS, return ret);
  std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> indicesTensorPtr(indices, aclDestroyTensor);

  // 3. Call the CANN operator library API. Replace with the specific API name.
  uint64_t workspaceSize = 0;
  aclOpExecutor* executor;
  // Call the first-phase API of aclnnMoeFusedTopk.
  ret = aclnnMoeFusedTopkGetWorkspaceSize(x,
                                       addNum,
                                       mappingNum,
                                       mappingTable,
                                       groupNum,
                                       groupTopk,
                                       topN,
                                       topK,
                                       activateType,
                                       isNorm,
                                       scale,
                                       enableExpertMapping,
                                       y,
                                       indices,
                                       &workspaceSize,
                                       &executor);
  CHECK_FREE_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMoeFusedTopkGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  // Apply for device memory based on the workspaceSize calculated by the first-phase API.
  void* workspaceAddr = nullptr;
  std::unique_ptr<void, aclError (*)(void *)> workspaceAddrPtr(nullptr, aclrtFree);
  if (workspaceSize > 0) {
    ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_FREE_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
    workspaceAddrPtr.reset(workspaceAddr);
  }
  // Call the aclnnMoeFusedTopk second-phase API.
  ret = aclnnMoeFusedTopk(workspaceAddr, workspaceSize, executor, stream);
  CHECK_FREE_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMoeFusedTopk failed. ERROR: %d\n", ret); return ret);

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

  // 5. Obtain the output values and copy the results from the device memory to the host side. Modify this based on the specific API definition.
  auto size = GetShapeSize(yShape);
  std::vector<float> yData(size, 0);
  ret = aclrtMemcpy(yData.data(), yData.size() * sizeof(yData[0]), yDeviceAddr,
                    size * sizeof(yData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
  CHECK_FREE_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("out result[%ld] is: %f\n", i, yData[i]);
  }

  return ACL_SUCCESS;
}

int main() {
  // 1. (Boilerplate) Initialize the device/stream. Refer to the external API list.
  // Enter the actual device ID for deviceId.
  int32_t deviceId = 0;
  aclrtStream stream;
  auto ret = aclnnMoeFusedTopkTest(deviceId, stream);
  CHECK_FREE_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMoeFusedTopkTest failed. ERROR: %d\n", ret); return ret);

  Finalize(deviceId, stream);
  return 0;
}
This site uses cookies to improve your browsing experience. By continuing to browse this site, you agree to our use of cookies.View Details