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

TL;DR

FlashRNN addresses the demand for fast, state-tracking sequence models on modern GPUs by introducing hardware-aware optimizations for traditional RNNs. It combines fused CUDA and Triton kernels with a head-wise, block-diagonal recurrence and an automatic ConstrINT-based tiling tuner to adapt to GPU cache and register constraints. The approach yields substantial speedups (up to tens of times faster than vanilla PyTorch LSTMs) and enables large hidden sizes, while preserving state-tracking capabilities demonstrated on parity extrapolation tasks. By providing open-source implementations and tooling, FlashRNN aims to make state-tracking RNNs a practical option for research and applications.

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. Our open-source kernels and the optimization library are released here to boost research in the direction of state-tracking enabled RNNs and sequence modeling: https://github.com/NX-AI/flashrnn

Figures (15)

• 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 Kernels (forward) for maximum hidden sizes, with two 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.
• 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.
• Figure 3: LSTM Runtime for different batch sizes (B) on a NVIDIA H100. We use 12 heads with head dimension 64 and sequence length 1024. Left: Forward pass. Right: Forward + backward pass.
• Figure 4: LSTM Runtime for different batch sizes (B) on a NVIDIA H100. We use one head with head dimension 768 and sequence length 1024. Left: Forward pass. Right: Forward + backward pass.
• Figure 5: JIT Optimization procedure for first kernel call. RNN parameters and GPU hardware info are processed by ConstrINT for a feasible kernel parametrization. Since register use cannot fully be predicted in advance, register use is iteratively optimized with feedback from the compiler. Subsequently, the kernel is cached as well as the intermediate optimization solutions.