nx_f32_dit_block_tiny.nx
buildroot/runtime/nx_f32_dit_block_tiny.nx
about
nx_f32_dit_block_tiny.nx -- first sovereign f32 DiT (Diffusion Transformer) BLOCK, assembled from the
gated bricks. The repeating unit of the Z-Image diffusion backbone (Peebles&Xie 2023 DiT; adaLN-Zero).
sd-server -> Nishi migration milestone (the DiT half's analogue of nx_f32_vae_decode_tiny). Per token:
x1 = x + gate1 * Wo @ Attention( Wq,Wk,Wv @ adaLN_mod( RMSNorm(x), scale1, shift1 ) )
out = x1 + gate2 * SwiGLU( adaLN_mod( RMSNorm(x1), scale2, shift2 ) )
Composes the gated organs: nx_f32_rmsnorm + nx_f32_adaln (modulate + gated residual) + nx_f32_attention +
nx_f32_silu (in the inline SwiGLU) + inline linear projections. Single head v1 (head_dim = D). The adaLN
scale/shift/gate come from the conditioning MLP (caller-supplied; adaLN-Zero starts gate=0 => identity).
x,out: flat *i64 f32 bits [n_tokens, D]. W{q,k,v,o}: [D,D]. SwiGLU W1,W3: [d_ff,D], W2: [D,d_ff].
sc/sh/gt{1,2}: [D]. gamma: [D] (RMSNorm scale). Scratch is allocated internally (fork-per-use organ).
license_tier: ORIGINAL
dependencies 7 imports · 1 importers
imports: nx_syscalls.nxnx_f32.nxnx_f32_cvt.nxnx_f32_rmsnorm.nxnx_f32_attention.nxnx_f32_adaln.nxnx_f32_activations.nx
imported by: nx_f32_rectflow_denoise.nx
call flow from main pre-order; caps 40 nodes / depth 6 declared; ↻ = already shown
structs
| none |
consts
| 21 | const K_MAGIC_100000: i64 = 100000 |
functions
| 24 | func ditb_linear(inp: *i64, n_tokens: i64, Din: i64, W: *i64, Dout: i64, out: *i64) -> i64 |
| 41 | func ditb_swiglu(inp: *i64, n_tokens: i64, D: i64, d_ff: i64, W1: *i64, W3: *i64, W2: *i64, out: *i64, h: *i64) -> i64 |
| 66 | func nx_f32_dit_block_tiny(x: *i64, n_tokens: i64, D: i64, d_ff: i64, gamma: *i64, eps: i64, |
| 114 | func main() -> i64 |