nx_nofloat_olmoe.nx source

1// nx_nofloat_olmoe.nx -- the faithful OLMoE transformer LAYER in deterministic no-float (MoE rung 4, 2// 2026-07-15). OLMoE-1B-7B blk shape (recon-verified): D=2048, 16 heads x head_dim 128, NO GQA, NO biases, 3// **QK-norm = RMSNorm over the FULL 2048-dim Q and K projections** (gamma[2048], the OLMoE/Qwen3 variant -- 4// applied AFTER projection, BEFORE RoPE), RoPE base 1e4, then MoE-FFN (64 experts top-8, lazy). Two 5// sublayers; the caller (gate/serve) does the residual adds + owns all scratch (lean frames, no big-frame 6// GP-fault). Composes ONLY proven primitives: rmsnorm_gamma_row / mm_out_in / rope_apply / fx_exp / silu / 7// qmul / dequant_to_q16 / nf_expert_byteoff. Deterministic by construction (exact integer sums, positional 8// tie-breaks). license_tier: ORIGINAL (lib: no main) 9import "nx_syscalls.nx" 10import "nx_nofloat_llm.nx" 11 12// ATTENTION sublayer -> writes attn projection (PRE-residual) to `out`. ap bundle: 13// [0]=x(T*D Q16) [1]=out(T*D) [2]=T [3]=D [4]=nh [5]=hd [6]=scale_q16 14// [7]=gAttn(D) [8]=Wq(D*D) [9]=Wk(D*D) [10]=Wv(D*D) [11]=Wo(D*D) 15// [12]=gQ(D) [13]=gK(D) [14]=freqs(hd/2) [15]=scr 16// scr layout (i64): xn[T*D] Q[T*D] K[T*D] V[T*D] sc[T] at[T] concat[T*D] 17func nolmoe_attn(ap: *i64) -> i64 { 18 let x: *i64 = ap[0] as *i64 19 let out: *i64 = ap[1] as *i64 20 let T: i64 = ap[2] 21 let D: i64 = ap[3] 22 let nh: i64 = ap[4] 23 let hd: i64 = ap[5] 24 let scale: i64 = ap[6] 25 let gA: *i64 = ap[7] as *i64 26 let Wq: *i64 = ap[8] as *i64 27 let Wk: *i64 = ap[9] as *i64 28 let Wv: *i64 = ap[10] as *i64 29 let Wo: *i64 = ap[11] as *i64 30 let gQ: *i64 = ap[12] as *i64 31 let gK: *i64 = ap[13] as *i64 32 let freqs: *i64 = ap[14] as *i64 33 let scr: *i64 = ap[15] as *i64 34 let xn: *i64 = scr 35 let Q: *i64 = ((scr as i64) + T*D*8) as *i64 36 let K: *i64 = ((Q as i64) + T*D*8) as *i64 37 let V: *i64 = ((K as i64) + T*D*8) as *i64 38 let sc: *i64 = ((V as i64) + T*D*8) as *i64 39 let at: *i64 = ((sc as i64) + T*8) as *i64 40 let concat: *i64 = ((at as i64) + T*8) as *i64 41 // per token: pre-norm -> QKV -> QK-norm -> RoPE 42 var t: i64 = 0 43 while t < T { 44 rmsnorm_gamma_row(x, gA, t*D, D, xn, t*D) 45 mm_out_in(((xn as i64) + t*D*8) as *i64, Wq, ((Q as i64) + t*D*8) as *i64, 1, D, D, 0) 46 mm_out_in(((xn as i64) + t*D*8) as *i64, Wk, ((K as i64) + t*D*8) as *i64, 1, D, D, 0) 47 mm_out_in(((xn as i64) + t*D*8) as *i64, Wv, ((V as i64) + t*D*8) as *i64, 1, D, D, 0) 48 // QK-norm: RMSNorm over the FULL D on Q and K (in place -- reads-then-writes each index, safe) 49 rmsnorm_gamma_row(Q, gQ, t*D, D, Q, t*D) 50 rmsnorm_gamma_row(K, gK, t*D, D, K, t*D) 51 var h: i64 = 0 52 while h < nh { 53 rope_apply(((Q as i64) + (t*D + h*hd)*8) as *i64, hd, t, freqs) 54 rope_apply(((K as i64) + (t*D + h*hd)*8) as *i64, hd, t, freqs) 55 h = h + 1 56 } 57 t = t + 1 58 } 59 // causal MHA 60 t = 0 61 while t < T { 62 var h: i64 = 0 63 while h < nh { 64 let qb: i64 = h*hd 65 var s: i64 = 0 66 while s <= t { 67 var dot: i64 = 0 68 var d: i64 = 0 69 while d < hd { dot = dot + (Q[t*D + qb + d] * K[s*D + qb + d]); d = d + 1 } 70 sc[s] = qmul(dot >> 16, scale) 71 s = s + 1 72 } 73 var mmax: i64 = sc[0] 74 var j: i64 = 1 75 while j <= t { if sc[j] > mmax { mmax = sc[j] } j = j + 1 } 76 var sum: i64 = 0 77 j = 0 78 while j <= t { let e: i64 = fx_exp(sc[j] - mmax); at[j] = e; sum = sum + e; j = j + 1 } 79 if sum < 1 { sum = 1 } 80 j = 0 81 while j <= t { at[j] = (at[j] << 16) / sum; j = j + 1 } 82 var d2: i64 = 0 83 while d2 < hd { 84 var acc: i64 = 0 85 s = 0 86 while s <= t { acc = acc + (at[s] * V[s*D + qb + d2]); s = s + 1 } 87 concat[t*D + qb + d2] = acc >> 16 88 d2 = d2 + 1 89 } 90 h = h + 1 91 } 92 t = t + 1 93 } 94 // output projection 95 t = 0 96 while t < T { mm_out_in(((concat as i64) + t*D*8) as *i64, Wo, ((out as i64) + t*D*8) as *i64, 1, D, D, 0); t = t + 1 } 97 return 0 98} 99 100// MoE-FFN sublayer (per-token, LAZY top-K expert fetch) -> writes MoE output (PRE-residual) to `out`. 101// mp bundle: [0]=x(T*D Q16) [1]=out(T*D) [2]=T [3]=D [4]=ff [5]=E [6]=K [7]=gFfn(D) [8]=Wr(E*D Q16) 102// [9]=buf [10]=eb(ptr to [bG,tyG,bU,tyU,bD,tyD]) [11]=scr [12]=cnt(E) 103// scr layout: xn[D] r[E] sel[K] w[K] gs[ff*D] us[ff*D] ds[D*ff] g[ff] u[ff] hh[ff] eo[D] 104func nolmoe_moe(mp: *i64) -> i64 { 105 let x: *i64 = mp[0] as *i64 106 let out: *i64 = mp[1] as *i64 107 let T: i64 = mp[2] 108 let D: i64 = mp[3] 109 let ff: i64 = mp[4] 110 let E: i64 = mp[5] 111 let K: i64 = mp[6] 112 let gF: *i64 = mp[7] as *i64 113 let Wr: *i64 = mp[8] as *i64 114 let buf: *u8 = mp[9] as *u8 115 let eb: *i64 = mp[10] as *i64 116 let scr: *i64 = mp[11] as *i64 117 let cnt: *i64 = mp[12] as *i64 118 let bG: i64 = eb[0] 119 let tyG: i64 = eb[1] 120 let bU: i64 = eb[2] 121 let tyU: i64 = eb[3] 122 let bD: i64 = eb[4] 123 let tyD: i64 = eb[5] 124 let xn: *i64 = scr 125 let r: *i64 = ((xn as i64) + D*8) as *i64 126 let sel: *i64 = ((r as i64) + E*8) as *i64 127 let w: *i64 = ((sel as i64) + K*8) as *i64 128 let gs: *i64 = ((w as i64) + K*8) as *i64 129 let us: *i64 = ((gs as i64) + ff*D*8) as *i64 130 let ds: *i64 = ((us as i64) + ff*D*8) as *i64 131 let g: *i64 = ((ds as i64) + D*ff*8) as *i64 132 let u: *i64 = ((g as i64) + ff*8) as *i64 133 let hh: *i64 = ((u as i64) + ff*8) as *i64 134 let eo: *i64 = ((hh as i64) + ff*8) as *i64 135 var t: i64 = 0 136 while t < T { 137 rmsnorm_gamma_row(x, gF, t*D, D, xn, 0) 138 mm_out_in(xn, Wr, r, 1, D, E, 0) 139 // top-K by router logit (strict >, lowest index on ties) 140 var ki: i64 = 0 141 while ki < K { 142 var best: i64 = 0 - 1 143 var bestv: i64 = 0 144 var e: i64 = 0 145 while e < E { 146 var taken: i64 = 0 147 var q: i64 = 0 148 while q < ki { if sel[q] == e { taken = 1 } q = q + 1 } 149 if taken == 0 { if best < 0 { best = e; bestv = r[e] } else { if r[e] > bestv { best = e; bestv = r[e] } } } 150 e = e + 1 151 } 152 sel[ki] = best 153 ki = ki + 1 154 } 155 var m: i64 = r[sel[0]] 156 var i: i64 = 1 157 while i < K { if r[sel[i]] > m { m = r[sel[i]] } i = i + 1 } 158 var sum: i64 = 0 159 i = 0 160 while i < K { let ev: i64 = fx_exp(r[sel[i]] - m); w[i] = ev; sum = sum + ev; i = i + 1 } 161 if sum < 1 { sum = 1 } 162 i = 0 163 while i < K { w[i] = (w[i] << 16) / sum; i = i + 1 } 164 var d: i64 = 0 165 while d < D { out[t*D + d] = 0; d = d + 1 } 166 i = 0 167 while i < K { 168 let e2: i64 = sel[i] 169 cnt[e2] = cnt[e2] + 1 170 dequant_to_q16(buf, bG + nf_expert_byteoff(tyG, e2*ff*D), tyG, ff*D, gs) 171 dequant_to_q16(buf, bU + nf_expert_byteoff(tyU, e2*ff*D), tyU, ff*D, us) 172 dequant_to_q16(buf, bD + nf_expert_byteoff(tyD, e2*D*ff), tyD, D*ff, ds) 173 mm_out_in(xn, gs, g, 1, D, ff, 0) 174 mm_out_in(xn, us, u, 1, D, ff, 0) 175 var j: i64 = 0 176 while j < ff { hh[j] = qmul(silu(g[j]), u[j]); j = j + 1 } 177 mm_out_in(hh, ds, eo, 1, ff, D, 0) 178 d = 0 179 while d < D { out[t*D + d] = out[t*D + d] + qmul(w[i], eo[d]); d = d + 1 } 180 i = i + 1 181 } 182 t = t + 1 183 } 184 return 0 185}