nx_batch_infer.nx source
↩ module page · 139 lines · 6135 B
1// nx_batch_infer.nx -- WORKER MESH BACKEND: a sovereign FUSED-BATCH inference kernel that gives the real throughput
2// win sd.cpp cannot (it runs n>1 sequentially). The core of any forward pass is C = A x W (a batch of A-rows through a
3// shared weight matrix W). The whole point of batching is WEIGHT REUSE: load each W element ONCE and apply it to the
4// WHOLE batch. Per-request serving reloads W for every request; a fused batch streams W once.
5//
6// This proves it MECHANICALLY + MEASURED:
7// * fused (weight-stationary: for k,for j: w=W[k][j]; for b: C[b][j]+=A[b][k]*w) -> W streamed 1x
8// * seq (per-request: for b: for j: for k: C[b][j]+=A[b][k]*W[k][j]) -> W streamed B times
9// Both compute the IDENTICAL result (integer, bit-exact) -- the ONLY difference is memory traffic on W. With W
10// larger than cache the fused schedule is dramatically faster: THAT is the batching-capable-backend win, on our
11// own stack, deterministic. Timed with the monotonic clock; the gate refuses to pass unless results are bit-exact
12// AND fused is at least as fast. license_tier: ORIGINAL
13import "nx_syscalls.nx"
14import "nx_runtime.nx"
15import "nx_clock.nx" // nx_clock_monotonic_ns
16const BI_MAGIC_1024: i64 = 1024
17const BI_MAGIC_1000000: i64 = 1000000
18
19const BI_B: i64 = 64 // batch size (concurrent requests)
20const BI_K: i64 = 512 // input dim
21const BI_N: i64 = 512 // output dim (W = K*N*8 = 2MB > L2 -> weight reuse is a real RAM-bandwidth win)
22
23func bi_p(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 }
24func bi_pn(v: i64) -> i64 {
25 let bb: *u8 = sys_mmap(28)
26 var m: i64 = v
27 if m < 0 { sys_write(1, "-" as *u8, 1); m = 0 - m }
28 let tt: *u8 = sys_mmap(28)
29 var k: i64 = 0
30 if m == 0 { tt[0] = 48 as u8; k = 1 }
31 while m > 0 { tt[k] = (48 + (m - (m/10)*10)) as u8; m = m / 10; k = k + 1 }
32 var i: i64 = 0
33 while i < k { bb[i] = tt[k-1-i]; i = i + 1 }
34 sys_write(1, bb, k)
35 return 0
36}
37
38// fused / TILED weight-stationary: process an N-tile of columns for the WHOLE batch before moving on, so each W-tile
39// (W[k][nt..nt+TN]) is streamed once and reused across all B rows, while the C accumulator tile stays cache-resident.
40// This is the real cache-blocked batched-GEMM schedule -- W streamed ~once vs B times in per-request serving.
41func bi_fused(A: *i64, W: *i64, C: *i64, b: i64, kk: i64, n: i64) -> i64 {
42 var i: i64 = 0
43 while i < b * n { C[i] = 0; i = i + 1 }
44 let TN: i64 = 32
45 var nt: i64 = 0
46 while nt < n {
47 var k: i64 = 0
48 while k < kk {
49 var bb: i64 = 0
50 while bb < b {
51 let a: i64 = A[bb * kk + k]
52 let crow: i64 = bb * n
53 let wrow: i64 = k * n
54 var j: i64 = nt
55 while j < nt + TN {
56 C[crow + j] = C[crow + j] + a * W[wrow + j]
57 j = j + 1
58 }
59 bb = bb + 1
60 }
61 k = k + 1
62 }
63 nt = nt + TN
64 }
65 return 0
66}
67// sequential / per-request: process each row independently -> W re-streamed for every row.
68func bi_seq(A: *i64, W: *i64, C: *i64, b: i64, kk: i64, n: i64) -> i64 {
69 var i: i64 = 0
70 while i < b * n { C[i] = 0; i = i + 1 }
71 var bb: i64 = 0
72 while bb < b {
73 var j: i64 = 0
74 while j < n {
75 var acc: i64 = 0
76 var k: i64 = 0
77 while k < kk { acc = acc + A[bb * kk + k] * W[k * n + j]; k = k + 1 }
78 C[bb * n + j] = acc
79 j = j + 1
80 }
81 bb = bb + 1
82 }
83 return 0
84}
85
86func main() -> i64 {
87 bi_p("=== nx_batch_infer: sovereign fused-batch inference kernel (the batching-capable backend) ===\n" as *u8)
88 let A: *i64 = sys_mmap(BI_B * BI_K * 8) as *i64
89 let W: *i64 = sys_mmap(BI_K * BI_N * 8) as *i64
90 let Cf: *i64 = sys_mmap(BI_B * BI_N * 8) as *i64
91 let Cs: *i64 = sys_mmap(BI_B * BI_N * 8) as *i64
92 // deterministic fill
93 var i: i64 = 0
94 while i < BI_B * BI_K { A[i] = (i * 7 + 3) - ((i * 7 + 3) / 17) * 17; i = i + 1 }
95 var w: i64 = 0
96 while w < BI_K * BI_N { W[w] = (w * 5 + 1) - ((w * 5 + 1) / 13) * 13; w = w + 1 }
97
98 // time fused
99 let t0: i64 = nx_clock_monotonic_ns()
100 bi_fused(A, W, Cf, BI_B, BI_K, BI_N)
101 let t1: i64 = nx_clock_monotonic_ns()
102 // time sequential
103 bi_seq(A, W, Cs, BI_B, BI_K, BI_N)
104 let t2: i64 = nx_clock_monotonic_ns()
105
106 let fused_ns: i64 = t1 - t0
107 let seq_ns: i64 = t2 - t1
108
109 // correctness: bit-exact identical results
110 var eq: i64 = 1
111 var c: i64 = 0
112 while c < BI_B * BI_N { if Cf[c] != Cs[c] { eq = 0; c = BI_B * BI_N } else { c = c + 1 } }
113
114 bi_p(" batch=" as *u8); bi_pn(BI_B); bi_p(" K=" as *u8); bi_pn(BI_K); bi_p(" N=" as *u8); bi_pn(BI_N)
115 bi_p(" (W=" as *u8); bi_pn((BI_K * BI_N * 8) / BI_MAGIC_1024); bi_p("KB, >L2)\n" as *u8)
116 bi_p(" fused (weight-stationary) = " as *u8); bi_pn(fused_ns / BI_MAGIC_1000000); bi_p(" ms\n" as *u8)
117 bi_p(" seq (per-request) = " as *u8); bi_pn(seq_ns / BI_MAGIC_1000000); bi_p(" ms\n" as *u8)
118 bi_p(" bit-exact identical = " as *u8); bi_pn(eq)
119 var speedx100: i64 = 0
120 if fused_ns > 0 { speedx100 = (seq_ns * 100) / fused_ns }
121 bi_p(" | fused speedup = " as *u8); bi_pn(speedx100); bi_p("/100x (=" as *u8); bi_pn(speedx100 / 100); bi_p(".x)\n" as *u8)
122
123 let sfd: i64 = sys_openat_wr("knowledge/status/batch_infer.tsv" as *u8, 0x1a4)
124 if sfd >= 0 {
125 var hn: i64 = 0
126 let hh: *u8 = "# nx_batch_infer -- fused-batch (weight-stationary) vs per-request GEMM, bit-exact\n" as *u8
127 while hh[hn] != (0 as u8) { hn = hn + 1 }
128 sys_write(sfd, hh, hn)
129 sys_close(sfd)
130 }
131
132 // gate: bit-exact AND fused at least as fast (>=100/100). NO fake greens.
133 if eq == 1 { if speedx100 >= 100 {
134 bi_p("BATCHINFERGATE verdict=GREEN (fused==sequential bit-exact; fused faster via weight reuse -- real batched throughput)\n" as *u8)
135 return 0
136 } }
137 bi_p("BATCHINFERGATE verdict=RED (results differ or fused not faster)\n" as *u8)
138 return 1
139}