nx_specdec_gate.nx source
↩ module page · 209 lines · 8414 B
1// nx_specdec_gate.nx -- MEASURED gate for sovereign SPECULATIVE DECODING
2// (prompt-lookup draft + multi-token greedy verify) on the REAL model.
3//
4// KAT pure drafter: match/no-match/most-recent-match
5// EQUIV-1 chunked prefill (ONE m=n forward) == sequential prefill
6// (n m=1 forwards): first 8 greedy tokens identical
7// EQUIV-2 spec_decode_greedy == plain_greedy: full 24-token sequence
8// identical (proves verify rows + cache truncation + offset-
9// causal positions)
10// SPEED same EQUIV-2 runs timed: us_plain vs us_spec + acceptance
11// stats (echo-y prompt = lookup-friendly; honest workload note)
12//
13// license_tier: ORIGINAL expect_exit: 0
14
15import "nx_syscalls.nx"
16import "nx_tier.nx"
17import "nx_le.nx"
18import "nx_bpe.nx"
19import "nx_gguf.nx"
20import "nx_gguf_load.nx"
21import "nx_gguf_meta.nx"
22import "nx_f32.nx"
23import "nx_f32_kv_cache.nx"
24import "nx_f32_lazy_weight.nx"
25import "nx_f32_llama_block.nx"
26import "nx_f32_llama_block_v4.nx"
27import "nx_f32_llama_stack_v4.nx"
28import "nx_f32_llama_layer_lazy_load.nx"
29import "nx_f32_llm.nx"
30import "nx_f32_llm_v4.nx"
31import "nx_f32_llm_read_dims.nx"
32import "nx_f32_bpe_load.nx"
33import "nx_f32_llm_special_tokens.nx"
34import "nx_f32_sampler.nx"
35import "nx_prng.nx"
36import "nx_reasoning.nx"
37import "nx_specdec.nx"
38
39func sg_w(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 }
40func sg_wn(v: i64) -> i64 {
41 var m: i64 = v
42 if m < 0 { sg_w("-" as *u8); m = 0 - m }
43 let t: *u8 = sys_mmap(28)
44 var k: i64 = 0
45 if m == 0 { t[0] = 48 as u8; k = 1 }
46 while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
47 let o: *u8 = sys_mmap(28)
48 var i: i64 = 0
49 while i < k { o[i] = t[k - 1 - i]; i = i + 1 }
50 sys_write(1, o, k)
51 return 0
52}
53
54func main() -> i64 {
55 // ---- KAT: pure drafter ----------------------------------------
56 let c1: *i64 = sys_mmap(16 * 8) as *i64
57 c1[0]=1; c1[1]=2; c1[2]=3; c1[3]=4; c1[4]=1; c1[5]=2; c1[6]=3
58 let d1: *i64 = sys_mmap(8 * 8) as *i64
59 let n1: nx_int = nx_spec_draft_lookup(c1, 7, 6, d1)
60 if n1 != 4 { return 11 }
61 if d1[0] != 4 { return 11 }
62 if d1[1] != 1 { return 11 }
63 if d1[2] != 2 { return 11 }
64 if d1[3] != 3 { return 11 }
65
66 let c2: *i64 = sys_mmap(8 * 8) as *i64
67 c2[0]=1; c2[1]=2; c2[2]=3; c2[3]=4; c2[4]=5
68 let n2: nx_int = nx_spec_draft_lookup(c2, 5, 6, d1)
69 if n2 != 0 { return 12 }
70
71 let c3: *i64 = sys_mmap(16 * 8) as *i64
72 c3[0]=7; c3[1]=8; c3[2]=9; c3[3]=5; c3[4]=7; c3[5]=8; c3[6]=10; c3[7]=7; c3[8]=8
73 let n3: nx_int = nx_spec_draft_lookup(c3, 9, 3, d1)
74 if n3 < 1 { return 13 }
75 if d1[0] != 10 { return 13 } // most-recent match wins (p=4, not p=0)
76 sg_w("SG KAT drafter match/no-match/recency OK\n" as *u8)
77
78 // ---- load the real model once ----------------------------------
79 let path: *u8 = "/tmp/nx_real_model.gguf" as *u8
80 let len_out: *i64 = sys_mmap(8) as *i64
81 let buf: *u8 = sys_read_file(path, len_out)
82 if buf == (0 as *u8) { return 10 }
83 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
84 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { return 20 }
85 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
86 let out_err: *i64 = sys_mmap(8) as *i64
87 if nx_f32_llm_read_dims_from_gguf(buf, len_out[0], hdr, model, out_err) != NX_FLD_OK { return 30 }
88 if nx_f32_llm_load_weights_v4_from_gguf(buf, hdr, model, out_err) != NX_FLV4_OK { return 40 }
89 let vocab: *NxBpeVocab = nx_bpe_vocab_new(67108864, 262144, 524288)
90 let nt2: *i64 = sys_mmap(8) as *i64
91 let nm: *i64 = sys_mmap(8) as *i64
92 if nx_f32_bpe_load_from_gguf(buf, len_out[0], hdr, vocab, nt2, nm, out_err) != NX_FBL_OK { return 50 }
93 let eos: nx_int = nx_f32_llm_read_eos(buf, len_out[0], hdr)
94 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(
95 model.n_layers, model.n_kv_heads, 160, model.head_dim)
96
97 let rc: *NxReasonCfg = nx_reason_cfg_alloc()
98 rc.model = model
99 rc.vocab = vocab
100 rc.cache = cache
101 rc.max_new = 24
102 rc.inv_temp_f32 = 0x3FA00000
103 rc.top_k = 40
104 rc.eps = 0x358637BD
105 rc.attn_scale = 0x3E000000
106 rc.rope_log_base = 0x415D0EAB
107 rc.eos = eos
108 rc.im_start = 151644
109 rc.im_end = 151645
110
111 // ---- EQUIV-1: chunked prefill == sequential prefill ------------
112 let q1: *u8 = "What is 47 plus 38? Answer with only the number." as *u8
113 let toks: *i64 = sys_mmap(512 * 8) as *i64
114 let ntq: nx_int = nx_reason_build_chat_toks(rc, q1, 49, toks)
115
116 // path A: sequential m=1 prefill + 8 greedy tokens.
117 nx_f32_kv_cache_reset(cache)
118 let lg: *i64 = sys_mmap(model.vocab_size * 8) as *i64
119 var pf: nx_int = 0
120 while pf < ntq {
121 let one: *i64 = ((toks as i64) + pf * 8) as *i64
122 if nx_f32_llm_forward_v4(model, one, 1, cache, rc.eps, rc.attn_scale,
123 rc.rope_log_base, 1, lg) != NX_FLV4_OK { return 60 }
124 pf = pf + 1
125 }
126 let seqA: *i64 = sys_mmap(8 * 8) as *i64
127 var tA: nx_int = nx_f32_sampler_argmax(lg, model.vocab_size)
128 var ia: nx_int = 0
129 let oneA: *i64 = sys_mmap(8) as *i64
130 while ia < 8 {
131 seqA[ia] = tA as i64
132 oneA[0] = tA as i64
133 if nx_f32_llm_forward_v4(model, oneA, 1, cache, rc.eps, rc.attn_scale,
134 rc.rope_log_base, 1, lg) != NX_FLV4_OK { return 60 }
135 tA = nx_f32_sampler_argmax(lg, model.vocab_size)
136 ia = ia + 1
137 }
138
139 // path B: plain_greedy (chunked prefill inside).
140 let stB: *NxSpecStats = nx_spec_stats_alloc()
141 let seqB: *i64 = sys_mmap(16 * 8) as *i64
142 let nB: nx_int = nx_spec_plain_greedy(rc, toks, ntq, 8, seqB, stB)
143 if nB < 8 {
144 // stop token inside 8 -- compare only nB then
145 sg_w("SG note: greedy stopped at " as *u8); sg_wn(nB as i64); sg_w(" tokens\n" as *u8)
146 }
147 var cmpn: nx_int = 8
148 if nB < 8 { cmpn = nB }
149 var ei: nx_int = 0
150 while ei < cmpn {
151 if seqA[ei] != seqB[ei] {
152 sg_w("EQUIV-1 MISMATCH at " as *u8); sg_wn(ei as i64)
153 sg_w(" seq=" as *u8); sg_wn(seqA[ei])
154 sg_w(" chunk=" as *u8); sg_wn(seqB[ei]); sg_w("\n" as *u8)
155 return 61
156 }
157 ei = ei + 1
158 }
159 sg_w("SG EQUIV-1 chunked-prefill == sequential-prefill (8 tokens) OK\n" as *u8)
160
161 // ---- EQUIV-2 + SPEED: spec == plain on an echo-y prompt --------
162 let q2: *u8 = "Repeat exactly: one two three four five six seven eight nine ten" as *u8
163 let ntq2: nx_int = nx_reason_build_chat_toks(rc, q2, 64, toks)
164
165 let stP: *NxSpecStats = nx_spec_stats_alloc()
166 let seqP: *i64 = sys_mmap(32 * 8) as *i64
167 let t0: i64 = sys_now_us()
168 let nP: nx_int = nx_spec_plain_greedy(rc, toks, ntq2, 24, seqP, stP)
169 let usP: i64 = sys_now_us() - t0
170 if nP < 0 { return 62 }
171
172 let stS: *NxSpecStats = nx_spec_stats_alloc()
173 let seqS: *i64 = sys_mmap(40 * 8) as *i64
174 let t1: i64 = sys_now_us()
175 let nS: nx_int = nx_spec_decode_greedy(rc, toks, ntq2, 24, seqS, stS)
176 let usS: i64 = sys_now_us() - t1
177 if nS < 0 { return 63 }
178
179 if nP != nS {
180 sg_w("EQUIV-2 LENGTH MISMATCH plain=" as *u8); sg_wn(nP as i64)
181 sg_w(" spec=" as *u8); sg_wn(nS as i64); sg_w("\n" as *u8)
182 return 64
183 }
184 var qi: nx_int = 0
185 while qi < nP {
186 if seqP[qi] != seqS[qi] {
187 sg_w("EQUIV-2 MISMATCH at " as *u8); sg_wn(qi as i64)
188 sg_w(" plain=" as *u8); sg_wn(seqP[qi])
189 sg_w(" spec=" as *u8); sg_wn(seqS[qi]); sg_w("\n" as *u8)
190 return 65
191 }
192 qi = qi + 1
193 }
194 sg_w("SG EQUIV-2 spec == plain BIT-IDENTICAL n=" as *u8); sg_wn(nP as i64); sg_w("\n" as *u8)
195
196 sg_w("SPEED plain_us=" as *u8); sg_wn(usP)
197 sg_w(" fw=" as *u8); sg_wn(stP.n_forwards)
198 sg_w(" spec_us=" as *u8); sg_wn(usS)
199 sg_w(" fw=" as *u8); sg_wn(stS.n_forwards)
200 sg_w(" drafted=" as *u8); sg_wn(stS.n_drafted)
201 sg_w(" accepted=" as *u8); sg_wn(stS.n_accepted)
202 sg_w("\n" as *u8)
203 var usS2: i64 = usS
204 if usS2 < 1 { usS2 = 1 }
205 sg_w("spec_vs_plain_x100=" as *u8); sg_wn(usP * 100 / usS2); sg_w("\n" as *u8)
206 sg_w("NOTE acceptance is workload-dependent (echo/list text accepts; free-form does not)\n" as *u8)
207 sg_w("SPECDEC_GATE DONE\n" as *u8)
208 return 0
209}