aclnnGroupedMatmulFinalizeRouting

📄 View Source Code

Applicable Products

Function

A fused operator of GroupedMatmul and MoeFinalizeRouting. It performs a combine operation on the output of GroupedMatmul based on indices.

Function Prototype

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

aclnnStatus aclnnGroupedMatmulFinalizeRoutingGetWorkspaceSize(
  const aclTensor *x,
  aclTensor       *w,
  const aclTensor *scaleOptional,
  const aclTensor *biasOptional,
  const aclTensor *pertokenScaleOptional,
  const aclTensor *groupListOptional,
  const aclTensor *sharedInputOptional,
  const aclTensor *logitOptional,
  const aclTensor *rowIndexOptional,
  int64_t          dtype,
  float            sharedInputWeight,
  int64_t          sharedInputOffset,
  bool             transposeX,
  bool             transposeW,
  int64_t          groupListType,
  aclTensor       *y,
  uint64_t        *workspaceSize,
  aclOpExecutor   **executor)

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

aclnnGroupedMatmulFinalizeRoutingGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
  x	Input	Input x (left matrix).	-	INT8	ND	(m, k), where m ranges from [1, 16*1024*8] and k supports 2048.	-
  w	Input	Input weight (right matrix).	-	INT4	ND	Supports three-dimensional shape. When the input is INT32, the dimensions are (e, k, n / 8). When the input is converted to INT4, the dimensions are (e, k, n). e ranges from [1, 256], k supports 2048, and n supports 7168.	-
  scaleOptional	Input	Scale factor in quantization parameters, a per-channel quantization parameter.	-	INT64	ND	The shape supports three dimensions (e, 1, n), where e and n are consistent with those of w.	-
  biasOptional	Input	Offset of the matrix.	-	FLOAT32	ND	Shape supports two dimensions, with dimensions (e, n), where e and n are consistent with e and n of w.	-
  offsetOptional	Input	Offset for asymmetric quantization.	-	FLOAT32	ND	Supports a three-dimensional shape with dimensions (e, 1, n), where e and n are consistent with those of w.	-
  antiquantScaleOptional	Input	Scaling factor for pseudo-quantization.	Not enabled currently.	FLOAT32	ND		-
  antiquantOffsetOptional	Input	Offset for pseudo-quantization.	Not enabled currently.	FLOAT32	ND		-
  pertokenScaleOptional	Input	Dequantization parameter for matrix computation.		FLOAT32	ND	The shape supports one dimension, with the dimension being (m), where m is consistent with the m of x.	-
  groupListOptional	Input	Distribution of matmul sizes along the grouping axis for input and output.		INT64	ND	Shape supports one dimension, with the dimension being (e), where e is consistent with the e of w.	-
  sharedInputOptional	Input	Output of the shared expert in MoE computation, which needs to be combined with the output of the MoE expert.		BF16	ND	The shape supports two dimensions, with dimensions (bsdp, n). bsdp must be less than or equal to batchSize/e, and n must be consistent with n of w.	-
  logitOptional	Input	Logit values of MoE experts for each token.		FLOAT32	ND	The shape supports one dimension, with the dimension being (m), where m is consistent with m of x.	-
  rowIndexOptional	Input	Combines the MoE expert outputs according to this rowIndex, where the values serve as indices for scatter-add during combination.		INT64	ND	The shape supports one dimension, with the dimension being (m), where m is consistent with m of x.	-
  dtype	Input	Output type for computation: 0: FLOAT32; 1: FLOAT16; 2: BFLOAT16. Currently only 0 is supported.		INT64			-
  sharedInputWeight	Input	Coefficient for combining the shared expert and MoE expert. sharedInput is first multiplied by this parameter, then accumulated with the MoE expert result.		FLOAT32			-
  sharedInputOffset	Input	Offset of the shared expert output in the total output.		INT64			-
  transposeX	Input	Indicates whether to transpose the left matrix. Only false is supported.		BOOL			-
  transposeW	Input	Whether the right matrix is transposed. Only false is supported.		BOOL			-
  groupListType	Input	Group mode: Set to 0 for cumsum mode (prefix sum), or 1 for count mode.		INT64			-
  y	Output	Output result.	The shape is the same as self.	FLOAT32	ND	0-8	√
  workspaceSize	Output	Size of the workspace to be applied for on the Device side.	-	-	-	-	-
  executor	Output	Returns the operator 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	Error Code	Description
  ACLNN_ERR_PARAM_NULLPTR	161001	The input x, w, scaleOptional, biasOptional, or y is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The data type or data format of x, w, scaleOptional, biasOptional, offsetOptional, antiQuantScaleOptional, antiQuantOffsetOptional, perTokenScaleOptional, groupListOptional, sharedInputOptional, logitOptional, rowIndexOptional, sharedInputWeight, sharedInputOffset, transposeX, transposeW, or y is not within the supported range.
  		The shape of x, w, scaleOptional, biasOptional, offsetOptional, antiQuantScaleOptional, antiQuantOffsetOptional, perTokenScaleOptional, groupListOptional, sharedInputOptional, logitOptional, rowIndexOptional, or y does not meet the validation condition.
  		The shape of x, w, scaleOptional, biasOptional, offsetOptional, antiQuantScaleOptional, antiQuantOffsetOptional, perTokenScaleOptional, groupListOptional, sharedInputOptional, logitOptional, rowIndexOptional, or y is an empty tensor.
  		The value range of dtype, sharedInputOffset, transposeX, transposeW, or groupListType does not meet the condition.

aclnnGroupedMatmulFinalizeRouting

• 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 through the first-phase API aclnnGroupedMatmulFinalizeRoutingGetWorkspaceSize.
  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:

  • aclnnGroupedMatmulFinalizeRouting defaults to a non-deterministic implementation. You can enable deterministic computation by calling aclrtCtxSetSysParamOpt.

• Supported Types in Pseudo-Quantization Scenarios

  The input and output support the following data type combinations:

  x	w	scale	bias	pertokenScale	groupList	sharedInput	logit	rowIndex	out
  INT8	INT4	INT64	FLOAT32	FLOAT32	INT64	BFLOAT16	FLOAT32	INT64	FLOAT32

  • In this scenario, scaleOptional represents the result of offline fusion of per-channel and per-group.

  • In this scenario, biasOptional represents the auxiliary result of offline computation. The value must be $8 \times w \times scaleOptional$, and is accumulated along the first dimension.

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_permute.h"
#include "aclnnop/aclnn_grouped_matmul_finalize_routing.h"
#include "aclnnop/aclnn_trans_matmul_weight.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: 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);

    // Compute the strides of a 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;
}

template <typename T>
int CreateAclTensorWeight(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr,
                      aclDataType dataType, aclTensor **tensor)
{
    auto size = static_cast<uint64_t>(GetShapeSize(shape));

    const aclIntArray *mat2Size = aclCreateIntArray(shape.data(), shape.size());
    auto ret = aclnnCalculateMatmulWeightSizeV2(mat2Size, dataType, &size);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCalculateMatmulWeightSizeV2 failed. ERROR: %d\n", ret);
              return ret);
    size *= sizeof(T);

    // Call aclrtMalloc to allocate memory on the device side.
    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);

    // Compute the strides of a 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];
    }

    std::vector<int64_t> storageShape;
    storageShape.push_back(GetShapeSize(shape));

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

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

    // 2. Construct the input and output. Customize the construction based on the API.
    int64_t m = 8;
    int64_t k = 2048;
    int64_t n = 7168;
    int64_t e = 1;
    int64_t batch = 8;
    int64_t bsdp = 1;
    int64_t dtype = 0;
    float shareInputWeight = 1.0;
    int64_t sharedInputOffset = 0;
    bool transposeX = false;
    bool transposeW = false;
    int64_t groupListType = 1;

    std::vector<int64_t> xShape = {m, k};
    std::vector<int64_t> wShape = {e, k, n / 8};
    std::vector<int64_t> scaleShape = {e, 1, n};
    std::vector<int64_t> biasShape = {e, n};
    std::vector<int64_t> pertokenScaleShape = {m};
    std::vector<int64_t> groupListShape = {e};
    std::vector<int64_t> sharedInputShape = {bsdp, n};
    std::vector<int64_t> logitShape = {m};
    std::vector<int64_t> rowIndexShape = {m};
    std::vector<int64_t> outShape = {batch, n};

    void *xDeviceAddr = nullptr;
    void *wDeviceAddr = nullptr;
    void *biasDeviceAddr = nullptr;
    void *scaleDeviceAddr = nullptr;
    void *pertokenScaleDeviceAddr = nullptr;
    void *groupListDeviceAddr = nullptr;
    void *sharedInputDeviceAddr = nullptr;
    void *logitDeviceAddr = nullptr;
    void *rowIndexDeviceAddr = nullptr;
    void *outDeviceAddr = nullptr;

    aclTensor* x = nullptr;
    aclTensor* w = nullptr;
    aclTensor* bias = nullptr;
    aclTensor* groupList = nullptr;
    aclTensor* scale = nullptr;
    aclTensor* pertokenScale = nullptr;
    aclTensor* sharedInput = nullptr;
    aclTensor* logit = nullptr;
    aclTensor* rowIndex = nullptr;
    aclTensor* out = nullptr;

    std::vector<int8_t> xHostData(GetShapeSize(xShape));
    std::vector<int32_t> wHostData(GetShapeSize(wShape));
    std::vector<int64_t> scaleHostData(GetShapeSize(scaleShape));
    std::vector<float> biasHostData(GetShapeSize(biasShape));
    std::vector<float> pertokenScaleHostData(GetShapeSize(pertokenScaleShape));
    std::vector<int64_t> groupListHostData(GetShapeSize(groupListShape));
    std::vector<uint16_t> sharedInputHostData(GetShapeSize(sharedInputShape));
    std::vector<int64_t> logitHostData(GetShapeSize(logitShape));
    std::vector<float> rowIndexHostData(GetShapeSize(rowIndexShape));
    std::vector<float> outHostData(GetShapeSize(outShape));
    // Assign a value to groupList.
    groupListHostData[0] = 8;
    // Create the x aclTensor.
    ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_INT8, &x);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> xTensorPtr(x, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void *)> xDeviceAddrPtr(xDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the w aclTensor of int32_t, which will be converted to int4 later.
    ret = CreateAclTensorWeight(wHostData, wShape, &wDeviceAddr, aclDataType::ACL_INT32, &w);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> wTensorPtr(w, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void *)> wDeviceAddrPtr(wDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create scale aclTensor.
    ret = CreateAclTensor(scaleHostData, scaleShape, &scaleDeviceAddr, aclDataType::ACL_INT64, &scale);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> scaleTensorPtr(scale, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void *)> scaleDeviceAddrPtr(scaleDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create bias aclTensor.
    ret = CreateAclTensor(biasHostData, biasShape, &biasDeviceAddr, aclDataType::ACL_FLOAT, &bias);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> biasTensorPtr(bias, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void *)> biasDeviceAddrPtr(biasDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create pertokenScale aclTensor.
    ret = CreateAclTensor(pertokenScaleHostData, pertokenScaleShape, &pertokenScaleDeviceAddr, aclDataType::ACL_FLOAT, &pertokenScale);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> pertokenScaleTensorPtr(pertokenScale, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void *)> pertokenScaleDeviceAddrPtr(pertokenScaleDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create groupList aclTensor.
    ret = CreateAclTensor(groupListHostData, groupListShape, &groupListDeviceAddr, aclDataType::ACL_INT64, &groupList);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> groupListTensorPtr(groupList, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void *)> groupListDeviceAddrPtr(groupListDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create sharedInput aclTensor.
    ret = CreateAclTensor(sharedInputHostData, sharedInputShape, &sharedInputDeviceAddr, aclDataType::ACL_BF16, &sharedInput);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> sharedInputTensorPtr(sharedInput, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void *)> sharedInputDeviceAddrPtr(sharedInputDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the logit aclTensor.
    ret = CreateAclTensor(logitHostData, logitShape, &logitDeviceAddr, aclDataType::ACL_FLOAT, &logit);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> logitTensorPtr(logit, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void *)> logitDeviceAddrPtr(logitDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the rowIndex aclTensor.
    ret = CreateAclTensor(rowIndexHostData, rowIndexShape, &rowIndexDeviceAddr, aclDataType::ACL_INT64, &rowIndex);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> rowIndexTensorPtr(rowIndex, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void *)> rowIndexDeviceAddrPtr(rowIndexDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the out aclTensor.
    ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor *)> outTensorPtr(out, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void *)> outDeviceAddrPtr(outDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

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

    // Call the first-phase API of aclnnGroupedMatmulFinalizeRouting.
    workspaceSize = 0;
    ret = aclnnGroupedMatmulFinalizeRoutingGetWorkspaceSize(x, w, scale, bias, pertokenScale, groupList, sharedInput, logit, rowIndex, dtype, shareInputWeight, sharedInputOffset, transposeX, transposeW, groupListType, out, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnGroupedMatmulFinalizeRoutingGetWorkspaceSize failed. ERROR: %d\n", ret);
              return ret);
    // Apply for device memory based on the workspaceSize calculated by the first-phase API.
    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 aclnnGroupedMatmulFinalizeRouting.
    ret = aclnnGroupedMatmulFinalizeRouting(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnGroupedMatmulFinalizeRouting failed. ERROR: %d\n", ret); return ret);

    // 4. (Boilerplate) Synchronize and 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 memory. Modify based on the specific API definition.
    auto size = GetShapeSize(outShape);
    std::vector<float> resultData(size, 0);
    ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr,
                      size * sizeof(resultData[0]), 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[%lld] is: %f\n", i, resultData[i]);
    }

    // 6. Release aclTensor resources. Modify based on the specific API definition.
    aclDestroyTensor(x);
    aclDestroyTensor(w);
    aclDestroyTensor(scale);
    aclDestroyTensor(bias);
    aclDestroyTensor(pertokenScale);
    aclDestroyTensor(groupList);
    aclDestroyTensor(sharedInput);
    aclDestroyTensor(logit);
    aclDestroyTensor(rowIndex);
    aclDestroyTensor(out);

    // 7. Release device resources. Modify based on the specific API definition.
    aclrtFree(xDeviceAddr);
    aclrtFree(wDeviceAddr);
    aclrtFree(scaleDeviceAddr);
    aclrtFree(biasDeviceAddr);
    aclrtFree(pertokenScaleDeviceAddr);
    aclrtFree(groupListDeviceAddr);
    aclrtFree(sharedInputDeviceAddr);
    aclrtFree(logitDeviceAddr);
    aclrtFree(rowIndexDeviceAddr);
    aclrtFree(outDeviceAddr);

    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}