nx_reasoning.nx source
↩ module page · 322 lines · 12663 B
1// nx_reasoning.nx -- sovereign TEST-TIME COMPUTE organ (self-consistency).
2//
3// The no.1 2025/26 frontier gap on the LLM census (nx_swcompare_gapmap llm,
4// momentum=11): spend MORE inference compute per question to get a BETTER
5// answer. This organ implements the foundational primitive:
6//
7// Self-consistency majority vote (Wang et al. 2023, ICLR): sample N
8// diverse completions at temperature T, extract each final answer,
9// return the mode. Accuracy rises with N on tasks where the error
10// mode is "diverse wrong answers, consistent right answer".
11//
12// Generation goes through the CHATML TEMPLATE at the TOKEN level
13// (<|im_start|> injected as ids -- the nx_f32_llm_chat_test.nx proven
14// path). MEASURED 2026-07-09: Qwen2.5-Instruct on a RAW continuation
15// prompt ("47 + 38 =") emits garbage even with correct digit-split
16// tokenization (verified by nx_tok_probe); the instruct model is only
17// in-distribution INSIDE its template. Composes the EXISTING sampler
18// (Hinton-2015 temperature + Fan-2018 top-k) + forward_v4 -- ZERO new
19// model machinery. The answers array is the composition point for
20// future best-of-N verifier re-ranking.
21//
22// genealogy_id: wang_2023_self_consistency
23// lineage_id: substrate_reasoning_v2 (v1 = raw-continuation, measured-hostile)
24
25import "nx_syscalls.nx"
26import "nx_tier.nx"
27import "nx_prng.nx"
28import "nx_bpe.nx"
29import "nx_f32_kv_cache.nx"
30import "nx_f32_llm.nx"
31import "nx_f32_llm_v4.nx"
32import "nx_f32_sampler.nx"
33
34// Sentinel: no integer found in the generated text.
35const NX_REASON_NONE: i64 = 0 - 999999999
36
37// ===== Config bundle (keeps every call far under the 16-arg cap) =====
38
39struct NxReasonCfg {
40 model: *NxF32LlamaModel,
41 vocab: *NxBpeVocab,
42 cache: *NxF32KVCache,
43 max_new: nx_int,
44 inv_temp_f32: i64, // 1/T f32 bits (sampled path only)
45 top_k: nx_int, // sampled path only
46 eps: i64,
47 attn_scale: i64,
48 rope_log_base: i64,
49 eos: nx_int,
50 im_start: nx_int, // ChatML <|im_start|> id (Qwen2: 151644)
51 im_end: nx_int // ChatML <|im_end|> id (Qwen2: 151645)
52}
53
54const NX_REASON_CFG_BYTES: nx_int = 96 // 12 fields * 8
55
56func nx_reason_cfg_alloc() -> *NxReasonCfg {
57 return sys_mmap(NX_REASON_CFG_BYTES) as *NxReasonCfg
58}
59
60// append byte-level-encoded text into toks at off; return count added.
61func _nxr_enc(vocab: *NxBpeVocab, text: *u8, tlen: nx_int, toks: *i64, off: nx_int) -> nx_int {
62 let tmp: *i64 = sys_mmap(256 * 8) as *i64
63 let cnt: nx_int = nx_bpe_encode_bytelevel(vocab, text, tlen, tmp)
64 var i: nx_int = 0
65 while i < cnt { toks[off + i] = tmp[i]; i = i + 1 }
66 return cnt
67}
68
69// ===== One CHAT generation (greedy or sampled) ======================
70// Builds "<|im_start|>user\n{q}<|im_end|>\n<|im_start|>assistant\n" as
71// TOKEN IDS, resets the KV cache, prefills token-by-token, then decodes
72// up to rc.max_new tokens. prng_state == null -> pure argmax (greedy,
73// deterministic). prng_state != null -> temperature + top-k sampling.
74// Stops at <|im_end|> or eos. Returns bytes written to out_bytes.
75
76// Build the ChatML token sequence for one user turn into toks; returns count.
77// Exposed so other test-time organs (nx_specdec) reuse the exact template.
78
79func nx_reason_build_chat_toks(rc: *NxReasonCfg, qtext: *u8, qlen: nx_int,
80 toks: *i64) -> nx_int {
81 var nt: nx_int = 0
82 toks[nt] = rc.im_start as i64; nt = nt + 1
83 let e1: nx_int = _nxr_enc(rc.vocab, "user\n" as *u8, 5, toks, nt)
84 nt = nt + e1
85 let e2: nx_int = _nxr_enc(rc.vocab, qtext, qlen, toks, nt)
86 nt = nt + e2
87 toks[nt] = rc.im_end as i64; nt = nt + 1
88 let e3: nx_int = _nxr_enc(rc.vocab, "\n" as *u8, 1, toks, nt)
89 nt = nt + e3
90 toks[nt] = rc.im_start as i64; nt = nt + 1
91 let e4: nx_int = _nxr_enc(rc.vocab, "assistant\n" as *u8, 10, toks, nt)
92 nt = nt + e4
93 return nt
94}
95
96func nx_reason_chat_gen(rc: *NxReasonCfg, qtext: *u8, qlen: nx_int,
97 prng_state: *i64, out_bytes: *u8,
98 out_cap: nx_int) -> nx_int {
99 nx_f32_kv_cache_reset(rc.cache)
100
101 let toks: *i64 = sys_mmap(512 * 8) as *i64
102 let nt: nx_int = nx_reason_build_chat_toks(rc, qtext, qlen, toks)
103
104 let model: *NxF32LlamaModel = rc.model
105 let logits: *i64 = sys_mmap((nt + 2) * model.vocab_size * 8) as *i64
106
107 // CHUNKED prefill: ONE m=nt forward instead of nt sequential m=1
108 // forwards (weight bytes read once for the whole prompt). Proven
109 // token-identical to sequential prefill by nx_specdec_gate EQUIV-1
110 // (2026-07-09); offset-causal positions per nx_f32_attn_multi.
111 if nx_f32_llm_forward_v4(model, toks, nt, rc.cache, rc.eps,
112 rc.attn_scale, rc.rope_log_base, 1,
113 logits) != NX_FLV4_OK { return 0 - 1 }
114 // decode loop reads row 0 of `logits`; move the last prefill row there.
115 let lrow: *i64 = ((logits as i64) + (nt - 1) * model.vocab_size * 8) as *i64
116 var mv: nx_int = 0
117 while mv < model.vocab_size { logits[mv] = lrow[mv]; mv = mv + 1 }
118
119 // Decode loop: sample or argmax, stop at im_end/eos, decode bytes.
120 var no: nx_int = 0
121 let nb1: *i64 = sys_mmap(8) as *i64
122 let db: *u8 = sys_mmap(64)
123 var step: nx_int = 0
124 while step < rc.max_new {
125 var nid: nx_int = 0
126 if prng_state == (0 as *i64) {
127 nid = nx_f32_sampler_argmax(logits, model.vocab_size)
128 } else {
129 nid = nx_f32_sampler_sample_top_k(logits, model.vocab_size,
130 rc.top_k, rc.inv_temp_f32,
131 prng_state)
132 }
133 if nid == rc.im_end { step = rc.max_new } else {
134 if rc.eos >= 0 { if nid == rc.eos { step = rc.max_new } }
135 }
136 if step < rc.max_new {
137 nb1[0] = nid as i64
138 let dn: nx_int = nx_bpe_decode_bytelevel(rc.vocab, nb1, 1, db)
139 var bi: nx_int = 0
140 while bi < dn {
141 if no < out_cap { out_bytes[no] = db[bi]; no = no + 1 }
142 bi = bi + 1
143 }
144 if nx_f32_llm_forward_v4(model, nb1, 1, rc.cache, rc.eps,
145 rc.attn_scale, rc.rope_log_base, 1,
146 logits) != NX_FLV4_OK { return 0 - 1 }
147 step = step + 1
148 }
149 }
150 return no
151}
152
153// ===== Answer extraction ===========================================
154// First integer in the generated text (digit run, up to 12 digits).
155// Mechanical + format-agnostic: works for "85", " 85.", "47 + 38 = 85".
156// NOTE for arithmetic echo-style answers ("47 + 38 = 85") use
157// nx_reason_extract_last_int instead -- the FIRST int is the echoed
158// operand, the LAST is the answer.
159
160func nx_reason_extract_int(text: *u8, n: nx_int) -> i64 {
161 var i: nx_int = 0
162 while i < n {
163 let c: i64 = text[i] as i64
164 if c >= 48 {
165 if c <= 57 {
166 var v: i64 = 0
167 var k: nx_int = i
168 var nd: i64 = 0
169 while k < n {
170 let d: i64 = text[k] as i64
171 if d < 48 { break }
172 if d > 57 { break }
173 if nd < 12 { v = v * 10 + (d - 48); nd = nd + 1 }
174 k = k + 1
175 }
176 return v
177 }
178 }
179 i = i + 1
180 }
181 return NX_REASON_NONE
182}
183
184// LAST integer in the text (the natural "final answer" position for
185// chat answers that echo the question, e.g. "47 + 38 = 85").
186
187func nx_reason_extract_last_int(text: *u8, n: nx_int) -> i64 {
188 var last: i64 = NX_REASON_NONE
189 var i: nx_int = 0
190 while i < n {
191 let c: i64 = text[i] as i64
192 var is_d: i64 = 0
193 if c >= 48 { if c <= 57 { is_d = 1 } }
194 if is_d == 1 {
195 // consume the whole digit run; record it as the latest int.
196 var v: i64 = 0
197 var nd: i64 = 0
198 while i < n {
199 let d: i64 = text[i] as i64
200 if d < 48 { break }
201 if d > 57 { break }
202 if nd < 12 { v = v * 10 + (d - 48); nd = nd + 1 }
203 i = i + 1
204 }
205 last = v
206 } else {
207 i = i + 1
208 }
209 }
210 return last
211}
212
213// ===== Majority vote (the self-consistency reducer) ================
214// Mode over answers, ignoring NX_REASON_NONE. Ties resolve to the
215// EARLIEST-seen answer (deterministic). All-sentinel -> sentinel.
216
217func nx_reason_majority(answers: *i64, n: nx_int) -> i64 {
218 var best: i64 = NX_REASON_NONE
219 var best_cnt: i64 = 0
220 var i: nx_int = 0
221 while i < n {
222 if answers[i] != NX_REASON_NONE {
223 var cnt: i64 = 0
224 var j: nx_int = 0
225 while j < n {
226 if answers[j] == answers[i] { cnt = cnt + 1 }
227 j = j + 1
228 }
229 if cnt > best_cnt { best_cnt = cnt; best = answers[i] }
230 }
231 i = i + 1
232 }
233 return best
234}
235
236// ===== Self-consistency: N chat samples -> majority ================
237// Runs n_samples stochastic generations (seed0 + s*7919, prime-spaced),
238// stores each text at texts + s*text_stride (lens_out[s] = byte count),
239// each LAST-int answer in answers_out[s]. Returns the majority answer.
240
241func nx_reason_selfconsist(rc: *NxReasonCfg, qtext: *u8, qlen: nx_int,
242 n_samples: nx_int, seed0: i64,
243 answers_out: *i64, texts: *u8,
244 text_stride: nx_int, lens_out: *i64) -> i64 {
245 var s: nx_int = 0
246 while s < n_samples {
247 let prng: *i64 = sys_mmap(8) as *i64
248 nx_prng_init(prng, seed0 + (s as i64) * 7919)
249 let ob: *u8 = ((texts as i64) + (s as i64) * (text_stride as i64)) as *u8
250 let nb: nx_int = nx_reason_chat_gen(rc, qtext, qlen, prng, ob, text_stride)
251 var nbc: i64 = nb as i64
252 if nbc < 0 { nbc = 0 }
253 lens_out[s] = nbc
254 let av: i64 = nx_reason_extract_last_int(ob, nbc as nx_int)
255 answers_out[s] = av
256 s = s + 1
257 }
258 return nx_reason_majority(answers_out, n_samples)
259}
260
261// ===== Question builders (data-driven) =============================
262
263func _nxr_put_u(buf: *u8, off: i64, v: i64) -> i64 {
264 let t: *u8 = sys_mmap(24)
265 var m: i64 = v
266 var k: i64 = 0
267 if m == 0 { t[0] = 48 as u8; k = 1 }
268 while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
269 var i: i64 = 0
270 while i < k { buf[off + i] = t[k - 1 - i]; i = i + 1 }
271 return off + k
272}
273
274func _nxr_put_s(buf: *u8, off: i64, s: *u8) -> i64 {
275 var o: i64 = off
276 var i: i64 = 0
277 while s[i] != (0 as u8) { buf[o] = s[i]; o = o + 1; i = i + 1 }
278 return o
279}
280
281// "What is A plus B?" (op_code 0=plus 1=minus 2=times) -- the natural
282// chat phrasing for an instruct model. Returns byte length.
283
284func nx_reason_build_q(buf: *u8, a: i64, op_code: i64, b: i64) -> nx_int {
285 var o: i64 = _nxr_put_s(buf, 0, "What is " as *u8)
286 o = _nxr_put_u(buf, o, a)
287 if op_code == 0 { o = _nxr_put_s(buf, o, " plus " as *u8) }
288 if op_code == 1 { o = _nxr_put_s(buf, o, " minus " as *u8) }
289 if op_code == 2 { o = _nxr_put_s(buf, o, " times " as *u8) }
290 o = _nxr_put_u(buf, o, b)
291 buf[o] = 63 as u8; o = o + 1 // '?'
292 return o as nx_int
293}
294
295// Short-answer variant: appends an only-the-number instruction so the
296// instruct model skips the narration chain ("To find the sum of 47 and
297// 38, you add them..." never reached the answer in 16 tokens -- measured
298// 2026-07-09). Keeps generations short AND the final answer extractable.
299
300func nx_reason_build_q_short(buf: *u8, a: i64, op_code: i64, b: i64) -> nx_int {
301 var o: i64 = nx_reason_build_q(buf, a, op_code, b) as i64
302 o = _nxr_put_s(buf, o, " Answer with only the number." as *u8)
303 return o as nx_int
304}
305
306// RAW continuation builder ("A <op> B =", no trailing space) -- kept
307// for BASE models. MEASURED-HOSTILE for Qwen2.5-INSTRUCT (2026-07-09):
308// off-template continuation emits garbage; use nx_reason_chat_gen.
309
310func nx_reason_build_arith_prompt(buf: *u8, a: i64, op_code: i64, b: i64) -> nx_int {
311 var o: i64 = _nxr_put_u(buf, 0, a)
312 buf[o] = 32 as u8; o = o + 1
313 var oc: i64 = 43 // '+'
314 if op_code == 1 { oc = 45 } // '-'
315 if op_code == 2 { oc = 42 } // '*'
316 buf[o] = oc as u8; o = o + 1
317 buf[o] = 32 as u8; o = o + 1
318 o = _nxr_put_u(buf, o, b)
319 buf[o] = 32 as u8; o = o + 1
320 buf[o] = 61 as u8; o = o + 1 // '=' (no trailing space)
321 return o as nx_int
322}