状态（Status）： Draft
作者（Authors）： @QianZH97
创建日期（Created）： 2026-09-16
更新日期（Updated）： 2026-09-17
相关 Issue/PR： #123、#128

1. 概述

1.1 简介

本提案旨在将 GPU 三方仓（recsys-examples、FBGEMM）中沉淀的典型推荐领域 Triton 算子迁移至 ops-rec 的 experimental/triton/atk_test 目录，为每个算子补充完善的 NPU / GPU 双端 ATK 精度与性能测试用例，并在 NPU 侧通过 autotune 等策略完成性能优化与基线固化。本提案同时定义了「源算子 → 双端 ATK 用例 → 性能优化 → 基线固化」的可持续迁移规范，后续新增三方 Triton 算子按此模板持续迁移。

1.2 动机

使用场景 / 用例

推荐系统（含生成式推荐）模型训练与推理链路大量依赖社区 Triton 算子，典型算子包括 jagged/dense 不规则张量转换、变长 Flash Attention 及其反向（dk / dv / dbias）等。这些算子上游主要源自 NVIDIA 推荐生态的两大三方仓：ecsys-examples 和 FBGEMM。将这类算子迁移至 ops-rec，可补齐昇腾平台推荐场景的 Triton 算子能力，并复用到 RecSDK 等的推荐模型链路中。

当前痛点

问题	说明
目录规范缺失	atk_test 下各算子目录组织形式、NPU/GPU 两侧文件归属缺乏统一定义
迁移流程不完整	缺少「GPU 原版算子 → NPU 适配（含 autotune 等优化）→ ATK 精度/性能用例」的完整链路定义，人工迁移靠口头约定，无法规模化复制
来源无追踪	算子来源于 recsys-examples / FBGEMM，缺少统一的来源（repo + commit + 文件行号）登记，升级追源与合规审计困难

必要性与用户价值

若不立项，三方 Triton 算子迁移将停留在「每次迁移凭经验手工组织、无规范可依、无基线可查」的状态：新算子落地成本高、命名与目录组织杂乱、性能优化收益不可量化。完成后可形成一套可复制的迁移模板，显著降低后续算子迁移成本，并为每个算子沉淀永久可追踪的双端精度基准与性能基线。

不做此提案的影响

昇腾推荐场景的 Triton 算子生态适配持续依靠零散人工迁移，无法规模化吸收社区算子红利；算子性能优化（autotune 等）缺少统一的迁移规范与可量化的性能基线。

1.3 目标

目标：

• 建立统一的三方 Triton 算子 ATK 测试迁移目录规范与命名约定，后续迁移算子全部按模板落盘。
• 每个算子具备：NPU 端性能优化后的 kernel（autotune 等策略 + 精度用例 + 性能基线）；GPU 端原版 kernel（精度对比基准）。
• 性能优化收益可量化（autotune 等策略的收益对比 GPU 基线 / 非优化基线），性能基线随算子一并固化入库。

非目标（边界）：

• 本次不承接算子性能目标的绝对承诺（如「必须快于 GPU 多少倍」），仅要求性能基线固化与优化收益可量化。
• 本次不实现全自动化迁移 Agent，迁移仍为半人工流程（规范 + 模板 + 清单），自动化工具可后续独立 RFC。

约束说明：

• ATK 测试依赖 NPU 端 + GPU 端双机环境（GPU 端作为精度基准 server，如果和 GPU 对比精度不满足要求，则转换为双精度即以 CPU 为准）。

2. 用例分析

用例	功能点	关键性能指标	安全隐私	DFX 要求	使用限制/约束
精度对比用例（ATK accuracy）	基于 generate_<op_name>.py + <op_name>.yaml 生成全组合精度用例，NPU 端 kernel 与 GPU 端原版 kernel 逐用例对比	无独立性能要求；为性能测试提供正确性前提	用例纯用随机/确定性生成的张量数据，不涉敏感数据，无需额外安全管控	可测试性：用例可重复生成、口径统一，由 atk case / atk node 自动化执行；可靠性：NPU/GPU 结果一致	精度阈值按各算子的 dtype/shape 场景约定；若与 GPU 对比精度不满足要求，则转换为双精度、以 CPU 为准
性能对比用例（ATK performance_device）	复用 <op_name>_performance.json 固化 shape，NPU 端 autotune 优化 kernel 与 GPU 端对照，读取性能数据	每次运行产出可重复的设备侧性能数据（ms/us），与 GPU 基线对比量化收益	同上	可维护性：性能测试用例固化随算子入库；可测试性：性能用例可由 atk node 重复执行对比收益	性能数据受设备负载影响
端到端算子调用（ATK node/server）	NPU 客户端 + GPU server 双端拉起精度/性能任务	双端链路稳定可复用	无敏感数据	兼容性：GPU/NPU kernel 接口签名一致（差异仅 autotune 等策略）	需 GPU 端 atk server 与 NPU 端 atk node 双机可用

3. 方案设计

3.1 总体方案

3.1.1 整体设计思路

以「原版 kernel 保留 + NPU 优化 kernel 独立文件 + 双端 ATK 用例」为核心思路：GPU 侧保留原版 Triton kernel 作为精度/性能的对照基准（不优化、不做 autotune），NPU 侧提供适配昇腾硬件并经 autotune 等策略优化后的 kernel，两者通过一致的 ATK API 封装接入 ATK 框架，实现跨端精度对比与性能基准测试。

3.1.2 迁移目录规范

在 experimental/triton/atk_test 下，以 Triton 算子名命名文件夹，文件夹内包含 npu_atk_test 与 gpu_atk_test 两个子目录：

experimental/triton/atk_test/
└── <op_name>/                        # 以 Triton 算子名命名
    ├── README.md                     # 算子来源、功能、约束、目录说明、ATK 用例使用指南
    ├── npu_atk_test/
    │   ├── <op_name>.py              # NPU 适配优化后的 Triton 算子 kernel（autotune 等性能优化策略后）
    │   ├── <op_name>_api.py          # NPU 端 ATK 自定义 API 封装（register + BaseApi）
    │   ├── <op_name>.yaml            # NPU 端 ATK 精度测试用例定义文件
    │   ├── generate_<op_name>.py     # 精度测试用例数据生成脚本
    │   └── <op_name>_performance.json # 固化的性能测试用例（性能基线数据）
    └── gpu_atk_test/
        ├── <op_name>.py              # 原版 Triton 算子 kernel（迁移来源，GPU 侧原样部署）
        └── <op_name>_api.py          # GPU 端 ATK 自定义 API 封装

各产物职责说明：

产物	归属	说明
<op_name>.py (kernel)	NPU / GPU	GPU 侧为原版 kernel（对齐迁移来源，不改性能参数）；NPU 侧为适配优化后 kernel（以 @triton.autotune 等包装，num_warps / BLOCK_* 等参数自动选优）
<op_name>_api.py	NPU / GPU	ATK 自定义 API 执行文件：@register 注册算子、继承 BaseApi，实现 run/compare 逻辑，调用同目录 kernel。GPU 侧与 NPU 侧均必须放置
<op_name>.yaml	NPU	ATK 精度测试用例定义文件（name / api_type / triton_api_type / generate / standard 等字段，可声明 triton_name 指定调用原版 kernel 名）
generate_<op_name>.py	NPU	精度测试用例数据生成脚本（采用确定性值生成策略，限制用例入参shape）
<op_name>_performance.json	NPU	固化的性能测试用例（performance_device 任务的输入
README.md	算子目录	记录来源（repo + commit + 文件/行号）、算子功能、输入约束、目录文件说明、ATK 端到端用例流程

命名统一约定：kernel / api / yaml / generate 脚本 / performance json 均与算子名 <op_name> 同名；yaml 中 name / api_type / triton_api_type 与算子名对齐，必要时通过 triton_name 显式关联原版 kernel。

约束说明：GPU 侧不放置 yaml / generate 脚本 / performance json（GPU 不承担数据生成与性能固化职责）；NPU 侧为精度与性能测试的唯一数据生成方与基线持有方。

3.1.3 源算子与产物放置

• gpu_atk_test/ 放置原版 Triton 算子文件以及 ATK 调用测试的 api 文件。
• npu_atk_test/ 放置：
  1. 精度测试用例生成相关文件（<op_name>.yaml + generate_<op_name>.py）；
  2. 固化的性能测试用例（<op_name>_performance.json，性能基线数据随算子一并入库）；
  3. 性能优化（autotune 等策略）后的 Triton 算子文件（<op_name>.py，BLOCK_* / num_warps 等策略自动选优）；
  4. ATK 调用测试的 api 文件（<op_name>_api.py），同时作为 NPU 侧精度/性能任务的执行入口。

3.1.4 里程碑

阶段	内容	交付物
M1：算子迁移 + autotune 调优	按「来源确认 → GPU 原版 kernel 落盘 gpu_atk_test → NPU 适配优化 kernel 落盘 npu_atk_test → 精度/性能用例生成与固化 → README 记录来源」逐算子迁移，NPU 端 kernel 以 autotune 完成基础调优	每算子双端 ATK 用例 + 含 autotune 基础调优的 NPU kernel + README
M2：NPU 算子深度性能调优	对 NPU 端算子开展进一步性能调优（如访存优化、计算粒度拆解、配置策略细化等），联动更新性能基线	性能优化后的 NPU kernel + 更新的性能基线

3.2 技术选型

备选方案	优势	劣势	不选择理由
方案 1：双端闭环（选定）——NPU 侧保留优化 kernel + GPU 侧保留原版 kernel，双端 ATK 对比	1. GPU 原版 kernel 是稳定的精度/性能对照基准，可比性强 2. 目录清晰、天然支持跨端对比 3. 单算子模板可复制	需要双机（GPU + NPU）ATK 环境	收益显著大于双机成本，且推荐生态双机环境已是常态
方案 2：仅 NPU 单端 + PyTorch 参考实现对比	单机即可，无 GPU 依赖	缺少跨端可比性；参考实现往往非原版 kernel，无法体现原版优化水平；无法复用 GPU 基线性能数据	无法满足「与 GPU 社区算子对标」的目标，对比口径不统一
方案 3：全自动化迁移 Agent	人力成本最低	需额外开发 Agent + 精度自动回溯链路，当前阶段 arbitrary_func 类自动化流程尚未沉淀为通用能力	本提案聚焦规范 + 用例 + 优化基线，自动化可作为后续独立 RFC

3.3 功能与性能设计

3.3.1 NPU 端 kernel 性能优化

NPU 侧以 @triton.autotune 包装 kernel，将 BLOCK_T / BLOCK_L / BLOCK_D / num_warps 等策略参数纳入自动选优，grid 从 META 中读取 BLOCK_T 等动态计算：

# NPU 端示例（autotune 策略）：
# 以变长 Flash Attention 反向 dk kernel 为例
@triton.autotune(configs=_JAGGED_DENSE_FLASH_ATTENTION_BWD_DK_AUTOTUNE_CONFIGS, key=["..."] )
def _jagged_dense_flash_attention_bwd_dk_kernel(...):
    ...

优化策略要点：

• 适用产品：autotune 优化主要针对 Ascend 950PR & 950DT 系列产品，不对其它 SOC 变体（如 A2 / A3）做差异化 autotune 配置。
• configs 覆盖：类似 BLOCK_T ∈ {16, 32, 64, 128}、BLOCK_L ∈ {16, 32}、num_warps ∈ {1, 2, 4, 8} 等组合，覆盖典型 shape。
• 可复现性：autotune 结果后固定最佳 config，性能基线数据基于固定 config 产出，保证基线可复现、可比。
• GPU 侧不对齐：GPU 侧保留原版 kernel（显式传参或无新增 autotune），保证对照为社区原版实现。

3.3.2 数据模型（ATK 用例文件）

文件	关键字段 / 内容
<op_name>.yaml	name、api（pytorch）、version、api_type、triton_api_type、triton_name、generate、dtype_numbers、outputs、standard（如 acc: single_bm）等
generate_<op_name>.py	确定性值生成策略（不依赖随机数）：如 jagged 算子的 seq_offsets 单调递增整数、ind_offsets 可选非负整数等
<op_name>_performance.json	性能测试用例 shape 列表

3.3.3 影响范围

仅影响 ops-rec 仓库 experimental/triton/atk_test/ 下的目录结构与文件，不涉及 AscendC 静态算子、Python 包 ops_rec/ 打包逻辑、构建系统 build.sh / CMake。

3.4 安全隐私与 DFX 设计

3.4.1 安全隐私设计

本提案所有 ATK 用例数据均为本地生成的随机 / 确定性张量，不涉及敏感业务数据、无越权访问与数据泄露风险，仅需遵循 ops-rec 现有安全规范（含 Copyright (c) Huawei Technologies Co., Ltd. 许可头、Apache-2.0 License）。

3.4.2 DFX 设计

兼容性：

• GPU / NPU kernel 保持 API 签名一致，性能差异仅限 autotune 等策略参数；接口 1:1 兼容迁移来源（FBGEMM / recsys-examples）。

可维护性：

• 目录与命名统一，新增算子按模板落盘；基线 json 随算子入库，可随时更新。
• 遵循 ops-rec 编码规范，核心 kernel 逻辑附必要说明（README + API 注释）。

可测试性：

• NPU 端精度 / 性能用例均为自动化 ATK 任务（atk case / atk node 执行），GPU 端原版 kernel 作为对比基准；用例生成脚本确定性可重放。

可靠性：

• NPU 端与 GPU 端结果在精度阈值内一致；性能基线随算子固化，供 M2 深度调优阶段对比优化收益。

3.5 编程与调用设计

3.5.1 编程模型基本设计

开发环境设计：

类别	具体内容
硬件环境	NPU 端：Ascend 950PR & 950DT 系列产品（以算子实际支持为准）；GPU 端：任意 NVIDIA 卡（作为 ATK 精度/性能基准 server）
软件环境	NPU端: CANN ≥ 9.1（配置 set_env.sh 环境变量）、PyTorch + torch_npu、triton + triton_ascend、ATK 工具; GPU 端: triton、ATK 工具
开发 / 调试工具链	atk case / atk server / atk node 命令

开发约束：

• 算子 kernel 采用 Triton DSL 编写，遵循原版 kernel 的输入输出约束（如 jagged 算子的 seq_offsets 单调递增、首 0 尾 sum_len）。
• 迁移算子需登记来源（repo + commit + 文件/行号），携带 Apache-2.0 许可头。

可验收设计：

• 每个算子：精度用例全通过（阈值按 dtype/shape 约定）+ 性能基线固化 + README 迁移来源可查。
• 里程碑 M1 / M2 的交付物按「3.1.4 里程碑」验收。

3.5.2 接口定义与设计

3.5.2.1 <op_name>.yaml（ATK 精度用例定义文件）

• 接口描述：声明算子精度测试的用例定义、生成脚本名、标准配置。
• 接口原型：YAML 文件，位于 npu_atk_test/<op_name>.yaml。

输入参数：

参数名称	输入/输出	类型	描述	取值范围
name	输入	string	yaml 用例名，与算子名 <op_name> 一致	算子命名规范
api	输入	string	API 框架类型	pytorch
api_type	输入	string	用例 API 类型，与算子名对齐	算子命名规范
triton_api_type	输入	string	Triton API 类型	算子命名规范
triton_name	输入（可选）	string	显式关联的原版 kernel 名	Triton kernel 命名
generate	输入	string	数据生成脚本名（不带后缀），如 generate_<op_name>	与算子同名
dtype_numbers	输入	int	用例数量	≥ 1
outputs	输入	int	输出张量个数	≥ 0
standard	输入	结构体	标准配置，如 acc: single_bm	见 ATK 框架规范

• 返回参数：无直接返回；运行 atk case -f <op_name>.yaml -p generate_<op_name>.py 后于 result/<op_name>/json/ 生成 all_<op_name>.json。
• 异常处理：字段校验失败时 ATK 框架报错，需保证 generate 指向的生成脚本存在。
• 约束说明：NPU 侧放置；GPU 侧不放置。
• 变更说明：新增算子时随算子新增。
• 调用参考代码：

name: <op_name>
api: pytorch
version: v2.1
api_type: <op_name>
triton_api_type: <op_name>
triton_name: <triton原版kernel名>
generate: generate_<op_name>
dtype_numbers: 50
outputs: 0
standard:
  acc: single_bm

3.5.2.2 <op_name>_api.py（ATK 自定义 API 封装）

• 接口描述：ATK 自定义 API 执行文件，注册算子并封装 run/compare 逻辑，供精度/性能任务调用。
• 接口原型：Python 模块，位于 npu_atk_test/ 与 gpu_atk_test/，分别调用对应目录的 kernel。

输入参数：

参数名称	输入/输出	类型	描述	取值范围
device	输入	torch.device	运行设备	npu:0 / cuda:0
input tensor	输入	torch.Tensor	算子输入张量（由 ATK 框架注入）	按算子约束

• 返回参数：ATK run 结果 dict / 输出张量（与 standard 配置关联的比对结果）。
• 异常处理：设备不可用时回退（如 torch.npu.is_available() 判断）。
• 约束说明：NPU / GPU 双端必须同时提供；导入同目录 kernel。
• 变更说明：新增算子时双端各一份。
• 调用参考代码：

import torch
from atk.configs.dataset_config import InputDataset
from atk.tasks.api_execute import register
from atk.tasks.api_execute.base_api import BaseApi
from <op_name> import <op_name>_kernel  # 导入同目录 kernel

@register
class <OpName>Api(BaseApi):
    def __init__(self, dataset: InputDataset):
        super().__init__(dataset)
        # 解析 yaml / dataset 提供输入
    def run(self):
        # 调用 <op_name>_kernel 完成计算，返回输出
        ...

3.5.2.3 端到端 ATK 流程命令

生成精度测试用例

atk case -f <op_name>.yaml -p generate_<op_name>.py

生成的精度用例位于 result/<op_name>/json/，文件名为 all_<op_name>.json。

远端 GPU 服务器起监听命令

atk server --devices 0 --plugin_path /path/to/api文件所在文件夹

精度测试

atk node --backend npu --devices 0 node --backend gpu -h {gpu端ip} -p {gpu端监听端口} --devices 0 \
    task -c all_<op_name>.json --task accuracy -p <op_name>_api.py

注：如需保存性能对比数据则需在指令后加 --save_data profile。

性能测试

atk node --backend npu --devices 0 node --backend gpu -h {gpu端ip} -p {gpu端监听端口} --devices 0 \
    task -c <op_name>_performance.json --task performance_device -p <op_name>_api.py

3.5.3 编程手册设计

• 章节规划：迁移目录规范与命名约定、单算子迁移步骤（来源确认 → 双端 kernel → 精度/性能用例 → 基线固化 → README）、ATK 命令说明、常见问题（FAQ）。
• 发布形式：随 ops-rec 仓库维护，先沉淀在 experimental/triton/atk_test/ 的算子 README。

4. 测试设计

• 单元测试：单算子 kernel 的功能正确性（NPU 端 kernel 与参考实现对比）。
• 集成测试（ATK 精度对比）：atk case 生成全组合精度用例 + atk node 拉取 accuracy 任务，NPU 端与 GPU 端原版 kernel 对比，阈值按 dtype/shape 约定；若与 GPU 对比精度不满足要求，则转换为双精度、以 CPU 为准。
• 端到端测试（ATK 性能对比）：performance_device 任务基于 <op_name>_performance.json 读取性能数据，量化 M1 / M2 性能优化策略的收益。M2 阶段深度调优后同步更新性能基线。
• 验收用例：覆盖各算子典型推荐 shape（多 shape + 边界 shape），dtype 覆盖 fp16/fp32/bf16（以算子支持为准）。

5. 缺点和风险

风险类型	具体描述	应对措施
性能回退风险	autotune / 深度调优在部分 shape 上选优不稳定，导致与 GPU 基线相比性能不达标	固定最佳 config 入库（性能基线）；必要时为边缘 shape 提供策略开关
复杂度提升风险	双端目录 + 双机环境增加维护成本	严格按模板落盘，README 收敛单算子说明，不引入额外抽象
兼容性风险	GPU 侧保留的原版 kernel 与 NPU 侧优化 kernel 接口签名需保持一致	迁移时强制 API 签名一致（差异仅 autotune 策略参数），ATK 精度用例兜底
安全风险	无	无
命名兼容性影响	新增算子采用 <op_name>_performance.json、与算子同名 yaml/api 等新命名，需保证与既有引用/脚本兼容	迁移算子 README 记录文件用途与生成命令，README 模板随每个迁移算子落地同步维护
实现成本	算子迁移（M1）+ 深度性能调优（M2），涉及人力投入	分阶段实现，M1 先行落地迁移模板；M2 按算子逐个深度调优
版本兼容性	迁移基于的 FBGEMM / recsys-examples commit 升级后 kernel 可能变化	README 记录 commit，升级时 diff 追踪

6. 现有技术

参考项目 / 社区设计：

• FBGEMM（pytorch/FBGEMM）：fbgemm_gpu 内提供大量 jagged/dense Triton kernel（如 triton_jagged_tensor_ops.py、triton_jagged_dense_flash_attention.py），是 jagged 类算子迁移的主要上游，本提案保留其原版 kernel 作为 GPU 侧基准。
• recsys-examples：推荐模型参考实现仓，提供 gate 等推荐场景 Triton 算子。

借鉴与差异：

参考项目	借鉴点	差异点
FBGEMM / recsys-examples	1. 算子来源与 kernel 签名 2. 推荐场景的典型 shape/用例	1. 运行平台为昇腾 NPU（Triton JIT） 2. 增加 autotune 等昇腾性能优化 3. 双端（GPU↔NPU）ATK 精度/性能对比体系
CANN ops-transformer 测试体系	算子与用例同仓、修改即覆盖	基于 ATK 框架双端对比，而非纯 NPU 单端 UT/ST

7. 未解决问题

autotune 配置：autotune 优化主要针对 Ascend 950PR & 950DT 系列产品进行 config 覆盖与选优，不对其它 SOC 变体（如 A2 / A3）做差异化 autotune 配置（见 3.3.1）。

附录

1. FBGEMM 仓库：https://github.com/pytorch/FBGEMM
2. recsys-examples 仓库：https://github.com/nvidia/recsys-examples
3. ops-rec 仓库：https://gitcode.com/Ascend/ops-rec
4. ATK 测试框架参考：ops-rec experimental/triton/atk_test 下现有算子 README

术语表：

术语	定义
ATK	ATK 是一款面向算子接口的端到端测试工具，支持测试用例生成、算子任务执行、精度与性能结果导出等能力。
Jagged Tensor	不规则（变长）张量，以 lengths / offsets 描述变长序列，推荐领域 embedding/序列特征常用
FBGEMM	Facebook GEneral Matrix Multiplication 库，提供 CPU/GPU 优化的稀疏算子与 Triton kernel，本提案主要上游
Triton DSL	Triton 编程语言，GPU/NPU 上编写高性能 kernel 的类 CUDA DSL
autotune	Triton 自动调优机制，基于 config 列表自动选择最优 block/num_warps 等参数

欢迎加入社区，感谢您对社区的贡献 🎉!