nx_f32_linear_attention.nx
buildroot/runtime/nx_f32_linear_attention.nx
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
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
structs
| none |
consts
| none |
functions
| 17 | func la_phi(x: i64) -> i64 |
| 23 | func nx_f32_linear_attention(Q: *i64, K: *i64, V: *i64, n_tokens: i64, head_dim: i64, out: *i64) -> i64 |
| 67 | func la_close(x: i64, e: i64, tol: i64) -> i64 |
| 72 | func main() -> i64 |