nx_nofloat_simd_dot_probe_gate.nx source
↩ module page · 138 lines · 7966 B
1// nx_nofloat_simd_dot_probe_gate.nx -- DECISIVE probe for Stage 2 (SIMD decode). Before rebuilding the whole
2// decode around __i16x16_madd (vpmaddwd), answer three questions on representative data:
3// (1) does the W8A8 SIMD dot MATH work end-to-end (i16-lane pack2 store + vpmaddwd + i32x8 hsum + rescale)?
4// (2) how big is the quantization ERROR vs the exact i64 dot (the precision cost of the fork)?
5// (3) is it actually FASTER than the scalar i64 dot (the whole point)?
6// If error is small AND it's faster, wire into the real decode + measure token faithfulness. Else keep i32 lossless.
7// NOTE: nx_cc has NO i16 type -> i16 lanes are *u8 buffers with manual LE pack2 (as nx_nofloat_gemm does).
8// No hw writes (Rule 26). expect_exit: 0 license_tier: ORIGINAL
9import "nx_syscalls.nx"
10import "nx_nofloat_gemm.nx" // pack2 + i32x8_hsum + the __i16x16_madd intrinsic contract
11
12func sp(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
13func sn(v: i64) -> i64 { let b: *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 as u8;k=1} while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} var i: i64=0; while i<k{b[i]=t[k-1-i];i=i+1} sys_write(1,b,k); return 0 }
14func iabs(v: i64) -> i64 { if v<0 { return 0-v } return v }
15// int32x8 horizontal sum. *i32 loads SIGN-extend natively since the 2026-07-10 sext compiler fix (this fn
16// carried a manual branchless sign-extend during the workaround era; removed in lockstep with the bless).
17func hsum_sx(acc: *u8) -> i64 {
18 let p: *i32 = acc as *i32
19 var s: i64=0; var i: i64=0
20 while i<8 { s=s+(p[i] as i64); i=i+1 }
21 return s
22}
23
24// exact scalar dot: sum(x[k]*w[k]) >> shift (x=Q24, w=Q16 -> mirrors the real matmul reduction)
25func exact_dot(x: *i64, w: *i64, n: i64, shift: i64) -> i64 {
26 var s: i64=0; var k: i64=0
27 while k<n { s=s+(x[k]*w[k]); k=k+1 }
28 return s>>shift
29}
30// W8A8 SIMD dot: quantize x and w to int8-in-i16 (per-vector scales), vpmaddwd accumulate, staged rescale.
31// (sx*sw) ~2e9, *(raw ~1e7) ~2e16 -> stays in i64, then >>shift.
32func simd_w8a8_dot(x: *i64, w: *i64, xi: *u8, wi: *u8, acc: *u8, n: i64, shift: i64) -> i64 {
33 var xmax: i64=0; var wmax: i64=0; var k: i64=0
34 while k<n { let ax: i64=iabs(x[k]); if ax>xmax { xmax=ax } let aw: i64=iabs(w[k]); if aw>wmax { wmax=aw } k=k+1 }
35 var sx: i64=xmax/127; if sx<1 { sx=1 }
36 var sw: i64=wmax/127; if sw<1 { sw=1 }
37 k=0
38 while k<n { pack2(xi, k, x[k]/sx); pack2(wi, k, w[k]/sw); k=k+1 }
39 let z: *i64 = acc as *i64; z[0]=0; z[1]=0; z[2]=0; z[3]=0
40 let xb: i64 = xi as i64; let wb: i64 = wi as i64
41 k=0
42 while k<n { __i16x16_madd(acc, (xb+k*2) as *u8, (wb+k*2) as *u8); k=k+16 }
43 let raw: i64 = hsum_sx(acc)
44 return ((sx*sw)*raw)>>shift
45}
46
47func main() -> i64 {
48 sp("SIMD W8A8 DOT PROBE -- does vpmaddwd give a fast+faithful-enough dot vs the exact i64 dot?\n\n" as *u8)
49 let N: i64=896
50 let x: *i64 = sys_mmap(N*8) as *i64
51 let w: *i64 = sys_mmap(N*8) as *i64
52 let xi: *u8 = sys_mmap(N*2)
53 let wi: *u8 = sys_mmap(N*2)
54 let acc: *u8 = sys_mmap(64)
55 var k: i64=0
56 while k<N {
57 let a: i64 = ((k*2654435761 + 1013904223) % 33554432) - 16777216 // Q24 in [-2^24, 2^24)
58 let b: i64 = ((k*40503 + 12345) % 40000) - 20000 // Q16 in [-20000, 20000)
59 x[k]=a; w[k]=b; k=k+1
60 }
61
62 var pass: i64=0; var ttl: i64=0
63 var worst_ppm: i64=0
64 var sh: i64=16
65 while sh<=24 {
66 let ex: i64 = exact_dot(x, w, N, sh)
67 let sd: i64 = simd_w8a8_dot(x, w, xi, wi, acc, N, sh)
68 var ppm: i64=0
69 if ex!=0 { ppm = (iabs(sd-ex)*1000000)/iabs(ex) }
70 if ppm>worst_ppm { worst_ppm=ppm }
71 sp(" shift="); sn(sh); sp(": exact="); sn(ex); sp(" simd="); sn(sd); sp(" rel-err="); sn(ppm); sp(" ppm\n")
72 sh=sh+4
73 }
74 // (1) i16 store path: manual LE read of stored lane 7 matches x[7]/sx (proves pack2 store + SIMD read same bytes)
75 ttl=ttl+1
76 var xmax: i64=0; k=0; while k<N { let ax: i64=iabs(x[k]); if ax>xmax { xmax=ax } k=k+1 }
77 var sx: i64=xmax/127; if sx<1 { sx=1 }
78 let want7: i64 = x[7]/sx
79 var man: i64 = (xi[14] as i64) | ((xi[15] as i64)<<8) // slot 7 * 2 bytes = offset 14
80 if man >= 32768 { man = man - 65536 }
81 sp(" i16 pack2 store slot7: want="); sn(want7); sp(" manualLE="); sn(man); sp(": ")
82 if man==want7 { pass=pass+1; sp("PASS\n") } else { sp("FAIL\n") }
83 // (2) error under a threshold on representative data
84 ttl=ttl+1; sp(" worst rel-err <= 20000 ppm (2%): "); if worst_ppm<=20000 { pass=pass+1; sp("PASS ("); sn(worst_ppm); sp(" ppm)\n") } else { sp("HIGH ("); sn(worst_ppm); sp(" ppm)\n") }
85
86 // (3) REAL decode-kernel speed: a full matmul (OUTD rows x N). Weights pre-quantized ONCE (as at cache build);
87 // per matmul the activation is quantized ONCE then reused across all rows -> this is what decode actually does.
88 let OUTD: i64=896 // e.g. the q-proj / o-proj width (single-token T=1)
89 let Wm: *i64 = sys_mmap(OUTD*N*8) as *i64 // full weight matrix (i64 Q16)
90 let Wm8: *u8 = sys_mmap(OUTD*N*2) // pre-quantized i16-lane weights
91 let sw_r: *i64 = sys_mmap(OUTD*8) as *i64 // per-row weight scale
92 let dst_s: *i64 = sys_mmap(OUTD*8) as *i64
93 let dst_v: *i64 = sys_mmap(OUTD*8) as *i64
94 var o: i64=0
95 while o<OUTD {
96 var kk: i64=0; var rm: i64=0
97 while kk<N { let bv: i64=(((o*7+kk)*40503 + 12345) % 40000) - 20000; Wm[o*N+kk]=bv; if iabs(bv)>rm { rm=iabs(bv) } kk=kk+1 }
98 var s: i64=rm/127; if s<1 { s=1 }
99 sw_r[o]=s
100 kk=0; while kk<N { pack2(((Wm8 as i64)+o*N*2) as *u8, kk, Wm[o*N+kk]/s); kk=kk+1 }
101 o=o+1
102 }
103 let REPS: i64=200
104 // scalar i64 matmul
105 let t0: i64=sys_now_ms()
106 var r: i64=0
107 while r<REPS { o=0; while o<OUTD { dst_s[o]=exact_dot(x, ((Wm as i64)+o*N*8) as *i64, N, 16); o=o+1 } r=r+1 }
108 let t1: i64=sys_now_ms()
109 // SIMD W8A8 matmul: quantize activation ONCE per matmul, then madd each row against pre-quant weights
110 r=0
111 while r<REPS {
112 var xmx: i64=0; k=0; while k<N { let ax: i64=iabs(x[k]); if ax>xmx { xmx=ax } k=k+1 }
113 var sxq: i64=xmx/127; if sxq<1 { sxq=1 }
114 k=0; while k<N { pack2(xi, k, x[k]/sxq); k=k+1 }
115 o=0
116 while o<OUTD {
117 let z: *i64 = acc as *i64; z[0]=0; z[1]=0; z[2]=0; z[3]=0
118 let xb: i64=xi as i64; let wb: i64=(Wm8 as i64)+o*N*2
119 k=0; while k<N { __i16x16_madd(acc, (xb+k*2) as *u8, (wb+k*2) as *u8); k=k+16 }
120 dst_v[o]=((sxq*sw_r[o])*hsum_sx(acc))>>16
121 o=o+1
122 }
123 r=r+1
124 }
125 let t2: i64=sys_now_ms()
126 let ex_ms: i64=t1-t0; let sd_ms: i64=t2-t1
127 sp("\n REAL matmul ("); sn(OUTD); sp("x"); sn(N); sp(", "); sn(REPS); sp(" reps): scalar-i64="); sn(ex_ms); sp(" ms, simd-w8a8="); sn(sd_ms); sp(" ms")
128 if sd_ms>0 { sp(" (=> "); sn((ex_ms*100)/sd_ms); sp("/100x)\n") } else { sp("\n") }
129 // matmul argmax must still agree (the thing that picks the token)
130 var bs: i64=0; var bsv: i64=dst_s[0]; o=1; while o<OUTD { if dst_s[o]>bsv { bsv=dst_s[o]; bs=o } o=o+1 }
131 var bv2: i64=0; var bvv: i64=dst_v[0]; o=1; while o<OUTD { if dst_v[o]>bvv { bvv=dst_v[o]; bv2=o } o=o+1 }
132 ttl=ttl+1; sp(" matmul argmax agrees (scalar row "); sn(bs); sp(" == simd row "); sn(bv2); sp("): "); if bs==bv2 { pass=pass+1; sp("PASS\n") } else { sp("DIFF (quant flipped the argmax on this synthetic data)\n") }
133 ttl=ttl+1; sp(" SIMD matmul faster than scalar: "); if sd_ms<ex_ms { pass=pass+1; sp("PASS\n") } else { sp("NO\n") }
134
135 sp("\nNX-NOFLOAT-SIMD-DOT-PROBE-GATE passed "); sn(pass); sp("/"); sn(ttl)
136 if pass==ttl { sp(" verdict=GREEN (W8A8 SIMD dot viable -> wire into decode + measure token faithfulness)\n"); return 0 }
137 sp(" verdict=AMBER (read the numbers -> decide)\n"); return 0
138}