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

↩ module page · 222 lines · 10758 B
1// nx_nofloat_olmoe_gate.nx -- MoE rung 4 (2026-07-15): a FAITHFUL OLMoE transformer LAYER on REAL blk.0
2// weights, end to end in deterministic no-float. Attention (pre-norm -> QKV -> QK-norm over full 2048 ->
3// RoPE base 1e4 -> 16-head causal MHA -> Wo -> residual) then MoE-FFN (ffn_norm -> lazy top-8 SwiGLU mix ->
4// residual). T=2 tokens (proves per-token routing + causal attention). Teeth:
5//   T1 DETERMINISM: whole-layer output byte-identical on repeat
6//   T2 REAL TRANSFORMATION: layer output differs from input on >half the cells (both tokens)
7//   T3 PER-TOKEN ROUTING: the two tokens select DIFFERENT expert sets (real router, real per-token MoE)
8//   T4 QK-NORM MATTERS (neg-control): skipping q_norm (gamma=1) changes the output -- proves QK-norm is live
9//   T5 LAZY FETCH: exactly K=8 experts dequanted per token (cnt sum == T*K), attention output nonzero
10// Requires /home/elderwesto/nx_stage/nx_moe_model.gguf. ~4.2GB mapped (zero-copy). Return from main.
11// No hw writes (Rule 26). expect_exit: 0   license_tier: ORIGINAL
12import "nx_syscalls.nx"
13import "nx_tier.nx"
14import "nx_le.nx"
15import "nx_tensor.nx"
16import "nx_gguf.nx"
17import "nx_gguf_load.nx"
18import "nx_gguf_meta.nx"
19import "nx_nofloat_llm.nx"
20import "nx_nofloat_arch.nx"
21import "nx_nofloat_olmoe.nx"
22import "nx_gate_verdict.nx"
23import "nx_stage_path.nx"
24
25func og_w(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 }
26func og_n(v: i64) -> i64 {
  var m: i64 = v
  if m < 0 { og_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
38}
39func og_slen(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } return n }
40func og_det(i: i64) -> i64 { return ((i * 2654435761) % 8191) - 4095 }
41// dequant a full 2-D weight tensor by name into an i64 buffer (attn Q/K/V/O are small enough to materialize).
42func og_load(buf: *u8, hdr: *NxGgufHeader, name: *u8, out: *i64, nvals: i64) -> i64 {
  let idx: nx_int = nx_gguf_find_tensor(hdr, name, og_slen(name))
  if idx < 0 { return 0 - 1 }
  let ti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, idx)
  let hl: *NxGgufHeader = hdr
  return dequant_to_q16(buf, hl.data_off + ti.offset, ti.ggml_type, nvals, out)
48}
49func og_expbase(buf: *u8, hdr: *NxGgufHeader, name: *u8, eb: *i64, slot: i64) -> i64 {
  let idx: nx_int = nx_gguf_find_tensor(hdr, name, og_slen(name))
  if idx < 0 { return 0 - 1 }
  let ti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, idx)
  let hl: *NxGgufHeader = hdr
  eb[slot] = hl.data_off + ti.offset
  eb[slot + 1] = ti.ggml_type
  return 0
57}
58
59func main() -> i64 {
  og_w("=== NX-NOFLOAT-OLMOE -- faithful OLMoE layer (QK-norm attn + lazy MoE-FFN) on real blk.0 ===\n" as *u8)
  let ln: *i64 = sys_mmap(8) as *i64
  ln[0] = 0
  let buf: *u8 = sys_map_file(sp_guarded("NOFLOAT-OLMOE-GATE" as *u8, "nx_moe_model.gguf" as *u8), ln)
  if (buf as i64) == 0 { og_w("MODEL ABSENT\n" as *u8); return 1 }
  let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
  if nx_gguf_parse(buf, ln[0], hdr) != NX_GGUF_OK { og_w("PARSE FAIL\n" as *u8); return 1 }
  let cfg: *i64 = sys_mmap(16*8) as *i64
  let arch: *u8 = sys_mmap(48)
  if nac_read_config(buf, ln[0], hdr, cfg, arch) != 0 { og_w("ARCH FAIL\n" as *u8); return 1 }
  let D: i64 = cfg[0]
  let nh: i64 = cfg[2]
  let hd: i64 = cfg[4]
  let scale: i64 = cfg[8]
  let ff: i64 = cfg[7]
  let E: i64 = 64
  let K: i64 = 8
  og_w("[cfg] D=" as *u8); og_n(D); og_w(" heads=" as *u8); og_n(nh); og_w(" hd=" as *u8); og_n(hd)
  og_w(" scale=" as *u8); og_n(scale); og_w(" ff=" as *u8); og_n(ff); og_w(" base=" as *u8); og_n(cfg[9]); og_w("\n" as *u8)
79
  // per-layer weights (materialize the small attn set; experts stay lazy)
  let gA: *i64 = sys_mmap(D*8) as *i64
  let gQ: *i64 = sys_mmap(D*8) as *i64
  let gK: *i64 = sys_mmap(D*8) as *i64
  let gF: *i64 = sys_mmap(D*8) as *i64
  let Wq: *i64 = sys_mmap(D*D*8) as *i64
  let Wk: *i64 = sys_mmap(D*D*8) as *i64
  let Wv: *i64 = sys_mmap(D*D*8) as *i64
  let Wo: *i64 = sys_mmap(D*D*8) as *i64
  let Wr: *i64 = sys_mmap(E*D*8) as *i64
  og_load(buf, hdr, "blk.0.attn_norm.weight" as *u8, gA, D)
  og_load(buf, hdr, "blk.0.attn_q_norm.weight" as *u8, gQ, D)
  og_load(buf, hdr, "blk.0.attn_k_norm.weight" as *u8, gK, D)
  og_load(buf, hdr, "blk.0.ffn_norm.weight" as *u8, gF, D)
  og_load(buf, hdr, "blk.0.attn_q.weight" as *u8, Wq, D*D)
  og_load(buf, hdr, "blk.0.attn_k.weight" as *u8, Wk, D*D)
  og_load(buf, hdr, "blk.0.attn_v.weight" as *u8, Wv, D*D)
  og_load(buf, hdr, "blk.0.attn_output.weight" as *u8, Wo, D*D)
  og_load(buf, hdr, "blk.0.ffn_gate_inp.weight" as *u8, Wr, E*D)
  let eb: *i64 = sys_mmap(8*8) as *i64
  og_expbase(buf, hdr, "blk.0.ffn_gate_exps.weight" as *u8, eb, 0)
  og_expbase(buf, hdr, "blk.0.ffn_up_exps.weight" as *u8, eb, 2)
  og_expbase(buf, hdr, "blk.0.ffn_down_exps.weight" as *u8, eb, 4)
103
  let freqs: *i64 = sys_mmap((hd/2)*8) as *i64
  rope_freqs_base(freqs, hd, cfg[11])           // base 1e4 rope (cfg[11]=ln(base))
106
  let T: i64 = 2
  let x: *i64 = sys_mmap(T*D*8) as *i64
  var i: i64 = 0
  while i < T*D { x[i] = og_det(i*7 + 13); i = i + 1 }
111
  let attnout: *i64 = sys_mmap(T*D*8) as *i64
  let xmid: *i64 = sys_mmap(T*D*8) as *i64
  let moeout: *i64 = sys_mmap(T*D*8) as *i64
  let out1: *i64 = sys_mmap(T*D*8) as *i64
  let out2: *i64 = sys_mmap(T*D*8) as *i64
  let ascr: *i64 = sys_mmap((4*T*D + 2*T + T*D + 64)*8) as *i64
  let mscr: *i64 = sys_mmap((D + E + 2*K + 2*ff*D + D*ff + 3*ff + D + 64)*8) as *i64
  let cnt: *i64 = sys_mmap(E*8) as *i64
  let selmask: *i64 = sys_mmap(T*8) as *i64
121
  let ap: *i64 = sys_mmap(16*8) as *i64
  let mp: *i64 = sys_mmap(16*8) as *i64
124
  // ---- run the full layer (attn + residual, moe + residual) into `outp` ----
  // (inlined twice for determinism; second run into out2)
  var run: i64 = 0
  while run < 2 {
      var outp: *i64 = out1
      if run == 1 { outp = out2 }
      var cz: i64 = 0
      while cz < E { cnt[cz] = 0; cz = cz + 1 }
      ap[0]=x as i64; ap[1]=attnout as i64; ap[2]=T; ap[3]=D; ap[4]=nh; ap[5]=hd; ap[6]=scale
      ap[7]=gA as i64; ap[8]=Wq as i64; ap[9]=Wk as i64; ap[10]=Wv as i64; ap[11]=Wo as i64
      ap[12]=gQ as i64; ap[13]=gK as i64; ap[14]=freqs as i64; ap[15]=ascr as i64
      nolmoe_attn(ap)
      var d: i64 = 0
      while d < T*D { xmid[d] = x[d] + attnout[d]; d = d + 1 }
      mp[0]=xmid as i64; mp[1]=moeout as i64; mp[2]=T; mp[3]=D; mp[4]=ff; mp[5]=E; mp[6]=K
      mp[7]=gF as i64; mp[8]=Wr as i64; mp[9]=buf as i64; mp[10]=eb as i64; mp[11]=mscr as i64; mp[12]=cnt as i64
      nolmoe_moe(mp)
      d = 0
      while d < T*D { outp[d] = xmid[d] + moeout[d]; d = d + 1 }
      run = run + 1
  }
  // record per-token expert masks from a fresh routed pass
  var mt: i64 = 0
  while mt < T {
      var cz2: i64 = 0
      while cz2 < E { cnt[cz2] = 0; cz2 = cz2 + 1 }
      mp[0]=((xmid as i64) + mt*D*8); mp[1]=((moeout as i64) + mt*D*8); mp[2]=1
      nolmoe_moe(mp)
      var msk: i64 = 0
      cz2 = 0
      while cz2 < E { if cnt[cz2] > 0 { msk = msk | (1 << cz2) } cz2 = cz2 + 1 }
      selmask[mt] = msk
      og_w("  tok" as *u8); og_n(mt); og_w(" expert-mask=" as *u8); og_n(msk); og_w("\n" as *u8)
      mt = mt + 1
  }
  mp[2] = T
161
  // ---- teeth ----
  var pass: i64 = 0
  var ttl: i64 = 0
  ttl = ttl + 1
  var repm: i64 = 0
  i = 0
  while i < T*D { if out1[i] != out2[i] { repm = repm + 1 } i = i + 1 }
  og_w("  T1 whole-layer determinism (repeat byte-identical, mism " as *u8); og_n(repm); og_w("): " as *u8)
  if repm == 0 { pass = pass + 1; og_w("PASS\n" as *u8) } else { og_w("FAIL\n" as *u8) }
  ttl = ttl + 1
  var neq: i64 = 0
  i = 0
  while i < T*D { if out1[i] != x[i] { neq = neq + 1 } i = i + 1 }
  og_w("  T2 real transformation (out!=in on " as *u8); og_n(neq); og_w("/" as *u8); og_n(T*D); og_w("): " as *u8)
  if neq > T*D/2 { pass = pass + 1; og_w("PASS\n" as *u8) } else { og_w("FAIL\n" as *u8) }
  ttl = ttl + 1
  og_w("  T3 per-token routing (tok0 mask " as *u8); og_n(selmask[0]); og_w(" != tok1 mask " as *u8); og_n(selmask[1]); og_w("): " as *u8)
  if selmask[0] != selmask[1] { pass = pass + 1; og_w("PASS\n" as *u8) } else { og_w("FAIL\n" as *u8) }
  ttl = ttl + 1
  // T4 QK-norm neg-control: run attention with gQ = all-ones (Q16) -> attn output must differ
  let ones: *i64 = sys_mmap(D*8) as *i64
  i = 0
  while i < D { ones[i] = 65536; i = i + 1 }
  let attn2: *i64 = sys_mmap(T*D*8) as *i64
  ap[1]=attn2 as i64
  ap[12]=ones as i64
  nolmoe_attn(ap)
  ap[1]=attnout as i64
  ap[12]=gQ as i64
  var qkdiff: i64 = 0
  i = 0
  while i < T*D { if attn2[i] != attnout[i] { qkdiff = qkdiff + 1 } i = i + 1 }
  og_w("  T4 QK-norm matters (skip q_norm -> attn differs on " as *u8); og_n(qkdiff); og_w(" cells): " as *u8)
  if qkdiff > 0 { pass = pass + 1; og_w("PASS\n" as *u8) } else { og_w("FAIL\n" as *u8) }
  ttl = ttl + 1
  // T5 lazy fetch count == T*K on a fresh full pass + attn nonzero
  var cz3: i64 = 0
  while cz3 < E { cnt[cz3] = 0; cz3 = cz3 + 1 }
  mp[0]=xmid as i64; mp[1]=moeout as i64; mp[2]=T
  nolmoe_moe(mp)
  var fetched: i64 = 0
  cz3 = 0
  while cz3 < E { fetched = fetched + cnt[cz3]; cz3 = cz3 + 1 }
  var anz: i64 = 0
  i = 0
  while i < T*D { if attnout[i] != 0 { anz = anz + 1 } i = i + 1 }
  og_w("  T5 lazy fetch total=" as *u8); og_n(fetched); og_w(" (want " as *u8); og_n(T*K); og_w("), attn nonzero " as *u8); og_n(anz); og_w(": " as *u8)
  var ok5: i64 = ((fetched == T*K) as i64) & ((anz > T*D/2) as i64)
  if ok5 == 1 { pass = pass + 1; og_w("PASS\n" as *u8) } else { og_w("FAIL\n" as *u8) }
211
  og_w("NX-NOFLOAT-OLMOE-GATE passed " as *u8); og_n(pass); og_w("/" as *u8); og_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-OLMOE-GATE" as *u8, ctr__dry, "real OLMoE layer: QK-norm attention + per-token lazy MoE, deterministic integer)" as *u8)
  sys_exit(rc__dry)
  return rc__dry
222}
nx_nofloat_olmoe_gate.nx source

`nx_nofloat_olmoe_gate.nx` source
