nx_f32_llm_probe.nx source
↩ module page · 167 lines · 8147 B
1// nx_f32_llm_probe.nx -- SOVEREIGN FORWARD PROBE: the tool that turns "why is the LLM forward garbage?"
2// into a one-run solve. Loads the real model, runs ONE token through layer-0 stage by stage, and prints
3// tensor HEALTH after every stage: NaN/Inf count + max biased exponent (magnitude proxy: 127=|x|~1,
4// 134=~128, 142=~32k, 150=~8M, 255=NaN/Inf) + the first 3 raw f32 words. The FIRST stage whose exponent
5// explodes or whose naninf>0 (relative to a sane transformer activation, maxexp~127-140) is the bug.
6// No gen loop, no sampler -- pure forward dissection. Reuses the REAL functions (rmsnorm/lazy_matmul/
7// attn_with_cache/silu) so it measures the real path, not a reimplementation.
8// expect_exit: 0 license_tier: ORIGINAL module: nishi-core.llm.probe
9import "nx_syscalls.nx"
10import "nx_tier.nx"
11import "nx_gguf.nx"
12import "nx_gguf_load.nx"
13import "nx_gguf_meta.nx"
14import "nx_f32.nx"
15import "nx_silu.nx"
16import "nx_f32_rmsnorm.nx"
17import "nx_f32_kv_cache.nx"
18import "nx_f32_lazy_weight.nx"
19import "nx_f32_attn_cached.nx"
20import "nx_f32_llama_block_v4.nx"
21import "nx_f32_llama_layer_lazy_load.nx"
22import "nx_f32_llm_v4.nx"
23import "nx_f32_llm_read_dims.nx"
24const K_MAGIC_9707: i64 = 9707
25
26func 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 }
27func pr_num(v: i64) -> i64 {
28 let bb: *u8=sys_mmap(28); var m: i64=v; let t: *u8=sys_mmap(28); var k: i64=0; if m==0 { t[0]=48 as u8; k=1 }
29 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
30}
31func pr_hex8(v: i64) -> i64 {
32 let buf: *u8 = sys_mmap(8); var i: i64 = 0
33 while i < 8 {
34 let nib: i64 = (v >> ((7-i)*4)) & 0xF
35 if nib < 10 { buf[i] = (48+nib) as u8 } else { buf[i] = (87+nib) as u8 }
36 i = i + 1
37 }
38 sys_write(1, buf, 8); return 0
39}
40// tensor health line: naninf count + max biased exponent + first 3 raw f32 words.
41func probe(label: *u8, vec: *i64, n: nx_int) -> i64 {
42 var naninf: i64 = 0
43 var maxexp: i64 = 0
44 var i: nx_int = 0
45 while i < n {
46 let raw: i64 = vec[i] & 0xFFFFFFFF
47 let exp: i64 = (raw >> 23) & 0xFF
48 if exp == 0xFF { naninf = naninf + 1 }
49 if exp > maxexp { maxexp = exp }
50 i = i + 1
51 }
52 pr_puts(" " as *u8); pr_puts(label)
53 pr_puts(" n=" as *u8); pr_num(n as i64)
54 pr_puts(" naninf=" as *u8); pr_num(naninf)
55 pr_puts(" maxexp=" as *u8); pr_num(maxexp)
56 pr_puts(" v0=" as *u8); pr_hex8(vec[0])
57 pr_puts(" v1=" as *u8); pr_hex8(vec[1])
58 pr_puts(" v2=" as *u8); pr_hex8(vec[2])
59 pr_puts("\n" as *u8)
60 return 0
61}
62
63func main() -> i64 {
64 let path: *u8 = "/tmp/nx_real_model.gguf" as *u8
65 let len_out: *i64 = sys_mmap(8) as *i64
66 let buf: *u8 = sys_read_file(path, len_out)
67 if buf == (0 as *u8) { pr_puts("PROBE: no model at /tmp/nx_real_model.gguf\n" as *u8); return 10 }
68
69 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
70 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { return 20 }
71
72 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
73 let out_err: *i64 = sys_mmap(8) as *i64
74 if nx_f32_llm_read_dims_from_gguf(buf, len_out[0], hdr, model, out_err) != NX_FLD_OK { return 30 }
75 if nx_f32_llm_load_weights_v4_from_gguf(buf, hdr, model, out_err) != NX_FLV4_OK { return 40 }
76
77 let hidden: nx_int = model.hidden_dim
78 let q_dim: nx_int = model.n_heads * model.head_dim
79 let kv_dim: nx_int = model.n_kv_heads * model.head_dim
80 let ffn: nx_int = model.ffn_dim
81 let eps: i64 = 0x358637BD
82 let attn_scale: i64 = 0x3E000000
83
84 pr_puts("=== FORWARD PROBE (layer 0, single token) ===\n" as *u8)
85 pr_puts("dims: hidden=" as *u8); pr_num(hidden as i64)
86 pr_puts(" q_dim=" as *u8); pr_num(q_dim as i64)
87 pr_puts(" kv_dim=" as *u8); pr_num(kv_dim as i64)
88 pr_puts(" ffn=" as *u8); pr_num(ffn as i64)
89 pr_puts(" n_layers=" as *u8); pr_num(model.n_layers as i64)
90 pr_puts(" vocab=" as *u8); pr_num(model.vocab_size as i64); pr_puts("\n" as *u8)
91 pr_puts("gguf: data_off=" as *u8); pr_num(hdr.data_off)
92 pr_puts(" file_size=" as *u8); pr_num(len_out[0])
93 pr_puts(" n_tensors=" as *u8); pr_num(hdr.n_tensors as i64); pr_puts("\n" as *u8)
94
95 // Stage 0: embed lookup for a fixed valid token.
96 let tok: nx_int = K_MAGIC_9707
97 let emb: *i64 = sys_mmap(hidden * 8) as *i64
98 var d: nx_int = 0
99 while d < hidden { emb[d] = model.embed_weights[tok * hidden + d]; d = d + 1 }
100 probe("embed " as *u8, emb, hidden)
101
102 let layer: *NxF32LlamaLayerLazy = model.layers[0] as *NxF32LlamaLayerLazy
103
104 // Stage 1: attn RMSNorm.
105 let nrm: *i64 = sys_mmap(hidden * 8) as *i64
106 nx_f32_rmsnorm(emb, layer.gamma_attn, hidden, eps, nrm)
107 probe("attn_norm " as *u8, nrm, hidden)
108
109 // Stage 2: Q/K/V projections (tests lazy_matmul on the REAL Q5_0/Q8_0 weights).
110 let Q: *i64 = sys_mmap(q_dim * 8) as *i64
111 let K: *i64 = sys_mmap(kv_dim * 8) as *i64
112 let V: *i64 = sys_mmap(kv_dim * 8) as *i64
113 nx_f32_lazy_matmul(nrm, layer.W_q, Q, 1, hidden, q_dim); probe("Q (proj) " as *u8, Q, q_dim)
114 nx_f32_lazy_matmul(nrm, layer.W_k, K, 1, hidden, kv_dim); probe("K (proj) " as *u8, K, kv_dim)
115 nx_f32_lazy_matmul(nrm, layer.W_v, V, 1, hidden, kv_dim); probe("V (proj) " as *u8, V, kv_dim)
116
117 // Stage 3: attention (pos 0, self only -> out ~ V).
118 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(model.n_layers, model.n_kv_heads, 64, model.head_dim)
119 let concat: *i64 = sys_mmap(q_dim * 8) as *i64
120 let va: nx_int = nx_f32_attn_with_cache(Q, K, V, 1, model.n_heads, model.n_kv_heads, model.head_dim, cache, 0, 1, attn_scale, concat)
121 probe("attn_out " as *u8, concat, q_dim)
122
123 // Stage 4: output projection + residual.
124 let proj: *i64 = sys_mmap(hidden * 8) as *i64
125 nx_f32_lazy_matmul(concat, layer.W_o, proj, 1, q_dim, hidden); probe("attn_proj " as *u8, proj, hidden)
126 let xmid: *i64 = sys_mmap(hidden * 8) as *i64
127 var i2: nx_int = 0
128 while i2 < hidden { xmid[i2] = nx_f32_add(emb[i2], proj[i2]); i2 = i2 + 1 }
129 probe("x_mid (res) " as *u8, xmid, hidden)
130
131 // Stage 5: FFN RMSNorm + SwiGLU.
132 let fnrm: *i64 = sys_mmap(hidden * 8) as *i64
133 nx_f32_rmsnorm(xmid, layer.gamma_ffn, hidden, eps, fnrm)
134 probe("ffn_norm " as *u8, fnrm, hidden)
135 let gate: *i64 = sys_mmap(ffn * 8) as *i64
136 let up: *i64 = sys_mmap(ffn * 8) as *i64
137 nx_f32_lazy_matmul(fnrm, layer.W_gate, gate, 1, hidden, ffn); probe("gate " as *u8, gate, ffn)
138 nx_f32_lazy_matmul(fnrm, layer.W_up, up, 1, hidden, ffn); probe("up " as *u8, up, ffn)
139 let hb: *i64 = sys_mmap(ffn * 8) as *i64
140 var j: nx_int = 0
141 while j < ffn { hb[j] = nx_f32_mul(nx_f32_silu(gate[j]), up[j]); j = j + 1 }
142 probe("silu*up " as *u8, hb, ffn)
143 let dproj: *i64 = sys_mmap(hidden * 8) as *i64
144 nx_f32_lazy_matmul(hb, layer.W_down, dproj, 1, ffn, hidden); probe("ffn_down " as *u8, dproj, hidden)
145
146 // === FULL forward (all layers + final norm + lm-head) -> logits, argmax ===
147 let rope_base: i64 = 0x415D0EAB
148 let cache2: *NxF32KVCache = nx_f32_kv_cache_alloc(model.n_layers, model.n_kv_heads, 64, model.head_dim)
149 let toks: *i64 = sys_mmap(8) as *i64
150 toks[0] = tok as i64
151 let logits: *i64 = sys_mmap(model.vocab_size * 8) as *i64
152 let fv: nx_int = nx_f32_llm_forward_v4(model, toks, 1, cache2, eps, attn_scale, rope_base, 1, logits)
153 pr_puts("full forward_v4 rc=" as *u8); pr_num(fv as i64); pr_puts("\n" as *u8)
154 probe("LOGITS " as *u8, logits, model.vocab_size)
155 var best: nx_int = 0
156 var bestv: i64 = logits[0]
157 var ai: nx_int = 1
158 while ai < model.vocab_size {
159 if nx_f32_gt(logits[ai], bestv) == 1 { bestv = logits[ai]; best = ai }
160 ai = ai + 1
161 }
162 pr_puts(" >>> argmax token id = " as *u8); pr_num(best as i64)
163 pr_puts(" (winning logit bits=" as *u8); pr_hex8(bestv); pr_puts(")\n" as *u8)
164
165 pr_puts("=== PROBE DONE: layer-0 healthy => not a blowup. Check logits sanity + argmax token above ===\n" as *u8)
166 return 0
167}