Bigstones

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

5 min read

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

출처: GitHub - RightNow-AI/AutoKernel

📌 핵심 요약

  • Andrej Karpathy 의 Autoresearch 개념을 GPU 커널 최적화에 적용한 도구.
  • PyTorch 모델을 입력하면 Triton 또는 CUDA C++ 커널을 자동으로 최적화.
  • 모델의 병목 커널을 프로파일링 → 추출 → 최적화 → 검증 하는 완전 자동화 파이프라인 제공.
  • Amdahl 의 법칙 기반 오케스트레이션으로 우선순위 조정, 하룻밤에 수백 회 실험 가능.
  • KernelBench 와 통합되어 체계적인 최적화 공간 탐색 지원.

⚙️ 작동 원리 (Workflow)

AutoKernel 은 PyTorch 모델을 입력받아 다음 단계를 자동으로 수행합니다.

  1. 프로파일링 (Profiling): GPU 시간 기준 병목 커널 식별 (profile.py).
  2. 추출 (Extraction): 상위 병목 커널을 독립적인 Triton/CUDA 커널로 분리 (extract.py).
  3. 자동 최적화 루프:
    • program.md 지침에 따라 kernel.py 수정.
    • bench.py 를 통해 5 단계 정확성 검사 및 roofline 분석 수행.
    • 수정 → 벤치마크 → 유지/되돌리기 반복.
  4. 검증 및 보고: 정확성 검증 후 전체 모델 속도 향상 보고 (verify.py).

💡 효율성: 각 실험은 약 90 초 소요 (시간당 약 40 회, 밤새 약 320 회 실험 가능).

🧰 주요 구성 요소 (Scripts)

스크립트역할
profile.pyGPU 시간 기준 커널 순위화
extract.py상위 병목 커널 추출
bench.py5 단계 정확성 및 성능 벤치마크
orchestrate.pyAmdahl 의 법칙 기반 다중 커널 스케줄링
verify.py전체 모델 검증 및 속도 향상 보고
program.md6 단계 최적화 전략, 충돌 처리, 의사결정 프레임워크 포함

🧩 지원 커널 및 예제 모델

지원 커널 (9 종)

  • matmul, softmax, layernorm, rmsnorm, flash_attention, fused_mlp, cross_entropy, rotary_embedding, reduce
  • 각 커널은 PyTorch 기준 구현 (reference.py) 과 Triton/CUDA 초기 버전 포함.

예제 모델

  • GPT-2, LLaMA (7B 포함), BERT-base, 사용자 정의 모델
  • HuggingFace 모델 연동 가능 (uv sync --extra models)

📊 KernelBench 통합

  • Stanford Scaling Intelligence Lab 의 KernelBench 와 통합.
  • AI 생성 GPU 커널의 표준 벤치마크 수행.
  • 문제당 50~300 회 반복 실험 을 통해 단발성 생성이 아닌 체계적 탐색 가능.

KernelBench 도구

  • bridge.py: 문제 로드 및 초기 커널 생성
  • bench_kb.py: 정확성·성능 평가
  • scorer.py: 전체 레벨 점수 계산
  • program_kb.md: KernelBench 전용 에이전트 지침

📤 HuggingFace Hub 내보내기

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

🏗️ 설계 원칙

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

📝 프로젝트 배경

  • 영감: Andrej Karpathy 의 autoresearch 개념 (LLM 연구용 자율 AI 에이전트 구조).
  • 기반: KernelBench 통합은 Stanford Scaling Intelligence Lab 연구 기반.
  • 개발: RightNow AI 의 Forge 팀 개발.
  • 라이선스: MIT 라이선스 공개.

🔗 관련 태그

GPU 커널최적화 Triton CUDA PyTorch AutoResearch AI 성능개선 RightNowAI

Graph View