TritonBench: Benchmarking Large Language Model Capabilities for Generating Triton Operators

要約

効率的なGPUカーネルを構築するために設計された高レベルのPythonのような言語であるTritonは、その携帯性、柔軟性、およびアクセシビリティのために、深い学習フレームワークに広く採用されています。
ただし、プログラミングと並列の最適化には、Triton開発者からのかなりの試行錯誤が必要です。
従来のコード生成の大規模な言語モデル（LLMS）の進歩にもかかわらず、これらのモデルは、その仕様とGPUプログラミングの複雑さの認識が欠けているため、正確でパフォーマンスが最適化されたTritonコードを生成するのに苦労しています。
さらに重要なことは、トリトンに合わせた体系的な評価が緊急に必要であることです。
この作業では、Triton Operator Generationの最初の包括的なベンチマークであるTritonbenchを紹介します。
Tritonbenchには、GitHubの184の実世界オペレーターのキュレーションセットと、Pytorchインターフェイスに沿ったオペレーターのコレクションの2つの評価チャネルがあります。
機能的正しさの優先順位付けの従来のコードベンチマークとは異なり、Tritonbenchは、業界アプリケーションに合わせた広く展開されたGPUの効率性能もプロファイルします。
私たちの研究は、現在の最先端のコードLLMSが効率的なTritonオペレーターを生成するのに苦労しており、高性能コード生成の大きなギャップを強調していることを明らかにしています。
Tritonbenchはhttps://github.com/thunlp/tritonbenchで入手できます。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google