CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation

CUDA Agent 은 확장 가능한 데이터 합성 파이프라인, 자동 검증 및 프로파일링이 가능한 기술 증강 CUDA 개발 환경, 그리고 안정적인 학습을 위한 강화학습 기법을 통해 커널 최적화 전문성을 함양하여 torch.compile 및 최첨단 상용 모델들을 압도하는 성능으로 CUDA 커널 생성 분야에서 새로운 최고 기록을 달성했습니다.

핵심

🚀 CUDA 에이전트: GPU 의 '수퍼 엔지니어'를 만든 AI

이 논문은 **"인공지능이 어떻게 컴퓨터 칩 (GPU) 을 위한 초고속 코드를 직접 설계하고 최적화할 수 있게 되었는지"**에 대한 이야기입니다.

기존의 AI 는 글을 쓰거나 코드를 짜는 데는 훌륭했지만, 컴퓨터의 속도를 극한으로 끌어올리는 **'GPU 커널 (핵심 연산 코드)'**을 만드는 데는 아직 부족했습니다. 마치 일반인에게는 훌륭한 요리사지만, 미슐랭 3 성 레스토랑의 주방장처럼 정교한 장비를 다루는 데는 서툰 것과 비슷하죠.

이 연구팀은 **'CUDA 에이전트 (CUDA Agent)'**라는 새로운 시스템을 만들어, AI 가 스스로 배우고 실전 훈련을 거쳐 세계적인 GPU 엔지니어가 되도록 했습니다.

---

🧩 1. 왜 이게 중요한가요? (배경)

현대 AI(예: 챗봇, 이미지 생성 AI) 가 돌아가려면 거대한 **GPU(그래픽 카드)**가 필요합니다. 하지만 이 GPU 를 100% 효율적으로 쓰게 하려면, 하드웨어의 미세한 구조를 완벽히 이해하는 전문가 수준의 코드를 직접 짜야 합니다.

• 기존 방식: 사람이 직접 코드를 짭니다. (매우 어렵고 시간이 걸림)
• 기존 AI 방식: AI 가 코드를 짜주지만, 성능이 나빠서 자동으로 최적화해주는 도구 (torch.compile 등) 보다 못합니다.
• 이 연구의 목표: AI 가 사람보다 더 똑똑하고 빠른 GPU 코드를 스스로 짜게 만드는 것.

---

🛠️ 2. 어떻게 만들었나요? (세 가지 핵심 비법)

연구팀은 AI 를 단순히 "코딩하게" 하는 것이 아니라, 실제 엔지니어처럼 일하게 만들었습니다.

① 📚 "수퍼 교재" 만들기 (데이터 합성 파이프라인)

AI 가 배우려면 좋은 문제집이 필요합니다. 하지만 좋은 GPU 최적화 문제는 너무 드뭅니다.

• 비유: 요리 학교에 학생을 보내려면, 다양한 재료를 섞어 새로운 요리를 만들어내는 수천 가지의 레시피 문제집이 필요합니다.
• 방법: AI 가 이미 있는 기본 연산 (재료) 들을 섞어서 새로운 복잡한 작업 (요리) 을 자동으로 만들어냈습니다. 그리고 이 문제들이 실제로 실행 가능한지, 너무 쉬운지 어려운지 자동으로 검사해서 좋은 문제만 남겼습니다.

② 🏗️ "안전한 실습장" 만들기 (스킬 강화 에이전트 환경)

AI 가 코드를 짜고 실행해보는 과정이 필요합니다. 하지만 AI 가 실수하면 컴퓨터가 터지거나, 점수를 조작할 수도 있습니다.

• 비유: 요리 실습장에 가는데, **도구 사용법 (칼질, 불 조절)**을 정확히 알려주고, 위생 검사관이 상주하며, 점수 조작을 막는 안전한 공간입니다.
• 방법:
  • AI 에게 "이 코드를 짜고, 실행해보고, 속도를 재고, 고쳐라"라는 **단계별 가이드 (스킬)**를 주었습니다.
  • AI 가 점수를 조작하지 못하도록 보안 장벽을 쳤습니다.
  • AI 가 코드를 짜면, 자동으로 실행해서 **"이게 원래보다 5% 이상 빠르냐?"**를 체크하고 점수를 줍니다.

③ 🧠 "안정적인 학습법" (강화 학습 알고리즘)

AI 가 처음부터 실전 훈련을 하면, 코드가 엉망이 되면서 학습이 붕괴되곤 합니다.

• 비유: 수영을 배우는데, 바로 깊은 바다로 뛰어들면 익사합니다. 먼저 **수영장 (단순 학습)**에서 기본기를 다지고, 그다음 **파도 치는 바다 (복잡한 실전)**로 나가는 것이 좋습니다.
• 방법:
  • 워밍업: 먼저 간단한 코딩 문제만 풀게 해서 기본기를 다졌습니다.
  • 거부 학습 (RFT): 엉뚱한 답을 내놓는 AI 의 행동을 걸러내고, 좋은 답만 남게 훈련시켰습니다.
  • 비평가 (Critic) 훈련: AI 가 "이 코드는 잘됐어, 저 코드는 나빠"라고 스스로 판단할 수 있게 비평가 AI도 함께 훈련시켜, 학습이 흔들리지 않게 했습니다.

---

🏆 3. 결과는 어땠나요? (성공 스토리)

이 'CUDA 에이전트'는 KernelBench라는 유명한 시험에서 놀라운 성과를 냈습니다.

• 기존 AI vs 이 AI:

  • 기존 AI (Claude, Gemini 등) 는 코드를 짜는 데는 잘했지만, 속도 최적화에서는 자동 도구 (torch.compile) 보다 느린 경우가 많았습니다.
  • CUDA 에이전트는 100% 가까운 확률로 정답을 맞췄고, 기존 자동 도구보다 최대 2 배 이상 빠른 코드를 만들어냈습니다.
  • 특히 가장 어려운 문제 (Level 3) 에서도 기존 최고의 상용 모델들보다 약 40% 더 좋은 성능을 보였습니다.

• 실제 사례:

  • 레벨 1 (단순): 행렬 곱셈을 할 때, 불필요한 계산을 아예 없애버려 속도를 73 배나 높였습니다. (마치 불필요한 계산을 줄여서 택시 대신 자전거를 탄 것처럼 빠름)
  • 레벨 3 (복잡): 복잡한 신경망 (ResNet) 을 최적화했을 때, 여러 단계의 작업을 하나로 합쳐서 3.5 배 더 빠르게 만들었습니다.

---

💡 4. 핵심 요약 (한 줄로 정리)

"이 연구는 AI 에게 'GPU 최적화'라는 어려운 기술을 가르치기 위해, 수천 개의 문제집을 만들고, 안전한 실습장을 제공하며, 단계별로 훈련시켜서, 이제 AI 가 사람 엔지니어보다 더 빠르고 똑똑한 GPU 코드를 직접 설계할 수 있게 만들었습니다."

이 기술이 발전하면, 앞으로 우리가 AI 모델을 쓸 때 더 빠르고, 더 저렴하며, 더 효율적인 하드웨어를 자동으로 만들어주는 시대가 올 것입니다. 마치 요리사가 직접 재료를 다듬고 조리법을 고안해내듯, AI 가 컴퓨터의 속도를 직접 조절하는 시대가 온 것입니다. 🚀

기술 요약

CUDA Agent: 대규모 에이전트 강화학습을 통한 고성능 CUDA 커널 생성 기술 요약

이 논문은 현대 딥러닝 인프라의 핵심인 GPU 커널 최적화 문제를 해결하기 위해 제안된 CUDA Agent에 대한 연구입니다. 기존 대규모 언어 모델 (LLM) 이 일반적인 프로그래밍에서는 탁월한 성능을 보이지만, CUDA 커널 생성 및 최적화 분야에서는 torch.compile과 같은 자동화 컴파일러 시스템보다 성능이 떨어지는 한계를 극복하고자 합니다.

1. 문제 정의 (Problem)

• CUDA 커널 최적화의 난이도: 고성능 CUDA 커널 개발은 GPU 미세 아키텍처에 대한 깊은 이해와 정교한 프로파일링 도구를 요구하는 매우 전문적인 작업입니다.
• 기존 LLM 의 한계: 기존 LLM 기반 CUDA 생성 접근법들은 훈련 없이 휴리스틱을 적용하거나, 고정된 다중 턴 실행 피드백 루프 내에서 미세 조정 (Fine-tuning) 하는 방식에 그쳤습니다. 이러한 방법들은 모델의 내재적인 CUDA 최적화 능력을 근본적으로 향상시키지 못해 성능 향상에 한계가 있었습니다.
• 경쟁력 부족: 현재 LLM 들은 torch.compile 같은 자동 최적화 도구조차 능가하지 못하며, 인간 전문가 수준에도 미치지 못합니다.

2. 방법론 (Methodology)

CUDA Agent 는 데이터, 에이전트 환경, 강화학습 (RL) 알고리즘이라는 세 가지 상호 보완적인 차원에서 시스템을 확장하여 구축되었습니다.

2.1 확장 가능한 데이터 합성 파이프라인 (Scalable Data Synthesis)

• 시드 문제 크롤링: PyTorch 및 Transformers 라이브러리에서 기본 연산자 (Seed Operators) 를 추출합니다.
• LLM 기반 조합적 합성: LLM 을 활용하여 여러 연산자를 결합 (Fusion) 한 새로운 복잡한 태스크를 생성합니다. 이는 개별 연산자를 최적화하는 것을 넘어, 중간 메모리 할당을 제거하고 레지스터/공유 메모리 제약을 고려한 통합 최적화 공간을 탐색하게 합니다.
• 엄격한 필터링: 실행 가능성, 결정론적 특성, 비자명성 (Non-trivial), 그리고 적절한 실행 시간 (1ms~100ms) 을 기준으로 데이터를 선별하여 6,000 개의 고품질 학습 데이터셋 (CUDA-Agent-Ops-6K) 을 구성했습니다.

2.2 기술 통합 에이전트 루프 (Skill-Integrated Agent Loop)

• 에이전트 환경: OpenHands 프레임워크를 기반으로 하며, Bash, 파일 편집, 프로파일링 도구 등을 포함한 표준 셸 유틸리티를 제공합니다.
• CUDA 코딩 기술 (Skill): 모델에게 SKILL.md 를 통해 CUDA 커널 최적화의 표준 워크플로우 (성능 분석 -> 커널 구현 -> 컴파일 및 검증 -> 반복 최적화) 를 명시적으로 학습시킵니다.
• 보상 해킹 방지:
  • 평가 스크립트의 파일 권한 통제 (수정 불가).
  • torch.nn.functional 의 폴백 호출 금지 (순수 CUDA 커널만 성능 향상을 내야 함).
  • 5 가지 무작위 입력에 대한 정합성 검증.
  • 프로파일링 파이프라인의 동기화 및 평균화를 통한 측정 노이즈 제거.

2.3 안정적인 RL 학습을 위한 알고리즘 개선

• 다단계 워밍업 전략 (Multi-stage Warm-up): 초기 RL 학습에서 모델 성능이 급격히 붕괴되는 원인 (도메인 분포 불일치 및 중요도 샘플링 비율의 불안정성) 을 해결하기 위해 다음과 같은 전략을 도입했습니다.
  1. Single-Turn RL: 베이스 모델에 대해 단일 턴 RL 을 수행하여 CUDA 생성 능력을 초기화합니다.
  2. RFT (Rejection Fine-Tuning): 에이전트 트래젝토리를 수집하여 품질이 낮은 (보상이 음수이거나 비효율적인) 샘플을 제거하고, 고품질 데이터로 에이전트 모델을 지도 학습합니다.
  3. Value Pretraining: 크리틱 (Critic) 모델을 초기화하여 다중 턴 상호작용 상태의 가치 함수를 정확히 예측하도록 합니다.
• 강건한 보상 설계: 단순한 속도 향상 (Speed-up) 비율 대신, 정합성 (Correctness) 과 성능 마일스톤 (5% 이상 향상 등) 에 기반한 이산적 보상 (Discrete Reward) 을 사용하여 학습의 안정성을 높였습니다.

3. 주요 기여 (Key Contributions)

1. CUDA Agent 시스템 제안: 확장 가능한 데이터 합성, 기술 통합 에이전트 환경, 안정적 RL 기법을 결합하여 LLM 의 CUDA 커널 코딩 및 최적화 능력을 체계적으로 향상시킨 대규모 에이전트 RL 시스템입니다.
2. SOTA 성능 달성: KernelBench 벤치마크에서 torch.compile 을 능가하는 성능을 기록했으며, 특히 난이도 Level-3 에서 Claude Opus 4.5 나 Gemini 3 Pro 와 같은 최상위 상용 모델보다 약 40% 더 높은 성능을 보였습니다.
3. 학습 안정성 해결: 대규모 컨텍스트 (128k 토큰) 와 긴 상호작용 (최대 200 턴) 을 가진 에이전트 RL 학습에서 발생하는 불안정성 문제를 RFT 와 Value Pretraining 을 통해 해결했습니다.

4. 실험 결과 (Results)

KernelBench (Level 1~3) 에서의 주요 결과는 다음과 같습니다:

지표	Level 1 (간단)	Level 2 (중간)	Level 3 (난이도 높음)
Pass Rate (정합성)	100%	100%	94%
Faster Rate vs. Compile	99%	100%	90%
Speed-up (Compile 대비)	1.87x	2.80x	1.52x

• 전반적 성과: CUDA Agent 는 Level 1, 2, 3 에서 각각 torch.compile 대비 100%, 100%, 92% 의 속도 향상 (Faster Rate) 을 기록했습니다.
• 경쟁 모델 비교: 가장 어려운 Level-3 설정에서 Claude Opus 4.5 와 Gemini 3 Pro 보다 약 40% 더 높은 성능을 달성했습니다.
• 컴파일러 우위: 정적 컴파일러인 torch.compile 이 처리하기 어려운 연산자 융합 (Operator Fusion) 및 복잡한 메모리 접근 패턴에서 에이전트 RL 기반의 학습된 최적화 정책이 우월함을 입증했습니다.

5. 의의 및 결론 (Significance)

• 하드웨어 인식 최적화의 자동화: CUDA Agent 는 LLM 을 단순한 코드 생성기를 넘어, 하드웨어 아키텍처를 이해하고 실행 기반 피드백을 통해 성능을 최적화하는 능동적인 시스템 최적화 에이전트로 변모시켰습니다.
• 컴파일러 대체 가능성: 전통적인 컴파일러 기반 최적화 (torch.compile) 를 능가하는 성능을 보여줌으로써, LLM 기반 커널 생성이 고성능 GPU 컴퓨팅 분야에서 경쟁력 있는 대안, 혹은 우월한 대안이 될 수 있음을 시사합니다.
• 미래 전망: 구조화된 환경과 신뢰할 수 있는 실행 기반 보상을 갖춘 에이전트 시스템은 GPU 컴퓨팅뿐만 아니라 성능이 중요한 소프트웨어 개발 전반의 자동화에 새로운 길을 열 것으로 기대됩니다.

이 연구는 대규모 강화학습과 에이전트 기술을 활용하여 GPU 커널 최적화라는 고난도 전문 영역을 자동화하는 데 있어 중요한 이정표가 되었습니다.