Phase 3 — Attention & Fusion Kernels

Duration target: ~10 weeks (Months 4–6, overlapping the last month of Phase 2). Prerequisites: Phase 2 (Triton fluency, GPU hardware model, Nsight discipline). Do not start Phase 3 until you can write a fused softmax in Triton from memory.

Why attention dominates this phase

Look at any real LLM inference profile. On a Llama 3 70B decode step on H100:

• Attention (Q/K/V matmul + attention math + output projection) is ~ 30–50% of latency

• MLP (up + down projections) is ~ 30–40%

• Norm/activation/quant/comm is the rest

But zoom into attention itself and it is architecturally special:

1. Naïve attention is HBM-bound. The N×N score matrix must be materialized in HBM if you compute softmax the textbook way. For long contexts this becomes catastrophic.

2. FlashAttention removes that HBM traffic by tiling + online softmax. Memory traffic goes from O(N²d) to O(N²d²/M) where M is SMEM size.

3. Attention is the workload where all your Phase 2 skills converge: memory hierarchy, tiling, tensor cores, async pipelines, numerics discipline, fusion.

4. The frontier changes yearly. FA1 (2022) → FA2 (2023) → FA3 (2024, Hopper) → FA4 (March 2026, Blackwell). Each is a masterclass in the current generation’s primitives.

If you can hand-write FlashAttention 2 in Triton, verify its numerics, and profile it, you can hold any inference-engineering conversation. This is Phase 3’s target.

Mental model — the three lenses

1. Online softmax is the mathematical trick that makes tiled attention possible. It maintains a running max and running sum and rescales when a new tile arrives. Understand this on paper first.

2. The block-sparse view. All modern attention (causal, sliding-window, paged, radix) is just a mask pattern on which (Q_tile, KV_tile) pairs to compute. Sparsity/pattern awareness is the difference between mediocre and SOTA.

3. Prefill vs decode are different kernels.

   • Prefill: many queries × many keys. Compute-bound. FlashAttention wins.

   • Decode: 1 (or few) queries × many keys per batch element. Memory-bound. FlashDecoding + paged attention win. Everything in FlashInfer flows from this distinction.

Exit criteria (do not leave this phase until all pass)

• Triton FA2 forward you wrote matches torch.nn.functional.scaled_dot_product_attention within numerical tolerance (bf16: atol=1e-2, rtol=1e-2) across at least 5 shape regimes: (B, H, N, D) = (1,8,512,64), (1,8,2048,64), (4,32,4096,128), (1,8,8192,128), (2,16,16384,64).

• Whiteboard derivation: you can derive the online softmax recurrence on paper without hints, and explain why FlashAttention’s memory traffic is O(N²d²/M) rather than O(N²d).

• Fusion literacy: you have implemented and benchmarked fused RMSNorm+residual, fused SwiGLU, and one of (dequant+GEMM) or (rotary+attention-input).

• CUDA graph experiment: you have replaced launch-overhead-dominated decode with a CUDA graph capture and measured the tokens/sec improvement.

• Numerics discipline: you know the difference between bf16 and fp8 tolerances, know why accumulate is in fp32, and can name at least 3 numerical failure patterns (NaN in softmax, all-zero output row, subnormal underflow).

• You have opinions: you can articulate the FA1→FA2→FA3→FA4 progression, know what FlashInfer solves that FlashAttention doesn’t, and know why FlashDecoding parallelizes across the KV dimension.

File map

File	Purpose
01_online_softmax.md	The math: derive it, worked example, why it’s the heart of FA
02_flashattention_lineage.md	FA1/2/3/4 with arxiv IDs and per-generation innovations
03_flashdecoding_flashinfer.md	Decode-specific attention, paged/ragged KV, FlashInfer API
04_writing_fa_in_triton.md	Step-by-step build of FA2 forward in Triton, bug catalog, numerics tests
05_fusion_thinking.md	Fused RMSNorm+residual, SwiGLU, dequant+GEMM, rotary+attn-input, torch.compile, CUDA graphs
06_numerics_discipline.md	bf16/fp8 tolerances, fp32 accumulate rule, failure smells
07_projects.md	Concrete Phase 3 projects with acceptance criteria

Reading order (recommended sequence)

1. 01_online_softmax.md — do the derivation on paper first.

2. Read the FA1 paper end-to-end (arXiv 2205.14135) while consulting 02_flashattention_lineage.md.

3. 04_writing_fa_in_triton.md — start the implementation.

4. FA2 paper (arXiv 2307.08691) once your FA1-ish Triton kernel works.

5. 03_flashdecoding_flashinfer.md — pivot to decode.

6. 05_fusion_thinking.md + 06_numerics_discipline.md — sharpen the toolkit.

7. 07_projects.md — ship the artifacts.

Non-goals for Phase 3

• You do NOT need to write a competitive FA3 in CUTLASS. Reading + understanding is enough.

• You do NOT need to implement FA4. Modal’s reverse-engineering blog is enough.

• You do NOT need to write your own paged attention from scratch — using vLLM’s PagedAttention and understanding its layout is enough.

• You do NOT need FP8 attention working. Understanding the numerics + reading FA3 §3 is enough.

Time-box everything. Frontier kernels take teams of PhDs quarters of work. Your job is to be able to read them fluently, not out-compete them alone in month 4.