nx_llm_layerlens_probe.nx source
↩ module page · 115 lines · 5200 B
1// nx_llm_layerlens_probe.nx -- SOVEREIGN logit-lens bisection (no external oracle). Prefills "The capital of
2// France is" (tokens 0..3), then runs the last token ("is") through the 24 layers one at a time, applying the
3// final RMSNorm + lm-head after EACH layer and printing the top token. In a healthy forward the prediction
4// sharpens toward the answer with depth; a layer where it DEGRADES or flips character localizes the residual bug.
5// Fully sovereign -- reuses block_v4/rmsnorm/matmul_t. expect_exit: 0
6import "nx_syscalls.nx"
7import "nx_tier.nx"
8import "nx_bpe.nx"
9import "nx_gguf.nx"
10import "nx_gguf_load.nx"
11import "nx_gguf_meta.nx"
12import "nx_f32.nx"
13import "nx_f32_rmsnorm.nx"
14import "nx_f32_matmul_t.nx"
15import "nx_f32_kv_cache.nx"
16import "nx_f32_lazy_weight.nx"
17import "nx_f32_llama_block.nx"
18import "nx_f32_llama_block_v4.nx"
19import "nx_f32_llama_stack_v4.nx"
20import "nx_f32_llama_layer_lazy_load.nx"
21import "nx_f32_llm.nx"
22import "nx_f32_llm_v4.nx"
23import "nx_f32_llm_read_dims.nx"
24import "nx_f32_bpe_load.nx"
25import "nx_f32_llm_special_tokens.nx"
26
27func pr_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
28func pr_num(v: i64) -> i64 { let bb: *u8=sys_mmap(28); var m: i64=v; if m<0{sys_write(1,"-" as *u8,1);m=0-m} let t: *u8=sys_mmap(28); var k: i64=0; if m==0{t[0]=48 as u8;k=1} while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} var i: i64=0; while i<k{bb[i]=t[k-1-i];i=i+1} sys_write(1,bb,k); return 0 }
29
30// argmax over f32 logits (raw bits) -- sign-magnitude compare.
31func amax(logits: *i64, V: nx_int) -> nx_int {
32 var best: nx_int = 0
33 var bv: i64 = logits[0] & 0xFFFFFFFF
34 var i: nx_int = 1
35 while i < V {
36 let lv: i64 = logits[i] & 0xFFFFFFFF
37 let sa: i64 = (bv >> 31) & 1
38 let sb: i64 = (lv >> 31) & 1
39 var gt: nx_int = 0
40 if sa == 1 { if sb == 1 { if lv < bv { gt = 1 } } else { gt = 1 } }
41 else { if sb == 0 { if lv > bv { gt = 1 } } }
42 if gt == 1 { bv = lv; best = i }
43 i = i + 1
44 }
45 return best
46}
47
48func main() -> i64 {
49 let path: *u8 = "/tmp/nx_real_model.gguf" as *u8
50 let len_out: *i64 = sys_mmap(8) as *i64
51 let buf: *u8 = sys_read_file(path, len_out)
52 if buf == (0 as *u8) { return 10 }
53 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
54 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { return 20 }
55 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
56 let oe: *i64 = sys_mmap(8) as *i64
57 if nx_f32_llm_read_dims_from_gguf(buf, len_out[0], hdr, model, oe) != NX_FLD_OK { return 30 }
58 if nx_f32_llm_load_weights_v4_from_gguf(buf, hdr, model, oe) != NX_FLV4_OK { return 40 }
59 let vocab: *NxBpeVocab = nx_bpe_vocab_new(67108864, 262144, 524288)
60 let nt: *i64 = sys_mmap(8) as *i64
61 let nm: *i64 = sys_mmap(8) as *i64
62 if nx_f32_bpe_load_from_gguf(buf, len_out[0], hdr, vocab, nt, nm, oe) != NX_FBL_OK { return 50 }
63
64 let eps: i64 = 0x358637BD
65 let attn_scale: i64 = 0x3E000000
66 let rope_base: i64 = 0x415D0EAB
67 let H: nx_int = model.hidden_dim
68 let V: nx_int = model.vocab_size
69 let NL: nx_int = model.n_layers
70
71 let toks: *i64 = sys_mmap(16 * 8) as *i64
72 toks[0]=785; toks[1]=6722; toks[2]=315; toks[3]=9625; toks[4]=374
73 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(NL, model.n_kv_heads, 64, model.head_dim)
74 let logits: *i64 = sys_mmap(V * 8) as *i64
75 // prefill tokens 0..3 (positions 0-3); cache seq_len -> 4
76 var p: nx_int = 0
77 while p < 4 {
78 let one: *i64 = (((toks as i64) + p * 8)) as *i64
79 if nx_f32_llm_forward_v4(model, one, 1, cache, eps, attn_scale, rope_base, 1, logits) != NX_FLV4_OK { return 60 }
80 p = p + 1
81 }
82
83 pr_puts("=== LOGIT-LENS by depth for last token 'is' (want convergence toward ' Paris') ===\n" as *u8)
84 // token 4 through layers, lens after each
85 let bufA: *i64 = sys_mmap(H * 8) as *i64
86 let bufB: *i64 = sys_mmap(H * 8) as *i64
87 var d: nx_int = 0
88 while d < H { bufA[d] = model.embed_weights[374 * H + d]; d = d + 1 }
89 let normed: *i64 = sys_mmap(H * 8) as *i64
90 let llog: *i64 = sys_mmap(V * 8) as *i64
91 var cur: *i64 = bufA
92 var use_a: nx_int = 1
93 var L: nx_int = 0
94 while L < NL {
95 let layer: *NxF32LlamaLayerLazy = model.layers[L] as *NxF32LlamaLayerLazy
96 var dst: *i64 = bufB
97 if use_a == 0 { dst = bufA }
98 nx_f32_llama_block_forward_v4(cur, 1, H, model.n_heads, model.n_kv_heads, model.head_dim, model.ffn_dim,
99 layer, cache, L, eps, attn_scale, rope_base, 1, dst)
100 cur = dst
101 use_a = 1 - use_a
102 // lens
103 nx_f32_rmsnorm(cur, model.gamma_out, H, eps, normed)
104 nx_f32_matmul_t(normed, model.lm_head, llog, 1, H, V)
105 let best: nx_int = amax(llog, V)
106 let tb: *i64 = sys_mmap(8) as *i64
107 tb[0] = best as i64
108 let db: *u8 = sys_mmap(64)
109 let dn: nx_int = nx_bpe_decode_bytelevel(vocab, tb, 1, db)
110 pr_puts(" L"); pr_num(L as i64); pr_puts(" -> tok="); pr_num(best as i64); pr_puts(" [")
111 sys_write(1, db, dn as i64); pr_puts("]\n" as *u8)
112 L = L + 1
113 }
114 return 0
115}