js-flash-attention-tiled-online-softmax-exactness
Done. I diagnosed the problem class, implemented the fix, verified it with Node, and wrote the self-contained write-up.
Deliverables
- ~/SOLUTION.md — the full markdown solution (title, root-cause analysis, exact code, verification).
- ~/flash_test.js — the runnable, verified implementation + test suite.
Root causes addressed
1. Broken online softmax — the denominator l and numerator acc must be rescaled by exp(m - mNew) whenever a block raises the running row max; the fix applies that recurrence exactly.
2. NaN on fully masked rows — exp(-Inf - -Inf) = NaN. Guarded by skipping blocks where mNew === -Infinity and emitting exact zeros when l === 0.
3. O(N²) memory — the kernel only allocates one blockSize score buffer plus a d-wide numerator, reused across rows/blocks; a tracker records the largest allocation.
Verification (all pass)
PASS exactness (span >= 1e6) span=3.555e+7
PASS max relative error <= 1e-9 err=0.000e+0
PASS peak scratch O(blockSize) peak=64 B=64 d=32
PASS peak independent of sequence length peak(N=4096)=32 peak(N=65536)=32
PASS fully masked row -> finite zeros (flash)
PASS fully masked row -> finite zeros (reference)
PASS masked-row exactness err=8.171e-16
PASS hand-computed case got=[6.69761549326657, 6.604769013466862]
Peak scratch stays max(blockSize, d) and is identical for N=4096 and N=65536, and fully masked rows return exactly 0 (finite) rather than NaN.
# Evidence - Problem class: js-flash-attention-tiled-online-softmax-exactness - Model: openrouter/deepseek/deepseek-v4.1-flash - Solved: 2026-09-23T10:09:57.143Z - Verification: solution produced by pi in sandbox; see signatures.json
{"description": "Implement a tiled flash-attention forward pass in plain JavaScript: iterate K/V in blocks, maintain the running row max m and running denominator l with online-softmax rescaling, and never materialize the full N-by-N score matrix. Prove exactness against a naive full-softmax reference (max relative error 1e-9) on random inputs that include rows whose entries span 1e6 in magnitude, and assert peak working memory stays O(block size) by tracking the largest array allocated. Also test that fully masked rows (all -Infinity scores) yield finite zeros instead of NaN.", "environment": "node20", "language": "js", "model": "openrouter/deepseek/deepseek-v4.1-flash", "problem_class": "js-flash-attention-tiled-online-softmax-exactness", "provider": "openrouter", "solved_at": "2026-09-23T10:09:57.144Z", "version": "20"}Done. I diagnosed the problem class, implemented the fix, verified it with Node, and wrote the self-contained write-up.
Deliverables
- ~/SOLUTION.md — the full markdown solution (title, root-cause analysis, exact code, verification).
- ~/flash_test.js — the runnable, verified implementation + test suite.
Root causes addressed
1. Broken online softmax — the denominator l and numerator acc must be rescaled by exp(m - mNew) whenever a block raises the running row max; the fix applies that recurrence exactly.
2. NaN on fully masked rows — exp(-Inf - -Inf) = NaN. Guarded by skipping blocks where mNew === -Infinity and emitting exact zeros when l === 0.
3. O(N²) memory — the kernel only allocates one blockSize score buffer plus a d-wide numerator, reused across rows/blocks; a tracker records the largest allocation.
Verification (all pass)
PASS exactness (span >= 1e6) span=3.555e+7
PASS max relative error <= 1e-9 err=0.000e+0
PASS peak scratch O(blockSize) peak=64 B=64 d=32
PASS peak independent of sequence length peak(N=4096)=32 peak(N=65536)=32
PASS fully masked row -> finite zeros (flash)
PASS fully masked row -> finite zeros (reference)
PASS masked-row exactness err=8.171e-16
PASS hand-computed case got=[6.69761549326657, 6.604769013466862]
Peak scratch stays max(blockSize, d) and is identical for N=4096 and N=65536, and fully masked rows return exactly 0 (finite) rather than NaN.
# Evidence - Problem class: js-flash-attention-tiled-online-softmax-exactness - Model: openrouter/deepseek/deepseek-v4.1-flash - Solved: 2026-09-23T10:09:57.143Z - Verification: solution produced by pi in sandbox; see signatures.json
{"description": "Implement a tiled flash-attention forward pass in plain JavaScript: iterate K/V in blocks, maintain the running row max m and running denominator l with online-softmax rescaling, and never materialize the full N-by-N score matrix. Prove exactness against a naive full-softmax reference (max relative error 1e-9) on random inputs that include rows whose entries span 1e6 in magnitude, and assert peak working memory stays O(block size) by tracking the largest array allocated. Also test that fully masked rows (all -Infinity scores) yield finite zeros instead of NaN.", "environment": "node20", "language": "js", "model": "openrouter/deepseek/deepseek-v4.1-flash", "problem_class": "js-flash-attention-tiled-online-softmax-exactness", "provider": "openrouter", "solved_at": "2026-09-23T10:09:57.144Z", "version": "20"}