FlashRNN: I/O-Aware Optimization of Traditional RNNs on modern hardware

Dear z31124,

thank you for your detailed review and all your questions and concerns. The speed comparison was conducted without torch.compile, as PyTorch would unroll the 1024 steps of the time loop, making it very slow to compile. In Appendix G.2, we added a comment on this, measuring the speed for short sequence lengths in exemplary tests on a RTX3090.

This work’s focus is on training, where the whole sequence is given in advance. There might be further optimizations for inference going step by step.

We have not explored the TVM compiler, because it is listed as „inference only“ backend in torch.compile() at the time of writing the paper and for our purposes we need Training&Inference backends. (See https://pytorch.org/docs/stable/torch.compiler.html). For torch.compile() we relied on the default backend „inductor“. Torch.compile() compiles the PyTorch code into Triton kernels. Since at the time of writing the paper torch.compile() was not able to compile the RNN in reasonable compile times, we decided to write our own custom Triton kernels for a comparison. We see that torch.compile is still in active development and hope that further generations will be more capable in compiling RNN workloads to efficient Triton kernels.

While we see the recurrent matrix caching as the main speed contribution, there are other optimizations made internally to reach the measured speeds. This includes coalesced memory accesses, use of shared memory for intermediate results, shared memory padding to avoid memory bank conflicts, and using synchronizations only at the points and levels where needed to ensure numerical correctness. Since the RNN architectures optimized in FlashRNN are not sequence parallelizable in their standard computational structure, any sequence parallelizable architecture can out-speed this given enough hardware parallelization and sequence length. In Appendix Section G, we briefly compare the computational and memory complexity of RNNs to Transformer architectures now.

FlashRNN does not intend to outperform FlashAttention in speed, nor in performance for tasks that do not need state tracking capabilities. Where these capabilities are needed however, it should serve as a library for speeding up research and investigating novel RNN architectures. While FlashRNN also falls into the category of kernel optimization, the approach here is very different to FlashAttention, as RNNs are not sequence parallelizable and typically have a large recurrent matrix that we are able to cache efficiently in our code. This is now stated as well in the benchmark introduction Section 6.1.

In Appendix Section H.1, we added now the attainable head dimensions for our fused kernels on different GPUs - using the automatic adaption provided by ConstrINT. For larger head dimensions one has to switch to the alternating kernel, which supports arbitrary head dimensions at the cost of typically lower speed.

Dear vHW4x01,

thanks for your nice review! Haste is very similar to the alternating version of our kernels, which is why we did not do a speed comparison in the first place. We have added Haste now to our speed measurements. It is slower than our alternating kernels as it is limited to float32 and float64 precision. Since haste has not been updated for around four years and it does not support float16 and bfloat16 precision. These precisions are integrated in FlashRNN and strongly benefit from TensorCore acceleration and memory speedups. Feel free to look also at the improvements, additions and clarifications made in response to the other reviewers, and to add further questions.

Dear 8hLb,

thank you for your review and comments on our paper. In the Appendix Section F we added the requested roofline plots using NVIDIA Nsight Compute Utils (ncu). There, one can see that the fused kernels have a higher arithmetic intensity compared to the alternating ones. The forward kernel also comes close to the closed-source cuDNN baseline, but there is still a gap the hardware limitations of the roofline.

We agree that comparing to pure PyTorch is unfair, but this is the typical entry point for researchers. We have tried torch.compile to speed up the vanilla PyTorch version, but this unrolls the whole recurrent computation graph along the sequence dimension, since there is no generic scan in PyTorch currently. This leads to excessively long compilation times, especially for deeper architectures.

As mentioned in Section 6.1 in the new revision of our paper, we compare to FlashAttention2 as a very common sequence modeling kernel. Attention however, does not have state tracking capabilities. It just serves as a baseline in terms of speed, when these capabilities are not needed.

We hope that our additional experiments and explanations could improve your impression over previous concerns. We are also happy to take in further comments or suggestions for improvement.

Dear fH1r04, thank you for your review and comments. Thanks for your detailed questions that we would like to address:

1. In the new revision of our paper, we added Appendix Section D to describe the higher level interaction of ConstrINT with the kernel compilation process.
2. Also, we provide a more detailed outline of the central algorithm in Appendix Section B.
3. Regarding the shared memory and registers used, we follow the strategy to use as much shared memory and registers as possible for the cached values (mostly the recurrent weight matrix). This was beneficial for the larger head dimensions of 768 and 1024, as otherwise the recurrent matrix would not fit at all. A more detailed tuning can be done by adapting the register use manually in the ConstrINT variables / kernel parameters.
4. The lack of synchronization does not impact accuracy in a numerical sense. One can view the different heads also as smaller recurrent neural networks that run in parallel, without any approximations. While there is no synchronization of the different heads at each time step, there is a synchronization of all outputs after computing the whole recurrent sequence. This is needed for a subsequent layer that potentially uses them in time-parallel fashion (e.g. in an MLP).
5. That is true. In Triton we do not have access to GPU registers, therefore we cannot keep the recurrent weights and the biases in registers manually. Instead we rely on the Triton compiler to manage SRAM and Register transfers. There is no significant gap in the design of the fused FlashRNN algorithm between CUDA and Triton. Both support the headwise parallelization and avoid multiple loading the recurrent weights and biases multiple times. However, as CUDA gives more fine grained control over GPU Hardware features such as the mentioned register level access and grid synchronization for example. Therefore, in our CUDA implementation we make use of these features to further optimize the kernels. Register level caching of the recurrent weights for example enables to avoid unnecessary loads from SRAM to registers in every time step. Grid synchronization allows to distribute one head to multiple thread blocks and accumulate the R s matrix multiply results over HBM. This enables the much larger head dimensions (up to 2688) in contrast to the limited size of 64 or 128 in Triton. In our revised paper we have added an Appendix E and Algorithm 5 which provide further details on the Triton implementation.
6. The minimal batch size is 8 as this is the minimal size of a TensorCore matrix multiplication in one outer dimension. This means that only 1/8th of the FLOPs are used for batch size one. However due to the order of magnitude higher FLOPs-speed of tensor cores this should overcompensate the efficiency drop. While the alternating kernel in Figure 4 does not have the limitation of the batch size, and may not suffer from the efficiency drop, it is still slower compared to the fused one.
7. We clarified our setup in Section 6.1 . The FlashAttention 2 baseline serves as a speed baseline of a common sequence modeling layer, which however does not have state tracking capabilities. The nn.LSTM baseline for the given hidden dimension is limited to one head (there is no multi-head version), and based on the cuDNN implementation of LSTM that is closed source. It is therefore not possible to try out new RNN architectures with nn.LSTM, but it can serve as a baseline of an industrially optimized kernel.
8. This is a setting Attention is commonly used in (Transformer), and RNN models can be used in as well. In the xLSTM paper, the sLSTM-only models were used in similar blocks. Our experiments show that training is feasible and not too much slower than attention based training. As previously mentioned, RNNs are currently the only hardware optimized sequence models with state tracking capabilities. Tasks that need those will therefore benefit from our kernels.
9. The tests in the code check our models for differences to a PyTorch baseline, where we see differences due to the low precision. We follow good practices, as in our kernels we use float32 accumulation internally, which potentially is not the case for the nn.LSTM baseline. This might be the cause for the difference in language model training performance between the FlashRNN and the nn.LSTM baseline. In addition, in Appendix H.6 we plot the numerical error of our CUDA fused kernel in bfloat16 compared to a vanilla PyTorch implementation in float64 over time. We see that the maximum errors stabilize in a modest range of 0.01 over time.

Thanks for all your detailed questions! We hope we could improve upon your previous impression of our paper, address most of your concerns and clarify the unclear points. We are happy to take further comments and suggestions!