Phase 3 Projects — concrete artifacts with acceptance criteria

Eight projects. Each one produces a code artifact + a written writeup + Nsight reports. Non-negotiable rule: if there’s no benchmark number and no numerics-vs-reference verification, the project isn’t done.

Target hardware: any Ampere+ GPU (H100 ideal, A100/L40S/4090 acceptable). Blackwell projects marked [BW].

---

P3.1 — Online softmax (pure Python + Triton)

Goal: own the recurrence at the bit level.

Steps:

1. Derive online softmax on paper (01_online_softmax.md).

2. Implement pure-Python online softmax; verify against torch.softmax.

3. Implement online-softmax attention (Q@K, softmax, @V) in pure Python; verify against F.scaled_dot_product_attention.

4. Implement online-softmax in Triton (a single row per program) as a warm-up before FA2.

Acceptance criteria:

• Numerics: atol=1e-6 at fp32 for pure-Python, atol=1e-3 at bf16 for Triton.

• Whiteboard-explain the tile recurrence including the V-update, from memory, in under 5 minutes.

Time-box: 2 days.

---

P3.2 — FlashAttention 2 forward in Triton (the flagship project)

Goal: the artifact that makes you an inference engineer.

Steps: follow 04_writing_fa_in_triton.md end-to-end.

Acceptance criteria:

• Numerics: matches F.scaled_dot_product_attention with atol=1e-2, rtol=1e-2 on bfloat16 inputs across these 5 shapes:

  • (1, 8, 512, 64)

  • (1, 8, 2048, 64)

  • (4, 32, 4096, 128)

  • (1, 8, 8192, 128)

  • (2, 16, 16384, 64)

• Both is_causal=True and is_causal=False pass.

• Autotuned. At least 4 configs, picked by triton.autotune.

• Performance: on H100, ≥ 50% of F.scaled_dot_product_attention(...,is_causal=True) throughput on (2, 16, 4096, 64). Stretch: ≥ 80%.

• Nsight report showing: SM utilization ≥ 70%, tensor-core utilization ≥ 50%, DRAM throughput characterized.

• Written README: architecture explanation + benchmarks table + Nsight metrics + known limitations.

Time-box: 1.5 weeks.

Stretch: implement backward using the stored log-sum-exp L.

---

P3.3 — Sliding-window and ALiBi variants

Goal: prove your FA kernel is a pattern you own, not a spell you copied.

Steps:

1. Add WINDOW_SIZE: tl.constexpr and change the causal mask to (offs_m >= offs_n) & (offs_m - offs_n < WINDOW_SIZE).

2. Add ALiBi slope bias: s += alibi_slope[head] * (offs_n - offs_m)[None, :].

3. Verify each against a reference SDPA + mask.

Acceptance criteria:

• Sliding-window matches Mistral-style windowed attention reference within tolerance.

• ALiBi matches BigScience/BLOOM’s ALiBi reference within tolerance.

Time-box: 1 day.

---

P3.4 — FlashDecoding-style decode kernel

Goal: feel the difference between prefill and decode.

Steps:

1. Take your FA2 kernel from P3.2.

2. Modify: at inference time N_Q = 1, and you want to split KV work across many CTAs.

3. Add a split factor G: launch grid (B, H, G) where each program computes attention over N_KV / G tokens and produces partial (m, d, o).

4. Write a second small kernel that merges the G partials with the online-softmax merge rule.

Acceptance criteria:

• Numerics: matches SDPA on decode shapes (1, 32, 1, 128) × (1, 32, 8192, 128).

• Perf: > 3× your P3.2 kernel used naively on decode (which should be badly underutilized).

• Ideally within 2× of flash_attn.flash_attn_with_kvcache (which uses FlashDecoding under the hood).

Time-box: 3 days.

---

P3.5 — Fusion trio: RMSNorm+residual, SwiGLU, Rotary+attn-input

Goal: internalize the fusion pattern.

Steps: implement all three from 05_fusion_thinking.md.

Acceptance criteria:

Kernel	Numerics vs eager	Perf vs eager	Perf vs torch.compile
Fused RMSNorm+residual	atol=1e-3, rtol=1e-3	≥ 1.8×	≥ 1.0×
Fused SwiGLU	atol=1e-3, rtol=1e-3	≥ 1.5×	≥ 1.0×
Rotary fused into attention	atol=5e-3, rtol=5e-3	≥ 1.05× vs unfused RoPE + FA2	—

Each kernel has an Nsight report and a do_bench benchmark output.

Time-box: 1 week.

---

P3.6 — CUDA Graph experiment (decode throughput)

Goal: understand the launch-overhead ceiling in decode.

Steps:

1. Load a small model (e.g., TinyLlama or Llama-3 1B) at bf16.

2. Measure decode tokens/sec at batch=1 with vanilla .forward().

3. Wrap decode with torch.compile(mode="reduce-overhead"). Measure.

4. Manually capture a CUDA graph over one decode step. Measure.

5. Combine torch.compile + explicit CUDA graph. Measure.

Acceptance criteria:

• Report a 4-row table: eager / compile / raw-graph / compile+graph.

• Reduce-overhead should give ≥ 1.5× at batch=1, small model.

• Explain (in the README) why the improvement is what it is: which fraction was launch overhead vs kernel-time.

Time-box: 2 days.

---

P3.7 — FlashInfer integration walkthrough

Goal: know the serving-time attention library.

Steps:

1. pip install flashinfer.

2. Build a BatchDecodeWithPagedKVCacheWrapper on a toy KV cache (2 sequences, 4 pages each).

3. Call plan() then run(). Time each.

4. Do the same with BatchPrefillWithRaggedKVCacheWrapper.

5. Read FlashInfer paper §3 and §4 while your code is running.

Acceptance criteria:

• Working example that runs.

• Written 1-page summary: what does FlashInfer’s JIT actually recompile? When does plan() need to be re-called? What is the workspace tensor for?

Time-box: 1 day.

---

P3.8 — The Phase 3 writeup

Goal: consolidate.

Steps: write a 2–4 page technical writeup covering:

1. Online softmax derivation (annotated with your worked example).

2. FA1 → FA2 → FA3 → FA4 progression, one paragraph each.

3. Your Triton FA2’s architecture: grid, tile shapes, autotune space, numerics, benchmark table, Nsight metrics.

4. Fusion patterns you implemented and their measured impact.

5. CUDA graph experiment results.

6. Open questions you didn’t get to (backward pass, FP8, FlashInfer custom masks).

Acceptance criteria:

• Fits in one markdown doc; readable by another inference engineer without asking questions.

• References cite arxiv IDs and URLs, not vibes.

Time-box: 2 days at end of Phase 3.

---

[BW] Optional Blackwell projects (only if you have B200 access)

P3.9 — Read Modal’s FA4 reverse-engineering blog

Read https://modal.com/blog/reverse-engineer-flash-attention-4 line-by-line. Take notes on: TMEM usage, 2-CTA MMA in backward, software exp, CuTe-DSL patterns.

Acceptance criteria: 1-page notes doc, key primitives named, one “aha” observation written down.

P3.10 — Watch Charles Frye’s FA4 lecture

GPU MODE lecture, Oct 2 2025: https://www.youtube.com/watch?v=ZIEq-WTquy4. Take notes.

---

Phase 3 exit study — self-check questions

Answer these without notes. If you can’t, that specific area is where to spend one more day.

1. Derive the online softmax recurrence including the V update. What is the invariant?

2. Why does FA1’s memory traffic scale as O(N²d²/M) rather than O(N²d)?

3. What are FA2’s three refinements over FA1, and which of them is most valuable at long context small batch?

4. What are FA3’s three innovations, and which Hopper primitives are they exploiting?

5. What are FA4’s five main changes, and which Blackwell primitives are they exploiting?

6. Why does FlashDecoding parallelize across KV dimension? What’s the merge rule?

7. What does FlashInfer’s plan() actually do? Why is it separate from run()?

8. In your Triton FA2 kernel, what’s in fp32 and what’s in bf16, and why?

9. Name three numerics failure smells and their likely causes.

10. When does torch.compile(mode="reduce-overhead") do the work of a hand-written kernel, and when does it fall short?

---

Phase 3 total time budget

Project	Days	Priority
P3.1 Online softmax	2	Critical
P3.2 Triton FA2 forward	10	Critical
P3.3 Sliding-window + ALiBi	1	Recommended
P3.4 FlashDecoding kernel	3	Critical
P3.5 Fusion trio	5	Critical
P3.6 CUDA Graph experiment	2	Critical
P3.7 FlashInfer walkthrough	1	Recommended
P3.8 Writeup	2	Critical
P3.9 FA4 blog notes [BW]	1	Optional
P3.10 FA4 lecture notes [BW]	0.5	Optional

Core budget: ~5–6 weeks. Buffer for the inevitable numerics bug hunt: 1–2 weeks. Total: 10 weeks, matching the phase.