TritonBench: Benchmarking Large Language Model Capabilities for Generating Triton Operators
这是什么
Triton 是 OpenAI 推出的 Python-like GPU kernel DSL,现在是 LLM 训练/推理栈里手写 kernel 的主流选择(FlashAttention 变体、各种 fused op 大量用 Triton 写)。但 LLM 生成 Triton 代码的能力此前没有系统评测——已有的 KernelBench 面向 CUDA/torch 编译,而 Triton 有自己的一套 block-level 编程模型和坑。
TritonBench 提供两个评测通道:TritonBench-G 从 95 个 100+ star 的 GitHub 仓库里人工清洗出 184 个真实 Triton 算子(补 wrapper、去重、分 d1-d5 五档难度),考察'照真实需求写算子';TritonBench-T 则是 166 个与 PyTorch 接口对齐的算子(含算子融合场景),考察'把 PyTorch 语义翻译成高效 Triton'。两个通道都带可执行的测试代码(平均每算子 3.6 个测试分支,用随机 tensor 做输入)。
与常规代码基准最大的差别是它不止测对错:对能跑对的算子,还在 A100 上测相对参考实现的 Speed Up,以及按专家标注的 memory bytes / FLOPs 算出实测吞吐占 A100 理论峰值的比例(GPU Efficiency)。ACL 2025 主会,S2 引用 78,是后续 Triton 代码生成工作(KernelLLM、各种 RL for kernel 工作)普遍引用的基线。
机制与做法
两个通道的构建
G 通道:爬 100+ star 的 Triton 相关仓库(95 个仓库、845 个 Python 文件),LLM prompt 过滤出 250 个含 Triton 代码的文件,再人工逐个修复缺失组件、拆分多算子文件、给纯 kernel 补 wrapper,用 CodeBERTScore 去重,最后为每个算子写自然语言 instruction(功能、函数名、输入输出示例,均经人工校验)。难度由 LLM 初标 + 两名专家复核,统计上 func 数/参数数/行数/token 数随难度单调上升。
T 通道:从 PyTorch 算子统计里采样,人工写出与 PyTorch 接口对齐的 Triton 参考实现,并额外构造 common/uncommon 算子的融合(fusion)组合——这是 Triton 相对 PyTorch 的主要收益来源(省掉中间读写)。另附 8k 条过滤/合成的 GitHub 训练语料(train_crawl 4024 + train_synth 4133),供微调或 RAG 用。
五个指标,重点在性能
评测链条:CodeBLEU 相似度(仅 G 通道)→ Call Accuracy(能否无报错跑起来)→ Execution Accuracy(输出与参考一致)→ 对跑对的算子测 Speed Up(t_ref/t_gen)和 GPU Efficiency(实测 GB/s、TFLOPS 占 A100 理论峰值的比例,取两者峰值)。
一个有意思的副产品:他们对 184 个人写的 G 通道参考算子本身也算了 GPU efficiency,平均只有 43.0%,19.6% 的算子低于 10%——专业 Triton 程序员写的 kernel 也大量没榨干硬件,说明这事本身就难,也提示'超过参考实现'并不等于'接近最优'。
主要实验设置
基线覆盖两类:通用大模型直接推理(GPT-4o、Claude-3.5-Sonnet、Qwen2.5-72B、DeepSeek-R1、GPT-o1),以及 7B 级代码模型(Qwen2.5-Coder、DeepSeek-Coder)在 8k Triton 语料上 SFT 后再测。zero-shot 和 one-shot(BM25 从训练语料检索最相近示例)各测一遍。
关键结果
- G 通道(真实 GitHub 算子)非常难:最好成绩是 GPT-o1 one-shot 23.91% execution accuracy,DeepSeek-R1 22.83%;zero-shot 下最好的也只有 15% 左右。
- 7B 代码模型(Qwen2.5-Coder、DeepSeek-Coder)zero-shot 在 G 通道 accuracy 为 0,SFT 8k Triton 语料后升到 5-12%——预训练语料里 Triton 太少,domain 数据立竿见影但天花板仍低。
- T 通道(PyTorch 对齐)相对容易:DeepSeek-R1 zero-shot execution accuracy 53.01%,GPT-o1 one-shot 43.37%;论文归因于 T 通道难度分布更均匀,而 G 通道集中在 d3/d4。
- 跑对的算子平均 Speed Up 大多在 1.0x 上下(R1 在 T 通道 one-shot 达 1.91x,主要来自算子融合省内存读写),即 LLM 目前基本只能'复刻'而非'优化'。
- one-shot 检索示例在 G 通道普遍带来约 10 个点的提升,且显著减少 Run&Logic 类错误,但会引入更多 Syntax / Name&Ref 错误——示例对逻辑结构有帮助,对 API 细节反而可能误导。
- 人写的参考算子平均 GPU efficiency 也只有 43%,近 1/5 低于 10%。
实证核查
与我们方向的关系
对课题组做'LLM 写 kernel / AI for systems'方向,这是比 KernelBench 更贴近现役训练栈的评测对象:Triton 是 vLLM、TorchInductor、各家训练框架 fused kernel 的通用 DSL。它的两通道设计(真实需求 vs PyTorch 语义翻译)和'正确性 + speedup + 占理论峰值比例'的三层指标,可以直接搬到自己的 kernel 生成/优化 agent 评测里。
两个可借鉴的教训:一是基准里连人写参考算子平均 GPU efficiency 都只有 43%,做'beat reference'类 RL reward 时要意识到 reference 本身水位不高;二是它暴露的坑(stdout 比对判等、参考 kernel 内存越界、复现分数对不上)正是自建 kernel 评测时最容易犯的错——判等应该在 tensor 层面做 allclose,参考实现要过 sanitizer。若要实际使用,优先考虑社区修复版 TritonBench_T_verified。