PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Lars Simon Zehnder

arXiv:2608.17641v2cs.LGcs.DCcs.PF

TL;DR

RL credit-assignment computations remain on the critical path because standard recurrences impose sequential work across rollout timesteps. rl-triton recasts seven estimators as one associative scan and delivers 1.6–5.70× full-call speedups over vectorized torch.compile in massively parallel simulation.

  • Problem

    RL credit-assignment recurrences impose O(T) sequential dependencies, while vectorized torch.compile still incurs intermediate HBM round-trips between scan stages.

  • Method

    rl-triton unifies seven credit-assignment algorithms as first-order linear recurrences evaluated by shared fused Triton associative scans with on-chip coefficient construction.

  • Results

    1.6–5.70× full-call speedups over the vectorized torch.compile baseline occur across all seven algorithms in the massively parallel simulation regime.

  • Takeaways & Limitations

    The shared on-chip scan provides a common high-performance implementation for diverse RL credit-assignment computations in massively parallel simulation.

  • Takeaways & Limitations

    The reported microbenchmark speedups are isolated-kernel gains and do not imply equivalent end-to-end training-throughput improvements.

Abstract

from arXiv · show

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$\times$ 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.

1 Introduction

The paper frames seven RL credit-assignment estimators as instances of one first-order linear recurrence, enabling a shared O(log T) associative-scan formulation instead of T-step sequential evaluation. It implements this formulation with fused Triton kernels that construct coefficients on-chip and explicitly handle episode boundaries, termination, truncation, and bootstrap values.

  • 1 Introduction: Credit-assignment quantities—including advantages, value targets, and return estimates—are computed at every training iteration to indicate which actions were good and by how much.The motivating estimators include GAE, V-Trace, Retrace(λ), and simpler return estimators.
  • 1 Introduction: A naive sequential implementation requires T serial steps despite available GPU parallelism, whereas an explicit associative scan restructures the recurrence to O(log T) depth.Wrapping sequential loops in torch.compile removes Python overhead but not their temporal dependency; the PyTorch scan implementation serves as the evaluation baseline.
  • 1 Introduction: Fused Triton kernels construct algorithm-specific recurrence coefficients from raw rollout inputs on-chip and evaluate all algorithms with the same scan operator and combine function.The five backward algorithms use the backward recurrence, while eligibility traces and episodic prefix sums use its forward-time mirror.
  • 1 Introduction: Seven standard RL credit-assignment algorithms are unified as instances of a single first-order linear recurrence evaluated with one associative scan.Backward estimators use At = αt + βt · At+1, while eligibility traces and episodic prefix sums use the forward-time mirror At = αt + βtAt−1.
  • 1 Introduction: The framework explicitly handles termination, truncation, and rollout-window bootstrap values, representing boundary bootstraps either in αT−1 or as a nonzero scan carry.These choices are algorithm-specific and are part of the unified kernel design.

2 Background

GAE credit assignment is a strict O(T) dependency chain whose GPU cost is dominated by repeated HBM traffic, which compilation alone cannot eliminate. Recasting the recurrence as an associative scan provides O(log T) depth, while fused Triton kernels reduce intermediate memory round-trips.

  • 2.1 Motivation: GAE [Schulman et al., 2016] forms a strict sequential chain of length T because each advantage depends on the next timestep.Computing A_0 requires A_1 through A_T−1.
  • 2.1 Motivation: For T=1024 and 128 environments, a naive GPU loop performs 1024 serial HBM round-trips despite parallel environment rows.The arithmetic is inexpensive; each timestep reads δ_t and γλ from HBM and writes A_t back.
  • 2.1 Motivation: torch.compile removes interpreter overhead but cannot break GAE’s true dependency, so the loop remains T sequential steps and is used as the Loop baseline.The vectorized torch.compile doubling-scan baseline instead restructures the recurrence to achieve O(log T) depth.
  • 2.2 Associative Scan: An associative scan computes all prefix reductions in O(log T) parallel steps, enabling the first-order recurrence A_t = α_t + β_t · A_{t+1} to be evaluated in parallel.The scan relies on an associative binary operator over recurrence tuples.
  • 2.3 Triton: Because credit-assignment kernels are memory-bandwidth bound, fused Triton kernels keep scan state on-chip and avoid intermediate HBM round-trips.Triton’s associative scan primitive expresses the tree reduction for NVIDIA and AMD GPU execution.

3 The Unified Linear Recurrence Framework

rl-triton expresses seven reinforcement-learning credit-assignment algorithms as forward or reverse applications of one associative affine scan, with fused kernels constructing recurrence coefficients directly from rollout tensors. The framework explicitly handles termination, truncation, and rollout-window boundaries without propagating credit across episodes or double-counting bootstrap values.

  • 3.1 Shared scan formulation: All seven algorithms share one associative scan operator, using reverse-time scans for five backward recurrences and forward-time scans for eligibility traces and episodic prefix sums.The same combine function supports both directions; forward scans require no successor-value bootstrap.
  • 3.1 Shared scan formulation: For sequences up to 131072 steps, fused kernels construct αt and βt on-chip from raw rollout tensors without materializing intermediate PyTorch tensors.Longer backward scans use a chunked unfused fallback, while Table 1 summarizes each algorithm’s coefficient mapping.
  • 3.2 Episode-boundary semantics: Termination and truncation both reset the scan carry, but termination removes successor bootstrapping whereas truncation inserts a caller-supplied continuation value without double-counting.Eligibility traces and episodic prefix sums reset at episode boundaries without value bootstrapping.
  • 3.2 Episode-boundary semantics: At a non-terminal rollout-window edge, continuation values enter αT−1, the boundary carry AT, or both, depending on the algorithm; a truncated final step instead uses truncation handling and zero carry.GAE, V-Trace, and Retrace use the supplied successor value locally, whereas TD(λ) and discounted returns distribute continuation between the local update and boundary carry.
  • 3.3 Algorithm-specific recurrences: The algorithm-specific mappings instantiate the shared recurrence across GAE, V-Trace, Retrace(λ), TD(λ), discounted returns, eligibility traces, and episodic prefix sums.Retrace shifts the decay to the next-step importance weight, while episodic prefix sums use a segmented forward scan with configurable boundary conventions.

4 Implementation

rl-triton maps each environment row to a Triton program that processes timestep blocks on-chip with a shared associative scan, while fused kernels construct recurrence coefficients in registers. Fusion reduces HBM traffic for common cases, but sequences longer than 131072 use an unfused chunked fallback and all kernels require float32 inputs.

  • Kernel design: Each environment row maps to one Triton program, with timestep blocks distributed across warps and processed in registers and on-chip memory.Parallelism is exposed across environment rows through the launch grid and across timesteps within each program.
  • Scan procedure: Fused kernels construct α_t and β_t in registers for sequences up to 131072, then apply an on-chip associative scan with O(log BLOCK SIZE) dependency depth.The generic backward kernel instead consumes precomputed α and β tensors, while both paths load masked timestep blocks and use the same scan procedure.
  • Fusion: For GAE, fusion reduces counted HBM accesses from 11 to 6 by eliminating δ and β materialization between preprocessing and scanning.The same traffic-reduction pattern applies to the other fused kernels.
  • Long sequences: Sequences longer than 131072 fall back to an unfused chunked PyTorch path for backward algorithms, while eligibility-trace and prefix-sum scans lack a chunked fallback.The chunked path materializes α and β as full tensors, and throughput is then dominated by global-memory bandwidth.
  • Numerical requirements: All kernels require float32 inputs because long scans accumulate numerical error, especially with bfloat16; mixed-precision pipelines should cast advantage inputs before scanning.The recommended pattern is float32 for the scan and bfloat16 for the policy forward pass.

5 Correctness Analysis

Associativity allows the recurrence tuples to be reduced in log2(T) dependent steps while preserving the sequential scan result. Setting β_t = 0 at episode boundaries severs backward and forward carry propagation, including for the two forward scans.

  • Associative reduction: The associative tree reduction reproduces the sequential recurrence in log2(T) dependent steps, yielding the same accumulated tuple regardless of grouping.Each combine preserves chronological ordering, and the resulting α component matches the exact sequential expansion across the window.
  • Episode boundaries: At a done boundary, β_t = 0 multiplies later accumulated terms by zero, so no credit propagates backward past t.The boundary result is (α_t, 0), leaving the scan at t unaffected by subsequent episodes.
  • Episode boundaries: For the two forward scans, the shifted boundary coefficient likewise sets β_t = 0 when d_{t−1} = 1, preventing carry from the preceding episode.Thus the tree reduction remains equivalent to the sequential recurrence across forward-scan episode boundaries.

6 Benchmarking Methodology

The evaluation uses rigorously warmed and repeated CUDA-event measurements, validates every kernel against an independent sequential reference, and compares all algorithms with a consistent doubling-scan baseline. The methodology also specifies the baseline’s architectural limitations and deployment environment.

  • Benchmark protocol: 50 timed iterations across 5 trials, using each trial’s median and the minimum trial median as the final measurement, follow 20 untimed warmups.Warmups absorb JIT compilation, autotuning, and first-touch allocation; CUDA events synchronize before and after each timed iteration.
  • Correctness validation: Every kernel is validated against an independent sequential reference implementation with atol=rtol=1e-4 before timing results are included.This guards against finite, plausible-looking credit-assignment errors that may evade NaN/Inf-only tests.
  • Execution environment: GPU benchmarks ran on RunPod cloud instances using torch==2.4.1+cu124 and Triton 3.0.0, covering the H100 and RTX 2000 Ada.The methodology reports the minimum of five independently warmed trial medians to reduce interference from frequency changes or OS scheduling.
  • Baseline design: All seven baselines use a consistent linear-space log2(T)-doubling associative scan, with suffix scans for backward algorithms and prefix scans for forward recurrences.A log-space alternative was rejected because repeated termination resets can cause float32 underflow and inf/nan outputs.
  • Baseline scope: The torch.compile comparison baseline is specific rather than an optimality claim: it uses 6–12 launches per call versus 1–2 for the rl-triton kernel.The launch-count difference reflects the inability of torch.compile/Inductor to fuse this elementwise-then-scan pattern as aggressively as hand-written Triton.

7 Results

rl-triton delivers 1.6–5.70× full-call speedups over the vectorized torch.compile baseline across seven credit-assignment algorithms, with larger gains generally emerging at longer sequence lengths. These isolated-kernel improvements do not translate directly into equivalent end-to-end training speedups because credit assignment may represent only a small fraction of total update time.

  • 7.1 Headline results: Truncation-aware speedups are smaller because rl-triton’s compile-time fast path benefits only the Triton kernel when truncations are absent, whereas torch.compile runs essentially the same graph in both cases.Retrace(λ) shows the smallest headline speedup because its heavier per-timestep work reduces the relative benefit of truncation-specific savings.
  • 7.2 Scaling behavior: GAE stays above 2.5× across the full shape sweep and reaches 7570× over the sequential loop at (512, 4096) on H100.The sweep covers twelve environment-count and sequence-length configurations, while the sequential loop incurs O(T) serial GPU operations and repeated launches and global-memory traffic.
  • 7.3 End-to-end impact: End-to-end PPO speedups are only ∼1.02× at (1024, 1024), versus 1.11–1.16× at (128, 128), because GAE accounts for less than 0.1% versus 10–15% of update time.These results follow Amdahl’s law, so the reported microbenchmarks should be interpreted as isolated-kernel gains rather than equivalent training-throughput improvements.
  • 7.2 Scaling behavior: rl-triton has lower full-call time at 31 of 40 evaluated shapes, with sequence length largely determining the crossover rather than environment count.Neither GPU consistently produces a larger Triton-to-torch.compile speedup, and the H100-versus-RTX ordering can reverse across shapes.

8 Related Work

rl-triton differs from existing RL training stacks by applying associative GPU scans to post-rollout scalar credit assignment, with explicit episode-boundary handling. It also connects this approach to parallel recurrence methods, IO-aware GPU kernels, and off-policy correction algorithms.

  • Parallel scans in ML: Associative scans parallelize linear recurrences in models such as S4 [Gu et al., 2022], S5 [Smith et al., 2023], LRU [Orvieto et al., 2023], and Mamba [Gu and Dao, 2023], whereas rl-triton targets post-rollout scalar credit assignment.LRUs have also been used as recurrent backbones for partially observable RL.
  • Off-policy correction: rl-triton implements V-Trace [Espeholt et al., 2018] and Retrace(λ) [Munos et al., 2016], whose recursive importance-weighted returns fit the affine-scan formulation.ACER [Wang et al., 2017] uses related truncated importance-sampling corrections with trust-region stabilization but is not currently implemented.

9 Discussion and Future Work

rl-triton’s recurrence may extend to long-context LLM post-training and spiking-neural-network training when scalar affine recurrences have sufficient parallel sequences, but current performance and functionality remain constrained by shape, sequence-length, precision, and kernel-design limitations. Future work targets broader rollout support, lower-overhead execution, expanded numerical formats, chunked scans, and additional differentiable RL kernels.

  • Broader Applicability: The recurrence may also support long-context LLM post-training and hard-reset spiking-neural-network BPTT when scalar affine recurrences have enough independent sequences for GPU parallelism.T-PPO and long-context PPO apply GAE-style token recurrences, while hard-reset spiking BPTT has comparable boundary semantics [Fan et al., 2025] [Gong et al., 2026] [Fang et al., 2023].
  • Limitations: Retrace(λ) slows below torch.compile above sequence length 2048 because of register-pressure effects.This regression is identified as a systems limitation requiring future correction.
  • Limitations: At 16 steps and 16,384 environments on H100, GAE’s device-only ratio falls below 1×, although its full-call ratio remains above 1×.GAE is the only one of the five measured algorithms showing this inversion, because launch and wrapper overhead dominate at this shape.
  • Limitations: Reported speedups depend on the torch.compile baseline, while GPU-margin direction varies by shape and the underlying H100-versus-RTX mechanism remains unidentified.A faster numerically stable PyTorch implementation would reduce the reported ratios; at launch-overhead-dominated shapes, several algorithms reverse which GPU has the larger margin.
  • Limitations: The library is limited to float32, scalar per-timestep βt, and flat fused scans of at most 131,072 steps, with incomplete long-sequence fallbacks.bfloat16 is deferred, per-dimension βt is unsupported, and eligibility-trace and episodic-prefix-sum forward scans lack chunked fallbacks beyond the sequence-length limit.
  • Future Work: Future work will add masked multi-turn and tool-use rollouts, fused loss operations, bfloat16 I/O, chunked scans, CUDA graphs, and differentiable PPO, GRPO, and KL-loss kernels.The plans also include fixing the Retrace register-pressure regression and extending support beyond the current rollout setting.

10 Conclusion

The seven credit-assignment algorithms share a common first-order linear recurrence evaluable with one associative scan operator. rl-triton implements each with a fused Triton kernel that constructs recurrence coefficients on-chip and keeps the O(log T) scan on-chip.

  • Unified recurrence framework: Seven RL credit-assignment algorithms share a common first-order linear recurrence structure evaluable with the same associative scan operator.The algorithms are GAE, V-Trace, Retrace(λ), TD(λ) returns, discounted returns, eligibility traces, and episodic prefix sums.
  • Kernel implementation: Each algorithm is implemented by its own fused Triton kernel, which constructs recurrence coefficients on-chip.
  • On-chip execution: The O(log T) associative scan is kept on-chip rather than following the sequential loop’s per-timestep HBM round-t…

A Four-Thread Scan Trace

The appendix details a reverse-chronological, four-thread associative scan that combines recurrence tuples in two fixed stride passes. After pass 2, each thread holds the correctly accumulated tuple for its position.

  • A Four-Thread Scan Trace: The scan loads four tuples in reverse chronological order, with thread i initially holding (α4−i, β4−i).Thus T1 = (α3, β3), T2 = (α2, β2), T3 = (α1, β1), and T4 = (α0, β0), where position 1 is timestep t=3 and position 4 is t=0.
  • A Four-Thread Scan Trace: Two fixed reduction passes suffice: stride 1 combines immediate neighbors, then stride 2 combines spans two positions apart.All stride-1 reads occur simultaneously before writes, so each combine uses the prior pass’s single-position tuples rather than updated spans.
  • A Four-Thread Scan Trace: After pass 1, threads 2–4 cover two adjacent positions through the associative combine, while thread 1 retains its initial tuple.For example, T1..2 = (α2 + β2α3, β3β2) and T3..4 = (α0 + β0α1, β1β0).
  • A Four-Thread Scan Trace: After pass 2, thread 4 covers all four positions and every thread holds the correct tuple for its own position.The loop terminates after log2(BLOCK SIZE) = log2 4 = 2 fixed passes, with no dynamic stopping criterion.

B Full-Grid Scaling: Remaining Algorithms

Table 6 reports full-grid results for the remaining five algorithms, none of which falls below 1× at any measured shape. Long-sequence RTX 2000 Ada speedups can substantially exceed H100 results, driven mainly by slower torch.compile baselines.

  • B Full-Grid Scaling: Remaining Algorithms: Table 6 reports the corresponding twelve-shape sweep for the remaining five algorithms, complementing Table 4’s GAE and Retrace(λ) results.Table 4 includes Retrace’s long-sequence regression analyzed in Section 7.3.
  • B Full-Grid Scaling: Remaining Algorithms: 5.8× on H100 versus 24.5× on RTX for discounted returns at (512, 4096), with neither result falling below 1× across measured shapes.The remaining five algorithms are covered by Table 6, and none falls below 1× at any measured shape.
  • B Full-Grid Scaling: Remaining Algorithms: 24.5× RTX speedup is driven mainly by the torch.compile baseline, which is 11× slower than on H100 at (512, 4096).For discounted returns, the baseline takes 2.62 ms on RTX versus 0.24 ms on H100, while the Triton kernel is only ∼2.6× slower.
Loading 2608.17641v2…