nx_f32_matmul_t.nx source
↩ module page · 242 lines · 11486 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// Serial column-tiled matmul_t. Same signature and same results as nx_f32_matmul_t, fewer weight reads.
123// ADDITIVE (rule 19): nothing is rewired by this file -- callers opt in, and the dispatcher change is a
124// separate deliberate act once the KAT has proven equality on the shapes that matter.
125func nx_f32_matmul_t_tiled(A: *i64, B: *i64, C: *i64,
126 m: nx_int, k: nx_int, n: nx_int) -> nx_int {
127 if m <= 0 { return NX_F32_MM_ERR_BAD_DIM }
128 if k <= 0 { return NX_F32_MM_ERR_BAD_DIM }
129 if n <= 0 { return NX_F32_MM_ERR_BAD_DIM }
130 mmt_tile_cols(A, B, C, m as i64, k as i64, n as i64, 0, n as i64)
131 return NX_F32_MM_OK
132}
133
134func nx_f32_matmul_t(A: *i64, B: *i64, C: *i64,
135 m: nx_int, k: nx_int, n: nx_int) -> nx_int {
136 if m <= 0 { return NX_F32_MM_ERR_BAD_DIM }
137 if k <= 0 { return NX_F32_MM_ERR_BAD_DIM }
138 if n <= 0 { return NX_F32_MM_ERR_BAD_DIM }
139
140 let total: i64 = (m as i64) * (n as i64)
141 // small -> serial (fork overhead not worth it)
142 if total * (k as i64) < MMT_PAR_MIN_MAC {
143 mmt_range(A, B, C, k, n, 0, total)
144 return NX_F32_MM_OK
145 }
146
147 // big -> SOVEREIGN fork-parallel: workers write flat-ranges into a shared temp, parent copies into C
148 let Ctmp: *i64 = sys_mmap_shared(total * 8) as *i64
149 var w: i64 = 0
150 while w < MMT_NW {
151 let pid: i64 = sys_fork()
152 if pid == 0 {
153 let lo: i64 = (w * total) / MMT_NW
154 let hi: i64 = ((w + 1) * total) / MMT_NW
155 mmt_range(A, B, Ctmp, k, n, lo, hi)
156 sys_exit(0)
157 }
158 w = w + 1
159 }
160 let stp: *i64 = sys_mmap(8) as *i64
161 var d: i64 = 0
162 while d < MMT_NW { sys_wait4(0 - 1, stp, 0); d = d + 1 }
163 var c: i64 = 0
164 while c < total { C[c] = Ctmp[c]; c = c + 1 }
165 return NX_F32_MM_OK
166}
167
168// Pooled matmul_t: flat-range bands on a caller-owned pool.
169// Bit-exact vs mmt_range (serial) on any data.
170func nx_f32_matmul_t_pool(pool: *NxThreadPool, A: *i64, B: *i64, C: *i64,
171 m: nx_int, k: nx_int, n: nx_int) -> nx_int {
172 if m <= 0 { return NX_F32_MM_ERR_BAD_DIM }
173 if k <= 0 { return NX_F32_MM_ERR_BAD_DIM }
174 if n <= 0 { return NX_F32_MM_ERR_BAD_DIM }
175
176 // ---- m>1 => COLUMN-TILED, WEIGHT-STATIONARY (2026-08-01) ---------------------------------------
177 // THE LAST HOP. mmt_tile_cols (the batching kernel) was built 07-31 with the measurement that proves
178 // the need -- nx_batchscale_kat at the real Qwen inner dim k=896 recorded us_per_row 4705/4056/4295/4191
179 // for m=1/2/4/8: FLAT, a 1.12x batch speedup where ~4x is expected -- but NOTHING CALLED IT. The pooled
180 // path below bands over FLAT CELLS and runs mmt_range, which walks output cells and therefore re-streams
181 // each weight column once PER ROW, so weight traffic is m*k*n instead of k*n. The forward only ever uses
182 // the pooled entry point, so the batching kernel was unreachable from production: a kernel with no
183 // caller is not an optimisation, it is a file. That is why prefill never amortised and why the sev-9
184 // rows measured per-token cost flat across a 15x range of m (606/673/567/593 ms/tok at m=8/32/64/125).
185 // Routing m>1 here converts prefill from memory-bound re-streaming to the compute-bound regime batching
186 // is supposed to reach; m=1 decode is left EXACTLY as it was, because at m=1 there is nothing to amortise
187 // (each weight is read once either way) and that path is already measured and proven.
188 // BIT-EXACT, not approximately equal: mmt_tile_cols accumulates kk ascending from f32 zero for every
189 // (i,j), the identical order mmt_range uses -- only the interleaving across i changes. So this is a
190 // memory-traffic change and NOT a numerics change, and nx_matmul_tile_kat asserts that equality directly.
191 // u00e2u02dcu2026A KERNEL WITH NO CALLER IS NOT AN OPTIMISATION, IT IS A FILE.
192 if m > 1 {
193 var cbands: i64 = pool.n_workers
194 if cbands > (n as i64) { cbands = n as i64 }
195 if cbands < 1 { cbands = 1 }
196 let cctxs: *u8 = sys_mmap(cbands * NX_MMT_CTX_BYTES)
197 let cdone: i64 = nx_pool_n_completed(pool)
198 var cb: i64 = 0
199 while cb < cbands {
200 let ccx: *NxMmtCtx = ((cctxs as i64) + cb * NX_MMT_CTX_BYTES) as *NxMmtCtx
201 ccx.a_ptr = A as i64
202 ccx.b_ptr = B as i64
203 ccx.c_ptr = C as i64
204 ccx.k_dim = k
205 ccx.n_dim = n
206 ccx.m_dim = m as i64
207 ccx.lo = (cb * (n as i64)) / cbands
208 ccx.hi = ((cb + 1) * (n as i64)) / cbands
209 nx_pool_submit(pool, _nx_mmt_tile_task, ccx as i64)
210 cb = cb + 1
211 }
212 let cwv: i64 = nx_pool_wait(pool, cdone + cbands)
213 sys_munmap(cctxs, cbands * NX_MMT_CTX_BYTES)
214 if cwv != 0 { return NX_F32_MM_ERR_BAD_DIM }
215 return NX_F32_MM_OK
216 }
217
218 let total: i64 = (m as i64) * (n as i64)
219 var bands: i64 = pool.n_workers
220 if bands > total { bands = total }
221 if bands < 1 { bands = 1 }
222
223 let ctxs: *u8 = sys_mmap(bands * NX_MMT_CTX_BYTES)
224 let done_before: i64 = nx_pool_n_completed(pool)
225 var b: i64 = 0
226 while b < bands {
227 let cx: *NxMmtCtx = ((ctxs as i64) + b * NX_MMT_CTX_BYTES) as *NxMmtCtx
228 cx.a_ptr = A as i64
229 cx.b_ptr = B as i64
230 cx.c_ptr = C as i64
231 cx.k_dim = k
232 cx.n_dim = n
233 cx.lo = (b * total) / bands
234 cx.hi = ((b + 1) * total) / bands
235 nx_pool_submit(pool, _nx_mmt_task, cx as i64)
236 b = b + 1
237 }
238 let wv: i64 = nx_pool_wait(pool, done_before + bands)
239 sys_munmap(ctxs, bands * NX_MMT_CTX_BYTES)
240 if wv != 0 { return NX_F32_MM_ERR_BAD_DIM }
241 return NX_F32_MM_OK
242}