Geak: Introducing Triton Kernel AI Agent & Evaluation Benchmarks

Motivation

• The framework targets the growing need for high-performance GPU kernels without manual tuning as workloads and hardware diversity increase.The motivation includes deployment on emerging hardware such as AMD Instinct MI300X GPUs.Source: §1 · p. 1

Contribution

• GEAK combines generation, evaluation, reflection, and optimization agents to iteratively improve Triton kernel correctness and performance.The framework uses inference-time compute scaling with frontier LLMs and requires only minimal task descriptions.Source: §1 · p. 2

Benchmarks

• The work releases TritonBench-revised and ROCm Triton benchmarks measuring both execution correctness and runtime performance.The suites contain 184 revised kernels and 30 real-world kernels from open-source AMD ROCm repositories.Source: §1 · p. 2

Results

• GEAK achieves up to 63.33% correct kernel generation, compared with less than 15% for direct prompting of strong LLMs.The reported rates are 54.89% on TritonBench-revised and 63.33% on the ROCm Triton benchmark.Source: §1 · p. 2

Contribution

• The authors also open-source the GEAK agent implementation and evaluation framework.Source: §1 · p. 2

Results

• GEAK-generated kernels deliver up to 2.59× average speedup over reference kernels on TritonBench-revised.Source: §1 · p. 2

Benchmarks

• Existing coding benchmarks commonly focus on functional correctness, whereas GPU-kernel evaluation also requires runtime and hardware-specific assessment.GPU-targeted benchmarks remain scarce despite the importance of hardware efficiency for large-scale deep learning.Source: §2 · pp. 2–3
• KernelBench and TritonBench provide domain-specific evaluation of correctness and runtime efficiency for GPU-kernel generation.These benchmarks expose the gap between general-purpose code generation and performant Triton code.Source: §2 · p. 3

LLMs for Code Generation

• General-purpose code models such as DeepSeek-Coder[11] and Qwen-Coder[12] perform strongly on broad coding benchmarks but have limited success on performance-critical low-level code.Source: §2 · p. 3
• KernelLLM is restricted to translating PyTorch into Triton, while GEAK accepts natural-language instructions and reference code.The paper also contrasts specialized models with the broader reasoning capabilities of GPT-4.1 and Claude 4 for constrained optimization tasks.Source: §2 · p. 3

Inference-Time Compute Scaling

• Inference-time compute scaling improves complex-task performance through additional reasoning passes, parallel generations, or longer contexts without additional training.GEAK combines this strategy with Reflexion-style feedback for iterative code refinement.Source: §2 · p. 3

Benchmark Construction

• The benchmarks contain 184 modified TritonBench kernels and 30 kernels drawn from open-source ROCm repositories.Source: §3 · p. 3
• The revised benchmark fixes AMD-GPU incompatibilities, missing test calls, inconsistent unit tests, and unreliable output comparisons.Corrections include tolerance-based tensor comparison and consistent random seeds.Source: §3.1 · p. 3
• The ROCm kernels were authored by Triton experts and adapted with frontier LLM assistance into the same evaluation format.This enables consistent comparison of agent accuracy and efficiency.Source: §3.2 · p. 4

Evaluation Metrics

• Call Accuracy measures the fraction of generated kernels that compile and run without errors.Source: §3.3 · p. 4
• Execution Accuracy measures the percentage of generated kernels passing all unit tests among kernels that compile successfully.Source: §3.3 · p. 4
• Speedup is the expected ratio of median reference latency to generated-kernel latency across unit tests, computed only for kernels passing all tests.Source: §3.3 · p. 4

4.1 Architecture

• The four modules are Generator, Reflector, Evaluator, and Optimizer, with evaluation cascading from functionality testing to latency and memory assessment.Failed functionality tests send error traces to the Reflector; passing code proceeds to performance evaluation.Source: §4.1 · p. 5

4.2 Modules

• One-shot prompting retrieves a code-similar Triton example, while knowledge injection supplies hardware-specific optimization guidance.The retrieved examples do not overlap with the benchmark datasets.Source: §4.2 · p. 5
• Reflexion uses execution error traces to diagnose failed kernels and iteratively refine them across multiple attempts.Source: §4.2 · p. 5
• Different underlying LLMs produce substantially different kernel-generation results, making model selection an important system factor.Source: §4.2 · p. 5
• The Optimizer ranks prior generations by performance and uses their code and metrics to propose more effective optimization directions.Source: §4.2 · p. 5
• A debugging-attempt limit prevents repeated reflection cycles from trapping the agent on the same bug; persistent failures trigger a new strategy and fresh code.Source: §4.2 · p. 5

4.3 Parallel Scaling

• Parallel independent runs diversify candidate kernels, and combining parallel with sequential scaling can improve performance estimates and outcomes.The experiments use pass@k for unbiased accuracy estimation.Source: §4.2 · p. 5

5.1 Main Results

• Direct prompting of frontier LLMs produces low execution accuracy, while one-shot prompting generally improves call accuracy, execution accuracy, and speedup over zero-shot prompting.Gemini 2.5 Pro is an exception for speedup: one-shot prompting lowers speedup despite improving accuracy.Source: §5.1 · pp. 5–6
• 54.89% execution accuracy and 2.59× speedup on TritonBench-revised, versus direct prompting’s limited results; ROCm reaches 63.33% execution accuracy and 0.92x speedup.These results are reported on MI300 and compare GEAK with direct LLM prompting.Source: §5.1 · p. 6, Table 3 · p. 6, Table 4 · p. 6

5.2.1 Sequential Scaling

• Increasing sequential iterations monotonically improves correctness, reaching 63.04% call accuracy and 44.02% execution accuracy at iter19, more than 3× the baseline.Accuracy continues improving beyond iter10, including gains from 56.52% to 63.04% in call accuracy and from 40.76% to 44.02% in execution accuracy.Source: §5.2.1 · p. 7, Table 5 · p. 7
• Sequential refinement improves runtime early: iter1 reaches 1.89× speedup, while later iterations remain above 1.4× from iter7 onward.The early gain is attributed to correcting suboptimal code generations.Source: §5.2.1 · p. 7, Table 5 · p. 7

5.2.2 Parallel Scaling

• Parallel scaling produces near-log-linear gains in call and execution accuracy and complements sequential refinement across MI250 and MI300.Together, the two axes provide a flexible way to tune the accuracy-performance tradeoff under a fixed compute budget and show consistent patterns across hardware.Source: §5.2.2 · pp. 7–8, Table 7 · p. 8

5.2.3 Effect of different Modules

• Knowledge injection raises call accuracy from 14.67% to 52.72% and execution accuracy from 8.70% to 20.11%, while adding one-shot reaches 54.35% and 27.17%.Adding the Optimizer to both modules yields 40.76% execution accuracy and the highest speedup, 1.45×.Source: §5.2.3 · pp. 8, 10, Table 9 · p. 10

• GEAK coordinates specialized agents to iteratively improve generated Triton code's correctness and performance without additional training.Source: §6 · p. 10
• The authors introduce TritonBench-revised and a new ROCm benchmark to evaluate functional correctness and runtime efficiency across diverse kernels.Source: §6 · p. 10
• The benchmarks and agent code are fully open-sourced to encourage community development of GPU kernels.Source: §6 · p. 10

• The case study evaluates a GEAK-generated flip kernel against a corresponding Triton expert-written kernel.Source: A Case study: Flip kernel · pp. 11–12
• The kernel processes 2D blocks by loading an (N, M) block, flipping each row horizontally, and storing the result.Source: A Case study: Flip kernel · p. 11
• GEAK computes flipped source pointers directly from row and flipped-column offsets before storing values at normal output positions.Source: A Case study: Flip kernel · p. 12
• The expert implementation defines a flip kernel over flattened input and output tensors with N rows and M columns.Source: A Case study: Flip kernel · p. 11

Expert written code limitations

• The expert implementation loads an entire N×M block into registers, flips it there, and stores it back, increasing bandwidth use and register pressure.Source: §A.1 · p. 12

GEAK generated code advantages

• GEAK's direct addressing reads flipped source positions and writes normal destinations in a single pass, avoiding intermediate whole-block storage.Source: §A.1 · p. 12
• Explicit boundary handling prevents out-of-bounds accesses, and row-major addressing maintains coalesced access because nearby threads read nearby locations.Source: §A.1 · p. 13

Performance Implications

• GEAK uses approximately 50% less memory bandwidth because it performs one read and write instead of two reads and writes.Source: §A.1 · p. 13
• Lower register pressure permits larger block sizes, while direct addressing improves cache friendliness and scalability for larger tensor dimensions.Source: §A.1 · p. 13

• Broad, well-crafted unit tests are necessary because insufficient coverage can allow partially correct or flawed generated code to pass.Source: The Critical Role of Robust Unit Testing · p. 13
• Higher statement, branch, or decision coverage increases confidence that tests exercise varied execution paths and reduces undetected failures.Source: The Critical Role of Robust Unit Testing · p. 13
• Automatically adding tests reduced flawed-code passing rates by up to 28.9 percentage points, while GPT-4o test suites averaged only 35.2% coverage.Source: The Critical Role of Robust Unit Testing · p. 13

• ROCm Triton Benchmark uses wide unit-test coverage across input conditions, edge cases, and reference outputs to assess functional and numerical correctness.Its broader testing is intended to increase confidence that passing kernels are genuinely correct.Source: Our Benchmark Test Coverage Comparison · p. 13
• TritonBench-G contains at most six unit tests per kernel, creating limited coverage and a risk of false positives.Kernels may pass the available tests while failing under untested conditions.Source: Our Benchmark Test Coverage Comparison · p. 13
• TritonBench-revised retains TritonBench-G's unchanged harness, so its narrow coverage may inflate apparent correctness for overfitted kernels.The benchmark is functionally consistent with TritonBench-G but remains vulnerable to test-case overfitting.Source: Our Benchmark Test Coverage Comparison · p. 13

• High-coverage benchmarks such as ROCm Triton should be preferred for AI-generated kernels, especially when subtle numerical errors matter.Source: Recommendations & Evaluation Scope · p. 13
• Narrow suites such as TritonBench-revised require explicit coverage caveats because pass rates of 50% or more may overstate general correctness.Source: Recommendations & Evaluation Scope · p. 14
• Combining broad coverage with transparent limitation reporting supports accurate and trustworthy evaluation of AI-generated GPU kernels.Source: Recommendations & Evaluation Scope · pp. 13–14