nx_f32_llama_block_v4.nx source
↩ module page · 389 lines · 14493 B
1// nx_f32_llama_block_v4.nx -- per-layer forward with lazy weights.
2//
3// Mirrors nx_f32_llama_block (v3) but accepts NxF32LazyWeight for
4// each of the 7 matmul weights, dispatching to f32 or Q4_K backends
5// per tensor. RMSNorm gammas remain as raw *i64 (they're 1D vectors,
6// always small enough that lazy dequant isn't needed).
7//
8// When real GGUF files have most weights as Q4_K (e.g. Qwen2.5-0.5B-
9// Instruct Q4_K_M), this forward avoids the eager-dequant memory
10// blow-up: each matmul streams the Q4_K bytes through nx_q4k_to_f32
11// one row at a time, peak working memory drops from O(k*n) to O(n).
12//
13// genealogy_id: vaswani_2017 + touvron_2023_llama +
14// ggml_format_taxonomy + tagged_dispatch
15// lineage_id: substrate_f32_llama_block_v4_lazy
16
17import "nx_syscalls.nx"
18import "nx_tier.nx"
19import "nx_f32.nx"
20import "nx_f32_rmsnorm.nx"
21import "nx_f32_matmul.nx"
22import "nx_f32_activations.nx"
23import "nx_f32_rope.nx"
24import "nx_f32_attn_multi.nx"
25import "nx_f32_kv_cache.nx"
26import "nx_f32_attn_cached.nx"
27import "nx_f32_lazy_weight.nx"
28import "nx_thread_pool.nx"
29import "nx_f32_llama_block.nx"
30const NX_MAGIC_2026: i64 = 2026
31
32const NX_BLK4_OK: nx_int = 0
33const NX_BLK4_ERR_BAD_DIM: nx_int = 1
34const NX_BLK4_ERR_NULL: nx_int = 2
35const NX_BLK4_ERR_CACHE: nx_int = 3
36const NX_BLK4_N_VERDICTS: nx_int = 4
37
38func nx_blk4_verdict_is_valid(v: nx_int) -> nx_int {
39 if v < 0 { return 0 }
40 if v >= NX_BLK4_N_VERDICTS { return 0 }
41 return 1
42}
43
44struct NxF32LlamaLayerLazy {
45 gamma_attn: *i64, // [hidden_dim]
46 gamma_ffn: *i64, // [hidden_dim]
47 W_q: *NxF32LazyWeight,
48 W_k: *NxF32LazyWeight,
49 W_v: *NxF32LazyWeight,
50 W_o: *NxF32LazyWeight,
51 W_gate: *NxF32LazyWeight,
52 W_up: *NxF32LazyWeight,
53 W_down: *NxF32LazyWeight,
54 bias_q: *i64, // [q_dim] -- Qwen2 attention_bias=true (0 ptr for Llama-style no-bias)
55 bias_k: *i64, // [kv_dim]
56 bias_v: *i64 // [kv_dim]
57}
58
59const NX_F32_LLAMA_LAYER_LAZY_BYTES: nx_int = 96 // 12 fields * 8
60
61func nx_f32_llama_layer_lazy_alloc() -> *NxF32LlamaLayerLazy {
62 return sys_mmap(NX_F32_LLAMA_LAYER_LAZY_BYTES) as *NxF32LlamaLayerLazy
63}
64
65// ===== Phase profiling (2026-07-08, off by default) ===============
66// Accumulates per-region us into 8 buckets so the decode forward can
67// be split without guessing. bp_t/bp_add are ~free when disabled
68// (one branch). Buckets: 0=ALLOC 1=RMSNORM 2=MATMUL 3=ROPE/BIAS
69// 4=ATTN 5=ACT 6=RESID.
70const BP_ALLOC: i64 = 0
71const BP_RMSNORM: i64 = 1
72const BP_MATMUL: i64 = 2
73const BP_ROPE: i64 = 3
74const BP_ATTN: i64 = 4
75const BP_ACT: i64 = 5
76const BP_RESID: i64 = 6
77
78static g_bp_en: i64
79static g_bp_acc: i64 // ptr to i64[8]
80
81func nx_blk4_prof_enable(on: i64) -> i64 {
82 g_bp_en = on
83 if g_bp_acc == 0 { g_bp_acc = sys_mmap(8 * 8) as i64 }
84 return 0
85}
86func nx_blk4_prof_reset() -> i64 {
87 if g_bp_acc != 0 {
88 let a: *i64 = g_bp_acc as *i64
89 var i: i64 = 0
90 while i < 8 { a[i] = 0; i = i + 1 }
91 }
92 return 0
93}
94func nx_blk4_prof_read(b: i64) -> i64 {
95 if g_bp_acc == 0 { return 0 }
96 let a: *i64 = g_bp_acc as *i64
97 return a[b]
98}
99func bp_t() -> i64 {
100 if g_bp_en == 0 { return 0 }
101 return sys_now_us()
102}
103func bp_add(b: i64, t0: i64) -> i64 {
104 if g_bp_en == 0 { return 0 }
105 let a: *i64 = g_bp_acc as *i64
106 a[b] = a[b] + (sys_now_us() - t0)
107 return 0
108}
109
110// ===== Pooled elementwise phases (2026-07-08) =====================
111// SwiGLU (silu*up) and RoPE were 20% + 17% of a decode token, both
112// SCALAR per-element transcendentals (software exp / sin+cos) running
113// serial. They are embarrassingly parallel -- band them on the
114// shared pool. Bit-exact (each output element/head computed wholly
115// in one band, no cross-element dependency).
116
117struct NxActCtx {
118 a_ptr: i64,
119 b_ptr: i64,
120 o_ptr: i64,
121 lo: i64,
122 hi: i64,
123 p0: i64,
124 p1: i64,
125 p2: i64,
126}
127const NX_ACT_CTX_BYTES: i64 = 64
128
129// Fused SwiGLU: out[j] = silu(gate[j]) * up[j] over [lo,hi).
130func _swiglu_task(ctx_i: i64) -> i64 {
131 let cx: *NxActCtx = ctx_i as *NxActCtx
132 let g: *i64 = cx.a_ptr as *i64
133 let u: *i64 = cx.b_ptr as *i64
134 let o: *i64 = cx.o_ptr as *i64
135 var j: i64 = cx.lo
136 while j < cx.hi {
137 o[j] = nx_f32_mul(nx_f32_silu(g[j]), u[j])
138 j = j + 1
139 }
140 return 0
141}
142
143func nx_blk4_swiglu_pool(gate: *i64, up: *i64, out: *i64, count: i64) -> i64 {
144 let pool: *NxThreadPool = nx_lw_shared_pool()
145 var bands: i64 = pool.n_workers
146 if bands > count { bands = count }
147 if bands < 1 { bands = 1 }
148 let ctxs: *u8 = sys_mmap(bands * NX_ACT_CTX_BYTES)
149 let per: i64 = (count + bands - 1) / bands
150 let done_before: i64 = nx_pool_n_completed(pool)
151 var b: i64 = 0
152 while b < bands {
153 let cx: *NxActCtx = ((ctxs as i64) + b * NX_ACT_CTX_BYTES) as *NxActCtx
154 cx.a_ptr = gate as i64
155 cx.b_ptr = up as i64
156 cx.o_ptr = out as i64
157 cx.lo = b * per
158 var hi: i64 = (b + 1) * per
159 if hi > count { hi = count }
160 cx.hi = hi
161 nx_pool_submit(pool, _swiglu_task, cx as i64)
162 b = b + 1
163 }
164 nx_pool_wait(pool, done_before + bands)
165 sys_munmap(ctxs, bands * NX_ACT_CTX_BYTES)
166 return 0
167}
168
169// RoPE over a band of heads. Each head is head_dim wide at
170// base + head*head_dim*8; pos + rope_log_base are shared. Bands
171// heads [lo,hi) so no head is split.
172func _rope_task(ctx_i: i64) -> i64 {
173 let cx: *NxActCtx = ctx_i as *NxActCtx
174 let base: i64 = cx.a_ptr
175 let head_dim: i64 = cx.p0
176 let pos: i64 = cx.p1
177 let rope_log_base: i64 = cx.p2
178 var h: i64 = cx.lo
179 while h < cx.hi {
180 let hv: *i64 = (base + h * head_dim * 8) as *i64
181 nx_f32_rope_apply_vector_neox(hv, head_dim, pos, rope_log_base)
182 h = h + 1
183 }
184 return 0
185}
186
187// Apply RoPE to n_heads consecutive heads at `base` (row for one
188// token), pool-banded. Serial-fallback for tiny head counts is the
189// band clamp (bands <= n_heads).
190func nx_blk4_rope_heads_pool(base: i64, n_heads: i64, head_dim: i64,
191 pos: i64, rope_log_base: i64) -> i64 {
192 let pool: *NxThreadPool = nx_lw_shared_pool()
193 var bands: i64 = pool.n_workers
194 if bands > n_heads { bands = n_heads }
195 if bands < 1 { bands = 1 }
196 let ctxs: *u8 = sys_mmap(bands * NX_ACT_CTX_BYTES)
197 let per: i64 = (n_heads + bands - 1) / bands
198 let done_before: i64 = nx_pool_n_completed(pool)
199 var b: i64 = 0
200 while b < bands {
201 let cx: *NxActCtx = ((ctxs as i64) + b * NX_ACT_CTX_BYTES) as *NxActCtx
202 cx.a_ptr = base
203 cx.p0 = head_dim
204 cx.p1 = pos
205 cx.p2 = rope_log_base
206 cx.lo = b * per
207 var hi: i64 = (b + 1) * per
208 if hi > n_heads { hi = n_heads }
209 cx.hi = hi
210 nx_pool_submit(pool, _rope_task, cx as i64)
211 b = b + 1
212 }
213 nx_pool_wait(pool, done_before + bands)
214 sys_munmap(ctxs, bands * NX_ACT_CTX_BYTES)
215 return 0
216}
217
218// v4 forward. 15 args (under the 16-arg limit).
219
220func nx_f32_llama_block_forward_v4(
221 x: *i64,
222 n_tokens: nx_int,
223 hidden_dim: nx_int,
224 n_heads: nx_int,
225 n_kv_heads: nx_int,
226 head_dim: nx_int,
227 ffn_dim: nx_int,
228 layer: *NxF32LlamaLayerLazy,
229 cache: *NxF32KVCache,
230 layer_idx: nx_int,
231 eps: i64,
232 attn_scale: i64,
233 rope_log_base: i64,
234 apply_rope: nx_int,
235 out: *i64) -> nx_int {
236
237 if n_tokens <= 0 { return NX_BLK4_ERR_BAD_DIM }
238 if hidden_dim <= 0 { return NX_BLK4_ERR_BAD_DIM }
239 if n_heads <= 0 { return NX_BLK4_ERR_BAD_DIM }
240 if n_kv_heads <= 0 { return NX_BLK4_ERR_BAD_DIM }
241 if head_dim <= 0 { return NX_BLK4_ERR_BAD_DIM }
242 if ffn_dim <= 0 { return NX_BLK4_ERR_BAD_DIM }
243 if n_heads * head_dim != hidden_dim { return NX_BLK4_ERR_BAD_DIM }
244 if layer == (0 as *NxF32LlamaLayerLazy) { return NX_BLK4_ERR_NULL }
245 if cache == (0 as *NxF32KVCache) { return NX_BLK4_ERR_NULL }
246 if x == (0 as *i64) { return NX_BLK4_ERR_NULL }
247 if out == (0 as *i64) { return NX_BLK4_ERR_NULL }
248
249 let q_dim: nx_int = n_heads * head_dim
250 let kv_dim: nx_int = n_kv_heads * head_dim
251
252 let _tp0: i64 = bp_t()
253 let attn_in: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
254 let Q: *i64 = sys_mmap(n_tokens * q_dim * 8) as *i64
255 let K_new: *i64 = sys_mmap(n_tokens * kv_dim * 8) as *i64
256 let V_new: *i64 = sys_mmap(n_tokens * kv_dim * 8) as *i64
257 let attn_concat: *i64 = sys_mmap(n_tokens * q_dim * 8) as *i64
258 let attn_proj: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
259 let x_mid: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
260 let ffn_in: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
261 let gate_raw: *i64 = sys_mmap(n_tokens * ffn_dim * 8) as *i64
262 let gate_act: *i64 = sys_mmap(n_tokens * ffn_dim * 8) as *i64
263 let up_buf: *i64 = sys_mmap(n_tokens * ffn_dim * 8) as *i64
264 let hidden_buf: *i64 = sys_mmap(n_tokens * ffn_dim * 8) as *i64
265 let ffn_proj: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
266 bp_add(BP_ALLOC, _tp0)
267
268 // ===== Attention sublayer =====
269 let _tp1: i64 = bp_t()
270 var t: nx_int = 0
271 while t < n_tokens {
272 let x_row: *i64 = (((x as i64) + t * hidden_dim * 8)) as *i64
273 let n_row: *i64 = (((attn_in as i64) + t * hidden_dim * 8)) as *i64
274 nx_f32_rmsnorm(x_row, layer.gamma_attn, hidden_dim, eps, n_row)
275 t = t + 1
276 }
277 bp_add(BP_RMSNORM, _tp1)
278
279 let _tp2: i64 = bp_t()
280 nx_f32_lazy_matmul(attn_in, layer.W_q, Q, n_tokens, hidden_dim, q_dim)
281 nx_f32_lazy_matmul(attn_in, layer.W_k, K_new, n_tokens, hidden_dim, kv_dim)
282 nx_f32_lazy_matmul(attn_in, layer.W_v, V_new, n_tokens, hidden_dim, kv_dim)
283 bp_add(BP_MATMUL, _tp2)
284
285 // Qwen2 attention bias (attention_bias=true): Q/K/V += bias, per token, BEFORE RoPE.
286 // Guarded so Llama-style no-bias models (bias ptr == 0) are unaffected. This missing
287 // additive term was the forward-correctness bug -- every matmul/rope/attn was bit-exact.
288 let _tp3: i64 = bp_t()
289 if (layer.bias_q as i64) != 0 {
290 var tb: nx_int = 0
291 while tb < n_tokens {
292 var iq: nx_int = 0
293 while iq < q_dim { Q[tb*q_dim+iq] = nx_f32_add(Q[tb*q_dim+iq], layer.bias_q[iq]); iq = iq + 1 }
294 var ik: nx_int = 0
295 while ik < kv_dim { K_new[tb*kv_dim+ik] = nx_f32_add(K_new[tb*kv_dim+ik], layer.bias_k[ik]); ik = ik + 1 }
296 var iv: nx_int = 0
297 while iv < kv_dim { V_new[tb*kv_dim+iv] = nx_f32_add(V_new[tb*kv_dim+iv], layer.bias_v[iv]); iv = iv + 1 }
298 tb = tb + 1
299 }
300 }
301
302 // RoPE: build the (cos,sin) table ONCE per position, then apply to
303 // every head with pure mul/add (2026-07-08). The trig is head-
304 // INDEPENDENT, so the old per-head apply recomputed exp+cos+sin ~16x
305 // per layer (RoPE was 123ms/token, mostly transcendentals). Bit-exact.
306 // Stays SERIAL (threading it per-head was a measured loss -- too fine-
307 // grained; nx_blk4_rope_heads_pool kept above, unwired).
308 if apply_rope != 0 {
309 let cache_before: nx_int = nx_f32_kv_cache_get_seq_len(cache)
310 let cs: *i64 = sys_mmap(head_dim * 8) as *i64
311 var t_r: nx_int = 0
312 while t_r < n_tokens {
313 let pos: nx_int = cache_before + t_r
314 nx_f32_rope_build_cs(cs, head_dim, pos, rope_log_base)
315 var h_q: nx_int = 0
316 while h_q < n_heads {
317 let qv: *i64 = (((Q as i64) + (t_r * q_dim + h_q * head_dim) * 8)) as *i64
318 nx_f32_rope_apply_cs_neox(qv, head_dim, cs)
319 h_q = h_q + 1
320 }
321 var h_kv: nx_int = 0
322 while h_kv < n_kv_heads {
323 let kv: *i64 = (((K_new as i64) + (t_r * kv_dim + h_kv * head_dim) * 8)) as *i64
324 nx_f32_rope_apply_cs_neox(kv, head_dim, cs)
325 h_kv = h_kv + 1
326 }
327 t_r = t_r + 1
328 }
329 sys_munmap(cs, head_dim * 8)
330 }
331 bp_add(BP_ROPE, _tp3)
332
333 let _tp4: i64 = bp_t()
334 let v_attn: nx_int = nx_f32_attn_with_cache(Q, K_new, V_new, n_tokens,
335 n_heads, n_kv_heads, head_dim,
336 cache, layer_idx, 1, attn_scale,
337 attn_concat)
338 if v_attn != NX_F32_AC_OK { return NX_BLK4_ERR_CACHE }
339 bp_add(BP_ATTN, _tp4)
340
341 let _tp5: i64 = bp_t()
342 nx_f32_lazy_matmul(attn_concat, layer.W_o, attn_proj, n_tokens, q_dim, hidden_dim)
343 bp_add(BP_MATMUL, _tp5)
344
345 let _tp6: i64 = bp_t()
346 var i: nx_int = 0
347 while i < n_tokens * hidden_dim {
348 x_mid[i] = __f32_add(x[i], attn_proj[i]) // hw SSE, bit-identical (NX_MAGIC_2026-07-10)
349 i = i + 1
350 }
351 bp_add(BP_RESID, _tp6)
352
353 // ===== FFN (SwiGLU) sublayer =====
354 let _tp7: i64 = bp_t()
355 var t2: nx_int = 0
356 while t2 < n_tokens {
357 let xm_row: *i64 = (((x_mid as i64) + t2 * hidden_dim * 8)) as *i64
358 let fi_row: *i64 = (((ffn_in as i64) + t2 * hidden_dim * 8)) as *i64
359 nx_f32_rmsnorm(xm_row, layer.gamma_ffn, hidden_dim, eps, fi_row)
360 t2 = t2 + 1
361 }
362 bp_add(BP_RMSNORM, _tp7)
363
364 let _tp8: i64 = bp_t()
365 nx_f32_lazy_matmul(ffn_in, layer.W_gate, gate_raw, n_tokens, hidden_dim, ffn_dim)
366 nx_f32_lazy_matmul(ffn_in, layer.W_up, up_buf, n_tokens, hidden_dim, ffn_dim)
367 bp_add(BP_MATMUL, _tp8)
368
369 let _tp9: i64 = bp_t()
370 nx_blk4_swiglu_pool(gate_raw, up_buf, hidden_buf, n_tokens * ffn_dim)
371 bp_add(BP_ACT, _tp9)
372
373 let _tp10: i64 = bp_t()
374 nx_f32_lazy_matmul(hidden_buf, layer.W_down, ffn_proj, n_tokens, ffn_dim, hidden_dim)
375 bp_add(BP_MATMUL, _tp10)
376
377 let _tp11: i64 = bp_t()
378 var k: nx_int = 0
379 while k < n_tokens * hidden_dim {
380 out[k] = __f32_add(x_mid[k], ffn_proj[k]) // hw SSE, bit-identical (NX_MAGIC_2026-07-10)
381 k = k + 1
382 }
383 bp_add(BP_RESID, _tp11)
384 // (dangling debug call removed 2026-07-09: the blk4_dump8/ffnmax
385 // "log2mag" definitions were stripped in the spam cleanup but this one
386 // call site was missed -- it broke EVERY build importing this block.)
387
388 return NX_BLK4_OK
389}