aclnnGroupedMatmulSwigluQuant

📄 View Source Code

Applicable Products

Function

• Description: Fuses GroupedMatmul, dquant, swiglu, and quant. For details, see Formula.

• Formula:

  Quantization Scenario A8W8 (A refers to the activation matrix, W refers to the weight matrix, and 8 refers to the INT8 data type):

  • 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_{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_8}^{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.

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

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

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

  • Output:

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

    • $Q\_scale∈\mathbb{R}^{M}$: Quantization scaling factor.

  • Computation Process

    • 1. Determine the tokens of the current group based on groupList[i], where $i \in [0,Len(groupList))$.

      Example: Assume groupList = [3, 4, 4, 6].

      The 0th right matrix W[0,:,:] corresponds to tokens x[0:3] at index positions [0, 3) (3 – 0 = 3 tokens in total), and corresponds to x_scale[0:3], w_scale[0], Q[0:3], and Q_scale[0:3].

      The 1st right matrix W[1,:,:] corresponds to tokens x[3:4] at index positions [3, 4) (4 – 3 = 1 token in total), and corresponds to x_scale[3:4], w_scale[1], Q[3:4], and Q_scale[3:4].

      The 2nd right matrix W[2,:,:] corresponds to tokens x[4:4] at index positions [4, 4) (4 – 4 = 0 tokens in total), and corresponds to x_scale[4:4], w_scale[2], Q[4:4], and Q_scale[4:4].

      The 3rd right matrix W[3,:,:] corresponds to tokens x[4:6] at index positions [4, 6) (6 – 4 = 2 tokens in total), and corresponds to x_scale[4:6], w_scale[3], Q[4:6], and Q_scale[4:6].

      Note: The parts not specified in groupList will not participate in the update. For example, when groupList = [12, 14, 18] and the shape of x is [30, N/2]:

      The shape of the first output Q is [30, N/2], where the part Q[18:,:] will not be updated or initialized, and its data is the original data at the time of memory allocation.

      Similarly, the shape of the second output Q_scale is [30], where the part Q_scale[18:] will not be updated or initialized, and its data is the original data at the time of memory allocation.

      That is, Q[:groupList[-1],:] and Q_scale[:groupList[-1]] of the output are the valid data portions.

    • 2. Perform the following computation based on the input parameters determined by the group:

      $C_{i} = (X_{i}\cdot W_{i} )\odot x\_scale_{i\,\text{Broadcast}} \odot w\_scale_{i\,\text{Broadcast}}$

      $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$

Function Prototype

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

aclnnStatus aclnnGroupedMatmulSwigluQuantGetWorkspaceSize(
  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 aclnnGroupedMatmulSwigluQuant(
  void          *workspace,
  uint64_t       workspaceSize,
  aclOpExecutor *executor,
  aclrtStream    stream)

aclnnGroupedMatmulSwigluQuantGetWorkspaceSize

• Parameters

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	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, shaped as (M, K).	√
  weight	Input	Weight matrix, corresponding to W in the formula.	INT32 is used for adaptation purposes. One INT32 value is interpreted as eight INT4 values. In the A8W4 scenario: shape supports 3 or 5 dimensions, corresponding to the ND and FRACTAL_NZ data formats respectively. In the A8W8 scenario: shape supports 5 dimensions, and the data format supports FRACTAL_NZ.	INT8, INT4, INT32	ND, FRACTAL_NZ	3, 5	✓
  bias	Input	An auxiliary matrix used in matrix multiplication, corresponding to weightAssistMatrix in the formula.	Effective 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 offset in the formula.	Reserved input. Not supported yet. Pass a null pointer.	FLOAT	ND	2	-
  weightScale	Input	Quantization scale of the right matrix, corresponding to w_scale in the formula.	The length of the first axis must equal the first axis dimension of weight, and the length of the last axis must equal the last axis of weight restored to ND format. A8W4 scenario: Supports 2D or 3D shape, with UINT64 data type. A8W8 scenario: Supports 2D shape, with FLOAT, FLOAT16, or BFLOAT16 data type.	UINT64, FLOAT, FLOAT16, BFLOAT16	ND	2, 3	√
  xScale	Input	Quantization scale of the left matrix, corresponding to x_scale in the formula.	Its length must equal 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 calculation process in the Function section.	INT64	ND	1	√
  Output	Output	Indicates the output quantization result, corresponding to Q in the formula.	-	INT8	ND	2	√
  outputScale	Output	Output quantization factor, corresponding to Q_scale in the formula.	-	FLOAT	ND	1	√
  outputOffset	Output	Asymmetric quantization offset of the output, corresponding to Q_offset in the formula.	Reserved input, not supported yet. Pass a null pointer.	FLOAT	ND	1	-
  workspaceSize	Output	Returns the workspace size that the user needs to apply for 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 x, weight, weightScale, xScale, groupList, output, and outputScale passed in are null pointers.
  ACLNN_ERR_PARAM_INVALID	161002	The data dimensions of x, weight, weightScale, xScale, groupList, output, and outputScale passed in do not meet the constraints.
  		The shapes of x, weight, weightScale, xScale, groupList, output, and outputScale passed in do not meet the constraint conditions.
  		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 is greater than the length of the first axis of weight.
  		The length of the N axis exceeds 10240.
  		In the A8W8 scenario, the length of the last axis of x is greater than or equal to 65536.
  		In the A8W4 scenario, the length of the last axis of x is greater than or equal to 20000.
  		In the A8W8 scenario, the data format of weight is not FRACTAL_NZ.

aclnnGroupedMatmulSwigluQuant

• Parameters

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

• aclnnGroupedMatmulSwigluQuant defaults to a deterministic implementation.

• A8W8 scenario

• 1. The last axis length of x cannot be greater than or equal to 65536.
  • 2. The data format of weight only supports FRACTAL_NZ.
  • 3. The length of the N axis cannot exceed 10240.

• A8W4 scenario

  • 1. The length of the last axis of x cannot be greater than or equal to 20000.
  • 2. The length of the N axis 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.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: 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, aclFormat formatType, aclTensor** tensor) {
    auto size = GetShapeSize(shape) * sizeof(T);
    // Call aclrtMalloc to apply for 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 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 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 aclCreateTensor 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. Refer to the acl API manual.
    // Fill in the deviceId based on your actual device.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = Init(deviceId, &stream);
    // Handle the check as needed.
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

    // 2. Construct the inputs and outputs, which need to be custom-constructed 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 the weightScale aclTensor.
    ret = CreateAclTensor(weightScaleHostData, weightScaleShape, &weightScaleDeviceAddr, aclDataType::ACL_FLOAT,  aclFormat::ACL_FORMAT_ND, &weightScale);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the xScale aclTensor.
    ret = CreateAclTensor(xScaleHostData, xScaleShape, &xScaleDeviceAddr, aclDataType::ACL_FLOAT,  aclFormat::ACL_FORMAT_ND, &xScale);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the groupList aclTensor.
    ret = CreateAclTensor(groupListHostData, groupListShape, &groupListDeviceAddr, aclDataType::ACL_INT64, aclFormat::ACL_FORMAT_ND, &groupList);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the output aclTensor.
    ret = CreateAclTensor(outputHostData, outputShape, &outputDeviceAddr, aclDataType::ACL_INT8, aclFormat::ACL_FORMAT_ND, &output);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the 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 aclnnGroupedMatmulSwigluQuant.
    ret = aclnnGroupedMatmulSwigluQuantGetWorkspaceSize(x, weight, nullptr, nullptr, weightScale, xScale,
                                                        groupList, output, outputScale, nullptr,
                                                        &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS,
    LOG_PRINT("aclnnGroupedMatmulSwigluQuantGetWorkspaceSize 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 aclnnGroupedMatmulSwigluQuant.
    ret = aclnnGroupedMatmulSwigluQuant(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS,
    LOG_PRINT("aclnnGroupedMatmulSwigluQuant failed. ERROR: %d\n", ret); return ret);

    // 4. (Boilerplate) Wait for task execution to finish synchronously.
    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 side memory to the Host side. 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;
}