nx_lowrank_kv_real.nx source
↩ module page · 366 lines · 18036 B
1// nx_lowrank_kv_real.nx -- M4c-1: is REAL Qwen2.5-0.5B KV actually low-rank? (the empirical foundation of the
2// low-rank-KV competition). Loads the sovereign-downloaded model, runs a real forward, then measures the
3// SINGULAR SPECTRUM of the real cached K: cumulative energy in the top-r directions vs total. If the top few
4// directions hold most of the energy, real KV IS low-rank -> our per-token compression (M3) helps real models.
5//
6// Reuses: the proven real-LM load+forward (nx_f32_llm_forward_v4, identical to nx_f32_llm_live_load_test) +
7// the M1 reorthogonalization lesson for the eigen-extraction. Arithmetic = nx_f32 (software, matches the LM
8// stack) + nx_f32_hw f32_of/f32_int for int<->f32. Sovereign: nx_cc->nxasm, no gcc. license_tier: ORIGINAL
9import "nx_syscalls.nx"
10import "nx_itoa_lib.nx" // shared MSB-first emitter (zero-alloc)
11import "nx_tier.nx"
12import "nx_gguf.nx"
13import "nx_gguf_load.nx"
14import "nx_gguf_meta.nx"
15import "nx_f32.nx"
16import "nx_f32_hw.nx"
17import "nx_f32_kv_cache.nx"
18import "nx_f32_lazy_weight.nx"
19import "nx_f32_llama_block.nx"
20import "nx_f32_llama_block_v4.nx"
21import "nx_f32_llama_stack_v4.nx"
22import "nx_f32_llama_layer_lazy_load.nx"
23import "nx_f32_llm.nx"
24import "nx_f32_llm_v4.nx"
25import "nx_f32_llm_read_dims.nx"
26
27func r_puts(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 }
28// MIGRATED to the shared emitter (debt 1785563586). The old body mmapped a scratch buffer
29// per call and never freed it. At PAGE granularity that is 4096B leaked PER CALL -- the
30// defect that took 28.5GB of a 36GB host in nx_ts_lumadiff (2MB input, ~3.66M calls).
31// nxi_* is MSB-first, allocates NOTHING, and emits identical bytes including the sign.
32func r_putn(v: i64) -> i64 { nxi_out(v); return 0 }
33func r_is_zero(x: i64) -> i64 { if (x & 0x7fffffff) == 0 { return 1 } return 0 }
34
35// ---- f32 eigen-spectrum (nx_f32 arithmetic + reorthogonalization) ----
36func r_dot(a: *i64, b: *i64, n: i64) -> i64 {
37 var s: i64 = 0
38 var i: i64 = 0
39 while i < n { s = nx_f32_add(s, nx_f32_mul(a[i], b[i])); i = i + 1 }
40 return s
41}
42func r_normalize(v: *i64, n: i64) -> i64 {
43 let ns: i64 = r_dot(v, v, n)
44 if r_is_zero(ns) == 1 { return -1 }
45 let nrm: i64 = nx_f32_sqrt(ns)
46 if r_is_zero(nrm) == 1 { return -1 }
47 let inv: i64 = nx_f32_div(f32_of(1), nrm)
48 var i: i64 = 0
49 while i < n { v[i] = nx_f32_mul(v[i], inv); i = i + 1 }
50 return 0
51}
52func r_gmatvec(G: *i64, v: *i64, d: i64, out: *i64) -> i64 {
53 var a: i64 = 0
54 while a < d {
55 var s: i64 = 0
56 var b: i64 = 0
57 while b < d { s = nx_f32_add(s, nx_f32_mul(G[a * d + b], v[b])); b = b + 1 }
58 out[a] = s
59 a = a + 1
60 }
61 return 0
62}
63func r_deflate1(v: *i64, basis: *i64, found: i64, d: i64) -> i64 {
64 var j: i64 = 0
65 while j < found {
66 var s: i64 = 0
67 var c: i64 = 0
68 while c < d { s = nx_f32_add(s, nx_f32_mul(v[c], basis[j * d + c])); c = c + 1 }
69 c = 0
70 while c < d { v[c] = nx_f32_sub(v[c], nx_f32_mul(s, basis[j * d + c])); c = c + 1 }
71 j = j + 1
72 }
73 return 0
74}
75func r_reorth(v: *i64, basis: *i64, found: i64, d: i64) -> i64 {
76 r_deflate1(v, basis, found, d); r_deflate1(v, basis, found, d); return 0
77}
78// one eigenvector via power iteration; returns eigenvalue lambda = v^T G v (f32); stores v in basis[found].
79func r_power_one(G: *i64, d: i64, basis: *i64, found: i64, v: *i64, w: *i64) -> i64 {
80 var c: i64 = 0
81 while c < d { v[c] = 0; c = c + 1 }
82 v[found % d] = f32_of(1)
83 r_reorth(v, basis, found, d)
84 if r_normalize(v, d) < 0 { return 0 }
85 var it: i64 = 0
86 while it < 64 {
87 r_gmatvec(G, v, d, w)
88 r_reorth(w, basis, found, d)
89 if r_is_zero(r_dot(w, w, d)) == 1 { it = 64 } else {
90 if r_normalize(w, d) < 0 { it = 64 } else {
91 c = 0
92 while c < d { v[c] = w[c]; c = c + 1 }
93 it = it + 1
94 }
95 }
96 }
97 c = 0
98 while c < d { basis[found * d + c] = v[c]; c = c + 1 }
99 r_gmatvec(G, v, d, w)
100 return r_dot(v, w, d) // lambda = v^T G v
101}
102
103// G = K^T K (kv_dim x kv_dim) from K[seq x kv_dim]
104func r_gram(K: *i64, seq: i64, d: i64, G: *i64) -> i64 {
105 var a: i64 = 0
106 while a < d {
107 var b: i64 = 0
108 while b < d {
109 var s: i64 = 0
110 var t: i64 = 0
111 while t < seq { s = nx_f32_add(s, nx_f32_mul(K[t * d + a], K[t * d + b])); t = t + 1 }
112 G[a * d + b] = s
113 b = b + 1
114 }
115 a = a + 1
116 }
117 return 0
118}
119func r_trace(G: *i64, d: i64) -> i64 {
120 var s: i64 = 0
121 var i: i64 = 0
122 while i < d { s = nx_f32_add(s, G[i * d + i]); i = i + 1 }
123 return s
124}
125
126// measure + print the cumulative top-r energy ratio (permille) of K's spectrum.
127func r_spectrum(K: *i64, seq: i64, d: i64, label: *u8) -> i64 {
128 let G: *i64 = sys_mmap(d * d * 8) as *i64
129 r_gram(K, seq, d, G)
130 let total: i64 = r_trace(G, d)
131 r_puts(label); r_puts(" kv_dim="); r_putn(d); r_puts(" seq="); r_putn(seq); r_puts(" top-r energy: ")
132 if r_is_zero(total) == 1 { r_puts("(zero K)\n"); return 0 }
133 let basis: *i64 = sys_mmap(8 * d * 8) as *i64
134 let v: *i64 = sys_mmap(d * 8) as *i64
135 let w: *i64 = sys_mmap(d * 8) as *i64
136 var cum: i64 = 0
137 var f: i64 = 0
138 while f < 8 {
139 let lam: i64 = r_power_one(G, d, basis, f, v, w)
140 cum = nx_f32_add(cum, lam)
141 // print at r = 1,2,4,8
142 if f == 0 { r_puts("r1="); r_putn(f32_int(nx_f32_mul(nx_f32_div(cum, total), f32_of(1000)))); r_puts(" ") }
143 if f == 1 { r_puts("r2="); r_putn(f32_int(nx_f32_mul(nx_f32_div(cum, total), f32_of(1000)))); r_puts(" ") }
144 if f == 3 { r_puts("r4="); r_putn(f32_int(nx_f32_mul(nx_f32_div(cum, total), f32_of(1000)))); r_puts(" ") }
145 if f == 7 { r_puts("r8="); r_putn(f32_int(nx_f32_mul(nx_f32_div(cum, total), f32_of(1000)))); r_puts("permille") }
146 f = f + 1
147 }
148 r_puts("\n")
149 return 0
150}
151
152// the nx_lowrank_rank_plan RULE applied to a real cache tensor: min r with cumulative energy >= thresh(permille)
153func r_rank_for_threshold(K: *i64, seq: i64, d: i64, thresh: i64) -> i64 {
154 let G: *i64 = sys_mmap(d * d * 8) as *i64
155 r_gram(K, seq, d, G)
156 let total: i64 = r_trace(G, d)
157 if r_is_zero(total) == 1 { return 0 }
158 let basis: *i64 = sys_mmap(seq * d * 8) as *i64
159 let v: *i64 = sys_mmap(d * 8) as *i64
160 let w: *i64 = sys_mmap(d * 8) as *i64
161 var cum: i64 = 0
162 var f: i64 = 0
163 while f < seq {
164 let lam: i64 = r_power_one(G, d, basis, f, v, w)
165 cum = nx_f32_add(cum, lam)
166 let ratio: i64 = f32_int(nx_f32_mul(nx_f32_div(cum, total), f32_of(1000)))
167 if ratio >= thresh { return f + 1 }
168 f = f + 1
169 }
170 return seq
171}
172
173func r_argmax(a: *i64, n: i64) -> i64 {
174 var bi: i64 = 0
175 var best: i64 = a[0]
176 var i: i64 = 1
177 while i < n {
178 if nx_f32_gt(a[i], best) == 1 { best = a[i]; bi = i }
179 i = i + 1
180 }
181 return bi
182}
183
184// compress K[seq x d] IN PLACE to rank r (per-token SVD projection -- the M3 compression on REAL cached KV).
185func r_compress_layer(K: *i64, seq: i64, d: i64, rank: i64) -> i64 {
186 let G: *i64 = sys_mmap(d * d * 8) as *i64
187 r_gram(K, seq, d, G)
188 let basis: *i64 = sys_mmap(rank * d * 8) as *i64
189 let vv: *i64 = sys_mmap(d * 8) as *i64
190 let ww: *i64 = sys_mmap(d * 8) as *i64
191 var f: i64 = 0
192 while f < rank { r_power_one(G, d, basis, f, vv, ww); f = f + 1 }
193 let coords: *i64 = sys_mmap(rank * 8) as *i64
194 var t: i64 = 0
195 while t < seq {
196 var j: i64 = 0
197 while j < rank {
198 var s: i64 = 0
199 var c: i64 = 0
200 while c < d { s = nx_f32_add(s, nx_f32_mul(K[t * d + c], basis[j * d + c])); c = c + 1 }
201 coords[j] = s
202 j = j + 1
203 }
204 var c2: i64 = 0
205 while c2 < d {
206 var s2: i64 = 0
207 j = 0
208 while j < rank { s2 = nx_f32_add(s2, nx_f32_mul(coords[j], basis[j * d + c2])); j = j + 1 }
209 K[t * d + c2] = s2
210 c2 = c2 + 1
211 }
212 t = t + 1
213 }
214 return 0
215}
216
217func r_copy(dst: *i64, src: *i64, n: i64) -> i64 { var i: i64 = 0; while i < n { dst[i] = src[i]; i = i + 1 } return 0 }
218
219// relative L2 error of logit vector la vs lb, in permille (whole vocab)
220func r_logit_rel(la: *i64, lb: *i64, n: i64) -> i64 {
221 var num: i64 = 0
222 var den: i64 = 0
223 var i: i64 = 0
224 while i < n {
225 let d: i64 = nx_f32_sub(la[i], lb[i])
226 num = nx_f32_add(num, nx_f32_mul(d, d))
227 den = nx_f32_add(den, nx_f32_mul(lb[i], lb[i]))
228 i = i + 1
229 }
230 if nx_f32_gt(den, 0) == 0 { return -1 }
231 return f32_int(nx_f32_mul(nx_f32_sqrt(nx_f32_div(num, den)), f32_of(1000)))
232}
233
234func main() -> i64 {
235 // ---- load the sovereign-downloaded model (identical to nx_f32_llm_live_load_test) ----
236 let path: *u8 = "/home/elderwesto/nx_stage/nx_real_model.gguf\x00"
237 let len_out: *i64 = sys_mmap(8) as *i64
238 let buf: *u8 = sys_read_file(path, len_out)
239 if buf == (0 as *u8) { r_puts("M4c: model read failed\n"); return 10 }
240 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
241 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { r_puts("M4c: gguf parse failed\n"); return 20 }
242 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
243 let oe: *i64 = sys_mmap(8) as *i64
244 if nx_f32_llm_read_dims_from_gguf(buf, len_out[0], hdr, model, oe) != NX_FLD_OK { r_puts("M4c: dims failed\n"); return 30 }
245 if nx_f32_llm_load_weights_v4_from_gguf(buf, hdr, model, oe) != NX_FLV4_OK { r_puts("M4c: weights failed\n"); return 40 }
246
247 let kv_dim: i64 = model.n_kv_heads * model.head_dim
248 r_puts("MODEL loaded: layers="); r_putn(model.n_layers); r_puts(" kv_heads="); r_putn(model.n_kv_heads)
249 r_puts(" head_dim="); r_putn(model.head_dim); r_puts(" kv_dim="); r_putn(kv_dim); r_puts("\n")
250
251 // ---- forward a 64-token prompt -> realistic-context KV cache (seq=8 too small to gauge real compressibility; operator cleared the slow forward 2026-06-17) ----
252 let seq: i64 = 64
253 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(model.n_layers, model.n_kv_heads, seq + 4, model.head_dim) // +4 headroom so the decode token (position seq) fits (was overflowing at max_seq=seq)
254 if cache == (0 as *NxF32KVCache) { r_puts("M4c: cache alloc failed\n"); return 50 }
255 let tokens: *i64 = sys_mmap(seq * 8) as *i64
256 var i: i64 = 0
257 while i < seq { tokens[i] = 1000 + i * 777; i = i + 1 } // DIVERSE real subwords (id>=256) -> rank-rich real cache (not rank-1)
258 let logits: *i64 = sys_mmap(seq * model.vocab_size * 8) as *i64
259 let eps: i64 = 0x322BCC77
260 let attn_scale: i64 = 0x3E000000
261 let rope_base: i64 = 0x4548F000
262 let ft0: i64 = sys_now_ms()
263 let v: nx_int = nx_f32_llm_forward_v4(model, tokens, seq, cache, eps, attn_scale, rope_base, 1, logits)
264 let ft1: i64 = sys_now_ms()
265 if v != NX_FLV4_OK { r_puts("M4c: forward failed v="); r_putn(v); r_puts("\n"); return 60 }
266 r_puts("FORWARD ok seq_len="); r_putn(nx_f32_kv_cache_get_seq_len(cache)); r_puts(" prefill(64tok x 24L) = "); r_putn(ft1 - ft0); r_puts("ms (SSE Q4_K matmul+dequant)\n\n")
267
268 r_puts("=== REAL Qwen2.5-0.5B KV singular spectrum (is real KV low-rank? higher permille @ small r = yes) ===\n")
269 // probe a few layers (shallow, middle, deep)
270 r_spectrum(nx_f32_kv_cache_get_K_layer(cache, 0), seq, kv_dim, " K layer0 " as *u8)
271 r_spectrum(nx_f32_kv_cache_get_K_layer(cache, 12), seq, kv_dim, " K layer12" as *u8)
272 r_spectrum(nx_f32_kv_cache_get_K_layer(cache, 23), seq, kv_dim, " K layer23" as *u8)
273 r_spectrum(nx_f32_kv_cache_get_V_layer(cache, 12), seq, kv_dim, " V layer12" as *u8)
274
275 r_puts("\n(seq=64 so rank<=64; r ratios = fraction of KV covariance energy in the top r of <=64 dirs)\n\n")
276
277 // ===== M4d: PRECISE per-LAYER rank SCHEDULE from the FULL measured spectrum (read-only; before M4c-2) =====
278 // Operator cleared the one-time slow forward. Apply the nx_lowrank_rank_plan rule (min r with cumulative
279 // covariance energy >= 90%) to EVERY layer x {K,V} -> the exact schedule, not just M4c-1's sparse anchors.
280 r_puts("=== M4d: data-driven per-LAYER rank schedule (rule: min r with cum energy >= 90%, seq=64 so r<=64) ===\n")
281 let THRESH: i64 = 900
282 var sumK: i64 = 0
283 var sumV: i64 = 0
284 let schedK: *i64 = sys_mmap(model.n_layers * 8) as *i64 // per-layer 90%-energy K rank (fed to M4c-2 validation)
285 let schedV: *i64 = sys_mmap(model.n_layers * 8) as *i64
286 var Lx: i64 = 0
287 while Lx < model.n_layers {
288 let rK: i64 = r_rank_for_threshold(nx_f32_kv_cache_get_K_layer(cache, Lx), seq, kv_dim, THRESH)
289 let rV: i64 = r_rank_for_threshold(nx_f32_kv_cache_get_V_layer(cache, Lx), seq, kv_dim, THRESH)
290 sumK = sumK + rK; sumV = sumV + rV
291 schedK[Lx] = rK; schedV[Lx] = rV
292 r_puts(" L"); r_putn(Lx); r_puts(": rK="); r_putn(rK); r_puts(" rV="); r_putn(rV); r_puts("\n")
293 Lx = Lx + 1
294 }
295 r_puts(" TOTAL over "); r_putn(model.n_layers); r_puts(" layers: sum rK="); r_putn(sumK); r_puts(" sum rV="); r_putn(sumV)
296 r_puts(" (avg K="); r_putn(sumK / model.n_layers); r_puts(" V="); r_putn(sumV / model.n_layers); r_puts(")\n")
297 let LCTX: i64 = 512
298 let full_cells: i64 = 2 * model.n_layers * LCTX * kv_dim
299 let uni_cells: i64 = model.n_layers * 2 * 4 * (LCTX + kv_dim)
300 let dd_cells: i64 = (sumK + sumV) * (LCTX + kv_dim)
301 r_puts(" VRAM @seq=512 (per-token-SVD cells): full="); r_putn(full_cells)
302 r_puts(" uniform-r4="); r_putn(uni_cells); r_puts(" ("); r_putn(full_cells / uni_cells); r_puts("x)")
303 r_puts(" data-driven="); r_putn(dd_cells); r_puts(" ("); r_putn(full_cells / dd_cells); r_puts("x)\n")
304 var m4d_ok: i64 = 1
305 if sumK < sumV { r_puts(" M4d G1 K compresses harder than V per-layer (sum rK<rV): PASS\n") } else { r_puts(" M4d G1: FAIL (K not < V)\n"); m4d_ok = 0 }
306 if dd_cells < full_cells { r_puts(" M4d G2 data-driven schedule < full VRAM: PASS\n") } else { r_puts(" M4d G2: FAIL\n"); m4d_ok = 0 }
307 if m4d_ok == 1 { r_puts(" M4d VERDICT=GREEN (precise per-layer schedule from the real full spectrum)\n\n") } else { r_puts(" M4d VERDICT=RED\n\n") }
308
309 // ===== M4c-2 COMPETITION: REAL model output preservation -- full KV vs low-rank KV vs ZERO-cache control.
310 // Diverse context -> rank-rich real cache. Metric = logit-distribution rel error vs the full-KV output.
311 // The ZERO-cache control proves the decode actually USES the cache (rel_zero must be >>0); then if r=4
312 // keeps rel small, low-rank KV genuinely preserves the real model output. (No forward surgery: cache
313 // backup/restore + incremental decode.) =====
314 // (uniform rank=4 removed -- M4c-2 now compresses to the per-layer M4d schedule schedK[L]/schedV[L])
315 let l_full: *i64 = sys_mmap(model.vocab_size * 8) as *i64
316 let l_r4: *i64 = sys_mmap(model.vocab_size * 8) as *i64
317 let l_zero: *i64 = sys_mmap(model.vocab_size * 8) as *i64
318 let tok9: *i64 = sys_mmap(8) as *i64
319 tok9[0] = 1000 + seq * 777
320 let total: i64 = model.n_layers * cache.max_seq_len * kv_dim
321 let bkK: *i64 = sys_mmap(total * 8) as *i64
322 let bkV: *i64 = sys_mmap(total * 8) as *i64
323
324 // (full) decode + backup the prefilled cache
325 cache.seq_len = seq
326 nx_f32_llm_forward_v4(model, tok9, 1, cache, eps, attn_scale, rope_base, 1, l_full)
327 r_copy(bkK, cache.cache_K, total)
328 r_copy(bkV, cache.cache_V, total)
329
330 // (a) r=4 low-rank KV
331 var L: i64 = 0
332 while L < model.n_layers {
333 r_compress_layer(nx_f32_kv_cache_get_K_layer(cache, L), seq, kv_dim, schedK[L]) // data-driven: this layer's 90%-energy K rank (M4d)
334 r_compress_layer(nx_f32_kv_cache_get_V_layer(cache, L), seq, kv_dim, schedV[L]) // data-driven: this layer's 90%-energy V rank (M4d)
335 L = L + 1
336 }
337 cache.seq_len = seq
338 nx_f32_llm_forward_v4(model, tok9, 1, cache, eps, attn_scale, rope_base, 1, l_r4)
339 let rel_r4: i64 = r_logit_rel(l_r4, l_full, model.vocab_size)
340
341 // (b) ZERO-cache control: restore, then zero the prefilled K/V
342 r_copy(cache.cache_K, bkK, total)
343 r_copy(cache.cache_V, bkV, total)
344 L = 0
345 while L < model.n_layers {
346 let Kp: *i64 = nx_f32_kv_cache_get_K_layer(cache, L)
347 let Vp: *i64 = nx_f32_kv_cache_get_V_layer(cache, L)
348 var z: i64 = 0
349 while z < seq * kv_dim { Kp[z] = 0; Vp[z] = 0; z = z + 1 }
350 L = L + 1
351 }
352 cache.seq_len = seq
353 nx_f32_llm_forward_v4(model, tok9, 1, cache, eps, attn_scale, rope_base, 1, l_zero)
354 let rel_zero: i64 = r_logit_rel(l_zero, l_full, model.vocab_size)
355
356 r_puts("=== M4c-2 COMPETITION: real Qwen output preservation (diverse context, "); r_putn(model.n_layers); r_puts(" layers) ===\n")
357 r_puts(" DATA-DRIVEN 90%-energy schedule (per-layer M4d ranks, avg K="); r_putn(sumK / model.n_layers); r_puts(" V="); r_putn(sumV / model.n_layers); r_puts("): logit rel error vs full = "); r_putn(rel_r4); r_puts("permille (uniform r=4 was ~615)\n")
358 r_puts(" ZERO-cache ctrl : logit rel error vs full = "); r_putn(rel_zero); r_puts("permille (>>0 = the decode USES the cache)\n")
359 if rel_zero > rel_r4 + 50 {
360 if rel_r4 < 100 { r_puts(" VERDICT: the data-driven 90% schedule PRESERVES the real output ("); r_putn(rel_r4); r_puts("permille; cache matters: zero->"); r_putn(rel_zero); r_puts(") => the rule WORKS end-to-end\n") }
361 else { r_puts(" the 90% schedule shifted output "); r_putn(rel_r4); r_puts("permille; "); if rel_r4 < 615 { r_puts("BEATS uniform-r4(~615) so data-driven sizing HELPS, but per-layer 90% still COMPOUNDS across 24 layers -> raise the per-layer threshold for tight preservation\n") } else { r_puts("does NOT beat uniform-r4(~615)\n") } }
362 }
363 else { r_puts(" INCONCLUSIVE: zero-cache barely changed the output (rel_zero="); r_putn(rel_zero); r_puts(") -> decode not cache-dependent as set up\n") }
364 sys_exit(0)
365 return 0
366}