AscendQuant

Applicability

Function Usage

Performs quantization by element. For example, quantize the half/float data type to the int8_t data type. The following is the formula, where round indicates rounding to the nearest even numbercast indicates the rounding mode.

• PER_TENSOR quantization: srcTensor corresponds to a quantization parameter, whose shape is [1].

  

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

  

• PER_TOKEN quantization: The elements in each group of tokens (m groups of tokens in the n direction) of srcTensor share a group of scale and offset. When the shape of srcTensor is [m, n], the shape of scale and offset is [m, 1]. offset is an optional input.
• 
• 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]. offset is an optional input.
  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/float4_e1m2 scenario (float4 scenario)
    • kDim = 0:

      

    • kDim = 1:

      

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

      

    • kDim = 1:

      

Principles

Figure 1 AscendQuant algorithm block diagram, where both scale and offset are scalars

Figure 2 AscendQuant algorithm block diagram, where both scale and offset are tensors

Figure 3 AscendQuant algorithm block diagram, where scale is a tensor and offset is a scalar

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

1. Precision conversion: If the input src, scale, or offset is of the float type, convert it to the half type.
2. Broadcast: If the input scale or offset is a vector, broadcast it to the same dimension as src.
3. Scale calculation: If src and scale are vectors, Mul calculation is performed. If scale is a scalar, Muls calculation is performed to obtain Tmp1.
4. Offset calculation: If Tmp1 and offset are vectors, Add calculation is performed. If offset is a scalar, Adds calculation is performed to obtain Tmp2.
5. Precision conversion: Convert Tmp2 from half to int8_t to obtain the output.

Figure 4 AscendQuant algorithm diagram in the PER_TOKEN and PER_GROUP scenarios where both scale and offset are tensors

Figure 5 AscendQuant algorithm diagram in the PER_TOKEN and PER_GROUP scenarios where scale is a tensor and offset is a scalar

The computation logic in the PER_TOKEN and PER_GROUP scenarios is as follows:

1. Data reading: Read the input src continuously, and use different modes to read the input scale and offset based on actual scenarios. For example, the input data is broadcast in the PER_TOKEN scenario, and gathered in the PER_GROUP scenario.
2. Precision conversion: Convert the data types of src, scale, and offset based on the input data type combination.
3. Computation: Perform the multiplication and addition operations on the data after type conversion.
4. Precision conversion: Convert the computation result to the dstT type as the final output.

Prototype

• dstTensor of the int8_t type
  • PER_TENSOR quantization:
    • Pass the temporary space through the sharedTmpBuffer input parameter.
      • All or part of the source operand tensors are involved in computation.

        1 2

        template <typename T, bool isReuseSource = false, const AscendQuantConfig& config = ASCEND_QUANT_DEFAULT_CFG> __aicore__ inline void AscendQuant(const LocalTensor<int8_t>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const float scale, const float offset, const uint32_t calCount)

      • All source operand tensors are involved in computation.

        1 2	template <typename T, bool isReuseSource = false, const AscendQuantConfig& config = ASCEND_QUANT_DEFAULT_CFG> __aicore__ inline void AscendQuant(const LocalTensor<int8_t>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const float scale, const float offset)

    • Allocate the temporary space through the API framework.
      • All or part of the source operand tensors are involved in computation.

        1 2	template <typename T, bool isReuseSource = false, const AscendQuantConfig& config = ASCEND_QUANT_DEFAULT_CFG> __aicore__ inline void AscendQuant(const LocalTensor<int8_t>& dstTensor, const LocalTensor<T>& srcTensor, const float scale, const float offset, const uint32_t calCount)

      • All source operand tensors are involved in computation.

        1 2	template <typename T, bool isReuseSource = false, const AscendQuantConfig& config = ASCEND_QUANT_DEFAULT_CFG> __aicore__ inline void AscendQuant(const LocalTensor<int8_t>& dstTensor, const LocalTensor<T>& srcTensor, const float scale, const float offset)

  • PER_CHANNEL quantization:
    • Pass the temporary space through the sharedTmpBuffer input parameter.
      • All or part of the source operand tensors are involved in computation.

        1 2

        template <typename T, bool isReuseSource = false, const AscendQuantConfig& config = ASCEND_QUANT_DEFAULT_CFG> __aicore__ inline void AscendQuant(const LocalTensor<int8_t>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const LocalTensor<T>& scaleTensor, const T offset, const uint32_t scaleCount, const uint32_t calCount)

        1 2

        template <typename T, bool isReuseSource = false, const AscendQuantConfig& config = ASCEND_QUANT_DEFAULT_CFG> __aicore__ inline void AscendQuant(const LocalTensor<int8_t>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const LocalTensor<T>& scaleTensor, const LocalTensor<T>& offsetTensor, const uint32_t scaleCount, const uint32_t offsetCount, const uint32_t calCount)

      • All source operand tensors are involved in computation.

        1 2

        template <typename T, bool isReuseSource = false, const AscendQuantConfig& config = ASCEND_QUANT_DEFAULT_CFG> __aicore__ inline void AscendQuant(const LocalTensor<int8_t>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const LocalTensor<T>& scaleTensor, const T offset)

        1 2

        template <typename T, bool isReuseSource = false, const AscendQuantConfig& config = ASCEND_QUANT_DEFAULT_CFG> __aicore__ inline void AscendQuant(const LocalTensor<int8_t>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const LocalTensor<T>& scaleTensor, const LocalTensor<T>& offsetTensor)

    • Allocate the temporary space through the API framework.
      • All or part of the source operand tensors are involved in computation.

        1 2

        template <typename T, bool isReuseSource = false, const AscendQuantConfig& config = ASCEND_QUANT_DEFAULT_CFG> __aicore__ inline void AscendQuant(const LocalTensor<int8_t>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<T>& scaleTensor, const T offset, const uint32_t scaleCount, const uint32_t calCount)

        1 2

        template <typename T, bool isReuseSource = false, const AscendQuantConfig& config = ASCEND_QUANT_DEFAULT_CFG> __aicore__ inline void AscendQuant(const LocalTensor<int8_t>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<T>& scaleTensor, const LocalTensor<T>& offsetTensor, const uint32_t scaleCount, const uint32_t offsetCount, const uint32_t calCount)

      • All source operand tensors are involved in computation.

        1 2	template <typename T, bool isReuseSource = false, const AscendQuantConfig& config = ASCEND_QUANT_DEFAULT_CFG> __aicore__ inline void AscendQuant(const LocalTensor<int8_t>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<T>& scaleTensor, const T offset)

        1 2	template <typename T, bool isReuseSource = false, const AscendQuantConfig& config = ASCEND_QUANT_DEFAULT_CFG> __aicore__ inline void AscendQuant(const LocalTensor<int8_t>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<T>& scaleTensor, const LocalTensor<T>& offsetTensor)

• dstTensor of non-fixed data type
  It is supported only by the Atlas 350 Accelerator Card.

  • PER_TENSOR quantization:
    • Pass the temporary space through the sharedTmpBuffer input parameter.
      • All or part of the source operand tensors are involved in computation.

        1 2	template <typename dstT, typename srcT, bool isReuseSource = false> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const float scale, const float offset, const uint32_t calCount)

      • All source operand tensors are involved in computation.

        1 2	template <typename dstT, typename srcT, bool isReuseSource = false> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const float scale, const float offset)

    • Allocate the temporary space through the API framework.
      • All or part of the source operand tensors are involved in computation.

        1 2	template <typename dstT, typename srcT, bool isReuseSource = false> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const float scale, const float offset, const uint32_t calCount)

      • All source operand tensors are involved in computation.

        1 2	template <typename dstT, typename srcT, bool isReuseSource = false> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const float scale, const float offset)

  • PER_CHANNEL quantization:
    • Pass the temporary space through the sharedTmpBuffer input parameter.
      • All or part of the source operand tensors are involved in computation.

        1 2

        template <typename dstT, typename srcT, bool isReuseSource = false> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const LocalTensor<srcT>& scaleTensor, const srcT offset, const uint32_t scaleCount, const uint32_t calCount)

        1 2

        template <typename dstT, typename srcT, bool isReuseSource = false> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const LocalTensor<srcT>& scaleTensor, const LocalTensor<srcT>& offsetTensor, const uint32_t scaleCount, const uint32_t offsetCount, const uint32_t calCount)

      • All source operand tensors are involved in computation.

        1 2	template <typename dstT, typename srcT, bool isReuseSource = false> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const LocalTensor<srcT>& scaleTensor, const srcT offset)

        1 2	template <typename dstT, typename srcT, bool isReuseSource = false> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const LocalTensor<srcT>& scaleTensor, const LocalTensor<srcT>& offsetTensor)

    • Allocate the temporary space through the API framework.
      • All or part of the source operand tensors are involved in computation.

        1 2	template <typename dstT, typename srcT, bool isReuseSource = false> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<srcT>& scaleTensor, const srcT offset, const uint32_t scaleCount, const uint32_t calCount)

        1 2

        template <typename dstT, typename srcT, bool isReuseSource = false> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<srcT>& scaleTensor, const LocalTensor<srcT>& offsetTensor, const uint32_t scaleCount, const uint32_t offsetCount, const uint32_t calCount)

      • All source operand tensors are involved in computation.

        1 2	template <typename dstT, typename srcT, bool isReuseSource = false> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<srcT>& scaleTensor, const srcT offset)

        1 2	template <typename dstT, typename srcT, bool isReuseSource = false> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<srcT>& scaleTensor, const LocalTensor<srcT>& offsetTensor)

  • PER_TOKEN or PER_GROUP quantization:
    • Pass the temporary space through the sharedTmpBuffer input parameter.
      • offset operand of tensor type

        1 2

        template <typename dstT, typename srcT, typename scaleT, bool isReuseSource = false, const AscendQuantConfig& config, const AscendQuantPolicy& policy> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const LocalTensor<scaleT>& scaleTensor, const LocalTensor<scaleT>& offsetTensor, const AscendQuantParam& para)

      • offset operand of scalar type

        1 2

        template <typename dstT, typename srcT, typename scaleT, bool isReuseSource = false, const AscendQuantConfig& config, const AscendQuantPolicy& policy> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<uint8_t>& sharedTmpBuffer, const LocalTensor<scaleT>& scaleTensor,const scaleT offset, const AscendQuantParam& para)

    • Allocate the temporary space through the API framework.
      • offset operand of tensor type

        1 2

        template <typename dstT, typename srcT, typename scaleT, bool isReuseSource = false, const AscendQuantConfig& config, const AscendQuantPolicy& policy> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<scaleT>& scaleTensor, const LocalTensor<scaleT>& offsetTensor, const AscendQuantParam& para)

      • offset operand of scalar type

        1 2

        template <typename dstT, typename srcT, typename scaleT, bool isReuseSource = false, const AscendQuantConfig& config, const AscendQuantPolicy& policy> __aicore__ inline void AscendQuant(const LocalTensor<dstT>& dstTensor, const LocalTensor<srcT>& srcTensor, const LocalTensor<scaleT>& scaleTensor, const scaleT offset, const AscendQuantParam& 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 GetAscendQuantMaxMinTmpSize API provided in GetAscendQuantMaxMinTmpSize.

Note that in the PER_TOKEN and PER_GROUP quantization scenarios, a temporary buffer is not required for the internal implementation, and the sharedTmpBuffer parameter is reserved in the corresponding API.

Parameters

1
2
3
4
5
6

struct AscendQuantConfig{
uint32_t calcCount = 0;
uint32_t offsetCount = 0;
uint32_t scaleCount = 0;
uint32_t workLocalSize = 0;
};

1

constexpr AscendQuantConfig ASCEND_QUANT_DEFAULT_CFG = {0, 0, 0, 0};

1
2
3
4
5

struct AscendQuantConfig {
        bool hasOffset;
        int32_t kDim = 1;
        RoundMode roundMode = RoundMode::CAST_RINT;
}

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 AscendQuantParam {
        uint32_t m;
        uint32_t n; 
        uint32_t calCount;  
        uint32_t groupSize = 0;
}

Returns

None

Constraints

• The source operand and destination operand can be reused.
• 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 scale is of the float type, its value range is still the range of values of the half type.
• The Atlas training product supports only PER_TENSOR quantization. It does not support PER_CHANNEL quantization.
• Only the Atlas 350 Accelerator Card supports the PER_TOKEN and PER_GROUP scenarios.
• In the PER_TOKEN and PER_GROUP scenarios, the size of data in the continuous compute direction (n direction) must be 32-byte aligned.

Examples

For a complete operator sample, see quant operator sample.

• The following is an example of calling the API in the PER_TENSOR quantization scenario:

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

  // dstLocal: tensor for storing the quantization computation result, with the shape of 1024 // srcLocal: tensor for storing the quantization computation input, with the shape of 1024 and the type of float or half // sharedTmpBuffer: tensor for storing the temporary buffer during quantization computation const float scale = 0.02; // quantization parameter const float offset = 0.9; // quantization parameter, dstLocal[i] = srcLocal[i] x scale + offset uint32_t calCount = 1022; // the first calCount elements of srcTensor involved in the computation // dstTensor of the int8_t data type, whose temporary space is passed through the sharedTmpBuffer input parameter AscendC::AscendQuant<srcType>(dstLocal, srcLocal, sharedTmpBuffer, scale, offset, calCount); // dstTensor of non-fixed data type AscendC::AscendQuant<dstType, srcType>(dstLocal, srcLocal, scale, offset, calCount);

  Result example:

  1 2 3 4

  Input (srcLocal): [-512. -511. -510. ... 509. 510. 511.] Input quantization parameter (scale): 0.02 Input quantization parameter (offset): 0.9 Output (dstLocal): [-9 -9 -9 ... 11 51. 51.1]

• The following is an example of calling the API in the PER_CHANNEL quantization scenario:

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

  // dstLocal: tensor for storing the quantization computation result, with the shape of 1024 // srcLocal: tensor for storing the quantization computation input, with the shape of 1024 and the type of float or half // scaleLocal: input tensor for storing quantization parameters // offsetLocal: input tensor for storing quantization parameters // sharedTmpBuffer: tensor for storing the temporary buffer during quantization computation uint32_t scaleCount = 64; // quantization parameter, which must be an integer multiple of 32 uint32_t offsetCount = 64; // quantization parameter, which must be the same as the value of scaleCount uint32_t calCount = 1022; // the first calCount elements of srcTensor involved in the computation // dstTensor of the int8_t data type, whose temporary space is passed through the sharedTmpBuffer input parameter AscendC::AscendQuant<srcType>(dstLocal, srcLocal, sharedTmpBuffer, scaleLocal, offsetLocal, scaleCount, offsetCount, calCount); // dstTensor of non-fixed data type AscendC::AscendQuant<srcType>(dstLocal, srcLocal, scaleLocal, offsetLocal, scaleCount, offsetCount, calCount);

  Result example:

  1 2 3 4

  Input (srcLocal): [-512. -511. -510. ... 509. 510. 511.] Input quantization parameter (scale): [0.02 0.02 0.02 ... 0.02] Input quantization parameter (offset): [1.01 1.02 1.03 ... 1.32] Output (dstLocal): [-9 -9 -9 ... 11 510. 511.]

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

• If roundMode is not set for the AscendQuantConfig parameter, use the default value RoundMode::CAST_RINT.

  1 2 3 4 5 6 7 8 9 10

  // Note that m and n must be passed from an external system. constexpr static bool isReuseSource = false; constexpr static AscendQuantConfig config = {has_offset, 1}; constexpr static AscendQuantPolicy policy = AscendQuantPolicy::PER_TOKEN; // The enumerated value can be modified to enable PER_GROUP. LocalTensor<uint8_t> sharedTmpBuffer = inQueue.AllocTensor<uint8_t>(); AscendQuantParam para; para.m = m; para.n = n; para.calCount = calCount; AscendQuant<dstType, srcType, scaleType, isReuseSource, config, policy>(dstLocal, srcLocal, sharedTmpBuffer, scaleLocal, offsetLocal, para);

• Set roundMode for the AscendQuantConfig parameter proactively.

  1 2 3 4 5 6 7 8 9 10

  // Note that m and n must be passed from an external system. constexpr static bool isReuseSource = false; constexpr static AscendQuantConfig config = {has_offset, 1, RoundMode::CAST_ROUND}; constexpr static AscendQuantPolicy policy = AscendQuantPolicy::PER_TOKEN; // The enumerated value can be modified to enable PER_GROUP. LocalTensor<uint8_t> sharedTmpBuffer = inQueue.AllocTensor<uint8_t>(); AscendQuantParam para; para.m = m; para.n = n; para.calCount = calCount; AscendQuant<dstType, srcType, scaleType, isReuseSource, config, policy>(dstLocal, srcLocal, sharedTmpBuffer, scaleLocal, offsetLocal, para);