TongmingLAIC/AKO4ALL

해결하는 문제

AKO4ALL은 GPU 커널 최적화의 번거로운 과정을 자동화합니다. 인간 전문가가 수동으로 프로파일링하고 코드를 재작성하여 성능을 극대화하는 대신, AI 에이전트를 사용해 커널의 속도를 반복적으로 향상시키면서도 정확성을 보장합니다.

작동 방식

이는 Claude Code와 같은 코딩 에이전트를 위한 "스킬"(프로토콜 문서)로 구현되었습니다. GPU 커널이 포함된 디렉터리에서 호출되면, 에이전트는 폐쇄 루프 사이클을 따릅니다: 작업 공간을 생성하고, 베이스라인을 검증하며, NVIDIA Nsight Compute (ncu)를 사용해 코드를 프로파일링하고, 커널을 재작성하며, 결과를 벤치마킹하고, 진행 상황을 로그에 기록합니다. 성능이 정체되면, 에이전트는 새로운 최적화 아이디어를 찾기 위해 웹 검색을 수행하거나, 코드를 재프로파일링하여 새로운 성능 저하 지점을 찾을 수 있습니다. Triton, CUDA, C++, TileLang 등 다양한 언어를 지원합니다.

대상 사용자

각각의 커널마다 대규모 인프라를 구축하지 않고도 최적화 프로세스를 자동화하고 싶은 고성능 GPU 커널을 다루는 개발자 및 연구자에게 적합합니다.

주요 특징

• 전문가 수준의 성능: NVIDIA B200에서 GQA 및 MLA 디코딩/프리필 연산자에 대해 FlashInfer의 전문가 커널보다 빠른 성능 향상을 입증했습니다.
• 언어 독립성: Triton, CUDA, C++, Python 등 다양한 언어로 작성된 커널을 최적화할 수 있으며, 필요 시 언어를 전환해 더 나은 성능을 찾을 수 있습니다.
• 통합 평가: 내장된 평가기 (KernelBench)와 에이전트가 메트릭을 조작하는 것을 방지하는 안티치트 메커니즘을 포함합니다.
• 즉시 통합: Claude Code용 단일 스킬로 작동하며, 최적화 실행을 시작하기 위해 최소한의 설정만 필요합니다.