ScalingIntelligence/KernelBench

LLM의 효율적인 GPU 커널 (kernel) 생성 능력을 평가하기 위한 벤치마크 및 환경

구체적으로, 우리는 LLM이 대상 GPU 상의 PyTorch 프로그램에 대해 정확하고 효율적인 CUDA / DSL 커널을 생성하도록 과제를 부여합니다.

arXiv | 블로그 포스트 | HuggingFace Dataset

최신 안정 버전은 main 브랜치에 제공됩니다. 우리는 지속적으로 저장소 (repo)를 업데이트하고 개선하고 있습니다.

Huggingface 데이터셋은 v0.1로 업데이트되었습니다.

이 저장소는 KernelBench를 위한 핵심 기능과 평가를 위한 사용하기 쉬운 스크립트 세트를 제공합니다. 이 저장소는 이 과제를 해결하는 복잡한 에이전트 스캐폴딩 (agentic scaffolds)을 제공하기 위한 것이 아닙니다. 실험을 위해 이 저장소를 클론 (cloning)하여 수정하거나, git 서브모듈 (submodule)로 사용하는 것을 권장합니다.

우리는 LLM이 PyTorch에 기술된 연산자 (operators)를 원하는 수준의 세분화 (granularity) 단계에 따라 CUDA 커널로 트랜스파일 (transpile) 하도록 문제를 구조화했습니다.

우리는 KernelBench를 4개의 레벨 카테고리로 구성했습니다:

Level 1 🧱: 단일 커널 연산자 (Single-kernel operators) (100개 문제) 신경망의 기초적인 구성 요소 (Convolutions, Matrix multiplies, Layer normalization)
Level 2 🔗: 단순 퓨전 패턴 (Simple fusion patterns) (100개 문제) 퓨전된 커널은 분리된 커널보다 빠를 것입니다 (Conv + Bias + ReLU, Matmul + Scale + Sigmoid)
Level 3 ⚛️: 전체 모델 아키텍처 (Full model architectures) (50개 문제) 전체 모델 아키텍처를 엔드투엔드 (end-to-end)로 최적화 (MobileNet, VGG, MiniGPT, Mamba)
Level 4 🤗: Hugging Face 레벨 - HuggingFace의 전체 모델 아키텍처를 최적화

우리는 cuda 외의 다른 DSL로도 KernelBench를 적극적으로 확장하고 있으며 (아래 참조), AMD GPU 지원도 확장하고 있습니다.

모델이 생성한 커널을 평가하기 위해, 우리는 다음 사항을 확인해야 합니다:

정확한가 ✅: 무작위 입력값에 대해 n_correctness 번 참조(reference) torch 연산자와 비교하여 확인합니다.
성능이 좋은가 ⏱️: 런타임 사이의 속도 향상 (speedup)을 측정하기 위해 n_trial 번 참조 torch 연산자와 비교합니다.

정확성 확인 및 타이밍 구현 방식에 대한 자세한 내용은 src/eval.py를, 평가 및 벤치마킹 가이드라인에 대한 참고 사항은 EVAL.md [WIP]를 확인하십시오.

우리는 단일 샘플 소스 코드를 참조 소스 코드와 비교하여 정확성을 확인하고 속도 향상 (speedup)을 계산할 수 있는 편리한 스크립트 scripts/run_and_check.py를 제공합니다. eval_mode=local 또는 eval_mode=modal을 설정하여 로컬 또는 원격에서 커널을 평가하는 데 사용할 수 있습니다.

정확성과 성능 둘 다 포착해야 하므로, 우리는 다음과 같은 지표 fast_p를 정의합니다:

fast_p: 정확하면서 동시에 임계값 p보다 큰 속도 향상을 보이는 작업의 비율; 속도 향상은 PyTorch 참조 실제 실행 시간 (wall-clock time) 대비 생성된 커널 시간의 비율로 계산됩니다.

속도 향상을 기준으로 필터링하는 이 지표를 설명하는 몇 가지 예시는 다음과 같습니다:

fast_1: LM이 생성한 커널이 정확하면서 PyTorch 베이스라인보다 더 빠른 작업의 비율
fast_2: LM이 생성한 커널이 정확하면서 PyTorch 베이스라인보다 최소 2배 더 빠른 작업의 비율
fast_0: LM이 생성한 커널이 정확한 작업의 비율 (정확도 (correctness rate)와 동일)

작업을 더 어렵게 만들기 위해 속도 향상 임계값 p를 높일 수 있습니다.

전체 벤치마크 성능을 계산하기 위해 scripts/greedy_analysis.py 스크립트를 제공합니다.
정확성과 성능 둘 다 포착해야 하므로, 우리는 지표 fast_p를 사용합니다: 정확하면서 동시에 임계값 p보다 큰 속도 향상을 보이는 작업의 비율; 속도 향상은 PyTorch 참조 실제 실행 시간 (wall-clock time) 대비 생성된 커널 시간의 비율로 계산됩니다.

저장소(repo)를 다음과 같은 구조로 구성했습니다:

KernelBench/
├── assets/
├── KernelBench/ # 벤치마크 데이터셋 파일
...

우리는 의존성 관리 (dependency management)를 위해 pyproject.toml 및 uv를 사용하는 방식으로 전환했습니다. 아직 설치하지 않았다면 uv를 설치하십시오.

# 기본 의존성 설치 (로컬 GPU 없이도 작동)
uv sync
# AMD ROCm 백엔드로 설치 (ROCm>=7.1 필요)
...

AMD GPU 즉 ROCm 백엔드 (ROCm>=7.1)의 경우, uv remove torch && uv add torch --index pytorch=https://download.pytorch.org/whl/rocm7.1를 추가해 주세요.

의존성 (dependencies)을 구성하기 위해 RoCm 호환 PyTorch를 사용하려면 uv add torch --index pytorch=https://download.pytorch.org/whl/rocm7.1를 추가해 주세요. ROCm 설정의 복잡성 때문에 Docker 이미지에서 실행하는 것을 권장합니다.

여전히 conda (python=3.10)를 사용하여 환경을 생성하고 requirements.txt로 의존성을 설치할 수 있습니다.

API 호출을 위해 litellm을 사용합니다. 우리의 .env.example을 따라 .env 파일을 생성하여 키를 설정해 주세요.

커널 (kernels) 실행 및 프로파일링 (profiling)에는 GPU가 필요합니다. 로컬에서 사용 가능한 GPU가 없는 경우, 클라우드 서버리스 (serverless) GPU 평가를 위해 Modal을 설정할 수 있습니다. 계정을 생성한 후 modal token new를 실행하여 Modal 토큰을 설정하세요. 그런 다음 generate_and_eval_single_sample_modal.py 스크립트를 사용하면 됩니다.

Google Colab을 통해 우리의 튜토리얼 노트북 (또한 notebooks/tutorial.ipynb에 있음)을 시도해 볼 수도 있습니다.

단일 문제로 시작하는 것이 더 쉽습니다. 이는 문제를 가져오고, 샘플을 생성하며, 샘플을 평가합니다.

# 예를 들어, huggingface에서 level 2의 40번 문제를 실행하고 생성을 위해 google gemini 2.5 flash를 사용합니다.
uv run python scripts/generate_and_eval_single_sample.py dataset_src=huggingface level=2 problem_id=40 server_type=google model_name=gemini/gemini-2.5-flash
# dataset_src는 "local" 또는 "huggingface"가 될 수 있습니다.
...

수정해야 할 수도 있는 사항

• 사용 중인 GPU에 따라 하드웨어를 반영하도록 gpu_arch 인자를 조정해야 할 수도 있습니다.
• precision=fp32를 통해 텐서 (tensor)의 정밀도 (precision)를 지정할 수 있습니다. 현재 보고된 모든 결과는 fp32이지만, fp16 및 bf16에 대한 지원을 추가했습니다.
• 또한 backend=cuda 이외의 다른 GPU 프로그래밍 언어도 지원합니다. 예를 들어, 간단히 backend=triton 또는 backend=hip을 지정하면 됩니다. 현재 우리는 프로그래밍 프레임워크 및 DSL을 사용하는 NVIDIA GPU를 지원합니다: cuda, triton, cute, tilelang, thunderkittens.

AMD GPU 참고 사항: hip 백엔드를 사용하세요. 현재 지원되는 gpu_arch는 gfx942, gfx950입니다.

ThunderKittens (TK)를 로컬에 설정하는 방법에 대한 참고 사항: backend=thunderkittens를 사용하려면

ThunderKittens (TK) 레포지토리를 git clone하고, 로컬 ThunderKittens 디렉토리를 가리키도록 다음과 같은 환경 변수를 설정해야 합니다: export THUNDERKITTENS_ROOT=<ThunderKittens 폴더 경로>

또한, 예시에 표시된 것과 같이 모든 ThunderKitten 프로그램은 tk_root = os.environ.get("THUNDERKITTENS_ROOT", "/root/ThunderKittens")를 포함해야 하며, 이를 통해 커널 (kernel)이 올바른 TK 프리미티브 (primitives)를 포함할 수 있게 됩니다. 추가로, 현재 TK는 BF16만 지원합니다.

포괄적인 옵션 세트를 확인하려면 설정 (config) 필드를 확인하십시오. 기본적으로 모델에 최소한의 정보와 함께 원샷 예시 (one-shot example)를 제공한다는 점에 유의하십시오. 다른 프롬프트 (prompt) 설정을 확인하거나 src/prompt_constructor_toml.py에서 직접 구성할 수 있습니다.

# 1. 응답을 생성하고 커널을 runs/{run_name} 디렉토리에 로컬로 저장
uv run python scripts/generate_samples.py run_name=test_hf_level_1 dataset_src=huggingface level=1 num_workers=50 server_type=deepseek model_name=deepseek-chat temperature=0
# 2. runs/{run_name} 디렉토리에 생성된 모든 커널에 대해 평가
...

성공률 (success rate), 타이밍 지표 (timing metric), 그리고 전반적인 벤치마크 성능인 fast_p를 계산하기 위해 평가 결과를 분석하는 scripts/benchmark_eval_analysis.py를 제공합니다.

uv run python scripts/benchmark_eval_analysis.py run_name=test_hf_level_1 level=1 hardware=L40S_matx3 baseline=baseline_time_torch

다른 하드웨어를 사용하는 경우, scripts/generate_baseline_time.py 스크립트로 베이스라인 시간 (baseline time)을 생성할 수 있습니다.

results/timing에 여러 세대에 걸친 다양한 NVIDIA GPU에 대한 몇 가지 참조용 베이스라인 시간을 제공하지만, 더 정확한 결과(클러스터 전력, 소프트웨어 버전 등이 모두 타이밍 결과에 영향을 미침)를 위해 직접 베이스라인 시간을 생성하는 것을 권장합니다. 자세한 내용은 results/timing/README.md를 참조하십시오.

우리가 추가할 계획인 기능들은 로드맵 (roadmap)을 확인해 보십시오. 이러한 방향에 대한 커뮤니티의 기여와 논의를 환영합니다.

또한 여러분의 프로젝트를 위해 KernelBench를 라이브러리로 사용할 수도 있습니다. 예를 들어:

from kernelbench import timing

, from kernelbench import eval as kb_eval

, 또는 from kernelbench.utils import set_gpu_arch

를 사용할 수 있습니다.

Harbor를 이용한 어댑터 (Adapter with Harbor)— 모델의 pass@1/k를 넘어 에이전트 성능 (agentic performance)에 대한 더 높은 처리량의 평가 (eval)와 더 풍부한 평가를 가능하게 하기 위해 Harbor와 통합하고 있습니다.(진행 중)

멀티턴 / 테스트 시간 스케일링 (Multi-Turn / Test-Time Scaling)— Caesar는 우리의 처리량 중심 멀티턴 추론 엔진 (ICML '25)으로, 논문의 반복적 정제 (iterative refinement) 실험에 사용되었습니다. 이는 생성 궤적 (generation trajectories)을 배치 (batch)로 실행하며, 순차적인 테스트 시간 스케일링 (test-time scaling)을 위해 턴(turn) 전반에 걸쳐 정확성, 실행 시간 및 프로파일링 (profiling) 신호를 피드백합니다.

강화학습 (Reinforcement Learning, RLVR)— kernelbench-tinker는 Thinking Machines Lab의 Tinker RL 라이브러리와의 엔드 투 엔드 (end-to-end) 통합입니다. 이 파이프라인은 정책 모델 (policy model)이 커널을 생성하게 하고, Modal을 통해 클라우드 GPU에서 이를 평가하며, 결과를 RL 보상 (rewards)으로 변환합니다 — 이는 GPU 커널 최적화에서 RLVR을 실험하기 위한 최소한의 플레이그라운드 (playground)입니다.

진화 탐색 (Evolutionary Search)— AlphaEvolve와 같은 진화 탐색은 최적화 문제에 대한 혁신적인 솔루션을 발견하는 데 유망함을 보여주었습니다. 우리는 OpenEvolve를 위한 통합 작업을 진행 중입니다. 곧 출시될 예정입니다.

루프라인 / 최대 가속 분석 (Roofline / Max Speedup Analysis)—(실험적, 작업 진행 중) simple-torchroofline은 PyTorch 프로그램에 대한 분석적 루프라인 (roofline) 분석을 제공하여, 대상 GPU에 대한 이론적 한계치 (speed-of-light, SoL) 연산 및 메모리 경계를 추정합니다 — 하드웨어가 필요하지 않습니다. 경험적 루프라인 분석을 위한 하드웨어 카운터 기반 프로파일링과 결합하면, 보고된 가속 (speedup)이 물리적으로 현실적인지 확인하는 데 도움이 됩니다.

출시 이후, 우리는 이 방향을 탐구하기 위해 KernelBench를 사용하는 연구자, 연구실 및 기업들로부터 많은 관심을 받아왔습니다. 우리는 KernelBench의 알려진 사용 사례와 자동화된 커널 생성 (automated kernel generations)을 향한 관련 노력들을 문서화했습니다. 만약 여러분이 KernelBench를 사용하고 있다면, 그에 대해 더 자세히 듣고 싶습니다!

면책 조항 (Disclaimer): KernelBench는 오픈 소스 (open-source) 평가 프레임워크 및 툴킷으로 설계되었습니다. KernelBench 팀은 개별 커널 (kernel) 또는 보고된 결과를 검토, 검증 또는 보증하지 않습니다. 사용자는 프레임워크를 사용하여 얻은 모든 결과를 독립적으로 검증할 책임이 있습니다. 커널 (kernel)의 벤치마킹 및 평가에 관한 더 자세한 안내는 EVAL.md를 확인해 주세요.

MIT. 자세한 내용은 LICENSE.md를 확인하세요.

@misc{ouyang2025kernelbenchllmswriteefficient,
title={KernelBench: Can LLMs Write Efficient GPU Kernels?},
author={Anne Ouyang and Simon Guo and Simran Arora and Alex L. Zhang and William Hu and Christopher Ré and Azalia Mirhoseini},
...

이 프로젝트를 가능하게 해준 GPU Mode, PyTorch, Modal Labs 및 더 넓은 오픈 소스 커뮤니티의 지원에 감사드립니다.