KernelBench: Can LLMs Write Efficient GPU Kernels?

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.

한 줄 요약

KernelBench는 LLM이 효율적인 GPU 커널을 생성할 수 있는지 평가하는 오픈소스 프레임워크로, 250개 PyTorch 워크로드에서 기능적 정확성과 속도 향상을 측정한다.

핵심 기여도

**KernelBench**: 250개 PyTorch ML 워크로드를 기반으로 LLM이 생성한 GPU 커널의 성능과 정확성을 평가하는 오픈소스 프레임워크를 제시.
**fast_p 메트릭**: 기능적으로 정확하고 기준 속도 향상(p) 이상을 달성한 커널 비율을 측정하는 새로운 평가 지표를 도입.
**실험 결과**: 최신 추론 모델(o1, DeepSeek-R1)이 PyTorch 베이스라인과 일치하는 경우는 20% 미만이며, 피드백을 활용한 반복적 개선 시 43%까지 향상 가능.
**CUDA 데이터 부족 문제**: CUDA는 오픈소스 코드 코퍼스에서 0.073%에 불과하여 모델 학습에 제한이 있음.

핵심 아이디어

KernelBench는 AI 엔지니어의 실제 작업 흐름을 반영하여 LLM이 GPU 커널을 생성하고 최적화하는 능력을 평가하는 환경을 구축했다. 모델은 PyTorch 기준 코드를 입력으로 받아, 어떤 연산을 최적화할지, 어떻게 최적화할지를 스스로 결정할 수 있다. 이는 컴파일러 피드백, 프로파일링 정보, 하드웨어 특성 등을 포함한 다양한 정보를 모델에 제공하는 것을 가능하게 한다.

핵심 통찰은 LLM이 단순히 코드를 생성하는 것이 아니라, **하드웨어 효율성**(예: 테일링, 퓨전)과 **알고리즘 최적화**(예: 스파스성 활용)를 고려한 커널을 생성할 수 있는지 여부이다. 그러나 CUDA는 훈련 데이터에서 상대적으로 희소(0.073%)하여, 모델이 정확한 커널을 생성하는 데 어려움이 있다.

기술적 접근법

**KernelBench 환경**: PyTorch 기준 코드를 입력으로 받아, LLM이 커널을 생성하고, 정확성과 성능을 자동으로 평가.
**fast_p 메트릭**: 기능적으로 정확하고 기준 속도 향상(p) 이상을 달성한 커널 비율을 측정.
**평가 방식**: 5개의 랜덤 입력으로 정확성을 검증, 반복 실행으로 성능을 평가.
**실험 모델**: OpenAI-o1, DeepSeek-R1 등 최신 추론 모델과, 피드백을 활용한 반복적 최적화 실험.
**하드웨어 정보 활용**: TFLOP, 대역폭, 테일링, 퓨전 기법 등 하드웨어 최적화 정보를 모델에 조건으로 제공.

주요 결과

**PyTorch 기준 대비 성능**: OpenAI-o1과 DeepSeek-R1은 PyTorch 기준과 일치하는 경우가 20% 미만.
**반복적 피드백 활용 시 향상**: fast_1 기준에서 12% → 43%, 36% → 72%, 12% → 18%로 향상.
**정확성 문제**: 모델은 실행 오류는 줄일 수 있지만, 기능적으로 정확한 코드를 생성하는 데 어려움.
**CUDA 데이터 부족**: CUDA는 The Stack v1.2에서 0.073%만 포함되어 있어, 모델 학습에 제한이 있음.

의의 및 한계

KernelBench는 LLM이 실제 AI 엔지니어링 환경에서 커널을 생성하고 최적화하는 능력을 평가하는 첫 번째 프레임워크로, **성능 향상과 에너지 절감**에 직접적인 영향을 미친다. 또한, PyTorch 기반으로 하드웨어 플랫폼에 따라 유연하게 평가 가능하며, **fast_p 메트릭은 시간이 지남에 따라 더 진보한 기준에 맞춰 조정 가능**하다.

하지만, **CUDA는 훈련 데이터에서 희소**하여 모델이 정확한 커널을 생성하는 데 어려움이 있으며, **복잡한 최적화**(예: 텐서 코어 wmma)를 시도하면 오류가 증가하는 trade-off가 존재한다. 또한, **기능적 정확성**(functional correctness)은 여전히 큰 도전 과제이다.

실용적 활용

KernelBench는 **AI 하드웨어 개발자**, **모델 최적화 엔지니어**, **LLM 기반 코드 생성 연구자**에게 유용한 평가 도구로 활용될 수 있다. 특히, **새로운 GPU 플랫폼에 맞춘 커널 최적화**, **LLM이 생성한 커널의 성능과 정확성 검증**, **하드웨어 특화 최적화 기법 연구**에 적용 가능하다.