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

• URL: http://arxiv.org/abs/2608.17641v2
• Date: Wed, 19 Aug 2026 11:29:40 GMT
• Title: rl-triton: High-Performance Triton GPU Kernels for Reinforcement Learning Credit Assignment
• Abstract summary: We present rl-triton, an open-source library of high-performance GPU kernels for reinforcement learning credit assignment.<n>The core contribution is a unified associative scan framework that recasts seven distinct estimation algorithms.<n> Benchmarks show a 1.6-5.70$times$ full-call speedup over a vectorized torch-compile baseline in the parallel simulation regime.
• Score: 0.0
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Xe-Forge: Multi-Stage LLM-Powered Kernel Optimization for Intel GPU [38.81440160993858]
  Porting deep learning algorithms to new hardware accelerators requires developers to apply the same low-level optimizations to every Triton kernel in their code-base.<n>We present Xe-Forge, a multi-stage LLM-powered pipeline that automates this process for Intel GPU.<n>A curated knowledge base encodes Intel GPU constraints that are absent from LLM training data, keeping the model within architecturally valid bounds.
  arXiv  Detail & Related papers  (2026-04-16T20:21:01Z)
• PolyKAN: Efficient Fused GPU Operators for Polynomial Kolmogorov-Arnold Network Variants [10.239332579225522]
  Kolmogorov-Arnold Networks (KANs) promise higher expressive capability and stronger interpretability than Multi-Layer Perceptron.<n>We present a GPU-accelerated operator library, named PolyKAN, which is the first general open-source implementation of KAN and its variants.
  arXiv  Detail & Related papers  (2025-11-18T19:05:16Z)
• Accelerated Evolving Set Processes for Local PageRank Computation [75.54334100808022]
  This work proposes a novel framework based on nested evolving set processes to accelerate Personalized PageRank computation.<n>We show that the time complexity of such localized methods is upper bounded by $mintildemathcalO(R2/epsilon2), tildemathcalO(m)$ to obtain an $epsilon$-approximation of the PPR vector.
  arXiv  Detail & Related papers  (2025-10-09T09:47:40Z)
• Near-Optimal Online Learning for Multi-Agent Submodular Coordination: Tight Approximation and Communication Efficiency [52.60557300927007]
  We present a $textbfMA-OSMA$ algorithm to transfer the discrete submodular problem into a continuous optimization.<n>We also introduce a projection-free $textbfMA-OSEA$ algorithm, which effectively utilizes the KL divergence by mixing a uniform distribution.<n>Our algorithms significantly improve the $(frac11+c)$-approximation provided by the state-of-the-art OSG algorithm.
  arXiv  Detail & Related papers  (2025-02-07T15:57:56Z)
• A User's Guide to $\texttt{KSig}$: GPU-Accelerated Computation of the Signature Kernel [12.111848705677138]
  The signature kernel is a positive definite kernel for sequential and temporal data.<n>In this chapter, we give a short introduction to $textttKSig$, a $textttScikit-Learn$ compatible Python package that implements various GPU-accelerated algorithms for computing signature kernels.
  arXiv  Detail & Related papers  (2025-01-13T09:11:13Z)
• Low Depth Phase Oracle Using a Parallel Piecewise Circuit [3.629687485125086]
  We explore the important task of applying a phase $exp(i,f(x))$ to a computational basis state $left| x right>$.<n>The closely related task of rotating a target qubit by an angle depending on $f(x)$ is also studied.<n>We find that strategies prioritising execution speed can achieve circuit depth as low as $O(logn+logS)$ for a register of $n$ qubits and a piecewise approximation of $S$ sections.
  arXiv  Detail & Related papers  (2024-09-06T19:57:13Z)
• TCNCA: Temporal Convolution Network with Chunked Attention for Scalable Sequence Processing [52.64837396100988]
  MEGA is a recent transformer-based architecture, which utilizes a linear recurrent operator whose parallel computation, based on the FFT, scales as $O(LlogL)$, with $L$ being the sequence length. We build upon their approach by replacing the linear recurrence with a special temporal convolutional network which permits larger receptive field size with shallower networks, and reduces the computational complexity to $O(L)$. We evaluate TCNCA on EnWik8 language modeling, long-range-arena (LRA) sequence classification, as well as a synthetic reasoning benchmark associative recall.
  arXiv  Detail & Related papers  (2023-12-09T16:12:25Z)
• Efficiently Learning One-Hidden-Layer ReLU Networks via Schur Polynomials [50.90125395570797]
  We study the problem of PAC learning a linear combination of $k$ ReLU activations under the standard Gaussian distribution on $mathbbRd$ with respect to the square loss. Our main result is an efficient algorithm for this learning task with sample and computational complexity $(dk/epsilon)O(k)$, whereepsilon>0$ is the target accuracy.
  arXiv  Detail & Related papers  (2023-07-24T14:37:22Z)
• ReSQueing Parallel and Private Stochastic Convex Optimization [59.53297063174519]
  We introduce a new tool for BFG convex optimization (SCO): a Reweighted Query (ReSQue) estimator for the gradient of a function convolved with a (Gaussian) probability density. We develop algorithms achieving state-of-the-art complexities for SCO in parallel and private settings.
  arXiv  Detail & Related papers  (2023-01-01T18:51:29Z)
• VersaGNN: a Versatile accelerator for Graph neural networks [81.1667080640009]
  We propose textitVersaGNN, an ultra-efficient, systolic-array-based versatile hardware accelerator. textitVersaGNN achieves on average 3712$times$ speedup with 1301.25$times$ energy reduction on CPU, and 35.4$times$ speedup with 17.66$times$ energy reduction on GPU.
  arXiv  Detail & Related papers  (2021-05-04T04:10:48Z)
• Asynchronous Advantage Actor Critic: Non-asymptotic Analysis and Linear Speedup [56.27526702716774]
  This paper revisits the A3C algorithm with TD(0) for the critic update, termed A3C-TD(0), with provable convergence guarantees. Under i.i.d. sampling, A3C-TD(0) obtains sample complexity of $mathcalO(epsilon-2.5/N)$ per worker to achieve $epsilon$ accuracy, where $N$ is the number of workers. Compared to the best-known sample complexity of $mathcalO(epsilon-2.5/N)$ for two
  arXiv  Detail & Related papers  (2020-12-31T09:07:09Z)
• Streaming Coresets for Symmetric Tensor Factorization [9.181791777532608]
  We show how to factorize tensors efficiently in the streaming setting. We introduce two novel algorithmic techniques: online filtering and kernelization. We show applications of our algorithms in learning single topic modeling.
  arXiv  Detail & Related papers  (2020-06-01T19:55:34Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.