nx_reasoning_paged.nx source

1// nx_reasoning_paged.nx -- PREFIX-SHARED self-consistency: the paged-KV 2// promotion of nx_reason_selfconsist. Prefill the ChatML prompt ONCE on a 3// root sequence, then FORK per sample (refcounted page-table copy; the 4// fork's first append copy-on-writes at most one tail block). MEASURED 5// 1.48x vs the reset-loop at N=3 (nx_paged_fwd_gate, real model, bytes 6// BIT-IDENTICAL per sample); the win grows with N since N-1 chunked 7// prefills are eliminated. 8// 9// genealogy_id: wang_2023_self_consistency + kwon_2023_pagedattention 10// lineage_id: substrate_reasoning_paged_v1 11 12import "nx_syscalls.nx" 13import "nx_tier.nx" 14import "nx_prng.nx" 15import "nx_bpe.nx" 16import "nx_f32_llm.nx" 17import "nx_f32_llm_v4.nx" 18import "nx_f32_sampler.nx" 19import "nx_reasoning.nx" 20import "nx_kvcache.nx" 21import "nx_f32_attn_paged.nx" 22import "nx_f32_llama_v4p.nx" 23import "nx_f32_llama_v4b.nx" 24 25// one sampled decode from pre-copied last-row logits on a paged seq. 26func _nxrp_sample_loop(rc: *NxReasonCfg, seq: *NxPagedSeq, logits: *i64, 27 seed: i64, out_bytes: *u8, out_cap: nx_int) -> nx_int { 28 let model: *NxF32LlamaModel = rc.model 29 let prng: *i64 = sys_mmap(8) as *i64 30 nx_prng_init(prng, seed) 31 let one: *i64 = sys_mmap(8) as *i64 32 let db: *u8 = sys_mmap(64) 33 var no: nx_int = 0 34 var step: nx_int = 0 35 while step < rc.max_new { 36 let nid: nx_int = nx_f32_sampler_sample_top_k(logits, model.vocab_size, 37 rc.top_k, rc.inv_temp_f32, prng) 38 var stop: nx_int = 0 39 if nid == rc.im_end { stop = 1 } 40 if rc.eos >= 0 { if nid == rc.eos { stop = 1 } } 41 if stop == 1 { return no } 42 one[0] = nid as i64 43 let dn: nx_int = nx_bpe_decode_bytelevel(rc.vocab, one, 1, db) 44 var bi: nx_int = 0 45 while bi < dn { 46 if no < out_cap { out_bytes[no] = db[bi]; no = no + 1 } 47 bi = bi + 1 48 } 49 if nx_f32_llm_forward_v4p(model, one, 1, seq, rc.eps, rc.attn_scale, 50 rc.rope_log_base, 1, logits) != NX_FLV4_OK { return 0 - 1 } 51 step = step + 1 52 } 53 return no 54} 55 56// Prefix-shared self-consistency: N samples over ONE shared prompt prefill. 57// Same contract as nx_reason_selfconsist (answers/texts/lens per sample, 58// returns the majority); pool must hold ~(prompt_blocks + N*2) blocks. 59 60func nx_reason_selfconsist_paged(rc: *NxReasonCfg, pool: *NxPagedPool, 61 qtext: *u8, qlen: nx_int, 62 n_samples: nx_int, seed0: i64, 63 answers_out: *i64, texts: *u8, 64 text_stride: nx_int, lens_out: *i64) -> i64 { 65 let model: *NxF32LlamaModel = rc.model 66 let vs: nx_int = model.vocab_size 67 68 let toks: *i64 = sys_mmap(512 * 8) as *i64 69 let nt: nx_int = nx_reason_build_chat_toks(rc, qtext, qlen, toks) 70 71 // ONE chunked prefill on the root sequence. 72 let seq0: *NxPagedSeq = nx_pkv_seq_new(pool, nt + rc.max_new + 32) 73 let logits: *i64 = sys_mmap((nt + 2) * vs * 8) as *i64 74 if nx_f32_llm_forward_v4p(model, toks, nt, seq0, rc.eps, rc.attn_scale, 75 rc.rope_log_base, 1, logits) != NX_FLV4_OK { 76 return NX_REASON_NONE 77 } 78 let base_row: *i64 = sys_mmap(vs * 8) as *i64 79 let srcrow: *i64 = ((logits as i64) + (nt - 1) * vs * 8) as *i64 80 var bc: nx_int = 0 81 while bc < vs { base_row[bc] = srcrow[bc]; bc = bc + 1 } 82 83 var s: nx_int = 0 84 while s < n_samples { 85 let fseq: *NxPagedSeq = nx_pkv_seq_fork(seq0) 86 var cw: nx_int = 0 87 while cw < vs { logits[cw] = base_row[cw]; cw = cw + 1 } 88 let ob: *u8 = ((texts as i64) + (s as i64) * (text_stride as i64)) as *u8 89 let nb: nx_int = _nxrp_sample_loop(rc, fseq, logits, 90 seed0 + (s as i64) * 7919, ob, text_stride) 91 var nbc: i64 = nb as i64 92 if nbc < 0 { nbc = 0 } 93 lens_out[s] = nbc 94 let av: i64 = nx_reason_extract_last_int(ob, nbc as nx_int) 95 answers_out[s] = av 96 nx_pkv_seq_free(fseq) 97 s = s + 1 98 } 99 nx_pkv_seq_free(seq0) 100 sys_munmap(base_row, vs * 8) 101 sys_munmap(logits, (nt + 2) * vs * 8) 102 return nx_reason_majority(answers_out, n_samples) 103} 104 105// ===== BATCHED prefix-shared self-consistency ======================== 106// Same contract; the N forks decode in LOCKSTEP ROUNDS -- one M-row 107// nx_f32_llm_forward_v4b per round advances every active fork (weight 108// bytes read once per round, not once per fork). Outputs BIT-IDENTICAL 109// to the sequential variant (nx_batched_gate, N=3, decode 1.47x); on top 110// of the shared prefill the stacked win vs a reset loop is ~2.4x at N=3 111// and grows with N. Early-stopping forks drop out of the batch (ragged 112// compaction); pool needs ~(prompt_blocks + N*2) blocks. 113 114func nx_reason_selfconsist_batched(rc: *NxReasonCfg, pool: *NxPagedPool, 115 qtext: *u8, qlen: nx_int, 116 n_samples: nx_int, seed0: i64, 117 answers_out: *i64, texts: *u8, 118 text_stride: nx_int, lens_out: *i64) -> i64 { 119 let model: *NxF32LlamaModel = rc.model 120 let vs: nx_int = model.vocab_size 121 122 let toks: *i64 = sys_mmap(512 * 8) as *i64 123 let nt: nx_int = nx_reason_build_chat_toks(rc, qtext, qlen, toks) 124 125 // ONE chunked prefill on the root sequence. 126 let seq0: *NxPagedSeq = nx_pkv_seq_new(pool, nt + rc.max_new + 32) 127 let plg: *i64 = sys_mmap((nt + 2) * vs * 8) as *i64 128 if nx_f32_llm_forward_v4p(model, toks, nt, seq0, rc.eps, rc.attn_scale, 129 rc.rope_log_base, 1, plg) != NX_FLV4_OK { 130 return NX_REASON_NONE 131 } 132 let base_row: *i64 = sys_mmap(vs * 8) as *i64 133 let srcrow: *i64 = ((plg as i64) + (nt - 1) * vs * 8) as *i64 134 var bc: nx_int = 0 135 while bc < vs { base_row[bc] = srcrow[bc]; bc = bc + 1 } 136 sys_munmap(plg, (nt + 2) * vs * 8) 137 138 // per-fork state + first sample from the shared base row. 139 let fseqs: *i64 = sys_mmap(n_samples * 8) as *i64 140 let prngs: *i64 = sys_mmap(n_samples * 8) as *i64 141 let curs: *i64 = sys_mmap(n_samples * 8) as *i64 142 let done: *i64 = sys_mmap(n_samples * 8) as *i64 143 let emitted: *i64 = sys_mmap(n_samples * 8) as *i64 144 let one: *i64 = sys_mmap(8) as *i64 145 let db: *u8 = sys_mmap(64) 146 var f: nx_int = 0 147 while f < n_samples { 148 let fs: *NxPagedSeq = nx_pkv_seq_fork(seq0) 149 fseqs[f] = fs as i64 150 let pr: *i64 = sys_mmap(8) as *i64 151 nx_prng_init(pr, seed0 + (f as i64) * 7919) 152 prngs[f] = pr as i64 153 lens_out[f] = 0 154 done[f] = 0 155 emitted[f] = 0 156 let nid0: nx_int = nx_f32_sampler_sample_top_k(base_row, vs, rc.top_k, 157 rc.inv_temp_f32, pr) 158 var stop0: nx_int = 0 159 if nid0 == rc.im_end { stop0 = 1 } 160 if rc.eos >= 0 { if nid0 == rc.eos { stop0 = 1 } } 161 if stop0 == 1 { done[f] = 1 } else { 162 one[0] = nid0 as i64 163 let dn: nx_int = nx_bpe_decode_bytelevel(rc.vocab, one, 1, db) 164 let ob: *u8 = ((texts as i64) + (f as i64) * (text_stride as i64)) as *u8 165 var bi: nx_int = 0 166 var no: i64 = 0 167 while bi < dn { 168 if no < text_stride { ob[no] = db[bi]; no = no + 1 } 169 bi = bi + 1 170 } 171 lens_out[f] = no 172 curs[f] = nid0 as i64 173 emitted[f] = 1 174 } 175 f = f + 1 176 } 177 178 // lockstep rounds: ONE batched forward advances every active fork. 179 let bat_ids: *i64 = sys_mmap(n_samples * 8) as *i64 180 let bat_seqs: *i64 = sys_mmap(n_samples * 8) as *i64 181 let bat_map: *i64 = sys_mmap(n_samples * 8) as *i64 182 let blg: *i64 = sys_mmap(n_samples * vs * 8) as *i64 183 var running: nx_int = 1 184 while running == 1 { 185 var Mb: nx_int = 0 186 var f2: nx_int = 0 187 while f2 < n_samples { 188 if done[f2] == 0 { 189 if emitted[f2] < rc.max_new { 190 bat_ids[Mb] = curs[f2] 191 bat_seqs[Mb] = fseqs[f2] 192 bat_map[Mb] = f2 193 Mb = Mb + 1 194 } else { 195 done[f2] = 1 196 } 197 } 198 f2 = f2 + 1 199 } 200 if Mb == 0 { running = 0 } else { 201 if nx_f32_llm_forward_v4b(model, bat_ids, Mb, bat_seqs, rc.eps, 202 rc.attn_scale, rc.rope_log_base, 1, 203 blg) != NX_FLV4_OK { running = 0 } else { 204 var r2: nx_int = 0 205 while r2 < Mb { 206 let fk: nx_int = bat_map[r2] as nx_int 207 let row: *i64 = ((blg as i64) + r2 * vs * 8) as *i64 208 let pr2: *i64 = prngs[fk] as *i64 209 let nid: nx_int = nx_f32_sampler_sample_top_k(row, vs, rc.top_k, 210 rc.inv_temp_f32, pr2) 211 var stop: nx_int = 0 212 if nid == rc.im_end { stop = 1 } 213 if rc.eos >= 0 { if nid == rc.eos { stop = 1 } } 214 if stop == 1 { done[fk] = 1 } else { 215 one[0] = nid as i64 216 let dn2: nx_int = nx_bpe_decode_bytelevel(rc.vocab, one, 1, db) 217 let ob2: *u8 = ((texts as i64) + (fk as i64) * (text_stride as i64)) as *u8 218 var bi2: nx_int = 0 219 var no2: i64 = lens_out[fk] 220 while bi2 < dn2 { 221 if no2 < text_stride { ob2[no2] = db[bi2]; no2 = no2 + 1 } 222 bi2 = bi2 + 1 223 } 224 lens_out[fk] = no2 225 curs[fk] = nid as i64 226 emitted[fk] = emitted[fk] + 1 227 } 228 r2 = r2 + 1 229 } 230 } 231 } 232 } 233 234 var f3: nx_int = 0 235 while f3 < n_samples { 236 nx_pkv_seq_free(fseqs[f3] as *NxPagedSeq) 237 let ob3: *u8 = ((texts as i64) + (f3 as i64) * (text_stride as i64)) as *u8 238 let av: i64 = nx_reason_extract_last_int(ob3, lens_out[f3] as nx_int) 239 answers_out[f3] = av 240 f3 = f3 + 1 241 } 242 nx_pkv_seq_free(seq0) 243 sys_munmap(base_row, vs * 8) 244 sys_munmap(blg, n_samples * vs * 8) 245 return nx_reason_majority(answers_out, n_samples) 246}