CUDA 커널 생성에 특화된 확산 LLM DICE

초록

DICE는 확산 기반 대형 언어 모델(dLLM)을 CUDA 커널 생성에 맞게 설계한 시스템으로, 고성능 CuKe 데이터셋과 두 단계로 구성된 BiC‑RL(커널 인필링 → 전체 생성) 강화학습 프레임워크를 활용한다. 1.7B, 4B, 8B 규모의 모델은 KernelBench에서 기존 AR‑LLM 및 기존 dLLM을 크게 앞서며, 컴파일 성공률과 실행 속도 모두에서 새로운 최고 수준을 기록한다.

상세 분석

본 논문은 코드 생성에 있어 비선형적인 리팩터링과 전역 구조 설계가 요구되는 CUDA 커널을 대상으로, 기존의 순차적 토큰 예측 방식인 Autoregressive(AR) 모델의 한계를 지적한다. AR 모델은 토큰을 왼쪽에서 오른쪽으로 하나씩 생성하므로, 긴 코드와 복잡한 연산 흐름을 다룰 때 높은 지연과 비효율을 초래한다. 반면 확산 기반 대형 언어 모델(dLLM)은 전체 시퀀스를 마스킹하고 여러 차례 디노이징 과정을 통해 전역 정보를 동시에 반영할 수 있어, 코드의 구조적 완전성을 빠르게 확보한다.

논문은 이러한 dLLM의 장점을 CUDA 커널에 적용하기 위해 두 가지 핵심 기여를 제시한다. 첫째, 고성능 CUDA 커널을 집중적으로 수집·정제한 CuKe 데이터셋을 구축하였다. 기존 ConCuR 데이터셋에 속도 향상 2× 이상의 엄격한 필터를 적용해 1,425개의 확실히 가속된 PyTorch‑CUDA 쌍을 확보하고, Attention·MLP 등 실제 AI 모델에서 흔히 사용되는 복합 구조 291개를 추가로 수집해 총 6,303개의 고품질 샘플을 만든 점이 눈에 띈다. 둘째, BiC‑RL이라는 두 단계 강화학습 프레임워크를 설계했다. 초기 단계인 커널 인필링에서는 프리픽스·코어·서픽스라는 3‑파트 구조로 커널을 분해해, 모델이 핵심 연산 로직에 집중하도록 유도한다. 이때 보상은 컴파일 성공 여부와 실행 정확성에 기반하며, ‘속임수’(예: 프롬프트 예시를 그대로 복제하거나 호출 로직을 누락하는 행위)를 방지한다. 두 번째 단계인 전체 커널 생성에서는 완전한 CUDA 구현을 요구함으로써 인필링 단계에서 학습한 핵심 로직을 실제 프로그램에 통합하도록 한다.

모델 아키텍처는 블록 디퓨전(Block Diffusion) 방식을 채택한다. 시퀀스를 일정 길이 블록으로 나누어 블록 내부는 비자율적으로 병렬 디코딩하고, 블록 간에는 AR 방식으로 순서를 유지한다. 이는 KV‑cache 재활용을 가능하게 하여 메모리 효율성을 높이고, 토큰당 지연을 크게 감소시킨다. 실험에서는 1.7B, 4B, 8B 규모의 DICE 모델을 각각 학습시켰으며, KernelBench 벤치마크에서 실행 정확도(Exec)와 속도 향상(fast 1, fast 2) 지표 모두에서 기존 AR‑LLM(예: Gemini‑3‑Pro, DeepSeek‑Coder) 및 기존 dLLM(예: Gemini‑Diffusion)보다 현저히 우수한 결과를 보였다. 특히 8B 모델은 14 × 이상의 속도 향상을 달성하면서도 100 % 컴파일 성공률을 유지했다.

이와 같이 DICE는 데이터 효율성, 구조적 보상 설계, 블록 디퓨전 기반 토큰 병렬화라는 세 축을 결합해, 고성능 GPU 코드를 자동 생성하는 새로운 패러다임을 제시한다. 향후 연구에서는 더 큰 파라미터 스케일, 멀티‑GPU 분산 디퓨전, 그리고 다른 저수준 언어(예: OpenCL, SYCL)로의 확장을 기대한다.