AutoKernel: GPU 커널 자동 연구 시스템

• Karpathy의 Autoresearch를 GPU 커널 최적화에 적용
• PyTorch 모델을 입력하면 자동으로 Triton 또는 CUDA C++ 커널을 최적화하는 GPU 커널 자동 연구 도구
• 모델의 병목 커널을 프로파일링·추출·최적화·검증하는 완전 자동화 파이프라인 제공
• Amdahl의 법칙 기반 오케스트레이션으로 우선순위를 조정하며, 각 실험은 약 90초 소요되어 밤새 수백 회 반복 가능
• KernelBench와 통합되어 250개 이상 문제에 대해 50~300회 반복 실험을 수행, 단발성 생성보다 체계적 탐색 가능
• Triton과 CUDA C++ 이중 백엔드, 정확성 우선 검증, 단일 파일 수정 구조로 효율적이고 재현 가능한 커널 최적화 지원

작동 원리

• AutoKernel은 PyTorch 모델을 입력받아 다음 단계를 수행
  • 프로파일링으로 GPU 병목 커널 식별
  • 추출하여 독립적인 Triton 또는 CUDA C++ 커널로 분리
  • 자동 최적화 루프를 통해 수정·벤치마크·유지/되돌리기 반복
  • 정확성 검증 후 전체 속도 향상 보고
• program.md에 포함된 지침을 기반으로 kernel.py를 수정하며, bench.py를 통해 5단계 정확성 검사와 roofline 분석을 수행
• 각 실험은 약 90초, 시간당 약 40회, 하룻밤 동안 약 320회 실험 가능

실행 및 구성 요소

• 필수 환경: NVIDIA GPU(H100/A100/RTX 4090), Python 3.10+, uv 패키지
• 주요 스크립트 구성
  • profile.py: GPU 시간 기준 커널 순위화
  • extract.py: 상위 병목 커널 추출
  • bench.py: 5단계 정확성 및 성능 벤치마크
  • orchestrate.py: Amdahl의 법칙 기반 다중 커널 스케줄링
  • verify.py: 전체 모델 검증 및 속도 향상 보고
• program.md는 6단계 최적화 전략, 충돌 처리, 의사결정 프레임워크를 포함해 장시간 자동 실행 가능

지원 커널 및 예제 모델

• 9종 커널 유형 지원: matmul, softmax, layernorm, rmsnorm, flash_attention, fused_mlp, cross_entropy, rotary_embedding, reduce
• 각 커널은 PyTorch 기준 구현(reference.py), Triton/CUDA 초기 버전(kernels/, kernels/cuda/) 포함
• 예제 모델: GPT-2, LLaMA(7B 포함), BERT-base, 사용자 정의 모델
• HuggingFace 모델도 uv sync --extra models로 연동 가능

KernelBench 통합

• KernelBench(Stanford Scaling Intelligence Lab)와 통합되어 AI 생성 GPU 커널의 표준 벤치마크 수행
• AutoKernel은 문제당 50~300회 반복 실험으로 최적화 공간을 체계적으로 탐색
• 주요 도구
  • bridge.py: 문제 로드 및 초기 커널 생성
  • bench_kb.py: 정확성·성능 평가
  • scorer.py: 전체 레벨 점수 계산
  • program_kb.md: KernelBench 전용 에이전트 지침

HuggingFace Hub 내보내기

• 최적화된 커널을 HuggingFace Hub로 내보내고, get_kernel()로 간단히 불러오기 가능
• export_hf.py로 CUDA 커널 내보내기 및 업로드 지원

설계 원칙

• Triton + CUDA C++ 이중 백엔드: Triton은 빠른 반복, CUDA는 최고 성능 제공
• 정확성 우선: PyTorch 결과와 일치하지 않으면 즉시 되돌림
• Amdahl의 법칙 기반 오케스트레이션으로 전체 성능 기여도에 따라 우선순위 결정
• 단일 파일 수정 구조(kernel.py)로 변경 추적과 복구 용이
• TSV 로그(results.tsv)로 실험 결과를 단순·가독성 있게 기록

결과 기록 형식

• results.tsv에 각 실험의 번호, 커널 유형, 처리량(TFLOPS), 지연시간, GPU 피크 대비 비율, PyTorch 대비 속도 향상, 정확성, VRAM 사용량, 설명이 기록됨

프로젝트 배경

• Andrej Karpathy의 autoresearch 개념에서 영감을 받아, LLM 연구용 자율 AI 에이전트 구조를 GPU 커널 최적화에 적용
• KernelBench 통합은 Stanford Scaling Intelligence Lab의 연구를 기반으로 하며, AutoKernel은 단발성 생성 대신 반복적 최적화를 수행
• 프로젝트는 RightNow AI의 Forge 팀이 개발, MIT 라이선스로 공개됨