nx_specdec.nx source
↩ module page · 238 lines · 9897 B
1// nx_specdec.nx -- sovereign SPECULATIVE DECODING (prompt-lookup drafted,
2// greedy-verified; census OPPORTUNITY "Speculative decoding", momentum 3).
3//
4// Decode is MEMORY-bound (weights re-read per token). Speculation turns k
5// sequential m=1 forwards into ONE m=k+1 forward: DRAFT k tokens cheaply,
6// VERIFY them in a single multi-token pass (weight bytes read ONCE for the
7// whole block -- the same amortization that makes prefill cheap). With
8// GREEDY acceptance (argmax match) the output is BIT-IDENTICAL to plain
9// greedy decode: pure measured speedup, zero quality change.
10//
11// Draft source here = PROMPT-LOOKUP (n-gram continuation from the model's
12// own context; Saxena 2023 "prompt lookup decoding", the draft-model-free
13// member of the speculative family: Leviathan 2023 / Chen 2023; EAGLE-class
14// trained drafts can replace it later behind the same verify loop).
15// Acceptance is WORKLOAD-DEPENDENT: echo/quote/list-heavy text accepts big
16// chunks; free-form text accepts little (then cost ~= plain decode plus the
17// occasional rejected block).
18//
19// Cache contract: the verify forward appends k+1 rows; on partial accept the
20// cache is TRUNCATED back to the accepted prefix (seq_len rewind -- rows are
21// plain arrays, later appends overwrite). Positions/causality inside the
22// multi-token pass are the offset-causal q_pos = total_k - n_q + i already
23// implemented by nx_f32_attn_multi (proven by nx_specdec_gate EQUIV rows).
24//
25// genealogy_id: leviathan_2023_speculative + saxena_2023_prompt_lookup
26// lineage_id: substrate_specdec_v1
27
28import "nx_syscalls.nx"
29import "nx_tier.nx"
30import "nx_bpe.nx"
31import "nx_f32_kv_cache.nx"
32import "nx_f32_llm.nx"
33import "nx_f32_llm_v4.nx"
34import "nx_f32_sampler.nx"
35import "nx_reasoning.nx"
36
37const NX_SPEC_NGRAM_MAX: nx_int = 3 // longest suffix n-gram to match
38const NX_SPEC_NGRAM_MIN: nx_int = 2 // 1-grams match spuriously; stop at 2
39const NX_SPEC_K: nx_int = 6 // max drafted tokens per round
40
41// ===== Prompt-lookup drafter (pure ints; KAT-covered) ==============
42// Match the last g tokens of ctx (g = NGRAM_MAX..NGRAM_MIN) at an earlier
43// position; on the MOST RECENT match, draft the k tokens that followed it.
44// Returns n_drafted (0 = no match).
45
46func nx_spec_draft_lookup(ctx: *i64, n_ctx: nx_int, k: nx_int,
47 out_draft: *i64) -> nx_int {
48 if n_ctx < NX_SPEC_NGRAM_MIN + 1 { return 0 }
49 var g: nx_int = NX_SPEC_NGRAM_MAX
50 while g >= NX_SPEC_NGRAM_MIN {
51 if n_ctx > g {
52 // scan candidate positions RIGHT-to-LEFT (most recent first);
53 // candidate match ends at p+g-1, must end before the suffix.
54 var p: nx_int = n_ctx - g - 1
55 while p >= 0 {
56 var hit: nx_int = 1
57 var j: nx_int = 0
58 while j < g {
59 if ctx[p + j] != ctx[n_ctx - g + j] { hit = 0; j = g } else { j = j + 1 }
60 }
61 if hit == 1 {
62 // draft the tokens that FOLLOWED the match (all history;
63 // may overlap the suffix region -- that is fine, it is
64 // still past text). p <= n_ctx-g-1 so >= 1 draftable.
65 var nd: nx_int = 0
66 var src: nx_int = p + g
67 while nd < k {
68 if src >= n_ctx { nd = k } else {
69 out_draft[nd] = ctx[src]
70 nd = nd + 1
71 src = src + 1
72 }
73 }
74 let actual: nx_int = src - (p + g)
75 if actual > 0 { return actual }
76 return 0
77 }
78 p = p - 1
79 }
80 }
81 g = g - 1
82 }
83 return 0
84}
85
86// ===== Stats bundle ================================================
87
88struct NxSpecStats {
89 n_forwards: i64, // model forward calls (any m)
90 n_rows: i64, // total token-rows pushed through forwards
91 n_drafted: i64,
92 n_accepted: i64,
93 n_emitted: i64
94}
95
96const NX_SPEC_STATS_BYTES: nx_int = 40
97
98func nx_spec_stats_alloc() -> *NxSpecStats {
99 return sys_mmap(NX_SPEC_STATS_BYTES) as *NxSpecStats
100}
101
102func _spec_trunc(cache: *NxF32KVCache, new_len: nx_int) -> i64 {
103 cache.seq_len = new_len
104 return 0
105}
106
107func _spec_stop_tok(rc: *NxReasonCfg, t: nx_int) -> nx_int {
108 if t == rc.im_end { return 1 }
109 if rc.eos >= 0 { if t == rc.eos { return 1 } }
110 return 0
111}
112
113// ===== Plain greedy (m=1 steps) -- the equivalence baseline ========
114// Emits TOKEN IDS to out_toks; returns count. Chunked prefill (one
115// m=n_prompt forward) then token-by-token decode.
116
117func nx_spec_plain_greedy(rc: *NxReasonCfg, toks: *i64, nt: nx_int,
118 max_new: nx_int, out_toks: *i64,
119 st: *NxSpecStats) -> nx_int {
120 nx_f32_kv_cache_reset(rc.cache)
121 let model: *NxF32LlamaModel = rc.model
122 let logits: *i64 = sys_mmap((nt + 2) * model.vocab_size * 8) as *i64
123 if nx_f32_llm_forward_v4(model, toks, nt, rc.cache, rc.eps,
124 rc.attn_scale, rc.rope_log_base, 1,
125 logits) != NX_FLV4_OK { return 0 - 1 }
126 st.n_forwards = st.n_forwards + 1
127 st.n_rows = st.n_rows + (nt as i64)
128 let last_row: *i64 = ((logits as i64) + (nt - 1) * model.vocab_size * 8) as *i64
129 var t: nx_int = nx_f32_sampler_argmax(last_row, model.vocab_size)
130 var no: nx_int = 0
131 let one: *i64 = sys_mmap(8) as *i64
132 while no < max_new {
133 if _spec_stop_tok(rc, t) == 1 { return no }
134 out_toks[no] = t as i64
135 no = no + 1
136 if no >= max_new { return no }
137 one[0] = t as i64
138 if nx_f32_llm_forward_v4(model, one, 1, rc.cache, rc.eps,
139 rc.attn_scale, rc.rope_log_base, 1,
140 logits) != NX_FLV4_OK { return 0 - 1 }
141 st.n_forwards = st.n_forwards + 1
142 st.n_rows = st.n_rows + 1
143 t = nx_f32_sampler_argmax(logits, model.vocab_size)
144 }
145 return no
146}
147
148// ===== Speculative greedy decode ===================================
149// Same contract as plain_greedy; output token sequence is BIT-IDENTICAL
150// (greedy acceptance). ctx (prompt+emitted) feeds the lookup drafter.
151
152func nx_spec_decode_greedy(rc: *NxReasonCfg, toks: *i64, nt: nx_int,
153 max_new: nx_int, out_toks: *i64,
154 st: *NxSpecStats) -> nx_int {
155 nx_f32_kv_cache_reset(rc.cache)
156 let model: *NxF32LlamaModel = rc.model
157 let vs: nx_int = model.vocab_size
158 let logits: *i64 = sys_mmap((nt + NX_SPEC_K + 2) * vs * 8) as *i64
159
160 // ctx = prompt ++ emitted (drafter's source).
161 let ctx: *i64 = sys_mmap((nt + max_new + NX_SPEC_K + 2) * 8) as *i64
162 var n_ctx: nx_int = 0
163 var ci: nx_int = 0
164 while ci < nt { ctx[ci] = toks[ci]; ci = ci + 1 }
165 n_ctx = nt
166
167 // chunked prefill.
168 if nx_f32_llm_forward_v4(model, toks, nt, rc.cache, rc.eps,
169 rc.attn_scale, rc.rope_log_base, 1,
170 logits) != NX_FLV4_OK { return 0 - 1 }
171 st.n_forwards = st.n_forwards + 1
172 st.n_rows = st.n_rows + (nt as i64)
173 let last_row: *i64 = ((logits as i64) + (nt - 1) * vs * 8) as *i64
174 var t: nx_int = nx_f32_sampler_argmax(last_row, vs)
175
176 let draft: *i64 = sys_mmap(NX_SPEC_K * 8) as *i64
177 let feed: *i64 = sys_mmap((NX_SPEC_K + 1) * 8) as *i64
178 var no: nx_int = 0
179
180 while no < max_new {
181 if _spec_stop_tok(rc, t) == 1 { return no }
182 out_toks[no] = t as i64
183 ctx[n_ctx] = t as i64
184 n_ctx = n_ctx + 1
185 no = no + 1
186 if no >= max_new { return no }
187
188 var kd: nx_int = nx_spec_draft_lookup(ctx, n_ctx, NX_SPEC_K, draft)
189 if kd > max_new - no { kd = max_new - no }
190 st.n_drafted = st.n_drafted + (kd as i64)
191
192 // feed = [t, d1..dkd]: ONE forward, kd+1 rows.
193 feed[0] = t as i64
194 var fi: nx_int = 0
195 while fi < kd { feed[fi + 1] = draft[fi]; fi = fi + 1 }
196 let s_before: nx_int = nx_f32_kv_cache_get_seq_len(rc.cache)
197 if nx_f32_llm_forward_v4(model, feed, kd + 1, rc.cache, rc.eps,
198 rc.attn_scale, rc.rope_log_base, 1,
199 logits) != NX_FLV4_OK { return 0 - 1 }
200 st.n_forwards = st.n_forwards + 1
201 st.n_rows = st.n_rows + ((kd + 1) as i64)
202
203 // accept the longest prefix of drafts matching the model's argmax.
204 var n_acc: nx_int = 0
205 var j: nx_int = 0
206 var stop_hit: nx_int = 0
207 while j < kd {
208 let row_j: *i64 = ((logits as i64) + j * vs * 8) as *i64
209 let mj: nx_int = nx_f32_sampler_argmax(row_j, vs)
210 if mj != (draft[j] as nx_int) { j = kd } else {
211 // draft j confirmed by the model
212 st.n_accepted = st.n_accepted + 1
213 if _spec_stop_tok(rc, draft[j] as nx_int) == 1 {
214 stop_hit = 1
215 j = kd
216 } else {
217 out_toks[no] = draft[j]
218 ctx[n_ctx] = draft[j]
219 n_ctx = n_ctx + 1
220 no = no + 1
221 n_acc = n_acc + 1
222 if no >= max_new { stop_hit = 2; j = kd } else { j = j + 1 }
223 }
224 }
225 }
226 if stop_hit == 1 { return no }
227 if stop_hit == 2 { return no } // max_new mid-accept: done
228
229 // truncate cache to prefix [.. t + accepted drafts].
230 _spec_trunc(rc.cache, s_before + 1 + n_acc)
231
232 // bonus token = the model's own argmax at the first mismatch row
233 // (or after the last accepted draft) -- becomes next round's t.
234 let row_b: *i64 = ((logits as i64) + n_acc * vs * 8) as *i64
235 t = nx_f32_sampler_argmax(row_b, vs)
236 }
237 return no
238}