nx_nofloat_muon_ns.nx source

1// nx_nofloat_muon_ns.nx -- SOVEREIGN no-float Newton-Schulz orthogonalization = the core of the Muon 2// optimizer (K3 F405 gap; operator 2026-07-19 "beyond SOTA, novel"). Muon = MomentUm Orthogonalized by 3// Newton-Schulz (Keller Jordan). NS quintic iteration drives a matrix toward its orthogonal polar factor: 4// X0 = G / ||G||_F (Frobenius-normalize so spectral norm <= 1) 5// X_{k+1} = a*X + b*(A*X) + c*(A*A*X), A = X*X^T, (a,b,c) = (3.4445, -4.7750, 2.0315) 6// After ~5 iters X*X^T -> I. Done in PURE INTEGER Q16 (matmul accumulate-then-shift + integer isqrt) = 7// BIT-EXACT DETERMINISTIC orthogonalization -- a float training stack CANNOT reproduce this bit-for-bit 8// (non-associative accumulation); that is the novel exceed. Grounded: kellerjordan.github.io/posts/muon. 9// license_tier: ORIGINAL No hw writes (Rule 26). 10import "nx_syscalls.nx" 11 12const MN_Q: i64 = 65536 // Q16 one 13const MN_QBITS: i64 = 16 14const MN_NMAX: i64 = 8 // max matrix dim supported 15// NS quintic coeffs in Q16 (Muon): a=3.4445, b=-4.7750, c=2.0315 16const MN_A: i64 = 225755 17const MN_B: i64 = 0 - 312934 18const MN_C: i64 = 133136 19const MN_ITERS: i64 = 5 20 21// integer sqrt of a non-negative i64 (floor). Newton iteration. 22func mn_isqrt(v: i64) -> i64 { 23 if v <= 0 { return 0 } 24 var x: i64 = v 25 var y: i64 = (x + 1) / 2 26 while y < x { x = y; y = (x + v / x) / 2 } 27 return x 28} 29// C[n x n] = A[n x n] * B[n x n] in Q16 (accumulate i64, >>16). A,B,C distinct buffers. 30func mn_mm(a: *i64, b: *i64, c: *i64, n: i64) -> i64 { 31 var i: i64 = 0 32 while i < n { 33 var j: i64 = 0 34 while j < n { 35 var s: i64 = 0 36 var k: i64 = 0 37 while k < n { s = s + a[i*n+k] * b[k*n+j]; k = k + 1 } 38 c[i*n+j] = s >> MN_QBITS 39 j = j + 1 40 } 41 i = i + 1 42 } 43 return 0 44} 45// C = A^T (n x n) 46func mn_transpose(a: *i64, c: *i64, n: i64) -> i64 { 47 var i: i64 = 0 48 while i < n { var j: i64 = 0; while j < n { c[j*n+i] = a[i*n+j]; j = j + 1 } i = i + 1 } 49 return 0 50} 51// out = pa*A + pb*B + pc*C (scalars in Q16, matrices in Q16) elementwise, n x n 52func mn_comb(a: *i64, b: *i64, c: *i64, out: *i64, pa: i64, pb: i64, pc: i64, n: i64) -> i64 { 53 var i: i64 = 0 54 let nn: i64 = n * n 55 while i < nn { 56 out[i] = ((pa * a[i]) >> MN_QBITS) + ((pb * b[i]) >> MN_QBITS) + ((pc * c[i]) >> MN_QBITS) 57 i = i + 1 58 } 59 return 0 60} 61// Newton-Schulz orthogonalize G (n x n, Q16) -> X (n x n, Q16), `iters` NS steps. 62// Muon NS = APPROXIMATE orthogonalization: singular values -> ~1 (fixed points 0.87/1.26), not exactly 1. 63func mn_orthogonalize_it(g: *i64, x: *i64, n: i64, iters: i64) -> i64 { 64 let nn: i64 = n * n 65 // Frobenius norm ||G||_F = sqrt(sum g^2). g in Q16 -> g^2 in Q32; sum in i64; sqrt -> Q16 norm. 66 var ss: i64 = 0 67 var i: i64 = 0 68 while i < nn { let v: i64 = g[i] >> 8; ss = ss + v * v; i = i + 1 } // (g>>8)^2 = g^2 >> 16 = Q16-scaled sumsq 69 let fro: i64 = mn_isqrt(ss) // ~ ||G||_F in Q8-ish; normalize by ratio 70 if fro <= 0 { return 0 } 71 // X0 = G / ||G||_F : x = g * Q / (fro<<8) (fro was sqrt of Q16-scaled -> scale back) 72 let denom: i64 = fro << 8 73 i = 0 74 while i < nn { x[i] = (g[i] << MN_QBITS) / denom; i = i + 1 } 75 // iteration buffers 76 let xt: *i64 = sys_mmap(MN_NMAX*MN_NMAX*8) as *i64 77 let aM: *i64 = sys_mmap(MN_NMAX*MN_NMAX*8) as *i64 78 let ax: *i64 = sys_mmap(MN_NMAX*MN_NMAX*8) as *i64 79 let aax: *i64 = sys_mmap(MN_NMAX*MN_NMAX*8) as *i64 80 var it: i64 = 0 81 while it < iters { 82 mn_transpose(x, xt, n) // X^T 83 mn_mm(x, xt, aM, n) // A = X X^T 84 mn_mm(aM, x, ax, n) // A X 85 mn_mm(aM, ax, aax, n) // A A X 86 mn_comb(x, ax, aax, x, MN_A, MN_B, MN_C, n) // X = aX + b(AX) + c(AAX) 87 it = it + 1 88 } 89 return 0 90} 91// default 5-iteration Muon NS 92func mn_orthogonalize(g: *i64, x: *i64, n: i64) -> i64 { return mn_orthogonalize_it(g, x, n, MN_ITERS) } 93// orthogonality error metric: sum |(X X^T)[i][j] - I[i][j]| in Q16 (0 = perfectly orthogonal) 94func mn_ortho_err(x: *i64, n: i64) -> i64 { 95 let xt: *i64 = sys_mmap(MN_NMAX*MN_NMAX*8) as *i64 96 let g: *i64 = sys_mmap(MN_NMAX*MN_NMAX*8) as *i64 97 mn_transpose(x, xt, n) 98 mn_mm(x, xt, g, n) 99 var e: i64 = 0 100 var i: i64 = 0 101 while i < n { 102 var j: i64 = 0 103 while j < n { 104 var t: i64 = g[i*n+j] 105 if i == j { t = t - MN_Q } 106 if t < 0 { t = 0 - t } 107 e = e + t 108 j = j + 1 109 } 110 i = i + 1 111 } 112 return e 113}