이 논문은 현대 딥러닝 인프라의 핵심 구성 요소인 GPU 커널의 최적화 문제를 다룹니다. NVIDIA의 CUDA 아키텍처는 널리 사용되지만, 고성능 CUDA 커널 개발 및 최적화는 깊은 하드웨어 전문 지식을 요구하는 매우 어려운 작업입니다. 대규모 언어 모델(LLM)이 일반적인 소프트웨어 개발 작업에서 뛰어난 성능을 보였음에도 불구하고, 기존의 CUDA 커널 생성 접근 방식은 torch.compile과 같은 컴파일러 기반 시스템에 비해 경쟁력이 낮았습니다. 기존 연구들은 훈련 없는(training-free) 코드 개선이나 고정된 다중 턴(multi-turn) 실행-피드백 루프 내 모델 미세 조정을 사용했지만, 이는 모델의 본질적인 CUDA 최적화 능력을 근본적으로 향상시키지 못하여 성능 이득이 제한적이었습니다.

이러한 한계를 극복하기 위해, 본 논문은 CUDA Agent라는 대규모 Agentic 강화 학습(RL) 시스템을 제안합니다. CUDA Agent는 세 가지 핵심 구성 요소를 통해 LLM의 CUDA 커널 코딩 및 최적화 능력을 체계적으로 향상시킵니다.

1. 확장 가능한 데이터 합성 파이프라인 (Scalable Data Synthesis Pipeline):
고품질 CUDA 커널 훈련 데이터의 부족 문제를 해결하기 위해, CUDA Agent는 강화 학습에 필요한 방대하고 다양한 참조 연산자(operator) 코퍼스를 생성하는 확장 가능한 데이터 수집 파이프라인을 개발했습니다. 이는 다음 세 단계로 구성됩니다:

• Seed Problem Crawling: torch 및 transformers 라이브러리에서 PyTorch로 구현된 참조 연산자(예: matmul, relu, conv2d)를 수집하여 포괄적인 시드(seed) 문제 세트를 구축합니다. 각 연산자는 초기화 및 forward 메서드를 가진 Python 클래스로 표현됩니다.
• Combinatorial Problem Construction: LLM을 활용하여 최대 5개의 연산자 클래스를 샘플링하고, 이를 순차적으로 결합하여 단일 연산 계층으로 융합(fused)된 복합 연산자 태스크를 합성합니다. 이 과정은 PyTorch 연산자만을 대상으로 하여 transformers 라이브러리의 고수준 모듈은 제외합니다. 이러한 융합 태스크는 단순히 개별 연산자를 최적화하는 것보다 더 복잡한 최적화 문제를 제시하여, 공유 메모리(SMEM) 및 점유율(occupancy) 제약 등을 고려한 통합된 병렬 매핑 및 데이터 레이아웃이 필요합니다.
• Problems Filtering: 실행 기반 피드백을 통해 너무 쉽거나 어려운 문제를 걸러내는 엄격한 데이터 선택 과정을 거칩니다. 필터링 기준은 다음과 같습니다: (1) Eager 및 Compile 모드에서 성공적인 실행, (2) 재현성을 위한 내재적 무작위성(stochasticity) 배제, (3) 리워드 해킹(reward hacking) 방지를 위한 입력 변화에 따른 출력 유효성 검증, (4) Eager 모드 실행 시간 1ms에서 100ms 범위 제한 (자명하거나 과도하게 무거운 태스크 제외), (5) KernelBench 테스트 케이스와의 높은 유사성 배제.

이 과정을 통해 6,000개의 고품질 샘플로 구성된 CUDA-Agent-Ops-6K 데이터셋이 구축되었습니다.

2. 스킬 통합 에이전트 환경 (Skill-Augmented CUDA Development Environment):
CUDA Agent는 OpenHands 프레임워크와 유사하게 설계된 Agent Loop를 사용합니다. LLM은 BashTool, GlobTool, MultiEditTool, TodoWriteTool 등 표준 셸 유틸리티를 제공받습니다. 여기에 CUDA 코딩에 특화된 지침(SKILL.md)과 도구(예: 프로파일링 도구)가 Agent Skills 형식으로 통합됩니다. Agent Loop는 사고(reasoning), 행동 실행(action execution), 관찰(observation)을 교차하며 반복적인 코딩, 디버깅, 성능 최적화를 가능하게 하는 ReAct-스타일 패러다임을 따릅니다.
특히, 리워드 신호의 신뢰성을 보장하기 위해 견고한 리워드 스케줄링 및 리워드 해킹 방지 노력을 기울였습니다.

• 견고한 리워드 스케줄링: 기존의 순수 스피드업(speedup) 기반 리워드 대신, 정확성(correctness)과 실행 지연 시간(latency)을 함께 최적화하는 정규화된 리워드 스키마를 제안합니다. 리워드 $r$은 다음과 같이 정의됩니다:

$r = \begin{cases} -1 & \text{if correctness check fails} \\ 3 & \text{if } b(t, t_{\text{eager}}) \land b(t, t_{\text{compile}}) \\ 2 & \text{if } b(t, t_{\text{eager}}) \\ 1 & \text{otherwise} \end{cases}$
여기서 $t$는 생성된 커널의 런타임, $t_{\text{eager}}$는 PyTorch의 eager 구현 런타임, $t_{\text{compile}}$은 torch.compile 버전 런타임입니다. $b(t, t_0) = \mathbb{I} [(t_0 - t)/t_0 > 5\%]$는 기준선 $t_0$ 대비 5% 이상의 유의미한 속도 향상을 나타냅니다.

• 리워드 해킹 방지: (1) 정확성 검증 및 성능 프로파일링 스크립트 보호 (파일 권한 제어), (2) torch.nn.functional의 폴백(fallback) 구현 호출 금지 (성능 향상이 순수하게 생성된 CUDA 커널에서 나오도록 보장), (3) 5개의 무작위 입력에 대한 커널 출력 유효성 검증 (기능적 정확성 보장), (4) 정밀한 프로파일링 파이프라인 (측정 노이즈 감소), (5) 웹 검색 또는 외부 정보 검색 도구 미제공 (순수하게 로컬 실행 환경에서 솔루션 도출) 등의 제약 조건을 통해 리워드 해킹을 방지합니다.

3. 안정적인 훈련을 위한 강화 학습 알고리즘 기술 (RL Algorithmic Techniques for Stable Training):
초기 RL 시도에서 훈련 불안정성 문제를 겪었으며, 그 원인이 기본 모델의 사전 학습된 분포와 CUDA 커널 코딩에 필요한 데이터 분포 간의 심각한 불일치(domain distribution mismatch)임을 파악했습니다. 이 불일치는 낮은 확률의 토큰 샘플링과 FP16/BF16 정밀도 문제로 인한 중요도 샘플링 비율($\rho_t(\theta)$)의 불안정성으로 이어졌습니다.
이를 해결하기 위해 다음의 웜업(warm-up) 전략을 제안합니다:

• Single-Turn Warm-up: 기본 모델에 대해 단일 턴 RL을 먼저 수행하여 CUDA 커널 생성 능력을 향상시킵니다.
• Actor Initialization (Rejection Fine-Tuning, RFT): 단일 턴 RL을 통해 생성된 에이전트 궤적(trajectory)을 수집합니다. 이후 다음 기준에 따라 고품질 롤아웃(rollout)만을 선별하는 Rejection Fine-Tuning(RFT)을 적용하여 액터 모델($\pi_\theta$)을 초기화합니다: (1) 양수 리워드(R > 0)를 달성한 궤적만 유지, (2) 비효율적이거나 유효하지 않은 행동(예: 중복 다중 턴 루프, 도구 호출 스키마 위반)을 보이는 궤적 폐기. 필터링된 궤적은 다음 목적 함수로 액터를 최적화하는 데 사용됩니다:

$L_{\text{RFT}}(\theta) = -\mathbb{E}_{\tau \sim \mathcal{D}'} \left[ \sum_{t=1}^{T} \log \pi_{\theta}(a_t | s_t, a_{<t}) \right]$

• Critic Initialization (Value Pretraining): 상태 시퀀스와 해당 최종 리워드로 구성된 샘플링된 궤적 데이터를 활용하여 비평가(critic) 네트워크를 사전 학습합니다. GAE(Generalized Advantage Estimation) [19]를 사용하여 목표 값($V^{\text{targ}}_t$)을 계산합니다. 비평가 파라미터 $\phi$는 다음 평균 제곱 오차를 최소화함으로써 최적화됩니다:

$L_{\text{VP}}(\phi) = \frac{1}{2} \mathbb{E}_{\tau \sim \mathcal{D}} \left[ \frac{1}{T} \sum_{t=0}^{T-1} (V_{\phi}(s_t) - V^{\text{targ}}_t)^2 \right]$
이러한 초기화 단계는 RL 훈련을 200 스텝까지 안정적으로 유지하며 일관된 리워드 성장을 가능하게 합니다.
최종 RL 알고리즘은 PPO [18]를 사용하여 액터 모델 $\pi_\theta$를 최적화하며, 클리핑된 surrogate 목적 함수를 사용합니다:
$L_{\text{CLIP}}(\theta) = \mathbb{E}_{\tau \sim \mathcal{D}} \left[ \frac{1}{T} \sum_{t=0}^{T-1} \min \left( \rho_t(\theta) \hat{A}_t, \text{clip}(\rho_t(\theta), 1 - \epsilon_{\text{lower}}, 1 + \epsilon_{\text{higher}}) \hat{A}_t \right) \right]$
여기서 $\rho_t(\theta) = \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta_{\text{old}}}(a_t|s_t)}$는 중요도 샘플링 비율이며, $\hat{A}_t$는 어드밴티지(advantage), $\epsilon_{\text{lower}}=0.2$, $\epsilon_{\text{higher}}=0.28$입니다.

CUDA Agent는 Seed1.6 [20] MoE 모델을 기본 모델로 사용하며, KernelBench [17]의 Level 1부터 Level 3까지의 250개 태스크를 통해 평가되었습니다. Claude Opus 4.5, Gemini 3 Pro, GLM 4.6, Kimi K2 등 최첨단 모델들과 비교되었습니다.
결과적으로 CUDA Agent는 KernelBench에서 SOTA 성능을 달성하여, Level-1, Level-2, Level-3에서 torch.compile 대비 각각 100%, 100%, 92%의 빠른 속도(Faster Rate)를 보였습니다. 특히 가장 어려운 Level-3 설정에서는 Claude Opus 4.5 및 Gemini 3 Pro와 같은 강력한 독점 모델보다 약 40% 더 나은 성능을 보여주었습니다. 이는 LLM 기반 커널 생성이 기존 컴파일러 기반 최적화에 비해 경쟁력 있고 종종 우월한 대안임을 입증합니다.