nx_q8_coldwarm.nx source
↩ module page · 96 lines · 3696 B
1// nx_q8_coldwarm.nx -- decides the post-futex decode-matmul direction:
2// COMPUTE-bound (kernel work: AVX2/multi-acc pays) vs COLD-STREAMING-bound
3// (fewer bytes / batch tokens pays). Measures the SAME production Q8_0
4// matmul two ways at FULL-MODEL working-set scale (~209MB, like the real
5// forward's 24 distinct layers read once/token):
6// WARM = one weight buffer, REPS reps (cache-resident after pass 1).
7// COLD = NB DISTINCT buffers (aggregate >> LLC), one matmul each, never
8// reused (the forward's true access pattern).
9// COLD/WARM ~1 -> compute-bound (kernel is the lever).
10// COLD/WARM >>1 -> cold DRAM streaming (fewer bytes / batch is the lever).
11// license_tier: ORIGINAL expect_exit: 0
12import "nx_syscalls.nx"
13import "nx_tier.nx"
14import "nx_le.nx"
15import "nx_f32.nx"
16import "nx_f32_cvt.nx"
17import "nx_thread_pool.nx"
18import "nx_f32_lazy_weight.nx"
19import "nx_fmt.nx"
20
21const MK: i64 = 896
22const NG: i64 = 4864 // W_gate width; one buf = 896*4864*34/32 = 4.63MB
23const Q8B: i64 = 34
24const Q8V: i64 = 32
25const NB: i64 = 48 // 48 * 4.63MB = 222MB ~ full model, >> LLC
26
27func cw_nl() -> i64 { fmt_puts("\n" as *u8); return 0 }
28func cw_lcg(s: i64) -> i64 { var v: i64 = s * 1103515245 + 12345; v = v & 2147483647; return v }
29func cw_weight(seed: i64) -> *NxF32LazyWeight {
30 let bpr: i64 = (MK / Q8V) * Q8B
31 let w: *u8 = sys_mmap(NG * bpr)
32 var s: i64 = seed
33 var r: i64 = 0
34 while r < NG {
35 var b: i64 = 0
36 while b < MK / Q8V {
37 let off: i64 = r * bpr + b * Q8B
38 w[off + 1] = 0x2C as u8
39 var q: i64 = 0
40 while q < 32 { s = cw_lcg(s); w[off + 2 + q] = ((s % 17) - 8) as u8; q = q + 1 }
41 b = b + 1
42 }
43 r = r + 1
44 }
45 return nx_f32_lazy_weight_new_q8_0(w, 0, NG, MK)
46}
47func cw_rep(tag: *u8, us: i64, calls: i64) -> i64 {
48 var u: i64 = us
49 if u < 1 { u = 1 }
50 let flop: i64 = 2 * MK * NG * calls
51 let bytes: i64 = (NG * (MK / Q8V) * Q8B) * calls
52 fmt_puts(tag)
53 fmt_puts(" per_call_us="); fmt_putn(us / calls)
54 fmt_puts(" MFLOPs="); fmt_putn(flop / u)
55 fmt_puts(" MBps="); fmt_putn(bytes / u)
56 cw_nl()
57 return 0
58}
59
60func main() -> i64 {
61 let A: *i64 = sys_mmap(MK * 8) as *i64
62 var s: i64 = 4242
63 var i: i64 = 0
64 while i < MK { s = cw_lcg(s); A[i] = nx_i32_to_f32((s % 9) - 4); i = i + 1 }
65 let C: *i64 = sys_mmap(NG * 8) as *i64
66 nx_lw_shared_pool()
67
68 // allocate NB distinct weights up front (touch each once so mmap is
69 // committed; they then get EVICTED as we allocate the rest -> cold).
70 let bufs: *i64 = sys_mmap(NB * 8) as *i64
71 var b0: i64 = 0
72 while b0 < NB { bufs[b0] = cw_weight(700 + b0 * 13) as i64; b0 = b0 + 1 }
73
74 // COLD: one matmul per DISTINCT buffer (each ~4.6MB, aggregate 222MB
75 // >> LLC -> every touch streams from DRAM). This is the forward.
76 let tc: i64 = sys_now_us()
77 var rc: i64 = 0
78 while rc < NB { nx_f32_lazy_matmul(A, bufs[rc] as *NxF32LazyWeight, C, 1, MK, NG); rc = rc + 1 }
79 let usc: i64 = sys_now_us() - tc
80 cw_rep("COLD(222MB sweep)" as *u8, usc, NB)
81
82 // WARM: same COUNT of matmuls but ONE resident buffer (cache-hot).
83 let lw: *NxF32LazyWeight = bufs[0] as *NxF32LazyWeight
84 nx_f32_lazy_matmul(A, lw, C, 1, MK, NG) // prime
85 let tw: i64 = sys_now_us()
86 var rw: i64 = 0
87 while rw < NB { nx_f32_lazy_matmul(A, lw, C, 1, MK, NG); rw = rw + 1 }
88 let usw: i64 = sys_now_us() - tw
89 cw_rep("WARM(1 resident) " as *u8, usw, NB)
90
91 var uw: i64 = usw
92 if uw < 1 { uw = 1 }
93 fmt_puts("COLD_over_WARM_x100="); fmt_putn(usc * 100 / uw); cw_nl()
94 fmt_puts("Q8_COLDWARM DONE"); cw_nl()
95 return 0
96}