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

↩ module page · 708 lines · 33570 B
1// nx_reader_squad_f32_gate.nx -- R3e of the NEURAL PASSAGE READER arc: TRAIN-ON-F32 (quantize-after comes
2// next). The Q16 lane closed honestly: 6 configurations all entered the uniform-collapse absorbing state at
3// E=24 (see nx_reader_squad_gate). f32 has no quantization floor and no bit-identical-states trap, so the
4// SAME architecture + data + recipe should train. This gate = the Q16 gate ported to the ta_* f32 tensor
5// autograd (transformer ops added 2026-07-09, identities verbatim from the GREEN nfa_* lib).
6//   T1 reader held-out span-F1 > question-blind ablation + 100 (liar-kill)
7//   T2 reader held-out exact-start > 60 permille (chance ~11)
8//   T3 train loss falls > 40% (real learning, no collapse)
9//   T4 deterministic re-eval
10// Consumes: reader_rows.bin ('dt'), semppmi_v1.bin (vocab), embed_v1.bin (SGNS Q10 -> f32 RAW magnitude).
11// expect_exit: 0   license_tier: ORIGINAL   Sovereign: nx_autograd_tensor (f32) + nx_syscalls.
12import "nx_autograd_tensor.nx"
13import "nx_thread_pool.nx"        // RF_PAR: data-parallel batch training across the 16-worker pool
14import "nx_hw.nx"
15import "nx_syscalls.nx"
16
17func rf_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
18func rf_pn(v: i64) -> i64 { let b: *u8=sys_mmap(28); var x: i64=v; if x<0{b[0]=45;sys_write(1,b,1);x=0-x} if x==0{b[0]=48;sys_write(1,b,1);return 0} var d: i64=0; var y: i64=x; while y>0{d=d+1;y=y/10} var i: i64=d-1; y=x; while i>=0{b[i]=(48+(y%10)) as u8;y=y/10;i=i-1} sys_write(1,b,d); return 0 }
19func rf_ck(name: *u8, c: i64) -> i64 { if c==1 { rf_puts("  PASS " as *u8) } else { rf_puts("  FAIL " as *u8) } rf_puts(name); rf_puts("\n" as *u8); return c }
20
21const RF_E: i64 = 24
22const RF_F: i64 = 48
23const RF_MAXT: i64 = 96
24const RF_MAXQ: i64 = 20
25const RF_MAXG: i64 = 8
26// RF_XL flips the whole gate between the small (998-row reader_rows.bin) and the XL (4771-row
27// reader_rows_xl.bin) datasets -- ONE flag, no duplication. RF_NROWS sizes the row arena for either.
28const RF_XL: i64 = 1
29// RF_PAR: run the batch of 8 examples DATA-PARALLEL on the thread pool (per-example private arenas; grads
30// reduced in mb-order => BIT-IDENTICAL to serial => same loss curve, verified vs the serial 8353 reference).
31// Proven correct in nx_ta_parallel_grad_gate. ~5x faster to the verdict on the soft-float bottleneck.
32const RF_PAR: i64 = 1
33const RF_BATCH: i64 = 8
34// RF_CAUSAL: attention masking. 0 = BIDIRECTIONAL (BERT-style), 1 = causal (autoregressive, from the nano-LM).
35// THEORY TESTED + REFUTED (2026-07-10): I hypothesized causal masking was the bug behind the +9-over-ablation
36// gap. Retrained bidirectional (RF_CAUSAL=0): F1=79 vs causal F1=81 -- a WASH (noise-level, slightly WORSE).
37// => masking is NOT the bug; at E=24/1-layer/1-head with static hash embeddings, attention can't extract much
38// question signal regardless of masking. The gap is a CAPACITY/PARADIGM wall, not a wiring error. Kept =1
39// (the marginally-better measured baseline). Real fix = pretrained model (sovereign Qwen), not from-scratch tweaks.
40const RF_CAUSAL: i64 = 1
41const RF_NROWS: i64 = 4800
42const RF_MAXV: i64 = 16384
43const RF_MAPN: i64 = 32768
44const FOFF_EMB: i64 = 0
45const FOFF_WQ: i64 = 393216
46const FOFF_WK: i64 = 393792
47const FOFF_WV: i64 = 394368
48const FOFF_WO: i64 = 394944
49const FOFF_WG: i64 = 395520
50const FOFF_WU: i64 = 396672
51const FOFF_WD: i64 = 397824
52const FOFF_US: i64 = 398976
53const FOFF_UE: i64 = 399000
54const RF_NP: i64  = 399024
55
56// dataset path helpers -- RF_XL selects small vs XL; used by main + the checkpoint fns (no arg-threading).
57func rf_rows_path() -> *u8 { if RF_XL == 1 { return "knowledge/index/reader_rows_xl.bin" as *u8 } return "knowledge/index/reader_rows.bin" as *u8 }
58func rf_ck_path() -> *u8 { if RF_XL == 1 { return "knowledge/index/rsqf32xl_ckpt.bin" as *u8 } return "knowledge/index/rsqf32_ckpt.bin" as *u8 }
59func rf_ck_part_path() -> *u8 { if RF_XL == 1 { return "knowledge/index/rsqf32xl_ckpt.bin.part" as *u8 } return "knowledge/index/rsqf32_ckpt.bin.part" as *u8 }
60func rf_persist_path() -> *u8 { if RF_XL == 1 { return "knowledge/index/reader_neural_f32xl.bin" as *u8 } return "knowledge/index/reader_neural_f32.bin" as *u8 }
61
62func rf_ehash(buf: *u8, off: i64, len: i64) -> i64 {
  var h1: i64 = 5381
  var h2: i64 = 77245
  var i: i64 = 0
  while i < len { let c: i64 = buf[off+i] as i64; h1 = (h1*33 + c) % 1073741789; h2 = (h2*131 + c) % 1073741783; i = i + 1 }
  return h1 * 1073741783 + h2
68}
69func rf_bsearch(a: *i64, n: i64, v: i64) -> i64 {
  var lo: i64 = 0
  var hi: i64 = n - 1
  while lo <= hi { let mid: i64 = (lo+hi)/2; if a[mid] == v { return mid } if a[mid] < v { lo = mid+1 } else { hi = mid-1 } }
  return 0-1
74}
75func rf_lc(c: i64) -> i64 { if c >= 65 { if c <= 90 { return c + 32 } } return c }
76func rf_isal(c: i64) -> i64 { if c >= 97 { if c <= 122 { return 1 } } if c >= 48 { if c <= 57 { return 1 } } return 0 }
77func rf_tok(buf: *u8, n: i64, hout: *i64, cap: i64) -> i64 {
  let wb: *u8 = sys_mmap(64)
  var nt: i64 = 0
  var i: i64 = 0
  while i < n {
      let c: i64 = rf_lc(buf[i] as i64)
      if rf_isal(c) == 1 {
          var wl: i64 = 0
          var j: i64 = i
          var live: i64 = 1
          while live == 1 {
              if j >= n { live = 0 } else {
                  let cj: i64 = rf_lc(buf[j] as i64)
                  if rf_isal(cj) == 0 { live = 0 } else { if wl < 48 { wb[wl] = cj as u8; wl = wl + 1 } j = j + 1 }
              }
          }
          if nt < cap { hout[nt] = rf_ehash(wb, 0, wl); nt = nt + 1 }
          i = j
      } else { i = i + 1 }
  }
  return nt
98}
99func rf_vid(map: *i64, vhash: *i64, vst: *i64, h: i64) -> i64 {
  var slot: i64 = h & (RF_MAPN - 1)
  if slot < 0 { slot = 0 - slot }
  var probe: i64 = 0
  while probe < RF_MAPN {
      let e: i64 = map[slot]
      if e == 0 {
          if vst[0] >= RF_MAXV - 1 { return 1 }
          vst[0] = vst[0] + 1
          let id: i64 = vst[0]
          vhash[id] = h
          map[slot] = id
          return id
      }
      if vhash[e] == h { return e }
      slot = (slot + 1) & (RF_MAPN - 1)
      probe = probe + 1
  }
  return 1
118}
119
120// forward on the ta_ tape (f32). out[0]=loss node, out[1]=Ls, out[2]=Le, out[3]=Xg leaf, out[4..12]=leaves.
121// rtab = precomputed RoPE cos/sin table (nx_f32_sin/cos mmap PER CALL -> table-driven or OOM).
122func rf_fwd(tape: *i64, vals: *i64, st: *i64, W: *i64, Xg: *i64, tgts: *i64, tgte: *i64, T: i64, useattn: i64, out: *i64, rtab: *i64) -> i64 {
  st[0] = 0
  st[1] = 0
  let scale: i64 = nx_f32_div(TA_F32_ONE, nx_f32_sqrt(nx_i32_to_f32(RF_E)))
  let nXg: i64 = ta_leaf(tape, vals, st, T, RF_E, Xg, 0)
  let nWq: i64 = ta_leaf(tape, vals, st, RF_E, RF_E, W, FOFF_WQ)
  let nWk: i64 = ta_leaf(tape, vals, st, RF_E, RF_E, W, FOFF_WK)
  let nWv: i64 = ta_leaf(tape, vals, st, RF_E, RF_E, W, FOFF_WV)
  let nWo: i64 = ta_leaf(tape, vals, st, RF_E, RF_E, W, FOFF_WO)
  let nWg: i64 = ta_leaf(tape, vals, st, RF_E, RF_F, W, FOFF_WG)
  let nWu: i64 = ta_leaf(tape, vals, st, RF_E, RF_F, W, FOFF_WU)
  let nWd: i64 = ta_leaf(tape, vals, st, RF_F, RF_E, W, FOFF_WD)
  let nUs: i64 = ta_leaf(tape, vals, st, 1, RF_E, W, FOFF_US)
  let nUe: i64 = ta_leaf(tape, vals, st, 1, RF_E, W, FOFF_UE)
  var nH: i64 = nXg
  if useattn == 1 {
      let nXn: i64 = ta_rmsnorm_rows(tape, vals, st, nXg)
      let nQ: i64 = ta_matmul(tape, vals, st, nXn, nWq)
      let nK: i64 = ta_matmul(tape, vals, st, nXn, nWk)
      let nV: i64 = ta_matmul(tape, vals, st, nXn, nWv)
      let nQr: i64 = ta_rope_tab(tape, vals, st, nQ, rtab)
      let nKr: i64 = ta_rope_tab(tape, vals, st, nK, rtab)
      let nS: i64 = ta_matmul_nt(tape, vals, st, nQr, nKr)
      let nSs: i64 = ta_cmul(tape, vals, st, nS, scale)
      let nA: i64 = ta_softmax_rows(tape, vals, st, nSs, RF_CAUSAL)   // RF_CAUSAL=0 => BIDIRECTIONAL reader
      let nO: i64 = ta_matmul(tape, vals, st, nA, nV)
      let nOp: i64 = ta_matmul(tape, vals, st, nO, nWo)
      nH = ta_vadd(tape, vals, st, nXg, nOp)
  }
  let nHn: i64 = ta_rmsnorm_rows(tape, vals, st, nH)
  let nG: i64 = ta_matmul(tape, vals, st, nHn, nWg)
  let nU2: i64 = ta_matmul(tape, vals, st, nHn, nWu)
  let nSg: i64 = ta_silu(tape, vals, st, nG)
  let nHs: i64 = ta_hadamard(tape, vals, st, nSg, nU2)
  let nD: i64 = ta_matmul(tape, vals, st, nHs, nWd)
  let nY: i64 = ta_vadd(tape, vals, st, nH, nD)
  let nYn: i64 = ta_rmsnorm_rows(tape, vals, st, nY)
  let nLs: i64 = ta_matmul_nt(tape, vals, st, nUs, nYn)
  let nLe: i64 = ta_matmul_nt(tape, vals, st, nUe, nYn)
  let nCs: i64 = ta_softce_rows(tape, vals, st, nLs, tgts)
  let nCe: i64 = ta_softce_rows(tape, vals, st, nLe, tgte)
  let nRoot: i64 = ta_vadd(tape, vals, st, nCs, nCe)
  out[0] = nRoot
  out[1] = nLs
  out[2] = nLe
  out[3] = nXg
  out[4] = nWq
  out[5] = nWk
  out[6] = nWv
  out[7] = nWo
  out[8] = nWg
  out[9] = nWu
  out[10] = nWd
  out[11] = nUs
  out[12] = nUe
  return 0
178}
179
180// accumulate grads into flat G (f32 adds). Dense weights by layout; embedding rows scattered by token id.
181func rf_grab(tape: *i64, grads: *i64, out: *i64, G: *i64, ids: *i64, T: i64) -> i64 {
  let nXg: i64 = out[3]
  let gx: i64 = tape[7*nXg+6]
  var t: i64 = 0
  while t < T {
      let id: i64 = ids[t]
      var e: i64 = 0
      while e < RF_E { G[FOFF_EMB + id*RF_E + e] = nx_f32_add(G[FOFF_EMB + id*RF_E + e], grads[gx + t*RF_E + e]); e = e + 1 }
      t = t + 1
  }
  var k: i64 = 0
  while k < 9 {
      var sz: i64 = 576
      var of: i64 = 0
      if k == 0 { of = FOFF_WQ }
      if k == 1 { of = FOFF_WK }
      if k == 2 { of = FOFF_WV }
      if k == 3 { of = FOFF_WO }
      if k == 4 { sz = 1152; of = FOFF_WG }
      if k == 5 { sz = 1152; of = FOFF_WU }
      if k == 6 { sz = 1152; of = FOFF_WD }
      if k == 7 { sz = 24; of = FOFF_US }
      if k == 8 { sz = 24; of = FOFF_UE }
      let nid: i64 = out[4+k]
      let goff: i64 = tape[7*nid+6]
      var i: i64 = 0
      while i < sz { G[of+i] = nx_f32_add(G[of+i], grads[goff+i]); i = i + 1 }
      k = k + 1
  }
  return 0
211}
212
213// ---- data-parallel batch training (RF_PAR): one example per pool worker, private arenas ----
214struct RfCtx {
  w_ptr:     i64,
  ids_ptr:   i64,
  tt:        i64,
  tape_ptr:  i64,
  vals_ptr:  i64,
  grads_ptr: i64,
  st_ptr:    i64,
  xg_ptr:    i64,
  out_ptr:   i64,
  ge_ptr:    i64,
  rtab_ptr:  i64,
  useattn:   i64,
  tg_ptr:    i64,
  te_ptr:    i64,
  loss_ptr:  i64,
230}
231const RF_CTX_BYTES: i64 = 120
232
233// pool worker: gather this example's embeddings, forward, backward, grab grads into its PRIVATE Ge, store loss.
234// Reads W + rtab (shared, read-only) and writes only its private arenas -> thread-safe (proven pattern).
235func rf_par_task(ctx_i: i64) -> i64 {
  let c: *RfCtx = ctx_i as *RfCtx
  let W: *i64 = c.w_ptr as *i64
  let ids: *i64 = c.ids_ptr as *i64
  let T: i64 = c.tt
  let tape: *i64 = c.tape_ptr as *i64
  let vals: *i64 = c.vals_ptr as *i64
  let grads: *i64 = c.grads_ptr as *i64
  let st: *i64 = c.st_ptr as *i64
  let Xg: *i64 = c.xg_ptr as *i64
  let out: *i64 = c.out_ptr as *i64
  let Ge: *i64 = c.ge_ptr as *i64
  let rtab: *i64 = c.rtab_ptr as *i64
  let tgts: *i64 = c.tg_ptr as *i64
  let tgte: *i64 = c.te_ptr as *i64
  let loss: *i64 = c.loss_ptr as *i64
  var t: i64 = 0
  while t < T { var e: i64 = 0; while e < RF_E { Xg[t*RF_E+e] = W[FOFF_EMB + ids[t]*RF_E + e]; e = e + 1 } t = t + 1 }
  rf_fwd(tape, vals, st, W, Xg, tgts, tgte, T, c.useattn, out, rtab)
  loss[0] = ta_val(tape, vals, out[0], 0)
  ta_backward(tape, vals, grads, st[0], out[0])
  var z: i64 = 0
  while z < RF_NP { Ge[z] = TA_F32_ZERO; z = z + 1 }
  rf_grab(tape, grads, out, Ge, ids, T)
  return 0
260}
261
262func rf_init(W: *i64) -> i64 {
  var i: i64 = 0
  while i < 576 { W[FOFF_WQ+i] = ta_constf((((i*13)%9)-4), 32); W[FOFF_WK+i] = ta_constf((((i*17)%9)-4), 32); W[FOFF_WV+i] = ta_constf((((i*19)%9)-4), 32); W[FOFF_WO+i] = ta_constf((((i*23)%9)-4), 32); i = i + 1 }
  i = 0
  while i < 1152 { W[FOFF_WG+i] = ta_constf((((i*29)%9)-4), 40); W[FOFF_WU+i] = ta_constf((((i*31)%9)-4), 40); W[FOFF_WD+i] = ta_constf((((i*37)%9)-4), 40); i = i + 1 }
  i = 0
  while i < 24 { W[FOFF_US+i] = ta_constf((((i*5+1)%7)-3), 8); W[FOFF_UE+i] = ta_constf((((i*11+2)%7)-3), 8); i = i + 1 }
  return 0
270}
271
272func rf_span_f1(ids: *i64, a: i64, b: i64, gs: i64, ge: i64) -> i64 {
  let pl: i64 = b - a + 1
  let gl: i64 = ge - gs + 1
  if pl <= 0 { return 0 }
  if gl <= 0 { return 0 }
  var common: i64 = 0
  let used: *i64 = sys_mmap(RF_MAXG*8) as *i64
  var u: i64 = 0
  while u < gl { used[u] = 0; u = u + 1 }
  var p: i64 = a
  while p <= b {
      var q2: i64 = 0
      var got: i64 = 0
      while q2 < gl {
          if got == 0 { if used[q2] == 0 { if ids[gs+q2] == ids[p] { used[q2] = 1; common = common + 1; got = 1 } } }
          q2 = q2 + 1
      }
      p = p + 1
  }
  if common == 0 { return 0 }
  return (2*common*1000)/(pl+gl)
293}
294
295func rf_eval(tape: *i64, vals: *i64, st: *i64, W: *i64, dat: *i64, meta: *i64, r0: i64, r1: i64, useattn: i64, res: *i64, rtab: *i64) -> i64 {
  let out: *i64 = sys_mmap(16*8) as *i64
  let Xg: *i64 = sys_mmap(RF_MAXT*RF_E*8) as *i64
  let tgts: *i64 = sys_mmap(8) as *i64
  let tgte: *i64 = sys_mmap(8) as *i64
  var f1sum: i64 = 0
  var exact: i64 = 0
  var n: i64 = 0
  var r: i64 = r0
  while r < r1 {
      let T: i64 = meta[r*4+0]
      let qn: i64 = meta[r*4+1]
      let gs: i64 = meta[r*4+2]
      let ge: i64 = meta[r*4+3]
      let ids: *i64 = (dat as i64 + r*RF_MAXT*8) as *i64
      var t: i64 = 0
      while t < T { var e: i64 = 0; while e < RF_E { Xg[t*RF_E+e] = W[FOFF_EMB + ids[t]*RF_E + e]; e = e + 1 } t = t + 1 }
      tgts[0] = gs
      tgte[0] = ge
      rf_fwd(tape, vals, st, W, Xg, tgts, tgte, T, useattn, out, rtab)
      let nLs: i64 = out[1]
      let nLe: i64 = out[2]
      var s: i64 = qn + 1
      var bv: i64 = ta_val(tape, vals, nLs, qn+1)
      var j: i64 = qn + 2
      while j < T { let v: i64 = ta_val(tape, vals, nLs, j); if nx_f32_gt(v, bv) == 1 { bv = v; s = j } j = j + 1 }
      var en: i64 = s
      var ev: i64 = ta_val(tape, vals, nLe, s)
      var lim: i64 = s + RF_MAXG
      if lim > T { lim = T }
      j = s + 1
      while j < lim { let v2: i64 = ta_val(tape, vals, nLe, j); if nx_f32_gt(v2, ev) == 1 { ev = v2; en = j } j = j + 1 }
      f1sum = f1sum + rf_span_f1(ids, s, en, gs, ge)
      if s == gs { exact = exact + 1 }
      n = n + 1
      r = r + 1
  }
  if n == 0 { res[0] = 0; res[1] = 0; return 0 }
  res[0] = f1sum / n
  res[1] = exact * 1000 / n
  return 0
336}
337
338// ---- CRASH-TOLERANT CHECKPOINTING (the box gets VM-restarted by sibling sessions; per-epoch checkpoints make
339//      any interruption cost <=1 epoch and progress accumulate monotonically). Layout (i64 cells):
340//      ck[0]=magic (20262), ck[1]=reader epochs done, ck[2]=ablation epochs done, ck[3]=nw,
341//      ck[4]=reader FIRST-epoch loss milli, ck[5]=reader LAST-epoch loss milli (T3 survives resume),
342//      ck[6..6+NP) = W (reader), ck[6+NP..6+2NP) = W2 (ablation). W/W2 point INTO ck -> flush = one write. ----
343const RF_CKMAGIC: i64 = 20262
344const RF_CKCELLS: i64 = 6 + RF_NP + RF_NP
345
346func rf_ck_flush(ck: *i64) -> i64 {
  // ATOMIC checkpoint: write .part, then rename over the live file. A VM kill mid-write (the box gets
  // sibling-restarted ~15min) must never leave a TORN checkpoint -- a torn one resumes into garbage weights
  // and trains on silently. rename on the same fs is atomic => the live ckpt is always old-complete or
  // new-complete, never partial.
  let fd: i64 = sys_openat_wr(rf_ck_part_path(), 0x1a4)
  if fd < 0 { return 0 }
  let nb: i64 = RF_CKCELLS * 8
  var w: i64 = 0
  while w < nb { let x: i64 = sys_write(fd, (ck as i64 + w) as *u8, nb - w); if x <= 0 { w = nb } else { w = w + x } }
  sys_close(fd)
  if w < nb { return 0 }
  sys_renameat(rf_ck_part_path(), rf_ck_path())
  return 1
360}
361
362func rf_ck_load(ck: *i64, nw: i64) -> i64 {
  let fd: i64 = sys_openat_rd(rf_ck_path())
  if fd < 0 { return 0 }
  let nb: i64 = RF_CKCELLS * 8
  var got: i64 = 0
  var r: i64 = 1
  while r > 0 { if got >= nb { r = 0 } else { r = sys_read(fd, (ck as i64 + got) as *u8, nb - got); if r > 0 { got = got + r } } }
  sys_close(fd)
  if got < nb { return 0 }
  if ck[0] != RF_CKMAGIC { return 0 }
  if ck[3] != nw { return 0 }
  return 1
374}
375
376// train from start-epoch ep0 (checkpoint-resume); saves the checkpoint after EVERY epoch (ckslot 1=reader,
377// 2=ablation). res0[0]=final-epoch mean loss milli, res0[1]=first-SEEN-epoch mean (for the fall test).
378func rf_train(tape: *i64, vals: *i64, grads: *i64, st: *i64, W: *i64, dat: *i64, meta: *i64, ntrain: i64, useattn: i64, EP: i64, res0: *i64, ck: *i64, ckslot: i64, ep0: i64, rtab: *i64) -> i64 {
  let out: *i64 = sys_mmap(16*8) as *i64
  let Xg: *i64 = sys_mmap(RF_MAXT*RF_E*8) as *i64
  let tgts: *i64 = sys_mmap(8) as *i64
  let tgte: *i64 = sys_mmap(8) as *i64
  let G: *i64 = sys_mmap(RF_NP*8) as *i64
  let lr0: i64 = ta_constf(1, 32)                     // 0.03125 = R3a's proven rate
  let clip: i64 = nx_i32_to_f32(4)
  let nclip: i64 = nx_f32_neg(clip)
  let inv8: i64 = ta_constf(1, 8)
  // RF_PAR: pool + RF_BATCH private per-example arena sets (allocated ONCE, reused across epochs/batches)
  var pool: *NxThreadPool = 0 as *NxThreadPool
  let ptapes: *i64 = sys_mmap(RF_BATCH*4096*7*8) as *i64
  let pvals: *i64 = sys_mmap(RF_BATCH*262144*8) as *i64
  let pgrads: *i64 = sys_mmap(RF_BATCH*262144*8) as *i64
  let psts: *i64 = sys_mmap(RF_BATCH*2*8) as *i64
  let pxgs: *i64 = sys_mmap(RF_BATCH*RF_MAXT*RF_E*8) as *i64
  let pouts: *i64 = sys_mmap(RF_BATCH*16*8) as *i64
  let pges: *i64 = sys_mmap(RF_BATCH*RF_NP*8) as *i64
  let plosses: *i64 = sys_mmap(RF_BATCH*8) as *i64
  let ptg: *i64 = sys_mmap(RF_BATCH*2*8) as *i64
  let pctxs: *u8 = sys_mmap(RF_BATCH*RF_CTX_BYTES)
  if RF_PAR == 1 { pool = nx_pool_new(0, 0) }
  var first: i64 = 1
  var ep: i64 = ep0
  while ep < EP {
      var lsum: i64 = TA_F32_ZERO
      var nb: i64 = 0
      var r: i64 = 0
      while r + 8 <= ntrain {
          var z: i64 = 0
          while z < RF_NP { G[z] = TA_F32_ZERO; z = z + 1 }
          if RF_PAR == 1 {
              // ---- PARALLEL: submit RF_BATCH examples to the pool, each on private arenas ----
              let db: i64 = nx_pool_n_completed(pool)
              var mb: i64 = 0
              while mb < RF_BATCH {
                  let rr: i64 = r + mb
                  ptg[mb*2] = meta[rr*4+2]
                  ptg[mb*2+1] = meta[rr*4+3]
                  let cx: *RfCtx = ((pctxs as i64) + mb*RF_CTX_BYTES) as *RfCtx
                  cx.w_ptr     = W as i64
                  cx.ids_ptr   = (dat as i64 + rr*RF_MAXT*8)
                  cx.tt        = meta[rr*4+0]
                  cx.tape_ptr  = (ptapes as i64 + mb*4096*7*8)
                  cx.vals_ptr  = (pvals as i64 + mb*262144*8)
                  cx.grads_ptr = (pgrads as i64 + mb*262144*8)
                  cx.st_ptr    = (psts as i64 + mb*2*8)
                  cx.xg_ptr    = (pxgs as i64 + mb*RF_MAXT*RF_E*8)
                  cx.out_ptr   = (pouts as i64 + mb*16*8)
                  cx.ge_ptr    = (pges as i64 + mb*RF_NP*8)
                  cx.rtab_ptr  = rtab as i64
                  cx.useattn   = useattn
                  cx.tg_ptr    = (ptg as i64 + mb*2*8)
                  cx.te_ptr    = (ptg as i64 + (mb*2+1)*8)
                  cx.loss_ptr  = (plosses as i64 + mb*8)
                  nx_pool_submit(pool, rf_par_task, cx as i64)
                  mb = mb + 1
              }
              nx_pool_wait(pool, db + RF_BATCH)
              // reduce per-example grads (mb order = bit-identical to serial) + sum losses
              mb = 0
              while mb < RF_BATCH {
                  let Ge: *i64 = (pges as i64 + mb*RF_NP*8) as *i64
                  var zz: i64 = 0
                  while zz < RF_NP { G[zz] = nx_f32_add(G[zz], Ge[zz]); zz = zz + 1 }
                  lsum = nx_f32_add(lsum, plosses[mb])
                  mb = mb + 1
              }
          } else {
              var mb: i64 = 0
              while mb < RF_BATCH {
                  let rr: i64 = r + mb
                  let T: i64 = meta[rr*4+0]
                  let gs: i64 = meta[rr*4+2]
                  let ge: i64 = meta[rr*4+3]
                  let ids: *i64 = (dat as i64 + rr*RF_MAXT*8) as *i64
                  var t: i64 = 0
                  while t < T { var e: i64 = 0; while e < RF_E { Xg[t*RF_E+e] = W[FOFF_EMB + ids[t]*RF_E + e]; e = e + 1 } t = t + 1 }
                  tgts[0] = gs
                  tgte[0] = ge
                  rf_fwd(tape, vals, st, W, Xg, tgts, tgte, T, useattn, out, rtab)
                  let nRoot: i64 = out[0]
                  lsum = nx_f32_add(lsum, ta_val(tape, vals, nRoot, 0))
                  ta_backward(tape, vals, grads, st[0], nRoot)
                  rf_grab(tape, grads, out, G, ids, T)
                  mb = mb + 1
              }
          }
          // decay: small=24/40 (48-ep schedule); XL=8/12 (16-ep, 7x more data/epoch converges faster)
          var d1: i64 = 24
          var d2: i64 = 40
          if RF_XL == 1 { d1 = 8; d2 = 12 }
          var elr: i64 = lr0
          if ep >= d1 { elr = nx_f32_mul(lr0, ta_constf(1, 2)) }
          if ep >= d2 { elr = nx_f32_mul(lr0, ta_constf(1, 4)) }
          var w: i64 = 0
          while w < RF_NP {
              var gv: i64 = nx_f32_mul(G[w], inv8)
              if nx_f32_gt(gv, clip) == 1 { gv = clip }
              if nx_f32_lt(gv, nclip) == 1 { gv = nclip }
              W[w] = nx_f32_sub(W[w], nx_f32_mul(elr, gv))
              w = w + 1
          }
          nb = nb + 1
          r = r + 8
      }
      let mean: i64 = nx_f32_div(lsum, nx_i32_to_f32(nb*8))
      let mm: i64 = ta_f32_to_milli(mean)
      if first == 1 { res0[1] = mm; first = 0 }
      res0[0] = mm
      rf_puts("    epoch "); rf_pn(ep); rf_puts("  mean_loss_milli="); rf_pn(mm); rf_puts("\n" as *u8)
      if ckslot == 1 { if ep == 0 { ck[4] = mm } ck[5] = mm }
      ck[ckslot] = ep + 1
      rf_ck_flush(ck)
      ep = ep + 1
  }
  return 0
496}
497
498func main() -> i64 {
  rf_puts("nx_reader_squad_f32_gate (R3e: TRAIN-ON-F32 -- same reader, ta_* tensor autograd; Q16 collapse trap does not exist here)\n" as *u8)
  var pass: i64 = 0
  var total: i64 = 0
502
  let fd: i64 = sys_openat_rd(rf_rows_path())
  if fd < 0 { rf_puts("RED -- reader rows file missing (run: nx_qabench dt / dx)\n" as *u8); return 1 }
  let cap: i64 = 16777216
  let raw: *u8 = sys_mmap(cap)
  var got: i64 = 0
  var rr2: i64 = 1
  while rr2 > 0 { rr2 = sys_read(fd, (raw as i64 + got) as *u8, cap - got); if rr2 > 0 { got = got + rr2 } }
  sys_close(fd)
  if got < 16 { rf_puts("RED -- reader_rows.bin empty\n" as *u8); return 1 }
512
  let pfd: i64 = sys_openat_rd("knowledge/index/semppmi_v1.bin" as *u8)
  if pfd < 0 { rf_puts("RED -- semppmi_v1.bin missing\n" as *u8); return 1 }
  let pcap: i64 = 2097152
  let pblob: *u8 = sys_mmap(pcap)
  var pgot: i64 = 0
  var pr: i64 = 1
  while pr > 0 { if pgot >= pcap { pr = 0 } else { pr = sys_read(pfd, (pblob as i64 + pgot) as *u8, pcap - pgot); if pr > 0 { pgot = pgot + pr } } }
  sys_close(pfd)
  let phi: *i64 = (pblob as i64 + 8) as *i64
  let nv: i64 = phi[0]
  if pgot < 32 + nv*8 { rf_puts("RED -- semppmi vh truncated\n" as *u8); return 1 }
  let vh: *i64 = (pblob as i64 + 32) as *i64
  let efd: i64 = sys_openat_rd("knowledge/index/embed_v1.bin" as *u8)
  if efd < 0 { rf_puts("RED -- embed_v1.bin missing\n" as *u8); return 1 }
  let ecap: i64 = 67108864
  let eblob: *u8 = sys_mmap(ecap)
  var egot: i64 = 0
  var er: i64 = 1
  while er > 0 { if egot >= ecap { er = 0 } else { er = sys_read(efd, (eblob as i64 + egot) as *u8, ecap - egot); if er > 0 { egot = egot + er } } }
  sys_close(efd)
  let ehi: *i64 = (eblob as i64 + 8) as *i64
  let env: i64 = ehi[0]
  let edim: i64 = ehi[1]
  if edim != RF_E { rf_puts("RED -- embed dim mismatch\n" as *u8); return 1 }
  if env != nv { rf_puts("RED -- embed/vocab nv mismatch\n" as *u8); return 1 }
  let EQ: *i64 = (eblob as i64 + 24) as *i64
  rf_puts("  SGNS loaded: vocab nv="); rf_pn(nv); rf_puts("  dim="); rf_pn(edim); rf_puts("\n" as *u8)
540
  let dat: *i64 = sys_mmap(RF_NROWS*RF_MAXT*8) as *i64
  let meta: *i64 = sys_mmap(RF_NROWS*4*8) as *i64
  let qh: *i64 = sys_mmap(64*8) as *i64
  let ch: *i64 = sys_mmap(4096*8) as *i64
  let gh: *i64 = sys_mmap(64*8) as *i64
  let vhash: *i64 = sys_mmap(RF_MAXV*8) as *i64
  let vmap: *i64 = sys_mmap(RF_MAPN*8) as *i64
  let vst: *i64 = sys_mmap(8) as *i64
  vst[0] = 0
  var nrows: i64 = 0
  var off: i64 = 8
  var scanned: i64 = 0
  while off + 24 < got {
      if nrows >= RF_NROWS { off = got } else {
          let hp: *i64 = (raw as i64 + off) as *i64
          let qlen: i64 = hp[0]
          if qlen < 0 { off = got } else { if qlen > 4000 { off = got } else {
              let qp: *u8 = (raw as i64 + off + 8) as *u8
              let hp2: *i64 = (raw as i64 + off + 8 + qlen) as *i64
              let clen: i64 = hp2[0]
              let cp: *u8 = (raw as i64 + off + 16 + qlen) as *u8
              let hp3: *i64 = (raw as i64 + off + 16 + qlen + clen) as *i64
              let alen: i64 = hp3[0]
              let ap: *u8 = (raw as i64 + off + 24 + qlen + clen) as *u8
              off = off + 24 + qlen + clen + alen
              scanned = scanned + 1
              var qn: i64 = rf_tok(qp, qlen, qh, RF_MAXQ)
              let cn0: i64 = rf_tok(cp, clen, ch, 4096)
              let gn: i64 = rf_tok(ap, alen, gh, RF_MAXG)
              var cn: i64 = RF_MAXT - qn - 1
              if cn > cn0 { cn = cn0 }
              if qn >= 3 { if gn >= 1 { if cn >= 8 {
                  var gs: i64 = 0-1
                  var c: i64 = 0
                  while c + gn <= cn {
                      if gs < 0 {
                          var m: i64 = 1
                          var k: i64 = 0
                          while k < gn { if ch[c+k] != gh[k] { m = 0; k = gn } else { k = k + 1 } }
                          if m == 1 { gs = c }
                      }
                      c = c + 1
                  }
                  if gs >= 0 {
                      let T: i64 = qn + 1 + cn
                      let ids: *i64 = (dat as i64 + nrows*RF_MAXT*8) as *i64
                      var t: i64 = 0
                      while t < qn { let lid: i64 = rf_vid(vmap, vhash, vst, qh[t]); ids[t] = lid; t = t + 1 }
                      ids[qn] = 0
                      t = 0
                      while t < cn { let lid2: i64 = rf_vid(vmap, vhash, vst, ch[t]); ids[qn+1+t] = lid2; t = t + 1 }
                      meta[nrows*4+0] = T
                      meta[nrows*4+1] = qn
                      meta[nrows*4+2] = qn + 1 + gs
                      meta[nrows*4+3] = qn + 1 + gs + gn - 1
                      nrows = nrows + 1
                  }
              } } }
          } }
      }
  }
  let nw: i64 = vst[0]
  rf_puts("  scanned="); rf_pn(scanned); rf_puts("  usable="); rf_pn(nrows); rf_puts("  vocab nw="); rf_pn(nw); rf_puts("\n" as *u8)
  if nrows < 200 { rf_puts("RED -- too few usable rows\n" as *u8); return 1 }
  var ntrain: i64 = (nrows * 4) / 5
  rf_puts("  train="); rf_pn(ntrain); rf_puts("  held-out="); rf_pn(nrows - ntrain); rf_puts("\n" as *u8)
607
  // ---- checkpoint buffer: W/W2 live INSIDE it; resume if a valid checkpoint exists ----
  let ck: *i64 = sys_mmap(RF_CKCELLS*8) as *i64
  let W: *i64 = (ck as i64 + 6*8) as *i64
  let W2: *i64 = (ck as i64 + (6 + RF_NP)*8) as *i64
  let resumed: i64 = rf_ck_load(ck, nw)
  if resumed == 1 {
      rf_puts("  CHECKPOINT RESUMED: reader epochs done="); rf_pn(ck[1]); rf_puts("  ablation done="); rf_pn(ck[2]); rf_puts("\n" as *u8)
  } else {
      ck[0] = RF_CKMAGIC
      ck[1] = 0
      ck[2] = 0
      ck[3] = nw
      ck[4] = 0
      ck[5] = 0
      // embeddings: SGNS Q10 -> f32 RAW magnitude (f32 has no absorbing trap; true init)
      let q10: i64 = nx_i32_to_f32(1024)
      var hits: i64 = 0
      var wv: i64 = 1
      while wv <= nw {
          let r3: i64 = rf_bsearch(vh, nv, vhash[wv])
          if r3 >= 0 {
              var e2: i64 = 0
              while e2 < RF_E { W[FOFF_EMB + wv*RF_E + e2] = nx_f32_div(nx_i32_to_f32(EQ[r3*RF_E + e2]), q10); e2 = e2 + 1 }
              hits = hits + 1
          } else {
              var e3: i64 = 0
              while e3 < RF_E { W[FOFF_EMB + wv*RF_E + e3] = ta_constf((((wv*7+e3*3)%11)-5), 16); e3 = e3 + 1 }
          }
          wv = wv + 1
      }
      var e4: i64 = 0
      while e4 < RF_E { W[FOFF_EMB + e4] = ta_constf((((e4*5+2)%9)-4), 12); e4 = e4 + 1 }
      rf_puts("  SGNS coverage: "); rf_pn(hits); rf_puts("/"); rf_pn(nw); rf_puts("\n" as *u8)
      rf_init(W)
      var cw: i64 = 0
      while cw < RF_NP { W2[cw] = W[cw]; cw = cw + 1 }
      rf_init(W2)
      rf_ck_flush(ck)
  }
647
  let tape: *i64 = sys_mmap(4096*7*8) as *i64
  let vals: *i64 = sys_mmap(262144*8) as *i64
  let grads: *i64 = sys_mmap(262144*8) as *i64
  let st: *i64 = sys_mmap(2*8) as *i64
  var EP: i64 = 48              // small: 48-ep (664 rows undertrained at 12). XL: 16-ep (4771 rows, 7x more
  if RF_XL == 1 { EP = 16 }     // examples/epoch -> converges in fewer epochs; checkpoint lets us stop early.
  // RoPE table: ONE build (sin/cos mmap per call; 96*12*2 calls once = fine, per-element in the loop = OOM)
  let rtab: *i64 = sys_mmap((2 + 2*RF_MAXT*(RF_E/2))*8) as *i64
  ta_rope_build_tab(rtab, RF_MAXT, RF_E/2)
  rf_puts("  RoPE table built ("); rf_pn(RF_MAXT); rf_puts("x"); rf_pn(RF_E/2); rf_puts(")\n" as *u8)
658
  rf_puts("  training READER (f32, SGNS init, R3a recipe; checkpointed per epoch)...\n" as *u8)
  let resL: *i64 = sys_mmap(2*8) as *i64
  resL[1] = 0
  if ck[1] < EP { rf_train(tape, vals, grads, st, W, dat, meta, ntrain, 1, EP, resL, ck, 1, ck[1], rtab) }
  let resA: *i64 = sys_mmap(2*8) as *i64
  rf_eval(tape, vals, st, W, dat, meta, ntrain, nrows, 1, resA, rtab)
  rf_puts("  READER held-out: span-F1="); rf_pn(resA[0]); rf_puts("  exact-start="); rf_pn(resA[1]); rf_puts(" permille\n" as *u8)
666
  rf_puts("  training ABLATION (question-blind, same everything; checkpointed)...\n" as *u8)
  let resL2: *i64 = sys_mmap(2*8) as *i64
  if ck[2] < EP { rf_train(tape, vals, grads, st, W2, dat, meta, ntrain, 0, EP, resL2, ck, 2, ck[2], rtab) }
  let resB: *i64 = sys_mmap(2*8) as *i64
  rf_eval(tape, vals, st, W2, dat, meta, ntrain, nrows, 0, resB, rtab)
  rf_puts("  ABLATION held-out: span-F1="); rf_pn(resB[0]); rf_puts("  exact-start="); rf_pn(resB[1]); rf_puts(" permille\n" as *u8)
673
  total = total + 1
  if resA[0] > resB[0] + 100 { pass = pass + 1; rf_ck("T1 reader beats question-blind ablation by >100 F1 (liar-kill)" as *u8, 1) } else { rf_ck("T1 reader beats question-blind ablation by >100 F1 (liar-kill)" as *u8, 0) }
  total = total + 1
  if resA[1] > 60 { pass = pass + 1; rf_ck("T2 reader exact-start > 60 permille" as *u8, 1) } else { rf_ck("T2 reader exact-start > 60 permille" as *u8, 0) }
  total = total + 1
  var t3: i64 = 0
  if ck[4] > 0 { if ck[5]*10 < ck[4]*6 { t3 = 1 } }        // first vs last reader-epoch loss FROM THE CHECKPOINT (resume-safe)
  if t3 == 1 { pass = pass + 1; rf_ck("T3 train loss falls >40% (no collapse)" as *u8, 1) } else { rf_ck("T3 train loss falls >40% (no collapse)" as *u8, 0) }
682
  let resC: *i64 = sys_mmap(2*8) as *i64
  rf_eval(tape, vals, st, W, dat, meta, ntrain, nrows, 1, resC, rtab)
  total = total + 1
  var t4: i64 = 0
  if resC[0] == resA[0] { if resC[1] == resA[1] { t4 = 1 } }
  if t4 == 1 { pass = pass + 1; rf_ck("T4 deterministic re-eval" as *u8, 1) } else { rf_ck("T4 deterministic re-eval" as *u8, 0) }
689
  let pf: i64 = sys_openat_wr(rf_persist_path(), 0x1a4)
  if pf >= 0 {
      let mg: *i64 = sys_mmap(32) as *i64
      let mgb: *u8 = mg as *u8
      mgb[0]=78 as u8; mgb[1]=88 as u8; mgb[2]=78 as u8; mgb[3]=82 as u8; mgb[4]=51 as u8; mgb[5]=0 as u8; mgb[6]=0 as u8; mgb[7]=0 as u8
      mg[1] = nw
      mg[2] = RF_E
      mg[3] = RF_F
      sys_write(pf, mg as *u8, 32)
      sys_write(pf, W as *u8, RF_NP*8)
      sys_close(pf)
      rf_puts("  persisted trained f32 reader -> knowledge/index/reader_neural_f32.bin (NXNR3)\n" as *u8)
  }
703
  rf_puts("---- nx_reader_squad_f32_gate: passed "); rf_pn(pass); rf_puts(" / "); rf_pn(total); rf_puts("\n" as *u8)
  if pass == total { rf_puts("READER R3e GREEN -- the f32-trained attention reader LEARNS + GENERALIZES on real SQuAD where the question-blind reader cannot; quantize-after is next.\n" as *u8); return 0 }
  rf_puts("RED -- R3e not fully passed (see numbers)\n" as *u8)
  return 1
708}
nx_reader_squad_f32_gate.nx source

`nx_reader_squad_f32_gate.nx` source
