🤖 machine learning

AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis

AutoMegaKernel은 자기 개선 에이전트 루프를 통해 추론급 GPU에서 베이스라인 대비 최대 1.72배의 속도 향상을 달ras면서, 다양한 NVIDIA 아키텍처 전반에 걸쳐 데드락 없는 실행을 보장하며 HuggingFace Llama 계열 모델을 위한 단일 지속성 CUDA 메가커널을 자동으로 합성하는 정적 검사 기반의 에이전트 하네스입니다.

원저자: Jaber Jaber, Osama Jaber

게시일 2026-06-09

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Jaber Jaber, Osama Jaber

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 거대한 규모의 복잡한 공장(대규모 언어 모델)을 운영하여 단 하나의 제품(텍스트 한 단어)을 생산하려고 한다고 상상해 보십시오.

과거의 방식 (표준 AI):
현재 이 공장을 운영하는 것은 작업자를 창고로 보내 특정 도구를 가져오게 하고, 이를 다시 조립 라인으로 가져와 아주 작은 작업을 수행한 뒤, 다음 도구를 위해 다시 작업자를 창고로 보내는 과정을 단어 하나를 만들 때마다 수십 번 반복하는 것과 같습니다.

문제점: 작업자는 만드는 시간보다 창고를 왔다 갔다 하는 데(메모리에서 데이터를 기다리는 시간) 대부분의 시간을 소비합니다. 이를 "대역폭 제한(bandwidth-bound)"이라고 합니다. 공장은 얼마나 빨리 만드느냐가 아니라, 작업자가 얼마나 빨리 걷느냐에 의해 제한됩니다.
비용: 작업자가 도구를 가져오기 위해 멈출 때마다 미세한 지연(실행 지연, launch latency)이 발생합니다. 매 단계마다 이 과정이 더해지면 전체적인 속도가 느려집니다.

새로운 방식 (AutoMegaKernel):
이 논문의 저자들은 **AutoMega-Kernel (AMK)**이라는 시스템을 구축했습니다. AMK는 규칙을 완전히 바꾸는 혁신적인 새로운 공장 관리자라고 생각하면 됩니다.

1. "원샷(One-Shot)" 공장

작업자를 왔다 갔다 하게 하는 대신, AMK는 단 **한 번의 실행(single launch)**으로 전체 작업을 완료하도록 공장을 구성합니다.

비유: 작업자가 도구를 하나씩 가져오는 대신, 모든 기계에 작업자들이 배치되어 있다고 상상해 보십시오. 그들은 모두 직접적인 내부 파이프라인으로 연결되어 있습니다. "시작" 버튼을 누르면, 아무도 바닥을 떠나 창고로 갈 필요 없이 전체 조립 라인이 처음부터 끝까지 하나의 연속적인 움직임으로 실행됩니다.
결과: 이는 "걷는 시간"과 "멈춤-시작"의 지연을 제거하여, 이론적으로 프로세스를 훨씬 빠르게 만듭니다.

2. "안전 검사관" (정적 검증 에이전트)

모든 사람이 동시에 작업하는 공장을 짓는 것은 위험합니다. 만약 작업자들이 완벽하게 조율되지 않는다면, 서로 충돌하거나(경쟁 상태, race conditions) 도구가 도착하기만을 영원히 기다리는 상황(데드락, deadlocks)이 발생할 수 있습니다.

혁신: 보통 이를 구현하려면 인간 전문가가 코드를 수동으로 작성해야 하는데, 이는 매우 어렵고 오류가 발생하기 쉽습니다. AMK는 AI 에이전트를 사용하여 공장 레이아웃을 자동으로 설계합니다.
안전망: 공장이 가동되기 전, **변경 불가능한 고정된 안전 검사관(검증기)**이 AI의 설계도를 점검합니다. 이 검사관은 다음 사항을 수학적으로 증명합니다:
- 아무도 영원히 기다리며 갇히지 않을 것임.
- 두 명의 작업자가 동시에 같은 도구를 사용하려 하지 않을 것임.
- 작업 순서가 완벽함.
주장: 저자들은 이 검사관을 7,160개의 "까다롭거나" 결함이 있는 설계도에 대해 테스트했습니다. 검사관은 모든 위험한 설계도를 잡아내어 공장이 시작되기 전에 거부했습니다. 단 한 번도 위험한 계획을 통과시킨 적이 없습니다.

3. "만능 어댑터" (자기 재타겟팅, Self-Retargeting)

보통 특정 건물(특정 컴퓨터 칩)에 맞춰 설계된 공장은 다른 건물에서는 작동하지 않습니다.

마법: AMK는 "자기 재타겟팅"이 가능합니다. 설계도를 한 번만 작성하면, 시스템이 인간이 코드를 새로 쓸 필요 없이 다양한 유형의 컴퓨터 칩(예: RTX 5090, A100, H100)에서 완벽하게 작동하도록 자동으로 적응합니다.

4. "경량화" 버전 (양자화, Quantization)

이 시스템은 "더 가벼운" 도구를 사용하여 공장을 운영하는 방법도 찾아냈습니다.

비유: 작업자들이 무거운 풀사이즈 도구(고정밀 수학)를 들고 가는 대신, 가볍고 컴팩트한 도구(int8 또는 int4 정밀도)를 사용합니다.
이점: 도구가 가벼워지면 한 번에 더 많은 도구를 운반할 수 있고, 작업자들도 더 빠르게 움직일 수 있습니다.
- Int8: 무거운 도구와 똑같이 잘 작동하면서도(무손실) 12% 더 빠릅니다.
- Int4: 훨씬 더 빠르게 움직이지만(걷는 시간을 2배 이상 단축), 최종 결과물이 약간 덜 완벽할 수 있습니다(손실 발생). 하지만 여전히 읽을 수 있는 수준입니다.

5. 솔직한 결과 (현실적인 분석)

저자들은 이 시스템이 어디에서 승리하고 어디에서 패배하는지에 대해 매우 투명하게 밝히고 있습니다.

승리하는 지점: "추론급(Inference-Class)" 칩(L4, L40S, A10G 및 소비자용 RTX 5090 등)에서, AMK 시스템은 단어를 하나씩 생성할 때 현재 업계 표준(cuBLAS)보다 빠릅니다. 경쟁 모델보다 1.1배에서 1.3배 더 빠릅니다.
- 이유: 현재의 시스템들을 괴롭히는 "왔다 갔다 하는" 지연을 성공적으로 제거했기 때문입니다.
패배하는 지점: 거대하고 빠른 "학습급(Training-Class)" 칩(A100 및 H100 등)에서는 현재 버전이 기존 표준 시스템보다 더 느립니다.
- 이유: 작업자들 사이의 "안전 조율(동기화)"에 약간의 시간이 소요됩니다. 가장 빠른 칩에서는 이 미세한 지연이 병목 현상이 됩니다. 저자들은 이를 솔직하게 인정합니다: "우리는 병목 지점을 찾아냈으며, 이에 대해 정직하게 밝힙니다."

요약

AutoMegaKernel은 AI 모델을 실행하기 위한 "원샷" 공장을 설계하는 데 AI 에이전트를 사용하는 시스템입니다. 이 시스템은 설계가 절대 실패하지 않도록 보장하는 엄격한 안전 검사관을 포함하고 있습니다. 또한 다양한 컴퓨터 칩에서 자동으로 작동하도록 설계되었습니다.

핵-이득: 현재 소비자용 및 데이터센터용 칩에서 텍스트를 한 단어씩 생성하는 데 있어 가장 빠른 방식 중 하나이며, 불필요한 지연을 제거함으로써 기존 도구들을 앞질렀습니다.
한계: 아직 완벽하지는 않습니다. 절대적인 속도를 자랑하는 슈퍼컴퓨터급 칩에서는 조율 오버헤드로 인해 기존의 거물들보다 약간 느리지만, 저자들은 이 시스템이 안전하고 정확하며 스스로 개선될 수 있음을 입증했습니다.

기술 요약: AutoMegaKernel (AMK)

문제 정의

배치 크기 1에서의 단일 스트림 자기회귀(autoregressive) 디코딩은 근본적으로 대역폭 제한(bandwidth-bound)을 받습니다. 지연 시간의 하한선은 총 가중치 바이트와 HBM 대역폭의 비율( $t_{min} = \text{weight\_bytes} / \text{HBM\_bandwidth}$ )에 의해 결정됩니다. 기존 실행 프레임워크(예: PyTorch, cuBLAS)는 모든 연산마다 별도의 커널을 실행하기 때문에, 상당한 CPU 런칭 오버헤드를 유발하고 연산 사이에 활성화 값(activations)이 HBM을 거쳐 왕복하도록 강제하여 이 하한선에 도달하지 못합니다. CUDA Graphs가 런칭 오버헤드는 상쇄할 수 있지만, 연산 간의 경계와 HBM 왕복 문제는 여전히 유지됩니다.

기존의 "메가커널(megakernel)" 접근 방식(예: Mirage Persistent Kernel)이나 수작업으로 최적화된 커널들은 전체 순전파(forward pass)를 하나의 지속적(persistent) 커널로 융합하여 지연 시간 하한선 문제를 해결합니다. 그러나 이러한 시스템들은 **신뢰성(trust)**과 **이식성(portability)**이 부족합니다. 이들은 종종 정적 안전 게이트(static safety gates)를 누락하여 데드락(deadlock)이나 레이스 컨디션(race condition)의 위험을 초래하며, 특정 모델이나 아키텍처를 수동으로 타겟팅해야 하고, 자율 에이전트가 시스템을 안전하게 최적화할 수 있는 구조화된 인터페이스를 제공하지 않습니다.

방법론

AutoMegaKernel (AMK)은 수작업으로 작성된 CUDA 없이 HuggingFace Llama 계열 모델을 단일 지속 협력 커널(single persistent cooperative kernel)로 컴파일하는 자율 에이전트 기반 시스템입니다. 이 시스템은 4개의 레이어와 2개의 오토리서치(autoresearch) 루프로 구성됩니다:

Layer 0 (The VM): 수작업으로 작성되고 고정된, 아키텍처별 지속 메가커널(vm/scheduler.cu)입니다. 이는 각 Streaming Multiprocessor (SM)당 하나의 협력적 스레드블록(cooperative threadblock)을 실행합니다. 각 블록은 로컬 큐를 소유하며, wait → dispatch → signal 명령어를 실행하며 전역 토폴로지 순서를 따라갑니다. 동기화를 위해 락(lock) 대신 단조 증가 카운터(monotonic counters)에 의존합니다.
Layer 1 (Instructions): 격리된 상태에서 정확성이 검증된 ABI 준수 마이크로 커널(GEMV, attention, RMSNorm 등)입니다.
Layer 2 (Scheduler): 에이전트가 수정할 수 있는 표면(edit surface)입니다. 모델을 임포트하여 타입이 지정된 Schedule IR(단조 증가 카운터로 동기화된 task-DAG)로 낮추고, ScheduleConfig 객체를 통해 구조적 구성(타이링, 융합, SM 할당, 파이프라이닝 깊이 등)을 제안합니다.
Layer 3 (Roadmap): 미래의 역동성(예: continuous batching)을 위한 플레이스홀더입니다.

안전 게이트 (The Safety Gate):
핵심 혁신은 정적 검사 검증기(schedule/ir.py의 validate())입니다. 어떤 GPU 런칭 이전에, 검증기는 Schedule IR에 대해 일련의 정적 그래프 체크를 수행하여 다음을 인증합니다:

데드락 프리(Deadlock-freedom): 프로듀서-컨슈머 DAG가 비순환(acyclic)인지, 대기 임계값( $1 \le t \le \#producers$ )이 충족 가능한지, 그리고 per-SM 큐 순서가 유효한지 확인합니다.
레이스 프리(Race-freedom): 공유 카운터가 올바르게 조인되는지(모든 프로듀서가 완료된 후 컨슈머가 멀티 프로듀서 카운터를 기다려야 함), 그리고 메모리 액세스가 happens-before 관계를 준수하는지 확인합니다.
정확성(Correctness): 출력 도달 가능성 및 참조 무결성을 검증합니다.

만약 스케줄이 이러한 검사를 통과하지 못하면 즉시 거부됩니다. 이는 에이전트가 제안한 안전하지 않은 스케일이 GPU를 멈추게 하는 것을 방지합니다.

오토리서치 루프 (Autoresearch Loop):
AMK는 ScheduleConfig 변경을 반복적으로 제안하고, 이를 검증하며, CPU 레퍼런스와 정확성 게이트를 대조하고, 지연 시간을 측정하며, 변경 사항을 유지하거나 되돌리는 무인(unattended), 에이전트 구동 루프를 포함합니다. 이 루프는 자신의 베이스라인보다 메가커널을 자율적으로 개선합니다.

주요 기여

정적으로 검증된 Schedule IR: 신뢰할 수 있는 베이스 위에서 정적 그래프 체크를 통해 데드락 및 레이스 프리를 인증하는 고정된 검증기를 갖추고 있습니다. 이는 7,160개의 적대적으로 구성된 스케줄(그 중 6,091개는 안전하지 않음)에 대해 **제로 False-accept(오탐 없는 거부)**를 달 기록했으며, 360개의 실제 모델 로우어링(lowering)은 모두 수용했습니다.
자동 생성 커버리지: AMK는 수작업 CUDA 없이 10개의 지원 모델 전부에 대해 10/10의 정확한 메가커널을 자동 생성합니다(SmolLM2-135M 및 TinyLlama-1.1B와 같은 실제 체크포인트 포함). 또한 지원되지 않는 변형(예: 편향된 프로젝션)을 성공적으로 거부하거나 그 사각지대를 정직하게 문서화합니다.
자기 재타겟팅(Self-Retargeting): 단일 소스 코드베이스가 sm_80 (A100), sm_90 (H100), sm_120 (RTX 5090)을 위한 올바른 메가커널을 생성하며, 티처 포싱(teacher-forced) 결과와 비트 단위로 일치하거나 매우 타이트한 fp32 허용 오차( $\le 3.8 \times 10^{-5}$ ) 내에서 일치합니다.
양자화된 메가커널: 동일한 IR 경로로부터 int8(그리디-로스리스, 1.12배 빠름) 및 int4(로스형, 2.42배 감소된 가중치 트래픽) 웨이트 온리(weight-only) 양자화 커널을 자동 생성합니다.
에이전트 구동 컴파일러: 시스템은 정확성을 구조적 불변량(invariant)으로 취급하여, 에이전트가 런타임 중단 위험 없이 안전하게 스케줄 공간을 탐색할 수 있게 합니다.
정직한 성능 연구: 논문은 인퍼런스급 및 트레이닝급 GPU 전반에 걸친 결과를 보고하며, 시스템이 우위를 점하는 지점과 뒤처지는 지점을 명시적으로 공개합니다.

실험 결과

정확성 및 이식성:

RTX 5090, A100, H100에서 AMK는 HuggingFace의 그리디 디코드(greedy decode)를 토큰 단위로 일치시키며, SmolLM2-135M 체크포인트에 대해 티처 포싱 퍼플렉시티(perplexity)를 $2.5 \times 10^{-7}$ 수준까지 일치시킵니다.
검증기는 스트레스 테스트에서 6,091개의 안전하지 않은 스케줄을 모두 거부하여, 0%의 False-accept rate를 입증했습니다.

인퍼런스급 GPU에서의 성능 (승리 지점):

RTX 5090 (소비자용): 자동 튜닝된 int8 웨이트 온리 (W8A16) 메가커널은 배치-1 디코드 시 CUDA-graphed cuBLAS (bf16)보다 1.19–1.23× 더 빠릅니다. 이는 정밀도 비대칭 비교(int8 vs. bf16)이며, 승리의 원인은 커널 자체의 속도가 아니라 더 적은 가중치 바이트(0.61배)를 스트리밍하는 데 있습니다.
데이터센터 인퍼런스 플릿: int8 메가커널은 인퍼런스급 플릿 전반에서 CUDA-graphed cuBLAS bf16보다 빠릅니다:
- NVIDIA L4: 최대 1.33× 빠름 (모델 크기에 따라 증가).
- L40S: 1.25–1.27× 빠름.
- A10G: 규모에 따라 최대 1.08× 빠름.
자기 개선: 오토리서치 루프는 RTX 5090에서 기본 설정 대비 메가커널을 1.25–1.72× 자율적으로 개선했습니다.

트레이닝급 GPU에서의 성능 (패배 지점):

고대역폭 트레이닝급 GPU(A100, H100)에서 자동 생성된 메가커널은 CUDA-graphed cuBLAS 및 vLLM에 뒤처집니다.
- H100에서 AMK는 vLLM의 기본 CUDA-graph 경로보다 1.65배 느렸습니다.
- A100에서 최적화된 bf16 GEMV는 측정된 HBM 피크의 **약 63%**에 도달했습니다 (cuBLAS의 ~90% 대비).
진단: 병목 지점은 per-tile cross-SM 동기화 비용(cuBLAS는 지불하지 않는 고정 비용)입니다. 고대역폭 실리콘에서는 이 동기화 오버헤드가 지배적인 반면, 인퍼런스급 GPU에서는 융합 및 양자화를 통한 대역폭 절감 효과가 이 비용을 상쇄합니다.

의의 및 주장

본 논문은 AutoMegaKernel이 "원시 속도"에서 체계적 안전성과 이식성으로의 전환을 나타낸다고 주장합니다. 주요 기여는 자율 에이전트가 전체 모델 스케줄 공간을 안전하게 탐색할 수 있도록 하는 정적으로 검증된 컴파일러 아키텍처입니다.

신뢰성: GPU가 멈출 위험을 대신하여, 실행 전 안전하지 않은 스케줄을 거부하는 정적 안전 게이트를 제공합니다.
범용성: 수동 CUDA 튜닝 없이도 다양한 아키텍처(sm_80/90/120)와 모델 패밀리에 대해 자기 재타겟팅을 달성합니다.
정직함: 저자들은 성능 결과를 환경 의존적(regime-dependent)으로 명확히 프레임화했습니다. 시스템은 융합과 양자화를 활용해 메모리 트래픽을 줄임으로써 인퍼런스급 GPU에서 승리하지만, 현재는 더 미세한(coarser-grained) 스케줄링이 필요한 동기화 오버헤드로 인해 트레이닝급 GPU에서는 패배합니다.

결론적으로, 현재의 커널 품질이 고성능 트레이닝 하드웨어에서 수작업 최적화 베이스라인에 미치지 못할지라도, **정적으로 검증된 하네스(statically-checked harness)**는 단순한 마이크로 커널 개선을 넘어 더 거친 동기화(coarser synchronization)나 더 깊은 파이프라이닝과 같은 미래의 개선을 안전하게 자동화할 수 있는 필수적인 토대를 제공합니다.