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

섹션별 상세

01

CUDA Agent는 데이터 합성, 에이전트 환경, 강화학습 알고리즘의 세 가지 핵심 구성 요소로 이루어진 대규모 에이전트 강화학습 시스템이다.

02

데이터 합성 단계에서는 PyTorch와 Transformer 라이브러리에서 시드 연산자를 수집하고, LLM을 통해 최대 5개의 연산자를 조합한 6,000개의 고품질 학습 데이터셋(CUDA-Agent-Ops-6K)을 구축했다.

03

에이전트 환경은 ReAct 스타일의 워크플로우를 따르며, SKILL.md 가이드를 기반으로 코딩, 컴파일, 디버깅, 프로파일링을 반복 수행하여 torch.compile 대비 5% 이상의 성능 향상을 목표로 한다.

04

강화학습 파이프라인은 단일 턴 PPO 웜업, Rejection Fine-Tuning(RFT)을 통한 액터 초기화, 가치 사전 학습을 통한 크리틱 초기화 단계를 거쳐 최대 150턴의 긴 문맥에서도 안정적인 학습을 지원한다.

05

KernelBench 평가 결과, 가장 어려운 Level-3 작업에서 torch.compile 대비 90%의 Faster Rate와 1.52배의 기하 평균 속도 향상을 기록하며 최신 상용 모델들을 압도했다.

쿠다(CUDA): — NVIDIA GPU에서 병렬 컴퓨팅을 수행하기 위한 아키텍처 및 프로그래밍 모델로, 딥러닝 연산 가속을 위한 커널 최적화의 핵심 기술이다.
근사 정책 최적화(PPO): — 에이전트의 정책을 안정적으로 업데이트하기 위해 변화량을 제한하는 강화학습 알고리즘으로, 학습 과정의 급격한 성능 붕괴를 방지한다.
토치 컴파일(torch.compile): — PyTorch 코드를 자동으로 최적화된 커널로 변환하여 실행 속도를 높이는 컴파일러 기능으로, 본 연구에서 성능 비교의 주요 기준점이 된다.
거부 미세 조정(Rejection Fine-Tuning): — 생성된 결과물 중 품질이 낮은 데이터를 제외하고 우수한 데이터만 선별하여 모델을 미세 조정하는 기법으로, 정책의 효율성과 안정성을 높인다.
기하 평균 속도 향상(Geometric Mean Speed-up): — 여러 작업의 속도 향상 비율을 기하 평균으로 계산하여 전체적인 성능 개선 정도를 나타내는 지표로, 산술 평균보다 이상치의 영향을 덜 받는다.