FlashRNN: I/O-Aware Optimization of Traditional RNNs on modern hardware
RNN;LSTM;Sequence Models;GPU;Hardware Optimization
Yes
No ethics review needed.
Please see the weakness section above.
- The paper is easy to follow and well organized.
- Hardware aware fine tuning
- Comprehensive benchmarking
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.
- 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?
Yes
No ethics review needed.
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.
+ 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.
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.
- Certain critical procedures are somewhat unclear and would benefit from further clarification.
- The experimental section could be strengthened with additional details and improvements.
Yes
No ethics review needed.
Can you elaborate on the similarities and differences of the HASTE implementation and your approach. Did you benchmark against HASTE? If not, why not?
- 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)
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.
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.
Yes
No ethics review needed.
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.
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.
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.
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.
We optimize traditional RNNs with memory mixing close to transformer levels on modern hardware.
@inproceedings{
anonymous2024flashrnn,
title={Flash{RNN}: I/O-Aware Optimization of Traditional {RNN}s on modern hardware},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=l0ZzTvPfTw},
note={under review}
}
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 will open-source our kernels and the optimization library to boost research in the direction of state-tracking enabled RNNs and sequence modeling.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
RNN LSTM Sequence Models GPU Hardware Optimization
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/ceae65ae3fa63817b31d7c5609b3dc388e9f789c.pdf
infrastructure, software libraries, hardware, systems, etc.
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
/attachment/267c94c372c76215cac66f956a2aae64d0138eea.zip
FlashRNN: I/O-Aware Optimization of Traditional RNNs on modern hardware
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission