nx_f32_dit_block_linear.nx
buildroot/runtime/nx_f32_dit_block_linear.nx
about
nx_f32_dit_block_linear.nx -- SANA-style DiT block using LINEAR attention (the efficient Z-Image DiT unit).
From DiT-IC/SANA (ingested 2026-07-01): a DiT block with O(n) linear attention instead of O(n^2) softmax.
adaLN(x; sc1,sh1) -> Q/K/V -> linear-attn -> Wo -> x1 = x + gt1*ao
adaLN(x1; sc2,sh2) -> SwiGLU(Wg,Wu,Wd) -> out = x1 + gt2*ffn
adaLN-Zero gate (gt1=gt2=0) => out == x bit-exact, proving the WIRING independent of attn/FFN internals
(linear-attn + SwiGLU are separately gated). Single-head (head_dim=D). Linear-attn inlined; SiLU via lib.
license_tier: ORIGINAL
dependencies 5 imports · 0 importers
imports: nx_syscalls.nxnx_f32.nxnx_f32_div.nxnx_f32_cvt.nxnx_f32_activations.nx
imported by: nobody (leaf or entry point)
call flow from main pre-order; caps 40 nodes / depth 6 declared; ↻ = already shown
structs
| none |
consts
| none |
functions
| 15 | func dbl_phi(v: i64) -> i64 |
| 22 | func dbl_linear(inp: *i64, W: *i64, out: *i64, n: i64, id: i64, od: i64) -> i64 |
| 39 | func dbl_modulate(x: *i64, sc: *i64, sh: *i64, out: *i64, n: i64, D: i64) -> i64 |
| 52 | func dbl_linattn(Q: *i64, K: *i64, V: *i64, n: i64, D: i64, out: *i64, S: *i64, z: *i64) -> i64 |
| 88 | func dbl_gate_res(x: *i64, y: *i64, gt: *i64, out: *i64, n: i64, D: i64) -> i64 |
| 98 | func nx_f32_dit_block_linear(x: *i64, n: i64, D: i64, d_ff: i64, |
| 140 | func main() -> i64 |