aclnnDynamicDualLevelMxQuant

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 target data type being FLOAT4. Quantization is performed only on the last axis, with all preceding axes processed jointly. The input is divided into multiple data blocks based on the given level0BlockSize. First-level quantization is performed on each data block, and the quantization scale level0ScaleOut is output. Then, the result of the first-level quantization is used as the new input and divided into multiple data blocks based on the given level1BlockSize. Second-level quantization is performed on each data block, and the quantization scale level1ScaleOut is output. Data type conversion is performed according to round_mode to obtain the quantization result yOut. For details, see the diagram.
Formula:
- Group the input x on the last axis by $k_0$ = level0BlockSize elements. A group of $k_0$ elements $\{\{x_i\}_{i=1}^{k_0}\}$ is dynamically quantized into $\{level0Scale, \{temp_i\}_{i=1}^{k_0}\}$ , where $k_0$ = level0BlockSize. Then, group temp on the last axis by $k_1$ = level1BlockSize elements. A group of $k_1$ elements $\{\{temp_i\}_{i=1}^{k_1}\}$ is dynamically quantized into $\{level1Scale, \{y_i\}_{i=1}^{k_1}\}$ , where $k_1$ = level1BlockSize.
$input\_max_i = max_i(abs(x_i))$ $level0Scale = input\_max_i / (FP4\_E2M1\_MAX)$ $temp_i = cast\_to\_x\_type(x_i / level0Scale), \space i\space from\space 1\space to\space level0BlockSize$ $shared\_exp = floor(log_2(max_i(|temp_i|))) - emax$ $level1Scale = 2^{shared\_exp}$ $y_i = cast\_to\_FP4\_E2M1(temp_i/level1Scale, round\_mode), \space i\space from\space 1\space to\space level1BlockSize$
- The quantized $y_{i}$ forms the output yOut according to the position of the corresponding $x_{i}$ , level0Scale forms the output level0ScaleOut according to the groups corresponding to the last axis, and level1Scale forms the output level1ScaleOut according to the groups corresponding to the last axis.
- max_i represents the maximum value in the i-th group.
- emax: The exponent of the maximum normal number of the corresponding data type.
  
  DataType emax
  FLOAT4_E2M1 2

DataType	emax
FLOAT4_E2M1	2

Function Prototype

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

aclnnStatus aclnnDynamicDualLevelMxQuantGetWorkspaceSize(
  const aclTensor *x,
  const aclTensor *smoothScaleOptional,
  char            *roundModeOptional,
  int64_t          level0BlockSize,
  int64_t          level1BlockSize,
  const aclTensor *yOut,
  const aclTensor *level0ScaleOut,
  const aclTensor *level1ScaleOut,
  uint64_t        *workspaceSize,
  aclOpExecutor   **executor)

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

aclnnDynamicDualLevelMxQuantGetWorkspaceSize

Parameters:

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-Contiguous Tensor
x (aclTensor*)	Input	Represents the input x, corresponding to x_i in the formula.	The last dimension of x must be an even number. Empty tensors are not supported.	FLOAT16, BFLOAT16	ND	1-7	√
smoothScalesOptional (aclTensor*)	Input	Indicates the optional input smoothScaleOptional.	Currently, this feature is not supported. Only nullptr is supported as input.	Same as the input x.	ND	1	√
roundModeOptional (char*)	Input	Indicates the data conversion mode, corresponding to round_mode in the formula.	Supports {"rint", "round", "floor"}; Defaults to "rint".	STRING	-	-	-
level0BlockSize (int64_t)	Input	Represents the block_size of the first-level quantization, corresponding to level0BlockSize in the formula.	The input range is {512}.	INT64	-	-	-
level1BlockSize (int64_t)	Input	Indicates the block_size of the second-level quantization, which corresponds to level1BlockSize in the formula.	The input range is {32}.	INT64	-	-	-
yOut (aclTensor*)	Output	Indicates the corresponding result after quantizing the input x, corresponding to y_i in the formula.	The shape is the same as the input x. Empty tensors are not supported.	FLOAT4_E2M1	ND	1-7	√
level0ScaleOut (aclTensor*)	Output	Represents the scale of the first-level quantization, corresponding to level0Scale in the formula.	The value of shape on the last axis is the value of the last axis of x divided by level0BlockSize and rounded up. Empty tensors are not supported.	FLOAT32	ND	1-7	√
level1ScaleOut (aclTensor*)	Output	Represents the scale of the second-level quantization, corresponding to level1Scale in the formula.	The size of shape is dim of x + 1. The values of shape on the last two axes are ((ceil(x.shape[-1] / level1Blocksize) + 2 - 1) / 2, 2), and even padding is applied to them, with the padding value being 0. Empty tensors are not supported.	FLOAT8_E8M0	ND	1-8	√
workspaceSize	Output	Returns the workspace size to be applied for on the Device side.	-	-	-	-	-
executor	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 contains a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data type or data format of x, smoothScaleOptional, yOut, level0ScaleOut, or level1ScaleOut is not within the supported range.
		The shape of x, smoothScaleOptional, yOut, level0ScaleOut, or level1ScaleOut does not meet the validation conditions.
		roundModeOptional, level0BlockSize, or level1BlockSize does not match the currently supported values.
ACLNN_ERR_RUNTIME_ERROR	361001	The current platform is not within the supported platform range.

aclnnDynamicDualLevelMxQuant

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

The shape constraints for x, level0ScaleOut, and level1ScaleOut are described as follows:
- rank(level1ScaleOut) = rank(x) + 1.
- level0ScaleOut.shape[-1] = ceil(x.shape[-1] / level0Blocksize).
- level1ScaleOut.shape[-2] = (ceil(x.shape[-1] / level1Blocksize) + 2 - 1) / 2.
- level1ScaleOut.shape[-1] = 2.
- Other dimensions are the same as those of the input x.
Deterministic computation description: aclnnDynamicDualLevelMxQuant defaults to a deterministic implementation.

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_dual_level_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);
    // Copy host-side data to device-side 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 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];
    }

    // Create an aclTensor by calling the aclCreateTensor API.
    *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 aclnnDynamicDualLevelMxQuantTest(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. Custom construction is required based on the API.
    std::vector<int64_t> xShape = {1, 512};
    std::vector<int64_t> smoothScaleOptionalShape = {1};
    std::vector<int64_t> yOutShape = {1, 512};
    std::vector<int64_t> level0ScaleOutShape = {1, 1};
    std::vector<int64_t> level1ScaleOutShape = {1, 8, 2};
    void* xDeviceAddr = nullptr;
    void* smoothScaleOptionalDeviceAddr = nullptr;
    void* yOutDeviceAddr = nullptr;
    void* level0ScaleOutDeviceAddr = nullptr;
    void* level1ScaleOutDeviceAddr = nullptr;
    aclTensor* x = nullptr;
    aclTensor* smoothScaleOptional = nullptr;
    aclTensor* yOut = nullptr;
    aclTensor* level0ScaleOut = nullptr;
    aclTensor* level1ScaleOut = nullptr;

    // Corresponding BF16 values (0->0, 16640->8, 17024->64, 17408->512).
    std::vector<uint16_t> xHostData(512, 16640);
    std::vector<uint16_t> smoothScaleOptionalHostData = {0};
    // Values corresponding to float4_e2m1 (0->0, 72->4, 96->32, 120->256).
    std::vector<uint8_t> yOutHostData(512, 0);
    // Values corresponding to float32 (0->0).
    std::vector<float> level0ScaleOutHostData = {{0}};
    //Values corresponding to float8_e8m0 (128->2).
    std::vector<std::vector<std::vector<uint8_t>>> level1ScaleOutHostData(1, std::vector<std::vector<uint8_t>>(8, std::vector<uint8_t>(2, 0)));
    const char* roundModeOptional = "rint";
    int64_t level0Blocksize = 512;
    int64_t level1Blocksize = 32;

    // 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 smoothScaleOptional aclTensor.
    ret = CreateAclTensor(smoothScaleOptionalHostData, smoothScaleOptionalShape, &smoothScaleOptionalDeviceAddr, aclDataType::ACL_BF16, &smoothScaleOptional);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> smoothScaleOptionalTensorPtr(smoothScaleOptional, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void*)> smoothScaleOptionalDeviceAddrPtr(smoothScaleOptionalDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create yOut aclTensor.
    ret = CreateAclTensor(yOutHostData, yOutShape, &yOutDeviceAddr, aclDataType::ACL_FLOAT4_E2M1, &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 level0ScaleOut aclTensor.
    ret = CreateAclTensor(level0ScaleOutHostData, level0ScaleOutShape, &level0ScaleOutDeviceAddr, aclDataType::ACL_FLOAT, &level0ScaleOut);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> level0ScaleOutTensorPtr(level0ScaleOut, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void*)> level0ScaleOutDeviceAddrPtr(level0ScaleOutDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create level1ScaleOut aclTensor.
    ret = CreateAclTensor(level1ScaleOutHostData, level1ScaleOutShape, &level1ScaleOutDeviceAddr, aclDataType::ACL_FLOAT8_E8M0, &level1ScaleOut);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> level1ScaleOutTensorPtr(level1ScaleOut, aclDestroyTensor);
    std::unique_ptr<void, aclError (*)(void*)> level1ScaleOutDeviceAddrPtr(level1ScaleOutDeviceAddr, aclrtFree);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

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

    // Call the first-phase API of aclnnDynamicDualLevelMxQuant.
    ret = aclnnDynamicDualLevelMxQuantGetWorkspaceSize(x, smoothScaleOptional, (char*)roundModeOptional, level0Blocksize, level1Blocksize, yOut, level0ScaleOut, level1ScaleOut, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnDynamicDualLevelMxQuantGetWorkspaceSize failed. ERROR: %d\n", ret);
              return ret);
    // Apply for 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 aclnnDynamicDualLevelMxQuant.
    ret = aclnnDynamicDualLevelMxQuant(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnDynamicDualLevelMxQuant failed. ERROR: %d\n", ret); return ret);

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

    // Obtain the output value. Copy the result from the device side memory to the host side. Modify this based on the specific API definition.
    auto size = GetShapeSize(yOutShape) / 2;
    std::vector<uint8_t> yOutData(
        size, 0);  // In C language, fp4 data cannot be printed directly. It needs to be read out using uint8 and manually converted to fp4 through binary conversion.
    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(level0ScaleOutShape);
    std::vector<float> level0ScaleOutData(
        size, 0);
    ret = aclrtMemcpy(level0ScaleOutData.data(), level0ScaleOutData.size() * sizeof(level0ScaleOutData[0]), level0ScaleOutDeviceAddr,
                      size * sizeof(level0ScaleOutData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy level0ScaleOut from device to host failed. ERROR: %d\n", ret);
              return ret);
    for (int64_t i = 0; i < size; i++) {
        LOG_PRINT("level0ScaleOut[%ld] is: %f\n", i, level0ScaleOutData[i]);
    }
    size = GetShapeSize(level1ScaleOutShape);
    std::vector<uint8_t> level1ScaleOutData(
        size, 0);  // In C language, fp8 data cannot be printed directly. Read it as uint8 and manually convert it to fp8 through binary conversion.
    ret = aclrtMemcpy(level1ScaleOutData.data(), level1ScaleOutData.size() * sizeof(level1ScaleOutData[0]), level1ScaleOutDeviceAddr,
                      size * sizeof(level1ScaleOutData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy level1ScaleOut from device to host failed. ERROR: %d\n", ret);
              return ret);
    for (int64_t i = 0; i < size; i++) {
        LOG_PRINT("level1ScaleOut[%ld] is: %d\n", i, level1ScaleOutData[i]);
    }
    return ACL_SUCCESS;
}

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

    Finalize(deviceId, stream);
    return 0;
}