# Flash Attention 2 — backward pass notes

Implementation log for the second half of task #47. Forward landed last
session (`webgpu/attention_fa2.wgsl` + `tests/test_fa2_parity.mjs` +
`docs/performance/fa2_forward_notes.md`); this is the runway for backward.

This delivery: **algorithm-level parity verified in Node**
(`tests/test_fa2_backward_parity.mjs`). The WGSL kernel + ops.ts wiring
are left for a follow-up session because of a real binding-budget
problem documented below.

## The math (FA2 backward with recompute)

Standard attention backward starts from `dO = ∂loss/∂ctx` and produces
`dQ, dK, dV`. Naive version uses the cached attention matrix `P`. The
FA2 trick: with the forward's saved `L = m_final + log(l_final)`
(log-sum-exp per Q row), we reconstruct `P = exp(S − L)` from `q/k`
without ever reading the full `[B, H, T, T]` attn matrix.

Per `(batch, head, q_row)`:

```
D[q_row] = ⟨dO[q_row], O[q_row]⟩          // scalar; O = ctx from forward
for t2 in causal range (t2 ≤ q_row):
    S    = q[q_row] · k[t2] · scale
    P    = exp(S − L[q_row])               // recompute
    dP   = ⟨dO[q_row], V[t2]⟩
    dS   = P · (dP − D[q_row])
    dQ[q_row] += dS · k[t2] · scale        // own thread
    dK[t2]    += dS · q[q_row] · scale     // others contribute too
    dV[t2]    += P  · dO[q_row]            // others contribute too
```

`dQ` is owned by `q_row`'s thread — no contention. `dK` and `dV` are
gathered: every `q_row ≥ t2` contributes to the same `k[t2]` and `v[t2]`.
That contention shapes the kernel design (see below).

## Parity verification — what's already done

`tests/test_fa2_backward_parity.mjs` mirrors the planned WGSL kernel
in plain JS. Six shapes including the masked-boundary case (T=20),
the Behemoth-shaped head dim (hd=64), and the regime-crossover case
(T=256). All 18 checks (6 shapes × dQ/dK/dV) pass within 1 ULP of
the naive cached-attn reference:

```
$ node tests/test_fa2_backward_parity.mjs
ok   fa2 dQ vs naive [B=1 T=16 C=32 H=4 hd=8]    maxAbs=2.98e-8
ok   fa2 dK vs naive [B=1 T=16 C=32 H=4 hd=8]    maxAbs=2.98e-8
ok   fa2 dV vs naive [B=1 T=16 C=32 H=4 hd=8]    maxAbs=2.38e-7
...
ok   fa2 dV vs naive [B=1 T=256 C=64 H=2 hd=32]  maxAbs=2.98e-7
ALL PASS
```

`dV` shows higher absolute error because it sums `T` contributions
per output element; 2.98e-7 at T=256 is ~T·ε for f32 — clean.

## Why the WGSL kernel isn't shipped this turn

The natural binding layout for the backward kernel is **7 storage
buffers**:

```
g0 = q       (read)
g1 = k       (read)
g2 = v       (read)
g3 = L       (read)     ← new in forward, saved per (b, h, t1)
g4 = dO      (read)
g5 = ctx     (read)     ← needed to compute D[t1] = ⟨dO, O⟩
g? = dQ/dK/dV  (write)
```

The shared bind layout in `webgpu/ops.ts` exposes exactly six storage
buffers + one uniform. Three workable resolutions, none of which
deserve to be rushed:

1. **Pre-pass kernel** computes `D` once and stores it in a small
   `[B, H, T]` buffer. The main backward kernel then drops `ctx` and
   reads `D` instead — fits in six bindings cleanly. Two kernels total.
2. **Buffer aliasing trick** — overwrite `g4 = dO` with `dq` in place.
   Each thread reads its full `dO[t1, :]` row into private registers
   before the K-block loop, then writes `dq[t1, :]` to the same slot.
   No race because the buffer is per-thread (per-Q-row). One kernel,
   but the aliasing rule (`ops.ts` already comments that WebGPU forbids
   aliasing the same buffer in two writable bindings of one bind group)
   means we need a slightly different bind-group construction path.
3. **Two-kernel split** by output role: `fa2_backward_dq` (one workgroup
   per Q tile, dQ-only — no contention) and `fa2_backward_dkv`
   (one workgroup per K tile, walking all Q rows ≥ k_row, dK + dV).
   This is the FA2 paper's actual implementation choice. Avoids float
   atomics entirely; uses the natural data-flow split (each kernel's
   threads each own one row of their output buffer). Likely the right
   long-term shape.

Picking between (1) and (3) is a 20-minute design decision; doing it
right is another half-day of careful WGSL plus the parity gate
(`tests/test_webgpu_train.mjs` end-to-end, 5% drift bar). Today's
delivery pins the math so whichever path the next session picks, the
algorithm itself is verified.

## What the next session needs to do

1. **Decide on (1) vs (3) above.** My recommendation: (3), because it
   matches the published FA2 backward and parallelises naturally.
2. **Modify forward** to save `L = m + log(l)` into a new `[B, H, T]`
   buffer (small) so backward can read it. Today the kernel saves
   `m_final` and `l_final` to `var<workgroup>` arrays — we just expose
   them to a host-visible storage buffer.
3. **Write the WGSL kernel(s)** following the plan in
   `test_fa2_backward_parity.mjs`. The Node mirror's loops are the
   exact WGSL structure modulo cooperative loading of K/V tiles into
   shared memory.
4. **Live WGSL compile check** via Playwright + a real WebGPU adapter,
   same pattern as `tests/test_fa2_compile.mjs` did for the forward.
5. **Integration in `ops.ts`**: add an `attentionBackwardFA2` method
   that calls the new kernel(s); route `attentionBackward` to it
   conditionally on whether `attentionForward` used the FA2 path
   (today both are conditional on `hd ≤ 64`).
6. **End-to-end gate**: `tests/test_webgpu_train.mjs` — drift < 5%.
7. **Drop the second `attn` pass from the forward kernel** once
   backward is on the recompute path. Saves the O(B·H·T²) global
   memory write that the FA1-style backward currently needs. This is
   the other half of the FA2 memory win.
8. **Bench**: one Mega-preset (ctx=512) step with FA1 forward + FA1
   backward vs FA2 forward + FA2 backward. Document in
   `browser/devlog.html`. **Single shot**, per AGENTS.md "Safety
   rules" — no loops.

## Caveats specific to the backward path

- **The L = -∞ row.** For Q rows where every K position is masked
  (shouldn't happen with causal attention as long as `q_row < T`, but
  defensive code wins): `L = -∞` means `exp(S − L) = exp(+∞)` blows up
  to NaN. The forward already writes zero `ctx` for those lanes; the
  backward needs the same guard (skip rows where `l == 0`).
- **D = ⟨dO, O⟩ as scalar.** Cheap, but needs `O` (= forward ctx)
  available. If we go with split (3), `fa2_backward_dq` can compute
  D from O+dO on the fly; `fa2_backward_dkv` doesn't need D directly
  (it needs `dS` which factors through D, so it might need a small
  D buffer either pre-computed or duplicated in shared memory).
- **No dropout, no relative position bias.** Same scope as the forward.

## Where this lives in the codebase

```
docs/performance/fa2_forward_notes.md          — forward writeup (last session)
docs/performance/fa2_backward_notes.md         — this file
tests/test_fa2_parity.mjs          — forward parity (last session)
tests/test_fa2_backward_parity.mjs — backward parity (this session)
webgpu/attention_fa2.wgsl          — forward kernel (last session)
webgpu/attention_fa2_backward.wgsl — backward kernel (next session)
webgpu/ops.ts                      — integration site (next session)
```