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

Lars Simon Zehnder

cs.LGAug 18, 2026

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

Authors: Lars Simon Zehnder

Organizations: Independent Researcher Faro, Portugal

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( $λ$ ), 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.

Figures & tables

Algorithm	$\boldsymbol{\alpha_{t}}$	$\boldsymbol{\beta_{t}}$	Bootstrap
GAE	$r_{t}+\gamma(1-d^{\mathrm{term}}_{t})V(s_{t+1})-V(s_{t})$	$\gamma\lambda(1-d_{t})$	$A_{T}=0$
V-Trace	$\rho_{t}\!\left(r_{t}+\gamma(1-d^{\mathrm{term}}_{t})V(s_{t+1})-V(s_{t})\right)$	$\gamma c_{t}(1-d_{t})$	$\Delta_{T}=0$
Retrace( $\lambda$ )	$r_{t}+\gamma(1-d^{\mathrm{term}}_{t})\mathbb{E}_{\pi}[Q(s_{t+1},\cdot)]-Q(s_{t},a_{t})$	$\gamma c_{t+1}(1-d_{t})$	$0$
TD( $\lambda$ )	$r_{t}+\gamma(1-\lambda)(1-d^{\mathrm{term}}_{t})V(s_{t+1})$	$\gamma\lambda(1-d_{t})$	$G^{\lambda}_{T}$
Disc. Returns	$r_{t}$	$\gamma(1-d_{t})$	$G_{T}$
Elig. Traces (fwd)	$\nabla_{\mathbf{w}}\hat{V}(s_{t},\mathbf{w}_{t})$	$\gamma\lambda(1-d_{t-1})$	$\mathbf{z}_{-1}=\mathbf{0}$

Table 1: Algorithm-specific coefficients for the shared affine scan, applied in reverse or forward time as appropriate.

d^{\mathrm{term}}_{t},d^{\mathrm{trunc}}_{t}\in\{0,1\}

flag termination and truncation;

d_{t}=d^{\mathrm{term}}_{t}\vee d^{\mathrm{trunc}}_{t}

is the flag used in

\beta_{t}

, combining both.

\Delta_{t}=v_{t}-V(s_{t})

is the V-Trace value delta (Section 3.3.2 ). The Bootstrap column gives each algorithm’s scan boundary carry (Section 3.2 explains why it is

0

for GAE/V-Trace/Retrace(

\lambda

) but nonzero for TD(

\lambda

)/discounted returns). Eligibility traces use vector notation since the trace is vector-valued; see Section 9.2 for per-call scalar handling.

Fused pipeline (single Triton kernel)
Step	Reads	Writes
Load $r$ , $v$ (read twice at offsets rev and rev+1 ), $d$ , $bs$ ; compute $\alpha$ , $\beta$ in-register; write $A$	5	1
Total	6
Unfused pipeline (three GPU operations)
Compute $\delta=r+\gamma v_{t+1}(1{-}d)-v_{t}$ ( $\delta$ is GAE’s $\alpha_{t}$ ; $v_{t+1}$ reads $bs$ at the boundary column, a 5th tensor)	5	1
Compute $\beta=\gamma\lambda(1{-}d)$	1	1

Table 2: HBM accesses for GAE (with bootstrap values, no truncations).

r

: rewards;

v

: values;

d

: terminateds;

bs

: bootstrap values;

\delta

: TD residuals;

\beta

: decay scalars;

A

: advantages.

	H100 80GB HBM3			RTX 2000 Ada
Algorithm	Loop	Plain	With trunc.	Loop	Plain	With trunc.
GAE	317 $\times$	2.36 $\times$	1.7 $\times$	281 $\times$	2.57 $\times$	1.6 $\times$
V-Trace	96 $\times$	2.78 $\times$	2.7 $\times$	86 $\times$	3.46 $\times$	3.3 $\times$
Retrace( $\lambda$ )	72 $\times$	1.91 $\times$	1.9 $\times$	16 $\times$	1.62 $\times$	1.6 $\times$
$\lambda$ -returns	267 $\times$	2.85 $\times$	2.6 $\times$	251 $\times$	5.14 $\times$	4.4 $\times$
Discounted returns	167 $\times$	2.70 $\times$	2.6 $\times$	171 $\times$	5.70 $\times$	4.6 $\times$

Table 3: Production-regime speedup (

\texttt{num\_envs}{=}4096

\texttt{seq\_len}{=}128

), full-call wall time, v0.1.3 release measurements. Loop is the Triton kernel’s speedup over the naive uncompiled sequential-loop baseline. Plain and With trunc. are its speedup over the vectorized torch.compile baseline (Section 6.2 ). Eligibility traces and prefix sum have no with-truncations path.

Figure 1: Launch/wrapper overhead and device time for GAE (plain path, H100) across the twelve shapes in Table 4 . Top : launch/wrapper overhead as a fraction of full-call time for rl-triton and the torch.compile baseline, decreasing from

\sim

88–95% at the smallest shapes to

\sim

11–37% at the largest. Bottom : full-call time decomposed into launch/wrapper and device components. rl-triton ’s launch/wrapper cost remains

\sim

0.026–0.035 ms, while the baseline’s is 1.4–3.1

\times

larger.

num_envs	seq_len	GAE	Retrace	GAE	Retrace
		H100 80GB HBM3		RTX 2000 Ada
64	512	2.7 $\times$	2.2 $\times$	2.7 $\times$	2.0 $\times$
128	1024	3.2 $\times$	2.3 $\times$	3.1 $\times$	1.8 $\times$
256	1024	3.3 $\times$	2.3 $\times$	3.0 $\times$	1.4 $\times$
512	128	2.8 $\times$	2.0 $\times$	2.5 $\times$	1.8 $\times$
512	512	2.9 $\times$	2.2 $\times$	2.8 $\times$	1.4 $\times$

Table 4: Full shape sweep, plain path, full-call speedup vs. the vectorized torch.compile baseline (v0.1.3 release measurements).

Hidden size	num_envs	GAE share / E2E vs. loop	GAE share / E2E vs. scan
(1024, 1024)	4096	${\sim}2.2\%$ / ${\sim}1.02\times$	$<0.1\%$ / ${\sim}1.00\times$ (within noise)
(128, 128)	16384	${\sim}10$ – $15\%$ / ${\sim}1.11$ – $1.16\times$	not measured at this configuration †

Table 5: GAE share of a PPO update and end-to-end Triton speedup over each baseline, by policy hidden size (

\texttt{seq\_len}{=}128

Figure 2: Full-call GAE performance of rl-triton and PufferLib in the classic on-policy regime on an H100. rl-triton has lower runtime across all 20 evaluated shapes.

Figure 3: Full-call GAE performance of rl-triton and PufferLib in the massively parallel simulation regime on an H100. rl-triton has lower runtime at 11 of 20 shapes, including every configuration with

\texttt{seq\_len}\geq 32

except

(\texttt{num\_envs},\texttt{seq\_len})=(32768,32)

; PufferLib is faster at all configurations with

\texttt{seq\_len}\leq 16

. The crossover shifts to longer sequences as num_envs increases, reflecting PufferLib’s low overhead for very short per-thread sequential scans.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

num_envs	seq_len	V-Trace	$\lambda$ -ret.	Disc. ret.	Elig.	Prefix	V-Trace	$\lambda$ -ret.	Disc. ret.	Elig.	Prefix
		H100 80GB HBM3					RTX 2000 Ada
64	512	3.0 $\times$	3.0 $\times$	3.0 $\times$	2.8 $\times$	2.8 $\times$	2.9 $\times$	2.9 $\times$	2.9 $\times$	2.6 $\times$	2.5 $\times$
128	1024	3.0 $\times$	3.3 $\times$	3.3 $\times$	2.7 $\times$	2.7 $\times$	2.9 $\times$	3.1 $\times$	3.1 $\times$	2.6 $\times$	2.5 $\times$
256	1024	3.0 $\times$	3.2 $\times$	3.2 $\times$	2.8 $\times$	2.8 $\times$	2.9 $\times$	3.6 $\times$	3.3 $\times$	2.6 $\times$	2.6 $\times$
512	128	2.9 $\times$	2.9 $\times$	2.8 $\times$	2.5 $\times$	2.4 $\times$	2.8 $\times$	2.7 $\times$	2.6 $\times$	2.3 $\times$	2.2 $\times$
512	512	3.1 $\times$	3.1 $\times$	3.1 $\times$	2.8 $\times$	2.7 $\times$	3.0 $\times$	3.3 $\times$	3.3 $\times$	2.6 $\times$	2.5 $\times$

Appendix

Table 6: Full twelve-shape sweep for V-Trace, TD(

\lambda

\lambda

-returns, discounted returns, eligibility traces, and episodic prefix sum. Values are full-call speedups over the vectorized torch.compile baseline using the v0.1.3 release measurements.

Credit Assignment in RL

Credit assignment in reinforcement learning studies how to trace rewards and failures back to the actions, states, or agents that caused them.

See all 218 papers

daVinci-kernel: Co-Evolving Skill Selection, Summarization, and Utilization via RL for GPU Kernel Optimization
Jun 15, 2026Dayuan Fu, Mohan Jiang, Tongyu Wang +5RL for Code Generation GPU Kernel Optimization
1SII · 2SJTU · 3GAIR
TRACE: Turn-level Reward Assignment via Credit Estimation for Long-Horizon Agents
Jul 15, 2026Leitian Tao, Baolin Peng, Wenlin Yao +5Credit Assignment in RL Step-Level Credit Assignment in RL
University of Wisconsin–Madison · Microsoft Research
RTMC: Step-Level Credit Assignment via Rollout Trees
Apr 13, 2026Tao Wang, Suhang Zheng, Xiaoxiao XuStep-Level Credit Assignment in RL Agentic RL
Alibaba Group

Jun 15, 2026J/K move · Enter open · S save

daVinci-kernel: Co-Evolving Skill Selection, Summarization, and Utilization via RL for GPU Kernel Optimization

Dayuan Fu, Mohan Jiang, Tongyu Wang, Dian Yang+4

1SII · 2SJTU · 3GAIR

RL for Code Generation GPU Kernel Optimization Large Language Model-Guided Optimization GPU Kernel Generation RL Post-Training

Open paper AskPDF arXiv

GPU kernel optimization represents a paradigm where functional correctness is assumed and execution efficiency is the objective. We present daVinci-kernel, a reinforcement learning framework that couples skill discovery with skill exploitation through a dynamically evolving skill library. daVinci-kernel jointly trains three agents sharing one LLM backbone: a Skill Selection Agent that retrieves relevant techniques via BM25 and LLM reranking, a Policy Agent that generates multi-turn CUDA/Triton kernels conditioned on selected skills, and a Skill Summary Agent that distills successful rollouts into reusable skills. Candidate skills are added only after execution-based verification confirms reproducible speedups. All three agents share a single LLM backbone, are initialized via a structured SFT cold start on diversity-filtered data, and are then jointly optimized end-to-end with multi-turn REINFORCE and per-agent advantage estimation. On KernelBench, daVinci-kernel-14B achieves 37.2%, 70.6%, and 32.2% on Level 1, Level 2, and Level 3 under the Fast

_1

threshold, outperforming the strongest prior RL-trained model, Dr. Kernel-14B.

Abstract

Credit Assignment in RL

Explore similar work

daVinci-kernel: Co-Evolving Skill Selection, Summarization, and Utilization via RL for GPU Kernel Optimization

TRACE: Turn-level Reward Assignment via Credit Estimation for Long-Horizon Agents

RTMC: Step-Level Credit Assignment via Rollout Trees