cann_ops_nn.swiglu_group【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn产品支持情况Ascend 950PR/Ascend 950DT 支持Atlas A3 训练系列产品/Atlas A3 推理系列产品 不支持Atlas A2 训练系列产品/Atlas A2 推理系列产品 不支持Atlas 200I/500 A2 推理产品 不支持Atlas 推理系列产品 不支持Atlas 训练系列产品 不支持功能说明接口功能对输入x执行 SwiGLU 激活支持可选的逐 token 权重和 count 模式分组。底层封装aclnnSwigluGroup。计算公式令输入x最后一维大小为DA x[..., :D/2]B x[..., D/2:]。$$ y silu(A) \times B \frac{A}{1 e^{-A}} \times B $$当clamp_limit 0时先对A和B做截断当weight非空时对结果逐 token 乘以weight。函数原型torch.ops.cann_ops_nn.swiglu_group( x, *, weightNone, group_indexNone, clamp_limit-1.0, ) - Tensor参数说明参数名参数类型可选/必选描述数据类型维度(shape)xTensor必选SwiGLU 输入最后一维会被均分为两部分。torch.float16、torch.bfloat16、torch.float321-8维最后一维为偶数weightTensor可选逐 token 权重非空时乘到 SwiGLU 结果上。torch.float32元素个数等于x除最后一维外的元素个数group_indexTensor可选count 模式分组 token 数。torch.int641维clamp_limitfloat可选激活前截断阈值默认-1.0表示不启用截断。--返回值说明参数名参数类型描述数据类型维度(shape)yTensorSwiGLU 激活结果。与x相同与x相同但最后一维为x.shape[-1] // 2约束说明该接口支持单算子模式和 TorchAir 图模式调用。x、weight、group_index均需为 NPU Tensor可选 Tensor 可以传None。x的最后一维必须为偶数。不支持空 Tensor 和非连续 Tensor。group_index表示每个 group 的 token 数所有元素需大于等于 0。确定性计算默认支持确定性计算。调用说明单算子模式调用import torch import torch_npu import cann_ops_nn.ops x torch.randn(8, 512, dtypetorch.float16).npu() weight torch.ones(8, dtypetorch.float32).npu() group_index torch.tensor([4, 4], dtypetorch.int64).npu() y torch.ops.cann_ops_nn.swiglu_group( x, weightweight, group_indexgroup_index, clamp_limit-1.0) print(y.shape, y.dtype)图模式torchair调用import torch import torch_npu import torchair import cann_ops_nn.ops class Model(torch.nn.Module): def forward(self, x): return torch.ops.cann_ops_nn.swiglu_group(x, clamp_limit-1.0) model torch.compile(Model().npu(), backendtorchair.get_npu_backend(), dynamicFalse) x torch.randn(8, 512, dtypetorch.float16).npu() y model(x)【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考