nx_moe.nx source
↩ module page · 202 lines · 7429 B
1// nx_moe.nx -- sovereign MIXTURE-OF-EXPERTS sparse FFN (the last census
2// OPPORTUNITY axis "Mixture of experts (sparse)", momentum 3; Mixtral /
3// Switch-class, Shazeer 2017 / Fedus 2021 / Jiang 2024).
4//
5// Per token: router logits = x @ W_router^T -> softmax -> TOP-K experts
6// -> renormalize the k probs -> out = sum_k p_i * SwiGLU_FFN_expert_i(x).
7// Only k of n_experts FFNs run per token (the sparsity that lets total
8// parameters scale past per-token compute). Experts accumulate in
9// SELECTION ORDER (prob-descending; ties -> lowest index) --
10// deterministic, documented, and mirrored by the gate's oracle.
11//
12// Composes the EXISTING primitives: nx_f32_matmul_t (weights transposed,
13// the LLM convention), nx_f32_softmax, nx_f32_silu. Weights here are
14// plain F32 buffers; binding real MoE GGUF expert tensors (and the lazy/
15// quantized paths) is the documented follow-on once an MoE model is on
16// disk. Touch counters per expert make the sparsity MEASURABLE (gated),
17// not asserted.
18//
19// genealogy_id: shazeer_2017_moe + fedus_2021_switch + jiang_2024_mixtral
20// lineage_id: substrate_moe_v1
21
22import "nx_syscalls.nx"
23import "nx_tier.nx"
24import "nx_f32.nx"
25import "nx_f32_div.nx"
26import "nx_f32_softmax.nx"
27import "nx_f32_matmul_t.nx"
28import "nx_f32_activations.nx"
29
30const NX_MOE_OK: nx_int = 0
31const NX_MOE_ERR_NULL: nx_int = 1
32const NX_MOE_ERR_BAD_DIM: nx_int = 2
33
34// ===== Layer =========================================================
35// experts_* are POINTER ARRAYS (one *i64 weight buffer per expert, stored
36// as i64): gate/up = [ffn_dim, hidden_dim] (transposed), down =
37// [hidden_dim, ffn_dim] (transposed). W_router = [n_experts, hidden_dim].
38
39struct NxMoeLayer {
40 n_experts: nx_int,
41 top_k: nx_int,
42 hidden_dim: nx_int,
43 ffn_dim: nx_int,
44 W_router: *i64,
45 experts_gate: *i64,
46 experts_up: *i64,
47 experts_down: *i64,
48 touch: *i64 // per-expert FFN-run counters (sparsity gate)
49}
50
51const NX_MOE_LAYER_BYTES: nx_int = 72
52
53func nx_moe_layer_alloc(n_experts: nx_int, top_k: nx_int,
54 hidden_dim: nx_int, ffn_dim: nx_int) -> *NxMoeLayer {
55 if n_experts <= 0 { return 0 as *NxMoeLayer }
56 if top_k <= 0 { return 0 as *NxMoeLayer }
57 if top_k > n_experts { return 0 as *NxMoeLayer }
58 if hidden_dim <= 0 { return 0 as *NxMoeLayer }
59 if ffn_dim <= 0 { return 0 as *NxMoeLayer }
60 let L: *NxMoeLayer = sys_mmap(NX_MOE_LAYER_BYTES) as *NxMoeLayer
61 L.n_experts = n_experts
62 L.top_k = top_k
63 L.hidden_dim = hidden_dim
64 L.ffn_dim = ffn_dim
65 L.W_router = sys_mmap(n_experts * hidden_dim * 8) as *i64
66 L.experts_gate = sys_mmap(n_experts * 8) as *i64
67 L.experts_up = sys_mmap(n_experts * 8) as *i64
68 L.experts_down = sys_mmap(n_experts * 8) as *i64
69 L.touch = sys_mmap(n_experts * 8) as *i64
70 var e: nx_int = 0
71 while e < n_experts {
72 let wg: *i64 = sys_mmap(ffn_dim * hidden_dim * 8) as *i64
73 let wu: *i64 = sys_mmap(ffn_dim * hidden_dim * 8) as *i64
74 let wd: *i64 = sys_mmap(hidden_dim * ffn_dim * 8) as *i64
75 L.experts_gate[e] = wg as i64
76 L.experts_up[e] = wu as i64
77 L.experts_down[e] = wd as i64
78 e = e + 1
79 }
80 return L
81}
82
83// ===== Routing =======================================================
84// logits -> softmax -> top-k (prob-descending; ties -> lowest index) ->
85// renormalize. out_idx[k], out_p[k] (f32 bits).
86
87func nx_moe_route(L: *NxMoeLayer, x_row: *i64,
88 out_idx: *i64, out_p: *i64) -> nx_int {
89 if L == (0 as *NxMoeLayer) { return NX_MOE_ERR_NULL }
90 let ne: nx_int = L.n_experts
91 let logits: *i64 = sys_mmap(ne * 8) as *i64
92 nx_f32_matmul_t(x_row, L.W_router, logits, 1, L.hidden_dim, ne)
93 let probs: *i64 = sys_mmap(ne * 8) as *i64
94 nx_f32_softmax(logits, ne, probs)
95
96 // top-k by repeated argmax-with-exclusion (ne is small).
97 let used: *i64 = sys_mmap(ne * 8) as *i64
98 var kk: nx_int = 0
99 while kk < L.top_k {
100 var best: nx_int = 0 - 1
101 var bestv: i64 = 0
102 var e: nx_int = 0
103 while e < ne {
104 if used[e] == 0 {
105 var take: nx_int = 0
106 if best < 0 { take = 1 } else {
107 if nx_f32_gt(probs[e], bestv) == 1 { take = 1 }
108 }
109 if take == 1 { best = e; bestv = probs[e] }
110 }
111 e = e + 1
112 }
113 used[best] = 1
114 out_idx[kk] = best as i64
115 out_p[kk] = bestv
116 kk = kk + 1
117 }
118
119 // renormalize the selected probs to sum 1.
120 var psum: i64 = 0
121 var s: nx_int = 0
122 while s < L.top_k { psum = nx_f32_add(psum, out_p[s]); s = s + 1 }
123 var s2: nx_int = 0
124 while s2 < L.top_k {
125 let pr: i64 = nx_f32_div(out_p[s2], psum)
126 out_p[s2] = pr
127 s2 = s2 + 1
128 }
129 return NX_MOE_OK
130}
131
132// ===== One expert's SwiGLU FFN (exported: the gate's oracle reuses it) ==
133// y[hidden] = ( silu(x@Wg^T) * (x@Wu^T) ) @ Wd^T
134
135func nx_moe_expert_ffn(L: *NxMoeLayer, e: nx_int, x_row: *i64,
136 y_out: *i64) -> nx_int {
137 let fd: nx_int = L.ffn_dim
138 let hd: nx_int = L.hidden_dim
139 let wg: *i64 = L.experts_gate[e] as *i64
140 let wu: *i64 = L.experts_up[e] as *i64
141 let wd: *i64 = L.experts_down[e] as *i64
142 let gate_raw: *i64 = sys_mmap(fd * 8) as *i64
143 let up_raw: *i64 = sys_mmap(fd * 8) as *i64
144 let hidden: *i64 = sys_mmap(fd * 8) as *i64
145 nx_f32_matmul_t(x_row, wg, gate_raw, 1, hd, fd)
146 nx_f32_matmul_t(x_row, wu, up_raw, 1, hd, fd)
147 var j: nx_int = 0
148 while j < fd {
149 let sg: i64 = nx_f32_silu(gate_raw[j])
150 let hv: i64 = nx_f32_mul(sg, up_raw[j])
151 hidden[j] = hv
152 j = j + 1
153 }
154 nx_f32_matmul_t(hidden, wd, y_out, 1, fd, hd)
155 sys_munmap(gate_raw, fd * 8)
156 sys_munmap(up_raw, fd * 8)
157 sys_munmap(hidden, fd * 8)
158 L.touch[e] = L.touch[e] + 1
159 return NX_MOE_OK
160}
161
162// ===== MoE forward: [m, hidden] -> [m, hidden] ======================
163
164func nx_moe_forward(L: *NxMoeLayer, X: *i64, m: nx_int,
165 out: *i64) -> nx_int {
166 if L == (0 as *NxMoeLayer) { return NX_MOE_ERR_NULL }
167 if X == (0 as *i64) { return NX_MOE_ERR_NULL }
168 if out == (0 as *i64) { return NX_MOE_ERR_NULL }
169 if m <= 0 { return NX_MOE_ERR_BAD_DIM }
170 let hd: nx_int = L.hidden_dim
171 let idx: *i64 = sys_mmap(L.top_k * 8) as *i64
172 let prb: *i64 = sys_mmap(L.top_k * 8) as *i64
173 let y: *i64 = sys_mmap(hd * 8) as *i64
174
175 var t: nx_int = 0
176 while t < m {
177 let x_row: *i64 = ((X as i64) + t * hd * 8) as *i64
178 let o_row: *i64 = ((out as i64) + t * hd * 8) as *i64
179 let vr: nx_int = nx_moe_route(L, x_row, idx, prb)
180 if vr != NX_MOE_OK { return vr }
181 var d0: nx_int = 0
182 while d0 < hd { o_row[d0] = 0; d0 = d0 + 1 }
183 var kk: nx_int = 0
184 while kk < L.top_k {
185 let e: nx_int = idx[kk] as nx_int
186 nx_moe_expert_ffn(L, e, x_row, y)
187 var d: nx_int = 0
188 while d < hd {
189 let contrib: i64 = nx_f32_mul(prb[kk], y[d])
190 let acc: i64 = nx_f32_add(o_row[d], contrib)
191 o_row[d] = acc
192 d = d + 1
193 }
194 kk = kk + 1
195 }
196 t = t + 1
197 }
198 sys_munmap(idx, L.top_k * 8)
199 sys_munmap(prb, L.top_k * 8)
200 sys_munmap(y, hd * 8)
201 return NX_MOE_OK
202}