aclnnGroupedMatmulSwigluQuantWeightNZ

• Definition:

  • ⋅ denotes matrix multiplication.

  • ⊙ denotes element-wise multiplication.

  • $\left \lfloor x\right \rceil$ denotes rounding x to the nearest integer.

  • $\mathbb{Z_8} = \{ x \in \mathbb{Z} | −128≤x≤127 \}$

  • $\mathbb{Z_4} = \{ x \in \mathbb{Z} | −8≤x≤7 \}$

  • $\mathbb{Z_{32}} = \{ x \in \mathbb{Z} | -2147483648≤x≤2147483647 \}$

• Input:

  • $X∈\mathbb{Z_8}^{M \times K}$: activation matrix (left matrix), where M is the total number of tokens and K is the feature dimension.

  • $W∈\mathbb{Z_4}^{E \times K \times N}$: Grouped weight matrix (right matrix), where E is the number of experts, K is the feature dimension, and N is the output dimension.

  • $weightAssistMatrix∈\mathbb{R}^{E \times N}$: Auxiliary matrix used during matrix multiplication (see the computation process for generating the auxiliary matrix below).

  • $w\_scale∈\mathbb{R}^{E \times K\_group\_num \times N}$: Per-channel scale factor of the grouped weight matrix (right matrix), where E is the number of experts, K_group_num is the number of groups along the K axis, and N is the output dimension.

  • $x\_scale∈\mathbb{R}^{M}$: Per-token scale factor of the activation matrix (left matrix), where M is the total number of tokens.

  • $groupList∈\mathbb{N}^{E}$: Group index list of cumsum.

• Output:

  • $Q∈\mathbb{Z_8}^{M \times N / 2}$: quantized output matrix.

  • $Q\_scale∈\mathbb{R}^{M}$: quantization scale factor.

• Computation Process

  • 1. Determine the tokens of the current group based on groupList[i], where $i \in [0,Len(groupList))$.
    • The grouping logic is the same as that of A8W8.
  • 2. The computation process for generating the auxiliary matrix (weightAssistMatrix) (note that part of the weightAssistMatrix computation is generated offline as input, not completed inside the operator):

    • For per-channel quantization (where $w\_scale$ is 2-dimensional):

      $weightAssistMatrix_{i} = 8 × w\_scale × Σ_{k=0}^{K-1} weight[:,k,:]$

    • For per-group quantization (where $w\_scale$ is 3-dimensional):

      $weightAssistMatrix_{i} = 8 × Σ_{k=0}^{K-1} (weight[:,k,:] × w\_scale[:, ⌊k/num\_per\_group⌋, :])$

      NOTE

$num\_per\_group = K // K\_group\_num$

  - 3. Perform the following computation based on the input parameters determined by grouping:

    - 3.1. Split the left matrix $\mathbb{Z_8}$ into high and low 4-bit parts $\mathbb{Z_4}$:

      $X\_high\_4bits_{i} = \lfloor \frac{X_{i}}{16} \rfloor$
      $X\_low\_4bits_{i} = X_{i} \& 0x0f - 8$

    - 3.2. Enable per-channel or per-group quantization during matrix multiplication.

      per-channel:

      $C\_high_{i} = (X\_high\_4bits_{i} \cdot W_{i}) \odot w\_scale_{i}$

      $C\_low_{i} = (X\_low\_4bits_{i} \cdot W_{i}) \odot w\_scale_{i}$

      per-group:

      $C\_high_{i} = \\ Σ_{k=0}^{K-1}((X\_high\_4bits_{i}[:, k * num\_per\_group : (k+1) * num\_per\_group] \cdot W_{i}[k *   num\_per\_group : (k+1) * num\_per\_group, :]) \odot w\_scale_{i}[k, :] )$

      $C\_low_{i} = \\ Σ_{k=0}^{K-1}((X\_low\_4bits_{i}[:, k * num\_per\_group : (k+1) * num\_per\_group] \cdot W_{i}[k *   num\_per\_group : (k+1) * num\_per\_group, :]) \odot w\_scale_{i}[k, :] )$

    - 3.3. Restore the high/low-bit matrix multiplication results to the overall result.

      $C_{i} = (C\_high_{i} * 16 + C\_low_{i} + weightAssistMatrix_{i}) \odot x\_scale_{i}$

      $C_{i,act}, gate_{i} = split(C_{i})$

      $S_{i}=Swish(C_{i,act})\odot gate_{i}$     where $Swish(x)=\frac{x}{1+e^{-x}}$

  - 3. Quantize the output result.

    $Q\_scale_{i} = \frac{max(|S_{i}|)}{127}$

    $Q_{i} = \left\lfloor \frac{S_{i}}{Q\_scale_{i}} \right\rceil$
</details>

Function Prototype

Each operator uses a two-phase API. You must first call the "aclnnGroupedMatmulSwigluQuantWeightNZGetWorkspaceSize" API to obtain the workspace size required for computation and the executor that contains the operator computation process, and then call the "aclnnGroupedMatmulSwigluQuantWeightNZ" API to perform the computation.

aclnnStatus aclnnGroupedMatmulSwigluQuantWeightNZGetWorkspaceSize(
  const aclTensor *x,
  const aclTensor *weight,
  const aclTensor *bias,
  const aclTensor *offset,
  const aclTensor *weightScale,
  const aclTensor *xScale,
  const aclTensor *groupList,
  aclTensor       *output,
  aclTensor       *outputScale,
  aclTensor       *outputOffset,
  uint64_t        *workspaceSize,
  aclOpExecutor  **executor)

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

aclnnGroupedMatmulSwigluQuantWeightNZGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-Contiguous Tensor
  x	Input	Left matrix, corresponding to X in the formula.	If the shape is [M, K], K must be less than 65536.	INT8	ND	2	√
  weight	Input	Weight matrix, corresponding to W in the formula.	Note that this API ignores the data format of weight and forcibly treats it as FRACTAL_NZ format.	INT8, INT4, INT32 (INT32 is for adaptation purposes; one INT32 value is actually interpreted as eight INT4 values).	FRACTAL_NZ	5	√
  bias	Input	Specifies the auxiliary matrix used in matrix multiplication, corresponding to weightAssistMatrix in the formula.	Takes effect only in the A8W4 scenario. A null pointer must be passed in the A8W8 scenario.	FLOAT	ND	2	√
  offset	Input.	Offset for per-channel asymmetric dequantization, corresponding to the offset in the formula.	Reserved input. Not supported yet. Pass a null pointer.	-	-	-	-
  weightScale	Input	Quantization scale factor of the right matrix, corresponding to w_scale in the formula.	The first axis length must be equal to the first axis dimension of weight, and the last axis length must be the same as the last axis of weight restored to ND format.	FLOAT, FLOAT16, BFLOAT16	ND	2 (per-channel) or 3 (per-group)	√
  xScale	Input	Quantization factor of the left matrix, corresponding to x_scale in the formula.	Its length must be equal to the first axis dimension of x.	FLOAT	ND	1	√
  groupList	Input	Indicates the number of tokens participating in the computation for each group, corresponding to groupList in the formula.	The length must be equal to the first axis dimension of weight. The last value in groupList constrains the valid portion of the output data. For details, see the computation process in Function.	INT64	ND	1	√
  output	Output	Output quantization result, corresponding to Q in the formula.	-	INT8	ND	2	√
  outputScale	Output	Output quantization scale factor, corresponding to Q_scale in the formula.	-	FLOAT	ND	1	√
  outputOffset	Output	Offset of the output asymmetric quantization, corresponding to Q_offset in the formula.	Reserved input. Not supported currently. Pass a null pointer.	-	-	-	-
  workspaceSize	Output	Returns the workspace size to be applied for on the NPU 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	Error Code	Description
  ACLNN_ERR_PARAM_NULLPTR	161001	The input x, weight, weightScale, xScale, groupList, output, or outputScale is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The data dimension of the input x, weight, weightScale, xScale, groupList, output, or outputScale does not meet the constraints.
  		The shape of the input x, weight, weightScale, xScale, groupList, output, or outputScale does not meet the constraints.
  		The format of the input x, weight, weightScale, xScale, groupList, output, or outputScale data does not meet the constraints.
  		The number of elements in groupList exceeds the first axis length of weight.
  		The N axis length exceeds 10240.
  		In the A8W8 scenario, the last axis length of x is greater than or equal to 65536.
  		In the A8W4 scenario, the last axis length of x is greater than or equal to 20000.

aclnnGroupedMatmulSwigluQuantWeightNZ

• Parameters

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

• Return value

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

Constraints

• Deterministic computation:

• aclnnGroupedMatmulSwigluQuantWeightNZ defaults to a deterministic implementation.

• A8W8 scenario (A refers to the activation matrix (left matrix), W refers to the weight matrix (right matrix), and 8 refers to the INT8 data type)

• 1. The last axis length of x must not be greater than or equal to 65536.
  • 2. The N axis length cannot exceed 10240.

• A8W4 scenario (A refers to the activation matrix (left matrix), W refers to the weight matrix (right matrix), and 4 refers to the INT4 data type)

  • 1. The last axis length of x cannot be greater than or equal to 20000.
  • 2. The N axis length cannot exceed 10240.

Example

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

#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_grouped_matmul_swiglu_quant_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(const std::vector<T>& hostData, const std::vector<int64_t>& shape,
                    void** deviceAddr, aclDataType dataType, aclFormat formatType, 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 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);

    // Calculate 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, formatType,
                            shape.data(), shape.size(), *deviceAddr);
    return 0;
}

int main() {
    // 1. (Boilerplate) Initialize the device/stream. See the ACL API manual.
    // Fill in the 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 == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

    // 2. Construct the inputs and outputs. Customize the construction based on the API.
    int64_t E = 4;
    int64_t M = 8192;
    int64_t N = 4096;
    int64_t K = 7168;
    std::vector<int64_t> xShape = {M, K};
    std::vector<int64_t> weightShape = {E, N / 32 ,K / 16, 16, 32};
    std::vector<int64_t> weightScaleShape = {E, N};
    std::vector<int64_t> xScaleShape = {M};
    std::vector<int64_t> groupListShape = {E};
    std::vector<int64_t> outputShape = {M, N / 2};
    std::vector<int64_t> outputScaleShape = {M};

    void* xDeviceAddr = nullptr;
    void* weightDeviceAddr = nullptr;
    void* weightScaleDeviceAddr = nullptr;
    void* xScaleDeviceAddr = nullptr;
    void* groupListDeviceAddr = nullptr;
    void* outputDeviceAddr = nullptr;
    void* outputScaleDeviceAddr = nullptr;

    aclTensor* x = nullptr;
    aclTensor* weight = nullptr;
    aclTensor* weightScale = nullptr;
    aclTensor* xScale = nullptr;
    aclTensor* groupList = nullptr;
    aclTensor* output = nullptr;
    aclTensor* outputScale = nullptr;

    std::vector<int8_t> xHostData(M * K, 0);
    std::vector<int8_t> weightHostData(E * N * K, 0);
    std::vector<float> weightScaleHostData(E * N, 0);
    std::vector<float> xScaleHostData(M, 0);
    std::vector<int64_t> groupListHostData(E, 0);
    std::vector<int8_t> outputHostData(M * N / 2, 0);
    std::vector<float> outputScaleHostData(M, 0);

    // Create the x aclTensor.
    ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr,  aclDataType::ACL_INT8, aclFormat::ACL_FORMAT_ND, &x);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the weight aclTensor.
    ret = CreateAclTensor(weightHostData, weightShape, &weightDeviceAddr,  aclDataType::ACL_INT8, aclFormat::ACL_FORMAT_FRACTAL_NZ, &weight);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create weightScale aclTensor.
    ret = CreateAclTensor(weightScaleHostData, weightScaleShape, &weightScaleDeviceAddr, aclDataType::ACL_FLOAT,  aclFormat::ACL_FORMAT_ND, &weightScale);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create xScale aclTensor.
    ret = CreateAclTensor(xScaleHostData, xScaleShape, &xScaleDeviceAddr, aclDataType::ACL_FLOAT,  aclFormat::ACL_FORMAT_ND, &xScale);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create groupList aclTensor.
    ret = CreateAclTensor(groupListHostData, groupListShape, &groupListDeviceAddr, aclDataType::ACL_INT64, aclFormat::ACL_FORMAT_ND, &groupList);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create output aclTensor.
    ret = CreateAclTensor(outputHostData, outputShape, &outputDeviceAddr, aclDataType::ACL_INT8, aclFormat::ACL_FORMAT_ND, &output);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create outputScale aclTensor.
    ret = CreateAclTensor(outputScaleHostData, outputScaleShape, &outputScaleDeviceAddr, aclDataType::ACL_FLOAT, aclFormat::ACL_FORMAT_ND, &outputScale);
    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 aclnnGroupedMatmulSwigluQuantWeightNZ.
    ret = aclnnGroupedMatmulSwigluQuantWeightNZGetWorkspaceSize(x, weight, nullptr, nullptr, weightScale, xScale,
                                                        groupList, output, outputScale, nullptr,
                                                        &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS,
    LOG_PRINT("aclnnGroupedMatmulSwigluQuantWeightNZGetWorkspaceSize 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 aclnnGroupedMatmulSwigluQuantWeightNZ.
    ret = aclnnGroupedMatmulSwigluQuantWeightNZ(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS,
    LOG_PRINT("aclnnGroupedMatmulSwigluQuantWeightNZ 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. Modify based on the specific API definition.
    auto size = GetShapeSize(outputShape);
    std::vector<int8_t> out1Data(size, 0);
    ret = aclrtMemcpy(out1Data.data(), out1Data.size() * sizeof(out1Data[0]), outputDeviceAddr,
                        size * sizeof(out1Data[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, out1Data[j]);
    }
    size = GetShapeSize(outputScaleShape);
    std::vector<float> out2Data(size, 0);
    ret = aclrtMemcpy(out2Data.data(), out2Data.size() * sizeof(out2Data[0]), outputScaleDeviceAddr,
                        size * sizeof(out2Data[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: %f\n", j, out2Data[j]);
    }
    // 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
    aclDestroyTensor(x);
    aclDestroyTensor(weight);
    aclDestroyTensor(weightScale);
    aclDestroyTensor(xScale);
    aclDestroyTensor(groupList);
    aclDestroyTensor(output);
    aclDestroyTensor(outputScale);

    // 7. Release device resources. Modify based on the specific API definition.
    aclrtFree(xDeviceAddr);
    aclrtFree(weightDeviceAddr);
    aclrtFree(weightScaleDeviceAddr);
    aclrtFree(xScaleDeviceAddr);
    aclrtFree(groupListDeviceAddr);
    aclrtFree(outputDeviceAddr);
    aclrtFree(outputScaleDeviceAddr);
    if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}