nx_paged_kv_gate.nx source
↩ module page · 222 lines · 9874 B
1// nx_paged_kv_gate.nx -- MEASURED gate for the sovereign PAGED KV-cache
2// (nx_kvcache + nx_f32_attn_paged). PURE (no model, fast):
3//
4// KAT pool alloc/refcount/free; fork refcounts; exhaustion fail-fast;
5// truncate unref
6// EQUIV paged attention == contiguous attention BIT-EXACT over 3
7// append rounds (5-token prefill, 1-token decode, 3-token chunk)
8// x 2 layers x GQA heads
9// FORK fork after prefill; diverge A/RA vs B/RB; copy-on-append fires
10// (refcounts return to 1); BOTH forks bit-exact vs fresh
11// contiguous replays of their histories
12// REWIND truncate A back to the fork point, append RB -> bit-exact ==
13// B's outputs (the speculative-decode rewind contract)
14//
15// license_tier: ORIGINAL expect_exit: 0
16
17import "nx_syscalls.nx"
18import "nx_tier.nx"
19import "nx_f32.nx"
20import "nx_f32_cvt.nx"
21import "nx_f32_kv_cache.nx"
22import "nx_f32_attn_cached.nx"
23import "nx_kvcache.nx"
24import "nx_f32_attn_paged.nx"
25import "nx_fmt.nx"
26
27const PG_NL: nx_int = 2 // layers
28const PG_NH: nx_int = 4 // heads
29const PG_NKV: nx_int = 2 // kv heads
30const PG_HD: nx_int = 8 // head dim
31// q_dim = 32, kv_dim = 16
32
33func pg_nl2() -> i64 { fmt_puts("\n" as *u8); return 0 }
34func pg_lcg(s: i64) -> i64 { var v: i64 = s * 1103515245 + 12345; v = v & 2147483647; return v }
35
36// fill n f32 slots with small ints (exact regime), advancing *state.
37func pg_fill(p: *i64, n: nx_int, state: *i64) -> i64 {
38 var s: i64 = state[0]
39 var i: nx_int = 0
40 while i < n {
41 s = pg_lcg(s)
42 let vv: i64 = nx_i32_to_f32((s % 17) - 8)
43 p[i] = vv
44 i = i + 1
45 }
46 state[0] = s
47 return 0
48}
49
50func pg_cmp(a: *i64, b: *i64, n: nx_int) -> i64 {
51 var i: nx_int = 0
52 while i < n { if a[i] != b[i] { return 0 } i = i + 1 }
53 return 1
54}
55
56// run one append-round on BOTH paths; compare bit-exact. Returns 0 ok.
57func pg_round(cc: *NxF32KVCache, ps: *NxPagedSeq,
58 Q: *i64, K: *i64, V: *i64, n_tok: nx_int,
59 scale: i64, oc: *i64, op: *i64) -> i64 {
60 let ve: nx_int = nx_pkv_ensure_append(ps, n_tok)
61 if ve != NX_PKV_OK { return 1 }
62 var L: nx_int = 0
63 while L < PG_NL {
64 let vc: nx_int = nx_f32_attn_with_cache(Q, K, V, n_tok, PG_NH, PG_NKV,
65 PG_HD, cc, L, 1, scale, oc)
66 if vc != NX_F32_AC_OK { return 2 }
67 let vp: nx_int = nx_f32_attn_with_paged(Q, K, V, n_tok, PG_NH, PG_NKV,
68 PG_HD, ps, L, 1, scale, op)
69 if vp != NX_F32_AP_OK { return 3 }
70 if pg_cmp(oc, op, n_tok * PG_NH * PG_HD) != 1 { return 4 }
71 L = L + 1
72 }
73 nx_f32_kv_cache_advance(cc, n_tok)
74 nx_pkv_advance(ps, n_tok)
75 return 0
76}
77
78// replay rounds on a paged seq ONLY (for fork divergence), collecting the
79// last round's outputs per layer into op (layer-major reuse: compare per
80// layer inside). Returns 0 ok.
81func pg_round_paged_only(ps: *NxPagedSeq, Q: *i64, K: *i64, V: *i64,
82 n_tok: nx_int, scale: i64, op: *i64) -> i64 {
83 let ve: nx_int = nx_pkv_ensure_append(ps, n_tok)
84 if ve != NX_PKV_OK { return 1 }
85 var L: nx_int = 0
86 while L < PG_NL {
87 let vp: nx_int = nx_f32_attn_with_paged(Q, K, V, n_tok, PG_NH, PG_NKV,
88 PG_HD, ps, L, 1, scale,
89 ((op as i64) + L * n_tok * PG_NH * PG_HD * 8) as *i64)
90 if vp != NX_F32_AP_OK { return 2 }
91 L = L + 1
92 }
93 nx_pkv_advance(ps, n_tok)
94 return 0
95}
96
97func pg_round_contig_only(cc: *NxF32KVCache, Q: *i64, K: *i64, V: *i64,
98 n_tok: nx_int, scale: i64, oc: *i64) -> i64 {
99 var L: nx_int = 0
100 while L < PG_NL {
101 let vc: nx_int = nx_f32_attn_with_cache(Q, K, V, n_tok, PG_NH, PG_NKV,
102 PG_HD, cc, L, 1, scale,
103 ((oc as i64) + L * n_tok * PG_NH * PG_HD * 8) as *i64)
104 if vc != NX_F32_AC_OK { return 2 }
105 L = L + 1
106 }
107 nx_f32_kv_cache_advance(cc, n_tok)
108 return 0
109}
110
111func main() -> i64 {
112 let kv_dim: nx_int = PG_NKV * PG_HD
113 let q_dim: nx_int = PG_NH * PG_HD
114 let scale: i64 = 0x3E800000 // 0.25 f32 (any fixed scale; both paths share)
115
116 // ---- KAT: pool mechanics ---------------------------------------
117 let p0: *NxPagedPool = nx_pkv_pool_new(3, PG_NL, kv_dim)
118 let b0: i64 = nx_pkv_alloc_block(p0)
119 let b1: i64 = nx_pkv_alloc_block(p0)
120 let b2: i64 = nx_pkv_alloc_block(p0)
121 if b0 != 0 { return 11 }
122 if b1 != 1 { return 11 }
123 if b2 != 2 { return 11 }
124 let b3: i64 = nx_pkv_alloc_block(p0)
125 if b3 != 0 - 1 { return 12 } // exhaustion fail-fast
126 nx_pkv_unref_block(p0, b1)
127 let b4: i64 = nx_pkv_alloc_block(p0)
128 if b4 != 1 { return 13 } // freed id recycles
129 fmt_puts("PKV KAT pool alloc/exhaust/recycle OK"); pg_nl2()
130
131 // ---- shared deterministic inputs -------------------------------
132 let st: *i64 = sys_mmap(8) as *i64
133 st[0] = 20260709
134 let Q1: *i64 = sys_mmap(5 * q_dim * 8) as *i64
135 let K1: *i64 = sys_mmap(5 * kv_dim * 8) as *i64
136 let V1: *i64 = sys_mmap(5 * kv_dim * 8) as *i64
137 pg_fill(Q1, 5 * q_dim, st); pg_fill(K1, 5 * kv_dim, st); pg_fill(V1, 5 * kv_dim, st)
138 let Q2: *i64 = sys_mmap(1 * q_dim * 8) as *i64
139 let K2: *i64 = sys_mmap(1 * kv_dim * 8) as *i64
140 let V2: *i64 = sys_mmap(1 * kv_dim * 8) as *i64
141 pg_fill(Q2, 1 * q_dim, st); pg_fill(K2, 1 * kv_dim, st); pg_fill(V2, 1 * kv_dim, st)
142 let Q3: *i64 = sys_mmap(3 * q_dim * 8) as *i64
143 let K3: *i64 = sys_mmap(3 * kv_dim * 8) as *i64
144 let V3: *i64 = sys_mmap(3 * kv_dim * 8) as *i64
145 pg_fill(Q3, 3 * q_dim, st); pg_fill(K3, 3 * kv_dim, st); pg_fill(V3, 3 * kv_dim, st)
146
147 // ---- EQUIV: 3 rounds paged == contiguous ------------------------
148 let cc: *NxF32KVCache = nx_f32_kv_cache_alloc(PG_NL, PG_NKV, 64, PG_HD)
149 let pool: *NxPagedPool = nx_pkv_pool_new(16, PG_NL, kv_dim)
150 let sA: *NxPagedSeq = nx_pkv_seq_new(pool, 64)
151 let oc: *i64 = sys_mmap(5 * q_dim * 8) as *i64
152 let op: *i64 = sys_mmap(5 * q_dim * 8) as *i64
153 let r1: i64 = pg_round(cc, sA, Q1, K1, V1, 5, scale, oc, op)
154 if r1 != 0 { return 20 + r1 }
155 let r2: i64 = pg_round(cc, sA, Q2, K2, V2, 1, scale, oc, op)
156 if r2 != 0 { return 30 + r2 }
157 let r3: i64 = pg_round(cc, sA, Q3, K3, V3, 3, scale, oc, op)
158 if r3 != 0 { return 40 + r3 }
159 fmt_puts("PKV EQUIV paged==contiguous BIT-EXACT (5+1+3 rows, 2 layers, GQA) OK"); pg_nl2()
160
161 // ---- FORK + COW divergence --------------------------------------
162 // fresh pool/seq: prefill 5, fork, diverge.
163 let pool2: *NxPagedPool = nx_pkv_pool_new(16, PG_NL, kv_dim)
164 let fA: *NxPagedSeq = nx_pkv_seq_new(pool2, 64)
165 let ccA: *NxF32KVCache = nx_f32_kv_cache_alloc(PG_NL, PG_NKV, 64, PG_HD)
166 let rA1: i64 = pg_round(ccA, fA, Q1, K1, V1, 5, scale, oc, op)
167 if rA1 != 0 { return 50 + rA1 }
168
169 let fB: *NxPagedSeq = nx_pkv_seq_fork(fA)
170 if pool2.refcnt[fA.pt[0]] != 2 { return 61 } // shared prompt block
171
172 // divergent rows RA (2 tokens) vs RB (2 tokens)
173 let QA: *i64 = sys_mmap(2 * q_dim * 8) as *i64
174 let KA: *i64 = sys_mmap(2 * kv_dim * 8) as *i64
175 let VA: *i64 = sys_mmap(2 * kv_dim * 8) as *i64
176 pg_fill(QA, 2 * q_dim, st); pg_fill(KA, 2 * kv_dim, st); pg_fill(VA, 2 * kv_dim, st)
177 let QB: *i64 = sys_mmap(2 * q_dim * 8) as *i64
178 let KB: *i64 = sys_mmap(2 * kv_dim * 8) as *i64
179 let VB: *i64 = sys_mmap(2 * kv_dim * 8) as *i64
180 pg_fill(QB, 2 * q_dim, st); pg_fill(KB, 2 * kv_dim, st); pg_fill(VB, 2 * kv_dim, st)
181
182 let opA: *i64 = sys_mmap(PG_NL * 2 * q_dim * 8) as *i64
183 let opB: *i64 = sys_mmap(PG_NL * 2 * q_dim * 8) as *i64
184 let rdA: i64 = pg_round_paged_only(fA, QA, KA, VA, 2, scale, opA)
185 if rdA != 0 { return 70 + rdA }
186 let rdB: i64 = pg_round_paged_only(fB, QB, KB, VB, 2, scale, opB) // COW fires here
187 if rdB != 0 { return 80 + rdB }
188 if fA.pt[0] == fB.pt[0] { return 62 } // must have diverged
189 if pool2.refcnt[fA.pt[0]] != 1 { return 63 }
190 if pool2.refcnt[fB.pt[0]] != 1 { return 63 }
191
192 // ground truth: fresh contiguous replays.
193 let gA: *NxF32KVCache = nx_f32_kv_cache_alloc(PG_NL, PG_NKV, 64, PG_HD)
194 let gB: *NxF32KVCache = nx_f32_kv_cache_alloc(PG_NL, PG_NKV, 64, PG_HD)
195 let sink: *i64 = sys_mmap(PG_NL * 5 * q_dim * 8) as *i64
196 let g1: i64 = pg_round_contig_only(gA, Q1, K1, V1, 5, scale, sink)
197 if g1 != 0 { return 90 }
198 let g2: i64 = pg_round_contig_only(gB, Q1, K1, V1, 5, scale, sink)
199 if g2 != 0 { return 90 }
200 let ocA: *i64 = sys_mmap(PG_NL * 2 * q_dim * 8) as *i64
201 let ocB: *i64 = sys_mmap(PG_NL * 2 * q_dim * 8) as *i64
202 let g3: i64 = pg_round_contig_only(gA, QA, KA, VA, 2, scale, ocA)
203 if g3 != 0 { return 91 }
204 let g4: i64 = pg_round_contig_only(gB, QB, KB, VB, 2, scale, ocB)
205 if g4 != 0 { return 91 }
206 if pg_cmp(opA, ocA, PG_NL * 2 * q_dim) != 1 { return 92 }
207 if pg_cmp(opB, ocB, PG_NL * 2 * q_dim) != 1 { return 93 }
208 fmt_puts("PKV FORK+COW both forks BIT-EXACT vs contiguous replays OK"); pg_nl2()
209
210 // ---- REWIND (speculative truncation contract) --------------------
211 let tr: nx_int = nx_pkv_seq_truncate(fA, 5)
212 if tr != NX_PKV_OK { return 94 }
213 let opR: *i64 = sys_mmap(PG_NL * 2 * q_dim * 8) as *i64
214 let rdR: i64 = pg_round_paged_only(fA, QB, KB, VB, 2, scale, opR)
215 if rdR != 0 { return 95 }
216 if pg_cmp(opR, ocB, PG_NL * 2 * q_dim) != 1 { return 96 }
217 fmt_puts("PKV REWIND truncate-then-append BIT-EXACT (spec-decode contract) OK"); pg_nl2()
218
219 fmt_puts("LIAR-KILL pool-mechanics=1 equiv-bitexact=1 fork-cow=1 rewind=1"); pg_nl2()
220 fmt_puts("PAGED_KV_GATE DONE"); pg_nl2()
221 return 0
222}