Research map / Architectures and efficiency / RNNs and normalization
FlashRNN: I/O-Aware Optimization of Traditional RNNs on modern hardware
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces FlashRNN, hardware-optimized CUDA and Triton kernels for traditional RNNs (LSTM, GRU, sLSTM) with fused matrix-multiply and pointwise ops, head-wise block-diagonal recurrence, and an integer CSP solver (ConstrINT) for auto-tuning tile sizes.
Problem
Traditional RNNs with state tracking capabilities are slow due to strictly sequential processing and inefficient GPU kernels (alternating matmul and pointwise kernels, limited precision support), so they lag Transformers in practical speed.
Approach
Fuse the recurrent matrix multiplication and pointwise state updates into a single persistent GPU kernel that caches the recurrent weights in registers/SRAM, combined with block-diagonal (multi-head) recurrent matrices and an automatic tile-size optimizer (ConstrINT) that solves integer CSPs with equality, inequality and divisibility constraints modeling hardware limits.
Key results
- CUDA fused kernels achieve up to 50x speed-up over vanilla PyTorch RNN implementation on H100 (as claimed in the abstract and conclusion).
- Fused kernels support head dimensions up to ~2688 on H100 versus Triton limits of 128 (forward) and 64 (backward), i.e. about 40x larger hidden sizes than the Triton implementation.
- Parity task extrapolation: Elman, GRU, LSTM and sLSTM reach 1.00 accuracy while Transformer (0.52), Mamba (0.56) and mLSTM (0.54) fail.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use CUDA fused kernels instead of CUDA alternating kernels (ablation) | FlashRNN LSTM, 1 head · SlimPajama training | Train Time: 10.8 -> 9.9 h | Helped |
| Use CUDA fused kernels instead of CUDA alternating kernels (ablation) | FlashRNN LSTM, 1 head · SlimPajama validation | Val PPL: 21.9 -> 22.1 | Hurt |
| Use CUDA fused kernels instead of CUDA alternating kernels (ablation) | FlashRNN LSTM, 12 heads · SlimPajama training | Train Time: 9.6 -> 5.9 h | Helped |
| Use CUDA fused kernels instead of CUDA alternating kernels (ablation) | FlashRNN LSTM, 12 heads · SlimPajama validation | Val PPL: 22.1 -> 22.2 | Hurt |
| Use CUDA fused kernels instead of CUDA alternating kernels (ablation) | FlashRNN sLSTM, 1 head · SlimPajama training | Train Time: 10.9 -> 10.1 h | Helped |
| Use CUDA fused kernels instead of CUDA alternating kernels (ablation) | FlashRNN sLSTM, 1 head · SlimPajama validation | Val PPL: 21.4 -> 21.3 | Helped |
| Use CUDA fused kernels instead of CUDA alternating kernels (ablation) | FlashRNN sLSTM, 12 heads · SlimPajama training | Train Time: 9.7 -> 6.8 h | Helped |
| Use CUDA fused kernels instead of CUDA alternating kernels (ablation) | FlashRNN sLSTM, 12 heads · SlimPajama validation | Val PPL: 21.8 -> 21.7 | Helped |
| Replace attention with FlashRNN LSTM or sLSTM layers | 165M-parameter language model on SlimPajama · SlimPajama training | Train Time: Transformer 2.9; FlashRNN LSTM/sLSTM 5.9–10.9 | Hurt |
| Replace attention with FlashRNN LSTM or sLSTM layers | 165M-parameter language model on SlimPajama · SlimPajama validation | Val PPL: 17.9 -> 21.3–22.2 | Hurt |
| Use FlashRNN LSTM instead of PyTorch nn.LSTM | LSTM language model, 1 head · SlimPajama training | Train Time: 4.5 -> 9.9 h | Hurt |
| Use FlashRNN LSTM instead of PyTorch nn.LSTM | LSTM language model, 1 head · SlimPajama validation | Val PPL: 25.8 -> 22.1 | Helped |
Surprising findings
- nn.LSTM (cuDNN) is about 3x faster than the FlashRNN CUDA fused kernel for a single head of dimension 768 at that size.
- FlashRNN LSTM converged faster than the PyTorch cuDNN LSTM in language modeling despite only small numerical differences in single kernel calls.
Limitations
- Triton implementation limited to head dimension 128 (forward) and 64 (backward) due to lack of grid synchronization.
- Fused CUDA kernel limited by register/SRAM size; large head sizes require alternating kernels.
- RNNs still slower than attention in language modeling (about 25% for equal head dimension, 140% for a single head).
Experimental setup
Datasets: SlimPajama synthetic parity task data
Benchmarks: SlimPajama validation perplexity parity sequence extrapolation
Metrics: runtime (ms) speed-up over PyTorch validation perplexity training time extrapolation accuracy
Baselines: PyTorch vanilla LSTM nn.LSTM (cuDNN) HASTE FlashAttention2 Transformer Mamba mLSTM
Builds on: HASTE FlashAttention xLSTM (sLSTM) Triton
Closest papers by meaning
- Accelerating Linear Recurrent Neural Networks for the Edge with Unstructured Sparsity (ICML 2025 · similarity 0.86)
- Tiled Flash Linear Attention: More Efficient Linear RNN and xLSTM Kernels (NeurIPS 2025 · similarity 0.86)
- Second-order forward-mode optimization of recurrent neural networks for neuroscience (NeurIPS 2024 · similarity 0.86)
- Harnessing Manycore Processors with Distributed Memory for Accelerated Training of Sparse and Recurrent Models (AAAI 2024 · similarity 0.86)
- Sequence Accumulation and Beyond: Infinite Context Length on Single GPU and Large Clusters (AAAI 2025 · similarity 0.86)
- FlashMoE: Fast Distributed MoE in a Single Kernel (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.