股票 
板块 
用户 
资讯 算子开发一直是大模型Infra的基本需求，然而一个高性能算子的开发对开发人员的要求很高，需要对硬件架构十分熟悉，且开发周期长。随着AI Coding能力变强，AI算子生成成为可能。本文基于昇腾端到端训练了一个30B的模型进行算子生成任务，以算子的torch实现作为输入，生成NPU-Triton算子代码。在生产级的算子验收标准下，模型生成的Triton算子通过率在Kernelbench level 2达到29%，是Drkernel开源模型（通过率9%）的3倍以上。相关代码已在verl-ascend-recipe仓库提交PR，后文附上链接。
生产级精度校验标准及
Triton-only Hacking检测
自动化的数据合成以及后期RL训练需要online的算子验证，NPU的沙箱开发是基础的工作。我们的Sandbox基于开源的KernelGym(https://github.com/hkust-nlp/KernelGYM)，在昇腾上打通了算子编译、正确性及性能提升的自动化校验。
NPU-Triton的自动化校验
Sandbox的校验逻辑是根据输入的算子torch代码以及Triton代码，开启一个进程然后独占一张卡并执行两套算子代码。根据两者的输出比较精度，同时根据在NPU上实跑得到的时延来计算加速比。昇腾的环境需要在启动时设置CANN/ATB环境变量路径，我们将路径注入到子进程的运行环境中避免环境带来的错误反馈。
精度校验标准及Hacking检测
与KernelGym不同，本次的算子生成任务对齐生产级标准。我们发现开源的校验标准较为宽松，相对误差和绝对误差的阈值是1.00E-02，且生成的算子代码允许部分调用torch接口。而我们则要求模型生成Triton-only的kernel代码，所有调用了计算类torch接口的实现会被判定为作弊。正确性的校验标准我们参考了https://github.com/Just-it/AscendOpGenAgent/tree/br_430/skills/triton/kernel-verifier/scripts，会根据不同Dtype制定不同的标准，譬如bf16类型的数据阈值是2^{-7}，要严于1.00E-02。
性能数据采集
我们注意到KernelGym在计算加速比时采用的是打点计时得到的时间，这并不准确。我们在其中添加了profiling采集kernel侧耗时的方式来计算加速比。
NPU-Triton训练数据配方
由于开源的权重在昇腾上无法生成高质量的算子，如果我们直接进行RL训练，则会由于一直采样不到正样本而导致实验失败。我们需要先进行SFT让模型具备一些基础的昇腾算子生成能力。因此我们搭建了数据合成的Pipeline合成高质量多轮轨迹数据SFT数据集。
数据合成Pipeline
我们通过 LLM 多轮对话 + 沙箱验证的方式，自动循环五轮将 torch 算子代码改写为 Triton 算子代码，Pipeline如下图所示。
图1.数据合成pipeline
第一轮冷启动，LLM首次接触PyTorch代码，独立完成分析和改写。从第二轮起，LLM基于沙箱的客观验证反馈做有针对性的改进。沙箱反馈包含正确性、性能加速比和作弊检测三个维度的信号，LLM根据反馈自行判断当前问题，选择修复或优化策略。达到预先设定的最大轮数上限时停止，让LLM在通过后仍有机会继续优化性能。我们目前共进行五轮对话。
数据质量评估
在数据蒸馏的外部模型选择上，我们选取了DeepSeek系列模型，开源的Triton数据集则选取了领先的闭源模型进行蒸馏。经过上节的数据合成过程之后，我们把蒸馏得到的Triton算子进行通过率和加速比的统计。
采用v1-standard标准评估的结果在下图右边，可以发现GPT 5的质量明显高于DeepSeek系列。但在我们的流程中，我们按照自己改进的沙箱（即v2-standard）进行了数据合成。由于多轮交互中会有沙箱的反馈结果，高标准的反馈促使了模型往正确的方向进行优化，评估结果如下图左边所示，DeepSeek系列质量高于GPT 5。证明了我们沙箱改进的价值，以及得到了一批NPU上高质量的Triton算子轨迹数据集。
图2.数据质量对比。图中橙色为DeepSeekV4-Pro的结果，绿色为DeepSeekV3.2的结果，蓝色是从开源Triton数据集提取出来的结果。柱状图反映通过率，折线图反映加速比。
提示词激发模型潜力
我们主要选取Qwen3-30B-Thinking为基础模型进行后训练。考虑到训练效率，我们把序列长度控制在64K以内,因此我们主要合成了5轮的数据，所用的 SFT提示词如下图：
图3.数据合成及SFT训练提示词
我们采用异步RL框架进行多轮的RL训练。Rollout的过程进行3轮的采样，每轮生成Triton算子经过沙箱验证之后继续生成。提示词和SFT稍有区别，如下图所示，我们在提示词中显示提到了一些模型常犯的错误，以及提了几点NPU Triton实现的特有约束，避免模型犯低级错误。
图4.RL训练提示词
生产级验收标准下，3倍通过率提升
后训练效果
在评测过程中，我们发现生成算子的通过率和硬件型号、CANN环境版本等相关。评测时我们分别在Atlas 800T A3和Atlas 900 A2 PoD硬件上测试，结果如下，经过SFT与RL之后，训练得到的权重均能在通过率和Fast@1 / 1.2指标上有所提升，例如Pass@8通过率是基线的2.5倍以上，性能指标也能体现模型具有写出优于torch eager算子的能力。
图5.Atlas 900 A2 PoD评测结果
图6.Atlas 800T A3评测结果
NPU-Triton业界对比
我们选取了业界SOTA的Drkernel工作开放的开源8B和14B权重，在业界设备上补测结果。由于开源沙箱在加速比的计算采用了打点计时的方式，这里我们并不比较性能的数据。开源标准的评测结果如下表所示：
KernelBench评测Pass@8通过率对比（开源标准）
模型权重
Level
8B
14B
Ours
SFT
1
36%
55%
65%
2
18%
37%
94%
RL
1
57%
86%
89%
2
73%
95%
96%
此外，我们也做了生产级标准（即v2-standard）的对比，结果如下表所示：
KernelBench评测Pass@8通过率对比（生产级标准）
模型权重
Level
8B
14B
Ours
SFT
1
7%
18%
50%
2
1%
1%
24%
RL
1
8%
40%
60%
2
1%
9%
29%
以上两张表可以发现，SFT与RL后训练的权重均能超越开源的模型，其中KernelBench level 2的通过率超过了14B的3倍以上。
Test-time Scaling
我们将Triton算子生成建模为树状搜索 + 上下文窗口的推理过程。核心思想是从树状搜索、深度搜索和STTS窗口三个维度，通过配置参数node_num、max_rounds、window_size独立控制。实验配置如下：
图7.Test-time scaling实验结果
我们发现实验通过率与加速收益随搜索次数提升而增长；且增加搜索次数能够带来一定收益，但提升幅度会逐渐减小。对于优化难度较高的算子，可以适当增加搜索次数以获得更好的结果。
开源链接
本文实践的训练全流程代码已开源在verl-ascend-recipe仓库。
SFT训练部分代码开源链接：
https://github.com/verl-project/verl-ascend-recipe/pull/48
RL训练代码开源链接：
https://github.com/verl-project/verl-ascend-recipe/pull/49
数据集和权重链接预计7月底开源：
https://modelers.cn/datasets/zhouchuansai/Triton-Training_SFT_Datasets
https://modelers.cn/models/zhouchuansai/Qwen3-30B-Triton-Training
◆往期推荐◆
匠心时刻丨MindStudio Agent：让精度调试告别“专家瓶颈”
评论 (0)