PaperlayerView exact revision on arXiv

Source-linked AI summary

KernelBench: Can LLMs Write Efficient GPU Kernels?

Anne Ouyang, Simon Guo, Simran Arora, Alex L. Zhang, William Hu, Christopher Ré, Azalia Mirhoseini

arXiv:2502.10517v1cs.LGcs.AIcs.PFcs.SE

TL;DR

Writing efficient GPU kernels is difficult and hardware-dependent, motivating a benchmark for whether language models can automate correct, performant kernel generation. The paper introduces KernelBench and fast_p to evaluate this capability across diverse PyTorch workloads. Frontier models remain below the PyTorch baseline on most tasks, while execution and profiling feedback improves iterative refinement.

  • Problem

    Efficient GPU kernels are important for AI performance, cost, and energy, but developing them remains challenging across diverse architectures and hardware platforms.

  • Method

    The paper introduces KernelBench, an evaluation-only framework with 250 PyTorch workloads, automated correctness and performance testing, iterative feedback, and the fast_p metric.

  • Results

    OpenAI-o1 and DeepSeek-R1 matched the PyTorch Eager baseline on < 20% of tasks out of the box, while feedback-based refinement improved reported fast1 results.

  • Takeaways & Limitations

    KernelBench is designed to remain useful as new AI workloads and hardware platforms emerge, with benchmark success intended to map to practical production value.

  • Takeaways & Limitations

    The evaluation is limited to GPUs, and iterative-refinement benefits depend on the quality of the base model.

Abstract

from arXiv · show

Efficient GPU kernels are crucial for building performant machine learning architectures, but writing them is a time-consuming challenge that requires significant expertise; therefore, we explore using language models (LMs) to automate kernel generation. We introduce KernelBench, an open-source framework for evaluating LMs' ability to write fast and correct kernels on a suite of 250 carefully selected PyTorch ML workloads. KernelBench represents a real-world engineering environment and making progress on the introduced benchmark directly translates to faster practical kernels. We introduce a new evaluation metric fast_p, which measures the percentage of generated kernels that are functionally correct and offer a speedup greater than an adjustable threshold p over baseline. Our experiments across various state-of-the-art models and test-time methods show that frontier reasoning models perform the best out of the box but still fall short overall, matching the PyTorch baseline in less than 20% of the cases. While we show that results can improve by leveraging execution and profiling feedback during iterative refinement, KernelBench remains a challenging benchmark, with its difficulty increasing as we raise speedup threshold p.

1 Introduction

KernelBench frames LM kernel generation as an attempt to automate a difficult, hardware-dependent engineering task. It evaluates whether models can produce correct and performant kernels while exposing workflow information and iterative feedback.

  • Motivation: Efficient GPU kernels matter for AI performance, cost, and energy, but kernel development remains difficult across rapidly changing architectures and hardware.The paper highlights delayed optimization for important workloads such as FlashAttention.
  • Framework: KernelBench provides an environment that automates the AI engineer’s workflow and supports diverse algorithms, languages, hardware platforms, and evaluation information.The environment can incorporate compiler, profiling, hardware-specific, and execution information.
  • Benchmark scope: The benchmark spans individual operations, fused sequences of operations, and end-to-end ML architectures.These levels test optimization from foundational operators through multi-operation fusion and complete models.
  • Results: < 20% of tasks were matched to the PyTorch Eager baseline by OpenAI-o1 and DeepSeek-R1 out of the box.Generated kernels suffered from execution errors, functional-correctness issues, and limited platform-specific optimization.
  • Results: Language models showed nascent potential for algorithmic kernel improvements through sparsity, operator fusion, and hardware features.Conditioning on hardware information and optimization demonstrations revealed more such instances.
  • Iterative refinement: 12%, 36%, and 12% fast1 results improved to 43%, 72%, and 18%, respectively, after multiple refinements using execution and profiler feedback.The paper presents feedback-driven refinement as a way to reduce execution errors and discover faster solutions.

2 Related Works

Prior kernel-programming systems either require substantial human engineering or focus on narrower code-generation settings. KernelBench instead targets wall-clock efficiency across diverse, real-world deep-learning workloads.

  • Kernel libraries and compilers: Hardware-specific libraries such as cuDNN, CUTLASS, and Apple MLX demand substantial human engineering effort.ThunderKittens and Triton broaden access to fast kernels but still require human programming.
  • LM code generation: Earlier LM code-generation work emphasized functional code, while later studies explored algorithmic and asymptotic efficiency.KernelBench focuses specifically on wall-clock efficiency.
  • HPC code generation: KernelBench curates 250 diverse kernels from modern deep-learning workloads rather than translating arbitrary C++ or generating only standard low-level kernels.Many selected workloads lack existing human-written implementations, making successful solutions immediately relevant to real workloads.

3 KernelBench: A Framework for AI Kernel Generation

KernelBench is an evaluation-only framework for testing language models on correct and fast kernel generation across 250 AI workloads. It combines automated correctness and wall-clock performance evaluation with a unified speedup metric and multiple workload levels.

  • Task design: KernelBench contains 250 tasks spanning single operations, operator sequences, and full ML architectures.The levels contain 100, 100, and 50 tasks, respectively.
  • Task format: Each task gives the LM PyTorch reference code and prescribed input tensors, then requires an optimized ModelNew implementation.The LM may choose which operations to optimize and use libraries or techniques such as fusion, tiling, CUDA, PTX, Triton, or CUTLASS.
  • Motivation: The benchmark targets practical kernel-generation problems whose solutions can directly produce real-world impact.The selected workloads include modern architectures whose performant kernels may require algorithmic changes beyond compiler scope.
  • Evaluation approach: KernelBench is automatically verifiable without ground-truth kernels by comparing outputs on 5 random inputs and measuring repeated wall-clock execution times.Correctness compares Model with ModelNew, while performance compares their execution times.
  • Metric design: The evaluation compares language models using a unified metric that jointly captures correctness and performance.The threshold p can be adjusted to evaluate different speedup requirements.
  • Metric design: fast_p is the fraction of tasks that are functionally correct and achieve a speedup greater than threshold p over PyTorch wall-clock time.The speedup is the ratio of PyTorch time to generated-kernel time; fast_0 equals the correctness rate.

4 KernelBench Baseline Evaluation

KernelBench evaluates off-the-shelf LM-generated kernels for correctness and speed against PyTorch baselines, revealing substantial failures and limited speedups. Reasoning models generally perform better, but performance varies across hardware.

  • Evaluation setup: The one-shot evaluation gives LMs a PyTorch optimization task, generates ModelNew by greedy decoding, and profiles kernels on an NVIDIA L40S GPU.The benchmark reports fast_p across all problems and compares generated kernels with PyTorch baselines.
  • Speed performance: Fewer than 20% of tasks show LM-generated kernels faster than PyTorch Eager on average.Table 1 reports fast1 as the percentage of problems where generated kernels beat the PyTorch Eager and torch.compile baselines.
  • Metric difficulty: Increasing the fast_p speedup threshold raises evaluation difficulty because fast0 counts correct kernels while fast1 requires more than 1× PyTorch speedup.The metric distribution is examined as the threshold p increases.
  • Failure modes: Reasoning models produce fewer execution failures, but all models struggle similarly with functional correctness.Failure categories include compiler and runtime failures, memory violations, and output shape or value mismatches.
  • Speed performance: At p = 1, fewer than 15% of functionally correct LM-generated kernels outperform PyTorch across KernelBench levels.Reasoning-based LMs generally provide more speedups than other evaluated models.
  • Hardware variation: One-shot kernels show hardware-dependent performance, with DeepSeek R1 Level 2 fast1 reaching 36% on L40S versus 47% on A10G.Level 1 speedups transfer more similarly across GPUs, whereas Level 2 exhibits larger variation.

5 Analysis of Model Capabilities

KernelBench case studies show that repeated sampling and iterative feedback can improve kernel quality, while model capabilities remain uneven and hardware-aware optimization is difficult. Gains depend strongly on the base model and the feedback or prompting strategy.

  • Repeated Sampling: Repeated sampling raises fast1 across KernelBench levels, with DeepSeek-V3 reaching 37% at k = 100 on Level 2 versus 4% one-shot.High-temperature sampling explores more solutions, but provides limited benefit when a model has very low task-solving probability.
  • Iterative Refinement: Execution and profiler feedback improve iterative generations, increasing DeepSeek-R1 Level 2 fast1 from 36% to 72% after 10 turns.The framework supplies compiler, correctness, timing, and profiling signals during refinement.
  • Comparing Test-Time Methods: With a fixed budget of 10 calls, iterative refinement outperforms repeated sampling in 5 of 6 comparisons, but effectiveness depends on base-model quality.DeepSeek-V3 consistently benefits more from repeated sampling than Llama 3.1 70B, while feedback does not always help every model.
  • Hardware-aware In-Context Examples: Few-shot optimization examples increase execution failures overall, despite prompting o1 to use tiling on 77% of Level 1 GEMM variants.Among correct solutions, some generations use shared-memory management and other hardware-oriented optimizations.
  • Hardware-aware In-Context Examples: Models rarely produce hardware-optimized kernels, although R1 and o1 generate occasional correct outliers that are ≥2× faster than baseline.R1 attempts WMMA instructions for approximately 50% of Level 1 matrix multiplication problems, but most attempts fail to compile.

6 Discussion

KernelBench is presented as a framework for evaluating LM-driven kernel optimization with practical relevance beyond a fixed benchmark. The discussion identifies substantial room for improvement and several directions for extending models, abstractions, and hardware coverage.

  • Conclusion: KernelBench evaluates diverse models and approaches while analyzing their strengths, limitations, and opportunities for improvement.The paper’s stated contributions are the framework and its broad evaluation of LM-driven kernel optimization.
  • Conclusion: The fastp metric can evolve with new workloads and increasingly advanced baselines, while PyTorch-based tasks remain portable across hardware releases.The authors connect benchmark progress with production value, including lower costs and energy use.
  • Opportunities for Future Work: The study reports significant room for improvement in currently available models.This conclusion motivates the proposed future work on training data, reasoning, code abstractions, and broader hardware coverage.
  • Opportunities for Future Work: Future work includes better fine-tuning and reasoning, more high-quality CUDA data, alternative programming abstractions, and evaluation on non-GPU accelerators.Suggested abstractions include ThunderKittens, CUTLASS, and Triton, which may simplify generation and tensor-core use.

Ethics Statement

The paper frames efficient GPU kernels as relevant to computational cost and energy use, while its benchmark avoids human-subject and privacy concerns. KernelBench relies on publicly available repository code rather than proprietary or private data.

  • Ethics Statement: Optimized GPU kernels can reduce computational costs and energy use in large-scale machine learning workloads.The paper links these savings to reduced environmental impact and more energy-efficient AI systems.
  • Ethics Statement: KernelBench involves no human studies or user data and uses only publicly available GitHub repositories.The paper therefore reports no privacy concerns from human participants or private code.

A KernelBench Task Example

A KernelBench task pairs a fixed-shape PyTorch reference module with a model-generated implementation that may include custom inline CUDA. Evaluation checks randomized-input correctness and compares repeated forward-pass timing against the reference.

  • Task Structure: Each task defines a Model class with init and forward functions, fixed input shapes, and randomly generated initialization and execution inputs.The task specification supports evaluating a PyTorch reference architecture and a generated replacement.
  • PyTorch Reference Architecture: The example task performs matrix multiplication, with inputs shaped (M, K) and (K, N) and output shaped (M, N).The reference architecture is a simple PyTorch module that computes C = A * B.
  • Generated Implementation: The generated implementation must provide kernel code and integration code, commonly using inline CUDA compiled through PyTorch extensions.The evaluation runs the generated forward pass as a PyTorch operator.
  • Evaluation Environment: Evaluations run on an NVIDIA L40S using Python 3.10, PyTorch 2.5.0+cu124, and CUDA 12.4 unless otherwise stated.The default GPU has Ada Lovelace architecture, 48 GB of HBM memory, and 300W power.
  • Evaluation: Correctness uses five randomized inputs, while performance averages 100 warmed-up CUDA-event timing trials for each forward pass.Speedup is computed as TModel/TModelNew and compared with threshold p for fastp scoring.
  • Correctness Evaluation: Five-input random testing is a practical heuristic for numerical kernel equivalence, not a formal proof of program equivalence.In one 100-kernel experiment, 50 were correct, while the remainder showed mismatches, shape errors, runtime errors, or compilation errors.

B.3 Distribution of Model Performance for One-Shot Baseline

Across KernelBench levels, correctly generated kernels usually do not exceed PyTorch Eager speed, although some Level 1 cases produce substantial outliers. Reasoning-optimized models perform best overall, while model scale alone does not ensure strong kernel generation.

  • Performance distribution: Level 1 produces the largest outliers, including one kernel with more than 10x speedup.The highlighted case applies a diagonal-matrix optimization by multiplying each vector element directly with the corresponding row of the matrix.
  • Model comparison: Reasoning-optimized models OpenAI-o1 and DeepSeek-R1 perform best out of the box across all workload levels, especially Level 2 fusion tasks.Llama 3.1 models perform poorly regardless of whether they have 405B or 70B parameters.
  • Baseline comparison: PyTorch torch.compile is generally faster than PyTorch Eager on Levels 2 and 3 because of graph-level optimizations such as operator fusion.On Level 1, torch.compile can be slower because runtime overhead is significant for small kernels; backend variability motivates focusing analysis on PyTorch Eager.
  • Evaluation setup: The one-shot baseline uses minimal prompting, greedy decoding at temperature = 0, and in-context examples to test models’ out-of-the-box kernel-generation ability.The prompt permits replacing selected operators, leaving others unchanged, and exploiting fusion opportunities; other experiments use different sampling or feedback settings.

D.2 Kernel Fusion

Generated kernels show isolated successes from fusion and hardware-aware optimization, but overall Level 2 outputs still exhibit insufficient fusion. Iterative feedback can improve performance and correctness, yet some kernels remain persistently incorrect.

  • Level 1: 2.9x speedup was achieved on Level 1 Problem 87 for a GeLU reference in torch.
  • Level 1: 1.3x speedup was achieved on Level 1 Problem 29 for a SoftSign reference in torch.
  • Level 2: Overall Level 2 generated kernels showed insufficient fusion, although most Level 2 problems can be expressed in a single fused kernel.
  • Level 3: 1.9x speedup resulted on Level 3 Problem 49 after fusing scaling, masked fill, and ReLU.The optimization did not fuse other operations.
  • Level 1: 2.8x speedup was achieved on Level 1 Problem 96 for Cosine Similarity Loss.
  • Iterative refinement: 8× speedup from the initial generation was achieved through repeated refinement using compiler, shared-memory, and texture-cache optimizations.The model first parallelized across an output-channel dimension, then used shared memory and ldg texture-cache loads.
  • Iterative refinement: One Level 2 example became functionally correct after execution-feedback refinement, but remained non-performant.
  • Iterative refinement: At a ten-turn budget, DeepSeek-R1 generated functional kernels for over 90% of Level 1 and 2 tasks, while remaining failures were usually functional-incorrectness errors.Execution feedback was especially effective for fixing execution errors, whereas correctness feedback was less granular.

F Few Shot Experiment

Few-shot prompting supplies examples of optimization techniques such as fusion, tiling, recompute, and asynchrony. The resulting kernels improve over one-shot generation in selected workloads, especially through tiling and shared-memory I/O management.

  • Prompting: The few-shot experiment provides examples of fusion, tiling, recompute, and asynchrony during kernel generation.
  • Evaluation: The experiment compares Section 4.1 baselines with few-shot prompting using fast0, fast1, and generated-kernel character length across levels.
  • Matrix multiplication: 77% of Level 1 matrix multiplication problems achieved a speedup over the one-shot baseline through tiling.
  • Matrix multiplication: Level 1 matrix-multiplication performance is compared between the Section 4.1 baseline and few-shot prompting.
  • Level 2: Few-shot Level 2 kernels outperformed PyTorch Eager in problems using aggressive shared-memory I/O management.
  • Level 2: Level 2 performance is compared for problems whose few-shot kernels outperform PyTorch Eager.

G.1 Evaluation across different hardware

KernelBench evaluates generated kernels across NVIDIA GPUs with different micro-architectures and capabilities. It compares fast1 speedups against PyTorch Eager and examines how hardware-specific context affects generation.

  • Hardware setup: KernelBench spans NVIDIA GPUs with different micro-architectures and capabilities, using hardware specifications such as memory, power, TFLOPS, and bandwidth.These specifications are listed in Table 13.
  • Hardware setup: The same kernels are evaluated across hardware, with fast1 measured against the PyTorch Eager baseline on each platform.Table 14 reports cross-hardware comparisons for kernels generated without hardware-specific information.
  • Evaluation: Speedup plots for Levels 1 and 2 use PyTorch Eager as the reference and mark y = 1.0 as the fast1 cutoff.The plots show per-problem speedups across GPUs on a logarithmic scale.
  • Hardware-aware generation: Providing hardware specifications in context produces fast1 comparisons across GPUs, models, and workload levels.Table 15 summarizes speedup over Torch Eager when hardware specifications are supplied during generation.
  • Hardware-aware generation: DeepSeek-R1 begins generating kernels with warp matrix multiply-accumulate instructions when given H100 specifications for Level 1 Problem 1.Figure 10 shows the resulting CUDA kernel.

H.1 Single-shot Experiments: Batched Kernel Generation

KernelBench uses a highly parallelized system to generate, compile, execute, profile, and inspect many kernels. Its iterative platform overlaps computation while isolating GPU execution and handling faulty generations.

  • Single-shot system: KernelBench separates high-volume evaluation into inference, CPU pre-compilation, and GPU evaluation stages.The system parallelizes generation and compilation, while GPU timing evaluates one kernel at a time per device.
  • Iterative refinement: Iterative refinement models kernel generation, pre-compilation, execution, and profiling as finite-state-machine stages driven by environment feedback.Transitions can vary across experiment setups.
  • Iterative refinement: A pipelined multiprocessing design overlaps generation and execution across a node with eight available GPUs.Only kernel execution requires acquiring a GPU, while CPU processes handle independent refinement tasks.
  • Iterative refinement: GPU acquisition uses a separate orchestrator process and semaphores to maximize resource throughput on limited devices.The orchestrator minimizes process control over GPUs during execution and evaluation.
  • Resource management: Kernels are pre-compiled with nvcc on CPUs, and GPU execution and profiling are the only finite-state-machine stages using GPUs.GPU waiting is identified as the orchestrator’s primary bottleneck, motivating maximized device occupancy.
  • Robustness and inspection: The system captures CUDA illegal memory accesses and deadlocks by releasing and spawning processes instead of crashing the orchestrator.KernelBench also provides an interface for visualizing generated kernels and comparing techniques and configurations.
Loading 2502.10517v1…