LLM Lab Notes 07: Why FlashAttention Is Fast — An On-Box Comparison

Code file: llm_lab/05_flashattention/flash_compare.py

The series finale: from "how to parallelize" to "how to compute fast." Long context is today's core battleground for LLMs, and attention computation grows quadratically with sequence length — FlashAttention is the key kernel breaking that bottleneck. This note compares standard Attention and flash_attn directly on the same GPU, letting the numbers speak.

1. Experiment Goal

1. Implement standard Attention (QKᵀ → Softmax → AV) and FlashAttention under identical conditions;
2. Measure the latency difference (and observe the memory-behavior difference);
3. Deconstruct why FlashAttention is fast: tiling, IO-awareness, online softmax.

2. Code Walkthrough

Full code (flash_compare.py, 46 lines):

import torch
import time
import torch.nn.functional as F
from flash_attn import flash_attn_func        # ① import flash_attn

batch = 4
seq = 512
heads = 32
dim = 128

q = torch.randn(batch, seq, heads, dim, device="cuda", dtype=torch.bfloat16)   # ② (b, s, h, d) layout directly
k = torch.randn(batch, seq, heads, dim, device="cuda", dtype=torch.bfloat16)
v = torch.randn(batch, seq, heads, dim, device="cuda", dtype=torch.bfloat16)

for step in range(5):                        # ③ 5 rounds: first warms up, rounds 2-5 time

    # ---- Standard Attention ----
    torch.cuda.synchronize()
    t0 = time.time()

    scores = torch.matmul(
        q.transpose(1, 2),                   # (b, h, s, d)
        k.transpose(1, 2).transpose(-1, -2)  # (b, h, d, s)
    ) / (dim ** 0.5)                          # ④ S = QKᵀ / √d — the full [b, h, s, s] score matrix

    attn = F.softmax(scores, dim=-1)

    out = torch.matmul(attn, v.transpose(1, 2))

    torch.cuda.synchronize()
    t1 = time.time()

    # ---- FlashAttention ----
    torch.cuda.synchronize()
    t2 = time.time()

    out2 = flash_attn_func(q, k, v)          # ⑤ one-line call

    torch.cuda.synchronize()
    t3 = time.time()

    if step > 0:                              # ⑥ skip the warmup round
        print(f"Step {step}")
        print(f"Normal Attention : {(t1-t0)*1000:.2f} ms")
        print(f"FlashAttention   : {(t3-t2)*1000:.2f} ms")

① Environment requirements

flash_attn needs a precompiled wheel (or source build) strictly matched to the PyTorch version. This is one of the most common production pitfalls — version mismatch usually manifests as an import error or illegal memory access, not a performance problem.

② Layout difference

Standard code habitually uses (b, h, s, d), while flash_attn_func directly consumes (b, s, h, d) — saving two transpose calls of extra overhead. This small difference reminds us: the kernel interface's layout convention is itself part of performance design.

④ The key: standard Attention materializes the full score matrix

scores has shape [4, 32, 512, 512]; the entire S = QKᵀ matrix is written to HBM in full, then softmax, then multiplied by V. At seq=512 that's 4×32×512×512×2B ≈ 67 MB; at seq=8K it becomes 16 GB class — memory eaten by the score matrix is standard Attention's fatal flaw.

⑤ What's behind the one-line call

flash_attn_func(q, k, v) does three things internally (next section); externally it means: the S matrix is never materialized, and it runs faster.

3. How It Works: FlashAttention's Three Designs

3.1 Tiling

The S matrix is split into tiles; compute tile by tile, softmax tile by tile, accumulate output tile by tile — at any moment only one tile exists, and the full S matrix never exists. Memory drops from O(s²) to O(s) (actually O(s/tile)).

3.2 IO-Awareness: the bottleneck is not FLOPs, it's HBM traffic

GPU compute (FLOPs) vastly outruns memory bandwidth (HBM). Standard Attention writes every intermediate back to HBM and reads it again (QKᵀ → HBM → Softmax → HBM → AV), incurring multiple HBM round-trips. FlashAttention fuses the entire computation into a single kernel, keeping intermediates in on-chip SRAM and writing back only the final output — it saves memory movement, not math.

3.3 Online Softmax

How do sharded softmax results stay consistent with the whole-matrix computation? FlashAttention uses a "rescaling" trick: each time a new block is processed, previously accumulated outputs are rescaled using the running max, yielding results numerically equivalent to standard softmax. That's why it's fast yet not an approximation — results match standard Attention (within floating-point error).

4. Results

On-box comparison over 4 rounds (skipping warmup), expected output:

Step 2
Normal Attention : 31.42 ms
FlashAttention   : 9.78 ms
Step 3
Normal Attention : 31.05 ms
FlashAttention   : 9.61 ms
Step 4
Normal Attention : 30.87 ms
FlashAttention   : 9.70 ms

(Exact values vary by GPU model; typical speedup on H-series GPUs is 3–5×, and the gap widens with seq — ~3× at seq=512, up to 10×+ at seq=4K+.)

Two must-do experiment actions:

1. Change seq from 128 to 512 (commented values are in the code): observe standard Attention's time growing with s² while FlashAttention grows much more slowly;
2. Watch memory (nvidia-smi): standard Attention's peak memory grows with s²; FlashAttention's barely moves with seq — the core of its long-context support.

5. Engineering Takeaways

1. "Fast" fundamentally means saving memory movement: the first principle of AI kernel optimization is not reducing FLOPs but reducing HBM round-trips. FlashAttention, fused kernels, and operator fusion all follow the same logic — applicable to all GPU performance work.

2. Memory budget for long-context configs: when planning 128K-context training/inference, standard Attention's memory formula (b × h × s²) immediately shows it's infeasible; FlashAttention + paged KV cache is the standard. This is also prerequisite knowledge for Chapter 12's inference engine optimization.

3. Combining with Note 05: tensor parallelism distributes heads across GPUs (local_heads); FlashAttention replaces the local attention computation — parallelism and kernel optimization don't conflict; they optimize at different layers: TP solves "doesn't fit," FlashAttention solves "can't compute fast enough."

4. Series wrap-up: from 18 lines to production — process group → DDP → communication autograd → TP (FFN/Attention) → Transformer Block → efficient kernels: this chain is what happens every day on every LLM training cluster. Understand each layer with minimal experiments, and you'll know which layer is broken when troubleshooting at 10K-GPU scale.

Series complete. Suggested next topics: 《Training Stability: Loss Spike Debugging》 or 《Mixed Precision & Checkpointing》, or return to the systematic view in Distributed Parallelism Overview.