FlashRNN: I/O-Aware Optimization of Traditional RNNs on modern hardware
Authors: Korbinian Pöppel
Maximilian Beck
Sepp Hochreiter
Venue: ICLR 2025 Poster (2025)
Keywords: RNN
LSTM
Sequence Models
GPU
Hardware Optimization
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: We optimize traditional RNNs with memory mixing close to transformer levels on modern hardware.

Abstract: While Transformers and other sequence-parallelizable neural network architectures seem like the current state of the art in sequence modeling, they specifically lack state-tracking capabilities. These are important for time-series tasks and logical reasoning. Traditional RNNs like LSTMs and GRUs, as well as modern variants like sLSTM do have these capabilities at the cost of strictly sequential processing. While this is often seen as a strong limitation, we show how fast these networks can get with our hardware-optimization FlashRNN in Triton and CUDA, optimizing kernels to the register level on modern GPUs. We extend traditional RNNs with a parallelization variant that processes multiple RNNs of smaller hidden state in parallel, similar to the head-wise processing in Transformers. To enable flexibility on different GPU variants, we introduce a new optimization framework for hardware-internal cache sizes, memory and compute handling. It models the hardware in a setting using polyhedral-like constraints, including the notion of divisibility. This speeds up the solution process in our ConstrINT library for general integer constraint satisfaction problems (integer CSPs).
We show that our kernels can achieve 50x speed-ups over a vanilla PyTorch implementation and allow 40x larger hidden sizes compared to our Triton implementation. We have open-sourced our kernels and the optimization library to boost research in the direction of state-tracking enabled RNNs and sequence modeling here: https://github.com/NX-AI/flashrnn

Published as a conference paper at ICLR 2025
FLASHRNN: I/O-AWARE OPTIMIZATION OF TRADI- TIONAL RNNS ON MODERN HARDWARE
Korbinian Pöppel & Maximilian Beck & Sepp Hochreiter Johannes Kepler University, NXAI Lab and NXAI GmbH
Altenberger Str. 69 A-4040 Linz, Austria
(poeppel, beck, hochreit jku,
ABSTRACT
While Transformers and other sequence-parane parallelizzalleeee neural network architec- tures seem like the current state of the art in sequence modeling. they specifi- cally lack state-tracking capabilities. These are important for time-series tasks and logical reasoning. Traditional RNNs like LSTMs and GRUs, well as mod- em variants like sLSTM do have these capabilities at the cost of strictly sequential processing. While this is often seen strong limitation, these networks can get with our hardware-optinizztion FlashRNN in Triton and CUDA, optimizing kernels to the register level on modern GPUs. We extend traditional RNNs with a parallclization variant that processes multiple RNNs of smaller hidden state in parallel, similar to the head-wise processing in Transform- ers. To enable flexibility on different GPU variants, we introduce a new optimiza- tion framework for hardware-internal cache sizes, memory and compute handling. It models the hardware in a setting using polyhedral-like constraints, including the notion of divisibility. This speeds up the solution process in our ConstrINT library for general integer constraint satisfaction problems (integer CSPs). We show that our kernels can achieve 50x speed-ups over a vanilla PyTorch imple- mentation and allow 40x larger hidden sizes compared to our Triton implemen- tation. We have open-sourced our kernels and the optimization library to boost research in the direction of tracking enabled RNNs and sequence modeling here: htttps///ithub. com/NX AIfflashrnn
show how fast
INTRODUCTION
Sequence models are at the core of many applications like time-series modeling, natural language processing, text, audio and video models, and predictions for physical systems based ODEs or PDEs (Vaswani et al.) 2017 Degrave et al. 2022, Nearing et al. 2024). While there are modern eeguence-parallelizable architectures like the Transformer (Vaswani et al. 2017). Mamba (Gu & Dao 2023) or mLSTM (Beck et al. 2024). these lack the state-tracking capabilities (Merrill et al., 2024 Merrill Sabharwal 2023 Delétang et al. 2023) of traditional RNNs like LSTM (Hochreiter & Schmidhuber 1997). GRU (Cho et al. 2014) and sLSTM (Beck et al. 2024).
Traditional RNNs include recurrent connection or memory mixing, that connects the previous hidden state in non-linear way to the current state update and this way mixes the states of different memory cells. While the sequence has to be processed step by step. computed hidden states and the recurrent matrix weights can stay cached, enabling large speed optimization. In this work, we introduce FlashRNN as a generic hardware-optimizzed library for these RNN-style architectures.
Our library facilitates research in the direction of state-tracking enabled RNN architectures, in two ways: Firstly, enables easier and more efficient use recent RNN-aachittcttrrs like sLSTM (Beck) et al. 2024). This includes the notion of block-diagonal recurrent matrices that can speed up net- works while lowering the number of parameters. Secondly, it can be easily extended to novel RNN- like architecture variants, it supports generic state and gate numbers per cell. The LSTM (Hochre- iiter & Schnidhuber [997 Gers et al.) 1999). with its two states and four gates (we consider the cell update fourth "gaae" for simplicity here), can be implemented easy simple Elman-RNN

---

Published as a conference paper at ICLR 2025
with one state and one gate (Elman 1990). sLSTM with its three states and four gates (Beck al. 2024).
To realize the shown speed-ups. fuse the recurrent matrix multipliattiiiiiiinnnn part with the point activation part, both wrapped the sequential loop into one kernel. This can be used on different GPUs and with different state/gate variants, operations automatically based on the models' hidden sizes and the cache and register sizes of the hardware.
library optimizes internal memory sizes and
For the auto-optimization we introduce an integer constraint satisfaction library ConstrINT. With this library, one model generic integer CSP problems with equality, inequality and divisibil- ity constraints as these can model size constraints on modern hardware with specific tensor-core, register and SRAM memory sizes.
GPU
Rh
Rh
degister
GPU SRAM
GPU HBM
Main CPU
Co
Go
Memory
Fused
Alternating
Sequence Loop
Write HBM
Pointwise Non-Linearity
Matrix Multiplication Addition
Kernel Boundary
Read from HBM
Figure 1: FlashRNN Kernel overview: Left: Basic Memory Hierarchy in modern GPUs. Center: Fused Kernel (forward) leveraging all caching options for maximal speed. Right: Alternating Ker- nels (forward) for maximum hidden sizes, with kernel calls per time step. The colors show the caching level of the different tensors, the batch dimension is depicted to the right (except for R), the hidden gate dimension vertically.
RELATED WORK
Hardware-aware algorithms and their open-source implementations common sequence modeling primitives have been focused primarily around the Transformer architecture (Vaswani et al.) 2017 and its attention operation because of its ubiquity in language modeling. FlashAttention (Dao et al. 2022) introduced an IO-aware attention algorithm and CUDA implementation that uses tiling duce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM, and achieves significant memory savings. FlashAttention2 (Dao 2024) improves FlashAttention with better work partitioning and the additional parallelization over the sequence di- mension. FlashAttention3 (Shah et al.) (024) takes advantage new capabilites, such as asynchrony and FP8 low precision support of the recent NVIDA Hopper GPU generation.
Recently, novel sequence models taking inspiration of Linear Attention (Katharopoulos et al. 2020) have shown promising performance compared Transformer Attention (Beck et al. 2024 Yang ctt al. 2024: Dao Gul 2024). Yang et al. (2024) provide an hardware-efficient algorithm and implementation in Triton for Gated Linear Attention that trades off memory movement against par- allelizability and show that it faster than Flastontention2.
Traditional RNNs like LSTMs (Hochreiter & Schmidhuber 1997) GRUs (Cho et al. 2014)
still widely used in many applications, such as for example time series modeling or reinforcement learning (Nearing al. 2024 Degrave al. 2022 Many of these applications rely optimized
2

---

Published as a conference paper at ICLR 2025
closed-source implementations of these RNN operations such as in the NVIDIA cuDNN library, which integrated in PyTorch. [Sharvi 2020 provide an open-source altemaatve in CUDA for specific LSTM and GRU variants in their HASTE library, which served as inspiration for this work. HASTE is limited in speed due to a sequence of alternating calls of matrix multiplication and point- wise kernels, well as its limitation to higher (but slower) precision.
Our work FlashRNN overcomes this limitation by fusing the recurrent matrix multiplication with the pointwise operations into a single persistent kernel with custom caching of the recurrent weights in registers, FlashRNN also supports the bfloat dtype and block-diagonal recurrent matrices. By open-sourcing our CUDA and Triton kernels we aim to enable researchers to quickly reach similar speeds compared to optimized closed source libraries.
3 GENERIC RECURRENT NEURAL NETWORK ARCHITECTURE WITH MEMORY MIXING
A generic RNN architecture that aim optimize has N, states s( Rd, and Ng gates (or pre-activations) goو) E Rdd, with d being the embedding dimension or hidden size of the RNN. For example the LSTM (Hochreiter & Schmidhuber 1997) has N, 2 states and Ng 4 gates.
Each gate receives an input (+1) E Rdd As learnable parameters, the gates have recurrent matrix Rú) E Raxd that models the dependency on the previous hidden state s, and bias b(1) Rd The state sequence of the RNN is then defined as:
(0) 9: xx + RD0) s +
(4) p(i)
( (8()) PE(I..N,)
( SE(1.NN))
with a point-wise / element-wise function p(o) that does not mix different cells along the vector dimension (unlike the recurrent weight). In Appendix A we show how this generic formulation translates to the most common RNN variants.
Usually for these networks, the input is modified with another weight matrix W. We omit this here as it can be moved outside of the basic kernels. In the common training setting, where the whole sequence is given as input, the weight matrix W can be applied parallel to all timesteps before processing sequence in the RNN. Our runtime experiments in Section 6.1 show that this operation has only marginal impact on the overall runtime.
GENERIC GRADIENT FOR BACK BACK-PROOE PROPAGATION THROUGH TIME
In back paco-propagttion through time Mozer 1995), the backward pass of this RNN architecture can be recursively defined as well. The backward pass reads:
Op!))
(9)))
8gو (
kE(1..N))
SE(L..N,)) 88,
(3)
0g!
Op(!)
(((
(
(1)
JE(L..N))
8s (4)
kE (.)
(1) +
Σ
R
if i=
(
(t)
(E((I..Ng)
(4)
The structure of the gradient shows that, also for the backward pass, we have an alternation of wosnt-wise operations (left) and matrix multiplication (right).
The input gradient is equal to the gate gradients, the bias gradient the sum of the input gradients and the recurrent weight matrix gradient is the time-wise sum of the outer product of gate gradients
3

---

Published as a conference paper at ICLR 2025
with the state values:
)x) (1) 5gg (1)
(5)
8b = Σ 5g. (1)
(6)
8R Σ ög. St
(3) (0)
4.1 VANISHING AND EXPLODING GRADIENTS AND GRADIENT MODIFICATIONS
For neural network be stably trainable, there must not be exploding gradients, also vanishing gradients should be prohibited for long context sequence modeling (Hochreiter & Schmidhuber, 1997]. Still, for the generic structure of Equations [3] there can be exploding components: Firstly, one more eigenvalues of the point-wise function Jacobian can be greater than one in magnitude. This can be mitigated by proper choice of the point-wise function. Secondly, the combination
of recurrent matrix and gate gradients with the gradient app()) R could have singular values of
وg;5)
magnitude > 1. This case cannot be excluded directly, as the recurrent matrix consists of trainable, typically unconstrained weights. However, for practical training this is rarely a limitation.
In our library, we implement a simple approach for mitigating this at the cost of additional gradient noise, clipping the gradient values on a scalar level after each time step. Specifically, we clip the be
term containing the recurrent matrix within pre-defined magnitude. The gradients cut to zero, leading to typically worse convergence at the henefit of faster training, as the recurrent matrix part in Equation iss cut zero for the backward pass.
4.2 HEAD-WISE PARALLELIZATION
When increasing the size of a neural network, typically the width, i.c. the embedding dimension or hidden size is increased. Vaswani et al. (201)) found that for the attention operation it is beneficial to linearly project the input embedding vectors into multiple smaller input vectors, the so called heads, and then perform attention each of these small vectors in parallel. This parallelization primitive enables also efficient implementations on GPUs, since each head can be computed in different thread blocks of the GPU (Dao al., 2022 in parallel (see also Section 5.1)
Many more recent architectures also rely on this hesd parallelization primitive (Beck et al, 2024 Yang et al. 20244] Dao & Gu 2024), where the embedding or hidden vector of dimension is split into Nheads heads of smaller dimension dhead = d/Nneads each of which is processed independently inside the sequential part. In FlashRNN, apply this primitive to traditional RNNs by dividing the recurrent matrix R into multiple blocks or heads Rhend E Rddeadd Xdhead rendering the recurrent matrix R block-diagonal matrix.
5 HARDWARE- IMPLEMENTATION
5.1 GPU-ACCLERATED COMPUTING
Modern compute hardware the form of GPUs enables massive parallelization and accelerated matrix multiplication. This means that both point-wise (scalar) operations can be parallelized and matrix multiplications have good support via BLAS-like libraries (Lawson et al. 1979: Thakkar
al. 2023). as used for RNN training workloads as defined above.
Execution Model Specifically, a modern GPU consists of larger computational super-units (i.e. streaming multiprocessors (SMs) that have some faster memory attached to them. There are three levels of memory, the large HBM that allows global random access from all computational units at the cost of low speed (still fast compared to CPU RAM access), the SRAM that is attached to one computational super-unit and the registers that tied to smallest computational unit (i.e. thread). One anper-unit usually supports up to 1024 threads in parallel (with varying register sizes), which typically referred to as a block thread block. Multiple blacks executed in parallel mul-
4

---

Published as a conference paper at ICLR 2025
? An NVIDIA H100, for example, consists of 132 streaming multiprocessors, with 256 KB SRAM per SM and a SRAM bandwidth of around 33 TB/s (Spector et al., 2024), compared to the up to 3 TB/s for access to the 80 GB of HBM. Starting from the NVIDIA Ampere Architecture and newer, there is hardware acceleration for asynchronous loading and SRAM interconnection, which we did not utilize in this work. [ Beyond the memory levels, a computational super-unit allows for
tiple super-units are called the grid.
matrix multiplication (e.g. via Tensor- Cores, "wmma" operation). Typically, it is divided into sub-umits (warps) of certain number of threads (NVIDIA: 32) that act as one for matrix multiplication. There are certain size limita- tions for this acceleration, which have to be considered in the kemne optimization process. For a NVIDIA H100, this means that only minimal matrices of sizes 32x 1688, 16x 6 16 or 8x16x32 can be multiplied for the low-precision 6 or f dtypes, larger matrix multiplications have to be composed of those, by parallelization along the outer dimensions and summation along the accumulating dimension.
Performance measures The specific limitation of computational load falls into two regimes: Being compute-bound or being memory-bound. In the former case, the arithmetic intensity is high, there are many compute operations per loaded byte and therefore, the main limitation is the compu- tational part. In the latter case, arithmetic intensity is low and the bottleneck is the memory access to load inputs and store outputs (Williams al. 2009). Small operations, like applying an activation function in parallel are typically memory bound and should be grouped together into a fused kernel.
Fused Kernels To minimize HBM memory accesses, one combines multiple arithmetic operations in one GPU kernel. A kernel is set of instructions on the GPU. which is executed in parallel on its parts. Only within the execution of one kernel SRAM and registers are kept and can serve as a cache. Therefore, for memory-boume operations it is helpful to fuse multiple arithmetic operations into one kernel to leverage these lower cache levels. While compilers can fuse point-wise operations, an alternation of both point-wis, computations and matrix multiplication is non-trivial.
5.2 FLASHRNN KERNELS
As the RNN operations of Equations and 3 are a sequential alternation between matrix multi- plication and pointwise non-linearities, there simple speed up variant that optimizes these two primitives separately. Our library implements this variant, in the alternating backend. This enables arbitrarily large head dimensions (to the limits of HBM GPU memory). Also, a vanilla PyTorch implementation relying on auto-grads will work in this primitive, but for every time step separate state is saved for the backward pass, leading to inefficiencies beyond memory accesses. Moving the time-loop into CUDA can already give large speedups over the vanilla Py Torch implementation.
The downside of the alternating implementation is that there are no I/O optimizations beyond a single time step. For every time step, the current input and last state, as well as the recurrent matrix and the biases have to be re-loaded. However, both the recurrent matrix R and the biases remain the same for the whole time loop and the previous states stay in memory as they the previous time step. Since the structure of the computation remains the same over the time steps, one can even store most of these values in registers. Registers have the highest memory bandwidth and, while they can only be used within the lowest computation unit (threads), their total size on a GPU is comparable to the SRAM (both 256 KB per SM on H100).
computed in
To reach the maximum speed, we implement FlashRNN fused kernels that store the recurrent matrix R and the biases b in registers (and SRAM if register memory exceeded). The matrix multiplica- tion results are stored and accumulated in shared memory (or HBM if SRAM sizes are exceeded). In the forward pass, the computations are mainly tiled along the gate dimension (or the dimension of the new hidden states). This way, we use the maximum amount of memory along the previous state dimension. This dimension is the accumulating dimension of the recurrent matrix multiplica- tion. For the backward pass, the computations are typically tiled along the previous state gradient dimension, such that the gate dimension, which is accumulated over, is minimally tiled. Algorithm shows simplified representation of the forward pass in pseudo-code and in Appendix Section B this algorithm is shown in more detail.
tepppreeeeeettpppeeee-eaeeee hopper
5

---

Published as a conference paper at ICLR 2025
Algorithm FlashRNN-fused forward pass
All states tiled along threads (single ALU) in Warps (for e.g. Matrix Multiplication) in block (SRAM level, streaming multiprocessor) and blocks in the grid (multiple streaming multiprocessors) additionally there can be looping levels where the parallelization resolved to simple loop. Dimensions are: b: batch, t: time, g: gates, s/: previous/new state
Require: Recurrent matrix Rys. inputs Etbg. biases bg
Require: Initial states Spbs
Load Rys. by to registers and SRAM
for l, in LB do
Load 80bs registers for / E 0..7 do
for Matrix Tiles in Registers do
Calculate and Accumulate Matrix product Ythg = R,g the along $
end for
for Matrix Tiles in SRAM do
Load Matrix Tile of Rgs
(0) Calculate and Accumulate Matrix product Ytbg Rgs ths along end for
Accumulate MatMul results Ytby along in shared memory (Write, Load and Sum) state dimension too big for SRAM then Accumulate MatMul results Ytbg along in HBM (Write, Grid Sync, Load, Sum)
end if
Sum Gate inputs Biby with Ytby and biases by gates gtby
Compute Point-wise Function 8 ths" = ((stss' 9rbg) with aligned states s' and gates Write gates for backward pass and new states HBM
Grid-Level Sync (for new states to be available across the whole grid)
end for end for
5.3 TRITON IMPLEMENTATION
With FlashRNN we also implement Triton variant of the fused FlashRNN kernel. Triton is domain specific language and compiler for parallel programming that provides a Python-based en- vironment for writing custom GPU kernels.
For the Triton kernel we parallelize the computation over two dimensions the batch dimension and the head dimension. See Appendix [ for detailed description of the Triton implementation in Al Algorithm 5 of the FlashRNN algorithm. As described in section 4.2 we partition the embedding dimension into multiple heads and compute each head in parallel in different programs (or thread blocks) with no synchronization in between these programs. In Triton, each program (which corre- sponds to thread block CUDA) will hold its recurrent weight matrix Rhead and bias bread in SRAM. In contrast to CUDA, Triton gives access registers the GPU. Therefore, cannot apply the custom caching strategy of the fused CUDA kernels and instead rely on Triton for manag- ing the shared memory and register cache. Additionally, there is (grid) synchronization between programs in Triton, which makes it impossible to communicate values between different programs HBM. In section 6.1 find that this poses limitation the maximum head dimension of 128 for the forward pass and 64 for the backward pass on a NVIDIA H100 GPU.
The recurrent matrix multiplication in equation I and 3 implemented with Triton's matrix multiply operation t1. dot, which gives an interface to the Tensor Core units on GPUs. In Triton, the minimum block size of these matrix multiplies 16x16, which gives a limit on the minimum batch size. In practice, we enable smaller batch sizes by padding zeros at the cost of efficiency.
https://tritonn11ae http://ttrittonnnn...ro

---

Published as a conference paper at ICLR 2025
5.4 AUTOMATIC TUNING OF TILING AND LOOPING DIMENSIONS
While Algorithm| | describes the algorithmic behaviour, the tile, block and grid sizes and loop itera- tions depend on the specific hardware architecture, i.e. the number of computational super-units (streaming multiprocessors). the SRAM per super-unit, the sizes of matrix-multiphae maltiiixuulippiaattion units, threads (warps and threads) per super-unit and the number of registers per thread. On NVIDIA H100s (and most other NVIDIA GPUs), there is a varying amount of registers per thread, depend- ing on the block size used, while the total number of registers per SM is physically fixed.
These physical constraints can now be reformulated as equalities, inequalities and divisibility con- straints inside an integer constraint satisfaction problem (integer CSP). Typically this optimization is done via polyhedral constraint optimization in compilers (Baghdadi et al. 2018). For solving these constraints FlashRNN, we implement an efficient solver ConstrINT Python for general integer CSPs going over large number ranges and including the notion of divisibility constraints, which are needed to model the minimal matrix sizes.
For more details on the solution algorithm, see Appendix Section
EXPERIMENTS
In Section 6.1 we benchmark the runtime of our FlashRNN kernels and compare against the LSTM and Attention implementations provided in PyTorch, In Section(6.2) we measure training time with FlashRNN kernels on language modeling. Finally, in Section 6.3 we confirm that traditional RNNs like LSTM and sLSTM implemented in FlashRNN can solve state tracking problems.
6.1 RUNTIME BENCHMARK
We evaluate the runtime of backends of our FlashRNN library that implement an LSTM:
CUDA fused: CUDA implementation that fuses matrix multiplication and pointwise oper- ations of the LSTM in single kernel that is persistent over all time iterations.
CUDA alternating: CUDA implementation that implements the time loop in C++ and alternates between matrix multiply kernel and LSTM pointwise kernel.
Triton fused: Triton implementation that fuses matrix multiplication and pointwise opera- tions similar to CUDA fused.
Vanilla Py Torch: Py implementation of the LSTT operation with our custom back- ward pass implementation, which is faster than the PyTorch autograd backward pass. We do not use torch. compile due to very long compile times.
We compare our backends two references from Py PyTorch and the haste library (Sharvil 2020
FlashAttention2: PyTorch Attention with FlashAttentions backend. Note that FlashAt- tention2 is not a recurrent operation and can be parallelized across batch, head, and se- quence dimension on the GPU. FlashAttention2 does not fall into the category of RNNs, which FlashRNN aims to speed up, and is not able to solve state tracking tasks. Therefore, in our benchmarks it should be seen as a widely adopted reference to better interpret the runtimes instead of a direct bascline that we aim to outperform.
nn.LSTM: Py PoTor LSTM with NVIDIA cuDNN as backend. In contrast to our FlashRNN LSTM, nn. L.STM also integrates the gate pre-activation computation into the function call (not kernel call), which we do not (see Section (3). In Section [H.4 in the ap- pendix, we provide a comparison to the combination of linear layer and our FlashRNN LSTM kernel with nn LSTM. Moreover, nn. LSTM does not support multiple heads on the embedding dimension as described in Section 4.2 nn. LS TM always uses a single head.
haste: The haste library is implementation of LSTM and GRU and variations in CUDA, using alternating kernels between pointwise and matrix multiplication operations.
nn funcct ona sca ed__
dot Dot_product_actentiion

---

Published as a conference paper at ICLR 2025
Its last release was in 2020, with compilation support for Ampere or later architectures in the standard setting" It solely supports float32 and float64 precision and does not have multi-head option.
Setup. We assess the impact of the input dimensions batch size (B), sequence length (T) and head dimension (DH) and number of heads (NH). The number of heads together with the head dimension give the embedding dimension d NH DH. Except for PyTorch nn .ITMM we run all runtime experiments with bfloat precision. For nn. LSTM we use float16 precision, since this precision yielded the fastest runtimes. For every runtime measurement we do 25 warmup iterations and then report the average across 1000 iterations on NVIDIA H100 GPUs. We use Py PyTorch 2.4 and with CUDA version 12.4 for our experiments. Further details and additional experiments can be found in Section H in the appendix.
Head dimension. We measure the runtime of all of our FlashRNN kernels and our two references FlashAttention22 and Py PyTorch nn. LSTM for different head dimensions. We fix the embedding di- mension d' NH x DH to 768 and vary the head dimension from 16 to 768. We use batch size 16 and sequence length 1024. In Figure Σ we report the runtime of each the forward pass only on the left and the forward combined with the backward pass. FlashAttentions does not allow for head dimension larger than 256, due shared memory limitation. The PyTorch mn. LSTM does not sup- port multiple heads or blockdiagonal recurrent matrices. Therefore, we only report the runtime for a single head of dimension 768, including the gate pre-activation computation. At this dimension, nn L. LM is about times faster than CUDA fused. The Triton kernels are limited to head dimen- sion 128 and 64, but are about two times faster than CUDA fused for small head dimensions 16 and 32. The fused CUDA kernels support all head dimensions up to 768 (actually more, see Appendix Section H. and are about two to three times faster than the alternating kernels.
Vanilla PyTorch
1S.L Hoste 1SM fy32
Vanila PyTorch FlushAteentin??
CUDA fooed
FluehAttention?
CUDA fioeed CUDA altemating
Haste 1TM fp32
CUDA altematingg
160
352.00 05
146.04
16rAr
470.44
A90 81
600
43
()
Time (ms))
Time (osc)
12-
10:2
10-
20
13:4
11
3333
10
2.50
5:0
11
15
11
966
DOSMOO0M
odial
DOSMOOSM
DH-16
DH-64
DH-128 NH#
De 256
DH-768
DH-16
DH
DH-64
DH-128
De
DH-768
Nit:00
Nm::2
NIH=12
Nis
N111
Nit:00
NH::2
NHH=12
NH=6
Ni
N1:1
Figure 2: LSTM Runtime for different head dimensions (DH) and number of heads (NH) on a NVIDIA H100. Overall embedding dimension is fixed at 768. We use batch size 16 and sequence length 1024. Left: Forward pass. Right: Forward backward pass.
Batch size. We measure the runtime of all LSTM kernels while varying the batch size (B) from 2 to 256 at sequence length 1024. F shows the results for NH=12 heads with head dimension DH=64. The CUDA fused backend is optimized for smaller batch sizes and shows 2x speed up over the alternating backend for batch sizes up to 32. For larger batch sizes than 128 CUDA alternating is faster. Figure 4 shows the results for a single head with head dimension DH=768. At this head dimension CUDA fused is still faster than CUDA alternating up to batch size 32. For larger batch sizes, CUDA alternating more than two times faster. Comparing the PyTorch nn LSTM, we find for medium batch sizes from 8 to 64 it is about 2-3 times faster than and CUDA fused and for larger batch sizes about about 30% faster than CUDA alternating.
htttps:////g http:///ttttubccoc
/hatte
8

---

Published as a conference paper at ICLR 2025
fized CUDA foxed
alternating
alternating
CUDA Foved
20.86
74.45
E0
12.17
Time (ms)
Time lmms
10:00
5/38014
10
1:23
18.13051
ta Pa. 21
2.1
2.7
3,
2,1
3.00
10
33
99
28
B-2
B-16
B-64
B-128
B-256
B-2
B-16
B-64
B-128
B-256
Figure 3: LSTM Runtime for different batch sizes (B) on NVIDIA H100. We use 12 heads with head dimension 64 and sequence length 1024. Left: Forward pass. Right: Forward + backward pass.
CUDA fused
PT LSL
Hoste LSTM
CUDA fused
PT m.LSTM
Haste LSTM fp32
CUDA altemsting
CUDA alternating
104.44
221.8
197.60
70.12
150
Time (mos)
60
133
40
80
50
20
16
2135
60
13.08
20:28 40
25.15
776
3.1
14
9
20
115
B-2
B=8
B-16
B-
B-64
B-128
B-256
B=2
B:8
B:16
BB-22
B-64
B=128
B-256
Figure 4: LSTM Runtime for different batch sizes (B) on a NVIDIA H100. We use one head (dhead = 768) and sequence length 1024. Left: Forward pass. Right: Forward + backward pass.
Additional Runtime Experiments. In section in the appendix, we include experiments on varying sequence lengths. We see the expected linear runtime scaling for our N kernels and validate that the above findings transfer to longer sequences. In addition, section H.4 we compare the FlashRNN LSTM kernel in combination with a linear layer that computes the gate pre-activations externally to the PyTorch nn. LSTM baseline, which integrates the gate pre-activation computation. We find that the gate pre-activation computation has only marginal impact on the overall runtime. we provide all runtime results also for the sLSTM (Beck et al. 2024).
Finally, in section H.5
6.2 LANGUAGE MODELING
Even though we do no expect traditional RNNs to outperform Transformers, the language model- ing setting serves as an important benchmark for speed on larger scales. Here, we train models at the 165M parameter scale for a Llama-style Transformer without weight tying, i.e. 12 Transformer blocks with Pre-LayerNorm and Swish-Gated MLP after the attention layer. We replace attention with FlashRNN LSTM and sLSTM layers for a speed comparison. The results show a slowdown of roughly 25 % over attention for equal head dimensions or 140 % for one RNN head, see Table 6.2 (H100) and Appendix Table [ (A1)).. In our experiments, we also compare to the cuDNN imple- mentation of LSTM integrated in PyTorch Morch.nn.LSTM) While it's integration into PyTorch is considerably faster, there are numerical differences to the FlashRNN implementation. With same initialization, FlashRNN LSTMs converge faster in language experiments (both 6 and

---

Published as a conference paper at ICLR 2025
Model
Heads
Param. (M)
Train Time (h)
Median Step (s)
Val PPL (val)
LSTM CUDA fused
190
9.9
0.535
22.1
LSTM CUDA altern.
190
10.8
0.575
21.9
LSTM PT nn LSTM 1
190
4.5
0.285
25.8
LSTM CUDA fused
12
164
5.9
0.325
22.2
LSTM CUDA altern.
12
164
9.6
0.511
22.1
sLSTM CUDA fused
190
10.1
0.543
21.3
SLSTM CUDA altern.
1
190
10.9
0.577
21.4
sLSTM CUDA fused
12
164
6.8
0.342
21.7
SLSTM CUDA altern.
12
164
9.7
0,509
21.8
Transformer
12
162
2.9
0.190
17.9
Table 1: 165M Model training on 15B tokens of SlimPajama on 8xH100s with two gradient accu- mulation steps.
float32), even though the differences in single kernel call are the expected levels of numerical precision. This deviation should be investigated further and suggests the use of FlashRNN even for the established LSTM architecture. We provide an analysis our kernel precision compared float64 baseline in section H..6
For larger models, we expect local batch sizes to be smaller and the effective speed difference for fused kernels to be higher compared to the alternating version as measured in Section 6.1
6.3 STATE TRACKING TASK
To show state tracking capabilities of traditional RNNs contrast to Transformers and State Space Models experimentally, we train implementation the Parity task and evaluate longer quences to measure extrapolation capabilities (Zhou et al. 2024). This serves as a litmus test for state tracking capabilities (Merrill et al. 2024).
Model Acc (Ext.)
Transformer 0.52
Mamba 0.56
mLSTM
Elman 1.00
GRU 1.00
LSTM
sLSTM 1.00
0.54
1.00
Table 2: Parity Task in Sequence Extrapolation: Transformers, State Space Models and mLSTM fails at this task (close to random chance 0.5), while traditional recurrent models can learn to extrapolate. Extrapolation accuracies are averaged over three seeds for the best learning rate.
7 CONCLUSION
The FlashRNN library serves as a fast and extendable implementation of traditional RNNs with a recurrent connection memory mixing. It extends RNNs with the multi-head paradigm introduced by Beck et al. (2024) for sLSTM. FlashRNN provides a speed-up of up to 50x over vanilla Py Torch implementations of RNNs and may recurrent connection.
backbone for future RNN architectures that have
FlashRNN implements two variants, an alternating version switching between point-wise and matrix-multi maadii-multplliaation kernels and fused implementation optimizing memory transfers, while using hardware-optimizer matrix- mutiix-muliplliattion. The second leads to further 3-4x speed-up over the alternating option for small batch sizes. The implementation auto-optimizes for different cache levels via the ConstrINT library custom library solving general integer con- straint satisfaction problems with equality, inequality and divisibility constraints. This library may be re-used for other optimization problems regarding cache sizes on hardware platforms and beyond.
internal sizes
We show that with FlashRNN, traditional RNNs are not too far in speed from Transformers in practice, even though they are not parallelizable along the sequence dimension. In the future, it may be optimized to leverage asynchronous memory operations and SRAe connections recent hardware features that promise further speed ups not realized in this work.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
thank Markus Spanring, Maxim Milakov, Pieter-Jan Hoedt, Günter Klambauer and Fabian Pais- cher for helpful discussions and feedback.
We acknowledge EuroHPC Joint Undertaking for awarding us access to Karolina at IT Czech Republic, McluXina LuxProvide, Luxembourg, Leonardo at CINECA, Italy and LUMI at CSC, Finland. The ELLIS Unit Linz, the LIT Al Lab, the Institute for Machine Learning, are supported by the Federal State Upper Austria. This research was funded in whole in part by the Austrian Science Fund (FWF) (10.55776/COE12), We thank the projects INCONTROL-RL (FFG- 881064), PRIMAL (FFG-873979), S3AI (FFG-872172), DL for GranularFlow (FFG-87 1002),, EPILEPSIA (FFG-892171), FWF AIRI FG 9-N (10.55776/FG9) Al4Grreenteeatingrridd (FFG- 899943), INTEGRATE (FFG-892418), ELISE (H2020-ICT-2019 (HORIZON-CL6- 2021--22---2222 1ILII--LAA---0-00111 We thank NXAI GmbH. AudiJKKU Deep Learning Cen- ter, TGW LOGISTICS GROUP GMBH, Silicon Austria Labs (SAL), FILL Gesellschaft mbH, Any- line GmbH, Google, ZF Friedrichshafen AG, Robert Bosch GmbH, UCB Biopharma SRL, Merek Healthcare KGaÅ, Verbund AG, GLS (Univ. Waterloo), Software Competence Center Hagenberg GmbH, Borealis AG, TÜV Austria, Frauscher Sensonic, TRUMPF and the NVIDIA Corporation.
ID: 951847), Stars4Waters
ETHICS STATEMENT
We use open dataset (SlimPajama) that uses publicly crawled internet data for Language Model training. Our implementation speeds up a certain class of Machine Learning models. This may reduce the environmental impact of the research field, in case these architectures remain important in future research. Also, it may speed up development of Machine Learning research in the direction of recurrent sequence modeling with state tracking capabilities. The further implications of these impacts may or may not be a benefit for society.
REPRODUCIBILITY STATEMENT
We provide the source code for your implementations along with this paper. The detailed training setup for speed tests is described in Section 6.1 For Language Modeling this setup description provided in Appendix Section D and uses the open SlimPajama dataset, for the parity task exper- iments Appendix Section K the training and test data can be synthetically generated using the mentioned distributions.
The observed deviations in language model training compared to the PyTorch LSTT based on cuDNN should be further investigated. The results on A100 and H100, as well as across differ- ent kernels are within the expected small-scale numerical deviations.
The code is released here: https: /tt/t/gttth//t www/XX-AI/Elaatttte library in fackirm//aaouunn///sssinnnn as single-file file implementation and optional caching.
with the ConstrINT
REFERENCES
R. Baghdadi, J. Ray. M. B. Romdhane, E. Del Sozzo, A. Akkas, Y. Zhang, P. Suriana, S. Kamil, and S. Amarasinghe. Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code, December 2018. URL htth:ttt/axxiii org / abs/ /104. 10694 arXiv:1804.1069 [cs].
M. Beck, K. Pöppel, M. Spanring, A. Auer, O. Prudnikova, M. Kopp, G. Klambauer, J. Brandstetter, and S. Hochreiter. XL.SMM: Extended Long Short-Term Memory, May 2024. URI. http: //arxiv. wwwxii/oorg/bb//40055.5 04517 arXiv:2405.04517 [cs, stat].
K. Cho, B. van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation. In Alessandro Moschitti, Bo Pang, and Walter Daelemans (eds.). Proceedings of the 2014 Con- ference Empirical Methods in Natural Language Processing (EMNLP), pp. 1724-172511 Doha, Qatar, October 2014. Association for Computational Linguistics. doi: 10.3115/v 1179. URL htt.lttts://// antto ogt.oog//ooo-11111111g
11

---

Published as a conference paper at ICLR 2025
T. Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, 2024. URL https:// openreview net/ wtwp:/ormmmmmmdmmmmmmmm
T. Dao and A. Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. Forty-first International Conference on Machine Learning, 2024. URL https:://penreeeeeeeeeer// /ttp:e/wwwweeeeeeeeeoeoeeeeeeeee 2o-ounnnid-ttttCCR11
net/ Eorum?
T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré. FlashAttention: Fast and memory-efficien exact attention with IO-awareness. Advances in Neural Information Processing Systems (NeurIPS), 2022.
J. Degrave, F. Felici, J. Buchli, M. Neunert, B. Tracey, F. Carpanese, T. Ewalds, R. Hafner, A. Ab- dolmaleki, de Casas, C. Donner, L.. Fritz, C. Galperti, A. Huber, Keeling, M. Tsimpoukelli, J. Kay, A. Merle, J. Moret, S. Noury, F. Pesamosca, D. Pfau, O. Sauter, C. Sommariva, S. Coda, B. Duval, A. Fasoli, P. Kohli, K. Kavukcuoglu, D. Hassabis, and M. Riedmiller. Magnetic con- trol of tokamak plasmas through deep reinforcement learning, Nature, 602:414-419, 2022. doi: 10.1038/s41586-02 11--0022-103--9
G. Delétang, A. Ruoss, J. Grau-Moya, T. Genewein, I.. K. Wenliang, E. Catt, C. Cundy, M. Hutter, S. Legg, J. Veness, and P. A. Ortega. Neural networks and the chomsky hierarchy. In Eleventh International Conference on Learning Representations, 2023.
Elman. Finding Structure in Time. Cognitive Science, 14(2):179-211, March 1990. ISSN 0364-
0213, 1551-6709. doi: 10. 10.1007// wiley. wtt.:cccmmcoooooo 110.0111011771115510000000000000
URL https://onl inel Sieeliben ibbrrr.
F.A. Gers, J. Schmidhuber, and Cummins. Learning to forget: continual prediction with LSTM. In 9th International Conference on Artificial Neural Networks: ICANN '99, vol- ume 1999, pp. 850-855, Edinburgh, UK, 1999. IEE. ISBN 978-0-85296-721-8. doi:
10. 10.1099//p:199911218 URL wttpt::://// //dggt//////aaaaa Sehhggtal-librrrye wew.ary..hiiiee... www.eeiieeettroooteeetttt....... conferences/10: 11t.0011999pp11990000 1tt.ee/eenee/...119//ppppppppp
A. Gu and T. Dao. Mamba: Linear- Time Sequence Modeling with Selective State Spaces, December 2023. htt::tt::///// //axivv http::////wwwiw...//aaaaas/e/e/ org/ /231 C0752 arXiv:2312.00752 [cs].
S. Hochreiter and J. Schmidhuber. Long Short-Term Memory.
Neural Computation, 9(8):
doi: 10.1162/neco.199 1997.9. 8.1735. URL https::/ wwwww mi tatiipsssssii journals org/ doi www.sss..... 1162/neco. 1997. 9.7.1.35.13
1735-1780, November 1997. ISSN 0899-7667, 1530-888X.
A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are mns:: Fast autoregressive transformers with linear attention. In Proceedings of the International Conference on Machine Learning (ICML), 2020.
C. Lawson R. J. Hanson, D. Kincaid, and F. Krogh. Basic Linear Algebra Subprograms for Fortran Usage. ACM Transactions on Mathematical Software, 5(3):308-335 September 1979. ISSN 0098-3500, 1557-7295. doi: 10.1145/358441.355888 URL https: acm. acm.org/ org/
1145100.145/20
355847
A. K. Mackworth. Consistency in networks of relations. Artificial Intelligence, 8(1)999-118, February 1977. ISSN 00043702. doi: 10101660004-300477790000007 URL https://
linkinghub. e elsevier com http:///w..........///////////10
W. Merrill and A. Sabharwal. The Parallelism Tradeoff: Limitations of Log-recision Transformers. Transactions of the Association for Computational Linguistics, 11:531-545, 06 2023. ISSN 2307- 387X. doi: 10.
/ 2 00562
W. Merrill, J. Petty, and A. Sabharwal. The illusion of state in state-space models. In Forty-first International Conference on Machine Learning, 2024. URL https::/openreeview..t Eorum (ildd=22goJJJpl
M. Mozer. A focused backpropagation algorithm for temporal pattern recognition. Complex Sys- tems, 3, (1 1995.
12

---

Published as a conference paper at ICLR 2025
G. Nearing, D. Cohen, V. Dube, M. Gauch, O. Gilon, S. Harrigan, A. Hassidim, D. Klotz, F. Kratzert, A. Metzger, S. Nevo, F. Pappenberger, C. Prudhomme, G. Shalev, S. Shenzis, T. Y. Tekalign, D. Weitzner, and Y. M. B. Kosko. Global prediction of extreme floods in ungauged watersheds. Nature, 627:559-563, 2024. doi: 10.1088844156660244074445555
J. Shah, G. Bikshandi, Y. Zhang, V. Thakkar, Ramani, and T. Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision, 2024. URL https://axxiv.org/
abs/ fos/2407.0m 08608
N. Sharvil. Haste: fast, simple, and open mn library. htttp:://github.com// haste/ Jan 2020.
com/
B. Spector, A. Singhal, Arora, and C. Ré. GPUs Go Brr, 2024. edw/b/du//bbl bl wwtpbl/g/////aaa-eoo//// 20210022--0552222222
URL https://
hazyresearch.stanne
V. Thakkar, P. Ramani, C. Cecka, A. Shivam, H. Lu, Yan, J. Kosaian, M. Hoemmen, H. Wu, A. Kerr, M. Nicely, D. Merrill, D. Blasig, F. Qiao, P. Majcher, P. Springer, M. Hohnerbach, Wang, and M. Gupta. Cutlass, 2023. https:::://git wiigttttu......... wwt.ccm///////// NVIDIA/cutlasa tree/v3.Cm
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N Gomez, L. Kaiser, and Polosukhin. Attention is All you Need. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL htttps:/pproceseeee wtw.:eeeedddnn.................. www.enee..ceeeeeeeeeeee fi 3lle33f5e22334 3o Papes
I. Guyon, U. Von Luxburg,
https//eepppppp..eeeeeeeeeeeeeee files//apppe///017//
S. Williams, A. Waterman, and D. Patterson. Roofline: an insightful visual performance model for multicore architectures. Communications of the ACM, 52(4):65-776 76, April 2009. ISSN 0001- 0782, 1557-7317. doi: 10. 144/1499866. URL https:ttpss. :/ htt//aaccm.mm//// acm. org dwo//o.......// 1145/1498761 1498785
Publisher: Association for Computing Machinery (ACM).
S. Yang, B. Wang. Y. Shen, R. Panda, and Y. Kim. Gated linear attention transformers with hardware-efficient training. In Forty-first International Conference Machine Learning, 2024.
net forum? URL http:::://ppeeeeiiiiiitee wtw.leeww.neettttnmiieeeeee ditaummmdddaaaaaxxuu13
H. Zhou, A. Bradley, E. Littwin, N. Razin, O. Saremi, J. M. Susskind, S. Bengio, and P. Nakkiran. What algorithms can transformers learn? study in length generalization. In The Twelfth Interna- tional Conference on Learning Representations, 2024. URL
Forum?id=Age IunmHXX
13

---

Published as a conference paper at ICLR 2025
A RNN VARIANTS WITH MEMORY MIXING RECURRENT CONNECTIONS MODELED FLASHRNN
Elman RNNs
Elman 1990) Number of states: 1, Number of gates:
(0) tanh (g." (0)
(8)
Here, we omit as well a possible post-processing of the sequence that does not inter-mix states of different time steps. This can as well be parallelized for offlin training.
LSTM
(Hochreiter & Schmidhuber 1997 Gers al. 1999) Number of states: 2, Number of
gates: 4
States: h, (0) hidden state, C₁ s (1) cell state
(3) Gates: 2, = g: cell input, ft = 9r forget gate, ét = input gate, o, - gi gi output gate
(1)
(2)
h: (o:) tanh (c+)
(9)
G = (ft) 1 1 + a (é.) tanh (z+)
(10)
GRU (Cho et al. 2014) Number of states: 1, Number gates: 4 (in the definition of this paper)
States: 8 hidden state
Gates: 2, = g cell imput, T2 = g; forget gate, n, = 9: input gate, O, = gi output gate Here,
(0)
the nu, gate is not dependent the previous state, whereas the 9: gate is not dependent on the input. This behavior can be modeled in FlashRNN as well.
h4 = a ( h + (1 - (zx)) tanh (n: + a (r:) tanh(g:))
(11)
sLSTM
(Beck et 2024) Number of states: 4, Number of gates: States: h (0) hidden
state, = cell state, s (2) normalizer state, m, si stabilizer state
Gates: 2, g: (0)
cell input, fig. (1) forget gate, i, = g input gate, o, g. output gate
(3)
h o (ot) C
(12)
C₁ exp((log a (ft) + m,-1 m²) C- + exp (i. me) tanh (zt)
(13)
n2 exp((log 0 (fi) + m,-1 mi) n2- + exp (i: m.)
(14)
m, max (log (fi) + m4-1,2t)
(15)
B FLASHRNN ALGORITHM IN DETAIL
In the following algorithm, Rabx((ed) means this is
matrix tile of size b (c- dd, where a is an additional outer index (typically time / or states 8), which denotes that this is used as separate outer dimension. The merged dimension (c d) is shown merged as it is used in matrix multiplications, but is split (typically into gates) in the pointwise function. The dimensions are: time, s: states, g: gates, b: batch dimension, d: head dimension (abbreviated from dhead). For accumulation of the recurrent matrix product there are two matrix dimensions: The dimension along the previous state (of total size d) and the dimension along the new gates ğ of total size g d, since for every of the d RNN cells there are gates. Multiple heads are parallelized either sequentially or over multiple blocks in the grid Bи but we omit this for clarity here.
axis A happens an elementary tile size within warp E multiple warps WA in a thread block, multiple blocks BA within the grid, and sequentially via a loop LA. The total size has satisfy SA FA WA BA LA- The typical elementary size FA usable for matrix multiplications bfloat16 is EA E {8, 16, 32] for an outer dimension and EA = 16 for the accumulating dimension. The number of elementary tiles TA
Tiling along
S. ΣA WA x BA x LA- To optimize speed internally, we use memory padding to minimize memory bank conflicts and coalesced memory loading.
14

---

Published as a conference paper at ICLR 2025
Algorithm 2 FlashRNN-fusee forward pass
Tiling across blocks should kept minimal along accumulating dimensions, and can be extended along parallelizing dimensions (here gate dimension). So ideally Bs < Bg. Indices b, 5, implicitly updated from loop indices [p. Is. lo incorporating the respective warp/block indices.
Require: Recurrent matrix R E (d-)) inputs R(bx((d)) biases E Rd4
Require: Initial state ( RR.1,bdd
Require: Tiling dimensions for the grid [Ba. Bs, BB Bh] and block size [32 x Wo Ws. WB]
Divide R¹ into [Ts. = [Ls x Ws x Bs, LG x We x Bcl tiles R E IR 16x (16or 32) with
E (1..Tsh. 9 E (1.Ta) along the state (first) and gate (second) dimension
Divide the bias b into L x Wa] tiles along the gate dimension as by E (1)
Load tiles Rgg. by from HBM into registers (Ls, La] / [La] per warp) and potentially SRAM across multiple thread blocks [Bs, Bg, BB] for 1B in {1..LB) do
Load from initial state s (0) batch tile bs for t E 0..T - do
for La E 1. do
Initialize MatMul result E R Ea with zero in registers.
11.LL. for Lg (reg) do
Load state matrix tile Sorba from HBM
Calculate and Accumulate Matrix Product y = y + 8prs R.g along
end for
if L (reg)
< Ls then for Is L.LL)A do
Load state matrix tile Sulss from HBM
Load recurrent matrix tile RT from SRAM
Calculate and Accumulate Matrix Product y y + Sudda R.g. along 5.
end for
end if
Store MatMul result y in SRAM
Block Level Sync
for wy in 1..Ws - do
Load other MatMul result ÿ:
Accumulate MatMul result = y:
end for
Block Level Sync
Store MatMul result y in SRAM
if Bs > then
Reorder tiling here for coalescing memory access and optimal work partitioning
Store MatMul result y in HBM
Grid Level Sync
end if
# Reorder tiling here with a block for one thread per point-wise op.
Load Gate inputs 2rig from HBM
Load MatMul result y from SRAM + by + +
Add Lggg
for by
E 2. B. do Load other MatMul result ÿs from HBM + y:
Add
end for
Point-wise Update 814165 16. P(s with aligned states s and gates ğ Write out gates 9:bg to HBM for backward pass
Write out new states 84165 to HBM
end for
Grid-Level Syne (for new states to be available across the whole grid)
end for end for
15

---

Published as a conference paper at ICLR 2025
C CONSTRINT RESOLUTION ALGORITHMS
To model the hardware constraints, we define IntegerVariables, c.g. variable describing a tiling size in the FlashRNN algorithm tiprocessor. These can attain a set of numbers (domain), e.g. initially a large range for a so far unconstrained tiling size or certain value for constant. These variables can be composed to terms via addition and multiplication, and these terms can be constrained via equalities, inequalities and divisibility constraints.
constant that defines the total SRAM for one streaming mul-
Algorithm 3 ConstrINT Resolution
Require: Input Constants/ Variables
Require: Resolution Variables with Heuristic
Require: Equality, Incquality and Divisibility Constraints
Generate intermediiatebaccgroout variables for terms that propagate constraints Reach are consistency via "Global ARC-Reduce" if any variable has empty domain [y| = 0 then
return "No Solution viable" end if
Sort values for each Resolution Variable via Heuristic while any Resolution Variable has domain Dv| 1 (not fixed) do
Choose Variable via Heuristic, Increase Index Count
Lowest Order Variable has empty domain then return "No Solution viable" end
Set Variable Value via Heuristic Reach are consistency via Global ARC-Reduce if any variable has empty domain |Dv = 0 then Backtrack
end if end while return Solution
Algorithm ConstrINT Global ARC-Reduce
Require: Expression Parse Tree of Constraints and Variables Status=Not Converged
while Change in Root IntegerVariable or NotConverged do
Propagate Restrictions to SubTerms of Expression for Sub-Term in Koott-xpression do
> Top-Down Application of Constraints
Apply Global ARC-Reduce on Sub-Term - Get changes end for
if any change in values for Sub-Term then Propagate Restriction from Sub-Terma to Root Variable
> Bottom- Up Application of Constraints
Status=Not Converged
else
Status=Converged
end
end while
return change in Root IntegerVariable
Specific resolution variables additionally have a heuristic added that defines the behaviour of iter- ation for choosing among possible values. If the domain of all resolution variables is reduced to single number, this a solution. The heuristic gives an order these variables and for each variable, if smaller larger values are expected to result in "better" solution. This helps optimization as there might be many possible solutions, but certain ones promise most speed-ups (e.g. using most TensorCores).
At the lowest level, a term is composed two Integer Variables (or intermediate variables), so con- straints it propagate down to the summand factor IntegeVVariable Equality, Inequality
16

---

Published as a conference paper at ICLR 2025
and Divisibility constraints propagate the contained terms well. For example, since all numbers are strictly positive the upper bound on a sum of two IntegerVariables Variables applies to both the summands minus one. Applying the constraints iteratively upwards and downwards in the expression parse tree until convergence (i.e. no change for any variable) leads to an arc-consistent state, which we call "Global ARC- The binary "ARC-Reduce" algorithm is part of the "AC-3", constraint satisfaction problem solver for a more general setting (Mackworth) 1977). An arc-consistent state might still have no solution, it is merely super-set of all possible solutions. Based on the heuristic the ConstrINT algorithm applies a depth-firs first tree search with "Global ARC-Reduce" application at each step and backtracking for an empty solution domain.
D CONSTRINT KERNEL OPTIMIZATION
First Kernel Call
RNN Sizes
GPUInfo Kernel
(Heads, HeadDim...)
SRAM, Register
ConstriNT Optimizer
Kernel Parameters Registers exceed
(Tilling Register use)
limits
Compiler NVCC
Test Execution
Executable Kernel
Figure 5: JIT Optimization procedure for first kernel call. RNN parameters and GPU hardware info are processed by ConstrINT for feasible kernel parametrization. Since register cannot fully be predicted in advance, register use is iteratively optimized with feedback from the compiler. Subsequently, the kernel is cached as well the intermediate optimization solutions.
The fused kernel described in Section 5.2 and in more detail in Appendix Section B has certain external parameters, which have to be set correctly for the kernels to be runnable and fast. The main constraints here are the size limitations of registers and SRAM. For a large hidden or head dimension d, e.g. d' = 768, the recurrent weight matrix for an LSTM has the size 4 x 768 x 768 x 2 x 4 MB. However, for an H100 GPU, the register file size is 256 B and the SRAM/ shared memory is up to K per SM/ / block. Therefore, this matrix needs to be sharded over multiple SMs in coopera- tive grid group that can synchronize on the grid level. In particular, the different variables defined in
Appendix Section B EB, WB, BB, LB, Eg. Wo, BG, LG, Es, Ws. Bs, L L (SRAM) imply cer- tain sizes of registers and SRAM via a polynomial function. ConstrINT optimizes these variables to fit within the boundaries of the hardware and to achieve a reasonable speed, by ordering the variables and applying heuristic their values. For example, the gate dimension ğ in the forward pass is pure parallelization, whereas the state dimension 8 is accumulated over. Accumulation necessitates additional memory operations and synchronization that make the execution slower, which is not the case for the purely parallel dimension. Therefore the BG variable is maximized, while Bs is mini- mized during the constraint satisfaction solution search.
(reg)
Furthermore, it is priori not clear how many registers can be used by the kernel for storing ad- ditional variables. Therefore, ConstrINT is used in a feedback loop together with the compiler performing binary search for the largest attainable register size.
Given a more detailed understanding of the kernel as shown in Section B ConstrINT variables can be fixed to different values for manual optimization of specific kernel size. ConstrINT will auto- matically optimize all other variables using the hardware limits and given heuristics.
An example, where this is used in the standard version is the block size (threads/ warps per thread block). While the maximum could be 1024 on NVIDIA GPUs, we set this manually to a fourth. This is usually faster, while not restricting the usable memory.
17

---

(5) Published as a conference paper at ICLR 2025
E DETAILS ON TRITON IMPLEMENTATION
Algorithm provides details on the Triton implementation. shows the computation for a single program or thread block, which computes one head of dimension d and block Bg of the batch dimension b. We run a grid of (Rhead x Bg of these programs in parallel for FlashRNN forward pass in Triton. We load the recurrent weights R and biases only once from HBM to SRAM and keep them in SRAM throughout the time loop.
b
On a higher level the main differences to the CUDA implementation in algorithm are that in CUDA force the kernel keep the recurrent
multiple thread blocks for single head and weights R in registers instead of SRAM. One can see this difference for example in the kernel launch grid, which parallelizes only over number of heads and blocks of batch size in Triton, while it has two more parallelization dimensions in CUDA (see Algorithm T.
Algorithm Triton FlashRNN Forward Pass
(5) E Rd for Require: Recurrent weights R(i) E Rdxdd biases b()) E Rd for gates j and inputs x.
gates j and timesteps 1..T;
Initial states $ E Rd
Load R() E Rdd biases b6) E Rdd
Load initial states 8p E Rd
for timestep f 1..T.d. do
Load inputs .r!")
Compute gate preactivations 9t 2t + & + b0)
(i) (1)
(0)
Compute pointwise operations s = pli) (f8/-1be {g. }3
Store states s(()
if Store output gates then
Store gates (5)
end if 1 s (i)
9:-1 gr (t) end for
())
(8) return States 8 gates 9
F COMPUTATIONAL COMPLEXITY
Traditional RNNs go over the sequence step by step, while applying recurrent matrix multiplica- tion and pointwise activation function at each step. For the back-propagation in time, all past state values usually stored. In this paper, we implement the head-wise notion limiting the recurrent matrix to a block diagonal form. The computational complexity is therefore: OT Mheade &head), with head size dheadd Rheads the number of heads and T the sequence length. The matrix vector product at each step is the dominant computational factor (for large head sizes). For inference, the memory needed is defined by the state of the RNN and is O(Theads dRead) in size.
In contrast, Attention computes a weighted sum over past inputs at each step, with the weight de- fined by the sofumax over scalar products between query and key vectors. This leads computa- tional complexity OT Nheads dhead). The space complexity is O(T Nheads dheed) (Vaswani et al.) 2017). In conclusion, RNNs are more compressive, while their computational complexity is higher when computing only few steps. For training with BPTT the space complexity of RNNs matches the one of Attention, as all past states have be stored.
G ROoFLINE ANALYSIS
As mentioned in Section 5.1 kernel speed is fundamentally limited by two factors: computation and memory bandwidth. This is usually visualized in the Roofline-Plot, showing the position of a kernel in terms of its computation throughput and arithmetic intensity. We NVIDIA NSight Compute,
18

---

Published as a conference paper at ICLR 2025
to analyse this for the alternating and fused FlashRNN LSTM kernel compared to the nn.LSTM (cuDNN) baseline on H100-SXM:
Roofline Analysis NVIDIA Nsight Compute
1013
Porfamancc (FLOP/4)
10
iITTT EW
FW
10
Puintwse
Postwies BIV
10t0
Memory Bandwidth
Peak Compute Performance
10
Intensity (FLOP/Byte)
Figure 6: LSTM Kernels in the roofline plot measured with NVIDIA NSight Compute plotting arithmetic intensity to the right, computation speed to the top. Alternating kernels show lower arithmetic intensity and performance than the fused kernels. The fused backward kernel might still be optimized compared to the nn.LSTM baseline. RNNs in general are still deep in the memory bound regime of low arithmetic intensity. The peak performance is the scalar performance limit for float32 FLOPs.
H ADDITIONAL BENCHMARK EXPERIMENTS
H.1 FUSED KERNEL LIMITS
Since the fused CUDA kernel FlashRNN based on keeping the recurrent memory matrix in registers and shared memory, there is limit on the maximal head size corresponding the size of the R matrix. As ConstrINT can solve these constraints automatically, there is no additional overhead other than setting this number and let it check if it works. Here, the hardware limits become visible an RTX 3090 and A40 have 128KB of SRAM compared to 192 KB of an 1 and 228 KB of an H100 For the LSTM fused kernels (forward backward), get the following attainable head dimensions (greater than 1280):
RTX3090: [1280, 1312, 1344, 1440, 1536, 1600, 1632, 1728, 1824]
A40: [1280. 1312, 1344, 1440, 1536. 1600, 1632, 1728, 1824]
A100: [1280, 1312, 1344, 1376, 1408, 1440. 1472, 1504, 1536, 1568, 1600, 1632, 1664, 1696. 1728, 1760, 1792, 1824, 1920, 2016, 2080, 2112, 2304]
H100: [1280, 1312, 1344, 1376, 1408, 1440, 1472, 1504, 1536, 1568, 1600, 1632 1664, 1696, 1728, 1760, 1792, 1824, 1856. 1888, 1920. 1952, 1984, 2016, 2048, 2080, 2112, 2208, 2240, 2304, 2400, 2496, 2560, 2688)
19

---

Published as a conference paper at ICLR 2025
For larger head sizes the alternating kernels can to be used, since these not restricted in the head dimension.
H.2 TORCH.COMPILE BASELINE
Since torch. compil le seems unroll the vanilla PyTorch implementation of kernels, long sequence lengths take very long compilation times. Exemplary tests for small sequence length 64 took minutes compile, while being only about times faster than the vanilla PyTorch imple- mentation without torch. compile. For comparison our fused kernels are up to 50 times faster.
H.3 LSTM SEQUENCE LENGTH RUNTIME EXPERIMENTS
We confirm that the findings from Section 6.1 hold true also for varying sequence lengths from 256 to 2048. We fix the batch size to 16 and measure the runtime for 12 heads with head dimension 64 (see Figure and single head with head dimension 768 (see Figure ( these experiments, we see the expected linear scaling of the runtime of all LSTM kernels for increasing sequence lengths. The previous findings transfer across sequence lengths.
Triton fazed
CUDA alterrating
Rochhlttertion2
Triton Fased
CUDA alternttngg
Flach/ttteriio22
CUDA fuxed
CUDA faxed
40
35.00
15
Tima (mi)
Time (mi)
20
10.02
6.19
4.25
10.58
31
6:60
1.2
0.53
0.08
1.0
3.00
1.18
2.15
0.15
0.29
Tm256
T=512
T=1024
T=2048
Tm
T=512
T=1024
T=2048
Figure 7: LSTM Runtime for different sequence lengths (T) on a NVIDIA H100. We use 12 heads with head dimension 64 and batch size 16. Left: Forward pass. Right: Forward + backward pass.
CUIA lused CUDA altomazing
PT LSTM
Haste LSTM lo32
CUDA fused CUDA alternating
PT mLLSTM
Haste LSTM fo32
Time (ned)
Tima (m63)
50-
12.54
30-
27.55
9.89
20
20
1.0
1.63
11.
15:40
13.4
2.45
2.00
0.85
0.52
T-1024
T=2048
Figure 8: LSTM Runtime for different sequence lengths (T) NVIDIA H100. We use one head with head dimension and batch size 16. Left: Forward pass. Right: Forward + backward pass.
20

---

Published as a conference paper at ICLR 2025
H.4 FLASHRNN WITH EXTERNAL GATE PRE-ACTIVATION COMPUTATION
In Figure 9 we compare the kernel runtimes of the CUDA alternating and CUDA fused kernel with and without the external gate preactivation. w/ Linear denotes with external gate preactivation computation via a linear layer. The impact the gate preactivation computation marginal com- pared to the overall kernel runtime.
CUDA fused
CUDA Linear
CUDA fused
CUDA altemazing Linear
CUDA fused w/ Linse
PT m5LS
CUDA fined w/ Linear
PT n5.LSTM
CUDA altornatngg
alternting
31,722.0
404.8
30
Tome [m))
60
15
1.5
3.0
40
10
539.74
28 789 4577
31.45
11:4..16
3.08
9
Figure 9: LSTM Runtime for different batch sizes (B) on a NVIDIA H100. We use one head with head dimension 768. We compare the kernel runtime with and without the gate preactivation matrix multiplication. Left: Forward pass. Right: Forward backward pass.
H.5 SLSTM RUNTIME EXPERIMENTS
In Figures 10
12
13 and
14
show the results of the experiments from Section 6.1 for the
sLSTM (Beck et al. 2024).
Tritos fused CUDA fused
CUDA, altemating Vanila PyTorch
Triton fused CUDA fused
CUDA altemating
Vanila PyTorch
207 49.44
346.9
251.05
699
671.38
602.8
15
Time (m))
25
25.73
24.05
10
1.29
15
10.75
35
2,1
NH-48
NH-24
NH-12 D41:64
NHw6
N
NH-1
NH-48
NH+24
NHw
NH+m6
NH
NHw1 DI1:768
DII:16
Di:: 32
D1:
D51
D68
DH::1
DI::12
Di 64
DIf:128
D51:250
Figure 10: sLSTM Runtime for different head dimensions (DH) and number of heads (NH) on NVIDIA H100. Overall embedding dimension is fixed at 768. We use batch size 16 and sequence length 1024, Left: Forward pass. Right: Forward backward pass.
21

---

Published as a conference paper at ICLR 2025
altemanting
fused
fused
alternating
P21
60
Time (mi))
5:82
0.00
40
23.08
23.20
23.60
21.71
20
13.71
2,50
10
B-16
B-64
B-256
B-2
6-16
8-128
B-256
Figure 11: sLSTM Runtime for different batch sizes (B) on a NVIDIA H100, at 12 heads with head dimension 64 and sequence length 1024. Left: Forward pass. Right: Forward + backward pass.
CUDA fused
CUDA altemating
CUDA fused
CUDA altomating
50
Tiha (mq)
30
1000
91.75
80
1.23
15 16
7.16
60
10
12.08
11:22
40
24:87
25.00
30 03100
4.53
3.44
3.47
455
20
10.1
101
B-2
Bm8
B-16
Bm32
Bw64
Be128
Bm256
B=2
8m6
8-16
Bm 8m32
Bm64
Bm128
Bm256
Figure 12: sLSTM Runtime for different batch sizes (B) on a NVIDIA H100, one head with head dimension 768 and sequence length 1024. Left: Forward pass. Right: Forward backward pass.
Triton
altemating
Trion foeed
CUDA alternatnng
20
40,55
40
15
Timee (m8)
10
3.36
207
15.05
35:23
9-101
2.25
2.50
5.99
5.72
4.7
T-256
T=2048
Figure 13: sLSTM Runtime for different sequence lengths (T) on a NVIDIA H100. We use 12 heads with head dimension 64 and batch size 16. Left: Forward pass. Right: Forward + backward pass.
22

---

Published as a conference paper at ICLR 2025
CUDA
altemating
fused
altemating
50
40
15
Time (mh))
11.0
30
20.99
28.05
10
20
14.29
10
7.31
T-256
T-512
T-1024
T-2048
T-256
T-512
T-1024
T-2048
Figure 14: sLSTM Runtime for different sequence lengths (T) on NVIDIA H100. We use one head with head dimension 768 and batch size 16. Left: Forward pass. Right: Forward + backward pass.
23

---

Published as a conference paper at ICLR 2025
H.6 NUMERICAL ERROR ANALYSIS
Hidden State
50th percentile
90th percentile
0.020
100th percentlle
0.015
Absolute Ererr
السهرر
0.010
0.005
0.000
100
200
300
400
500
Time Step
Numerical error the CUDA fused kernel in 6 compared to a vanilla Py PyTocch baseline in float64 over the sequence length. For an RNN, would assume accumulation of errors over multiple steps.
Figure 15:
In Figure 15 we plot the numerical deviations in the LSTM hidden states (i.e. the outputs) over time. We compare CUDA fused kernel in 6 (the default setting) vanilla PyTorch implementation in float64. For this experiment we use single example with sequence length 512 and single head with head dimension 768. We use random normal distribution generate the weights, biases and inputs.
We plot the 50th, 90th and 100th percentiles of the absolute errors of the LSTM hidden state output- many per timestep. Percentiles are computed over the head dimension of 768.
There exist maximum deviations of about 0.01. but this error stabilizes over time.
I LANGUAGE MODEL TRAINING ON A100s
Model
Heads
Param. (M)
Train Time (h)
Median Step (s)
Val PPL (val)
LSTM CUDA fused
190
15.4
1.699
22.1
LSTM CUDA altemn.
I
190
15.1
1.685
22.1
LSTM PT nn LSTM 1
190
6.6
0.730
25.9
sLSTM CUDA fused
1
190
15.3
1.707
21.3
sLSTM CUDA altern.
190
15.6
1.720
21.3
LSTM CUDA fused
12
164
7.8
0.820
22.3
LSTM CUDA altern.
12
164
7.7
0.809
22.3
sLSTM CUDA fused
12
164
8.0
0.865
21.7
sLSTM CUDA altern.
12
164
8.0
0.852
21.7
Transformer
12
162
6.3
0.688
18.0
Table 3: 165M Model training on 15B tokens of SlimPajama on 8xA 100s.
24

---

Published as a conference paper at ICLR 2025
J LANGUAGE TRAINING DETAILS
All models are roughly at 165 M parameter scale, that means 12 Transformer blocks (post-up pro- jection), with a swish-gated MLP and embedding dimension 768. The Transformer uses RoPE embeddings, whereas the other models do not use any additional positional information. We train with context length 1024 and global batch size of 512, resulting in roughly 30 k training steps for 15 B tokens of the SlimPajama dataset. We use the GPT-2 tokenizer and learning rate 1e-3 with lin- ear warmup over 750 steps and cosine decay to Ie-4 over 30k steps. We use PyTorch in version 2.4.0 and CUDA 12.1 for A100 and 12.4 for H100s. The training uses PyTorch FSDP in the NO SHARD setting (DDP) with Automated Mixed Precision using 6 and float32 for accumulations. For the A100 experiments,
one node of eight A100-SXM (80GB) GPUs and local batch size of 64, For H100-SXM we reduce the local batch size to 32 and use gradient accumulation errors, though they have the same HBM size (80 GB).
steps due For the language model trainings, we see more spikes in the training step times for FlashRNN mod- els compared the PyTorch implementations, which should be investigated further.
K EXPERIMENTAL DETAILS PARITY TASK IN SEQUENCE EXTRAPOLATION
For the parity task we train the parity task with varying training sequence lengths up 40. For the reported validation, we evaluate on sequence lengths between 40 and 256. Sequence lengths are uniformly sampled in the respective ranges. We train three seeds for learning rates {1e-2, 1e-3, lc-4] and choose best learning rates. We train for up to 100k steps with batch size 256 with linear warmup of 10k and cosine decay 10 % of the peak learning rate. Elman networks and LSTM cannot reach 100 % accuracy on sequence extrapolation for the smallest learning rate. All models reach low losses and high accuracies on the training set.
25


Reviews:
[{"review_id": "PuZ7XZgZqN", "reviewer": "Reviewer_8hLb", "summary": "This paper introduces FlashRNN, which optimizes traditional RNNs on modern GPUs. It divides RNNs into smaller pieces and parallelize over it to increase efficiency on GPUs. The model auto-adjusts its internal cache settings using a constraint satisfaction problem solver, enabling further performance optimization based on hardware. The experimental results show that the proposed solution is 50x faster than a vanilla PyTorch implementation.", "questions": "Please see the weakness section above.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is easy to follow and well organized.\n- Hardware aware fine tuning\n- Comprehensive benchmarking", "weaknesses": "- The paper emphasizes the register level CUDA optimization. It'd be great if the authors can show how the implementation is close to the roofline through profiling.\n- I believe comparing to pytorch is not fair as pytorch has overhead to launch its internal CUDA kernel.\n- Several fusion can be achieved through compilation. How is it different? Also, not using compilation seems to be unfair comparison.\n- Why FlashAttention2 matters in the RNN?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fNBlykMrOA", "reviewer": "Reviewer_fH1r", "summary": "While Transformers are state-of-the-art in sequence modeling, they lack state-tracking abilities. RNN families, like LSTMs and GRUs, can perform state tracking but at the expense of sequential processing. To address this, this paper introduces FlashRNN, which optimizes kernel operations to leverage GPU registers and SRAM. Additionally, it incorporates ConstrINT, an auto-tuning optimization tool that operates under hardware constraints, such as cache and memory limits, by solving an integer constraint satisfaction problem. Experimental results demonstrate the effectiveness and efficiency of FlashRNN.", "questions": "1. The interaction between ConstrINT and other components in FlashRNN is somewhat unclear. Providing an overall architecture or workflow for FlashRNN would be helpful.\n2. Algorithm 5.1 outlines the kernel fusion procedure but lacks a detailed explanation. A further illustration of this algorithm would be beneficial.\n3. In Algorithm 5.1, line 242 mentions “Load RtbgsR_{tbgs}Rtbgs​, BtgB_{tg}Btg​ to registers and SRAM.” How many registers and how much SRAM can be used for storing these values? Utilizing all available registers and SRAM may not be ideal, as they may also be needed for storing intermediate results during computation. If this strategy is used, it could impact performance. Detailed clarification is needed.\n4. Line 299 states that multiple heads are computed in parallel in different programs without synchronization. Could this lack of synchronization impact accuracy?\n5. Lines 301 to 303 indicate that custom caching strategies cannot be applied due to restrictions on GPU register access in Triton. Is there a significant gap between the intended design in Algorithm 5.1 and the Triton implementation? How large is this gap?\n6. Line 310 mentions that FlashRNN pads with zeros to support small batch sizes, at the cost of efficiency. How much efficiency is lost due to this padding?\n7. The evaluation setup is somewhat confusing. Are all four backends (CUDA fused, CUDA alternating, Triton fused, and Vanilla PyTorch) being compared with the two baselines, FlashAttention2 and nn.LSTM? If so, it appears that FlashRNN’s performance is lower than the baselines in most configurations shown in Figures 1 to 4. Additionally, why is FlashRNN compared with FlashAttention2?\n8. In the experiments, lines 467 to 468 mention replacing attention blocks with FlashRNN LSTM for speed comparisons. Does this substitution provide a meaningful comparison?\n9. None of the experiments include accuracy comparisons, which appears to be a limitation.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "+ Hardware-aware optimization for accelerating RNN model training and inference is promising.\n+ It claims that FlashRNN will be open-sourced, which helps further research in this direction.", "weaknesses": "- Certain critical procedures are somewhat unclear and would benefit from further clarification.\n- The experimental section could be strengthened with additional details and improvements.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vLrSO6UQQ1", "reviewer": "Reviewer_HW4x", "summary": "The paper optimizes traditional RNNs on GPUs. They introduce optimized kernels and an optimization framework. One challenge for efficient and scalable implementations that the authors tackle is the requirement of strict sequential processing of RNNs. They introduce a parallelization variant that processes multiple RNNs of smaller hidden state in parallel, inspired by the head-wise processing in transformers. They implement two variants, an alternating version switching between point-wise and matrix-multiplication kernels and a fused implementation, optimizing memory transfers while using hardware-optimized matrix-multiplication. The second leads to a further 3-4x speed-up over the alternating option for small batch sizes. Their kernels can achieve 50x speed-ups over a vanilla PyTorch implementation and allow 40x larger hidden sizes compared to their Triton implementation. \nTo enable flexibility on different GPU variants, they also introduce an optimization framework for hardware-internal cache sizes, memory and compute handling.", "questions": "Can you elaborate on the similarities and differences of the HASTE implementation and your approach. Did you benchmark against HASTE? If not, why not?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is overall well-written, uses clear language and structure. \n- it tackles an important problem of overcoming a fundamental limitation important for further scaling rnns. This would enable state-intensive tasks where state continuity and stepwise dependency matter, which might be less naturally handled by Transformers\n- They compare to relevant Backends (e.g., cuDNN)", "weaknesses": "They mention the HASTE library and that they overcome their limitations. However, I think they don't benchmark directly against HASTE which would be relevant.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "byfiAHb4Ta", "reviewer": "Reviewer_z311", "summary": "FlashRNN is a hardware-optimized library that significantly accelerates traditional RNNs, like LSTMs and GRUs, using optimized kernels for GPUs. The paper introduces kernel-level optimizations that fuse matrix multiplications and activation functions, achieving up to 50x speed-ups compared to standard implementations. FlashRNN also supports parallel processing through a \"multi-head\" approach and uses the ConstrINT framework to model hardware constraints. This library enables faster and more efficient RNN performance, particularly in state-tracking tasks.", "questions": "A few questions and feedback:\n\nWas the profiling conducted with torch.compile optimizations? Have you considered trying TVM auto-tuning for further optimization?\n\nSince LSTM's arithmetic intensity is mostly memory-bound, the proposed method seems to focus primarily on fusion operations. It’s unclear if other techniques to tackle bandwidth bottlenecks for inference were explored. Given that Transformers were specifically designed to overcome LSTM limitations, this aspect could use more clarification.  The paper would benefit from a more detailed comparison of computational complexity scaling between FlashRNN and Transformers, particularly how each scales with respect to input size and hidden states.\n\nWhile it’s good that the paper addresses why LSTM optimization remains relevant compared to Transformers, it is not fully convincing why FlashRNN would outperform FlashAttention or how FlashRNN differs fundamentally from applying FlashAttention-like methods to RNNs. \n\nHas the method been tested on GPUs with less SRAM? It would be valuable to understand how FlashRNN performs in environments with different hardware constraints. What are the options if the weight matrix (W) cannot be cached? Addressing this scenario would help clarify how the method adapts to different hardware limitations.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "FlashRNN offers performance improvements by focusing on kernel fusion techniques that optimize LSTM and GRU operations on modern GPUs. The library’s focus on fusing matrix multiplication and activation functions into a single operation significantly reduces memory access overhead, a critical bottleneck in memory-bound computations. The integration of the ConstrINT framework for hardware-aware optimization ensures that the method adapts well to various GPU architectures. Additionally, the open-source nature of FlashRNN encourages further exploration and development in state-tracking tasks, where RNNs still excel. The paper’s effort to justify LSTM optimization versus Transformers provides valuable insight into the relevance of RNNs in certain domains.", "weaknesses": "FlashRNN has some limitations, particularly in its use cases, as its advantages are more pronounced for state-tracking tasks and less clear for applications where parallelizable models like Transformers excel, which may limit its broader appeal. Additionally, its focus on GPU-based optimizations reduces its applicability for users working with other hardware, such as TPUs or CPUs. Despite performance improvements, RNNs still face inherent scalability issues due to sequential processing, and FlashRNN doesn't fully address the parallelization advantage held by Transformers. Lastly, the paper notes numerical precision differences in training compared to implementations like PyTorch's cuDNN, which may need further investigation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)