RightNow-AI/autokernel

해결하는 문제

빠른 GPU 커널(신경망을 실행하는 저수준 코드)을 작성하는 것은 시간이 많이 걸리고 깊은 전문 지식이 필요합니다. AutoKernel은 코딩 AI 에이전트를 모델에 투입하여 이를 자동화합니다. 가장 느린 부분을 찾아 별도의 커널로 추출하고, 더 빨라질 때까지 최적화를 계속 시도합니다. 각 작업을 수동으로 튜닝하는 대신, 모델을 제공하기만 하면 측정된 속도 향상 수치가 포함된 최적화된 커널을 돌려받을 수 있습니다.

작동 방식

PyTorch 모델을 제공하면 AutoKernel은 다음 파이프라인을 실행합니다:

1. Profile(프로파일링): PyTorch의 프로파일러를 사용하여 가장 많은 시간을 소비하는 GPU 커널의 순위를 매깁니다.
2. Extract(추출): 각 병목 현상을 독립적인 Triton 또는 CUDA C++ 커널 파일로 추출합니다.
3. Optimize(최적화): AI 코딩 에이전트(Claude 또는 Codex 등)가 일련의 지침(program.md)을 읽고, 커널 파일을 편집하며, 고정된 벤치마크를 실행합니다. 정확성 검사를 통과하고 속도가 향상된 경우에만 변경 사항을 유지합니다. 이 루프를 계속 반복하며, 각 실험은 약 90초가 소요됩니다.
4. Verify(검증): 모든 병목 현상을 최적화한 후, 새로운 커널을 모델에 다시 삽입하고 전체적인 엔드투엔드 속도 향상 수치를 보고합니다.

벤치마크에는 성능 측정 전 5단계 정확성 검사(스모크 테스트, 형상 스윕, 수치 안정성, 결정론적 동작, 엣지 케이스)가 포함됩니다. 오케스트레이터는 암달의 법칙을 사용하여 전체 모델에 가장 큰 영향을 미치는 커널을 우선적으로 처리합니다.

대상 사용자

PyTorch 모델의 GPU 추론이나 학습 속도를 높이고 싶지만 CUDA나 Triton 커널을 직접 작성하고 싶지 않은 개발자 및 연구원. NVIDIA GPU(H100/A100/RTX 4090에서 테스트됨)와 Python 3.10+가 필요합니다. 즉시 실행 가능한 모델 정의(GPT-2, LLaMA, BERT)가 포함되어 있어 외부 모델 라이브러리 없이도 바로 시도해 볼 수 있습니다.

주요 특징

• 자율 최적화: 밤새 감독 없이 수백 번의 실험을 실행하며, 파일을 편집하고 성공적인 변경 사항만 유지합니다.
• 듀얼 백엔드: Triton(빠른 반복)과 네이티브 CUDA C++(최대 성능) 커널을 모두 생성하며 동일한 인터페이스를 공유합니다.
• 9가지 커널 유형: matmul, softmax, layernorm, flash attention, fused MLP 등 일반적인 딥러닝 연산을 다룹니다.
• KernelBench 통합: AI 생성 GPU 커널을 위한 표준 벤치마크와 연동되어 일회성 추측이 아닌 반복적인 정교화 과정을 거칩니다.
• 정확성 우선: 5단계 정확성 테스트를 통과하지 못한 커널은 즉시 이전 상태로 되돌려지므로 속도 향상이 항상 안전합니다.
• 간편한 출력: 모든 실험은 일반 results.tsv 파일에 기록되어 쉽게 추적할 수 있습니다.