Quantize

Applicability

Function Usage

Performs quantization computation by element to convert high-precision data into low-precision data. This API supports four types of quantization policies: PER_TENSOR, PER_CHANNEL, PER_TOKEN, and PER_GROUP. Each type of quantization policy supports the configuration of the rounding mode. This API supports only two-dimensional input data.

Quantize has similar function to AscendQuant. Quantize extends the function of configuring the rounding mode in the PER_TENSOR and PER_CHANNEL quantization scenarios. Therefore, you are advised to use this API.

• PER_TENSOR quantization: srcTensor corresponds to a quantization parameter. The shape of scale and offset is [1].

  

• PER_CHANNEL quantization: The shape of srcTensor is [m, n]. Each channel dimension corresponds to a quantization parameter. The shape of scale and offset is [1, n].

  

• PER_TOKEN quantization: The elements in each group of tokens (m groups of tokens in the n direction) of srcTensor share a quantization parameter. When the shape of srcTensor is [m, n], the shape of scale and offset is [m, 1].

  

• PER_GROUP quantization: The compute direction of groups is defined as k. In the k direction, each groupSize elements of srcTensor shares a group of scale and offset. When the shape of srcTensor 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].
  Based on the output 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)
    • kDim = 0:

      

    • kDim = 1:

      

  • int8_t/hifloat8_t/fp8_e5m2_t/fp8_e4m3fn_t scenario (b8 scenario)
    • kDim = 0:

      

    • kDim = 1:

      

Prototype

• Pass the temporary space through the sharedTmpBuffer input parameter.

  1 2

  template <const QuantizeConfig& config, typename DstT, typename SrcT, typename ScaleT, typename OffsetT> __aicore__ inline void Quantize(const LocalTensor<DstT>& dstTensor, const LocalTensor<SrcT>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const ScaleT& scale, const OffsetT& offset, const QuantizeParams& params)

• Allocate the temporary space through the API framework.

  1 2	template <const QuantizeConfig& config, typename DstT, typename SrcT, typename ScaleT, typename OffsetT> __aicore__ inline void Quantize(const LocalTensor<DstT>& dstTensor, const LocalTensor<SrcT>& srcTensor,const ScaleT& scale, const OffsetT& offset, const QuantizeParams& params)

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

Parameters

1
2
3
4
5
6

struct QuantizeConfig {
    QuantizePolicy policy;
    bool hasOffset;
    RoundMode roundMode = RoundMode::CAST_RINT;
    int32_t kDim = 1;
 }

1
2
3
4
5

struct QuantizeParams {
        uint32_t m;
        uint32_t n;
        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.
• The data volume in the continuous computation direction (n direction) must be 32-byte aligned.
• In the PER_GROUP float4 scenario, offset is not supported. In this scenario, the hasOffset parameter in the template parameter config must be set to false.

Examples

• PER_CHANNEL, PER_TOKEN, and PER_GROUP modes

  1 2 3 4 5 6 7 8 9 10 11 12

  constexpr static QuantizePolicy tokenPolicy = QuantizePolicy::PER_TOKEN; constexpr static QuantizePolicy channelPolicy = QuantizePolicy::PER_CHANNEL; constexpr static QuantizePolicy groupPolicy = QuantizePolicy::PER_GROUP; // The PER_TOKEN mode is used as an example. Enable the offset function and set the rounding mode to CAST_ROUND. kDim is valid only in the PER_GROUP scenario, indicating that the compute direction of groups is n. constexpr static QuantizeConfig config = {tokenPolicy, true, RoundMode::CAST_ROUND, 1}; QuantizeParams params; // m and n are external input parameters, indicating the number of elements in the m and n directions involved in srcLocal. params.m = m; params.n = n; params.groupSize = n; // This is valid only in the PER_GROUP scenario. It indicates that all elements in the n direction share a group of scale and offset. // dstLocal is a LocalTensor of the int8_t type, and srcLocal, scale, and offset are a LocalTensor of the half type. Quantize<config>(dstLocal, srcLocal, scale, offset, params);

• PER_TENSOR mode

  1 2 3 4 5 6 7 8 9 10

  constexpr static QuantizePolicy tensorPolicy = QuantizePolicy::PER_TENSOR; // Enable the offset function and set the rounding mode to CAST_ROUND. constexpr static QuantizeConfig config = {tensorPolicy, true, RoundMode::CAST_ROUND, -1}; QuantizeParams params; // m and n are external input parameters, indicating the number of elements in the m and n directions involved in srcLocal. params.m = m; params.n = n; params.groupSize = 0; // This is valid only in the PER_GROUP scenario. // dstLocal is a LocalTensor of the int8_t type, srcLocal is a LocalTensor of the half type, and scale and offset are scalars of the half type. Quantize<config>(dstLocal, srcLocal, scale, offset, params);

Result example:

Input (srcLocal):
[-4.4, 2.5, -2.9, -3.1, -1.5, -4.8, 1.8, 3.5, 4.5, 1.1, -2.7, 0.5, ... 1.6]
Input (scale vector):
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ... 1]
Input (scale scalar):
[1]
Input (offset vector):
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ... 0]
Input (offset scalar):
[0]
Output (dstLocal):
[-4, 3, -3, -3, -1, -5, 2, 4, 5, 1, -3, 1, ... 2]