arXiv:2502.14752cs.CLcs.LG2025-02ACL被引 78

首个专为Triton代码生成设计的基准,评估大模型生成高效GPU算子能力。

TritonBench: Benchmarking Large Language Model Capabilities for Generating Triton Operators

  • 构建包含184个真实算子的双通道评测集,覆盖GitHub与PyTorch接口
  • 发现当前顶级大模型在生成高效Triton代码上表现不佳,效率差距显著
  • 聚焦工业级GPU性能,填补高性能代码生成评估空白,适合框架开发者

Triton是一种面向GPU内核开发的高级Python类语言,因其可移植性、灵活性和易用性被广泛应用于深度学习框架。然而,编写和并行优化仍需大量试错。尽管大语言模型在常规代码生成方面取得进展,但它们在生成准确且性能优化的Triton代码时表现不佳,原因在于缺乏对Triton规范及GPU编程复杂性的认知。更关键的是,亟需针对Triton的系统性评估。本文提出TritonBench,首个全面的Triton算子生成基准。该基准包含两个评测通道:从GitHub精选的184个真实算子,以及与PyTorch接口对齐的算子集合。不同于传统代码基准仅关注功能正确性,TritonBench还针对广泛应用的工业级GPU进行效率性能分析。研究揭示,当前最先进的代码大模型在生成高效Triton算子方面存在明显短板,暴露出高性能代码生成的巨大差距。TritonBench将开源发布于https://github.com/thunlp/TritonBench。

原文摘要 · Abstract (English)

Triton, a high-level Python-like language designed for building efficient GPU kernels, is widely adopted in deep learning frameworks due to its portability, flexibility, and accessibility. However, programming and parallel optimization still require considerable trial and error from Triton developers. Despite advances in large language models (LLMs) for conventional code generation, these models struggle to generate accurate, performance-optimized Triton code, as they lack awareness of its specifications and the complexities of GPU programming. More critically, there is an urgent need for systematic evaluations tailored to Triton. In this work, we introduce TritonBench, the first comprehensive benchmark for Triton operator generation. TritonBench features two evaluation channels: a curated set of 184 real-world operators from GitHub and a collection of operators aligned with PyTorch interfaces. Unlike conventional code benchmarks prioritizing functional correctness, TritonBench also profiles efficiency performance on widely deployed GPUs aligned with industry applications. Our study reveals that current state-of-the-art code LLMs struggle to generate efficient Triton operators, highlighting a significant gap in high-performance code generation. TritonBench will be available at https://github.com/thunlp/TritonBench.

Triton大模型代码生成性能优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。