nx_matmul_t_pool_gate.nx source
↩ module page · 170 lines · 6408 B
1// nx_matmul_t_pool_gate.nx -- gate + MEASUREMENT for the pooled
2// transposed matmul (nx_f32_matmul_t_pool) that threads lm_head.
3//
4// lm_head is the forward's biggest single matmul: at decode it is
5// nx_f32_matmul_t(A[1,896], B[vocab,896], C[1,vocab], 1, 896, 151936)
6// = 136M MACs, and it ran the SCALAR SERIAL mmt_range every token
7// (fork was disabled in the organ). This gate runs that EXACT shape
8// serial vs pooled so the printed serial_us IS lm_head's real
9// per-token cost, and proves the threaded result bit-identical.
10//
11// BIT-EXACT ON ANY DATA: flat-range banding computes each C[idx]
12// wholly inside one band with the identical accumulation order as
13// serial -- so pool == serial is bit-exact regardless of magnitude.
14// Values are still kept small-int-exact (finite, no NaN) so the
15// bit-compare is clean.
16//
17// Checks (6):
18// 1 small prime shape (m=3, n=37): pool == serial bit-exact
19// (band-boundary correctness -- 37 is not a multiple of 16)
20// 2 serial oracle OK at the real lm_head decode shape
21// 3 pooled result bit-exact vs serial at the real shape
22// 4 auto-worker pool (nx_pool_new(0)) == serial
23// 5 pool reuse (second call, new A) == fresh serial
24// 6 SPEEDUP: pooled >= floor x serial at the real shape
25// (serial_us printed = lm_head's real per-token cost)
26//
27// lineage_id: matmul_t_pool_gate_v1
28
29import "nx_f32_matmul_t.nx"
30import "nx_f32_cvt.nx"
31import "nx_fmt.nx"
32
33const TG_K: i64 = 896 // Qwen2.5-0.5B hidden
34const TG_VOCAB: i64 = 151936 // real vocab
35const TG_FLOOR_X100: i64 = 300
36
37// 8 exact small-int f32 values (finite; |v| <= 2 so k=896 dot stays
38// far under 2^24 -- exact regime, clean bit-compare).
39static G_TBL: i64
40func _tbl() -> *i64 {
41 if G_TBL == 0 {
42 let t: *i64 = sys_mmap(8 * 8) as *i64
43 t[0] = nx_i32_to_f32(1)
44 t[1] = nx_i32_to_f32(0 - 1)
45 t[2] = nx_i32_to_f32(2)
46 t[3] = nx_i32_to_f32(0)
47 t[4] = nx_i32_to_f32(0 - 2)
48 t[5] = nx_i32_to_f32(1)
49 t[6] = nx_i32_to_f32(0 - 1)
50 t[7] = nx_i32_to_f32(2)
51 G_TBL = t as i64
52 }
53 return G_TBL as *i64
54}
55
56func t_fill(p: *i64, count: i64, seed: i64) -> i64 {
57 let tb: *i64 = _tbl()
58 var i: i64 = 0
59 while i < count {
60 p[i] = tb[(i + seed) & 7]
61 i = i + 1
62 }
63 return 0
64}
65
66func t_poison(p: *i64, count: i64) -> i64 {
67 var i: i64 = 0
68 while i < count { p[i] = 0 - 777777; i = i + 1 }
69 return 0
70}
71
72func t_same(a: *i64, b: *i64, count: i64) -> i64 {
73 var i: i64 = 0
74 while i < count { if a[i] != b[i] { return 0 } i = i + 1 }
75 return 1
76}
77
78func t_nl() -> i64 { fmt_puts("\n" as *u8); return 0 }
79
80func main() -> i64 {
81 var pass: i64 = 0
82
83 // ---- 1: small prime shape, band-boundary correctness ----
84 let SM: i64 = 3
85 let SN: i64 = 37
86 let sa: *i64 = sys_mmap(SM * TG_K * 8) as *i64
87 let sb: *i64 = sys_mmap(SN * TG_K * 8) as *i64
88 let sc: *i64 = sys_mmap(SM * SN * 8) as *i64
89 let sp: *i64 = sys_mmap(SM * SN * 8) as *i64
90 t_fill(sa, SM * TG_K, 1)
91 t_fill(sb, SN * TG_K, 5)
92 nx_f32_matmul_t(sa, sb, sc, SM, TG_K, SN)
93 let poolS: *NxThreadPool = nx_pool_new(4, 0)
94 t_poison(sp, SM * SN)
95 nx_f32_matmul_t_pool(poolS, sa, sb, sp, SM, TG_K, SN)
96 if t_same(sp, sc, SM * SN) != 1 { fmt_puts("MMTG 1 SMALL-BAND FAIL"); t_nl(); return 11 }
97 nx_pool_shutdown(poolS)
98 fmt_puts("MMTG 1 SMALL-BAND EXACT OK"); t_nl()
99 pass = pass + 1
100
101 // ---- real lm_head decode shape ----
102 let A: *i64 = sys_mmap(TG_K * 8) as *i64
103 fmt_puts("allocating lm_head B (~1.09GB)..."); t_nl()
104 let B: *i64 = sys_mmap(TG_VOCAB * TG_K * 8) as *i64
105 let Cs: *i64 = sys_mmap(TG_VOCAB * 8) as *i64
106 let Cp: *i64 = sys_mmap(TG_VOCAB * 8) as *i64
107 t_fill(A, TG_K, 1)
108 t_fill(B, TG_VOCAB * TG_K, 3)
109
110 // ---- 2: serial oracle = lm_head's real per-token cost ----
111 let t0: i64 = sys_now_us()
112 let v2: nx_int = nx_f32_matmul_t(A, B, Cs, 1, TG_K, TG_VOCAB)
113 let serial_us: i64 = sys_now_us() - t0
114 if v2 != NX_F32_MM_OK { fmt_puts("MMTG 2 SERIAL FAIL"); t_nl(); return 12 }
115 fmt_puts("MMTG 2 SERIAL OK lm_head_serial_us="); fmt_putn(serial_us); t_nl()
116 pass = pass + 1
117
118 // ---- 3: pooled (fixed 8) bit-exact ----
119 let pool8: *NxThreadPool = nx_pool_new(8, 0)
120 t_poison(Cp, TG_VOCAB)
121 let v3: nx_int = nx_f32_matmul_t_pool(pool8, A, B, Cp, 1, TG_K, TG_VOCAB)
122 var ok3: i64 = 0
123 if v3 == NX_F32_MM_OK { ok3 = t_same(Cp, Cs, TG_VOCAB) }
124 if ok3 != 1 { fmt_puts("MMTG 3 POOL8-EXACT FAIL"); t_nl(); return 13 }
125 nx_pool_shutdown(pool8)
126 fmt_puts("MMTG 3 POOL8 EXACT OK"); t_nl()
127 pass = pass + 1
128
129 // ---- 4: auto-worker pool + speedup timing ----
130 let poolA: *NxThreadPool = nx_pool_new(0, 0)
131 t_poison(Cp, TG_VOCAB)
132 let t1: i64 = sys_now_us()
133 let v4: nx_int = nx_f32_matmul_t_pool(poolA, A, B, Cp, 1, TG_K, TG_VOCAB)
134 let pool_us: i64 = sys_now_us() - t1
135 var ok4: i64 = 0
136 if v4 == NX_F32_MM_OK { ok4 = t_same(Cp, Cs, TG_VOCAB) }
137 if ok4 != 1 { fmt_puts("MMTG 4 POOL-AUTO FAIL"); t_nl(); return 14 }
138 fmt_puts("MMTG 4 POOL-AUTO EXACT OK workers="); fmt_putn(poolA.n_workers); t_nl()
139 pass = pass + 1
140
141 // ---- 5: pool reuse, new A ----
142 t_fill(A, TG_K, 4)
143 let v5s: nx_int = nx_f32_matmul_t(A, B, Cs, 1, TG_K, TG_VOCAB)
144 if v5s != NX_F32_MM_OK { return 15 }
145 t_poison(Cp, TG_VOCAB)
146 let v5: nx_int = nx_f32_matmul_t_pool(poolA, A, B, Cp, 1, TG_K, TG_VOCAB)
147 var ok5: i64 = 0
148 if v5 == NX_F32_MM_OK { ok5 = t_same(Cp, Cs, TG_VOCAB) }
149 nx_pool_shutdown(poolA)
150 if ok5 != 1 { fmt_puts("MMTG 5 REUSE FAIL"); t_nl(); return 15 }
151 fmt_puts("MMTG 5 REUSE EXACT OK"); t_nl()
152 pass = pass + 1
153
154 // ---- 6: speedup verdict ----
155 var us_s: i64 = serial_us
156 if us_s < 1 { us_s = 1 }
157 var us_p: i64 = pool_us
158 if us_p < 1 { us_p = 1 }
159 let macs: i64 = TG_K * TG_VOCAB
160 fmt_puts("serial_mflops="); fmt_putn(2 * macs / us_s); t_nl()
161 fmt_puts("pool_us="); fmt_putn(us_p); fmt_puts(" mflops="); fmt_putn(2 * macs / us_p); t_nl()
162 let sx100: i64 = us_s * 100 / us_p
163 fmt_puts("lmhead_speedup_x100="); fmt_putn(sx100); t_nl()
164 if sx100 < TG_FLOOR_X100 { fmt_puts("MMTG 6 SPEEDUP FAIL"); t_nl(); return 16 }
165 fmt_puts("MMTG 6 SPEEDUP OK"); t_nl()
166 pass = pass + 1
167
168 fmt_puts("MATMUL_T_POOL_GATE "); fmt_putn(pass); fmt_puts("/6 GREEN"); t_nl()
169 return 0
170}