nx_gpu_serve_gate.nx source
↩ module page · 151 lines · 10692 B
1// nx_gpu_serve_gate.nx -- #23 KEYSTONE: the SOVEREIGN serve-side GPU client, gated byte-exact.
2// Loads the CODER model (nx_coder_model.gguf), generates NGEN greedy tokens on the proven CPU path
3// (sequential decode_step_kv_cached_i8 + i32 head -- the serve math), then repeats the SAME generation
4// via the GPU logits server on unix socket /home/elderwesto/nx_stage/nx_gpu.sock: per token the client
5// sends [pos: i64][x1: ne i64 embedding row] and receives [logits: vocab i64], argmaxes on CPU
6// (lowest-id tie-break = head_argmax convention). GATE: token ids IDENTICAL CPU vs GPU for the
7// prompt-next + all NGEN steps. The C server is the BOOTSTRAP ORACLE (doctrine: 3rd parties =
8// benchmarks); THIS client is the sovereign integration shape for nx_nofloat_serve_core.
9// Requires the server running: /tmp/gls --serve. expect_exit: 0 license_tier: ORIGINAL
10import "nx_syscalls.nx"
11import "nx_tier.nx"
12import "nx_le.nx"
13import "nx_tensor.nx"
14import "nx_gguf.nx"
15import "nx_gguf_load.nx"
16import "nx_gguf_meta.nx"
17import "nx_nofloat_llm.nx"
18import "nx_nofloat_tok.nx"
19import "nx_gate_verdict.nx"
20import "nx_stage_path.nx"
21
22func gs_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
23func gs_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 }
24func gs_slen(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} return n }
25// write-all / read-all over the socket (partial transfers WILL happen on the 1.2MB logits)
26func gs_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 }
27func gs_rall(fd: i64, buf: *u8, count: i64) -> i64 { var off: i64=0; while off<count { let r: i64=sys_read(fd, ((buf as i64)+off) as *u8, count-off); if r<=0 { return 0-1 } off=off+r } return 0 }
28// connect to the unix-domain GPU socket; -1 on failure
29func gs_connect(path: *u8) -> i64 {
30 let fd: i64 = sys_socket(1, 1, 0)
31 if fd < 0 { return 0-1 }
32 let sa: *u8 = sys_mmap(128)
33 sa[0]=1 as u8
34 sa[1]=0 as u8
35 let plen: i64 = gs_slen(path)
36 var i: i64=0
37 while i<plen { sa[2+i]=path[i]; i=i+1 }
38 let rc: i64 = sys_connect(fd, sa, 2+plen+1)
39 if rc < 0 { return 0-1 }
40 return fd
41}
42// argmax over an i64 logit vector, LOWEST id on ties (matches head_argmax_cached_i32)
43func gs_argmax(lgv: *i64, vocab: i64) -> i64 { var best: i64=lgv[0]; var bid: i64=0; var v: i64=1; while v<vocab { if lgv[v]>best { best=lgv[v]; bid=v } v=v+1 } return bid }
44
45func main() -> i64 {
46 gs_puts("GPU-SERVE-GATE: sovereign socket client vs CPU reference (CODER model, greedy, byte-exact ids)\n\n" as *u8)
47 let path: *u8 = sp_path("nx_coder_model.gguf" as *u8, sys_mmap(SP_PATH_MAX))
48 sp_skip_unless("GPU-SERVE-GATE" as *u8, path)
49 let len_out: *i64 = sys_mmap(8) as *i64; len_out[0]=0
50 let buf: *u8 = sys_read_file(path, len_out)
51 if buf == (0 as *u8) { gs_puts("MODEL ABSENT\n" as *u8); return 1 }
52 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
53 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { gs_puts("PARSE FAIL\n" as *u8); return 1 }
54 let ne: i64=896; let qd: i64=896; let kvd: i64=128; let fd: i64=4864; let MAXT: i64=16; let NGEN: i64=6
55 let voff: *i64=sys_mmap(8) as *i64; let vty: *i64=sys_mmap(8) as *i64
56 var mfirst: i64=0; var nm_c: i64=0; var vfirst: i64=0; var vocab: i64=0
57 let km: *u8="tokenizer.ggml.merges\x00" as *u8; let kt: *u8="tokenizer.ggml.tokens\x00" as *u8
58 if nx_gguf_meta_find(buf, len_out[0], hdr, km, gs_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]) }
59 if nx_gguf_meta_find(buf, len_out[0], hdr, kt, gs_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]) }
60 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
61 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)
62 if ie<0 { gs_puts("NO EMBD\n" as *u8); return 1 } if io<0 { gs_puts("NO HEAD\n" as *u8); return 1 } if inn<0 { gs_puts("NO NORM\n" as *u8); return 1 }
63 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
64 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
65 let gout: *i64=sys_mmap(ne*8) as *i64; load_named_q16(buf, hdr, nn, 18, gout, ne)
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 let kvc: *i64=sys_mmap(48*8) as *i64
71 var kl: i64=0; 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 }
72 let nmbuf: *u8=sys_mmap(64); let freqs: *i64=sys_mmap(32*8) as *i64; rope_freqs(freqs, 64)
73 let tmp: *i64=sys_mmap(64*256*8) as *i64
74 let x1: *i64=sys_mmap(ne*8) as *i64
75 let h1: *i64=sys_mmap(ne*8) as *i64
76 let normed: *i64=sys_mmap(ne*8) as *i64
77 let idout: *i64=sys_mmap(8) as *i64; let lgout: *i64=sys_mmap(8) as *i64
78 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
79 let cfgF: *i64=sys_mmap(4*8) as *i64; cfgF[1]=ne; cfgF[2]=fd
80 let input: *u8="The capital of France is\x00" as *u8
81 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
82 let nprompt: i64=tk_bpe_encode(buf, mfirst, nm_c, vfirst, vocab, input, gs_slen(input), tokp, tokl, ids)
83 if nprompt<1 { gs_puts("ENCODE FAIL\n" as *u8); return 1 }
84 gs_puts("prompt tokens: " as *u8); gs_num(nprompt); gs_puts(" building i8+head caches...\n" as *u8)
85 let wc8: *i64=sys_mmap(24*8) as *i64
86 nf_dequant_all_layers_i8(buf, hdr, wc8, 24, ne, qd, kvd, fd)
87 let hcache: *i32=nf_dequant_head_all_i32(buf, oh_base, oh_ty, vocab, ne)
88 if (hcache as i64)==0 { gs_puts("HEAD OOM\n" as *u8); return 1 }
89 let hcp: *i64=sys_mmap(6*8) as *i64
90 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
91 // ---- CPU reference: sequential cached prefill + NGEN greedy (the serve math) ----
92 let ids_cpu: *i64=sys_mmap(MAXT*8) as *i64
93 var i: i64=0
94 while i<nprompt { ids_cpu[i]=ids[i]; dequant_row(buf, te_base, te_ty, ids[i], ne, x1, tmp); decode_step_kv_cached_i8(buf, hdr, x1, h1, wc8, sb, nmbuf, freqs, kvc, i, cfgA, cfgF, 24); i=i+1 }
95 rmsnorm_gamma_row_q24(h1, gout, 0, ne, normed, 0)
96 var next: i64=head_argmax_cached_i32(hcp)
97 var T: i64=nprompt
98 ids_cpu[T]=next; T=T+1
99 var g: i64=1
100 while g<NGEN {
101 let pos: i64=T-1
102 dequant_row(buf, te_base, te_ty, ids_cpu[pos], ne, x1, tmp)
103 decode_step_kv_cached_i8(buf, hdr, x1, h1, wc8, sb, nmbuf, freqs, kvc, pos, cfgA, cfgF, 24)
104 rmsnorm_gamma_row_q24(h1, gout, 0, ne, normed, 0)
105 next=head_argmax_cached_i32(hcp)
106 ids_cpu[T]=next; T=T+1; g=g+1
107 }
108 gs_puts("CPU ids: " as *u8); var pi: i64=nprompt; while pi<T { gs_num(ids_cpu[pi]); gs_puts(" " as *u8); pi=pi+1 } gs_puts("\n" as *u8)
109 // ---- GPU path: same generation via the logits socket ----
110 let sock: *u8 = sp_path("nx_gpu.sock" as *u8, sys_mmap(SP_PATH_MAX))
111 let cfd: i64=gs_connect(sock)
112 if cfd<0 { gs_puts("GPU SERVER DOWN (start /tmp/gls --serve)\nverdict=RED\n" as *u8); return 1 }
113 let lgv: *i64=sys_mmap(vocab*8) as *i64
114 let posbuf: *i64=sys_mmap(8) as *i64
115 let ids_gpu: *i64=sys_mmap(MAXT*8) as *i64
116 var j: i64=0; while j<nprompt { ids_gpu[j]=ids[j]; j=j+1 }
117 var Tg: i64=nprompt
118 let t0: i64=sys_now_ms()
119 var pos2: i64=0
120 while Tg<T {
121 // feed tokens 0..Tg-1 progressively; request logits at each pos; only argmax when at the frontier
122 if pos2>=Tg { gs_puts("SEQ ERROR\n" as *u8); return 1 }
123 dequant_row(buf, te_base, te_ty, ids_gpu[pos2], ne, x1, tmp)
124 posbuf[0]=pos2
125 if gs_wall(cfd, posbuf as *u8, 8)<0 { gs_puts("SEND FAIL\nverdict=RED\n" as *u8); return 1 }
126 if gs_wall(cfd, x1 as *u8, ne*8)<0 { gs_puts("SEND FAIL\nverdict=RED\n" as *u8); return 1 }
127 if gs_rall(cfd, lgv as *u8, vocab*8)<0 { gs_puts("RECV FAIL\nverdict=RED\n" as *u8); return 1 }
128 if pos2==Tg-1 { let nx: i64=gs_argmax(lgv, vocab); ids_gpu[Tg]=nx; Tg=Tg+1 }
129 pos2=pos2+1
130 }
131 let t1: i64=sys_now_ms()
132 gs_puts("GPU ids: " as *u8); pi=nprompt; while pi<Tg { gs_num(ids_gpu[pi]); gs_puts(" " as *u8); pi=pi+1 } gs_puts("\n" as *u8)
133 // ---- compare + decode text ----
134 var pass: i64=0; var ttl: i64=0
135 pi=nprompt
136 while pi<T { ttl=ttl+1; if ids_gpu[pi]==ids_cpu[pi] { pass=pass+1 } pi=pi+1 }
137 gs_puts("completion: '" as *u8)
138 pi=nprompt; while pi<Tg { let o2: i64=tk_decode_off(buf, vfirst, ids_gpu[pi]); 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) } pi=pi+1 }
139 gs_puts("'\n" as *u8)
140 gs_puts("gpu wall: " as *u8); gs_num(t1-t0); gs_puts(" ms for " as *u8); gs_num(T-1); gs_puts(" fwd steps (incl 1.2MB logits/step over socket)\n" as *u8)
141 gs_puts("NX-GPU-SERVE-GATE passed " as *u8); gs_num(pass); gs_puts("/" as *u8); gs_num(ttl)
142 // MIGRATED onto nx_gate_verdict by nx_gate_dry_apply (D001, minimal form): every check
143 // row above is untouched, so the PASS/FAIL vector cannot change; only the hand-rolled
144 // verdict emission is replaced by the ONE shared base class. Proven by nx_gate_migrate verify.
145 let ctr__dry: *i64 = gv_ctr()
146 ctr__dry[0] = pass
147 ctr__dry[1] = ttl
148 let rc__dry: i64 = gv_verdict("GPU-SERVE-GATE" as *u8, ctr__dry, "sovereign socket client reproduces the CPU serve byte-exact -- the serve can ride the GPU)" as *u8)
149 sys_exit(rc__dry)
150 return rc__dry
151}