nx_f32_linear_attention.nx source
↩ module page · 103 lines · 4336 B
1// nx_f32_linear_attention.nx -- sovereign O(n) linear attention (SANA-style), the DiT VRAM/speed lever.
2//
3// From DiT-IC/SANA (ingested 2026-07-01): replace softmax attention (O(n^2 d), materializes the n×n matrix)
4// with LINEAR attention (O(n d^2), no n×n matrix): with a positive feature map phi(x)=relu(x)+1,
5// S = Σ_j phi(K_j) ⊗ V_j (d_k × d_v) z = Σ_j phi(K_j) (d_k)
6// out_i = ( phi(Q_i) @ S ) / ( phi(Q_i) · z )
7// For n_tokens >> head_dim this is a big compute/VRAM cut -> smaller-VRAM, faster, more-hardware Z-Image DiT.
8// Gate: uniform K/Q -> out = mean(V) (matches softmax's uniform case); + a non-uniform hand-computed check.
9// Q,K,V,out: flat *i64 f32 bits [n_tokens, head_dim].
10// license_tier: ORIGINAL
11import "nx_syscalls.nx"
12import "nx_f32.nx"
13import "nx_f32_div.nx"
14import "nx_f32_cvt.nx"
15
16// phi(x) = relu(x) + 1 (positive feature map). relu via the f32 sign bit.
17func la_phi(x: i64) -> i64 {
18 var r: i64 = x
19 if (x & 0x80000000) != 0 { r = 0 } // negative -> 0
20 return nx_f32_add(r, nx_i32_to_f32(1))
21}
22
23func nx_f32_linear_attention(Q: *i64, K: *i64, V: *i64, n_tokens: i64, head_dim: i64, out: *i64) -> i64 {
24 let S: *i64 = sys_mmap(head_dim * head_dim * 8) as *i64
25 let z: *i64 = sys_mmap(head_dim * 8) as *i64
26 var a: i64 = 0
27 while a < head_dim * head_dim { S[a] = 0; a = a + 1 }
28 a = 0
29 while a < head_dim { z[a] = 0; a = a + 1 }
30
31 // accumulate KV-summary
32 var j: i64 = 0
33 while j < n_tokens {
34 a = 0
35 while a < head_dim {
36 let pk: i64 = la_phi(K[j * head_dim + a])
37 z[a] = nx_f32_add(z[a], pk)
38 var b: i64 = 0
39 while b < head_dim {
40 S[a * head_dim + b] = nx_f32_add(S[a * head_dim + b], nx_f32_mul(pk, V[j * head_dim + b]))
41 b = b + 1
42 }
43 a = a + 1
44 }
45 j = j + 1
46 }
47
48 // per-query
49 var i: i64 = 0
50 while i < n_tokens {
51 var den: i64 = 0
52 a = 0
53 while a < head_dim { den = nx_f32_add(den, nx_f32_mul(la_phi(Q[i * head_dim + a]), z[a])); a = a + 1 }
54 var b: i64 = 0
55 while b < head_dim {
56 var num: i64 = 0
57 a = 0
58 while a < head_dim { num = nx_f32_add(num, nx_f32_mul(la_phi(Q[i * head_dim + a]), S[a * head_dim + b])); a = a + 1 }
59 out[i * head_dim + b] = nx_f32_div(num, den)
60 b = b + 1
61 }
62 i = i + 1
63 }
64 return 0
65}
66
67func la_close(x: i64, e: i64, tol: i64) -> i64 {
68 if (nx_f32_sub(x, e) & 0x7FFFFFFF) < tol { return 1 }
69 return 0
70}
71
72func main() -> i64 {
73 let tol: i64 = nx_f32_div(nx_i32_to_f32(1), nx_i32_to_f32(100)) // 0.01
74 let Q: *i64 = sys_mmap(4 * 8) as *i64
75 let K: *i64 = sys_mmap(4 * 8) as *i64
76 let V: *i64 = sys_mmap(4 * 8) as *i64
77 let out: *i64 = sys_mmap(4 * 8) as *i64
78
79 // ---- gate 1: uniform K/Q (all 0) -> out = mean(V) = [4,6] ----
80 var i: i64 = 0
81 while i < 4 { Q[i] = 0; K[i] = 0; i = i + 1 }
82 V[0] = nx_i32_to_f32(2); V[1] = nx_i32_to_f32(4); V[2] = nx_i32_to_f32(6); V[3] = nx_i32_to_f32(8)
83 nx_f32_linear_attention(Q, K, V, 2, 2, out)
84 if la_close(out[0], nx_i32_to_f32(4), tol) != 1 { return 10 }
85 if la_close(out[1], nx_i32_to_f32(6), tol) != 1 { return 11 }
86 if la_close(out[2], nx_i32_to_f32(4), tol) != 1 { return 12 }
87 if la_close(out[3], nx_i32_to_f32(6), tol) != 1 { return 13 }
88
89 // ---- gate 2: non-uniform. Q=K=[[1,0],[0,1]], V=[[2,4],[6,8]] ----
90 // phi(K0)=[2,1],phi(K1)=[1,2]; S=[[10,16],[14,20]]; z=[3,3]
91 // out0 = phi(Q0=[1,0])=[2,1] -> num=[34,52]/den=9 = [34/9,52/9]
92 // out1 = phi(Q1=[0,1])=[1,2] -> num=[38,56]/den=9 = [38/9,56/9]
93 Q[0] = nx_i32_to_f32(1); Q[1] = nx_i32_to_f32(0); Q[2] = nx_i32_to_f32(0); Q[3] = nx_i32_to_f32(1)
94 K[0] = nx_i32_to_f32(1); K[1] = nx_i32_to_f32(0); K[2] = nx_i32_to_f32(0); K[3] = nx_i32_to_f32(1)
95 nx_f32_linear_attention(Q, K, V, 2, 2, out)
96 let nine: i64 = nx_i32_to_f32(9)
97 if la_close(out[0], nx_f32_div(nx_i32_to_f32(34), nine), tol) != 1 { return 20 }
98 if la_close(out[1], nx_f32_div(nx_i32_to_f32(52), nine), tol) != 1 { return 21 }
99 if la_close(out[2], nx_f32_div(nx_i32_to_f32(38), nine), tol) != 1 { return 22 }
100 if la_close(out[3], nx_f32_div(nx_i32_to_f32(56), nine), tol) != 1 { return 23 }
101
102 return 0
103}