swiglu_group_quant

产品支持情况

• Ascend 950PR/Ascend 950DT：支持

• Atlas A3 训练系列产品/Atlas A3 推理系列产品：不支持

• Atlas A2 训练系列产品/Atlas A2 推理系列产品：不支持

• Atlas 200I/500 A2 推理产品：不支持

• Atlas 推理系列产品：不支持

• Atlas 训练系列产品：不支持

功能说明

• 接口功能：

  对输入x执行SwiGLU激活后进行分组低比特量化，支持Block FP8、MX FP8、MX FP4、HiFloat8静态量化和HiFloat8动态量化。底层封装 aclnnSwigluGroupQuant。

• 计算流程：

  1. 将x的最后一维均分为A、B两部分，计算y_origin = silu(A) * B。
  2. 当weight非空时，对y_origin逐token乘以weight。
  3. 按quant_mode对y_origin量化，输出y和y_scale。
  4. 当output_origin=True时，额外返回量化前的y_origin；否则返回shape为[0]的占位Tensor。

函数原型

cann_ops_nn.swiglu_group_quant(
    x,
    *,
    weight=None,
    group_index=None,
    scale=None,
    dst_type=291,
    quant_mode=0,
    block_size=0,
    round_scale=False,
    clamp_limit=-1.0,
    dst_type_max=15.0,
    output_origin=False,
) -> (Tensor, Tensor, Tensor)

参数说明

参数名	参数类型	可选/必选	描述	数据类型	维度(shape)
x	Tensor	必选	SwiGLU输入，最后一维会被均分为两部分。	torch.float16、torch.bfloat16、torch.float32	1-7维
weight	Tensor	可选	逐token权重，非空时乘到量化前结果上。	torch.float32	1维，元素个数等于 x 除最后一维外的元素个数
group_index	Tensor	可选	count模式分组token数。	torch.int64	1维
scale	Tensor	可选	HiFloat8静态量化使用的scale。	torch.float32	1维
dst_type	int	可选	目标量化类型的torch dtype编码，默认291。	-	-
quant_mode	int	可选	量化模式，支持0、1、2、3。	-	-
block_size	int	可选	量化块大小，0表示使用模式默认值。	-	-
round_scale	bool	可选	MX量化是否将scale舍入为2的幂。	-	-
clamp_limit	float	可选	SwiGLU计算前的截断阈值，默认-1.0表示不启用截断。	-	-
dst_type_max	float	可选	HiFloat8动态量化计算 scale时使用的最大有限值。	-	-
output_origin	bool	可选	是否返回量化前的SwiGLU结果。	-	-

quant_mode 与 dst_type

quant_mode	含义	dst_type 支持值	y 的 torch dtype	y_scale 的 torch dtype
0	Block FP8	23/291表示float8_e5m2；24/292表示float8_e4m3fn	torch.float8_e5m2或torch.float8_e4m3fn	torch.float32
1	MX FP8 / MX FP4	23、24、291、292、296、297	FP8使用torch FP8 dtype；FP4在eager/图模式（aclgraph）路径使用torch.uint8打包存储。TorchAir GE图模式中296输出类型为torch.float4_e2m1fn_x2，297输出类型为torch.uint8	torch.float8_e8m0fnu
2	HiFloat8 静态量化	290表示HiFloat8	torch.uint8	torch.float32，shape为[0]
3	HiFloat8 动态量化	290表示HiFloat8	torch.uint8	torch.float32

dst_type 编码说明

dst_type	对应类型	来源	下发到 GE/ACL 的类型	适用 quant_mode
23	torch.float8_e5m2	PyTorch原生dtype int值	DT_FLOAT8_E5M2 / ACL_FLOAT8_E5M2	0、1
24	torch.float8_e4m3fn	PyTorch原生dtype int值	DT_FLOAT8_E4M3FN / ACL_FLOAT8_E4M3FN	0、1
291	torch_npu.float8_e5m2，语义同torch.float8_e5m2	torch_npu扩展dtype编码	DT_FLOAT8_E5M2 / ACL_FLOAT8_E5M2	0、1
292	torch_npu.float8_e4m3fn，语义同torch.float8_e4m3fn	torch_npu扩展dtype编码	DT_FLOAT8_E4M3FN / ACL_FLOAT8_E4M3FN	0、1
290	torch_npu.hifloat8	torch_npu扩展dtype编码	DT_HIFLOAT8 / ACL_HIFLOAT8	2、3
296	torch_npu.float4_e2m1fn_x2	torch_npu扩展dtype编码	DT_FLOAT4_E2M1 / ACL_FLOAT4_E2M1	1
297	torch_npu.float4_e1m2fn_x2	torch_npu扩展dtype编码	DT_FLOAT4_E1M2 / ACL_FLOAT4_E1M2	1

说明：quant_mode=2/3为HiFloat8模式，实际下发为DT_HIFLOAT8/ACL_HIFLOAT8。graph_convert会把上表中的torch dtype编码转换为GE dtype attr，ACLNN路径会把对应编码转换为aclDataType后调用底层接口。

返回值说明

参数名	参数类型	描述	数据类型	维度(shape)
y	Tensor	量化输出。	参见quant_mode 与 dst_type	FP8/HiFloat8为x.shape[:-1] + [D/2]；FP4为x.shape[:-1] + [ceil((D/2)/2)]，在D可被256整除时等价于x.shape[:-1] + [D/4]
y_scale	Tensor	量化scale输出。	参见quant_mode 与 dst_type	quant_mode=0为x.shape[:-1] + [ceil((D/2)/128)]；quant_mode=1为x.shape[:-1] + [ceil(ceil((D/2)/32)/2), 2]；quant_mode=2为[0]；quant_mode=3为group_index.shape或[1]
y_origin	Tensor	量化前SwiGLU结果或占位Tensor。	与x相同	output_origin=True时为x.shape[:-1] + [D/2]，否则为[0]

其中 D = x.shape[-1]。

约束说明

• 该接口支持单算子模式和TorchAir图模式调用。
• x、weight、group_index、scale均需为NPU Tensor；可选Tensor可以传 None。
• 输入x的rank 必须大于0，最后一维D必须大于等于256且能被256整除。
• dst_type支持FP8、FP4和HiFloat8对应的torch dtype编码，详见dst_type 编码说明。
• quant_mode=0时仅支持FP8输出，dst_type支持23、24、291、292，block_size支持0或128。
• quant_mode=1时支持FP8/FP4 输出，dst_type支持23、24、291、292、296、297，block_size支持0或32，round_scale必须为True。
• quant_mode=2时支持HiFloat8静态量化输出，需传入scale，dst_type、block_size和round_scale不生效，实际下发HiFloat8。
• quant_mode=3时支持HiFloat8动态量化输出，不使用scale，dst_type、block_size和round_scale不生效，实际下发 HiFloat8。
• dst_type为296或297，即FLOAT4_E2M1或FLOAT4_E1M2时，必须使用quant_mode=1。
• y_scale的数据类型必须与quant_mode匹配：Block FP8为torch.float32，MX为torch.float8_e8m0fnu，HiFloat8为torch.float32。
• quant_mode=3时，group_index可用于MoE场景的分组动态量化；y_scale的shape为group_index.shape，未传group_index时为[1]。
• clamp_limit不启用时使用默认占位值-1.0；启用时必须大于0。
• quant_mode=0和quant_mode=1时，output_origin仅支持False。
• group_index中的元素值须大于等于0。
• 不支持空Tensor和非连续Tensor。

确定性计算

默认支持确定性计算。

调用说明

• 单算子模式调用：

  import torch
  import torch_npu
  import cann_ops_nn.ops
  
  x = torch.randn(8, 512, dtype=torch.float16).npu()
  y, y_scale, y_origin = torch.ops.cann_ops_nn.swiglu_group_quant(
      x,
      dst_type=291,
      quant_mode=0,
      block_size=128,
  )
  

  MX FP4 模式：

  y, y_scale, y_origin = torch.ops.cann_ops_nn.swiglu_group_quant(
      x,
      dst_type=296,
      quant_mode=1,
      block_size=32,
      round_scale=True,
  )
  

• 图模式（torchair）调用：

  import torch
  import torch_npu
  import torchair
  import cann_ops_nn.ops
  
  class Model(torch.nn.Module):
      def forward(self, x):
          y, y_scale, _ = torch.ops.cann_ops_nn.swiglu_group_quant(
              x,
              dst_type=291,
              quant_mode=0,
              block_size=128,
          )
          return y, y_scale
  
  model = torch.compile(Model().npu(), backend=torchair.get_npu_backend(), dynamic=False)
  x = torch.randn(8, 512, dtype=torch.float16).npu()
  y, y_scale = model(x)