Efficient Attention: Flash Attention, Sliding Window, Sparse

Three families of techniques have emerged to make attention practical at scale. They differ in whether they preserve exact attention or trade some accuracy for efficiency.

Flash Attention (exact, memory-efficient)

Flash Attention (Dao et al. 2022) computes the same attention output as the standard formula, but never materializes the full (n × n) score matrix in GPU HBM. Instead, it tiles the computation: load a block of Q, a block of K and V, compute the partial softmax with online normalization, write the partial result, repeat. This trades a bit of recompute for an enormous reduction in memory bandwidth. Flash Attention 2 added better parallelism; Flash Attention 3 supports FP8 and warp specialization for H100, hitting ~740 TFLOPs/s.

Sliding-window attention (approximate, local)

Restrict each token to attend to a local window of W tokens. O(n × W) instead of O(n²). Used by Mistral 7B, Mixtral, and Gemma 3. The model loses direct long-range attention but gains practical long-context efficiency. Often combined with a few "global" attention layers or "sink" tokens to preserve some long-range signal.

Sparse attention (approximate, structured)

Compute attention only over a structured sparse pattern: e.g., a diagonal band, plus a few global tokens, plus random connections (Longformer, BigBird). Used by Phi-3-small. Gets you O(n × log n) or O(n) at the cost of approximate output.

Code

Flash Attention is just a kernel call·python

import torch
import torch.nn.functional as F

# In modern PyTorch (>=2.0), F.scaled_dot_product_attention
# automatically dispatches to Flash Attention 2 or 3
# when the inputs are eligible (CUDA, FP16/BF16, etc).

Q = torch.randn(2, 32, 8192, 128, device='cuda', dtype=torch.float16)
K = torch.randn(2, 32, 8192, 128, device='cuda', dtype=torch.float16)
V = torch.randn(2, 32, 8192, 128, device='cuda', dtype=torch.float16)

with torch.backends.cuda.sdp_kernel(enable_flash=True):
    out = F.scaled_dot_product_attention(Q, K, V, is_causal=True)
# Same output as a textbook implementation, but never instantiates
# the full (8192, 8192) score matrix per head.

Sliding-window mask·python

def sliding_window_mask(seq_len, window):
    # Allow each position to attend to itself and `window` previous positions
    i = torch.arange(seq_len)[:, None]
    j = torch.arange(seq_len)[None, :]
    return (j > i) | (j < i - window + 1)   # True where masked

Efficient Attention: Flash Attention, Sliding Window, Sparse

Flash Attention (exact, memory-efficient)

Sliding-window attention (approximate, local)

Sparse attention (approximate, structured)

Code

External links

Exercise

Progress

Comments 0