Are LLM-Generated GPU Kernels Production-Ready? A Trace-Driven Benchmark and Optimization Agent

2026-07-15

AI Summary

This work addresses the limitation of existing GPU kernel generation benchmarks, which often diverge from real-world production workloads and thus fail to accurately evaluate the performance of large-model-generated kernels. To bridge this gap, we introduce Atrex-Bench—the first benchmark constructed from full-cluster, real inference traces with GPU time–weighted metrics—and present Atrex-Kernel-Agent, an optimization agent that integrates performance profiling with knowledge-driven strategies. Leveraging trace-driven sampling, iterative evaluation, an optimization discard mechanism, and a hierarchical knowledge base comprising 298 reference kernels and 244 documentation artifacts, our approach significantly enhances both correctness and efficiency of generated kernels. Experiments show that state-of-the-art coding agents achieve only ~10% of hardware roofline performance on Atrex-Bench, whereas our method eliminates reliance on PyTorch fallbacks and produces kernels that match or even surpass hand-tuned baselines.

Abstract

Existing GPU kernel generation benchmarks draw problems from synthetic or curated sources that diverge from deployed workloads. We present Atrex-Bench, a benchmark whose 30 operators and 440 shapes are sampled directly from full-cluster production inference traces of compute-limited, memory-rich GPUs. Each problem carries an importance weight derived from its share of observed GPU time, weighted by application card-hours and computed separately for the serving phases in which it runs, together with a per-problem roofline ceiling, so the aggregate score emphasizes the kernels that consume the most serving time. Evaluating six frontier coding agents on Atrex-Bench shows that even the best vanilla model reaches only ${\sim}10\%$ of the hardware roofline on production operators; and correctness alone overstates capability, since much of the apparent pass rate comes from PyTorch fallbacks rather than kernels the model wrote. To close this gap, we co-release Atrex-Kernel-Agent (AKA), a profile-driven kernel-optimization agent that combines iterative measure-revise search, optimization dropout for escaping stalled search contexts, and a layered GPU-optimization knowledge base (298 reference-kernel files and 244 optimization-knowledge documents, plus external upstream reference projects for API/ISA lookup). In a controlled case study, the agent converts zero-FlyDSL fallbacks into real kernels that match or exceed hand-tuned production baselines.

Problem

• LLM-generated GPU kernels
• production readiness
• trace-driven benchmark
• kernel optimization
• performance evaluation

Innovation

• trace-driven benchmark
• GPU kernel generation
• optimization agent
• profile-driven search
• production-ready evaluation