nx_f32_llm_run_v2.nx source
↩ module page · 215 lines · 8254 B
1// nx_f32_llm_run_v2.nx -- production runner with full sampler config.
2//
3// Wraps the v1 runner pattern but accepts a sampler config struct
4// covering temperature + top-k + top-p + repetition penalty, and
5// maintains a recent-tokens ring buffer that feeds the penalty
6// across decode steps.
7//
8// Composes:
9// nx_bpe_encode / nx_bpe_decode
10// nx_f32_llm_forward
11// nx_f32_sampler_apply_repetition_penalty
12// nx_f32_sampler_sample_top_p / top_k / temp / argmax
13// nx_f32_kv_cache_alloc/append/advance
14//
15// genealogy_id: standard_decode_loop_with_sampler_pipeline
16// lineage_id: substrate_f32_llm_run_v2
17
18import "nx_syscalls.nx"
19import "nx_tier.nx"
20import "nx_bpe.nx"
21import "nx_f32.nx"
22import "nx_f32_kv_cache.nx"
23import "nx_f32_llm.nx"
24import "nx_f32_sampler.nx"
25const NX_MAGIC_16384: i64 = 16384
26
27const NX_FRN2_OK_BASE: nx_int = 0
28const NX_FRN2_ERR_NULL: nx_int = -1
29const NX_FRN2_ERR_BAD_DIM: nx_int = -2
30const NX_FRN2_ERR_FORWARD: nx_int = -3
31
32struct NxF32SamplerCfg {
33 inv_temp_f32: i64, // 1/temperature in f32 bits (0x3F800000 = 1.0)
34 top_k: nx_int, // 0 = disabled
35 top_p_q14: i64, // 0 = disabled, NX_MAGIC_16384 = no truncation
36 penalty_f32: i64, // 0x3F800000 = 1.0 (disabled)
37 recent_cap: nx_int // ring-buffer size (e.g. 64); 0 = no penalty
38}
39
40const NX_F32_SAMPLER_CFG_BYTES: nx_int = 40 // 5 fields * 8
41
42func nx_f32_sampler_cfg_alloc() -> *NxF32SamplerCfg {
43 let cfg: *NxF32SamplerCfg = sys_mmap(NX_F32_SAMPLER_CFG_BYTES) as *NxF32SamplerCfg
44 cfg.inv_temp_f32 = 0x3F800000 // 1.0
45 cfg.top_k = 0
46 cfg.top_p_q14 = 0
47 cfg.penalty_f32 = 0x3F800000 // 1.0 (disabled)
48 cfg.recent_cap = 0
49 return cfg
50}
51
52// Sampler dispatch: choose best primitive based on cfg.
53
54func _frn2_pick(cfg: *NxF32SamplerCfg, logits: *i64, vocab_size: nx_int,
55 prng_state: *i64) -> nx_int {
56 if prng_state == (0 as *i64) {
57 return nx_f32_sampler_argmax(logits, vocab_size)
58 }
59 if cfg.top_p_q14 > 0 {
60 return nx_f32_sampler_sample_top_p(logits, vocab_size, cfg.top_p_q14,
61 cfg.inv_temp_f32, prng_state)
62 }
63 if cfg.top_k > 0 {
64 return nx_f32_sampler_sample_top_k(logits, vocab_size, cfg.top_k,
65 cfg.inv_temp_f32, prng_state)
66 }
67 return nx_f32_sampler_argmax(logits, vocab_size)
68}
69
70// Apply repetition penalty IF configured.
71
72func _frn2_apply_penalty(cfg: *NxF32SamplerCfg, logits: *i64, vocab_size: nx_int,
73 recent: *i64, n_recent: nx_int) -> nx_int {
74 // Skip if penalty is exactly 1.0 (default / disabled).
75 if cfg.penalty_f32 == 0x3F800000 { return 0 }
76 if n_recent <= 0 { return 0 }
77 nx_f32_sampler_apply_repetition_penalty(logits, vocab_size,
78 recent, n_recent, cfg.penalty_f32)
79 return 0
80}
81
82// Push a token id onto a ring buffer. head/n_recent are caller-maintained.
83
84func _frn2_recent_push(recent: *i64, recent_cap: nx_int,
85 head_ptr: *i64, n_recent_ptr: *i64,
86 token_id: nx_int) -> nx_int {
87 if recent_cap <= 0 { return 0 }
88 let h: nx_int = head_ptr[0] as nx_int
89 recent[h] = token_id as i64
90 head_ptr[0] = ((h + 1) - ((h + 1) / recent_cap) * recent_cap) as i64
91 let nr: nx_int = n_recent_ptr[0] as nx_int
92 if nr < recent_cap {
93 n_recent_ptr[0] = (nr + 1) as i64
94 }
95 return 0
96}
97
98// v2 runner. 15 args.
99
100func nx_f32_llm_run_v2(model: *NxF32LlamaModel,
101 vocab: *NxBpeVocab,
102 cache: *NxF32KVCache,
103 prompt: *u8,
104 n_prompt: nx_int,
105 max_new_tokens: nx_int,
106 cfg: *NxF32SamplerCfg,
107 eps: i64,
108 attn_scale: i64,
109 rope_log_base: i64,
110 apply_rope: nx_int,
111 prng_state: *i64,
112 eos_token_id: nx_int,
113 out_bytes: *u8,
114 out_bytes_cap: nx_int) -> nx_int {
115
116 if model == (0 as *NxF32LlamaModel) { return NX_FRN2_ERR_NULL }
117 if vocab == (0 as *NxBpeVocab) { return NX_FRN2_ERR_NULL }
118 if cache == (0 as *NxF32KVCache) { return NX_FRN2_ERR_NULL }
119 if cfg == (0 as *NxF32SamplerCfg) { return NX_FRN2_ERR_NULL }
120 if out_bytes_cap <= 0 { return NX_FRN2_ERR_BAD_DIM }
121 if max_new_tokens < 0 { return NX_FRN2_ERR_BAD_DIM }
122
123 // ===== 1. Tokenize prompt =====
124 let prompt_tokens: *i64 = sys_mmap(n_prompt * 8) as *i64
125 let n_prompt_tok: nx_int = nx_bpe_encode(vocab, prompt, n_prompt, prompt_tokens)
126 if n_prompt_tok < 0 { return NX_FRN2_ERR_BAD_DIM }
127 if n_prompt_tok == 0 { return NX_FRN2_ERR_BAD_DIM }
128
129 // ===== 2. Recent-tokens ring buffer (for repetition penalty) =====
130 var recent_cap: nx_int = cfg.recent_cap
131 if recent_cap < 0 { recent_cap = 0 }
132 let recent: *i64 = sys_mmap((recent_cap + 1) * 8) as *i64
133 let recent_head: *i64 = sys_mmap(8) as *i64
134 let n_recent: *i64 = sys_mmap(8) as *i64
135 recent_head[0] = 0
136 n_recent[0] = 0
137
138 // Seed recent window with the prompt tokens (most recent up to cap).
139 if recent_cap > 0 {
140 var pp: nx_int = 0
141 while pp < n_prompt_tok {
142 _frn2_recent_push(recent, recent_cap, recent_head, n_recent,
143 prompt_tokens[pp] as nx_int)
144 pp = pp + 1
145 }
146 }
147
148 // ===== 3. Allocate logits scratch =====
149 let logits: *i64 = sys_mmap(n_prompt_tok * model.vocab_size * 8) as *i64
150 let logits1: *i64 = sys_mmap(model.vocab_size * 8) as *i64
151
152 // ===== 4. Prefill =====
153 let v_pre: nx_int = nx_f32_llm_forward(model, prompt_tokens, n_prompt_tok,
154 cache, eps, attn_scale, rope_log_base,
155 apply_rope, logits)
156 if v_pre != NX_F32_LLM_OK { return NX_FRN2_ERR_FORWARD }
157
158 let last_row: *i64 = (((logits as i64) +
159 (n_prompt_tok - 1) * model.vocab_size * 8)) as *i64
160
161 var n_emitted: nx_int = 0
162 let next_buf: *i64 = sys_mmap(8) as *i64
163 let one_byte_buf: *u8 = sys_mmap(64)
164
165 // ===== 5. Decode loop =====
166 var step: nx_int = 0
167 var current_logits: *i64 = last_row
168 while step < max_new_tokens {
169 // Apply repetition penalty IN PLACE before sampling.
170 _frn2_apply_penalty(cfg, current_logits, model.vocab_size,
171 recent, n_recent[0] as nx_int)
172
173 // Sample.
174 let next_id: nx_int = _frn2_pick(cfg, current_logits, model.vocab_size,
175 prng_state)
176
177 // EOS check.
178 if eos_token_id >= 0 {
179 if next_id == eos_token_id { step = max_new_tokens }
180 }
181 if step < max_new_tokens {
182 // Detokenize.
183 next_buf[0] = next_id as i64
184 let nb: nx_int = nx_bpe_decode(vocab, next_buf, 1, one_byte_buf)
185 if nb > 0 {
186 var bi: nx_int = 0
187 while bi < nb {
188 if n_emitted >= out_bytes_cap {
189 bi = nb
190 step = max_new_tokens
191 } else {
192 out_bytes[n_emitted] = one_byte_buf[bi]
193 n_emitted = n_emitted + 1
194 bi = bi + 1
195 }
196 }
197 }
198
199 if step < max_new_tokens {
200 // Push next_id into recent window before next iteration.
201 _frn2_recent_push(recent, recent_cap, recent_head, n_recent, next_id)
202
203 // Forward decode 1 token.
204 let v_d: nx_int = nx_f32_llm_forward(model, next_buf, 1, cache,
205 eps, attn_scale, rope_log_base,
206 apply_rope, logits1)
207 if v_d != NX_F32_LLM_OK { return NX_FRN2_ERR_FORWARD }
208 current_logits = logits1
209 }
210 step = step + 1
211 }
212 }
213
214 return n_emitted
215}