LEAP: Lean Environment-Feedback via Adaptive Pruning for Code RL in GPU Kernel Generation

2026-08-03

AI Summary

This work addresses the challenges of sparse rewards, high compilation latency, and signal dilution across multi-round feedback in reinforcement learning for low-level system programming tasks such as CUDA kernel generation. The authors propose LEAP, a novel framework that integrates difficulty-conditioned pruning (DCP) with a critic-free multi-round reinforcement learning paradigm. By dynamically pruning tasks to focus on high-value, complex samples and employing a ranking-based pairwise tournament reward mechanism, LEAP enables adaptive optimization without manual hyperparameter tuning. The framework combines Group Relative Policy Optimization (GRPO) with hardware-aware multi-round environment feedback, maintaining strong initial-generation performance while significantly improving robustness and convergence speed during iterative refinement—offering an efficient and practical reinforcement learning solution for low-level code generation.

Abstract

Post-training large language models (LLMs) via reinforcement learning (RL) has significantly advanced code generation capabilities. To bypass the heavy memory footprint of critic networks, current state-of-the-art frameworks leverage critic-free paradigms like Group Relative Policy Optimization (GRPO) tied to rule-based verification sandboxes. However, applying these frameworks to low-level systems programming, such as CUDA kernel generation-presents severe challenges: binary pass/fail rewards introduce severe signal sparsity, while multi-turn environmental feedback loops suffer from prohibitive compilation latencies and reward dilution across trajectories. In this work, we introduce LEAP (Lean Environment-Feedback via Adaptive Pruning), a scalable and computationally efficient multi-turn RL framework optimized for low-level hardware accelerator alignment. LEAP features Difficulty-Conditioned Pruning (DCP), a dynamic gating mechanism that adaptively cuts off simple and overly catastrophic tasks from multi-turn expansion, focusing resource-heavy compilation and hardware exploration exclusively on high-value, complex tasks. To fully operationalize these paths without manual hyperparameter engineering, we propose a Rank-Based Reward formulation. By deriving scale-free relative advantages from pairwise tournament outcomes within the GRPO rollout group, our method inherently penalizes token inefficiency on simple prompts while maximizing learning gradients on challenging distributions. Empirical evaluations show that LEAP achieves superior first-turn proficiency and robust multi-turn debugging resilience while converging faster than unpruned multi-turn baselines, establishing a practical paradigm for low-level code RL.

Problem

• reinforcement learning
• code generation
• CUDA kernel
• reward sparsity
• compilation latency

Innovation

• Adaptive Pruning
• Difficulty-Conditioned Pruning
• Rank-Based Reward
• Multi-turn RL
• CUDA Kernel Generation