The Hitchhiker's Guide — Phase 08: Inference Optimization at Depth

"FlashAttention is the most important algorithmic contribution to LLM inference since the transformer itself." — Tri Dao

Section 1: FlashAttention — The Complete Derivation

Why Standard Attention is IO-Bound

Standard scaled dot-product attention:

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

The naive implementation materializes the $n \times n$ attention matrix $S = QK^T / \sqrt{d_k}$ in HBM (GPU DRAM), then the softmax of $S$, then the final product with $V$:

Q [n×d] → HBM
K [n×d] → HBM  
V [n×d] → HBM
S [n×n] ← COMPUTE: QKᵀ → WRITE to HBM  ← n² memory
P [n×n] ← READ S from HBM → softmax → WRITE to HBM  ← n² memory
O [n×d] ← READ P from HBM → PV → WRITE to HBM

Total HBM reads/writes: $O(n^2 d + n^2) = O(n^2)$ — quadratic in sequence length.

At $n=8192$, $d=128$: $8192^2 \times 2$ bytes = 134 MB of intermediate attention matrices. On A100 with 2 TB/s bandwidth: $134 \text{ MB} / 2 \text{ TB/s} = 67$ µs just for memory traffic — before any compute.

FlashAttention: IO-Optimal Attention

The key insight: we don't need to store the full attention matrix. We can compute the attention output tile-by-tile, using only SRAM (on-chip memory), by tracking running statistics.

Online softmax: given the recurrence, we can compute $\text{softmax}(x_1, \ldots, x_n)$ by processing blocks:

For tile $i$ of the logit vector (size $B_c$ each):

$$m_i = \max(m_{i-1},\ \text{rowmax}(S_i))$$ $$l_i = e^{m_{i-1} - m_i} l_{i-1} + \text{rowsum}(e^{S_i - m_i})$$ $$O_i = \text{diag}(e^{m_{i-1} - m_i}) O_{i-1} + e^{S_i - m_i} V_i$$

At the end: $O = \text{diag}(l_n)^{-1} O_n$

This is correct because: $$\text{softmax}(S)V = \frac{\sum_j e^{S_j} V_j}{\sum_j e^{S_j}} = \frac{e^{-m} \sum_j e^{S_j - m} V_j}{e^{-m} \sum_j e^{S_j - m}}$$

The $e^{-m}$ cancels — we only need the accumulated numerator $O$ and denominator $l$, tracked with running maximum $m$ for numerical stability.

FlashAttention Tiling Strategy

Split $Q$ into row blocks of $B_r$ rows. For each $Q$ block:

• Load $Q_i \in \mathbb{R}^{B_r \times d}$ to SRAM (stays for entire inner loop)
• Iterate over $K, V$ blocks of $B_c$ columns:
  • Load $K_j \in \mathbb{R}^{B_c \times d}$, $V_j \in \mathbb{R}^{B_c \times d}$ to SRAM
  • Compute $S_{ij} = Q_i K_j^T / \sqrt{d_k} \in \mathbb{R}^{B_r \times B_c}$
  • Update running $(m_i, l_i, O_i)$ using the online softmax recurrence
• Write final $O_i$ to HBM

SRAM requirement: $B_r \times d + 2 \times B_c \times d$ for $Q$, $K$, $V$ blocks. For $d=128$ and $B_r = B_c = 64$: $64 \times 128 \times 4 \times 3 = 98$ KB — fits in A100 SRAM (192 KB per SM).

HBM reads/writes: $O$ and each $K, V$ block read once = $O(nd)$ — linear in sequence length.

FlashAttention in Triton — Key Implementation Details

@triton.jit
def flash_attn_fwd_kernel(
    Q_ptr, K_ptr, V_ptr, Out_ptr,
    L_ptr,          # logsumexp for backward: L[i] = m[i] + log(l[i])
    stride_qm, stride_qk,
    stride_km, stride_kk,
    stride_vm, stride_vk,
    stride_om, stride_ok,
    n_heads, seq_len, head_dim,
    scale: tl.constexpr,
    BLOCK_M: tl.constexpr,   # B_r: rows of Q per block
    BLOCK_N: tl.constexpr,   # B_c: columns of K/V per block
    HEAD_DIM: tl.constexpr,
):
    # Which Q block this program handles
    start_m = tl.program_id(0)
    off_h = tl.program_id(1)    # head index
    off_b = tl.program_id(2)    # batch index

    # Initialize accumulators
    m_i = tl.full([BLOCK_M], float('-inf'), dtype=tl.float32)
    l_i = tl.full([BLOCK_M], 0.0, dtype=tl.float32)
    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)

    # Load Q block (stays in SRAM for entire inner loop)
    q_offsets = ...
    q = tl.load(Q_ptr + q_offsets)

    # Iterate over K/V blocks
    for start_n in range(0, seq_len, BLOCK_N):
        k = tl.load(K_ptr + k_offsets)
        v = tl.load(V_ptr + v_offsets)
        
        # Compute attention scores
        s = tl.dot(q, tl.trans(k)) * scale  # [BLOCK_M, BLOCK_N]

        # Online softmax update
        m_new = tl.maximum(m_i, tl.max(s, axis=1))
        alpha = tl.exp(m_i - m_new)   # rescale factor for old accumulator
        p = tl.exp(s - m_new[:, None]) # softmax numerators for this block
        
        # Update accumulator
        acc = acc * alpha[:, None] + tl.dot(p.to(tl.float16), v)
        
        # Update running stats
        l_i = l_i * alpha + tl.sum(p, axis=1)
        m_i = m_new
    
    # Finalize
    acc = acc / l_i[:, None]
    
    # Store logsumexp for backward
    L = m_i + tl.log(l_i)
    tl.store(L_ptr + l_offsets, L)
    tl.store(Out_ptr + out_offsets, acc.to(tl.float16))

Section 2: Speculative Decoding — Mathematical Proof of Losslessness

The Setup

Let $p(x)$ = target model distribution over next token
Let $q(x)$ = draft model distribution over next token

The draft model proposes tokens $x_1, \ldots, x_k$. The target model evaluates all $k+1$ positions in one forward pass (getting $p(x | \text{prefix})$ for each position).

The Acceptance Criterion

For token $x_i$ proposed by the draft:

• Accept with probability $\min\left(1, \frac{p(x_i)}{q(x_i)}\right)$
• If rejected: sample from corrected distribution $p'(x) = \text{norm}(\max(0, p(x) - q(x)))$

Theorem: The output distribution equals $p(x)$ (the target model's distribution).

Proof: Let $A$ = event that token is accepted.

$$P(\text{output} = x) = P(A) \cdot q(x) \cdot \frac{p(x)}{q(x)} + P(\text{reject}) \cdot p'(x)$$

where $P(A) = \sum_x \min(q(x), p(x))$ and $P(\text{reject}) = 1 - P(A) = \sum_x \max(0, q(x) - p(x))$.

The correction term $p'(x) = \frac{\max(0, p(x) - q(x))}{\sum_x \max(0, p(x) - q(x))}$ ensures:

$$P(\text{output} = x) = \min(p(x), q(x)) + \frac{\max(0, p(x) - q(x))}{\sum_x \max(0, p(x)-q(x))} \cdot \sum_x \max(0, q(x) - p(x))$$

With $\sum_x \max(0, q(x)-p(x)) = \sum_x \max(0, p(x)-q(x))$ (since both $p$ and $q$ sum to 1):

$$P(\text{output} = x) = \min(p(x), q(x)) + \max(0, p(x) - q(x)) = p(x) \quad \square$$

Expected Speedup Formula

Let $\alpha$ = acceptance rate per token position (assumed constant for simplicity).

With $k$ draft tokens, the expected number of tokens accepted before first rejection: $$E[\text{tokens accepted}] = \sum_{i=1}^{k} i \cdot \alpha^{i-1}(1-\alpha) + k \cdot \alpha^k = \frac{1 - \alpha^{k+1}}{1 - \alpha}$$

Each target model forward pass verifies $k$ tokens at once (same cost as generating 1 token in standard decoding). Speedup:

$$\text{Speedup} = \frac{E[\text{tokens per target call}]}{1} = \frac{1 - \alpha^{k+1}}{1 - \alpha}$$

For $\alpha = 0.7$, $k = 5$: speedup $= (1 - 0.7^6)/(1 - 0.7) = (1 - 0.118)/0.3 = 2.94\times$

Section 3: Continuous Batching & Paged Attention

Why Static Batching Wastes GPU

Static batching: pack $B$ requests into one batch, run until the longest request finishes. If requests have lengths [10, 100, 1000, 2000], then the GPU is processing all 4 until the 2000-token request finishes — 99.5% of GPU cycles are spent waiting for the longest request.

Continuous batching (iteration-level scheduling): at each decoding step, re-schedule. When any request finishes, immediately insert a new waiting request into the batch. GPU utilization approaches 100%.

Paged Attention — KV Cache Without Waste

Standard KV cache: pre-allocate max_seq_len × d_model × 2 × num_layers memory per request. For max_len=4096 and LLaMA-7B: 4096 × 128 × 32 × 2 × 32 × 2 bytes = 1 GB per request. With 8 concurrent requests: 8 GB — used even if all requests are only 100 tokens.

Paged attention: divide the KV cache into fixed-size pages (e.g., 16 tokens per page). Each request gets a page table mapping logical sequence positions to physical pages in the cache pool. Pages are allocated on demand, freed when the request completes.

KV Cache Pool: 1000 pages × 16 tokens × 128 dim × 32 layers
                                        ↑
Page Table for Request A:   [page 7, page 23, page 156, ...]
Page Table for Request B:   [page 1, page 44, ...]

During attention computation, the paged attention kernel uses the page table as an indirection:

# Instead of: K_cache[seq_idx] 
# Paged:       K_cache[block_table[seq_idx // block_size], seq_idx % block_size]

Section 4: Interview Pitfalls

Section 5: Resources

1. FlashAttention paper — Tri Dao et al., "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" (NeurIPS 2022)
2. FlashAttention-2 — Tri Dao, "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning" (ICLR 2024)
3. FlashAttention-3 — Jay Shah et al., "FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision" (2024)
4. Speculative decoding — Leviathan et al., "Fast Inference from Transformers via Speculative Decoding" (ICML 2023)
5. vLLM / PagedAttention — Kwon et al., "Efficient Memory Management for Large Language Model Serving with PagedAttention" (SOSP 2023)
6. Continuous batching — Orca blog post — https://www.usenix.org/conference/osdi22/presentation/yu
7. Triton documentation — https://triton-lang.org/main/index.html
8. Dao-AILab/flash-attention GitHub — read the Triton reference implementation