nx_genblock.nx source
↩ module page · 665 lines · 27273 B
1// nx_genblock.nx -- shared DiT block primitives for the sovereign gen engine.
2//
3// Extracted from nx_gen_blockrun once a second consumer appeared (nx_gen_ditchain, which loops
4// these over N layers). Copying them would have been the third instance of this lane's most
5// expensive recurring mistake -- two implementations of one op drift, and the drift is silent
6// because both produce finite plausible numbers.
7//
8// Everything here is layer-agnostic and model-agnostic: tensor names, dimensions and worker
9// counts are arguments. Hot loops use the HARDWARE __f32_* intrinsics, never the nx_f32_*
10// software IEEE-754 twins (measured 9-10x slower in this lane).
11// license_tier: ORIGINAL
12
13import "nx_syscalls.nx"
14import "nx_le.nx"
15import "nx_f32.nx"
16import "nx_f32_div.nx"
17import "nx_f32_cvt.nx"
18import "nx_f16.nx"
19import "nx_f32_exp.nx"
20import "nx_f32_activations.nx"
21import "nx_strconv.nx"
22import "nx_genfix.nx"
23import "nx_genver.nx"
24import "nx_genweights.nx"
25
26import "nx_genarch.nx"
27
28// ---- ARCHITECTURE: BOUND AT RUNTIME, NOT COMPILED IN ----------------------------------------
29// Operator: *"z image turbo is the mix of all the parts ... just make sure that they are hot
30// swappable with other models"*. A `const` head count is a model welded into the binary: swapping
31// the checkpoint then produces a picture computed with the WRONG geometry and no error anywhere.
32// ★★★★★★ AN ARCHITECTURE CONSTANT IN THE CODE IS A MODEL THAT CANNOT BE SWAPPED.
33//
34// These start at 0 and every consumer must call br_arch_bind() first. ZERO IS DELIBERATE: the
35// tempting default is the Z-Image geometry, but then a Flux checkpoint that forgot to bind runs to
36// completion against 30 heads of 128 and returns plausible noise.
37// ★★★★★ A DEFAULT THAT IS ALSO A VALID ANSWER FOR ONE MODEL CANNOT SIGNAL "NOBODY SET THIS".
38static BR_HEAD_DIM: i64
39static BR_N_HEADS: i64
40static BR_CHUNKS: i64
41
42// Derive the geometry from the weights themselves and bind it. 0 = ok, negative = refuse.
43// The GGUF here carries `kv 0`, so tensor SHAPES are the model's only self-description -- which is
44// why nx_genarch probes rather than reads, and why it refuses instead of defaulting.
45func br_arch_bind(gw: *i64) -> i64 {
46 let a: *i64 = sys_mmap(NX_ARCH_SLOTS * 8 + 64) as *i64
47 if nx_arch_probe(gw, a) != 0 { return 0 - 1 }
48 if a[NX_ARCH_HEAD_DIM] <= 0 { return 0 - 2 }
49 if a[NX_ARCH_N_HEADS] <= 0 { return 0 - 3 }
50 if a[NX_ARCH_N_CHUNKS] <= 0 { return 0 - 4 }
51 BR_HEAD_DIM = a[NX_ARCH_HEAD_DIM]
52 BR_N_HEADS = a[NX_ARCH_N_HEADS]
53 BR_CHUNKS = a[NX_ARCH_N_CHUNKS]
54 return 0
55}
56// 1 if bound. Callers check this rather than trusting that someone upstream did.
57func br_arch_ready() -> i64 {
58 if BR_HEAD_DIM <= 0 { return 0 }
59 if BR_N_HEADS <= 0 { return 0 }
60 if BR_CHUNKS <= 0 { return 0 }
61 return 1
62}
63
64func br_puts(s: *u8) -> i64 {
65 var n: i64 = 0
66 while s[n] != (0 as u8) { n = n + 1 }
67 return sys_write(1, s, n)
68}
69func br_strlen(s: *u8) -> i64 {
70 var n: i64 = 0
71 while s[n] != (0 as u8) { n = n + 1 }
72 return n
73}
74
75// ---- Q8_0 weight, pre-decoded scales -------------------------------------------------
76// Returns the raw block bytes; fills `scales` with one f32 per 32-value block.
77func br_load_q8(model: *u8, name: *u8, in_dim: i64, out_dim: i64, scales_out: *i64) -> *u8 {
78 let nl: i64 = br_strlen(name)
79 let nblk: i64 = in_dim / 32
80 let bytes: i64 = nblk * 34 * out_dim
81 let w: *u8 = nx_genfix_load_raw(model, name, nl, bytes)
82 if (w as i64) == 0 { return 0 as *u8 }
83 var o: i64 = 0
84 while o < out_dim {
85 var b: i64 = 0
86 while b < nblk {
87 scales_out[o * nblk + b] = nx_f16_to_f32(nx_le_read_u16(w, (o * nblk + b) * 34))
88 b = b + 1
89 }
90 o = o + 1
91 }
92 return w
93}
94
95// ---- Q8_0 matmul: out[t][o] = sum_i x[t][i] * W[o][i], forked over the output band ----
96func br_mm_band(w: *u8, scales: *i64, x: *u8, out: *u8,
97 rows: i64, in_dim: i64, out_dim: i64, o0: i64, o1: i64) -> i64 {
98 let nblk: i64 = in_dim / 32
99 var t: i64 = 0
100 while t < rows {
101 let ab: i64 = (x as i64) + t * in_dim * 4
102 var o: i64 = o0
103 while o < o1 {
104 let qb: i64 = (w as i64) + o * nblk * 34
105 let sb: i64 = o * nblk
106 var acc: i64 = 0
107 var b: i64 = 0
108 while b < nblk {
109 let r: i64 = __f32_i8dot32a((qb + b * 34 + 2) as *u8, (ab + b * 128) as *u8)
110 acc = __f32_add(acc, __f32_mul(scales[sb + b], r))
111 b = b + 1
112 }
113 nx_le_write_u32(out, (t * out_dim + o) * 4, acc)
114 o = o + 1
115 }
116 t = t + 1
117 }
118 return 0
119}
120
121func br_matmul(w: *u8, scales: *i64, x: *u8, out: *u8,
122 rows: i64, in_dim: i64, out_dim: i64, nw: i64) -> i64 {
123 if nw <= 1 {
124 br_mm_band(w, scales, x, out, rows, in_dim, out_dim, 0, out_dim)
125 return 0
126 }
127 let pids: *i64 = sys_mmap(nw * 8 + 64) as *i64
128 var k: i64 = 0
129 while k < nw {
130 let o0: i64 = k * out_dim / nw
131 let o1: i64 = (k + 1) * out_dim / nw
132 let pid: i64 = sys_fork()
133 if pid == 0 {
134 br_mm_band(w, scales, x, out, rows, in_dim, out_dim, o0, o1)
135 sys_exit(0)
136 }
137 pids[k] = pid
138 k = k + 1
139 }
140 let st: *i64 = sys_mmap(64) as *i64
141 k = 0
142 while k < nw { sys_wait4(pids[k], st, 0); k = k + 1 }
143 return 0
144}
145
146// ---- RMSNorm (optionally * (1+scale)) ------------------------------------------------
147func br_rmsnorm(x: *u8, w: *u8, mods: *i64, out: *u8, rows: i64, d: i64, eps: i64) -> i64 {
148 let dinv: i64 = nx_f32_div(nx_i32_to_f32(1), nx_i32_to_f32(d))
149 var t: i64 = 0
150 while t < rows {
151 let base: i64 = t * d
152 var ss: i64 = 0
153 var i: i64 = 0
154 while i < d {
155 let v: i64 = nx_le_read_u32(x, (base + i) * 4)
156 ss = __f32_add(ss, __f32_mul(v, v))
157 i = i + 1
158 }
159 let inv: i64 = nx_f32_div(nx_i32_to_f32(1),
160 nx_f32_sqrt(__f32_add(__f32_mul(ss, dinv), eps)))
161 i = 0
162 while i < d {
163 var v: i64 = __f32_mul(__f32_mul(nx_le_read_u32(x, (base + i) * 4), inv),
164 nx_le_read_u32(w, i * 4))
165 if (mods as i64) != 0 { v = __f32_mul(v, mods[i]) }
166 nx_le_write_u32(out, (base + i) * 4, v)
167 i = i + 1
168 }
169 t = t + 1
170 }
171 return 0
172}
173
174// ---- gated residual: out = inner * tanh(gate) + residual -------------------------------
175func br_gate_resid(inner: *u8, tg: *i64, resid: *u8, out: *u8, rows: i64, d: i64) -> i64 {
176 var t: i64 = 0
177 while t < rows {
178 var i: i64 = 0
179 while i < d {
180 let f: i64 = t * d + i
181 let v: i64 = __f32_add(__f32_mul(nx_le_read_u32(inner, f * 4), tg[i]),
182 nx_le_read_u32(resid, f * 4))
183 nx_le_write_u32(out, f * 4, v)
184 i = i + 1
185 }
186 t = t + 1
187 }
188 return 0
189}
190
191
192// ---- SDPA over a band of heads ---------------------------------------------------------
193// Heads are fully independent -- separate q/k/v slices, separate output columns -- so this is a
194// clean parallel axis needing no reduction. It was the last scalar single-threaded stage in the
195// block and dominated the wall time even though the matmuls around it were tuned.
196// The q.k dot uses __f32x8_fma + one hsum, the same kernel shape that beat __f32x8_dot 17x in
197// the projections; the v accumulation stays scalar because it is an AXPY (scalar times vector),
198// which the available intrinsics do not cover.
199func br_sdpa_band(qrp: *u8, krp: *u8, qkv: *u8, aop: *u8,
200 T: i64, D: i64, QKV: i64, head_dim: i64, n_heads: i64,
201 scale: i64, h0: i64, h1: i64) -> i64 {
202 let sc: *i64 = sys_mmap(T * 8) as *i64
203 let acc: *i64 = sys_mmap(head_dim * 8) as *i64
204 let fa: *u8 = sys_mmap(64)
205 let az: *i64 = fa as *i64
206 let nch: i64 = head_dim / 8
207 var hh: i64 = h0
208 while hh < h1 {
209 var tq: i64 = 0
210 while tq < T {
211 let qrow: i64 = (hh * T + tq) * head_dim
212 let qb: i64 = (qrp as i64) + qrow * 4
213 var mx: i64 = 0
214 var tk: i64 = 0
215 while tk < T {
216 let kb: i64 = (krp as i64) + ((hh * T + tk) * head_dim) * 4
217 az[0] = 0
218 az[1] = 0
219 az[2] = 0
220 az[3] = 0
221 var ch: i64 = 0
222 while ch < nch {
223 __f32x8_fma(fa, (qb + ch * 32) as *u8, (kb + ch * 32) as *u8)
224 ch = ch + 1
225 }
226 let sv: i64 = __f32_mul(__f32x8_hsum(fa), scale)
227 sc[tk] = sv
228 if tk == 0 { mx = sv } else { if nx_f32_lt(mx, sv) == 1 { mx = sv } }
229 tk = tk + 1
230 }
231 let nmx: i64 = mx ^ 0x80000000
232 var sum: i64 = 0
233 tk = 0
234 while tk < T {
235 let e: i64 = nx_f32_exp(__f32_add(sc[tk], nmx))
236 sc[tk] = e
237 sum = __f32_add(sum, e)
238 tk = tk + 1
239 }
240 let sinv: i64 = nx_f32_div(nx_i32_to_f32(1), sum)
241 var d2: i64 = 0
242 while d2 < head_dim { acc[d2] = 0; d2 = d2 + 1 }
243 tk = 0
244 while tk < T {
245 let wgt: i64 = __f32_mul(sc[tk], sinv)
246 let vrow: i64 = tk * QKV + (2 * n_heads + hh) * head_dim
247 d2 = 0
248 while d2 < head_dim {
249 acc[d2] = __f32_add(acc[d2], __f32_mul(wgt, nx_le_read_u32(qkv, (vrow + d2) * 4)))
250 d2 = d2 + 1
251 }
252 tk = tk + 1
253 }
254 d2 = 0
255 while d2 < head_dim {
256 nx_le_write_u32(aop, (tq * D + hh * head_dim + d2) * 4, acc[d2])
257 d2 = d2 + 1
258 }
259 tq = tq + 1
260 }
261 hh = hh + 1
262 }
263 return 0
264}
265
266
267// ---- elementwise SwiGLU over a band ----------------------------------------------------
268// silu() goes through a SOFTWARE exp, so this loop is not cheap despite being "just elementwise":
269// measured at ~1.2s of a 6.6s block while the matmuls beside it were fully tuned.
270// ★ AN ELEMENTWISE STAGE IS NOT AUTOMATICALLY A CHEAP STAGE.
271func br_swiglu_band(a: *u8, b: *u8, out: *u8, i0: i64, i1: i64) -> i64 {
272 var i: i64 = i0
273 while i < i1 {
274 nx_le_write_u32(out, i * 4,
275 __f32_mul(nx_f32_silu(nx_le_read_u32(a, i * 4)),
276 nx_le_read_u32(b, i * 4)))
277 i = i + 1
278 }
279 return 0
280}
281
282// Fork a [0,n) range across nw workers and wait. The range carries no cross-element dependency,
283// so bands need no reduction.
284func br_fork_range(kind: i64, a: *u8, b: *u8, out: *u8, n: i64, nw: i64) -> i64 {
285 if nw <= 1 {
286 br_swiglu_band(a, b, out, 0, n)
287 return 0
288 }
289 let pids: *i64 = sys_mmap(nw * 8 + 64) as *i64
290 var k: i64 = 0
291 while k < nw {
292 let i0: i64 = k * n / nw
293 let i1: i64 = (k + 1) * n / nw
294 let pid: i64 = sys_fork()
295 if pid == 0 { br_swiglu_band(a, b, out, i0, i1); sys_exit(0) }
296 pids[k] = pid
297 k = k + 1
298 }
299 let st: *i64 = sys_mmap(64) as *i64
300 k = 0
301 while k < nw { sys_wait4(pids[k], st, 0); k = k + 1 }
302 return 0
303}
304
305
306// ---- GGUF-backed Q8_0 weight: raw blocks IN PLACE, scales decoded once ------------------
307// No copy: the blocks are read straight out of the mapped model file. That is the whole point of
308// reading the GGUF rather than a dumped fixture -- the engine touches only the pages it uses.
309func br_gw_q8(gw: *i64, name: *u8, in_dim: i64, out_dim: i64, scales_out: *i64) -> *u8 {
310 let idx: i64 = nx_gw_find(gw, name, br_strlen(name))
311 if idx < 0 { return 0 as *u8 }
312 if nx_gw_type(gw, idx) != NX_GW_TYPE_Q8_0 { return 0 as *u8 }
313 // Refuse a shape that disagrees with the caller rather than reading the wrong stride.
314 if nx_gw_dim0(gw, idx) != in_dim { return 0 as *u8 }
315 if nx_gw_dim1(gw, idx) != out_dim { return 0 as *u8 }
316 let w: *u8 = nx_gw_data(gw, idx)
317 let nblk: i64 = in_dim / 32
318 var o: i64 = 0
319 while o < out_dim {
320 var b: i64 = 0
321 while b < nblk {
322 scales_out[o * nblk + b] = nx_f16_to_f32(nx_le_read_u16(w, (o * nblk + b) * 34))
323 b = b + 1
324 }
325 o = o + 1
326 }
327 return w
328}
329
330// GGUF-backed small tensor as packed f32 (norms are stored quantized but the engine wants f32).
331func br_gw_f32(gw: *i64, name: *u8, n: i64) -> *u8 {
332 let idx: i64 = nx_gw_find(gw, name, br_strlen(name))
333 if idx < 0 { return 0 as *u8 }
334 let out: *u8 = sys_mmap(n * 4 + 64)
335 if nx_gw_to_f32_packed(gw, idx, out, n) != 0 { return 0 as *u8 }
336 return out
337}
338
339
340// ---- build "model.diffusion_model.layers.<n>.<suffix>" ---------------------------------
341// Layer-indexed names are what make an N-block chain possible without N sets of fixtures.
342// prefix is "layers" / "context_refiner" / "noise_refiner" -- the three block stacks share one
343// tensor layout, so one name builder and one block function serve all of them.
344func br_name_p(out: *u8, prefix: *u8, layer: i64, suffix: *u8) -> *u8 {
345 let pre: *u8 = "model.diffusion_model." as *u8
346 var o0: i64 = 0
347 var i0: i64 = 0
348 while pre[i0] != (0 as u8) { out[o0] = pre[i0]; o0 = o0 + 1; i0 = i0 + 1 }
349 i0 = 0
350 while prefix[i0] != (0 as u8) { out[o0] = prefix[i0]; o0 = o0 + 1; i0 = i0 + 1 }
351 out[o0] = 0x2E; o0 = o0 + 1
352 let dec0: *u8 = sys_mmap(32)
353 let nd0: i64 = nx_strconv_format_i64(layer, dec0)
354 var k0: i64 = 0
355 while k0 < nd0 { out[o0] = dec0[k0]; o0 = o0 + 1; k0 = k0 + 1 }
356 out[o0] = 0x2E; o0 = o0 + 1
357 i0 = 0
358 while suffix[i0] != (0 as u8) { out[o0] = suffix[i0]; o0 = o0 + 1; i0 = i0 + 1 }
359 out[o0] = 0
360 return out
361}
362
363func br_name(out: *u8, layer: i64, suffix: *u8) -> *u8 {
364 let pre: *u8 = "model.diffusion_model.layers." as *u8
365 var o: i64 = 0
366 var i: i64 = 0
367 while pre[i] != (0 as u8) { out[o] = pre[i]; o = o + 1; i = i + 1 }
368 let dec: *u8 = sys_mmap(32)
369 let nd: i64 = nx_strconv_format_i64(layer, dec)
370 var k: i64 = 0
371 while k < nd { out[o] = dec[k]; o = o + 1; k = k + 1 }
372 out[o] = 0x2E; o = o + 1
373 i = 0
374 while suffix[i] != (0 as u8) { out[o] = suffix[i]; o = o + 1; i = i + 1 }
375 out[o] = 0
376 return out
377}
378
379// ---- adaLN modulation vector for a layer, computed from the GGUF ------------------------
380// adaln[o] = sum_i t_emb[i] * W[o][i] + bias[o]
381// Computing this per layer removes the last per-layer fixture dependency: a chain then needs only
382// the layer-0 input, t_emb and the rope table.
383func br_adaln(gw: *i64, prefix: *u8, layer: i64, t_emb: *u8, out: *u8, embed_dim: i64, width: i64) -> i64 {
384 let nm: *u8 = sys_mmap(256)
385 br_name_p(nm, prefix, layer, "adaLN_modulation.0.weight" as *u8)
386 let wi: i64 = nx_gw_find(gw, nm, br_strlen(nm))
387 if wi < 0 { return 0 - 1 }
388 if nx_gw_dim0(gw, wi) != embed_dim { return 0 - 2 }
389 if nx_gw_dim1(gw, wi) != width { return 0 - 3 }
390 let w: *u8 = nx_gw_data(gw, wi)
391 let nblk: i64 = embed_dim / 32
392
393 let bm: *u8 = sys_mmap(256)
394 br_name_p(bm, prefix, layer, "adaLN_modulation.0.bias" as *u8)
395 let bias: *u8 = br_gw_f32(gw, bm, width)
396 if (bias as i64) == 0 { return 0 - 4 }
397
398 let ab: i64 = t_emb as i64
399 var o: i64 = 0
400 while o < width {
401 let qb: i64 = (w as i64) + o * nblk * 34
402 var acc: i64 = 0
403 var b: i64 = 0
404 while b < nblk {
405 let d32: i64 = nx_f16_to_f32(nx_le_read_u16(w, (o * nblk + b) * 34))
406 let r: i64 = __f32_i8dot32a((qb + b * 34 + 2) as *u8, (ab + b * 128) as *u8)
407 acc = __f32_add(acc, __f32_mul(d32, r))
408 b = b + 1
409 }
410 nx_le_write_u32(out, o * 4, __f32_add(acc, nx_le_read_u32(bias, o * 4)))
411 o = o + 1
412 }
413 return 0
414}
415
416
417
418
419// ---- f32 matmul (for F16-stored weights) ------------------------------------------------
420// Not every tensor in a "Q8_0" model is Q8_0: the final layer and the embedders are stored F16.
421// br_gw_q8 REFUSES those by type rather than misreading 2-byte halves as 34-byte blocks, so they
422// need their own path -- dequantize once to packed f32, then the same __f32x8_fma kernel.
423// ★ A MODEL FILE'S NAME IS NOT ITS TYPE SYSTEM.
424func br_mmf32_band(w: *u8, x: *u8, out: *u8,
425 rows: i64, in_dim: i64, out_dim: i64, o0: i64, o1: i64) -> i64 {
426 let fa: *u8 = sys_mmap(64)
427 let az: *i64 = fa as *i64
428 let nch: i64 = in_dim / 8
429 var t: i64 = 0
430 while t < rows {
431 let ab: i64 = (x as i64) + t * in_dim * 4
432 var o: i64 = o0
433 while o < o1 {
434 let wb: i64 = (w as i64) + o * in_dim * 4
435 az[0] = 0
436 az[1] = 0
437 az[2] = 0
438 az[3] = 0
439 var ch: i64 = 0
440 while ch < nch {
441 __f32x8_fma(fa, (ab + ch * 32) as *u8, (wb + ch * 32) as *u8)
442 ch = ch + 1
443 }
444 var acc: i64 = __f32x8_hsum(fa)
445 var i: i64 = nch * 8
446 while i < in_dim {
447 acc = __f32_add(acc, __f32_mul(nx_le_read_u32(x, t * in_dim * 4 + i * 4),
448 nx_le_read_u32(w, o * in_dim * 4 + i * 4)))
449 i = i + 1
450 }
451 nx_le_write_u32(out, (t * out_dim + o) * 4, acc)
452 o = o + 1
453 }
454 t = t + 1
455 }
456 return 0
457}
458
459func br_matmul_f32(w: *u8, x: *u8, out: *u8, rows: i64, in_dim: i64, out_dim: i64, nw: i64) -> i64 {
460 if nw <= 1 { br_mmf32_band(w, x, out, rows, in_dim, out_dim, 0, out_dim); return 0 }
461 let pids: *i64 = sys_mmap(nw * 8 + 64) as *i64
462 var k: i64 = 0
463 while k < nw {
464 let o0: i64 = k * out_dim / nw
465 let o1: i64 = (k + 1) * out_dim / nw
466 let pid: i64 = sys_fork()
467 if pid == 0 { br_mmf32_band(w, x, out, rows, in_dim, out_dim, o0, o1); sys_exit(0) }
468 pids[k] = pid
469 k = k + 1
470 }
471 let st: *i64 = sys_mmap(64) as *i64
472 k = 0
473 while k < nw { sys_wait4(pids[k], st, 0); k = k + 1 }
474 return 0
475}
476
477// ---- QK-norm + interleaved RoPE, straight into the post-RoPE layout ----------------------
478// Reads q and k out of the PACKED qkv (heads 0..n-1 and n..2n-1) and writes
479// [head_dim, L, n_heads] -- RoPE permutes the head and token axes, so the destination stride is
480// deliberately different from the source's and both are written out rather than shared.
481func br_qk_rope(qkv: *u8, wqn: *u8, wkn: *u8, pe: *u8, qrp: *u8, krp: *u8,
482 T: i64, D: i64, QKV: i64, eps: i64) -> i64 {
483 let half: i64 = BR_HEAD_DIM / 2
484 let hdinv: i64 = nx_f32_div(nx_i32_to_f32(1), nx_i32_to_f32(BR_HEAD_DIM))
485 var l: i64 = 0
486 while l < T {
487 var h: i64 = 0
488 while h < BR_N_HEADS {
489 var side: i64 = 0
490 while side < 2 {
491 var wnorm: *u8 = wqn
492 var dst: *u8 = qrp
493 var hbase: i64 = 0
494 if side == 1 { wnorm = wkn; dst = krp; hbase = BR_N_HEADS }
495 let src: i64 = l * QKV + (hbase + h) * BR_HEAD_DIM
496 var ss: i64 = 0
497 var d: i64 = 0
498 while d < BR_HEAD_DIM {
499 let v: i64 = nx_le_read_u32(qkv, (src + d) * 4)
500 ss = __f32_add(ss, __f32_mul(v, v))
501 d = d + 1
502 }
503 let inv: i64 = nx_f32_div(nx_i32_to_f32(1),
504 nx_f32_sqrt(__f32_add(__f32_mul(ss, hdinv), eps)))
505 let orow: i64 = (h * T + l) * BR_HEAD_DIM
506 var j: i64 = 0
507 while j < half {
508 let a0: i64 = __f32_mul(__f32_mul(nx_le_read_u32(qkv, (src + 2 * j) * 4), inv),
509 nx_le_read_u32(wnorm, (2 * j) * 4))
510 let a1: i64 = __f32_mul(__f32_mul(nx_le_read_u32(qkv, (src + 2 * j + 1) * 4), inv),
511 nx_le_read_u32(wnorm, (2 * j + 1) * 4))
512 var r: i64 = 0
513 while r < 2 {
514 let pb: i64 = ((l * half + j) * 2 + r) * 2
515 let v: i64 = __f32_add(__f32_mul(a0, nx_le_read_u32(pe, pb * 4)),
516 __f32_mul(a1, nx_le_read_u32(pe, (pb + 1) * 4)))
517 nx_le_write_u32(dst, (orow + 2 * j + r) * 4, v)
518 r = r + 1
519 }
520 j = j + 1
521 }
522 side = side + 1
523 }
524 h = h + 1
525 }
526 l = l + 1
527 }
528 return 0
529}
530
531// ---- ONE FULL DiT BLOCK, layer-indexed, weights straight from the GGUF -------------------
532//
533// Shared by nx_gen_blockrun (one layer, graded) and nx_gen_ditchain (N layers, chained). The
534// SEQUENCING lives here exactly once: two copies of an op order drift silently, because both
535// still produce finite plausible activations.
536//
537// scr[] holds pre-allocated scratch so an N-layer chain does not re-map ~200MB per layer:
538// 0 h_attn 1 qkv 2 qrp 3 krp 4 aop 5 atto 6 an2 7 mid
539// 8 h_ffn 9 fw1 10 fw3 11 fact 12 fdn 13 fn2
540// 14 s_qkv 15 s_out 16 s_w1 17 s_w3 18 s_w2 (scale tables)
541// 19 adaln 20 mod_msa 21 tg_msa 22 mod_mlp 23 tg_mlp
542const BR_SCR_SLOTS: i64 = 24
543
544func br_block(gw: *i64, prefix: *u8, layer: i64, modulation: i64, cur: *u8, outb: *u8, scr: *i64,
545 T: i64, D: i64, QKV: i64, FD: i64, NW: i64,
546 pe: *u8, t_emb: *u8, eps: i64) -> i64 {
547 // Refuse rather than compute with geometry nobody set. Unbound would mean head_dim 0, which
548 // silently degenerates every loop below into zero iterations and returns an all-zero tensor --
549 // a result that looks like a working pipeline producing a black image.
550 if br_arch_ready() == 0 { return 0 - 99 }
551 let nm: *u8 = sys_mmap(256)
552
553 let h_attn: *u8 = scr[0] as *u8
554 let qkv: *u8 = scr[1] as *u8
555 let qrp: *u8 = scr[2] as *u8
556 let krp: *u8 = scr[3] as *u8
557 let aop: *u8 = scr[4] as *u8
558 let atto: *u8 = scr[5] as *u8
559 let an2: *u8 = scr[6] as *u8
560 let mid: *u8 = scr[7] as *u8
561 let h_ffn: *u8 = scr[8] as *u8
562 let fw1: *u8 = scr[9] as *u8
563 let fw3: *u8 = scr[10] as *u8
564 let fact: *u8 = scr[11] as *u8
565 let fdn: *u8 = scr[12] as *u8
566 let fn2: *u8 = scr[13] as *u8
567 let s_qkv: *i64 = scr[14] as *i64
568 let s_out: *i64 = scr[15] as *i64
569 let s_w1: *i64 = scr[16] as *i64
570 let s_w3: *i64 = scr[17] as *i64
571 let s_w2: *i64 = scr[18] as *i64
572 let adaln: *u8 = scr[19] as *u8
573 let mod_msa: *i64 = scr[20] as *i64
574 let tg_msa: *i64 = scr[21] as *i64
575 let mod_mlp: *i64 = scr[22] as *i64
576 let tg_mlp: *i64 = scr[23] as *i64
577
578 // weights for THIS layer
579 let W_qkv: *u8 = br_gw_q8(gw, br_name_p(nm, prefix, layer, "attention.qkv.weight" as *u8), D, QKV, s_qkv)
580 if (W_qkv as i64) == 0 { return 0 - 10 }
581 let W_out: *u8 = br_gw_q8(gw, br_name_p(nm, prefix, layer, "attention.out.weight" as *u8), D, D, s_out)
582 if (W_out as i64) == 0 { return 0 - 11 }
583 let W_w1: *u8 = br_gw_q8(gw, br_name_p(nm, prefix, layer, "feed_forward.w1.weight" as *u8), D, FD, s_w1)
584 if (W_w1 as i64) == 0 { return 0 - 12 }
585 let W_w3: *u8 = br_gw_q8(gw, br_name_p(nm, prefix, layer, "feed_forward.w3.weight" as *u8), D, FD, s_w3)
586 if (W_w3 as i64) == 0 { return 0 - 13 }
587 let W_w2: *u8 = br_gw_q8(gw, br_name_p(nm, prefix, layer, "feed_forward.w2.weight" as *u8), FD, D, s_w2)
588 if (W_w2 as i64) == 0 { return 0 - 14 }
589 let wn1: *u8 = br_gw_f32(gw, br_name_p(nm, prefix, layer, "attention_norm1.weight" as *u8), D)
590 if (wn1 as i64) == 0 { return 0 - 15 }
591 let wn2: *u8 = br_gw_f32(gw, br_name_p(nm, prefix, layer, "attention_norm2.weight" as *u8), D)
592 if (wn2 as i64) == 0 { return 0 - 16 }
593 let wf1: *u8 = br_gw_f32(gw, br_name_p(nm, prefix, layer, "ffn_norm1.weight" as *u8), D)
594 if (wf1 as i64) == 0 { return 0 - 17 }
595 let wf2: *u8 = br_gw_f32(gw, br_name_p(nm, prefix, layer, "ffn_norm2.weight" as *u8), D)
596 if (wf2 as i64) == 0 { return 0 - 18 }
597 let wqn: *u8 = br_gw_f32(gw, br_name_p(nm, prefix, layer, "attention.q_norm.weight" as *u8), BR_HEAD_DIM)
598 if (wqn as i64) == 0 { return 0 - 19 }
599 let wkn: *u8 = br_gw_f32(gw, br_name_p(nm, prefix, layer, "attention.k_norm.weight" as *u8), BR_HEAD_DIM)
600 if (wkn as i64) == 0 { return 0 - 20 }
601
602 // adaLN only exists on modulated blocks. context_refiner runs UNMODULATED: no adaLN tensor,
603 // no scale, and a PLAIN residual add instead of a tanh gate. Feeding it identity scale (1) and
604 // identity gate (1) reproduces that exactly through the same code path, so the two variants
605 // cannot drift apart the way two hand-written block functions would.
606 let one: i64 = nx_i32_to_f32(1)
607 var i0: i64 = 0
608 if modulation != 0 {
609 let ra: i64 = br_adaln(gw, prefix, layer, t_emb, adaln, 256, BR_CHUNKS * D)
610 if ra != 0 { return 0 - 30 }
611 while i0 < D {
612 mod_msa[i0] = __f32_add(one, nx_le_read_u32(adaln, (0 * D + i0) * 4))
613 tg_msa[i0] = nx_f32_tanh(nx_le_read_u32(adaln, (1 * D + i0) * 4))
614 mod_mlp[i0] = __f32_add(one, nx_le_read_u32(adaln, (2 * D + i0) * 4))
615 tg_mlp[i0] = nx_f32_tanh(nx_le_read_u32(adaln, (3 * D + i0) * 4))
616 i0 = i0 + 1
617 }
618 }
619 if modulation == 0 {
620 while i0 < D {
621 mod_msa[i0] = one
622 tg_msa[i0] = one
623 mod_mlp[i0] = one
624 tg_mlp[i0] = one
625 i0 = i0 + 1
626 }
627 }
628
629 br_rmsnorm(cur, wn1, mod_msa, h_attn, T, D, eps)
630 br_matmul(W_qkv, s_qkv, h_attn, qkv, T, D, QKV, NW)
631 br_qk_rope(qkv, wqn, wkn, pe, qrp, krp, T, D, QKV, eps)
632
633 let scale: i64 = nx_f32_div(nx_i32_to_f32(1), nx_f32_sqrt(nx_i32_to_f32(BR_HEAD_DIM)))
634 var sw: i64 = NW
635 if sw > BR_N_HEADS { sw = BR_N_HEADS }
636 if sw <= 1 { br_sdpa_band(qrp, krp, qkv, aop, T, D, QKV, BR_HEAD_DIM, BR_N_HEADS, scale, 0, BR_N_HEADS) }
637 if sw > 1 {
638 let sp: *i64 = sys_mmap(sw * 8 + 64) as *i64
639 var sk: i64 = 0
640 while sk < sw {
641 let a0: i64 = sk * BR_N_HEADS / sw
642 let a1: i64 = (sk + 1) * BR_N_HEADS / sw
643 let pid: i64 = sys_fork()
644 if pid == 0 { br_sdpa_band(qrp, krp, qkv, aop, T, D, QKV, BR_HEAD_DIM, BR_N_HEADS, scale, a0, a1); sys_exit(0) }
645 sp[sk] = pid
646 sk = sk + 1
647 }
648 let st: *i64 = sys_mmap(64) as *i64
649 sk = 0
650 while sk < sw { sys_wait4(sp[sk], st, 0); sk = sk + 1 }
651 }
652
653 br_matmul(W_out, s_out, aop, atto, T, D, D, NW)
654 br_rmsnorm(atto, wn2, 0 as *i64, an2, T, D, eps)
655 br_gate_resid(an2, tg_msa, cur, mid, T, D)
656
657 br_rmsnorm(mid, wf1, mod_mlp, h_ffn, T, D, eps)
658 br_matmul(W_w1, s_w1, h_ffn, fw1, T, D, FD, NW)
659 br_matmul(W_w3, s_w3, h_ffn, fw3, T, D, FD, NW)
660 br_fork_range(0, fw1, fw3, fact, T * FD, NW)
661 br_matmul(W_w2, s_w2, fact, fdn, T, FD, D, NW)
662 br_rmsnorm(fdn, wf2, 0 as *i64, fn2, T, D, eps)
663 br_gate_resid(fn2, tg_mlp, mid, outb, T, D)
664 return 0
665}