CUDA 에이전트, LLM을 훈련해 GPU 커널을 작성, torch.compile 능가하는 속도

서론

바이트댄스 Seed 팀과 칭화대학교 AI 산업 연구원(AIR)이 공동으로 CUDA Agent를 출시했다. 이는 대규모 에이전트 강화학습 시스템으로, 언어 모델이 고성능 CUDA 커널을 생성하도록 가르치는 것을 목표로 한다.

이 작업은 AI가 생성하는 저수준 코드에서 오랫동안 존재해 온 약점을 겨냥한다. 최첨단 언어 모델은 컴파일 가능하고 올바르게 실행되는 CUDA 코드를 생성할 수 있는 경우가 많지만, 프로덕션급 GPU 워크로드에는 정확성만으로는 충분하지 않다. 생성된 커널은 고도로 최적화된 컴파일러 생성 구현과도 경쟁할 수 있어야 한다.

CUDA Agent는 바로 이 두 번째 문제를 중심으로 구축되었다. 유효한 CUDA 코드를 생성하는 것뿐만 아니라, 유의미한 런타임 개선이 달성될 때까지 커널을 반복적으로 프로파일링하고 검증하며 최적화하는 방법을 학습하는 것이다.

올바른 CUDA 코드만으로는 충분하지 않은 이유

에이전트 강화학습을 도입하기 전에 Seed1.6 기반 모델은 KernelBench에서 이미 상당히 강력한 CUDA 코딩 능력을 보여주었다.

연구자들이 사용한 250개의 KernelBench 레벨 1~3 작업에서 Seed1.6은 74.0%의 통과율을 달성했다. 즉, 대부분의 벤치마크 작업에 대해 기능적으로 올바른 솔루션을 생성할 수 있었다는 뜻이다.

하지만 성능은 또 다른 이야기였다.

기반 모델이 생성한 커널 중 27.2% 만이 torch.compile보다 빨랐으며, 컴파일러 기준선 대비 기하 평균 속도는 **0.69×**에 불과했다.

모델	통과율	`torch.compile`보다 빠름	`torch.compile` 대비 기하 평균 속도
Seed1.6 기반 모델	74.0%	27.2%	0.69×
CUDA Agent	98.8%	96.8%	2.11×

이 격차는 자동 GPU 커널 생성의 주요 과제를 부각시킨다.

언어 모델은 CUDA 구문을 이해할 수 있지만, 비효율적인 메모리 접근, 과도한 커널 실행, 최적이 아닌 타일링 선택, 불필요한 중간 결과 생성 등을 여전히 만들어낼 수 있다.

성능 프로파일링

CUDA 최적화 지침
제한된 샌드박스 환경
실제 실행 결과에 연동된 보상 신호

SKILL.md 파일은 또한 벤치마크 결과를 왜곡할 수 있는 지름길을 제한한다. 예를 들어, 에이전트는 커스텀 커널 구현에서 임의의 PyTorch 연산으로 단순히 대체할 수 없다.

이것은 강화학습 시스템이 의도된 최적화 문제를 해결하지 않고 보상을 극대화하는 방법을 찾을 수 있기 때문에 중요하다.

PPO 훈련이 에이전트를 장기 최적화 궤적으로 확장

연구자들은 **근접 정책 최적화(PPO)**를 사용하여 CUDA Agent를 훈련시켰다.

주요 과제 중 하나는 GPU 최적화가 본질적으로 장기간에 걸친 작업이라는 점이다. 모델은 최적의 결과에 도달하기 위해 여러 차례의 코드 편집, 컴파일, 디버깅, 프로파일링 및 추가 최적화를 거쳐야 할 수 있다.

따라서 논문은 학습을 안정화하기 위해 다양한 컨텍스트 길이와 훈련 단계를 채택했다.

기반 모델은 Seed1.6으로, 총 2,300억 개의 파라미터와 230억 개의 활성 파라미터를 가진 혼합 전문가 모델이다.

에이전트 강화학습의 경우:

컨텍스트 창은 131,072토큰이다.
훈련 재생은 최대 150라운드의 에이전트 상호작용을 허용한다.
평가는 최대 200라운드의 에이전트 상호작용을 허용한다.
모델은 150스텝의 RL을 훈련한다.
Actor와 critic 모델은 전체 장기 최적화 전에 다단계 워밍업 전략을 사용한다.

이는 단순히 CUDA 프롬프트와 답변 쌍에 대해 모델을 미세 조정하는 것보다 훨씬 복잡하다.

목표는 실행 피드백과의 반복적 상호작용을 통해 모델이 자체 커널을 개선하는 방법을 학습하도록 하는 것이다.

샌드박스가 컴파일과 GPU 프로파일링을 분리

실행 속도가 보상의 일부이므로 안정적인 성능 측정이 중요하다.

따라서 연구자들은 CPU-GPU 리소스가 분리된 샌드박스 아키텍처를 구축했다.

Docker 기반 터미널 샌드박스는 컴파일과 같은 CPU 중심 작업을 처리하고, 검증과 성능 프로파일링은 128개의 NVIDIA H20 GPU를 포함하는 전용 GPU 샌드박스 풀에 배정된다.

이러한 분리에는 이중 목적이 있다.

첫째, 컴파일과 에이전트 상호작용을 GPU 벤치마킹과 격리한다. 둘째, 전용 GPU 할당은 동시 워크로드 간의 간섭을 줄여 지연 시간 측정을 더욱 안정적으로 만든다.

이것은 강화학습에 중요하다. 타이밍 측정이 노이즈의 영향을 받으면 모델은 신뢰할 수 없는 보상 신호를 받고 잘못된 최적화 동작을 학습할 수 있기 때문이다.

CUDA Agent, 98.8% 통과율 달성, 96.8% 작업에서 컴파일러 능가

최종 시스템은 KernelBench 레벨 1~3의 250개 작업에서 평가되었다.

CUDA Agent는 다음을 달성했다:

98.8%의 전체 통과율
PyTorch 즉시 실행보다 98.4% 빠름
torch.compile보다 96.8% 빠름
즉시 실행 대비 2.60×의 기하 평균 가속
torch.compile 대비 2.11×의 기하 평균 가속

이러한 결과는 Seed1.6 출발점 대비 큰 개선이다.

모델은 정확성 테스트를 통과하는 코드를 생성하는 데 더 나아졌을 뿐만 아니라, 생성된 커널이 컴파일러 기준선을 능가할 가능성도 높아졌다.

KernelBench 난이도별 결과

세 가지 KernelBench 난이도 모두에서 성능 개선은 강력하게 유지되었다.

KernelBench 난이도	통과율	`torch.compile`보다 빠른 비율	`torch.compile` 대비 기하 평균 가속
레벨 1	100.0%	97.0%	1.87×
레벨 2	100.0%	100.0%	2.80×
레벨 3	94.0%	90.0%	1.52×

레벨 2의 성과가 특히 두드러진다.

이러한 작업은 퓨전과 메모리 이동을 통한 최적화 기회가 있는 연산자 시퀀스를 포함하며, 이는 정적 컴파일러의 휴리스틱 전략이 항상 효과적으로 활용하지 못하는 부분이다.

논문 저자들은 반복적 학습 최적화 프로그램이 특정 하드웨어를 겨냥한 메모리 접근 방식, 타일링 및 퓨전 선택을 포함한 더 넓은 구현 전략 공간을 탐색할 수 있다고 주장한다.

훈련 데이터 곧 공개 예정

현재 전체 훈련 모델 가중치는 이 프로젝트의 공개 릴리스에 포함되지 않았다.

하지만 연구자들은 훈련 스택의 몇 가지 중요한 구성 요소를 이미 공개했다.

무료 · We0 AI 사용하기

몇 분 만에 쇼케이스 사이트를 만들고 리드를 늘리세요

아이디어를 한 문장으로 입력하면 We0 AI가 쇼케이스 사이트, 페이지, CMS를 생성하고 출시 후 고객과 트래픽 확보를 돕습니다.

무료 등록을 위한 하나의 완전한 프로젝트 생성

하나의 완전한 생성 흐름을 시도하고 첫 번째 프로젝트 초안을 빠르게 보는 데 가장 적합합니다.

무제-1

공들여 나열한 것

채우다

반지름

생성 중

CUDA-Agent-Ops-6K

CUDA-Agent-Ops-6K 데이터셋은 6,000개의 합성 연산자 수준 훈련 작업을 포함한다.

구축 프로세스는 다음과 같다:

torch, transformers 등의 라이브러리에서 시드 연산자를 수집한다.
LLM을 사용하여 여러 연산자를 더 복잡한 퓨전 작업으로 결합한다.
실행 결과 기반 검사를 통해 생성된 작업을 필터링한다.

필터링 단계는 무작위성이 있거나, 너무 단순하거나, 유효하지 않거나, 평가 작업과 너무 유사한 경우를 제거한다.

이 데이터셋은 Hugging Face에 게시되었으며 CC BY 4.0 라이선스를 사용한다.

SKILL.md와 에이전트 환경

GitHub 저장소에는 CUDA용 SKILL.md 지침 파일과 에이전트 작업 환경도 포함되어 있다.

이러한 자료는 최적화 워크플로우, 사용 가능한 도구, 제약 조건 및 에이전트를 부트스트랩하기 위한 실행 환경 설정을 문서화한다.

보상 메커니즘과 훈련 방식

논문과 저장소에는 보상 설계, 워밍업 단계, critic 초기화 및 PPO 기반 훈련 방식이 장기 최적화 프로세스를 안정화하기 위해 설명되어 있다.

이는 최종 벤치마크 점수를 재현하는 것뿐만 아니라 시스템 프로그래밍 에이전트 훈련에 관심이 있는 연구자들에게 특히 유용하다.

이것이 KernelBench 자체를 넘어 중요한 이유

CUDA 커널 최적화는 현대 AI 인프라의 상당 부분을 뒷받침한다.

더 빠른 커널은 다음을 향상시킬 수 있다:

모델 훈련 처리량
LLM 추론 지연 시간
GPU 활용률
서비스 비용 효율성
실시간 AI 파이프라인
고성능 수치 연산 워크로드

AIBase의 원문 기사는 AI에서의 중요성을 지적했다.

인프라, 대규모 언어 모델 추론 서비스, 자율주행, 그리고 고빈도 매매 등의 분야에서의 잠재적 응용——이러한 분야에서는 미세한 지연 시간 차이조차 결정적일 수 있다.

CUDA Agent의 성과가 전통적인 컴파일러가 구식이 되었다는 것을 의미하지는 않는다.

torch.compile은 여전히 강력한 자동 기준선이며, CUDA Agent의 평가 자체가 통제된 벤치마크 환경과 특정 GPU 환경에 의존한다.

이 연구가 보여주는 결론은 범위가 더 좁지만 여전히 중요하다: 충분한 실행 피드백과 전용 강화 학습의 지원을 받아, 언어 모델은 대다수 테스트 GPU 커널 작업에서 정적 컴파일러 기준선을 능가하는 최적화 전략을 학습할 수 있다.

이는 하위 성능 엔지니어링을 단순한 일회성 코드 생성이 아닌, 에이전트 기반 학습의 합리적인 영역으로 전환할 수 있음을 의미한다.

자주 묻는 질문

CUDA Agent란 무엇인가?

CUDA Agent는 바이트댄스 Seed, 칭화대학교 지능산업연구원(AIR) 및 그들의 공동 SIA-Lab이 개발한 대규모 에이전트 강화 학습 시스템이다. 이 시스템은 언어 모델이 반복적으로 CUDA 커널을 작성, 검증, 분석 및 최적화하도록 훈련한다.

CUDA Agent는 어떤 모델을 기반으로 하는가?

이 연구는 Seed1.6을 기반 모델로 사용한다. 논문에서는 이를 총 파라미터 수 2300억 개, 활성 파라미터 수 230억 개의 혼합 전문가 모델로 설명한다.

KernelBench란 무엇인가?

KernelBench는 언어 모델이 PyTorch 프로그램을 위해 정확하고 효율적인 GPU 커널을 생성할 수 있는지 평가하는 오픈 벤치마크이다. CUDA Agent는 KernelBench 레벨 1부터 레벨 3까지를涵盖하는 250개 작업에서 평가되었다.

CUDA Agent는 `torch.compile`보다 얼마나 빠른가?

전체 벤치마크에서 CUDA Agent는 torch.compile 대비 기하 평균 2.11배의 가속을 달성했다. 생성된 커널은 평가 작업의 **96.8%**에서 컴파일러 기준선보다 빨랐다.

CUDA Agent는 강화 학습을 사용하는가?

그렇다. 이 시스템은 PPO를 사용하며, 다단계 사전 학습(pre-warmup), 가치 모델 사전 훈련, 견고한 보상 설계 및 장시간 다중 턴 에이전트 궤적을 결합한다.

CUDA Agent 최적화 궤적 한 번은 얼마나 오래 실행될 수 있는가?

훈련 롤아웃은 최대 150개의 에이전트 턴을 허용하며, 평가는 최대 200개 턴을 허용한다. 에이전트 훈련 컨텍스트 창은 131,072개의 토큰이다.

CUDA Agent는 오픈 소스인가?

이 프로젝트는 GitHub 저장소, 에이전트 환경, SKILL.md, 훈련 레시피 및 CUDA-Agent-Ops-6K 데이터셋을 공개했다. 완전히 훈련된 모델 가중치는 현재 공개 배포 목록에 포함되어 있지 않다.

CUDA-Agent-Ops-6K란 무엇인가?

CUDA-Agent-Ops-6K는 6,000개의 합성 연산자 수준 CUDA 훈련 작업을 포함하는 데이터셋이다. 이는 대규모 에이전트 강화 학습을 위해 특별히 설계되었으며, 작업이 실행 가능하고, 결정적이며, 사소하지 않고, KernelBench 평가 세트와 분리되도록 필터링 단계를 포함한다.

요약

CUDA Agent는 대규모 언어 모델이 생성한 시스템 코드의 구체적인 약점을 해결한다: 생성된 CUDA 코드는 정확할 뿐만 아니라 실제로 컴파일러가 생성한 대안보다 빨라야 한다.

Seed1.6을 시작으로, 연구진은 CUDA 전용 에이전트 환경, 실행 피드백, 견고한 보상 설계 및 장기 지평 PPO 훈련을 통해 통과율을 74.0%에서 98.8%로, torch.compile보다 빨라진 비율을 27.2%에서 96.8%로 끌어올렸다.

이 프로젝트는 또한 6,000개 작업의 훈련 데이터셋, SKILL.md, 에이전트 환경 자료 및 훈련 레시피를 공개하여, 연구자들이 학습된 GPU 최적화 방향에 대한 후속 작업을 진행할 수 있는 구체적인 기반을 제공한다.

CUDA Agent의 주요 기여는 성능 엔지니어링 자체가 일회성 코드 생성으로 근사하는 것이 아니라, 반복적인 에이전트 루프를 통해 학습될 수 있음을 입증했다는 점이다.