Train diffusion language models to generate optimized CUDA kernels using bi-phase reinforcement learning. First phase masks and regenerates core kernel logic with provided scaffolding to prevent PyTorch shortcuts. Second phase enables end-to-end generation. Leverage diffusion's global context awareness for non-sequential code generation.
Train diffusion language models to generate optimized CUDA kernels using bi-phase reinforcement learning. First phase masks and regenerates core kernel logic with provided scaffolding to prevent PyTorch shortcuts. Second phase enables end-to-end generation. Leverage diffusion's global context awareness for non-sequential code generation.
DICE: Diffusion LLMs Excel at Generating CUDA Kernels
Problem Context
Generating correct, efficient CUDA kernels requires understanding memory access patterns, thread synchronization, and hardware constraints. Autoregressive models struggle with the global coordination needed for optimal kernels. Diffusion models excel because they can revise entire sequences, refining global properties like memory coalescing and instruction-level parallelism. However, they require specialized RL training to avoid deceptive solutions (e.g., using PyTorch functions instead of custom kernels).
Core Concept
DICE (Diffusion for Imperative Code Execution) uses bi-phase RL training:
Kernel infilling: Models learn correct CUDA logic with scaffolding (prefix/suffix provided), preventing shortcuts
End-to-end generation: Models progress to full kernel implementation with invocation logic
This leverages diffusion's iterative refinement to handle the multi-constraint optimization problem of kernel generation.
Architecture Overview
CuKe dataset: 6,303 high-performance CUDA kernels with verified 2.0× speedups
Kernel infilling stage: Learn core logic with fixed prefix/suffix via RL
End-to-end stage: Full kernel generation including invocation and bounds checks
classEndToEndKernelGeneration:
"""Generate complete kernels including invocation code."""def__init__(self, model, verifier: KernelVerifier):
self.model = model
self.verifier = verifier
defgenerate_full_kernel(
self,
problem_description: str,
max_length: int = 500,
num_samples: int = 1) -> List[str]:
"""
Generate complete CUDA kernel from problem statement.
Args:
problem_description: What the kernel should compute
max_length: Maximum code length
num_samples: Number of samples to generate
Returns:
kernels: List of generated kernel codes
"""
prompt = f"""
Generate an optimized CUDA kernel for the following problem:
{problem_description}
Provide:
1. __global__ kernel function definition
2. Host code to launch the kernel
3. Memory management
CUDA kernel:
"""
samples = []
for _ inrange(num_samples):
kernel, _ = self.model.generate_with_logprobs(
prompt, max_tokens=max_length, temperature=0.7
)
samples.append(kernel)
return samples
defevaluate_kernels(
self,
kernels: List[str],
problem: str,
baseline: str = None) -> List[Dict]:
"""
Evaluate multiple generated kernels.
Returns:
metrics: Compilation success, speedup, etc.
"""
results = []
for kernel in kernels:
success, error = self.verifier.compile_kernel(kernel)
speedup = self.verifier.measure_speedup(kernel, baseline) if success else0.0
reward = self.verifier.compute_reward(kernel, baseline)
results.append({
'compiles': success,
'error': error,
'speedup': speedup,
'reward': reward
})
return results
Step 5: Bi-phase training
deftrain_dice_bi_phase(
model,
dataset: CUDaKernelDataset,
optimizer,
verifier: KernelVerifier,
num_epochs: int = 10,
phase_transition_epoch: int = 5,
device: str = 'cuda'):
"""
Train DICE in two phases: infilling then end-to-end.
Args:
phase_transition_epoch: Switch to end-to-end training at this epoch
"""
infill_trainer = KernelInfillingRL(model, optimizer, verifier)
e2e_trainer = EndToEndKernelGeneration(model, verifier)
for epoch inrange(num_epochs):
if epoch < phase_transition_epoch:
# Phase 1: Kernel infillingprint(f"Epoch {epoch + 1}: Infilling phase")
infill_examples = dataset.get_infill_examples(num_examples=50)
avg_loss = infill_trainer.training_step(infill_examples)
print(f" Loss: {avg_loss:.4f}")
else:
# Phase 2: End-to-end generationprint(f"Epoch {epoch + 1}: End-to-end phase")
problems = dataset.get_problems()[:10]
for problem in problems:
kernels = e2e_trainer.generate_full_kernel(
problem['statement'],
num_samples=4
)
results = e2e_trainer.evaluate_kernels(
kernels, problem, baseline=problem['baseline']
)
avg_reward = sum(r['reward'] for r in results) / len(results)
print(f" Problem: avg_reward={avg_reward:.4f}")
return model
Practical Guidance
When to use: Generating CUDA kernels, optimized system code, or other complex imperative programs where global structure matters more than sequential coherence
Hyperparameters:
infilling_phase_epochs: 3-5 (build fundamental skills)
e2e_phase_epochs: 5-10 (refinement)
max_kernel_length: 300-500 tokens
num_generation_samples: 4-8 (ensembling before compilation)
Verified rewards via compilation and performance testing
Handles global properties (memory coalescing) better than autoregressive
Common pitfalls:
Infilling phase too short → models not learning kernel patterns
Compilation reward not strict enough → deceptive solutions
Not measuring actual speedup → accepting slow kernels
Phase transition too abrupt → catastrophic forgetting
Scaling: Dataset curation and verification are bottlenecks; consider synthetic kernel generation.
Reference
Paper: https://arxiv.org/abs/2602.11715
Related work: Code generation, diffusion models, system optimization, program synthesis
Benchmarks: KernelBench, custom CUDA kernel correctness and performance
Dataset: CuKe (6,303 verified kernels)
