aclnnQuantMatmulAlltoAll

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	×

Function

Description: Fuses quantized Matmul computation, Permute (to ensure contiguous addresses after communication), and All-to-All Communication. Computation precedes communication. Supports K-C quantization and mx Quantization Mode.
Formula: Assume the shape of x1 is (BS, H1), the shape of x2 is (H1, H2), and rankSize is the number of NPUs.
- Atlas A2 training products/Atlas A2 inference products:
  - K-C Quantization Scenario: $computeOut = (x1 @ x2) * x1Scale * x2Scale + bias \\ permutedOut = computeOut.view(BS, rankSize, H2 / rankSize).permute(1, 0, 2) \\ output = AlltoAll(permutedOut).view(rankSize * BS, H2 / rankSize)$
- Ascend 950PR/Ascend 950DT:
  - K-C Quantization Scenario:
    $computeOut = (x1 @ x2 + bias) * x1Scale * x2Scale \\ permutedOut = computeOut.view(BS, rankSize, H2 / rankSize).permute(1, 0, 2) \\ output = AlltoAll(permutedOut).view(rankSize * BS, H2 / rankSize)$
  - mx quantization scenario:
    $computeOut = (x1* x1Scale)@(x2* x2Scale) + bias \\ permutedOut = computeOut.view(BS, rankSize, H2 / rankSize).permute(1, 0, 2) \\ output = AlltoAll(permutedOut).view(rankSize * BS, H2 / rankSize)$

Function Prototype

Each operator is divided into a two-phase API. The "aclnnQuantMatmulAlltoAllGetWorkspaceSize" API must be called first to obtain the workspace size required for computation and the executor that encapsulates the operator computation flow. Then, the "aclnnQuantMatmulAlltoAll" API is called to perform the computation.

aclnnStatus aclnnQuantMatmulAlltoAllGetWorkspaceSize(
const aclTensor*   x1,
const aclTensor*   x2,
const aclTensor*   biasOptional,
const aclTensor*   x1Scale,
const aclTensor*   x2Scale,
const aclTensor*   commScaleOptional,
const aclTensor*   x1OffsetOptional,
const aclTensor*   x2OffsetOptional,
const aclIntArray* alltoAllAxesOptional,
const char*        group,
int64_t            x1QuantMode,
int64_t            x2QuantMode,
int64_t            commQuantMode,
int64_t            commQuantDtype,
int64_t            groupSize,
bool               transposeX1,
bool               transposeX2,
const aclTensor*   output,
uint64_t*          workspaceSize,
aclOpExecutor**    executor);

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

aclnnQuantMatmulAlltoAllGetWorkspaceSize

Parameters:

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
x1	Input	The left matrix input of the fused operator, corresponding to x1 in the formula.	This input serves as the left matrix input for the MatMul computation. The supported data types vary by device model. For details, see Constraints.	FLOAT8_E4M3FN, FLOAT8_E5M2, INT8	ND	2D, with shape (BS, H1).	x
x2	Input	Represents the right matrix input of the fused operator, corresponding to x2 in the formula.	Serves directly as the right matrix input for MatMul computation. The constraints on data types and non-contiguous tensors vary by device model. For details, see Constraints.	FLOAT8_E4M3FN, FLOAT8_E5M2, INT8	ND	2-dimensional, with shape (H1, H2).	√
biasOptional	Optional Input	Bias accumulated after matrix multiplication, corresponding to bias in the formula.	Data type restrictions vary by device model. For details, see Constraints.	FLOAT16, BFLOAT16, FLOAT32	ND	1D, with shape (H2)	x
x1Scale	Input	Quantization coefficient of the left matrix.	Corresponds to x1Scale in the formula.	FLOAT32, FLOAT8_E8M0	ND	1D or 3D. The shape is (BS) in the K-C quantization scenario, and (BS, ceil(H1/64), 2) in the mx quantization scenario.	x
x2Scale	Input	Quantization coefficient of the right matrix.	Corresponds to x2Scale in the formula.	FLOAT32, FLOAT8_E8M0	ND	1D/3D. The shape is (H2) in K-C quantization scenarios, and (H2, ceil(H1/64), 2) in mx quantization scenarios.	x
commScaleOptional	Optional Input	Quantization coefficient for low-bit communication.	Reserved parameter. Low-bit communication is not supported yet.	-	-	-	-
x1OffsetOptional	Optional Input	Quantization offset of the left matrix.	Reserved parameter. Not supported yet.	-	-	-	-
x2OffsetOptional	Optional Input	Quantization bias of the right matrix.	Reserved parameter. Not supported yet.	-	-	-	-
alltoAllAxesOptional	Optional Input	Direction of AlltoAll and Permute data exchange.	Supports empty or [-1, -2]. When empty, defaults to [-1, -2], indicating that the input is transformed from (BS, H2) to (BS*rankSize, H2/rankSize).	aclIntArray* (element type INT64).	-	1D, with shape (2).	-
Group	Input	Specifies the string that identifies the column group, i.e., the communication domain name. Obtain commName through the Hccl API HcclGetCommName and use it as this parameter.	The string length must be in the range (0, 128).	STRING	-	-	-
x1QuantMode	Input	Quantization mode of the left matrix.	The value restrictions vary by device model. For details, see Constraints.	INT	-	-	-
x2QuantMode	Input	Quantization mode of the right matrix.	The value restrictions vary by device model. For details, see Constraints.	INT	-	-	-
commQuantMode	Input	Quantization mode for low-bit communication.	Reserved parameter. Currently only supports configuration as 0, indicating no quantization.	INT	-	-	-
commQuantDtype	Input	Quantization type for low-bit communication.	Reserved parameter. Currently only supports -1, indicating ACL_DT_UNDEFINED.	INT	-	-	-
groupSize	Input	Quantization group size for Matmul computation in three dimensions.	Specifies the groupSize input, which is composed of three values—groupSizeM, groupSizeN, and groupSizeK—for the three dimensions. Each value occupies 16 bits, and together they occupy the lower 48 bits of the int64_t groupSize (the upper 16 bits of groupSize are invalid). The formula is: groupSize = groupSizeK \| groupSizeN << 16 \| groupSizeM << 32. In the mx quantization scenario, it is configured as 4295032864, which corresponds to [1, 1, 32] derived from the formula. In other quantization scenarios, it defaults to 0 and the value does not take effect.	INT	-	-	-
transposeX1	Input	Indicates whether the left matrix is transposed.	Configuring it to True is currently not supported.	bool	-	-	-
transposeX2	Input	Indicates whether the right matrix has been transposed.	When set to True, the right matrix shape is (H2, H1). Must be set to True in mx quantization mode.	bool	-	-	-
output	Output	The final computation result.		FLOAT16, BFLOAT16, FLOAT32	ND	2D, with shape (BS*rankSize, H2/rankSize).	x
workspaceSize	Output	Returns the workspace size to be allocated on the Device side.		UINT64	-	-	-
executor	Output	Returns the operator executor, which contains the computation flow of the operator.		aclOpExecutor*	-	-	-

The mapping between the enumeration values of x1QuantMode, x2QuantMode, and commQuantMode and the Quantization Mode is as follows:

0: no quantization
1: pertensor
2: perchannel
3: pertoken
4: pergroup
5: perblock
6: mx quantization
7: pertoken dynamic quantization

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	A required input or output tensor is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data type of an input or output is not within the supported range.
		An input tensor is empty.
		alltoAllAxesOptional is invalid.
		transposeX1 is true.
		The communication domain length is invalid.
		The input/output tensor dimensions are invalid.
		The input/output format is a private format.

aclnnQuantMatmulAlltoAll

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 aclnnQuantMatmulAlltoAllGetWorkspaceSize.
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

Defaults to a deterministic computation.
Number of NPU cards (rankSize) varies by device model:
- Atlas A2 training products/Atlas A2 inference products: Supports 2, 4, or 8 cards.
- Ascend 950PR/Ascend 950DT: Supports 2, 4, 8, or 16 cards.
The variable H2 used in the shape in Parameters must be divisible by the number of NPU cards.
The values of BS * rankSize and H2 must not exceed 2147483647 (INT32_MAX). The value of BS must not be less than 1, and the value of H2 must not be less than 2.
Empty tensors are not supported.
Support for non-contiguous tensors varies by device model:
- Atlas A2 training products/Atlas A2 inference products: No non-contiguous tensors are supported.
- Ascend 950PR/Ascend 950DT: Only x2 is supported as a non-contiguous tensor. Other non-contiguous tensors are not supported.
The input x1, x2, x1Scale, x2Scale, and output must not be null pointers, and
- Atlas A2 training products/Atlas A2 inference products: biasOptional does not support null pointers.

The data types, data dimensions, and quantization modes of the input and output of this operator have different restrictions depending on the device model:

Atlas A2 training products/Atlas A2 inference products:
- Quantization Mode:
  - Currently supported: K-C quantization, left matrix perToken quantization with x1QuantMode=3, right matrix perChannel quantization with x2QuantMode=2.
  - The bias is added after quantization.
- Type constraints:
  - The supported data type combinations for input/output are:
    - K-C quantization:
      
      x1 x2 biasOptional output
      INT8 INT8 FLOAT16 FLOAT16
      INT8 INT8 FLOAT32 FLOAT16
      INT8 INT8 BFLOAT16 BFLOAT16
      INT8 INT8 FLOAT32 BFLOAT16
- Dimensional Constraints:
  - H1 supports only the range [1, 65535].

x1	x2	biasOptional	output
INT8	INT8	FLOAT16	FLOAT16
INT8	INT8	FLOAT32	FLOAT16
INT8	INT8	BFLOAT16	BFLOAT16
INT8	INT8	FLOAT32	BFLOAT16

Ascend 950PR/Ascend 950DT:

Quantization Mode:
- Currently supported: K-C quantization, left matrix perToken quantization, x1QuantMode=3, right matrix perChannel quantization, x2QuantMode=2; mx quantization, left matrix mx quantization, x1QuantMode=6, right matrix mx quantization, x2QuantMode=6.
- The bias offset is added before quantization.

Type constraints:

biasOptional can be empty.

The supported input/output data type combinations are:

K-C Quantization:

x1	x2	biasOptional	output	x1QuantDtype	x2QuantDtype	x1ScaleOptional	x2Scale
FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT32	FLOAT16	3	2	FLOAT32	FLOAT32
FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT32	BFLOAT16	3	2	FLOAT32	FLOAT32
FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT32	FLOAT32	3	2	FLOAT32	FLOAT32
FLOAT8_E4M3FN	FLOAT8_E5M2	FLOAT32	FLOAT16	3	2	FLOAT32	FLOAT32
FLOAT8_E4M3FN	FLOAT8_E5M2	FLOAT32	BFLOAT16	3	2	FLOAT32	FLOAT32
FLOAT8_E4M3FN	FLOAT8_E5M2	FLOAT32	FLOAT32	3	2	FLOAT32	FLOAT32
FLOAT8_E5M2	FLOAT8_E4M3FN	FLOAT32	FLOAT16	3	2	FLOAT32	FLOAT32
FLOAT8_E5M2	FLOAT8_E4M3FN	FLOAT32	BFLOAT16	3	2	FLOAT32	FLOAT32
FLOAT8_E5M2	FLOAT8_E4M3FN	FLOAT32	FLOAT32	3	2	FLOAT32	FLOAT32
FLOAT8_E5M2	FLOAT8_E5M2	FLOAT32	FLOAT16	3	2	FLOAT32	FLOAT32
FLOAT8_E5M2	FLOAT8_E5M2	FLOAT32	BFLOAT16	3	2	FLOAT32	FLOAT32
FLOAT8_E5M2	FLOAT8_E5M2	FLOAT32	FLOAT32	3	2	FLOAT32	FLOAT32

mx Quantization:

x1	x2	biasOptional	output	x1QuantDtype	x2QuantDtype	x1ScaleOptional	x2Scale
FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT32	FLOAT16	6	6	FLOAT8_E8M0	FLOAT8_E8M0
FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT32	BFLOAT16	6	6	FLOAT8_E8M0	FLOAT8_E8M0
FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT32	FLOAT32	6	6	FLOAT8_E8M0	FLOAT8_E8M0
FLOAT8_E4M3FN	FLOAT8_E5M2	FLOAT32	FLOAT16	6	6	FLOAT8_E8M0	FLOAT8_E8M0
FLOAT8_E4M3FN	FLOAT8_E5M2	FLOAT32	BFLOAT16	6	6	FLOAT8_E8M0	FLOAT8_E8M0
FLOAT8_E4M3FN	FLOAT8_E5M2	FLOAT32	FLOAT32	6	6	FLOAT8_E8M0	FLOAT8_E8M0
FLOAT8_E5M2	FLOAT8_E4M3FN	FLOAT32	FLOAT16	6	6	FLOAT8_E8M0	FLOAT8_E8M0
FLOAT8_E5M2	FLOAT8_E4M3FN	FLOAT32	BFLOAT16	6	6	FLOAT8_E8M0	FLOAT8_E8M0
FLOAT8_E5M2	FLOAT8_E4M3FN	FLOAT32	FLOAT32	6	6	FLOAT8_E8M0	FLOAT8_E8M0
FLOAT8_E5M2	FLOAT8_E5M2	FLOAT32	FLOAT16	6	6	FLOAT8_E8M0	FLOAT8_E8M0
FLOAT8_E5M2	FLOAT8_E5M2	FLOAT32	BFLOAT16	6	6	FLOAT8_E8M0	FLOAT8_E8M0
FLOAT8_E5M2	FLOAT8_E5M2	FLOAT32	FLOAT32	6	6	FLOAT8_E8M0	FLOAT8_E8M0

Dimension Constraints:
- The H1 range supports only [1, 65535].
- In mx quantization scenarios, x2 must be transposed, with shape (H2, H1) and transposeX2 set to True.

The communication-computation fusion operator does not support concurrent calls, and different communication-computation fusion operators also do not support concurrent calls.
Cross-supernode communication is not supported. Only intra-supernode communication is supported.

Example

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

Note: This sample code calls some HCCL collective communication library APIs: HcclGetCommName, HcclCommInitAll, and HcclCommDestroy. See HCCL API (C).

Atlas A2 training products/Atlas A2 inference products:

#include <thread>
#include <iostream>
#include <string>
#include <cstring>
#include <vector>
#include <acl/acl.h>
#include <hccl/hccl.h>
#include "aclnn/opdev/fp16_t.h"
#include "aclnnop/aclnn_quant_matmul_allto_all.h"

int ndev = 2;

#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);
    // Allocate device memory using 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 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);
    // Compute the strides of the 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 using aclCreateTensor.
    *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 launchOneThreadQuantMatmulAlltoAll(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> x1ScaleShape = {32};
    std::vector<int64_t> x2ScaleShape = {128};
    std::vector<int64_t> outShape = {32, 128};
    void *x1DeviceAddr = nullptr;
    void *x2DeviceAddr = nullptr;
    void *biasDeviceAddr = nullptr;
    void *x1ScaleDeviceAddr = nullptr;
    void *x2ScaleDeviceAddr = nullptr;
    void *outDeviceAddr = nullptr;
    aclTensor *x1 = nullptr;
    aclTensor *x2 = nullptr;
    aclTensor *bias = nullptr;
    aclTensor *x1Scale = nullptr;
    aclTensor *x2Scale = nullptr;
    aclTensor *out = nullptr;

    int64_t x1QuantMode = 3;
    int64_t x2QuantMode = 2;
    int64_t commQuantMode = 0;
    int64_t commQuantDtype = -1;
    int64_t groupSize = 0;

    int64_t a2aAxes[2] = {-1, -2};
    aclIntArray* alltoAllAxesOptional = aclCreateIntArray(a2aAxes, static_cast<uint64_t>(2));
    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 x1ScaleShapeSize = GetShapeSize(x1ScaleShape);
    long long x2ScaleShapeSize = GetShapeSize(x2ScaleShape);
    long long outShapeSize = GetShapeSize(outShape);
    std::vector<int8_t> x1HostData(x1ShapeSize, 1);
    std::vector<int8_t> x2HostData(x2ShapeSize, 1);
    std::vector<op::fp16_t> biasHostData(biasShapeSize, 1);
    std::vector<float> x1ScaleHostData(x1ScaleShapeSize, 1);
    std::vector<float> x2ScaleHostData(x2ScaleShapeSize, 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);
    ret = CreateAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_INT8, &x2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(biasHostData, biasShape, &biasDeviceAddr, aclDataType::ACL_FLOAT16, &bias);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(x1ScaleHostData, x1ScaleShape, &x1ScaleDeviceAddr, aclDataType::ACL_FLOAT, &x1Scale);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(x2ScaleHostData, x2ScaleShape, &x2ScaleDeviceAddr, aclDataType::ACL_FLOAT, &x2Scale);
    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 = aclnnQuantMatmulAlltoAllGetWorkspaceSize(x1, x2, bias, x1Scale, x2Scale, nullptr, nullptr, nullptr,
                                                  alltoAllAxesOptional, hcom_name, x1QuantMode, x2QuantMode,
                                                  commQuantMode, commQuantDtype, groupSize, false, false,
                                                  out, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS,
            LOG_PRINT("aclnnQuantMatmulAlltoAllGetWorkspaceSize 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 = aclnnQuantMatmulAlltoAll(workspaceAddr, workspaceSize, executor, args.stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnQuantMatmulAlltoAll failed. ERROR: %d\n", ret); return ret);
    //(Boilerplate) Synchronously wait for the 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 aclnnQuantMatmulAlltoAll 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 (x1Scale != nullptr) {
        aclDestroyTensor(x1Scale);
    }
    if (x2Scale != nullptr) {
        aclDestroyTensor(x2Scale);
    }
    if (out != nullptr) {
        aclDestroyTensor(out);
    }
    if (x1DeviceAddr != nullptr) {
        aclrtFree(x1DeviceAddr);
    }
    if (x2DeviceAddr != nullptr) {
        aclrtFree(x2DeviceAddr);
    }
    if (biasDeviceAddr != nullptr) {
        aclrtFree(biasDeviceAddr);
    }
    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[]) {
    // This sample is implemented based on Atlas A2 and must run on Atlas A2.
    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(&launchOneThreadQuantMatmulAlltoAll, std::ref(args[rankId])));
    }
    for (uint32_t rankId = 0; rankId < ndev; rankId++) {
        threads[rankId]->join();
    }
    aclFinalize();
    return 0;
}

Ascend 950PR/Ascend 950DT:

#include <thread>
#include <iostream>
#include <string>
#include <cstring>
#include <vector>
#include <acl/acl.h>
#include <hccl/hccl.h>
#include "aclnnop/aclnn_quant_matmul_allto_all.h"

int ndev = 2;

#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 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 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);
    return 0;
}

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

int launchOneThreadQuantMatmulAlltoAll(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> x1ScaleShape = {32};
    std::vector<int64_t> x2ScaleShape = {128};
    std::vector<int64_t> outShape = {32 * ndev, 128 / ndev};
    void *x1DeviceAddr = nullptr;
    void *x2DeviceAddr = nullptr;
    void *biasDeviceAddr = nullptr;
    void *x1ScaleDeviceAddr = nullptr;
    void *x2ScaleDeviceAddr = nullptr;
    void *outDeviceAddr = nullptr;
    aclTensor *x1 = nullptr;
    aclTensor *x2 = nullptr;
    aclTensor *bias = nullptr;
    aclTensor *x1Scale = nullptr;
    aclTensor *x2Scale = nullptr;
    aclTensor *out = nullptr;

    int64_t x1QuantMode = 3;
    int64_t x2QuantMode = 2;
    int64_t commQuantMode = 0;
    int64_t commQuantDtype = -1;
    int64_t groupSize = 0;

    int64_t a2aAxes[2] = {-1, -2};
    aclIntArray* alltoAllAxesOptional = aclCreateIntArray(a2aAxes, static_cast<uint64_t>(2));
    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 x1ScaleShapeSize = GetShapeSize(x1ScaleShape);
    long long x2ScaleShapeSize = GetShapeSize(x2ScaleShape);
    long long outShapeSize = GetShapeSize(outShape);
    std::vector<int16_t> x1HostData(x1ShapeSize, 1);
    std::vector<int16_t> x2HostData(x2ShapeSize, 1);
    std::vector<int16_t> biasHostData(biasShapeSize, 1);
    std::vector<int16_t> x1ScaleHostData(x1ShapeSize, 1);
    std::vector<int16_t> x2ScaleHostData(x2ShapeSize, 1);
    std::vector<int16_t> outHostData(outShapeSize, 0);
    // Create a tensor.
    ret = CreateAclTensor(x1HostData, x1Shape, &x1DeviceAddr, aclDataType::ACL_FLOAT8_E4M3FN, &x1);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_FLOAT8_E4M3FN, &x2);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(biasHostData, biasShape, &biasDeviceAddr, aclDataType::ACL_FLOAT, &bias);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(x1ScaleHostData, x1ScaleShape, &x1ScaleDeviceAddr, aclDataType::ACL_FLOAT, &x1Scale);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    ret = CreateAclTensor(x2ScaleHostData, x2ScaleShape, &x2ScaleDeviceAddr, aclDataType::ACL_FLOAT, &x2Scale);
    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 = aclnnQuantMatmulAlltoAllGetWorkspaceSize(x1, x2, bias, x1Scale, x2Scale, nullptr, nullptr, nullptr,
                                                   alltoAllAxesOptional, hcom_name, x1QuantMode, x2QuantMode,
                                                   commQuantMode, commQuantDtype, groupSize, false, false,
                                                   out, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS,
            LOG_PRINT("aclnnQuantMatmulAlltoAllGetWorkspaceSize 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 = aclnnQuantMatmulAlltoAll(workspaceAddr, workspaceSize, executor, args.stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnQuantMatmulAlltoAll 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 aclnnQuantMatmulAlltoAll 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 (x1Scale != nullptr) {
        aclDestroyTensor(x1Scale);
    }
    if (x2Scale != nullptr) {
        aclDestroyTensor(x2Scale);
    }
    if (out != nullptr) {
        aclDestroyTensor(out);
    }
    if (x1DeviceAddr != nullptr) {
        aclrtFree(x1DeviceAddr);
    }
    if (x2DeviceAddr != nullptr) {
        aclrtFree(x2DeviceAddr);
    }
    if (biasDeviceAddr != nullptr) {
        aclrtFree(biasDeviceAddr);
    }
    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[]) {
    // This sample is implemented based on Atlas A5 and must run on Atlas A5.
    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 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(&launchOneThreadQuantMatmulAlltoAll, std::ref(args[rankId])));
    }
    for (uint32_t rankId = 0; rankId < ndev; rankId++) {
        threads[rankId]->join();
    }
    aclFinalize();
    return 0;
}