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" 12 13const KS_QBITS: i64 = 16 14const KS_Q: i64 = 65536 15const KS_LMAX: i64 = 8 16const KS_CELLMAX: i64 = 256 // t*d cap per layer buffer 17 18func ks_mean(h: *i64, td: i64) -> i64 { 19 var s: i64 = 0 20 var i: i64 = 0 21 while i < td { s = s + h[i]; i = i + 1 } 22 return s / td 23} 24func ks_isqrt(v: i64) -> i64 { 25 if v <= 0 { return 0 } 26 var x: i64 = v 27 var y: i64 = (x + 1) / 2 28 while y < x { x = y; y = (x + v / x) / 2 } 29 return x 30} 31// RMSNorm (pre-norm transformers keep activations bounded at depth): h[i] <- h[i] / RMS(h), unit RMS. 32// derivation: RMS_q16 ~ isqrt(sum(h>>8)^2 / td) * 256 ; h_norm[i] = (h[i]<<8) / isqrt(...) (unit-RMS Q16). 33func ks_rmsnorm(h: *i64, td: i64) -> i64 { 34 var ss: i64 = 0 35 var i: i64 = 0 36 while i < td { let v: i64 = h[i] >> 8; ss = ss + v * v; i = i + 1 } 37 let q: i64 = ks_isqrt(ss / td) 38 if q <= 0 { return 0 } 39 i = 0 40 while i < td { h[i] = (h[i] << 8) / q; i = i + 1 } 41 return 0 42} 43// N-layer K3 stack with AttnRes cross-layer residual. weights tied across layers. out: t x d. 44func 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 { 45 let td: i64 = t * d 46 // layer output buffers h[0..layers] 47 let hbase: *i64 = sys_mmap((KS_LMAX+1) * KS_CELLMAX * 8) as *i64 48 let means: *i64 = sys_mmap((KS_LMAX+1) * 8) as *i64 49 // h[0] = x 50 var i: i64 = 0 51 while i < td { hbase[i] = x[i]; i = i + 1 } 52 means[0] = ks_mean(hbase, td) 53 let r: *i64 = sys_mmap(KS_CELLMAX * 8) as *i64 54 var l: i64 = 1 55 while l <= layers { 56 // AttnRes: content-weighted combine of h[0..l-1] 57 let prev: i64 = l - 1 58 // scores sc[s] = |mean_s * mean_prev| >> 16, +1 to avoid all-zero 59 var sumsc: i64 = 0 60 let sc: *i64 = sys_mmap((KS_LMAX+1)*8) as *i64 61 var s: i64 = 0 62 while s < l { 63 var v: i64 = (means[s] * means[prev]) >> KS_QBITS 64 if v < 0 { v = 0 - v } 65 v = v + 1 66 sc[s] = v 67 sumsc = sumsc + v 68 s = s + 1 69 } 70 // r = sum_s (sc[s]/sumsc) * h[s] 71 var c: i64 = 0 72 while c < td { r[c] = 0; c = c + 1 } 73 s = 0 74 while s < l { 75 let w: i64 = (sc[s] << KS_QBITS) / sumsc 76 let hs: *i64 = ((hbase as i64) + s * KS_CELLMAX * 8) as *i64 77 c = 0 78 while c < td { r[c] = r[c] + ((w * hs[c]) >> KS_QBITS); c = c + 1 } 79 s = s + 1 80 } 81 // h[l] = block(r) 82 let hl: *i64 = ((hbase as i64) + l * KS_CELLMAX * 8) as *i64 83 k3_block(r, wq, wk, wv, wo, w1, w2, hl, t, d, dff) 84 ks_rmsnorm(hl, td) // pre-norm keeps activations bounded at depth (DeepSeek-V3 style) 85 means[l] = ks_mean(hl, td) 86 l = l + 1 87 } 88 // out = h[layers] 89 let hL: *i64 = ((hbase as i64) + layers * KS_CELLMAX * 8) as *i64 90 i = 0 91 while i < td { out[i] = hL[i]; i = i + 1 } 92 return 0 93} 94// re-expose: get layer-l input residual r (for the cross-layer test) -- runs the AttnRes combine only 95func 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 { 96 return k3_stack(x, wq, wk, wv, wo, w1, w2, layers, t, d, dff, out) 97}