custom-npu_swiglu_group_quant【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√功能说明在SwiGlu激活函数后添加量化操作实现输入x的SwiGluQuant计算。根据quant_mode不同量化分组大小有所差异quant_mode1或3时groupSize固定为128quant_mode2时groupSize固定为32。计算过程见计算公式。计算公式$$ Y_{tmp} SwiGLU(x) Swish(A)*B $$$$ scalerow_max(abs(Y_{tmp}))/dstTypeScale $$$$ Y Cast(Mul(Y_{tmp}, Scale)) $$ 其中A表示输入x的前半部分B表示输入x的后半部分。函数原型custom.npu_swiglu_group_quant(Tensor x, *, Tensor? topk_weightNone, Tensor? group_indexNone, ScalarType dst_type, int quant_mode1, int group_size128, bool round_scaleFalse, bool ue8m0_scaleFalse, bool output_originFalse, int group_list_type0, float clamp_value0.0) - (Tensor, Tensor, Tensor)参数说明说明bbatch size表示输入样本批量大小、ssequence length表示输入样本序列长度、dhead dimension表示注意力头的维度数、T表示bs合轴后的大小、G表示MoE场景下的分组数量。xTensor必选参数输入tensor公式中的输入 x 。不支持非连续数据格式支持ND数据类型支持float16和bfloat16shape为[ b, s, d ]或[ T, d ]T为bs合轴后的大小。*代表其之前的参数是位置相关的必须按照顺序输入属于必选参数其之后的参数是键值对赋值与位置无关属于可选参数不传入会使用默认值。topk_weightTensor可选MoE场景下的topk权重tensor用于对SwiGLU输出进行加权。数据类型支持float32shape为[ T, 1 ]。默认为None表示不进行加权操作。group_indexTensor可选MoE场景下分组索引tensor用于指定每个group的token个数。数据类型支持int64shape为[G]其中G表示分组数量。该Tensor的数值总和代表输入x中的有效Token数。默认值为None表示非MoE场景此时所有Token都有效。dst_typeScalarType可选量化输出数据类型支持torch.float8_e5m2和torch.float8_e4m3fn。默认为torch.float8_e4m3fn。quant_modeint, 可选量化模式取值范围为1、2、3。1PerGroup fp8量化量化输出scale为float32类型groupSize固定为128。2MX fp8量化量化输出scale为float8_e8m0类型groupSize固定为32。3PerGroup fp8量化增强模式groupSize固定为128。scale类型由ue8m0_scale参数决定ue8m0_scaleFalse时scale为float32类型与quant_mode1功能相同ue8m0_scaleTrue时scale为float8_e8m0类型。相比quant_mode1该模式还支持round_scale、output_origin等额外参数。group_sizeint, 可选量化分组大小。quant_mode1或3时固定为128quant_mode2时固定为32。默认为128。round_scalebool, 可选是否对scale进行舍入处理。默认为False。仅在quant_mode3时有效。ue8m0_scalebool, 可选是否使用float8_e8m0格式输出scale。默认为Falsescale为float32类型设置为True时scale为float8_e8m0类型。仅在quant_mode3时有效。output_originbool, 可选是否输出SwiGLU计算后的原始fp16/bf16结果。默认为False。仅在quant_mode3时有效。group_list_typeint, 可选分组列表类型固定为0默认为0表示count模式。在count模式下group_index的每个元素代表对应分组的元素个数。在quant_mode1、2、3时均生效。clamp_valuefloat, 可选SwiGLU激活函数的截断值。当clamp_value0时对激活值进行截断A部分只有最大值约束截断到不超过clamp_valueB部分截断到[-clamp_value, clamp_value]。默认为0.0表示不进行截断。返回值说明yTensor输出tensor量化后的value。不支持非连续数据格式支持ND数据类型支持float8_e4m3fn和float8_e5m2shape为[ b, s, d/2 ]或[ T, d/2 ]。当使用group_index时前有效Token数行的数据有效其余为填充数据。scaleTensor输出tensor量化后的scale。不支持非连续数据格式支持ND。当quant_mode1或quant_mode3且ue8m0_scaleFalse时数据类型为float32当quant_mode2或quant_mode3且ue8m0_scaleTrue时数据类型为float8_e8m0。yOriginTensor输出tensorSwiGLU计算后的原始结果。仅在quant_mode3且output_originTrue时有效输出其他情况下返回空tensor。数据类型与输入x相同float16或bfloat16shape为[ b, s, d/2 ]或[ T, d/2 ]。约束说明shape 字段取值范围约束 | quant_mode | d最后一维约束 | |------------|------------------| | 1 或 3 | 必须是256的倍数 | | 2 | 必须是128的倍数 |quant_mode 取值范围为1、2、3。具体说明详见参数说明中的quant_mode字段。round_scale、ue8m0_scale、output_origin参数仅在quant_mode3时有效。group_list_type固定为0表示count模式在quant_mode1、2、3时均生效。group_index的数值总和必须不超过输入x的第一维大小。该接口支持推理场景下使用。该接口支持aclgraph入图。该接口与PyTorch配合使用时需要保证CANN相关包与PyTorch相关包的版本匹配。调用示例详见 test_npu_swiglu_group_quant.py【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考