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

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

왜 중요한가

핵심 기여

대규모 CUDA 학습 데이터셋 합성 파이프라인 구축

PyTorch 연산자를 조합하고 LLM을 활용해 복합 연산 과제를 생성하는 파이프라인을 통해 6,000개의 고품질 학습 데이터인 CUDA-Agent-Ops-6K를 구축했다.

기술 증강 에이전트 루프 및 환경 설계

SKILL.md를 통해 표준화된 워크플로우를 제공하고, 자동화된 검증 및 프로파일링 도구를 통합하여 에이전트가 스스로 코드를 수정하고 최적화할 수 있는 환경을 구현했다.

안정적인 에이전트 강화학습을 위한 다단계 웜업 전략

도메인 불일치로 인한 모델 붕괴를 방지하기 위해 Rejection Fine-Tuning(RFT)과 가치망 사전 학습을 포함한 다단계 초기화 기법을 도입하여 학습 안정성을 확보했다.

KernelBench 벤치마크 SOTA 달성

가장 어려운 Level-3 과제에서 Claude 4.5 및 Gemini 3 Pro 대비 약 40% 높은 성능을 기록했으며, torch.compile 대비 최대 100%의 성능 향상 비율을 보였다.

핵심 아이디어 이해하기

기존 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와 같은 막대한 컴퓨팅 리소스가 필요하므로 일반적인 연구 환경에서의 재현성이 제한될 수 있다.

실무 활용

딥러닝 모델의 추론 및 학습 속도를 개선하고자 하는 엔지니어들이 수동으로 CUDA 커널을 작성하는 대신 이 시스템을 활용해 최적화된 코드를 자동 생성할 수 있다.

복합 연산자(Fused Operators)를 위한 고성능 맞춤형 CUDA 커널 생성
기존 PyTorch 모델의 특정 병목 지점을 하드웨어 특화 커널로 교체하여 가속화
ResNet BasicBlock과 같은 표준 신경망 블록의 메모리 접근 패턴 최적화 및 레이아웃 변환

코드 공개 여부: 공개

코드 저장소 보기

키워드

CUDA(쿠다)Reinforcement Learning(강화학습)Agent(에이전트)Kernel Optimization(커널 최적화)LLM(대형 언어 모델)PPO(근사 정책 최적화)

용어 해설

CUDA 커널(CUDA Kernel): — GPU에서 병렬로 실행되는 함수 단위이다. 딥러닝 연산의 가장 기초적인 실행 블록으로, 메모리 접근 방식과 연산 순서에 따라 전체 모델의 성능이 결정되는 핵심 요소이다.
연산자 퓨전(Operator Fusion): — 여러 개의 독립적인 연산을 하나의 커널로 합치는 최적화 기법이다. 중간 데이터를 전역 메모리에 썼다 읽는 과정을 생략하여 메모리 대역폭 낭비를 줄이고 실행 속도를 높인다.
근사 정책 최적화(PPO): — 강화학습에서 정책을 안정적으로 업데이트하기 위한 알고리즘이다. 이전 정책과 너무 큰 차이가 나지 않도록 변화 폭을 제한(clipping)하여 학습의 안정성을 보장한다.
거부 파인튜닝(Rejection Fine-Tuning): — 모델이 생성한 여러 결과물 중 높은 보상을 받은 성공적인 궤적만 선별하여 다시 학습시키는 기법이다. 강화학습 초기 단계에서 모델이 올바른 방향으로 수렴하도록 돕는다.
커널벤치(KernelBench): — LLM이 생성한 CUDA 커널의 정확성과 성능을 평가하기 위한 벤치마크 데이터셋이다. 단순한 수식부터 복잡한 신경망 블록까지 다양한 난이도의 연산 과제를 포함한다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

CUDA Agent, CUDA 커널 최적화 성능 입증

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

CUDA Agent, CUDA 커널 최적화 성능 입증