nx_f32_matmul_t.nx source
↩ module page · 334 lines · 17112 B
1// nx_f32_matmul_t.nx -- matmul where B is laid out per ggml convention (B[k,j] at B[k + j*K]).
2// Same C[i,j] = sum_k A[i,k]*B[k,j] math. Inner MAC = sovereign SSE (__f32_add/__f32_mul, bit-identical IEEE f32).
3//
4// SOVEREIGN MULTICORE (2026-06-18, operator "no 3rd party from the hardware rung up"): for big matmuls this
5// forks MMT_NW workers (sys_fork), each computing a flat output-range into a SHARED temp (sys_mmap_shared),
6// parent joins (sys_wait4) and copies the temp into C. Drop-in: same signature + output, C may be private (the
7// shared temp bridges). Pure syscalls -- NO libvulkan/CUDA/Mesa/pthread. Pattern proven in nx_par_matmul (6.78x,
8// bit-identical). Small matmuls stay serial (fork overhead not worth it).
9//
10// genealogy_id: standard_matmul + ggml_dim_0_fast_layout + sovereign_fork_parallel
11// lineage_id: substrate_f32_matmul_t_v2_parallel
12import "nx_syscalls.nx"
13import "nx_tier.nx"
14import "nx_f32.nx"
15import "nx_f32_matmul.nx"
16import "nx_thread_pool.nx"
17
18const MMT_NW: i64 = 8 // worker forks
19const MMT_PAR_MIN_MAC: i64 = 1000000000000000 // fork DISABLED: per-matmul fork was net-SLOWER in the forward (1300+ forks of a 491MB proc > the parallelism gain). Pool (fork-once) is the fix. Serial path (mmt_range) is the proven-correct SSE serial.
20
21// ===== POOL (fork-once) variant, 2026-07-08 =======================
22//
23// The "Pool (fork-once) is the fix" note above, delivered. Threads
24// mmt_range on the SHARED nx_thread_pool by FLAT output-range bands
25// (idx = i*n+j maps 1:1 to output cells; a band is a contiguous flat
26// range, so no cell is split and pool == serial is BIT-EXACT on any
27// data -- same per-cell accumulation order). Caller-owned pool
28// (spawn-once, one per process): THE fix for lm_head, which ran
29// mmt_range serial every token (136M MACs, the forward's biggest
30// single matmul, measured 2026-07-08). Delta-wait so pool reuse is
31// safe; single-submitter contract (the forward drives from one
32// thread).
33
34struct NxMmtCtx {
35 a_ptr: i64,
36 b_ptr: i64,
37 c_ptr: i64,
38 k_dim: i64,
39 n_dim: i64,
40 lo: i64,
41 hi: i64,
42 m_dim: i64,
43}
44
45const NX_MMT_CTX_BYTES: i64 = 64
46
47// Pool task for the m>1 COLUMN-TILED kernel. lo/hi are a COLUMN range [j0,j1), not flat cell indices --
48// a band must own entire columns or it re-reads them, which is the whole defect being fixed.
49func _nx_mmt_tile_task(ctx_i: i64) -> i64 {
50 let cx: *NxMmtCtx = ctx_i as *NxMmtCtx
51 return mmt_tile_cols(cx.a_ptr as *i64, cx.b_ptr as *i64, cx.c_ptr as *i64,
52 cx.m_dim, cx.k_dim, cx.n_dim, cx.lo, cx.hi)
53}
54
55func _nx_mmt_task(ctx_i: i64) -> i64 {
56 let cx: *NxMmtCtx = ctx_i as *NxMmtCtx
57 return mmt_range(cx.a_ptr as *i64, cx.b_ptr as *i64, cx.c_ptr as *i64,
58 cx.k_dim, cx.n_dim, cx.lo, cx.hi)
59}
60
61// compute the flat output-range C[lo..hi) (flat idx = i*n + j); B is ggml column-major (B[kk + j*k]).
62func mmt_range(A: *i64, B: *i64, C: *i64, k: nx_int, n: nx_int, lo: i64, hi: i64) -> i64 {
63 var i: i64 = lo / n
64 var j: i64 = lo % n
65 var idx: i64 = lo
66 while idx < hi {
67 var sum: i64 = 0
68 var kk: nx_int = 0
69 while kk < k { sum = __f32_add(sum, __f32_mul(A[i * k + kk], B[kk + j * k])); kk = kk + 1 }
70 C[idx] = sum
71 idx = idx + 1
72 j = j + 1
73 if j >= n { j = 0; i = i + 1 }
74 }
75 return 0
76}
77
78// ---- m>1 COLUMN-TILED KERNEL (2026-07-31) -----------------------------------------------------
79// THE DEFECT IT FIXES, MEASURED not guessed: mmt_range walks OUTPUT CELLS, so for every cell (i,j) it
80// streams the whole column B[*,j] again. With m rows that means EACH COLUMN OF W IS READ m TIMES.
81// nx_batchscale_kat measured the consequence directly at the real Qwen inner dim (k=896): us_per_row
82// 4705/4056/4295/4191 for m=1/2/4/8 -- FLAT, a batch speedup of 1.12x where batching should give ~4x
83// (external: Graviton3 Llama3-8B-4bit 45.5 tok/s at b=1 -> 184.8 at b=8). Flat per-row cost IS the
84// signature of re-streaming the weights, and it is why prefill never amortised.
85//
86// THE FIX: hoist the column. For each j, load B[kk + j*k] ONCE and accumulate it against ALL m rows.
87// Weight traffic drops from m*k*n to k*n; A is re-read instead, but A is m*k which is negligible beside
88// the k*n weight matrix (at k=896,n=4864,m=8 that is 7k cells vs 4.3M).
89//
90// ★BIT-EXACT BY CONSTRUCTION, and this is the load-bearing property: for any given (i,j) the accumulation
91// ORDER over kk is IDENTICAL to mmt_range -- sum starts at f32 zero and adds kk ascending. Only the
92// INTERLEAVING across i changes, and f32 addition order within a cell is untouched. So this is a memory
93// -traffic change, NOT a numerics change, and a bit-exact equality check against mmt_range is a legitimate
94// oracle rather than a tolerance comparison. nx_matmul_tile_kat asserts exactly that.
95//
96// Banded by COLUMN RANGE [j0,j1) rather than by flat cell index, because the whole point is that a band
97// must own entire columns -- a flat cell band would split a column across workers and re-read it anyway.
98func mmt_tile_cols(A: *i64, B: *i64, C: *i64, m: i64, k: i64, n: i64, j0: i64, j1: i64) -> i64 {
99 if m <= 0 { return 0 }
100 let acc: *i64 = sys_mmap(m * 8) as *i64
101 var j: i64 = j0
102 while j < j1 {
103 var z: i64 = 0
104 while z < m { acc[z] = 0; z = z + 1 }
105 var kk: i64 = 0
106 while kk < k {
107 let bv: i64 = B[kk + j * k]
108 var r: i64 = 0
109 while r < m {
110 acc[r] = __f32_add(acc[r], __f32_mul(A[r * k + kk], bv))
111 r = r + 1
112 }
113 kk = kk + 1
114 }
115 var w: i64 = 0
116 while w < m { C[w * n + j] = acc[w]; w = w + 1 }
117 j = j + 1
118 }
119 return 0
120}
121
122// ---- 4x4 REGISTER-BLOCKED matmul_t (2026-09-15, the cross-encoder throughput rung) ----------------------
123// WHY. mmt_range and mmt_tile_cols both move m*n*k*8 bytes per matmul (every weight column once per row, or
124// every activation element once per column), and mmt_range's one accumulator per cell is a DEPENDENT f32 add
125// chain, so a core waits on the previous add before the next. One cross-encoder forward (T 300..512, D 384,
126// I 1536, 6 layers) therefore streamed ~30 GB through the caches at ~1.8 GMAC/s on one core, and eight
127// forwards side by side ran NO faster than one -- measured 2026-09-15 on the laptop: 5 s per forward at
128// 8-wide (thread pool AND fork-join, identically) against 1.6 s serial, because the eight shared one memory
129// system. Blocking 4 rows x 4 columns keeps 16 sums live per kk step: each loaded A value serves 4 columns and
130// each loaded B value serves 4 rows (traffic divided by 4), and the 16 sums are independent (no chain).
131// BIT-EXACT WITH mmt_range BY CONSTRUCTION: every cell still accumulates kk ascending from f32 zero in its own
132// sum; only the interleaving ACROSS cells changes, so nx_f32_matmul_block_gate asserts equality, never a
133// tolerance. ADDITIVE: callers opt in (the cross-encoder does); the dispatcher is untouched.
134const MMT_BR: i64 = 4
135const MMT_BC: i64 = 4
136// rows [i0,i1) against every column; the 4x4 interior is spelled with named sums, the ragged edge (m or n not a
137// multiple of 4) through a small accumulator array with the identical per-cell order
138func mmt_block(A: *i64, B: *i64, C: *i64, m: i64, k: i64, n: i64, i0: i64, i1: i64) -> i64 {
139 let acc: *i64 = sys_mmap(MMT_BR * MMT_BC * 8) as *i64
140 var i: i64 = i0
141 while i < i1 {
142 var rb: i64 = MMT_BR
143 if i1 - i < rb { rb = i1 - i }
144 var j: i64 = 0
145 while j < n {
146 var cb: i64 = MMT_BC
147 if n - j < cb { cb = n - j }
148 if rb == MMT_BR { if cb == MMT_BC {
149 let ia0: i64 = i * k
150 let ia1: i64 = ia0 + k
151 let ia2: i64 = ia1 + k
152 let ia3: i64 = ia2 + k
153 let ib0: i64 = j * k
154 let ib1: i64 = ib0 + k
155 let ib2: i64 = ib1 + k
156 let ib3: i64 = ib2 + k
157 var s00: i64 = 0; var s01: i64 = 0; var s02: i64 = 0; var s03: i64 = 0
158 var s10: i64 = 0; var s11: i64 = 0; var s12: i64 = 0; var s13: i64 = 0
159 var s20: i64 = 0; var s21: i64 = 0; var s22: i64 = 0; var s23: i64 = 0
160 var s30: i64 = 0; var s31: i64 = 0; var s32: i64 = 0; var s33: i64 = 0
161 var kk: i64 = 0
162 while kk < k {
163 let a0: i64 = A[ia0 + kk]; let a1: i64 = A[ia1 + kk]; let a2: i64 = A[ia2 + kk]; let a3: i64 = A[ia3 + kk]
164 let b0: i64 = B[ib0 + kk]; let b1: i64 = B[ib1 + kk]; let b2: i64 = B[ib2 + kk]; let b3: i64 = B[ib3 + kk]
165 s00 = __f32_add(s00, __f32_mul(a0, b0)); s01 = __f32_add(s01, __f32_mul(a0, b1)); s02 = __f32_add(s02, __f32_mul(a0, b2)); s03 = __f32_add(s03, __f32_mul(a0, b3))
166 s10 = __f32_add(s10, __f32_mul(a1, b0)); s11 = __f32_add(s11, __f32_mul(a1, b1)); s12 = __f32_add(s12, __f32_mul(a1, b2)); s13 = __f32_add(s13, __f32_mul(a1, b3))
167 s20 = __f32_add(s20, __f32_mul(a2, b0)); s21 = __f32_add(s21, __f32_mul(a2, b1)); s22 = __f32_add(s22, __f32_mul(a2, b2)); s23 = __f32_add(s23, __f32_mul(a2, b3))
168 s30 = __f32_add(s30, __f32_mul(a3, b0)); s31 = __f32_add(s31, __f32_mul(a3, b1)); s32 = __f32_add(s32, __f32_mul(a3, b2)); s33 = __f32_add(s33, __f32_mul(a3, b3))
169 kk = kk + 1
170 }
171 let ic0: i64 = i * n + j
172 let ic1: i64 = ic0 + n
173 let ic2: i64 = ic1 + n
174 let ic3: i64 = ic2 + n
175 C[ic0] = s00; C[ic0 + 1] = s01; C[ic0 + 2] = s02; C[ic0 + 3] = s03
176 C[ic1] = s10; C[ic1 + 1] = s11; C[ic1 + 2] = s12; C[ic1 + 3] = s13
177 C[ic2] = s20; C[ic2 + 1] = s21; C[ic2 + 2] = s22; C[ic2 + 3] = s23
178 C[ic3] = s30; C[ic3 + 1] = s31; C[ic3 + 2] = s32; C[ic3 + 3] = s33
179 } }
180 if rb < MMT_BR || cb < MMT_BC {
181 var z: i64 = 0
182 while z < MMT_BR * MMT_BC { acc[z] = 0; z = z + 1 }
183 var kq: i64 = 0
184 while kq < k {
185 var r: i64 = 0
186 while r < rb {
187 let av: i64 = A[(i + r) * k + kq]
188 var c: i64 = 0
189 while c < cb { acc[r * MMT_BC + c] = __f32_add(acc[r * MMT_BC + c], __f32_mul(av, B[kq + (j + c) * k])); c = c + 1 }
190 r = r + 1
191 }
192 kq = kq + 1
193 }
194 var r2: i64 = 0
195 while r2 < rb { var c2: i64 = 0; while c2 < cb { C[(i + r2) * n + j + c2] = acc[r2 * MMT_BC + c2]; c2 = c2 + 1 } r2 = r2 + 1 }
196 }
197 j = j + cb
198 }
199 i = i + rb
200 }
201 sys_munmap(acc as *u8, MMT_BR * MMT_BC * 8)
202 return 0
203}
204// Serial 4x4-blocked matmul_t. Same signature and bit-identical results to nx_f32_matmul_t; a quarter of the traffic.
205func nx_f32_matmul_t_blocked(A: *i64, B: *i64, C: *i64,
206 m: nx_int, k: nx_int, n: nx_int) -> nx_int {
207 if m <= 0 { return NX_F32_MM_ERR_BAD_DIM }
208 if k <= 0 { return NX_F32_MM_ERR_BAD_DIM }
209 if n <= 0 { return NX_F32_MM_ERR_BAD_DIM }
210 mmt_block(A, B, C, m as i64, k as i64, n as i64, 0, m as i64)
211 return NX_F32_MM_OK
212}
213
214// Serial column-tiled matmul_t. Same signature and same results as nx_f32_matmul_t, fewer weight reads.
215// ADDITIVE (rule 19): nothing is rewired by this file -- callers opt in, and the dispatcher change is a
216// separate deliberate act once the KAT has proven equality on the shapes that matter.
217func nx_f32_matmul_t_tiled(A: *i64, B: *i64, C: *i64,
218 m: nx_int, k: nx_int, n: nx_int) -> nx_int {
219 if m <= 0 { return NX_F32_MM_ERR_BAD_DIM }
220 if k <= 0 { return NX_F32_MM_ERR_BAD_DIM }
221 if n <= 0 { return NX_F32_MM_ERR_BAD_DIM }
222 mmt_tile_cols(A, B, C, m as i64, k as i64, n as i64, 0, n as i64)
223 return NX_F32_MM_OK
224}
225
226func nx_f32_matmul_t(A: *i64, B: *i64, C: *i64,
227 m: nx_int, k: nx_int, n: nx_int) -> nx_int {
228 if m <= 0 { return NX_F32_MM_ERR_BAD_DIM }
229 if k <= 0 { return NX_F32_MM_ERR_BAD_DIM }
230 if n <= 0 { return NX_F32_MM_ERR_BAD_DIM }
231
232 let total: i64 = (m as i64) * (n as i64)
233 // small -> serial (fork overhead not worth it)
234 if total * (k as i64) < MMT_PAR_MIN_MAC {
235 mmt_range(A, B, C, k, n, 0, total)
236 return NX_F32_MM_OK
237 }
238
239 // big -> SOVEREIGN fork-parallel: workers write flat-ranges into a shared temp, parent copies into C
240 let Ctmp: *i64 = sys_mmap_shared(total * 8) as *i64
241 var w: i64 = 0
242 while w < MMT_NW {
243 let pid: i64 = sys_fork()
244 if pid == 0 {
245 let lo: i64 = (w * total) / MMT_NW
246 let hi: i64 = ((w + 1) * total) / MMT_NW
247 mmt_range(A, B, Ctmp, k, n, lo, hi)
248 sys_exit(0)
249 }
250 w = w + 1
251 }
252 let stp: *i64 = sys_mmap(8) as *i64
253 var d: i64 = 0
254 while d < MMT_NW { sys_wait4(0 - 1, stp, 0); d = d + 1 }
255 var c: i64 = 0
256 while c < total { C[c] = Ctmp[c]; c = c + 1 }
257 return NX_F32_MM_OK
258}
259
260// Pooled matmul_t: flat-range bands on a caller-owned pool.
261// Bit-exact vs mmt_range (serial) on any data.
262func nx_f32_matmul_t_pool(pool: *NxThreadPool, A: *i64, B: *i64, C: *i64,
263 m: nx_int, k: nx_int, n: nx_int) -> nx_int {
264 if m <= 0 { return NX_F32_MM_ERR_BAD_DIM }
265 if k <= 0 { return NX_F32_MM_ERR_BAD_DIM }
266 if n <= 0 { return NX_F32_MM_ERR_BAD_DIM }
267
268 // ---- m>1 => COLUMN-TILED, WEIGHT-STATIONARY (2026-08-01) ---------------------------------------
269 // THE LAST HOP. mmt_tile_cols (the batching kernel) was built 07-31 with the measurement that proves
270 // the need -- nx_batchscale_kat at the real Qwen inner dim k=896 recorded us_per_row 4705/4056/4295/4191
271 // for m=1/2/4/8: FLAT, a 1.12x batch speedup where ~4x is expected -- but NOTHING CALLED IT. The pooled
272 // path below bands over FLAT CELLS and runs mmt_range, which walks output cells and therefore re-streams
273 // each weight column once PER ROW, so weight traffic is m*k*n instead of k*n. The forward only ever uses
274 // the pooled entry point, so the batching kernel was unreachable from production: a kernel with no
275 // caller is not an optimisation, it is a file. That is why prefill never amortised and why the sev-9
276 // rows measured per-token cost flat across a 15x range of m (606/673/567/593 ms/tok at m=8/32/64/125).
277 // Routing m>1 here converts prefill from memory-bound re-streaming to the compute-bound regime batching
278 // is supposed to reach; m=1 decode is left EXACTLY as it was, because at m=1 there is nothing to amortise
279 // (each weight is read once either way) and that path is already measured and proven.
280 // BIT-EXACT, not approximately equal: mmt_tile_cols accumulates kk ascending from f32 zero for every
281 // (i,j), the identical order mmt_range uses -- only the interleaving across i changes. So this is a
282 // memory-traffic change and NOT a numerics change, and nx_matmul_tile_kat asserts that equality directly.
283 // u00e2u02dcu2026A KERNEL WITH NO CALLER IS NOT AN OPTIMISATION, IT IS A FILE.
284 if m > 1 {
285 var cbands: i64 = pool.n_workers
286 if cbands > (n as i64) { cbands = n as i64 }
287 if cbands < 1 { cbands = 1 }
288 let cctxs: *u8 = sys_mmap(cbands * NX_MMT_CTX_BYTES)
289 let cdone: i64 = nx_pool_n_completed(pool)
290 var cb: i64 = 0
291 while cb < cbands {
292 let ccx: *NxMmtCtx = ((cctxs as i64) + cb * NX_MMT_CTX_BYTES) as *NxMmtCtx
293 ccx.a_ptr = A as i64
294 ccx.b_ptr = B as i64
295 ccx.c_ptr = C as i64
296 ccx.k_dim = k
297 ccx.n_dim = n
298 ccx.m_dim = m as i64
299 ccx.lo = (cb * (n as i64)) / cbands
300 ccx.hi = ((cb + 1) * (n as i64)) / cbands
301 nx_pool_submit(pool, _nx_mmt_tile_task, ccx as i64)
302 cb = cb + 1
303 }
304 let cwv: i64 = nx_pool_wait(pool, cdone + cbands)
305 sys_munmap(cctxs, cbands * NX_MMT_CTX_BYTES)
306 if cwv != 0 { return NX_F32_MM_ERR_BAD_DIM }
307 return NX_F32_MM_OK
308 }
309
310 let total: i64 = (m as i64) * (n as i64)
311 var bands: i64 = pool.n_workers
312 if bands > total { bands = total }
313 if bands < 1 { bands = 1 }
314
315 let ctxs: *u8 = sys_mmap(bands * NX_MMT_CTX_BYTES)
316 let done_before: i64 = nx_pool_n_completed(pool)
317 var b: i64 = 0
318 while b < bands {
319 let cx: *NxMmtCtx = ((ctxs as i64) + b * NX_MMT_CTX_BYTES) as *NxMmtCtx
320 cx.a_ptr = A as i64
321 cx.b_ptr = B as i64
322 cx.c_ptr = C as i64
323 cx.k_dim = k
324 cx.n_dim = n
325 cx.lo = (b * total) / bands
326 cx.hi = ((b + 1) * total) / bands
327 nx_pool_submit(pool, _nx_mmt_task, cx as i64)
328 b = b + 1
329 }
330 let wv: i64 = nx_pool_wait(pool, done_before + bands)
331 sys_munmap(ctxs, bands * NX_MMT_CTX_BYTES)
332 if wv != 0 { return NX_F32_MM_ERR_BAD_DIM }
333 return NX_F32_MM_OK
334}