`nx_f32_adaln.nx` source

↩ module page · 108 lines · 4583 B
1// nx_f32_adaln.nx -- software-f32 adaptive LayerNorm (adaLN / adaLN-Zero) modulation, the DiT-specific brick.
2//
3// sd-server -> Nishi migration: this is what turns a plain transformer block into a DIFFUSION transformer
4// (DiT) block (Peebles & Xie 2023 "Scalable Diffusion Models with Transformers"; Z-Image uses it). A small
5// MLP maps the conditioning vector (timestep + pooled text) to per-block modulation params
6// {shift, scale, gate} for each sublayer; the block then applies:
7//
8//   modulate(norm(x), shift, scale) = norm(x) * (1 + scale) + shift     (broadcast over tokens)
9//   x = x + gate * sublayer(modulate(...))                              (gated residual; adaLN-Zero starts gate=0)
10//
11// This file owns the two elementwise ops (modulate + gated-residual). The conditioning MLP that produces
12// shift/scale/gate is a plain linear (nx_f32_matmul) the caller supplies. Composes nx_f32_add/mul + cvt.
13// x, out, res, h: flat *i64 f32 bits [n_tokens, D]. scale, shift, gate: [D] f32 (per-feature, token-broadcast).
14// license_tier: ORIGINAL
15import "nx_syscalls.nx"
16import "nx_f32.nx"
17import "nx_f32_cvt.nx"
18
19const NX_F32ADALN_OK:  i64 = 0
20const NX_F32ADALN_ERR: i64 = 1
21
22// modulate: out[t][d] = x[t][d] * (1 + scale[d]) + shift[d]
23func nx_f32_adaln_modulate(x: *i64, n_tokens: i64, D: i64, scale: *i64, shift: *i64, out: *i64) -> i64 {
24    if n_tokens <= 0 { return NX_F32ADALN_ERR }
25    if D <= 0 { return NX_F32ADALN_ERR }
26    let one: i64 = nx_i32_to_f32(1)
27    var t: i64 = 0
28    while t < n_tokens {
29        let base: i64 = t * D
30        var d: i64 = 0
31        while d < D {
32            let sp1: i64 = nx_f32_add(one, scale[d])               // 1 + scale[d]
33            out[base + d] = nx_f32_add(nx_f32_mul(x[base + d], sp1), shift[d])
34            d = d + 1
35        }
36        t = t + 1
37    }
38    return NX_F32ADALN_OK
39}
40
41// gated residual: out[t][d] = res[t][d] + gate[d] * h[t][d]   (adaLN-Zero: gate starts at 0 -> identity init)
42func nx_f32_adaln_gate(res: *i64, h: *i64, n_tokens: i64, D: i64, gate: *i64, out: *i64) -> i64 {
43    if n_tokens <= 0 { return NX_F32ADALN_ERR }
44    if D <= 0 { return NX_F32ADALN_ERR }
45    var t: i64 = 0
46    while t < n_tokens {
47        let base: i64 = t * D
48        var d: i64 = 0
49        while d < D {
50            out[base + d] = nx_f32_add(res[base + d], nx_f32_mul(gate[d], h[base + d]))
51            d = d + 1
52        }
53        t = t + 1
54    }
55    return NX_F32ADALN_OK
56}
57
58// ===== Self-test (inline gate) -- all invariants are bit-exact f32 =====
59//   (a) modulate identity: scale=0, shift=0 -> out == x
60//   (b) modulate: scale=1 -> out == 2x ;  shift=k -> out == x + k
61//   (c) gate-Zero: gate=0 -> out == res ;  gate=1 -> out == res + h
62func main() -> i64 {
63    let x:   *i64 = sys_mmap(16 * 8) as *i64
64    let out: *i64 = sys_mmap(16 * 8) as *i64
65    let res: *i64 = sys_mmap(16 * 8) as *i64
66    let hh:  *i64 = sys_mmap(16 * 8) as *i64
67    let sc:  *i64 = sys_mmap(8 * 8) as *i64
68    let sh:  *i64 = sys_mmap(8 * 8) as *i64
69    let gt:  *i64 = sys_mmap(8 * 8) as *i64
70    let one: i64 = nx_i32_to_f32(1)
71
72    // n_tokens=2, D=2. x = [1,2,3,4]
73    var i: i64 = 0
74    while i < 4 { x[i] = nx_i32_to_f32(i + 1); i = i + 1 }
75
76    // (a) identity modulate
77    sc[0] = 0; sc[1] = 0; sh[0] = 0; sh[1] = 0
78    let va: i64 = nx_f32_adaln_modulate(x, 2, 2, sc, sh, out)
79    if va != NX_F32ADALN_OK { return 10 }
80    i = 0
81    while i < 4 { if out[i] != x[i] { return 20 } i = i + 1 }
82
83    // (b) scale=1 -> 2x ; then shift=5 (scale 0) -> x+5
84    sc[0] = one; sc[1] = one; sh[0] = 0; sh[1] = 0
85    nx_f32_adaln_modulate(x, 2, 2, sc, sh, out)
86    if out[0] != nx_i32_to_f32(2) { return 30 }   // 1*2
87    if out[1] != nx_i32_to_f32(4) { return 31 }   // 2*2
88    if out[3] != nx_i32_to_f32(8) { return 32 }   // 4*2
89    sc[0] = 0; sc[1] = 0; sh[0] = nx_i32_to_f32(5); sh[1] = nx_i32_to_f32(5)
90    nx_f32_adaln_modulate(x, 2, 2, sc, sh, out)
91    if out[0] != nx_i32_to_f32(6) { return 33 }   // 1+5
92    if out[3] != nx_i32_to_f32(9) { return 34 }   // 4+5
93
94    // (c) gated residual
95    i = 0
96    while i < 4 { res[i] = nx_i32_to_f32(i + 1); hh[i] = nx_i32_to_f32((i + 1) * 10); i = i + 1 }
97    gt[0] = 0; gt[1] = 0
98    let vc: i64 = nx_f32_adaln_gate(res, hh, 2, 2, gt, out)
99    if vc != NX_F32ADALN_OK { return 40 }
100    i = 0
101    while i < 4 { if out[i] != res[i] { return 41 } i = i + 1 }   // gate 0 -> unchanged
102    gt[0] = one; gt[1] = one
103    nx_f32_adaln_gate(res, hh, 2, 2, gt, out)
104    if out[0] != nx_i32_to_f32(11) { return 42 }   // 1 + 10
105    if out[3] != nx_i32_to_f32(44) { return 43 }   // 4 + 40
106
107    return 0
108}
nx_f32_adaln.nx source

`nx_f32_adaln.nx` source
