nx_f32_llama_block.nx source
↩ module page · 238 lines · 9692 B
1// nx_f32_llama_block.nx -- bits-up Llama-class transformer block
2// forward pass (v3). The full per-layer integration.
3//
4// Composes:
5// nx_f32_rmsnorm
6// nx_f32_matmul (for W_q, W_k, W_v, W_o, W_gate, W_up, W_down)
7// nx_f32_rope_apply_vector (optional, per query/kv head per token)
8// nx_f32_attn_with_cache (the new KV-cached attention kernel)
9// nx_f32_silu (FFN SwiGLU gate activation)
10// nx_f32_mul (element-wise gate * up)
11// nx_f32_add (residual additions)
12//
13// Layer structure (canonical Llama / Llama-2 / Qwen2 form):
14// x_in = x // input
15// attn_in = RMSNorm(x_in, gamma_attn, eps)
16// Q = matmul(attn_in, W_q) // [n_tokens, q_dim]
17// K_new = matmul(attn_in, W_k) // [n_tokens, kv_dim]
18// V_new = matmul(attn_in, W_v) // [n_tokens, kv_dim]
19// if apply_rope:
20// for each token t, pos = cache.seq_len + t:
21// for each query head h:
22// RoPE(Q[t, h*head_dim:(h+1)*head_dim], pos)
23// for each kv head h:
24// RoPE(K_new[t, h*head_dim:(h+1)*head_dim], pos)
25// attn_concat = attn_with_cache(Q, K_new, V_new, ..., cache, layer_idx)
26// attn_proj = matmul(attn_concat, W_o) // [n_tokens, hidden_dim]
27// x_mid = x_in + attn_proj // residual
28//
29// ffn_in = RMSNorm(x_mid, gamma_ffn, eps)
30// gate_raw = matmul(ffn_in, W_gate) // [n_tokens, ffn_dim]
31// up = matmul(ffn_in, W_up) // [n_tokens, ffn_dim]
32// for each element: hidden = SiLU(gate_raw) * up
33// ffn_proj = matmul(hidden, W_down) // [n_tokens, hidden_dim]
34// out = x_mid + ffn_proj // residual
35//
36// The CALLER advances cache.seq_len ONCE per forward pass after
37// invoking this block for every layer in the stack.
38//
39// Weight layout convention (matches matmul convention):
40// matmul(A [M,K], B [K,N]) -> [M,N]
41// W_q: [hidden_dim, q_dim], attn_in [n_tokens, hidden_dim] @ W_q
42// W_k: [hidden_dim, kv_dim]
43// W_v: [hidden_dim, kv_dim]
44// W_o: [q_dim, hidden_dim]
45// W_gate: [hidden_dim, ffn_dim]
46// W_up: [hidden_dim, ffn_dim]
47// W_down: [ffn_dim, hidden_dim]
48//
49// Note: real Llama GGUF stores transposed weights but the loader
50// transposes them on load (or the matmul does column-walk). v1
51// uses the canonical [M,K]@[K,N]=[M,N] form so the math is clean.
52// A separate transposing layer can adapt GGUF on demand.
53//
54// genealogy_id: vaswani_2017_attention + touvron_2023_llama +
55// zhang_sennrich_2019_rmsnorm + su_2021_rope +
56// ainslie_2023_gqa + shazeer_2020_swiglu
57// lineage_id: substrate_f32_llama_block_v3
58
59import "nx_syscalls.nx"
60import "nx_tier.nx"
61import "nx_f32.nx"
62import "nx_f32_rmsnorm.nx"
63import "nx_f32_matmul.nx"
64import "nx_f32_activations.nx"
65import "nx_f32_rope.nx"
66import "nx_f32_attn_multi.nx"
67import "nx_f32_kv_cache.nx"
68import "nx_f32_attn_cached.nx"
69
70const NX_F32_BLK_OK: nx_int = 0
71const NX_F32_BLK_ERR_BAD_DIM: nx_int = 1
72const NX_F32_BLK_ERR_NULL: nx_int = 2
73const NX_F32_BLK_ERR_CACHE: nx_int = 3
74const NX_F32_BLK_N_VERDICTS: nx_int = 4
75
76func nx_f32_blk_verdict_is_valid(v: nx_int) -> nx_int {
77 if v < 0 { return 0 }
78 if v >= NX_F32_BLK_N_VERDICTS { return 0 }
79 return 1
80}
81
82// Per-layer weights bundled into a single struct.
83
84struct NxF32LlamaLayer {
85 gamma_attn: *i64, // [hidden_dim]
86 gamma_ffn: *i64, // [hidden_dim]
87 W_q: *i64, // [hidden_dim, q_dim]
88 W_k: *i64, // [hidden_dim, kv_dim]
89 W_v: *i64, // [hidden_dim, kv_dim]
90 W_o: *i64, // [q_dim, hidden_dim]
91 W_gate: *i64, // [hidden_dim, ffn_dim]
92 W_up: *i64, // [hidden_dim, ffn_dim]
93 W_down: *i64 // [ffn_dim, hidden_dim]
94}
95
96const NX_F32_LLAMA_LAYER_BYTES: nx_int = 72 // 9 fields * 8
97
98// Allocate a zero-initialized layer struct. Caller fills in weights.
99
100func nx_f32_llama_layer_alloc() -> *NxF32LlamaLayer {
101 return sys_mmap(NX_F32_LLAMA_LAYER_BYTES) as *NxF32LlamaLayer
102}
103
104// Block forward. out and x may NOT alias.
105// 15 args -- under the 16-arg limit.
106
107func nx_f32_llama_block_forward_v3(
108 x: *i64,
109 n_tokens: nx_int,
110 hidden_dim: nx_int,
111 n_heads: nx_int,
112 n_kv_heads: nx_int,
113 head_dim: nx_int,
114 ffn_dim: nx_int,
115 layer: *NxF32LlamaLayer,
116 cache: *NxF32KVCache,
117 layer_idx: nx_int,
118 eps: i64,
119 attn_scale: i64,
120 rope_log_base: i64,
121 apply_rope: nx_int,
122 out: *i64) -> nx_int {
123
124 if n_tokens <= 0 { return NX_F32_BLK_ERR_BAD_DIM }
125 if hidden_dim <= 0 { return NX_F32_BLK_ERR_BAD_DIM }
126 if n_heads <= 0 { return NX_F32_BLK_ERR_BAD_DIM }
127 if n_kv_heads <= 0 { return NX_F32_BLK_ERR_BAD_DIM }
128 if head_dim <= 0 { return NX_F32_BLK_ERR_BAD_DIM }
129 if ffn_dim <= 0 { return NX_F32_BLK_ERR_BAD_DIM }
130 if n_heads * head_dim != hidden_dim { return NX_F32_BLK_ERR_BAD_DIM }
131 if layer == (0 as *NxF32LlamaLayer) { return NX_F32_BLK_ERR_NULL }
132 if cache == (0 as *NxF32KVCache) { return NX_F32_BLK_ERR_NULL }
133 if x == (0 as *i64) { return NX_F32_BLK_ERR_NULL }
134 if out == (0 as *i64) { return NX_F32_BLK_ERR_NULL }
135
136 let q_dim: nx_int = n_heads * head_dim
137 let kv_dim: nx_int = n_kv_heads * head_dim
138
139 // Scratch buffers.
140 let attn_in: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
141 let Q: *i64 = sys_mmap(n_tokens * q_dim * 8) as *i64
142 let K_new: *i64 = sys_mmap(n_tokens * kv_dim * 8) as *i64
143 let V_new: *i64 = sys_mmap(n_tokens * kv_dim * 8) as *i64
144 let attn_concat: *i64 = sys_mmap(n_tokens * q_dim * 8) as *i64
145 let attn_proj: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
146 let x_mid: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
147 let ffn_in: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
148 let gate_raw: *i64 = sys_mmap(n_tokens * ffn_dim * 8) as *i64
149 let gate_act: *i64 = sys_mmap(n_tokens * ffn_dim * 8) as *i64
150 let up_buf: *i64 = sys_mmap(n_tokens * ffn_dim * 8) as *i64
151 let hidden_buf: *i64 = sys_mmap(n_tokens * ffn_dim * 8) as *i64
152 let ffn_proj: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
153
154 // ===== Attention sublayer =====
155 // RMSNorm per token.
156 var t: nx_int = 0
157 while t < n_tokens {
158 let x_row: *i64 = (((x as i64) + t * hidden_dim * 8)) as *i64
159 let n_row: *i64 = (((attn_in as i64) + t * hidden_dim * 8)) as *i64
160 nx_f32_rmsnorm(x_row, layer.gamma_attn, hidden_dim, eps, n_row)
161 t = t + 1
162 }
163
164 // Q, K_new, V_new = matmul(attn_in, W_*).
165 nx_f32_matmul(attn_in, layer.W_q, Q, n_tokens, hidden_dim, q_dim)
166 nx_f32_matmul(attn_in, layer.W_k, K_new, n_tokens, hidden_dim, kv_dim)
167 nx_f32_matmul(attn_in, layer.W_v, V_new, n_tokens, hidden_dim, kv_dim)
168
169 // Optional RoPE per token position per head.
170 if apply_rope != 0 {
171 let cache_before: nx_int = nx_f32_kv_cache_get_seq_len(cache)
172 var t_r: nx_int = 0
173 while t_r < n_tokens {
174 let pos: nx_int = cache_before + t_r
175 var h_q: nx_int = 0
176 while h_q < n_heads {
177 let qv: *i64 = (((Q as i64) + (t_r * q_dim + h_q * head_dim) * 8)) as *i64
178 nx_f32_rope_apply_vector(qv, head_dim, pos, rope_log_base)
179 h_q = h_q + 1
180 }
181 var h_kv: nx_int = 0
182 while h_kv < n_kv_heads {
183 let kv: *i64 = (((K_new as i64) + (t_r * kv_dim + h_kv * head_dim) * 8)) as *i64
184 nx_f32_rope_apply_vector(kv, head_dim, pos, rope_log_base)
185 h_kv = h_kv + 1
186 }
187 t_r = t_r + 1
188 }
189 }
190
191 // Cached multi-head attention.
192 let v_attn: nx_int = nx_f32_attn_with_cache(Q, K_new, V_new, n_tokens,
193 n_heads, n_kv_heads, head_dim,
194 cache, layer_idx, 1, attn_scale,
195 attn_concat)
196 if v_attn != NX_F32_AC_OK { return NX_F32_BLK_ERR_CACHE }
197
198 // attn_proj = matmul(attn_concat, W_o).
199 nx_f32_matmul(attn_concat, layer.W_o, attn_proj, n_tokens, q_dim, hidden_dim)
200
201 // Residual: x_mid = x + attn_proj.
202 var i: nx_int = 0
203 while i < n_tokens * hidden_dim {
204 x_mid[i] = nx_f32_add(x[i], attn_proj[i])
205 i = i + 1
206 }
207
208 // ===== FFN (SwiGLU) sublayer =====
209 var t2: nx_int = 0
210 while t2 < n_tokens {
211 let xm_row: *i64 = (((x_mid as i64) + t2 * hidden_dim * 8)) as *i64
212 let fi_row: *i64 = (((ffn_in as i64) + t2 * hidden_dim * 8)) as *i64
213 nx_f32_rmsnorm(xm_row, layer.gamma_ffn, hidden_dim, eps, fi_row)
214 t2 = t2 + 1
215 }
216
217 nx_f32_matmul(ffn_in, layer.W_gate, gate_raw, n_tokens, hidden_dim, ffn_dim)
218 nx_f32_matmul(ffn_in, layer.W_up, up_buf, n_tokens, hidden_dim, ffn_dim)
219
220 // gate_act = SiLU(gate_raw); hidden = gate_act * up.
221 var j: nx_int = 0
222 while j < n_tokens * ffn_dim {
223 gate_act[j] = nx_f32_silu(gate_raw[j])
224 hidden_buf[j] = nx_f32_mul(gate_act[j], up_buf[j])
225 j = j + 1
226 }
227
228 nx_f32_matmul(hidden_buf, layer.W_down, ffn_proj, n_tokens, ffn_dim, hidden_dim)
229
230 // Residual: out = x_mid + ffn_proj.
231 var k: nx_int = 0
232 while k < n_tokens * hidden_dim {
233 out[k] = nx_f32_add(x_mid[k], ffn_proj[k])
234 k = k + 1
235 }
236
237 return NX_F32_BLK_OK
238}