# Optimizing Attention Computation for Long Sequences

Standard self-attention scales **O(n²)** in both time and memory with sequence length *n*, which becomes prohibitive for long sequences. Here are the main strategies to address this:

---

## 1. **Sparse Attention Patterns**
Instead of attending to all tokens, attend to a subset:
- **Local/Sliding Window** — each token attends only to nearby tokens (e.g., Longformer, Mistral's sliding window)
- **Strided/Dilated** — attend to every *k*-th token to capture long-range dependencies
- **Combined** — mix local + global attention (e.g., Longformer adds global tokens, BigBird adds random attention)

**Complexity:** O(n · w) where *w* is window size

## 2. **FlashAttention**
An **IO-aware** exact attention algorithm that:
- Tiles the computation to fit in GPU SRAM
- Avoids materializing the full N×N attention matrix in HBM
- Fuses the softmax, masking, and dropout into a single kernel

```python
# PyTorch 2.0+ built-in
from torch.nn.functional import scaled_dot_product_attention
output = scaled_dot_product_attention(Q, K, V, is_causal=True)  # auto-selects FlashAttention
```
**Complexity:** Still O(n²) FLOPs, but **O(n)** memory and much faster wall-clock time.

## 3. **Linear / Kernel Attention**
Approximate softmax(QKᵀ)V by decomposing it:
- **Key idea:** Use feature maps φ so that φ(Q)·φ(K)ᵀ ≈ softmax(QKᵀ)
- Compute as φ(Q) · (φ(K)ᵀ · V) — right-to-left associativity makes it **O(n · d²)**
- Examples: **Linear Transformers**, **Performers** (Random Feature Attention)

## 4. **Multi-Query & Grouped-Query Attention (MQA/GQA)**
Reduce the KV cache size, which is the bottleneck at inference:
- **MQA:** All heads share one K and V projection
- **GQA:** Groups of heads share K/V (used in Llama 2/3, Gemma)

This doesn't change asymptotic complexity but drastically reduces **memory bandwidth** and **KV cache size**.

## 5. **KV Cache Optimization**
For autoregressive inference:
- **Paged Attention** (vLLM) — manages KV cache like virtual memory pages, reducing fragmentation
- **KV Cache Quantization** — store K/V in FP8/INT8
- **Sliding Window Cache** — only keep the last *w* tokens in cache

## 6. **Ring Attention / Sequence Parallelism**
Distribute the sequence across multiple devices:
- Each device holds a chunk of the sequence
- K/V blocks are passed in a ring between devices
- Enables sequences of **millions of tokens** across a cluster

## 7. **Hierarchical / Chunked Approaches**
- **Chunked cross-attention** (Retrieval-augmented models)
- **Mamba / State Space Models** — replace attention entirely with O(n) recurrent computation
- **MEGABYTE** — patch-level then byte-level transformers

---

## Quick Decision Guide

| Scenario | Recommended Approach |
|---|---|
| Training, GPU-bound | **FlashAttention** (exact, fastest wall-clock) |
| Very long context (100K+) training | FlashAttention + **Ring Attention** |
| Inference, KV cache is bottleneck | **GQA** + **Paged Attention** + KV quantization |
| Need true subquadratic scaling | **Sparse attention** or **SSMs** (Mamba) |
| Resource-constrained | **Linear attention** approximations |

---

## Practical Starting Point

For most use cases, start with **FlashAttention** (it's exact and requires no architectural changes), then layer on GQA and sparse patterns if you still hit limits:

```python
# Hugging Face Transformers — enable FlashAttention 2
model = AutoModelForCausalLM.from_pretrained(
    "model-name",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
)
```

Would you like me to go deeper into any of these approaches?