DICE: 대규모 언어 모델의 확산 기법으로 CUDA 커널 생성 성능 향상

초록

확산 대형 언어 모델(dLLM)은 토큰 병렬 생성 능력으로 인해 자기회귀(AR) LLM의 매력적인 대안으로 부상하고 있습니다. 이러한 패러다임은 전체적인 구조 계획과 비순차적 정교화가 중요한 코드 생성에 특히 적합합니다. 이러한 잠재력에도 불구하고, dLLM을 CUDA 커널 생성에 맞춰 조정하는 것은 높은 전문성과 고품질 학습 데이터의 심각한 부족으로 인해 여전히 어려운 과제로 남아 있습니다. 이러한 문제를 해결하기 위해 우리는 고성능 CUDA 커널에 최적화된 증강 지도 미세 조정 데이터 세트인 CuKe를 구축했습니다. 이를 기반으로 CUDA 커널 채우기 단계와 종단간 CUDA 커널 생성 단계로 구성된 2단계 선별 강화 학습(BiC-RL) 프레임워크를 제안합니다. 이 학습 프레임워크를 활용하여 우리는 1.7B, 4B, 8B의 세 가지 매개변수 규모를 아우르는 CUDA 커널 생성을 위해 설계된 확산 대형 언어 모델 시리즈인 DICE를 소개합니다. KernelBench에 대한 광범위한 실험 결과, DICE는 동등 규모의 자기회귀 및 확산 LLM을 모두 크게 능가하며 CUDA 커널 생성 분야의 새로운 최첨단 기술을 확립했습니다.

English

Diffusion large language models (dLLMs) have emerged as a compelling alternative to autoregressive (AR) LLMs, owing to their capacity for parallel token generation. This paradigm is particularly well-suited for code generation, where holistic structural planning and non-sequential refinement are critical. Despite this potential, tailoring dLLMs for CUDA kernel generation remains challenging, obstructed not only by the high specialization but also by the severe lack of high-quality training data. To address these challenges, we construct CuKe, an augmented supervised fine-tuning dataset optimized for high-performance CUDA kernels. On top of it, we propose a bi-phase curated reinforcement learning (BiC-RL) framework consisting of a CUDA kernel infilling stage and an end-to-end CUDA kernel generation stage. Leveraging this training framework, we introduce DICE, a series of diffusion large language models designed for CUDA kernel generation, spanning three parameter scales, 1.7B, 4B, and 8B. Extensive experiments on KernelBench demonstrate that DICE significantly outperforms both autoregressive and diffusion LLMs of comparable scale, establishing a new state-of-the-art for CUDA kernel generation.

DICE: 대규모 언어 모델의 확산 기법으로 CUDA 커널 생성 성능 향상

DICE: Diffusion Large Language Models Excel at Generating CUDA Kernels

초록

Support