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.
Figures & tables
Algorithm
αt
βt
Bootstrap
GAE
rt+γ(1−dtterm)V(st+1)−V(st)
γλ(1−dt)
AT=0
V-Trace
ρt(rt+γ(1−dtterm)V(st+1)−V(st))
γct(1−dt)
ΔT=0
Retrace( λ )
rt+γ(1−dtterm)Eπ[Q(st+1,⋅)]−Q(st,at)
γct+1(1−dt)
0
TD( λ )
rt+γ(1−λ)(1−dtterm)V(st+1)
γλ(1−dt)
GTλ
Disc. Returns
rt
γ(1−dt)
GT
Elig. Traces (fwd)
∇wV^(st,wt)
γλ(1−dt−1)
z−1=0
Table 1: Algorithm-specific coefficients for the shared affine scan, applied in reverse or forward time as appropriate. dtterm,dttrunc∈{0,1} flag termination and truncation; dt=dtterm∨dttrunc is the flag used in βt , combining both. Δt=vt−V(st) 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( λ ) but nonzero for TD( λ )/discounted returns). Eligibility traces use vector notation since the trace is vector-valued; see Section 9.2 for per-call scalar handling.
Step
Reads
Writes
Fused pipeline (single Triton kernel)
Load r , v (read twice at offsets rev and rev+1 ), d , bs ; compute α , β in-register; write A
5
1
Total
6
Unfused pipeline (three GPU operations)
Compute δ=r+γvt+1(1−d)−vt ( δ is GAE’s αt ; vt+1 reads bs at the boundary column, a 5th tensor)
5
1
Compute β=γλ(1−d)
1
1
Table 2: HBM accesses for GAE (with bootstrap values, no truncations). r : rewards; v : values; d : terminateds; bs : bootstrap values; δ : TD residuals; β : decay scalars; A : advantages.
H100 80GB HBM3
RTX 2000 Ada
Algorithm
Loop
Plain
With trunc.
Loop
Plain
With trunc.
GAE
317 ×
2.36 ×
1.7 ×
281 ×
2.57 ×
1.6 ×
V-Trace
96 ×
2.78 ×
2.7 ×
86 ×
3.46 ×
3.3 ×
Retrace( λ )
72 ×
1.91 ×
1.9 ×
16 ×
1.62 ×
1.6 ×
λ -returns
267 ×
2.85 ×
2.6 ×
251 ×
5.14 ×
4.4 ×
Discounted returns
167 ×
2.70 ×
2.6 ×
171 ×
5.70 ×
4.6 ×
Table 3: Production-regime speedup ( num_envs=4096 , 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 ∼ 88–95% at the smallest shapes to ∼ 11–37% at the largest. Bottom : full-call time decomposed into launch/wrapper and device components. rl-triton ’s launch/wrapper cost remains ∼ 0.026–0.035 ms, while the baseline’s is 1.4–3.1 × larger.
H100 80GB HBM3
RTX 2000 Ada
num_envs
seq_len
GAE
Retrace
GAE
Retrace
64
512
2.7 ×
2.2 ×
2.7 ×
2.0 ×
128
1024
3.2 ×
2.3 ×
3.1 ×
1.8 ×
256
1024
3.3 ×
2.3 ×
3.0 ×
1.4 ×
512
128
2.8 ×
2.0 ×
2.5 ×
1.8 ×
512
512
2.9 ×
2.2 ×
2.8 ×
1.4 ×
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
∼2.2% / ∼1.02×
<0.1% / ∼1.00× (within noise)
(128, 128)
16384
∼10 – 15% / ∼1.11 – 1.16×
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 ( 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 seq_len≥32 except (num_envs,seq_len)=(32768,32) ; PufferLib is faster at all configurations with seq_len≤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
H100 80GB HBM3
RTX 2000 Ada
num_envs
seq_len
V-Trace
λ -ret.
Disc. ret.
Elig.
Prefix
V-Trace
λ -ret.
Disc. ret.
Elig.
Prefix
64
512
3.0 ×
3.0 ×
3.0 ×
2.8 ×
2.8 ×
2.9 ×
2.9 ×
2.9 ×
2.6 ×
2.5 ×
128
1024
3.0 ×
3.3 ×
3.3 ×
2.7 ×
2.7 ×
2.9 ×
3.1 ×
3.1 ×
2.6 ×
2.5 ×
256
1024
3.0 ×
3.2 ×
3.2 ×
2.8 ×
2.8 ×
2.9 ×
3.6 ×
3.3 ×
2.6 ×
2.6 ×
512
128
2.9 ×
2.9 ×
2.8 ×
2.5 ×
2.4 ×
2.8 ×
2.7 ×
2.6 ×
2.3 ×
2.2 ×
512
512
3.1 ×
3.1 ×
3.1 ×
2.8 ×
2.7 ×
3.0 ×
3.3 ×
3.3 ×
2.6 ×
2.5 ×
Appendix
Table 6: Full twelve-shape sweep for V-Trace, TD( λ )/ λ -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.
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 Fast1 threshold, outperforming the strongest prior RL-trained model, Dr. Kernel-14B.
Multi-turn agents solve complex tasks through extended sequences of tool interactions before producing a final answer, making credit assignment a fundamental challenge during post-training. Outcome rewards provide reliable supervision for short-horizon reasoning, but become sparse and high-variance as trajectories grow to tens or hundreds of tool calls. They can also be misleading: a failed rollout may contain many useful actions that move the agent closer to the goal, yet outcome-only training assigns them the same negative advantage as the eventual mistake. We propose TRACE (Turn-level Reward Assignment via Credit Estimation), a dense credit-assignment method for agentic reinforcement learning. TRACE represents rollouts as state transitions at tool-call boundaries, obtains gold-answer log-probabilities from a frozen reference model, transforms them into log-ratio state values, and derives per-action rewards as Temporal-Difference changes in those values. This requires no additional critic or process-label training, and its one-step log-ratio TD component telescopes across redundant tool calls. On long-horizon complex search, TRACE substantially improves base-model tool-use ability using pure RL, without a cold-start supervised fine-tuning stage, an agentic mid-training stage, or training on live-web data. On the closed-web BrowseComp-Plus benchmark, it raises Qwen3-4B from 7.2 to 35.6 and Qwen3-30B-A3B from 8.4 to 42.6. The learned search behavior also transfers to open-web benchmarks, and the learning curves show earlier improvement and faster convergence during RL training.
Leitian Tao, Baolin Peng, Wenlin Yao +5
University of Wisconsin–Madison · Microsoft Research
Multi-step agentic reinforcement learning benefits from fine-grained credit assignment, yet existing approaches offer limited options: critic-free methods like GRPO assign a uniform advantage to every action in a trajectory, while learned value networks introduce notable overhead and can be fragile under sparse rewards. We observe that group rollouts targeting the same problem often traverse overlapping intermediate states, implicitly forming a tree whose branches diverge at successive decision points. Building on this insight, we introduce Rollout-Tree Monte Carlo (RTMC) advantage estimation, which aggregates return statistics across rollouts sharing a common state to produce per-step Q-values and advantages--without any learned critic. A state-action signature system compresses raw interaction histories into compact, comparable representations, making cross-rollout state matching tractable. On SWE-bench Verified, RTMC improves pass@1 by 3.2 percentage points over GRPO.