code wiki / (root) / nx_f32_linear_attention.nx

nx_f32_linear_attention.nx

buildroot/runtime/nx_f32_linear_attention.nx

4336 B103 linesdepth 4pulls 6 transitivereach 0 importersview sourcekind tooltopic f32
docsdependenciesstructsconstsfunctions

about

nx_f32_linear_attention.nx -- sovereign O(n) linear attention (SANA-style), the DiT VRAM/speed lever. From DiT-IC/SANA (ingested 2026-07-01): replace softmax attention (O(n^2 d), materializes the n×n matrix) with LINEAR attention (O(n d^2), no n×n matrix): with a positive feature map phi(x)=relu(x)+1, S = Σ_j phi(K_j) ⊗ V_j (d_k × d_v) z = Σ_j phi(K_j) (d_k) out_i = ( phi(Q_i) @ S ) / ( phi(Q_i) · z ) For n_tokens >> head_dim this is a big compute/VRAM cut -> smaller-VRAM, faster, more-hardware Z-Image DiT. Gate: uniform K/Q -> out = mean(V) (matches softmax's uniform case); + a non-uniform hand-computed check. Q,K,V,out: flat *i64 f32 bits [n_tokens, head_dim]. license_tier: ORIGINAL

dependencies 4 imports · 0 importers

nx_syscalls.nx nx_f32.nx nx_f32_div.nx nx_f32_cvt.nx nx_f32_linear_attention.nx

imports: nx_syscalls.nxnx_f32.nxnx_f32_div.nxnx_f32_cvt.nx

imported by: nobody (leaf or entry point)

call flow from main pre-order; caps 40 nodes / depth 6 declared; ↻ = already shown

main 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_i32_to_f32 sys_mmap nx_f32_linear_attention sys_mmap ↻ la_phi nx_f32_add nx_f32_classify ↻ nx_f32_sign ↻ nx_f32_mant_field ↻ nx_f32_exp_field ↻ nx_i32_to_f32 ↻ nx_f32_add ↻ nx_f32_mul nx_f32_classify ↻ nx_f32_sign ↻ nx_f32_mant_field ↻ nx_f32_exp_field ↻ nx_f32_div ↻ la_close nx_f32_sub nx_f32_add ↻ nx_f32_neg

structs

none

consts

none

functions

17func la_phi(x: i64) -> i64
23func nx_f32_linear_attention(Q: *i64, K: *i64, V: *i64, n_tokens: i64, head_dim: i64, out: *i64) -> i64
67func la_close(x: i64, e: i64, tol: i64) -> i64
called by 1: main calls 1: nx_f32_sub
72func main() -> i64