aclnnGroupedMatmulWeightNz

Supported Products

Function

• Description: Implements grouped matrix multiplication, supporting non-uniform matrix dimension sizes across multiple groups. The basic function is matrix multiplication, for example, $y_i[m_i,n_i]=x_i[m_i,k_i] times weight_i[k_i,n_i], i=1...g$, where $g$ indicates the number of groups and $m_i$, $k_i$, and $n_i$ define the shapes for each group. Both inputs and outputs are of the aclTensorList type, with the following functions:

  • K-axis grouping: $k_i$ varies across groups, while $m_i$ and $n_i$ remain the same for each group. In this case, $x_i$ and $weight_i$ can be concatenated along the K-axis.
  • M-axis grouping: $k_i$ remains the same for each group. In this case, $weight_i$ and $y_i$ can be concatenated along the N-axis.

  New features compared with GroupedMatmulV5:

  • The input weight data supports the AI processor-affinity layout format (FRACTAL_NZ).
  • The quantGroupSize parameter is added. It is an integer that indicates the quantization group size in per-group mode. If per-group quantization is not involved, set this parameter to 0.

• Formulas:

  • Non-quantization scenario:$$y_i=x_i \times weight_i + bias_i$$

  • Quantization scenario (without perTokenScaleOptional):

    • x in INT8 and bias in INT32

      $$y_i=(x_i \times weight_i + bias_i) * scale_i + offset_i$$

    • x in INT8 and bias in BFLOAT16/FLOAT16/FLOAT32, without offset

      $$y_i=(x_i \times weight_i) * scale_i + bias_i$$

  • Quantization scenario (with perTokenScaleOptional):

    • x in INT8 and bias in INT32

      $$y_i=(x_i \times weight_i + bias_i) * scale_i * per\_token\_scale_i$$

    • x in INT8 and bias in BFLOAT16/FLOAT16/FLOAT32

      $$y_i=(x_i \times weight_i) * scale_i * per\_token\_scale_i + bias_i$$

  • Quantization scenario (MX quantization, without bias or activation layer):

    $$y_i=(x_i \times per\_token\_scale_i) * (weight_i \times scale_i)$$
  • Dequantization scenario:$$y_i=(x_i \times weight_i + bias_i) * scale_i$$
  • Fake-quantization scenario:$$y_i=x_i \times (weight_i + antiquant\_offset_i) * antiquant\_scale_i + bias_i$$

Prototype

Each operator has two-phase API calls. First, aclnnGroupedMatmulWeightNzGetWorkspaceSize is called to obtain the input parameters and compute the required workspace size based on the process. Then, aclnnGroupedMatmulWeightNz is called to perform computation.

aclnnStatus aclnnGroupedMatmulWeightNzGetWorkspaceSize(
    const aclTensorList *x, 
    const aclTensorList *weight, 
    const aclTensorList *biasOptional, 
    const aclTensorList *scaleOptional, 
    const aclTensorList *offsetOptional, 
    const aclTensorList *antiquantScaleOptional, 
    const aclTensorList *antiquantOffsetOptional, 
    const aclTensorList *perTokenScaleOptional, 
    const aclTensor     *groupListOptional, 
    const aclTensorList *activationInputOptional, 
    const aclTensorList *activationQuantScaleOptional, 
    const aclTensorList *activationQuantOffsetOptional, 
    int64_t              splitItem, 
    int64_t              groupType, 
    int64_t              groupListType, 
    int64_t              actType, 
    aclIntArray         *tuningConfigOptional, 
    int64_t              quantGroupSize, 
    aclTensorList       *out, 
    aclTensorList       *activationFeatureOutOptional, 
    aclTensorList       *dynQuantScaleOutOptional, 
    uint64_t            *workspaceSize, 
    aclOpExecutor      **executor)

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

aclnnGroupedMatmulWeightNzGetWorkspaceSize

• Parameters

  Name	Input/Output	Description	Usage Notes	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
  x	Input	Input x in the formula.	The maximum length is 128.	FLOAT16, BFLOAT16, INT8, INT41, INT321, FLOAT8_E4M3FN2	ND	-	-
  weight	Input	Weight in the formula.	The maximum length is 128. The Ascend-affinity data layout format (NZ) is supported.	FLOAT16, BFLOAT16, INT8, INT4, INT32, FLOAT32, FLOAT4_E2M12	ND, FRACTAL_NZ	-	-
  biasOptional	Optional input	Bias in the formula.	Same length as weight.	FLOAT16, FLOAT32, INT32, BFLOAT16	ND	-	-
  scaleOptional	Optional input	Scale in the formula, indicating the scale factor for quantization parameters.	Generally, the length is the same as the weight length. For details about the constraints, see Constraints.	UINT641, BFLOAT161, FLOAT32	ND	-	-
  offsetOptional	Optional input	Offset in the formula, indicating the offset for quantization parameters.	Same length as weight.	FLOAT32	ND	-	-
  antiquantScaleOptional	Optional input	antiquant_scale in the formula, indicating the scale factor for fake-quantization parameters.	Same length as weight. For details about the constraints, see Constraints.	FLOAT16, BFLOAT161, FLOAT8_E8M02	ND	-	-
  antiquantOffsetOptional	Optional input	antiquant_offset in the formula, indicating the offset for fake-quantization parameters.	Same length as weight. For details about the constraints, see Constraints.	FLOAT16, BFLOAT161	ND	-	-
  perTokenScaleOptional	Optional input	per_token_scale in the formula, indicating the scale factor (introduced by x quantization) for quantization parameters.	It is valid in scenarios where x, weight, and out are all single-tensor. For details about the constraints, see Constraints.	FLOAT32, FLOAT8_E8M02	ND	-	-
  groupListOptional	Optional input	Matmul size distribution along the grouping axis for inputs and outputs.	Input data in different formats based on groupListType. Note that when the length of the output TensorList is 1, the last value of this parameter constrains the valid portion of the output.	INT64	ND	-	-
  activationInputOptional	Optional input	Backward input of the activation function.	Currently, only nullptr can be passed.	-	-	-	-
  activationQuantScaleOptional	Optional input	-	Currently, only nullptr can be passed.	-	-	-	-
  activationQuantOffsetOptional	Optional input	-	Currently, only nullptr can be passed.	-	-	-	-
  splitItem	Input	Specifies whether to perform tensor splitting on the output.	0/1 indicates that the output is multi-tensor; 2/3 indicates that the output is single-tensor.	INT64	-	-	-
  groupType	Input	Axis to be grouped.	-1: no grouping; 0: M-axis grouping; 1: N-axis grouping; 2: K-axis grouping. For details about the constraints, see Constraints.	INT64	-	-	-
  groupListType	Input	Grouping mode of the groupList input.	0: Cumulative sum result; 1: size of each group; 2: [groupIdx, groupSize]. For details about the constraints, see Constraints.	INT64	-	-	-
  actType	Input	Activation function type.	The value ranges from 0 to 5. For details about the constraints, see Constraints.	INT64	-	-	-
  tuningConfigOptional	Optional input	The first value indicates the expected number of tokens to be processed by each expert, which is used to optimize tiling.	Compatible with earlier versions. If this parameter is not used, do not pass it (that is, pass nullptr).	INT64	-	-	-
  quantGroupSize	Input	Group size for per-group quantization.	If per-group quantization is not involved, set this parameter to `0`.	INT64	-	-	-
  out	Output	Output y in the formula.	The maximum length is 128.	FLOAT16, BFLOAT16, INT8, FLOAT32, INT32	ND	-	-
  activationFeatureOutOptional	Output	Input data of the activation function.	Currently, only nullptr can be passed.	-	-	-	-
  dynQuantScaleOutOptional	Output	-	Currently, only nullptr can be passed.	-	-	-	-
  workspaceSize	Output	Size of the workspace required to be allocated on the device.	-	-	-	-	-
  executor	Output	Operator executor, containing the operator computation process.	-	-	-	-	-

  • Atlas A2 training products/Atlas A2 inference products and Atlas A3 training products/Atlas A3 inference products:
    • The superscript "1" in the "Data Type" column of the table above indicates data types that are supported by the products, and the superscript "2" indicates data types that are not supported by the products.
    • The weight format can be converted from ND to NZ by calling aclnnCalculateMatmulWeightSizeV2 and aclnnTransMatmulWeight. When the input data type is INT32, each INT32 value is treated as eight INT4 values inside the API.
  • Atlas inference products:
    • Only FLOAT16 is supported. weight supports only the FRACTAL_NZ format, which needs to be converted through auxiliary APIs.
    • Quantization/asymmetric quantization parameters such as scaleOptional and offsetOptional are not supported, and need to be passed as null pointers.
    • groupType supports only M-axis grouping (0). The value of actType can only be 0. tuningConfigOptional is not supported.

• Return

  aclnnStatus status code. For details, see aclnn Return Codes.

  The first-phase API implements input parameter validation. The following errors may be thrown:

  Return	Error Code	Description
  ACLNN_ERR_PARAM_NULLPTR	161001	1. The required input, output, or attribute is passed as a null pointer.
  		2. The input weight contains elements that are null pointers.
  		3. The input x contains elements that are null pointers, while the corresponding elements in the output out are non-null pointers.
  		4. The input x contains elements that are non-null pointers, while the corresponding elements in the output out are null pointers.
  ACLNN_ERR_PARAM_INVALID	161002	1. The data type or data format of x, weight, biasOptional, scaleOptional, offsetOptional, antiquantScaleOptional, antiquantOffsetOptional, groupListOptional, or out is not supported.
  		2. The weight length is greater than 128; the bias is not null but its length is not equal to the weight length.
  		3. The dimension of groupListOptional is 1.
  		4. When splitItem is set to 2 or 3, the length of out is not 1.
  		5. When splitItem is set to 0 or 1, the length of out or groupListOptional is not equal to the weight length.
  		6. An element in the input parameter tuningConfigOptional is negative or exceeds the number of rows (M) in x.

aclnnGroupedMatmulWeightNz

• Parameters

  Parameter	Input/Output	Description
  workspace	Input	Address of the workspace to be allocated on the device.
  workspaceSize	Input	Size of the workspace to be allocated on the device, which is obtained by calling aclnnGroupedMatmulWeightNzGetWorkspaceSize.
  executor	Input	Operator executor, containing the operator computation process.
  stream	Input	Stream for executing the task.

• Return

  aclnnStatus status code. For details, see aclnn Return Codes.

Constraints

• Deterministic computation:
  • aclnnGroupedMatmulWeightNz defaults to deterministic implementation.
• Common constraints
  • If groupListOptional is passed: when groupListType is 0, groupListOptional must be a non-negative, monotonically non-decreasing sequence; when groupListType is 1, groupListOptional must be a non-negative sequence and its length cannot be 1; when groupListType is 2, the second column of groupListOptional must be a non-negative sequence and its length cannot be 1.
  • The size of each dimension for every tensor in x and weight, after 32-byte alignment, should be less than the maximum value of INT32 (2147483647).

Calling Example

The following example is for reference only. For details, see Compilation and Running Sample.

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

#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 shapeSize = 1;
  for (auto i : shape) {
    shapeSize *= i;
  }
  return shapeSize;
}

int Init(int32_t deviceId, aclrtStream* stream) {
  // (Boilerplate) Initialize resources.
  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_New(const std::vector<int64_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.
  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 the data on the host to the memory on the device.
  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;
}

template <typename T>
int CreateAclTensor(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.
  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 the data on the host to the memory on the device.
  std::vector<T> hostData(size, 0);
  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 CreateAclTensorList(const std::vector<std::vector<int64_t>>& shapes, void** deviceAddr,
                        aclDataType dataType, aclTensorList** tensor) {
  int size = shapes.size();
  aclTensor* tensors[size];
  for (int i = 0; i < size; i++) {
    int ret = CreateAclTensor<uint16_t>(shapes[i], deviceAddr + i, dataType, tensors + i);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
  }
  *tensor = aclCreateTensorList(tensors, size);
  return ACL_SUCCESS;
}

template <typename T>
int CreateAclTensorNz(const std::vector<T> &hostData, const std::vector<std::vector<int64_t>> &shapes, 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];
  }
int64_t E = shape[0];
int64_t K = shape[1];
int64_t N = shape[2];
std::vector<int64_t> shapeNz = {E, N/64, K/16, 16, 64};

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

template <typename T>
int CreateAclTensorListNz(const std::vector<std::vector> &hostData, const std::vector<std::vector<int64_t>> &shapes, void **deviceAddr,
                        aclDataType dataType, aclTensorList **tensor)
{
  int size = shapes.size();
  aclTensor * tensors[size];
  for (int i = 0; i < size; i++) {
    int ret = CreateAclTensorNz<T>(hostData[i], shapes[i], deviceAddr + i, dataType, tensors + i);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
  }
  *tensor = aclCreateTensorList(tensors, size);
  return ACL_SUCCESS;
}

int main() {
  // 1. (Boilerplate) Initialize the device and stream. For details, see the ACL API manual.
  // Set the device ID in use.
  int32_t deviceId = 0;
  aclrtStream stream;
  auto ret = Init(deviceId, &stream);
  // Customize error handling based on your requirements.
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

  // 2. Construct inputs and outputs based on API definitions.
  std::vector<std::vector<int64_t>> xShape = {{512, 256}};
  std::vector<std::vector<int64_t>> weightShape= {{2, 256, 256}};
  std::vector<std::vector<int64_t>> yShape = {{512, 256}};
  std::vector<int64_t> groupListShape = {{2}};
  std::vector<int64_t> groupListData = {256, 512};
  void* xDeviceAddr[1];
  void* weightDeviceAddr[1];
  void* yDeviceAddr[1];
  void* groupListDeviceAddr;
  aclTensorList* x = nullptr;
  aclTensorList* weight = nullptr;
  aclTensorList* bias = nullptr;
  aclTensor* groupedList = nullptr;
  aclTensorList* scale = nullptr;
  aclTensorList* offset = nullptr;
  aclTensorList* antiquantScale = nullptr;
  aclTensorList* antiquantOffset = nullptr;
  aclTensorList* perTokenScale = nullptr;
  aclTensorList* activationInput = nullptr;
  aclTensorList* activationQuantScale = nullptr;
  aclTensorList* activationQuantOffset = nullptr;
  aclTensorList* out = nullptr;
  aclTensorList* activationFeatureOut = nullptr;
  aclTensorList* dynQuantScaleOut = nullptr;
  int64_t splitItem = 3;
  int64_t groupType = 0;
  int64_t groupListType = 0;
  int64_t actType = 0;
  std::vector<int8_t> wHostData(GetShapeSize(weightShape));
  // Create a tuningconfig aclIntArray.
  std::vector<int64_t> tuningConfigData = {512};
  aclIntArray *tuningConfig = aclCreateIntArray(tuningConfigData.data(), 1);

  // Create an x aclTensorList.
  ret = CreateAclTensorList(xShape, xDeviceAddr, aclDataType::ACL_BF16, &x);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create a weight aclTensorList.
  ret = CreateAclTensorListNz(wHostData, weightShape, weightDeviceAddr, aclDataType::ACL_BF16, &weight);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create a y aclTensorList.
  ret = CreateAclTensorList(yShape, yDeviceAddr, aclDataType::ACL_BF16, &out);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create a group_list aclTensor.
  ret = CreateAclTensor_New<int64_t>(groupListData, groupListShape, &groupListDeviceAddr, aclDataType::ACL_INT64, &groupedList);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  uint64_t workspaceSize = 0;
  aclOpExecutor* executor;

  // 3. Call the CANN operator library API.
  // Call the first-phase API of aclnnGroupedMatmulWeightNz.
  ret = aclnnGroupedMatmulWeightNzGetWorkspaceSize(x, weight, bias, scale, offset, antiquantScale, antiquantOffset, perTokenScale, groupedList, activationInput, activationQuantScale, activationQuantOffset, splitItem, groupType, groupListType, actType, tuningConfig, 0, out, activationFeatureOut, dynQuantScaleOut, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnGroupedMatmulWeightNzGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  // Allocate device memory based on 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 aclnnGroupedMatmulWeightNz.
  ret = aclnnGroupedMatmulWeightNz(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnGroupedMatmulWeightNz failed. ERROR: %d\n", ret); return ret);

  // 4. (Boilerplate) Wait until the task execution is 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 the code based on the API definition.
  for (int i = 0; i < 1; i++) {
    auto size = GetShapeSize(yShape[i]);
    std::vector<uint16_t> resultData(size, 0);
    ret = aclrtMemcpy(resultData.data(), size * sizeof(resultData[0]), yDeviceAddr[i],
                      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 j = 0; j < size; j++) {
        LOG_PRINT("result[%ld] is: %d\n", j, resultData[j]);
    }
  }

  // 6. Release aclTensors and aclScalars. Modify the code based on the API definition.
  aclDestroyTensorList(x);
  aclDestroyTensorList(weight);
  aclDestroyTensorList(bias);
  aclDestroyTensorList(out);

  // 7. Release device resources. Modify the code based on the API definition.
  for (int i = 0; i < 1; i++) {
    aclrtFree(xDeviceAddr[i]);
    aclrtFree(weightDeviceAddr[i]);
    aclrtFree(biasDeviceAddr[i]);
    aclrtFree(yDeviceAddr[i]);
  }
  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();
  return 0;
}