torch_npu.npu_dynamic_quant

功能描述

• 算子功能：为输入的张量进行pre-token对称动态量化。
• 计算公式：
  • 如果smooth_scales不存在：

    

  • 如果smooth_scales存在：

    

接口原型

torch_npu.npu_dynamic_quant(Tensor x, Tensor? smooth_scales=None) ->(Tensor, Tensor)

参数说明

• x：Device侧的Tensor类型，需要进行量化的源数据张量，必选输入，数据类型支持FLOAT16、BFLOAT16，数据格式支持ND，支持非连续的Tensor。

• smooth_scales：Device侧的Tensor类型，对x进行scales的张量，可选输入，数据类型支持FLOAT16、BFLOAT16，数据格式支持ND，支持非连续的Tensor。

输出说明

两个Tensor类型的输出：

• 第一个输出为Tensor，量化后的输出y，shape和输入x一致，数据类型支持INT8。
• 第二个输出为Tensor，用于后面反量化的scale，shape为x剔除最后一维，数据类型支持FLOAT32。

约束说明

• 该融合算子仅在推理场景使用。
• 输入x的维度必须大于1，针对如下产品，x的最后一个维度大小不超过11264：
  • Atlas A2训练系列产品/Atlas 800I A2推理产品
• 输入smooth_scales必须是关键字传参，即调用接口时需要显式地写smooth_scales=xxx。

支持的型号

• Atlas A2训练系列产品/Atlas 800I A2推理产品

调用示例

• 只有一个输入x

  import torch
  import torch_npu
   
  x = torch.rand((3, 3), dtype = torch.float16).to("npu")
  output, scale = torch_npu.npu_dynamic_quant(x)
  print(output)
  print(scale)

• 使用smooth_scales输入

  import torch
  import torch_npu
   
  x = torch.rand((3, 3), dtype = torch.float16).to("npu")
  smooth_scales = torch.rand((3,), dtype = torch.float16).to("npu")
  output, scale = torch_npu.npu_dynamic_quant(x, smooth_scales=smooth_scales)
  print(output)
  print(scale)