nx_lowrank_attn.nx source

1// nx_lowrank_attn.nx -- M3: low-rank KV through a REAL softmax-attention forward (end-to-end quality proof). 2// 3// module: nishi-core.ai.lowrank_attn capability: VRAM_REDUCTION (per-token low-rank KV, measured end-to-end) 4// 5// M0/M1 measured COVARIANCE fidelity (K^T K) of an FD sketch. But softmax attention needs PER-TOKEN scores 6// Q.K[i], which an FD covariance sketch cannot give. The per-token low-rank compression that DOES work is 7// Eigen-Attention / SVD-projection: find the top-r right singular basis V_r [r x d] of K (top-r eigenvectors 8// of K^T K), store per-token coords[i] = K[i] @ V_r^T (n x r) + the shared basis (r x d), reconstruct 9// K[i] ~= coords[i] @ V_r. This PRESERVES per-token identity -> usable in real attention. The eigen-extraction 10// REUSES the M1 lesson: reorthogonalize (Gram-Schmidt twice) or the dominant direction leaks under f32. 11// 12// PROOF: compare the attention OUTPUT (after softmax + value mixing) full-KV vs compressed-KV. For genuinely 13// rank-2 K/V, compressing to r>=2 reconstructs exactly -> output error ~0 (VRAM cut, quality kept end-to-end). 14// Sovereign: nx_f32_hw (SSE arith) + nx_f32_sqrt + nx_f32_softmax (reused). no gcc. HONEST: reference-scale d=8. 15import "nx_syscalls.nx" 16import "nx_f32_hw.nx" 17import "nx_f32.nx" 18import "nx_f32_softmax.nx" 19const LA_MAGIC_1103515245: i64 = 1103515245 20const LA_MAGIC_12345: i64 = 12345 21const LA_MAGIC_2048: i64 = 2048 22const LA_MAGIC_4096: i64 = 4096 23const LA_MAGIC_99991: i64 = 99991 24const LA_MAGIC_11003: i64 = 11003 25const LA_MAGIC_22013: i64 = 22013 26const LA_MAGIC_33029: i64 = 33029 27const LA_MAGIC_44039: i64 = 44039 28const LA_MAGIC_55049: i64 = 55049 29const LA_MAGIC_60061: i64 = 60061 30const LA_MAGIC_70067: i64 = 70067 31const LA_MAGIC_80071: i64 = 80071 32 33const LA_POWER_ITERS: i64 = 64 34 35// ---- stdout ---- 36func la_strlen(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } return n } 37func g_puts(s: *u8) -> i64 { sys_write(1, s, la_strlen(s)); return 0 } 38func g_putn(v: i64) -> i64 { 39 if v == 0 { sys_write(1, "0" as *u8, 1); return 0 } 40 var m: i64 = v 41 if m < 0 { sys_write(1, "-" as *u8, 1); m = 0 - m } 42 let d: *u8 = sys_mmap(24); var k: i64 = 0 43 while m > 0 { d[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 } 44 var j: i64 = k - 1 45 while j >= 0 { sys_write(1, ((d as i64)+j) as *u8, 1); j = j - 1 } 46 return 0 47} 48 49// ---- f32 bit predicate ---- 50func la_is_zero(x: i64) -> i64 { if (x & 0x7fffffff) == 0 { return 1 } return 0 } 51 52func f32frac(num: i64, den: i64) -> i64 { return f32_div(f32_of(num), f32_of(den)) } 53 54// ---- pseudo-random stream (high bits) ---- 55func la_lcg(st: *i64) -> i64 { st[0] = st[0] * LA_MAGIC_1103515245 + LA_MAGIC_12345; return (st[0] >> 20) & 0xfff } 56 57// genuinely rank-r matrix M[n x d] = C[n x r] @ B[r x d], zero-mean random factors. 58func la_build_rankr(M: *i64, n: i64, d: i64, r: i64, seed: i64) -> i64 { 59 let st: *i64 = sys_mmap(8) as *i64 60 st[0] = seed 61 let C: *i64 = sys_mmap(n * r * 8) as *i64 62 let B: *i64 = sys_mmap(r * d * 8) as *i64 63 var t: i64 = 0 64 while t < n * r { C[t] = f32frac(la_lcg(st) - LA_MAGIC_2048, LA_MAGIC_4096); t = t + 1 } 65 t = 0 66 while t < r * d { B[t] = f32frac(la_lcg(st) - LA_MAGIC_2048, LA_MAGIC_4096); t = t + 1 } 67 var i: i64 = 0 68 while i < n { 69 var c: i64 = 0 70 while c < d { 71 var acc: i64 = 0 72 var j: i64 = 0 73 while j < r { acc = f32_add(acc, f32_mul(C[i * r + j], B[j * d + c])); j = j + 1 } 74 M[i * d + c] = acc 75 c = c + 1 76 } 77 i = i + 1 78 } 79 return 0 80} 81 82// ---- f32 vector ops (hw SSE arith + sw sqrt) ---- 83func la_dot(a: *i64, b: *i64, n: i64) -> i64 { 84 var s: i64 = 0 85 var i: i64 = 0 86 while i < n { s = f32_add(s, f32_mul(a[i], b[i])); i = i + 1 } 87 return s 88} 89 90func la_normalize(v: *i64, n: i64) -> i64 { 91 let ns: i64 = la_dot(v, v, n) 92 if la_is_zero(ns) == 1 { return -1 } 93 let nrm: i64 = nx_f32_sqrt(ns) 94 if la_is_zero(nrm) == 1 { return -1 } 95 let inv: i64 = f32_div(f32_of(1), nrm) 96 var i: i64 = 0 97 while i < n { v[i] = f32_mul(v[i], inv); i = i + 1 } 98 return 0 99} 100 101// G @ v (G is d x d symmetric) 102func la_gmatvec(G: *i64, v: *i64, d: i64, out: *i64) -> i64 { 103 var a: i64 = 0 104 while a < d { 105 var s: i64 = 0 106 var b: i64 = 0 107 while b < d { s = f32_add(s, f32_mul(G[a * d + b], v[b])); b = b + 1 } 108 out[a] = s 109 a = a + 1 110 } 111 return 0 112} 113 114// reorthogonalize v against basis[0..found-1] (each d long) -- Gram-Schmidt TWICE (M1 lesson). 115func la_deflate_once(v: *i64, basis: *i64, found: i64, d: i64) -> i64 { 116 var j: i64 = 0 117 while j < found { 118 var s: i64 = 0 119 var c: i64 = 0 120 while c < d { s = f32_add(s, f32_mul(v[c], basis[j * d + c])); c = c + 1 } 121 c = 0 122 while c < d { v[c] = f32_sub(v[c], f32_mul(s, basis[j * d + c])); c = c + 1 } 123 j = j + 1 124 } 125 return 0 126} 127func la_reorth(v: *i64, basis: *i64, found: i64, d: i64) -> i64 { 128 la_deflate_once(v, basis, found, d); la_deflate_once(v, basis, found, d); return 0 129} 130 131// top-r eigenvectors of G [d x d] -> basis[r x d] (power iteration + reorth deflation). 132func la_topr_basis(G: *i64, d: i64, r: i64, basis: *i64) -> i64 { 133 let v: *i64 = sys_mmap(d * 8) as *i64 134 let w: *i64 = sys_mmap(d * 8) as *i64 135 var f: i64 = 0 136 while f < r { 137 var c: i64 = 0 138 while c < d { v[c] = 0; c = c + 1 } 139 v[f % d] = f32_of(1) 140 la_reorth(v, basis, f, d) 141 if la_normalize(v, d) < 0 { 142 c = 0 143 while c < d { basis[f * d + c] = 0; c = c + 1 } 144 f = f + 1 145 } else { 146 var it: i64 = 0 147 while it < LA_POWER_ITERS { 148 la_gmatvec(G, v, d, w) 149 la_reorth(w, basis, f, d) 150 if la_normalize(w, d) < 0 { it = LA_POWER_ITERS } else { 151 c = 0 152 while c < d { v[c] = w[c]; c = c + 1 } 153 it = it + 1 154 } 155 } 156 c = 0 157 while c < d { basis[f * d + c] = v[c]; c = c + 1 } 158 f = f + 1 159 } 160 } 161 return 0 162} 163 164// G = M^T M [d x d] 165func la_gram(M: *i64, n: i64, d: i64, G: *i64) -> i64 { 166 var a: i64 = 0 167 while a < d { 168 var b: i64 = 0 169 while b < d { 170 var s: i64 = 0 171 var i: i64 = 0 172 while i < n { s = f32_add(s, f32_mul(M[i * d + a], M[i * d + b])); i = i + 1 } 173 G[a * d + b] = s 174 b = b + 1 175 } 176 a = a + 1 177 } 178 return 0 179} 180 181// compress M[n x d] to rank r (per-token SVD projection) and reconstruct into Mout[n x d]. 182// stored size = n*r (coords) + r*d (basis); returned via the reconstruction (proof of fidelity). 183func la_compress_reconstruct(M: *i64, n: i64, d: i64, r: i64, Mout: *i64) -> i64 { 184 let G: *i64 = sys_mmap(d * d * 8) as *i64 185 let basis: *i64 = sys_mmap(r * d * 8) as *i64 186 la_gram(M, n, d, G) 187 la_topr_basis(G, d, r, basis) 188 // per token: coords[j] = M[i] . basis[j]; Mout[i] = sum_j coords[j]*basis[j] 189 var i: i64 = 0 190 while i < n { 191 var j: i64 = 0 192 let coords: *i64 = sys_mmap(r * 8) as *i64 193 while j < r { 194 var s: i64 = 0 195 var c: i64 = 0 196 while c < d { s = f32_add(s, f32_mul(M[i * d + c], basis[j * d + c])); c = c + 1 } 197 coords[j] = s 198 j = j + 1 199 } 200 var c2: i64 = 0 201 while c2 < d { 202 var s2: i64 = 0 203 j = 0 204 while j < r { s2 = f32_add(s2, f32_mul(coords[j], basis[j * d + c2])); j = j + 1 } 205 Mout[i * d + c2] = s2 206 c2 = c2 + 1 207 } 208 i = i + 1 209 } 210 return 0 211} 212 213// single-query attention: out[d] = softmax(Q.K[i]*scale) . V 214func la_attention(Q: *i64, Km: *i64, Vm: *i64, n: i64, d: i64, scale: i64, out: *i64) -> i64 { 215 let scores: *i64 = sys_mmap(n * 8) as *i64 216 let weights: *i64 = sys_mmap(n * 8) as *i64 217 var i: i64 = 0 218 while i < n { 219 var s: i64 = 0 220 var c: i64 = 0 221 while c < d { s = f32_add(s, f32_mul(Q[c], Km[i * d + c])); c = c + 1 } 222 scores[i] = f32_mul(s, scale) 223 i = i + 1 224 } 225 nx_f32_softmax(scores, n, weights) 226 var c3: i64 = 0 227 while c3 < d { 228 var s3: i64 = 0 229 i = 0 230 while i < n { s3 = f32_add(s3, f32_mul(weights[i], Vm[i * d + c3])); i = i + 1 } 231 out[c3] = s3 232 c3 = c3 + 1 233 } 234 return 0 235} 236 237// relative L2 error of a vs reference b, in permille 238func la_rel_permille(a: *i64, b: *i64, d: i64) -> i64 { 239 var num: i64 = 0 240 var den: i64 = 0 241 var i: i64 = 0 242 while i < d { 243 let diff: i64 = f32_sub(a[i], b[i]) 244 num = f32_add(num, f32_mul(diff, diff)) 245 den = f32_add(den, f32_mul(b[i], b[i])) 246 i = i + 1 247 } 248 if la_is_zero(den) == 1 { return 0 } 249 let rel: i64 = nx_f32_sqrt(f32_div(num, den)) 250 return f32_int(f32_mul(rel, f32_of(1000))) 251} 252 253// run attention with K,V compressed to rank r; return output rel-error vs full-KV reference out_full. 254func la_run_compressed(Q: *i64, K: *i64, V: *i64, n: i64, d: i64, r: i64, scale: i64, out_full: *i64) -> i64 { 255 let Ka: *i64 = sys_mmap(n * d * 8) as *i64 256 let Va: *i64 = sys_mmap(n * d * 8) as *i64 257 la_compress_reconstruct(K, n, d, r, Ka) 258 la_compress_reconstruct(V, n, d, r, Va) 259 let outc: *i64 = sys_mmap(d * 8) as *i64 260 la_attention(Q, Ka, Va, n, d, scale, outc) 261 return la_rel_permille(outc, out_full, d) 262} 263 264// approximately low-rank: dominant rank-`base` + small full-rank noise tail = what REAL KV looks like. 265func la_build_approx(M: *i64, n: i64, d: i64, base: i64, noise_num: i64, seed: i64) -> i64 { 266 la_build_rankr(M, n, d, base, seed) 267 let st: *i64 = sys_mmap(8) as *i64; st[0] = seed + LA_MAGIC_99991 268 let nf: i64 = f32frac(noise_num, 1000) 269 var i: i64 = 0 270 while i < n * d { M[i] = f32_add(M[i], f32_mul(nf, f32frac(la_lcg(st) - LA_MAGIC_2048, LA_MAGIC_4096))); i = i + 1 } 271 return 0 272} 273 274func main() -> i64 { 275 let d: i64 = 8 276 let n: i64 = 64 277 g_puts("=== M3: low-rank KV through a REAL softmax-attention forward (end-to-end output error) ===\n") 278 g_puts(" d="); g_putn(d); g_puts(" n="); g_putn(n); g_puts(" compress K,V to rank r (per-token SVD projection)\n\n") 279 280 let scale: i64 = f32_div(f32_of(1), nx_f32_sqrt(f32_of(d))) 281 282 // genuinely rank-2 K and V; random query Q 283 let K: *i64 = sys_mmap(n * d * 8) as *i64 284 let V: *i64 = sys_mmap(n * d * 8) as *i64 285 la_build_rankr(K, n, d, 2, LA_MAGIC_11003) 286 la_build_rankr(V, n, d, 2, LA_MAGIC_22013) 287 let Q: *i64 = sys_mmap(d * 8) as *i64 288 let st: *i64 = sys_mmap(8) as *i64; st[0] = LA_MAGIC_33029 289 var c: i64 = 0 290 while c < d { Q[c] = f32frac(la_lcg(st) - LA_MAGIC_2048, LA_MAGIC_4096); c = c + 1 } 291 292 // full-KV attention reference 293 let out_full: *i64 = sys_mmap(d * 8) as *i64 294 la_attention(Q, K, V, n, d, scale, out_full) 295 296 // sweep r over rank-2 K,V 297 g_puts(" rank-2 K,V (true rank=2; compress to r):\n") 298 let full_entries: i64 = 2 * n * d 299 var err_r1: i64 = 0 300 var err_r2: i64 = 0 301 let rs: *i64 = sys_mmap(8 * 8) as *i64 302 rs[0] = 1; rs[1] = 2; rs[2] = 3; rs[3] = 4 303 var idx: i64 = 0 304 while idx < 4 { 305 let r: i64 = rs[idx] 306 let e: i64 = la_run_compressed(Q, K, V, n, d, r, scale, out_full) 307 let comp_entries: i64 = 2 * n * r + 2 * r * d 308 g_puts(" r="); g_putn(r); g_puts(" KV entries "); g_putn(full_entries); g_puts("->"); g_putn(comp_entries) 309 g_puts(" ("); g_putn(full_entries / comp_entries); g_puts("x) attn_out_err="); g_putn(e); g_puts("permille") 310 if r >= 2 { g_puts(" [r>=rank: preserved]") } else { g_puts(" [r<rank: lossy]") } 311 g_puts("\n") 312 if r == 1 { err_r1 = e } 313 if r == 2 { err_r2 = e } 314 idx = idx + 1 315 } 316 317 // neg-control: genuinely FULL-rank K,V compressed to r=4 -> output must degrade 318 let Kf: *i64 = sys_mmap(n * d * 8) as *i64 319 let Vf: *i64 = sys_mmap(n * d * 8) as *i64 320 la_build_rankr(Kf, n, d, d, LA_MAGIC_44039) 321 la_build_rankr(Vf, n, d, d, LA_MAGIC_55049) 322 let outf_full: *i64 = sys_mmap(d * 8) as *i64 323 la_attention(Q, Kf, Vf, n, d, scale, outf_full) 324 let err_full: i64 = la_run_compressed(Q, Kf, Vf, n, d, 4, scale, outf_full) 325 g_puts("\n full-rank K,V compressed to r=4 (neg-control) attn_out_err="); g_putn(err_full); g_puts("permille\n\n") 326 327 // --- scale-d sweep (rank-2 K,V @ r=2): output fidelity holds while VRAM ratio GROWS with d 328 // (retires the d=8-reference caveat; the per-token SVD path has NO d limit, unlike the FD sketch) --- 329 g_puts(" scale-d sweep (rank-2 K,V, compress to r=2):\n") 330 let ds: *i64 = sys_mmap(8 * 8) as *i64 331 ds[0] = 8; ds[1] = 32; ds[2] = 64; ds[3] = 128 332 var err_d128: i64 = 0 333 var di: i64 = 0 334 while di < 4 { 335 let dd: i64 = ds[di] 336 let Kd: *i64 = sys_mmap(n * dd * 8) as *i64 337 let Vd: *i64 = sys_mmap(n * dd * 8) as *i64 338 la_build_rankr(Kd, n, dd, 2, LA_MAGIC_11003 + dd) 339 la_build_rankr(Vd, n, dd, 2, LA_MAGIC_22013 + dd) 340 let Qd: *i64 = sys_mmap(dd * 8) as *i64 341 let std: *i64 = sys_mmap(8) as *i64; std[0] = LA_MAGIC_33029 + dd 342 var cc: i64 = 0 343 while cc < dd { Qd[cc] = f32frac(la_lcg(std) - LA_MAGIC_2048, LA_MAGIC_4096); cc = cc + 1 } 344 let scaled: i64 = f32_div(f32_of(1), nx_f32_sqrt(f32_of(dd))) 345 let of: *i64 = sys_mmap(dd * 8) as *i64 346 la_attention(Qd, Kd, Vd, n, dd, scaled, of) 347 let ed: i64 = la_run_compressed(Qd, Kd, Vd, n, dd, 2, scaled, of) 348 let fe: i64 = 2 * n * dd 349 let ce: i64 = 2 * n * 2 + 2 * 2 * dd 350 g_puts(" d="); g_putn(dd); g_puts(" KV entries "); g_putn(fe); g_puts("->"); g_putn(ce) 351 g_puts(" ("); g_putn(fe / ce); g_puts("x) attn_out_err="); g_putn(ed); g_puts("permille\n") 352 if dd == 128 { err_d128 = ed } 353 di = di + 1 354 } 355 g_puts("\n") 356 357 // --- approx-low-rank sweep (d=64, dominant rank-2 + noise tail = realistic KV; compress r=2) --- 358 g_puts(" approx-low-rank sweep (d=64, dominant rank-2 + noise tail = realistic KV; r=2):\n") 359 let noises: *i64 = sys_mmap(8 * 8) as *i64 360 noises[0] = 0; noises[1] = 50; noises[2] = 150; noises[3] = 400 361 var err_n5: i64 = 0 362 var ni: i64 = 0 363 while ni < 4 { 364 let nz: i64 = noises[ni] 365 let Ka: *i64 = sys_mmap(n * 64 * 8) as *i64 366 let Va: *i64 = sys_mmap(n * 64 * 8) as *i64 367 la_build_approx(Ka, n, 64, 2, nz, LA_MAGIC_60061 + nz) 368 la_build_approx(Va, n, 64, 2, nz, LA_MAGIC_70067 + nz) 369 let Qa: *i64 = sys_mmap(64 * 8) as *i64 370 let sta: *i64 = sys_mmap(8) as *i64; sta[0] = LA_MAGIC_80071 + nz 371 var cc: i64 = 0 372 while cc < 64 { Qa[cc] = f32frac(la_lcg(sta) - LA_MAGIC_2048, LA_MAGIC_4096); cc = cc + 1 } 373 let sca: i64 = f32_div(f32_of(1), nx_f32_sqrt(f32_of(64))) 374 let ofa: *i64 = sys_mmap(64 * 8) as *i64 375 la_attention(Qa, Ka, Va, n, 64, sca, ofa) 376 let ea: i64 = la_run_compressed(Qa, Ka, Va, n, 64, 2, sca, ofa) 377 g_puts(" noise="); g_putn(nz); g_puts("permille of signal -> attn_out_err="); g_putn(ea); g_puts("permille\n") 378 if nz == 50 { err_n5 = ea } 379 ni = ni + 1 380 } 381 g_puts("\n") 382 383 // --- gate --- 384 var pass: i64 = 0 385 var fail: i64 = 0 386 if err_r2 < 50 { g_puts(" T1 rank-2 KV @ r=2: attention output PRESERVED end-to-end (<50permille): PASS ("); g_putn(err_r2); g_puts(")\n"); pass = pass + 1 } 387 else { g_puts(" T1 output preserved: FAIL ("); g_putn(err_r2); g_puts(")\n"); fail = fail + 1 } 388 if full_entries > (2 * n * 2 + 2 * 2 * d) { g_puts(" T2 VRAM reduced at r=2 ("); g_putn(full_entries); g_puts("->"); g_putn(2 * n * 2 + 2 * 2 * d); g_puts(" entries): PASS\n"); pass = pass + 1 } 389 else { g_puts(" T2 VRAM reduced: FAIL\n"); fail = fail + 1 } 390 if err_full > err_r2 * 3 + 20 { g_puts(" T3 neg-control: full-rank output degrades >> rank-2 (metric real): PASS\n"); pass = pass + 1 } 391 else { g_puts(" T3 neg-control: FAIL\n"); fail = fail + 1 } 392 if err_r1 > err_r2 { g_puts(" T4 graceful: under-ranking (r=1<rank) loses more than r=2: PASS\n"); pass = pass + 1 } 393 else { g_puts(" T4 graceful: FAIL\n"); fail = fail + 1 } 394 if err_d128 < 50 { g_puts(" T5 fidelity holds at d=128 (output err<50permille, ratio grows): PASS ("); g_putn(err_d128); g_puts(")\n"); pass = pass + 1 } 395 else { g_puts(" T5 scale-d fidelity: FAIL ("); g_putn(err_d128); g_puts(")\n"); fail = fail + 1 } 396 if err_n5 < 150 { g_puts(" T6 approx-low-rank (5% noise tail) output err acceptable (<150permille): PASS ("); g_putn(err_n5); g_puts(")\n"); pass = pass + 1 } 397 else { g_puts(" T6 approx-low-rank: FAIL ("); g_putn(err_n5); g_puts(")\n"); fail = fail + 1 } 398 399 g_puts("\n PASS="); g_putn(pass); g_puts("/6 ") 400 if fail == 0 { g_puts("VERDICT=GREEN (low-rank KV preserves REAL attention output; VRAM cut, quality kept end-to-end)\n"); sys_exit(0); return 0 } 401 g_puts("VERDICT=RED\n"); sys_exit(1); return 1 402}