The Flash Attention 2 path was easier to finish after the algorithm existed in plain JavaScript first.
The forward pass used one workgroup per batch, head, and query tile. It walked keys and values in blocks while keeping the online-softmax state in registers. The backward pass recomputed attention from q, k, and a saved log-sum-exp value instead of reading a cached attention matrix.
At the recorded Mega-class shape, removing that cached-matrix traffic avoided about 67 MB of global memory movement per layer per step. On the Medium preset, the end-to-end comparison moved from 6.8 seconds for WASM SIMD to 0.7 seconds for WebGPU with the new forward and backward paths. Loss drift was 2.5%.
Both halves were checked against a naive Node reference before the WGSL implementation was written. That separated algorithm errors from shader and integration errors. The shader became a translation of known math rather than the first place the math was tested.
The implementation record is at https://posttrainllm.com/devlog.
Top comments (0)