NpuOpsTransformerExt
【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法,提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer
简介 | Overview
该目录包含用于Qwen3-Next 推理的 NPU 融合算子实现,目前包括:
- rmsnormgated融合算子
- Gated Delta Network (GDN)融合算子
这些算子以 PyTorch Extension 的形式实现,并在安装后注册到torch.ops命名空间中供框架调用。
核心组件 | Core Components
本模块主要包含以下关键组件:
gated_delta_net/<op_dir>/算子实现目录,主要包含:<op_name>.cpp:算子调用文件。op_kernel/:算子 Kernel 具体实现代码。
gated_delta_net/<op_dir>/CMakeLists.txt算子编译配置文件。npu_ops_transformer_ext/npu_ops_transformer_ext/npu_ops_def.cpp算子接口注册文件。
环境要求 | Prerequisites
- Python ≥ 3.8
- CANN Ascend Toolkit
- PyTorch ≥ 2.1.0
- torch_npu (PyTorchAdapter)
上述依赖的安装与环境配置请参考 Qwen3-Next README。
安装步骤 | Installation
- 进入算子目录,安装依赖:
pip install -r requirements.txt- 从源码构建
.whl包:
python -m build --wheel -n- 安装构建好的
.whl包:
pip install dist/*.whl --force-reinstall --no-deps- (可选)如果需要重新编译,建议先清理编译缓存:
python setup.py clean算子调用 | Usage
完成编译并安装.whl包后,自定义算子会注册到torch.ops命名空间中,可通过如下方式调用:
import torch import npu_ops_transformer_ext # 调用自定义算子 out = torch.ops.npu_ops_transformer_ext.my_ops(input)其中:
npu_ops_transformer_ext为算子注册的 namespacemy_ops为具体算子名称(在npu_ops_def.cpp中定义)
当前模块包含的算子示例:
torch.ops.npu_ops_transformer_ext.recurrent_gated_delta_rule(...) torch.ops.npu_ops_transformer_ext.mambav2_rmsnormgated(...)具体输入参数格式请参考对应算子的实现代码。
【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法,提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考