rl-triton: High-Performance Triton GPU Kernels for Reinforcement Learning Credit Assignment

arXiv:2608.17641v2 Announce Type: new Abstract: We present rl-triton, an open-source library of high-performance GPU kernels for reinforcement learning credit assignment, implemented in Triton. The core contribution is a unified associative scan framework that recasts seven distinct RL estimation algorithms - Generalized Advantage Estimation (GAE), V-Trace, Retrace(lambda), TD(lambda) returns, discounted returns, eligibility traces, and episodic prefix sums - as instances of a single first-order linear recurrence solved in O(log T) parallel steps. All algorithms share the same associative scan operator, with algorithm-specific fused Triton kernels constructing their recurrence coefficients on-chip. We verify the associative operator algebraically and define the treatment of terminated and truncated episodes explicitly. Benchmarks show a 1.6-5.70imes full-call speedup over a vectorized torch-compile baseline in the massively parallel simulation regime (thousands of environments, short rollouts). The reported range covers all seven algorithms on both GPUs, both with and without per-step truncation handling. For most algorithms, speedups increase at longer sequence lengths, as the baseline requires more scan stages as log T grows, each adding an intermediate HBM round-trip. The library is available at https://github.com/simonsays1980/rl-triton.

Related

• Nautilus: An Auto-Scheduling Tensor Compiler for Efficient Tiled GPU Kernels
• Towards Automated Kernel Generation in the Era of LLMs
• Kernel Foundry: A Diagnosis-driven Evolutionary Kernel Optimizer with Multi-Experts
• TideRL: Boosting Agentic RL Goodput with Readiness-Aware Scheduling

Source: arXiv cs.LG | 2026-08-19