Sparse Attention Patterns: Local, Strided, Hash-Based, and Learnable Sparsity

Full self-attention computes a score between every pair of tokens in the sequence. For a sequence of length $n$, this produces an $n \times n$ attention matrix, requiring $O(n^2)$ time and memory. At $n = 2048$ (GPT-2 era), this was manageable. At $n = 32768$ (Llama 3), it costs $32768^2 = 1.07 \times 10^9$ entries per head per layer. At $n = 131072$ (Llama 3.1 extended), it costs $1.72 \times 10^{10}$ entries. At $n = 1000000$ (the frontier), it costs $10^{12}$ entries. The quadratic scaling is a hard wall.

Sparse attention exploits an empirical observation: in trained transformers, most attention weights are near zero. The attention matrix is sparse in practice, even though the computation is dense. Sparse attention methods predefine which token pairs can attend to each other, skipping the rest. This reduces the cost from $O(n^2)$ to $O(n \cdot k)$ where $k$ is the number of tokens each position attends to.

This post covers every major sparse attention pattern, derives their complexity, provides implementations, and explains why FlashAttention changed the calculus — making dense attention fast enough that sparse methods only win at extreme context lengths.

The Sparsity Observation

1.1 Attention Weight Distribution

After training a standard transformer, examine the attention weights (post-softmax values) across all heads and layers. The distribution is highly concentrated:

• The top 10% of attention weights capture 80-95% of the total mass
• Most off-diagonal entries are less than $\frac{1}{n}$ (the uniform attention baseline)
• The pattern varies by layer: early layers attend locally, late layers attend globally

This means the $O(n^2)$ attention computation produces mostly near-zero values. Sparse attention avoids computing these near-zero entries entirely.

1.2 Formalizing Sparsity

Define a sparsity pattern $S \subseteq \{1, \ldots, n\} \times \{1, \ldots, n\}$ as the set of allowed attention connections. For full attention, $|S| = n^2$. For sparse attention, $|S| = O(n \cdot k)$ where $k \ll n$.

The sparse attention computation:

$\text{Attn}(Q, K, V)_i = \sum_{j \in S(i)} \frac{\exp(q_i \cdot k_j / \sqrt{d_k})}{\sum_{j' \in S(i)} \exp(q_i \cdot k_{j'} / \sqrt{d_k})} v_j$

where $S(i) = \{j : (i, j) \in S\}$ is the set of keys that position $i$ is allowed to attend to.

The softmax normalization is over the sparse set $S(i)$ only. This is mathematically different from full attention followed by zeroing out entries — the normalization denominator changes.

Local (Windowed) Attention

2.1 Definition

Local attention restricts each token to attend to a fixed window of $W$ neighboring tokens:

$S_{\text{local}}(i) = \{j : |i - j| \leq W/2\}$

For a window size $W$, each token attends to at most $W$ keys. The total number of attention entries is $O(n \cdot W)$.

For causal (autoregressive) models, the window is one-sided:

$S_{\text{causal-local}}(i) = \{j : \max(0, i - W + 1) \leq j \leq i\}$

2.2 Complexity

• Time: $O(n \cdot W \cdot d_k)$ for the QK computation, $O(n \cdot W \cdot d_v)$ for the AV multiplication
• Memory: $O(n \cdot W)$ for the attention weights
• Speedup over full: $\frac{n}{W}$

For $n = 131072$ and $W = 4096$: speedup = 32x. For $n = 1000000$ and $W = 4096$: speedup = 244x.

2.3 Limitation

Local attention cannot capture long-range dependencies. If two tokens are more than $W$ positions apart, they cannot directly attend to each other. Information can only flow long-range through multiple layers, with each layer propagating information by $W$ positions. To propagate information across a sequence of length $n$, you need at least $\lceil n/W \rceil$ layers.

For $n = 131072$ and $W = 4096$: 32 layers needed for information to propagate end-to-end. For $n = 1000000$ and $W = 4096$: 244 layers needed. Most models have 32-80 layers, so local attention alone cannot support very long contexts.

2.4 Implementation

import torch
import torch.nn.functional as F
import math

def local_attention(q, k, v, window_size, causal=True):
    """
    Local windowed attention with causal masking.

    Args:
        q: (B, H, S, D) queries
        k: (B, H, S, D) keys
        v: (B, H, S, D) values
        window_size: number of past tokens to attend to
        causal: if True, only attend to past tokens
    Returns:
        output: (B, H, S, D) attention output
    """
    B, H, S, D = q.shape
    scale = 1.0 / math.sqrt(D)

    output = torch.zeros_like(q)

    for i in range(S):
        # Define the window for position i
        if causal:
            start = max(0, i - window_size + 1)
            end = i + 1
        else:
            start = max(0, i - window_size // 2)
            end = min(S, i + window_size // 2 + 1)

        # Extract keys and values in the window
        k_window = k[:, :, start:end, :]   # (B, H, W_eff, D)
        v_window = v[:, :, start:end, :]   # (B, H, W_eff, D)

        # Compute attention scores for position i
        q_i = q[:, :, i:i+1, :]            # (B, H, 1, D)
        scores = torch.matmul(q_i, k_window.transpose(-2, -1)) * scale
        weights = F.softmax(scores, dim=-1)
        output[:, :, i:i+1, :] = torch.matmul(weights, v_window)

    return output

The efficient implementation tiles the computation into blocks and processes each block as a dense attention over a $W \times W$ submatrix:

def local_attention_blocked(q, k, v, window_size, causal=True):
    """
    Blocked local attention -- more GPU-friendly.

    Process attention in blocks of size window_size.
    Each block attends to itself and the previous block.
    """
    B, H, S, D = q.shape
    scale = 1.0 / math.sqrt(D)
    W = window_size

    # Pad sequence to multiple of window_size
    pad = (W - S % W) % W
    if pad > 0:
        q = F.pad(q, (0, 0, 0, pad))
        k = F.pad(k, (0, 0, 0, pad))
        v = F.pad(v, (0, 0, 0, pad))

    S_padded = q.shape[2]
    n_blocks = S_padded // W

    # Reshape into blocks: (B, H, n_blocks, W, D)
    q_blocks = q.view(B, H, n_blocks, W, D)
    k_blocks = k.view(B, H, n_blocks, W, D)
    v_blocks = v.view(B, H, n_blocks, W, D)

    outputs = []
    for block_idx in range(n_blocks):
        q_block = q_blocks[:, :, block_idx]  # (B, H, W, D)

        # Current block attends to itself + previous block
        if block_idx == 0:
            k_ctx = k_blocks[:, :, 0]
            v_ctx = v_blocks[:, :, 0]
        else:
            k_ctx = torch.cat([
                k_blocks[:, :, block_idx - 1],
                k_blocks[:, :, block_idx]
            ], dim=2)  # (B, H, 2W, D)
            v_ctx = torch.cat([
                v_blocks[:, :, block_idx - 1],
                v_blocks[:, :, block_idx]
            ], dim=2)

        # Dense attention within the block context
        scores = torch.matmul(q_block, k_ctx.transpose(-2, -1)) * scale

        if causal:
            # Build causal mask for this block
            ctx_len = k_ctx.shape[2]
            mask = torch.ones(W, ctx_len, dtype=torch.bool,
                              device=q.device)
            for qi in range(W):
                global_qi = block_idx * W + qi
                for ki in range(ctx_len):
                    if block_idx == 0:
                        global_ki = ki
                    else:
                        global_ki = (block_idx - 1) * W + ki
                    if global_ki > global_qi:
                        mask[qi, ki] = False
            scores = scores.masked_fill(~mask, float('-inf'))

        weights = F.softmax(scores, dim=-1)
        out_block = torch.matmul(weights, v_ctx)
        outputs.append(out_block)

    output = torch.stack(outputs, dim=2).view(B, H, S_padded, D)
    return output[:, :, :S, :]  # Remove padding

Strided Attention

3.1 Definition

Strided attention allows each token to attend to every $k$-th token in the sequence, providing global coverage with $O(n/k)$ connections per position:

$S_{\text{stride}}(i) = \{j : j \mod k = i \mod k\}$

Alternatively, a fixed-stride pattern:

$S_{\text{stride}}(i) = \{j : j \mod k = 0\} \cup \{i\}$

This ensures every position attends to the same set of “landmark” positions (every $k$-th token), plus itself.

3.2 Complexity

• Entries per position: $n/k$
• Total entries: $O(n^2/k)$
• Memory: $O(n^2/k)$

For $k = 128$, this is a 128x reduction from full attention.

3.3 The Sparse Transformer (Child et al., 2019)

The Sparse Transformer combines local and strided attention in a two-head pattern:

• Head pattern A: Local attention with window $W = \sqrt{n}$
• Head pattern B: Strided attention with stride $k = \sqrt{n}$

Together, any two tokens can communicate in at most 2 hops: token $i$ sends information to the nearest landmark via Head A’s local window, and the landmark sends information to any other token via Head B’s stride.

Total entries per position: $O(\sqrt{n})$ from each pattern, so $O(\sqrt{n})$ total. Total computation: $O(n\sqrt{n})$.

def strided_attention_mask(seq_len, stride):
    """
    Create a strided attention mask.
    Each position attends to every stride-th position.
    """
    mask = torch.zeros(seq_len, seq_len, dtype=torch.bool)

    for i in range(seq_len):
        # Attend to every stride-th position
        for j in range(0, seq_len, stride):
            mask[i, j] = True
        # Always attend to self
        mask[i, i] = True

    return mask

def combined_sparse_attention(q, k, v, window_size, stride):
    """
    Sparse Transformer style: union of local + strided patterns.
    """
    B, H, S, D = q.shape
    scale = 1.0 / math.sqrt(D)

    # Build combined mask: local OR strided
    local_mask = torch.zeros(S, S, dtype=torch.bool, device=q.device)
    stride_mask = torch.zeros(S, S, dtype=torch.bool, device=q.device)

    for i in range(S):
        # Local: attend to window
        start = max(0, i - window_size + 1)
        local_mask[i, start:i+1] = True

        # Strided: attend to every stride-th position up to i
        for j in range(0, i + 1, stride):
            stride_mask[i, j] = True

    combined_mask = local_mask | stride_mask  # Union

    # Compute attention with combined mask
    scores = torch.matmul(q, k.transpose(-2, -1)) * scale
    scores = scores.masked_fill(~combined_mask, float('-inf'))
    weights = F.softmax(scores, dim=-1)
    return torch.matmul(weights, v)

BigBird: Local + Random + Global

4.1 The Three Components

BigBird (Zaheer et al., 2020) combines three attention patterns:

1. Local attention: Each token attends to $W$ neighbors (same as section 2)
2. Random attention: Each token attends to $R$ randomly selected tokens
3. Global attention: $G$ designated tokens attend to (and are attended by) all tokens

The combined pattern:

$S_{\text{BigBird}}(i) = S_{\text{local}}(i) \cup S_{\text{random}}(i) \cup S_{\text{global}}$

4.2 Why This Combination Works

BigBird’s theoretical contribution is proving that this combination is a universal approximator of sequence functions, while pure local attention is not. The key insight is from random graph theory:

A random graph with $O(\log n)$ random edges per node is connected with high probability (Erdos-Renyi). Adding $R = O(\log n)$ random attention edges per token ensures that information can flow between any two tokens in $O(\log n / \log \log n)$ hops, even if they are far apart.

The global tokens serve as “hubs” that aggregate and broadcast information. With $G$ global tokens, the information pathway between any two tokens $i$ and $j$ is:

$i \to \text{global token} \to j$

This is 2 hops, regardless of the distance between $i$ and $j$.

4.3 Complexity

• Local: $n \cdot W$ entries
• Random: $n \cdot R$ entries
• Global: $2 \cdot n \cdot G$ entries (global tokens attend to all, and all attend to global tokens)
• Total: $O(n \cdot (W + R + G))$, which is $O(n)$ when $W$, $R$, $G$ are constants

Typical values: $W = 64$, $R = 3$, $G = 2$. Total entries per position: $\sim 69$.

4.4 Implementation

def bigbird_attention_mask(seq_len, window_size, n_random, n_global):
    """
    Create BigBird attention mask: local + random + global.

    Args:
        seq_len: sequence length
        window_size: local attention window
        n_random: number of random attention connections per token
        n_global: number of global tokens (first n_global tokens)
    """
    mask = torch.zeros(seq_len, seq_len, dtype=torch.bool)

    for i in range(seq_len):
        # 1. Local attention
        start = max(0, i - window_size // 2)
        end = min(seq_len, i + window_size // 2 + 1)
        mask[i, start:end] = True

        # 2. Random attention
        candidates = list(range(seq_len))
        candidates.remove(i)
        random_targets = torch.randperm(len(candidates))[:n_random]
        for idx in random_targets:
            mask[i, candidates[idx]] = True

        # 3. Global attention (first n_global tokens)
        mask[i, :n_global] = True    # Every token attends to global tokens
        mask[:n_global, i] = True    # Global tokens attend to every token

    return mask

# Example: BigBird mask for short sequence
mask = bigbird_attention_mask(
    seq_len=32, window_size=8, n_random=3, n_global=2
)
# Density: mask.float().mean() shows fraction of entries that are non-zero
density = mask.float().mean().item()
print(f"Attention density: {density:.2%}")  # Much less than 100%

Longformer: Local + Global Sentinels

5.1 Architecture

Longformer (Beltagy et al., 2020) simplifies BigBird by removing random attention and using task-specific global tokens:

1. Local attention: Sliding window of size $W$ for all tokens
2. Global attention: Selected tokens (e.g., [CLS], question tokens in QA) have full attention

The global tokens are not fixed — they are chosen based on the task:

• Classification: [CLS] token is global
• Question answering: all question tokens are global
• Summarization: specific sentinel tokens are global

5.2 Implementation Details

Longformer uses different projections for local and global attention:

• Local attention uses $Q_l, K_l, V_l$ projections
• Global attention uses $Q_g, K_g, V_g$ projections

This doubles the parameter count for attention weights but allows the model to learn different attention patterns for local vs global contexts.

class LongformerAttention(nn.Module):
    def __init__(self, d_model, n_heads, window_size):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.window_size = window_size

        # Local attention projections
        self.Q_local = nn.Linear(d_model, d_model, bias=False)
        self.K_local = nn.Linear(d_model, d_model, bias=False)
        self.V_local = nn.Linear(d_model, d_model, bias=False)

        # Global attention projections (separate parameters)
        self.Q_global = nn.Linear(d_model, d_model, bias=False)
        self.K_global = nn.Linear(d_model, d_model, bias=False)
        self.V_global = nn.Linear(d_model, d_model, bias=False)

        self.out_proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, global_mask):
        """
        Args:
            x: (B, S, D) input
            global_mask: (B, S) bool tensor, True for global tokens
        """
        B, S, D = x.shape
        H = self.n_heads
        dk = self.d_k
        W = self.window_size
        scale = 1.0 / math.sqrt(dk)

        # Compute all projections
        q_l = self.Q_local(x).view(B, S, H, dk).transpose(1, 2)
        k_l = self.K_local(x).view(B, S, H, dk).transpose(1, 2)
        v_l = self.V_local(x).view(B, S, H, dk).transpose(1, 2)

        q_g = self.Q_global(x).view(B, S, H, dk).transpose(1, 2)
        k_g = self.K_global(x).view(B, S, H, dk).transpose(1, 2)
        v_g = self.V_global(x).view(B, S, H, dk).transpose(1, 2)

        output = torch.zeros_like(q_l)

        for b in range(B):
            global_indices = global_mask[b].nonzero(as_tuple=True)[0]
            n_global = len(global_indices)

            for i in range(S):
                if global_mask[b, i]:
                    # Global token: full attention using global projections
                    q_i = q_g[b, :, i:i+1, :]
                    scores = torch.matmul(q_i, k_g[b].transpose(-2, -1))
                    scores = scores * scale
                    weights = F.softmax(scores, dim=-1)
                    output[b, :, i:i+1, :] = torch.matmul(weights, v_g[b])
                else:
                    # Local token: window + global tokens
                    start = max(0, i - W // 2)
                    end = min(S, i + W // 2 + 1)
                    local_indices = torch.arange(start, end,
                                                  device=x.device)

                    # Combine local window indices with global indices
                    all_indices = torch.cat([
                        local_indices, global_indices
                    ]).unique()

                    q_i = q_l[b, :, i:i+1, :]
                    k_ctx = k_l[b, :, all_indices, :]
                    v_ctx = v_l[b, :, all_indices, :]

                    scores = torch.matmul(q_i, k_ctx.transpose(-2, -1))
                    scores = scores * scale
                    weights = F.softmax(scores, dim=-1)
                    output[b, :, i:i+1, :] = torch.matmul(weights, v_ctx)

        output = output.transpose(1, 2).contiguous().view(B, S, D)
        return self.out_proj(output)

Hash-Based Attention (Reformer)

6.1 Locality-Sensitive Hashing (LSH)

Reformer (Kitaev et al., 2020) uses locality-sensitive hashing to identify which key-query pairs will have high attention scores, then only computes attention within hash buckets.

The core idea: if $q_i$ and $k_j$ have high dot product (high attention score), they point in similar directions. A hash function that maps similar vectors to the same bucket will group together the query-key pairs that matter.

LSH for angular similarity uses random hyperplane projections:

$h(x) = \text{sign}(Rx)$

where $R \in \mathbb{R}^{b \times d}$ is a random matrix and $b$ is the number of hash bits. Vectors pointing in similar directions will have the same sign pattern with high probability.

6.2 The Reformer Attention Algorithm

1. Set $Q = K$ (shared QK attention — Reformer uses this to ensure queries and keys are in the same space)
2. Hash all queries/keys: $h(q_i)$ for all $i$
3. Sort tokens by hash bucket
4. Within each bucket, compute full attention
5. Use multiple hash rounds to reduce the chance of missing important pairs

def lsh_attention(q, v, n_hashes=8, n_buckets=64):
    """
    Simplified LSH attention (Reformer style).
    Uses shared QK (q serves as both query and key).

    Args:
        q: (B, H, S, D) queries (also used as keys)
        v: (B, H, S, D) values
        n_hashes: number of hash rounds
        n_buckets: number of hash buckets
    """
    B, H, S, D = q.shape
    scale = 1.0 / math.sqrt(D)

    # Accumulate attention from multiple hash rounds
    all_outputs = torch.zeros_like(q)
    all_log_weights = torch.full((B, H, S, 1), float('-inf'),
                                  device=q.device)

    for round_idx in range(n_hashes):
        # Random projection for this hash round
        random_proj = torch.randn(D, n_buckets // 2, device=q.device)
        # Project and hash: use sign of projection, concatenate
        # with negation for balanced buckets
        proj = torch.matmul(q, random_proj)  # (B, H, S, n_buckets//2)
        hash_codes = torch.argmax(
            torch.cat([proj, -proj], dim=-1), dim=-1
        )  # (B, H, S) bucket assignments

        # For each bucket, compute attention among its members
        output_round = torch.zeros_like(q)

        for bucket_id in range(n_buckets):
            # Find tokens in this bucket
            bucket_mask = (hash_codes == bucket_id)  # (B, H, S)

            for b in range(B):
                for h in range(H):
                    indices = bucket_mask[b, h].nonzero(as_tuple=True)[0]
                    if len(indices) == 0:
                        continue

                    q_bucket = q[b, h, indices]  # (bucket_size, D)
                    v_bucket = v[b, h, indices]  # (bucket_size, D)

                    # Full attention within bucket
                    scores = torch.matmul(
                        q_bucket, q_bucket.transpose(-2, -1)
                    ) * scale

                    # Causal mask within bucket
                    bucket_size = len(indices)
                    causal = torch.tril(
                        torch.ones(bucket_size, bucket_size,
                                   device=q.device, dtype=torch.bool)
                    )
                    # Map back to original positions for causal ordering
                    for qi in range(bucket_size):
                        for ki in range(bucket_size):
                            if indices[ki] > indices[qi]:
                                causal[qi, ki] = False

                    scores = scores.masked_fill(~causal, float('-inf'))
                    weights = F.softmax(scores, dim=-1)
                    out = torch.matmul(weights, v_bucket)

                    output_round[b, h, indices] = out

        all_outputs += output_round

    # Average over hash rounds
    return all_outputs / n_hashes

6.3 Complexity Analysis

• Hash computation: $O(n \cdot d \cdot b)$ per round, where $b$ is the number of hash bits
• Sorting by bucket: $O(n \log n)$
• Attention within buckets: if each bucket has $\sim n / n_{\text{buckets}}$ tokens, the total attention cost is $n_{\text{buckets}} \cdot (n / n_{\text{buckets}})^2 = n^2 / n_{\text{buckets}}$
• With $n_{\text{buckets}} = O(n / \log n)$: total cost is $O(n \log n)$

Multiple hash rounds (typically 4-8) multiply the cost by a constant factor.

6.4 Limitations

1. Shared QK requirement: Reformer ties $Q = K$ to ensure queries and keys are in the same hash space. This removes a degree of freedom from the attention mechanism.
2. Sorting overhead: Sorting by hash bucket is $O(n \log n)$ and not GPU-friendly (irregular memory access patterns).
3. Bucket size variance: Some buckets may have many tokens, others few. This creates load imbalance on GPUs.
4. Approximation quality: LSH is probabilistic. Important attention pairs may be missed if they fall in different buckets. Multiple rounds mitigate this but increase cost.

Learnable Sparsity

7.1 The Idea

Instead of hand-designing the sparsity pattern, let the model learn which tokens to attend to. Several approaches exist:

Routing-based: Use a lightweight scoring network to predict which keys are relevant for each query, then only compute full attention for the top-$k$ pairs.

Threshold-based: Compute a cheap approximation of attention scores (e.g., using low-rank projections) and only compute full attention for pairs above a threshold.

7.2 Top-k Sparse Attention

For each query $q_i$, compute a cheap relevance score for all keys, then select the top $k$ and compute full attention only over those:

class TopKSparseAttention(nn.Module):
    def __init__(self, d_model, n_heads, top_k=256):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.top_k = top_k

        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

        # Low-rank scoring network for cheap relevance estimation
        self.score_rank = 32
        self.score_q = nn.Linear(d_model, self.score_rank, bias=False)
        self.score_k = nn.Linear(d_model, self.score_rank, bias=False)

    def forward(self, x, causal=True):
        B, S, D = x.shape
        H = self.n_heads
        dk = self.d_k
        k_select = min(self.top_k, S)
        scale = 1.0 / math.sqrt(dk)

        # Full projections (computed for all positions)
        q = self.W_q(x).view(B, S, H, dk).transpose(1, 2)
        k = self.W_k(x).view(B, S, H, dk).transpose(1, 2)
        v = self.W_v(x).view(B, S, H, dk).transpose(1, 2)

        # Cheap scoring: low-rank dot product to estimate relevance
        score_q = self.score_q(x)  # (B, S, rank)
        score_k = self.score_k(x)  # (B, S, rank)
        cheap_scores = torch.matmul(
            score_q, score_k.transpose(-2, -1)
        )  # (B, S, S)

        # Apply causal mask to cheap scores
        if causal:
            causal_mask = torch.triu(
                torch.ones(S, S, device=x.device, dtype=torch.bool),
                diagonal=1
            )
            cheap_scores = cheap_scores.masked_fill(causal_mask,
                                                     float('-inf'))

        # Select top-k keys for each query
        _, top_indices = cheap_scores.topk(k_select, dim=-1)
        # top_indices: (B, S, k_select)

        # Gather selected keys and values
        output = torch.zeros(B, H, S, dk, device=x.device)

        for b in range(B):
            for i in range(S):
                idx = top_indices[b, i]  # (k_select,)
                k_sel = k[b, :, idx, :]  # (H, k_select, dk)
                v_sel = v[b, :, idx, :]  # (H, k_select, dk)
                q_i = q[b, :, i:i+1, :]  # (H, 1, dk)

                scores = torch.matmul(
                    q_i, k_sel.transpose(-2, -1)
                ) * scale  # (H, 1, k_select)
                weights = F.softmax(scores, dim=-1)
                output[b, :, i:i+1, :] = torch.matmul(weights, v_sel)

        output = output.transpose(1, 2).contiguous().view(B, S, D)
        return self.W_o(output)

7.3 Complexity

• Cheap scoring: $O(n^2 r)$ where $r$ is the scoring rank (32-64, much smaller than $d$)
• Top-k selection: $O(n^2 \log k)$ (can be done with partial sort)
• Sparse attention: $O(n k d)$
• Total: $O(n^2 r + nkd)$

If $r \ll d$ and $k \ll n$, this is cheaper than full attention’s $O(n^2 d)$. But the cheap scoring step is still $O(n^2)$ in the sequence length, just with a much smaller constant.

7.4 The Fundamental Challenge

Learnable sparsity faces a chicken-and-egg problem: to decide which tokens to attend to, you need some information about the token representations, but the representations depend on the attention output. Most approaches use the representations from the previous layer or a cheap approximation, which introduces a one-step lag.

In practice, learnable sparsity has not been widely adopted because:

1. The scoring overhead partially offsets the savings
2. The top-k selection is not differentiable (requires straight-through estimators or Gumbel-softmax)
3. The irregular memory access patterns from gathered indices are slow on GPUs

Why Sparse Attention Lost to FlashAttention

8.1 The IO Bottleneck

The key insight of FlashAttention (Dao et al., 2022): standard attention is bottlenecked by memory IO, not computation. The attention matrix of size $O(n^2)$ must be written to and read from GPU HBM (high-bandwidth memory). The actual matrix multiplications are fast; the memory transfers are slow.

FlashAttention never materializes the full $n \times n$ attention matrix. Instead, it tiles the computation into blocks that fit in GPU SRAM (on-chip memory, ~20MB on A100) and computes attention one block at a time, accumulating the output using the online softmax trick.

The result: FlashAttention computes mathematically exact full attention with:

• $O(n^2 d / M)$ HBM accesses (where $M$ is SRAM size), compared to $O(n^2 + n^2 d)$ for standard attention
• No $O(n^2)$ intermediate storage
• 2-4x wall-clock speedup over standard attention

8.2 The Break-Even Point

Sparse attention reduces FLOPs from $O(n^2 d)$ to $O(nkd)$. FlashAttention does not reduce FLOPs — it still computes $O(n^2 d)$ FLOPs — but it reduces memory IO by a factor of $O(M/d)$ where $M$ is SRAM size.

On an A100 GPU:

• SRAM: ~20MB, which holds ~5 million float16 values
• HBM bandwidth: 2 TB/s
• Compute: 312 TFLOPS (float16)
• Arithmetic intensity needed to saturate compute: $312 / 2 = 156$ FLOPS/byte

FlashAttention’s tiled computation achieves high arithmetic intensity because it reuses data in SRAM. Sparse attention has lower FLOPs but worse memory access patterns (irregular gather/scatter operations), which reduces its effective throughput.

The crossover point: sparse attention with $k = n/4$ (75% sparsity) is faster than FlashAttention only when $n$ is large enough that the FLOP reduction outweighs the memory efficiency loss. Empirically, this crossover is around $n = 65536$ to $n = 131072$ for common sparsity patterns.

8.3 The Current Landscape

As of 2025, the practical situation is:

• Context length up to 32K: FlashAttention (dense) is faster than all sparse methods. This covers GPT-4, Claude, and most production LLMs.
• Context length 32K-128K: Hybrid approaches win. FlashAttention with sliding window (local attention) for most layers, full attention for a few layers.
• Context length above 128K: Sparse attention or linear attention is necessary. Ring attention (distributing the sequence across GPUs) combined with local attention is the current approach.

Llama 3.1 with 128K context uses a hybrid: local attention with $W = 32768$ in most layers, with a few layers using full attention for global coverage. This matches the BigBird-style reasoning (local for most, global for some) but implemented with FlashAttention’s IO-efficient kernels.

Implementation: Local Attention with Causal Mask (Production Quality)

Here is a production-oriented local attention implementation that works with FlashAttention-style tiling:

import torch
import torch.nn.functional as F
import math

class SlidingWindowAttention(nn.Module):
    """
    Sliding window (local) attention with causal masking.
    Compatible with standard transformer architectures.

    For use in models targeting long contexts (32K+) where
    full attention is too expensive but FlashAttention alone
    is not enough.
    """

    def __init__(self, d_model, n_heads, window_size=4096):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.window_size = window_size

        self.W_qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

    def _build_sliding_window_mask(self, seq_len, device):
        """Build a causal sliding window mask."""
        # Start with causal mask
        mask = torch.tril(torch.ones(seq_len, seq_len,
                                      device=device, dtype=torch.bool))
        # Apply window: zero out positions beyond window_size
        for i in range(seq_len):
            if i >= self.window_size:
                mask[i, :i - self.window_size + 1] = False
        return mask

    def forward(self, x):
        B, S, D = x.shape
        H = self.n_heads
        dk = self.d_k
        scale = 1.0 / math.sqrt(dk)

        # Fused QKV projection
        qkv = self.W_qkv(x).view(B, S, 3, H, dk)
        q, k, v = qkv.unbind(dim=2)
        q = q.transpose(1, 2)  # (B, H, S, dk)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # Compute full attention scores
        scores = torch.matmul(q, k.transpose(-2, -1)) * scale

        # Apply sliding window + causal mask
        mask = self._build_sliding_window_mask(S, x.device)
        scores = scores.masked_fill(~mask, float('-inf'))

        weights = F.softmax(scores, dim=-1)
        output = torch.matmul(weights, v)

        output = output.transpose(1, 2).contiguous().view(B, S, D)
        return self.W_o(output)

# Test: verify that local attention produces reasonable outputs
torch.manual_seed(42)
model = SlidingWindowAttention(d_model=512, n_heads=8, window_size=64)
x = torch.randn(2, 256, 512)

with torch.no_grad():
    out = model(x)

print(f"Input shape:  {x.shape}")
print(f"Output shape: {out.shape}")
print(f"Output mean:  {out.mean():.6f}")
print(f"Output std:   {out.std():.6f}")

# Verify the attention mask
mask = model._build_sliding_window_mask(256, torch.device('cpu'))
# Position 100 should attend to positions 37-100 (window=64)
assert mask[100, 36] == False  # Too far back
assert mask[100, 37] == True   # Start of window
assert mask[100, 100] == True  # Self-attention
assert mask[100, 101] == False # Future (causal)
print("Mask verification passed")

Summary

Sparse attention trades the completeness of full attention for reduced computation. The major patterns:

Pattern	Entries per token	Long-range	Key insight
Local	$W$	No	Locality bias matches natural language
Strided	$n/k$	Yes	Landmarks provide global coverage
BigBird	$W + R + G$	Yes	Random edges ensure connectivity
Longformer	$W + G$	Yes	Task-specific global tokens
Reformer	$O(\log n)$ expected	Probabilistic	Similar vectors hash together
Learnable	$k$ (chosen)	Yes	Model selects relevant tokens

The practical outcome: FlashAttention made dense attention fast enough for context lengths up to 32K-64K by optimizing memory IO instead of reducing FLOPs. Above 128K tokens, sparse patterns (particularly sliding window with global sentinel tokens) remain necessary. The winning architecture for long-context models is a hybrid: mostly local attention with a few full-attention layers, implemented with IO-efficient kernels.