← TinyGPT · docs · devlog · roadmap · speedup
source: docs/fa2_forward_notes.md · view on GitHub ↗

Flash Attention 2 — forward pass notes

Implementation log for task #47 (FA2 in WGSL), forward-only. Backward will be a separate session — see “What’s still needed” at the bottom.

Files in this change:

The algorithm

Standard causal scaled-dot-product attention is

S = Q Kᵀ / √d                      // [T, T]
S = causal_mask(S)                 // S[i, j] = -inf for j > i
P = softmax(S, axis=-1)            // [T, T]
O = P V                            // [T, d]

The existing attn_fused_sv in train.wgsl collapses softmax and the P·V matmul into a single pass but still allocates the full row of scores (array<f32, 1024>) in per-invocation private memory. At ctx=1024 that’s 4 KB per thread, which Apple WebGPU spills to global memory.

FA2 keeps Q in shared memory, walks K and V in blocks, and maintains the softmax state online — never materialising a full row of scores.

Notation

Loop (per (b, h, Q-tile))

  1. Load Qtile[Br][hd] into workgroup-shared memory, cooperatively.
  2. Each thread lane ∈ [0, Br) initialises m_i = −∞, l_i = 0, O_i = 0 for its Q row.
  3. For each K block kj:
    • Causal skip. If kj·Bc > q_end_of_tile, every score in this and all later blocks is −∞; break.
    • Cooperatively load Ktile[Bc][hd] and Vtile[Bc][hd].
    • Each thread, against its Q row:
      • Compute S[jj] = Q · Ktile[jj] · scale for jj ∈ [0, Bc), writing −∞ whenever the position is out-of-range or strictly above the causal diagonal.
      • m_block = max_j S[j], m_new = max(m_i, m_block).
      • α = exp(m_i − m_new); rescale O_i *= α, l_i *= α.
      • For each valid jj: p = exp(S[jj] − m_new); l_i += p; O_i += p · Vtile[jj].
      • Update m_i = m_new.
  4. After the loop, O_i /= l_i and write to ctx[b, q_row].
  5. Compatibility second pass. Re-walk K blocks; for each valid (q_row, t2) compute the score and write attn[b, h, q_row, t2] = exp(S − m_final) / l_final so the existing backward kernels still find the attention matrix where they expect it.

The math of the online merge is the standard FA1/FA2 identity:

softmax(concat(A, B)) ↔ softmax-merge of the per-block (m_A, l_A) and (m_B, l_B).

I’m doing FA2’s loop order (outer = Q-tile, inner = K-block) rather than FA1’s (outer = K-block, inner = Q-tile). The FA2 order means each thread only writes its O_i to global memory once, at the very end — no atomic accumulation across workgroups.

Why this is better than attn_fused_sv

Two real wins, one that’s a wash for now.

  1. Private-memory pressure goes from O(T) to O(hd). attn_fused_sv declares var sc: array<f32, 1024> per invocation; FA2 only holds m_i, l_i, and O_i[hd] per thread (so 1 + 1 + hd floats, ≤ 66 floats at hd=64 vs. 1024 in the old kernel). On Apple drivers, the old version provably spills; FA2 fits in registers.

  2. K and V are read once per Q tile, not once per Q row. attn_fused_sv has one thread per Q row and each thread re-reads every K and V row independently. FA2 has 16 threads per Q tile cooperatively loading K/V into shared memory — that’s a 16× reduction in global K/V reads. This is where the FA2 paper’s “I/O complexity” win comes from.

  3. What does NOT improve until ctx ≥ 256: for tiny ctx (Tiny preset has ctx=32, Trial has ctx=64), the existing kernel already fits comfortably and the dispatch overhead of a per-tile workgroup dominates. Expected crossover is around ctx=128–256, with the gap widening at Mega (ctx=512) and the projected Behemoth path (ctx ≥ 1024).

Per AGENTS.md “Safety rules for heavy GPU / compile loops” I did not run a benchmark sweep to measure the crossover. That’s a separate single-shot measurement session once the kernel is integrated.

Parity verification

tests/test_fa2_parity.mjs checks the algorithm (not the WGSL — see caveats) against the naive softmax(causal(QKᵀ/√d)) V reference, in plain JS. Six shapes, including:

Result: all twelve checks pass, with max |ctx − ctx_ref| ≈ 5.96e-8 (≈ 1 f32 ULP — comes from the order in which exp(·) · v[d] terms are summed) and max |attn − attn_ref| = 0 (the second-pass attn writeback uses the same summation order as the naive reference, so it’s bit-exact).

$ node tests/test_fa2_parity.mjs
ok   fa2 ctx   [B=1 T=16 C=32 H=4 hd=8]   maxAbsDiff=5.96e-8
...
ok   fa2 attn  [B=1 T=256 C=64 H=2 hd=32] maxAbsDiff=0.00e+0
ALL PASS

Why not a live WGSL parity dispatch

Node doesn’t ship WebGPU; every existing WGSL parity test in this repo runs through a headed Chromium + Playwright (browser/src/webgpu-test.ts driven by browser/webgpu_test.mjs). Spinning that harness up just to add one assertion would also require touching ops.ts to expose the new kernel — and per the partition for this turn, both of those files are off-limits (another agent is working in parallel).

So the live WGSL parity assertion will go into webgpu-test.ts in the integration session, alongside the ops.ts wiring. The block to add mirrors the existing // --- attention (stage 3) --- section:

// Add to browser/src/webgpu-test.ts after the existing attention check.
{
  const B = 2, T = 48, C = 24, H = 3;     // exercise multi-block walk
  const q = rand(B*T*C), k = rand(B*T*C), v = rand(B*T*C);
  const qt = GpuTensor.fromData(dev, q);
  const kt = GpuTensor.fromData(dev, k);
  const vt = GpuTensor.fromData(dev, v);
  const fa2 = ops.attentionForwardFA2(qt, kt, vt, B, T, C, H);   // new method
  const ref = ops.attentionForward(qt, kt, vt, B, T, C, H);      // existing FA1-style
  check("fa2 forward ctx",  maxError(await fa2.ctx.download(),  await ref.ctx.download())  < 1e-4, "");
  check("fa2 forward attn", maxError(await fa2.attn.download(), await ref.attn.download()) < 1e-4, "");
}

Manual smoke test in the meantime: open the page in Chrome with WebGPU enabled, hit the existing “WebGPU kernel parity tests” route (/webgpu-test.html) after wiring, and confirm ALL PASS. The end-to-end gate stays tests/test_webgpu_train.mjs (50 WASM steps vs 50 WebGPU steps with the same seed, <5% loss drift).

Caveats / known issues

What’s still needed for full FA2

This delivery is forward only. To finish FA2 the next session needs:

  1. Backward kernel with recomputation. The standard FA2 backward recomputes attention on the fly from Q, K, V and the saved log-sum-exp L = m_final + log(l_final). We’d write L instead of the full attn matrix in the forward, then the backward walks K/V blocks twice: once to recompute P for the dV contribution, once to compute dS and accumulate dQ / dK. The existing attn_dscores, attn_dq, attn_dk, attn_dv kernels can be retired.
  2. Drop the second attn writeback pass. Once (1) lands, the FA2 forward stops writing attn[B,H,T,T] entirely — saves O(B·H·T²) of global memory, which is the other half of the FA2 memory win.
  3. Integration into ops.ts. Add attentionForwardFA2 (and once backward lands, switch attentionForward itself to it). Dispatch: workgroup count x = ceil(T / 16), y = B · H.
  4. End-to-end gate. Re-run tests/test_webgpu_train.mjs and the webgpu-test.ts overfit gate (150 training steps on a tiny model driving loss from ~ln(256) to <0.5) after integration.
  5. Single-shot benchmark. One Mega-preset (ctx=512) step on FA1 vs FA2; the win should be visible in step time. Document it like the vec4 writeup in browser/devlog.html. Do not loop — AGENTS.md “Safety rules” applies.

If the next agent gets a NaN out of the backward, the first thing to check is the L = m + log(l) reconstruction at extremely-negative m: when an entire row was masked (shouldn’t happen for causal attention with q_row < T, but the boundary tile has q_row ≥ T lanes that produce l = 0), L = -∞ and downstream exp(s − L) blows up. The forward kernel already returns zero ctx for those lanes; the backward needs to match.