aclnnQuantMatmulAllReduceV3

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 later matching version. Otherwise, errors such as BUS ERROR may occur.

Function

Description: The aclnnQuantMatmulAllReduceV3 API extends the functionality of the aclnnQuantMatmulAllReduceV2 API by adding support for low-bit communication: the MatMul computation result undergoes AllToAll communication, ReduceSum computation, AllGather communication, and dequant dequantization in sequence, replacing the original flow of dequant and pertoken computation followed by AllReduce communication. It supports pertensor, perchannel, and pertoken quantization modes.
Formula:

into the following three scenarios:
- Scenario 1: After performing MatMul on the quantized inputs x1 and x2, dequantization is performed, followed by an Add operation with x3, and finally an AllReduce computation.
  $output= AllReduce(dequantScale*(x1_{int8}@x2_{int8} + bias_{int32}) + x3)$
- Scenario 2: After performing MatMul on the quantized inputs x1 and x2, dequantization and per-token computation are performed, followed by an Add operation with x3, and finally an AllReduce computation.
  $output= AllReduce(dequantScale * pertokenScaleOptional * (x1_{int8}@x2_{int8} + biasOptional_{int32}) + x3Optional)$
- Scenario 3: Perform MatMul, dequant, and pertoken computation on the quantized inputs x1 and x2, then perform an Add operation with x3, apply perchannel quantization to the output, then perform AllToAll communication, perform ReduceSum computation on the first communication result, then perform AllGather communication, and finally dequant the second communication result to obtain the final output.
  $matmulAddOutPut = (dequantScale * pertokenScaleOptional * (x1_{int8}@x2_{int8} + biasOptional_{int32}) + x3Optional);$ $alltoallOutPut_{int8} = AllToAll(matmulAddOutPut / commQuantScale1Optional);$ $reduceSumOutPut_{int8} = (add(alltoallOutPut_{int8}) * (commQuantScale1Optional / commQuantScale2Optional));$ $outPut = (AllGather(reduceSumOutPut_{int8}) * commQuantScale2Optional);$

Function Prototype

Each operator uses a two-phase API. You must first call the aclnnQuantMatmulAllReduceV3GetWorkspaceSize API to obtain the workspace size required for computation and an executor that encapsulates the operator computation process, and then call the aclnnQuantMatmulAllReduceV3 API to perform the computation.

aclnnStatus aclnnQuantMatmulAllReduceV3GetWorkspaceSize(
    const aclTensor *x1,
    const aclTensor *x2,
    const aclTensor *biasOptional,
    const aclTensor *x3Optional,
    const aclTensor *dequantScale,
    const aclTensor *pertokenScaleOptional,
    const aclTensor *commQuantScale1Optional,
    const aclTensor *commQuantScale2Optional,
    const char      *group,
    const char      *reduceOp,
    int64_t          commTurn,
    int64_t          streamMode,
    const aclTensor *output,
    uint64_t        *workspaceSize,
    aclOpExecutor  **executor)

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

aclnnQuantMatmulAllReduceV3GetWorkspaceSize

Parameters

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (shape)	Non-Contiguous Tensor
x1	Input	Left matrix for MatMul computation, i.e., x1 in the formula.	Only two-dimensional or three-dimensional input is supported in the current version. Supports the non-transpose scenario.	INT8	ND	2-3	×
x2	Input	Right matrix of the MatMul computation, i.e., x2 in the formula.	Supports transpose/non-transpose scenarios. In ND format, supports non-contiguous tensors with the last two axes transposed. Other non-contiguous tensors are not supported.	INT8	ND, FRACTAL_NZ	2	×
biasOptional	Input	Bias offset, i.e., biasOptional in the formula.	Only one-dimensional input is supported in the current version.	INT32	ND	1	√
x3Optional	Input	The add operation after MatMul computation, i.e., x3Optional in the formula.	The dimensions are consistent with output. Currently, only the BFLOAT16 output scenario is supported, and only non-empty input is supported.	FLOAT16, BFLOAT16	ND	2	√
dequantScale	Input	Dequantization coefficient after MatMul computation, i.e., dequantScale in the formula.	The shape is (1) in the pertensor scenario, and (n)/(1, n) in the perchannel scenario. When the output is BFLOAT16, directly pass the dequantScale of BFLOAT16 type to this API. When the output is FLOAT16, if pertokenScale is not empty, you can directly pass the dequantScale of FLOAT32 type to this API; if pertokenScale is empty, you need to call the aclnn API of the TransQuantParamV2 operator in advance to convert dequantScale to the INT64/UINT64 data type.	INT64, UINT64, FLOAT32, BFLOAT16	ND	2	√
pertokenScaleOptional	Input	Pertoken dequantization coefficient after MatMul computation, i.e., pertokenScaleOptional in the formula.	When x1 is (b, s, k), the shape is (b*s); when x1 is (m, k), the shape is (m).	FLOAT32	ND	2	√
commQuantScale1Optional	Input	Per-channel quantization scale after the matmulAdd computation, i.e., commQuantScale1Optional in the formula.	When x2 is (k, n), the shape can be (n) or (1, n).	BFLOAT16, FLOAT16	ND	2	√
commQuantScale2Optional	Input	Per-channel quantization scale after the AllGather computation, i.e., commQuantScale2Optional in the formula.	When x2 is (k, n), the shape can be (n) or (1, n).	BFLOAT16, FLOAT16	ND	2	√
group	Input	Name of the communication domain.	Obtained through the Hccl interface "extern HcclResult HcclGetCommName(HcclComm comm, char* commName);", where commName is the group.	String	-	-	-
reduceOp	Input	Specifies the reduce operation type.	Currently, only "sum" is supported.	String	-	-	-
commTurn	Input	Number of communication data splits, i.e., total data volume divided by single communication volume.	The current version only supports input 0.	INT64	-	-	-
streamMode	Input	Enumeration of stream modes.	Currently only the enumeration value 1 is supported.	INT64	-	-	-
output	Output	Result of the MatMul computation and AllReduce communication, i.e., the output in the formula.	The number of dimensions of output is the same as that of x1.	FLOAT16, BFLOAT16	ND	2-3	√
workspaceSize	Output	Returns the workspace size to be requested on the Device side.	-	-	-	-	-
executor	Output	Returns the operator executor, which contains the operator computation flow.	-	-	-	-	-

Atlas A2 training products/Atlas A2 inference products:
- The data format of input x2 supports ND and FRACTAL_NZ formats. The shape rules for the input are as follows:
  - When the data format of x2 is FRACTAL_NZ, the current version only supports four-dimensional input. Use aclnnCalculateMatmulWeightSizeV2 and aclnnTransMatmulWeight to complete the conversion from ND to NZ. Non-contiguous tensors only support the transpose scenario.
  - When the data format of x2 is ND, the current version only supports two-dimensional input.
Ascend 950PR/Ascend 950DT:
- The data format of input x2 only supports the ND format (the current version only supports two-dimensional input).

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: 1149px"><colgroup>
<col style="width: 282px">
<col style="width: 120px">
<col style="width: 747px">
</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>The passed x1, x2, dequantScale, or output is a null pointer.</td>
</tr>
<tr>
    <td rowspan="3">ACLNN_ERR_PARAM_INVALID</td>
    <td rowspan="3">161002</td>
    <td>The data type of x1, x2, biasOptional, dequantScale, pertokenScaleOptional, x3Optional, commQuantScale1Optional, commQuantScale2Optional, or output is not within the supported range.</td>
</tr>
<tr>
    <td>streamMode is not within the valid range.</td>
</tr>
<tr>
    <td>The shape of x1, x2, biasOptional, dequantScale, pertokenScaleOptional, x3Optional, commQuantScale1Optional, commQuantScale2Optional, or output does not meet the constraint requirements.</td>
</tr>
</tbody>
</table>

aclnnQuantMatmulAllReduceV3

Parameters

Parameter	Input/Output	Description
workspace	Input	Workspace memory address applied for on the Device side.
workspaceSize	Input	Workspace size applied for on the Device side, obtained from the first-phase API `aclnnQuantMatmulAllReduceV3GetWorkspaceSize`.
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:
- Atlas A2 training products/Atlas A2 inference products: aclnnQuantMatmulAllReduceV3 defaults to a non-deterministic implementation. Deterministic computation can be enabled by setting the HCCL_DETERMINISTIC environment variable to true.
- Ascend 950PR/Ascend 950DT: aclnnQuantMatmulAllReduceV3 defaults to a deterministic implementation.
MC2 is disabled in incremental scenarios and enabled in full scenarios.
The input x1 can be 2D or 3D and must not be an empty tensor, with shape (b, s, k) or (m, k). x2 must be 2D and must not be an empty tensor, with shape (k, n). The k axes must satisfy the input requirements of the mm operator, meaning the k axes must be equal.
The size of m must not exceed 2147483647. The size of the last dimension of x1 and x2 must not exceed 65535. The last dimension of x1 refers to k, and the last dimension of x2 refers to k when transposed or n when not transposed.
The passed x1, x2, dequantScale, or output must not be null pointers.
The data types and data formats of x1, x2, dequantScale, output, biasOptional (when not null), and x3 (when not null) must fall within the supported ranges.
If the output type is FLOAT16: when pertokenScale is null, the type of dequantScale is INT64 or UINT64; when pertokenScale is not null, the type of dequantScale is FLOAT32. If the output type is BFLOAT16, the type of dequantScale is BFLOAT16 and the type of x3 is BFLOAT16.
The commQuantScale1 and commQuantScale2 passed in must both be null pointers or both be non-null pointers. If both commQuantScale1 and commQuantScale2 are non-null pointers, the shapes of the two quantization parameters must be consistent, their types must be consistent with the operator output type, and the input on each card must be consistent.
When the shape of x1 is (b, s, k), the shape of pertokenScaleOptional is (b*s); when the shape of x1 is (m, k), the shape of pertokenScaleOptional is (m).
Only HCCS all-mesh networking is supported.
- Atlas A2 training products/Atlas A2 inference products: 1, 2, 4, and 8 cards are supported.
- Ascend 950PR/Ascend 950DT: 1, 2, 4, 8, 16, 32, and 64 cards are supported.
Atlas A2 training products/Atlas A2 inference products: The computation-communication fusion MC2 operator in a model supports only the same communication domain.
int8 low-bit communication provides performance benefits only in communication-bound scenarios. It is not recommended to enable int8 low-bit communication in computation-bound scenarios, meaning that it is not recommended to input commQuantScale1 and commQuantScale2.
Empty tensor support:
- Empty tensors are not supported.

Example

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

NOTE The sample code in this example calls some HCCL collective communication library APIs: HcclGetCommName, HcclCommInitAll, and HcclCommDestroy. For details, see HCCL API (C).

Atlas A2 training products/Atlas A2 inference products , Ascend 950PR/Ascend 950DT:

#include <iostream>
#include <vector>
#include <thread>
#include "hccl/hccl.h"
#include "aclnn/opdev/fp16_t.h"
#include "aclnnop/aclnn_trans_matmul_weight.h"
#include "aclnnop/aclnn_quant_matmul_all_reduce_v3.h"

int ndev = 2;

#define ACL_CHECK(ret)                                                                                     \
    do {                                                                                                   \
        auto retcode = ret;                                                                                \
        if (retcode != ACL_SUCCESS) {                                                                      \
            printf("[ERROR] acl interface return err %s:%d, retcode: %d \n", __FILE__, __LINE__, retcode); \
            return retcode;                                                                                \
        }                                                                                                  \
    } while (0)

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

struct Args {
    uint32_t rankId;
    HcclComm hcclComm;
    aclrtStream stream;
    aclrtContext context;
    std::string format;
  };

template<typename T>
int CreateWeightNzAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr,
                            aclDataType dataType, aclTensor **tensor, Args &args) {
    auto size = GetShapeSize(shape) * sizeof(T);
    const aclIntArray *mat2Size = aclCreateIntArray(shape.data(), shape.size());
    auto ret = aclnnCalculateMatmulWeightSizeV2(mat2Size, ACL_INT8, &size);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCalculateMatmulWeightSizeV2 failed. ERROR: %d\n", ret); return    ret);

    // Call aclrtMalloc to apply for device memory.
    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 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, aclFormat::ACL_FORMAT_ND,
                              shape.data(), shape.size(), *deviceAddr);

    uint64_t transWorkspaceSize;
    aclOpExecutor *executor;
    void *transWorkspaceAddr = nullptr;
    ret = aclnnTransMatmulWeightGetWorkspaceSize(*tensor, &transWorkspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS && transWorkspaceSize > 0,
              printf("[ERROR] aclnnTransMatmulWeightGetWorkspaceSize failed. ret = %d \n", ret); return ret);
    ACL_CHECK(aclrtMalloc(&transWorkspaceAddr, transWorkspaceSize, ACL_MEM_MALLOC_HUGE_FIRST));
    ret = aclnnTransMatmulWeight(transWorkspaceAddr, transWorkspaceSize, executor, args.stream);
    CHECK_RET(ret == ACL_SUCCESS, printf("[ERROR] aclnnTransMatmulWeight failed. ret = %d \n", ret);return ret);
    ACL_CHECK(aclrtSynchronizeStreamWithTimeout(args.stream, 20000));

    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);
    // Allocate device memory by calling aclrtMalloc.
    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 data to device memory by calling 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 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];
    }
    // Create an aclTensor by calling aclCreateTensor.
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                              shape.data(), shape.size(), *deviceAddr);
    return 0;
}

int launchOneThreadQuantMatmulAllReduce(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> pertokenScaleShape = {32};
    std::vector<int64_t> commQuantScale1Shape = {128};
    std::vector<int64_t> commQuantScale2Shape = {128};
    std::vector<int64_t> x3Shape = {32, 128};
    std::vector<int64_t> outShape = {32, 128};
    void *x1DeviceAddr = nullptr;
    void *x2DeviceAddr = nullptr;
    void *biasDeviceAddr = nullptr;
    void *dequantScaleDeviceAddr = nullptr;
    void *pertokenScaleDeviceAddr = nullptr;
    void *commQuantScale1DeviceAddr = nullptr;
    void *commQuantScale2DeviceAddr = nullptr;
    void *x3DeviceAddr = nullptr;
    void *outDeviceAddr = nullptr;
    aclTensor *x1 = nullptr;
    aclTensor *x2 = nullptr;
    aclTensor *bias = nullptr;
    aclTensor *dequantScale = nullptr;
    aclTensor *pertokenScale = nullptr;
    aclTensor *commQuantScale1 = nullptr;
    aclTensor *commQuantScale2 = nullptr;
    aclTensor *x3 = nullptr;
    aclTensor *out = nullptr;

    int64_t commTurn = 0;
    int64_t streamMode = 1;
    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 pertokenScaleShapeSize = GetShapeSize(pertokenScaleShape);
    long long commQuantScale1ShapeSize = GetShapeSize(commQuantScale1Shape);
    long long commQuantScale2ShapeSize = GetShapeSize(commQuantScale2Shape);
    long long x3ShapeSize = GetShapeSize(x3Shape);
    long long outShapeSize = GetShapeSize(outShape);

    std::vector<int8_t> x1HostData(x1ShapeSize, 1);
    std::vector<int8_t> x2HostData(x2ShapeSize, 1);
    std::vector<int32_t> biasHostData(biasShapeSize, 1);
    std::vector<float> dequantScaleHostData(dequantScaleShapeSize, 1);
    std::vector<float> pertokenScaleHostData(pertokenScaleShapeSize, 1);
    std::vector<op::fp16_t> commQuantScale1HostData(commQuantScale1ShapeSize, 1);
    std::vector<op::fp16_t> commQuantScale2HostData(commQuantScale2ShapeSize, 1);
    std::vector<op::fp16_t> x3HostData(x3ShapeSize, 1);
    std::vector<op::fp16_t> outHostData(outShapeSize, 0);
    // Create a tensor.
    ret = CreateAclTensor(x1HostData, x1Shape, &x1DeviceAddr, aclDataType::ACL_INT8, &x1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    if (args.format == "NZ") {
        ret = CreateWeightNzAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_INT8, &x2, args);
    } else {
        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_FLOAT, &dequantScale);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(pertokenScaleHostData, pertokenScaleShape, &pertokenScaleDeviceAddr,
                          aclDataType::ACL_FLOAT, &pertokenScale);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(commQuantScale1HostData, commQuantScale1Shape, &commQuantScale1DeviceAddr,
                          aclDataType::ACL_FLOAT16, &commQuantScale1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(commQuantScale2HostData, commQuantScale2Shape, &commQuantScale2DeviceAddr,
                          aclDataType::ACL_FLOAT16, &commQuantScale2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(x3HostData, x3Shape, &x3DeviceAddr, aclDataType::ACL_FLOAT16, &x3);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT16, &out);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Call the first-phase API.
    ret = aclnnQuantMatmulAllReduceV3GetWorkspaceSize(x1, x2, bias, x3, dequantScale, pertokenScale,
                                                      commQuantScale1, commQuantScale2, hcom_name,
                                                      "sum", commTurn, streamMode, out,
                                                      &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS,
              LOG_PRINT("aclnnQuantMatmulAllReduceV3GetWorkspaceSize 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 = aclnnQuantMatmulAllReduceV3(workspaceAddr, workspaceSize, executor, args.stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnQuantMatmulAllReduceV3 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 aclnnQuantMatmulAllReduceV3 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 (pertokenScale != nullptr) {
        aclDestroyTensor(pertokenScale);
    }
    if (commQuantScale1 != nullptr) {
        aclDestroyTensor(commQuantScale1);
    }
    if (commQuantScale2 != nullptr) {
        aclDestroyTensor(commQuantScale2);
    }
    if (x3 != nullptr) {
        aclDestroyTensor(x3);
    }
    if (out != nullptr) {
        aclDestroyTensor(out);
    }
    if (x1DeviceAddr != nullptr) {
        aclrtFree(x1DeviceAddr);
    }
    if (x2DeviceAddr != nullptr) {
        aclrtFree(x2DeviceAddr);
    }
    if (biasDeviceAddr != nullptr) {
        aclrtFree(biasDeviceAddr);
    }
    if (dequantScaleDeviceAddr != nullptr) {
        aclrtFree(dequantScaleDeviceAddr);
    }
    if (pertokenScaleDeviceAddr != nullptr) {
        aclrtFree(pertokenScaleDeviceAddr);
    }
    if (commQuantScale1DeviceAddr != nullptr) {
        aclrtFree(commQuantScale1DeviceAddr);
    }
    if (commQuantScale2DeviceAddr != nullptr) {
        aclrtFree(commQuantScale2DeviceAddr);
    }
    if (x3DeviceAddr != nullptr) {
        aclrtFree(x3DeviceAddr);
    }
    if (outDeviceAddr != nullptr) {
        aclrtFree(outDeviceAddr);
    }
    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(&launchOneThreadQuantMatmulAllReduce, std::ref(args[rankId])));
    }
    for (uint32_t rankId = 0; rankId < ndev; rankId++) {
        threads[rankId]->join();
    }
    aclFinalize();
    return 0;
}