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:

webgpu/attention_fa2.wgsl — the new kernel (entry point fa2_forward).
tests/test_fa2_parity.mjs — algorithm-level parity vs the naive reference.

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

Br = Bc = 16 — Q-row tile size and K/V-row block size.
m_i — running per-row max log-score.
l_i — running per-row softmax denominator.
O_i — running per-row partial output (Σ exp(s − m) · v, length hd).
α = exp(m_old − m_new) — rescale factor when the max changes.

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

Load Qtile[Br][hd] into workgroup-shared memory, cooperatively.
Each thread lane ∈ [0, Br) initialises m_i = −∞, l_i = 0, O_i = 0 for its Q row.
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.
After the loop, O_i /= l_i and write to ctx[b, q_row].
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.

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.
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.
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:

T = 16 (one full Q tile, T == Br, T % Bc == 0)
T = 8 (single partial tile)
T = 20 (boundary K block has the mask)
T = 256 (ctx=256, the regime where FA2 starts to pay off)
hd = 64 (Behemoth-shaped head dim, equal to MAX_HD)

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

MAX_HD = 64. Workgroup memory is sized at compile time, so we bound hd at 64. Every current preset is well within that (byte-tinygpt-v0 = 32, small/Behemoth = 64). Bumping this needs both a constant change and a re-verification against device.limits.maxComputeWorkgroupStorageSize — at 64 we sit at 3 × 16 × 64 × 4 = 12 KB of workgroup memory; at 128 we’d hit 24 KB, which is over Apple’s 16 KB default on some drivers.
T does not need to be a multiple of 16. Out-of-range columns get -∞ scores and are dropped by the softmax. Verified by the T=20 case.
The second attn pass costs roughly 1× the score work over again. That’s the price of keeping the existing backward kernels unchanged. Once FA2 backward is shipped (next session), this pass can be deleted and the kernel becomes purely log-sum-exp + ctx out.
No dropout, no relative-position bias. This project’s model uses neither, so neither is implemented.
Subgroup primitives not used. A future pass could use subgroupShuffle to skip the workgroup barrier between K-block loads and score computation, but the basic kernel doesn’t need it. Existing subgroup-using kernels live in train_sg.wgsl — see the pattern there.

What’s still needed for full FA2

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

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.
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.
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.
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.
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.