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

arXiv: https://arxiv.org/abs/2608.17641 · HTML: https://mlsanity.com/paper/2608.17641/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 24 of its 37 references that are in the map.
- First posted on arXiv 18 Aug 2026.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 24 of 24)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning](https://mlsanity.com/paper/2108.10470/) (2021)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Simplified State Space Layers for Sequence Modeling](https://mlsanity.com/paper/2208.04933/) (2022)
- [Resurrecting Recurrent Neural Networks for Long Sequences](https://mlsanity.com/paper/2303.06349/) (2023)
- [Safe and Efficient Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1606.02647/) (2016)
- [Sample Efficient Actor-Critic with Experience Replay](https://mlsanity.com/paper/1611.01224/) (2016)
- [SpikingJelly: An open-source machine learning infrastructure platform for spike-based intelligence](https://mlsanity.com/paper/2310.16620/) (2023)
- [Training Agents Inside of Scalable World Models](https://mlsanity.com/paper/2509.24527/) (2025)
- [Sample Factory: Egocentric 3D Control from Pixels at 100000 FPS with Asynchronous Reinforcement Learning](https://mlsanity.com/paper/2006.11751/) (2020)
- [Robust Autonomy Emerges from Self-Play](https://mlsanity.com/paper/2502.03349/) (2025)
- [Nocturne: a scalable driving benchmark for bringing multi-agent learning one step closer to the real world](https://mlsanity.com/paper/2206.09889/) (2022)
- [GPUDrive: Data-driven, multi-agent driving simulation at 1 million FPS](https://mlsanity.com/paper/2408.01584/) (2024)
- [True Online Temporal-Difference Learning](https://mlsanity.com/paper/1512.04087/) (2015)
- [Truncated Proximal Policy Optimization](https://mlsanity.com/paper/2506.15050/) (2025)
- [Segmental Advantage Estimation: Enhancing PPO for Long-Context LLM Training](https://mlsanity.com/paper/2601.07320/) (2026)
- [UniLab: A Heterogeneous Architecture for Robot RL Beyond GPU-Dominant Paradigms](https://mlsanity.com/paper/2605.30313/) (2026)

## Abstract (arXiv)

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(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.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 log T grows, each adding an intermediate HBM round-trip. The library is available at https://github.com/simonsays1980/rl-triton.