nx_nofloat_k3stack.nx source

1// nx_nofloat_k3stack.nx -- MULTI-LAYER K3 stack with ATTENTION RESIDUALS in no-float (operator 2026-07-19 2// "logically integrated, mature"). Stacks the K3 block at DEPTH, and closes the K3-specific Attention- 3// Residuals gap: instead of a plain x+f(x) residual, each layer's input is a CONTENT-WEIGHTED ATTENTION 4// over ALL preceding layer outputs (arXiv 2603.15031: "learned input-dependent attention over preceding 5// layers... earlier-layer representations stay reachable deeper in the stack"). Minimal faithful AttnRes: 6// score_s = |mean(h_s) . mean(h_{l-1})| ; w_s = score_s / sum ; r_l = sum_{s<l} w_s h_s ; h_l = block(r_l) 7// Pure integer Q16, deterministic AT DEPTH -- a float stack drifts per-layer and compounds; ours is bit-exact 8// through N layers. Composes nx_nofloat_k3block (the K3 block) + weight-tying across layers. 9// license_tier: ORIGINAL No hw writes (Rule 26). 10import "nx_nofloat_k3block.nx" 11import "nx_syscalls.nx" 12import "nx_vecmath.nx" 13 14const KS_QBITS: i64 = 16 15const KS_Q: i64 = 65536 16const KS_LMAX: i64 = 8 17const KS_CELLMAX: i64 = 256 // t*d cap per layer buffer 18 19func ks_mean(h: *i64, td: i64) -> i64 { 20 var s: i64 = 0 21 var i: i64 = 0 22 while i < td { s = s + h[i]; i = i + 1 } 23 return s / td 24} 25func ks_isqrt(v: i64) -> i64 { return vm_isqrt(v) } 26// RMSNorm (pre-norm transformers keep activations bounded at depth): h[i] <- h[i] / RMS(h), unit RMS. 27// derivation: RMS_q16 ~ isqrt(sum(h>>8)^2 / td) * 256 ; h_norm[i] = (h[i]<<8) / isqrt(...) (unit-RMS Q16). 28func ks_rmsnorm(h: *i64, td: i64) -> i64 { 29 var ss: i64 = 0 30 var i: i64 = 0 31 while i < td { let v: i64 = h[i] >> 8; ss = ss + v * v; i = i + 1 } 32 let q: i64 = ks_isqrt(ss / td) 33 if q <= 0 { return 0 } 34 i = 0 35 while i < td { h[i] = (h[i] << 8) / q; i = i + 1 } 36 return 0 37} 38// N-layer K3 stack with AttnRes cross-layer residual. weights tied across layers. out: t x d. 39func k3_stack(x: *i64, wq: *i64, wk: *i64, wv: *i64, wo: *i64, w1: *i64, w2: *i64, layers: i64, t: i64, d: i64, dff: i64, out: *i64) -> i64 { 40 let td: i64 = t * d 41 // layer output buffers h[0..layers] 42 let hbase: *i64 = sys_mmap((KS_LMAX+1) * KS_CELLMAX * 8) as *i64 43 let means: *i64 = sys_mmap((KS_LMAX+1) * 8) as *i64 44 // h[0] = x 45 var i: i64 = 0 46 while i < td { hbase[i] = x[i]; i = i + 1 } 47 means[0] = ks_mean(hbase, td) 48 let r: *i64 = sys_mmap(KS_CELLMAX * 8) as *i64 49 var l: i64 = 1 50 while l <= layers { 51 // AttnRes: content-weighted combine of h[0..l-1] 52 let prev: i64 = l - 1 53 // scores sc[s] = |mean_s * mean_prev| >> 16, +1 to avoid all-zero 54 var sumsc: i64 = 0 55 let sc: *i64 = sys_mmap((KS_LMAX+1)*8) as *i64 56 var s: i64 = 0 57 while s < l { 58 var v: i64 = (means[s] * means[prev]) >> KS_QBITS 59 if v < 0 { v = 0 - v } 60 v = v + 1 61 sc[s] = v 62 sumsc = sumsc + v 63 s = s + 1 64 } 65 // r = sum_s (sc[s]/sumsc) * h[s] 66 var c: i64 = 0 67 while c < td { r[c] = 0; c = c + 1 } 68 s = 0 69 while s < l { 70 let w: i64 = (sc[s] << KS_QBITS) / sumsc 71 let hs: *i64 = ((hbase as i64) + s * KS_CELLMAX * 8) as *i64 72 c = 0 73 while c < td { r[c] = r[c] + ((w * hs[c]) >> KS_QBITS); c = c + 1 } 74 s = s + 1 75 } 76 // h[l] = block(r) 77 let hl: *i64 = ((hbase as i64) + l * KS_CELLMAX * 8) as *i64 78 k3_block(r, wq, wk, wv, wo, w1, w2, hl, t, d, dff) 79 ks_rmsnorm(hl, td) // pre-norm keeps activations bounded at depth (DeepSeek-V3 style) 80 means[l] = ks_mean(hl, td) 81 l = l + 1 82 } 83 // out = h[layers] 84 let hL: *i64 = ((hbase as i64) + layers * KS_CELLMAX * 8) as *i64 85 i = 0 86 while i < td { out[i] = hL[i]; i = i + 1 } 87 return 0 88} 89// re-expose: get layer-l input residual r (for the cross-layer test) -- runs the AttnRes combine only 90func k3_stack_out_layer(x: *i64, wq: *i64, wk: *i64, wv: *i64, wo: *i64, w1: *i64, w2: *i64, layers: i64, t: i64, d: i64, dff: i64, out: *i64) -> i64 { 91 return k3_stack(x, wq, wk, wv, wo, w1, w2, layers, t, d, dff, out) 92}