GPU 與強化學習系統
rl-triton Unifies Seven Reinforcement Learning Return Computations as GPU Scans, Delivering Up to a 5.70× Speedup for Full Calls
The open-source rl-triton library rewrites credit-assignment recurrences such as GAE, V-Trace, and Retrace as Triton associative scans, reducing sequential dependency depth from O(T) to O(log T). In tests with 4,096 environments and 128 steps, it was 1.6× to 5.70× faster than an already vectorized torch.compile baseline, although end-to-end PPO speedup for large policies was only about 1.02×.
rl-triton has released a set of fused GPU kernels written in Triton for the often-overlooked credit-assignment stage of reinforcement learning training. The library covers GAE, V-Trace, Retrace(λ), TD(λ) returns, discounted returns, eligibility traces, and episodic prefix sums. The author notes that all seven algorithms can be expressed as the first-order linear recurrence `A_t = α_t + β_t A_{t+1}` and then transformed into an associative scan using the same combining operation, replacing the original step-by-step O(T) dependency chain with O(log T) parallel stages.
Algorithm-specific kernels construct α and β directly in registers from rewards, values, done flags, and importance ratios, then complete the intermediate scan on-chip. This avoids writing results back to HBM at every doubling stage. The implementation also distinguishes episode termination, time-limit truncation, and rollout-window boundaries, preventing bootstrap values from being incorrectly zeroed out. These details are particularly important for the numerical correctness of GAE and off-policy methods.
The author benchmarked v0.1.3 on an H100 80GB and an RTX 2000 Ada using 4,096 parallel environments and 128 steps per rollout. Compared with a `torch.compile` baseline that already uses an O(log T) scan, full-function-call speedups across the seven operations ranged from 1.6× to 5.70×. Discounted returns on the RTX 2000 Ada achieved the highest result, while GAE on the H100 reached 2.36×. This comparison is more informative than benchmarking against an uncompiled step-by-step loop.
However, kernel microbenchmarks do not translate directly into overall training throughput. With 1024×1024 policy hidden layers, GAE accounts for only about 2.2% of a PPO update, yielding an end-to-end improvement of roughly 1.02×. Smaller 128×128 policies achieve 1.11× to 1.16×. The current release also supports only FP32; Retrace becomes slower than the compiled baseline at a sequence length of 4,096 because of register pressure, and the two forward scans have no fallback beyond 131,072 steps. Engineering teams should first measure the share of training time spent on credit assignment in their own workloads before deciding whether to integrate the library.