nx_q4k_simd2_gate.nx source
↩ module page · 203 lines · 11977 B
1// nx_q4k_simd2_gate.nx -- the REFEREE for __q4k_unpack32s (the nx_cc intrinsic) and nx_q4k_dot_simd2 (2026-09-02, LM4c).
2// The intrinsic is proven at TWO levels. (1) LANE LEVEL: on packed bytes whose nibbles and scales are chosen so every
3// expected lane is known, out[k] must equal (byte_k & 15) * sc_lo and out[32+k] must equal (byte_k >> 4) * sc_hi for all
4// 64 lanes, including the extreme lane value 15*63 = 945 (no lane carry) and a zero scale; and a run over 5000 pseudo-random
5// byte blocks must agree lane-for-lane with the scalar spread+multiply nx_q4k_dot_simd_lib already trusts (ds_spread4).
6// (2) DOT LEVEL: on 200 pseudo-random Q4_K super-block rows with real-shaped d, dmin, scales and mins, nx_q4k_dot_simd2 must
7// return EXACTLY what nx_q4k_dot_simd returns (both are exact integer functions of the same bytes) and exactly what the
8// scalar reference nx_q4k_dot_row_col returns. Neg-controls: a corrupted lane must NOT match (the lane tooth can fail) and a
9// flipped nibble byte must change the dot (the dot tooth can fail). A timing line reports simd2 vs simd over the same rows
10// (a number, never a verdict -- the decode benchmark is the speed ruler).
11// nx_q4k_simd2_gate (no args; needs a compiler that knows __q4k_unpack32s, else it does not build)
12// license_tier: ORIGINAL. Writes nothing. No hw writes.
13import "nx_syscalls.nx"
14import "nx_tier.nx"
15import "nx_le.nx"
16import "nx_tensor.nx"
17import "nx_gguf.nx"
18import "nx_gguf_load.nx"
19import "nx_dequant_iter.nx"
20import "nx_q4k_matmul.nx"
21import "nx_clock.nx"
22import "nx_gate_verdict.nx"
23import "nx_q4k_dot_simd_lib.nx"
24import "nx_q4k_dot_simd2_lib.nx"
25import "nx_thread_pool.nx"
26import "nx_nofloat_llm.nx"
27import "nx_nofloat_q4k.nx" // the pooled fused GEMM for the scaling probe
28
29const G2_LCG_A: i64 = 6364136223846793005
30const G2_LCG_C: i64 = 1442695040888963407
31const G2_LCG_MASK: i64 = 0x7FFFFFFFFFFFFFFF
32const G2_LCG_SEED: i64 = 424242424242
33const G2_BYTES: i64 = 32
34const G2_LANES: i64 = 64
35const G2_LANE_MAX: i64 = 945 // 15 * 63
36const G2_SC_MAX: i64 = 63
37const G2_RAND_BLOCKS: i64 = 5000
38const G2_ROWS: i64 = 200
39const G2_NBLK: i64 = 14 // 3584-wide row: the 7B attention shape
40const G2_SB: i64 = 144
41const G2_COL_MAX: i64 = 32767
42const G2_TIMING_ITERS: i64 = 2000
43const G2_QBUF_BYTES: i64 = 128
44const G2_NIBBLE_MASK: i64 = 15
45const G2_NIB_SHIFT: i64 = 4
46const G2_SC_HI_SHIFT: i64 = 16
47
48static g_g2_st: i64
49func g2_lcg() -> i64 { g_g2_st = (g_g2_st*G2_LCG_A + G2_LCG_C) & G2_LCG_MASK; return g_g2_st >> 16 }
50
51// read lane k (i16, sign-extended) of a packed i16 buffer
52func g2_lane(p: *u8, k: i64) -> i64 { var v: i64 = nx_le_read_u16(p, k*2); if v >= 32768 { v = v - 65536 } return v }
53
54// lane-level check of one unpack: 1 when all 64 lanes equal the arithmetic expectation
55func g2_lanes_ok(qs: *u8, out: *u8, sclo: i64, schi: i64) -> i64 {
56 var ok: i64 = 1
57 var k: i64 = 0
58 while k < G2_BYTES {
59 let b: i64 = qs[k] & 0xFF
60 if g2_lane(out, k) != (b & G2_NIBBLE_MASK) * sclo { ok = 0 }
61 if g2_lane(out, k + G2_BYTES) != (b >> G2_NIB_SHIFT) * schi { ok = 0 }
62 k = k + 1
63 }
64 return ok
65}
66
67// the scalar path the shipping lib trusts: spread + multiply, written to a second buffer for lane comparison
68func g2_scalar_unpack(qs: *u8, out: *i64, sclo: i64, schi: i64) -> i64 {
69 var j: i64 = 0
70 while j < 4 {
71 let w: i64 = nx_le_read_u64(qs, j * 8)
72 let lo: i64 = w & DS_M_NIB
73 let hi: i64 = (w >> DS_NIB_BITS) & DS_M_NIB
74 out[j + j] = ds_spread4(lo) * sclo
75 out[j + j + 1] = ds_spread4(lo >> DS_HALF_BITS) * sclo
76 out[8 + j + j] = ds_spread4(hi) * schi
77 out[8 + j + j + 1] = ds_spread4(hi >> DS_HALF_BITS) * schi
78 j = j + 1
79 }
80 return 0
81}
82func g2_bufs_equal(a: *u8, b: *u8, n: i64) -> i64 { var i: i64 = 0; while i < n { if a[i] != b[i] { return 0 } i = i + 1 } return 1 }
83
84// one pseudo-random Q4_K super-block row (nblk super-blocks) with real-shaped fields
85func g2_fill_row(buf: *u8, nblk: i64) -> i64 {
86 var b: i64 = 0
87 while b < nblk {
88 let base: i64 = b * G2_SB
89 // d, dmin: random f16 with a small positive exponent (values ~ 2^-14 .. 2^3), sign bit clear
90 let dexp: i64 = 1 + (g2_lcg() % 18); let dman: i64 = g2_lcg() % 1024
91 let d16: i64 = (dexp << 10) | dman
92 buf[base] = (d16 & 0xFF) as u8; buf[base+1] = (d16 >> 8) as u8
93 let mexp: i64 = 1 + (g2_lcg() % 14); let mman: i64 = g2_lcg() % 1024
94 let m16: i64 = (mexp << 10) | mman
95 buf[base+2] = (m16 & 0xFF) as u8; buf[base+3] = (m16 >> 8) as u8
96 // scales: every 6-bit sub-block scale is drawn NON-ZERO so the neg-control's flipped nibble is always
97 // observable (under a zero scale the flip is arithmetically invisible -- measured: 1 row of 200 drew a
98 // zero and the control read as a miss). A zero scale is a legitimate Q4_K value; T3 covers it lane-level.
99 // bytes 0..3: sc(0..3) low 6 bits in 1..63 with random high bits; bytes 4..7: mins, any; bytes 8..11:
100 // low nibble (sc 4..7 low bits) in 1..15 with a random high nibble (mins 4..7 high bits).
101 var s: i64 = 0
102 while s < 4 { buf[base+4+s] = ((1 + (g2_lcg() % 63)) | ((g2_lcg() % 4) << 6)) as u8; s = s + 1 }
103 while s < 8 { buf[base+4+s] = (g2_lcg() % 256) as u8; s = s + 1 }
104 while s < 12 { buf[base+4+s] = ((1 + (g2_lcg() % 15)) | ((g2_lcg() % 16) << 4)) as u8; s = s + 1 }
105 var q: i64 = 0
106 while q < 128 { buf[base+16+q] = (g2_lcg() % 256) as u8; q = q + 1 }
107 b = b + 1
108 }
109 return 0
110}
111
112func main() -> i64 {
113 gv_head("NX-Q4K-SIMD2-GATE: the compiler-emitted Q4_K unpack-and-scale, lane-exact and dot-exact against the scalar rulers" as *u8)
114 let ctr: *i64 = gv_ctr()
115 g_g2_st = G2_LCG_SEED
116 let qs: *u8 = sys_mmap(G2_BYTES)
117 let out: *u8 = sys_mmap(G2_QBUF_BYTES)
118 let ref: *i64 = sys_mmap(G2_QBUF_BYTES) as *i64
119
120 // T1 extreme lanes: every byte 0xFF, scales 63/63 -> every lane 945 (no carry into the neighbour)
121 var k: i64 = 0; while k < G2_BYTES { qs[k] = 255 as u8; k = k + 1 }
122 __q4k_unpack32s(qs, out, G2_SC_MAX | (G2_SC_MAX << G2_SC_HI_SHIFT))
123 gv_check("T1 all-ones bytes at the maximum scale give 945 in all 64 lanes (no lane carry)" as *u8, g2_lanes_ok(qs, out, G2_SC_MAX, G2_SC_MAX), ctr)
124 gv_check("T1 fixture-reached: lane 0 reads 945 and lane 63 reads 945" as *u8, (g2_lane(out, 0) == G2_LANE_MAX) & (g2_lane(out, 63) == G2_LANE_MAX), ctr)
125 // T2 distinct nibbles and asymmetric scales: byte k = (k & 15) | ((15 - (k & 15)) << 4), sc_lo=7 sc_hi=63
126 k = 0; while k < G2_BYTES { qs[k] = ((k & G2_NIBBLE_MASK) | ((15 - (k & G2_NIBBLE_MASK)) << G2_NIB_SHIFT)) as u8; k = k + 1 }
127 __q4k_unpack32s(qs, out, 7 | (G2_SC_MAX << G2_SC_HI_SHIFT))
128 gv_check("T2 distinct nibbles with asymmetric scales land in the right half at the right multiple" as *u8, g2_lanes_ok(qs, out, 7, G2_SC_MAX), ctr)
129 // T3 a zero scale zeroes its half and leaves the other half intact
130 __q4k_unpack32s(qs, out, 0 | (5 << G2_SC_HI_SHIFT))
131 gv_check("T3 sc_lo=0 zeroes the low half while the high half is scaled by 5" as *u8, g2_lanes_ok(qs, out, 0, 5), ctr)
132 // T4 neg-control: corrupt one output lane after the unpack -> the lane check must FAIL
133 __q4k_unpack32s(qs, out, 7 | (G2_SC_MAX << G2_SC_HI_SHIFT))
134 out[10] = (out[10] ^ (1 as u8)) as u8
135 gv_check("neg-control-T4 a corrupted lane is detected by the lane tooth (the tooth can fail)" as *u8, g2_lanes_ok(qs, out, 7, G2_SC_MAX) == 0, ctr)
136 // T5 5000 random blocks: intrinsic lanes == scalar spread+multiply lanes, byte for byte
137 var agree: i64 = 0; var tried: i64 = 0
138 var r: i64 = 0
139 while r < G2_RAND_BLOCKS {
140 k = 0; while k < G2_BYTES { qs[k] = (g2_lcg() % 256) as u8; k = k + 1 }
141 let slo: i64 = g2_lcg() % (G2_SC_MAX + 1); let shi: i64 = g2_lcg() % (G2_SC_MAX + 1)
142 __q4k_unpack32s(qs, out, slo | (shi << G2_SC_HI_SHIFT))
143 g2_scalar_unpack(qs, ref, slo, shi)
144 tried = tried + 1
145 if g2_bufs_equal(out, ref as *u8, G2_QBUF_BYTES) == 1 { agree = agree + 1 }
146 r = r + 1
147 }
148 gv_puts(" T5 random blocks tried=" as *u8); gv_num(tried); gv_puts(" agree=" as *u8); gv_num(agree); gv_puts("\n" as *u8)
149 gv_check("T5 intrinsic == scalar spread+multiply on every one of 5000 random blocks (denominator bound)" as *u8, (agree == tried) & (tried == G2_RAND_BLOCKS), ctr)
150
151 // T6 dot level: 200 random Q4_K rows, simd2 == simd == scalar row_col, plus a neg-control
152 let W: *u8 = sys_mmap(G2_NBLK * G2_SB)
153 let IN: i64 = G2_NBLK * 256
154 let col: *i64 = sys_mmap(IN * 8) as *i64
155 let col16: *i64 = sys_mmap(IN * 2) as *i64
156 let scpre: *i64 = sys_mmap(G2_NBLK * 8 * 8) as *i64
157 let qpk: *i64 = sys_mmap(64) as *i64
158 let qhi: *i64 = sys_mmap(64) as *i64
159 let qbuf: *i64 = sys_mmap(G2_QBUF_BYTES) as *i64
160 let acc: *i64 = sys_mmap(32) as *i64
161 let it: *NxQ4KBlockIter = nx_q4k_iter_alloc()
162 var eq_simd: i64 = 0; var eq_scalar: i64 = 0; var rows: i64 = 0; var neg_changed: i64 = 0
163 var tsimd: i64 = 0; var tsimd2: i64 = 0
164 r = 0
165 while r < G2_ROWS {
166 g2_fill_row(W, G2_NBLK)
167 var i: i64 = 0
168 while i < IN { col[i] = (g2_lcg() % (2*G2_COL_MAX + 1)) - G2_COL_MAX; i = i + 1 }
169 var jj: i64 = 0
170 while jj < IN / 4 { col16[jj] = ds_pack4(col[jj*4], col[jj*4+1], col[jj*4+2], col[jj*4+3]); jj = jj + 1 }
171 nx_q4k_sc_precompute(col, G2_NBLK, scpre)
172 let t0: i64 = nx_clock_monotonic_ns()
173 let d1: i64 = nx_q4k_dot_simd(W, 0, G2_NBLK, col16, qpk, qhi, acc, scpre)
174 let t1: i64 = nx_clock_monotonic_ns()
175 let d2: i64 = nx_q4k_dot_simd2(W, 0, G2_NBLK, col16, qbuf, acc, scpre)
176 let t2: i64 = nx_clock_monotonic_ns()
177 let d0: i64 = nx_q4k_dot_row_col(W, 0, G2_NBLK, col, it)
178 tsimd = tsimd + (t1 - t0); tsimd2 = tsimd2 + (t2 - t1)
179 rows = rows + 1
180 if d2 == d1 { eq_simd = eq_simd + 1 }
181 if d2 == d0 { eq_scalar = eq_scalar + 1 }
182 // neg-control: flip one nibble byte of the row and the dot must move
183 W[16 + (r % 128)] = (W[16 + (r % 128)] ^ (1 as u8)) as u8
184 let d3: i64 = nx_q4k_dot_simd2(W, 0, G2_NBLK, col16, qbuf, acc, scpre)
185 if d3 != d2 { neg_changed = neg_changed + 1 }
186 r = r + 1
187 }
188 gv_puts(" T6 rows=" as *u8); gv_num(rows); gv_puts(" simd2==simd " as *u8); gv_num(eq_simd); gv_puts(" simd2==scalar " as *u8); gv_num(eq_scalar); gv_puts(" neg_changed=" as *u8); gv_num(neg_changed); gv_puts("\n" as *u8)
189 gv_check("T6 simd2 dot == simd dot on every one of 200 random real-shaped rows" as *u8, (eq_simd == rows) & (rows == G2_ROWS), ctr)
190 gv_check("T6b simd2 dot == scalar row_col dot on every row (the estate's original ruler)" as *u8, (eq_scalar == rows) & (rows == G2_ROWS), ctr)
191 gv_check("neg-control-T6c flipping one nibble byte changes the simd2 dot on every row (the dot tooth can fail)" as *u8, neg_changed == rows, ctr)
192 // timing: a number, not a verdict
193 var k2: i64 = 0; let ta: i64 = nx_clock_monotonic_ns()
194 while k2 < G2_TIMING_ITERS { nx_q4k_dot_simd(W, 0, G2_NBLK, col16, qpk, qhi, acc, scpre); k2 = k2 + 1 }
195 let tb: i64 = nx_clock_monotonic_ns()
196 k2 = 0
197 while k2 < G2_TIMING_ITERS { nx_q4k_dot_simd2(W, 0, G2_NBLK, col16, qbuf, acc, scpre); k2 = k2 + 1 }
198 let tc: i64 = nx_clock_monotonic_ns()
199 gv_puts(" timing over " as *u8); gv_num(G2_TIMING_ITERS); gv_puts(" rows of 14 super-blocks: simd_ns=" as *u8); gv_num(tb - ta); gv_puts(" simd2_ns=" as *u8); gv_num(tc - tb)
200 if (tc - tb) > 0 { gv_puts(" speedup_x100=" as *u8); gv_num((tb - ta) * 100 / (tc - tb)) }
201 gv_puts("\n" as *u8)
202 return gv_verdict("nx_q4k_simd2_gate" as *u8, ctr, "the compiler-emitted unpack is lane-exact against the scalar spread and dot-exact against both incumbent rulers on random real-shaped rows; timing is reported, the decode bench is the speed ruler" as *u8)
203}