Geak: Introducing Triton Kernel AI Agent & Evaluation Benchmarks

TL;DR

This paper introduces GEAK, an AI framework that generates high-performance Triton GPU kernels for AMD hardware, demonstrating significant improvements over baseline methods in correctness and speed, and providing an evaluation suite for Triton kernels.

Contribution

The paper presents GEAK, a novel AI-driven framework for generating efficient Triton GPU kernels for AMD GPUs, along with an evaluation suite for benchmarking such kernels.

Findings

GEAK outperforms baseline LLM prompts and Reflexion pipelines in correctness and speed.

Achieves correctness up to 63% and speedups up to 2.59x on evaluation benchmarks.

Demonstrates potential for democratizing expert-level GPU kernel development.

Abstract

The demand for AI-generated GPU kernels is rapidly growing, influenced by the need for scalable, hardware-optimized solutions in both industry and academia. As deep learning workloads grow in complexity and diversity, it is imperative to automate low-level kernel development to meet performance and productivity demands. Major cloud providers, semiconductor companies, and research institutions are now investing heavily in AI-driven code generation for GPUs, aiming to reduce manual optimization efforts while achieving near-expert performance on hardware like AMD MI300X. The Triton language, a Python-based DSL for GPU programming, has emerged as a popular target for such AI-generated kernels due to its balance of performance and ease-of-coding. In this work, we present an evaluation suite for Triton-based GPU kernels and GEAK (Generating Efficient AI-centric GPU Kernels)-a framework that leverages cutting-edge LLMs to generate performant Triton code specifically for AMD GPUs, including the AMD MI300X and MI250. GEAK leverages inference-time compute scaling to produce Triton-based GPU kernels using a reasoning loop adapted from Reflexion-style feedback mechanisms. On two evaluation benchmarks, GEAK significantly outperformed the baselines of directly prompting frontier LLMs as well as Reflexion-based generation pipelines by achieving correctness up to $63$% and execution speed up of up to $2.59$X. These results highlight the promise of GEAK-like agentic code generation for accelerating the adoption of diverse hardware platforms and democratizing access to expert-level kernel performance.