FlashAttention, Linear Attention & Long Context: Efficient Transformer Attention Explained

The transformer architecture has revolutionized machine learning since its introduction in 2017, powering everything from large language models to protein structure prediction. At the heart of this architecture lies the attention mechanism—a powerful but computationally expensive operation that allows models to weigh the importance of different parts of an input sequence. As we push toward longer context windows and more efficient inference, understanding and optimizing attention has become one of the most active areas of research in deep learning.

The Standard Attention Mechanism

The standard scaled dot-product attention computes a weighted sum of values based on the similarity between queries and keys. Given input matrices $Q$ (queries), $K$ (keys), and $V$ (values), the attention output is computed as:

$\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

where $d_k$ is the dimension of the keys, and the scaling factor $\frac{1}{\sqrt{d_k}}$ prevents the dot products from growing too large in magnitude.

The Quadratic Bottleneck

The computational complexity of standard attention is $O(n^2 d)$, where $n$ is the sequence length and $d$ is the embedding dimension. More critically, the memory complexity is also $O(n^2)$ because we must materialize the full $n \times n$ attention matrix before applying softmax.

This quadratic scaling creates severe practical limitations:

• Memory constraints: A sequence of length 32,768 with 16 attention heads requires storing attention matrices totaling over 17 billion elements in fp16
• Compute bottleneck: The matrix multiplication $QK^T$ dominates computation time for long sequences
• IO overhead: Moving large attention matrices between high-bandwidth memory (HBM) and on-chip SRAM becomes the primary bottleneck on modern GPUs

For context, processing a 100,000 token sequence—roughly the length of a novel—would require materializing attention matrices with 10 billion elements per layer per head. This is simply intractable with standard attention.

Linear Attention: Removing the Quadratic Dependency

Linear attention reformulates the attention computation to achieve $O(n)$ complexity by leveraging the associativity of matrix multiplication. The key insight is that we can avoid explicitly computing the $n \times n$ attention matrix.

The Kernel Trick

Linear attention replaces the softmax with a feature map $\phi(\cdot)$ applied separately to queries and keys:

$\text{LinAttn}(Q,K,V) = \frac{\phi(Q)\phi(K)^TV}{\phi(Q)\phi(K)^T\mathbf{1}}$

where $\mathbf{1}$ is a vector of ones for normalization.

The crucial observation is that we can reorder the computation. Instead of computing $(\phi(Q)\phi(K)^T)V$ which requires the $n \times n$ intermediate matrix, we compute $\phi(Q)(\phi(K)^TV)$. Since $\phi(K)^TV$ produces a $d \times d$ matrix (where $d$ is the feature dimension), the overall complexity becomes $O(nd^2)$—linear in sequence length.

Common Feature Maps

Several feature maps have been proposed:

• ELU-based: $\phi(x) = \text{elu}(x) + 1$
• ReLU: $\phi(x) = \text{ReLU}(x)$
• Random Fourier features: Approximating the softmax kernel with randomized trigonometric features
• Identity with normalization: Simply using normalized queries and keys

The SLAB architecture introduces a simplified linear attention mechanism that combines efficient attention patterns with progressive re-parameterized batch normalization. This work demonstrates that carefully designed linear attention can match or exceed standard attention performance while maintaining computational efficiency on resource-constrained devices.

The Expressiveness Trade-off

Pure linear attention removes the softmax nonlinearity, which fundamentally changes what the attention mechanism can express. Softmax attention produces a proper probability distribution over positions, enabling sharp, selective attention patterns. Linear attention tends to produce smoother, more diffuse attention patterns.

This expressiveness gap motivated the development of more sophisticated variants that attempt to recover the modeling power of softmax attention while maintaining linear complexity.

Gated Linear Attention: The Best of Both Worlds

Gated Linear Attention (GLA) represents a significant advancement in efficient attention mechanisms, achieving strong performance while enabling hardware-efficient training. The key innovation is introducing data-dependent gating that allows the model to selectively forget or retain information.

The Recurrent Formulation

GLA can be understood through its recurrent form. At each time step $t$, we maintain a hidden state $S_t \in \mathbb{R}^{d_k \times d_v}$—a 2D matrix rather than the 1D vector of traditional RNNs:

$S_t = G_t \odot S_{t-1} + k_t^T v_t$

$o_t = q_t S_t$

where:

• $G_t$ is a data-dependent gating matrix that controls forgetting
• $k_t$ and $v_t$ are the key and value vectors at position $t$
• $q_t$ is the query vector
• $\odot$ denotes element-wise multiplication

The gating mechanism allows GLA to selectively decay the hidden state, addressing a key limitation of pure linear attention where all historical information accumulates without any forgetting mechanism.

The Parallel Formulation

For efficient training on GPUs, which excel at parallel computation, GLA can be reformulated in a chunkwise parallel manner. The sequence is divided into chunks, and within each chunk, a parallel attention-like computation is performed:

$O_{\text{intra}} = (Q K^T \odot M) V$

where $M$ is a causal mask incorporating the cumulative gate values. Between chunks, the hidden state is propagated recurrently. This hybrid approach enables:

• Parallel training: Chunks are processed with matrix operations that fully utilize GPU tensor cores
• Linear memory: The cross-chunk state is constant-size regardless of sequence length
• Efficient inference: The recurrent form enables $O(1)$ per-token generation

Hardware-Efficient Training

The GLA paper introduces a training algorithm that is explicitly designed around GPU memory hierarchies. Modern GPUs have multiple levels of memory:

1. HBM (High Bandwidth Memory): Large capacity (40-80GB) but relatively slow access
2. SRAM (on-chip): Small capacity (tens of MB) but extremely fast access

The algorithm minimizes data movement between HBM and SRAM by:

• Chunked computation: Processing the sequence in chunks that fit in SRAM
• Fused operations: Combining multiple operations into single GPU kernels to avoid intermediate writes to HBM
• Materialization-free backward pass: Computing gradients without storing the full attention matrix

This IO-aware design is critical for practical efficiency. Even with $O(n)$ computational complexity, an algorithm can be slower than $O(n^2)$ attention if it requires excessive memory transfers.

Practical Considerations and Trade-offs

When to Use Each Variant

Standard softmax attention remains the best choice when:

• Sequence lengths are moderate (under 4,096 tokens)
• Maximum expressiveness is required
• You're using optimized implementations like FlashAttention

Linear attention is preferable when:

• Extremely long sequences are required
• Inference latency is critical
• Memory is severely constrained

Gated linear attention offers a strong middle ground:

• Near-transformer quality on language modeling
• Efficient training and inference
• Good performance on tasks requiring selective memory

The Role of Specialized Hardware

The efficiency of attention mechanisms is deeply intertwined with hardware design. Tensor cores on modern GPUs are optimized for matrix multiplications of specific sizes. Attention implementations must be designed with these constraints in mind.

The trend toward longer context windows (128K tokens and beyond) is driving both algorithmic innovation and hardware co-design. We're seeing:

• Custom attention kernels: Hand-tuned implementations for specific attention patterns
• Sparse attention accelerators: Hardware support for block-sparse attention patterns
• Linear attention in hardware: Specialized units for recurrent-style computations

Looking Ahead

The landscape of efficient attention continues to evolve rapidly. Several directions show particular promise:

Hybrid architectures that combine different attention mechanisms at different layers or for different heads, leveraging the strengths of each approach.

Learned efficiency where models learn to allocate computation dynamically based on input complexity, using full attention only where necessary.

Hardware-software co-design as new accelerators are developed specifically for efficient attention patterns, enabling algorithms that would be impractical on current hardware.

The quadratic attention bottleneck that once seemed fundamental to transformers is increasingly becoming an engineering challenge rather than a theoretical limitation. As these efficient alternatives mature, we can expect to see transformers processing ever-longer sequences while becoming more accessible for deployment on edge devices and in resource-constrained environments.

Understanding these mechanisms—from the mathematical foundations of linear attention to the hardware-aware optimizations of gated variants—is essential for anyone working at the frontier of modern machine learning systems.