aclnnQuantMatmulAllReduceV2

📄 View Source Code

Applicable Products

ProductSupported
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 the driver firmware package and the CANN package are both of the matching 8.0.RC2 version or a matching later version. Otherwise, errors such as BUS ERROR will occur.

Function

  • Description: The aclnnQuantMatmulAllReduceV2 API is a functional extension of the aclnnQuantMatmulAllReduce API, adding support for pertoken quantization. aclnnQuantMatmulAllReduceV2 supports pertensor, perchannel, and pertoken quantization modes.

  • Formula:

    There are two cases:

    • Case 1: After MatMul computation on the quantized input parameters x1 and x2, dequant computation is performed, followed by an Add operation with x3, and finally an AllReduce computation.

      output=AllReduce(dequantScale∗(x1int8@x2int8+biasint32)+x3)output= AllReduce(dequantScale*(x1_{int8}@x2_{int8} + bias_{int32}) + x3)
    • Case 2: After MatMul computation on the quantized input parameters x1 and x2, dequant and pertoken computations are performed, followed by an Add operation with x3, and finally an AllReduce computation.

      output=AllReduce(dequantScale∗pertokenScaleOptional∗(x1int8@x2int8+biasOptionalint32)+x3Optional)output= AllReduce(dequantScale * pertokenScaleOptional * (x1_{int8}@x2_{int8} + biasOptional_{int32}) + x3Optional)

Function Prototype

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

aclnnStatus aclnnQuantMatmulAllReduceV2GetWorkspaceSize(
    const aclTensor  *x1,
    const aclTensor  *x2,
    const aclTensor  *biasOptional,
    const aclTensor  *x3Optional,
    const aclTensor  *dequantScale,
    const aclTensor  *pertokenScaleOptional,
    const char       *group,
    const char       *reduceOp,
    int64_t           commTurn,
    int64_t           streamMode,
    const aclTensor  *output,
    uint64_t         *workspaceSize,
    aclOpExecutor   **executor)
aclnnStatus aclnnQuantMatmulAllReduceV2(
    void              *workspace,
    uint64_t           workspaceSize,
    aclOpExecutor     *executor,
    const aclrtStream  stream)

aclnnQuantMatmulAllReduceV2GetWorkspaceSize

  • Parameters

    Parameter Input/Output Description Instruction Data Type Data Format Dimension (Shape) Non-Contiguous Tensor
    x1 Input The left matrix of the matrix multiplication, i.e., x1 in the formula.
    • Currently, only two-dimensional or three-dimensional input is supported.
    • Supports the non-transpose scenario.
    INT8 ND 2-3 ×
    x2 Input The right matrix of the MatMul computation, i.e., x2 in the formula.
    • The current version only supports two-dimensional input.
    • Transpose and non-transpose scenarios are supported.
    • In ND format, non-contiguous tensors with the last two axes transposed are supported; other non-contiguous tensors are not supported.
    INT8 ND, FRACTAL_NZ 2 ×
    biasOptional Input biasOptional in the formula. The current version supports only one-dimensional input. INT32 ND 0-1 √
    x3Optional Input The add operation after the matrix multiplication, i.e., x3Optional in the formula. The shape is consistent with the shape after the matrix multiplication. 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) or (1, n) in the PerChannel scenario.
    • When the output is BFLOAT16, directly pass the BFLOAT16 dequantScale to this API.
    • When the output is FLOAT16, if pertokenScaleOptional is not empty, you can directly pass the FLOAT32 dequantScale to this API; if pertokenScaleOptional 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 Per-token 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 √
    group Input Communication domain name. Obtained through the Hccl API "extern HcclResult HcclGetCommName(HcclComm comm, char* commName);", where commName is the group. String - - -
    reduceOp Input Type of the reduce operation. Only "sum" is supported in the current version. String - - -
    commTurn Input Number of communication data splits, i.e., total data volume / single communication volume. Only input 0 is supported in the current version. INT64 - - -
    streamMode Input Enumeration of the stream mode. Only the enumeration value 1 is supported in the current version. INT64 - - -
    output Output Result of the MatMul computation and AllReduce communication, i.e., output in the formula. The dimensions of output are the same as those 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 (only two-dimensional input is supported in the current version) and FRACTAL_NZ format (only four-dimensional input is supported in the current version). When the data format of x2 is FRACTAL_NZ, use aclnnCalculateMatmulWeightSizeV2 and aclnnTransMatmulWeight to complete the conversion from ND data format to NZ data format. Non-contiguous tensors only support the transpose scenario.
    • Ascend 950PR/Ascend 950DT:

      • The data format of input x2 supports only the ND format (only two-dimensional input is supported in the current version).
  • 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 Return Code Description
    ACLNN_ERR_PARAM_NULLPTR 161001 x1, x2, dequantScale, or output is a null pointer.
    ACLNN_ERR_PARAM_INVALID 161002 The data type of x1, x2, biasOptional, dequantScale, pertokenScaleOptional, x3Optional, or output is not within the supported range.
    streamMode is not within the valid range.
    The shape of x1, x2, biasOptional, dequantScale, pertokenScaleOptional, x3Optional, or output does not meet the constraint requirements.

aclnnQuantMatmulAllReduceV2

  • 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 through the first-phase API aclnnQuantMatmulAllReduceV2GetWorkspaceSize.
    executor Input Operator executor, which 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: aclnnQuantMatmulAllReduceV2 defaults to a non-deterministic implementation. Deterministic computation can be enabled by setting the HCCL_DETERMINISTIC environment variable to true.

    • Ascend 950PR/Ascend 950DT: aclnnQuantMatmulAllReduceV2 defaults to a deterministic implementation.

  • MC2 is disabled in incremental scenarios and enabled in full-scale scenarios.

  • The input x1 can be two-dimensional or three-dimensional and must be a non-empty tensor, with its shape being (b, s, k) or (m, k). x2 must be two-dimensional and a non-empty tensor. Its shape is (k, n), and the k axis 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.

  • If bias is non-empty, its shape is (n). If x3 is non-empty, its shape must be the same as that of output.

  • When the shape of input x1 is (b, s, k), the shape of output is (b, s, n). When the shape of input x1 is (m, k), the shape of output is (m, n).

  • The passed x1, x2, dequantScale, or output must not be null pointers.

  • The data types and data formats of x1, x2, dequantScale, output, bias (when not empty), and x3 (when not empty) must be within the supported range.

  • If the output type is FLOAT16, when pertokenScaleOptional is empty, the type of dequantScale is INT64 or UINT64; when pertokenScaleOptional is not empty, the type of dequantScale is FLOAT32.

  • If the output type is BFLOAT16, the type of dequantScale is BFLOAT16.

  • 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 the HCCS all-mesh networking is supported.

    • Atlas A2 training products/Atlas A2 inference products: Supports 1, 2, 4, and 8 cards.

    • Ascend 950PR/Ascend 950DT: Supports 1, 2, 4, 8, 16, 32, and 64 cards.

  • Atlas A2 training products/Atlas A2 inference products: The MC2 operator for general-purpose computation fusion in a model supports only the same communication domain.

  • 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. 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_v2.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);
    
        // Apply for device memory by calling aclrtMalloc.
        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 data from the host side to the 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);
    
        // Compute 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 aclCreateTensor 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);
        // Call aclrtMalloc to allocate 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);
        // Call aclrtMemcpy to copy host-side data to the device-side 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 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 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> 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 *x3DeviceAddr = nullptr;
        void *outDeviceAddr = nullptr;
        aclTensor *x1 = nullptr;
        aclTensor *x2 = nullptr;
        aclTensor *bias = nullptr;
        aclTensor *dequantScale = nullptr;
        aclTensor *pertokenScale = 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 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> x3HostData(x3ShapeSize, 1);
        std::vector<op::fp16_t> outHostData(outShapeSize, 0);
        // Create the 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(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 = aclnnQuantMatmulAllReduceV2GetWorkspaceSize(x1, x2, bias, x3, dequantScale, pertokenScale,
                                                        hcom_name, "sum", commTurn, streamMode, out,
                                                        &workspaceSize, &executor);
        CHECK_RET(ret == ACL_SUCCESS,
                LOG_PRINT("aclnnQuantMatmulAllReduceV2GetWorkspaceSize 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 = aclnnQuantMatmulAllReduceV2(workspaceAddr, workspaceSize, executor, args.stream);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnQuantMatmulAllReduceV2 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 aclnnQuantMatmulAllReduceV2 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 (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 (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 multiple threads.
        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;
    }