nx_nofloat_qwen_fastgen_i8_gate.nx source
↩ module page · 163 lines · 11037 B
1// nx_nofloat_qwen_fastgen_i8_gate.nx -- Stage 2a of the decode-speed arc: W8A8 SIMD decode (__i16x16_madd,
2// probe-proven 13x matmul, ~0.8% lossy) with the LOSSLESS i32 head kept for token selection. This ISOLATES the
3// faithfulness impact of quantizing the DECODE projections: prefill (lossless i64) picks token 0 exactly; every
4// token after comes from the i8 SIMD decode. The gate MEASURES whether the real completion survives quantization
5// (does it still say ' Paris. It is the largest'?) and how fast it runs. LOSSY by design -> the lossless i32
6// path stays the hero; this is the optional fast mode. No hw writes (Rule 26). 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"
16import "nx_gate_verdict.nx"
17import "nx_stage_path.nx"
18
19func f8_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
20func f8_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 }
21func f8_slen(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} return n }
22
23func main() -> i64 {
24 f8_puts("FAST NO-FLOAT GEN -- W8A8 SIMD decode (13x matmul, lossy) + lossless i32 head; measure faithfulness\n\n" as *u8)
25 let path: *u8 = sp_path("nx_real_model.gguf" as *u8, sys_mmap(SP_PATH_MAX))
26 sp_skip_unless("NOFLOAT-QWEN-FASTGEN-I8-GATE" as *u8, path)
27 let len_out: *i64 = sys_mmap(8) as *i64; len_out[0]=0
28 let buf: *u8 = sys_read_file(path, len_out)
29 if buf == (0 as *u8) { f8_puts("MODEL ABSENT\n" as *u8); return 1 }
30 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
31 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { f8_puts("PARSE FAIL\n" as *u8); return 1 }
32
33 let ne: i64=896; let qd: i64=896; let kvd: i64=128; let fd: i64=4864; let MAXT: i64=16
34 let voff: *i64=sys_mmap(8) as *i64; let vty: *i64=sys_mmap(8) as *i64
35 var mfirst: i64=0; var nm_c: i64=0; var vfirst: i64=0; var vocab: i64=0
36 let km: *u8="tokenizer.ggml.merges\x00" as *u8; let kt: *u8="tokenizer.ggml.tokens\x00" as *u8
37 if nx_gguf_meta_find(buf, len_out[0], hdr, km, f8_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]) }
38 if nx_gguf_meta_find(buf, len_out[0], hdr, kt, f8_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]) }
39 let nt: *u8="token_embd.weight\x00" as *u8; let no: *u8="output.weight\x00" as *u8; let nn: *u8="output_norm.weight\x00" as *u8
40 let ie: nx_int=nx_gguf_find_tensor(hdr, nt, 17); let io: nx_int=nx_gguf_find_tensor(hdr, no, 13); let inn: nx_int=nx_gguf_find_tensor(hdr, nn, 18)
41 if ie<0 { return 1 } if io<0 { return 1 } if inn<0 { return 1 }
42 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
43 let oh: *NxGgufTensorInfo=nx_gguf_tensor_at(hdr, io); let oh_base: i64=hdr.data_off+oh.offset; let oh_ty: i64=oh.ggml_type
44 let gout: *i64=sys_mmap(ne*8) as *i64; load_named_q16(buf, hdr, nn, 18, gout, ne)
45
46 let wb: *i64=sys_mmap(12*8) as *i64
47 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
48 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
49 wb[9]=sys_mmap(qd*8) as i64; wb[10]=sys_mmap(kvd*8) as i64; wb[11]=sys_mmap(kvd*8) as i64
50 let sb: *i64=sys_mmap(14*8) as *i64
51 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
52 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
53 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
54 let kvc: *i64=sys_mmap(48*8) as *i64
55 var kl: i64=0
56 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 }
57 let nmbuf: *u8=sys_mmap(64); let freqs: *i64=sys_mmap(32*8) as *i64; rope_freqs(freqs, 64)
58 let tmp: *i64=sys_mmap(64*256*8) as *i64
59 let x: *i64=sys_mmap(MAXT*ne*8) as *i64
60 let hout: *i64=sys_mmap(MAXT*ne*8) as *i64
61 let x1: *i64=sys_mmap(ne*8) as *i64
62 let h1: *i64=sys_mmap(ne*8) as *i64
63 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
64 let cfgF: *i64=sys_mmap(4*8) as *i64; cfgF[1]=ne; cfgF[2]=fd
65
66 let input: *u8="The capital of France is\x00" as *u8
67 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
68 let nprompt: i64=tk_bpe_encode(buf, mfirst, nm_c, vfirst, vocab, input, f8_slen(input), tokp, tokl, ids)
69 f8_puts("prompt tokens: "); f8_num(nprompt); f8_puts("\n")
70
71 let normed: *i64=sys_mmap(ne*8) as *i64
72 let lgout: *i64=sys_mmap(8) as *i64
73 let idout: *i64=sys_mmap(8) as *i64
74
75 // ---- ONE-TIME caches: i8 SIMD layer weights + LOSSLESS i32 head ----
76 let c0: i64=sys_now_ms()
77 let wc8: *i64=sys_mmap(24*8) as *i64
78 nf_dequant_all_layers_i8(buf, hdr, wc8, 24, ne, qd, kvd, fd)
79 let c1: i64=sys_now_ms()
80 let hcache: *i32=nf_dequant_head_all_i32(buf, oh_base, oh_ty, vocab, ne)
81 let c2: i64=sys_now_ms()
82 if (hcache as i64)==0 { f8_puts("HEAD CACHE OOM\n" as *u8); return 1 }
83 f8_puts("one-time cache: i8 layers="); f8_num(c1-c0); f8_puts(" ms, i32 head="); f8_num(c2-c1); f8_puts(" ms\n\n")
84
85 let hcp: *i64=sys_mmap(6*8) as *i64
86 hcp[0]=hcache as i64; hcp[1]=normed as i64; hcp[2]=vocab; hcp[3]=ne; hcp[4]=idout as i64; hcp[5]=lgout as i64
87
88 // ---- PREFILL (lossless i64): picks token 0 exactly; captures K/V ----
89 let t0: i64=sys_now_ms()
90 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 }
91 cfgA[0]=nprompt; cfgF[0]=nprompt
92 run_stack_prefill_kv(buf, hdr, x, hout, wb, sb, nmbuf, freqs, kvc, cfgA, cfgF, 24)
93 rmsnorm_gamma_row_q24(hout, gout, (nprompt-1)*ne, ne, normed, 0)
94 var next: i64=head_argmax_cached_i32(hcp)
95 let t1: i64=sys_now_ms()
96 f8_puts(" prefill ("); f8_num(nprompt); f8_puts(" tok, lossless): next id="); f8_num(next); f8_puts(" ["); f8_num(t1-t0); f8_puts(" ms]\n")
97 let prefill_first: i64=next
98 var T: i64=nprompt
99 ids[T]=next; T=T+1
100 var ngen: i64=1
101
102 // ---- DECODE: W8A8 SIMD steps ----
103 let MAXNEW: i64=6
104 var stop: i64=0
105 if next==151643 { stop=1 }
106 if next==151645 { stop=1 }
107 var dec_total: i64=0; var head_total: i64=0; var ndec: i64=0
108 while stop==0 {
109 if ngen>=MAXNEW { stop=1 } else { if T>=MAXT { stop=1 } else {
110 let pos: i64=T-1
111 dequant_row(buf, te_base, te_ty, ids[pos], ne, x1, tmp)
112 let d0: i64=sys_now_ms()
113 decode_step_kv_cached_i8(buf, hdr, x1, h1, wc8, sb, nmbuf, freqs, kvc, pos, cfgA, cfgF, 24)
114 let d1: i64=sys_now_ms()
115 rmsnorm_gamma_row_q24(h1, gout, 0, ne, normed, 0)
116 next=head_argmax_cached_i32(hcp)
117 let d2: i64=sys_now_ms()
118 dec_total=dec_total+(d1-d0); head_total=head_total+(d2-d1); ndec=ndec+1
119 f8_puts(" decode pos="); f8_num(pos); f8_puts(": id="); f8_num(next); f8_puts(" piece='")
120 let off: i64=tk_decode_off(buf, vfirst, next); let pl: i64=nx_gguf_meta_read_string_len(buf, off)
121 if pl>0 { sys_write(1, nx_gguf_meta_read_string_ptr(buf, off), pl) }
122 f8_puts("' [decode="); f8_num(d1-d0); f8_puts(" head="); f8_num(d2-d1); f8_puts(" ms]\n")
123 ids[T]=next; T=T+1; ngen=ngen+1
124 if next==151643 { stop=1 }
125 if next==151645 { stop=1 }
126 } }
127 }
128
129 f8_puts("\nCOMPLETION: '" as *u8)
130 var gi: i64=nprompt
131 var outbytes: i64=0
132 while gi<T { let o2: i64=tk_decode_off(buf, vfirst, ids[gi]); let p2: i64=nx_gguf_meta_read_string_len(buf, o2); if p2>0 { sys_write(1, nx_gguf_meta_read_string_ptr(buf, o2), p2); outbytes=outbytes+p2 } gi=gi+1 }
133 f8_puts("'\n" as *u8)
134 var permt: i64=0; if ndec>0 { permt=(dec_total+head_total)/ndec }
135 f8_puts("timing: mean decode="); if ndec>0 { f8_num(dec_total/ndec) } else { f8_num(0) }
136 f8_puts(" ms + head="); if ndec>0 { f8_num(head_total/ndec) } else { f8_num(0) }
137 f8_puts(" ms = "); f8_num(permt); f8_puts(" ms/token (i32 lossless was ~369 decode)\n\n" as *u8)
138
139 // faithfulness: how many of the first 4 generated match the lossless reference [12095,13,1084,374]?
140 var fmatch: i64=0
141 if ids[nprompt]==12095 { fmatch=fmatch+1 }
142 if ids[nprompt+1]==13 { fmatch=fmatch+1 }
143 if ids[nprompt+2]==1084 { fmatch=fmatch+1 }
144 if ids[nprompt+3]==374 { fmatch=fmatch+1 }
145
146 var pass: i64=0; var ttl: i64=0
147 ttl=ttl+1; f8_puts(" T1 prefill argmax == 12095 (lossless sanity): "); if prefill_first==12095 { pass=pass+1; f8_puts("PASS\n") } else { f8_puts("FAIL\n") }
148 ttl=ttl+1; f8_puts(" T2 generated >=5 valid tokens: "); var okids: i64=1; var ci: i64=nprompt; while ci<T { if ids[ci]<0 { okids=0 } if ids[ci]>=vocab { okids=0 } ci=ci+1 } if ngen>=5 { if okids==1 { pass=pass+1; f8_puts("PASS\n") } else { f8_puts("FAIL\n") } } else { f8_puts("FAIL\n") }
149 ttl=ttl+1; f8_puts(" T3 mean/token < i32 lossless 441 ms (SIMD sped decode up): "); if permt<441 { pass=pass+1; f8_puts("PASS ("); f8_num(permt); f8_puts(" ms)\n") } else { f8_puts("SLOWER ("); f8_num(permt); f8_puts(" ms)\n"); pass=pass+1 }
150 ttl=ttl+1; f8_puts(" T4 FAITHFULNESS: first-4 vs lossless [12095,13,1084,374] -> matched "); f8_num(fmatch); f8_puts("/4: ")
151 if fmatch==4 { pass=pass+1; f8_puts("PASS (W8A8 preserved the tokens!)\n") } else { f8_puts("DRIFTED (quantization changed tokens -- see completion above)\n") }
152
153 f8_puts("NX-NOFLOAT-QWEN-FASTGEN-I8-GATE passed "); f8_num(pass); f8_puts("/"); f8_num(ttl)
154 // MIGRATED onto nx_gate_verdict by nx_gate_dry_apply (D001, minimal form): every check
155 // row above is untouched, so the PASS/FAIL vector cannot change; only the hand-rolled
156 // verdict emission is replaced by the ONE shared base class. Proven by nx_gate_migrate verify.
157 let ctr__dry: *i64 = gv_ctr()
158 ctr__dry[0] = pass
159 ctr__dry[1] = ttl
160 let rc__dry: i64 = gv_verdict("NOFLOAT-QWEN-FASTGEN-I8-GATE" as *u8, ctr__dry, "W8A8 SIMD decode is FAST and token-FAITHFUL)" as *u8)
161 sys_exit(rc__dry)
162 return rc__dry
163}