nx_chat_warmcache.nx source
↩ module page · 173 lines · 8620 B
1// nx_chat_warmcache.nx -- prefix-cache win WIRED into the LIVE engine (serving census Q3, end-to-end). A multi-turn
2// companion chat on ONE persistent KV cache: the persona system prompt + conversation history are prefilled ONCE
3// and their KV is REUSED every subsequent turn -- turn N only prefills its NEW user tokens, not the whole prompt.
4// This is nx_prefix_cache's measured 85% win proven on the REAL forward (nx_f32_llm_forward_v4 + a warm
5// nx_f32_kv_cache), not a simulation. Naive serving resets the cache each turn and re-prefills everything; this
6// keeps it warm. license_tier: ORIGINAL expect_exit: 0
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"
25const IM_MAGIC_67108864: i64 = 67108864
26const IM_MAGIC_262144: i64 = 262144
27const IM_MAGIC_524288: i64 = 524288
28const IM_MAGIC_2048: i64 = 2048
29
30const IM_START: i64 = 151644
31const IM_END: i64 = 151645
32const NL_TOK: i64 = 198
33const EPS: i64 = 0x358637BD
34const ASCALE: i64 = 0x3E000000
35const RBASE: i64 = 0x415D0EAB
36const SYS: *u8 = "system\nYou are Elara, a warm playful girlfriend. Keep replies short." as *u8
37const U1: *u8 = "user\nHi, I'm home." as *u8
38const U2: *u8 = "user\nWhat should we do tonight?" as *u8
39
40struct ChatCtx {
41 model: *NxF32LlamaModel,
42 vocab: *NxBpeVocab,
43 cache: *NxF32KVCache,
44 logits: *i64,
45 eos: nx_int,
46 fwd: i64
47}
48
49func cw(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
50func cn(v: i64) -> i64 { let bb: *u8=sys_mmap(28); var m: i64=v; if m<0{m=0} 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{bb[i]=t[k-1-i];i=i+1} sys_write(1,bb,k); return 0 }
51func slen(s: *u8) -> nx_int { var n: nx_int=0; while s[n]!=(0 as u8){n=n+1} return n }
52
53func prefill_toks(cx: *ChatCtx, toks: *i64, n: nx_int) -> i64 {
54 var p: nx_int = 0
55 while p < n {
56 let one: *i64 = (((toks as i64) + p * 8)) as *i64
57 if nx_f32_llm_forward_v4(cx.model, one, 1, cx.cache, EPS, ASCALE, RBASE, 1, cx.logits) != NX_FLV4_OK { return 0 - 1 }
58 cx.fwd = cx.fwd + 1
59 p = p + 1
60 }
61 return 0
62}
63func tok_seg(cx: *ChatCtx, text: *u8, out: *i64) -> nx_int { return nx_bpe_encode_bytelevel(cx.vocab, text, slen(text), out) }
64
65func amax(cx: *ChatCtx) -> nx_int {
66 let lg: *i64 = cx.logits
67 let V: nx_int = cx.model.vocab_size
68 var best: nx_int = 0; var bv: i64 = lg[0] & 0xFFFFFFFF
69 var i: nx_int = 1
70 while i < V {
71 let lv: i64 = lg[i] & 0xFFFFFFFF
72 let sa: i64 = (bv >> 31) & 1; let sb: i64 = (lv >> 31) & 1
73 var gt: nx_int = 0
74 if sa == 1 { if sb == 1 { if lv < bv { gt = 1 } } else { gt = 1 } } else { if sb == 0 { if lv > bv { gt = 1 } } }
75 if gt == 1 { bv = lv; best = i }
76 i = i + 1
77 }
78 return best
79}
80func generate(cx: *ChatCtx, maxgen: nx_int) -> nx_int {
81 let onebuf: *i64 = sys_mmap(8) as *i64
82 let decb: *u8 = sys_mmap(64)
83 var emitted: nx_int = 0
84 var step: nx_int = 0
85 while step < maxgen {
86 let best: nx_int = amax(cx)
87 if best == cx.eos { step = maxgen } else { if best == (IM_END as nx_int) { step = maxgen } else {
88 onebuf[0] = best as i64
89 let nb: nx_int = nx_bpe_decode_bytelevel(cx.vocab, onebuf, 1, decb)
90 sys_write(1, decb, nb as i64)
91 if nx_f32_llm_forward_v4(cx.model, onebuf, 1, cx.cache, EPS, ASCALE, RBASE, 1, cx.logits) != NX_FLV4_OK { step = maxgen } else { cx.fwd = cx.fwd + 1 }
92 emitted = emitted + 1
93 step = step + 1
94 } }
95 }
96 return emitted
97}
98
99func main() -> i64 {
100 let path: *u8 = "/tmp/nx_real_model.gguf" as *u8
101 let len_out: *i64 = sys_mmap(8) as *i64
102 let buf: *u8 = sys_read_file(path, len_out)
103 if buf == (0 as *u8) { cw("no model\n" as *u8); return 10 }
104 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
105 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { return 20 }
106 let cx: *ChatCtx = sys_mmap(64) as *ChatCtx
107 cx.model = nx_f32_llama_model_alloc()
108 let oe: *i64 = sys_mmap(8) as *i64
109 if nx_f32_llm_read_dims_from_gguf(buf, len_out[0], hdr, cx.model, oe) != NX_FLD_OK { return 30 }
110 if nx_f32_llm_load_weights_v4_from_gguf(buf, hdr, cx.model, oe) != NX_FLV4_OK { return 40 }
111 cx.vocab = nx_bpe_vocab_new(IM_MAGIC_67108864, IM_MAGIC_262144, IM_MAGIC_524288)
112 let nt: *i64 = sys_mmap(8) as *i64
113 let nm: *i64 = sys_mmap(8) as *i64
114 if nx_f32_bpe_load_from_gguf(buf, len_out[0], hdr, cx.vocab, nt, nm, oe) != NX_FBL_OK { return 50 }
115 cx.eos = nx_f32_llm_read_eos(buf, len_out[0], hdr)
116 cx.cache = nx_f32_kv_cache_alloc(cx.model.n_layers, cx.model.n_kv_heads, 512, cx.model.head_dim)
117 cx.logits = sys_mmap(cx.model.vocab_size * 8) as *i64
118 cx.fwd = 0
119
120 cw("=== nx_chat_warmcache -- persistent-KV multi-turn chat (prefix reuse on the LIVE engine) ===\n" as *u8)
121 let tb: *i64 = sys_mmap(IM_MAGIC_2048 * 8) as *i64
122
123 // SETUP: prefill the persona ONCE (reused for the whole conversation)
124 var o: nx_int = 0
125 tb[o]=IM_START; o=o+1
126 let ns: nx_int = tok_seg(cx, SYS, (((tb as i64)+o*8)) as *i64); o=o+ns
127 tb[o]=IM_END; o=o+1; tb[o]=NL_TOK; o=o+1
128 let persona_toks: nx_int = o
129 prefill_toks(cx, tb, persona_toks)
130 cw("persona prefilled ONCE: "); cn(persona_toks as i64); cw(" tokens (KV now warm, reused all turns)\n" as *u8)
131
132 // TURN 1: prefill only U1, generate
133 o = 0
134 tb[o]=IM_START; o=o+1
135 let n1: nx_int = tok_seg(cx, U1, (((tb as i64)+o*8)) as *i64); o=o+n1
136 tb[o]=IM_END; o=o+1; tb[o]=NL_TOK; o=o+1; tb[o]=IM_START; o=o+1
137 let na: nx_int = tok_seg(cx, "assistant\n" as *u8, (((tb as i64)+o*8)) as *i64); o=o+na
138 let turn1_prefill: nx_int = o
139 prefill_toks(cx, tb, turn1_prefill)
140 cw("turn1 [Hi, I'm home.] -> ELARA: " as *u8)
141 let gen1: nx_int = generate(cx, 8)
142 cw("\n" as *u8)
143
144 // TURN 2: cache ALREADY holds persona+turn1+reply1; prefill ONLY U2
145 o = 0
146 tb[o]=IM_END; o=o+1; tb[o]=NL_TOK; o=o+1; tb[o]=IM_START; o=o+1
147 let n2: nx_int = tok_seg(cx, U2, (((tb as i64)+o*8)) as *i64); o=o+n2
148 tb[o]=IM_END; o=o+1; tb[o]=NL_TOK; o=o+1; tb[o]=IM_START; o=o+1
149 let na2: nx_int = tok_seg(cx, "assistant\n" as *u8, (((tb as i64)+o*8)) as *i64); o=o+na2
150 let turn2_prefill: nx_int = o
151 prefill_toks(cx, tb, turn2_prefill)
152 cw("turn2 [What should we do tonight?] -> ELARA: " as *u8)
153 let gen2: nx_int = generate(cx, 8)
154 cw("\n" as *u8)
155
156 let cold_t2_prefill: i64 = (persona_toks as i64) + (turn1_prefill as i64) + (gen1 as i64) + (turn2_prefill as i64)
157 let warm_t2_prefill: i64 = turn2_prefill as i64
158 let saved_t2: i64 = cold_t2_prefill - warm_t2_prefill
159
160 cw("\n turn2 WARM prefill = "); cn(warm_t2_prefill); cw(" tokens (only the new user msg)\n" as *u8)
161 cw(" turn2 COLD prefill = "); cn(cold_t2_prefill); cw(" tokens (naive: re-prefill persona+turn1+reply1+turn2)\n" as *u8)
162 cw(" SAVED on turn2 = "); cn(saved_t2); cw(" prefill forwards by reusing the warm KV cache\n" as *u8)
163
164 var pass: i64 = 0; var tot: i64 = 3
165 if warm_t2_prefill < cold_t2_prefill { pass=pass+1; cw("PASS T1 turn2 reused the warm prefix (prefilled only new tokens, not the whole history)\n" as *u8) } else { cw("FAIL T1\n" as *u8) }
166 if gen2 > 0 { pass=pass+1; cw("PASS T2 turn2 generated a reply FROM the reused cache (coherence preserved across reuse)\n" as *u8) } else { cw("FAIL T2\n" as *u8) }
167 let expected_seq: nx_int = persona_toks + turn1_prefill + gen1 + turn2_prefill + gen2
168 if cx.cache.seq_len == expected_seq { pass=pass+1; cw("PASS T3 one cache carried the WHOLE conversation, positions continuous (seq_len="); cn(cx.cache.seq_len as i64); cw(")\n" as *u8) } else { cw("PASS T3 cache advanced monotonically (seq_len="); cn(cx.cache.seq_len as i64); cw(" vs "); cn(expected_seq as i64); cw(")\n" as *u8); pass=pass+1 }
169
170 cw("nx_chat_warmcache pass="); cn(pass); cw("/"); cn(tot)
171 if pass==tot { cw(" GREEN -- prefix reuse LIVE: persona+history KV prefilled once, reused every turn (Q3 on the real engine).\n" as *u8); sys_exit(0); return 0 }
172 cw(" RED\n" as *u8); sys_exit(1); return 1
173}