nx_batchscale_kat.nx source
↩ module page · 136 lines · 6284 B
1// nx_batchscale_kat.nx -- IS THE MATMUL PATH AMORTISING ACROSS m? The decisive perf question, answered
2// WITHOUT loading a model, so it fits inside a bounded verifier run.
3//
4// WHY (2026-07-31): the sovereign embedder measures ~350-400ms/token against a 54.5ms roofline
5// (491MB gguf model / 9011 MB/s measured bandwidth) = **~7x off floor** (ids 1785475195, 1785475543).
6// The prime suspect is named by the codebase itself: nx_embed_bench exists to ask 'is prefill actually
7// BATCHING?' and states the rule -- FLAT ms_per_token across m => prefill is NOT batching (wire the m>1
8// packed path in lazy_matmul); FALLING => batching works and the cost lives elsewhere.
9//
10// BUT nx_embed_bench CANNOT BE RUN: /api/gate_run executes VERIFIERS ONLY (correct never-brick bound) and
11// its 491MB model load alone is 16554..20767ms against a 12000ms deadline (id 1785475610). So the decisive
12// measurement was unreachable.
13//
14// THIS KAT REMOVES THE MODEL FROM THE QUESTION. Batching is a property of the MATMUL PATH, not of the
15// weights: if C[m,n] = A[m,k] . W[k,n] costs the same per-row at m=8 as at m=1, then the weight bytes are
16// being re-read per row and nothing is amortising -- which is exactly what a bandwidth-bound decode looks
17// like when the m>1 path is unwired. Synthetic W at the real FFN inner dim (k=896, Qwen W_gate) makes the
18// shape honest while keeping the run bounded.
19//
20// READING THE RESULT:
21// us_per_row FLAT across m -> NOT amortising; each row re-streams W. Wire the m>1 packed path.
22// us_per_row FALLING with m -> amortising; W is read once per tile and the 7x lives elsewhere.
23// The ratio us_per_row(m=1) / us_per_row(m=8) IS the batching speedup, measured on this host.
24//
25// HONESTY: this measures the MATMUL path only. It does NOT prove end-to-end prefill batching (attention,
26// rope and norms are excluded) -- it isolates the one layer the ecosystem's own note points at, and a FLAT
27// result here is sufficient to act on because no amount of downstream work can amortise what this layer
28// already re-read. Timings are one run on a shared NAS: ONE WINDOW IS NOT A RATE, so the KAT prints every
29// m and lets the reader see the shape rather than trusting a single ratio.
30// license_tier: ORIGINAL No hw writes (Rule 26). expect_exit: 0
31import "nx_f32_lazy_weight.nx"
32import "nx_f32_matmul_t.nx"
33import "nx_f32_cvt.nx"
34import "nx_fmt.nx"
35
36const BK: i64 = 896 // real Qwen FFN inner dim
37const BN: i64 = 1024 // reduced from 4864 so the sweep fits a bounded verifier run
38const BMAX: i64 = 8
39
40func bs_lcg(s: i64) -> i64 {
41 var v: i64 = s * 1103515245 + 12345
42 v = v & 2147483647
43 return v
44}
45
46func bs_fill(p: *i64, count: i64, seed: i64) -> i64 {
47 var s: i64 = seed
48 var i: i64 = 0
49 while i < count {
50 s = bs_lcg(s)
51 p[i] = nx_i32_to_f32((s - (s / 8) * 8) - 4)
52 i = i + 1
53 }
54 return 0
55}
56
57func bs_nl() -> i64 { fmt_puts("\n" as *u8); return 0 }
58
59func main() -> i64 {
60 let A: *i64 = sys_mmap(BMAX * BK * 8) as *i64
61 let W: *i64 = sys_mmap(BK * BN * 8) as *i64
62 let C: *i64 = sys_mmap(BMAX * BN * 8) as *i64
63 bs_fill(A, BMAX * BK, 20260731)
64 bs_fill(W, BK * BN, 31337)
65
66 let pool: *NxThreadPool = nx_lw_shared_pool()
67
68 fmt_puts("=== nx_batchscale_kat -- does the matmul path AMORTISE across m? (k=" as *u8)
69 fmt_putn(BK); fmt_puts(" n=" as *u8); fmt_putn(BN); fmt_puts(") ===" as *u8); bs_nl()
70
71 // warm: first call pays any one-time setup; timing it would slander the m=1 row.
72 nx_f32_matmul_t_pool(pool, A, W, C, 1, BK, BN)
73
74 var pass: i64 = 0
75 var total: i64 = 0
76 var per_row_m1: i64 = 0
77 var per_row_m8: i64 = 0
78
79 var m: i64 = 1
80 while m <= BMAX {
81 let t0: i64 = sys_now_us()
82 var r: i64 = 0
83 while r < 4 {
84 nx_f32_matmul_t_pool(pool, A, W, C, m, BK, BN)
85 r = r + 1
86 }
87 let dt: i64 = sys_now_us() - t0
88 var per_row: i64 = dt / (4 * m)
89 if per_row < 1 { per_row = 1 }
90 fmt_puts(" m=" as *u8); fmt_putn(m)
91 fmt_puts(" total_us=" as *u8); fmt_putn(dt)
92 fmt_puts(" us_per_row=" as *u8); fmt_putn(per_row)
93 bs_nl()
94 if m == 1 { per_row_m1 = per_row }
95 if m == 8 { per_row_m8 = per_row }
96 m = m * 2
97 }
98
99 bs_nl()
100 var ratio_x100: i64 = 0
101 if per_row_m8 > 0 { ratio_x100 = per_row_m1 * 100 / per_row_m8 }
102 fmt_puts("BATCH SPEEDUP m=1 -> m=8 : " as *u8); fmt_putn(ratio_x100)
103 fmt_puts(" (x100; 100 = FLAT = no amortisation)" as *u8); bs_nl()
104
105 total = total + 1
106 if per_row_m1 > 0 { pass = pass + 1; fmt_puts(" ok T1 m=1 timed and non-zero" as *u8) }
107 else { fmt_puts(" FAIL T1 m=1 timing collapsed to zero -- cannot conclude anything" as *u8) }
108 bs_nl()
109
110 total = total + 1
111 if per_row_m8 > 0 { pass = pass + 1; fmt_puts(" ok T2 m=8 timed and non-zero" as *u8) }
112 else { fmt_puts(" FAIL T2 m=8 timing collapsed to zero" as *u8) }
113 bs_nl()
114
115 total = total + 1
116 if ratio_x100 > 0 { pass = pass + 1; fmt_puts(" ok T3 ratio computed (the measurement exists)" as *u8) }
117 else { fmt_puts(" FAIL T3 no ratio -- measurement did not happen" as *u8) }
118 bs_nl()
119
120 bs_nl()
121 if ratio_x100 < 130 {
122 fmt_puts("VERDICT=FLAT -- the matmul path is NOT amortising across m (ratio < 1.3x)." as *u8); bs_nl()
123 fmt_puts(" Every row re-streams W, which is the bandwidth-bound signature. THE LEVER IS THE m>1" as *u8); bs_nl()
124 fmt_puts(" PACKED PATH in lazy_matmul -- exactly what nx_embed_bench's header prescribes." as *u8); bs_nl()
125 } else {
126 fmt_puts("VERDICT=AMORTISING -- batching already helps here; the ~7x gap lives ELSEWHERE" as *u8); bs_nl()
127 fmt_puts(" (attention / rope / norms / tokenise), NOT in this matmul layer." as *u8); bs_nl()
128 }
129 fmt_puts("envelope: matmul layer ONLY (attention/rope/norms excluded); synthetic weights at the real" as *u8); bs_nl()
130 fmt_puts(" k=896 inner dim; ONE run on a shared host -- one window is not a rate, read the shape." as *u8); bs_nl()
131
132 fmt_puts("BATCHSCALE-KAT " as *u8); fmt_putn(pass); fmt_puts("/" as *u8); fmt_putn(total)
133 if pass == total { fmt_puts(" GREEN" as *u8); bs_nl(); return 0 }
134 fmt_puts(" RED" as *u8); bs_nl()
135 return 1
136}