AscendAntiQuant

Applicability

Function Usage

Performs fake quantization by element. For example, apply fake quantization to convert the int8_t data type to the half data type. The formulas are as follows:

• PER_CHANNEL scenario (quantization by channel)
  • Input transposing disabled

    groupSize = src.shape[0]/offset.shape[0]

    dst[i][j] = scale[i/groupSize][j] x (src[i][j] + offset[i/groupSize][j])

  • Input transposing enabled

    groupSize = src.shape[1]/offset.shape[1]

    dst[i][j] = scale[i][j/groupSize] x (src[i][j] + offset[i][j/groupSize])

• PER_TENSOR scenario (quantization by tensor)

  dst[i][j] = scale x (src[i][j] + offset)

• PER_TOKEN scenario (quantization by token)

  

• PER_GROUP scenario (quantization by group)

  Based on the input data type, PER_GROUP is classified into two scenarios: fp4x2_e2m1_t/fp4x2_e1m2_t scenario (referred to as the float4 scenario) and int8_t/hifloat8_t/fp8_e5m2_t/fp8_e4m3fn_t scenario (referred to as the b8 scenario).

  • fp4x2_e2m1_t/fp4x2_e1m2_t scenario (float4 scenario)
    • Configurable groupSize API
      The compute direction of a group is defined as k. In the k direction, each groupSize element of src shares scale. When the shape of src is [m, n], if kDim is 0, k is in the m direction, and the shape of scale is [(m + groupSize – 1)/groupSize, n]; if kDim is 1, k is in the n direction, and the shape of scale is [m, (n + groupSize – 1)/groupSize]. If isTranspose is set to True, src, scale, and dst are transposed matrices.

      • k is the m direction, indicating that the i axis in the formula is the compute direction of the group: (kDim is 0 and isTranspose is False) or (kDim is 1 and isTranspose is True).

        dst[i][j] = scale[i/groupSize][j] x src[i][j]

      • k is the n direction, indicating that the j axis in the formula is the compute direction of the group: (kDim is 0 and isTranspose is True) or (kDim is 1 and isTranspose is False).

        dst[i][j] = scale[i][j / groupSize] * src[i][j]

    • groupSize fixed at 32
      If isTranspose is set to True, src, scale, and dst are transposed matrices.

      • Input transposing disabled (isTranspose = False)

        dst[i][j] = scale[i/groupSize][j] x src[i][j]

      • Input transposing enabled (isTranspose = True)

        dst[i][j] = scale[i][j/groupSize] x src[i][j]

  • int8_t/hifloat8_t/fp8_e5m2_t/fp8_e4m3fn_t scenario (b8 scenario)
    The compute direction of a group is defined as k. In the k direction, each groupSize element of src shares a group of scale and offset. When the shape of src is [m, n], if kDim is 0, k is in the m direction, and the shape of scale and offset is [(m + groupSize – 1)/groupSize, n]; if kDim is 1, k is in the n direction, and the shape of scale and offset is [m, (n + groupSize – 1)/groupSize]. offset is an optional input. If isTranspose is set to True, src, scale, and dst are transposed matrices.

    • k is the m direction, indicating that the i axis in the formula is the compute direction of the group: (kDim is 0 and isTranspose is False) or (kDim is 1 and isTranspose is True).

      

    • k is the n direction, indicating that the j axis in the formula is the compute direction of the group: (kDim is 0 and isTranspose is True) or (kDim is 1 and isTranspose is False).

      

Principles

Figure 1 AscendAntiQuant algorithm block diagram

The preceding figure shows the algorithm block diagram of AscendAntiQuant in typical scenarios. The computation process is divided into the following steps, all of which are performed on vectors:

1. Precision conversion: Convert the input src to the half type.
2. Offset calculation: Perform Add calculation when offset is a vector, and perform Adds calculation when offset is a scalar.
3. Scale calculation: Perform Mul calculation when scale is a vector, and perform Muls calculation when scale is a scalar.

Figure 2 AscendAntiQuant algorithm block diagram when isTranspose is False and the output is of the bfloat16 type

On the Atlas A2 training product / Atlas A2 inference product , when the output is of the bfloat16 type, the computation process is divided into the following steps:

1. src precision conversion: Convert the input src to the half type, and then convert it to the float type and store it in tmp1.
2. Offset precision conversion: When the input offset is a vector, it is converted to the float type and stored in tmp2; when the input offset is a scalar, it is converted to the float type through Cast.
3. Offset calculation: When the input offset is a vector, perform the Add calculation with tmp2. When it is a scalar, perform the Adds calculation.
4. Scale precision conversion: When the input scale is a vector, it is converted to the float type and stored in tmp2; when the input scale is a scalar, it is converted to the float type through Cast.
5. Scale calculation: When the input scale is a vector, perform the Mul calculation with tmp2. When it is a scalar, perform the Muls calculation.
6. dst precision conversion: Convert tmp1 to the bf16 type.

Prototype

• Pass the temporary space through the sharedTmpBuffer input parameter.
  • PER_CHANNEL scenario (quantization by channel)

    1 2

    template <typename InputDataType, typename OutputDataType, bool isTranspose> __aicore__ inline void AscendAntiQuant(const LocalTensor<OutputDataType>& dst, const LocalTensor<InputDataType>& src, const LocalTensor<OutputDataType>& offset, const LocalTensor<OutputDataType>& scale, const LocalTensor<uint8_t>& sharedTmpBuffer, const uint32_t k, const AntiQuantShapeInfo& shapeInfo = {})

  • PER_CHANNEL scenario (quantization by channel without offset)

    1 2

    template <typename InputDataType, typename OutputDataType, bool isTranspose> __aicore__ inline void AscendAntiQuant(const LocalTensor<OutputDataType>& dst, const LocalTensor<InputDataType>& src, const LocalTensor<OutputDataType>& scale, const LocalTensor<uint8_t>& sharedTmpBuffer, const uint32_t k, const AntiQuantShapeInfo& shapeInfo = {})

  • PER_TENSOR scenario (quantization by tensor)

    1 2

    template <typename InputDataType, typename OutputDataType, bool isTranspose> __aicore__ inline void AscendAntiQuant(const LocalTensor<OutputDataType>& dst, const LocalTensor<InputDataType>& src, const OutputDataType offset, const OutputDataType scale, const LocalTensor<uint8_t>& sharedTmpBuffer, const uint32_t k, const AntiQuantShapeInfo& shapeInfo = {})

  • PER_TENSOR scenario (quantization by tensor without offset)

    1 2

    template <typename InputDataType, typename OutputDataType, bool isTranspose> __aicore__ inline void AscendAntiQuant(const LocalTensor<OutputDataType> &dst, const LocalTensor<InputDataType> &src, const OutputDataType scale, const LocalTensor<uint8_t> &sharedTmpBuffer, const uint32_t k, const AntiQuantShapeInfo& shapeInfo = {})

  • PER_GROUP float4 scenario (quantization by group)
    It is supported only by the Atlas 350 Accelerator Card.

    1 2

    template <typename InputDataType, typename OutputDataType, bool isTranspose> __aicore__ inline void AscendAntiQuant(const LocalTensor<OutputDataType>& dst, const LocalTensor<InputDataType>& src, const LocalTensor<fp8_e8m0_t>& scale, const LocalTensor<uint8_t>& sharedTmpBuffer, const uint32_t k, const AntiQuantShapeInfo& shapeInfo = {})

  • PER_TOKEN b8/float4 (quantization by token) or PER_GROUP b8/float4 scenario (quantization by group)
    It is supported only by the Atlas 350 Accelerator Card.

    1 2

    template <typename dstT, typename srcT, typename scaleT, const AscendAntiQuantConfig& config, const AscendAntiQuantPolicy& policy> __aicore__ inline void AscendAntiQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<scaleT>& scaleTensor, const LocalTensor<scaleT>& offsetTensor,const LocalTensor<uint8_t>& sharedTmpBuffer, const AscendAntiQuantParam& para)

• Allocate the temporary space through the API framework.
  • PER_CHANNEL scenario

    1 2

    template <typename InputDataType, typename OutputDataType, bool isTranspose> __aicore__ inline void AscendAntiQuant(const LocalTensor<OutputDataType>& dst, const LocalTensor<InputDataType>& src, const LocalTensor<OutputDataType>& offset, const LocalTensor<OutputDataType>& scale, const uint32_t k, const AntiQuantShapeInfo& shapeInfo = {})

  • PER_TENSOR scenario

    1 2

    template <typename InputDataType, typename OutputDataType, bool isTranspose> __aicore__ inline void AscendAntiQuant(const LocalTensor<OutputDataType>& dst, const LocalTensor<InputDataType>& src, const OutputDataType offset, const OutputDataType scale, const uint32_t k, const AntiQuantShapeInfo& shapeInfo = {})

  • PER_GROUP float4 scenario (groupSize fixed at 32)
    It is supported only by the Atlas 350 Accelerator Card.

    1 2	template <typename InputDataType, typename OutputDataType, bool isTranspose> __aicore__ inline void AscendAntiQuant(const LocalTensor<OutputDataType>& dst, const LocalTensor<InputDataType>& src, const LocalTensor<fp8_e8m0_t>& scale, const uint32_t k, const AntiQuantShapeInfo& shapeInfo = {})

  • PER_TOKEN b8/float4 or PER_GROUP b8/float4 scenario (configurable groupSize)
    It is supported only by the Atlas 350 Accelerator Card.

    1 2

    template <typename dstT, typename srcT, typename scaleT, const AscendAntiQuantConfig& config, const AscendAntiQuantPolicy& policy> __aicore__ inline void AscendAntiQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<scaleT>& scaleTensor, const LocalTensor<scaleT>& offsetTensor,const AscendAntiQuantParam& para)

Due to the complex mathematical computation involved in the internal implementation of this API, extra temporary space is required to store intermediate variables generated during computation. The temporary space can be allocated through the API framework or passed by developers through the sharedTmpBuffer input parameter.

• When the API framework is used for temporary space allocation, developers do not need to allocate the space, but must reserve the required size for the temporary space.

• When the sharedTmpBuffer input parameter is used for passing the temporary space, the tensor serves as the temporary space. In this case, the API framework is not required for temporary space allocation. This enables developers to manage the sharedTmpBuffer space and reuse the buffer after calling the API, so that the buffer is not repeatedly allocated or deallocated, improving the flexibility and buffer utilization.

If the API framework is used, developers must reserve the temporary space. If sharedTmpBuffer is used, developers must allocate space for sharedTmpBuffer. To obtain the size of the temporary space (BufferSize) to be reserved, use the API provided in GetAscendAntiQuantMaxMinTmpSize.

Parameters

1
2
3
4
5

struct AscendAntiQuantConfig {
        bool hasOffset;
        bool isTranspose;
        int32_t kDim = 1;
 };

1
2
3
4
5
6
7
8

enum class AscendQuantPolicy : int32_t {
        PER_TENSOR // Reserved.
        PER_CHANNEL // Reserved.
        PER_TOKEN // PER_TOKEN scenario
        PER_GROUP // PER_GROUP scenario
        PER_CHANNEL_PER_GROUP // Reserved.
        PER_TOKEN_PER_GROUP // Reserved.
}

1
2
3
4
5
6

struct AntiQuantShapeInfo {
    uint32_t offsetHeight{0};  // Offset height
    uint32_t offsetWidth{0};  // Offset width
    uint32_t scaleHeight{0};  // Scale height
    uint32_t scaleWidth{0};  // Scale width
};

1
2
3
4
5
6

struct AscendAntiQuantParam {
        uint32_t m;
        uint32_t n;
        uint32_t calCount;
        uint32_t groupSize = 0;
}

Returns

None

Constraints

• The source operand address must not overlap the destination operand address.
• For details about the operand address alignment requirements, see General Address Alignment Restrictions.
• The length of the data involved in computation of the input and output operands must be 32-byte aligned.
• When inputs are transposed, k must be 32-byte aligned.
• Before calling the API, ensure that the size of the input data is correct, and the size and shape of offset and scale are correct.
• Currently, only the Atlas 350 Accelerator Card supports the PER_TOKEN and PER_GROUP scenarios.
• In the PER_TOKEN b8/float4 and PER_GROUP b8/float4 scenarios, the size of data in the continuous compute direction (n direction) must be 32-byte aligned.

Examples

For a complete call example, see AscendAntiQuant operator sample.

1
2
3
4
5
6
7
8
9

// dstLocal: result tensor
// srcLocal: quantized input
// offsetLocal: offset parameter
// scaleLocal: scaling parameter
// sharedTmpBuffer: temporary buffer managed by developers, used to store intermediate variables during internal computation
// k: length of the k axis
// shapeInfo: shape information of the offsetLocal and scaleLocal tensors
AscendC::AntiQuantShapeInfo shapeInfo = {1, elementCountOfOffset, 1, elementCountOfOffset};
AscendC::AscendAntiQuant<InputType, OutType, false>(dstLocal, srcLocal, offsetLocal, scaleLocal, sharedTmpBuffer, k, shapeInfo);

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

Input data src (shape: [2, 64], non-transposing scenario):
[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1]
offset (shape: [1, 64]):
[2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2.
 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2.
 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2.]
scale (shape: [1, 64]):
[3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3.
 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3.
 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3. 3.]
Output dstLocal (shape: [2, 64]):
[9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9.
 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9.
 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9.
 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9.
 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9. 9.
 9. 9. 9. 9. 9. 9. 9. 9.]

The following is an example of calling the API in the PER_TOKEN b8 and PER_GROUP b8 scenarios:

1
2
3
4
5
6
7
8
9

// Note that m and n must be passed from an external system.
constexpr static bool isReuseSource = false;
constexpr static AscendAntiQuantConfig config = {has_offset, has_transpose, -1};
constexpr static AscendAntiQuantPolicy policy = AscendAntiQuantPolicy::PER_TOKEN;
AscendAntiQuantParam para;
para.m = m;
para.n = n;
para.calCount = calCount;
AscendAntiQuant<dstType, srcType, scaleType, config, policy>(dstLocal, srcLocal, scaleLocal, offsetLocal, para);