nx_f32_llm_chat_test.nx source
↩ module page · 128 lines · 5757 B
1// nx_f32_llm_chat_test.nx -- SOVEREIGN CHAT HARNESS: the decisive coherence test. Qwen2.5-Instruct is trained
2// for the chat template; a raw prompt is out-of-distribution and drifts. This builds the REAL template with
3// special-token ids (<|im_start|>=151644, <|im_end|>=151645) -> "<|im_start|>user\n{msg}<|im_end|>\n
4// <|im_start|>assistant\n" -> prefills token-by-token (memory-safe) -> greedy-generates -> decodes (byte-level
5// unmap). If the engine is correct, this produces a coherent English answer. Proves engine-correct vs forward-bug.
6// expect_exit: 0 license_tier: ORIGINAL module: nishi-core.llm.chat
7import "nx_syscalls.nx"
8import "nx_tier.nx"
9import "nx_bpe.nx"
10import "nx_gguf.nx"
11import "nx_gguf_load.nx"
12import "nx_gguf_meta.nx"
13import "nx_f32.nx"
14import "nx_f32_kv_cache.nx"
15import "nx_f32_lazy_weight.nx"
16import "nx_f32_llama_block.nx"
17import "nx_f32_llama_block_v4.nx"
18import "nx_f32_llama_stack_v4.nx"
19import "nx_f32_llama_layer_lazy_load.nx"
20import "nx_f32_llm.nx"
21import "nx_f32_llm_v4.nx"
22import "nx_f32_llm_read_dims.nx"
23import "nx_f32_bpe_load.nx"
24import "nx_f32_llm_special_tokens.nx"
25
26func pr_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
27func pr_num(v: i64) -> i64 {
28 let bb: *u8=sys_mmap(28); var m: i64=v; let t: *u8=sys_mmap(28); var k: i64=0; if m==0 { t[0]=48 as u8; k=1 }
29 while m>0 { t[k]=(48+(m%10)) as u8; m=m/10; k=k+1 } var i: i64=0; while i<k { bb[i]=t[k-1-i]; i=i+1 } sys_write(1,bb,k); return 0
30}
31// append byte-level-encoded text into toks at off; return count added.
32func ch_enc(vocab: *NxBpeVocab, text: *u8, tlen: nx_int, toks: *i64, off: nx_int) -> nx_int {
33 let tmp: *i64 = sys_mmap(256 * 8) as *i64
34 let cnt: nx_int = nx_bpe_encode_bytelevel(vocab, text, tlen, tmp)
35 var i: nx_int = 0
36 while i < cnt { toks[off + i] = tmp[i]; i = i + 1 }
37 return cnt
38}
39func ch_argmax(logits: *i64, vocab: nx_int) -> nx_int {
40 var best: nx_int = 0
41 var bestv: i64 = logits[0]
42 var i: nx_int = 1
43 while i < vocab {
44 if nx_f32_gt(logits[i], bestv) == 1 { bestv = logits[i]; best = i }
45 i = i + 1
46 }
47 return best
48}
49
50func main() -> i64 {
51 let path: *u8 = "/tmp/nx_real_model.gguf" as *u8
52 let len_out: *i64 = sys_mmap(8) as *i64
53 let buf: *u8 = sys_read_file(path, len_out)
54 if buf == (0 as *u8) { return 10 }
55 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
56 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { return 20 }
57 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
58 let out_err: *i64 = sys_mmap(8) as *i64
59 if nx_f32_llm_read_dims_from_gguf(buf, len_out[0], hdr, model, out_err) != NX_FLD_OK { return 30 }
60 if nx_f32_llm_load_weights_v4_from_gguf(buf, hdr, model, out_err) != NX_FLV4_OK { return 40 }
61 let vocab: *NxBpeVocab = nx_bpe_vocab_new(67108864, 262144, 524288)
62 let nt2: *i64 = sys_mmap(8) as *i64
63 let nm: *i64 = sys_mmap(8) as *i64
64 if nx_f32_bpe_load_from_gguf(buf, len_out[0], hdr, vocab, nt2, nm, out_err) != NX_FBL_OK { return 50 }
65 let eos: nx_int = nx_f32_llm_read_eos(buf, len_out[0], hdr)
66 pr_puts("DIMS n_heads=" as *u8); pr_num(model.n_heads as i64)
67 pr_puts(" n_kv_heads=" as *u8); pr_num(model.n_kv_heads as i64)
68 pr_puts(" head_dim=" as *u8); pr_num(model.head_dim as i64)
69 pr_puts(" hidden=" as *u8); pr_num(model.hidden_dim as i64)
70 pr_puts(" eos=" as *u8); pr_num(eos as i64); pr_puts("\n" as *u8)
71
72 let eps: i64 = 0x358637BD
73 let attn_scale: i64 = 0x3E000000
74 let rope_base: i64 = 0x415D0EAB
75
76 // Build "<|im_start|>user\n{msg}<|im_end|>\n<|im_start|>assistant\n"
77 let toks: *i64 = sys_mmap(512 * 8) as *i64
78 var nt: nx_int = 0
79 toks[nt] = 151644; nt = nt + 1
80 nt = nt + ch_enc(vocab, "user\n" as *u8, 5, toks, nt)
81 nt = nt + ch_enc(vocab, "What is the capital of France?" as *u8, 30, toks, nt)
82 toks[nt] = 151645; nt = nt + 1
83 nt = nt + ch_enc(vocab, "\n" as *u8, 1, toks, nt)
84 toks[nt] = 151644; nt = nt + 1
85 nt = nt + ch_enc(vocab, "assistant\n" as *u8, 10, toks, nt)
86
87 pr_puts("template tokens (n=" as *u8); pr_num(nt as i64); pr_puts("): " as *u8)
88 var pt: nx_int = 0
89 while pt < nt { pr_puts(" " as *u8); pr_num(toks[pt]); pt = pt + 1 }
90 pr_puts("\n" as *u8)
91
92 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(model.n_layers, model.n_kv_heads, 128, model.head_dim)
93 let logits: *i64 = sys_mmap(model.vocab_size * 8) as *i64
94
95 // Prefill token-by-token (memory-safe).
96 var pf: nx_int = 0
97 while pf < nt {
98 let one: *i64 = (((toks as i64) + pf * 8)) as *i64
99 if nx_f32_llm_forward_v4(model, one, 1, cache, eps, attn_scale, rope_base, 1, logits) != NX_FLV4_OK { return 60 }
100 pf = pf + 1
101 }
102
103 // Greedy generate up to 24 tokens; stop at <|im_end|> or eos.
104 let outbuf: *u8 = sys_mmap(2048)
105 var no: nx_int = 0
106 let nb1: *i64 = sys_mmap(8) as *i64
107 let db: *u8 = sys_mmap(64)
108 var step: nx_int = 0
109 while step < 24 {
110 let nid: nx_int = ch_argmax(logits, model.vocab_size)
111 if nid == 151645 { step = 24 } else {
112 if eos >= 0 { if nid == eos { step = 24 } }
113 }
114 if step < 24 {
115 nb1[0] = nid as i64
116 let dn2: nx_int = nx_bpe_decode_bytelevel(vocab, nb1, 1, db)
117 var bi: nx_int = 0
118 while bi < dn2 { if no < 2040 { outbuf[no] = db[bi]; no = no + 1 } bi = bi + 1 }
119 if nx_f32_llm_forward_v4(model, nb1, 1, cache, eps, attn_scale, rope_base, 1, logits) != NX_FLV4_OK { return 70 }
120 step = step + 1
121 }
122 }
123
124 pr_puts("USER: What is the capital of France?\nASSISTANT: " as *u8)
125 sys_write(1, outbuf, no as i64)
126 pr_puts("\n" as *u8)
127 return no as i64
128}