cann_ops_nn.swiglu_group_quant【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn产品支持情况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。计算流程将x的最后一维均分为A、B两部分计算y_origin silu(A) * B。当weight非空时对y_origin逐 token 乘以weight。按quant_mode对y_origin量化输出y和y_scale。当output_originTrue时额外返回量化前的y_origin否则返回 shape 为[0]的占位 Tensor。函数原型torch.ops.cann_ops_nn.swiglu_group_quant( x, *, weightNone, group_indexNone, scaleNone, dst_type291, quant_mode0, block_size0, round_scaleFalse, clamp_limit-1.0, dst_type_max15.0, output_originFalse, ) - (Tensor, Tensor, Tensor)参数说明参数名参数类型可选/必选描述数据类型维度(shape)xTensor必选SwiGLU 输入最后一维会被均分为两部分。torch.float16、torch.bfloat16、torch.float321-7维weightTensor可选逐 token 权重非空时乘到量化前结果上。torch.float321维元素个数等于x除最后一维外的元素个数group_indexTensor可选count 模式分组 token 数。torch.int641维scaleTensor可选HiFloat8 静态量化使用的 scale。torch.float321维dst_typeint可选目标量化类型的 torch dtype 编码默认291。--quant_modeint可选量化模式支持0、1、2、3。--block_sizeint可选量化块大小0表示使用模式默认值。--round_scalebool可选MX 量化是否将 scale 舍入为 2 的幂。--clamp_limitfloat可选SwiGLU 计算前的截断阈值默认-1.0表示不启用截断。--dst_type_maxfloat可选HiFloat8 动态量化计算 scale 时使用的最大有限值。--output_originbool可选是否返回量化前的 SwiGLU 结果。--quant_mode 与 dst_typequant_mode含义dst_type支持值y的 torch dtypey_scale的 torch dtype0Block FP823/291表示float8_e5m224/292表示float8_e4m3fntorch.float8_e5m2或torch.float8_e4m3fntorch.float321MX FP8 / MX FP423、24、291、292、296、297FP8 使用 torch FP8 dtypeFP4 在 eager/ACL Graph 路径使用torch.uint8打包存储。TorchAir GE 图模式中296输出类型为torch.float4_e2m1fn_x2297输出类型为torch.uint8torch.float8_e8m0fnu2HiFloat8 静态量化建议使用290表示 HiFloat8当前 torch 接口实际下发 HIFLOAT8torch.uint8torch.float32shape 为[0]3HiFloat8 动态量化建议使用290表示 HiFloat8当前 torch 接口实际下发 HIFLOAT8torch.uint8torch.float32dst_type 编码说明dst_type对应类型来源下发到 GE/ACL 的类型适用quant_mode23torch.float8_e5m2PyTorch 原生 dtype int 值DT_FLOAT8_E5M2/ACL_FLOAT8_E5M20、124torch.float8_e4m3fnPyTorch 原生 dtype int 值DT_FLOAT8_E4M3FN/ACL_FLOAT8_E4M3FN0、1291torch_npu.float8_e5m2语义同torch.float8_e5m2torch_npu 扩展 dtype 编码DT_FLOAT8_E5M2/ACL_FLOAT8_E5M20、1292torch_npu.float8_e4m3fn语义同torch.float8_e4m3fntorch_npu 扩展 dtype 编码DT_FLOAT8_E4M3FN/ACL_FLOAT8_E4M3FN0、1290torch_npu.hifloat8torch_npu 扩展 dtype 编码DT_HIFLOAT8/ACL_HIFLOAT82、3296torch_npu.float4_e2m1fn_x2torch_npu 扩展 dtype 编码DT_FLOAT4_E2M1/ACL_FLOAT4_E2M11297torch_npu.float4_e1m2fn_x2torch_npu 扩展 dtype 编码DT_FLOAT4_E1M2/ACL_FLOAT4_E1M21说明quant_mode2/3为 HiFloat8 模式实际下发为DT_HIFLOAT8/ACL_HIFLOAT8。graph_convert 会把上表中的 torch dtype 编码转换为 GE dtype attrACLNN 路径会把对应编码转换为aclDataType后调用底层接口。返回值说明参数名参数类型描述数据类型维度(shape)yTensor量化输出。见quant_mode 与 dst_typeFP8/HiFloat8 为x.shape[:-1] [D/2]FP4 为x.shape[:-1] [ceil((D/2)/2)]在D可被 256 整除时等价于x.shape[:-1] [D/4]y_scaleTensor量化 scale 输出。见quant_mode 与 dst_typequant_mode0为x.shape[:-1] [ceil((D/2)/128)]quant_mode1为x.shape[:-1] [ceil(ceil((D/2)/32)/2), 2]quant_mode2为[0]quant_mode3为group_index.shape或[1]y_originTensor量化前 SwiGLU 结果或占位 Tensor。与x相同output_originTrue时为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_mode0时仅支持 FP8 输出dst_type支持23、24、291、292block_size支持0或128。quant_mode1时支持 FP8/FP4 输出dst_type支持23、24、291、292、296、297block_size支持0或32round_scale必须为True。quant_mode2时支持 HiFloat8 静态量化输出需传入scaledst_type、block_size和round_scale不生效实际下发 HiFloat8。quant_mode3时支持 HiFloat8 动态量化输出不使用scaledst_type、block_size和round_scale不生效实际下发 HiFloat8。dst_type为296或297即FLOAT4_E2M1或FLOAT4_E1M2时必须使用quant_mode1。y_scale的数据类型必须与quant_mode匹配Block FP8 为torch.float32MX 为torch.float8_e8m0fnuHiFloat8 为torch.float32。quant_mode3时group_index可用于 MoE 场景的分组动态量化y_scale的 shape 为group_index.shape未传group_index时为[1]。clamp_limit不启用时使用默认占位值-1.0启用时必须大于 0。quant_mode0和quant_mode1时output_origin仅支持False。不支持空 Tensor 和非连续 Tensor。确定性计算默认支持确定性计算。调用说明单算子模式调用import torch import torch_npu import cann_ops_nn.ops x torch.randn(8, 512, dtypetorch.float16).npu() y, y_scale, y_origin torch.ops.cann_ops_nn.swiglu_group_quant( x, dst_type291, quant_mode0, block_size128, )MX FP4 模式y, y_scale, y_origin torch.ops.cann_ops_nn.swiglu_group_quant( x, dst_type296, quant_mode1, block_size32, round_scaleTrue, )图模式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_type291, quant_mode0, block_size128, ) return y, y_scale model torch.compile(Model().npu(), backendtorchair.get_npu_backend(), dynamicFalse) x torch.randn(8, 512, dtypetorch.float16).npu() y, y_scale model(x)【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考