Research map / Architectures and efficiency / Attention and transformers
Tiled Flash Linear Attention: More Efficient Linear RNN and xLSTM Kernels
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.
Proposes Tiled Flash Linear Attention (TFLA), a two-level sequence-parallel kernel algorithm for linear RNNs/mLSTM enabling arbitrary chunk sizes, plus mLSTMsig with sigmoid input gate that is faster at equal language modeling performance.
Problem
Chunkwise-parallel linear RNN kernels (e.g., Flash Linear Attention) are limited in chunk size by GPU SRAM, forcing many intermediate states into HBM, which lowers arithmetic intensity and increases memory use and IO for long-context training.
Approach
Add a second level of sequence parallelism by tiling the intra-chunk matrix multiplications along the sequence dimension, so chunk size can grow without limit; additionally replace the exponential input gate of mLSTM with a sigmoid gate, removing the normalizer and max states to reduce computation.
Key results
- TFLA mLSTMsig forward kernel is over 30% faster than the mLSTMexp TFLA forward kernel.
- mLSTMexp TFLA xl_chunk kernels are about 25% faster than limit_chunk kernels.
- In training, TFLA kernels are faster than FlashAttention 3 for longer sequences and more than 2x faster than Mamba 2 kernels for all sequence lengths.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the exponential input gate with a sigmoid input gate (ablation) | mLSTMexp · DCLM language modeling, up to 1.4B parameters | perplexity | No clear effect |
| Use the TFLA xl_chunk kernel instead of the limit_chunk kernel (ablation) | mLSTMexp · DCLM validation, context length 4096 | perplexity: 21.03 -> 21.18; 16.55 -> 16.80; 13.20 -> 13.87* | Mixed |
| Initialize input gate biases to -10 (ablation) | mLSTMexp and mLSTMsig · DCLM language-model training | training stability | Helped |
| Use TFLA xl_chunk instead of limit_chunk (ablation) | mLSTMexp kernels · Kernel runtime benchmark on NVIDIA H100 | runtime: about 25% faster | Helped |
| Replace mLSTMexp with mLSTMsig (ablation) | TFLA forward kernel · Kernel runtime benchmark on NVIDIA H100 | forward-pass runtime: over 30% faster | Helped |
| Use TFLA mLSTM kernels | mLSTM · Training kernel benchmark; longer sequences | runtime | Helped |
| Use TFLA mLSTM kernels instead of Mamba 2 kernels | mLSTM · Training kernel benchmark; all sequence lengths | runtime: more than 2x faster | Helped |
| Increase the TFLA chunk size (ablation) | mLSTMsig xl_chunk kernel · Kernel runtime benchmark | runtime | Mixed |
Surprising findings
- Initializing input gate biases at large negative values (e.g., -10) mitigates gradient norm spikes for both mLSTM variants.
- mLSTMsig without normalizer shows the same transfer behavior as mLSTMexp with normalizer and max state.
Limitations
- Some head dimensions show gradient norm or loss spikes for xl_chunk kernels affecting final loss.
- Current kernels do not reach the theoretical runtime-optimal chunk size, suggesting unexploited performance potential.
- Language modeling experiments limited to models up to 1.4B parameters.
Experimental setup
Datasets: DCLM
Benchmarks: validation perplexity on DCLM TFLA kernel runtime on H100
Metrics: validation perplexity kernel runtime arithmetic intensity FLOPs
Baselines: FlashAttention-2 FlashAttention-3 Mamba-2 GLA FLA kernels LightningAttention-2 Llama2-style Transformer
Builds on: Flash Linear Attention FlashAttention-2 xLSTM Gated Linear Attention
Closest papers by meaning
- Various Lengths, Constant Speed: Efficient Language Modeling with Lightning Attention (ICML 2024 · similarity 0.89)
- Gated Linear Attention Transformers with Hardware-Efficient Training (ICML 2024 · similarity 0.89)
- FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (ICLR 2024 · similarity 0.87)
- xLSTM: Extended Long Short-Term Memory (NeurIPS 2024 · similarity 0.87)
- FlashMask: Efficient and Rich Mask Extension of FlashAttention (ICLR 2025 · similarity 0.87)
- Short-Long Convolutions Help Hardware-Efficient Linear Attention to Focus on Long Sequences (ICML 2024 · similarity 0.87)
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.