CANN ops-nn 算子实践situ_glu 的 SiTU 门控线性单元接口详解与 NPU 加速原理【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn导读cann_ops_nn.situ_glu是 CANN ops-nn 算子库提供的 SiTU 门控线性单元SiTU Gated Linear Unit激活函数接口面向 Ascend NPU 训练与推理场景。本文以 torchapi_situ_glu.md 为核心完整讲解该接口的产品支持矩阵、计算公式、参数语义、调用方式与约束并结合仓库中的算子定义、形状推导、tiling 与内核实现剖析其在 NPU 上切分—激活—乘门控的底层计算链路帮助读者在 PyTorch 单算子模式与 TorchAir 图模式下正确、高效地使用该算子。产品支持情况根据 torchapi_situ_glu.md 中的产品支持矩阵situ_glu的支持情况如下产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持该支持矩阵与源码中的平台能力校验保持一致在 situ_glu_tiling.cpp 的IsCapable()中仅当socVersion_为ASCEND910_93、ASCEND910B、ASCEND950之一时才返回true算子定义侧 situ_glu_def.cpp 也只向ascend910b、ascend910_93、ascend950三个平台注册了 AICore 配置对应 op_host/config 目录下的ascend910b、ascend910_93、ascend950三份 binary 配置 json。功能说明接口功能SiTU 门控线性单元SiTU Gated Linear Unit激活函数对输入张量x沿指定维度dim切分为门控gate与上路径up两半按 SiTU 公式计算输出。PyTorch 层接口cann_ops_nn.situ_glu底层封装aclnnSituGlu算子接口对应文档见 aclnnSituGlu.md最终落到 AICore 内核执行。计算公式对给定的输入张量x按以下步骤计算 SiTU 输出y合轴基于dim进行合轴合轴后维度为[pre, cut]其中cut必须为偶数。令h cut // 2。前后切分根据activate_left对x进行切分activate_leftTrue默认gate x[..., :h]up x[..., h:]activate_leftFalsegate x[..., h:]up x[..., :h]SiTU 计算根据beta、linear_beta进行计算$$ situ_a \beta \cdot \tanh\left(\frac{gate}{\beta}\right) \cdot \text{sigmoid}(gate) $$当linear_beta 0时对 up 路径施加有界化线性 tanh 变换$$ up linear_beta \cdot \tanh\left(\frac{up}{linear_beta}\right) $$当linear_beta 0时up直接透传。输出$$ y situ_a \cdot up $$输出形状输出张量y的维度数量与x一致dim轴上的大小为x的一半其他维度与x相同。以上计算流程在 AICore 内核 situ_glu.hpp 的SituCore中得到逐指令印证先Muls乘1/beta、Tanh、再乘回beta得到beta * tanh(gate/beta)随后Sigmoid(gate)并与前者Mul得到situ_a当linearBeta_ 0时对 up 走同样的 tanh 有界化变换后与situ_a相乘否则直接Mul(y, y, up, n)透传 up。函数原型cann_ops_nn.situ_glu( x, *, dim-1, beta1.0, linear_beta0.0, activate_leftTrue ) - Tensor该签名在 situ_glu.py 中以torch.library的 schema 形式注册为situ_glu(Tensor x, *, int dim-1, float beta1.0, float linear_beta0.0, bool activate_leftTrue) - Tensor并通过OpBuilder动态加载 C 扩展situ_glu.cpp见 csrc/situ_glu.cpp完成算子调用。参数说明参数名参数类型可选/必选描述数据类型维度(shape)xTensor必选SiTU 输入dim维度会被均分为门控gate与上路径up两部分。float16、float32、bfloat161-8 维dim维度为偶数dimint可选切分维度取值范围[-x.dim(), x.dim()-1]默认-1。--betafloat可选SiTU 门控部分的缩放系数控制 tanh 非线性强度。建议为非零正值默认1.0。--linear_betafloat可选up 路径线性 tanh 的缩放系数。大于 0 时对 up 施加有界化变换小于等于 0 时 up 直接透传。默认0.0。--activate_leftbool可选切分x时 gate 是否为前半部分。True表示 gate 为前半、up 为后半False表示 gate 为后半、up 为前半。默认True。--补充说明来源于源码校验逻辑x仅支持 float32 / float16 / bfloat16 三种数据类型且要求 ND 格式。该约束在 situ_glu_def.cppxDtype {DT_FLOAT, DT_FLOAT16, DT_BF16}、xFormat FORMAT_ND和 situ_glu_tiling.cpp 的SUPPORT_DTYPE中均有校验。dim支持负索引内部会统一转换为正数cutDim_ cutDim_ 0 ? (cutDim_ xDims_) : cutDim_越界时在 infershape 与 tiling 阶段都会报错。beta与linear_beta的默认值在 situ_glu_def.cpp 中定义为DEFAULT_BETA 1.0、DEFAULT_LINEAR_BETA 0.0。dim维度大小必须为偶数否则在 Python 侧 situ_glu.py 与 infershape 侧 situ_glu_infershape.cpp 都会抛出 should be even / not divisible by 2 类错误。返回值说明参数名参数类型描述数据类型维度(shape)yTensorSiTU 激活结果。与x一致与x相同但dim维度大小为x.shape[dim] // 2输出形状由形状推导逻辑确定在 situ_glu_infershape.cpp 中输出yShape先复制xShape再仅将splitDim一维设为x.shape[dim] / 2输出数据类型与输入保持一致InferDataTypeForSituGlu。Python 侧 situ_glu.py 的_situ_glu_shape也实现了相同的 shape 推导并注册为 Meta 实现供动态 shape 场景使用。约束说明该接口支持训练、推理场景下使用。该接口支持单算子模式和 TorchAir 图模式调用。dim非 last 维时若尾部维度组合不满足半行 32B 对齐会回退到 Long-H 路径功能正确但性能略降。最后一条约束可以结合 tiling 源码深入理解在 situ_glu_tiling.cpp 的CountMaxPair()中tiling 会判断halfRowAligned (dimH * dsize) % 32 0即半行字节数是否 32B 对齐。当半行不满足 32B 对齐或 UB 一次装不下整行dimH时isLongH_被置 1内核走 Long-H 连续单块拷贝路径由DataCopyPad处理 sub-block 对齐否则走 Short-H 行 stride 拷贝路径。两种路径对应 situ_glu.hpp 中的CopyInShortH按行跳步拷贝 gate/up 两半与CopyInLongH连续拷贝后再切半CopyOut同样区分两种模式。确定性计算该接口默认支持确定性计算。这意味着在相同输入与参数下多次调用得到的计算结果保持一致适合对可复现性有要求的训练与推理流水。调用示例单算子模式调用单算子模式通过cann_ops_nn.situ_glu直接调用前提是环境已安装并初始化torch_npuimport torch import torch_npu import cann_ops_nn x torch.randn(2, 32, dtypetorch.float32).npu() # 默认参数调用 y cann_ops_nn.situ_glu(x) print(y.shape, y.dtype) # torch.Size([2, 16]) torch.float32 # 指定全部参数 y cann_ops_nn.situ_glu(x, dim-1, beta1.0, linear_beta0.0, activate_leftTrue) print(y.shape, y.dtype)示例中x的 shape 为(2, 32)默认dim-1表示在最后一维大小为 32上切分因此输出 shape 为(2, 16)。若把dim指定为0则按第一维切分输出 shape 将变为(1, 32)。注意dim维度大小必须为偶数否则调用会直接报错。TorchAir 图模式调用图模式通过torch.compile配合torchair的 NPU 后端完成构图与编译import torch import torch_npu import torchair import cann_ops_nn class Model(torch.nn.Module): def forward(self, x): return cann_ops_nn.situ_glu(x, dim-1, beta1.0) model torch.compile(Model().npu(), backendtorchair.get_npu_backend(), dynamicFalse) x torch.randn(2, 32, dtypetorch.float32).npu() y model(x)图模式场景下算子会经由 graph_convert_situ_glu.py 完成图中算子转换并以 situ_glu_proto.h 定义的算子 IR 形式参与构图。底层实现从 PyTorch 调用到 AICore 内核调用链路situ_glu的完整调用链可归纳为Python 层cann_ops_nn.situ_glu(...)→OpBuilder动态加载 C 扩展situ_glu.py形状推导Meta 实现_situ_glu_shape计算输出 shapedim校验与偶数校验在此完成situ_glu.pyHost 层算子定义situ_glu_def.cpp声明输入/输出/属性并注册平台infershapesitu_glu_infershape.cpp推导输出形状与数据类型Tiling 层situ_glu_tiling.cpp 根据平台 AIV 核数、UB 内存大小与输入 shape 计算切分参数coreNumAll、dimBatchSize、dim2H、isLongH、ubMaxPair等并申请 16MB workspaceKernel 层AICore 内核 situ_glu.cpp 按ORIG_DTYPE_X编译为 float / half / bfloat16 三种实例执行 situ_glu.hpp 中的分块计算。内核计算细节合轴与分核tiling 将x合轴为[dimBatchSize, dim2H]dimBatchSize为dim之前各维乘积dim2H为dim及之后各维乘积再由CalTilingParamsitu_glu.hpp按 batch 在多个 AIV 核之间切分核内再按 UB 容量分 loop 处理。gate/up 切分CopyIn一次搬运 gate 与 up 两段到 xQueue前半与后半各占一块 bufferCompute中根据activateLeft_选择 gate 与 up 的偏移situ_glu.hpp。混合精度路径fp16 / bf16 输入先Cast到 float32 计算无损加宽算完再Cast回原类型其中 bf16 回写必须使用CAST_RINT舍入模式源码注释明确指出CAST_NONE在非 A2 平台如 Ascend950 上会产生错误结果见 situ_glu.hpp。FP32 原生路径float32 输入直接在原 buffer 上以 float 视图计算无需 Cast减少一次数据搬运。参考资源接口文档torchapi_situ_glu.md、底层 aclnn 接口文档 aclnnSituGlu.md算子 README含参数表与调用方式汇总README.md算子定义与推导situ_glu_def.cpp、situ_glu_infershape.cppTiling 实现situ_glu_tiling.cpp、situ_glu_tiling.hAICore 内核situ_glu.cpp、situ_glu.hppPyTorch 扩展situ_glu.py、csrc/situ_glu.cpp、graph_convert_situ_glu.py调用示例test_aclnn_situ_glu.cppaclnn API 方式测试用例test_situ_glu_network.py网络级 pytest、executor_aclnnSituGlu.pyST 执行器反向算子训练场景配套situ_glu_grad 目录及其接口文档 torchapi_situ_glu_grad.md【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-nn 算子实践：situ_glu 的 SiTU 门控线性单元接口详解与 NPU 加速原理
CANN ops-nn 算子实践situ_glu 的 SiTU 门控线性单元接口详解与 NPU 加速原理【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn导读cann_ops_nn.situ_glu是 CANN ops-nn 算子库提供的 SiTU 门控线性单元SiTU Gated Linear Unit激活函数接口面向 Ascend NPU 训练与推理场景。本文以 torchapi_situ_glu.md 为核心完整讲解该接口的产品支持矩阵、计算公式、参数语义、调用方式与约束并结合仓库中的算子定义、形状推导、tiling 与内核实现剖析其在 NPU 上切分—激活—乘门控的底层计算链路帮助读者在 PyTorch 单算子模式与 TorchAir 图模式下正确、高效地使用该算子。产品支持情况根据 torchapi_situ_glu.md 中的产品支持矩阵situ_glu的支持情况如下产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持该支持矩阵与源码中的平台能力校验保持一致在 situ_glu_tiling.cpp 的IsCapable()中仅当socVersion_为ASCEND910_93、ASCEND910B、ASCEND950之一时才返回true算子定义侧 situ_glu_def.cpp 也只向ascend910b、ascend910_93、ascend950三个平台注册了 AICore 配置对应 op_host/config 目录下的ascend910b、ascend910_93、ascend950三份 binary 配置 json。功能说明接口功能SiTU 门控线性单元SiTU Gated Linear Unit激活函数对输入张量x沿指定维度dim切分为门控gate与上路径up两半按 SiTU 公式计算输出。PyTorch 层接口cann_ops_nn.situ_glu底层封装aclnnSituGlu算子接口对应文档见 aclnnSituGlu.md最终落到 AICore 内核执行。计算公式对给定的输入张量x按以下步骤计算 SiTU 输出y合轴基于dim进行合轴合轴后维度为[pre, cut]其中cut必须为偶数。令h cut // 2。前后切分根据activate_left对x进行切分activate_leftTrue默认gate x[..., :h]up x[..., h:]activate_leftFalsegate x[..., h:]up x[..., :h]SiTU 计算根据beta、linear_beta进行计算$$ situ_a \beta \cdot \tanh\left(\frac{gate}{\beta}\right) \cdot \text{sigmoid}(gate) $$当linear_beta 0时对 up 路径施加有界化线性 tanh 变换$$ up linear_beta \cdot \tanh\left(\frac{up}{linear_beta}\right) $$当linear_beta 0时up直接透传。输出$$ y situ_a \cdot up $$输出形状输出张量y的维度数量与x一致dim轴上的大小为x的一半其他维度与x相同。以上计算流程在 AICore 内核 situ_glu.hpp 的SituCore中得到逐指令印证先Muls乘1/beta、Tanh、再乘回beta得到beta * tanh(gate/beta)随后Sigmoid(gate)并与前者Mul得到situ_a当linearBeta_ 0时对 up 走同样的 tanh 有界化变换后与situ_a相乘否则直接Mul(y, y, up, n)透传 up。函数原型cann_ops_nn.situ_glu( x, *, dim-1, beta1.0, linear_beta0.0, activate_leftTrue ) - Tensor该签名在 situ_glu.py 中以torch.library的 schema 形式注册为situ_glu(Tensor x, *, int dim-1, float beta1.0, float linear_beta0.0, bool activate_leftTrue) - Tensor并通过OpBuilder动态加载 C 扩展situ_glu.cpp见 csrc/situ_glu.cpp完成算子调用。参数说明参数名参数类型可选/必选描述数据类型维度(shape)xTensor必选SiTU 输入dim维度会被均分为门控gate与上路径up两部分。float16、float32、bfloat161-8 维dim维度为偶数dimint可选切分维度取值范围[-x.dim(), x.dim()-1]默认-1。--betafloat可选SiTU 门控部分的缩放系数控制 tanh 非线性强度。建议为非零正值默认1.0。--linear_betafloat可选up 路径线性 tanh 的缩放系数。大于 0 时对 up 施加有界化变换小于等于 0 时 up 直接透传。默认0.0。--activate_leftbool可选切分x时 gate 是否为前半部分。True表示 gate 为前半、up 为后半False表示 gate 为后半、up 为前半。默认True。--补充说明来源于源码校验逻辑x仅支持 float32 / float16 / bfloat16 三种数据类型且要求 ND 格式。该约束在 situ_glu_def.cppxDtype {DT_FLOAT, DT_FLOAT16, DT_BF16}、xFormat FORMAT_ND和 situ_glu_tiling.cpp 的SUPPORT_DTYPE中均有校验。dim支持负索引内部会统一转换为正数cutDim_ cutDim_ 0 ? (cutDim_ xDims_) : cutDim_越界时在 infershape 与 tiling 阶段都会报错。beta与linear_beta的默认值在 situ_glu_def.cpp 中定义为DEFAULT_BETA 1.0、DEFAULT_LINEAR_BETA 0.0。dim维度大小必须为偶数否则在 Python 侧 situ_glu.py 与 infershape 侧 situ_glu_infershape.cpp 都会抛出 should be even / not divisible by 2 类错误。返回值说明参数名参数类型描述数据类型维度(shape)yTensorSiTU 激活结果。与x一致与x相同但dim维度大小为x.shape[dim] // 2输出形状由形状推导逻辑确定在 situ_glu_infershape.cpp 中输出yShape先复制xShape再仅将splitDim一维设为x.shape[dim] / 2输出数据类型与输入保持一致InferDataTypeForSituGlu。Python 侧 situ_glu.py 的_situ_glu_shape也实现了相同的 shape 推导并注册为 Meta 实现供动态 shape 场景使用。约束说明该接口支持训练、推理场景下使用。该接口支持单算子模式和 TorchAir 图模式调用。dim非 last 维时若尾部维度组合不满足半行 32B 对齐会回退到 Long-H 路径功能正确但性能略降。最后一条约束可以结合 tiling 源码深入理解在 situ_glu_tiling.cpp 的CountMaxPair()中tiling 会判断halfRowAligned (dimH * dsize) % 32 0即半行字节数是否 32B 对齐。当半行不满足 32B 对齐或 UB 一次装不下整行dimH时isLongH_被置 1内核走 Long-H 连续单块拷贝路径由DataCopyPad处理 sub-block 对齐否则走 Short-H 行 stride 拷贝路径。两种路径对应 situ_glu.hpp 中的CopyInShortH按行跳步拷贝 gate/up 两半与CopyInLongH连续拷贝后再切半CopyOut同样区分两种模式。确定性计算该接口默认支持确定性计算。这意味着在相同输入与参数下多次调用得到的计算结果保持一致适合对可复现性有要求的训练与推理流水。调用示例单算子模式调用单算子模式通过cann_ops_nn.situ_glu直接调用前提是环境已安装并初始化torch_npuimport torch import torch_npu import cann_ops_nn x torch.randn(2, 32, dtypetorch.float32).npu() # 默认参数调用 y cann_ops_nn.situ_glu(x) print(y.shape, y.dtype) # torch.Size([2, 16]) torch.float32 # 指定全部参数 y cann_ops_nn.situ_glu(x, dim-1, beta1.0, linear_beta0.0, activate_leftTrue) print(y.shape, y.dtype)示例中x的 shape 为(2, 32)默认dim-1表示在最后一维大小为 32上切分因此输出 shape 为(2, 16)。若把dim指定为0则按第一维切分输出 shape 将变为(1, 32)。注意dim维度大小必须为偶数否则调用会直接报错。TorchAir 图模式调用图模式通过torch.compile配合torchair的 NPU 后端完成构图与编译import torch import torch_npu import torchair import cann_ops_nn class Model(torch.nn.Module): def forward(self, x): return cann_ops_nn.situ_glu(x, dim-1, beta1.0) model torch.compile(Model().npu(), backendtorchair.get_npu_backend(), dynamicFalse) x torch.randn(2, 32, dtypetorch.float32).npu() y model(x)图模式场景下算子会经由 graph_convert_situ_glu.py 完成图中算子转换并以 situ_glu_proto.h 定义的算子 IR 形式参与构图。底层实现从 PyTorch 调用到 AICore 内核调用链路situ_glu的完整调用链可归纳为Python 层cann_ops_nn.situ_glu(...)→OpBuilder动态加载 C 扩展situ_glu.py形状推导Meta 实现_situ_glu_shape计算输出 shapedim校验与偶数校验在此完成situ_glu.pyHost 层算子定义situ_glu_def.cpp声明输入/输出/属性并注册平台infershapesitu_glu_infershape.cpp推导输出形状与数据类型Tiling 层situ_glu_tiling.cpp 根据平台 AIV 核数、UB 内存大小与输入 shape 计算切分参数coreNumAll、dimBatchSize、dim2H、isLongH、ubMaxPair等并申请 16MB workspaceKernel 层AICore 内核 situ_glu.cpp 按ORIG_DTYPE_X编译为 float / half / bfloat16 三种实例执行 situ_glu.hpp 中的分块计算。内核计算细节合轴与分核tiling 将x合轴为[dimBatchSize, dim2H]dimBatchSize为dim之前各维乘积dim2H为dim及之后各维乘积再由CalTilingParamsitu_glu.hpp按 batch 在多个 AIV 核之间切分核内再按 UB 容量分 loop 处理。gate/up 切分CopyIn一次搬运 gate 与 up 两段到 xQueue前半与后半各占一块 bufferCompute中根据activateLeft_选择 gate 与 up 的偏移situ_glu.hpp。混合精度路径fp16 / bf16 输入先Cast到 float32 计算无损加宽算完再Cast回原类型其中 bf16 回写必须使用CAST_RINT舍入模式源码注释明确指出CAST_NONE在非 A2 平台如 Ascend950 上会产生错误结果见 situ_glu.hpp。FP32 原生路径float32 输入直接在原 buffer 上以 float 视图计算无需 Cast减少一次数据搬运。参考资源接口文档torchapi_situ_glu.md、底层 aclnn 接口文档 aclnnSituGlu.md算子 README含参数表与调用方式汇总README.md算子定义与推导situ_glu_def.cpp、situ_glu_infershape.cppTiling 实现situ_glu_tiling.cpp、situ_glu_tiling.hAICore 内核situ_glu.cpp、situ_glu.hppPyTorch 扩展situ_glu.py、csrc/situ_glu.cpp、graph_convert_situ_glu.py调用示例test_aclnn_situ_glu.cppaclnn API 方式测试用例test_situ_glu_network.py网络级 pytest、executor_aclnnSituGlu.pyST 执行器反向算子训练场景配套situ_glu_grad 目录及其接口文档 torchapi_situ_glu_grad.md【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考