nx_q8_0_simd_gate.nx source
↩ module page · 164 lines · 5702 B
1// nx_q8_0_simd_gate.nx -- does the SIMD __f32_i8dot32 fused Q8_0 dequant-dot
2// WIN? Q8_0 has NO nibble unpack (int8 directly), so the intrinsic applies
3// cleanly. lm_head is Q8_0 (136M values, 143ms in the forward, memory-bound).
4// Threaded Q8_0-fused vs F32-pool at ~lm_head shape (m=1, k=896, big n).
5// Ratio host-noise-tolerant. NO forward changes.
6//
7// 1 Q8_0-fused == F32-pool bit-exact (exact-int regime)
8// 2 threaded ratio: WIN prints if >= 1.0x
9//
10// lineage_id: q8_0_simd_gate_v1
11import "nx_syscalls.nx"
12import "nx_tier.nx"
13import "nx_le.nx"
14import "nx_f32.nx"
15import "nx_f32_cvt.nx"
16import "nx_q8_0_to_f32.nx"
17import "nx_thread_pool.nx"
18import "nx_f32_matmul_t.nx"
19import "nx_fmt.nx"
20
21const QK: i64 = 896
22const QN: i64 = 32768
23const QREPS: i64 = 60
24const Q8B: i64 = 34
25const Q8V: i64 = 32
26
27struct Q8Ctx { qbuf: i64, aptr: i64, cptr: i64, k: i64, n: i64, jlo: i64, jhi: i64 }
28const Q8CTX_BYTES: i64 = 56
29
30func q_lcg(s: i64) -> i64 { var v: i64 = s * 1103515245 + 12345; v = v & 2147483647; return v }
31
32func q_fill_weight(w: *u8, seed: i64) -> i64 {
33 let bpr: i64 = (QK / Q8V) * Q8B
34 var s: i64 = seed
35 var r: i64 = 0
36 while r < QN {
37 var b: i64 = 0
38 while b < QK / Q8V {
39 let off: i64 = r * bpr + b * Q8B
40 w[off + 0] = 0x00 as u8
41 w[off + 1] = 0x2C as u8 // f16 d ~ 0.0625
42 var q: i64 = 0
43 while q < 32 { s = q_lcg(s); w[off + 2 + q] = ((s % 17) - 8) as u8; q = q + 1 } // int8 -8..8
44 b = b + 1
45 }
46 r = r + 1
47 }
48 return 0
49}
50func q_fill_a(p: *i64, count: i64, seed: i64) -> i64 {
51 var s: i64 = seed
52 var i: i64 = 0
53 while i < count { s = q_lcg(s); p[i] = nx_i32_to_f32((s % 9) - 4); i = i + 1 }
54 return 0
55}
56
57func q8_0_dot(qbuf: *u8, qoff: i64, Apk: *u8, k: i64) -> i64 {
58 let nblk: i64 = k / Q8V
59 let qpb: i64 = qbuf as i64
60 let apb: i64 = Apk as i64
61 var acc: i64 = 0
62 var b: i64 = 0
63 while b < nblk {
64 let boff: i64 = qoff + b * Q8B
65 let d_f32: i64 = nx_f16_to_f32(nx_le_read_u16(qbuf, boff))
66 let raw: i64 = __f32_i8dot32((qpb + boff + 2) as *u8, (apb + b * Q8V * 4) as *u8)
67 acc = __f32_add(acc, __f32_mul(d_f32, raw))
68 b = b + 1
69 }
70 return acc
71}
72
73func q8_task(ctx_i: i64) -> i64 {
74 let cx: *Q8Ctx = ctx_i as *Q8Ctx
75 let bpr: i64 = (cx.k / Q8V) * Q8B
76 let qbuf: *u8 = cx.qbuf as *u8
77 let Apk: *u8 = cx.aptr as *u8
78 let C: *i64 = cx.cptr as *i64
79 var j: i64 = cx.jlo
80 while j < cx.jhi {
81 C[j] = q8_0_dot(qbuf, j * bpr, Apk, cx.k)
82 j = j + 1
83 }
84 return 0
85}
86
87func q8_pool_matmul(pool: *NxThreadPool, qbuf: *u8, Apk: *u8, C: *i64, k: i64, n: i64) -> i64 {
88 var bands: i64 = pool.n_workers
89 if bands > n { bands = n }
90 if bands < 1 { bands = 1 }
91 let ctxs: *u8 = sys_mmap(bands * Q8CTX_BYTES)
92 let cpb: i64 = (n + bands - 1) / bands
93 let done_before: i64 = nx_pool_n_completed(pool)
94 var b: i64 = 0
95 while b < bands {
96 let cx: *Q8Ctx = ((ctxs as i64) + b * Q8CTX_BYTES) as *Q8Ctx
97 cx.qbuf = qbuf as i64; cx.aptr = Apk as i64; cx.cptr = C as i64
98 cx.k = k; cx.n = n; cx.jlo = b * cpb
99 var jhi: i64 = (b + 1) * cpb
100 if jhi > n { jhi = n }
101 cx.jhi = jhi
102 nx_pool_submit(pool, q8_task, cx as i64)
103 b = b + 1
104 }
105 nx_pool_wait(pool, done_before + bands)
106 sys_munmap(ctxs, bands * Q8CTX_BYTES)
107 return 0
108}
109
110func q_nl() -> i64 { fmt_puts("\n" as *u8); return 0 }
111
112func main() -> i64 {
113 let bpr: i64 = (QK / Q8V) * Q8B
114 let W: *u8 = sys_mmap(QN * bpr)
115 let A: *i64 = sys_mmap(QK * 8) as *i64
116 let WF: *i64 = sys_mmap(QN * QK * 8) as *i64
117 let Cref: *i64 = sys_mmap(QN * 8) as *i64
118 let Cfus: *i64 = sys_mmap(QN * 8) as *i64
119 q_fill_weight(W, 20260708)
120 q_fill_a(A, QK, 4242)
121 var mr: i64 = 0
122 while mr < QN { nx_q8_0_to_f32(W, mr * bpr, QK, ((WF as i64) + mr * QK * 8) as *i64); mr = mr + 1 }
123 let Apk: *u8 = sys_mmap(QK * 4)
124 var pj: i64 = 0
125 while pj < QK {
126 let bits: i64 = A[pj]
127 Apk[pj*4+0] = bits as u8; Apk[pj*4+1] = (bits>>8) as u8; Apk[pj*4+2] = (bits>>16) as u8; Apk[pj*4+3] = (bits>>24) as u8
128 pj = pj + 1
129 }
130
131 let pool: *NxThreadPool = nx_pool_new(0, 0)
132
133 nx_f32_matmul_t_pool(pool, A, WF, Cref, 1, QK, QN)
134 q8_pool_matmul(pool, W, Apk, Cfus, QK, QN)
135 var ok1: i64 = 1
136 var c: i64 = 0
137 while c < QN { if Cref[c] != Cfus[c] { ok1 = 0; c = QN } else { c = c + 1 } }
138 if ok1 != 1 { fmt_puts("Q8T 1 EXACT FAIL"); q_nl(); nx_pool_shutdown(pool); return 11 }
139 fmt_puts("Q8T 1 F32POOL==Q8FUSEDPOOL EXACT OK"); q_nl()
140
141 let t0: i64 = sys_now_us()
142 var r0: i64 = 0
143 while r0 < QREPS { nx_f32_matmul_t_pool(pool, A, WF, Cref, 1, QK, QN); r0 = r0 + 1 }
144 let us_f32: i64 = sys_now_us() - t0
145 let t1: i64 = sys_now_us()
146 var r1: i64 = 0
147 while r1 < QREPS { q8_pool_matmul(pool, W, Apk, Cfus, QK, QN); r1 = r1 + 1 }
148 let us_q8: i64 = sys_now_us() - t1
149 nx_pool_shutdown(pool)
150
151 var uf: i64 = us_f32
152 if uf < 1 { uf = 1 }
153 var uq: i64 = us_q8
154 if uq < 1 { uq = 1 }
155 let macs: i64 = QK * QN * QREPS
156 fmt_puts("F32_pool_us="); fmt_putn(us_f32); fmt_puts(" mflops="); fmt_putn(2 * macs / uf); q_nl()
157 fmt_puts("Q8_0_fused_pool_us="); fmt_putn(us_q8); fmt_puts(" mflops="); fmt_putn(2 * macs / uq); q_nl()
158 let sx100: i64 = uf * 100 / uq
159 fmt_puts("q8_vs_f32_threaded_x100="); fmt_putn(sx100); q_nl()
160 if sx100 >= 100 { fmt_puts("VERDICT: SIMD Q8_0 dequant-dot WINS -- wire lm_head to Q8_0") } else { fmt_puts("VERDICT: still slower") }
161 q_nl()
162 fmt_puts("Q8_0_SIMD_GATE DONE"); q_nl()
163 return 0
164}