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

📄 arXiv: 2608.01804v1 📥 PDF

作者: Tankun Li, Zhi Chen, Yaohua Tang

分类: cs.LG, cs.AI

发布日期: 2026-08-03


💡 一句话要点

提出LEAP以解决低级代码生成中的反馈稀疏问题

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 强化学习 代码生成 GPU内核 困难条件修剪 多轮反馈 任务选择 性能优化

📋 核心要点

  1. 现有的强化学习框架在低级代码生成中面临信号稀疏和编译延迟等挑战,影响了性能和效率。
  2. LEAP通过困难条件修剪(DCP)机制,动态聚焦于高价值复杂任务,提升了资源利用效率。
  3. 实验证明,LEAP在首轮生成能力和多轮调试上表现优异,收敛速度较未修剪基线显著提升。

📝 摘要(中文)

通过强化学习(RL)对大型语言模型(LLMs)进行后训练,显著提升了代码生成能力。然而,现有方法在低级系统编程(如CUDA内核生成)中面临挑战:二进制的通过/失败奖励导致信号稀疏,而多轮环境反馈循环则受到编译延迟和奖励稀释的影响。为此,本文提出了LEAP(Lean Environment-Feedback via Adaptive Pruning),一个可扩展且计算高效的多轮RL框架,专为低级硬件加速器优化。LEAP采用困难条件修剪(DCP)动态机制,专注于高价值复杂任务,避免简单和灾难性任务的资源浪费。通过基于排名的奖励公式,我们的方法在简单提示上惩罚令牌低效,同时在挑战性分布上最大化学习梯度。实验证明,LEAP在首轮表现和多轮调试的韧性上优于未修剪的基线,且收敛速度更快。

🔬 方法详解

问题定义:本文旨在解决低级代码生成中由于二进制奖励导致的信号稀疏和多轮反馈循环中的编译延迟问题。现有方法在处理这些问题时效率低下,无法有效利用资源。

核心思路:LEAP的核心思想是通过困难条件修剪(DCP)动态机制,优先处理高价值复杂任务,从而避免简单任务的资源浪费,提升整体学习效率。

技术框架:LEAP的整体架构包括多个模块,首先通过环境反馈收集任务信息,然后应用DCP机制进行任务筛选,最后通过基于排名的奖励公式优化学习过程。

关键创新:LEAP的主要创新在于引入了困难条件修剪(DCP)和基于排名的奖励机制,这与现有的无评论网络方法形成了鲜明对比,显著提高了低级代码生成的效率和效果。

关键设计:在设计中,DCP机制动态调整任务难度,确保资源集中在高价值任务上;基于排名的奖励公式通过对比赛结果计算相对优势,优化了学习过程,减少了简单提示的低效令牌使用。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,LEAP在首轮生成能力上优于未修剪的基线,且在多轮调试中表现出更强的韧性。具体而言,LEAP的收敛速度比未修剪的多轮基线快,显示出在复杂任务处理上的优势。

🎯 应用场景

LEAP的研究成果在低级系统编程、GPU内核生成等领域具有广泛的应用潜力。其高效的反馈机制和任务选择策略能够显著提升代码生成的质量和效率,推动相关技术的发展,尤其是在需要高性能计算的应用场景中。

📄 摘要(原文)

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.