LEAP: GPU 커널 생성을 위한 코드 강화 학습에서 적응적 가지치기를 통한 효율적인 환경-피드백

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

강화 학습(RL)을 통해 대규모 언어 모델(LLM)을 추가 학습하면 코드 생성 능력이 크게 향상됩니다. 현재 최첨단 프레임워크는 비평 네트워크의 과도한 메모리 사용량을 피하기 위해 그룹 상대 정책 최적화(GRPO)와 같은 비평 기반 패러다임을 활용하며, 규칙 기반 검증 환경과 결합합니다. 그러나 이러한 프레임워크를 CUDA 커널 생성과 같은 저수준 시스템 프로그래밍에 적용하는 것은 심각한 어려움을 야기합니다. 이진 형태의 합격/불합격 보상은 신호 희소성을 초래하고, 다중 단계 환경 피드백 루프는 엄청난 컴파일 지연 및 경로 전체에 걸친 보상 희석 문제를 발생시킵니다. 본 연구에서는 저수준 하드웨어 가속기 최적화를 위한 확장 가능하고 계산 효율적인 다중 단계 강화 학습 프레임워크인 LEAP(Lean Environment-Feedback via Adaptive Pruning)을 소개합니다. LEAP는 난이도에 따른 가지치기(DCP, Difficulty-Conditioned Pruning)라는 동적 게이팅 메커니즘을 특징으로 하며, 이는 간단하거나 심각한 오류를 유발하는 작업을 다중 단계 확장에서 적응적으로 제거하여, 리소스 집약적인 컴파일 및 하드웨어 탐색을 고가치 복잡한 작업에만 집중하도록 합니다. 이러한 기능을 수동 하이퍼파라미터 조정 없이 구현하기 위해, 순위 기반 보상(Rank-Based Reward) 방식을 제안합니다. GRPO 롤아웃 그룹 내의 쌍별 토너먼트 결과를 통해 스케일과 무관한 상대적 이점을 도출함으로써, 본 방법은 간단한 프롬프트에서 발생하는 토큰 비효율성을 줄이는 동시에 어려운 분포에 대한 학습 기울기를 최대화합니다. 실험 결과는 LEAP가 가지치기가 없는 다중 단계 기준 모델보다 더 빠른 수렴 속도를 보이며, 첫 번째 단계의 성능이 우수하고 강력한 다중 단계 디버깅 복원력을 갖춘다는 것을 보여줍니다. 이러한 결과를 통해 저수준 코드 강화 학습을 위한 실용적인 패러다임을 제시합니다.

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.