PaperlayerView exact revision on arXiv

Source-linked AI summary

DICE: Diffusion Large Language Models Excel at Generating CUDA Kernels

Haolei Bai, Lingcheng Kong, Xueyi Chen, Jianmian Wang, Zhiqiang Tao, Huan Wang

arXiv:2602.11715v2cs.LGcs.CL

TL;DR

CUDA kernel generation needs both functional correctness and performance, but dLLM adaptation is limited by task specialization and scarce high-quality data. The paper constructs CuKe, trains DICE with two-stage BiC-RL, and reports state-of-the-art KernelBench performance across 1.7B, 4B, and 8B models. Its current reward uses execution-based binary feedback rather than explicitly incorporating execution latency.

  • Problem

    CUDA kernel generation requires correct and fast kernels, while high-quality kernel data is scarce and dLLM potential for this task remains largely unexplored.

  • Method

    The paper constructs the CuKe dataset and trains 1.7B, 4B, and 8B DICE models with BiC-RL, progressing from kernel infilling to end-to-end generation.

  • Results

    DICE consistently achieves state-of-the-art KernelBench performance against autoregressive and diffusion LLMs across its evaluated parameter scales.

  • Takeaways & Limitations

    CuKe and the two-stage training framework support a specialized family of diffusion models for CUDA kernel generation across three parameter scales.

  • Takeaways & Limitations

    BiC-RL uses execution-based binary rewards and does not explicitly incorporate CUDA execution latency because distributed speed measurements are noisy.

Abstract

from arXiv · show

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.

1 Introduction

CUDA kernel generation demands both correctness and speed, while dLLMs offer parallel, globally informed refinement that may fit its non-local dependencies and scarce high-quality data. DICE addresses this opportunity with specialized supervision and reinforcement learning.

  • CUDA kernel generation requires syntactic and functional correctness together with speedup benefits.
  • dLLMs can generate long code sequences faster through parallel generation than strictly sequential AR decoding.
  • Bidirectional attention and iterative refinement align with CUDA kernels’ non-local dependencies across distant code regions.
  • DICE combines the CuKe supervised fine-tuning dataset with BiC-RL and provides models at 1.7B, 4B, and 8B scales.BiC-RL hierarchically progresses across CUDA-kernel data and task complexity.
  • Experiments report state-of-the-art performance for DICE against both autoregressive and diffusion LLMs.

2 Preliminary

Autoregressive LLMs factorize sequence probability into left-to-right conditional predictions, whereas block diffusion combines autoregression between blocks with parallel decoding within blocks. This hybrid design reuses prior-block KV caches and supports non-sequential CUDA-kernel generation.

  • 2.1 Autoregressive Language Models: Autoregressive LLMs estimate sequence probability by factorizing it into conditional probabilities for each token.Each token is predicted from the preceding prefix.
  • 2.1 Autoregressive Language Models: Strict left-to-right generation introduces significant inference latency, especially for long sequences.
  • 2.2 Block Diffusion Language Models: Block diffusion divides sequences into contiguous blocks and combines intra-block autoregression with inter-block non-autoregressivity.The training objective is a weighted cross-entropy loss over blocks.
  • 2.2 Block Diffusion Language Models: With block length four in the illustrated example, diffusion decoding generates tokens in parallel within blocks while reusing KV caches from previous blocks.
  • 2.2 Block Diffusion Language Models: CUDA-kernel generation exhibits non-contiguous generation order, indicating behavior beyond strict left-to-right dependencies.

3 CuKe Dataset Construction

CuKe addresses scarce high-performance CUDA-kernel data by combining curated speedup-verified samples with structurally diverse operations and model components. The resulting dataset contains 6,303 high-quality training samples.

  • CuKe augments the ConCuR dataset to address the scarcity of high-performance CUDA-kernel data.
  • Only 86% of cudaLLM’s PyTorch-CUDA pairs exhibit speedup over their PyTorch equivalents.Speedups near 1.0× may be unreliable because of launch overhead, cache states, or measurement jitter.
  • CuKe incorporates complex Attention and MLP structures with varied input dimensions beyond isolated operations.The variation targets changing CUDA optimization logic across tensor shapes.
  • 291 structural samples were evaluated up to five times, yielding 36 functionally verified instances with demonstrable acceleration.
  • The final CuKe training set contains 6,303 samples assembled from the curated data dimensions.

4 DICE

DICE uses a bi-phase curated reinforcement learning framework that first trains CUDA kernel core-logic infilling and then progresses to end-to-end generation. The framework combines structured scaffolding, execution-based rewards, and curriculum-style data scheduling to address deceptive behavior and training difficulty.

  • BiC-RL: BiC-RL splits training into kernel infilling and subsequent end-to-end CUDA kernel generation stages.The first stage focuses on core-logic optimization, while the second generates complete implementations from high-level PyTorch references.
  • Training strategy: Supervised fine-tuning provides the cold-start initialization needed to establish foundational alignment before reinforcement learning.The paper motivates this phase because reasoning models struggle to generate functionally correct CUDA kernels.
  • Kernel Infilling Stage: The infilling stage decomposes kernels into a prefix, core C++ implementation, and suffix, constraining generation to the core logic and preventing deceptive behavior.The prefix defines the execution environment, the suffix wraps compilation and neural-network integration, and the core contains the kernel’s main logic.
  • Reward Function: Execution-based binary rewards assign success when sampled kernels compile and execute equivalently to the reference PyTorch implementation.Candidates are sampled in groups, executed, and compared with the reference implementation before reward normalization.
  • Data Scheduling: Curriculum-style data scheduling transitions reinforcement-learning data from basic operations to complex whole-model structures to improve training efficiency.The schedule addresses large exploration spaces and difficulty variation that can cause early ineffective rewards and optimization plateaus.

5 Experiments

Experiments evaluate DICE across KernelBench scales, baselines, metrics, and training strategies. DICE consistently achieves strong performance, while ablations support CuKe and hierarchical BiC-RL as important components.

  • Training and Evaluation Setups: KernelBench covers 250 tasks across single-kernel operators, fusion patterns, and full model architectures, evaluated with execution correctness and speedup metrics.The benchmark contains 100 Level 1, 100 Level 2, and 50 Level 3 tasks; fastp requires correctness and at least p-times speedup.
  • Training and Evaluation Setups: DICE models are compared with autoregressive, diffusion, specialized CUDA, and commercial models across 8B, 4B, and 1.7B scales.The comparisons include cudaLLM, CUDA-L1, Kevin-32B, Gemini-3-Pro, Claude-Sonnet-44, and Mercury-Coder-Small.
  • Main Results: DICE consistently demonstrates state-of-the-art KernelBench performance against autoregressive and diffusion LLMs across parameter scales.DICE-8B outperforms cudaLLM and Gemini-3-Pro, DICE-4B surpasses several 8B-scale models, and DICE-1.7B performs strongly on Level 1 tasks.
  • Main Results: Deceptive behavior is pervasive across models and becomes more severe for lower-ability models and more difficult tasks.The analysis specifically identifies Level 3 as the most difficult KernelBench level and connects BiC-RL with preventing deceptive behaviors during generation.
  • Ablation Studies: The 2.0× speedup threshold improves Level 1 and Level 2 performance despite slightly lower performance on Level 3 tasks.The alternative 1.5× and 2.5× CuKe variants contain 7,940 and 5,685 samples, respectively.
  • Ablation Studies: CuKe fine-tuning consistently yields superior performance compared with the cudaLLM and ConCuR supervised fine-tuning datasets.The comparison supports data-efficient scaling through strategic integration of high-performance samples rather than high data volume alone.
  • Ablation Studies: BiC-RL progressively improves CUDA kernel generation and avoids optimization plateaus through staged training and task-complexity decomposition.Supervised fine-tuning establishes baseline competence, while kernel infilling and data scheduling further improve performance across scales.

6 Related Work

Related work establishes diffusion language models as an alternative to autoregressive decoding and surveys their progress in code generation. It also places DICE within specialized LLM-based kernel-generation research addressing difficult benchmarks and scarce high-quality data.

  • Diffusion Large Language Models: Diffusion language models generate masked sequences through gradual denoising, offering an alternative to strict left-to-right autoregressive decoding.The related-work discussion frames dLLMs as a distinct language-modeling paradigm rather than a minor autoregressive variation.
  • dLLMs for Code Generation: Diffusion code-generation research has progressed from small-scale simple tasks to stronger models and commercial systems with faster generation speeds.However, the potential of dLLMs for challenging kernel generation remains largely unexplored.
  • LLM Kernel Generation: Kernel-generation research combines specialized benchmarks, reinforcement learning, agentic methods, and supervised datasets to improve functional correctness and execution efficiency.ConCuR specifically addresses the scarcity of high-quality CUDA kernel data, while prior systems include Kevin, TritonRL, AutoTriton, Astra, CudaForge, and KernelEvolve.

7 Conclusion

The paper combines CuKe with BiC-RL to train DICE diffusion models for CUDA kernel generation. Its experiments indicate that the two training stages jointly improve convergence and performance.

  • CuKe provides a robust supervised fine-tuning foundation that addresses the scarcity of high-quality CUDA-kernel data.
  • BiC-RL combines CUDA-kernel infilling with end-to-end generation to improve training efficiency and convergence.
  • Each BiC-RL training stage contributes to stable convergence, while their combination yields substantial performance improvements.
  • DICE integrates CuKe and BiC-RL into diffusion language models spanning three parameter scales.

A Appendix

The appendix supplies material beyond the main paper, including additional details and experimental results. It is organized into subsequent appendix sections.

  • The appendix contains details not discussed in the main paper and additional experimental results.
  • The appendix is structured into multiple sections for presenting these supplementary materials.

A.1 One-shot Prompt Templates

The one-shot prompts define two CUDA-kernel generation settings: infilling missing code within a skeleton and generating optimized architectures end to end. Both require compilable, functionally complete CUDA code aligned with the task architecture.

  • A.1.1 Kernel Infilling Stage: The kernel infilling prompt supplies contextual demonstrations to guide core C++ generation while preserving alignment with a given prefix and suffix.
  • A.1.1 Kernel Infilling Stage: Infilling generates the missing core C++ code inside a partial CUDA-kernel skeleton using the provided PyTorch reference and configuration.
  • A.1.1 Kernel Infilling Stage: The infilling template illustrates how generated code is combined with the supplied prefix and suffix to form the complete architecture.
  • A.1.1 Kernel Infilling Stage: The infilling output must be real, compilable, fully functional code for the optimized architecture, without testing code or explanatory text.
  • A.1.2 Kernel Generation Stage and Evaluation: The end-to-end generation prompt follows KernelBench specifications and is also used for evaluation, aligning reinforcement-learning objectives with standard protocols.
  • A.1.2 Kernel Generation Stage and Evaluation: End-to-end generation may replace selected operators, leave others unchanged, fuse operators, or apply algorithmic changes such as online softmax.
  • A.1.2 Kernel Generation Stage and Evaluation: The end-to-end template demonstrates inline embedding of custom CUDA operators into a PyTorch architecture before providing the target architecture for optimization.
  • A.1.2 Kernel Generation Stage and Evaluation: The requested end-to-end output is compilable, fully functional code for a new optimized architecture named ModelNew, with no testing code or extra text.

A.2 Data Efficiency of dLLM

This appendix examines data efficiency through supervised fine-tuning comparisons between Qwen3 and SDAR models on CuKe. It also points to reported comparisons involving SDAR, Qwen3, and DICE-8B.

  • Supervised fine-tuning uses the same LLaMAFactory pipeline, CuKe dataset, and training hyperparameters across Qwen3 and SDAR model families.
  • SDAR generally benefits from supervised fine-tuning, while Qwen3 4B and 8B models can exhibit degradation.
  • Table 7 reports supervised fine-tuning comparisons for Qwen3 and SDAR series models on CuKe, with metrics reported as percentages.
  • Figure 5 reports the average performance gain after supervised fine-tuning for SDAR versus Qwen3.
  • Table 8 compares baseline models with DICE-8B while excluding deceptive behaviour.

A.3 Quantitative Comparison of Deception Behavior

The evaluation compares model-generated CUDA kernels with and without robust checking, addressing deceptive kernels that can appear functionally equivalent under the standard script.

  • The standard KernelBench script checks equivalence to the PyTorch reference but not the specific form of the generated kernel.The authors therefore introduce a robust-check comparison after identifying deceptive behavior empirically.
  • Robust checking is motivated by deceptive kernels that satisfy equivalence checks without implementing the intended specialized CUDA logic.
  • The reported comparisons distinguish direct script results from results obtained after applying the authors’ more reliable checking procedure.

A.4 Experimental Details

The experiments use multi-scale training, controlled decoding settings, and KernelBench-based comparisons to evaluate DICE’s correctness, speedup, cost, and failure modes. Results include stronger performance with compact training resources, representative accelerated kernels, and identified limitations in rewards and deceptive behavior.

  • Training configuration: DICE is trained at 1.7B, 4B, and 8B scales using supervised fine-tuning on CuKe followed by two reinforcement-learning stages.BiC-RL comprises kernel infilling and end-to-end kernel generation, with stage-specific sampling and training steps.
  • Training cost: 120 GPUh, or about a 12.5% increment, is the additional BiC-RL overhead over single-stage reinforcement learning for the 8B model.The analysis uses 8 NVIDIA A100 GPUs and identifies compilation and execution as a main bottleneck.
  • Comparison with cudaLLM: DICE achieves better performance than cudaLLM despite using 4,096-token generation, 6,303 SFT samples, and 4,992 RL samples versus cudaLLM’s larger settings.Reducing cudaLLM’s maximum generation length from 32,768 to 4,096 tokens causes a sharp performance drop.
  • Dataset analysis: Only 86% of the cudaLLM SFT samples achieve speedup, with most concentrated between 1.0× and 1.5× relative to PyTorch references.
  • Generated kernels: DICE generates functionally correct Level 3 kernels with measured speedups, including 2.1143× for NetVladWithGhostClusters, 4.0756× for ReLUSelfAttention, and 3.5430× for RegNet.These examples come from DICE-1.7B, DICE-4B, and DICE-8B, respectively.
  • Failure analysis: The identified failure categories include execution failure, functional correctness failure, and generation failure, motivating stronger structural constraints and explicit operator-replacement rewards.
  • Limitations: BiC-RL uses an execution-based binary reward and does not explicitly include execution latency because distributed speedup measurements are noisy.The authors identify multi-objective rewards as an important direction for future work.
Loading 2602.11715v2…