nx_q4k_matmul_test.nx source
↩ module page · 130 lines · 4995 B
1// nx_q4k_matmul_test.nx -- smoke + bit-exact KAT for fused Q4_K dot.
2//
3// Two test classes:
4// A) Verdict gate + Q20 → Q10 helper unit checks
5// B) KAT: build a Q4_K block + col vector, compute fused dot,
6// compare bit-exactly to "dequant whole block then dot" path
7
8import "nx_syscalls.nx"
9import "nx_tier.nx"
10import "nx_le.nx"
11import "nx_gguf.nx"
12import "nx_gguf_load.nx"
13import "nx_dequant_iter.nx"
14import "nx_q4k_matmul.nx"
15
16// Build the same canonical Q4_K block we use in nx_gguf_load_q4k_test.
17// d=1.0, dmin=0, sc[0..3]={1,1,1,1} sc[4]=2 sc[5..7]=0, all mins=0,
18// nibble byte 0 = 0x32 (sb 0, l 0), nibble byte 64 = 0x32 (sb 4, l 0),
19// all other nibbles 0.
20func _build_canon_block(buf: *u8, off: i64) -> nx_int {
21 nx_le_write_u16(buf, off + 0, 0x3C00)
22 nx_le_write_u16(buf, off + 2, 0x0000)
23 buf[off + 4]=0x01; buf[off + 5]=0x01
24 buf[off + 6]=0x01; buf[off + 7]=0x01
25 buf[off + 8]=0x00; buf[off + 9]=0x00
26 buf[off + 10]=0x00; buf[off + 11]=0x00
27 buf[off + 12]=0x02; buf[off + 13]=0x00
28 buf[off + 14]=0x00; buf[off + 15]=0x00
29 var zi: nx_int = 0
30 while zi < 128 {
31 buf[off + 16 + zi] = 0
32 zi = zi + 1
33 }
34 buf[off + 16 + 0] = 0x32
35 buf[off + 16 + 64] = 0x32
36 return 0
37}
38
39func main() -> i64 {
40 // ----- A) Verdict gate -----
41 var vi: nx_int = 0
42 while vi < NX_Q4KM_N_VERDICTS {
43 if nx_q4km_verdict_is_valid(vi) != 1 { return 5 + vi }
44 vi = vi + 1
45 }
46
47 // Normalize (Q34 -> Q10) unit checks. The fused dot is now Q34
48 // (Q24 weight x Q10 col); the helper divides by 2^24 (round-half
49 // away from zero, C truncation).
50 // 2^24 * 1024 (Q34) -> 1024 (Q10 = 1.0)
51 if nx_q4km_q20_to_q10(16777216 * 1024) != 1024 { return 10 }
52 if nx_q4km_q20_to_q10(0) != 0 { return 11 }
53 if nx_q4km_q20_to_q10(0 - 16777216 * 1024) != (0 - 1024) { return 12 }
54 // Rounding: 2^23 - 1 rounds to 0 ; 2^23 rounds to 1
55 if nx_q4km_q20_to_q10(8388607) != 0 { return 13 }
56 if nx_q4km_q20_to_q10(8388608) != 1 { return 14 }
57
58 // ----- B) Build canonical Q4_K block at off 0 -----
59 let buf: *u8 = sys_mmap(512)
60 _build_canon_block(buf, 0)
61
62 // ggml layout + Q24 super-scale (d=1.0 -> d_q24 = 16777216). The
63 // dequantized block has only 3 non-zero values now:
64 // v[0] = 33554432 (16777216 * sc[0]=1 * 2 - 0) low nibble, byte 0
65 // v[32] = 50331648 (16777216 * sc[1]=1 * 3 - 0) high nibble, byte 0
66 // v[128] = 67108864 (16777216 * sc[4]=2 * 2 - 0) low nibble, byte 64
67
68 // Build a col vector [256] Q10, activations aligned to the nonzero
69 // OUTPUT positions: col[0]=1.0=1024, col[32]=2.0=2048, col[128]=0.5=512.
70 // Expected fused dot (Q34 = Q24 weight x Q10 col):
71 // 33554432*1024 + 50331648*2048 + 67108864*512
72 // = 34359738368 + 103079215104 + 34359738368
73 // = 171798691840
74 let col: *i64 = sys_mmap(256 * 8) as *i64
75 var ci: nx_int = 0
76 while ci < 256 { col[ci] = 0; ci = ci + 1 }
77 col[0] = 1024
78 col[32] = 2048
79 col[128] = 512
80
81 let it: *NxQ4KBlockIter = nx_q4k_iter_alloc()
82 let dot_q20: i64 = nx_q4k_dot_row_col(buf, 0, 1, col, it)
83 if dot_q20 != 171798691840 { return 30 }
84
85 // Normalize Q34 -> Q10: (171798691840 + 2^23) >> 24 = 10240 (=10.0).
86 let dot_q10: i64 = nx_q4km_q20_to_q10(dot_q20)
87 if dot_q10 != 10240 { return 31 }
88
89 // ----- C) Cross-check vs "dequant then dot" path -----
90 //
91 // Dequant the whole block, then do a manual dot, then compare.
92 let weights_q10: *i64 = sys_mmap(256 * 8) as *i64
93 nx_gguf_dequant_q4_k(buf, 0, 256, weights_q10)
94 var manual_dot: i64 = 0
95 var mi: nx_int = 0
96 while mi < 256 {
97 manual_dot = manual_dot + weights_q10[mi] * col[mi]
98 mi = mi + 1
99 }
100 if manual_dot != dot_q20 { return 40 }
101
102 // ----- D) Multi-block: 2 blocks in a row -----
103 //
104 // Repeat the same canon block at offset 144. Second block's
105 // nonzero outputs are at 256+{0,32,128} (v[256]=33554432,
106 // v[288]=50331648, v[384]=67108864). Align activations for block 2
107 // at col[256]=1024, col[288]=2048 (so block 2 contributes its first
108 // two nonzero values; nothing at 384).
109 let buf2: *u8 = sys_mmap(512)
110 _build_canon_block(buf2, 0)
111 _build_canon_block(buf2, NX_GL_Q4_K_BPB)
112 let col2: *i64 = sys_mmap(512 * 8) as *i64
113 var ci2: nx_int = 0
114 while ci2 < 512 { col2[ci2] = 0; ci2 = ci2 + 1 }
115 col2[0] = 1024
116 col2[32] = 2048
117 col2[128] = 512
118 col2[256] = 1024 // block 2's v[256]
119 col2[288] = 2048 // block 2's v[288]
120 // Expected: first block contribution + second block contribution (Q34)
121 // block 1: 171798691840 (from B above)
122 // block 2: 33554432*1024 + 50331648*2048 = 34359738368 + 103079215104
123 // = 137438953472
124 // total: 309237645312
125 let it2: *NxQ4KBlockIter = nx_q4k_iter_alloc()
126 let dot2: i64 = nx_q4k_dot_row_col(buf2, 0, 2, col2, it2)
127 if dot2 != 309237645312 { return 50 }
128
129 return 0
130}