nx_flash_attention.nx

nx_flash_attention.nx -- tiled IO-aware attention (Dao 2022). Realises the F-001 row in VRAM_OPTIMIZATION_REALISTIC_TRACKING.md: 2-3 GB working-memory savings at ZERO quality loss. HONEST about what transfers vs what doesn't on the i64 substrate: CUDA FlashAttention wins on: * Memory savings: NEVER materialise full n*n score matrix * Latency: kernel fusion + async memory hides DRAM latency * Numerical stability: online softmax via running max On i64 substrate without SIMD/GPU: * Memory savings: SAME -- algorithmic, transfers fully. For n=1024 / head_dim=64: saves ~8 MB per head per layer. Across 16 heads * 40 layers = ~5 GB of working memory. * Latency: substrate has no kernel fusion (every op is a NishiLang function call); per-block compute is *slower* than naive per-row compute on small inputs. Win flips back POSITIVE for n > ~1024 because cache locality dominates. * Numerical stability: SAME -- online softmax is algorithm, not hardware. Bottom line: ship for the memory win, not the latency win (yet). When SIMD lands the latency story flips to "FlashAttention native always wins." Algorithm (one query row at a time): m = -INF running row-max l = 0 running normaliser sum O = zero vector (head_dim) running output for each KV block [j_start, j_end): S_block[j] = Q[i, :] @ K[j, :] (per j in block) m_block = max(S_block) m_new = max(m, m_block) rescale = exp(m - m_new) <= 1 P_block[j] = exp(S_block[j] - m_new) l = l * rescale + sum(P_block)