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

↩ module page · 294 lines · 14996 B
1// nx_nofloat_olmoe_forward_gate.nx -- MoE rung 5 (2026-07-15): a COMPLETE OLMoE forward -- 16 stacked
2// real layers + token embedding + final norm + LM head -> a REAL next-token from the real 64-expert model,
3// in deterministic no-float. Lazy per-layer weight load (attn set into REUSED buffers; experts lazy) over
4// the zero-copy mmap. Composes every eaten-debt piece: full-quant dequant, exact scale, real rope base 1e4,
5// sys_map_file, nolmoe_attn/nolmoe_moe (rung 4). Teeth:
6//   T1 TOKENIZER round-trip: encode(prompt) -> decode == prompt (our BPE self-consistent on OLMoE's gpt2 vocab)
7//   T2 NON-DEGENERATE logits: argmax logit strictly > min logit (the head discriminates, not all-equal)
8//   T3 DETERMINISM: the full forward's predicted id + top logit byte-identical on repeat
9//   T4 ROUTING ALIVE AT DEPTH: >=2 distinct expert-masks seen across the 16 layers' last-token routing
10//   T5 PREDICTION real: a valid non-negative vocab id, decoded + printed (honest: OUR segmentation, uniform
11//      Q16 -- HF-canonical BPE + q24 precision + H2H = rung 6, deliberately not claimed here)
12// Requires the OLMoE gguf. Heavy (~1-3 min: 16 layers x lazy dequant + 50304-row head). Return from main.
13// No hw writes (Rule 26). expect_exit: 0   license_tier: ORIGINAL
14import "nx_syscalls.nx"
15import "nx_tier.nx"
16import "nx_le.nx"
17import "nx_tensor.nx"
18import "nx_gguf.nx"
19import "nx_gguf_load.nx"
20import "nx_gguf_meta.nx"
21import "nx_nofloat_llm.nx"
22import "nx_nofloat_tok.nx"
23import "nx_nofloat_arch.nx"
24import "nx_nofloat_olmoe.nx"
25import "nx_gate_verdict.nx"
26
27func of_w(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 }
28func of_n(v: i64) -> i64 {
  var m: i64 = v
  if m < 0 { of_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
40}
41func of_slen(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } return n }
42// load blk.L.<suffix> (full 2-D tensor) into out (Q16)
43func of_load_blk(buf: *u8, hdr: *NxGgufHeader, L: i64, suffix: *u8, out: *i64, nvals: i64) -> i64 {
  let nm: *u8 = sys_mmap(64)
  build_name(nm, L, suffix)
  let idx: nx_int = nx_gguf_find_tensor(hdr, nm, of_slen(nm))
  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)
51}
52func of_expbase_blk(buf: *u8, hdr: *NxGgufHeader, L: i64, suffix: *u8, eb: *i64, slot: i64) -> i64 {
  let nm: *u8 = sys_mmap(64)
  build_name(nm, L, suffix)
  let idx: nx_int = nx_gguf_find_tensor(hdr, nm, of_slen(nm))
  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
62}
63
64func main() -> i64 {
  of_w("=== NX-NOFLOAT-OLMOE-FORWARD -- full 16-layer OLMoE forward -> real next-token (no-float) ===\n" as *u8)
  let ln: *i64 = sys_mmap(8) as *i64
  ln[0] = 0
  let buf: *u8 = sys_map_file("/home/elderwesto/nx_stage/nx_moe_model.gguf" as *u8, ln)
  if (buf as i64) == 0 { of_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 { of_w("PARSE FAIL\n" as *u8); return 1 }
  let hloc: *NxGgufHeader = hdr
  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 { of_w("ARCH FAIL\n" as *u8); return 1 }
  let D: i64 = cfg[0]
  let NL: i64 = cfg[1]
  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
84
  // tokenizer meta
  let voff: *i64 = sys_mmap(8) as *i64
  let vty: *i64 = sys_mmap(8) as *i64
  var mfirst: i64 = 0
  var nm_c: i64 = 0
  var vfirst: i64 = 0
  var vocab: i64 = 0
  let km: *u8 = "tokenizer.ggml.merges" as *u8
  let kt: *u8 = "tokenizer.ggml.tokens" as *u8
  if nx_gguf_meta_find(buf, ln[0], hdr, km, of_slen(km), voff, vty) == NX_GMETA_OK { nm_c = nx_gguf_meta_array_count(buf, voff[0]); mfirst = nx_gguf_meta_array_first_elt_off(buf, voff[0]) }
  if nx_gguf_meta_find(buf, ln[0], hdr, kt, of_slen(kt), voff, vty) == NX_GMETA_OK { vocab = nx_gguf_meta_array_count(buf, voff[0]); vfirst = nx_gguf_meta_array_first_elt_off(buf, voff[0]) }
  of_w("[cfg] olmoe D=" as *u8); of_n(D); of_w(" layers=" as *u8); of_n(NL); of_w(" heads=" as *u8); of_n(nh); of_w(" ff=" as *u8); of_n(ff); of_w(" vocab=" as *u8); of_n(vocab); of_w(" base=" as *u8); of_n(cfg[9]); of_w("\n" as *u8)
97
  // ---- T1: tokenizer round-trip ----
  let prompt: *u8 = "The capital of France is" as *u8
  let plen: i64 = of_slen(prompt)
  let MAXT: i64 = 32
  let ids: *i64 = sys_mmap(MAXT*8) as *i64
  let tokp: *i64 = sys_mmap(512*8) as *i64
  let tokl: *i64 = sys_mmap(512*8) as *i64
  let np: i64 = tk_bpe_encode(buf, mfirst, nm_c, vfirst, vocab, prompt, plen, tokp, tokl, ids)
  of_w("[tok] '" as *u8); of_w(prompt); of_w("' -> " as *u8); of_n(np); of_w(" ids:" as *u8)
  var ti0: i64 = 0
  while ti0 < np { of_w(" " as *u8); of_n(ids[ti0]); ti0 = ti0 + 1 }
  of_w("\n" as *u8)
  // decode concat == prompt bytes?
  let rt: *u8 = sys_mmap(256)
  var rn: i64 = 0
  var di: i64 = 0
  while di < np {
      let off: i64 = tk_decode_off(buf, vfirst, ids[di])
      let pl: i64 = nx_gguf_meta_read_string_len(buf, off)
      let pp: *u8 = nx_gguf_meta_read_string_ptr(buf, off)
      var b: i64 = 0
      while b < pl { rt[rn] = pp[b]; rn = rn + 1; b = b + 1 }
      di = di + 1
  }
  // OLMoE gpt2 byte-level: a leading space becomes 0xC4 0xA0; compare the DECODED-BYTE-LEVEL form to the
  // byte-level-encoded prompt for a fair self-consistency check (encode's own byte map -> decode inverts it).
  var t1ok: i64 = 1
  // reconstruct expected byte-level string of the prompt
  let exp: *u8 = sys_mmap(256)
  var en: i64 = 0
  var pi: i64 = 0
  while pi < plen { en = en + tk_byte_char(prompt[pi] as i64, ((exp as i64) + en) as *u8); pi = pi + 1 }
  if rn != en { t1ok = 0 } else { var c: i64 = 0; while c < rn { if rt[c] != exp[c] { t1ok = 0 } c = c + 1 } }
131
  // ---- embed prompt ----
  let te_i: nx_int = nx_gguf_find_tensor(hdr, "token_embd.weight" as *u8, 17)
  let on_i: nx_int = nx_gguf_find_tensor(hdr, "output_norm.weight" as *u8, 18)
  let oh_i: nx_int = nx_gguf_find_tensor(hdr, "output.weight" as *u8, 13)
  if te_i < 0 { return 1 } if on_i < 0 { return 1 } if oh_i < 0 { return 1 }
  let te: *NxGgufTensorInfo = nx_gguf_tensor_at(hloc, te_i)
  let te_base: i64 = hloc.data_off + te.offset
  let te_ty: i64 = te.ggml_type
  let oh: *NxGgufTensorInfo = nx_gguf_tensor_at(hloc, oh_i)
  let oh_base: i64 = hloc.data_off + oh.offset
  let oh_ty: i64 = oh.ggml_type
  let gOut: *i64 = sys_mmap(D*8) as *i64
  let on: *NxGgufTensorInfo = nx_gguf_tensor_at(hloc, on_i)
  dequant_to_q16(buf, hloc.data_off + on.offset, on.ggml_type, D, gOut)
146
  let tmp: *i64 = sys_mmap(64*256*8) as *i64
  let x: *i64 = sys_mmap(MAXT*D*8) as *i64
  var ei: i64 = 0
  while ei < np { dequant_row(buf, te_base, te_ty, ids[ei], D, ((x as i64) + ei*D*8) as *i64, tmp); ei = ei + 1 }
151
  // ---- per-layer weight buffers (reused) + scratch ----
  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
  let eb: *i64 = sys_mmap(8*8) as *i64
  let freqs: *i64 = sys_mmap((hd/2)*8) as *i64
  rope_freqs_base(freqs, hd, cfg[11])
  let attnout: *i64 = sys_mmap(MAXT*D*8) as *i64
  let xmid: *i64 = sys_mmap(MAXT*D*8) as *i64
  let moeout: *i64 = sys_mmap(MAXT*D*8) as *i64
  let ascr: *i64 = sys_mmap((4*MAXT*D + 2*MAXT + MAXT*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 ap: *i64 = sys_mmap(16*8) as *i64
  let mp: *i64 = sys_mmap(16*8) as *i64
  let lastmask: *i64 = sys_mmap(MAXT*8) as *i64   // per-layer last-token expert mask
174
  let normed: *i64 = sys_mmap(D*8) as *i64
  let hrow: *i64 = sys_mmap(D*8) as *i64
  let pred: *i64 = sys_mmap(8) as *i64
  let toplg: *i64 = sys_mmap(8) as *i64
  let minlg: *i64 = sys_mmap(8) as *i64
180
  var distinct_masks: i64 = 0
  var run: i64 = 0
  var predA: i64 = 0
  var topA: i64 = 0
  while run < 2 {
      // fresh embed each run (x is overwritten by layers)
      ei = 0
      while ei < np { dequant_row(buf, te_base, te_ty, ids[ei], D, ((x as i64) + ei*D*8) as *i64, tmp); ei = ei + 1 }
      var seenmask: i64 = 0
      var L: i64 = 0
      while L < NL {
          of_load_blk(buf, hdr, L, ".attn_norm.weight\x00" as *u8, gA, D)
          of_load_blk(buf, hdr, L, ".attn_q_norm.weight\x00" as *u8, gQ, D)
          of_load_blk(buf, hdr, L, ".attn_k_norm.weight\x00" as *u8, gK, D)
          of_load_blk(buf, hdr, L, ".ffn_norm.weight\x00" as *u8, gF, D)
          of_load_blk(buf, hdr, L, ".attn_q.weight\x00" as *u8, Wq, D*D)
          of_load_blk(buf, hdr, L, ".attn_k.weight\x00" as *u8, Wk, D*D)
          of_load_blk(buf, hdr, L, ".attn_v.weight\x00" as *u8, Wv, D*D)
          of_load_blk(buf, hdr, L, ".attn_output.weight\x00" as *u8, Wo, D*D)
          of_load_blk(buf, hdr, L, ".ffn_gate_inp.weight\x00" as *u8, Wr, E*D)
          of_expbase_blk(buf, hdr, L, ".ffn_gate_exps.weight\x00" as *u8, eb, 0)
          of_expbase_blk(buf, hdr, L, ".ffn_up_exps.weight\x00" as *u8, eb, 2)
          of_expbase_blk(buf, hdr, L, ".ffn_down_exps.weight\x00" as *u8, eb, 4)
          ap[0]=x as i64; ap[1]=attnout as i64; ap[2]=np; 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 < np*D { xmid[d] = x[d] + attnout[d]; d = d + 1 }
          var cz: i64 = 0
          while cz < E { cnt[cz] = 0; cz = cz + 1 }
          mp[0]=xmid as i64; mp[1]=moeout as i64; mp[2]=np; 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 < np*D { x[d] = xmid[d] + moeout[d]; d = d + 1 }
          // record last-token routing mask for this layer (re-route just the last token)
          var cz2: i64 = 0
          while cz2 < E { cnt[cz2] = 0; cz2 = cz2 + 1 }
          mp[0]=((xmid as i64) + (np-1)*D*8); mp[1]=moeout as i64; 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 }
          if run == 0 { lastmask[L] = msk }
          L = L + 1
      }
      // final norm on the last token + streamed head argmax
      rmsnorm_gamma_row(x, gOut, (np-1)*D, D, normed, 0)
      var best: i64 = 0
      var bestv: i64 = 0 - 9223372036854775807
      var worst: i64 = 9223372036854775807
      var v: i64 = 0
      while v < vocab {
          dequant_row(buf, oh_base, oh_ty, v, D, hrow, tmp)
          var s: i64 = 0
          var kk: i64 = 0
          while kk < D { s = s + (normed[kk] * hrow[kk]); kk = kk + 1 }
          let lg: i64 = s >> 16
          if lg > bestv { bestv = lg; best = v }
          if lg < worst { worst = lg }
          v = v + 1
      }
      if run == 0 { predA = best; topA = bestv; toplg[0] = bestv; minlg[0] = worst; pred[0] = best }
      if run == 1 { pred[1] = best; toplg[1] = bestv }
      run = run + 1
  }
  // distinct masks across layers
  var L2: i64 = 0
  while L2 < NL {
      var seen: i64 = 0
      var L3: i64 = 0
      while L3 < L2 { if lastmask[L3] == lastmask[L2] { seen = 1 } L3 = L3 + 1 }
      if seen == 0 { distinct_masks = distinct_masks + 1 }
      L2 = L2 + 1
  }
257
  of_w("[fwd] predicted next id=" as *u8); of_n(predA); of_w(" piece='" as *u8)
  let po: i64 = tk_decode_off(buf, vfirst, predA)
  let ppl: i64 = nx_gguf_meta_read_string_len(buf, po)
  if ppl > 0 { sys_write(1, nx_gguf_meta_read_string_ptr(buf, po), ppl) }
  of_w("'  top_logit=" as *u8); of_n(topA); of_w(" min_logit=" as *u8); of_n(minlg[0]); of_w(" distinct-layer-masks=" as *u8); of_n(distinct_masks); of_w("\n" as *u8)
263
  var pass: i64 = 0
  var ttl: i64 = 0
  ttl = ttl + 1
  of_w("  T1 tokenizer round-trip (encode->decode == byte-level prompt): " as *u8)
  if t1ok == 1 { pass = pass + 1; of_w("PASS\n" as *u8) } else { of_w("FAIL\n" as *u8) }
  ttl = ttl + 1
  of_w("  T2 non-degenerate logits (top " as *u8); of_n(topA); of_w(" > min " as *u8); of_n(minlg[0]); of_w("): " as *u8)
  if topA > minlg[0] { pass = pass + 1; of_w("PASS\n" as *u8) } else { of_w("FAIL\n" as *u8) }
  ttl = ttl + 1
  var detok: i64 = 0
  if pred[0] == pred[1] { if toplg[0] == toplg[1] { detok = 1 } }
  of_w("  T3 full-forward determinism (pred+logit identical on repeat): " as *u8)
  if detok == 1 { pass = pass + 1; of_w("PASS\n" as *u8) } else { of_w("FAIL\n" as *u8) }
  ttl = ttl + 1
  of_w("  T4 routing alive at depth (distinct layer masks " as *u8); of_n(distinct_masks); of_w(">=2): " as *u8)
  if distinct_masks >= 2 { pass = pass + 1; of_w("PASS\n" as *u8) } else { of_w("FAIL\n" as *u8) }
  ttl = ttl + 1
  of_w("  T5 prediction is a valid vocab id: " as *u8)
  if predA >= 0 { if predA < vocab { pass = pass + 1; of_w("PASS\n" as *u8) } else { of_w("FAIL\n" as *u8) } } else { of_w("FAIL\n" as *u8) }
283
  of_w("NX-NOFLOAT-OLMOE-FORWARD-GATE passed " as *u8); of_n(pass); of_w("/" as *u8); of_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-FORWARD-GATE" as *u8, ctr__dry, "real 16-layer OLMoE forward -> next-token, deterministic no-float integer)" as *u8)
  sys_exit(rc__dry)
  return rc__dry
294}
nx_nofloat_olmoe_forward_gate.nx source

`nx_nofloat_olmoe_forward_gate.nx` source
