nx_q5_0_threaded_gate.nx source
↩ module page · 220 lines · 7754 B
1// nx_q5_0_threaded_gate.nx -- THE DECISIVE test: does THREADING flip the
2// fused Q5_0 dequant-dot from a loss to a win?
3//
4// Single-thread (nx_q5_0_fused_dot_gate) the fused kernel LOST 2x: compute-
5// bound (dequant overhead > MAC), memory not saturated. But at 16 cores the
6// F32 matmul saturates RAM bandwidth (16 x ~2GB/s > ~20-40GB/s), so reading
7// 11.6x fewer bytes (Q5_0 3MB vs F32 34.9MB) could win THERE even with a
8// scalar dequant. Both paths POOL-threaded, SSE dot (__f32_mul/__f32_add),
9// same run => ratio is host-noise-tolerant. THIS decides whether quantized-
10// read is viable on our CPU substrate or needs SIMD-dequant intrinsics first.
11// NO forward changes.
12//
13// Checks:
14// 1 Q5_0-fused-pool == F32-pool (mmt_range), bit-exact
15// 2 threaded ratio: fused vs F32 at FFN shape (report; not a hard floor --
16// this is a DECISION measurement, GREEN prints the verdict either way)
17//
18// lineage_id: q5_0_threaded_gate_v1
19
20import "nx_syscalls.nx"
21import "nx_tier.nx"
22import "nx_le.nx"
23import "nx_f32.nx"
24import "nx_f32_cvt.nx"
25import "nx_q5_0_to_f32.nx"
26import "nx_thread_pool.nx"
27import "nx_f32_matmul_t.nx"
28import "nx_fmt.nx"
29
30const QK: i64 = 896
31const QN: i64 = 4864
32const QREPS: i64 = 120
33const Q5B: i64 = 22
34const Q5V: i64 = 32
35
36struct Q5Ctx {
37 qbuf: i64,
38 aptr: i64,
39 cptr: i64,
40 k: i64,
41 n: i64,
42 jlo: i64,
43 jhi: i64,
44}
45const Q5CTX_BYTES: i64 = 56
46
47func q_lcg(s: i64) -> i64 { var v: i64 = s * 1103515245 + 12345; v = v & 2147483647; return v }
48
49func q_fill_weight(w: *u8, seed: i64) -> i64 {
50 let bpr: i64 = (QK / Q5V) * Q5B
51 var s: i64 = seed
52 var r: i64 = 0
53 while r < QN {
54 var b: i64 = 0
55 while b < QK / Q5V {
56 let off: i64 = r * bpr + b * Q5B
57 w[off + 0] = 0x00 as u8
58 w[off + 1] = 0x2C as u8
59 var i: i64 = 0
60 while i < 4 { s = q_lcg(s); w[off + 2 + i] = (s & 255) as u8; i = i + 1 }
61 var q: i64 = 0
62 while q < 16 { s = q_lcg(s); w[off + 6 + q] = (s & 255) as u8; q = q + 1 }
63 b = b + 1
64 }
65 r = r + 1
66 }
67 return 0
68}
69func q_fill_a(p: *i64, count: i64, seed: i64) -> i64 {
70 var s: i64 = seed
71 var i: i64 = 0
72 while i < count { s = q_lcg(s); p[i] = nx_i32_to_f32((s % 9) - 4); i = i + 1 }
73 return 0
74}
75
76// Build the 80-byte __q5_unpack32 mask constants ONCE (module static).
77static G_Q5_CONSTS: i64
78func q5_consts() -> *u8 {
79 if G_Q5_CONSTS == 0 {
80 let cc: *u8 = sys_mmap(80)
81 var i: i64 = 0
82 while i < 16 { cc[i] = 0x0F as u8; i = i + 1 } // c_0F
83 i = 0
84 while i < 8 { cc[16+i] = 0 as u8; i = i + 1 }
85 while i < 16 { cc[16+i] = 1 as u8; i = i + 1 } // pshuf_lo
86 i = 0
87 while i < 8 { cc[32+i] = 2 as u8; i = i + 1 }
88 while i < 16 { cc[32+i] = 3 as u8; i = i + 1 } // pshuf_hi
89 i = 0
90 while i < 8 { cc[48+i] = (1 << i) as u8; cc[48+8+i] = (1 << i) as u8; i = i + 1 } // bitmask
91 i = 0
92 while i < 16 { cc[64+i] = 0x10 as u8; i = i + 1 } // c_10
93 G_Q5_CONSTS = cc as i64
94 }
95 return G_Q5_CONSTS as *u8
96}
97
98// FULLY SIMD Q5_0 dequant-dot: __q5_unpack32 (SSE nibble+qh unpack -> 32 int8)
99// then __f32_i8dot32 (SSE convert+dot), x d. qh at block+2, qs at block+6 are
100// contiguous -> qhqs = block+2.
101func q5_0_dot(qbuf: *u8, qoff: i64, Apk: *u8, k: i64, i8scr: *u8) -> i64 {
102 let nblk: i64 = k / Q5V
103 let apb: i64 = Apk as i64
104 let qpb: i64 = qbuf as i64
105 let consts: *u8 = q5_consts()
106 var acc: i64 = 0
107 var b: i64 = 0
108 while b < nblk {
109 let boff: i64 = qoff + b * Q5B
110 let d_f32: i64 = nx_f16_to_f32(nx_le_read_u16(qbuf, boff + 0))
111 __q5_unpack32((qpb + boff + 2) as *u8, i8scr, consts)
112 let raw: i64 = __f32_i8dot32(i8scr, (apb + b * Q5V * 4) as *u8)
113 acc = __f32_add(acc, __f32_mul(d_f32, raw))
114 b = b + 1
115 }
116 return acc
117}
118
119func q5_task(ctx_i: i64) -> i64 {
120 let cx: *Q5Ctx = ctx_i as *Q5Ctx
121 let i8scr: *u8 = sys_mmap(Q5V)
122 let bpr: i64 = (cx.k / Q5V) * Q5B
123 let qbuf: *u8 = cx.qbuf as *u8
124 let Apk: *u8 = cx.aptr as *u8
125 let C: *i64 = cx.cptr as *i64
126 var j: i64 = cx.jlo
127 while j < cx.jhi {
128 C[j] = q5_0_dot(qbuf, j * bpr, Apk, cx.k, i8scr)
129 j = j + 1
130 }
131 sys_munmap(i8scr, Q5V)
132 return 0
133}
134
135func q5_pool_matmul(pool: *NxThreadPool, qbuf: *u8, Apk: *u8, C: *i64, k: i64, n: i64) -> i64 {
136 var bands: i64 = pool.n_workers
137 if bands > n { bands = n }
138 if bands < 1 { bands = 1 }
139 let ctxs: *u8 = sys_mmap(bands * Q5CTX_BYTES)
140 let cpb: i64 = (n + bands - 1) / bands
141 let done_before: i64 = nx_pool_n_completed(pool)
142 var b: i64 = 0
143 while b < bands {
144 let cx: *Q5Ctx = ((ctxs as i64) + b * Q5CTX_BYTES) as *Q5Ctx
145 cx.qbuf = qbuf as i64
146 cx.aptr = Apk as i64
147 cx.cptr = C as i64
148 cx.k = k
149 cx.n = n
150 cx.jlo = b * cpb
151 var jhi: i64 = (b + 1) * cpb
152 if jhi > n { jhi = n }
153 cx.jhi = jhi
154 nx_pool_submit(pool, q5_task, cx as i64)
155 b = b + 1
156 }
157 nx_pool_wait(pool, done_before + bands)
158 sys_munmap(ctxs, bands * Q5CTX_BYTES)
159 return 0
160}
161
162func q_nl() -> i64 { fmt_puts("\n" as *u8); return 0 }
163
164func main() -> i64 {
165 let bpr: i64 = (QK / Q5V) * Q5B
166 let W: *u8 = sys_mmap(QN * bpr)
167 let A: *i64 = sys_mmap(QK * 8) as *i64
168 let WF: *i64 = sys_mmap(QN * QK * 8) as *i64
169 let Cref: *i64 = sys_mmap(QN * 8) as *i64
170 let Cfus: *i64 = sys_mmap(QN * 8) as *i64
171 q_fill_weight(W, 20260708)
172 q_fill_a(A, QK, 4242)
173 var mr: i64 = 0
174 while mr < QN { nx_q5_0_to_f32(W, mr * bpr, QK, ((WF as i64) + mr * QK * 8) as *i64); mr = mr + 1 }
175 // pack A (i64-slot f32) -> contiguous 4-byte f32 for __f32_i8dot32.
176 let Apk: *u8 = sys_mmap(QK * 4)
177 var pj: i64 = 0
178 while pj < QK {
179 let bits: i64 = A[pj]
180 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
181 pj = pj + 1
182 }
183
184 let pool: *NxThreadPool = nx_pool_new(0, 0)
185
186 // ---- 1: bit-exact (F32-pool vs Q5_0-fused-pool) ----
187 nx_f32_matmul_t_pool(pool, A, WF, Cref, 1, QK, QN)
188 q5_pool_matmul(pool, W, Apk, Cfus, QK, QN)
189 var ok1: i64 = 1
190 var c: i64 = 0
191 while c < QN { if Cref[c] != Cfus[c] { ok1 = 0; c = QN } else { c = c + 1 } }
192 if ok1 != 1 { fmt_puts("Q5T 1 EXACT FAIL"); q_nl(); nx_pool_shutdown(pool); return 11 }
193 fmt_puts("Q5T 1 F32POOL==Q5FUSEDPOOL EXACT OK"); q_nl()
194
195 // ---- 2: threaded speed ----
196 let t0: i64 = sys_now_us()
197 var r0: i64 = 0
198 while r0 < QREPS { nx_f32_matmul_t_pool(pool, A, WF, Cref, 1, QK, QN); r0 = r0 + 1 }
199 let us_f32: i64 = sys_now_us() - t0
200
201 let t1: i64 = sys_now_us()
202 var r1: i64 = 0
203 while r1 < QREPS { q5_pool_matmul(pool, W, Apk, Cfus, QK, QN); r1 = r1 + 1 }
204 let us_q5: i64 = sys_now_us() - t1
205 nx_pool_shutdown(pool)
206
207 var uf: i64 = us_f32
208 if uf < 1 { uf = 1 }
209 var uq: i64 = us_q5
210 if uq < 1 { uq = 1 }
211 let macs: i64 = QK * QN * QREPS
212 fmt_puts("F32_pool_us="); fmt_putn(us_f32); fmt_puts(" mflops="); fmt_putn(2 * macs / uf); q_nl()
213 fmt_puts("Q5_0_fused_pool_us="); fmt_putn(us_q5); fmt_puts(" mflops="); fmt_putn(2 * macs / uq); q_nl()
214 let sx100: i64 = uf * 100 / uq
215 fmt_puts("q5_vs_f32_threaded_x100="); fmt_putn(sx100); q_nl()
216 if sx100 >= 100 { fmt_puts("VERDICT: THREADING FLIPS IT -- fused Q5_0 WINS (memory-bound at 16c) -> wire it") } else { fmt_puts("VERDICT: still compute-bound even threaded -> needs SIMD dequant intrinsics (CPU-Marlin)") }
217 q_nl()
218 fmt_puts("Q5_0_THREADED_GATE DONE"); q_nl()
219 return 0
220}