CANN Qwen3-Next NPU算子扩展
2026/5/9 15:54:36 网站建设 项目流程

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

本模块主要包含以下关键组件:

  1. gated_delta_net/<op_dir>/算子实现目录,主要包含:

    • <op_name>.cpp:算子调用文件。
    • op_kernel/:算子 Kernel 具体实现代码。
  2. gated_delta_net/<op_dir>/CMakeLists.txt算子编译配置文件。

  3. 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

  1. 进入算子目录,安装依赖:
pip install -r requirements.txt
  1. 从源码构建.whl包:
python -m build --wheel -n
  1. 安装构建好的.whl包:
pip install dist/*.whl --force-reinstall --no-deps
  1. (可选)如果需要重新编译,建议先清理编译缓存:
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为算子注册的 namespace
  • my_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),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询