AutoMegaKernel: 자가 재타겟팅 Megakernel 합성을 위한 정적 검증 에이전트 하네스

AutoMegaKernel (AMK)은 HuggingFace Llama 계열 모델을 단일 지속적 협력 CUDA 커널(single persistent cooperative CUDA kernel)로 컴파일하여, 모델별로 수동 작성된 CUDA 없이 단 한 번의 런칭(launch)으로 전체 순전파(forward pass)를 실행합니다. 본 연구의 기여점은 순수 속도가 아닌 시스템 자체에 있습니다. 고정된 스케줄-IR 검증기(schedule-IR validator)는 정적 그래프 검사(static graph checks, 기계적 증명 아님)를 통해 데드락 프리(deadlock-freedom)와 레이스 프리(race-freedom)를 정적으로 인증하므로, 에이전트가 제안한 안전하지 않은 스케줄은 실행 전에 거부됩니다. 7,160개의 적대적 스케줄(이 중 6,091개는 안전하지 않음)을 대상으로 테스트했을 때, 잘못된 수락(false-accepts)은 0건이었으며 360개의 실제 로워링(lowering)은 모두 수락되었습니다. 동일한 소스로부터 하나의 코드베이스로 sm_80/sm_90/sm_120을 재타겟팅(retarget)할 수 있으며, 지원되는 10개 모델 모두에 대해 정확한 메가커널(megakernel)을 자동 생성합니다. 실제 SmolLM2-135M 체크포인트에서는 HuggingFace의 그리디 디코딩(greedy decode) 결과와 토큰 단위로 일치하는 결과를 재현했습니다 (perplexity 일치도 2.5e-7). 무인 상태로 에이전트가 구동하는 자동 연구 루프(autoresearch loop)는 자체 베이스라인 대비 메가커널을 스스로 개선합니다 (1.25-1.72x). 탐색을 통해 발견된 int8 (W8A16) 메가커널은 NVIDIA의 데이터센터 추론 플릿(inference fleet)에서 배치-1(batch-1) 디코딩 시 CUDA-graphed cuBLAS bf16보다 성능이 뛰어납니다: L4는 최대 1.33x, 현세대 L40S는 1.25-1.27x, A10G는 규모에 따라 최대 1.08x, 소비자용 RTX 5090은 1.19-1.23x의 성능 향상을 보였습니다. 성능 순위는 대역폭(bandwidth)의 단순한 함수가 아닙니다 (864 GB/s인 L40S가 600 GB/s인 A10G보다 빠름). 그 차이는 추론급(inference-class)과 학습급(training-class)의 차이에서 기인합니다. AMK는 대역폭이 높은 학습급인 A100/H100에서는 cuBLAS에 뒤처지며, 이 경우 하네스(harness)가 교차 SM 동기화(cross-SM-sync) 병목 현상을 국소화합니다. 저희는 이 격차를 있는 그대로 보고합니다. 이는 디코딩 위치 0에서의 정밀도 비대칭(precision-asymmetric, W8A16 vs bf16) 비교이며, 가장 큰 실제 체크포인트는 TinyLlama-1.1B입니다. 코드 및 하네스: https://github.com/RightNow-AI/AutoMegaKernel