nx_llm_forward_profile.nx source
↩ module page · 105 lines · 5070 B
1// nx_llm_forward_profile.nx -- clean per-decode-forward timing on the
2// REAL model, to find where a decode token's time actually goes.
3//
4// The 34-token bench is too noisy on this host (WSL bounces + multi-
5// session load) to attribute sub-20% changes. This times ONE
6// forward_v4 decode call at a time (no sampler/detok/penalty/loop
7// overhead), after a warmup prefill so the dequant cache + lm_head
8// pool are hot. Combined with the isolated lm_head number
9// (nx_matmul_t_pool_gate: ~143ms pooled at m=1,n=151936), it splits
10// the decode forward into lm_head vs the 24-layer block stack.
11//
12// Prints each step's us (watch for growth = KV-cache-length effect)
13// and the average. No vocab/bpe load (forward needs neither).
14//
15// expect_exit: 0
16
17import "nx_syscalls.nx"
18import "nx_itoa_lib.nx" // shared MSB-first emitter (zero-alloc)
19import "nx_tier.nx"
20import "nx_gguf.nx"
21import "nx_gguf_load.nx"
22import "nx_gguf_meta.nx"
23import "nx_f32.nx"
24import "nx_f32_kv_cache.nx"
25import "nx_f32_lazy_weight.nx"
26import "nx_f32_llama_block.nx"
27import "nx_f32_llama_block_v4.nx"
28import "nx_f32_llama_stack_v4.nx"
29import "nx_f32_llama_layer_lazy_load.nx"
30import "nx_f32_llm.nx"
31import "nx_f32_llm_v4.nx"
32import "nx_f32_llm_read_dims.nx"
33
34func pf_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
35// MIGRATED to the shared emitter (debt 1785563586). The old body mmapped a scratch buffer
36// per call and never freed it. At PAGE granularity that is 4096B leaked PER CALL -- the
37// defect that took 28.5GB of a 36GB host in nx_ts_lumadiff (2MB input, ~3.66M calls).
38// nxi_* is MSB-first, allocates NOTHING, and emits identical bytes including the sign.
39func pf_num(v: i64) -> i64 { nxi_out(v); return 0 }
40func pf_nl() -> i64 { pf_puts("\n" as *u8); return 0 }
41
42func main(argc: i64, argv: *i64) -> i64 {
43 var path: *u8 = "/volume1/homes/elderwesto/nx_bench/model/nx_real_model.gguf" as *u8
44 if argc > 1 { path = argv[1] as *u8 }
45 let len_out: *i64 = sys_mmap(8) as *i64
46 let buf: *u8 = sys_read_file(path, len_out)
47 if buf == (0 as *u8) { pf_puts("no model\n" as *u8); return 10 }
48 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
49 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { return 20 }
50 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
51 let out_err: *i64 = sys_mmap(8) as *i64
52 if nx_f32_llm_read_dims_from_gguf(buf, len_out[0], hdr, model, out_err) != NX_FLD_OK { return 30 }
53 if nx_f32_llm_load_weights_v4_from_gguf(buf, hdr, model, out_err) != NX_FLV4_OK { return 40 }
54
55 let eps: i64 = 0x358637BD
56 let attn_scale: i64 = 0x3E000000
57 let rope_base: i64 = 0x415D0EAB
58 let V: nx_int = model.vocab_size
59
60 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(model.n_layers, model.n_kv_heads, 256, model.head_dim)
61 let logits: *i64 = sys_mmap(V * 8) as *i64
62
63 // ---- Warmup prefill (5 tokens): fills dequant cache + warms the pool ----
64 let toks: *i64 = sys_mmap(16 * 8) as *i64
65 toks[0]=785; toks[1]=6722; toks[2]=315; toks[3]=9625; toks[4]=374
66 pf_puts("prefill (warms dequant cache + lm_head pool)...\n" as *u8)
67 var p: nx_int = 0
68 while p < 5 {
69 let one: *i64 = (((toks as i64) + p * 8)) as *i64
70 if nx_f32_llm_forward_v4(model, one, 1, cache, eps, attn_scale, rope_base, 1, logits) != NX_FLV4_OK { return 60 }
71 p = p + 1
72 }
73 pf_puts("packed weight cache used bytes="); pf_num(nx_lw_cache_used()); pf_nl()
74
75 // ---- Timed decode: one forward_v4 at a time, fixed feed token ----
76 let feed: *i64 = sys_mmap(8) as *i64
77 feed[0] = 785
78 let NS: nx_int = 6
79 nx_blk4_prof_enable(1)
80 nx_blk4_prof_reset()
81 let tstart: i64 = sys_now_us()
82 var s: nx_int = 0
83 while s < NS {
84 let t0: i64 = sys_now_us()
85 if nx_f32_llm_forward_v4(model, feed, 1, cache, eps, attn_scale, rope_base, 1, logits) != NX_FLV4_OK { return 61 }
86 let dt: i64 = sys_now_us() - t0
87 pf_puts(" decode step "); pf_num(s as i64); pf_puts(" us="); pf_num(dt); pf_nl()
88 s = s + 1
89 }
90 let ttot: i64 = sys_now_us() - tstart
91 nx_blk4_prof_enable(0)
92 pf_puts("avg_us_per_decode_forward="); pf_num(ttot / (NS as i64)); pf_nl()
93
94 // ---- Block phase breakdown (summed over NS decodes x 24 layers) ----
95 pf_puts("--- block phase totals over "); pf_num(NS as i64); pf_puts(" decodes (us; /"); pf_num(NS as i64); pf_puts(" = per-token) ---\n" as *u8)
96 pf_puts(" ALLOC(13 mmaps)="); pf_num(nx_blk4_prof_read(0)); pf_nl()
97 pf_puts(" RMSNORM="); pf_num(nx_blk4_prof_read(1)); pf_nl()
98 pf_puts(" MATMUL(7/layer)="); pf_num(nx_blk4_prof_read(2)); pf_nl()
99 pf_puts(" ROPE+BIAS="); pf_num(nx_blk4_prof_read(3)); pf_nl()
100 pf_puts(" ATTN(qk/sm/av)="); pf_num(nx_blk4_prof_read(4)); pf_nl()
101 pf_puts(" ACT(silu/mul)="); pf_num(nx_blk4_prof_read(5)); pf_nl()
102 pf_puts(" RESID="); pf_num(nx_blk4_prof_read(6)); pf_nl()
103 pf_puts("(lm_head is OUTSIDE the block, ~143000us/token pooled)\n" as *u8)
104 return 0
105}