AscendQuant

产品支持情况

功能说明

按元素做量化计算，比如将half/float数据类型量化为int8_t数据类型。计算公式如下，round表示四舍六入五成双取整，cast表示舍入模式：

• PER_TENSOR量化：整个srcTensor对应一个量化参数，量化参数的shape为[1]。

  

• PER_CHANNEL量化：srcTensor的shape为[m, n], 每个channel维度对应一个量化参数，量化参数的shape为[n]。

  

• PER_TOKEN量化：srcTensor的每组token（token为n方向，共有m组token）中的元素共享一组scale和offset参数，srcTensor的shape为[m, n]时，scale和offset的shape为[m, 1]。offset是可选输入。
• 
• PER_GROUP量化：这里定义group的计算方向为k方向，srcTensor在k方向上每groupSize个元素共享一组scale和offset。srcTensor的shape为[m, n]时，如果kDim=0，表示k是m方向，scale和offset的shape为[(m + groupSize - 1) / groupSize, n]；如果kDim=1，表示k是n方向，scale和offset的shape为[m，(n + groupSize - 1) / groupSize]。offset是可选输入。
  根据输出数据类型的不同，当前PER_GROUP分为两种场景：fp4x2_e2m1_t/fp4x2_e1m2_t场景（后续内容中简称为float4场景）和int8_t/hifloat8_t/fp8_e5m2_t/fp8_e4m3fn_t场景（后续内容中简称为b8场景）。

  • fp4x2_e2m1_t/float4_e1m2场景（float4场景）
    • kDim = 0:

      

    • kDim = 1:

      

  • int8_t/hifloat8_t/fp8_e5m2_t/fp8_e4m3fn_t场景（b8场景）
    • kDim=0：

      

    • kDim=1：

      

实现原理

图1 AscendQuant算法框图scale和offset都是scalar

图2 AscendQuant算法框图scale和offset都是Tensor

图3 AscendQuant算法框图scale是Tensor&offset是Scalar

如上图所示是AscendQuant内部算法框图，计算过程大致描述为如下几步，均在Vector上进行：

1. 精度转换：当输入的src，scale或者offset是float类型时，将其转换为half类型；
2. broadcast：当输入的scale或者offset是向量时，将其broadcast成和src相同维度；
3. 计算scale：当src和scale为向量时做Mul计算，当scale是scalar时做Muls计算，得到Tmp1；
4. 计算offset：当Tmp1和offset为向量时做Add计算，当offset是scalar时做Adds计算，得到Tmp2；
5. 精度转换：将Tmp2从half转换成int8_t类型，得到output。

图4 AscendQuant算法框图PER_TOKEN/PER__GROUP场景scale和offset都是tensor

图5 AscendQuant算法框图PER_TOKEN/PER__GROUP场景scale是tensor&offset是scalar

PER_TOKEN/PER_GROUP场景的计算逻辑如下：

1. 读取数据：连续读取输入src；根据不同的场景，对输入scale和offset，采用不同的读取方式；例如，PER_TOKEN场景做Broadcast处理，PER_GROUP场景做Gather处理；
2. 精度转换：根据不同输入的数据类型组合，对src/scale/offset进行相应的数据类型转换；
3. 计算：对类型转换后的数据做乘加操作；
4. 精度转换：将上述乘加操作得到的计算结果转换成dstT类型，得到最终输出。

函数原型

• dstTensor为int8_t数据类型
  • PER_TENSOR量化：
    • 通过sharedTmpBuffer入参传入临时空间
      • 源操作数Tensor全部/部分参与计算

        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)

      • 源操作数Tensor全部参与计算

        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)

    • 接口框架申请临时空间
      • 源操作数Tensor全部/部分参与计算

        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)

      • 源操作数Tensor全部参与计算

        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量化：
    • 通过sharedTmpBuffer入参传入临时空间
      • 源操作数Tensor全部/部分参与计算

        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)

      • 源操作数Tensor全部参与计算

        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)

    • 接口框架申请临时空间
      • 源操作数Tensor全部/部分参与计算

        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)

      • 源操作数Tensor全部参与计算

        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非固定数据类型
  仅支持Atlas 350 加速卡。

  • PER_TENSOR量化：
    • 通过sharedTmpBuffer入参传入临时空间
      • 源操作数Tensor全部/部分参与计算

        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)

      • 源操作数Tensor全部参与计算

        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)

    • 接口框架申请临时空间
      • 源操作数Tensor全部/部分参与计算

        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)

      • 源操作数Tensor全部参与计算

        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量化：
    • 通过sharedTmpBuffer入参传入临时空间
      • 源操作数Tensor全部/部分参与计算

        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)

      • 源操作数Tensor全部参与计算

        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)

    • 接口框架申请临时空间
      • 源操作数Tensor全部/部分参与计算

        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)

      • 源操作数Tensor全部参与计算

        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/PER_GROUP量化：
    • 通过sharedTmpBuffer入参传入临时空间
      • offset操作数类型为Tensor

        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操作数类型为scalar

        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)

    • 接口框架申请临时空间
      • offset操作数类型为Tensor

        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操作数类型为scalar

        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)

由于该接口的内部实现中涉及复杂的数学计算，需要额外的临时空间来存储计算过程中的中间变量。临时空间支持接口框架申请和开发者通过sharedTmpBuffer入参传入两种方式。

• 接口框架申请临时空间，开发者无需申请，但是需要预留临时空间的大小。

• 通过sharedTmpBuffer入参传入，使用该tensor作为临时空间进行处理，接口框架不再申请。该方式开发者可以自行管理sharedTmpBuffer内存空间，并在接口调用完成后，复用该部分内存，内存不会反复申请释放，灵活性较高，内存利用率也较高。

接口框架申请的方式，开发者需要预留临时空间；通过sharedTmpBuffer传入的情况，开发者需要为sharedTmpBuffer申请空间。临时空间大小BufferSize的获取方式如下：通过GetAscendQuantMaxMinTmpSize中提供的GetAscendQuantMaxMinTmpSize接口获取需要预留空间的范围大小。

需要注意的是，在PER_TOKEN/PER_GROUP量化场景，内部实现不需要临时空间Buffer，对应的接口中sharedTempBuffer为预留参数。

参数说明

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, // 预留参数，暂不支持
        PER_CHANNEL, // 预留参数，暂不支持
        PER_TOKEN, // 配置为PER_TOKEN场景
        PER_GROUP,  // 配置为PER__GROUP场景
        PER_CHANNEL_PER_GROUP, // 预留参数，暂不支持
        PER_TOKEN_PER_GROUP // 预留参数，暂不支持
}

1
2
3
4
5
6

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

返回值说明

无

约束说明

• 源操作数与目的操作数可以复用。
• 操作数地址对齐要求请参见通用地址对齐约束。
• 输入输出操作数参与计算的数据长度要求32B对齐。
• 当Scale为float类型时，其取值范围仍为half类型的取值范围。
• Atlas 训练系列产品 仅支持PER_TENSOR量化，不支持PER_CHANNEL量化。
• PER_TOKEN/PER_GROUP场景仅在Atlas 350 加速卡上支持。
• PER_TOKEN/PER_GROUP场景，连续计算方向（即n方向）的数据量要求32B对齐。

调用示例

完整的算子样例请参考quant算子样例。

• PER_TENSOR量化场景调用示例如下：

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

  // dstLocal: 存放量化计算的结果Tensor，shape为1024 // srcLocal: 存放量化计算的输入Tensor，shape为1024，类型为float/half // sharedTmpBuffer: 存放量化计算过程中临时缓存的Tensor const float scale = 0.02; // 量化参数 const float offset = 0.9; // 量化参数，dstLocal[i] = srcLocal[i] * scale + offset uint32_t calCount = 1022; // srcTensor的前calCount个于元素参与计算 // dstTensor为int8_t数据类型，通过sharedTmpBuffer入参传入临时空间 AscendC::AscendQuant<srcType>(dstLocal, srcLocal, sharedTmpBuffer, scale, offset, calCount); // dstTensor非固定数据类型，dstTensor非固定数据类型 AscendC::AscendQuant<dstType, srcType>(dstLocal, srcLocal, scale, offset, calCount);

  结果示例如下：

  1 2 3 4

  输入数据（srcLocal）: [-512. -511. -510. ... 509. 510. 511.] 输入量化参数（scale）: 0.02 输入量化参数（offset）: 0.9 输出数据（dstLocal）: [-9 -9 -9 ... 11 51. 51.1]

• PER_CHANNEL量化场景调用示例如下：

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

  // dstLocal: 存放量化计算的结果Tensor，shape为1024 // srcLocal: 存放量化计算的输入Tensor，shape为1024，类型为float/half // scaleLocal：存放量化参数的输入Tensor // offsetLocal：存放量化参数的输入Tensor // sharedTmpBuffer: 存放量化计算过程中临时缓存的Tensor uint32_t scaleCount = 64; // 量化参数，要求是32的整数倍 uint32_t offsetCount = 64; // 量化参数，要求与scaleCount相等 uint32_t calCount = 1022; // srcTensor的前calCount个于元素参与计算 // dstTensor为int8_t数据类型，通过sharedTmpBuffer入参传入临时空间 AscendC::AscendQuant<srcType>(dstLocal, srcLocal, sharedTmpBuffer, scaleLocal, offsetLocal, scaleCount, offsetCount, calCount); // dstTensor非固定数据类型，dstTensor非固定数据类型 AscendC::AscendQuant<srcType>(dstLocal, srcLocal, scaleLocal, offsetLocal, scaleCount, offsetCount, calCount);

  结果示例如下：

  1 2 3 4

  输入数据（srcLocal）: [-512. -511. -510. ... 509. 510. 511.] 输入量化参数（scale）: [0.02 0.02 0.02 ... 0.02] 输入量化参数（offset）: [1.01 1.02 1.03 ... 1.32] 输出数据（dstLocal）: [-9 -9 -9 ... 11 510. 511.]

PER_TOKEN/PER_GROUP场景调用示例如下。

• 未配置参数AscendQuantConfig的舍入模式roundMode，使用默认配置RoundMode::CAST_RINT。

  1 2 3 4 5 6 7 8 9 10

  // 注意m,n需从外部传入 constexpr static bool isReuseSource = false; constexpr static AscendQuantConfig config = {has_offset, 1}; constexpr static AscendQuantPolicy policy = AscendQuantPolicy::PER_TOKEN; // 可修改枚举值以使能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);

• 主动配置参数AscendQuantConfig的舍入模式roundMode。

  1 2 3 4 5 6 7 8 9 10

  // 注意m,n需从外部传入 constexpr static bool isReuseSource = false; constexpr static AscendQuantConfig config = {has_offset, 1, RoundMode::CAST_ROUND}; constexpr static AscendQuantPolicy policy = AscendQuantPolicy::PER_TOKEN; // 可修改枚举值以使能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);