code wiki / (root) / nx_zimage_qkv_verify.nx

nx_zimage_qkv_verify.nx source

↩ module page · 90 lines · 3376 B

1// nx_zimage_qkv_verify.nx -- SOVEREIGN Z-Image attention SUB-piece 2a: qkv projection, verified vs oracle. 2// 3// Small dedicated organ (the 285KB combined attention segfaults -> split into per-sub-step organs, each 4// small enough to dodge the codegen issue; piece #1 adaLN proved small organs verify cleanly). Computes 5// qkv = h_attn @ qkv_w^T [4,11520] sovereignly (f32) and verifies vs the oracle's dumped qkv. 6// license_tier: ORIGINAL 7import "nx_syscalls.nx" 8import "nx_le.nx" 9import "nx_f32.nx" 10import "nx_f32_div.nx" 11import "nx_f32_cvt.nx" 12import "nx_strconv.nx" 13const K_MAGIC_3840: i64 = 3840 14const K_MAGIC_11520: i64 = 11520 15 16func zqv_load(name: *u8, nl: i64, n_floats: i64) -> *u8 { 17 let base: *u8 = "/mnt/c/Users/elder/AppData/Local/Temp/claude/C--Users-elder/7be78b15-304c-449e-afe8-4d5bd7ddaa9c/scratchpad/zblk/" as *u8 18 let path: *u8 = sys_mmap(256) 19 var p: i64 = 0 20 var i: i64 = 0 21 while base[i] != 0 { path[p] = base[i]; p = p + 1; i = i + 1 } 22 i = 0 23 while i < nl { path[p] = name[i]; p = p + 1; i = i + 1 } 24 path[p] = 0x2E; p = p + 1 25 path[p] = 0x66; p = p + 1 26 path[p] = 0x33; p = p + 1 27 path[p] = 0x32; p = p + 1 28 path[p] = 0 29 let fd: i64 = sys_openat_rd(path) 30 if fd < 0 { return 0 as *u8 } 31 let bytes: i64 = n_floats * 4 32 let buf: *u8 = sys_mmap(bytes + 64) 33 var tot: i64 = 0 34 var go: i64 = 1 35 while go == 1 { 36 let r: i64 = sys_read(fd, ((buf as i64) + tot) as *u8, bytes - tot) 37 if r <= 0 { go = 0 } else { tot = tot + r; if tot >= bytes { go = 0 } } 38 } 39 sys_close(fd) 40 return buf 41} 42 43func main() -> i64 { 44 let D: i64 = K_MAGIC_3840 45 let NT: i64 = 4 46 let QKV: i64 = K_MAGIC_11520 47 let ha: *u8 = zqv_load("h_attn" as *u8, 6, NT * D) 48 let qkvw: *u8 = zqv_load("qkv_w" as *u8, 5, QKV * D) 49 let qkvg: *u8 = zqv_load("qkv" as *u8, 3, NT * QKV) 50 if (ha as i64) == 0 { return 30 } 51 if (qkvw as i64) == 0 { return 31 } 52 if (qkvg as i64) == 0 { return 32 } 53 54 let tolc: i64 = nx_f32_div(nx_i32_to_f32(2), nx_i32_to_f32(100)) // 0.02 55 var fails: i64 = 0 56 var first_bad: i64 = 0 - 1 57 var t: i64 = 0 58 while t < NT { 59 var o: i64 = 0 60 while o < QKV { 61 var acc: i64 = 0 62 var i: i64 = 0 63 let hb: i64 = t * D * 4 64 let wb: i64 = o * D * 4 65 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 } 66 let g: i64 = nx_le_read_u32(qkvg, (t * QKV + o) * 4) 67 var thr: i64 = tolc 68 let ag: i64 = g & 0x7FFFFFFF 69 if nx_f32_lt(thr, nx_f32_mul(tolc, ag)) == 1 { thr = nx_f32_mul(tolc, ag) } 70 if (nx_f32_sub(acc, g) & 0x7FFFFFFF) >= thr { fails = fails + 1; if first_bad < 0 { first_bad = t * QKV + o } } 71 o = o + 1 72 } 73 t = t + 1 74 } 75 76 let ofd: i64 = sys_openat_wr("/tmp/zqv.txt" as *u8, 0x1a4) 77 if ofd >= 0 { 78 let dec: *u8 = sys_mmap(32) 79 sys_write(ofd, "fails=" as *u8, 6) 80 let n1: i64 = nx_strconv_format_i64(fails, dec) 81 sys_write(ofd, dec, n1) 82 sys_write(ofd, " first_bad=" as *u8, 11) 83 let n2: i64 = nx_strconv_format_i64(first_bad, dec) 84 sys_write(ofd, dec, n2) 85 sys_write(ofd, "\n" as *u8, 1) 86 sys_close(ofd) 87 } 88 if fails > 0 { return 20 } 89 return 0 90}