nx_f32_llm_run_v3.nx source
↩ module page · 189 lines · 8131 B
1// nx_f32_llm_run_v3.nx -- lazy-aware autoregressive runner.
2//
3// Same API as nx_f32_llm_run_v2 (NxF32SamplerCfg + recent-tokens
4// ring + repetition penalty), but invokes nx_f32_llm_forward_v4
5// internally. Use with v4-loaded models (lazy Q4_K weights).
6//
7// genealogy_id: standard_decode_loop + lazy_dispatch
8// lineage_id: substrate_f32_llm_run_v3_lazy
9
10import "nx_syscalls.nx"
11import "nx_tier.nx"
12import "nx_bpe.nx"
13import "nx_f32.nx"
14import "nx_f32_kv_cache.nx"
15import "nx_f32_lazy_weight.nx"
16import "nx_f32_llama_block_v4.nx"
17import "nx_f32_llama_stack_v4.nx"
18import "nx_f32_llm.nx"
19import "nx_f32_llm_v4.nx"
20import "nx_f32_sampler.nx"
21import "nx_f32_llm_run_v2.nx"
22
23const NX_FRN3_OK_BASE: nx_int = 0
24const NX_FRN3_ERR_NULL: nx_int = -1
25const NX_FRN3_ERR_BAD_DIM: nx_int = -2
26const NX_FRN3_ERR_FORWARD: nx_int = -3
27
28// Prefill block size: bounds the prefill logits buffer at CHUNK*vocab
29// (~39MB for Qwen2.5's 151936 vocab) while cutting prefill forwards ~CHUNKx.
30const NX_FRN3_PF_CHUNK: nx_int = 32
31
32func _frn3_pick(cfg: *NxF32SamplerCfg, logits: *i64, vocab_size: nx_int,
33 prng_state: *i64) -> nx_int {
34 if prng_state == (0 as *i64) {
35 return nx_f32_sampler_argmax(logits, vocab_size)
36 }
37 if cfg.top_p_q14 > 0 {
38 return nx_f32_sampler_sample_top_p(logits, vocab_size, cfg.top_p_q14,
39 cfg.inv_temp_f32, prng_state)
40 }
41 if cfg.top_k > 0 {
42 return nx_f32_sampler_sample_top_k(logits, vocab_size, cfg.top_k,
43 cfg.inv_temp_f32, prng_state)
44 }
45 return nx_f32_sampler_argmax(logits, vocab_size)
46}
47
48func _frn3_apply_penalty(cfg: *NxF32SamplerCfg, logits: *i64, vocab_size: nx_int,
49 recent: *i64, n_recent: nx_int) -> nx_int {
50 if cfg.penalty_f32 == 0x3F800000 { return 0 }
51 if n_recent <= 0 { return 0 }
52 nx_f32_sampler_apply_repetition_penalty(logits, vocab_size,
53 recent, n_recent, cfg.penalty_f32)
54 return 0
55}
56
57func _frn3_recent_push(recent: *i64, recent_cap: nx_int,
58 head_ptr: *i64, n_recent_ptr: *i64,
59 token_id: nx_int) -> nx_int {
60 if recent_cap <= 0 { return 0 }
61 let h: nx_int = head_ptr[0] as nx_int
62 recent[h] = token_id as i64
63 head_ptr[0] = ((h + 1) - ((h + 1) / recent_cap) * recent_cap) as i64
64 let nr: nx_int = n_recent_ptr[0] as nx_int
65 if nr < recent_cap { n_recent_ptr[0] = (nr + 1) as i64 }
66 return 0
67}
68
69func nx_f32_llm_run_v3(model: *NxF32LlamaModel,
70 vocab: *NxBpeVocab,
71 cache: *NxF32KVCache,
72 prompt: *u8,
73 n_prompt: nx_int,
74 max_new_tokens: nx_int,
75 cfg: *NxF32SamplerCfg,
76 eps: i64,
77 attn_scale: i64,
78 rope_log_base: i64,
79 apply_rope: nx_int,
80 prng_state: *i64,
81 eos_token_id: nx_int,
82 out_bytes: *u8,
83 out_bytes_cap: nx_int) -> nx_int {
84
85 if model == (0 as *NxF32LlamaModel) { return NX_FRN3_ERR_NULL }
86 if vocab == (0 as *NxBpeVocab) { return NX_FRN3_ERR_NULL }
87 if cache == (0 as *NxF32KVCache) { return NX_FRN3_ERR_NULL }
88 if cfg == (0 as *NxF32SamplerCfg) { return NX_FRN3_ERR_NULL }
89 if out_bytes_cap <= 0 { return NX_FRN3_ERR_BAD_DIM }
90 if max_new_tokens < 0 { return NX_FRN3_ERR_BAD_DIM }
91
92 let prompt_tokens: *i64 = sys_mmap(n_prompt * 8) as *i64
93 let n_prompt_tok: nx_int = nx_bpe_encode_bytelevel(vocab, prompt, n_prompt, prompt_tokens)
94 if n_prompt_tok < 0 { return NX_FRN3_ERR_BAD_DIM }
95 if n_prompt_tok == 0 { return NX_FRN3_ERR_BAD_DIM }
96
97 var recent_cap: nx_int = cfg.recent_cap
98 if recent_cap < 0 { recent_cap = 0 }
99 let recent: *i64 = sys_mmap((recent_cap + 1) * 8) as *i64
100 let recent_head: *i64 = sys_mmap(8) as *i64
101 let n_recent: *i64 = sys_mmap(8) as *i64
102 recent_head[0] = 0
103 n_recent[0] = 0
104 if recent_cap > 0 {
105 var pp: nx_int = 0
106 while pp < n_prompt_tok {
107 _frn3_recent_push(recent, recent_cap, recent_head, n_recent,
108 prompt_tokens[pp] as nx_int)
109 pp = pp + 1
110 }
111 }
112
113 // BOUNDED-CHUNK prefill (2026-07-09): blocks of <=NX_FRN3_PF_CHUNK tokens,
114 // ONE forward per block -- weight bytes read once per BLOCK instead of once
115 // per token (~32x fewer prefill forwards = the TTFT win; chunked prefill
116 // proven token-identical to sequential by nx_specdec_gate EQUIV-1, and
117 // mixed chunk sizes by the paged gates' 5+1+3 rounds -- offset-causal
118 // q_pos handles any split). The logits buffer is FIXED at CHUNK*vocab
119 // (~39MB) -- the OLD all-at-once batched prefill got this wrong
120 // (n_prompt*vocab -> mmap fail on long prompts -> the -3 error); bounded
121 // chunks keep memory flat at ANY prompt length. Sequential 1-token
122 // prefill (the previous fix) was memory-safe but paid a full weight
123 // sweep PER TOKEN. After the last block its last row moves into
124 // logits1 (vocab*8), which the decode loop reuses unchanged.
125 let logits1: *i64 = sys_mmap(model.vocab_size * 8) as *i64
126 let logits_pf: *i64 = sys_mmap(NX_FRN3_PF_CHUNK * model.vocab_size * 8) as *i64
127 var pf: nx_int = 0
128 var last_n: nx_int = 0
129 while pf < n_prompt_tok {
130 var cn: nx_int = n_prompt_tok - pf
131 if cn > NX_FRN3_PF_CHUNK { cn = NX_FRN3_PF_CHUNK }
132 let chunk: *i64 = (((prompt_tokens as i64) + pf * 8)) as *i64
133 let v_pre: nx_int = nx_f32_llm_forward_v4(model, chunk, cn,
134 cache, eps, attn_scale, rope_log_base,
135 apply_rope, logits_pf)
136 if v_pre != NX_FLV4_OK { return NX_FRN3_ERR_FORWARD }
137 last_n = cn
138 pf = pf + cn
139 }
140 let lrow_pf: *i64 = (((logits_pf as i64) + (last_n - 1) * model.vocab_size * 8)) as *i64
141 var lcp: nx_int = 0
142 while lcp < model.vocab_size { logits1[lcp] = lrow_pf[lcp]; lcp = lcp + 1 }
143 sys_munmap(logits_pf, NX_FRN3_PF_CHUNK * model.vocab_size * 8)
144 let last_row: *i64 = logits1
145
146 var n_emitted: nx_int = 0
147 let next_buf: *i64 = sys_mmap(8) as *i64
148 let one_byte_buf: *u8 = sys_mmap(64)
149
150 var step: nx_int = 0
151 var current_logits: *i64 = last_row
152 while step < max_new_tokens {
153 _frn3_apply_penalty(cfg, current_logits, model.vocab_size,
154 recent, n_recent[0] as nx_int)
155 let next_id: nx_int = _frn3_pick(cfg, current_logits, model.vocab_size,
156 prng_state)
157 if eos_token_id >= 0 {
158 if next_id == eos_token_id { step = max_new_tokens }
159 }
160 if step < max_new_tokens {
161 next_buf[0] = next_id as i64
162 let nb: nx_int = nx_bpe_decode_bytelevel(vocab, next_buf, 1, one_byte_buf)
163 if nb > 0 {
164 var bi: nx_int = 0
165 while bi < nb {
166 if n_emitted >= out_bytes_cap {
167 bi = nb
168 step = max_new_tokens
169 } else {
170 out_bytes[n_emitted] = one_byte_buf[bi]
171 n_emitted = n_emitted + 1
172 bi = bi + 1
173 }
174 }
175 }
176 if step < max_new_tokens {
177 _frn3_recent_push(recent, recent_cap, recent_head, n_recent, next_id)
178 let v_d: nx_int = nx_f32_llm_forward_v4(model, next_buf, 1, cache,
179 eps, attn_scale, rope_log_base,
180 apply_rope, logits1)
181 if v_d != NX_FLV4_OK { return NX_FRN3_ERR_FORWARD }
182 current_logits = logits1
183 }
184 step = step + 1
185 }
186 }
187
188 return n_emitted
189}