nx_f32_transformer_test.nx source

1// nx_f32_transformer_test.nx -- integration smoke. 2// 3// Zero-weight identity check: with all weights = 0 and gamma = 0, 4// every matmul produces 0, every RMSNorm produces 0, both 5// residual additions reduce to "x + 0 = x", so output should 6// equal input unchanged. 7// 8// This is the structural-acceptance gate: proves the full chain 9// of RMSNorm + matmul + softmax + SiLU + residual composes 10// without crashing and preserves the identity at zero weights. 11 12import "nx_syscalls.nx" 13import "nx_tier.nx" 14import "nx_f32.nx" 15import "nx_f32_transformer.nx" 16 17func main() -> i64 { 18 // Verdict gate 19 var vi: nx_int = 0 20 while vi < NX_F32_TB_N_VERDICTS { 21 if nx_f32_tb_verdict_is_valid(vi) != 1 { return 5 + vi } 22 vi = vi + 1 23 } 24 25 let hidden_dim: nx_int = 4 26 let head_dim: nx_int = 4 27 let ffn_dim: nx_int = 8 28 29 // x = [1.0, 2.0, 3.0, 4.0] (f32 raw bits) 30 let x: *i64 = sys_mmap(hidden_dim * 8) as *i64 31 x[0] = 0x3F800000 // 1.0 32 x[1] = 0x40000000 // 2.0 33 x[2] = 0x40400000 // 3.0 34 x[3] = 0x40800000 // 4.0 35 36 // Save copy for post-check. 37 let x_orig: *i64 = sys_mmap(hidden_dim * 8) as *i64 38 var ci: nx_int = 0 39 while ci < hidden_dim { 40 x_orig[ci] = x[ci] 41 ci = ci + 1 42 } 43 44 // All weights and gammas = 0. 45 let gamma_attn: *i64 = sys_mmap(hidden_dim * 8) as *i64 46 let gamma_ffn: *i64 = sys_mmap(hidden_dim * 8) as *i64 47 let W_q: *i64 = sys_mmap(hidden_dim * head_dim * 8) as *i64 48 let W_k: *i64 = sys_mmap(hidden_dim * head_dim * 8) as *i64 49 let W_v: *i64 = sys_mmap(hidden_dim * head_dim * 8) as *i64 50 let W_o: *i64 = sys_mmap(head_dim * hidden_dim * 8) as *i64 51 let W_gate: *i64 = sys_mmap(hidden_dim * ffn_dim * 8) as *i64 52 let W_up: *i64 = sys_mmap(hidden_dim * ffn_dim * 8) as *i64 53 let W_down: *i64 = sys_mmap(ffn_dim * hidden_dim * 8) as *i64 54 55 // sys_mmap returns zeroed memory on this substrate; leave them. 56 57 let w: *NxF32TransformerBlockWeights = sys_mmap(NX_F32_TBW_BYTES) as *NxF32TransformerBlockWeights 58 w.gamma_attn = gamma_attn 59 w.W_q = W_q 60 w.W_k = W_k 61 w.W_v = W_v 62 w.W_o = W_o 63 w.gamma_ffn = gamma_ffn 64 w.W_gate = W_gate 65 w.W_up = W_up 66 w.W_down = W_down 67 68 // attn_scale = 1/sqrt(4) = 0.5 = 0x3F000000 69 let attn_scale: i64 = 0x3F000000 70 // eps = 1e-6 ~= 0x358637BD 71 let eps: i64 = 0x358637BD 72 73 let v: nx_int = nx_f32_transformer_block_forward( 74 x, hidden_dim, head_dim, ffn_dim, 75 w, attn_scale, eps) 76 if v != NX_F32_TB_OK { return 10 + v } 77 78 // With zero weights/gammas: 79 // attn_in = RMSNorm(x, 0) -> 0 (gamma=0 zeros) 80 // Q=K=V = matmul(0, 0) -> 0 81 // raw_score = dot(0,0) = 0; scores[0] = 0 * attn_scale = 0 82 // probs = softmax([0]) = [1.0] 83 // attn_out = probs[0] * V[i] = 1 * 0 = 0 84 // attn_proj = matmul(0, 0) = 0 85 // x = x + 0 = x (PRESERVED) 86 // 87 // ffn_in = RMSNorm(x, 0) = 0 88 // gate = SiLU(0) = 0 (since SiLU(0) = 0 exact) 89 // up = matmul(0, 0) = 0 90 // hidden = 0 * 0 = 0 91 // ffn_proj = matmul(0, 0) = 0 92 // x = x + 0 = x (PRESERVED) 93 94 // Verify x unchanged. 95 if x[0] != x_orig[0] { return 30 } 96 if x[1] != x_orig[1] { return 31 } 97 if x[2] != x_orig[2] { return 32 } 98 if x[3] != x_orig[3] { return 33 } 99 100 return 0 101}