nx_f32_transformer.nx source
↩ module page · 189 lines · 6826 B
1// nx_f32_transformer.nx -- bits-up f32 Llama-class transformer block
2// forward (single-token v1).
3//
4// L8 integration brick. Composes EVERY shipped primitive from this
5// session arc:
6//
7// nx_f32_rmsnorm (mul + add + sqrt + div + cvt)
8// nx_f32_matmul (mul + add)
9// nx_f32_softmax (exp + add + div + max)
10// nx_f32_silu (exp + mul + add + div)
11// nx_f32_mul/add (scalar L4)
12//
13// All bits-up. No libm, no compiler-builtin float, no copied code.
14//
15// Single-token, single-head v1. Multi-token + multi-head + GQA +
16// KV cache + RoPE are subsequent bricks (the structural scaffolding
17// for real Llama-7B+ inference; this brick proves the math chain).
18//
19// Per-block formula (Llama-style, RMSNorm + SwiGLU FFN):
20//
21// attn_in = RMSNorm(x, gamma_attn)
22// Q = matmul(attn_in, W_q) [1, head_dim]
23// K = matmul(attn_in, W_k)
24// V = matmul(attn_in, W_v)
25// scores = matmul(Q, K^T) * attn_scale [1, 1]
26// probs = softmax(scores)
27// attn_out = matmul(probs, V) [1, head_dim]
28// attn_proj = matmul(attn_out, W_o) [1, hidden_dim]
29// x = x + attn_proj
30//
31// ffn_in = RMSNorm(x, gamma_ffn)
32// gate = SiLU(matmul(ffn_in, W_gate)) [1, ffn_dim]
33// up = matmul(ffn_in, W_up) [1, ffn_dim]
34// hidden = gate * up [1, ffn_dim]
35// ffn_proj = matmul(hidden, W_down) [1, hidden_dim]
36// x = x + ffn_proj
37//
38// Weights bundle struct keeps the function signature under the
39// NishiLang 16-arg limit while making call-sites self-documenting.
40
41import "nx_syscalls.nx"
42import "nx_tier.nx"
43import "nx_f32.nx"
44import "nx_f32_rmsnorm.nx"
45import "nx_f32_matmul.nx"
46import "nx_f32_softmax.nx"
47import "nx_f32_activations.nx"
48
49const NX_F32_TB_OK: nx_int = 0
50const NX_F32_TB_ERR_BAD_DIM: nx_int = 1
51const NX_F32_TB_N_VERDICTS: nx_int = 2
52
53func nx_f32_tb_verdict_is_valid(v: nx_int) -> nx_int {
54 if v < 0 { return 0 }
55 if v >= NX_F32_TB_N_VERDICTS { return 0 }
56 return 1
57}
58
59// Weights bundle (all f32 raw bits).
60struct NxF32TransformerBlockWeights {
61 gamma_attn: *i64, // [hidden_dim]
62 W_q: *i64, // [hidden_dim, head_dim]
63 W_k: *i64, // [hidden_dim, head_dim]
64 W_v: *i64, // [hidden_dim, head_dim]
65 W_o: *i64, // [head_dim, hidden_dim]
66 gamma_ffn: *i64, // [hidden_dim]
67 W_gate: *i64, // [hidden_dim, ffn_dim]
68 W_up: *i64, // [hidden_dim, ffn_dim]
69 W_down: *i64 // [ffn_dim, hidden_dim]
70}
71
72const NX_F32_TBW_BYTES: nx_int = 72 // 9 * 8
73
74// Forward pass. In-place: x is BOTH input and output (1 x hidden_dim).
75//
76// Args (within the 16-arg limit):
77// x pointer to 1*hidden_dim f32 values (mutated in place)
78// hidden_dim d_model
79// head_dim per-head dim (single-head v1: == hidden_dim)
80// ffn_dim intermediate FFN dim
81// w weights bundle
82// attn_scale 1/sqrt(head_dim) f32 raw bits
83// eps_f32 RMSNorm epsilon (typically 1e-6)
84//
85// 7 args. Under the limit.
86
87func nx_f32_transformer_block_forward(
88 x: *i64,
89 hidden_dim: nx_int,
90 head_dim: nx_int,
91 ffn_dim: nx_int,
92 w: *NxF32TransformerBlockWeights,
93 attn_scale: i64,
94 eps_f32: i64) -> nx_int {
95
96 if hidden_dim <= 0 { return NX_F32_TB_ERR_BAD_DIM }
97 if head_dim <= 0 { return NX_F32_TB_ERR_BAD_DIM }
98 if ffn_dim <= 0 { return NX_F32_TB_ERR_BAD_DIM }
99
100 // ----- Scratch buffers (single token; small) -----
101 let attn_in: *i64 = sys_mmap(hidden_dim * 8) as *i64
102 let Q: *i64 = sys_mmap(head_dim * 8) as *i64
103 let K: *i64 = sys_mmap(head_dim * 8) as *i64
104 let V: *i64 = sys_mmap(head_dim * 8) as *i64
105 let scores: *i64 = sys_mmap(8) as *i64 // [1,1]
106 let probs: *i64 = sys_mmap(8) as *i64
107 let attn_out: *i64 = sys_mmap(head_dim * 8) as *i64
108 let attn_proj: *i64 = sys_mmap(hidden_dim * 8) as *i64
109 let ffn_in: *i64 = sys_mmap(hidden_dim * 8) as *i64
110 let gate: *i64 = sys_mmap(ffn_dim * 8) as *i64
111 let up: *i64 = sys_mmap(ffn_dim * 8) as *i64
112 let hidden: *i64 = sys_mmap(ffn_dim * 8) as *i64
113 let ffn_proj: *i64 = sys_mmap(hidden_dim * 8) as *i64
114
115 // ===== Attention sub-block =====================================
116
117 // 1) attn_in = RMSNorm(x, gamma_attn)
118 nx_f32_rmsnorm(x, w.gamma_attn, hidden_dim, eps_f32, attn_in)
119
120 // 2) Q,K,V projections. [1, hidden_dim] @ [hidden_dim, head_dim].
121 nx_f32_matmul(attn_in, w.W_q, Q, 1, hidden_dim, head_dim)
122 nx_f32_matmul(attn_in, w.W_k, K, 1, hidden_dim, head_dim)
123 nx_f32_matmul(attn_in, w.W_v, V, 1, hidden_dim, head_dim)
124
125 // 3) scores = Q . K (since K^T for n_tokens=1 is the same vector).
126 // scalar = dot(Q, K) * attn_scale.
127 let raw_score: i64 = nx_f32_dot(Q, K, head_dim)
128 scores[0] = nx_f32_mul(raw_score, attn_scale)
129
130 // 4) probs = softmax([raw_score])
131 // For a 1-element vector, softmax([s]) = [1.0] regardless of s.
132 // (e^s / e^s = 1.) But we run the full path for honest
133 // composition coverage.
134 nx_f32_softmax(scores, 1, probs)
135
136 // 5) attn_out = probs * V (probs is scalar; broadcast multiply).
137 var i: nx_int = 0
138 while i < head_dim {
139 attn_out[i] = nx_f32_mul(probs[0], V[i])
140 i = i + 1
141 }
142
143 // 6) attn_proj = matmul(attn_out, W_o)
144 // [1, head_dim] @ [head_dim, hidden_dim] -> [1, hidden_dim]
145 nx_f32_matmul(attn_out, w.W_o, attn_proj, 1, head_dim, hidden_dim)
146
147 // 7) x = x + attn_proj (residual)
148 var j: nx_int = 0
149 while j < hidden_dim {
150 x[j] = nx_f32_add(x[j], attn_proj[j])
151 j = j + 1
152 }
153
154 // ===== FFN sub-block ===========================================
155
156 // 8) ffn_in = RMSNorm(x, gamma_ffn)
157 nx_f32_rmsnorm(x, w.gamma_ffn, hidden_dim, eps_f32, ffn_in)
158
159 // 9) gate = SiLU(matmul(ffn_in, W_gate))
160 nx_f32_matmul(ffn_in, w.W_gate, gate, 1, hidden_dim, ffn_dim)
161 var g: nx_int = 0
162 while g < ffn_dim {
163 gate[g] = nx_f32_silu(gate[g])
164 g = g + 1
165 }
166
167 // 10) up = matmul(ffn_in, W_up)
168 nx_f32_matmul(ffn_in, w.W_up, up, 1, hidden_dim, ffn_dim)
169
170 // 11) hidden = gate * up (SwiGLU element-wise)
171 var h: nx_int = 0
172 while h < ffn_dim {
173 hidden[h] = nx_f32_mul(gate[h], up[h])
174 h = h + 1
175 }
176
177 // 12) ffn_proj = matmul(hidden, W_down)
178 // [1, ffn_dim] @ [ffn_dim, hidden_dim] -> [1, hidden_dim]
179 nx_f32_matmul(hidden, w.W_down, ffn_proj, 1, ffn_dim, hidden_dim)
180
181 // 13) x = x + ffn_proj (residual)
182 var k: nx_int = 0
183 while k < hidden_dim {
184 x[k] = nx_f32_add(x[k], ffn_proj[k])
185 k = k + 1
186 }
187
188 return NX_F32_TB_OK
189}