aclnnDynamicMxQuant

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: MX quantization with the destination data type being FLOAT4 or FLOAT8. On the given axis, for every blocksize number of elements, the quantization scale mxscale corresponding to this group is computed as the corresponding part of the output mxscaleOut. Then each element in this group is divided by mxscale and converted to the corresponding dstType according to round_mode, yielding the quantization result y as the corresponding part of the output yOut. When dstType is FLOAT8_E4M3FN or FLOAT8_E5M2, different algorithms for computing mxscale are specified based on the value of scaleAlg.
Formula:
- Scenario 1, when scaleAlg is 0:
  - Group the input x along the axis dimension by k = blocksize elements. Each group of k elements $\{\{V_i\}_{i=1}^{k}\}$ is dynamically quantized into $\{mxscale1, \{P_i\}_{i=1}^{k}\}$ , where k = blocksize.
  $shared\_exp = floor(log_2(max_i(|V_i|))) - emax \\ mxscale = 2^{shared\_exp}\\ P_i = cast\_to\_dst\_type(V_i/mxscale, round\_mode), \space i\space from\space 1\space to\space blocksize\\$
  - The quantized $P_{i}$ values form the output yOut according to the positions of the corresponding $V_{i}$ , and mxscale values form the output mxscaleOut according to the corresponding groups along the axis dimension.
  - emax: The exponent of the maximum normalized value of the corresponding data type.
    
    DataType emax
    FLOAT4_E2M1 2
    FLOAT4_E1M2 0
    FLOAT8_E4M3FN 8
    FLOAT8_E5M2 15
- Scenario 2: When scaleAlg is 1, only FP8 types are involved:
  - Divide the long vector into blocks, each with a length of k. Compute a block scale factor $S_{fp32}^b$ for each block individually, and then map all elements within the block to the target low-precision type FP8 using the same $S_{fp32}^b$ . If the last block has fewer than k elements, treat the missing values as 0 and process it as a complete block.
  - Find the maximum absolute value in the block:
    $Amax(D_{fp32}^b)=max(\{|d_{i}|\}_{i=1}^{k})$
  - Map FP32 to the representable range of the target data type FP8, where $Amax(DType)$ is the maximum value that the target precision can represent.
    $S_{fp32}^b = \frac{Amax(D_{fp32}^b)}{Amax(DType)}$
  - Convert the block scale factor $S_{fp32}^b$ to the representable scale value $S_{ue8m0}^b$ in FP8 format.
  - Extract the unbiased exponent $E_{int}^b$ and mantissa $M_{fixp}^b$ from the block floating-point scale factor $S_{fp32}^b$ .
  - To ensure no overflow during quantization, round up the exponent and keep it within the representable range of FP8:
    $E_{int}^b = \begin{cases} E_{int}^b + 1, & \text{if } S_{fp32}^b \text{ is a normal number, and } E_{int}^b < 254 \text{ and } M_{fixp}^b > 0 \\ E_{int}^b + 1, & \text{if } S_{fp32}^b \text{ is a subnormal number, and } M_{fixp}^b > 0.5 \\ E_{int}^b, & \text{otherwise} \end{cases}$
  - Calculate the block scale factor: $S_{ue8m0}^b=2^{E_{int}^b}$
  - Calculate the block conversion factor: $R_{fp32}^b=\frac{1}{fp32(S_{ue8m0}^b)}$
  - Apply to the final quantization step. For each element within a block, $d^i = DType(d_{fp32}^i \cdot R_{fp32}^n)$ . The final quantized output is $\left(S^b, [d^i]_{i=1}^k\right)$ , where $S^b$ represents the block scale factor, i.e., $S_{ue8m0}^b$ , and $[d^i]_{i=1}^k$ represents the quantized data within the block.

DataType	emax
FLOAT4_E2M1	2
FLOAT4_E1M2	0
FLOAT8_E4M3FN	8
FLOAT8_E5M2	15

Function Prototype

Each operator is divided into a two-phase API. You must first call the "aclnnDynamicMxQuantGetWorkspaceSize" API to obtain the workspace size required for computation and the executor that encapsulates the operator computation flow, and then call the "aclnnDynamicMxQuant" API to perform the computation.

aclnnStatus aclnnDynamicMxQuantGetWorkspaceSize(
  const aclTensor *x,
  int64_t          axis,
  char            *roundModeOptional,
  int64_t          dstType,
  int64_t          blocksize,
  int64_t          scaleAlg,
  const aclTensor *yOut,
  const aclTensor *mxscaleOut,
  uint64_t        *workspaceSize,
  aclOpExecutor   **executor)

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

aclnnDynamicMxQuantGetWorkspaceSize

Parameters

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
x (aclTensor*)	Input	Represents the input x, corresponding to Vi and di in the formula.	When the destination type is FLOAT4_E2M1 or FLOAT4_E1M2, the last dimension of x must be an even number.	FLOAT16, BFLOAT16	ND	1-7	√
axis (int64_t)	Input	Indicates the axis along which quantization is performed, corresponding to axis in the formula.	The value range is [-D, D-1], where D is the number of dimensions of x's shape.	INT64	ND	-	-
roundModeOptional (char*)	Input	Indicates the data conversion mode, which corresponds to round_mode in the formula.	When dstType is 40/41, {"rint", "floor", "round"} is supported. When dstType is 36/35, only {"rint"} is supported. When a null pointer is passed, the "rint" mode is used.	STRING	ND	-	-
dstType (int64_t)	input	Specifies the type of the converted output yOut, corresponding to DType in the formula.	The input range is {35, 36, 40, 41}, corresponding to the output yOut data types {35: FLOAT8_E5M2, 36: FLOAT8_E4M3FN, 40: FLOAT4_E2M1, 41: FLOAT4_E1M2}.	INT64	ND	-	-
blocksize (int64_t)	Input	Specifies the number of elements quantized each time, corresponding to blocksize in the formula.	Only multiples of 32 are supported. The value cannot be 0 and must not exceed 1024.	INT64	ND	-	-
scaleAlg (int64_t)	Input	Indicates the calculation method of mxscaleOut, corresponding to scaleAlg in the formula.	Supports values 0 and 1. A value of 0 represents scenario 1, and 1 represents scenario 2. When dstType is FLOAT4_E2M1/FLOAT4_E1M2, only the value 0 is supported.	INT64	ND	-	-
yOut（aclTensor*）	Output	Indicates the quantized result corresponding to the input x, which corresponds to Pi and di in the formula.	The shape is the same as that of the input x.	FLOAT4_E2M1, FLOAT4_E1M2, FLOAT8_E4M3FN, FLOAT8_E5M2	ND	1-7	√
mxscaleOut (aclTensor*)	Output	Quantization scale corresponding to each group, corresponding to mxscale and Sb in the formula.	The shape on the axis is the value of the corresponding axis of x divided by blocksize, rounded up, and then even-padded with a padding value of 0. When axis is not the last axis, the mxscaleOut output requires interleaving every two rows of data.	FLOAT8_E8M0	ND	1-8	√
workspaceSize（uint64_t）	Output	The workspace size to be applied for on the Device side.	-	-	-	-	-
executor (aclOpExecutor**)	Output	Returns the op executor, which contains the operator computation flow.	-	-	-	-	-

Return Value

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

The first-phase API performs input parameter validation and returns an error in the following scenarios:

Return Code	Error Code	Description
ACLNN_ERR_PARAM_NULLPTR	161001	x is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data type or data format of x, axis, roundModeOptional, dstType, blocksize, scaleAlg, yOut, or mxscaleOut is not within the supported range.
		The shape of x, yOut, or mxscaleOut does not meet the validation conditions.
		axis, roundModeOptional, dstType, blocksize, or scaleAlg does not match the currently supported values.
ACLNN_ERR_RUNTIME_ERROR	361001	The current platform is not within the supported platform range.

aclnnDynamicMxQuant

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 through the first-phase API aclnnDynamicMxQuantGetWorkspaceSize.
executor	Input	The op 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:
- aclnnDynamicMxQuant defaults to a deterministic implementation.
The shape constraints for x and mxscaleOut are as follows:
- rank(mxscaleOut) = rank(x) + 1.
- axis_change = axis if axis >= 0 else axis + rank(x).
- mxscaleOut.shape[axis_change] = (ceil(x.shape[axis] / blocksize) + 2 - 1) / 2.
- mxscaleOut.shape[-1] = 2.
- Other dimensions are the same as those of the input x.

Example

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

#include <iostream>
#include <memory>
#include <vector>

#include "acl/acl.h"
#include "aclnnop/aclnn_dynamic_mx_quant.h"

#define CHECK_RET(cond, return_expr) \
    do {                             \
        if (!(cond)) {               \
            return_expr;             \
        }                            \
    } while (0)

#define CHECK_FREE_RET(cond, return_expr) \
    do {                                  \
        if (!(cond)) {                    \
            Finalize(deviceId, stream);   \
            return_expr;                  \
        }                                 \
    } while (0)

#define LOG_PRINT(message, ...)         \
    do {                                \
        printf(message, ##__VA_ARGS__); \
    } while (0)

    int64_t
    GetShapeSize(const std::vector<int64_t>& shape)
{
    int64_t shapeSize = 1;
    for (auto i : shape) {
        shapeSize *= i;
    }
    return shapeSize;
}

int Init(int32_t deviceId, aclrtStream* stream)
{
    // Boilerplate for resource initialization.
    auto ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
    ret = aclrtSetDevice(deviceId);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
    ret = aclrtCreateStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
    return 0;
}

template <typename T>
int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
                    aclDataType dataType, aclTensor** tensor)
{
    auto size = GetShapeSize(shape) * sizeof(T);
    // Call aclrtMalloc to allocate device memory.
    auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
    // Call aclrtMemcpy to copy host data to device memory.
    ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);

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

void Finalize(int32_t deviceId, aclrtStream stream)
{
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
}

int aclnnDynamicMxQuantTest(int32_t deviceId, aclrtStream& stream)
{
    auto ret = Init(deviceId, &stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

    // 2. Construct inputs and outputs. Customize the construction based on the API interface.
    std::vector<int64_t> xShape = {1, 4};
    std::vector<int64_t> yOutShape = {1, 4};
    std::vector<int64_t> mxscaleOutShape = {1, 1, 2};
    void* xDeviceAddr = nullptr;
    void* yOutDeviceAddr = nullptr;
    void* mxscaleOutDeviceAddr = nullptr;
    aclTensor* x = nullptr;
    aclTensor* yOut = nullptr;
    aclTensor* mxscaleOut = nullptr;
    //Values corresponding to BF16: (0, 8, 64, 512).
    std::vector<uint16_t> xHostData = {0, 16640, 17024, 17408};
    //Values corresponding to float8_e4m3: (0, 4, 32, 256).
    std::vector<uint8_t> yOutHostData = {0, 72, 96, 120};
    //Value corresponding to float8_e8m0: (2).
    std::vector<uint8_t> mxscaleOutHostData = {{128, 0}};
    int64_t axis = -1;
    char* roundModeOptional = const_cast<char*>("rint");
    int64_t dstType = 36;
    int64_t blocksize = 32;
    int64_t scaleAlg = 0;
    // Create the x aclTensor.
    ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_BF16, &x);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> xTensorPtr(x, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void*)> xDeviceAddrPtr(xDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the yOut aclTensor.
    ret = CreateAclTensor(yOutHostData, yOutShape, &yOutDeviceAddr, aclDataType::ACL_FLOAT8_E4M3FN, &yOut);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> yOutTensorPtr(yOut, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void*)> yOutDeviceAddrPtr(yOutDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the mxscaleOut aclTensor.
    ret = CreateAclTensor(mxscaleOutHostData, mxscaleOutShape, &mxscaleOutDeviceAddr, aclDataType::ACL_FLOAT8_E8M0, &mxscaleOut);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> mxscaleOutTensorPtr(mxscaleOut, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void*)> mxscaleOutDeviceAddrPtr(mxscaleOutDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // Call the CANN operator library API. Replace with the actual API name.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;

    // Call the first-phase API of aclnnDynamicMxQuant.
    ret = aclnnDynamicMxQuantGetWorkspaceSize(x, axis, roundModeOptional, dstType, blocksize, scaleAlg, yOut, mxscaleOut, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnDynamicMxQuantGetWorkspaceSize failed. ERROR: %d\n", ret);
              return ret);
    // Allocate device memory based on the workspaceSize calculated by the first-phase API.
    void* workspaceAddr = nullptr;
    std::unique_ptr<void, aclError (*)(void*)> workspaceAddrPtr(nullptr, aclrtFree);
    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);
        workspaceAddrPtr.reset(workspaceAddr);
    }
    // Call the second-phase API of aclnnDynamicMxQuant.
    ret = aclnnDynamicMxQuant(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnDynamicMxQuant failed. ERROR: %d\n", ret); return ret);

    //(Boilerplate) Synchronously wait for the task execution to finish.
    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

    // Obtain the output value and copy the result from the device memory to the host. Modify this based on the specific API definition.
    auto size = GetShapeSize(yOutShape);
    std::vector<uint8_t> yOutData(
        size, 0);  // In C, fp4 data cannot be printed directly. Read it as uint8 and manually convert it to fp4 via binary.
    ret = aclrtMemcpy(yOutData.data(), yOutData.size() * sizeof(yOutData[0]), yOutDeviceAddr,
                      size * sizeof(yOutData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy yOut from device to host failed. ERROR: %d\n", ret);
              return ret);
    for (int64_t i = 0; i < size; i++) {
        LOG_PRINT("yOut[%ld] is: %d\n", i, yOutData[i]);
    }
    size = GetShapeSize(mxscaleOutShape);
    std::vector<uint8_t> mxscaleOutData(
        size, 0);  // In C, fp8 data cannot be printed directly. Read it as uint8 and manually convert it to fp8 via binary.
    ret = aclrtMemcpy(mxscaleOutData.data(), mxscaleOutData.size() * sizeof(mxscaleOutData[0]), mxscaleOutDeviceAddr,
                      size * sizeof(mxscaleOutData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy mxscaleOut from device to host failed. ERROR: %d\n", ret);
              return ret);
    for (int64_t i = 0; i < size; i++) {
        LOG_PRINT("mxscaleOut[%ld] is: %d\n", i, mxscaleOutData[i]);
    }
    return ACL_SUCCESS;
}

int main()
{
    // 1. (Boilerplate) Initialize the device/stream. Refer to the ACL API manual.
    // Fill in the deviceId based on your actual device.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = aclnnDynamicMxQuantTest(deviceId, stream);
    CHECK_FREE_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnDynamicMxQuantTest failed. ERROR: %d\n", ret); return ret);

    Finalize(deviceId, stream);
    return 0;
}