nx_nofloat_qwen_dqprobe_gate.nx source
↩ module page · 99 lines · 6749 B
1// nx_nofloat_qwen_dqprobe_gate.nx -- MEASUREMENT PROBE (not a milestone gate): isolate how much of a decode step
2// is spent RE-DEQUANTIZING the (unchanging) weights vs actual compute. decode_step_kv calls nf_load_layer_pool
3// every layer every token; the weights never change across steps, so if dequant dominates, a one-time dequant
4// cache is a LOSSLESS win. This prints: (A) one full decode_step_kv [dequant+compute], (B) a bare 24-layer
5// nf_load_layer_pool sweep [dequant only], and the derived compute floor (A-B). No hw writes (Rule 26).
6// expect_exit: 0 license_tier: ORIGINAL
7import "nx_syscalls.nx"
8import "nx_tier.nx"
9import "nx_le.nx"
10import "nx_tensor.nx"
11import "nx_gguf.nx"
12import "nx_gguf_load.nx"
13import "nx_gguf_meta.nx"
14import "nx_nofloat_llm.nx"
15import "nx_nofloat_tok.nx"
16
17func dp_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
18func dp_num(v: i64) -> i64 { let b: *u8=sys_mmap(28); var m: i64=v; if m<0{m=0-m;sys_write(1,"-" as *u8,1)} let t: *u8=sys_mmap(28); var k: i64=0; if m==0{t[0]=48 as u8;k=1} while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} var i: i64=0; while i<k{b[i]=t[k-1-i];i=i+1} sys_write(1,b,k); return 0 }
19func dp_slen(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} return n }
20
21func main() -> i64 {
22 dp_puts("DEQUANT-COST PROBE: is a decode step dominated by re-dequantizing the frozen weights?\n\n" as *u8)
23 let path: *u8 = "/home/elderwesto/nx_stage/nx_real_model.gguf\x00" as *u8
24 let len_out: *i64 = sys_mmap(8) as *i64; len_out[0]=0
25 let buf: *u8 = sys_read_file(path, len_out)
26 if buf == (0 as *u8) { dp_puts("MODEL ABSENT\n" as *u8); return 1 }
27 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
28 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { dp_puts("PARSE FAIL\n" as *u8); return 1 }
29
30 let ne: i64=896; let qd: i64=896; let kvd: i64=128; let fd: i64=4864; let MAXT: i64=16
31 let voff: *i64=sys_mmap(8) as *i64; let vty: *i64=sys_mmap(8) as *i64
32 var mfirst: i64=0; var nm_c: i64=0; var vfirst: i64=0; var vocab: i64=0
33 let km: *u8="tokenizer.ggml.merges\x00" as *u8; let kt: *u8="tokenizer.ggml.tokens\x00" as *u8
34 if nx_gguf_meta_find(buf, len_out[0], hdr, km, dp_slen(km), voff, vty)==NX_GMETA_OK { nm_c=nx_gguf_meta_array_count(buf, voff[0]); mfirst=nx_gguf_meta_array_first_elt_off(buf, voff[0]) }
35 if nx_gguf_meta_find(buf, len_out[0], hdr, kt, dp_slen(kt), voff, vty)==NX_GMETA_OK { vocab=nx_gguf_meta_array_count(buf, voff[0]); vfirst=nx_gguf_meta_array_first_elt_off(buf, voff[0]) }
36 let nt: *u8="token_embd.weight\x00" as *u8; let nn: *u8="output_norm.weight\x00" as *u8
37 let ie: nx_int=nx_gguf_find_tensor(hdr, nt, 17); let inn: nx_int=nx_gguf_find_tensor(hdr, nn, 18)
38 if ie<0 { return 1 } if inn<0 { return 1 }
39 let te: *NxGgufTensorInfo=nx_gguf_tensor_at(hdr, ie); let te_base: i64=hdr.data_off+te.offset; let te_ty: i64=te.ggml_type
40 let gout: *i64=sys_mmap(ne*8) as *i64; load_named_q16(buf, hdr, nn, 18, gout, ne)
41
42 let wb: *i64=sys_mmap(12*8) as *i64
43 wb[0]=sys_mmap(ne*8) as i64; wb[1]=sys_mmap(qd*ne*8) as i64; wb[2]=sys_mmap(kvd*ne*8) as i64; wb[3]=sys_mmap(kvd*ne*8) as i64; wb[4]=sys_mmap(ne*qd*8) as i64
44 wb[5]=sys_mmap(ne*8) as i64; wb[6]=sys_mmap(ne*fd*8) as i64; wb[7]=sys_mmap(ne*fd*8) as i64; wb[8]=sys_mmap(fd*ne*8) as i64
45 wb[9]=sys_mmap(qd*8) as i64; wb[10]=sys_mmap(kvd*8) as i64; wb[11]=sys_mmap(kvd*8) as i64
46 let sb: *i64=sys_mmap(14*8) as *i64
47 sb[0]=sys_mmap(MAXT*ne*8) as i64; sb[1]=sys_mmap(MAXT*qd*8) as i64; sb[2]=sys_mmap(MAXT*kvd*8) as i64; sb[3]=sys_mmap(MAXT*kvd*8) as i64; sb[4]=sys_mmap(MAXT*qd*8) as i64
48 sb[5]=sys_mmap(MAXT*8) as i64; sb[6]=sys_mmap(MAXT*8) as i64; sb[7]=sys_mmap(MAXT*ne*8) as i64; sb[8]=sys_mmap(MAXT*fd*8) as i64; sb[9]=sys_mmap(MAXT*fd*8) as i64
49 sb[10]=sys_mmap(MAXT*fd*8) as i64; sb[11]=sys_mmap(MAXT*ne*8) as i64; sb[12]=sys_mmap(MAXT*ne*8) as i64; sb[13]=sys_mmap(MAXT*ne*8) as i64
50 let kvc: *i64=sys_mmap(48*8) as *i64
51 var kl: i64=0
52 while kl<24 { kvc[2*kl]=sys_mmap(MAXT*kvd*8) as i64; kvc[2*kl+1]=sys_mmap(MAXT*kvd*8) as i64; kl=kl+1 }
53 let nmbuf: *u8=sys_mmap(64); let freqs: *i64=sys_mmap(32*8) as *i64; rope_freqs(freqs, 64)
54 let tmp: *i64=sys_mmap(64*256*8) as *i64
55 let x: *i64=sys_mmap(MAXT*ne*8) as *i64
56 let hout: *i64=sys_mmap(MAXT*ne*8) as *i64
57 let x1: *i64=sys_mmap(ne*8) as *i64
58 let h1: *i64=sys_mmap(ne*8) as *i64
59 let cfgA: *i64=sys_mmap(8*8) as *i64; cfgA[1]=ne; cfgA[2]=14; cfgA[3]=2; cfgA[4]=64; cfgA[5]=qd; cfgA[6]=kvd; cfgA[7]=8192
60 let cfgF: *i64=sys_mmap(4*8) as *i64; cfgF[1]=ne; cfgF[2]=fd
61
62 let input: *u8="The capital of France is\x00" as *u8
63 let ids: *i64=sys_mmap(MAXT*8) as *i64; let tokp: *i64=sys_mmap(MAXT*8) as *i64; let tokl: *i64=sys_mmap(MAXT*8) as *i64
64 let nprompt: i64=tk_bpe_encode(buf, mfirst, nm_c, vfirst, vocab, input, dp_slen(input), tokp, tokl, ids)
65
66 // ---- warm prefill so decode_step_kv attends over a valid cache ----
67 var ei: i64=0; while ei<nprompt { dequant_row(buf, te_base, te_ty, ids[ei], ne, ((x as i64)+ei*ne*8) as *i64, tmp); ei=ei+1 }
68 cfgA[0]=nprompt; cfgF[0]=nprompt
69 run_stack_prefill_kv(buf, hdr, x, hout, wb, sb, nmbuf, freqs, kvc, cfgA, cfgF, 24)
70 dp_puts("prefill done ("); dp_num(nprompt); dp_puts(" tok)\n\n");
71
72 let pos: i64=nprompt // decode the token right after the prompt
73 dequant_row(buf, te_base, te_ty, ids[nprompt-1], ne, x1, tmp)
74
75 // ---- (A) one FULL decode_step_kv = dequant + compute ----
76 let a0: i64=sys_now_ms()
77 decode_step_kv(buf, hdr, x1, h1, wb, sb, nmbuf, freqs, kvc, pos, cfgA, cfgF, 24)
78 let a1: i64=sys_now_ms()
79 let full_ms: i64=a1-a0
80
81 // ---- (B) bare 24-layer dequant sweep = dequant ONLY (no matmul/attn/ffn) ----
82 let b0: i64=sys_now_ms()
83 var L: i64=0
84 while L<24 { nf_load_layer_pool(buf, hdr, L, wb, ne, qd, kvd, fd); L=L+1 }
85 let b1: i64=sys_now_ms()
86 let dq_ms: i64=b1-b0
87
88 let compute_ms: i64=full_ms-dq_ms
89
90 dp_puts("(A) full decode_step_kv (dequant+compute): "); dp_num(full_ms); dp_puts(" ms\n")
91 dp_puts("(B) bare 24-layer dequant sweep only: "); dp_num(dq_ms); dp_puts(" ms\n")
92 dp_puts("=> derived COMPUTE floor (A-B): "); dp_num(compute_ms); dp_puts(" ms\n")
93 dp_puts("=> dequant fraction of a decode step: "); if full_ms>0 { dp_num((dq_ms*100)/full_ms) } else { dp_num(0) } dp_puts(" %\n\n")
94
95 // verdict is informational: GREEN if the measurement ran; the % tells us the lever.
96 dp_puts("PROBE COMPLETE. If dequant fraction is large, a one-time weight cache removes it LOSSLESSLY.\n")
97 dp_puts("NX-NOFLOAT-QWEN-DQPROBE-GATE verdict=GREEN (measurement captured)\n")
98 return 0
99}