nx_f32_gqa_attention.nx source
↩ module page · 103 lines · 4560 B
1// nx_f32_gqa_attention.nx -- software-f32 Grouped-Query multi-head attention (Ainslie 2023 GQA).
2//
3// sd-server -> Nishi migration: Qwen3-4B (the Z-Image text encoder + companion chat model) uses GQA with
4// 32 query heads sharing 8 kv heads (4:1, confirmed from the real GGUF: q_out 4096, kv_out 1024). This wraps
5// the gated single-head `nx_f32_attention`: for each query head h it gathers Q_h and the SHARED K/V of kv
6// head (h / group), runs scaled-dot-product attention, and scatters the head output back. This is the
7// attention core of a real Qwen (and DiT) transformer layer.
8//
9// Q,out: flat *i64 f32 bits [n_tokens, n_q_heads*head_dim]. K,V: [n_tokens, n_kv_heads*head_dim].
10// license_tier: ORIGINAL
11import "nx_syscalls.nx"
12import "nx_f32.nx"
13import "nx_f32_div.nx"
14import "nx_f32_cvt.nx"
15import "nx_f32_attention.nx"
16
17const NX_F32GQA_OK: i64 = 0
18const NX_F32GQA_ERR: i64 = 1
19
20func nx_f32_gqa_attention(Q: *i64, K: *i64, V: *i64, n_tokens: i64,
21 n_q_heads: i64, n_kv_heads: i64, head_dim: i64, scale: i64,
22 out: *i64) -> i64 {
23 if n_kv_heads <= 0 { return NX_F32GQA_ERR }
24 if n_q_heads - (n_q_heads / n_kv_heads) * n_kv_heads != 0 { return NX_F32GQA_ERR }
25 let group: i64 = n_q_heads / n_kv_heads
26 let qd: i64 = n_q_heads * head_dim
27 let kvd: i64 = n_kv_heads * head_dim
28 let Qh: *i64 = sys_mmap(n_tokens * head_dim * 8) as *i64
29 let Kh: *i64 = sys_mmap(n_tokens * head_dim * 8) as *i64
30 let Vh: *i64 = sys_mmap(n_tokens * head_dim * 8) as *i64
31 let Oh: *i64 = sys_mmap(n_tokens * head_dim * 8) as *i64
32 let sr: *i64 = sys_mmap(n_tokens * 8) as *i64
33 let pr: *i64 = sys_mmap(n_tokens * 8) as *i64
34 var h: i64 = 0
35 while h < n_q_heads {
36 let kv: i64 = h / group
37 // gather this head's Q and its shared K/V
38 var t: i64 = 0
39 while t < n_tokens {
40 var d: i64 = 0
41 while d < head_dim {
42 Qh[t * head_dim + d] = Q[t * qd + h * head_dim + d]
43 Kh[t * head_dim + d] = K[t * kvd + kv * head_dim + d]
44 Vh[t * head_dim + d] = V[t * kvd + kv * head_dim + d]
45 d = d + 1
46 }
47 t = t + 1
48 }
49 nx_f32_attention(Qh, Kh, Vh, n_tokens, head_dim, scale, Oh, sr, pr)
50 // scatter the head output back into the full [n_tokens, qd] tensor
51 t = 0
52 while t < n_tokens {
53 var d: i64 = 0
54 while d < head_dim { out[t * qd + h * head_dim + d] = Oh[t * head_dim + d]; d = d + 1 }
55 t = t + 1
56 }
57 h = h + 1
58 }
59 return NX_F32GQA_OK
60}
61
62// ===== Self-test (inline gate) ====================================
63// 2 query heads sharing 1 kv head (group=2), head_dim=2, n_tokens=2. With Q=0 every score is 0 ->
64// uniform softmax -> each head's output = mean over tokens of the (shared) kv head's V. Since both q
65// heads share kv head 0, BOTH heads produce the same mean(V) -> proves the grouping + the attention.
66func main() -> i64 {
67 let Q: *i64 = sys_mmap(16 * 8) as *i64
68 let K: *i64 = sys_mmap(16 * 8) as *i64
69 let V: *i64 = sys_mmap(16 * 8) as *i64
70 let out: *i64 = sys_mmap(16 * 8) as *i64
71 let scale: i64 = nx_i32_to_f32(1)
72 let tolb: i64 = nx_f32_div(nx_i32_to_f32(1), nx_i32_to_f32(100)) & 0x7FFFFFFF
73
74 // Q = 0 (all)
75 var i: i64 = 0
76 while i < 8 { Q[i] = 0; i = i + 1 }
77 // K arbitrary (Q=0 makes scores 0 regardless); K is [2 tokens, 1 kv head * 2] = 4 values
78 K[0] = nx_i32_to_f32(1); K[1] = nx_i32_to_f32(1); K[2] = nx_i32_to_f32(1); K[3] = nx_i32_to_f32(1)
79 // V = [[2,4],[6,8]] over the single kv head -> mean over tokens = [4,6]
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
82 let r: i64 = nx_f32_gqa_attention(Q, K, V, 2, 2, 1, 2, scale, out)
83 if r != NX_F32GQA_OK { return 10 }
84
85 // out is [2 tokens, 2 q-heads * 2] = 8 values; every (head, token) output should be [4,6]
86 let e4: i64 = nx_i32_to_f32(4)
87 let e6: i64 = nx_i32_to_f32(6)
88 var t: i64 = 0
89 while t < 2 {
90 var hh: i64 = 0
91 while hh < 2 {
92 let b0: i64 = out[t * 4 + hh * 2 + 0]
93 let b1: i64 = out[t * 4 + hh * 2 + 1]
94 if (nx_f32_sub(b0, e4) & 0x7FFFFFFF) >= tolb { return 20 }
95 if (nx_f32_sub(b1, e6) & 0x7FFFFFFF) >= tolb { return 21 }
96 hh = hh + 1
97 }
98 t = t + 1
99 }
100 // bad grouping
101 if nx_f32_gqa_attention(Q, K, V, 2, 3, 2, 2, scale, out) == NX_F32GQA_OK { return 30 }
102 return 0
103}