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

↩ module page · 205 lines · 7888 B
1// nx_zimage_attn_verify.nx -- SOVEREIGN Z-Image block piece #2: JointAttention, verified vs the oracle.
2//
3// z_image.hpp JointAttention: qkv = Linear(h); split Q|K|V each 30 heads x 128; per-head QK-RMSNorm(qn/kn);
4// axial RoPE; per-head SDPA (softmax(q.k^T / sqrt(128)) v); out = Linear. Reads oracle dumps h_attn (input),
5// qkv_w/out_w/qn/kn (weights), cosr/sinr (rope), and golden ao_out; computes sovereignly; verifies to tol.
6// Reference dumps are oracle-only; the ORGAN is pure Nishi (nx_f32_*, raw syscalls, static ELF).
7// license_tier: ORIGINAL
8import "nx_syscalls.nx"
9import "nx_le.nx"
10import "nx_f32.nx"
11import "nx_f32_div.nx"
12import "nx_f32_cvt.nx"
13import "nx_f32_exp.nx"
14import "nx_strconv.nx"
15const K_MAGIC_3840: i64 = 3840
16const K_MAGIC_11520: i64 = 11520
17const K_MAGIC_100000: i64 = 100000
18const K_MAGIC_2000: i64 = 2000
19
20func zav_load(name: *u8, nl: i64, n_floats: i64) -> *u8 {
  let base: *u8 = "/mnt/c/Users/elder/AppData/Local/Temp/claude/C--Users-elder/7be78b15-304c-449e-afe8-4d5bd7ddaa9c/scratchpad/zblk/" as *u8
  let path: *u8 = sys_mmap(256)
  var p: i64 = 0
  var i: i64 = 0
  while base[i] != 0 { path[p] = base[i]; p = p + 1; i = i + 1 }
  i = 0
  while i < nl { path[p] = name[i]; p = p + 1; i = i + 1 }
  path[p] = 0x2E; p = p + 1
  path[p] = 0x66; p = p + 1
  path[p] = 0x33; p = p + 1
  path[p] = 0x32; p = p + 1
  path[p] = 0
  let fd: i64 = sys_openat_rd(path)
  if fd < 0 { return 0 as *u8 }
  let bytes: i64 = n_floats * 4
  let buf: *u8 = sys_mmap(bytes + 64)
  var tot: i64 = 0
  var go: i64 = 1
  while go == 1 {
      let r: i64 = sys_read(fd, ((buf as i64) + tot) as *u8, bytes - tot)
      if r <= 0 { go = 0 } else { tot = tot + r; if tot >= bytes { go = 0 } }
  }
  sys_close(fd)
  return buf
45}
46
47// per-head RMSNorm(128) then axial-RoPE, in place on the i64 f32-bit array at qkv[off .. off+128]
48func zav_qknorm_rope(qkv: *i64, off: i64, w: *u8, cosr: *u8, sinr: *u8, t: i64, eps: i64) -> i64 {
  var ss: i64 = 0
  var d: i64 = 0
  while d < 128 { let v: i64 = qkv[off + d]; ss = nx_f32_add(ss, nx_f32_mul(v, v)); d = d + 1 }
  let ms: i64 = nx_f32_div(ss, nx_i32_to_f32(128))
  let rms: i64 = nx_f32_sqrt(nx_f32_add(ms, eps))
  d = 0
  while d < 128 { qkv[off + d] = nx_f32_div(nx_f32_mul(qkv[off + d], nx_le_read_u32(w, d * 4)), rms); d = d + 1 }
  var j: i64 = 0
  while j < 64 {
      let c: i64 = nx_le_read_u32(cosr, (t * 64 + j) * 4)
      let s: i64 = nx_le_read_u32(sinr, (t * 64 + j) * 4)
      let x0: i64 = qkv[off + 2 * j]
      let x1: i64 = qkv[off + 2 * j + 1]
      qkv[off + 2 * j] = nx_f32_sub(nx_f32_mul(x0, c), nx_f32_mul(x1, s))
      qkv[off + 2 * j + 1] = nx_f32_add(nx_f32_mul(x0, s), nx_f32_mul(x1, c))
      j = j + 1
  }
  return 0
67}
68
69func zstage(n: i64) -> i64 {
  let fd: i64 = sys_openat_wr("/tmp/zstage.txt" as *u8, 0x1a4)
  if fd >= 0 { let dec: *u8 = sys_mmap(16); let nd: i64 = nx_strconv_format_i64(n, dec); sys_write(fd, dec, nd); sys_write(fd, "\n" as *u8, 1); sys_close(fd) }
  return 0
73}
74
75func main() -> i64 {
  let D: i64 = K_MAGIC_3840
  let H: i64 = 30
  let HD: i64 = 128
  let NT: i64 = 4
  let QKV: i64 = K_MAGIC_11520
  let ha: *u8 = zav_load("h_attn" as *u8, 6, NT * D)
  let qkvw: *u8 = zav_load("qkv_w" as *u8, 5, QKV * D)
  let outw: *u8 = zav_load("out_w" as *u8, 5, D * D)
  let qn: *u8 = zav_load("qn" as *u8, 2, HD)
  let kn: *u8 = zav_load("kn" as *u8, 2, HD)
  let cosr: *u8 = zav_load("cosr" as *u8, 4, NT * 64)
  let sinr: *u8 = zav_load("sinr" as *u8, 4, NT * 64)
  let aog: *u8 = zav_load("ao_out" as *u8, 6, NT * D)
  if (ha as i64) == 0 { return 30 }
  if (qkvw as i64) == 0 { return 31 }
  if (aog as i64) == 0 { return 32 }
  if (outw as i64) == 0 { return 33 }
  if (qn as i64) == 0 { return 34 }
  if (cosr as i64) == 0 { return 35 }
  zstage(1)
96
  let eps: i64 = nx_f32_div(nx_i32_to_f32(1), nx_i32_to_f32(K_MAGIC_100000))
  let scale: i64 = nx_f32_div(nx_i32_to_f32(1), nx_f32_sqrt(nx_i32_to_f32(128)))
99
  // qkv = h_attn @ qkv_w^T   -> [NT, 11520] f32-bit
  let qkv: *i64 = sys_mmap(NT * QKV * 8) as *i64
  zstage((qkvw as i64) >> 28)            // marker: high bits of the qkvw pointer (valid mmap => large)
  var t: i64 = 0
  while t < NT {
      var o: i64 = 0
      while o < QKV {
          var acc: i64 = 0
          var i: i64 = 0
          let hb: i64 = t * D * 4
          let wb: i64 = o * D * 4
          while i < D { acc = nx_f32_add(acc, nx_f32_mul(nx_le_read_u32(ha, hb + i * 4), nx_le_read_u32(qkvw, wb + i * 4))); i = i + 1 }
          qkv[t * QKV + o] = acc
          if (o - (o / K_MAGIC_2000) * K_MAGIC_2000) == 0 { zstage(t * K_MAGIC_100000 + o) }
          o = o + 1
      }
      zstage(300 + t)
      t = t + 1
  }
119
  zstage(2)
  // qk-norm + rope on Q (offset 0) and K (offset 3840), per head
  t = 0
  while t < NT {
      var hd: i64 = 0
      while hd < H {
          zav_qknorm_rope(qkv, t * QKV + hd * HD, qn, cosr, sinr, t, eps)
          zav_qknorm_rope(qkv, t * QKV + D + hd * HD, kn, cosr, sinr, t, eps)
          hd = hd + 1
      }
      t = t + 1
  }
132
  zstage(3)
  // SDPA per head -> ao [NT, D]
  let ao: *i64 = sys_mmap(NT * D * 8) as *i64
  let sc: *i64 = sys_mmap(NT * 8) as *i64
  var hd2: i64 = 0
  while hd2 < H {
      var t1: i64 = 0
      while t1 < NT {
          var mx: i64 = 0
          var t2: i64 = 0
          while t2 < NT {
              var s: i64 = 0
              var d: i64 = 0
              while d < HD { s = nx_f32_add(s, nx_f32_mul(qkv[t1 * QKV + hd2 * HD + d], qkv[t2 * QKV + D + hd2 * HD + d])); d = d + 1 }
              s = nx_f32_mul(s, scale)
              sc[t2] = s
              if t2 == 0 { mx = s } else { if nx_f32_lt(mx, s) == 1 { mx = s } }
              t2 = t2 + 1
          }
          var sum: i64 = 0
          t2 = 0
          while t2 < NT { let e: i64 = nx_f32_exp(nx_f32_sub(sc[t2], mx)); sc[t2] = e; sum = nx_f32_add(sum, e); t2 = t2 + 1 }
          var d2: i64 = 0
          while d2 < HD {
              var a: i64 = 0
              t2 = 0
              while t2 < NT { a = nx_f32_add(a, nx_f32_mul(nx_f32_div(sc[t2], sum), qkv[t2 * QKV + D + D + hd2 * HD + d2])); t2 = t2 + 1 }
              ao[t1 * D + hd2 * HD + d2] = a
              d2 = d2 + 1
          }
          t1 = t1 + 1
      }
      hd2 = hd2 + 1
  }
167
  zstage(4)
  // out = ao @ out_w^T  ; verify vs ao_out golden
  let tolc: i64 = nx_f32_div(nx_i32_to_f32(5), nx_i32_to_f32(100))    // 0.05
  var fails: i64 = 0
  var first_bad: i64 = 0 - 1
  t = 0
  while t < NT {
      var o: i64 = 0
      while o < D {
          var acc: i64 = 0
          var i: i64 = 0
          let wb: i64 = o * D * 4
          while i < D { acc = nx_f32_add(acc, nx_f32_mul(ao[t * D + i], nx_le_read_u32(outw, wb + i * 4))); i = i + 1 }
          let g: i64 = nx_le_read_u32(aog, (t * D + o) * 4)
          var thr: i64 = tolc
          let ag: i64 = g & 0x7FFFFFFF
          if nx_f32_lt(thr, nx_f32_mul(tolc, ag)) == 1 { thr = nx_f32_mul(tolc, ag) }
          if (nx_f32_sub(acc, g) & 0x7FFFFFFF) >= thr { fails = fails + 1; if first_bad < 0 { first_bad = t * D + o } }
          o = o + 1
      }
      t = t + 1
  }
190
  let ofd: i64 = sys_openat_wr("/tmp/zattn.txt" as *u8, 0x1a4)
  if ofd >= 0 {
      let dec: *u8 = sys_mmap(32)
      sys_write(ofd, "fails=" as *u8, 6)
      let n1: i64 = nx_strconv_format_i64(fails, dec)
      sys_write(ofd, dec, n1)
      sys_write(ofd, " first_bad=" as *u8, 11)
      let n2: i64 = nx_strconv_format_i64(first_bad, dec)
      sys_write(ofd, dec, n2)
      sys_write(ofd, "\n" as *u8, 1)
      sys_close(ofd)
  }
  if fails > 0 { return 20 }
  return 0
205}
nx_zimage_attn_verify.nx source

`nx_zimage_attn_verify.nx` source
