Skip to content
Preprint

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

Aug 2026 · 0 citations · 37 references
Computer Science

TL;DR

This work introduces LEAP (Lean Environment-Feedback via Adaptive Pruning), a scalable and computationally efficient multi-turn RL framework optimized for low-level hardware accelerator alignment and proposes a Rank-Based Reward formulation, establishing a practical paradigm for low-level code RL.

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.

View source

Similar papers

Book Open access Aug 2026

Teaching LLMs to Write System Kernels for AI Accelerators: Post-Training, Reasoning, and Agentic Optimization

Writing high-performance system kernels for AI accelerators requires deep hardware expertise and low-level programming skill—a significant barrier for many researchers and engineers. This tutorial presents emerging methods that leverage large language models (LLMs) to automate kernel generation for GPUs and custom accelerators such as AWS Trainium. We cover the full pipeline: supervised fine-tuning on kernel corpora, reinforcement learning-based post-training (PPO, GRPO) with hardware-aware rewards, inference-time scaling for code reasoning, and multi-turn agentic systems for iterative kernel refinement. We also discuss key challenges unique to this domain, including reward hacking in compilation metrics, designing effective evaluation benchmarks for generated kernels, and strategies for data curation across heterogeneous hardware targets. The tutorial bridges the gap between LLM post-training research and systems programming, offering both foundational understanding and practical experience. We demonstrate how recent advances in reinforcement learning and inference-time compute scaling can be adapted to produce correct and performant kernels across diverse hardware backends. Attendees will gain hands-on experience generating Triton and NKI kernels using agentic LLM systems. Tutorial materials are available at https://neuron-science.github.io/llm_kernel_writing.

Youngsuk Park, R. Saha, Kaan Ozkara et al. · 0 citations
Preprint Jul 2026

Multi-turn RL with Structural and Performance Aware Rewards for CUDA Kernel Generation

Empirical findings suggest that CudaPerf significantly outperforms strong baselines, including Qwen-3-32B and CUDA Agent by achieving up to 5X and 3.32X improvements in speedup, and 17%&7% improvements in correctness, respectively.

Q. I. Mahmud, Nesreen K. Ahmed, Ali Jannesari · 0 citations
Preprint Aug 2026

MoE Proxy Models for Low-Cost Failure Reproduction and Diagnosis in LLM RL Post-Training

Reinforcement learning (RL) post-training of large language models (LLMs) is computationally intensive and involves complex system pipelines with substantial debugging overhead. In practice, factors such as framework adaptation, numerical precision, and operator implementation can cause failures, including gradient overflow and loss divergence. Reproducing such failures directly on large models requires considerable time and computational resources. This paper systematically analyzes failures encountered during large-scale RL training on the Huawei Ascend platform, summarizes representative failure types, and identifies three model-side factors relevant to fault reproduction. Based on these factors, we propose a proxy-model construction method for low-cost fault investigation and auxiliary diagnosis. It employs structure-preserving, clustering-based expert pruning to select representative experts while retaining the model's backbone architecture, routing mechanism, and basic task capabilities. Our experimental results show that the proxy models reduce accelerator requirements by 50%-87.5% and achieve up to a 33.3x reduction in per-step NPU-hour cost, while preserving major training dynamics and reproducing fault responses consistent with the original models. Overall, the proxy models can serve as low-cost surrogates for fault reproduction, targeted validation, and auxiliary diagnosis in RL post-training.

Yikai Wang, Chuansai Zhou, Yuhang Zhou et al. · 0 citations
Oct 2026

SynergyScale: Optimizing Offloading and Task Partitioning for Efficient Model Training

Deep neural networks (DNNs) with billions of parameters power many important applications, but their training is fundamentally constrained by the limited on-chip memory of GPUs. This memory wall forces training to rely on distributed execution or memory offloading, both of which introduce substantial inefficiencies. Existing offloading techniques can scale model size but often incur severe throughput degradation, while conventional distributed training suffers from poor hardware utilization due to limited cross-device bandwidth, leaving accelerator resources underused. We present SynergyScale, a software framework that addresses the GPU memory wall by jointly optimizing memory usage and computation efficiency. SynergyScale combines dynamic hierarchical memory offloading across GPUs, CPU, and secondary storage tiers to reduce GPU memory pressure with fine-grained multi-stream task partitioning that takes advantage of the concurrent execution capabilities of modern GPUs. The framework automatically selects offloading parameters that balance data movement overhead with parallel execution efficiency. Experimental results show that SynergyScale can train a 114.41B-parameter model on a single 80 GB NVIDIA A100 GPU and a 547.77B-parameter model on eight A100 GPUs, supporting models up to 29.5× larger than those enabled by existing approaches. By co-optimizing offloading and computation, SynergyScale achieves up to 4.5× higher training throughput than offloading-only methods, without altering training semantics, making billion-scale model training both efficient and cost-effective.

Xiaoyang Sun, Jie Xu, Zheng Wang · 0 citations
Open access Jul 2026

Validity-Aware Multi-Level Autotuning for Efficient Tensor Program Generation

MOCHA introduces two key ideas: a validity prediction model to filter out invalid configurations prior to expensive hardware profiling, and an advanced performance prediction model that leverages hidden features extracted during the compilation process to enhance autotuning efficiency and robustness.

Joohyoung Cha, Munyoung Lee, Jinse Kwon et al. · 0 citations
Preprint Aug 2026

Rollplex: Cross-Phase GPU Spatial Sharing for Vision Language Model Post-Training

Rollplex is presented, a runtime that decomposes the reference and training phase and moves the prefix computation into the rollout decode window and achieves speedup over serial colocation and disaggregation under the same GPU budget, while preserving the synchronous RL update.

Hanfeng Lu, Tianyu Feng, Suyi Li et al. · 0 citations