aclnnInplaceQuantMatmulAllReduceAddRmsNorm

NOTE

Applicable Products

Product	Supported
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	×

Note: When using this API, ensure that both the driver firmware package and the CANN package are of the matching 8.0.RC2 version or a matching later version. Otherwise, errors such as BUS ERROR may occur.

Function

Description: Completes the mm + all_reduce + add + rms_norm computation.
Formula:
$mm\_out = allReduce(dequantScale * (x1_{int8}@x2_{int8} + bias_{int32}))$ $y = mm\_out + residual$ $normOut = \frac{y}{RMS(y)} * gamma, RMS(y) = \sqrt{\frac{1}{d} \sum_{i=1}^{d} y_{i}^{2} + epsilon}$

Function Prototype

aclnnQuantMatmulAllReduceAddRmsNorm and aclnnInplaceQuantMatmulAllReduceAddRmsNorm implement the same functionality. The differences in usage are as follows. Select the appropriate operator based on your actual scenario.
- aclnnQuantMatmulAllReduceAddRmsNorm: Two output tensors, normOut and y, need to be created to store the computation results.
- aclnnInplaceQuantMatmulAllReduceAddRmsNorm: One output tensor normOut needs to be created. In non-in-place scenarios, the result stored in the newly created output tensor y is directly stored into the memory of the input tensor residual.
Each operator is divided into a two-phase API. You must first call the "aclnnInplaceQuantMatmulAllReduceAddRmsNormGetWorkspaceSize" API to obtain the workspace size required for computation and the executor that contains the operator computation process, and then call the "aclnnInplaceQuantMatmulAllReduceAddRmsNorm" API to perform the computation.

aclnnStatus aclnnInplaceQuantMatmulAllReduceAddRmsNormGetWorkspaceSize(
    const aclTensor *x1,
    const aclTensor *x2,
    const aclTensor *bias,
    const aclTensor *dequantScale,
    const aclTensor *residual,
    const aclTensor *gamma,
    double           epsilon,
    const char      *group,
    const char      *reduceOp,
    int64_t          commTurn,
    int64_t          streamMode,
    const aclTensor *normOut,
    uint64_t        *workspaceSize,
    aclOpExecutor  **executor)

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

aclnnInplaceQuantMatmulAllReduceAddRmsNormGetWorkspaceSize

Parameters:

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Shape	Non-contiguous Tensor
x1 (aclTensor)	Input	Left matrix for MatMul computation, i.e., x1 in the formula.	Supports empty tensors. Must have the same data type as x2. In the current version, only two-dimensional or three-dimensional inputs are supported.	INT8	ND	2-3	×
x2 (aclTensor)	Input	Right matrix of the MatMul computation, i.e., x2 in the formula.	Supports empty tensors. Must have the same data type as x1. In the current version, only two-dimensional input is supported, and both transposed and non-transposed scenarios are supported. Supports non-contiguous tensors in transposed scenarios.	INT8	ND	2	√
bias (aclTensor)	Input	Refers to bias in the formula.	Supports passing a null pointer. In the current version, only one-dimensional input is supported.	INT32	ND	1	√
dequantScale (aclTensor)	Input	Full quantization coefficient after MatMul computation, i.e., dequantScale in the formula.	Shape is (1) in the per-tensor scenario, and (n) or (1, n) in the per-channel scenario.	UINT64, INT64, BFLOAT16	ND	1-2	×
residual (aclTensor)	Input	Residual input of the AddRmsNorm fusion operator, i.e., residual in the formula.	In the in-place scenario, residual is used as the output address of y. The current version supports only three-dimensional input.	FLOAT16, BFLOAT16	ND	3	×
gamma (aclTensor)	Input	The RMS normalization computation input of the AddRmsNorm fused operator, i.e., gamma in the formula.	The current version supports only one-dimensional input.	FLOAT16, BFLOAT16	ND	1	×
epsilon (double)	Input	A double-precision value on the host side, used to prevent division-by-zero errors, i.e., epsilon in the formula.	The epsilon value falls within the range (0, 1).	Double	-	-	-
group (char)	Input	A string that identifies the communication domain on the host side. Specifies the communication domain name.	Obtained through the Hccl interface "extern HcclResult HcclGetCommName(HcclComm comm, char* commName);", where commName is the group.	String	-	-	-
reduceOp (char)	Input	A string on the host side that identifies the operation type, i.e., the reduce operation type.	Currently, only "sum" is supported.	String	-	-	-
commTurn (int64_t)	Input	Communication data split count, i.e., total data volume divided by single communication volume.	Only 0 is supported in the current version.	INT64	-	-	-
streamMode (int64_t)	Input	Enumeration of the stream mode.	Currently, only enumeration value 1 is supported.	INT64	-	-	-
normOut (aclTensor)	Output	Result of matrix multiplication + all-reduce + addition + RMS normalization, i.e., normOut in the formula.	Empty tensors are not supported. The data type is the same as that of the residual input.	FLOAT16, BFLOAT16	ND	2	√
workspaceSize (uint64_t)	Output	Returns the workspace size to be requested on the Device side.	-	-	-	-	-
executor (aclOpExecutor)	output	Returns the operator 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:

<table style="undefined;table-layout: fixed; width: 1030px"><colgroup>
<col style="width: 250px">
<col style="width: 130px">
<col style="width: 650px">
</colgroup>
<thead>
<tr>
    <th>Return Value</th>
    <th>Error Code</th>
    <th>Description</th>
</tr></thead>
<tbody>
<tr>
    <td>ACLNN_ERR_PARAM_NULLPTR</td>
    <td>161001</td>
    <td>x1, x2, residual, gamma, or dequantScale is a null pointer.</td>
</tr>
<tr>
    <td rowspan="3">ACLNN_ERR_PARAM_INVALID</td>
    <td rowspan="3">161002</td>
    <td>The shapes of x1, x2, and residual do not meet the constraint requirements.</td>
</tr>
<tr>
    <td>streamMode is not within the valid range.</td>
</tr>
<tr>
    <td>The size k of the last dimension of x1 is 0.</td>
</tr>
</tbody>
</table>

aclnnInplaceQuantMatmulAllReduceAddRmsNorm

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 from the first-phase API aclnnInplaceQuantMatmulAllReduceAddRmsNormGetWorkspaceSize.
executor	Input	Operator executor, which contains the operator computation flow.
stream	Input	Specifies the stream for executing the task.

Return value

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

Constraints

Deterministic computation:
- aclnnInplaceQuantMatmulAllReduceAddRmsNorm defaults to a non-deterministic implementation. Deterministic computation can be enabled by setting the HCCL_DETERMINISTIC environment variable to true.
The usage scenario is consistent with the fused operator aclnnQuantMatmulAllReduce: MC2 is disabled in incremental scenarios and enabled in full scenarios.
The input x1 can be two-dimensional or three-dimensional, with its shape being (b, s, k) or (m, k). x2 must be two-dimensional, with its shape being (k, n). The axes must satisfy the mm operator input requirements, with the k axes being equal. If bias is not null, bias is one-dimensional with its shape being (n).
The size of m does not exceed 2147483647. The size of the last dimension of x1 and x2 does not exceed 65535. The last dimension of x1 refers to k, and the last dimension of x2 refers to k during transpose or n during non-transpose.
The input residual must be three-dimensional with shape (b, s, n). When x1 is two-dimensional, (b*s) of residual equals m of x1. The input gamma has shape (n).
The output y and normOut have the same dimensions and data type as residual, with shape (b, s, n).
If the output residual type is FLOAT16, the type of dequantScale is UINT64 or INT64. If the output residual type is BFLOAT16, the type of dequantScale is BFLOAT16.
The data types of x1 and x2 are INT8. The data type of bias is INT32. The data types of residual, gamma, y, and normOut computation inputs must be consistent.
Only x2 matrix transposition/non-transposition is supported, and x1 matrix supports the non-transposed scenario.
1, 2, 4, or 8 cards are supported, and only HCCS link all-mesh networking is supported.
Empty tensors with (b*s) or n being 0 are supported, but empty tensors with k being 0 are not supported.
Atlas A2 training products/Atlas A2 inference products : In a model, the communication-computation fused MC2 operator supports only the same communication domain.

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 <thread>
#include <cstdio>
#include "hccl/hccl.h"
#include "aclnnop/aclnn_inplace_quant_matmul_all_reduce_add_rms_norm.h"

int ndev = 8;

#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;
}

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 device-side 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);
    // Copy host-side data to device-side memory using aclrtMemcpy.
    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, aclFormat::ACL_FORMAT_ND,
                            shape.data(), shape.size(), *deviceAddr);
    return 0;
}

struct Args {
    uint32_t rankId;
    HcclComm hcclComm;
    aclrtStream stream;
    aclrtContext context;
};

int launchOneThreadQuantMatmulAllReduceAddRmsNorm(Args &args) {
    int ret;
    ret = aclrtSetCurrentContext(args.context);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetCurrentContext failed. ERROR: %d\n", ret); return ret);
    char hcom_name[128];
    ret = HcclGetCommName(args.hcclComm, hcom_name);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetCommName failed. ret = %d \n", ret); return -1);
    LOG_PRINT("[INFO] rank %d hcom: %s stream: %p, context : %p\n", args.rankId, hcom_name, args.stream,
            args.context);

    std::vector<int64_t> x1Shape = {32, 64};
    std::vector<int64_t> x2Shape = {64, 128};
    std::vector<int64_t> biasShape = {128};
    std::vector<int64_t> dequantScaleShape = {128};
    std::vector<int64_t> residualShape = {1, 32, 128};
    std::vector<int64_t> gammaShape = {128};
    std::vector<int64_t> yShape = {1, 32, 128};
    std::vector<int64_t> normOutShape = {1, 32, 128};
    void *x1DeviceAddr = nullptr;
    void *x2DeviceAddr = nullptr;
    void *biasDeviceAddr = nullptr;
    void *dequantScaleDeviceAddr = nullptr;
    void *residualDeviceAddr = nullptr;
    void *gammaDeviceAddr = nullptr;
    void *yDeviceAddr = nullptr;
    void *normOutDeviceAddr = nullptr;
    aclTensor *x1 = nullptr;
    aclTensor *x2 = nullptr;
    aclTensor *bias = nullptr;
    aclTensor *dequantScale = nullptr;
    aclTensor *residual = nullptr;
    aclTensor *gamma = nullptr;
    aclTensor *y = nullptr;
    aclTensor *normOut = nullptr;

    int64_t commTurn = 0;
    int64_t streamMode = 1;
    double  epsilon = 0.000001;
    uint64_t workspaceSize = 0;
    aclOpExecutor *executor;
    void *workspaceAddr = nullptr;

    long long x1ShapeSize = GetShapeSize(x1Shape);
    long long x2ShapeSize = GetShapeSize(x2Shape);
    long long biasShapeSize = GetShapeSize(biasShape);
    long long dequantScaleShapeSize = GetShapeSize(dequantScaleShape);
    long long residualShapeSize = GetShapeSize(residualShape);
    long long gammaShapeSize = GetShapeSize(gammaShape);
    long long yShapeSize = GetShapeSize(yShape);
    long long normOutShapeSize = GetShapeSize(normOutShape);

    std::vector<int8_t> x1HostData(x1ShapeSize, 1);
    std::vector<int8_t> x2HostData(x2ShapeSize, 1);
    std::vector<int32_t> biasHostData(biasShapeSize, 1);
    std::vector<uint64_t> dequantScaleHostData(dequantScaleShapeSize, 1);
    std::vector<int16_t> residualHostData(residualShapeSize, 1);
    std::vector<int16_t> gammaHostData(gammaShapeSize, 1);
    std::vector<int16_t> yHostData(yShapeSize, 0);
    std::vector<int16_t> normOutHostData(normOutShapeSize, 0);
    // Create a tensor.
    ret = CreateAclTensor(x1HostData, x1Shape, &x1DeviceAddr, aclDataType::ACL_INT8, &x1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_INT8, &x2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(biasHostData, biasShape, &biasDeviceAddr, aclDataType::ACL_INT32, &bias);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(dequantScaleHostData, dequantScaleShape, &dequantScaleDeviceAddr,
                        aclDataType::ACL_UINT64, &dequantScale);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(residualHostData, residualShape, &residualDeviceAddr, aclDataType::ACL_FLOAT16, &residual);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(gammaHostData, gammaShape, &gammaDeviceAddr, aclDataType::ACL_FLOAT16, &gamma);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(yHostData, yShape, &yDeviceAddr, aclDataType::ACL_FLOAT16, &y);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(normOutHostData, normOutShape, &normOutDeviceAddr, aclDataType::ACL_FLOAT16, &normOut);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // Example of calling the aclnnInplaceQuantMatmulAllReduceAddRmsNorm API.
    // Call the first-phase API.
    ret = aclnnInplaceQuantMatmulAllReduceAddRmsNormGetWorkspaceSize(x1, x2, bias, dequantScale, residual, gamma, epsilon,
                                        hcom_name, "sum", commTurn, streamMode, normOut,
                                        &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS,
            LOG_PRINT("aclnnInplaceQuantMatmulAllReduceAddRmsNormGetWorkspaceSize 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.
    ret = aclnnInplaceQuantMatmulAllReduceAddRmsNorm(workspaceAddr, workspaceSize, executor, args.stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnInplaceQuantMatmulAllReduceAddRmsNorm failed. ERROR: %d\n", ret); return ret);
    //(Boilerplate) Synchronously wait for task execution to complete.
    ret = aclrtSynchronizeStreamWithTimeout(args.stream, 10000);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
    LOG_PRINT("device%d aclnnInplaceQuantMatmulAllReduceAddRmsNorm execute success \n", args.rankId);
    // Release device resources. Modify based on the specific API definition.

    if (x1 != nullptr) {
        aclDestroyTensor(x1);
    }
    if (x2 != nullptr) {
        aclDestroyTensor(x2);
    }
    if (bias != nullptr) {
        aclDestroyTensor(bias);
    }
    if (dequantScale != nullptr) {
        aclDestroyTensor(dequantScale);
    }
    if (residual != nullptr) {
        aclDestroyTensor(residual);
    }
    if (gamma != nullptr) {
        aclDestroyTensor(gamma);
    }
    if (y != nullptr) {
        aclDestroyTensor(y);
    }
    if (normOut != nullptr) {
        aclDestroyTensor(normOut);
    }
    if (x1DeviceAddr != nullptr) {
        aclrtFree(x1DeviceAddr);
    }
    if (x2DeviceAddr != nullptr) {
        aclrtFree(x2DeviceAddr);
    }
    if (biasDeviceAddr != nullptr) {
        aclrtFree(biasDeviceAddr);
    }
    if (dequantScaleDeviceAddr != nullptr) {
        aclrtFree(dequantScaleDeviceAddr);
    }
    if (residualDeviceAddr != nullptr) {
        aclrtFree(residualDeviceAddr);
    }
    if (gammaDeviceAddr != nullptr) {
        aclrtFree(gammaDeviceAddr);
    }
    if (yDeviceAddr != nullptr) {
        aclrtFree(yDeviceAddr);
    }
    if (normOutDeviceAddr != nullptr) {
        aclrtFree(normOutDeviceAddr);
    }
    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(args.stream);
    HcclCommDestroy(args.hcclComm);
    aclrtDestroyContext(args.context);
    aclrtResetDevice(args.rankId);
    return 0;
}

int main(int argc, char *argv[]) {
    int ret;
    int32_t devices[ndev];
    for (int i = 0; i < ndev; i++) {
        devices[i] = i;
    }
    HcclComm comms[128];
    ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
    // Initialize the collective communication domain.
    for (int i = 0; i < ndev; i++) {
        ret = aclrtSetDevice(devices[i]);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
    }
    ret = HcclCommInitAll(ndev, devices, comms);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("HcclCommInitAll failed. ERROR: %d\n", ret); return ret);
    Args args[ndev];
    aclrtStream stream[ndev];
    aclrtContext context[ndev];
    for (uint32_t rankId = 0; rankId < ndev; rankId++) {
        ret = aclrtSetDevice(rankId);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
        ret = aclrtCreateContext(&context[rankId], rankId);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateContext failed. ERROR: %d\n", ret); return ret);
        ret = aclrtCreateStream(&stream[rankId]);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
    }
    // Start multithreading.
    std::vector<std::unique_ptr<std::thread>> threads(ndev);
    for (uint32_t rankId = 0; rankId < ndev; rankId++) {
        args[rankId].rankId = rankId;
        args[rankId].hcclComm = comms[rankId];
        args[rankId].stream = stream[rankId];
        args[rankId].context = context[rankId];
        threads[rankId].reset(
                new(std::nothrow) std::thread(&launchOneThreadQuantMatmulAllReduceAddRmsNorm, std::ref(args[rankId])));
    }
    for (uint32_t rankId = 0; rankId < ndev; rankId++) {
        threads[rankId]->join();
    }
    aclFinalize();
    return 0;
}