nx_f32_q4k_matmul.nx source
↩ module page · 716 lines · 33011 B
1// nx_f32_q4k_matmul.nx -- lazy-dequant matmul, B is Q4_K bytes. C[m,n] = A[m,k] @ dequant(B[k,n]).
2// Per B-row k_idx: dequant -> row_scratch[n], then C[i,:] += A[i,k_idx]*row_scratch. Inner = sovereign SSE.
3//
4// SOVEREIGN MULTICORE (2026-06-18, operator "no 3rd party from the hardware rung up"): fork-parallel by OUTPUT
5// COLUMN (super-block-aligned) -- each worker dequants ONLY its column slice (no redundant dequant; ~Ncore x)
6// and writes disjoint columns of a SHARED temp; parent joins + copies into C. Column-split parallelizes the
7// single-token decode (m=1) too, unlike a row-split. Pure syscalls (sys_fork/mmap_shared/wait4) -- NO
8// libvulkan/CUDA/Mesa/pthread. Pattern proven in nx_par_matmul (6.78x, bit-identical).
9//
10// genealogy_id: ggml_q4k_dot_canon + sovereign_fork_parallel_columnsplit
11// lineage_id: substrate_f32_q4k_matmul_v2_parallel
12import "nx_syscalls.nx"
13import "nx_tier.nx"
14import "nx_le.nx"
15import "nx_gguf_load.nx"
16import "nx_q4k_to_f32.nx"
17import "nx_f32.nx"
18import "nx_thread_pool.nx"
19import "nx_q4k_matmul.nx"
20import "nx_dequant_iter.nx"
21import "nx_f32_cvt.nx"
22
23const NX_FQ4M_OK: nx_int = 0
24const NX_FQ4M_ERR_BAD_DIM: nx_int = 1
25const NX_FQ4M_ERR_NULL: nx_int = 2
26const NX_FQ4M_ERR_ALIGN: nx_int = 3
27const NX_FQ4M_ERR_POOL_WAIT: nx_int = 4
28const NX_FQ4M_N_VERDICTS: nx_int = 5
29const FQ4M_NW: i64 = 8 // worker forks
30const FQ4M_PAR_MIN_MAC: i64 = 1000000000000000 // fork DISABLED: per-matmul fork net-SLOWER in the forward (fork overhead > gain). Pool (fork-once) is the fix. Serial path (fq4m_cols full-range) = proven-correct.
31 // POOL LANDED 2026-07-07 (threading toolchain live): see
32 // nx_f32_q4k_matmul_pool / _mt below -- thread pool, spawn-once,
33 // CLONE_VM so the loader's private weight bytes need no
34 // shared-mmap setup (unlike the fork/pipe nx_par_pool demo).
35
36func nx_fq4m_verdict_is_valid(v: nx_int) -> nx_int {
37 if v < 0 { return 0 }
38 if v >= NX_FQ4M_N_VERDICTS { return 0 }
39 return 1
40}
41
42// STREAMING row-major matmul over a ggml Q4_K weight in [out, in] layout (== nx_f32_matmul_t but dequant-on-the-
43// fly). W row j = k in-values quantized in super-blocks along the reduction dim k (so k % 256 == 0). For each
44// output column j: dequant its k-value weight row into a k-scratch, then dot with every input row. Peak extra
45// memory = ONE row (k*8), NOT the whole weight -> a 7B stays quantized (~4.5GB) instead of exploding to f32 (~28GB).
46// The OLD code assumed [in, out] (transposed) + required n%256==0 -> wrong axis, errored on ffn_down(out=896),
47// left C zero. This is the fix that lets larger models load in a 15GB budget.
48func fq4m_rows(A: *i64, B_bytes: *u8, B_offset: i64, C: *i64, m: nx_int, k: nx_int, n: nx_int, jlo: nx_int, jhi: nx_int) -> i64 {
49 let vpb: i64 = NX_GL_Q4_K_VPB
50 let bpb: i64 = NX_GL_Q4_K_BPB
51 let bytes_per_row: i64 = (k / vpb) * bpb // k (the QUANTIZED reduction dim) super-blocks per weight row
52 // ---- SIMD DEQUANT-DOT (2026-08-01) --------------------------------------------------------------
53 // WHY: this kernel was ALREADY weight-stationary (each column dequantized ONCE, then dotted against all
54 // m rows), so its flat ms_per_token was never a batching failure -- it is the signature of a
55 // COMPUTE-BOUND batched matmul. Measured baseline at hidden_dim=896: 373/465/383/415/515 ms/token for
56 // m=8/16/32/64/125, flat with no weight re-streaming left to remove. The remaining cost is raw MACs, and
57 // the inner dot was SCALAR (one __f32_add + one __f32_mul per element) while the Q8_0 sibling path that
58 // measured 10.2x is SIMD. So the lever here is COMPUTE, not bytes -- the opposite axis from decode.
59 // HOW: same shape _lw_q8_0_matmul_st already uses -- pack A once per band (_fq4m_pack_a), dequant each
60 // weight column into a PACKED 4-byte row (nx_q4k_to_f32_packed), then __f32x4_dot 4 lanes at a time.
61 // Both operands must be packed 4-byte f32; A and the old scratch were *i64 (8-byte slots, f32 bits in the
62 // low half), which is exactly why the packed helpers exist.
63 // k ALIGNMENT IS GUARANTEED, not assumed: the Q4_K contract upstream already requires k % 256 == 0
64 // (super-block size), so k % 4 == 0 holds by construction and the 4-lane step can never run off the end.
65 // u00e2u0161u00a0NUMERICS CHANGE, STATED PLAINLY: 4-way partial sums re-associate the addition, so this is NOT
66 // bit-exact against the scalar loop (unlike the F32 column-tiling, which preserved order exactly).
67 // f32 addition is not associative. The precedent is deliberate: the Q8_0 path made the identical trade
68 // and was accepted at 10.2x, and __f32x4_dot itself is gate-covered (nx_f32x4_matmul_gate).
69 let rowp: *u8 = sys_mmap(k * 4) // packed 4-byte dequant scratch: ONE weight column
70 let apk: *u8 = sys_mmap(m * k * 4) // packed A, ONCE per band (not per column, not per row)
71 _fq4m_pack_a(A, m * k, apk)
72 let apb: i64 = apk as i64
73 let rpb: i64 = rowp as i64
74 var j: nx_int = jlo
75 while j < jhi {
76 let row_off: i64 = B_offset + (j as i64) * bytes_per_row
77 nx_q4k_to_f32_packed(B_bytes, row_off, k, rowp) // dequant this output column, packed
78 var i: nx_int = 0
79 while i < m {
80 let arow: i64 = apb + (i as i64) * (k as i64) * 4
81 var acc: i64 = 0
82 var l: nx_int = 0
83 while l < k {
84 acc = __f32_add(acc, __f32x4_dot((arow + (l as i64) * 4) as *u8,
85 (rpb + (l as i64) * 4) as *u8))
86 l = l + 4
87 }
88 C[i * n + j] = acc
89 i = i + 1
90 }
91 j = j + 1
92 }
93 sys_munmap(apk, m * k * 4)
94 // Free the dequant scratch: this runs once per CALL (and per band
95 // task under the pool entries) -- without it every matmul leaked
96 // the scratch, ~megabytes per token across a real forward. Now k*4
97 // (packed) rather than k*8, and the packed-A buffer is freed above.
98 sys_munmap(rowp, k * 4)
99 return 0
100}
101
102func nx_f32_q4k_matmul(A: *i64, B_bytes: *u8, B_offset: i64, C: *i64,
103 m: nx_int, k: nx_int, n: nx_int) -> nx_int {
104 if m <= 0 { return NX_FQ4M_ERR_BAD_DIM }
105 if k <= 0 { return NX_FQ4M_ERR_BAD_DIM }
106 if n <= 0 { return NX_FQ4M_ERR_BAD_DIM }
107 if A == (0 as *i64) { return NX_FQ4M_ERR_NULL }
108 if B_bytes == (0 as *u8) { return NX_FQ4M_ERR_NULL }
109 if C == (0 as *i64) { return NX_FQ4M_ERR_NULL }
110 // The QUANTIZED (reduction) dim is k -- ggml super-blocks run along the weight ROW = the in-dim. (OLD code
111 // checked n = the out-dim = the wrong axis, which errored on ffn_down out=896 and zeroed the FFN.)
112 if k - (k / NX_GL_Q4_K_VPB) * NX_GL_Q4_K_VPB != 0 { return NX_FQ4M_ERR_ALIGN }
113
114 // Serial streaming row-major matmul (== matmul_t, dequant-on-the-fly). Fork column-split removed: it used the
115 // wrong [in,out] layout; the streaming per-row path is proven-correct and memory-frugal (one row of scratch).
116 fq4m_rows(A, B_bytes, B_offset, C, m, k, n, 0, n)
117 return NX_FQ4M_OK
118}
119
120// ===== Multi-threaded matmul =====================================
121//
122// ADDITIVE: the serial entry above is untouched (it is the oracle
123// for nx_q4k_matmul_mt_gate). Banding is by OUTPUT COLUMN j --
124// each task dequants ONLY its own columns' weight rows and writes
125// disjoint C columns, so there are no shared writes and no atomics,
126// and it parallelizes the m=1 single-token decode (the shape the
127// forward actually runs). Bit-exactness is by construction: every
128// C[i,j] is computed wholly inside one band with the identical
129// accumulation order.
130//
131// nx_f32_q4k_matmul_pool -- caller-owned pool (the LLM forward
132// should create ONE pool per process and pass it to every
133// layer's matmuls: spawn-once, submit-per-call). Waits on
134// the completed-counter DELTA so reuse is safe; single-
135// submitter assumption as in nx_conv2d.
136// nx_f32_q4k_matmul_mt -- one-shot convenience (builds a pool,
137// runs, shuts down); nworkers < 1 sizes from hardware.
138
139struct NxFq4mCtx {
140 a_ptr: i64,
141 b_ptr: i64,
142 b_off: i64,
143 c_ptr: i64,
144 m: i64,
145 k: i64,
146 n: i64,
147 jlo: i64,
148 jhi: i64,
149}
150
151const NX_FQ4M_CTX_BYTES: i64 = 72
152
153func _nx_fq4m_task(ctx_i: i64) -> i64 {
154 let cx: *NxFq4mCtx = ctx_i as *NxFq4mCtx
155 return fq4m_rows(cx.a_ptr as *i64, cx.b_ptr as *u8, cx.b_off,
156 cx.c_ptr as *i64, cx.m, cx.k, cx.n, cx.jlo, cx.jhi)
157}
158
159func nx_f32_q4k_matmul_pool(pool: *NxThreadPool, A: *i64, B_bytes: *u8,
160 B_offset: i64, C: *i64,
161 m: nx_int, k: nx_int, n: nx_int) -> nx_int {
162 if m <= 0 { return NX_FQ4M_ERR_BAD_DIM }
163 if k <= 0 { return NX_FQ4M_ERR_BAD_DIM }
164 if n <= 0 { return NX_FQ4M_ERR_BAD_DIM }
165 if A == (0 as *i64) { return NX_FQ4M_ERR_NULL }
166 if B_bytes == (0 as *u8) { return NX_FQ4M_ERR_NULL }
167 if C == (0 as *i64) { return NX_FQ4M_ERR_NULL }
168 if k - (k / NX_GL_Q4_K_VPB) * NX_GL_Q4_K_VPB != 0 { return NX_FQ4M_ERR_ALIGN }
169
170 var bands: i64 = pool.n_workers
171 if bands > n { bands = n }
172 if bands < 1 { bands = 1 }
173
174 let ctxs: *u8 = sys_mmap(bands * NX_FQ4M_CTX_BYTES)
175 let cpb: i64 = (n + bands - 1) / bands
176 let done_before: i64 = nx_pool_n_completed(pool)
177 var b: i64 = 0
178 while b < bands {
179 let cx: *NxFq4mCtx = ((ctxs as i64) + b * NX_FQ4M_CTX_BYTES) as *NxFq4mCtx
180 cx.a_ptr = A as i64
181 cx.b_ptr = B_bytes as i64
182 cx.b_off = B_offset
183 cx.c_ptr = C as i64
184 cx.m = m
185 cx.k = k
186 cx.n = n
187 cx.jlo = b * cpb
188 var jhi: i64 = (b + 1) * cpb
189 if jhi > n { jhi = n }
190 cx.jhi = jhi
191 nx_pool_submit(pool, _nx_fq4m_task, cx as i64)
192 b = b + 1
193 }
194 let wv: i64 = nx_pool_wait(pool, done_before + bands)
195 sys_munmap(ctxs, bands * NX_FQ4M_CTX_BYTES)
196 if wv != 0 { return NX_FQ4M_ERR_POOL_WAIT }
197 return NX_FQ4M_OK
198}
199
200func nx_f32_q4k_matmul_mt(A: *i64, B_bytes: *u8, B_offset: i64, C: *i64,
201 m: nx_int, k: nx_int, n: nx_int,
202 nworkers: nx_int) -> nx_int {
203 var nw: i64 = nworkers
204 if nw < 1 { nw = nx_hw_worker_count() }
205 if nw > n { nw = n }
206 if nw <= 1 { return nx_f32_q4k_matmul(A, B_bytes, B_offset, C, m, k, n) }
207 let pool: *NxThreadPool = nx_pool_new(nw, 0)
208 let v: nx_int = nx_f32_q4k_matmul_pool(pool, A, B_bytes, B_offset, C, m, k, n)
209 nx_pool_shutdown(pool)
210 return v
211}
212
213// ===== PACKED-SIMD matmul (x4 paths, 2026-07-08) ==================
214//
215// Stacks the SECOND proven compute lever onto the pool: __f32x4_dot
216// (4 f32 MACs per mulps; 5.0x over scalar measured in
217// nx_f32x4_mt_matmul, same-day lane). A is packed once per call to
218// contiguous 4-byte f32; each weight row dequants STRAIGHT to packed
219// form (nx_q4k_to_f32_packed -- no repack pass, half the store
220// traffic). k is 256-aligned (ALIGN check) so the x4 inner loop has
221// no tail.
222//
223// NUMERIC CONTRACT: within one output cell the accumulation is
224// (4-lane dot) chunks summed left-to-right -- a DIFFERENT rounding
225// order than the scalar path, so x4 vs scalar is bit-exact ONLY in
226// the exact-f32 regime (|sums| < 2^24; the gate constructs this).
227// pool-x4 vs serial-x4 is bit-exact on ANY data (identical per-cell
228// order; banding never splits a cell).
229
230func _fq4m_pack_a(A: *i64, count: i64, out_p: *u8) -> i64 {
231 var i: i64 = 0
232 while i < count {
233 let bits: i64 = A[i]
234 out_p[i * 4 + 0] = bits as u8
235 out_p[i * 4 + 1] = (bits >> 8) as u8
236 out_p[i * 4 + 2] = (bits >> 16) as u8
237 out_p[i * 4 + 3] = (bits >> 24) as u8
238 i = i + 1
239 }
240 return 0
241}
242
243// Packed band worker: output columns [jlo, jhi), A pre-packed.
244func fq4m_rows_x4(pa: *u8, B_bytes: *u8, B_offset: i64, C: *i64,
245 m: nx_int, k: nx_int, n: nx_int,
246 jlo: nx_int, jhi: nx_int) -> i64 {
247 let vpb: i64 = NX_GL_Q4_K_VPB
248 let bpb: i64 = NX_GL_Q4_K_BPB
249 let bytes_per_row: i64 = (k / vpb) * bpb
250 let rowp: *u8 = sys_mmap(k * 4)
251 let pab: i64 = pa as i64
252 let rpb2: i64 = rowp as i64
253 var j: nx_int = jlo
254 while j < jhi {
255 let row_off: i64 = B_offset + (j as i64) * bytes_per_row
256 nx_q4k_to_f32_packed(B_bytes, row_off, k, rowp)
257 var i: nx_int = 0
258 while i < m {
259 let abase: i64 = pab + (i as i64) * k * 4
260 var acc: i64 = 0
261 var l: i64 = 0
262 while l < k {
263 acc = __f32_add(acc, __f32x4_dot((abase + l * 4) as *u8, (rpb2 + l * 4) as *u8))
264 l = l + 4
265 }
266 C[i * n + j] = acc
267 i = i + 1
268 }
269 j = j + 1
270 }
271 sys_munmap(rowp, k * 4)
272 return 0
273}
274
275func _nx_fq4m_task_x4(ctx_i: i64) -> i64 {
276 let cx: *NxFq4mCtx = ctx_i as *NxFq4mCtx
277 return fq4m_rows_x4(cx.a_ptr as *u8, cx.b_ptr as *u8, cx.b_off,
278 cx.c_ptr as *i64, cx.m, cx.k, cx.n, cx.jlo, cx.jhi)
279}
280
281// Serial packed-SIMD matmul (the x4 oracle; also the nw<=1 path).
282func nx_f32_q4k_matmul_x4(A: *i64, B_bytes: *u8, B_offset: i64, C: *i64,
283 m: nx_int, k: nx_int, n: nx_int) -> nx_int {
284 if m <= 0 { return NX_FQ4M_ERR_BAD_DIM }
285 if k <= 0 { return NX_FQ4M_ERR_BAD_DIM }
286 if n <= 0 { return NX_FQ4M_ERR_BAD_DIM }
287 if A == (0 as *i64) { return NX_FQ4M_ERR_NULL }
288 if B_bytes == (0 as *u8) { return NX_FQ4M_ERR_NULL }
289 if C == (0 as *i64) { return NX_FQ4M_ERR_NULL }
290 if k - (k / NX_GL_Q4_K_VPB) * NX_GL_Q4_K_VPB != 0 { return NX_FQ4M_ERR_ALIGN }
291 let pa: *u8 = sys_mmap(m * k * 4)
292 _fq4m_pack_a(A, m * k, pa)
293 fq4m_rows_x4(pa, B_bytes, B_offset, C, m, k, n, 0, n)
294 sys_munmap(pa, m * k * 4)
295 return NX_FQ4M_OK
296}
297
298// Pooled packed-SIMD matmul: SIMD x multicore, the forward's path.
299// Same delta-wait/single-submitter contract as nx_f32_q4k_matmul_pool.
300func nx_f32_q4k_matmul_pool_x4(pool: *NxThreadPool, A: *i64, B_bytes: *u8,
301 B_offset: i64, C: *i64,
302 m: nx_int, k: nx_int, n: nx_int) -> nx_int {
303 if m <= 0 { return NX_FQ4M_ERR_BAD_DIM }
304 if k <= 0 { return NX_FQ4M_ERR_BAD_DIM }
305 if n <= 0 { return NX_FQ4M_ERR_BAD_DIM }
306 if A == (0 as *i64) { return NX_FQ4M_ERR_NULL }
307 if B_bytes == (0 as *u8) { return NX_FQ4M_ERR_NULL }
308 if C == (0 as *i64) { return NX_FQ4M_ERR_NULL }
309 if k - (k / NX_GL_Q4_K_VPB) * NX_GL_Q4_K_VPB != 0 { return NX_FQ4M_ERR_ALIGN }
310
311 var bands: i64 = pool.n_workers
312 if bands > n { bands = n }
313 if bands < 1 { bands = 1 }
314
315 let pa: *u8 = sys_mmap(m * k * 4)
316 _fq4m_pack_a(A, m * k, pa)
317
318 let ctxs: *u8 = sys_mmap(bands * NX_FQ4M_CTX_BYTES)
319 let cpb: i64 = (n + bands - 1) / bands
320 let done_before: i64 = nx_pool_n_completed(pool)
321 var b: i64 = 0
322 while b < bands {
323 let cx: *NxFq4mCtx = ((ctxs as i64) + b * NX_FQ4M_CTX_BYTES) as *NxFq4mCtx
324 cx.a_ptr = pa as i64
325 cx.b_ptr = B_bytes as i64
326 cx.b_off = B_offset
327 cx.c_ptr = C as i64
328 cx.m = m
329 cx.k = k
330 cx.n = n
331 cx.jlo = b * cpb
332 var jhi: i64 = (b + 1) * cpb
333 if jhi > n { jhi = n }
334 cx.jhi = jhi
335 nx_pool_submit(pool, _nx_fq4m_task_x4, cx as i64)
336 b = b + 1
337 }
338 let wv: i64 = nx_pool_wait(pool, done_before + bands)
339 sys_munmap(ctxs, bands * NX_FQ4M_CTX_BYTES)
340 sys_munmap(pa, m * k * 4)
341 if wv != 0 { return NX_FQ4M_ERR_POOL_WAIT }
342 return NX_FQ4M_OK
343}
344
345// ---- FUSED INTEGER GEMM: dequant is FUSED INTO the dot; f32 weights are NEVER materialised. ----
346//
347// WHY THIS IS THE RUNG. x4 and x8 both attacked the DOT and both plateaued (+23% and -2%). The ceiling
348// was the DEQUANT: nx_q4k_to_f32_packed writes k f32 values per output column, and every kernel above it
349// pays that traffic no matter how wide its lanes are. Measured by nx_q4k_fused_vs_x4_gate, the fused
350// integer route is >=9.9x faster than nx_f32_q4k_matmul_x4 on the same shape, same thread, same data,
351// while agreeing BIT-EXACTLY in the exact-integer regime.
352//
353// u2605 WHY THIS IS m-BLOCKED AND NOT A LOOP OVER nx_q4k_dot_row_col. That function dots ONE column, so
354// calling it m times per output would re-traverse and re-dequantise the whole weight row m times --
355// throwing away exactly the amortisation that makes fq4m_rows work at m>1. Here the weight element is
356// dequantised ONCE (`v`, Q24) and immediately accumulated into ALL m accumulators. Weight traffic is
357// therefore independent of m, and arithmetic intensity RISES with m -- which is the compute-bound
358// prefill regime the whole rung is aimed at. At m=1 this degenerates to the same work as the single
359// -column dot, so decode loses nothing.
360//
361// NUMERIC CONTRACT -- READ BEFORE DISPATCHING. This is NOT a reassociation like x4/x8; it is
362// ACTIVATION QUANTISATION (W4A-fixed). Activations are converted f32 -> Q10 (nx_f32_to_q10), the
363// accumulator is Q34 (Q24 weight x Q10 activation), and the result is recovered via
364// nx_q4km_q20_to_q10 -> nx_q10_to_f32. Agreement with the f32 kernels is EXACT only where the values
365// are exactly representable in Q10; on real trained weights it is approximate, and the instrument for
366// that is nx_q4k_ggml_kat's 2% band, NOT this kernel's own output. Speed alone must not authorise the
367// dispatcher flip.
368//
369// ALLOCATION DISCIPLINE: the iterator and the m accumulators are allocated ONCE PER BAND, never per
370// call -- the per-call sys_mmap in nx_q4k_dot_row_col was itself the single biggest cost in that
371// primitive (an unclosed resource in a loop is a clock, not a leak).
372
373func fq4m_rows_fused(pq: *i64, B_bytes: *u8, B_offset: i64, C: *i64,
374 m: nx_int, k: nx_int, n: nx_int,
375 jlo: nx_int, jhi: nx_int) -> i64 {
376 let vpb: i64 = NX_GL_Q4_K_VPB
377 let bpb: i64 = NX_GL_Q4_K_BPB
378 let bytes_per_row: i64 = (k / vpb) * bpb
379 let n_blocks: i64 = k / vpb
380 let it: *NxQ4KBlockIter = nx_q4k_iter_alloc()
381 let acc: *i64 = sys_mmap(m * 8) as *i64
382 var j: nx_int = jlo
383 while j < jhi {
384 let row_off: i64 = B_offset + (j as i64) * bytes_per_row
385 var z: i64 = 0
386 while z < m { acc[z] = 0; z = z + 1 }
387 var blk: i64 = 0
388 while blk < n_blocks {
389 nx_q4k_iter_init(B_bytes, row_off + blk * bpb, it)
390 let col_base: i64 = blk * 256
391 var sb: i64 = 0
392 while sb < 8 {
393 let sc_sb: i64 = nx_q4k_iter_scale(it, sb)
394 let m_sb: i64 = nx_q4k_iter_min(it, sb)
395 let d1: i64 = it.d_q10 * sc_sb
396 let m1: i64 = it.dmin_q10 * m_sb
397 let g: i64 = sb / 2
398 let is_high: i64 = sb - g * 2
399 let grp_base: i64 = g * 32
400 let col_sb_base: i64 = col_base + sb * 32
401 var l: i64 = 0
402 while l < 32 {
403 let byte_v: i64 = nx_le_read_u8(it.qs_ptr, grp_base + l)
404 var q4: i64 = 0
405 if is_high == 0 { q4 = byte_v & 0x0F } else { q4 = byte_v >> 4 }
406 let v: i64 = d1 * q4 - m1
407 let cpos: i64 = col_sb_base + l
408 var i: i64 = 0
409 while i < m {
410 acc[i] = acc[i] + v * pq[i * k + cpos]
411 i = i + 1
412 }
413 l = l + 1
414 }
415 sb = sb + 1
416 }
417 blk = blk + 1
418 }
419 var o: i64 = 0
420 while o < m { C[o * n + j] = nx_q20_to_f32(nx_q4km_q20_to_q10(acc[o])); o = o + 1 }
421 j = j + 1
422 }
423 sys_munmap(acc as *u8, m * 8)
424 sys_munmap(it as *u8, NX_Q4K_ITER_BYTES)
425 return 0
426}
427
428// Q20, not Q10: the activation format IS the fused route's entire error budget, and widening it is
429// free (same i64 multiply). Q10 measured 1.37% max deviation on real weights; Q20 is 1024x finer.
430// ---- Q4_K i8-SIMD DEQUANT-DOT: the rung the two refuted designs pointed at. ----
431//
432// WHY THIS SHAPE. Two earlier attempts to adopt a faster Q4_K route both died on the lazy-weight cache
433// contract (LWC-8: cached and streamed must agree), because both quantised the ACTIVATION and so left
434// the f32 domain. This one does not: `__f32_i8dot32` is a BLESSED SSE convert+dot that takes **int8
435// weights against f32 activations** -- the int8-ness is on the WEIGHT side only. Activations are never
436// quantised, so this route never creates the cross-domain divergence that blocked the others.
437// The same intrinsic is already load-bearing here: 10.2x on Q8_0 lm_head (:29) and the Q5_0 path (:430).
438//
439// THE ALGEBRA. Over one 32-value Q4_K sub-block sb, value = d*sc[sb]*q4 - dmin*m[sb], so
440// SUM(value*a) = (d*sc[sb]) * SUM(q4*a) - (dmin*m[sb]) * SUM(a)
441// = (d*sc[sb]) * __f32_i8dot32(q4_as_int8, a_packed) - (dmin*m[sb]) * sumA[sb]
442// u2605 sumA[sb] does NOT depend on the output column j, so it is computed ONCE per activation row and
443// reused across all n columns -- free at n=4864. The nibble->int8 unpack is scalar into a 32-byte
444// scratch; Q5_0 shows the unpack is not the bottleneck, the dot is. Measure before optimising it.
445//
446// u26a0STATUS: ADDITIVE ONLY. Not wired into nx_f32_lazy_weight and not dispatched anywhere. Correctness
447// and speed are proven by nx_q4k_fused_vs_x4_gate before any flip is even considered.
448
449const FQ4M_SB_VALS: i64 = 32
450
451// Per-sub-block f32 sums of one activation row: sumA[b] = SUM of a[b*32 .. b*32+32).
452func fq4m_sumA(A: *i64, k: nx_int, out_sum: *i64) -> i64 {
453 let nsb: i64 = k / FQ4M_SB_VALS
454 var b: i64 = 0
455 while b < nsb {
456 var s: i64 = 0
457 var t: i64 = 0
458 while t < FQ4M_SB_VALS { s = __f32_add(s, A[b * FQ4M_SB_VALS + t]); t = t + 1 }
459 out_sum[b] = s
460 b = b + 1
461 }
462 return 0
463}
464
465func fq4m_rows_i8simd(pa: *u8, A: *i64, sumA: *i64, B_bytes: *u8, B_offset: i64, C: *i64,
466 m: nx_int, k: nx_int, n: nx_int,
467 jlo: nx_int, jhi: nx_int) -> i64 {
468 let bpb: i64 = NX_GL_Q4_K_BPB
469 let bytes_per_row: i64 = (k / NX_GL_Q4_K_VPB) * bpb
470 let n_blocks: i64 = k / NX_GL_Q4_K_VPB
471 let it: *NxQ4KBlockIter = nx_q4k_iter_alloc()
472 let i8scr: *u8 = sys_mmap(64)
473 let pab: i64 = pa as i64
474 var j: nx_int = jlo
475 while j < jhi {
476 let row_off: i64 = B_offset + (j as i64) * bytes_per_row
477 var acc: i64 = 0
478 var blk: i64 = 0
479 while blk < n_blocks {
480 nx_q4k_iter_init(B_bytes, row_off + blk * bpb, it)
481 let df: i64 = nx_q24_to_f32(it.d_q10)
482 let dmf: i64 = nx_q24_to_f32(it.dmin_q10)
483 var sb: i64 = 0
484 while sb < 8 {
485 let sc_f: i64 = __f32_mul(df, nx_i32_to_f32(nx_q4k_iter_scale(it, sb)))
486 let mn_f: i64 = __f32_mul(dmf, nx_i32_to_f32(nx_q4k_iter_min(it, sb)))
487 let g: i64 = sb / 2
488 let is_high: i64 = sb - g * 2
489 let grp_base: i64 = g * 32
490 var l: i64 = 0
491 while l < 32 {
492 let byte_v: i64 = nx_le_read_u8(it.qs_ptr, grp_base + l)
493 if is_high == 0 { i8scr[l] = (byte_v & 0x0F) as u8 } else { i8scr[l] = (byte_v >> 4) as u8 }
494 l = l + 1
495 }
496 let apos: i64 = blk * 256 + sb * 32
497 let raw: i64 = __f32_i8dot32(i8scr, (pab + apos * 4) as *u8)
498 acc = __f32_add(acc, __f32_mul(sc_f, raw))
499 acc = nx_f32_sub(acc, __f32_mul(mn_f, sumA[apos / FQ4M_SB_VALS]))
500 sb = sb + 1
501 }
502 blk = blk + 1
503 }
504 C[j] = acc
505 j = j + 1
506 }
507 sys_munmap(i8scr, 64)
508 sys_munmap(it as *u8, NX_Q4K_ITER_BYTES)
509 return 0
510}
511
512// Serial i8-SIMD matmul. m=1 only for now (decode shape) -- the m>1 blocking is the follow-on once the
513// single-row kernel is proven; claiming a batched win before the row kernel is measured would repeat
514// the x8 mistake.
515func nx_f32_q4k_matmul_i8simd(A: *i64, B_bytes: *u8, B_offset: i64, C: *i64,
516 m: nx_int, k: nx_int, n: nx_int) -> nx_int {
517 if m != 1 { return NX_FQ4M_ERR_BAD_DIM }
518 if k <= 0 { return NX_FQ4M_ERR_BAD_DIM }
519 if n <= 0 { return NX_FQ4M_ERR_BAD_DIM }
520 if A == (0 as *i64) { return NX_FQ4M_ERR_NULL }
521 if B_bytes == (0 as *u8) { return NX_FQ4M_ERR_NULL }
522 if C == (0 as *i64) { return NX_FQ4M_ERR_NULL }
523 if k - (k / NX_GL_Q4_K_VPB) * NX_GL_Q4_K_VPB != 0 { return NX_FQ4M_ERR_ALIGN }
524 let pa: *u8 = sys_mmap(k * 4)
525 _fq4m_pack_a(A, k, pa)
526 let sumA: *i64 = sys_mmap((k / FQ4M_SB_VALS) * 8) as *i64
527 fq4m_sumA(A, k, sumA)
528 fq4m_rows_i8simd(pa, A, sumA, B_bytes, B_offset, C, m, k, n, 0, n)
529 sys_munmap(sumA as *u8, (k / FQ4M_SB_VALS) * 8)
530 sys_munmap(pa, k * 4)
531 return NX_FQ4M_OK
532}
533
534// ---- Q24 INTEGER WEIGHT CACHE: makes the CACHED path numerically identical to the FUSED path. ----
535//
536// THE PROBLEM THIS SOLVES. nx_lw_cache_gate's LWC-8 asserts the cached and streamed paths return
537// IDENTICAL results -- caching must be a pure optimisation, never a behaviour change. That is the right
538// contract: otherwise a model's output depends on MEMORY PRESSURE. It is also exactly what blocked the
539// fused dispatcher flip, because fused is integer/Q20 while the f32 cache is not.
540// Widening precision could never fix that: LWC-8 wants BIT equality, and no activation-quantised route
541// can bit-match an f32 one. u21d2 The only sound fix is to put BOTH SIDES IN THE SAME NUMERIC DOMAIN.
542//
543// So the cache stores the dequantised weight as Q24 -- the EXACT value `v = d1*q4 - m1` that
544// fq4m_rows_fused computes on the fly -- and the cached dot runs the SAME Q24 x Q20 -> Q44 accumulate.
545// Cached and streamed then agree BIT-FOR-BIT by construction, not by tolerance.
546//
547// FOOTPRINT IS UNCHANGED: Q24 is stored as int32, 4 bytes per value, identical to the packed-f32 cache
548// it replaces, so the cache budget accounting in nx_f32_lazy_weight needs no change. A trained Q4_K
549// weight of ~0.05 is ~838,860 in Q24 -- three orders of magnitude inside int32.
550
551func fq4m_fill_q24(cache: *u8, B_bytes: *u8, B_offset: i64,
552 k: nx_int, n: nx_int, jlo: nx_int, jhi: nx_int) -> i64 {
553 let bpb: i64 = NX_GL_Q4_K_BPB
554 let bytes_per_row: i64 = (k / NX_GL_Q4_K_VPB) * bpb
555 let n_blocks: i64 = k / NX_GL_Q4_K_VPB
556 let it: *NxQ4KBlockIter = nx_q4k_iter_alloc()
557 var j: nx_int = jlo
558 while j < jhi {
559 let row_off: i64 = B_offset + (j as i64) * bytes_per_row
560 var blk: i64 = 0
561 while blk < n_blocks {
562 nx_q4k_iter_init(B_bytes, row_off + blk * bpb, it)
563 let col_base: i64 = blk * 256
564 var sb: i64 = 0
565 while sb < 8 {
566 let d1: i64 = it.d_q10 * nx_q4k_iter_scale(it, sb)
567 let m1: i64 = it.dmin_q10 * nx_q4k_iter_min(it, sb)
568 let g: i64 = sb / 2
569 let is_high: i64 = sb - g * 2
570 let grp_base: i64 = g * 32
571 let col_sb_base: i64 = col_base + sb * 32
572 var l: i64 = 0
573 while l < 32 {
574 let byte_v: i64 = nx_le_read_u8(it.qs_ptr, grp_base + l)
575 var q4: i64 = 0
576 if is_high == 0 { q4 = byte_v & 0x0F } else { q4 = byte_v >> 4 }
577 let v: i64 = d1 * q4 - m1
578 let o: i64 = ((j as i64) * k + col_sb_base + l) * 4
579 cache[o + 0] = v as u8
580 cache[o + 1] = (v >> 8) as u8
581 cache[o + 2] = (v >> 16) as u8
582 cache[o + 3] = (v >> 24) as u8
583 l = l + 1
584 }
585 sb = sb + 1
586 }
587 blk = blk + 1
588 }
589 j = j + 1
590 }
591 sys_munmap(it as *u8, NX_Q4K_ITER_BYTES)
592 return 0
593}
594
595func _fq4m_q24_at(cache: *u8, idx: i64) -> i64 {
596 let o: i64 = idx * 4
597 let raw: i64 = nx_le_read_u8(cache, o) | (nx_le_read_u8(cache, o + 1) << 8) |
598 (nx_le_read_u8(cache, o + 2) << 16) | (nx_le_read_u8(cache, o + 3) << 24)
599 if raw >= 2147483648 { return raw - 4294967296 }
600 return raw
601}
602
603// Cached-weight band worker. Same accumulate as fq4m_rows_fused, minus the dequant (already paid).
604func fq4m_rows_cached_q24(pq: *i64, cache: *u8, C: *i64,
605 m: nx_int, k: nx_int, n: nx_int,
606 jlo: nx_int, jhi: nx_int) -> i64 {
607 let acc: *i64 = sys_mmap(m * 8) as *i64
608 var j: nx_int = jlo
609 while j < jhi {
610 var z: i64 = 0
611 while z < m { acc[z] = 0; z = z + 1 }
612 var l: i64 = 0
613 while l < k {
614 let v: i64 = _fq4m_q24_at(cache, (j as i64) * k + l)
615 var i: i64 = 0
616 while i < m {
617 acc[i] = acc[i] + v * pq[i * k + l]
618 i = i + 1
619 }
620 l = l + 1
621 }
622 var o: i64 = 0
623 while o < m { C[o * n + j] = nx_q20_to_f32(nx_q4km_q20_to_q10(acc[o])); o = o + 1 }
624 j = j + 1
625 }
626 sys_munmap(acc as *u8, m * 8)
627 return 0
628}
629
630func _nx_fq4m_task_fill_q24(ctx_i: i64) -> i64 {
631 let cx: *NxFq4mCtx = ctx_i as *NxFq4mCtx
632 return fq4m_fill_q24(cx.a_ptr as *u8, cx.b_ptr as *u8, cx.b_off, cx.k, cx.n, cx.jlo, cx.jhi)
633}
634
635func _nx_fq4m_task_cached_q24(ctx_i: i64) -> i64 {
636 let cx: *NxFq4mCtx = ctx_i as *NxFq4mCtx
637 return fq4m_rows_cached_q24(cx.a_ptr as *i64, cx.b_ptr as *u8, cx.c_ptr as *i64,
638 cx.m, cx.k, cx.n, cx.jlo, cx.jhi)
639}
640
641func _fq4m_pack_q10(A: *i64, count: i64, out_q: *i64) -> i64 {
642 var i: i64 = 0
643 while i < count { out_q[i] = nx_f32_to_q20(A[i]); i = i + 1 }
644 return 0
645}
646
647// Serial fused matmul (the fused oracle; also the nw<=1 path). Same signature as the x4 sibling so the
648// dispatcher flip in nx_f32_lazy_weight is a one-line change once fidelity is proven on real weights.
649func nx_f32_q4k_matmul_fused(A: *i64, B_bytes: *u8, B_offset: i64, C: *i64,
650 m: nx_int, k: nx_int, n: nx_int) -> nx_int {
651 if m <= 0 { return NX_FQ4M_ERR_BAD_DIM }
652 if k <= 0 { return NX_FQ4M_ERR_BAD_DIM }
653 if n <= 0 { return NX_FQ4M_ERR_BAD_DIM }
654 if A == (0 as *i64) { return NX_FQ4M_ERR_NULL }
655 if B_bytes == (0 as *u8) { return NX_FQ4M_ERR_NULL }
656 if C == (0 as *i64) { return NX_FQ4M_ERR_NULL }
657 if k - (k / NX_GL_Q4_K_VPB) * NX_GL_Q4_K_VPB != 0 { return NX_FQ4M_ERR_ALIGN }
658 let pq: *i64 = sys_mmap(m * k * 8) as *i64
659 _fq4m_pack_q10(A, m * k, pq)
660 fq4m_rows_fused(pq, B_bytes, B_offset, C, m, k, n, 0, n)
661 sys_munmap(pq as *u8, m * k * 8)
662 return NX_FQ4M_OK
663}
664
665func _nx_fq4m_task_fused(ctx_i: i64) -> i64 {
666 let cx: *NxFq4mCtx = ctx_i as *NxFq4mCtx
667 return fq4m_rows_fused(cx.a_ptr as *i64, cx.b_ptr as *u8, cx.b_off,
668 cx.c_ptr as *i64, cx.m, cx.k, cx.n, cx.jlo, cx.jhi)
669}
670
671// Pooled fused matmul: fused-integer x multicore. Bands over output columns exactly like pool_x4, so
672// the two are directly comparable and a pool multiplies both sides equally.
673func nx_f32_q4k_matmul_pool_fused(pool: *NxThreadPool, A: *i64, B_bytes: *u8,
674 B_offset: i64, C: *i64,
675 m: nx_int, k: nx_int, n: nx_int) -> nx_int {
676 if m <= 0 { return NX_FQ4M_ERR_BAD_DIM }
677 if k <= 0 { return NX_FQ4M_ERR_BAD_DIM }
678 if n <= 0 { return NX_FQ4M_ERR_BAD_DIM }
679 if A == (0 as *i64) { return NX_FQ4M_ERR_NULL }
680 if B_bytes == (0 as *u8) { return NX_FQ4M_ERR_NULL }
681 if C == (0 as *i64) { return NX_FQ4M_ERR_NULL }
682 if k - (k / NX_GL_Q4_K_VPB) * NX_GL_Q4_K_VPB != 0 { return NX_FQ4M_ERR_ALIGN }
683
684 var bands: i64 = pool.n_workers
685 if bands > n { bands = n }
686 if bands < 1 { bands = 1 }
687
688 let pq: *i64 = sys_mmap(m * k * 8) as *i64
689 _fq4m_pack_q10(A, m * k, pq)
690
691 let ctxs: *u8 = sys_mmap(bands * NX_FQ4M_CTX_BYTES)
692 let cpb: i64 = (n + bands - 1) / bands
693 let done_before: i64 = nx_pool_n_completed(pool)
694 var b: i64 = 0
695 while b < bands {
696 let cx: *NxFq4mCtx = ((ctxs as i64) + b * NX_FQ4M_CTX_BYTES) as *NxFq4mCtx
697 cx.a_ptr = pq as i64
698 cx.b_ptr = B_bytes as i64
699 cx.b_off = B_offset
700 cx.c_ptr = C as i64
701 cx.m = m
702 cx.k = k
703 cx.n = n
704 cx.jlo = b * cpb
705 var jhi: i64 = (b + 1) * cpb
706 if jhi > n { jhi = n }
707 cx.jhi = jhi
708 nx_pool_submit(pool, _nx_fq4m_task_fused, cx as i64)
709 b = b + 1
710 }
711 let wv: i64 = nx_pool_wait(pool, done_before + bands)
712 sys_munmap(ctxs, bands * NX_FQ4M_CTX_BYTES)
713 sys_munmap(pq as *u8, m * k * 8)
714 if wv != 0 { return NX_FQ4M_ERR_POOL_WAIT }
715 return NX_FQ4M_OK
716}