PaperlayerView exact revision on arXiv

Source-linked AI summary

TritonBench: Benchmarking Large Language Model Capabilities for Generating Triton Operators

Jianling Li, Shangzhan Li, Zhenye Gao, Qi Shi, Yuxuan Li, Zefan Wang, Jiacheng Huang, Haojie Wang, Jianrong Wang, Xu Han, Zhiyuan Liu, Maosong Sun

arXiv:2502.14752v1cs.CLcs.LG

TL;DR

TritonBench addresses the lack of systematic, performance-aware evaluation for LLM-generated Triton operators, whose domain-specific GPU requirements remain difficult for current models. It introduces complementary real-world and PyTorch-aligned benchmark channels and finds that LLMs still struggle to generate accurate, efficient Triton code.

  • Problem

    LLMs struggle with Triton’s specifications and GPU-programming complexity, while existing evaluations insufficiently measure the accuracy and efficiency of generated operators.

  • Method

    TRITONBENCH combines 184 GitHub operators with PyTorch-aligned tasks and evaluates generated code for correctness and NVIDIA-GPU performance.

  • Results

    Current LLMs struggle with high-quality Triton generation; the highest execution accuracy reaches 23.91% on TRITONBENCH-G and 53.01% on TRITONBENCH-T.

  • Takeaways & Limitations

    TRITONBENCH provides a framework for advancing automated Triton operator generation while highlighting the need for accurate and performance-aware code generation.

  • Takeaways & Limitations

    Evaluations were conducted exclusively on the NVIDIA A100 GPU, with broader hardware architectures left for future work.

Abstract

from arXiv · show

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.

1 Introduction

Triton enables accessible GPU-kernel development but still demands difficult performance tuning, while LLMs remain inadequately evaluated for generating accurate, efficient Triton code. TRITONBENCH addresses this gap with complementary datasets, performance-aware evaluation, and broad experiments.

  • Triton simplifies GPU-kernel programming, but developers still manually manage memory access, parallel coordination, and hardware-specific optimization through trial and error.
  • LLMs perform strongly on general-purpose code yet struggle with Triton’s domain-specific specifications and GPU-programming intricacies.
  • TRITONBENCH combines 184 real-world GitHub operators in TRITONBENCH-G with PyTorch-aligned tasks in TRITONBENCH-T.
  • The benchmark evaluates generated operators using similarity, call and execution accuracy, speedup, and GPU efficiency, with expert GPU-performance assessment.
  • 23.91% is the highest execution accuracy on TRITONBENCH-G, versus 53.01% on TRITONBENCH-T; the best speedups are 1.56× and 1.91×, respectively.

2 Related Work

Prior coding benchmarks and LLMs largely emphasize general-purpose languages and functional correctness, leaving GPU-oriented DSL generation insufficiently assessed. TritonBench targets this gap with systematic evaluation of correctness and GPU performance.

  • Triton balances simplified GPU programming with explicit control over memory access and parallelism, yet developers must still tune parameters to exploit hardware capabilities.
  • Existing coding benchmarks primarily test functional correctness, while performance-oriented frameworks rarely target GPU code comprehensively.
  • LLMs achieve strong results on broad coding benchmarks but struggle with DSLs because specialized datasets are limited and DSL syntax, semantics, and parallel models are distinctive.

3 TRITONBENCH-G

TRITONBENCH-G curates real-world Triton operators, validates and annotates them, and evaluates generated code for correctness and GPU performance. Its dataset and expert ratings expose the practical difficulty of producing efficient Triton operators.

  • A Triton operator typically combines a GPU kernel for tensor addressing and parallel coordination with a Python wrapper that invokes it.
  • TRITONBENCH-G curates human-authored operators from GitHub repositories, then manually inspects, repairs, separates, wraps, deduplicates, and annotates them.
  • Difficulty levels d1 through d5 are assigned with LLM guidance and validated by two domain experts, with statistics covering functions, parameters, lines, and tokens.
  • Longer instructions deliberately provide richer context for nuanced requirements and better reflect real-world operator development.
  • 43.0% is the average GPU efficiency of TRITONBENCH-G operators, while 19.6% of professionally developed operators score below 10%.
  • Evaluation covers text similarity, call and execution accuracy, speedup, and GPU efficiency for generated operators.

4 TRITONBENCH-T

TRITONBENCH-T complements GitHub-derived operators with PyTorch-aligned tasks spanning common and uncommon GPU operators. Its construction and evaluation support a broader assessment of correctness and performance in Triton generation.

  • Task construction: TRITONBENCH-T selects GPU-interacting PyTorch operators by real-world usage frequency and fuses them into diverse operator-generation tasks.The final set contains 166 operators from PyTorch v2.6.0, each paired with a standard PyTorch call and documentation.
  • Task statistics: Its difficulty statistics cover five levels, d1–d5, with two domain experts validating LLM-guided categorizations.Average operator count, parameter count, mathematical-expression tokens, and description tokens generally increase with difficulty.
  • Benchmark comparison: TRITONBENCH-G emphasizes frequent real-world GitHub operators, whereas TRITONBENCH-T offers a more diverse PyTorch-derived distribution including uncommon operators.This contrast makes the two channels complementary rather than redundant.
  • Evaluation: Both channels evaluate correctness and performance, while TRITONBENCH-G additionally compares generated code with established implementations through similarity assessment.The shared metrics include Call Accuracy, Execution Accuracy, and Speed Up.
  • Benchmark comparison: The complementary benchmark designs enable a comprehensive and nuanced evaluation of Triton operator generation.TRITONBENCH-T broadens coverage beyond the highly frequent operations emphasized by the real-world channel.

5 Experiments

Experiments evaluate state-of-the-art specialized and general-purpose LLMs under zero-shot and one-shot settings. General-purpose models perform best overall, while the two benchmark channels differ substantially in difficulty and achievable accuracy.

  • Experimental setup: The experiments compare specialized open-source and general-purpose state-of-the-art LLMs under zero-shot and one-shot conditions.Domain-specific models receive supervised fine-tuning, and one-shot prompts use BM25 retrieval from the training corpus.
  • TRITONBENCH-G: 1.22× is DeepSeek-R1’s best reported GPU execution-time improvement on TRITONBENCH-G in the one-shot setting.Its zero-shot improvement is 1.11×.
  • Cross-channel comparison: Most models perform better on TRITONBENCH-T than on TRITONBENCH-G, likely because T has a more balanced difficulty distribution.TRITONBENCH-G is predominantly composed of higher-difficulty d3 and d4 operators.

6 Analysis

The analysis examines correctness across difficulty levels and categorizes compiler-reported execution failures. It finds widespread generation errors, model-specific error sensitivities, and differing effects of one-shot examples across the two channels.

  • TRITONBENCH-G difficulty: Most TRITONBENCH-G operators are generated incorrectly, although DeepSeek-R1 outperforms GPT-o1 and both improve substantially on d4 with one-shot examples.The d4 gains may reflect frequent Attention and Softmax operators that provide similar examples.
  • Error analysis: The 16 compiler-reported error types are grouped into four major categories for analysis.One-shot prompting increases Syntax and Name&Ref errors while reducing Attr&Type and Run&Logc errors for both leading models.
  • TRITONBENCH-T difficulty: DeepSeek-R1 generates more correct than incorrect TRITONBENCH-T operators, unlike the predominantly incorrect results observed for TRITONBENCH-G.This pattern is consistent with the smoother difficulty distribution in TRITONBENCH-T.
  • TRITONBENCH-T difficulty: One-shot prompting affects the models differently on TRITONBENCH-T: DeepSeek-R1 declines on d2–d4, whereas GPT-o1 improves on those subsets.The analysis associates these divergent trends with models’ differing use of the provided sample operators.
  • Error analysis: DeepSeek-R1 avoids Syntax errors entirely on TRITONBENCH-T, while GPT-o1 retains a high Syntax-error rate.Across both channels, one-shot prompting consistently reduces Run&Logc errors, while other error categories vary by model and setting.

7 Conclusion

TRITONBENCH is a dual-channel benchmark for Triton operator generation, combining real-world repository samples with PyTorch-aligned tasks. Its evaluation covers both functional accuracy and NVIDIA-GPU performance, and experiments show that current LLMs struggle to generate high-quality Triton operators.

  • Contribution: TRITONBENCH combines GitHub-derived real-world operators with complementary tasks aligned to PyTorch interfaces.The benchmark is designed specifically to evaluate LLM generation of Triton operators.
  • Evaluation: The evaluation framework measures both functional accuracy and performance on NVIDIA GPUs.This design addresses the need for Triton code that is not only correct but also performance-aware.
  • Findings: Experiments and detailed analyses find that current LLMs struggle to generate high-quality Triton operators.The conclusion underscores the need for further advances in accurate and performance-aware Triton code generation.

Limitations

The study evaluates TRITONBENCH exclusively on the NVIDIA A100 GPU, limiting its performance conclusions to that hardware setting.

  • The evaluations were conducted exclusively on the NVIDIA A100 GPU.The authors identify broader hardware coverage as future work.

A Training Corpus

The training corpus combines real-world Triton code from GitHub with synthetically generated data produced through compiler operations. The real-world component is cleaned, debugged, deduplicated against TRITONBENCH-G, and combined with synthetic samples for one-shot learning experiments.

  • The corpus contains real-world GitHub data and synthetically generated data produced through compiler operations.
  • GitHub-sourced Triton code undergoes basic cleaning, debugging, and CODEBERTSCORE-based filtering to remove samples highly similar to TRITONBENCH-G entries.The debugging is less rigorous than the TRITONBENCH-G methodology.
  • The synthetic component uses Ninetoothed4 to automate valid Triton-code synthesis from well-formed expressions.
  • The combined corpus serves as the foundational training dataset for experimental models in one-shot learning settings.

B Operator Performance Evaluation

Operator performance evaluation runs operators on increasingly large tensors, measures stabilized GPU execution time, and converts memory traffic and floating-point work into throughput and efficiency metrics. Performance curves show throughput rising before stabilizing as input dimensions grow.

  • Evaluation uses tensors with increasing dimensions, feeding each tensor sequentially into the operator.
  • Execution time is measured with triton.testing.do_bench after increasing warm-up and repetition times until measurements stabilize.Operators are often run hundreds of thousands of times for accurate timing.
  • Throughput is calculated as memory bandwidth divided by execution time in GB/s and floating-point operations divided by execution time in Tflops.
  • GPU efficiency compares measured GB/s and Tflops with the NVIDIA A100’s theoretical maximum performance.
  • As input dimensions increase, operator GB/s or Tflops rises and eventually stabilizes, indicating a bottleneck beyond a certain scale.

C Error Categories

TRITONBENCH organizes call and execution failures into four broad error groups and uses prompts to extract, describe, test, and assess Triton operators and their complexity.

  • Error Categories: Sixteen error types are grouped into Syntax, Attrb&Type, Name&Ref, and Run&Logc Errors.
  • Error Categories: ResultsError denotes inconsistency between the reference operator’s execution results and those of the generated operator.
  • Prompts: The workflow uses filtering, instruction, difficulty, and test-code prompts for operator extraction, description, scoring, and test generation.
  • Operator Extraction: Extraction retains @triton.jit kernels, their calling functions, and necessary imports while discarding files without Triton kernels.
  • Operator Description: Extracted operators receive concise English descriptions intended to support reproduction of their kernel and wrapper code.
  • Test Generation: Generated tests keep branch tests in one parameterless function beginning with test_, while tensor parameters must use GPU types.
  • Difficulty Assessment: The complexity score evaluates memory-layout and computation-scheduling complexity on a 1–5 scale using the ceiling of their average.
Loading 2502.14752v1…