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

Authors: Korbinian Pöppel, Maximilian Beck, Sepp Hochreiter
Primary area: infrastructure, software libraries, hardware, systems, etc.
Keywords: RNN, LSTM, Sequence Models, GPU, Hardware Optimization

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

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1729794723576

thread_type:
official_review

<thread_object>
creator: reviewer_z311

<content>
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.

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.

questions:
A few questions and feedback:

Was the profiling conducted with torch.compile optimizations? Have you considered trying TVM auto-tuning for further optimization?

Since 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.

While 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. 

Has 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730456497644

thread_type:
official_review

<thread_object>
creator: reviewer_hw4x

<content>
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. 
To enable flexibility on different GPU variants, they also introduce an optimization framework for hardware-internal cache sizes, memory and compute handling.

strengths:
- The paper is overall well-written, uses clear language and structure. 
- 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
- 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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691626956

thread_type:
official_review

<thread_object>
creator: reviewer_fh1r

<content>
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.

strengths:
+ Hardware-aware optimization for accelerating RNN model training and inference is promising.
+ 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.
- The experimental section could be strengthened with additional details and improvements.

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.
2. Algorithm 5.1 outlines the kernel fusion procedure but lacks a detailed explanation. A further illustration of this algorithm would be beneficial.
3. 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.
4. Line 299 states that multiple heads are computed in parallel in different programs without synchronization. Could this lack of synchronization impact accuracy?
5. 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?
6. 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?
7. 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?
8. In the experiments, lines 467 to 468 mention replacing attention blocks with FlashRNN LSTM for speed comparisons. Does this substitution provide a meaningful comparison?
9. None of the experiments include accuracy comparisons, which appears to be a limitation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730696204974

thread_type:
official_review

<thread_object>
creator: reviewer_8hlb

<content>
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.

strengths:
- The paper is easy to follow and well organized.
- Hardware aware fine tuning
- 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.
- I believe comparing to pytorch is not fair as pytorch has overhead to launch its internal CUDA kernel.
- Several fusion can be achieved through compilation. How is it different? Also, not using compilation seems to be unfair comparison.
- Why FlashAttention2 matters in the RNN?

questions:
Please see the weakness section above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732145310993

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank all the reviewers for their constructive comments and valuable feedback.

We are glad that the reviewers liked our work on hardware optimization for RNN architectures with state tracking capabilities.

We want to highlight the following major additions in the paper:
- Integration of the haste library [1] as an additional baseline
- Appendix Section B: FlashRNN Fused Algorithm in more detail
- Appendix Section D: ConstrINT Kernel Solution and Optimization Process
- Appendix Section E: Triton Algorithm
- Appendix Section F: Roofline Analysis
- Appendix Section G: Computational Complexity of RNNs compared to Attention
- Appendix Section H.1: Head Dimension Limits for different Hardware
- Appendix Section H.2: Note on torch.compile Experiments
- Appendix Section H.6: Numerical Error Analysis

We want to emphasize again that FlashRNN does not intend to compete against FlashAttention in terms of speed or against Transformers in Language modeling in its current form. However, Attention and current competitors do not have state tracking capabilities [2] that traditional RNNs do have, which is why we want to optimize these for speed on current GPU hardware to check if this is an actual limitation. We can show that, while still being slower than FlashAttention and Transformers, with our optimizations one can reach reasonable speeds that do neither limit experimentation with new RNN architectures nor moderate scaling to 165 M parameter sized models - reaching a ~50x speed up over vanilla PyTorch. 

[1] S. Nanavati 2020: Haste: a fast, simple, and open RNN library

[2] W. Merrill et al. 2024: The Illusion of State in State-Space Models
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes several techniques to improve the performance of RNN, such as adopting the multi-head architecture of Transformers in RNNs, optimized kernels for faster activation and backpropagation computation, and adaptation of said kernels for different hardware cache sizes.

Reviewers generally agreed the paper's major strengths are closing the implementation gap between RNNs and Transformer model (which can facilitate more experimentation and research into alternative model architectures), and the hardware-aware kernel library that can adapt to hardware of different cache sizes. At the same time, reviewers noted that the gap between RNNs and Transformers is not fully closed by this paper alone. Because it is not easy to improve the performance of RNNs (which has merit as a well-studied architecture) against the much larger volume of ongoing Transformer implementation optimization research, it is fair to say that multiple works in RNN optimization would be needed to change the state of the art, and that this paper makes a solid contribution towards that goal.