nx_f32_llama_block_test.nx source
↩ module page · 129 lines · 5162 B
1// nx_f32_llama_block_test.nx -- smoke for nx_f32_llama_block.nx.
2//
3// Identity-with-zero-weights test:
4// All weight matrices are zero (mmap default).
5// gamma_attn and gamma_ffn can be anything; with zero W_q/k/v
6// the attention projections are 0; with zero W_o attn_proj=0.
7// x_mid = x + 0 = x. FFN similarly produces 0 with zero W_*.
8// out = x_mid + 0 = x.
9//
10// This exercises the full forward pass (RMSNorm + matmul + RoPE
11// + attn-with-cache + SwiGLU + residual) and verifies the residual
12// path preserves x bit-exact when weights are zero.
13
14import "nx_syscalls.nx"
15import "nx_tier.nx"
16import "nx_f32.nx"
17import "nx_f32_rmsnorm.nx"
18import "nx_f32_matmul.nx"
19import "nx_f32_activations.nx"
20import "nx_f32_rope.nx"
21import "nx_f32_attn_multi.nx"
22import "nx_f32_kv_cache.nx"
23import "nx_f32_attn_cached.nx"
24import "nx_f32_llama_block.nx"
25
26func main() -> i64 {
27 var vi: nx_int = 0
28 while vi < NX_F32_BLK_N_VERDICTS {
29 if nx_f32_blk_verdict_is_valid(vi) != 1 { return 5 + vi }
30 vi = vi + 1
31 }
32
33 // Shape: hidden=4, n_heads=2, n_kv_heads=2, head_dim=2, ffn=8.
34 let hidden_dim: nx_int = 4
35 let n_heads: nx_int = 2
36 let n_kv_heads: nx_int = 2
37 let head_dim: nx_int = 2
38 let ffn_dim: nx_int = 8
39 let n_tokens: nx_int = 1
40 let q_dim: nx_int = n_heads * head_dim // 4
41 let kv_dim: nx_int = n_kv_heads * head_dim // 4
42
43 let layer: *NxF32LlamaLayer = nx_f32_llama_layer_alloc()
44 if layer == (0 as *NxF32LlamaLayer) { return 10 }
45
46 // All weights zero by default from sys_mmap. Allocate them.
47 layer.gamma_attn = sys_mmap(hidden_dim * 8) as *i64
48 layer.gamma_ffn = sys_mmap(hidden_dim * 8) as *i64
49 layer.W_q = sys_mmap(hidden_dim * q_dim * 8) as *i64
50 layer.W_k = sys_mmap(hidden_dim * kv_dim * 8) as *i64
51 layer.W_v = sys_mmap(hidden_dim * kv_dim * 8) as *i64
52 layer.W_o = sys_mmap(q_dim * hidden_dim * 8) as *i64
53 layer.W_gate = sys_mmap(hidden_dim * ffn_dim * 8) as *i64
54 layer.W_up = sys_mmap(hidden_dim * ffn_dim * 8) as *i64
55 layer.W_down = sys_mmap(ffn_dim * hidden_dim * 8) as *i64
56
57 // gamma = 1.0 (so RMSNorm produces meaningful output, but it's
58 // multiplied by zero W's downstream).
59 var gi: nx_int = 0
60 while gi < hidden_dim {
61 layer.gamma_attn[gi] = 0x3F800000
62 layer.gamma_ffn[gi] = 0x3F800000
63 gi = gi + 1
64 }
65
66 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(2, n_kv_heads, 8, head_dim)
67 if cache == (0 as *NxF32KVCache) { return 20 }
68
69 let x_in: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
70 let x_out: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
71 x_in[0] = 0x3F800000 // 1.0
72 x_in[1] = 0x40000000 // 2.0
73 x_in[2] = 0x40400000 // 3.0
74 x_in[3] = 0x40800000 // 4.0
75
76 let eps: i64 = 0x322BCC77 // ~1e-8
77 let attn_scale: i64 = 0x3F3504F3 // 1/sqrt(2)
78 let rope_base: i64 = 0x4548F000 // log(10000) approx (not used; apply_rope=0)
79
80 // Forward pass.
81 let v: nx_int = nx_f32_llama_block_forward_v3(
82 x_in, n_tokens, hidden_dim, n_heads, n_kv_heads, head_dim, ffn_dim,
83 layer, cache, 0, eps, attn_scale, rope_base, 0, x_out)
84 if v != NX_F32_BLK_OK { return 30 + v }
85
86 // Residual identity: out should equal x_in bit-exact since all
87 // projections are zero.
88 if x_out[0] != x_in[0] { return 40 }
89 if x_out[1] != x_in[1] { return 41 }
90 if x_out[2] != x_in[2] { return 42 }
91 if x_out[3] != x_in[3] { return 43 }
92
93 // Cache was appended for layer 0; advance and verify.
94 nx_f32_kv_cache_advance(cache, n_tokens)
95 if nx_f32_kv_cache_get_seq_len(cache) != 1 { return 50 }
96
97 // Second forward pass on same x_in (decode-style).
98 let x_out2: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
99 let v2: nx_int = nx_f32_llama_block_forward_v3(
100 x_in, n_tokens, hidden_dim, n_heads, n_kv_heads, head_dim, ffn_dim,
101 layer, cache, 0, eps, attn_scale, rope_base, 0, x_out2)
102 if v2 != NX_F32_BLK_OK { return 60 + v2 }
103
104 // Same identity result.
105 if x_out2[0] != x_in[0] { return 70 }
106 if x_out2[3] != x_in[3] { return 71 }
107
108 nx_f32_kv_cache_advance(cache, n_tokens)
109 if nx_f32_kv_cache_get_seq_len(cache) != 2 { return 80 }
110
111 // Test RoPE-enabled path: weights still zero so output is still
112 // identity (RoPE rotates Q/K but Q=K=0 so rotation is no-op).
113 let x_out3: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
114 let v3: nx_int = nx_f32_llama_block_forward_v3(
115 x_in, n_tokens, hidden_dim, n_heads, n_kv_heads, head_dim, ffn_dim,
116 layer, cache, 0, eps, attn_scale, rope_base, 1, x_out3)
117 if v3 != NX_F32_BLK_OK { return 90 + v3 }
118 if x_out3[0] != x_in[0] { return 100 }
119 if x_out3[3] != x_in[3] { return 101 }
120
121 // Bad-dim verdict.
122 let x_outE: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
123 let vE: nx_int = nx_f32_llama_block_forward_v3(
124 x_in, 0, hidden_dim, n_heads, n_kv_heads, head_dim, ffn_dim,
125 layer, cache, 0, eps, attn_scale, rope_base, 0, x_outE)
126 if vE != NX_F32_BLK_ERR_BAD_DIM { return 110 }
127
128 return 0
129}