code wiki / (root) / nx_f32_dit_block_tiny.nx

nx_f32_dit_block_tiny.nx

buildroot/runtime/nx_f32_dit_block_tiny.nx

7403 B169 linesdepth 7pulls 12 transitivereach 1 importersview sourcekind tooltopic f32
docsdependenciesstructsconstsfunctions

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

nx_syscalls.nx nx_f32.nx nx_f32_cvt.nx nx_f32_rmsnorm.nx nx_f32_attention.nx nx_f32_adaln.nx nx_f32_activations.nx nx_f32_dit_block_tiny.nx nx_f32_rectflow_denoise.nx

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

main sys_mmap nx_i32_to_f32 nx_f32_div nx_f32_classify nx_f32_exp_field nx_f32_mant_field nx_f32_sign nx_f32_mant_field ↻ nx_f32_exp_field ↻ nx_f32_dit_block_tiny sys_mmap ↻ nx_i32_to_f32 ↻ nx_f32_rmsnorm nx_i32_to_f32 ↻ nx_f32_sqrt nx_f32_classify ↻ nx_f32_sign ↻ nx_f32_exp_field ↻ nx_f32_mant_field ↻ _f32_isqrt nx_f32_adaln_modulate nx_i32_to_f32 ↻ nx_f32_add nx_f32_classify ↻ nx_f32_sign ↻ nx_f32_mant_field ↻ nx_f32_exp_field ↻ nx_f32_mul nx_f32_classify ↻ nx_f32_sign ↻ nx_f32_mant_field ↻ nx_f32_exp_field ↻ ditb_linear nx_f32_add ↻ nx_f32_mul ↻ nx_f32_attention nx_f32_add ↻ nx_f32_mul ↻ nx_f32_softmax

structs

none

consts

21const K_MAGIC_100000: i64 = 100000

functions

24func ditb_linear(inp: *i64, n_tokens: i64, Din: i64, W: *i64, Dout: i64, out: *i64) -> i64
41func ditb_swiglu(inp: *i64, n_tokens: i64, D: i64, d_ff: i64, W1: *i64, W3: *i64, W2: *i64, out: *i64, h: *i64) -> i64
66func nx_f32_dit_block_tiny(x: *i64, n_tokens: i64, D: i64, d_ff: i64, gamma: *i64, eps: i64,
114func main() -> i64