nx_nofloat_w4_gemm_gate.nx source
↩ module page · 161 lines · 8323 B
1// nx_nofloat_w4_gemm_gate.nx -- THE DECISIVE W4A16 EXPERIMENT (de-risks the inference-speed arc before any
2// core-engine change). Decode is memory-bound (baseline: i8 3x faster than i32 = 4x fewer weight bytes). W4A16
3// stores weights as 4-bit (0.5 B/weight = HALF of i8) -> less memory read, but needs extra 4-bit unpack. Q: does
4// the halved weight-read beat the unpack cost on the REAL decode shape (matvec: 1 x K activation * N x K weights)?
5// This gate builds both paths (i8-weight and 4-bit-weight, SAME __i16x16_madd inner product), checks correctness
6// (both approximate the int16 reference within quant error), and TIMES them. Verdict decides whether to build the
7// full W4 kernel (and whether it needs SIMD-unpack primitives). Isolated -- touches NO serve code. license_tier: ORIGINAL expect_exit: 0
8import "nx_syscalls.nx"
9
10func gp_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){ n=n+1 } sys_write(1,s,n); return 0 }
11func gp_n(v: i64) -> i64 {
12 if v==0 { sys_write(1,"0" as *u8,1); return 0 }
13 var m: i64=v; if m<0 { sys_write(1,"-" as *u8,1); m=0-m }
14 let d: *u8=sys_mmap(24); var k: i64=0
15 while m>0 { d[k]=(48+(m%10)) as u8; m=m/10; k=k+1 }
16 let o: *u8=sys_mmap(24); var wi: i64=0
17 while wi<k { o[wi]=d[k-1-wi]; wi=wi+1 }
18 sys_write(1,o,k)
19 return 0
20}
21func gp_iabs(x: i64) -> i64 { if x<0 { return 0-x } return x }
22func i32x8_hsum(acc: *u8) -> i64 { let p: *i32=acc as *i32; var s: i64=0; var i: i64=0; while i<8 { s=s+(p[i] as i64); i=i+1 } return s }
23func now_ms() -> i64 { let ts: *i64=sys_mmap(16) as *i64; sys_clock_gettime_mono(ts); return ts[0]*1000 + ts[1]/1000000 }
24
25// deterministic pseudo-random small int in [-lo,hi] (LCG; Date.now-free)
26func lcg(st: *i64) -> i64 { st[0] = (st[0]*6364136223846793005 + 1442695040888963407) & 0x7FFFFFFFFFFFFFFF; return st[0] }
27
28func main() -> i64 {
29 let st: *i64=sys_mmap(32) as *i64; st[0]=0; st[1]=0
30 gp_puts("=== nx_nofloat_w4_gemm_gate: does 4-bit weight-read beat i8 on the memory-bound decode matvec? ===\n\n" as *u8)
31
32 let K: i64 = 896 // Qwen2.5-0.5B hidden
33 let N: i64 = 4864 // FFN up-proj out-dim (a big, memory-bound weight matrix)
34 let seed: *i64 = sys_mmap(8) as *i64; seed[0]=1234567
35
36 // activation A: int16 [K], small values [-8,7]
37 let A: *u8 = sys_mmap(K*2)
38 var i: i64=0
39 while i<K { let v: i64 = (lcg(seed) & 15) - 8; A[i*2]=(v) as u8; A[i*2+1]=(v>>8) as u8; i=i+1 }
40
41 // weights W: int16 [N*K], values [-127,127] (fit i8 exactly -> the i8 path is bit-exact vs W).
42 let W: *i64 = sys_mmap(N*K*8) as *i64
43 i=0; while i<N*K { W[i] = (lcg(seed) % 255) - 127; i=i+1 }
44
45 // --- quantize each row of W to i8 (1 byte, exact here) and to 4-bit (0.5 byte, packed 2/byte) with a proper
46 // PER-ROW scale = ceil(maxabs/7) so codes land in [-8,7] (real Q4-style quant, low error). ---
47 let Wi8: *u8 = sys_mmap(N*K) // 1 byte/weight
48 let Wi4: *u8 = sys_mmap(N*K/2) // 0.5 byte/weight (two 4-bit codes per byte)
49 let scale4: *i64 = sys_mmap(N*8) as *i64
50 i=0
51 while i<N {
52 var k: i64=0
53 var maxabs: i64=0
54 while k<K { let a: i64 = gp_iabs(W[i*K+k]); if a>maxabs { maxabs=a } k=k+1 }
55 var sc: i64 = (maxabs + 6) / 7; if sc < 1 { sc = 1 }
56 scale4[i] = sc
57 k=0
58 while k<K {
59 let w: i64 = W[i*K+k]
60 Wi8[i*K+k] = (w & 255) as u8 // i8 exact (w in [-127,127])
61 var c4: i64 = (w + (sc/2)*(1)) / sc // round-ish to nearest code
62 if w < 0 { c4 = (w - (sc/2)) / sc }
63 if c4>7 { c4=7 } if c4<0-8 { c4=0-8 }
64 let bi: i64 = (i*K+k)
65 if (k & 1) == 0 { Wi4[bi/2] = (c4 & 15) as u8 } else { Wi4[bi/2] = (Wi4[bi/2] as i64 | ((c4 & 15) << 4)) as u8 }
66 k=k+1
67 }
68 i=i+1
69 }
70
71 // scratch: unpacked i16 weight row + SIMD accumulator
72 let wrow: *u8 = sys_mmap(K*2)
73 let acc: *u8 = sys_mmap(64)
74 let z: *i64 = acc as *i64
75 let Cref: *i64 = sys_mmap(N*8) as *i64
76 let C8: *i64 = sys_mmap(N*8) as *i64
77 let C4: *i64 = sys_mmap(N*8) as *i64
78
79 // reference: exact int16 matvec (A . W_row)
80 i=0
81 while i<N { var accr: i64=0; var k: i64=0
82 while k<K { let av: i64 = ((A[k*2] as i64) | ((A[k*2+1] as i64)<<8)); var as2: i64=av; if as2>=32768 { as2=as2-65536 }
83 accr = accr + as2*W[i*K+k]; k=k+1 }
84 Cref[i]=accr; i=i+1 }
85
86 // === i8 path: unpack i8->i16, SIMD madd ===
87 let t0: i64 = now_ms()
88 var rep: i64 = 0
89 while rep < 8 {
90 i=0
91 while i<N {
92 var k: i64=0
93 while k<K { var b: i64 = Wi8[i*K+k] as i64; if b>=128 { b=b-256 } wrow[k*2]=(b) as u8; wrow[k*2+1]=(b>>8) as u8; k=k+1 }
94 z[0]=0; z[1]=0; z[2]=0; z[3]=0
95 k=0
96 while k<K { __i16x16_madd(acc, (A as i64 + k*2) as *u8, (wrow as i64 + k*2) as *u8); k=k+16 }
97 C8[i]=i32x8_hsum(acc)
98 i=i+1
99 }
100 rep=rep+1
101 }
102 let t1: i64 = now_ms()
103
104 // === 4-bit path: Q4_K-STYLE SCALE-THE-SUM (the nx_q4k_dot_simd technique) -- madd the RAW signed codes
105 // (no per-weight multiply), then apply the per-row scale ONCE to the hsum. This is the fix my naive v1 missed. ===
106 var rep2: i64 = 0
107 while rep2 < 8 {
108 i=0
109 while i<N {
110 var k: i64=0
111 while k<K {
112 let bi: i64 = i*K+k
113 var nib: i64 = 0
114 if (k & 1) == 0 { nib = Wi4[bi/2] as i64 & 15 } else { nib = (Wi4[bi/2] as i64 >> 4) & 15 }
115 if nib >= 8 { nib = nib - 16 } // signed code -8..7 (NO *scale here)
116 wrow[k*2]=(nib) as u8; wrow[k*2+1]=(nib>>8) as u8
117 k=k+1
118 }
119 z[0]=0; z[1]=0; z[2]=0; z[3]=0
120 k=0
121 while k<K { __i16x16_madd(acc, (A as i64 + k*2) as *u8, (wrow as i64 + k*2) as *u8); k=k+16 }
122 C4[i]=i32x8_hsum(acc) * scale4[i] // apply the scale ONCE to the sum
123 i=i+1
124 }
125 rep2=rep2+1
126 }
127 let t2: i64 = now_ms()
128
129 let ms8: i64 = t1-t0
130 let ms4: i64 = t2-t1
131
132 // correctness: i8 must be EXACT vs ref (scale 1); i4 within quant error (mean |C4-Cref| small vs |Cref|)
133 var maxd8: i64=0; var sumabsref: i64=0; var sumabsd4: i64=0
134 i=0
135 while i<N {
136 let d8: i64 = gp_iabs(C8[i]-Cref[i]); if d8>maxd8 { maxd8=d8 }
137 sumabsref = sumabsref + gp_iabs(Cref[i])
138 sumabsd4 = sumabsd4 + gp_iabs(C4[i]-Cref[i])
139 i=i+1
140 }
141 let relerr4_permille: i64 = (1000*sumabsd4) / (sumabsref+1)
142
143 gp_puts("shape: decode matvec K=" as *u8); gp_n(K); gp_puts(" N=" as *u8); gp_n(N); gp_puts(" (x8 reps each)\n" as *u8)
144 gp_puts("i8-weight path : " as *u8); gp_n(ms8); gp_puts(" ms (1.0 byte/weight)\n" as *u8)
145 gp_puts("4bit-weight path: " as *u8); gp_n(ms4); gp_puts(" ms (0.5 byte/weight)\n" as *u8)
146 gp_puts("i8 exact vs ref: maxdiff=" as *u8); gp_n(maxd8); gp_puts(" (must be 0)\n" as *u8)
147 gp_puts("4bit rel-error : " as *u8); gp_n(relerr4_permille); gp_puts(" permille (quant loss, expect < 100)\n\n" as *u8)
148
149 // TEETH
150 if maxd8 == 0 { st[0]=st[0]+1; gp_puts(" PASS i8 path bit-exact vs int16 reference\n" as *u8) } else { st[1]=st[1]+1; gp_puts(" FAIL i8 path not exact\n" as *u8) }
151 if relerr4_permille < 120 { st[0]=st[0]+1; gp_puts(" PASS 4bit path within quant error (usable)\n" as *u8) } else { st[1]=st[1]+1; gp_puts(" FAIL 4bit quant error too high\n" as *u8) }
152 // the DECISIVE measurement (not a pass/fail -- it just reports the verdict)
153 gp_puts("\n=== VERDICT ===\n" as *u8)
154 if ms4 < ms8 { gp_puts("4-bit is FASTER than i8 (" as *u8); gp_n(ms8-ms4); gp_puts(" ms) -> memory-saving wins; W4 kernel is worth building.\n" as *u8) }
155 else { gp_puts("4-bit is NOT faster (" as *u8); gp_n(ms4-ms8); gp_puts(" ms slower) -> scalar unpack overhead dominates; W4 needs a SIMD 4-bit-unpack primitive (MARLIN-style fused dequant) before it pays.\n" as *u8) }
156
157 gp_puts("\nPASS=" as *u8); gp_n(st[0]); gp_puts(" FAIL=" as *u8); gp_n(st[1]); gp_puts("\n" as *u8)
158 if st[1]>0 { gp_puts("GATE RED\n" as *u8); return 1 }
159 gp_puts("GATE GREEN -- W4 correctness proven; the timing verdict scopes the next step (measured, not assumed)\n" as *u8)
160 return 0
161}