nx_f32_mha_multi.nx source
↩ module page · 140 lines · 4897 B
1// nx_f32_mha_multi.nx -- bits-up f32 multi-head + multi-token attention.
2//
3// L7 / L8 composition brick. Combines the head-decomposition pattern
4// from nx_f32_mha.nx (multi-head + GQA) with the multi-token causal
5// kernel from nx_f32_attn_multi.nx into the full attention forward.
6//
7// This is the actual attention kernel for:
8// - Prefill (n_tokens = prompt length, causal mask)
9// - KV-cache decode (n_tokens = 1 new + cache; future Task #13)
10// - Non-causal encoder attention (causal = 0)
11//
12// Algorithm:
13//
14// Q_all = matmul(x, W_q) [n_tokens, n_heads*head_dim]
15// K_all = matmul(x, W_k) [n_tokens, n_kv_heads*head_dim]
16// V_all = matmul(x, W_v) [n_tokens, n_kv_heads*head_dim]
17//
18// For each query head h in 0..n_heads:
19// kv_head = h / group_size
20// Slice Q_h, K_h, V_h from the column range [h*head_dim,
21// (h+1)*head_dim] for queries (or kv_head*head_dim for K,V).
22// attn_h = nx_f32_attn_multi(Q_h, K_h, V_h, n_tokens, n_tokens,
23// head_dim, causal, attn_scale)
24// Scatter attn_h back into attn_out_all at columns [h*head_dim,
25// (h+1)*head_dim].
26//
27// out = matmul(attn_out_all, W_o) [n_tokens, hidden_dim]
28//
29// Slicing is via copy in v1. v2 will use stride-aware kernels to
30// avoid the copy (perf lift queued).
31//
32// genealogy_id: vaswani_2017_mha + ainslie_2023_gqa + causal_mask_canon
33// lineage_id: substrate_f32_mha_multi_v1
34
35import "nx_syscalls.nx"
36import "nx_tier.nx"
37import "nx_f32.nx"
38import "nx_f32_matmul.nx"
39import "nx_f32_attn_multi.nx"
40
41const NX_F32_MHAM_OK: nx_int = 0
42const NX_F32_MHAM_ERR_BAD_DIM: nx_int = 1
43const NX_F32_MHAM_ERR_BAD_GQA: nx_int = 2
44const NX_F32_MHAM_N_VERDICTS: nx_int = 3
45
46func nx_f32_mham_verdict_is_valid(v: nx_int) -> nx_int {
47 if v < 0 { return 0 }
48 if v >= NX_F32_MHAM_N_VERDICTS { return 0 }
49 return 1
50}
51
52// 13 args -- under the 16-arg limit.
53
54func nx_f32_mha_multi_token(
55 x: *i64,
56 n_tokens: nx_int,
57 hidden_dim: nx_int,
58 n_heads: nx_int,
59 n_kv_heads: nx_int,
60 head_dim: nx_int,
61 W_q: *i64,
62 W_k: *i64,
63 W_v: *i64,
64 W_o: *i64,
65 causal: nx_int,
66 attn_scale: i64,
67 out: *i64) -> nx_int {
68
69 if n_tokens <= 0 { return NX_F32_MHAM_ERR_BAD_DIM }
70 if hidden_dim <= 0 { return NX_F32_MHAM_ERR_BAD_DIM }
71 if n_heads <= 0 { return NX_F32_MHAM_ERR_BAD_DIM }
72 if n_kv_heads <= 0 { return NX_F32_MHAM_ERR_BAD_DIM }
73 if head_dim <= 0 { return NX_F32_MHAM_ERR_BAD_DIM }
74 if n_heads * head_dim != hidden_dim {
75 return NX_F32_MHAM_ERR_BAD_DIM
76 }
77 if n_heads - (n_heads / n_kv_heads) * n_kv_heads != 0 {
78 return NX_F32_MHAM_ERR_BAD_GQA
79 }
80
81 let q_dim: nx_int = n_heads * head_dim
82 let kv_dim: nx_int = n_kv_heads * head_dim
83 let group_size: nx_int = n_heads / n_kv_heads
84
85 // Project all tokens.
86 let Q_all: *i64 = sys_mmap(n_tokens * q_dim * 8) as *i64
87 let K_all: *i64 = sys_mmap(n_tokens * kv_dim * 8) as *i64
88 let V_all: *i64 = sys_mmap(n_tokens * kv_dim * 8) as *i64
89 nx_f32_matmul(x, W_q, Q_all, n_tokens, hidden_dim, q_dim)
90 nx_f32_matmul(x, W_k, K_all, n_tokens, hidden_dim, kv_dim)
91 nx_f32_matmul(x, W_v, V_all, n_tokens, hidden_dim, kv_dim)
92
93 let attn_out_all: *i64 = sys_mmap(n_tokens * q_dim * 8) as *i64
94
95 // Per-head attention.
96 let Q_h: *i64 = sys_mmap(n_tokens * head_dim * 8) as *i64
97 let K_h: *i64 = sys_mmap(n_tokens * head_dim * 8) as *i64
98 let V_h: *i64 = sys_mmap(n_tokens * head_dim * 8) as *i64
99 let attn_h: *i64 = sys_mmap(n_tokens * head_dim * 8) as *i64
100
101 var h: nx_int = 0
102 while h < n_heads {
103 let kv_head: nx_int = h / group_size
104
105 // Gather Q_h, K_h, V_h slices across all tokens.
106 var t: nx_int = 0
107 while t < n_tokens {
108 var d: nx_int = 0
109 while d < head_dim {
110 Q_h[t * head_dim + d] = Q_all[t * q_dim + h * head_dim + d]
111 K_h[t * head_dim + d] = K_all[t * kv_dim + kv_head * head_dim + d]
112 V_h[t * head_dim + d] = V_all[t * kv_dim + kv_head * head_dim + d]
113 d = d + 1
114 }
115 t = t + 1
116 }
117
118 // Multi-token attention for this head.
119 nx_f32_attn_multi(Q_h, K_h, V_h, n_tokens, n_tokens, head_dim,
120 causal, attn_scale, attn_h)
121
122 // Scatter attn_h back into attn_out_all.
123 var t2: nx_int = 0
124 while t2 < n_tokens {
125 var d2: nx_int = 0
126 while d2 < head_dim {
127 attn_out_all[t2 * q_dim + h * head_dim + d2] =
128 attn_h[t2 * head_dim + d2]
129 d2 = d2 + 1
130 }
131 t2 = t2 + 1
132 }
133
134 h = h + 1
135 }
136
137 // Output projection.
138 nx_f32_matmul(attn_out_all, W_o, out, n_tokens, q_dim, hidden_dim)
139 return NX_F32_MHAM_OK
140}