融合 SwiGLU、Clamp 与动态量化的 AscendC 自定义算子解析：custom.npu_swiglu_clip_quant

融合 SwiGLU、Clamp 与动态量化的 AscendC 自定义算子解析custom.npu_swiglu_clip_quant【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer导读本文深入解析 CANN recipes 推理仓库中ops/ascendc目录下的融合算子custom.npu_swiglu_clip_quant它在 Swish 门控线性单元SwiGLU激活之后紧接 clamp 与动态量化将三段计算融合为单次 NPU 算子调用并支持按 token 分组使用不同的 clamp 缩放因子。读完本文你将掌握该算子的数学语义、全部输入输出参数与约束、Eager/Graph 两种调用方式以及从 OpDef 定义、Tiling 切分到 AscendC Kernel 与 PTA 前端的完整实现链路可直接在 Atlas A3 推理系列产品上复现运行。一、算子背景为什么需要融合 SwiGLU Clamp Quant在 LLM 推理的 MoE 场景中门控线性单元GLU类激活函数如 SwiGLU广泛应用于 FFN 计算。以 GPT-OSS、DeepSeek 系列等大模型为代表的网络中SwiGLU输出的数值范围较大直接送入低比特量化如 INT8会造成较大的精度损失。业界常见的做法是对激活输出做clamp截断限制其数值范围再乘以一个按组group配置的缩放因子group_alpha进一步收紧分布最后执行动态量化DynamicQuant按行/组计算 scale 并量化为 INT8。上述三段操作如果各自独立成算子会引入多次 GMGlobal Memory读写与 Kernel 启动开销。npu_swiglu_clip_quant将三者融合进一个自定义算子中一次完成激活 截断 按组缩放 动态量化是典型的计算密集型融合优化样例。该算子在仓库中被设计为AscendC自定义算子同时提供 PyTorchtorch_npu前端封装支持 Eager 模式与torch.compile图模式Graph 模式。二、功能说明与数学原理SwigluClipQuant的输入x在最后一维上被均分为前后两半A、B依次执行三步计算具体公式如下SwiGLU 激活$$ \textbf{swiglu}: swigluOut Swiglu(A)*B $$Clamp 截断按分组$$ \textbf{clip}: swigluClipOut Max(Min(swigluOut, Abs(ReduceMax(swigluOut)) * groupAlpha), -Abs(ReduceMax(swigluOut)) * groupAlpha) $$动态量化$$ \textbf{quant}: y, scale DynamicQuant(swigluClipOut, groupIndex) $$其中A表示输入x的前半部分B表示输入x的后半部分groupAlpha是与每个分组对应的 clamp 缩放因子通过group_alpha传入ReduceMax(swigluOut)沿最后一维H 维取最大值再乘以groupAlpha得到该组该行的 clamp 上界取负得到下界DynamicQuant依据groupIndex指定的分组方式按组计算量化 scale 并输出 INT8 结果。从仓库的 CPU 参考实现test_npu_swiglu_clip_quant.py可以更直观地还原语义先按group_index切分 token对每段做silu(x_left) * x_right再取abs的amax乘上group_alpha[index]作为 clamp 边界随后用max_values / 127作为动态 scale 完成除法量化并二次 clamp 到[-128, 127]最后四舍五入转 INT8。三、函数原型custom.npu_swiglu_clip_quant(Tensor x, Tensor group_index, Tensor group_alpha, *, bool activate_leftFalse, int quant_mode1, int clamp_mode1) - (Tensor, Tensor)原型要点x、group_index、group_alpha三个参数位于*之前是位置相关的必选参数必须按顺序传入*之后的activate_left、quant_mode、clamp_mode为键值对赋值的可选参数与位置无关不传时使用默认值。四、参数说明基础维度约定TokensNum表示传输的 Token 数自然数H表示嵌入向量的长度0groupNum表示group_index输入的长度0。参数类型必选/可选说明xTensor必选输入张量不支持非连续数据格式 ND数据类型bfloat16shape 为[TokensNum, H]group_indexTensor必选1D 张量不支持非连续数据格式 ND数据类型int64shape 为[groupNum]。当前仅支持 count 模式每个元素表示该分组包含的 Token 数非负整数即按 token 数量切分组group_alphaTensor必选1D 张量不支持非连续数据格式 ND数据类型float32shape 为[groupNum]。表示与各分组 Token 数对应的 clamp 缩放因子activate_leftbool可选是否对输入的左半部分做 SwiGLU 激活为false时对右半部分做激活。默认值Falsequant_modeint可选量化方式0静态量化1动态量化。默认值1当前仅支持动态量化clamp_modeint可选是否在指定分组 Token 数时使用group_alpha参数。默认值1使用group_alpha功能当前仅支持clamp_mode1源码印证swiglu_clip_quant_def.cppOpDef 中x约束为DT_BF16、group_index约束为DT_INT64、group_alpha约束为DT_FLOAT输出y为DT_INT8、scale为DT_FLOAT格式均为 ND属性activate_left为可选 bool默认 false、quant_mode为字符串属性默认dynamic、clamp_mode为 int 属性默认 1。五、返回值说明返回值类型说明yTensor公式中的输出y即量化后的输出张量数据类型int8数据格式 NDscaleTensor公式中的输出scale即量化 scale 参数数据类型float32数据格式 ND从 PTA 前端npu_swiglu_clip_quant.cpp可以看到输出 shape 的推导规则y与x的 shape 相同但最后一维减半H/2即 SwiGLU 输出宽度scale的 shape 为y去掉最后一维即[TokensNum]或[TokensNum, ...]去除 H 后的形状对应每行/每组一个 scale 的动态量化语义。六、约束说明使用该接口前请确认以下约束均来自算子文档与 Tiling 源码校验使用场景该接口仅支持推理场景使用。图模式支持该接口支持图模式Graph/torch.compile。版本匹配与 PyTorch 配合使用时需保证 CANN 相关包与 PyTorch 相关包的版本匹配。H 轴维度限制参数x的 H 轴满足H ≤ 10496且需要64 对齐规格不满足时算子会进行校验并报错。该限制源于 Tiling 中的inDimy_ % (BLOCK_SIZE * SWI_FACTOR) ! 0校验BLOCK_SIZE32、SWI_FACTOR2见 swiglu_clip_quant_tiling.cpp即输入最后一维必须能被 64 整除。group_index 越界保护group_index仅支持 count 模式需要保证group_index总和不超过x的TokensNum维度否则会出现越界访问。Kernel 侧对负值分组做了防护realDimx_ 0时忽略该组见 swiglu_clip_quant.h。输出垃圾数据输出y和scale中超过group_index总和的部分未做清理处理该部分内存为垃圾数据使用时需注意例如按group_index.sum()截取有效结果示例代码中正是用total_nums group_index.sum()做对比截断。量化模式quant_mode仅支持动态量化场景Tiling 中SUPPORT_QUANT_MODE {{dynamic, 1}}传入静态量化会报错见 swiglu_clip_quant_tiling.cpp。产品支持Atlas A3 推理系列产品此外 OpDef 中还注册了ascend910b对应 Atlas A2 系列等与ascend910_93的 AICore 配置A2 环境需通过-c ascend910b指定型号编译详见下文构建章节。七、调用示例Eager 与 Graph 双模式完整的可运行用例位于 test_npu_swiglu_clip_quant.py演示了 Eager 与torch.compileGraph 两种调用方式并与 CPU 参考实现做了精度比对。7.1 Eager 模式调用import torch import torch_npu import custom_ops DEVICE_ID 0 torch_npu.npu.set_device(int(DEVICE_ID)) m, n 5000, 4096 np.random.seed(0) x torch.tensor(np.random.uniform(-100000, 100000, (m, n))).to(torch.bfloat16) group_index torch.tensor([10, 1000, 2000, 500], dtypetorch.int64) group_alpha torch.tensor([0.12, 0.532, 0.927, 0.5], dtypetorch.float32) x x.to(npu:0) group_index group_index.to(npu:0) group_alpha group_alpha.to(npu:0) npu_y, npu_scale torch_npu.npu_swiglu_clip_quant( x, group_index, group_alpha, quant_mode1, clamp_mode1, activate_leftTrue)要点示例中输入x为[5000, 4096]的bfloat16张量H4096满足 64 对齐与H ≤ 10496group_index[10, 1000, 2000, 500]表示 4 个分组分别覆盖 10、1000、2000、500 个 token总和 3510 ≤ 5000group_alpha[0.12, 0.532, 0.927, 0.5]为各组对应的 clamp 缩放因子返回值npu_y的 shape 为[5000, 2048]H 减半npu_scale的 shape 为[5000]。7.2 Graph 模式调用import torch_npu import torchair import torch.nn as nn from torchair.configs.compiler_config import CompilerConfig class Network(nn.Module): def forward(self, x, group_index, group_alpha, quant_mode1, clamp_mode1, activate_leftTrue): y, scale torch_npu.npu_swiglu_clip_quant(x, group_index, group_alpha, quant_modequant_mode, clamp_modeclamp_mode, activate_leftactivate_left) return y, scale npu_mode Network().to(npu:0) config CompilerConfig() npu_backend torchair.get_npu_backend(compiler_configconfig) npu_mode torch.compile(npu_mode, fullgraphTrue, backendnpu_backend, dynamicFalse) npu_y, npu_scale npu_mode(x, group_index, group_alpha, quant_mode1, clamp_mode1, activate_leftTrue)Graph 模式依赖torchair提供的 GE 图转换能力通过register_fx_node_ge_converter注册的 converter见 npu_swiglu_clip_quant.py将torch.ops.custom.npu_swiglu_clip_quant.default节点转换为 GE 自定义算子节点SwigluClipQuant并携带activate_leftBool 属性、quant_mode字符串属性1 映射为dynamic0 映射为static、clamp_modeInt 属性。7.3 精度验证方式用例内置了 CPU 参考实现_swiglu_clip_quant与相对误差比对函数data_compare将 NPU 输出与 CPU 输出按group_index.sum()截取有效部分后比对统计相对误差超过阈值的比例与最大误差最终通过assert(compare_y[0] Pass and compare_scale[0] Pass)判定通过相对误差阈值0.01、通过比例要求95%、最大误差上限0.1。八、源码实现深度解析8.1 OpDef 算子定义swiglu_clip_quant_def.cpp 完成算子注册三个必选输入xBF16、group_indexINT64、group_alphaFLOAT均为 ND 格式两个必选输出yINT8、scaleFLOAT三个可选属性activate_left默认 false、quant_mode默认dynamic字符串、clamp_mode默认 1AICore 配置支持ascend910b与ascend910_93两个平台。8.2 Tiling 设计Host 侧Tiling 逻辑位于 swiglu_clip_quant_tiling.h 与 swiglu_clip_quant_tiling.cpp核心工作包括平台信息获取读取 AIV Core 数量与 UBUnified Buffer大小据此计算切分参数数据类型校验x仅支持bfloat16group_index仅支持int64group_alpha仅支持float32且要求group_index与group_alpha第一维长度一致属性校验quant_mode仅支持dynamic对应取值 1clamp_mode仅支持 1维度校验与切分要求x最后一维能被 64 整除BLOCK_SIZE * SWI_FACTOR并根据 UB 容量按切分行数UbFactorDimx× 输出列outDimy计算单块承载量预留 1024B UB 余量当inDimy为 2048 或 4096 时还有专门的性能切分因子PERFORMANCE_UB_FACTOR 4096 * 4tilingKey 组合通过有无分组、quantScale 类型、bias 类型、以及分组数大且每组 token 少groupNum 64且平均每组 token 数 ≤ 16等条件组合出不同的 tilingKey用于 Kernel 模板分派与性能优化分支。8.3 Kernel 实现Device 侧AscendC Kernel 位于 swiglu_clip_quant.h采用模板化设计template typename TBias, typename TQuantScale, typename TGroup, typename TXGm处理流程为Process按group_index遍历分组逐组计算起始偏移groupOffset对realDimx_ 0的分组做防护跳过ProcessSingleGroup / ProcessSingleGroupPerCore组内再做一次核间切分将realCount个 token 均摊到maxCoreNum个 AIV Core 上blockIdx_ realCoreDim的核参与计算最后一个核处理尾块blockDimxTailFactorComputeDequant将 BF16 输入Cast为 FP32CAST_NONE模式为后续计算做准备ComputeSwiGLU利用Copy按 act/gate 偏移切出A、B两半通过Muls(-1.0) → Exp → Adds(1.0) → Div → Mul的向量指令序列实现silu(A)*BSwiGLUComputeQuant / DynamicQuantquantMode 1时走动态量化内部使用dynamicQuantFactor 1.0 / 127.0与-1.0 / 127.0常量完成 INT8 对称量化CopyOut将 INT8 结果y与每行 FP32 的scale通过DataCopyPad写回 GM。数据流上算子使用TQueQuePosition::VECIN, 1输入队列与TQueQuePosition::VECOUT, 1输出队列配合TBufTPosition::VECCALC临时缓冲区在xActQueue_中复用同一块 UB 承载输入、激活与 scale 数据最大化缓存利用率。8.4 PTA 前端封装torch_npu 接入NPU 与 Meta 实现注册npu_swiglu_clip_quant.cpp 通过TORCH_LIBRARY_IMPL(custom, PrivateUse1, m)注册 NPU 侧实现内部调用aclnnSwigluClipQuant执行并通过TORCH_LIBRARY_IMPL(custom, Meta, m)注册 Meta 侧实现用于 shape 推导torch.compile静态分析需要输出张量由construct_swiglu_clip_quant_output_tensor完成 shape 推导与内存分配y最后一维减半scale去掉最后一维。GE Converternpu_swiglu_clip_quant.py 使用register_fx_node_ge_converter(torch.ops.custom.npu_swiglu_clip_quant.default)装饰器注册 FX 节点到 GE 图的转换将quant_mode整数映射为字符串属性1 →dynamic后通过torchair.ge.custom_op生成SwigluClipQuantGE 节点——这正是 Graph 模式torch.compile能正常成图的底层机制。九、构建与运行算子工程位于ops/ascendc使用仓库根目录下的 build.sh 进行编译常用命令如下# 编译全部算子A3 环境默认 bash build.sh # 指定 NPU 型号编译例如 Atlas A2 训练系列/A800I A2 推理等使用 ascend910b bash build.sh -c ascend910b # 仅编译 swiglu_clip_quant 单个算子 bash build.sh -n swiglu_clip_quant # 同时编译多个指定算子分号分隔 bash build.sh -n swiglu_clip_quant;gather_selection_kv_cache;hc_pre # 跳过版本校验拦截版本兼容性检查不通过时使用 bash build.sh --disable-check-compatible编译产物为自定义算子安装包CANN-custom_ops-linux.arch.run。安装后即可运行示例用例cd ops/ascendc/examples python3 test_npu_swiglu_clip_quant.py用例会依次执行 Eager 与 Graph 两个测试test_swiglu_clip_quant_eager/test_swiglu_clip_quant_graph并在两侧输出PTA eager BEGIN/FINISH、PTA graph BEGIN/FINISH日志最终通过 NPU 结果与 CPU 参考实现的比对断言。补充说明AscendC 算子默认基于 A3 环境编译包含 Atlas A3 训练系列产品/Atlas A3 推理系列产品其他环境需要在编译时通过-c ${soc_version}指定 NPU 型号。查看build.sh全部参数可执行bash build.sh --help。十、总结custom.npu_swiglu_clip_quant是 CANN recipes 中面向大模型推理的典型融合算子样例它把 SwiGLU 激活、按组 clamp 缩放与动态量化三段计算合并为一次 NPU Kernel 执行数学上实现y, scale DynamicQuant(Clamp(SwiGLU(x), ±|ReduceMax(SwiGLU(x))|·groupAlpha), groupIndex)的完整语义接口上通过group_indexcount 模式分组与group_alpha分组 clamp 因子支持按 token 分组粒度的灵活量化策略适配不同 token 分布下更精细的量化控制工程上从 OpDef 注册、Host 侧 Tiling 切分、Device 侧 AscendC Kernel 到 PTA 前端与 GE Converter 一应俱全同时支持 Eager 与torch.compileGraph 两种执行路径并附带 CPU 参考实现与精度比对用例可直接作为在 Atlas A3 推理系列产品上编写、编译与验证融合算子的完整参考模板。如需在自有模型中集成可参照 test_npu_swiglu_clip_quant.py 的调用方式将 FFN 中独立的SwiGLU → Clamp → Quant序列替换为该融合算子并确保输入满足文档中的维度、对齐与分组约束。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考