nx_q5_0_fused_dot_gate.nx source
↩ module page · 198 lines · 7118 B
1// nx_q5_0_fused_dot_gate.nx -- ISOLATED proof of the model-appropriate
2// SOTA decode lever: a FUSED Q5_0 dequant-dot (read the 22-byte block,
3// dequant in-register, dot) vs the CURRENT path (dequant the whole row
4// to F32 i64-slots, then dot).
5//
6// WHY: nx_gguf_typecensus shows our model is 79% Q5_0, but the loader
7// dequants it to F32 i64-slots (8 B/value) -- the memory-bound decode
8// matmul then reads 8 B/value when Q5_0 is 0.69 B/value native = ~11.6x
9// more memory than needed. SOTA (Marlin) keeps weights quantized and
10// dequants in-register. This gate proves the fused kernel bit-exact +
11// measures the memory-bound win at the real FFN shape. Ratio within one
12// run => host-noise-immune. NO forward changes (no " Paris" risk).
13//
14// Checks:
15// 1 fused q5_0 dot == dequant-to-f32-then-dot, bit-exact (whole matmul)
16// 2 fused reads ~11.6x fewer bytes; measure the speedup (floor 1.5x)
17//
18// lineage_id: q5_0_fused_dot_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_fmt.nx"
27
28const QK: i64 = 896 // hidden (reduction dim), div by 32
29const QN: i64 = 4864 // ffn (W_gate n)
30const QREPS: i64 = 60
31const QFLOOR_X100: i64 = 150
32
33// Q5_0: 22 bytes/32-values (d:2 f16, qh:4, qs:16).
34const Q5B: i64 = 22
35const Q5V: i64 = 32
36
37func q_lcg(s: i64) -> i64 {
38 var v: i64 = s * 1103515245 + 12345
39 v = v & 2147483647
40 return v
41}
42
43// Build QN rows of Q5_0, each QK values = (QK/32) blocks. Fixed small
44// f16 scale (0x2C00 ~= 0.0625) so dequant values stay finite; LCG qh/qs.
45func q_fill_weight(w: *u8, seed: i64) -> i64 {
46 let bpr: i64 = (QK / Q5V) * Q5B
47 var s: i64 = seed
48 var r: i64 = 0
49 while r < QN {
50 var b: i64 = 0
51 while b < QK / Q5V {
52 let off: i64 = r * bpr + b * Q5B
53 w[off + 0] = 0x00 as u8 // f16 d low
54 w[off + 1] = 0x2C as u8 // f16 d high (0x2C00 ~ 0.0625)
55 var i: i64 = 0
56 while i < 4 { s = q_lcg(s); w[off + 2 + i] = (s & 255) as u8; i = i + 1 }
57 var q: i64 = 0
58 while q < 16 { s = q_lcg(s); w[off + 6 + q] = (s & 255) as u8; q = q + 1 }
59 b = b + 1
60 }
61 r = r + 1
62 }
63 return 0
64}
65
66func q_fill_a(p: *i64, count: i64, seed: i64) -> i64 {
67 var s: i64 = seed
68 var i: i64 = 0
69 while i < count {
70 s = q_lcg(s)
71 p[i] = nx_i32_to_f32((s % 9) - 4)
72 i = i + 1
73 }
74 return 0
75}
76
77// FUSED Q5_0 dequant-dot: dot(A[0..k], dequant(qbuf[qoff..])) reading the
78// 22-byte blocks. Dequants each block into a 32-f32 scratch (L1) IN ORDER
79// so the accumulation order matches the oracle (bit-exact), while the
80// STREAMING read is the 22-byte block (0.69 B/val), not 8 B/val f32.
81func q5_0_dot(qbuf: *u8, qoff: i64, A: *i64, k: i64, vscr: *i64) -> i64 {
82 let nblk: i64 = k / Q5V
83 var acc: i64 = 0
84 var b: i64 = 0
85 while b < nblk {
86 let boff: i64 = qoff + b * Q5B
87 let d_f32: i64 = nx_f16_to_f32(nx_le_read_u16(qbuf, boff + 0))
88 let qh: i64 = nx_le_read_u32(qbuf, boff + 2)
89 var j: i64 = 0
90 while j < 16 {
91 let qs_byte: i64 = nx_le_read_u8(qbuf, boff + 6 + j)
92 let q5_lo: i64 = (qs_byte & 0x0F) | (((qh >> j) & 1) << 4)
93 let q5_hi: i64 = ((qs_byte >> 4) & 0x0F) | (((qh >> (j + 16)) & 1) << 4)
94 vscr[j] = nx_f32_mul(d_f32, nx_i32_to_f32(q5_lo - 16))
95 vscr[j + 16] = nx_f32_mul(d_f32, nx_i32_to_f32(q5_hi - 16))
96 j = j + 1
97 }
98 let abase: i64 = b * Q5V
99 var l: i64 = 0
100 while l < Q5V {
101 acc = nx_f32_add(acc, nx_f32_mul(A[abase + l], vscr[l]))
102 l = l + 1
103 }
104 b = b + 1
105 }
106 return acc
107}
108
109func q_nl() -> i64 { fmt_puts("\n" as *u8); return 0 }
110
111func main() -> i64 {
112 let bpr: i64 = (QK / Q5V) * Q5B
113 let W: *u8 = sys_mmap(QN * bpr)
114 let A: *i64 = sys_mmap(QK * 8) as *i64
115 let Cref: *i64 = sys_mmap(QN * 8) as *i64
116 let Cfus: *i64 = sys_mmap(QN * 8) as *i64
117 let rowf: *i64 = sys_mmap(QK * 8) as *i64 // oracle dequant scratch (per row)
118 let vscr: *i64 = sys_mmap(Q5V * 8) as *i64 // fused block scratch (32 f32)
119 q_fill_weight(W, 20260708)
120 q_fill_a(A, QK, 4242)
121
122 var pass: i64 = 0
123
124 // ---- 1: bit-exact (oracle = dequant-to-f32 then dot; fused = q5_0_dot) ----
125 var r: i64 = 0
126 while r < QN {
127 let roff: i64 = r * bpr
128 nx_q5_0_to_f32(W, roff, QK, rowf)
129 var acc: i64 = 0
130 var l: i64 = 0
131 while l < QK { acc = nx_f32_add(acc, nx_f32_mul(A[l], rowf[l])); l = l + 1 }
132 Cref[r] = acc
133 Cfus[r] = q5_0_dot(W, roff, A, QK, vscr)
134 r = r + 1
135 }
136 var ok1: i64 = 1
137 var c: i64 = 0
138 while c < QN { if Cref[c] != Cfus[c] { ok1 = 0; c = QN } else { c = c + 1 } }
139 if ok1 != 1 { fmt_puts("Q50 1 EXACT FAIL"); q_nl(); return 11 }
140 fmt_puts("Q50 1 FUSED==DEQUANT-THEN-DOT EXACT OK"); q_nl()
141 pass = pass + 1
142
143 // ---- 2: speed. The REAL forward stores the weight as F32 i64-slots
144 // (8 B/val = QN*QK*8 = 34.9MB) and the matmul reads ALL of it
145 // every token. Pre-materialize that ONCE (as the forward does at
146 // load), then time: F32 dot (reads 34.9MB) vs fused q5_0_dot
147 // (reads 3MB Q5_0). THIS is the memory-bound comparison. ----
148 let WF: *i64 = sys_mmap(QN * QK * 8) as *i64 // 34.9MB materialized F32 weight
149 var mr: i64 = 0
150 while mr < QN {
151 nx_q5_0_to_f32(W, mr * bpr, QK, ((WF as i64) + mr * QK * 8) as *i64)
152 mr = mr + 1
153 }
154 let t0: i64 = sys_now_us()
155 var rp0: i64 = 0
156 while rp0 < QREPS {
157 var rr: i64 = 0
158 while rr < QN {
159 let wrow: i64 = (WF as i64) + rr * QK * 8
160 let wp: *i64 = wrow as *i64
161 var acc: i64 = 0
162 var l: i64 = 0
163 while l < QK { acc = nx_f32_add(acc, nx_f32_mul(A[l], wp[l])); l = l + 1 }
164 Cref[rr] = acc
165 rr = rr + 1
166 }
167 rp0 = rp0 + 1
168 }
169 let us_cur: i64 = sys_now_us() - t0
170
171 let t1: i64 = sys_now_us()
172 var rp1: i64 = 0
173 while rp1 < QREPS {
174 var rr2: i64 = 0
175 while rr2 < QN {
176 Cfus[rr2] = q5_0_dot(W, rr2 * bpr, A, QK, vscr)
177 rr2 = rr2 + 1
178 }
179 rp1 = rp1 + 1
180 }
181 let us_fus: i64 = sys_now_us() - t1
182
183 var uc: i64 = us_cur
184 if uc < 1 { uc = 1 }
185 var uf: i64 = us_fus
186 if uf < 1 { uf = 1 }
187 let macs: i64 = QK * QN * QREPS
188 fmt_puts("current(dequant->f32->dot)_us="); fmt_putn(us_cur); fmt_puts(" mflops="); fmt_putn(2 * macs / uc); q_nl()
189 fmt_puts("fused(q5_0_dot)_us="); fmt_putn(us_fus); fmt_puts(" mflops="); fmt_putn(2 * macs / uf); q_nl()
190 let sx100: i64 = uc * 100 / uf
191 fmt_puts("fused_speedup_x100="); fmt_putn(sx100); q_nl()
192 if sx100 < QFLOOR_X100 { fmt_puts("Q50 2 SPEEDUP FAIL"); q_nl(); return 12 }
193 fmt_puts("Q50 2 SPEEDUP OK"); q_nl()
194 pass = pass + 1
195
196 fmt_puts("Q5_0_FUSED_DOT_GATE "); fmt_putn(pass); fmt_puts("/2 GREEN"); q_nl()
197 return 0
198}