PyPTO 算子设计中的 SwiGLU 激活原子模式(AT-07):从公式分解到融合算子实现 📅 发布时间:2026/9/19 16:39:03 👁 浏览次数: PyPTO 算子设计中的 SwiGLU 激活原子模式AT-07从公式分解到融合算子实现【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gymSwiGLUSwish-Gated Linear Unit是现代大模型 FFN/专家网络的标准激活结构其核心计算silu(gate) * value在 NPU 上既有纯 Vector 计算形态也有与 MatMul、量化深度融合的实现路径。本文以 PyPTO-Gym 仓库中算子设计模式库的 AT-07 卡片为骨架结合仓库内 FusedSwiGLU、GMM SwiGLU、GLMMoEFusion 三组真实实现与测试用例系统讲解 SwiGLU 的数学结构、两种 Vector 实现变体、前后向 kernel 写法、以及它与量化/分组 MatMul 的融合范式帮助算子设计者在编写 DESIGN.md 时快速复用该模式并避免精度与资源上的常见误区。一、模式定位AT-07 在计算模式库中的位置PyPTO-Gym 的算子设计工作流cannbot-skills/ops/pypto-op-design/SKILL.md要求拿到 SPEC 与 golden 之后先读 SK 索引 和 AT 索引再按计算依赖与适用条件选取候选卡片。AT-07 属于局部计算模式Atom类别被索引表归类为字段值IDAT-07名称SwiGLU Activationtagsactivationflow_patternV纯 VectorexamplesFusedSwiGLU、GMM SwiGLU、GLMMoEFusionflow_pattern: V 表示该模式在硬件计算流中只占用 Vector 单元通常紧跟在两个 MatMulgate 投影与 value 投影之后执行。在完整 FFN 链路中它扮演两个线性投影结果之间的逐元素门控这一角色是 AT-09 Quant Linear 等 MatMul 类模式的天然下游搭配。二、数学结构为什么是silu(gate) * valueSwiGLU 将输入分别送入两个线性投影得到门控路径gate与值路径value再按下式逐元素相乘y silu(gate) * value (gate * sigmoid(gate)) * value其中silu(x) x * sigmoid(x) x / (1 exp(-x))。相比传统 GELU/ReLU 类激活SwiGLU 通过可学习的门控路径对值路径做软开关在同等参数量下带来更强的表达能力因而被 Qwen、GutenOCR、GLM、DeepSeek 等主流模型族广泛采用。AT-07 卡片给出的接口契约如下输入gate: Tensor[M, N]— 门控路径FP32value: Tensor[M, N]— 值路径FP32输出y: Tensor[M, N]— 激活结果注意输入标注为 FP32即使上游 MatMul 以 BF16 输入、以 FP32 累加输出SwiGLU 的乘加链也建议在 FP32 下完成最后再一次性回落到目标精度BF16/INT8避免激活路径上的中间量化损失。三、两种 Vector 实现变体AT-07 卡片给出两条等价计算路径二者的区别仅在于 sigmoid 的展开方式变体 A — 使用 exp无 sigmoid 原语时sigmoid_gate div(ones, add(ones, exp(neg(gate)))) silu mul(gate, sigmoid_gate) y mul(silu, value)变体 B — 直接 silu 算子silu sigmoid(gate) * gate y mul(silu, value)选择依据仓库中两种写法均有实例见下文源码变体 A 在目标工具链/版本缺少sigmoid原语、或需要显式控制中间精度时使用。仓库 silu kernel reference 即给出完全一致的 exp 展开先exp(x * -1)加 1 得分母再x * div(ones, e, INTRINSIC)完成 silu且中间全程 FP32。参考实现还提示了两个关键实践batch 轴用pypto.loop切分、last-dim 计算轴整块处理。变体 B 在目标版本已提供pypto.sigmoid/pypto.silu原语时更简洁直观且显式乘法让精度控制更直接。无论哪个变体最终都是两个乘法和一个 sigmoid 展开计算强度低、完全落在 Vector 单元不需要 Cube 参与。四、源码级印证一FusedSwiGLU 前向 kernel仓库 fused_swiglu_impl.py 是 AT-07 变体 B 的完整生产实现其计算式与注释直接对应模式卡片y SiLU(x w_g b_g) * (x w_fc b_fc)实现要点与 AT-07 模式强对应动态 batch 支持x声明为pypto.Tensor([pypto.DYNAMIC, ...], pypto.DT_BF16)用tile_m 512沿 M 轴切分loop_count (m tile_m - 1) // tile_m计算循环轮数pypto.viewvalid_shape处理尾块——这正是 SKILL.md 中动态轴用框架循环遍历的落地方式。Cube 与 Vector 分时交替两个 MatMul 用pypto.matmul(x_tile, w_g, pypto.DT_FP32, extend_params{bias_tensor: b_g_fp32})完成FP32 累加、bias 融合进 MatMul随后切换pypto.set_vec_tile_shapes(128, 128)进入 Vector 段sigmoid - silu - 逐元素乘 - cast BF16。偏置先升 FP32b_g_fp32 pypto.cast(b_g, pypto.DT_FP32)在循环外完成避免每轮重复转换。JIT 配置pass_options设置vec_nbuffer_setting与cube_l1_reuse_settingruntime_options设置stitch_function_max_num与device_sched_mode用于控制 Vector 缓冲、Cube L1 复用与算子拼接调度。对应测试 test_fused_swiglu.py 以m220000, k512, n1024的大 M 形状验证golden 为gate x w_g b_g fc x w_fc b_fc gate_silu gate * torch.sigmoid(gate) y (gate_silu * fc).to(x.dtype)并以rtol0.0078125, atol0.0001做数值对齐验证了BF16 输入 FP32 中间 回写 BF16的精度路径。五、源码级印证二反传 kernel 的 silu 导数展开AT-07 卡片只覆盖前向但反传同样依赖 silu 的导数。仓库 fused_swiglu_grad_impl.py 给出三个反向 kernel其中反传的核心是 silu 导数silu(g) sigmoid(g) * (1 g * (1 - sigmoid(g)))实现中 sigmoid 用pypto.div(exp_g, 1.0 exp_g, precision_typepypto.PrecisionType.INTRINSIC)展开变体 A 的 exp 路线但分子直接取exp_g而非1避免一次除法随后dg dy * fc * silu(g)— gate 路径梯度dfc dy * silu(g)— value 路径梯度db用pypto.sum(dim0, keepdimTrue)沿 M 归约累加dw/dx分别通过a_transTrue与b_transTrue的 MatMul 完成三个 kernel 按tile_m 1024/2048分块循环dg/dfc由第一个 kernel 写出、被后两个 kernel 复用体现kernel 间中间张量落盘复用的典型设计。六、源码级印证三SwiGLU 与分组量化 MatMul 融合GMM SwiGLUAT-07 卡片 examples 中的 GMM SwiGLU 指仓库 gmm_swiglu_quant_impl.py一个 MXFP8 分组 scaled_matmul 的 kernel 里直接内联了 SwiGLU 与 per-token 量化。关键代码路径current_mm_out pypto.scaled_mm(x, weight, pypto.DT_FP32, scaled_x, scaled_weight) # 按列切分出 SwiGLU 的 value 与 gate value current_mm_out[:, : n_size // 2] gate current_mm_out[:, n_size // 2:] silu_value value * pypto.sigmoid(value) swiglu_out pypto.mul(silu_value, gate)两个要点直接呼应 AT-07列切分复用 MatMul 输出分组 MatMul 的输出按最后一维对半切成 value/gate省去两次独立投影这解释了模式描述中纯 V通常在两个 MatMul 之后的排布——在某些融合形态下两个 MatMul 甚至可合并为一次更宽的 MatMul。激活后立即量化swiglu_out先cast BF16再走abs - amax - 127/max - round - INT8的 per-token 量化链路产出x_int8与x_scale_quant形成MatMul → SwiGLU → 量化三级流水。测试 test_gmm_swiglu_quant.py 以m16, k512, n7168, group_list[7,9]的 FP8 分组场景验证golden 侧gmm_swiglu_quant_golden.py同样遵循silu(value) * gate结构量化输出用rtol1e-4严格对齐。七、源码级印证四完整 SwiGLU MLPGutenOCR-3B与 GLM MoE 融合7.1 三段式 SwiGLU MLP仓库 swiglu_mlp_impl.py 给出 SwiGLU 在 FFN 中的完整形态GutenOCR-3BH2048, I11008gate silu(gate_proj(x)) # Linear(H→I) up up_proj(x) # Linear(H→I) hidden gate * up # Element-wise mul output down_proj(hidden) # Linear(I→H)该文件同时提供动态 batch 版本pypto.silu(gate)原语写法与静态 batch 版本pypto.mul(gate, pypto.sigmoid(gate))展开写法对应 AT-07 的变体 B 与 A内置的精度对比以max diff 0.1为通过线性能对比则对 batch ∈ {1,4,8,16} 分别测量 Torch baseline 与 PyPTO 融合 kernel 的耗时。golden 参考见 swiglu_mlp_golden.pyF.silu(F.linear(...))标准实现。7.2 GLM MoE 共享专家量化 FFNAT-07 卡片 examples 中的 GLMMoEFusion 对应 test_glm_ffn_shared_expert_quant.py其参考路径展示了更贴近上线的组合output_w13 torch_npu.npu_quant_matmul(quantized_x, w13, w13_scale, ...) swiglu_out torch_npu.npu_swiglu(output_w13) # 单算子完成 silu(gate)*value quantized_x, x_scale torch_npu.npu_dynamic_quant(swiglu_out) output torch_npu.npu_quant_matmul(quantized_x, w2, ...)即 w13 合并投影 → SwiGLU → per-token 量化 → w2 投影的四步链路与 7.1 的三段式相比gate/up 两个投影被合并为一次更宽的 MatMulN 翻倍后对半切这正是第六节列切分思想的模型级版本。八、设计要点与常见陷阱综合模式卡片与上述实现设计 SwiGLU 相关算子时需注意精度走 FP32 中间无论变体 A/Bsilu 的乘加链都应保持在 FP32输出前再 castgate/value 输入为 FP32 契约的初衷即在此。sigmoid 原语可用性先行确认变体 Aexp 展开与变体 B直接 sigmoid/silu取决于目标版本是否提供原语设计时按 pypto-docs-search 工作流核对目标版本文档后再定。两个 MatMul 的合并机会gate/value 投影可合并为一次 N 翻倍的 MatMul再按列对半切分GMM SwiGLU、GLM w13 均为该模式可显著减少 Cube 调用次数代价是中间显存与寄存器占用上升需按 Tiling 约束核算。与量化的衔接SwiGLU 输出经常直接进入 per-token 量化127.0 / amax缩放 CAST_RINT舍入激活段精度直接决定量化误差属于 数据流约束 中需要重点标注的精度敏感点。Vector tile 配置参考实现中前向常用set_vec_tile_shapes(128, 128)Silu/MLP 场景也有(64, 512)、(64, 256)等组合属可调参数设计时标注为tunable而非固定常量。动态轴分块M 轴动态时用pypto.looppypto.view(valid_shape...)处理尾块bias 等小张量升 FP32 尽量移到循环外。九、总结AT-07 SwiGLU Activation 是纯 Vector、低计算强度、高融合潜力的经典原子模式前向只有 sigmoid 展开加两次逐元素乘法却能以列切分、算子融合两种方式与 MatMul/量化深度结合成为 FusedSwiGLU、GMM SwiGLU、GLMMoEFusion 三类算子共同的数学内核。设计者可将本文第四节至第七节的实现作为 DESIGN.md 的对照样板按 AT-07 卡片 的输入输出契约结合目标版本文档确认 sigmoid 原语可用性后直接复用仓库中的 tile 与精度配置快速完成可实现、可验证的设计。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考