nx_nofloat_muon.nx source

1// nx_nofloat_muon.nx -- the FULL Per-Head Muon optimizer LOOP in sovereign no-float (K3 F405; operator 2// 2026-07-19 "beyond SOTA, novel, evidence-driven"). Muon = MomentUm Orthogonalized by Newton-schulz 3// (Keller Jordan / Moonshot). One optimizer step on a weight matrix: 4// B_t = mu * B_{t-1} + G_t (heavy-ball momentum buffer) 5// U_t = NewtonSchulz(B_t) (orthogonalize the momentum -- the proven nx_nofloat_muon_ns core) 6// W_t = W_{t-1} - lr * U_t (apply the orthogonalized update) 7// PER-HEAD: attention weights are nheads independent hd x hd blocks; each is orthogonalized SEPARATELY 8// (that is the "Per-Head" in Per-Head Muon -- a per-head spectral step, not one global orthogonalization). 9// PURE INTEGER Q16 throughout -> BIT-EXACT DETERMINISTIC optimizer: a float Muon drifts (non-associative 10// gradient/momentum accumulation + float NS); ours is identical every machine, every run = the novel TRAIN 11// exceed. Composes nx_nofloat_muon_ns (mn_orthogonalize). license_tier: ORIGINAL No hw writes (Rule 26). 12import "nx_nofloat_muon_ns.nx" 13import "nx_syscalls.nx" 14 15const MU_Q: i64 = 65536 16 17// heavy-ball momentum update in place: buf[i] = (mu*buf[i])>>16 + grad[i] (MN_QBITS=16 from the NS lib) 18func mu_momentum(buf: *i64, grad: *i64, mu: i64, nn: i64) -> i64 { 19 var i: i64 = 0 20 while i < nn { buf[i] = ((mu * buf[i]) >> MN_QBITS) + grad[i]; i = i + 1 } 21 return 0 22} 23// apply the orthogonalized update: w[i] = w[i] - (lr*u[i])>>16 24func mu_apply(w: *i64, u: *i64, lr: i64, nn: i64) -> i64 { 25 var i: i64 = 0 26 while i < nn { w[i] = w[i] - ((lr * u[i]) >> MN_QBITS); i = i + 1 } 27 return 0 28} 29// one Muon step on a single n x n matrix. buf is the persistent momentum buffer; u is scratch (n x n). 30func mu_step(w: *i64, buf: *i64, grad: *i64, mu: i64, lr: i64, n: i64, u: *i64) -> i64 { 31 let nn: i64 = n * n 32 mu_momentum(buf, grad, mu, nn) // B = mu*B + G 33 mn_orthogonalize(buf, u, n) // U = NS(B) (the proven core) 34 mu_apply(w, u, lr, nn) // W = W - lr*U 35 return 0 36} 37// PER-HEAD Muon step: w/buf/grad are nheads contiguous hd x hd blocks. Each head orthogonalized separately. 38func mu_perhead_step(w: *i64, buf: *i64, grad: *i64, mu: i64, lr: i64, nheads: i64, hd: i64, u: *i64) -> i64 { 39 let blk: i64 = hd * hd 40 var h: i64 = 0 41 while h < nheads { 42 let off: i64 = h * blk 43 let wh: *i64 = ((w as i64) + off*8) as *i64 44 let bh: *i64 = ((buf as i64) + off*8) as *i64 45 let gh: *i64 = ((grad as i64) + off*8) as *i64 46 mu_step(wh, bh, gh, mu, lr, hd, u) 47 h = h + 1 48 } 49 return 0 50}