nx_f32_attention.nx source
↩ module page · 104 lines · 4763 B
1// nx_f32_attention.nx -- software-f32 multi-token scaled-dot-product attention (single head), the first
2// brick of the DiT diffusion-backbone half of the sd-server -> Nishi migration.
3//
4// `nx_f32_mha.nx` is SINGLE-TOKEN (autoregressive decode, n_tokens=1, degenerate softmax). The DiT needs
5// MULTI-TOKEN / bidirectional self-attention over all latent (+text) tokens -- exactly the follow-on its
6// header flags ("adds K^T transpose + multi-row softmax"). This is the standard prefill attention:
7//
8// scores[i][j] = dot(Q[i], K[j]) * scale (Q@K^T, scaled)
9// probs[i] = softmax(scores[i]) (over the j axis, numerically stable)
10// out[i] = sum_j probs[i][j] * V[j] (probs@V)
11//
12// Composes the gated primitives nx_f32_add/mul + nx_f32_softmax. Single-head; multi-head = caller loops
13// heads over the head_dim slices of Q/K/V (RoPE + QK-RMSNorm wrap this in the DiT block, next rung).
14// Q,K,V,out: flat *i64 f32 bits, [n_tokens, head_dim] row-major. score_row/prob_row: caller scratch [n_tokens].
15// license_tier: ORIGINAL
16import "nx_syscalls.nx"
17import "nx_f32.nx"
18import "nx_f32_div.nx"
19import "nx_f32_cvt.nx"
20import "nx_f32_softmax.nx"
21
22const NX_F32ATTN_OK: i64 = 0
23const NX_F32ATTN_ERR: i64 = 1
24
25func nx_f32_attention(Q: *i64, K: *i64, V: *i64, n_tokens: i64, head_dim: i64, scale: i64,
26 out: *i64, score_row: *i64, prob_row: *i64) -> i64 {
27 if n_tokens <= 0 { return NX_F32ATTN_ERR }
28 if head_dim <= 0 { return NX_F32ATTN_ERR }
29 var i: i64 = 0
30 while i < n_tokens {
31 // scores for query i against every key j
32 var j: i64 = 0
33 while j < n_tokens {
34 var acc: i64 = 0 // f32 +0.0
35 var d: i64 = 0
36 while d < head_dim {
37 acc = nx_f32_add(acc, nx_f32_mul(Q[i * head_dim + d], K[j * head_dim + d]))
38 d = d + 1
39 }
40 score_row[j] = nx_f32_mul(acc, scale)
41 j = j + 1
42 }
43 // softmax over the row (numerically stable, gated organ)
44 nx_f32_softmax(score_row, n_tokens, prob_row)
45 // out[i] = sum_j probs[i][j] * V[j]
46 var d2: i64 = 0
47 while d2 < head_dim {
48 var o: i64 = 0 // f32 +0.0
49 var jj: i64 = 0
50 while jj < n_tokens {
51 o = nx_f32_add(o, nx_f32_mul(prob_row[jj], V[jj * head_dim + d2]))
52 jj = jj + 1
53 }
54 out[i * head_dim + d2] = o
55 d2 = d2 + 1
56 }
57 i = i + 1
58 }
59 return NX_F32ATTN_OK
60}
61
62// ===== Self-test (inline gate) ====================================
63// f32 magnitude compare: |a - b| < tol <=> (nx_f32_sub(a,b) & 0x7FFFFFFF) < tol_bits.
64// (a) Q == 0 -> every score 0 -> softmax uniform 1/n -> out[i] == mean over tokens of V (for all i).
65// (b) n_tokens == 1 -> softmax([s]) == [1] -> out == V[0].
66// (c) bad dim -> ERR.
67func main() -> i64 {
68 let Q: *i64 = sys_mmap(16 * 8) as *i64
69 let K: *i64 = sys_mmap(16 * 8) as *i64
70 let V: *i64 = sys_mmap(16 * 8) as *i64
71 let out: *i64 = sys_mmap(16 * 8) as *i64
72 let sr: *i64 = sys_mmap(16 * 8) as *i64
73 let pr: *i64 = sys_mmap(16 * 8) as *i64
74 let scale: i64 = nx_i32_to_f32(1)
75 let tolb: i64 = nx_f32_div(nx_i32_to_f32(1), nx_i32_to_f32(100)) & 0x7FFFFFFF // 1e-2
76
77 // (a) n_tokens=2, head_dim=2. Q=0. V = [[2,4],[6,8]] -> mean over tokens = [4,6].
78 var i: i64 = 0
79 while i < 4 { Q[i] = 0; K[i] = nx_i32_to_f32(1); i = i + 1 }
80 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)
81 let va: i64 = nx_f32_attention(Q, K, V, 2, 2, scale, out, sr, pr)
82 if va != NX_F32ATTN_OK { return 10 }
83 let e4: i64 = nx_i32_to_f32(4)
84 let e6: i64 = nx_i32_to_f32(6)
85 if (nx_f32_sub(out[0], e4) & 0x7FFFFFFF) >= tolb { return 20 } // out[0][0]
86 if (nx_f32_sub(out[1], e6) & 0x7FFFFFFF) >= tolb { return 21 } // out[0][1]
87 if (nx_f32_sub(out[2], e4) & 0x7FFFFFFF) >= tolb { return 22 } // out[1][0]
88 if (nx_f32_sub(out[3], e6) & 0x7FFFFFFF) >= tolb { return 23 } // out[1][1]
89
90 // (b) n_tokens=1, head_dim=2. out == V[0] = [5,7].
91 Q[0] = nx_i32_to_f32(1); Q[1] = nx_i32_to_f32(1)
92 K[0] = nx_i32_to_f32(1); K[1] = nx_i32_to_f32(1)
93 V[0] = nx_i32_to_f32(5); V[1] = nx_i32_to_f32(7)
94 let vb: i64 = nx_f32_attention(Q, K, V, 1, 2, scale, out, sr, pr)
95 if vb != NX_F32ATTN_OK { return 30 }
96 if (nx_f32_sub(out[0], nx_i32_to_f32(5)) & 0x7FFFFFFF) >= tolb { return 31 }
97 if (nx_f32_sub(out[1], nx_i32_to_f32(7)) & 0x7FFFFFFF) >= tolb { return 32 }
98
99 // (c) bad dim
100 let vc: i64 = nx_f32_attention(Q, K, V, 0, 2, scale, out, sr, pr)
101 if vc == NX_F32ATTN_OK { return 40 }
102
103 return 0
104}