nx_gpu_export.nx source
↩ module page · 131 lines · 9508 B
1// nx_gpu_export.nx -- dump the no-float Qwen2.5-0.5B i8 caches + i32 head + output_norm gamma + freqs +
2// config + a sequential-decode CPU REFERENCE (input embeddings, final normed, next-token id) to a binary
3// that the GPU forward server loads. Reuses the PROVEN load/dequant lib (eat-the-debt: DUMP the built
4// cache, do NOT re-parse GGUF). One-shot offline tool. No hw writes (Rule 26). expect_exit: 0 license_tier: ORIGINAL
5import "nx_syscalls.nx"
6import "nx_tier.nx"
7import "nx_le.nx"
8import "nx_tensor.nx"
9import "nx_gguf.nx"
10import "nx_gguf_load.nx"
11import "nx_gguf_meta.nx"
12import "nx_nofloat_llm.nx"
13import "nx_nofloat_tok.nx"
14const K_MAGIC_65536: i64 = 65536
15const K_MAGIC_5108: i64 = 5108
16
17func ge_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 ge_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 ge_slen(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} return n }
20// read a u32 config value from GGUF metadata by key (like the config gate's cg_u32); -1 if absent
21func ge_u32(buf: *u8, flen: i64, hdr: *NxGgufHeader, key: *u8) -> i64 {
22 let voff: *i64 = sys_mmap(8) as *i64; let vty: *i64 = sys_mmap(8) as *i64
23 let r: nx_int = nx_gguf_meta_find(buf, flen, hdr, key, ge_slen(key), voff, vty)
24 if r != NX_GMETA_OK { return 0-1 }
25 return nx_gguf_meta_read_u32(buf, voff[0])
26}
27// write-all: loop until count bytes written (sys_write may be partial). 0 ok / -1 fail
28func ge_wall(fd: i64, buf: *u8, count: i64) -> i64 { var off: i64=0; while off<count { let w: i64=sys_write(fd, ((buf as i64)+off) as *u8, count-off); if w<=0 { return 0-1 } off=off+w } return 0 }
29func ge_w1(fd: i64, v: i64) -> i64 { let p: *i64=sys_mmap(8) as *i64; p[0]=v; return ge_wall(fd, p as *u8, 8) }
30
31func main() -> i64 {
32 ge_puts("GPU-EXPORT: dump i8 caches + i32 head + config + CPU ref for the GPU forward server\n\n" as *u8)
33 let path: *u8 = "/home/elderwesto/nx_stage/nx_coder_model.gguf\x00" as *u8
34 let len_out: *i64 = sys_mmap(8) as *i64; len_out[0]=0
35 let buf: *u8 = sys_read_file(path, len_out)
36 if buf == (0 as *u8) { ge_puts("MODEL ABSENT\n" as *u8); return 1 }
37 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
38 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { ge_puts("PARSE FAIL\n" as *u8); return 1 }
39 // dims from GGUF metadata (schema-first) -> handles ANY qwen2-schema model, not just the 0.5B
40 let ne: i64 = ge_u32(buf, len_out[0], hdr, "qwen2.embedding_length\x00" as *u8)
41 let nh: i64 = ge_u32(buf, len_out[0], hdr, "qwen2.attention.head_count\x00" as *u8)
42 let nkv: i64 = ge_u32(buf, len_out[0], hdr, "qwen2.attention.head_count_kv\x00" as *u8)
43 let nlayers: i64 = ge_u32(buf, len_out[0], hdr, "qwen2.block_count\x00" as *u8)
44 let fd: i64 = ge_u32(buf, len_out[0], hdr, "qwen2.feed_forward_length\x00" as *u8)
45 let hd: i64 = ne/nh; let qd: i64 = ne; let kvd: i64 = nkv*hd; let scale: i64 = K_MAGIC_65536/isqrt(hd); let MAXT: i64=16
46 ge_puts("config: ne="); ge_num(ne); ge_puts(" nh="); ge_num(nh); ge_puts(" nkv="); ge_num(nkv); ge_puts(" hd="); ge_num(hd); ge_puts(" fd="); ge_num(fd); ge_puts(" L="); ge_num(nlayers); ge_puts(" scale="); ge_num(scale); ge_puts("\n" as *u8)
47 if ne<1 { ge_puts("BAD CONFIG (not qwen2?)\n" as *u8); return 1 }
48 let voff: *i64=sys_mmap(8) as *i64; let vty: *i64=sys_mmap(8) as *i64
49 var mfirst: i64=0; var nm_c: i64=0; var vfirst: i64=0; var vocab: i64=0
50 let km: *u8="tokenizer.ggml.merges\x00" as *u8; let kt: *u8="tokenizer.ggml.tokens\x00" as *u8
51 if nx_gguf_meta_find(buf, len_out[0], hdr, km, ge_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]) }
52 if nx_gguf_meta_find(buf, len_out[0], hdr, kt, ge_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]) }
53 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
54 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)
55 if ie<0 { ge_puts("NO EMBD\n"); return 1 } if io<0 { ge_puts("NO HEAD\n"); return 1 } if inn<0 { ge_puts("NO NORM\n"); return 1 }
56 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
57 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
58 let gout: *i64=sys_mmap(ne*8) as *i64; load_named_q16(buf, hdr, nn, 18, gout, ne)
59 let freqs: *i64=sys_mmap((hd/2)*8) as *i64; rope_freqs(freqs, hd)
60 let cfgA: *i64=sys_mmap(8*8) as *i64; cfgA[1]=ne; cfgA[2]=nh; cfgA[3]=nkv; cfgA[4]=hd; cfgA[5]=qd; cfgA[6]=kvd; cfgA[7]=scale
61 let cfgF: *i64=sys_mmap(4*8) as *i64; cfgF[1]=ne; cfgF[2]=fd
62 ge_puts("building i8 layer cache...\n" as *u8)
63 let wc8: *i64=sys_mmap(nlayers*8) as *i64
64 nf_dequant_all_layers_i8(buf, hdr, wc8, nlayers, ne, qd, kvd, fd)
65 ge_puts("building i32 head cache...\n" as *u8)
66 let hcache: *i32=nf_dequant_head_all_i32(buf, oh_base, oh_ty, vocab, ne)
67 if (hcache as i64)==0 { ge_puts("HEAD OOM\n" as *u8); return 1 }
68 let sb: *i64=sys_mmap(14*8) as *i64
69 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
70 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
71 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
72 let kvc: *i64=sys_mmap(nlayers*2*8) as *i64
73 var kl: i64=0; while kl<nlayers { 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)
75 let tmp: *i64=sys_mmap(64*256*8) as *i64
76 let x1: *i64=sys_mmap(ne*8) as *i64
77 let h1: *i64=sys_mmap(ne*8) as *i64
78 let normed: *i64=sys_mmap(ne*8) as *i64
79 let idout: *i64=sys_mmap(8) as *i64; let lgout: *i64=sys_mmap(8) as *i64
80 let hcp: *i64=sys_mmap(6*8) as *i64
81 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
82 let input: *u8="The capital of France is\x00" as *u8
83 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
84 let nprompt: i64=tk_bpe_encode(buf, mfirst, nm_c, vfirst, vocab, input, ge_slen(input), tokp, tokl, ids)
85 if nprompt<1 { ge_puts("ENCODE FAIL\n" as *u8); return 1 }
86 ge_puts("prompt tokens: " as *u8); ge_num(nprompt); ge_puts("\n" as *u8)
87 // sequential decode = the CPU reference the GPU must reproduce (bit-identical to prefill via causal mask)
88 let emb: *i64=sys_mmap(MAXT*ne*8) as *i64
89 var i: i64=0
90 while i<nprompt {
91 dequant_row(buf, te_base, te_ty, ids[i], ne, x1, tmp)
92 var c: i64=0; while c<ne { emb[i*ne+c]=x1[c]; c=c+1 }
93 decode_step_kv_cached_i8(buf, hdr, x1, h1, wc8, sb, nmbuf, freqs, kvc, i, cfgA, cfgF, nlayers)
94 i=i+1
95 }
96 rmsnorm_gamma_row_q24(h1, gout, 0, ne, normed, 0)
97 let ref_next: i64=head_argmax_cached_i32(hcp)
98 ge_puts("CPU ref next-token id: " as *u8); ge_num(ref_next); ge_puts("\n" as *u8)
99 // ---- DUMP to binary ----
100 let outp: *u8="/home/elderwesto/nx_stage/nx_gpu_weights.bin\x00" as *u8
101 let fo: i64=sys_openat_wr(outp, 420)
102 if fo<0 { ge_puts("OPEN OUT FAIL\n" as *u8); return 1 }
103 ge_w1(fo, K_MAGIC_5108); ge_w1(fo, ne); ge_w1(fo, qd); ge_w1(fo, kvd); ge_w1(fo, fd)
104 ge_w1(fo, nh); ge_w1(fo, nkv); ge_w1(fo, hd); ge_w1(fo, scale); ge_w1(fo, vocab)
105 ge_w1(fo, nlayers); ge_w1(fo, hd/2); ge_w1(fo, nprompt); ge_w1(fo, ref_next)
106 ge_wall(fo, freqs as *u8, (hd/2)*8)
107 ge_wall(fo, gout as *u8, ne*8)
108 ge_wall(fo, normed as *u8, ne*8)
109 ge_wall(fo, emb as *u8, nprompt*ne*8)
110 var L: i64=0
111 while L<nlayers {
112 let s: *i64=wc8[L] as *i64
113 ge_wall(fo, s[0] as *u8, ne*8)
114 ge_wall(fo, s[1] as *u8, qd*ne*2); ge_wall(fo, s[2] as *u8, qd*8)
115 ge_wall(fo, s[3] as *u8, kvd*ne*2); ge_wall(fo, s[4] as *u8, kvd*8)
116 ge_wall(fo, s[5] as *u8, kvd*ne*2); ge_wall(fo, s[6] as *u8, kvd*8)
117 ge_wall(fo, s[7] as *u8, ne*qd*2); ge_wall(fo, s[8] as *u8, ne*8)
118 ge_wall(fo, s[9] as *u8, ne*8)
119 ge_wall(fo, s[10] as *u8, ne*fd*2); ge_wall(fo, s[11] as *u8, fd*8)
120 ge_wall(fo, s[12] as *u8, ne*fd*2); ge_wall(fo, s[13] as *u8, fd*8)
121 ge_wall(fo, s[14] as *u8, fd*ne*2); ge_wall(fo, s[15] as *u8, ne*8)
122 ge_wall(fo, s[16] as *u8, qd*8); ge_wall(fo, s[17] as *u8, kvd*8); ge_wall(fo, s[18] as *u8, kvd*8)
123 L=L+1
124 }
125 ge_wall(fo, hcache as *u8, vocab*ne*4)
126 sys_close(fo)
127 ge_puts("EXPORTED nx_gpu_weights.bin (nprompt=" as *u8); ge_num(nprompt); ge_puts(" ref_next=" as *u8); ge_num(ref_next); ge_puts(" vocab=" as *u8); ge_num(vocab); ge_puts(")\n" as *u8)
128 ge_puts("verdict=GREEN\n" as *u8)
129 sys_exit(0)
130 return 0
131}