nx_nofloat_qwen_kvgen_gate.nx source
↩ module page · 168 lines · 10807 B
1// nx_nofloat_qwen_kvgen_gate.nx -- KV-CACHED greedy generation through the sovereign no-float Qwen2.5-0.5B.
2// Prefill the 5-token prompt once (capturing every layer's K/V rows), then each new token is ONE single-token
3// decode_step_kv (flat cost in sequence length) instead of re-forwarding the whole growing sequence.
4// LOSSLESSNESS TEETH: the cached path must reproduce the uncached nx_nofloat_qwen_gen_gate tokens EXACTLY
5// (' Paris' '.' ' It' ' is' = 12095 13 1084 374) -- causal masking makes KV caching bit-exact, so any drift = bug.
6// 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 kg_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 kg_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 kg_slen(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} return n }
22
23// streamed LM head argmax over the full vocab for ONE normed row (Q24). Returns the id; logit -> lgout[0].
24func kg_head_argmax(buf: *u8, oh_base: i64, oh_ty: i64, normed: *i64, row: *i64, tmp: *i64, vocab: i64, ne: i64, lgout: *i64) -> i64 {
25 var best: i64=0
26 var bestv: i64=0-9223372036854775807
27 var v: i64=0
28 while v<vocab {
29 dequant_row(buf, oh_base, oh_ty, v, ne, row, tmp)
30 var s: i64=0; var k: i64=0; while k<ne { s=s+(normed[k]*row[k]); k=k+1 }
31 let lg: i64=s>>24
32 if lg>bestv { bestv=lg; best=v }
33 v=v+1
34 }
35 lgout[0]=bestv
36 return best
37}
38
39func main() -> i64 {
40 kg_puts("KV-CACHED NO-FLOAT GENERATION: prefill once, then flat single-token decode steps\n\n" as *u8)
41 let path: *u8 = sp_path("nx_real_model.gguf" as *u8, sys_mmap(SP_PATH_MAX))
42 sp_skip_unless("NOFLOAT-QWEN-KVGEN-GATE" as *u8, path)
43 let len_out: *i64 = sys_mmap(8) as *i64; len_out[0]=0
44 let buf: *u8 = sys_read_file(path, len_out)
45 if buf == (0 as *u8) { kg_puts("MODEL ABSENT\n" as *u8); return 1 }
46 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
47 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { kg_puts("PARSE FAIL\n" as *u8); return 1 }
48
49 let ne: i64=896; let qd: i64=896; let kvd: i64=128; let fd: i64=4864; let MAXT: i64=16
50 let voff: *i64=sys_mmap(8) as *i64; let vty: *i64=sys_mmap(8) as *i64
51 var mfirst: i64=0; var nm_c: i64=0; var vfirst: i64=0; var vocab: i64=0
52 let km: *u8="tokenizer.ggml.merges\x00" as *u8; let kt: *u8="tokenizer.ggml.tokens\x00" as *u8
53 if nx_gguf_meta_find(buf, len_out[0], hdr, km, kg_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]) }
54 if nx_gguf_meta_find(buf, len_out[0], hdr, kt, kg_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]) }
55 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
56 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)
57 if ie<0 { return 1 } if io<0 { return 1 } if inn<0 { return 1 }
58 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
59 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
60 let gout: *i64=sys_mmap(ne*8) as *i64; load_named_q16(buf, hdr, nn, 18, gout, ne)
61
62 let wb: *i64=sys_mmap(12*8) as *i64
63 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
64 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
65 wb[9]=sys_mmap(qd*8) as i64; wb[10]=sys_mmap(kvd*8) as i64; wb[11]=sys_mmap(kvd*8) as i64
66 let sb: *i64=sys_mmap(14*8) as *i64
67 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
68 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
69 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
70 // per-layer KV caches: 24 layers x (K + V), each MAXT x kvd
71 let kvc: *i64=sys_mmap(48*8) as *i64
72 var kl: i64=0
73 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 }
74 let nmbuf: *u8=sys_mmap(64); let freqs: *i64=sys_mmap(32*8) as *i64; rope_freqs(freqs, 64)
75 let tmp: *i64=sys_mmap(64*256*8) as *i64
76 let x: *i64=sys_mmap(MAXT*ne*8) as *i64
77 let hout: *i64=sys_mmap(MAXT*ne*8) as *i64
78 let x1: *i64=sys_mmap(ne*8) as *i64
79 let h1: *i64=sys_mmap(ne*8) as *i64
80 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
81 let cfgF: *i64=sys_mmap(4*8) as *i64; cfgF[1]=ne; cfgF[2]=fd
82
83 let input: *u8="The capital of France is\x00" as *u8
84 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
85 let nprompt: i64=tk_bpe_encode(buf, mfirst, nm_c, vfirst, vocab, input, kg_slen(input), tokp, tokl, ids)
86 kg_puts("prompt tokens: "); kg_num(nprompt); kg_puts("\n")
87
88 let normed: *i64=sys_mmap(ne*8) as *i64
89 let row: *i64=sys_mmap(ne*8) as *i64
90 let lgout: *i64=sys_mmap(8) as *i64
91 let idout: *i64=sys_mmap(8) as *i64
92 let hp: *i64=sys_mmap(8*8) as *i64
93 hp[0]=buf as i64; hp[1]=oh_base; hp[2]=oh_ty; hp[3]=normed as i64; hp[4]=vocab; hp[5]=ne; hp[6]=idout as i64; hp[7]=lgout as i64
94
95 // ---- PREFILL: full prompt once, capturing all layers' K/V ----
96 let t0: i64=sys_now_ms()
97 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 }
98 cfgA[0]=nprompt; cfgF[0]=nprompt
99 run_stack_prefill_kv(buf, hdr, x, hout, wb, sb, nmbuf, freqs, kvc, cfgA, cfgF, 24)
100 rmsnorm_gamma_row_q24(hout, gout, (nprompt-1)*ne, ne, normed, 0)
101 var next: i64=head_argmax_pool(hp)
102 let t1: i64=sys_now_ms()
103 kg_puts(" prefill ("); kg_num(nprompt); kg_puts(" tok): next id="); kg_num(next); kg_puts(" logit="); kg_num(lgout[0]); kg_puts(" ["); kg_num(t1-t0); kg_puts(" ms]\n")
104 let prefill_first: i64=next
105 var T: i64=nprompt
106 ids[T]=next; T=T+1
107 var ngen: i64=1
108
109 // ---- DECODE: one single-token step per new token, flat cost ----
110 let MAXNEW: i64=6
111 var stop: i64=0
112 if next==151643 { stop=1 }
113 if next==151645 { stop=1 }
114 var decms_total: i64=0
115 var ndec: i64=0
116 while stop==0 {
117 if ngen>=MAXNEW { stop=1 } else { if T>=MAXT { stop=1 } else {
118 let td0: i64=sys_now_ms()
119 let pos: i64=T-1 // the just-appended token's absolute position
120 dequant_row(buf, te_base, te_ty, ids[pos], ne, x1, tmp)
121 decode_step_kv(buf, hdr, x1, h1, wb, sb, nmbuf, freqs, kvc, pos, cfgA, cfgF, 24)
122 rmsnorm_gamma_row_q24(h1, gout, 0, ne, normed, 0)
123 next=head_argmax_pool(hp)
124 let td1: i64=sys_now_ms()
125 decms_total=decms_total+(td1-td0); ndec=ndec+1
126 kg_puts(" decode pos="); kg_num(pos); kg_puts(": next id="); kg_num(next); kg_puts(" logit="); kg_num(lgout[0]); kg_puts(" piece='")
127 let off: i64=tk_decode_off(buf, vfirst, next); let pl: i64=nx_gguf_meta_read_string_len(buf, off)
128 if pl>0 { sys_write(1, nx_gguf_meta_read_string_ptr(buf, off), pl) }
129 kg_puts("' ["); kg_num(td1-td0); kg_puts(" ms]\n")
130 ids[T]=next; T=T+1; ngen=ngen+1
131 if next==151643 { stop=1 }
132 if next==151645 { stop=1 }
133 } }
134 }
135
136 kg_puts("\nCOMPLETION: '" as *u8)
137 var gi: i64=nprompt
138 var outbytes: i64=0
139 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 }
140 kg_puts("'\n" as *u8)
141 kg_puts("timing: prefill="); kg_num(t1-t0); kg_puts(" ms; decode steps="); kg_num(ndec); kg_puts(" mean=")
142 if ndec>0 { kg_num(decms_total/ndec) } else { kg_num(0) }
143 kg_puts(" ms/token (uncached re-forward would GROW with T; cached is FLAT)\n\n" as *u8)
144
145 var pass: i64=0
146 var ttl: i64=0
147 ttl=ttl+1; kg_puts(" T1 cached prefill argmax == 12095 (ĠParis -- matches the proven uncached forward): "); if prefill_first==12095 { pass=pass+1; kg_puts("PASS\n") } else { kg_puts("FAIL\n") }
148 ttl=ttl+1; kg_puts(" T2 LOSSLESS KV: first 4 generated == uncached gen gate [12095,13,1084,374]: ")
149 var ok4: i64=1
150 if ids[nprompt]!=12095 { ok4=0 }
151 if ids[nprompt+1]!=13 { ok4=0 }
152 if ids[nprompt+2]!=1084 { ok4=0 }
153 if ids[nprompt+3]!=374 { ok4=0 }
154 if ok4==1 { pass=pass+1; kg_puts("PASS\n") } else { kg_puts("FAIL\n") }
155 ttl=ttl+1; kg_puts(" T3 generated "); kg_num(ngen); kg_puts(" tokens, all valid ids: "); 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; kg_puts("PASS\n") } else { kg_puts("FAIL\n") } } else { kg_puts("FAIL\n") }
156 ttl=ttl+1; kg_puts(" T4 completion decodes to non-empty text: "); if outbytes>0 { pass=pass+1; kg_puts("PASS\n") } else { kg_puts("FAIL\n") }
157
158 kg_puts("NX-NOFLOAT-QWEN-KVGEN-GATE passed "); kg_num(pass); kg_puts("/"); kg_num(ttl)
159 // MIGRATED onto nx_gate_verdict by nx_gate_dry_apply (D001, minimal form): every check
160 // row above is untouched, so the PASS/FAIL vector cannot change; only the hand-rolled
161 // verdict emission is replaced by the ONE shared base class. Proven by nx_gate_migrate verify.
162 let ctr__dry: *i64 = gv_ctr()
163 ctr__dry[0] = pass
164 ctr__dry[1] = ttl
165 let rc__dry: i64 = gv_verdict("NOFLOAT-QWEN-KVGEN-GATE" as *u8, ctr__dry, "KV-cached decode is LOSSLESS vs the uncached forward + flat per-token cost)" as *u8)
166 sys_exit(rc__dry)
167 return rc__dry
168}