GateGeluQuant 算子深度解析CANN ops-transformer 中 GeGLU 与 Per-Channel 量化融合 Kernel 的 Tiling 与实现原理【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读GateGeluQuant 是 CANN ops-transformer 仓库experimental/moe/gategelu_quant中提供的一个面向大模型 FFN 层的 NPU 融合算子它把 Gated GLU 结构中的 GELU 激活、门控逐元素乘法、激活值截断约束和 Per-Channel 动态量化输出 INT8四个操作融合进单个 AI Core Kernel直接服务于 LLaMA、Qwen、DeepSeek 等大模型 W8A8 量化推理部署。本文以该算子官方 README 为骨架结合仓库内 GATEGELU_QUANT2D_SBUF.cpp 源码、test_gategelu_quant.py 测试与 CMakeLists.txt 编译配置完整讲解其数学原理、参数语义、Tiling 策略、Kernel 计算流水线、编译与精度验证方法读完即可掌握该融合算子的设计与使用全貌。一、算子概览与产品支持情况GateGeluQuant 的定位是GeGLU Per-Channel Quantization的融合算子输入张量按列均分为 Gate 与 Value 两部分对 Gate 部分施加 GELU 激活函数后与 Value 部分逐元素相乘再经过可选的截断约束最后乘以缩放因子量化输出为 INT8。产品支持情况如下引自 README.md产品是否支持Atlas A2 训练系列产品是与之对应该算子的编译配置CMakeLists.txt中显式指定了目标 SoC 版本为Ascend910B1、核类型为VecCoreset_source_files_properties( ${GATEGELU_QUANT_NPU_SOURCES} PROPERTIES LANGUAGE CXX COMPILE_FLAGS --cce-soc-versionAscend910B1 --cce-soc-core-typeVecCore --cce-auto-sync -xcce )这意味着该 Kernel 面向 Ascend 910B 系列Atlas A2 训练系列的向量核VecCore编写使用 AscendC 编程模型实现。二、数学原理GeGLU Per-Channel 量化GLU门控线性单元Gated Linear Unit将输入在最后一维拆成两支一支经过激活函数后作为门另一支直接参与逐元素乘法。GateGeluQuant 中的门激活函数选用 GELU因此得到 GeGLU。设输入张量in的 shape 为(gbH, gbW)其中gbW 2 × hidden_size前半列[:, :W]为 Gate后半列[:, W:]为 Value且W gbW / 2。算子整体数学公式如下$$intermediate GELU(input[:, :W]) \odot input[:, W:]$$$$if \ constrait: \ intermediate Clamp(intermediate, -clampValue, clampValue)$$$$output Quantize(intermediate \times scale) \quad \in [-128, 127]$$其中⊙表示逐元素乘法Quantize表示四舍五入取整round并截断到 INT8 表示范围[-128, 127]scale是长度为gbW / 2的 Per-Channel 量化缩放因子即输出张量的每一列channel对应一个独立的 float32 缩放系数而非整个张量共用一个标量这正是Per-Channel量化的含义。注意constrait为 true 时的截断约束施加在乘以 scale 之前的 FP32 GeGLU 中间结果上其目的是在低精度量化前先行限制激活值幅值避免后续放大后溢出 INT8 表示范围详见下文 Kernel 计算流水线的步骤 4。三、参数说明下表完整继承自 README.md并补充了默认值、shape 关系等实操细节参数名输入/输出/属性描述数据类型数据格式in输入Gate 和 Value 拼接的输入张量shape 为(gbH, gbW)其中gbW 2 × hidden_size前半部分为 Gate后半部分为 Valuefloat16NDscale输入Per-Channel 量化缩放因子shape 为(gbW / 2, )float32NDout输出GeGLU 计算并量化后的 INT8 结果shape 为(gbH, gbW / 2)int8_tNDgbH输入输入张量的行数token 数量int64_t-gbW输入输入张量的列数Gate Value 拼接后的隐藏维度必须为偶数int64_t-constrait输入是否在量化前对 GeGLU 的 FP32 中间结果进行截断约束默认为 falsebool-clampValue输入截断约束的阈值配合constrait使用将值截断在[-clampValue, clampValue]内默认为 128.0float32-blockDim输入AI Core 的数量如 Ascend910B 为 40int64_t-stream输入Device 端的 streamAclrtStream-几个值得注意的细节三个张量的 shape 存在严格关联in为(gbH, gbW)scale为(gbW / 2,)out为(gbH, gbW / 2)即列方向经量化后缩减一半行数保持不变。gbW必须为偶数否则无法均分为 Gate / Value 两部分源码Init中直接以bkW_ gbW_ / 2作为输出列宽。constrait/clampValue与源码成员constrait_/clampValue_的默认值false / 128.0f完全一致见 GATEGELU_QUANT2D_SBUF.cpp注意约束一词在文档与源码中的拼写均为constrait。blockDim传入后并非无条件生效启动接口gateGeluQuant2dSBuf_lanuch中带有if (gbH blockDim) { blockDim gbH; }的保护逻辑GATEGELU_QUANT2D_SBUF.cpp即当 token 数行数少于 AI Core 数量时实际并行核数会被裁剪为行数避免空转核。四、约束说明输入张量的列数gbW必须为偶数以便均分为 Gate 和 Value 两部分。数据类型约束输入为 float16Scale 为 float32输出为 int8_t。从实现角度还可以补充两条推导约束其一tile 宽度在 UB 容量计算后需对齐到 64 元素粒度对应 128 字节对齐见下文 Tiling 策略其二输出是 INT8 定点结果因此在量化前以RoundMode::CAST_RINT四舍五入取整再以Mins/Maxs截断到[-128.0, 127.0]防止量化溢出。五、应用价值W8A8 量化推理中的 FFN 关键步骤在 LLaMA、Qwen、DeepSeek 等大语言模型的 W8A8 量化推理部署中GeGLU 激活及其输出量化是 FFN 层计算的核心步骤。若按朴素方式实现这一路径需要先写回 FP16/FP32 的 GeGLU 中间结果到 Global Memory再启动独立的量化 Kernel 读入、缩放、取整后写回 INT8中间会引入额外的显存带宽占用与 Kernel 启动开销。GateGeluQuant 的价值在于把GELU 激活、门控乘法、激活值截断及 Per-Channel 动态量化四个操作融合为一个 Kernel彻底消除了 FP16/FP32 中间结果的 Global Memory 写回与读入中间结果始终停留在片上 UB 中见下文 Compute 流水线从而大幅降低显存带宽压力和 Kernel 启动开销提升端到端推理性能以上效果描述引自 README.md 的价值/作用章节属于项目文档声明。从结构上看它位于 MoE / FFN 算子在experimental/moe目录下的算子族谱中与该目录中的biasgategelu、moegategeluclamp等门控激活类算子互为参照experimental/moe。六、设计方案一Tiling 策略分核策略按行gbH切分算子按行维度进行分核每个 Core 处理⌈gbH / blockNum⌉行数据向上取整。核心代码位于InitGATEGELU_QUANT2D_SBUF.cppbkLoop_ (int64_t)(gbH_ / blockNum_); if (gbH_ % blockNum_ ! 0) { bkLoop_ 1; }每个 Core 负责的行索引并非连续区间而是按核号跨步分布第i轮处理的行号是i * blockNum_ blockIdx_。行循环Process中通过if (i * blockNum_ blockIdx_ gbH_)判断当前行是否有效从而跳过因向上取整带来的尾部多余迭代GATEGELU_QUANT2D_SBUF.cpp。这种跨步分核方式使相邻核尽量处理相邻行配合 GELU 等逐行运算特性可均衡各核负载。分块策略按列bkW在 UB 容量约束下切 Tile在列方向上输出宽度bkW_ gbW / 2每个 Core 在每行内按 UB 可用容量继续切分 Tile计算单元素占用Init中按每个输出元素所需 Buffer估算字节占用temp——Gate、Value 各 1 个 half2 字节1 个 float scale4 字节1 个 int8 输出1 字节对应源码GATEGELU_QUANT2D_SBUF.cppint64_t temp BUFFER_NUM * bkH_ * sizeof(half) * 2; // Gate Value 输入 temp BUFFER_NUM * 2 * sizeof(half); temp BUFFER_NUM * sizeof(float); // scale temp BUFFER_NUM * bkH_ * sizeof(int8_t); // 输出求最大 Tile 宽度tlMaxW_ UB_MAX_BYTES / temp其中UB_MAX_BYTES 184 * 1024184KB见 GATEGELU_QUANT2D_SBUF.cpp随后tlMaxW_ tlMaxW_ / 64 * 64向下对齐到 64 元素以满足 128 字节对齐要求half 元素 2 字节64 × 2 128。确定实际 Tile 宽度tlW_GATEGELU_QUANT2D_SBUF.cpp若整行宽度bkW_不超过tlMaxW_则tlW_ bkW_单 Tile 处理整行否则先计算 Tile 数量向上取整再用tlW_ bkW_ / temp反推每个 Tile 的平均宽度最后tlW_ AlignUp(tlW_, 64)向上对齐到 64 元素。这种先定 Tile 数、再反推宽度的策略可保证 Tile 数量最少且各 Tile 尽量宽减少循环开销。尾部 Tile 处理tlTailW_ bkW_ % tlW_为尾部不足一个完整 Tile 的宽度tlAlignTailW_ AlignUp(tlTailW_, 64)为用于向量计算的对齐宽度向量指令要求对齐而实际有效宽度real_tlW仅用于搬入/搬出DataCopyPad按实际字节数搬运。tlLoop_ ceil(bkW_ / tlW_)给出每行内的 Tile 循环次数。七、设计方案二Kernel 侧设计整个 Kernel 采用Init Process两阶段结构其中 Process 内又分为数据搬入CopyIn、计算Compute、数据搬出CopyOut三步并使用单缓冲BUFFER_NUM 1机制即计算与搬入搬出不叠加流水无 double buffer 双缓冲乒乓简化了队列与同步管理。初始化阶段InitInit完成四类工作GATEGELU_QUANT2D_SBUF.cpp分核参数bkLoop_每个 Core 处理的行数、blockIdx_当前 Core 编号取自GetBlockIdx()并保存gbH_、gbW_、constrait_、clampValue_分块参数bkW_ gbW / 2基于 UB 容量计算tlMaxW_进而确定tlW_、tlTailW_、tlAlignTailW_、tlLoop_见上文分块策略此外还计算了按 32 对齐的整行宽度bkAlignW_作为地址对齐参考GM Tensor 映射建立inGm_half、scaleGm_float、outGm_int8三个 GlobalTensor分别绑定到in、scale、out的 GM 地址并声明元素个数队列初始化VECIN 输入队列inQueIn_Gate 与 Value 合并存放深度BUFFER_NUM、inQueScale_缩放因子深度 1VECOUT 输出队列outQueOut_INT8 输出深度BUFFER_NUM。计算流程ProcessProcess 外层按行循环FOR i 0 TO bkLoop_内层处理尾部 Tile 与完整 Tile 两种路径GATEGELU_QUANT2D_SBUF.cpp。结合源码各阶段细节如下CopyIn数据搬入从 GM 搬入拼接的 Gate 和 Value 数据到同一块本地内存in_local[0]存 Gatein_local[tlW_]存 Value——两次DataCopyPad的 GM 源地址相差bkW_个 half 元素即从 Gate 起始列跳到 Value 起始列行偏移为(i * blockNum_ blockIdx_) * (bkW_ * 2)同时将当前 Tile 对应的 float32scale数据搬入inQueScale_队列GATEGELU_QUANT2D_SBUF.cpp。Compute高度融合的量化计算流水线这是整个算子的核心共 8 步全部在 UB 内完成中间结果不落 GMGATEGELU_QUANT2D_SBUF.cppGelu(in_one, in_one)对 Gate 部分in_one_local原地计算 GELU 激活Mul(in_two, in_one, in_two)Gate已激活与 Valuein_two_local逐元素相乘得到 GeGLU 结果Cast(infloat_local, in_two, CAST_NONE)将 FP16 结果转为 FP32为高精度量化计算做准备[可选约束]若constrait_为 true则用Mins(infloat_local, clampValue_)与Maxs(infloat_local, -clampValue_)将 FP32 结果截断在[-clampValue_, clampValue_]内Mul(infloat_local, infloat_local, scale_local)乘以 Per-Channel 量化缩放因子按列对应逐元素广播Cast(infloat_local, infloat_local, CAST_RINT)以四舍五入round-to-nearest模式将 FP32 取整为整数仍以 FP32 格式存储Mins/Maxs将取整结果截断到[-128.0, 127.0]的 INT8 表示范围防止溢出Cast(in_one_local, infloat_local, CAST_NONE)与Cast(out_local, in_one_local, CAST_RINT)将结果从 FP32 经 FP16 中转最终转为 INT8 输出。需要说明的是步骤 8 的FP32 → FP16 → INT8两级 Cast 是为了利用向量指令的数据通路特性完成最终定点化属于实现层面的精度/指令权衡。CopyOut数据搬出将 INT8 计算结果从 UB 搬回 GM输出偏移量为offset (i * blockNum_ blockIdx_) * bkW_ j * tlW_GATEGELU_QUANT2D_SBUF.cpp与输入行的跨步映射一一对应。Kernel 入口与启动Kernel 采用extern C __global__ __aicore__导出参数为(gbH, gbW, in, scale, out, constrait, clampValue)通过gateGeluQuant2dSBuf_kernelblockDim, nullptr, stream启动GATEGELU_QUANT2D_SBUF.cpp。gategelu_quant_lanuch作为对外 C 接口返回启动状态其中包含上文提到的gbH blockDim时裁剪核数的保护。八、PyTorch 侧接入方式除了裸 Kernel 启动接口仓库还提供了 PyTorch 算子接入封装gategelu_quant_npuGATEGELU_QUANT2D_SBUF.cpp其要点包括使用TORCH_CHECK(torch_npu::utils::is_npu(...))校验in、scale、out三个张量均位于 NPU 设备上从输入张量推导gbH inTensor.size(0)、gbW inTensor.size(1)通过c10_npu::getCurrentNPUStream()获取当前 NPU stream以at_npu::native::OpCommand::RunOpApi(GategeluQuant, acl_call)异步执行 Kernel通过TORCH_LIBRARY_IMPL(ascend_ops, PrivateUse1, m)注册自定义算子gategelu_quant供torch.ops.ascend_ops.gategelu_quant(...)调用。结合测试脚本 test_gategelu_quant.py 中的调用示例NPU 侧最小调用方式为import torch import torch_npu import ascend_ops GBH, GBW, BLOCKDIM 4, 64, 4 input_npu torch.randn(GBH, GBW, dtypetorch.float16).npu() # (gbH, gbW) fp16 scale_npu torch.ones(GBW // 2, dtypetorch.float32).npu() # (gbW/2,) fp32 out_npu torch.empty(GBH, GBW // 2, dtypetorch.int8).npu() # (gbH, gbW/2) int8 torch.ops.ascend_ops.gategelu_quant(BLOCKDIM, input_npu, scale_npu, out_npu, False, 128.0) # 参数依次为: blockDim, in, scale, out, constrait, clampValue使用时需满足上文参数表的类型与 shape 约束且三张量必须在 NPU 上constraitFalse表示关闭截断约束测试即采用该配置。九、编译与精度验证编译集成该算子通过 CMakeLists.txt 以对象库形式加入构建以file(GLOB ...)收集目录内全部.cpp设置Ascend910B1/VecCore编译属性后创建gategelu_quant_objects对象库并挂接COMMON_COMPILE_OPTIONS与COMMON_INCLUDE_DIRS。其所在目录由 experimental/moe/CMakeLists.txt 遍历各子目录的CMakeLists.txt后统一add_subdirectory引入因此它是experimental/moe算子族整体构建链路上的一员。测试脚本与精度标准仓库提供了 CPU 参考实现对拍测试 test_gategelu_quant.py验证流程如下测试配置GBH4、GBW64、BLOCKDIM4输入用torch.randn(...) * 10生成 fp16 随机数据scale 用torch.ones(...) * 10构造constraitFalseCPU 参考实现gategelu_quant_cpu按列拆分in_one input[:, :GBW//2]、in_two input[:, GBW//2:]以 tanh 近似 GELU0.5x(1tanh(sqrt(2/π)(x0.044715x³)))计算激活、门控乘法、乘 scale、截断到 INT8 范围后round取整与算子公式逐一对齐结果对比同时统计绝对误差与基于 FP32 中间值的相对误差判定标准最大绝对误差 ≤ 1INT8 量化本身允许 1 个 LSB 的取整误差最大相对误差 ≤ 0.011%两者同时满足即判定CPU comparison test PASSED。该测试为验证算子数值正确性提供了可直接复现的对拍方法替换 shape / scale 取值即可扩展覆盖不同分核blockDim与gbH大小关系与尾部 TilebkW_ % tlW_ ! 0路径。十、小结GateGeluQuant 是 ops-transformer 仓库中一个典型的以访存优化为核心的融合算子范例按行跨步分核 按列 UB 容量分块的双层 Tiling 策略保证了多核并行度与片上存储的匹配8 步在片内完成的融合量化流水线GELU → 门控乘 → FP32 提升 → 可选截断 → 乘 scale → 取整 → INT8 范围截断 → 定点化使 FP16/FP32 中间结果始终不落 Global Memory而constrait/clampValue的运行时开关则为不同量化策略是否预截断激活值保留了灵活性。配合仓库内的 CPU 对拍测试与 Ascend910B1 编译配置该算子可以作为在 CANN 生态中实现GeGLU Per-Channel 量化融合 Kernel 的完整参考实现。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
长河君 · 2026/9/18 22:36:53 · 阅读 1.2k GateGeluQuant 算子深度解析：CANN ops-transformer 中 GeGLU 与 Per-Channel 量化融合 Kernel 的 Tiling 与实现原理
GateGeluQuant 算子深度解析CANN ops-transformer 中 GeGLU 与 Per-Channel 量化融合 Kernel 的 Tiling 与实现原理【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer导读GateGeluQuant 是 CANN ops-transformer 仓库experimental/moe/gategelu_quant中提供的一个面向大模型 FFN 层的 NPU 融合算子它把 Gated GLU 结构中的 GELU 激活、门控逐元素乘法、激活值截断约束和 Per-Channel 动态量化输出 INT8四个操作融合进单个 AI Core Kernel直接服务于 LLaMA、Qwen、DeepSeek 等大模型 W8A8 量化推理部署。本文以该算子官方 README 为骨架结合仓库内 GATEGELU_QUANT2D_SBUF.cpp 源码、test_gategelu_quant.py 测试与 CMakeLists.txt 编译配置完整讲解其数学原理、参数语义、Tiling 策略、Kernel 计算流水线、编译与精度验证方法读完即可掌握该融合算子的设计与使用全貌。一、算子概览与产品支持情况GateGeluQuant 的定位是GeGLU Per-Channel Quantization的融合算子输入张量按列均分为 Gate 与 Value 两部分对 Gate 部分施加 GELU 激活函数后与 Value 部分逐元素相乘再经过可选的截断约束最后乘以缩放因子量化输出为 INT8。产品支持情况如下引自 README.md产品是否支持Atlas A2 训练系列产品是与之对应该算子的编译配置CMakeLists.txt中显式指定了目标 SoC 版本为Ascend910B1、核类型为VecCoreset_source_files_properties( ${GATEGELU_QUANT_NPU_SOURCES} PROPERTIES LANGUAGE CXX COMPILE_FLAGS --cce-soc-versionAscend910B1 --cce-soc-core-typeVecCore --cce-auto-sync -xcce )这意味着该 Kernel 面向 Ascend 910B 系列Atlas A2 训练系列的向量核VecCore编写使用 AscendC 编程模型实现。二、数学原理GeGLU Per-Channel 量化GLU门控线性单元Gated Linear Unit将输入在最后一维拆成两支一支经过激活函数后作为门另一支直接参与逐元素乘法。GateGeluQuant 中的门激活函数选用 GELU因此得到 GeGLU。设输入张量in的 shape 为(gbH, gbW)其中gbW 2 × hidden_size前半列[:, :W]为 Gate后半列[:, W:]为 Value且W gbW / 2。算子整体数学公式如下$$intermediate GELU(input[:, :W]) \odot input[:, W:]$$$$if \ constrait: \ intermediate Clamp(intermediate, -clampValue, clampValue)$$$$output Quantize(intermediate \times scale) \quad \in [-128, 127]$$其中⊙表示逐元素乘法Quantize表示四舍五入取整round并截断到 INT8 表示范围[-128, 127]scale是长度为gbW / 2的 Per-Channel 量化缩放因子即输出张量的每一列channel对应一个独立的 float32 缩放系数而非整个张量共用一个标量这正是Per-Channel量化的含义。注意constrait为 true 时的截断约束施加在乘以 scale 之前的 FP32 GeGLU 中间结果上其目的是在低精度量化前先行限制激活值幅值避免后续放大后溢出 INT8 表示范围详见下文 Kernel 计算流水线的步骤 4。三、参数说明下表完整继承自 README.md并补充了默认值、shape 关系等实操细节参数名输入/输出/属性描述数据类型数据格式in输入Gate 和 Value 拼接的输入张量shape 为(gbH, gbW)其中gbW 2 × hidden_size前半部分为 Gate后半部分为 Valuefloat16NDscale输入Per-Channel 量化缩放因子shape 为(gbW / 2, )float32NDout输出GeGLU 计算并量化后的 INT8 结果shape 为(gbH, gbW / 2)int8_tNDgbH输入输入张量的行数token 数量int64_t-gbW输入输入张量的列数Gate Value 拼接后的隐藏维度必须为偶数int64_t-constrait输入是否在量化前对 GeGLU 的 FP32 中间结果进行截断约束默认为 falsebool-clampValue输入截断约束的阈值配合constrait使用将值截断在[-clampValue, clampValue]内默认为 128.0float32-blockDim输入AI Core 的数量如 Ascend910B 为 40int64_t-stream输入Device 端的 streamAclrtStream-几个值得注意的细节三个张量的 shape 存在严格关联in为(gbH, gbW)scale为(gbW / 2,)out为(gbH, gbW / 2)即列方向经量化后缩减一半行数保持不变。gbW必须为偶数否则无法均分为 Gate / Value 两部分源码Init中直接以bkW_ gbW_ / 2作为输出列宽。constrait/clampValue与源码成员constrait_/clampValue_的默认值false / 128.0f完全一致见 GATEGELU_QUANT2D_SBUF.cpp注意约束一词在文档与源码中的拼写均为constrait。blockDim传入后并非无条件生效启动接口gateGeluQuant2dSBuf_lanuch中带有if (gbH blockDim) { blockDim gbH; }的保护逻辑GATEGELU_QUANT2D_SBUF.cpp即当 token 数行数少于 AI Core 数量时实际并行核数会被裁剪为行数避免空转核。四、约束说明输入张量的列数gbW必须为偶数以便均分为 Gate 和 Value 两部分。数据类型约束输入为 float16Scale 为 float32输出为 int8_t。从实现角度还可以补充两条推导约束其一tile 宽度在 UB 容量计算后需对齐到 64 元素粒度对应 128 字节对齐见下文 Tiling 策略其二输出是 INT8 定点结果因此在量化前以RoundMode::CAST_RINT四舍五入取整再以Mins/Maxs截断到[-128.0, 127.0]防止量化溢出。五、应用价值W8A8 量化推理中的 FFN 关键步骤在 LLaMA、Qwen、DeepSeek 等大语言模型的 W8A8 量化推理部署中GeGLU 激活及其输出量化是 FFN 层计算的核心步骤。若按朴素方式实现这一路径需要先写回 FP16/FP32 的 GeGLU 中间结果到 Global Memory再启动独立的量化 Kernel 读入、缩放、取整后写回 INT8中间会引入额外的显存带宽占用与 Kernel 启动开销。GateGeluQuant 的价值在于把GELU 激活、门控乘法、激活值截断及 Per-Channel 动态量化四个操作融合为一个 Kernel彻底消除了 FP16/FP32 中间结果的 Global Memory 写回与读入中间结果始终停留在片上 UB 中见下文 Compute 流水线从而大幅降低显存带宽压力和 Kernel 启动开销提升端到端推理性能以上效果描述引自 README.md 的价值/作用章节属于项目文档声明。从结构上看它位于 MoE / FFN 算子在experimental/moe目录下的算子族谱中与该目录中的biasgategelu、moegategeluclamp等门控激活类算子互为参照experimental/moe。六、设计方案一Tiling 策略分核策略按行gbH切分算子按行维度进行分核每个 Core 处理⌈gbH / blockNum⌉行数据向上取整。核心代码位于InitGATEGELU_QUANT2D_SBUF.cppbkLoop_ (int64_t)(gbH_ / blockNum_); if (gbH_ % blockNum_ ! 0) { bkLoop_ 1; }每个 Core 负责的行索引并非连续区间而是按核号跨步分布第i轮处理的行号是i * blockNum_ blockIdx_。行循环Process中通过if (i * blockNum_ blockIdx_ gbH_)判断当前行是否有效从而跳过因向上取整带来的尾部多余迭代GATEGELU_QUANT2D_SBUF.cpp。这种跨步分核方式使相邻核尽量处理相邻行配合 GELU 等逐行运算特性可均衡各核负载。分块策略按列bkW在 UB 容量约束下切 Tile在列方向上输出宽度bkW_ gbW / 2每个 Core 在每行内按 UB 可用容量继续切分 Tile计算单元素占用Init中按每个输出元素所需 Buffer估算字节占用temp——Gate、Value 各 1 个 half2 字节1 个 float scale4 字节1 个 int8 输出1 字节对应源码GATEGELU_QUANT2D_SBUF.cppint64_t temp BUFFER_NUM * bkH_ * sizeof(half) * 2; // Gate Value 输入 temp BUFFER_NUM * 2 * sizeof(half); temp BUFFER_NUM * sizeof(float); // scale temp BUFFER_NUM * bkH_ * sizeof(int8_t); // 输出求最大 Tile 宽度tlMaxW_ UB_MAX_BYTES / temp其中UB_MAX_BYTES 184 * 1024184KB见 GATEGELU_QUANT2D_SBUF.cpp随后tlMaxW_ tlMaxW_ / 64 * 64向下对齐到 64 元素以满足 128 字节对齐要求half 元素 2 字节64 × 2 128。确定实际 Tile 宽度tlW_GATEGELU_QUANT2D_SBUF.cpp若整行宽度bkW_不超过tlMaxW_则tlW_ bkW_单 Tile 处理整行否则先计算 Tile 数量向上取整再用tlW_ bkW_ / temp反推每个 Tile 的平均宽度最后tlW_ AlignUp(tlW_, 64)向上对齐到 64 元素。这种先定 Tile 数、再反推宽度的策略可保证 Tile 数量最少且各 Tile 尽量宽减少循环开销。尾部 Tile 处理tlTailW_ bkW_ % tlW_为尾部不足一个完整 Tile 的宽度tlAlignTailW_ AlignUp(tlTailW_, 64)为用于向量计算的对齐宽度向量指令要求对齐而实际有效宽度real_tlW仅用于搬入/搬出DataCopyPad按实际字节数搬运。tlLoop_ ceil(bkW_ / tlW_)给出每行内的 Tile 循环次数。七、设计方案二Kernel 侧设计整个 Kernel 采用Init Process两阶段结构其中 Process 内又分为数据搬入CopyIn、计算Compute、数据搬出CopyOut三步并使用单缓冲BUFFER_NUM 1机制即计算与搬入搬出不叠加流水无 double buffer 双缓冲乒乓简化了队列与同步管理。初始化阶段InitInit完成四类工作GATEGELU_QUANT2D_SBUF.cpp分核参数bkLoop_每个 Core 处理的行数、blockIdx_当前 Core 编号取自GetBlockIdx()并保存gbH_、gbW_、constrait_、clampValue_分块参数bkW_ gbW / 2基于 UB 容量计算tlMaxW_进而确定tlW_、tlTailW_、tlAlignTailW_、tlLoop_见上文分块策略此外还计算了按 32 对齐的整行宽度bkAlignW_作为地址对齐参考GM Tensor 映射建立inGm_half、scaleGm_float、outGm_int8三个 GlobalTensor分别绑定到in、scale、out的 GM 地址并声明元素个数队列初始化VECIN 输入队列inQueIn_Gate 与 Value 合并存放深度BUFFER_NUM、inQueScale_缩放因子深度 1VECOUT 输出队列outQueOut_INT8 输出深度BUFFER_NUM。计算流程ProcessProcess 外层按行循环FOR i 0 TO bkLoop_内层处理尾部 Tile 与完整 Tile 两种路径GATEGELU_QUANT2D_SBUF.cpp。结合源码各阶段细节如下CopyIn数据搬入从 GM 搬入拼接的 Gate 和 Value 数据到同一块本地内存in_local[0]存 Gatein_local[tlW_]存 Value——两次DataCopyPad的 GM 源地址相差bkW_个 half 元素即从 Gate 起始列跳到 Value 起始列行偏移为(i * blockNum_ blockIdx_) * (bkW_ * 2)同时将当前 Tile 对应的 float32scale数据搬入inQueScale_队列GATEGELU_QUANT2D_SBUF.cpp。Compute高度融合的量化计算流水线这是整个算子的核心共 8 步全部在 UB 内完成中间结果不落 GMGATEGELU_QUANT2D_SBUF.cppGelu(in_one, in_one)对 Gate 部分in_one_local原地计算 GELU 激活Mul(in_two, in_one, in_two)Gate已激活与 Valuein_two_local逐元素相乘得到 GeGLU 结果Cast(infloat_local, in_two, CAST_NONE)将 FP16 结果转为 FP32为高精度量化计算做准备[可选约束]若constrait_为 true则用Mins(infloat_local, clampValue_)与Maxs(infloat_local, -clampValue_)将 FP32 结果截断在[-clampValue_, clampValue_]内Mul(infloat_local, infloat_local, scale_local)乘以 Per-Channel 量化缩放因子按列对应逐元素广播Cast(infloat_local, infloat_local, CAST_RINT)以四舍五入round-to-nearest模式将 FP32 取整为整数仍以 FP32 格式存储Mins/Maxs将取整结果截断到[-128.0, 127.0]的 INT8 表示范围防止溢出Cast(in_one_local, infloat_local, CAST_NONE)与Cast(out_local, in_one_local, CAST_RINT)将结果从 FP32 经 FP16 中转最终转为 INT8 输出。需要说明的是步骤 8 的FP32 → FP16 → INT8两级 Cast 是为了利用向量指令的数据通路特性完成最终定点化属于实现层面的精度/指令权衡。CopyOut数据搬出将 INT8 计算结果从 UB 搬回 GM输出偏移量为offset (i * blockNum_ blockIdx_) * bkW_ j * tlW_GATEGELU_QUANT2D_SBUF.cpp与输入行的跨步映射一一对应。Kernel 入口与启动Kernel 采用extern C __global__ __aicore__导出参数为(gbH, gbW, in, scale, out, constrait, clampValue)通过gateGeluQuant2dSBuf_kernelblockDim, nullptr, stream启动GATEGELU_QUANT2D_SBUF.cpp。gategelu_quant_lanuch作为对外 C 接口返回启动状态其中包含上文提到的gbH blockDim时裁剪核数的保护。八、PyTorch 侧接入方式除了裸 Kernel 启动接口仓库还提供了 PyTorch 算子接入封装gategelu_quant_npuGATEGELU_QUANT2D_SBUF.cpp其要点包括使用TORCH_CHECK(torch_npu::utils::is_npu(...))校验in、scale、out三个张量均位于 NPU 设备上从输入张量推导gbH inTensor.size(0)、gbW inTensor.size(1)通过c10_npu::getCurrentNPUStream()获取当前 NPU stream以at_npu::native::OpCommand::RunOpApi(GategeluQuant, acl_call)异步执行 Kernel通过TORCH_LIBRARY_IMPL(ascend_ops, PrivateUse1, m)注册自定义算子gategelu_quant供torch.ops.ascend_ops.gategelu_quant(...)调用。结合测试脚本 test_gategelu_quant.py 中的调用示例NPU 侧最小调用方式为import torch import torch_npu import ascend_ops GBH, GBW, BLOCKDIM 4, 64, 4 input_npu torch.randn(GBH, GBW, dtypetorch.float16).npu() # (gbH, gbW) fp16 scale_npu torch.ones(GBW // 2, dtypetorch.float32).npu() # (gbW/2,) fp32 out_npu torch.empty(GBH, GBW // 2, dtypetorch.int8).npu() # (gbH, gbW/2) int8 torch.ops.ascend_ops.gategelu_quant(BLOCKDIM, input_npu, scale_npu, out_npu, False, 128.0) # 参数依次为: blockDim, in, scale, out, constrait, clampValue使用时需满足上文参数表的类型与 shape 约束且三张量必须在 NPU 上constraitFalse表示关闭截断约束测试即采用该配置。九、编译与精度验证编译集成该算子通过 CMakeLists.txt 以对象库形式加入构建以file(GLOB ...)收集目录内全部.cpp设置Ascend910B1/VecCore编译属性后创建gategelu_quant_objects对象库并挂接COMMON_COMPILE_OPTIONS与COMMON_INCLUDE_DIRS。其所在目录由 experimental/moe/CMakeLists.txt 遍历各子目录的CMakeLists.txt后统一add_subdirectory引入因此它是experimental/moe算子族整体构建链路上的一员。测试脚本与精度标准仓库提供了 CPU 参考实现对拍测试 test_gategelu_quant.py验证流程如下测试配置GBH4、GBW64、BLOCKDIM4输入用torch.randn(...) * 10生成 fp16 随机数据scale 用torch.ones(...) * 10构造constraitFalseCPU 参考实现gategelu_quant_cpu按列拆分in_one input[:, :GBW//2]、in_two input[:, GBW//2:]以 tanh 近似 GELU0.5x(1tanh(sqrt(2/π)(x0.044715x³)))计算激活、门控乘法、乘 scale、截断到 INT8 范围后round取整与算子公式逐一对齐结果对比同时统计绝对误差与基于 FP32 中间值的相对误差判定标准最大绝对误差 ≤ 1INT8 量化本身允许 1 个 LSB 的取整误差最大相对误差 ≤ 0.011%两者同时满足即判定CPU comparison test PASSED。该测试为验证算子数值正确性提供了可直接复现的对拍方法替换 shape / scale 取值即可扩展覆盖不同分核blockDim与gbH大小关系与尾部 TilebkW_ % tlW_ ! 0路径。十、小结GateGeluQuant 是 ops-transformer 仓库中一个典型的以访存优化为核心的融合算子范例按行跨步分核 按列 UB 容量分块的双层 Tiling 策略保证了多核并行度与片上存储的匹配8 步在片内完成的融合量化流水线GELU → 门控乘 → FP32 提升 → 可选截断 → 乘 scale → 取整 → INT8 范围截断 → 定点化使 FP16/FP32 中间结果始终不落 Global Memory而constrait/clampValue的运行时开关则为不同量化策略是否预截断激活值保留了灵活性。配合仓库内的 CPU 对拍测试与 Ascend910B1 编译配置该算子可以作为在 CANN 生态中实现GeGLU Per-Channel 量化融合 Kernel 的完整参考实现。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考