nx_dequant_iter.nx source
↩ module page · 297 lines · 11194 B
1// nx_dequant_iter.nx -- per-block Q4_K dequant iterator.
2//
3// Architectural step toward conductor Phase C (dispatcher with fused
4// dequant+matmul). Instead of materializing a whole tensor's worth
5// of Q10 i64 storage, the dispatcher will hold a small per-block
6// iterator that amortizes the 8-scale-8-min unpack across the 256
7// value reads from that block. A matmul kernel asks the iterator
8// for individual values (or sub-block batches) on demand.
9//
10// Why this matters (HONESTLY, no perf claim attached):
11//
12// The existing nx_gguf_dequant_q4_k writes ALL 256 values of every
13// super-block into a flat output buffer. For Llama-7B's ~7B params
14// at 256 values/block that's ~27 million blocks materialized eagerly.
15// At ~11 ns/value (qemu-riscv64 measurement from nx_dequant_bench,
16// 2026-05-19) that's ~77 seconds of pure dequant if the substrate
17// ever materializes the full model.
18//
19// The iterator pattern (this brick) DOES NOT make per-value dequant
20// faster. It enables the SCAFFOLDING for a future dispatcher that
21// dequants only the blocks a given matmul tile reads. Whether the
22// dispatcher delivers net latency wins vs the materialize-everything
23// path depends on the matmul access pattern and remains to be
24// measured -- not claimed here.
25//
26// What this brick DOES claim, narrowly:
27// - One-time per-block setup: read f16 d + dmin, unpack 8 scales +
28// 8 mins from the 12-byte packed header (one constant-time pass)
29// - Constant-time per-value access: nx_q4k_iter_value(it, idx)
30// does index math + nibble extract + (d * sc) * q - dmin * m
31// - Bit-exact match against nx_gguf_dequant_q4_k for the same
32// block bytes (verified by smoke)
33//
34// genealogy_id: ggml_q4k_block_canon + iterator_pattern_canonical
35// lineage_id: substrate_dequant_iter_v1_q4k
36
37import "nx_syscalls.nx"
38import "nx_tier.nx"
39import "nx_le.nx"
40import "nx_gguf.nx"
41import "nx_gguf_load.nx"
42
43// ===== The iterator =================================================
44//
45// 18 i64 fields + 1 pointer = 152 bytes when laid out on the heap.
46// Allocate once per inference call; reinit per block.
47
48struct NxQ4KBlockIter {
49 d_q10: i64,
50 dmin_q10: i64,
51 sc_0: i64,
52 sc_1: i64,
53 sc_2: i64,
54 sc_3: i64,
55 sc_4: i64,
56 sc_5: i64,
57 sc_6: i64,
58 sc_7: i64,
59 m_0: i64,
60 m_1: i64,
61 m_2: i64,
62 m_3: i64,
63 m_4: i64,
64 m_5: i64,
65 m_6: i64,
66 m_7: i64,
67 qs_ptr: *u8
68}
69
70const NX_Q4K_ITER_BYTES: nx_int = 152 // 19 * 8
71
72const NX_Q4KI_OK: nx_int = 0
73const NX_Q4KI_ERR_BAD_BUF: nx_int = 1
74const NX_Q4KI_ERR_BAD_IDX: nx_int = 2
75const NX_Q4KI_N_VERDICTS: nx_int = 3
76
77func nx_q4ki_verdict_is_valid(v: nx_int) -> nx_int {
78 if v < 0 { return 0 }
79 if v >= NX_Q4KI_N_VERDICTS { return 0 }
80 return 1
81}
82
83// ===== Allocator ===================================================
84
85func nx_q4k_iter_alloc() -> *NxQ4KBlockIter {
86 return sys_mmap(NX_Q4K_ITER_BYTES) as *NxQ4KBlockIter
87}
88
89// ===== Init: read one Q4_K super-block's header into the iter =====
90//
91// buf: byte buffer holding the source GGUF (or arbitrary memory)
92// super_off: byte offset where the 144-byte super-block starts
93// it: caller-owned iterator to populate
94//
95// After init:
96// it.d_q10, it.dmin_q10 are the f16 super-scales decoded to Q24
97// (fields keep their historical *_q10 names but
98// now hold EXACT Q24 -- Q10/Q14 underflow/flip
99// sign on real trained-model super-scales)
100// it.sc_0..sc_7 are the 8 unpacked 6-bit sub-block scales
101// it.m_0..m_7 are the 8 unpacked 6-bit sub-block mins
102// it.qs_ptr points at bytes 16..143 of the super-block
103
104func nx_q4k_iter_init(buf: *u8, super_off: i64, it: *NxQ4KBlockIter) -> nx_int {
105 if buf == (0 as *u8) { return NX_Q4KI_ERR_BAD_BUF }
106
107 let d_raw: i64 = nx_le_read_u16(buf, super_off)
108 let dmin_raw: i64 = nx_le_read_u16(buf, super_off + 2)
109 it.d_q10 = _gguf_f16_to_q24(d_raw) // EXACT Q24 (field name historical)
110 it.dmin_q10 = _gguf_f16_to_q24(dmin_raw)
111
112 let scales_off: i64 = super_off + 4
113
114 // Low half: j in 0..4
115 // sc[j] = q[j] & 0x3F
116 // m[j] = q[j+4] & 0x3F
117 it.sc_0 = nx_le_read_u8(buf, scales_off + 0) & 0x3F
118 it.sc_1 = nx_le_read_u8(buf, scales_off + 1) & 0x3F
119 it.sc_2 = nx_le_read_u8(buf, scales_off + 2) & 0x3F
120 it.sc_3 = nx_le_read_u8(buf, scales_off + 3) & 0x3F
121 it.m_0 = nx_le_read_u8(buf, scales_off + 4) & 0x3F
122 it.m_1 = nx_le_read_u8(buf, scales_off + 5) & 0x3F
123 it.m_2 = nx_le_read_u8(buf, scales_off + 6) & 0x3F
124 it.m_3 = nx_le_read_u8(buf, scales_off + 7) & 0x3F
125
126 // High half: k = j - 4, j in 4..8
127 // sc[j] = ((q[k] >> 6) << 4) | (q[8+k] & 0x0F)
128 // m[j] = ((q[4+k] >> 6) << 4) | (q[8+k] >> 4)
129 let b_0: i64 = nx_le_read_u8(buf, scales_off + 0)
130 let b_1: i64 = nx_le_read_u8(buf, scales_off + 1)
131 let b_2: i64 = nx_le_read_u8(buf, scales_off + 2)
132 let b_3: i64 = nx_le_read_u8(buf, scales_off + 3)
133 let b_4: i64 = nx_le_read_u8(buf, scales_off + 4)
134 let b_5: i64 = nx_le_read_u8(buf, scales_off + 5)
135 let b_6: i64 = nx_le_read_u8(buf, scales_off + 6)
136 let b_7: i64 = nx_le_read_u8(buf, scales_off + 7)
137 let b_8: i64 = nx_le_read_u8(buf, scales_off + 8)
138 let b_9: i64 = nx_le_read_u8(buf, scales_off + 9)
139 let b_10: i64 = nx_le_read_u8(buf, scales_off + 10)
140 let b_11: i64 = nx_le_read_u8(buf, scales_off + 11)
141
142 it.sc_4 = ((b_0 >> 6) << 4) | (b_8 & 0x0F)
143 it.sc_5 = ((b_1 >> 6) << 4) | (b_9 & 0x0F)
144 it.sc_6 = ((b_2 >> 6) << 4) | (b_10 & 0x0F)
145 it.sc_7 = ((b_3 >> 6) << 4) | (b_11 & 0x0F)
146 it.m_4 = ((b_4 >> 6) << 4) | (b_8 >> 4)
147 it.m_5 = ((b_5 >> 6) << 4) | (b_9 >> 4)
148 it.m_6 = ((b_6 >> 6) << 4) | (b_10 >> 4)
149 it.m_7 = ((b_7 >> 6) << 4) | (b_11 >> 4)
150
151 it.qs_ptr = (buf as i64 + super_off + 16) as *u8
152
153 return NX_Q4KI_OK
154}
155
156// ===== Internal helpers ============================================
157
158func _q4k_iter_scale_at(it: *NxQ4KBlockIter, sb: nx_int) -> i64 {
159 if sb == 0 { return it.sc_0 }
160 if sb == 1 { return it.sc_1 }
161 if sb == 2 { return it.sc_2 }
162 if sb == 3 { return it.sc_3 }
163 if sb == 4 { return it.sc_4 }
164 if sb == 5 { return it.sc_5 }
165 if sb == 6 { return it.sc_6 }
166 if sb == 7 { return it.sc_7 }
167 return 0
168}
169
170func _q4k_iter_min_at(it: *NxQ4KBlockIter, sb: nx_int) -> i64 {
171 if sb == 0 { return it.m_0 }
172 if sb == 1 { return it.m_1 }
173 if sb == 2 { return it.m_2 }
174 if sb == 3 { return it.m_3 }
175 if sb == 4 { return it.m_4 }
176 if sb == 5 { return it.m_5 }
177 if sb == 6 { return it.m_6 }
178 if sb == 7 { return it.m_7 }
179 return 0
180}
181
182// ===== Public: get the (scale, min) pair for a sub-block =========
183//
184// Fused-dequant-matmul kernels need both per sub-block; this
185// promotes the previously-private helpers to a public API so the
186// kernel doesn't reach into underscore-prefixed internals.
187
188func nx_q4k_iter_scale_min(it: *NxQ4KBlockIter, sb: nx_int,
189 out_sc: *i64, out_m: *i64) -> nx_int {
190 if sb < 0 { return NX_Q4KI_ERR_BAD_IDX }
191 if sb >= 8 { return NX_Q4KI_ERR_BAD_IDX }
192 out_sc[0] = _q4k_iter_scale_at(it, sb)
193 out_m[0] = _q4k_iter_min_at(it, sb)
194 return NX_Q4KI_OK
195}
196
197// ===== Allocation-free scale/min accessors (hot-loop API) ==========
198//
199// WHY THESE EXIST: nx_q4k_iter_scale_min above returns through two
200// caller-supplied out-cells. A GEMM inner loop has nowhere to put those
201// cells, so nx_q4k_dot_row_col was allocating them with sys_mmap ON EVERY
202// CALL and never unmapping -- m*n page-rounded mappings leaked per matmul
203// (~9,700 at n=4864). These return the value directly, so a hot dot
204// kernel needs no scratch and performs no syscall at all.
205// Out-of-range sb yields 0, matching nx_q4k_iter_value's non-crashing
206// contract; callers that need to DISTINGUISH bad input keep using
207// nx_q4k_iter_scale_min, which still reports NX_Q4KI_ERR_BAD_IDX.
208
209func nx_q4k_iter_scale(it: *NxQ4KBlockIter, sb: nx_int) -> i64 {
210 if sb < 0 { return 0 }
211 if sb >= 8 { return 0 }
212 return _q4k_iter_scale_at(it, sb)
213}
214
215func nx_q4k_iter_min(it: *NxQ4KBlockIter, sb: nx_int) -> i64 {
216 if sb < 0 { return 0 }
217 if sb >= 8 { return 0 }
218 return _q4k_iter_min_at(it, sb)
219}
220
221// ===== Value-on-demand =============================================
222//
223// idx is in [0, 256) = the ggml OUTPUT position. Decodes idx under the
224// ggml layout (verified bit-exact): sub-block sb = idx/32, position
225// l = idx%32, group g = sb/2, source byte = qs[g*32 + l]; the nibble is
226// the LOW nibble when sb is even (sub-block 2g) and the HIGH nibble when
227// sb is odd (sub-block 2g+1). Value (Q24):
228// value_q24 = (d_q24 * sc[sb]) * q4 - (dmin_q24 * m[sb])
229//
230// Returns 0 for out-of-range idx (so a buggy caller doesn't crash).
231
232func nx_q4k_iter_value(it: *NxQ4KBlockIter, idx: nx_int) -> i64 {
233 if idx < 0 { return 0 }
234 if idx >= 256 { return 0 }
235
236 let sb: nx_int = idx / 32 // 0..8 output sub-block
237 let l: nx_int = idx - sb * 32 // 0..32 position in sub-block (idx % 32)
238 let g: nx_int = sb / 2 // 0..4 32-byte group
239 let is_high: nx_int = sb - g * 2 // 0 (low, sb=2g) or 1 (high, sb=2g+1)
240
241 let byte_v: i64 = nx_le_read_u8(it.qs_ptr, g * 32 + l)
242 var q4: i64 = 0
243 if is_high == 0 {
244 q4 = byte_v & 0x0F
245 } else {
246 q4 = byte_v >> 4
247 }
248
249 let sc_sb: i64 = _q4k_iter_scale_at(it, sb)
250 let m_sb: i64 = _q4k_iter_min_at(it, sb)
251 let d1: i64 = it.d_q10 * sc_sb
252 let m1: i64 = it.dmin_q10 * m_sb
253 return d1 * q4 - m1
254}
255
256// ===== Sub-block bulk dequant ======================================
257//
258// For callers that need ALL 32 values of one sub-block in a tight
259// loop (the typical matmul-tile pattern), this is more efficient
260// than 32 calls to nx_q4k_iter_value because the scale/min are
261// hoisted out.
262//
263// ggml layout (verified bit-exact): sub-block sb draws from 32-byte
264// group g = sb/2. Output position l (0..32, in order) = byte
265// qs[g*32 + l], LOW nibble when sb is even (sb=2g), HIGH nibble when sb
266// is odd (sb=2g+1). Output unit Q24.
267//
268// sb in 0..8. out_32 must hold 32 i64.
269
270func nx_q4k_iter_subblock_dequant(it: *NxQ4KBlockIter, sb: nx_int,
271 out_32: *i64) -> nx_int {
272 if sb < 0 { return NX_Q4KI_ERR_BAD_IDX }
273 if sb >= 8 { return NX_Q4KI_ERR_BAD_IDX }
274
275 let sc_sb: i64 = _q4k_iter_scale_at(it, sb)
276 let m_sb: i64 = _q4k_iter_min_at(it, sb)
277 let d1: i64 = it.d_q10 * sc_sb
278 let m1: i64 = it.dmin_q10 * m_sb
279
280 let g: nx_int = sb / 2
281 let is_high: nx_int = sb - g * 2 // 0 low, 1 high
282 let grp_base: i64 = g * 32
283
284 var l: nx_int = 0
285 while l < 32 {
286 let byte_v: i64 = nx_le_read_u8(it.qs_ptr, grp_base + l)
287 var q4: i64 = 0
288 if is_high == 0 {
289 q4 = byte_v & 0x0F
290 } else {
291 q4 = byte_v >> 4
292 }
293 out_32[l] = d1 * q4 - m1
294 l = l + 1
295 }
296 return NX_Q4KI_OK
297}