Block-Sparse & Linear Attention Kernels

Dense attention computes every token-to-token interaction. For many workloads, most of these interactions are unnecessary — a token 50K positions away rarely matters. Sparse and linear attention kernels exploit this structure for efficiency.

---

1. Why Sparse Attention?

Dense attention: $O(S^2)$ compute and memory. For S=128K, that’s 16 billion operations per head per layer.

But attention weights are often sparse in practice — most attention mass concentrates on:

• Nearby tokens (local context)
• A few globally important tokens (BOS, separator, key entities)
• Structured patterns (every-Nth token for summarization)

If we can skip the near-zero entries, we save compute proportional to the sparsity.

---

2. Block-Sparse Attention

The Idea

Instead of computing every position pair, define a mask at the block level: which blocks of Q attend to which blocks of K/V.

Block mask (16×16 blocks, each covers 256 tokens):

    K blocks: 0  1  2  3  4  5  6  7 ...
Q block 0:    ■  ■  □  □  □  □  □  □      ← causal, attend to self + neighbor
Q block 1:    ■  ■  ■  □  □  □  □  □
Q block 2:    □  ■  ■  ■  □  □  □  □      ← sliding window pattern
Q block 3:    □  □  ■  ■  ■  □  □  □
...

■ = compute attention (dense within block)
□ = skip (zero attention)

Kernel Implementation

Block-sparse attention uses the same FlashAttention tiling — but skips tiles where the mask says zero:

# Pseudocode: block-sparse FlashAttention
for each Q block i:
    for each K/V block j WHERE mask[i][j] == 1:  # skip zero blocks
        compute_attention_tile(Q_i, K_j, V_j)

If the mask has density $d$ (fraction of non-zero blocks), compute is $O(d \times S^2)$ instead of $O(S^2)$.

Common Sparse Patterns

Pattern	Description	Density	Use case
Causal	Lower triangle	50%	Standard LLM (already used everywhere)
Sliding window	Local band of width $w$	$w/S$	Mistral, local context models
Strided	Every k-th block	$1/k$	Longformer-style global + local
Custom	Data-dependent	Varies	Dynamic sparse attention

Sliding Window Attention (Mistral)

Mistral uses a sliding window of $W=4096$ tokens — each token attends only to the nearest 4K tokens:

$$\text{Attention}(q_i, K, V) = \text{softmax}\left(\frac{q_i K_{[i-W:i]}^\top}{\sqrt{d_k}}\right) V_{[i-W:i]}$$

Complexity: $O(S \times W)$ instead of $O(S^2)$. For S=128K, W=4K: $32\times$ reduction.

KV cache benefit: Only the most recent $W$ tokens’ KV cache is needed — old entries can be evicted. This caps KV cache at a fixed size regardless of conversation length.

---

3. Linear Attention

The Idea

Standard attention:

$$O = \text{softmax}(QK^\top) V \qquad O(S^2 d_k)$$

Linear attention removes softmax and uses a kernel trick:

$$O = \phi(Q) (\phi(K)^\top V) \qquad O(S d_k^2)$$

By computing $\phi(K)^\top V$ first (shape $d_k \times d_k$), then multiplying by $\phi(Q)$, we avoid the $S \times S$ matrix entirely.

The Tradeoff

	Dense attention	Linear attention
Compute	$O(S^2 d_k)$	$O(S d_k^2)$
Memory	$O(S^2)$ or $O(S)$ with FlashAttn	$O(d_k^2)$ fixed state
Quality	Exact	Approximate (no softmax)
Best when	$S < d_k^2 / d_k = d_k$ (always for $S > 128$)	Never competitive on quality for LLMs

In practice: No production LLM uses pure linear attention. The approximation quality gap is too large for tasks requiring precise recall.

Hybrid Approaches

State-space models (Mamba) and hybrid architectures (Jamba) use linear-time recurrence for most layers and standard attention for a few “retrieval” layers. This is covered in the Attention Variants chapter.

---

4. Kernel Implementation Challenges

Block-Sparse Challenges

Challenge	Solution
Irregular memory access	Precompute block indices, use gather
Load imbalance	Dynamic block scheduling
Warp utilization	Choose block size ≥ warp size (32)
Mask representation	Compressed sparse row (CSR) or block-level bitmap

FlashAttention + Sparsity

FlashAttention V2 natively supports block-sparse masks:

from flash_attn import flash_attn_func

# Custom block-sparse mask
# block_mask[i][j] = True if Q block i attends to K block j
output = flash_attn_func(Q, K, V, causal=False, block_mask=block_mask)

The kernel skips tiles where block_mask[i][j] = False — no wasted compute.

---

5. When to Use Which

Attention type	Best for	KV cache	Quality
Dense causal	Standard LLM serving	Full (grows with S)	Best
Sliding window	Long context with local bias	Capped at W	Good (loses very distant context)
Block-sparse	Domain-specific patterns	Depends on pattern	Good if pattern matches task
Linear	Infinite context streaming	Fixed size	Worst (no precise recall)
Hybrid (Mamba+Attn)	Long context with some retrieval	Mixed	Good balance

Default recommendation: Use dense causal attention with FlashAttention V2. Switch to sliding window (Mistral-style) only if context length exceeds what KV cache can support. Use hybrid architectures for 100K+ context applications.

---

Self-Check

• Explain why most attention weights are near-zero and how sparse attention exploits this
• Explain sliding window attention and its KV cache benefit
• Explain the order-of-operations trick that makes linear attention O(S dk²)
• Explain why linear attention isn’t used in production LLMs
• Choose the right attention type for a given sequence length and quality requirement

---

References

• Longformer: The Long-Document Transformer (Beltagy et al., 2020)
• Efficient Transformers: A Survey (Tay et al., 2022)
• Mistral 7B (Jiang et al., 2023)
• Mamba (Gu & Dao, 2023)