Transformer Based Linear Attention with Optimized GPU Kernel Implementation

2025/10/24 by Gerami, Armin, Duraiswami, Ramani
#Computation and Language (cs.CL) #FOS: Computer and information sciences #Machine Learning (cs.LG)

paper · doi:10.48550/arxiv.2510.21956

Abstract

The original softmax-based attention mechanism (regular attention) in the extremely successful Transformer architecture computes attention between N tokens, each embedded in a D-dimensional head, with a time complexity of O(N²D). Given the success of Transformers, improving their runtime during both training and inference is a popular research area. One such approach is the introduction of the linear attention (LA) mechanisms, which offers a linear time complexity of O(ND²) and have demonstrated comparable accuracy to regular attention. However, LA in practice lags behind its theoretical efficiency. We propose a novel method for LA's forward and backward passes, along with a highly-optimized CUDA implementation. Our approach outperforms the state-of-the-art by 3.3 times in speed and reduces memory consumption by 3.6 times. We validate these improvements in both single-layer and end-to-end settings by training a 1.4 billion parameter language model, which demonstrates similar expressivity to regular attention on major reasoning benchmarks.

Citations

• MiniMax-01: Scaling Foundation Models with Lightning Attention
• Gated Delta Networks: Improving Mamba2 with Delta Rule
• When Linear Attention Meets Autoregressive Decoding: Towards More Effective and Efficient Linearized Large Language Models
• Loki: Low-rank Keys for Efficient Sparse Attention
• Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality
• Simple linear attention language models balance the recall-throughput tradeoff
• The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry
• Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models
• Gated Linear Attention Transformers with Hardware-Efficient Training
• Mamba: Linear-Time Sequence Modeling with Selective State Spaces
• White-Box Transformers via Sparse Rate Reduction: Compression Is All There Is?
• FLatten Transformer: Vision Transformer using Focused Linear Attention
• TransNormerLLM: A Faster and Better Large Language Model with Improved TransNormer
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• A Note on Normalized Emergence Timing (in Pythia Language Model Evaluations)
• The Devil in Linear Transformer
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• cosFormer: Rethinking Softmax in Attention
• FMMformer: Efficient and Flexible Transformer via Decomposed Near-field and Far-field Attention
• Luna: Linear Unified Nested Attention
• RoFormer: Enhanced Transformer with Rotary Position Embedding
• Linear Transformers Are Secretly Fast Weight Programmers
• Measuring Massive Multitask Language Understanding
• Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
• Conv-Linformer: Boosting Linformer's Performance with Convolution in Small-Scale Settings
• PIQA: Reasoning about Physical Commonsense in Natural Language
• Efficient Attention: Attention with Linear Complexities
• Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge
• Attention Is All You Need
• EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction

Related

• Gated Linear Attention Transformers with Hardware-Efficient Training
• Lean Attention: Hardware-Aware Scalable Attention Mechanism for the Decode-Phase of Transformers
• MonarchAttention: Zero-Shot Conversion to Fast, Hardware-Aware Structured Attention
• Scaling Context Requires Rethinking Attention
• Stable, Fast and Accurate: Kernelized Attention with Relative Positional Encoding