ops-rec:基于昇腾平台的推荐系统高性能融合算子仓库

基于昇腾平台,面向推荐系统的高性能融合算子仓库

分支1Tags0
当前项目代码仓暂无内容

ops-rec

License

ℹ️ 简介

ops‑rec 是面向大规模推荐场景的高性能算子库,基于昇腾 CANN 平台深度定制,采用 AscendC 原生静态算子 + Triton JIT 动态可编程算子的双后端架构。算子库聚焦推荐模型核心算力瓶颈,覆盖特征预处理、多特征融合、矩阵计算、高阶非线性计算等高频场景,为推荐业务提供全链路硬件级加速支撑能力。

  • 双引擎异构算子架构:统一上层调用接口,同时支持 AscendC 极致静态高性能算子与 Triton JIT 动态灵活算子开发部署,兼顾极致性能与业务快速迭代能力,适配不同开发与上线诉求。
  • 推荐专属深度优化:针对推荐场景算子融合、HSTU 高阶复杂计算等核心性能瓶颈,完成昇腾硬件深度适配与定制重构。通过计算粒度拆解、访存合并、流水级并行优化,突破原生通用算子的性能上限,显著降低推荐模型延时、提升整机吞吐。
  • 极简轻量接入:屏蔽昇腾底层硬件、编译及平台适配细节,提供统一易用的 Python 接口,开箱即用,业务可低成本完成迁移,快速落地硬件加速收益。

⚙️ 功能介绍

功能 描述 文档
AscendC 静态高性能算子 面向推荐核心场景的高性能 AscendC 算子及其 aclnn 接口,覆盖特征预处理、多特征融合、矩阵计算、高阶非线性计算等 构建指南
Triton JIT 动态算子 与 ascendc/ 平级的 Triton 算子源码,支持动态可编程算子开发部署,兼顾快速迭代 算子贡献指南
统一 Python 调用接口 屏蔽 SOC 差异与编译细节,import ops_rec 自动按本机 SOC 加载对应算子实现 运行时 SOC 检测
多 SOC 构建与打包 一套源码编译全部 SOC 变体并合并为多 SOC whl,安装后自动选择当前芯片对应实现 编译产物结构

🚀 快速入门

环境要求

组件 要求
操作系统 Linux
Python >= 3.10
CANN >= 9.0.0,需执行 set_env.sh 配置环境变量
PyTorch + torch_npu 2.7.1 / 2.10.0
triton + triton-ascend 与 torch_npu 版本配套:torch_npu 2.7.1 建议配套 3.2.2,torch_npu 2.10.0 建议配套 3.6.0
构建依赖 gcc/g++(支持 C++17)、cmake、python3(含 torch/torch_npu);打包依赖 packaging

部署前请确保驱动和固件安装正确,且 npu-smi info 能识别芯片及设备状态。CANN 和 TorchNPU 的安装方式请参考:

安装 ops_rec

算子依赖 CATLASS 源码,编译前需要初始化 submodule:

git submodule update --init --recursive
# 默认编译三个 SOC 变体并安装(推荐):atlas_950、atlas_a2、atlas_a3
bash build.sh --install -j16

# 仅编译指定 SOC
bash build.sh --socs atlas_950 --install -j16

# 编译多个指定 SOC 并安装
bash build.sh --socs atlas_a2,atlas_a3 --install -j16

生成的 whl 位于 dist/,可直接 pip install dist/ops_rec-*.whl。多 SOC 产物合并打包说明见 构建指南

运行示例

# 运行所有 torch 示例
bash run_examples.sh --type torch

# 运行指定算子的 aclnn + torch 示例
bash run_examples.sh --ops ln_mul

run_examples.sh 支持 --type aclnn|torch|all--ops <list>,要求 ops_rec 已安装(bash build.sh --install)。

📖 使用指南

构建变体与硬件对应关系

build.sh 通过 --socs <list> 指定构建变体,英文逗号分隔。默认编译 atlas_950atlas_a2atlas_a3 三个变体(等价于 --socs atlas_950,atlas_a2,atlas_a3),各变体与芯片型号的对应关系如下:

构建变体 芯片型号 默认编译
atlas_950 Ascend950*
atlas_a2 Ascend910B*
atlas_a3 Ascend910_93*

运行时 import ops_rec 会通过 acl.get_soc_name() 自动检测芯片型号并按上述映射加载对应算子实现;也可通过环境变量 OPS_REC_FORCE_VARIANT 强制指定运行时变体。

完整参数说明见 构建指南

源码结构

ops-rec/                                  # 项目根目录
├── ascendc/                              # AscendC 算子源码
├── triton/                               # Triton 算子源码(与 ascendc/ 平级,按后端组织)
├── cmake/                                # CMake 模块(soc_config、toolchain 等)
├── docs/                                 # 开发指南和参考文档
├── examples/                             # 算子样例
├── experimental/                         # 实验性非正式算子
├── ops_rec/                              # pip 包主体
├── pre-commit/                           # pre-commit 配置
├── scripts/                              # 构建辅助脚本(merge_wheel.py 等)
├── tests/                                # UT 测试架构文件
├── third_party/                          # 三方软件
├── build.sh                              # 统一构建入口
├── run_examples.sh                       # 示例运行入口
└── CMakeLists.txt                        # 顶层 CMake 入口

本节展示的是源码仓库根目录结构。whl 安装后 Python 包 ops_rec/ 的内部布局见 构建指南

🛠️ 算子开发

新增算子时,按 <group>/<vendor> 目录约定放置代码即可。完整迁移流程、SOC 注册、torch_extension 写法、PR 模板等见 算子贡献指南

🛠️ 贡献指南

欢迎参与项目贡献,贡献流程和规范请参见 算子贡献指南

⚖️ 相关说明

🔹 构建指南
🔹 算子列表
🔹 算子贡献指南
🔹 Apache License 2.0 许可证

🤝 建议与交流

欢迎大家通过以下方式提出问题、交流讨论。

资源 说明
创建Issue 提交 Bug、需求或建议

🙏 致谢

感谢所有为 ops-rec 和昇腾开源社区贡献代码、文档、测试、问题反馈与技术建议的开发者。欢迎更多社区伙伴参与共建!

项目介绍

基于昇腾平台,面向推荐系统的高性能融合算子仓库

定制我的领域