CANN/ops-nn分组量化SwiGLU激活算子

SwigluGroupQuant

【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn

产品支持情况

功能说明

• 算子功能：融合实现SwiGLU激活和分组低比特量化，支持FP8和FP4输出。输入x的最后一维被均分为A和B，先计算silu(A) * B，再执行量化。

• 计算公式：

  $$ y_{tmp}=silu(A) \times B $$

  当传入clamp_limit时：

  $$ A=min(A, clamp_limit) $$

  $$ B=min(max(B, -clamp_limit), clamp_limit) $$

  当传入weight时，量化前执行：

  $$ y_{tmp}=y_{tmp} \times weight $$

  进行量化：

  $$ scale=row_max(abs(y_{tmp}))/dstTypeScale $$

  $$ y = Cast(Mul(y_{tmp}, 1/scale)) $$

  quant_mode为0时输出FP8类型的y和FLOAT32类型的y_scale；quant_mode为1时输出FP8/FP4类型的y和FLOAT8_E8M0类型的y_scale；quant_mode为2、3时输出HIFP8类型的y和FLOAT32类型的y_scale。

参数说明

约束说明

• 输入x的rank必须大于0，最后一维D必须大于等于256且能被256整除。
• dst_type支持HIFLOAT8、FLOAT8_E4M3FN、FLOAT8_E5M2、FLOAT4_E2M1、FLOAT4_E1M2。
• quant_mode=0时仅支持FP8输出，block_size支持0或128。
• quant_mode=1时支持FP8/FP4输出，block_size支持0或32，round_scale必须为true。
• quant_mode=2时支持HIFP8静态量化输出，dst_type、block_size和round_scale不生效。
• quant_mode=3时支持HIFP8动态量化输出，dst_type、block_size和round_scale不生效。
• dst_type为FLOAT4_E2M1或FLOAT4_E1M2时，必须使用quant_mode=1。
• dst_type为HIFLOAT8时，必须使用quant_mode=2或3。
• y_scale的数据类型必须与quant_mode匹配：Block FP8为FLOAT32，MX为FLOAT8_E8M0，HIFP8为FLOAT32。
• quant_mode=2或3时，group_index可用于MoE场景的分组量化，y_scale的shape为[G]（G为group数量）。
• clamp_limit不启用时使用默认占位值-1.0；启用时必须大于0。

调用说明

【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn