nx_matmul_t_pool_gate.nx
buildroot/runtime/nx_matmul_t_pool_gate.nx
about
nx_matmul_t_pool_gate.nx -- gate + MEASUREMENT for the pooled
transposed matmul (nx_f32_matmul_t_pool) that threads lm_head.
lm_head is the forward's biggest single matmul: at decode it is
nx_f32_matmul_t(A[1,896], B[vocab,896], C[1,vocab], 1, 896, 151936)
= 136M MACs, and it ran the SCALAR SERIAL mmt_range every token
(fork was disabled in the organ). This gate runs that EXACT shape
serial vs pooled so the printed serial_us IS lm_head's real
per-token cost, and proves the threaded result bit-identical.
BIT-EXACT ON ANY DATA: flat-range banding computes each C[idx]
wholly inside one band with the identical accumulation order as
serial -- so pool == serial is bit-exact regardless of magnitude.
Values are still kept small-int-exact (finite, no NaN) so the
bit-compare is clean.
Checks (6):
1 small prime shape (m=3, n=37): pool == serial bit-exact
(band-boundary correctness -- 37 is not a multiple of 16)
2 serial oracle OK at the real lm_head decode shape
3 pooled result bit-exact vs serial at the real shape
4 auto-worker pool (nx_pool_new(0)) == serial
5 pool reuse (second call, new A) == fresh serial
6 SPEEDUP: pooled >= floor x serial at the real shape
(serial_us printed = lm_head's real per-token cost)
lineage_id: matmul_t_pool_gate_v1
dependencies 3 imports · 0 importers
imports: nx_f32_matmul_t.nxnx_f32_cvt.nxnx_fmt.nx
imported by: nobody (leaf or entry point)
call flow from main pre-order; caps 40 nodes / depth 6 declared; ↻ = already shown
structs
| none |
consts
| 33 | const TG_K: i64 = 896 // Qwen2.5-0.5B hidden |
| 34 | const TG_VOCAB: i64 = 151936 // real vocab |
| 35 | const TG_FLOOR_X100: i64 = 300 |
functions
| 40 | func _tbl() -> *i64 |
| 56 | func t_fill(p: *i64, count: i64, seed: i64) -> i64 |
| 66 | func t_poison(p: *i64, count: i64) -> i64 called by 1: main |
| 72 | func t_same(a: *i64, b: *i64, count: i64) -> i64 called by 1: main |
| 78 | func t_nl() -> i64 { fmt_puts("\n" as *u8); return 0 } |
| 80 | func main() -> i64 |