nx_flash_attention.nx source
↩ module page · 273 lines · 10355 B
1// nx_flash_attention.nx -- tiled IO-aware attention (Dao 2022).
2//
3// Realises the F-001 row in VRAM_OPTIMIZATION_REALISTIC_TRACKING.md:
4// 2-3 GB working-memory savings at ZERO quality loss.
5//
6// HONEST about what transfers vs what doesn't on the i64 substrate:
7//
8// CUDA FlashAttention wins on:
9// * Memory savings: NEVER materialise full n*n score matrix
10// * Latency: kernel fusion + async memory hides DRAM latency
11// * Numerical stability: online softmax via running max
12//
13// On i64 substrate without SIMD/GPU:
14// * Memory savings: SAME -- algorithmic, transfers fully.
15// For n=1024 / head_dim=64: saves ~8 MB per head per layer.
16// Across 16 heads * 40 layers = ~5 GB of working memory.
17// * Latency: substrate has no kernel fusion (every op is a NishiLang
18// function call); per-block compute is *slower* than naive
19// per-row compute on small inputs. Win flips back POSITIVE for
20// n > ~1024 because cache locality dominates.
21// * Numerical stability: SAME -- online softmax is algorithm, not
22// hardware.
23//
24// Bottom line: ship for the memory win, not the latency win (yet).
25// When SIMD lands the latency story flips to "FlashAttention native
26// always wins."
27//
28// Algorithm (one query row at a time):
29//
30// m = -INF running row-max
31// l = 0 running normaliser sum
32// O = zero vector (head_dim) running output
33//
34// for each KV block [j_start, j_end):
35// S_block[j] = Q[i, :] @ K[j, :] (per j in block)
36// m_block = max(S_block)
37// m_new = max(m, m_block)
38// rescale = exp(m - m_new) <= 1
39// P_block[j] = exp(S_block[j] - m_new)
40// l = l * rescale + sum(P_block)
41// O = O * rescale + P_block @ V[j_start:j_end, :]
42// m = m_new
43//
44// out[i, :] = O / l
45//
46// In Q10 substrate semantics:
47// * rescale = nx_attn_exp_q10(m - m_new) (re-uses existing primitive)
48// * P_block values are Q10 themselves
49// * O accumulates Q10 * Q10 / Q10 products
50// * Final divide by l normalises back to Q10 probabilities
51//
52// genealogy_id: dao_2022_flashattention + dao_2023_flashattention_2 +
53// rabe_staats_2022_online_softmax + milakov_gimelshein_2018
54// lineage_id: substrate_flash_attention_v1
55
56// nx_safety_envelope:
57// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
58// sil_target: SIL1
59// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
60// verdict: NOT_YET_EVALUATED
61
62import "nx_syscalls.nx"
63import "nx_tier.nx"
64import "nx_tensor.nx"
65import "nx_attention.nx"
66const NX_MAGIC_10000: i64 = 10000
67
68const NX_FA_Q10: nx_int = 1024
69
70// ===== Sealed-enum: FlashAttnVerdict ==============================
71
72const NX_FA_OK: nx_int = 0
73const NX_FA_ERR_BAD_DTYPE: nx_int = 1
74const NX_FA_ERR_BAD_NDIM: nx_int = 2
75const NX_FA_ERR_SHAPE_MISMATCH: nx_int = 3
76const NX_FA_ERR_NOT_CONTIGUOUS: nx_int = 4
77const NX_FA_ERR_BAD_BLOCK_SIZE: nx_int = 5
78const NX_FA_N_VERDICTS: nx_int = 6
79
80func nx_fa_verdict_is_valid(v: nx_int) -> nx_int {
81 if v < 0 { return 0 }
82 if v >= NX_FA_N_VERDICTS { return 0 }
83 return 1
84}
85
86// ===== One-row online attention ===================================
87//
88// Process query row i against all KV positions in tiled blocks.
89// Outputs out_row[head_dim] (caller-supplied buffer).
90//
91// Substrate uses nx_attn_exp_q10 (already in nx_attention.nx) for
92// the rescale + per-element exp.
93
94func _fa_row(q: *NxTensor, k: *NxTensor, v: *NxTensor,
95 i: nx_int, scale_q10: nx_int, block_size: nx_int,
96 out_row: *i64) -> nx_int {
97 let d: nx_int = q.shape[1]
98 let n_kv: nx_int = k.shape[0]
99
100 let pq: *i64 = q.storage as *i64
101 let pk: *i64 = k.storage as *i64
102 let pv: *i64 = v.storage as *i64
103
104 // Use a sentinel "minus huge" for initial m; in Q10 with our exp
105 // lookup that clamps to 0 (exp(-large) ~ 0); choose -10000 as
106 // safely below all expected scores after scaling.
107 let M_NEG_INF: nx_int = 0 - NX_MAGIC_10000
108 var m_run: nx_int = M_NEG_INF
109 var l_run: nx_int = 0 // Q10 running sum
110
111 // Zero output row
112 var d_init: nx_int = 0
113 while d_init < d {
114 out_row[d_init] = 0
115 d_init = d_init + 1
116 }
117
118 // Per-block scratch
119 let s_block: *i64 = (sys_mmap(block_size * 8)) as *i64
120 let p_block: *i64 = (sys_mmap(block_size * 8)) as *i64
121
122 // Hoist the exp LUT: built ONCE per row, reused across every kv block's exp
123 // evaluations below (was one sys_mmap per exp call). Output byte-identical.
124 let fa_lut: *i64 = (sys_mmap(16 * 8)) as *i64
125 _attn_exp_lut_fill(fa_lut)
126
127 var j_start: nx_int = 0
128 while j_start < n_kv {
129 var j_end: nx_int = j_start + block_size
130 if j_end > n_kv { j_end = n_kv }
131 let block_n: nx_int = j_end - j_start
132
133 // Step 1: compute S_block[j] = Q[i, :] @ K[j, :] (scaled)
134 var m_block: nx_int = M_NEG_INF
135 var bj: nx_int = 0
136 while bj < block_n {
137 let j_abs: nx_int = j_start + bj
138 var acc: nx_int = 0
139 var t: nx_int = 0
140 while t < d {
141 acc = acc + pq[i * d + t] * pk[j_abs * d + t]
142 t = t + 1
143 }
144 // Apply Q10 scale
145 let scaled: nx_int = (acc * scale_q10) / NX_FA_Q10
146 s_block[bj] = scaled
147 if scaled > m_block { m_block = scaled }
148 bj = bj + 1
149 }
150
151 // Step 2: compute m_new = max(m_run, m_block)
152 var m_new: nx_int = m_run
153 if m_block > m_new { m_new = m_block }
154
155 // Step 3: rescale = exp(m_run - m_new)
156 let rescale: nx_int = _attn_exp_q10_lut(m_run - m_new, fa_lut)
157
158 // Step 4: P_block[j] = exp(S_block[j] - m_new), sum_p
159 var sum_p: nx_int = 0
160 var bp: nx_int = 0
161 while bp < block_n {
162 let e: nx_int = _attn_exp_q10_lut(s_block[bp] - m_new, fa_lut)
163 p_block[bp] = e
164 sum_p = sum_p + e
165 bp = bp + 1
166 }
167
168 // Step 5: l_new = l_run * rescale / Q10 + sum_p
169 l_run = (l_run * rescale) / NX_FA_Q10 + sum_p
170
171 // Step 6: O = O * rescale + P_block @ V[j_start..j_end, :]
172 var dd: nx_int = 0
173 while dd < d {
174 var o_old_scaled: nx_int = (out_row[dd] * rescale) / NX_FA_Q10
175 var p_at_v: nx_int = 0
176 var bv: nx_int = 0
177 while bv < block_n {
178 let j_abs2: nx_int = j_start + bv
179 p_at_v = p_at_v + p_block[bv] * pv[j_abs2 * d + dd]
180 bv = bv + 1
181 }
182 // p_block is Q10; divide
183 out_row[dd] = o_old_scaled + p_at_v / NX_FA_Q10
184 dd = dd + 1
185 }
186
187 m_run = m_new
188 j_start = j_end
189 }
190
191 // Final normalisation: out_row /= l_run (with Q10 multiplier
192 // bake-in -- but for simplicity we leave it as raw weighted sum
193 // and let caller normalise OR we divide here). Standard form:
194 // out = O / l
195 // Since l is in raw sum-of-exp space, just divide:
196 var df: nx_int = 0
197 while df < d {
198 if l_run > 0 {
199 out_row[df] = (out_row[df] * NX_FA_Q10) / l_run
200 }
201 df = df + 1
202 }
203 return NX_FA_OK
204}
205
206// ===== Public: full forward pass ===================================
207//
208// q: [n_q, d], k: [n_kv, d], v: [n_kv, d], out: [n_q, d]
209// scale_q10: Q10 1/sqrt(d_k) factor (caller pre-computes)
210// block_size: KV positions per tile. Typical 32-128. Must divide
211// into n_kv reasonably; we tolerate uneven last block.
212
213func nx_fa_forward(q: *NxTensor, k: *NxTensor, v: *NxTensor,
214 out: *NxTensor, scale_q10: nx_int,
215 block_size: nx_int) -> nx_int {
216 if q.dtype != NX_DT_I64 { return NX_FA_ERR_BAD_DTYPE }
217 if k.dtype != NX_DT_I64 { return NX_FA_ERR_BAD_DTYPE }
218 if v.dtype != NX_DT_I64 { return NX_FA_ERR_BAD_DTYPE }
219 if out.dtype != NX_DT_I64 { return NX_FA_ERR_BAD_DTYPE }
220 if q.ndim != 2 { return NX_FA_ERR_BAD_NDIM }
221 if k.ndim != 2 { return NX_FA_ERR_BAD_NDIM }
222 if v.ndim != 2 { return NX_FA_ERR_BAD_NDIM }
223 if out.ndim != 2 { return NX_FA_ERR_BAD_NDIM }
224 if q.shape[1] != k.shape[1] { return NX_FA_ERR_SHAPE_MISMATCH }
225 if v.shape[0] != k.shape[0] { return NX_FA_ERR_SHAPE_MISMATCH }
226 if v.shape[1] != q.shape[1] { return NX_FA_ERR_SHAPE_MISMATCH }
227 if out.shape[0] != q.shape[0] { return NX_FA_ERR_SHAPE_MISMATCH }
228 if out.shape[1] != q.shape[1] { return NX_FA_ERR_SHAPE_MISMATCH }
229 if nx_t_is_contiguous(q) == 0 { return NX_FA_ERR_NOT_CONTIGUOUS }
230 if nx_t_is_contiguous(k) == 0 { return NX_FA_ERR_NOT_CONTIGUOUS }
231 if nx_t_is_contiguous(v) == 0 { return NX_FA_ERR_NOT_CONTIGUOUS }
232 if nx_t_is_contiguous(out) == 0 { return NX_FA_ERR_NOT_CONTIGUOUS }
233 if block_size <= 0 { return NX_FA_ERR_BAD_BLOCK_SIZE }
234
235 let n_q: nx_int = q.shape[0]
236 let d: nx_int = q.shape[1]
237 let po: *i64 = out.storage as *i64
238
239 var i: nx_int = 0
240 while i < n_q {
241 let row_ptr: *i64 = ((po as nx_int) + i * d * NX_SIZEOF_NX_INT) as *i64
242 let rc: nx_int = _fa_row(q, k, v, i, scale_q10, block_size, row_ptr)
243 if rc != NX_FA_OK { return rc }
244 i = i + 1
245 }
246 return NX_FA_OK
247}
248
249// ===== Memory accounting (audit + dashboard) ======================
250//
251// Working memory at peak:
252// Standard attention: n_q * n_kv * 8 bytes (full S matrix)
253// FlashAttention: n_q * d * 8 bytes (per-row O accumulator)
254// + block_size * 8 bytes (per-block scratch)
255//
256// Returns (standard_bytes, flash_bytes) via out buffer [2 slots].
257
258func nx_fa_working_memory_bytes(n_q: nx_int, n_kv: nx_int, d: nx_int,
259 block_size: nx_int, out: *i64) -> nx_int {
260 out[0] = n_q * n_kv * 8 // standard
261 out[1] = n_q * d * 8 + block_size * 8 * 2 // flash (P + S scratch)
262 return 0
263}
264
265// ===== Memory savings ratio (Q10) =================================
266
267func nx_fa_savings_ratio_q10(n_q: nx_int, n_kv: nx_int, d: nx_int,
268 block_size: nx_int) -> nx_int {
269 let buf: *i64 = (sys_mmap(16)) as *i64
270 nx_fa_working_memory_bytes(n_q, n_kv, d, block_size, buf)
271 if buf[1] <= 0 { return 0 }
272 return (buf[0] * NX_FA_Q10) / buf[1]
273}