nx_paged_fwd_gate.nx source
↩ module page · 266 lines · 10974 B
1// nx_paged_fwd_gate.nx -- MEASURED gate for the PAGED forward chain
2// (nx_f32_llama_v4p) on the REAL model, and the payoff it exists for:
3// PREFIX-SHARED sampling (prefill once, fork N).
4//
5// EQUIV-G paged forward == contiguous forward: chunked prefill + 6
6// greedy tokens IDENTICAL (end-to-end: embed/rope/attn-paged/
7// COW/lm_head)
8// FORK-SC self-consistency N=3 (T=0.8, top-k 40, fixed seeds):
9// prefill ONCE on seq0 -> fork per sample vs the baseline
10// reset+re-prefill loop on the contiguous cache -- generated
11// BYTES BIT-IDENTICAL per sample
12// SPEED both loops timed: the fork path skips N-1 chunked prefills
13// HYGIENE freeing forks + root returns the pool to all-free
14//
15// license_tier: ORIGINAL expect_exit: 0
16
17import "nx_syscalls.nx"
18import "nx_tier.nx"
19import "nx_le.nx"
20import "nx_bpe.nx"
21import "nx_gguf.nx"
22import "nx_gguf_load.nx"
23import "nx_gguf_meta.nx"
24import "nx_f32.nx"
25import "nx_f32_kv_cache.nx"
26import "nx_f32_lazy_weight.nx"
27import "nx_f32_llama_block.nx"
28import "nx_f32_llama_block_v4.nx"
29import "nx_f32_llama_stack_v4.nx"
30import "nx_f32_llama_layer_lazy_load.nx"
31import "nx_f32_llm.nx"
32import "nx_f32_llm_v4.nx"
33import "nx_f32_llm_read_dims.nx"
34import "nx_f32_bpe_load.nx"
35import "nx_f32_llm_special_tokens.nx"
36import "nx_f32_sampler.nx"
37import "nx_prng.nx"
38import "nx_reasoning.nx"
39import "nx_kvcache.nx"
40import "nx_f32_attn_paged.nx"
41import "nx_f32_llama_v4p.nx"
42
43func fg_w(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 }
44func fg_wn(v: i64) -> i64 {
45 var m: i64 = v
46 if m < 0 { fg_w("-" as *u8); m = 0 - m }
47 let t: *u8 = sys_mmap(28)
48 var k: i64 = 0
49 if m == 0 { t[0] = 48 as u8; k = 1 }
50 while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
51 let o: *u8 = sys_mmap(28)
52 var i: i64 = 0
53 while i < k { o[i] = t[k - 1 - i]; i = i + 1 }
54 sys_write(1, o, k)
55 return 0
56}
57func fg_beq(a: *u8, b: *u8, n: i64) -> i64 {
58 var i: i64 = 0
59 while i < n { if a[i] != b[i] { return 0 } i = i + 1 }
60 return 1
61}
62
63// unified one-token step: mode 0 = contiguous cache, 1 = paged seq.
64func fg_step(mode: i64, h: i64, model: *NxF32LlamaModel, one: *i64,
65 logits: *i64, eps: i64, scale: i64, rope: i64) -> nx_int {
66 if mode == 0 {
67 return nx_f32_llm_forward_v4(model, one, 1, h as *NxF32KVCache,
68 eps, scale, rope, 1, logits)
69 }
70 return nx_f32_llm_forward_v4p(model, one, 1, h as *NxPagedSeq,
71 eps, scale, rope, 1, logits)
72}
73
74// sampled decode loop from post-prefill logits (lastrow pre-copied into
75// logits row 0); emits decoded bytes; stops at im_end/eos. 9 args.
76func fg_sample_loop(mode: i64, h: i64, rc: *NxReasonCfg, logits: *i64,
77 seed: i64, max_new: nx_int, out_bytes: *u8,
78 out_cap: nx_int, one: *i64) -> nx_int {
79 let model: *NxF32LlamaModel = rc.model
80 let prng: *i64 = sys_mmap(8) as *i64
81 nx_prng_init(prng, seed)
82 let db: *u8 = sys_mmap(64)
83 var no: nx_int = 0
84 var step: nx_int = 0
85 while step < max_new {
86 let nid: nx_int = nx_f32_sampler_sample_top_k(logits, model.vocab_size,
87 rc.top_k, rc.inv_temp_f32, prng)
88 var stop: nx_int = 0
89 if nid == rc.im_end { stop = 1 }
90 if rc.eos >= 0 { if nid == rc.eos { stop = 1 } }
91 if stop == 1 { return no }
92 one[0] = nid as i64
93 let dn: nx_int = nx_bpe_decode_bytelevel(rc.vocab, one, 1, db)
94 var bi: nx_int = 0
95 while bi < dn {
96 if no < out_cap { out_bytes[no] = db[bi]; no = no + 1 }
97 bi = bi + 1
98 }
99 if fg_step(mode, h, model, one, logits, rc.eps, rc.attn_scale,
100 rc.rope_log_base) != NX_FLV4_OK { return 0 - 1 }
101 step = step + 1
102 }
103 return no
104}
105
106func main() -> i64 {
107 // ---- load the real model once ----------------------------------
108 let path: *u8 = "/tmp/nx_real_model.gguf" as *u8
109 let len_out: *i64 = sys_mmap(8) as *i64
110 let buf: *u8 = sys_read_file(path, len_out)
111 if buf == (0 as *u8) { return 10 }
112 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
113 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { return 20 }
114 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
115 let out_err: *i64 = sys_mmap(8) as *i64
116 if nx_f32_llm_read_dims_from_gguf(buf, len_out[0], hdr, model, out_err) != NX_FLD_OK { return 30 }
117 if nx_f32_llm_load_weights_v4_from_gguf(buf, hdr, model, out_err) != NX_FLV4_OK { return 40 }
118 let vocab: *NxBpeVocab = nx_bpe_vocab_new(67108864, 262144, 524288)
119 let nt2: *i64 = sys_mmap(8) as *i64
120 let nm: *i64 = sys_mmap(8) as *i64
121 if nx_f32_bpe_load_from_gguf(buf, len_out[0], hdr, vocab, nt2, nm, out_err) != NX_FBL_OK { return 50 }
122 let eos: nx_int = nx_f32_llm_read_eos(buf, len_out[0], hdr)
123
124 let rc: *NxReasonCfg = nx_reason_cfg_alloc()
125 rc.model = model
126 rc.vocab = vocab
127 rc.cache = nx_f32_kv_cache_alloc(model.n_layers, model.n_kv_heads, 128, model.head_dim)
128 rc.max_new = 8
129 rc.inv_temp_f32 = 0x3FA00000
130 rc.top_k = 40
131 rc.eps = 0x358637BD
132 rc.attn_scale = 0x3E000000
133 rc.rope_log_base = 0x415D0EAB
134 rc.eos = eos
135 rc.im_start = 151644
136 rc.im_end = 151645
137
138 let kv_dim: nx_int = model.n_kv_heads * model.head_dim
139 let pool: *NxPagedPool = nx_pkv_pool_new(16, model.n_layers, kv_dim)
140
141 let q1: *u8 = "What is 47 plus 38? Answer with only the number." as *u8
142 let toks: *i64 = sys_mmap(512 * 8) as *i64
143 let ntq: nx_int = nx_reason_build_chat_toks(rc, q1, 49, toks)
144
145 let vs: nx_int = model.vocab_size
146 let lgC: *i64 = sys_mmap((ntq + 2) * vs * 8) as *i64
147 let lgP: *i64 = sys_mmap((ntq + 2) * vs * 8) as *i64
148 let one: *i64 = sys_mmap(8) as *i64
149
150 // ---- EQUIV-G: chunked prefill + 6 greedy, both paths -------------
151 nx_f32_kv_cache_reset(rc.cache)
152 if nx_f32_llm_forward_v4(model, toks, ntq, rc.cache, rc.eps, rc.attn_scale,
153 rc.rope_log_base, 1, lgC) != NX_FLV4_OK { return 60 }
154 let seqE: *NxPagedSeq = nx_pkv_seq_new(pool, 64)
155 if nx_f32_llm_forward_v4p(model, toks, ntq, seqE, rc.eps, rc.attn_scale,
156 rc.rope_log_base, 1, lgP) != NX_FLV4_OK { return 61 }
157 var gi: nx_int = 0
158 var tC: nx_int = 0
159 var tP: nx_int = 0
160 let rowC: *i64 = ((lgC as i64) + (ntq - 1) * vs * 8) as *i64
161 let rowP: *i64 = ((lgP as i64) + (ntq - 1) * vs * 8) as *i64
162 tC = nx_f32_sampler_argmax(rowC, vs)
163 tP = nx_f32_sampler_argmax(rowP, vs)
164 while gi < 6 {
165 if tC != tP {
166 fg_w("EQUIV-G MISMATCH at " as *u8); fg_wn(gi as i64)
167 fg_w(" contig=" as *u8); fg_wn(tC as i64)
168 fg_w(" paged=" as *u8); fg_wn(tP as i64); fg_w("\n" as *u8)
169 return 62
170 }
171 one[0] = tC as i64
172 if fg_step(0, rc.cache as i64, model, one, lgC, rc.eps, rc.attn_scale, rc.rope_log_base) != NX_FLV4_OK { return 63 }
173 if fg_step(1, seqE as i64, model, one, lgP, rc.eps, rc.attn_scale, rc.rope_log_base) != NX_FLV4_OK { return 64 }
174 tC = nx_f32_sampler_argmax(lgC, vs)
175 tP = nx_f32_sampler_argmax(lgP, vs)
176 gi = gi + 1
177 }
178 nx_pkv_seq_free(seqE)
179 fg_w("PFG EQUIV-G paged forward == contiguous forward (prefill + 6 greedy) OK\n" as *u8)
180
181 // ---- FORK-SC: prefill once + fork N vs reset loop --------------
182 let N: nx_int = 3
183 let lastrow_sz: i64 = vs * 8
184
185 // paged: ONE chunked prefill on seq0, snapshot last-row logits.
186 let tp0: i64 = sys_now_us()
187 let seq0: *NxPagedSeq = nx_pkv_seq_new(pool, 64)
188 if nx_f32_llm_forward_v4p(model, toks, ntq, seq0, rc.eps, rc.attn_scale,
189 rc.rope_log_base, 1, lgP) != NX_FLV4_OK { return 70 }
190 let base_row: *i64 = sys_mmap(lastrow_sz) as *i64
191 var bc: nx_int = 0
192 let srcrow: *i64 = ((lgP as i64) + (ntq - 1) * vs * 8) as *i64
193 while bc < vs { base_row[bc] = srcrow[bc]; bc = bc + 1 }
194 let txP: *u8 = sys_mmap(3 * 96)
195 let lnP: *i64 = sys_mmap(3 * 8) as *i64
196 var s: nx_int = 0
197 while s < N {
198 let fseq: *NxPagedSeq = nx_pkv_seq_fork(seq0)
199 var cw: nx_int = 0
200 while cw < vs { lgP[cw] = base_row[cw]; cw = cw + 1 }
201 let ob: *u8 = ((txP as i64) + s * 96) as *u8
202 let nb: nx_int = fg_sample_loop(1, fseq as i64, rc, lgP,
203 20260709 + (s as i64) * 7919, 8, ob, 96, one)
204 if nb < 0 { return 71 }
205 lnP[s] = nb as i64
206 nx_pkv_seq_free(fseq)
207 s = s + 1
208 }
209 let usP: i64 = sys_now_us() - tp0
210
211 // baseline: reset + full chunked re-prefill per sample.
212 let tb0: i64 = sys_now_us()
213 let txC: *u8 = sys_mmap(3 * 96)
214 let lnC: *i64 = sys_mmap(3 * 8) as *i64
215 var s2: nx_int = 0
216 while s2 < N {
217 nx_f32_kv_cache_reset(rc.cache)
218 if nx_f32_llm_forward_v4(model, toks, ntq, rc.cache, rc.eps, rc.attn_scale,
219 rc.rope_log_base, 1, lgC) != NX_FLV4_OK { return 72 }
220 let rowB: *i64 = ((lgC as i64) + (ntq - 1) * vs * 8) as *i64
221 var cw2: nx_int = 0
222 while cw2 < vs { lgC[cw2] = rowB[cw2]; cw2 = cw2 + 1 }
223 let ob2: *u8 = ((txC as i64) + s2 * 96) as *u8
224 let nb2: nx_int = fg_sample_loop(0, rc.cache as i64, rc, lgC,
225 20260709 + (s2 as i64) * 7919, 8, ob2, 96, one)
226 if nb2 < 0 { return 73 }
227 lnC[s2] = nb2 as i64
228 s2 = s2 + 1
229 }
230 let usC: i64 = sys_now_us() - tb0
231
232 // compare texts bit-identical per sample.
233 var s3: nx_int = 0
234 while s3 < N {
235 if lnP[s3] != lnC[s3] {
236 fg_w("FORK-SC LEN MISMATCH s=" as *u8); fg_wn(s3 as i64); fg_w("\n" as *u8)
237 return 74
238 }
239 let pa: *u8 = ((txP as i64) + s3 * 96) as *u8
240 let pb: *u8 = ((txC as i64) + s3 * 96) as *u8
241 if fg_beq(pa, pb, lnP[s3]) != 1 {
242 fg_w("FORK-SC BYTES MISMATCH s=" as *u8); fg_wn(s3 as i64); fg_w("\n" as *u8)
243 return 75
244 }
245 s3 = s3 + 1
246 }
247 fg_w("PFG FORK-SC forked samples == reset-loop samples BIT-IDENTICAL (N=3) OK\n" as *u8)
248 fg_w("SPEED forked_us=" as *u8); fg_wn(usP)
249 fg_w(" resetloop_us=" as *u8); fg_wn(usC)
250 var usP2: i64 = usP
251 if usP2 < 1 { usP2 = 1 }
252 fg_w(" resetloop_vs_forked_x100=" as *u8); fg_wn(usC * 100 / usP2)
253 fg_w("\n" as *u8)
254
255 // ---- HYGIENE: pool returns to all-free ----------------------------
256 nx_pkv_seq_free(seq0)
257 if pool.n_free != 16 {
258 fg_w("HYGIENE pool n_free=" as *u8); fg_wn(pool.n_free); fg_w(" want 16\n" as *u8)
259 return 80
260 }
261 fg_w("PFG HYGIENE pool all-free after seq frees OK\n" as *u8)
262
263 fg_w("LIAR-KILL equiv-forward=1 fork-bitident=1 hygiene=1\n" as *u8)
264 fg_w("PAGED_FWD_GATE DONE\n" as *u8)
265 return 0
266}