nx_kquant_test.nx source

1// nx_kquant_test.nx -- prove the QUALITY lever: finer-grained scales (the K-quant idea) cut the 2// quantization error on REALISTIC weights -- concentrated near zero with occasional OUTLIERS, the 3// real NN case. A naive block of 32 shares one scale, so one outlier inflates the step and blurs 4// the other 31 weights; sub-blocks of 8 isolate the outlier. Exit 0 only if: 5// (1) Q4 @ blk=8 has materially LOWER L2 error than Q4 @ blk=32 (the quality win is real); 6// (2) the cost is honest: blk=8 moves MORE bytes than blk=32 but still far less than fp16; 7// (3) so this is a real point on the quality/data frontier the Council balances. 8// license_tier: ORIGINAL 9 10import "nx_kquant.nx" 11import "nx_syscalls.nx" 12 13func kt_puts(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 } 14func kt_num(v: i64) -> i64 { let bb: *u8 = sys_mmap(28); var m: i64=v; if m<0 {m=0-m; sys_write(1,"-" as *u8,1)}; let t: *u8 = sys_mmap(28); var k: i64=0; if m==0 {t[0]=48;k=1}; while m>0 {t[k]=48+(m%10); m=m/10; k=k+1}; var i: i64=0; while i<k {bb[i]=t[k-1-i]; i=i+1}; sys_write(1, bb, k); return 0 } 15 16func main() -> i64 { 17 kt_puts("=== K-quant QUALITY lever: finer scales isolate outliers (realistic NN weights) ===\n" as *u8) 18 let R: i64 = 128 19 let C: i64 = 512 20 let w: *i64 = sys_mmap(8 * (R * C + 8)) as *i64 21 let x: *i64 = sys_mmap(8 * (C + 8)) as *i64 22 var rr: i64 = 0 23 while rr < R { 24 var cc: i64 = 0 25 while cc < C { 26 let idx: i64 = rr * C + cc 27 // concentrated ~[-24,24] (sum of 3 uniforms = bell-ish), with an outlier every 37th 28 var v: i64 = ((idx * 7 + 3) % 17 - 8) + ((idx * 5 + 1) % 17 - 8) + ((idx * 11 + 2) % 17 - 8) 29 if (idx % 37) == 0 { v = v * 6 } // outlier (real NN weights have them) 30 w[idx] = v 31 cc = cc + 1 32 } 33 rr = rr + 1 34 } 35 var c2: i64 = 0 36 while c2 < C { x[c2] = ((c2 * 13 + 5) % 17 - 8) + ((c2 * 3 + 7) % 17 - 8); c2 = c2 + 1 } 37 38 let codes: *i64 = sys_mmap(8 * (R * C + 8)) as *i64 39 let scales: *i64 = sys_mmap(8 * (R * 64 + 8)) as *i64 // enough for blk=8 (64 blocks/row) 40 let ye: *i64 = sys_mmap(8 * (R + 8)) as *i64 41 let y32: *i64 = sys_mmap(8 * (R + 8)) as *i64 42 let y8: *i64 = sys_mmap(8 * (R + 8)) as *i64 43 qlayer_matvec_exact(w, x, R, C, ye) 44 45 kq_layer_quantize(w, R, C, QMV_Q4MAX, 32, codes, scales) 46 kq_layer_matvec(codes, scales, x, R, C, QMV_Q4MAX, 32, y32) 47 let e32: i64 = qlayer_l2_relerr_permil(y32, ye, R) 48 49 kq_layer_quantize(w, R, C, QMV_Q4MAX, 8, codes, scales) 50 kq_layer_matvec(codes, scales, x, R, C, QMV_Q4MAX, 8, y8) 51 let e8: i64 = qlayer_l2_relerr_permil(y8, ye, R) 52 53 let b32: i64 = kq_bytes(R, C, 4, 32) 54 let b8: i64 = kq_bytes(R, C, 4, 8) 55 let bref: i64 = R * C * 2 56 kt_puts(" Q4 blk=32 L2 err permil = " as *u8); kt_num(e32); kt_puts(" bits/wt x100 = " as *u8); kt_num(kq_bits_x100(4, 32)); kt_puts("\n" as *u8) 57 kt_puts(" Q4 blk=8 L2 err permil = " as *u8); kt_num(e8); kt_puts(" bits/wt x100 = " as *u8); kt_num(kq_bits_x100(4, 8)); kt_puts("\n" as *u8) 58 kt_puts(" bytes: fp16=" as *u8); kt_num(bref); kt_puts(" blk32=" as *u8); kt_num(b32); kt_puts(" blk8=" as *u8); kt_num(b8); kt_puts("\n" as *u8) 59 if e8 < e32 { kt_puts(" -> finer scales cut error by permil " as *u8); kt_num(e32 - e8); kt_puts(" (quality win, more bytes)\n" as *u8) } 60 61 let r: *i64 = sys_mmap(8*8) as *i64 62 r[0] = 0; if e8 < e32 { r[0] = 1 } // finer reduces error 63 r[1] = 0; if e8 * 100 <= e32 * 80 { r[1] = 1 } // >=20% reduction (material) 64 r[2] = 0; if b8 > b32 { if b8 < bref { r[2] = 1 } } // honest cost, still compressed 65 var pass: i64 = 0; var i: i64 = 0 66 while i < 3 { pass = pass + r[i]; i = i + 1 } 67 kt_puts("----\n passed " as *u8); kt_num(pass); kt_puts("/3\n" as *u8) 68 if pass == 3 { kt_puts(" QUALITY LEVER PROVEN: finer scales (K-quant idea) trade bytes for materially better quality.\n" as *u8); sys_exit(0); return 0 } 69 kt_puts(" FAIL\n" as *u8); sys_exit(1); return 1 70}