Papers › rl-triton: High-Performance Triton GPU Kernels for Reinforcement Learning Credit Assignment
rl-triton: High-Performance Triton GPU Kernels for Reinforcement Learning Credit Assignment
Lars Simon Zehnder
Title, abstract, authors and date from arXiv's metadata (CC0); this paper is not in the Papers with Code archive (frozen 2025-07-28).
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(λ), TD(λ) returns, discounted returns, eligibility traces, and episodic prefix sums - as instances of a single first-order linear recurrence solved in O(logT) 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.70× 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 logT grows, each adding an intermediate HBM round-trip. The library is available at https://github.com/simonsays1980/rl-triton.
In Syntology Open this paper in Syntology's Atlas, the map of the papers in Syntology's graph and their citations.
Code
Reachability, where shown, is from one Syntology probe window (2026-09-16 to 2026-09-18); repositories not probed show nothing. GitHub stars are not tracked.
Code Syntology ran Syntology
Syntology holds the repository link but has not harvested or run code from it.
Results from the paper
The Papers with Code archive ends with its 2025-07-28 snapshot. This paper's arXiv identifier, 2608.17641, was issued in August 2026, after that date, so the archive has no leaderboard rows for it.
Report a problem or propose a change · a person checks every report against the paper or source before anything changes; decisions are listed on /corrections