nx_q4k_gemm_mt.nx source
↩ module page · 145 lines · 5112 B
1// nx_q4k_gemm_mt.nx -- multi-threaded integer GEMM: wire nx_thread into the Q4_K integer dot.
2//
3// sd-server -> Nishi migration (cross-hw census PARTIAL -> HAVE: threading). The GEMM is embarrassingly
4// parallel across output neurons: each worker thread computes a disjoint output range via nx_q4k_dot_row_col
5// (its own iterator), joined by a done-flag (matches the nx_thread_spawn smoke: raw *i64 ctx, spin-join).
6// Works on ANY multicore CPU. Verified bit-exact vs single-thread + measured speedup.
7// ctx (flat *i64, 9 slots): [0]buf [1]w_off [2]o_start [3]o_end [4]n_blocks [5]rstride [6]col [7]out [8]done
8// license_tier: ORIGINAL
9import "nx_syscalls.nx"
10import "nx_tier.nx"
11import "nx_le.nx"
12import "nx_strconv.nx"
13import "nx_tensor.nx"
14import "nx_gguf.nx"
15import "nx_gguf_load.nx"
16import "nx_gguf_meta.nx"
17import "nx_placement.nx"
18import "nx_gguf_load_lazy.nx"
19import "nx_q4k_matmul.nx"
20import "nx_dequant_iter.nx"
21import "nx_thread.nx"
22import "nx_clock.nx"
23
24func mt_worker(ctx_ptr: *u8) -> i64 {
25 let c: *i64 = ctx_ptr as *i64
26 let it: *NxQ4KBlockIter = nx_q4k_iter_alloc()
27 let bufp: *u8 = c[0] as *u8
28 let colp: *i64 = c[6] as *i64
29 let outp: *i64 = c[7] as *i64
30 let w_off: i64 = c[1]
31 let n_blocks: i64 = c[4]
32 let rstride: i64 = c[5]
33 var o: i64 = c[2]
34 while o < c[3] {
35 let dot: i64 = nx_q4k_dot_row_col(bufp, w_off + o * rstride, n_blocks, colp, it)
36 outp[o] = nx_q4km_q20_to_q10(dot)
37 o = o + 1
38 }
39 c[8] = 1
40 return 0
41}
42
43func mt_emit(fd: i64, key: *u8, kl: i64, v: i64) -> i64 {
44 let line: *u8 = sys_mmap(64)
45 var lo: i64 = 0
46 var i: i64 = 0
47 while i < kl { line[lo] = key[i]; lo = lo + 1; i = i + 1 }
48 line[lo] = 0x3D; lo = lo + 1
49 let dec: *u8 = sys_mmap(32)
50 let nd: i64 = nx_strconv_format_i64(v, dec)
51 var k: i64 = 0
52 while k < nd { line[lo] = dec[k]; lo = lo + 1; k = k + 1 }
53 line[lo] = 0x0A; lo = lo + 1
54 return sys_write(fd, line, lo)
55}
56
57func main() -> i64 {
58 let path: *u8 = "/mnt/c/Users/elder/elder-ai-platform/models/unified/text_encoder/Huihui-Qwen3-4B-Instruct-2507-abliterated-Q4_K_M.gguf" as *u8
59 let fd: i64 = sys_openat_rd(path)
60 if fd < 0 { return 30 }
61 let CAP: i64 = 1153433600
62 let buf: *u8 = sys_mmap(CAP)
63 var total: i64 = 0
64 var go: i64 = 1
65 while go == 1 {
66 let r: i64 = sys_read(fd, ((buf as i64) + total) as *u8, CAP - total)
67 if r <= 0 { go = 0 } else { total = total + r; if total >= CAP { go = 0 } }
68 }
69 sys_close(fd)
70 if total < 100000000 { return 31 }
71 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
72 if nx_gguf_parse(buf, total, hdr) != NX_GGUF_OK { return 40 }
73 let qi: nx_int = nx_gguf_find_tensor(hdr, "blk.0.attn_q.weight" as *u8, 19)
74 if qi < 0 { return 60 }
75 let ti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, qi)
76 let IN: i64 = ti.dim_0
77 let w_off: i64 = hdr.data_off + ti.offset
78 let n_blocks: i64 = IN / 256
79 let rstride: i64 = n_blocks * 144
80 let OD: i64 = 1024
81 if w_off + OD * rstride > total { return 63 }
82
83 let col: *i64 = sys_mmap(IN * 8) as *i64
84 var i: i64 = 0
85 while i < IN { col[i] = 1024 + (i - (i / 5) * 5) * 256; i = i + 1 }
86
87 // single-threaded baseline
88 let out_s: *i64 = sys_mmap(OD * 8) as *i64
89 let it0: *NxQ4KBlockIter = nx_q4k_iter_alloc()
90 let t0: i64 = nx_clock_monotonic_ns()
91 var o: i64 = 0
92 while o < OD { out_s[o] = nx_q4km_q20_to_q10(nx_q4k_dot_row_col(buf, w_off + o * rstride, n_blocks, col, it0)); o = o + 1 }
93 let t1: i64 = nx_clock_monotonic_ns()
94 let single_ns: i64 = t1 - t0
95
96 // multi-threaded
97 let N: i64 = 4
98 let per: i64 = OD / N
99 let out_m: *i64 = sys_mmap(OD * 8) as *i64
100 let ctxs: *i64 = sys_mmap(N * 8) as *i64
101 let t2: i64 = nx_clock_monotonic_ns()
102 var w: i64 = 0
103 while w < N {
104 let c: *i64 = sys_mmap(9 * 8) as *i64
105 c[0] = buf as i64
106 c[1] = w_off
107 c[2] = w * per
108 c[3] = (w + 1) * per
109 c[4] = n_blocks
110 c[5] = rstride
111 c[6] = col as i64
112 c[7] = out_m as i64
113 c[8] = 0
114 ctxs[w] = c as i64
115 let tid: i64 = nx_thread_spawn_fn(mt_worker, c as *u8, 1048576)
116 if tid <= 0 { return 70 }
117 w = w + 1
118 }
119 w = 0
120 while w < N {
121 let c: *i64 = ctxs[w] as *i64
122 while c[8] == 0 { nx_thread_yield() }
123 w = w + 1
124 }
125 let t3: i64 = nx_clock_monotonic_ns()
126 let mt_ns: i64 = t3 - t2
127
128 var mism: i64 = 0
129 o = 0
130 while o < OD { if out_m[o] != out_s[o] { mism = mism + 1 } o = o + 1 }
131
132 let ofd: i64 = sys_openat_wr("/tmp/zimg_gemm_mt.txt" as *u8, 0x1a4)
133 if ofd >= 0 {
134 mt_emit(ofd, "OD" as *u8, 2, OD)
135 mt_emit(ofd, "threads" as *u8, 7, N)
136 mt_emit(ofd, "single_ns" as *u8, 9, single_ns)
137 mt_emit(ofd, "mt_ns" as *u8, 5, mt_ns)
138 if mt_ns > 0 { mt_emit(ofd, "speedup_x100" as *u8, 12, single_ns * 100 / mt_ns) }
139 mt_emit(ofd, "mismatch" as *u8, 8, mism)
140 sys_close(ofd)
141 }
142 if mism > 0 { return 80 }
143 if mt_ns >= single_ns { return 81 }
144 return 0
145}