nx_attention.nx source
↩ module page · 352 lines · 13066 B
1// nx_attention.nx -- scaled dot-product attention kernel.
2//
3// THE kernel that activates everything we've shipped on the
4// sovereign-ML stack. Composes:
5//
6// nx_tensor -- typed shape contract
7// nx_blas_i64 -- gemm-shape Q @ K^T
8// nx_kv_cache -- O(n) autoregressive incremental
9// nx_sparse_tensor -- skip zero entries post-softmax
10// nx_quant_block -- 4-bit Q / K / V at storage layer (queued)
11// nx_numeric_oracle -- bit-exact vs reference verification
12//
13// Per Vaswani 2017:
14//
15// Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V
16//
17// On the i64 + Q10 substrate:
18// * Q, K, V are i64 tensors with shape [n_tokens, head_dim]
19// * Scores = Q @ K^T (raw integer dot products)
20// * Scale = 1 / sqrt(d_k) approximated via Q10 lookup or
21// constant the caller pre-computes
22// * Softmax in Q10: per-row, find max, subtract, exp via
23// piecewise approximation, normalise
24// * Output = weighted @ V
25//
26// For v1 we expose:
27// nx_attn_score_matrix Q @ K^T raw (i64)
28// nx_attn_scale_q10 apply Q10 scaling
29// nx_attn_softmax_row_q10 per-row softmax in Q10
30// nx_attn_apply_to_v scores @ V
31// nx_attn_forward full pass (compose the four above)
32//
33// Per the min-hardware-floor cardinal: every step is a separate
34// primitive so the caller can swap in sparse / quant / fused
35// variants at any layer.
36//
37// Softmax-in-Q10 honest accounting:
38// * exp(x) in Q10 uses a 16-entry lookup + linear interp; max
39// error ~3% relative for x in [-3, 0].
40// * Output rows sum to ~1024 (Q10 unity); rounding pads by up
41// to N (where N = sequence length) -- oracle epsilon should
42// accommodate this.
43//
44// genealogy_id: vaswani_2017_attention + flashattention_dao_2022 +
45// xformers_meta_2021 + reformer_kitaev_2020 +
46// linformer_wang_2020
47// lineage_id: substrate_attention_v1
48//
49// nx_safety_envelope:
50// intended_use: "Scaled dot-product attention (Vaswani 2017)
51// -- foundation for substrate transformer
52// primitives and downstream AI inference"
53// sil_target: SIL2 (AI correctness affects downstream
54// decisions; misuse class spans
55// quality-grader to medical)
56// asil_target: QM
57// dal_target: DAL C
58// evidence: [Vaswani_2017_canonical_basis, no_FP_in_index_math,
59// Q-scale_softmax_target,
60// deterministic_seed_path]
61// hazard_register: [bug-tape-attention-collapse-on-temperature-zero,
62// bug-tape-context-length-overflow,
63// bug-tape-positional-encoding-drift]
64// residual_risk: "Substrate primitive; alignment / robustness /
65// adversarial-input handling are upstream
66// responsibilities."
67// verdict: NOT_YET_EVALUATED
68
69import "nx_syscalls.nx"
70import "nx_tier.nx"
71import "nx_tensor.nx"
72import "nx_blas_i64.nx"
73const NX_MAGIC_1024: i64 = 1024
74const NX_MAGIC_7680: i64 = 7680
75
76const NX_ATTN_Q10: nx_int = 1024
77
78// ===== Sealed-enum: AttnVerdict ===================================
79
80const NX_ATTN_OK: nx_int = 0
81const NX_ATTN_ERR_BAD_DTYPE: nx_int = 1
82const NX_ATTN_ERR_BAD_NDIM: nx_int = 2
83const NX_ATTN_ERR_SHAPE_MISMATCH: nx_int = 3
84const NX_ATTN_ERR_NOT_CONTIGUOUS: nx_int = 4
85const NX_ATTN_N_VERDICTS: nx_int = 5
86
87func nx_attn_verdict_is_valid(v: nx_int) -> nx_int {
88 if v < 0 { return 0 }
89 if v >= NX_ATTN_N_VERDICTS { return 0 }
90 return 1
91}
92
93// ===== Step 1: raw score matrix Q @ K^T ===========================
94//
95// Q: [n_tokens, head_dim]
96// K: [n_kv, head_dim]
97// scores: [n_tokens, n_kv]
98//
99// Standard transposed-K matmul. We don't materialise K^T -- we walk
100// K row-by-row to compute scores[i, j] = dot(Q[i, :], K[j, :]).
101
102func nx_attn_score_matrix(q: *NxTensor, k: *NxTensor, scores: *NxTensor) -> nx_int {
103 if q.dtype != NX_DT_I64 { return NX_ATTN_ERR_BAD_DTYPE }
104 if k.dtype != NX_DT_I64 { return NX_ATTN_ERR_BAD_DTYPE }
105 if scores.dtype != NX_DT_I64 { return NX_ATTN_ERR_BAD_DTYPE }
106 if q.ndim != 2 { return NX_ATTN_ERR_BAD_NDIM }
107 if k.ndim != 2 { return NX_ATTN_ERR_BAD_NDIM }
108 if scores.ndim != 2 { return NX_ATTN_ERR_BAD_NDIM }
109 if q.shape[1] != k.shape[1] { return NX_ATTN_ERR_SHAPE_MISMATCH }
110 if scores.shape[0] != q.shape[0] { return NX_ATTN_ERR_SHAPE_MISMATCH }
111 if scores.shape[1] != k.shape[0] { return NX_ATTN_ERR_SHAPE_MISMATCH }
112 if nx_t_is_contiguous(q) == 0 { return NX_ATTN_ERR_NOT_CONTIGUOUS }
113 if nx_t_is_contiguous(k) == 0 { return NX_ATTN_ERR_NOT_CONTIGUOUS }
114 if nx_t_is_contiguous(scores) == 0 { return NX_ATTN_ERR_NOT_CONTIGUOUS }
115
116 let n_q: nx_int = q.shape[0]
117 let n_kv: nx_int = k.shape[0]
118 let d: nx_int = q.shape[1]
119 let pq: *i64 = q.storage as *i64
120 let pk: *i64 = k.storage as *i64
121 let ps: *i64 = scores.storage as *i64
122
123 var i: nx_int = 0
124 while i < n_q {
125 var j: nx_int = 0
126 while j < n_kv {
127 var acc: nx_int = 0
128 var t: nx_int = 0
129 while t < d {
130 acc = acc + pq[i * d + t] * pk[j * d + t]
131 t = t + 1
132 }
133 ps[i * n_kv + j] = acc
134 j = j + 1
135 }
136 i = i + 1
137 }
138 return NX_ATTN_OK
139}
140
141// ===== Step 2: scale by Q10 factor =================================
142//
143// scores[i, j] = (scores[i, j] * scale_q10) / Q10
144// Caller pre-computes scale_q10 = round(Q10 / sqrt(d_k))
145
146func nx_attn_scale_q10(scores: *NxTensor, scale_q10: nx_int) -> nx_int {
147 if scores.dtype != NX_DT_I64 { return NX_ATTN_ERR_BAD_DTYPE }
148 if nx_t_is_contiguous(scores) == 0 { return NX_ATTN_ERR_NOT_CONTIGUOUS }
149 let ps: *i64 = scores.storage as *i64
150 var k: nx_int = 0
151 while k < scores.numel {
152 ps[k] = (ps[k] * scale_q10) / NX_ATTN_Q10
153 k = k + 1
154 }
155 return NX_ATTN_OK
156}
157
158// ===== exp(x) approximation in Q10 ================================
159//
160// For softmax stability we exp(x - max(x)) so inputs are <= 0.
161// We use a 16-entry piecewise-linear lookup for x in [-8, 0] (Q10
162// scale). Entries are pre-computed Q10 values of exp(x/-1).
163//
164// exp(0) = 1024
165// exp(-0.5) = 621
166// exp(-1.0) = 376
167// exp(-1.5) = 228
168// exp(-2.0) = 138
169// exp(-2.5) = 84
170// exp(-3.0) = 51
171// exp(-3.5) = 31
172// exp(-4.0) = 19
173// exp(-4.5) = 11
174// exp(-5.0) = 7
175// exp(-5.5) = 4
176// exp(-6.0) = 3
177// exp(-6.5) = 2
178// exp(-7.0) = 1
179// exp(-7.5) = 1
180//
181// For x < -7.5 we clamp to 0 (negligible under Q10 precision).
182
183// Fill the 16-entry Q10 exp lookup table (a CONSTANT table) into lo_table.
184// Single source of the constants (DRY) for both entry points below.
185func _attn_exp_lut_fill(lo_table: *i64) -> i64 {
186 lo_table[0] = NX_MAGIC_1024; lo_table[1] = 621; lo_table[2] = 376; lo_table[3] = 228
187 lo_table[4] = 138; lo_table[5] = 84; lo_table[6] = 51; lo_table[7] = 31
188 lo_table[8] = 19; lo_table[9] = 11; lo_table[10] = 7; lo_table[11] = 4
189 lo_table[12] = 3; lo_table[13] = 2; lo_table[14] = 1; lo_table[15] = 1
190 return 0
191}
192
193// exp(x) in Q10 using a PRE-BUILT 16-entry table -- the HOT path. A caller that
194// fires this R*C times builds the table ONCE above its loop and passes it in,
195// eliminating the per-element sys_mmap. Interp is byte-identical to the per-call
196// form below (same bin / frac / lo_v / hi_v / return).
197func _attn_exp_q10_lut(x_q10: nx_int, lo_table: *i64) -> nx_int {
198 if x_q10 >= 0 { return NX_ATTN_Q10 }
199 let neg: nx_int = 0 - x_q10
200 if neg >= NX_MAGIC_7680 { return 0 } // > 7.5 in Q10
201 let bin: nx_int = neg / 512
202 let bin_lo: nx_int = bin * 512
203 let frac: nx_int = neg - bin_lo // 0..511
204 var lo_v: nx_int = lo_table[bin]
205 var hi_v: nx_int = 0
206 if bin + 1 < 16 { hi_v = lo_table[bin + 1] }
207 return lo_v + ((hi_v - lo_v) * frac) / 512
208}
209
210// Per-call form (cold / single-shot / test callers): build the table then
211// delegate. Output is byte-identical to the prior hand-inlined version.
212func _attn_exp_q10(x_q10: nx_int) -> nx_int {
213 let lo_table: *i64 = (sys_mmap(16 * 8)) as *i64
214 _attn_exp_lut_fill(lo_table)
215 return _attn_exp_q10_lut(x_q10, lo_table)
216}
217
218// ===== Step 3: per-row softmax in Q10 =============================
219//
220// For each row i of scores:
221// 1. find max_v = max(scores[i, :])
222// 2. shifted[j] = scores[i, j] - max_v (so all <= 0)
223// 3. exp_v[j] = _attn_exp_q10(shifted[j])
224// 4. sum_v = sum(exp_v)
225// 5. out[i, j] = (exp_v[j] * Q10) / sum_v (Q10 probability)
226
227func nx_attn_softmax_row_q10(scores: *NxTensor) -> nx_int {
228 if scores.dtype != NX_DT_I64 { return NX_ATTN_ERR_BAD_DTYPE }
229 if scores.ndim != 2 { return NX_ATTN_ERR_BAD_NDIM }
230 if nx_t_is_contiguous(scores) == 0 { return NX_ATTN_ERR_NOT_CONTIGUOUS }
231
232 let n_rows: nx_int = scores.shape[0]
233 let n_cols: nx_int = scores.shape[1]
234 let ps: *i64 = scores.storage as *i64
235
236 // Hoist the exp LUT: built ONCE here, reused across all n_rows*n_cols exp
237 // evaluations below (was one sys_mmap per element). Output byte-identical.
238 let lut: *i64 = (sys_mmap(16 * 8)) as *i64
239 _attn_exp_lut_fill(lut)
240
241 var i: nx_int = 0
242 while i < n_rows {
243 // max
244 var max_v: nx_int = ps[i * n_cols]
245 var j: nx_int = 1
246 while j < n_cols {
247 if ps[i * n_cols + j] > max_v { max_v = ps[i * n_cols + j] }
248 j = j + 1
249 }
250
251 // shifted exp + sum
252 var sum_v: nx_int = 0
253 var k: nx_int = 0
254 while k < n_cols {
255 let shifted: nx_int = ps[i * n_cols + k] - max_v
256 let e: nx_int = _attn_exp_q10_lut(shifted, lut)
257 ps[i * n_cols + k] = e
258 sum_v = sum_v + e
259 k = k + 1
260 }
261
262 // normalise
263 if sum_v <= 0 { sum_v = 1 }
264 var m: nx_int = 0
265 while m < n_cols {
266 ps[i * n_cols + m] = (ps[i * n_cols + m] * NX_ATTN_Q10) / sum_v
267 m = m + 1
268 }
269 i = i + 1
270 }
271 return NX_ATTN_OK
272}
273
274// ===== Step 4: apply attention weights to V ======================
275//
276// weights: [n_tokens, n_kv] (Q10 probabilities)
277// V: [n_kv, head_dim]
278// out: [n_tokens, head_dim]
279//
280// out[i, d] = sum_j(weights[i, j] * V[j, d]) / Q10
281
282func nx_attn_apply_to_v(weights: *NxTensor, v: *NxTensor, out: *NxTensor) -> nx_int {
283 if weights.dtype != NX_DT_I64 { return NX_ATTN_ERR_BAD_DTYPE }
284 if v.dtype != NX_DT_I64 { return NX_ATTN_ERR_BAD_DTYPE }
285 if out.dtype != NX_DT_I64 { return NX_ATTN_ERR_BAD_DTYPE }
286 if weights.ndim != 2 { return NX_ATTN_ERR_BAD_NDIM }
287 if v.ndim != 2 { return NX_ATTN_ERR_BAD_NDIM }
288 if out.ndim != 2 { return NX_ATTN_ERR_BAD_NDIM }
289 if weights.shape[1] != v.shape[0] { return NX_ATTN_ERR_SHAPE_MISMATCH }
290 if out.shape[0] != weights.shape[0] { return NX_ATTN_ERR_SHAPE_MISMATCH }
291 if out.shape[1] != v.shape[1] { return NX_ATTN_ERR_SHAPE_MISMATCH }
292 if nx_t_is_contiguous(weights) == 0 { return NX_ATTN_ERR_NOT_CONTIGUOUS }
293 if nx_t_is_contiguous(v) == 0 { return NX_ATTN_ERR_NOT_CONTIGUOUS }
294 if nx_t_is_contiguous(out) == 0 { return NX_ATTN_ERR_NOT_CONTIGUOUS }
295
296 let n_tok: nx_int = weights.shape[0]
297 let n_kv: nx_int = weights.shape[1]
298 let d: nx_int = v.shape[1]
299 let pw: *i64 = weights.storage as *i64
300 let pv: *i64 = v.storage as *i64
301 let po: *i64 = out.storage as *i64
302
303 var i: nx_int = 0
304 while i < n_tok {
305 var c: nx_int = 0
306 while c < d {
307 var acc: nx_int = 0
308 var j: nx_int = 0
309 while j < n_kv {
310 acc = acc + pw[i * n_kv + j] * pv[j * d + c]
311 j = j + 1
312 }
313 po[i * d + c] = acc / NX_ATTN_Q10
314 c = c + 1
315 }
316 i = i + 1
317 }
318 return NX_ATTN_OK
319}
320
321// ===== Full forward: scaled dot-product attention =================
322//
323// Caller provides Q, K, V (all [n_tokens, head_dim]); we allocate
324// scratch scores, do the four steps, write output.
325// scale_q10 is the pre-computed Q10 / sqrt(d_k) factor.
326//
327// Note: scratch tensor allocations live as long as the parent mmap
328// arena; v1 doesn't free them, future iteration adds an arena
329// allocator.
330
331func nx_attn_forward(q: *NxTensor, k: *NxTensor, v: *NxTensor,
332 out: *NxTensor, scale_q10: nx_int) -> nx_int {
333 let err: *i64 = (sys_mmap(8)) as *i64
334 let sh: *i64 = (sys_mmap(16)) as *i64
335 sh[0] = q.shape[0]
336 sh[1] = k.shape[0]
337 let scores: *NxTensor = nx_t_alloc(NX_DT_I64, sh, 2, err)
338
339 let rc1: nx_int = nx_attn_score_matrix(q, k, scores)
340 if rc1 != NX_ATTN_OK { return rc1 }
341
342 let rc2: nx_int = nx_attn_scale_q10(scores, scale_q10)
343 if rc2 != NX_ATTN_OK { return rc2 }
344
345 let rc3: nx_int = nx_attn_softmax_row_q10(scores)
346 if rc3 != NX_ATTN_OK { return rc3 }
347
348 let rc4: nx_int = nx_attn_apply_to_v(scores, v, out)
349 if rc4 != NX_ATTN_OK { return rc4 }
350
351 return NX_ATTN_OK
352}