nx_nofloat_moe.nx source
↩ module page · 122 lines · 5207 B
1// nx_nofloat_moe.nx -- no-float Mixture-of-Experts FFN (2026-07-15, MoE rung 1: the layer MECHANISM).
2// The July-2026 open-weight leaders (DeepSeek V4 / Qwen3.5 / Llama 4-5 / Kimi K2) are ALL sparse-MoE --
3// the arch census's #1 gap. The f32 TRAINER half exists (nx_f32_moe_train_gate 6/6, 07-10); this is the
4// INFERENCE half in pure Q16 integer, composed from the proven no-float primitives:
5// router logits r = x @ Wr^T ([out,in] GGUF convention, mm_out_in)
6// -> deterministic top-K (strict >, LOWEST index on ties -- positional, never order-dependent)
7// -> fx_exp softmax over the K selected (the attention softmax's own transcendental)
8// -> K x SwiGLU experts (silu(x@Wg) (*) x@Wu) @ Wd -- only the SELECTED experts are computed
9// -> weighted mix out = sum_i qmul(w_i, expert_i(x)).
10// cnt[] increments per COMPUTED expert = the sparsity instrument (K of E, measured not asserted).
11// Deterministic by construction: exact integer sums (order-free), positional tie-breaks, Q16 fx_exp.
12// mp = [0]=x(D) [1]=Wr(E*D) [2]=Wg(E blocks of H*D) [3]=Wu(E blocks of H*D) [4]=Wd(E blocks of D*H)
13// [5]=out(D) [6]=D [7]=H [8]=E [9]=K [10]=cnt(E) [11]=scr(i64 scratch >= E+2K+3H+D slots)
14// license_tier: ORIGINAL (lib: no main -- build standalone gives rc=102 by design)
15import "nx_syscalls.nx"
16import "nx_nofloat_llm.nx"
17
18// MoE BLOCK (rung 2): per-token pre-norm MoE-FFN with residual -- out[t] = x[t] + MoE(RMSNorm_gamma(x[t])).
19// PER-TOKEN routing (each token consults the router independently -- the property that makes MoE MoE at T>1).
20// bp = [0]=x(T*D) [1]=gamma(D) [2]=Wr [3]=Wg [4]=Wu [5]=Wd [6]=out(T*D) [7]=T [8]=D [9]=H [10]=E [11]=K
21// [12]=cnt(E, accumulates across tokens) [13]=scr (needs D extra slots for the normed token + rung-1 scr)
22func nmoe_block_forward(bp: *i64) -> i64 {
23 let x: *i64 = bp[0] as *i64
24 let gamma: *i64 = bp[1] as *i64
25 let out: *i64 = bp[6] as *i64
26 let T: i64 = bp[7]
27 let D: i64 = bp[8]
28 let scr: *i64 = bp[13] as *i64
29 let xn: *i64 = scr // D normed slots
30 let inner: *i64 = ((scr as i64) + D*8) as *i64 // rung-1 scratch (E+2K+3H+D+16 slots)
31 let mp: *i64 = ((inner as i64) + (bp[10] + 2*bp[11] + 3*bp[9] + D + 16)*8) as *i64 // 16-slot mp bundle after inner scr
32 var t: i64 = 0
33 while t < T {
34 rmsnorm_gamma_row(x, gamma, t*D, D, xn, 0)
35 mp[0] = xn as i64
36 mp[1] = bp[2]
37 mp[2] = bp[3]
38 mp[3] = bp[4]
39 mp[4] = bp[5]
40 mp[5] = ((out as i64) + t*D*8)
41 mp[6] = D
42 mp[7] = bp[9]
43 mp[8] = bp[10]
44 mp[9] = bp[11]
45 mp[10] = bp[12]
46 mp[11] = inner as i64
47 nmoe_forward(mp)
48 var d: i64 = 0
49 while d < D { out[t*D + d] = out[t*D + d] + x[t*D + d]; d = d + 1 }
50 t = t + 1
51 }
52 return 0
53}
54
55func nmoe_forward(mp: *i64) -> i64 {
56 let x: *i64 = mp[0] as *i64
57 let Wr: *i64 = mp[1] as *i64
58 let Wg: *i64 = mp[2] as *i64
59 let Wu: *i64 = mp[3] as *i64
60 let Wd: *i64 = mp[4] as *i64
61 let out: *i64 = mp[5] as *i64
62 let D: i64 = mp[6]
63 let H: i64 = mp[7]
64 let E: i64 = mp[8]
65 let K: i64 = mp[9]
66 let cnt: *i64 = mp[10] as *i64
67 let scr: *i64 = mp[11] as *i64
68 let r: *i64 = scr
69 let sel: *i64 = ((scr as i64) + E*8) as *i64
70 let w: *i64 = ((sel as i64) + K*8) as *i64
71 let g: *i64 = ((w as i64) + K*8) as *i64
72 let u: *i64 = ((g as i64) + H*8) as *i64
73 let h: *i64 = ((u as i64) + H*8) as *i64
74 let eo: *i64 = ((h as i64) + H*8) as *i64
75 // router logits
76 mm_out_in(x, Wr, r, 1, D, E, 0)
77 // deterministic top-K: strict max scan over unselected experts; lowest index wins ties
78 var ki: i64 = 0
79 while ki < K {
80 var best: i64 = 0 - 1
81 var bestv: i64 = 0
82 var e: i64 = 0
83 while e < E {
84 var taken: i64 = 0
85 var t: i64 = 0
86 while t < ki { if sel[t] == e { taken = 1 } t = t + 1 }
87 if taken == 0 {
88 if best < 0 { best = e; bestv = r[e] } else { if r[e] > bestv { best = e; bestv = r[e] } }
89 }
90 e = e + 1
91 }
92 sel[ki] = best
93 ki = ki + 1
94 }
95 // softmax over the selected K (max-subtracted, Q16)
96 var m: i64 = r[sel[0]]
97 var i: i64 = 1
98 while i < K { if r[sel[i]] > m { m = r[sel[i]] } i = i + 1 }
99 var sum: i64 = 0
100 i = 0
101 while i < K { let ev: i64 = fx_exp(r[sel[i]] - m); w[i] = ev; sum = sum + ev; i = i + 1 }
102 if sum < 1 { sum = 1 }
103 i = 0
104 while i < K { w[i] = (w[i] << 16) / sum; i = i + 1 }
105 // K selected experts only (sparsity), weighted mix
106 var d: i64 = 0
107 while d < D { out[d] = 0; d = d + 1 }
108 i = 0
109 while i < K {
110 let e2: i64 = sel[i]
111 cnt[e2] = cnt[e2] + 1
112 mm_out_in(x, ((Wg as i64) + e2*H*D*8) as *i64, g, 1, D, H, 0)
113 mm_out_in(x, ((Wu as i64) + e2*H*D*8) as *i64, u, 1, D, H, 0)
114 var j: i64 = 0
115 while j < H { h[j] = qmul(silu(g[j]), u[j]); j = j + 1 }
116 mm_out_in(h, ((Wd as i64) + e2*D*H*8) as *i64, eo, 1, H, D, 0)
117 d = 0
118 while d < D { out[d] = out[d] + qmul(w[i], eo[d]); d = d + 1 }
119 i = i + 1
120 }
121 return 0
122}