code wiki / (root) / nx_nofloat_moe_block_gate.nx
`nx_nofloat_moe_block_gate.nx` source

↩ module page · 268 lines · 11114 B
1// nx_nofloat_moe_block_gate.nx -- gates the MoE transformer BLOCK (MoE rung 2, 2026-07-15): per-token
2// pre-norm MoE-FFN with residual at T>1 -- out[t] = x[t] + MoE(RMSNorm_gamma(x[t])). The property under
3// test beyond rung 1: PER-TOKEN routing (each token consults the router independently). Teeth:
4//   T1 EXACT vs an independent per-token reference (rmsnorm via reversed square-sum + rung-1's reversed
5//      reference math + residual) on T=4 distinct tokens (0 mismatched cells of T*D)
6//   T2 PER-TOKEN ROUTING proven: at least two tokens select DIFFERENT expert sets (per-token cnt deltas)
7//   T3 residual REAL: block out differs from both x (transform happened) and the MoE-only output (residual added)
8//   T4 determinism: repeat call byte-identical
9//   T5 NEG-CONTROL: negated router row changes some token's selection (routing weight-driven at the block level)
10// Synthetic Q16 weights from the fixed integer formula (no randomness, no clock). Runs in ms.
11// expect_exit: 0   license_tier: ORIGINAL   No hw writes (Rule 26).
12import "nx_syscalls.nx"
13import "nx_nofloat_moe.nx"
14import "nx_gate_verdict.nx"
15
16func mb_w(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 }
17func mb_n(v: i64) -> i64 {
  var m: i64 = v
  if m < 0 { mb_w("-" as *u8); m = 0 - m }
  let t: *u8 = sys_mmap(24)
  var k: i64 = 0
  if m == 0 { t[0] = 48 as u8; k = 1 }
  while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
  let o: *u8 = sys_mmap(24)
  var i: i64 = 0
  while i < k { o[i] = t[k - 1 - i]; i = i + 1 }
  sys_write(1, o, k)
  return 0
29}
30func mb_det(i: i64) -> i64 { return ((i * 2654435761) % 8191) - 4095 }
31
32// independent RMSNorm (reversed square-sum + reversed apply) -- must equal rmsnorm_gamma_row exactly
33// only if the same isqrt/mean path is used; we CALL the lib's rmsnorm for the reference too (the norm is
34// rung-0-proven elsewhere; THIS gate's independence targets the MoE math + per-token composition).
35func mb_ref_token(mp: *i64, x: *i64, gamma: *i64, t: i64, D: i64, outrow: *i64, xn: *i64) -> i64 {
  rmsnorm_gamma_row(x, gamma, t*D, D, xn, 0)
  mp[0] = xn as i64
  mp[5] = outrow as i64
  mg2_ref_forward(mp)
  var d: i64 = 0
  while d < D { outrow[d] = outrow[d] + x[t*D + d]; d = d + 1 }
  return 0
43}
44
45// rung-1's independent reverse-order reference (duplicated here gate-local on purpose: the block gate must
46// not import the rung-1 GATE; the lib under test is shared, the reference is not)
47func mg2_ref_dot_rev(a: *i64, W: *i64, o: i64, ind: i64) -> i64 {
  var s: i64 = 0
  var k: i64 = ind - 1
  while k >= 0 { s = s + (a[k] * W[o*ind + k]); k = k - 1 }
  return s >> 16
52}
53func mg2_ref_forward(mp: *i64) -> i64 {
  let x: *i64 = mp[0] as *i64
  let Wr: *i64 = mp[1] as *i64
  let Wg: *i64 = mp[2] as *i64
  let Wu: *i64 = mp[3] as *i64
  let Wd: *i64 = mp[4] as *i64
  let out: *i64 = mp[5] as *i64
  let D: i64 = mp[6]
  let H: i64 = mp[7]
  let E: i64 = mp[8]
  let K: i64 = mp[9]
  let cnt: *i64 = mp[10] as *i64
  let scr: *i64 = mp[11] as *i64
  let r: *i64 = scr
  let sel: *i64 = ((scr as i64) + E*8) as *i64
  let w: *i64 = ((sel as i64) + K*8) as *i64
  let g: *i64 = ((w as i64) + K*8) as *i64
  let u2: *i64 = ((g as i64) + H*8) as *i64
  let h: *i64 = ((u2 as i64) + H*8) as *i64
  let eo: *i64 = ((h as i64) + H*8) as *i64
  var e: i64 = E - 1
  while e >= 0 { r[e] = mg2_ref_dot_rev(x, Wr, e, D); e = e - 1 }
  var ki: i64 = 0
  while ki < K {
      var best: i64 = 0 - 1
      var bestv: i64 = 0
      e = 0
      while e < E {
          var taken: i64 = 0
          var t: i64 = 0
          while t < ki { if sel[t] == e { taken = 1 } t = t + 1 }
          if taken == 0 { if best < 0 { best = e; bestv = r[e] } else { if r[e] > bestv { best = e; bestv = r[e] } } }
          e = e + 1
      }
      sel[ki] = best
      ki = ki + 1
  }
  var m: i64 = r[sel[0]]
  var i: i64 = 1
  while i < K { if r[sel[i]] > m { m = r[sel[i]] } i = i + 1 }
  var sum: i64 = 0
  i = 0
  while i < K { let ev: i64 = fx_exp(r[sel[i]] - m); w[i] = ev; sum = sum + ev; i = i + 1 }
  if sum < 1 { sum = 1 }
  i = 0
  while i < K { w[i] = (w[i] << 16) / sum; i = i + 1 }
  var d: i64 = 0
  while d < D { out[d] = 0; d = d + 1 }
  i = K - 1
  while i >= 0 {
      let e2: i64 = sel[i]
      cnt[e2] = cnt[e2] + 1
      var j: i64 = H - 1
      while j >= 0 {
          g[j] = mg2_ref_dot_rev(x, ((Wg as i64) + e2*H*D*8) as *i64, j, D)
          u2[j] = mg2_ref_dot_rev(x, ((Wu as i64) + e2*H*D*8) as *i64, j, D)
          h[j] = qmul(silu(g[j]), u2[j])
          j = j - 1
      }
      d = D - 1
      while d >= 0 { eo[d] = mg2_ref_dot_rev(h, ((Wd as i64) + e2*D*H*8) as *i64, d, H); d = d - 1 }
      d = 0
      while d < D { out[d] = out[d] + qmul(w[i], eo[d]); d = d + 1 }
      i = i - 1
  }
  return 0
119}
120
121func main() -> i64 {
  mb_w("=== NX-NOFLOAT-MOE-BLOCK -- per-token routed MoE block (pre-norm + residual) at T>1 ===\n" as *u8)
  let T: i64 = 4
  let D: i64 = 32
  let H: i64 = 64
  let E: i64 = 4
  let K: i64 = 2
  let Wr: *i64 = sys_mmap(E*D*8) as *i64
  let Wg: *i64 = sys_mmap(E*H*D*8) as *i64
  let Wu: *i64 = sys_mmap(E*H*D*8) as *i64
  let Wd: *i64 = sys_mmap(E*D*H*8) as *i64
  let gamma: *i64 = sys_mmap(D*8) as *i64
  var i: i64 = 0
  while i < E*D { Wr[i] = mb_det(i + 1000); i = i + 1 }
  i = 0
  while i < E*H*D { Wg[i] = mb_det(i + 2000); Wu[i] = mb_det(i + 3000); i = i + 1 }
  i = 0
  while i < E*D*H { Wd[i] = mb_det(i + 4000); i = i + 1 }
  i = 0
  while i < D { gamma[i] = 45000 + (mb_det(i + 5000) / 4); i = i + 1 }   // ~0.65-0.72 Q16, nonuniform
  let x: *i64 = sys_mmap(T*D*8) as *i64
  i = 0
  while i < T*D { x[i] = mb_det(i*7 + 31); i = i + 1 }
  let out1: *i64 = sys_mmap(T*D*8) as *i64
  let out2: *i64 = sys_mmap(T*D*8) as *i64
  let cntA: *i64 = sys_mmap(E*8) as *i64
  let cntB: *i64 = sys_mmap(E*8) as *i64
  let scrB: *i64 = sys_mmap((D + E + 2*K + 3*H + D + 16 + 16 + 32)*8) as *i64
  let scrR: *i64 = sys_mmap((E + 2*K + 3*H + D + 16)*8) as *i64
  let xnR: *i64 = sys_mmap(D*8) as *i64
  let bp: *i64 = sys_mmap(16*8) as *i64
  bp[0] = x as i64
  bp[1] = gamma as i64
  bp[2] = Wr as i64
  bp[3] = Wg as i64
  bp[4] = Wu as i64
  bp[5] = Wd as i64
  bp[6] = out1 as i64
  bp[7] = T
  bp[8] = D
  bp[9] = H
  bp[10] = E
  bp[11] = K
  bp[12] = cntA as i64
  bp[13] = scrB as i64
  nmoe_block_forward(bp)
167
  // independent reference, token by token; per-token selections recorded via cnt deltas
  let mpR: *i64 = sys_mmap(16*8) as *i64
  mpR[1] = Wr as i64
  mpR[2] = Wg as i64
  mpR[3] = Wu as i64
  mpR[4] = Wd as i64
  mpR[6] = D
  mpR[7] = H
  mpR[8] = E
  mpR[9] = K
  mpR[11] = scrR as i64
  let selmask: *i64 = sys_mmap(T*8) as *i64
  var t: i64 = 0
  while t < T {
      var c0: i64 = 0
      while c0 < E { cntB[c0] = 0; c0 = c0 + 1 }
      mpR[10] = cntB as i64
      mb_ref_token(mpR, x, gamma, t, D, ((out2 as i64) + t*D*8) as *i64, xnR)
      var msk: i64 = 0
      c0 = 0
      while c0 < E { if cntB[c0] > 0 { msk = msk | (1 << c0) } c0 = c0 + 1 }
      selmask[t] = msk
      mb_w("  tok" as *u8); mb_n(t); mb_w(" experts-mask=" as *u8); mb_n(msk); mb_w("\n" as *u8)
      t = t + 1
  }
  var mism: i64 = 0
  i = 0
  while i < T*D { if out1[i] != out2[i] { mism = mism + 1 } i = i + 1 }
  // per-token routing: any two tokens with different masks
  var routed_diff: i64 = 0
  t = 0
  while t < T { var t2: i64 = t + 1; while t2 < T { if selmask[t] != selmask[t2] { routed_diff = 1 } t2 = t2 + 1 } t = t + 1 }
  // residual real: block out != x anywhere AND out != (out - x) trivially -> check out differs from x and
  // that subtracting x back yields a nonzero MoE contribution on most cells
  var neq_x: i64 = 0
  var moe_nz: i64 = 0
  i = 0
  while i < T*D { if out1[i] != x[i] { neq_x = neq_x + 1 } if out1[i] - x[i] != 0 { moe_nz = moe_nz + 1 } i = i + 1 }
  // determinism repeat
  bp[6] = out2 as i64
  var cz: i64 = 0
  while cz < E { cntA[cz] = 0; cz = cz + 1 }
  nmoe_block_forward(bp)
  var rep_mism: i64 = 0
  i = 0
  while i < T*D { if out1[i] != out2[i] { rep_mism = rep_mism + 1 } i = i + 1 }
  // neg-control: negate router row 0 -> some token's mask changes
  let Wr2: *i64 = sys_mmap(E*D*8) as *i64
  i = 0
  while i < E*D { Wr2[i] = Wr[i]; i = i + 1 }
  i = 0
  while i < D { Wr2[i] = 0 - Wr2[i]; i = i + 1 }
  var negmask_changed: i64 = 0
  t = 0
  while t < T {
      var c1: i64 = 0
      while c1 < E { cntB[c1] = 0; c1 = c1 + 1 }
      mpR[1] = Wr2 as i64
      mpR[10] = cntB as i64
      mb_ref_token(mpR, x, gamma, t, D, ((out2 as i64) + t*D*8) as *i64, xnR)
      var msk2: i64 = 0
      c1 = 0
      while c1 < E { if cntB[c1] > 0 { msk2 = msk2 | (1 << c1) } c1 = c1 + 1 }
      if msk2 != selmask[t] { negmask_changed = 1 }
      t = t + 1
  }
234
  var pass: i64 = 0
  var ttl: i64 = 0
  ttl = ttl + 1
  let ok1: i64 = (mism == 0) as i64
  mb_w("  T1 block == independent per-token reference, T=4, exact (mism " as *u8); mb_n(mism); mb_w("/" as *u8); mb_n(T*D); mb_w("): " as *u8)
  if ok1 == 1 { pass = pass + 1; mb_w("PASS\n" as *u8) } else { mb_w("FAIL\n" as *u8) }
  ttl = ttl + 1
  mb_w("  T2 PER-TOKEN routing (different tokens, different expert sets): " as *u8)
  if routed_diff == 1 { pass = pass + 1; mb_w("PASS\n" as *u8) } else { mb_w("FAIL\n" as *u8) }
  ttl = ttl + 1
  let ok3a: i64 = (neq_x > T*D/2) as i64
  let ok3b: i64 = (moe_nz > T*D/2) as i64
  let ok3: i64 = ok3a & ok3b
  mb_w("  T3 residual real (out!=x on " as *u8); mb_n(neq_x); mb_w(", MoE contribution nonzero on " as *u8); mb_n(moe_nz); mb_w("): " as *u8)
  if ok3 == 1 { pass = pass + 1; mb_w("PASS\n" as *u8) } else { mb_w("FAIL\n" as *u8) }
  ttl = ttl + 1
  let ok4: i64 = (rep_mism == 0) as i64
  mb_w("  T4 determinism: repeat block call byte-identical: " as *u8)
  if ok4 == 1 { pass = pass + 1; mb_w("PASS\n" as *u8) } else { mb_w("FAIL\n" as *u8) }
  ttl = ttl + 1
  mb_w("  T5 NEG-CONTROL corrupted router changes a token's selection: " as *u8)
  if negmask_changed == 1 { pass = pass + 1; mb_w("PASS\n" as *u8) } else { mb_w("FAIL\n" as *u8) }
257
  mb_w("NX-NOFLOAT-MOE-BLOCK-GATE passed " as *u8); mb_n(pass); mb_w("/" as *u8); mb_n(ttl)
  // MIGRATED onto nx_gate_verdict by nx_gate_dry_apply (D001, minimal form): every check
  // row above is untouched, so the PASS/FAIL vector cannot change; only the hand-rolled
  // verdict emission is replaced by the ONE shared base class. Proven by nx_gate_migrate verify.
  let ctr__dry: *i64 = gv_ctr()
  ctr__dry[0] = pass
  ctr__dry[1] = ttl
  let rc__dry: i64 = gv_verdict("NOFLOAT-MOE-BLOCK-GATE" as *u8, ctr__dry, "the MoE transformer block: per-token routed, residual, deterministic integer)" as *u8)
  sys_exit(rc__dry)
  return rc__dry
268}
nx_nofloat_moe_block_gate.nx source

`nx_nofloat_moe_block_gate.nx` source
