CUDA Agent: 고성능 CUDA 커널 생성을 위한 대규모 에이전트 기반 강화학습 시스템

GPU 커널 최적화는 고도의 하드웨어 전문 지식이 필요한 영역으로, 기존 LLM은 이 분야에서 컴파일러보다 낮은 성능을 보였다. 이 논문은 대규모 데이터 합성 및 강화학습 기법을 통해 AI 에이전트가 전문가 수준의 CUDA 코드를 생성하고 최적화할 수 있음을 입증했다. 이는 딥러닝 인프라의 핵심인 커널 개발을 자동화하여 AI 모델의 실행 속도를 획기적으로 개선할 수 있는 가능성을 제시한다.

기존 LLM은 방대한 일반 코딩 데이터로 학습되었지만, CUDA와 같은 특수 하드웨어 프로그래밍 데이터는 전체의 0.01% 미만에 불과해 최적화 능력이 부족하다. 이는 마치 일반적인 글쓰기는 잘하지만 특정 기계의 상세 설계도는 그리지 못하는 것과 같다. CUDA Agent는 이 한계를 극복하기 위해 에이전트가 직접 코드를 실행하고 성능 지표를 확인하며 스스로 학습하는 '경험 기반 학습' 방식을 채택했다.

핵심 원리는 딥러닝의 기초인 강화학습(RL)을 코드 생성에 접목한 것이다. 단순히 정답 코드를 흉내 내는 것이 아니라, '어떤 코드가 GPU에서 더 빨리 돌아가는가'라는 보상 신호를 따라가도록 모델을 훈련시킨다. 이때 모델이 무작위로 코드를 짜다 길을 잃지 않도록, 성공적인 최적화 경로(궤적)를 먼저 학습시키고 성능 향상 단계마다 보상을 주는 정교한 가이드를 제공한다.

결과적으로 모델은 메모리 접근을 합치거나(Coalescing), 여러 연산을 하나로 묶는(Fusion) 등 하드웨어 전문가들이 사용하는 고난도 최적화 기법을 스스로 터득하게 된다. 이는 정적인 컴파일러가 미리 정해진 규칙대로만 최적화하는 것과 달리, AI가 문제의 수학적 구조를 이해하고 하드웨어 특성에 맞춰 유연하게 대응함을 의미한다.

전체 시스템은 데이터 합성, 에이전트 환경, 강화학습 알고리즘의 세 가지 축으로 구성된다. 데이터 합성 단계에서는 PyTorch 라이브러리에서 기본 연산자를 추출한 뒤, LLM이 이를 최대 5개까지 결합하여 새로운 복합 연산 문제를 만든다. 이후 실행 가능성, 결정론적 결과 여부, 연산량 등을 기준으로 엄격하게 필터링하여 6,000개의 학습 샘플을 확보한다.

에이전트 루프는 ReAct 패러다임을 따르며 Bash, 파일 편집, 프로파일링 도구 등을 기술(Skills) 형태로 사용한다. 보상 함수 r은 정확성 검증 실패 시 -1을 부여하고, 성공 시에는 생성된 커널의 실행 시간 t와 기준 모델(Eager, Compile)의 실행 시간 t_baseline을 비교한다. 구체적으로 b(t, t0) = I[(t0 - t)/t0 > 5%] 연산을 통해 5% 이상의 유의미한 속도 향상이 있을 때 더 높은 보상을 주어 성능 최적화를 유도한다.

학습 안정성을 위해 다단계 웜업을 수행한다. 먼저 단일 턴 RL로 기초 능력을 쌓은 뒤, 성공적인 궤적만 골라 Rejection Fine-Tuning(RFT)을 진행하여 Actor 모델을 초기화한다. 동시에 상태 시퀀스 s_t와 최종 보상 r을 입력으로 하여 가치망(Critic)을 사전 학습시킨다. 이후 PPO 알고리즘을 적용해 전체 에이전트 루프를 최적화하며, 이때 128k의 긴 컨텍스트와 최대 200회의 상호작용 턴을 지원하도록 설계했다.

KernelBench 벤치마크 평가 결과, CUDA Agent는 모든 난이도에서 압도적인 성능을 보였다. Level-1과 Level-2에서는 100%의 Faster Rate(기준 모델보다 빠른 커널 생성 비율)를 달성했으며, 가장 어려운 Level-3에서도 92%를 기록했다. 이는 정적 컴파일러인 torch.compile이 해결하지 못하는 복잡한 연산자 퓨전 시나리오에서도 AI가 더 효율적인 경로를 찾아냈음을 보여준다.

상용 모델과의 비교에서도 우위가 확인됐다. Claude 4.5와 Gemini 3 Pro는 약 90% 초반의 Pass Rate를 보였으나 실제 속도 향상 비율은 60%대에 머문 반면, CUDA Agent는 98.8%의 Pass Rate와 96.8%의 Faster Rate를 기록했다. 특히 Level-3에서 상용 모델들보다 약 40% 더 높은 성능을 기록하며 전문적인 최적화 도구로서의 가치를 입증했다.

Ablation Study를 통해 각 구성 요소의 중요성도 증명됐다. 에이전트 루프를 제거할 경우 성능이 급격히 하락했으며, RFT 초기화가 없을 경우 학습 17단계 만에 모델의 정책 엔트로피가 급증하며 성능이 붕괴되는 현상이 관찰됐다. 이는 복잡한 하드웨어 최적화 학습에서 단계별 가이드와 안정적인 초기화가 필수적임을 시사한다.

CUDA Agent는 Seed1.6(230B MoE) 모델을 기반으로 하며, 에이전트 학습을 위해 131,072 토큰의 긴 컨텍스트 윈도우를 활용한다. 시스템은 CPU와 GPU 자원이 분리된 Docker 기반 샌드박스 환경에서 작동하며, 128개의 NVIDIA H20 GPU 풀을 사용하여 대규모 병렬 프로파일링 및 학습을 수행한다.

학습의 핵심인 RFT 단계에서는 보상이 양수인 궤적만 필터링할 뿐만 아니라, 중복된 루프나 잘못된 도구 호출 패턴을 보이는 궤적을 제거하는 패턴 필터링을 병행한다. 이를 통해 모델이 효율적인 탐색 경로만 학습하도록 제약한다. Critic 모델은 Generalized Advantage Estimation(GAE)을 사용하여 Advantage를 계산하며, 사전 학습을 통해 Advantage 추정값의 분산을 줄여 PPO 학습의 수렴 속도를 높였다.

최적화 전략 측면에서 에이전트는 메모리 병합(Coalescing), 공유 메모리 타일링(Tiling), 벡터화된 로드(float4) 등 하드웨어 수준의 기법을 우선순위에 따라 적용하도록 훈련되었다. 특히 Level-3 과제에서는 BatchNorm 파라미터를 컨볼루션 가중치에 병합하거나 데이터 레이아웃을 NCHW에서 NHWC로 변경하는 등 고수준의 그래프 변환까지 수행하는 능력을 보여준다.

TVM과 같은 고도로 정교한 자동 튜닝 컴파일러 프레임워크와의 직접적인 성능 비교는 이루어지지 않았다. 또한 학습 과정에서 128개의 H20 GPU와 같은 막대한 컴퓨팅 리소스가 필요하므로 일반적인 연구 환경에서의 재현성이 제한될 수 있다.