nx_q4k_dot_simd.nx source
↩ module page · 165 lines · 6723 B
1// nx_q4k_dot_simd.nx -- SIMD-accelerated Q4_K fused dot (vpmaddwd), bit-exact vs nx_q4k_dot_row_col.
2//
3// Perf-exceed applied to the Qwen GEMM. Reformulation (exact integer algebra): per sub-block,
4// dot += Σ (d1*q - m0)*col = d1*(Σ q*col) - m0*(Σ col) = d1*sq - m0*sc
5// sq = Σ q*col via one vpmaddwd accumulate (q i16 0..15, col i16 Q10). Optimizations vs v1: read qs bytes
6// ONCE (extract lo+hi together), and take sc = Σ col from a PRECOMPUTED sub-block table (activation-only,
7// constant across the dot -- free in a GEMM). col: both Q10 i64 (unused now except pack) + i16-packed.
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_clock.nx"
22
23import "nx_q4k_dot_simd_lib.nx" // ds_pack4 / ds_hsum / nx_q4k_sc_precompute / nx_q4k_dot_simd (extracted 2026-09-02)
24
25func ds_emit(fd: i64, key: *u8, kl: i64, v: i64) -> i64 {
26 let line: *u8 = sys_mmap(64)
27 var lo: i64 = 0
28 var i: i64 = 0
29 while i < kl { line[lo] = key[i]; lo = lo + 1; i = i + 1 }
30 line[lo] = 0x3D; lo = lo + 1
31 let dec: *u8 = sys_mmap(32)
32 let nd: i64 = nx_strconv_format_i64(v, dec)
33 var k: i64 = 0
34 while k < nd { line[lo] = dec[k]; lo = lo + 1; k = k + 1 }
35 line[lo] = 0x0A; lo = lo + 1
36 return sys_write(fd, line, lo)
37}
38
39// Synthesise `nb` Q4_K blocks in place. A Q4_K block is 144 bytes: 2B d (fp16) + 2B dmin (fp16) +
40// 12B packed 6-bit scales + 128B of packed 4-bit quants. The PERF question -- is the SIMD dot faster
41// than the scalar one -- does not need real weights at all: both paths read the SAME bytes, so any
42// deterministic fill answers it exactly. This removes a 1.1GB model dependency from a timing bench.
43func ds_synth(buf: *u8, nb: i64) -> i64 {
44 var b: i64 = 0
45 while b < nb {
46 let sb: i64 = b * 144
47 buf[sb + 0] = 0 as u8
48 buf[sb + 1] = 60 as u8
49 buf[sb + 2] = 0 as u8
50 buf[sb + 3] = 52 as u8
51 var i: i64 = 0
52 while i < 12 { buf[sb + 4 + i] = ((b * 7 + i * 13) % 64) as u8; i = i + 1 }
53 i = 0
54 while i < 128 { buf[sb + 16 + i] = ((b * 31 + i * 17) % 256) as u8; i = i + 1 }
55 b = b + 1
56 }
57 return 0
58}
59
60// argv[1] = OPTIONAL gguf path (rule 17: never bake one). Absent or unopenable -> SYNTHETIC blocks, and
61// the run reports synthetic=1 so no reader can mistake a fixture number for a real-weights number.
62// The old build baked a /mnt/c laptop path and returned 30 on the NAS, writing zero bytes: a silent
63// failure that read as a hang. Results now also go to STDOUT, because nx_job_run captures stdout and a
64// file-only emitter looks empty even when it succeeded.
65func main(argc: i64, argv: *i64) -> i64 {
66 var buf: *u8 = 0 as *u8
67 var total: i64 = 0
68 var IN: i64 = 0
69 var w_off: i64 = 0
70 var synth: i64 = 1
71 if argc >= 2 {
72 let path: *u8 = argv[1] as *u8
73 let fd: i64 = sys_openat_rd(path)
74 if fd >= 0 {
75 let CAP: i64 = 1153433600
76 buf = sys_mmap(CAP)
77 var go: i64 = 1
78 while go == 1 {
79 let r: i64 = sys_read(fd, ((buf as i64) + total) as *u8, CAP - total)
80 if r <= 0 { go = 0 } else { total = total + r; if total >= CAP { go = 0 } }
81 }
82 sys_close(fd)
83 if total >= 100000000 {
84 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
85 if nx_gguf_parse(buf, total, hdr) == NX_GGUF_OK {
86 let qi: nx_int = nx_gguf_find_tensor(hdr, "blk.0.attn_q.weight" as *u8, 19)
87 if qi >= 0 {
88 let ti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, qi)
89 IN = ti.dim_0
90 w_off = hdr.data_off + ti.offset
91 synth = 0
92 }
93 }
94 }
95 }
96 }
97 var n_blocks: i64 = 0
98 if synth == 1 {
99 n_blocks = 14
100 IN = n_blocks * 256
101 total = n_blocks * 144
102 buf = sys_mmap(total)
103 ds_synth(buf, n_blocks)
104 w_off = 0
105 } else {
106 n_blocks = IN / 256
107 if w_off + n_blocks * 144 > total { return 63 }
108 }
109
110 let col: *i64 = sys_mmap(IN * 8) as *i64
111 let col_i16: *i64 = sys_mmap(IN * 2) as *i64
112 var i: i64 = 0
113 while i < IN { col[i] = 1024 + (i - (i / 5) * 5) * 256; i = i + 1 }
114 var jj: i64 = 0
115 while jj < IN / 4 {
116 col_i16[jj] = ds_pack4(col[jj * 4], col[jj * 4 + 1], col[jj * 4 + 2], col[jj * 4 + 3])
117 jj = jj + 1
118 }
119 let sc_pre: *i64 = sys_mmap(n_blocks * 8 * 8) as *i64
120 nx_q4k_sc_precompute(col, n_blocks, sc_pre)
121
122 let qpk: *i64 = sys_mmap(64) as *i64
123 let qhi: *i64 = sys_mmap(64) as *i64
124 let acc: *i64 = sys_mmap(32) as *i64
125 let it: *NxQ4KBlockIter = nx_q4k_iter_alloc()
126
127 let scalar_dot: i64 = nx_q4k_dot_row_col(buf, w_off, n_blocks, col, it)
128 let simd_dot: i64 = nx_q4k_dot_simd(buf, w_off, n_blocks, col_i16, qpk, qhi, acc, sc_pre)
129
130 let IT: i64 = 20000
131 let t0: i64 = nx_clock_monotonic_ns()
132 var s1: i64 = 0
133 var k: i64 = 0
134 while k < IT { s1 = s1 + nx_q4k_dot_row_col(buf, w_off, n_blocks, col, it); k = k + 1 }
135 let t1: i64 = nx_clock_monotonic_ns()
136 let scalar_ns: i64 = t1 - t0
137
138 let t2: i64 = nx_clock_monotonic_ns()
139 var s2: i64 = 0
140 k = 0
141 while k < IT { s2 = s2 + nx_q4k_dot_simd(buf, w_off, n_blocks, col_i16, qpk, qhi, acc, sc_pre); k = k + 1 }
142 let t3: i64 = nx_clock_monotonic_ns()
143 let simd_ns: i64 = t3 - t2
144
145 let ofd: i64 = sys_openat_wr("/tmp/q4k_simd.txt" as *u8, 0x1a4)
146 if ofd >= 0 {
147 ds_emit(ofd, "scalar_dot" as *u8, 10, scalar_dot)
148 ds_emit(ofd, "simd_dot" as *u8, 8, simd_dot)
149 ds_emit(ofd, "scalar_ns" as *u8, 9, scalar_ns)
150 ds_emit(ofd, "simd_ns" as *u8, 7, simd_ns)
151 if simd_ns > 0 { ds_emit(ofd, "speedup_x100" as *u8, 12, scalar_ns * 100 / simd_ns) }
152 sys_close(ofd)
153 }
154 ds_emit(1, "synthetic" as *u8, 9, synth)
155 ds_emit(1, "n_blocks" as *u8, 8, n_blocks)
156 ds_emit(1, "scalar_dot" as *u8, 10, scalar_dot)
157 ds_emit(1, "simd_dot" as *u8, 8, simd_dot)
158 ds_emit(1, "bit_exact" as *u8, 9, 1 - ((scalar_dot - simd_dot) * (scalar_dot - simd_dot)))
159 ds_emit(1, "scalar_ns" as *u8, 9, scalar_ns)
160 ds_emit(1, "simd_ns" as *u8, 7, simd_ns)
161 if simd_ns > 0 { ds_emit(1, "speedup_x100" as *u8, 12, scalar_ns * 100 / simd_ns) }
162 if simd_dot != scalar_dot { return 80 }
163 if simd_ns >= scalar_ns { return 81 }
164 return 0
165}