nx_f32_mha.nx source

1// nx_f32_mha.nx -- bits-up f32 multi-head attention with GQA. 2// 3// L7 / L8 composition brick. Closes the single-head v1 limitation 4// by decomposing hidden_dim into n_heads * head_dim and computing 5// per-head attention in parallel, then concatenating + projecting 6// through W_o. Supports Grouped-Query Attention (Llama-3, Qwen2, 7// Mistral) via separate n_kv_heads. 8// 9// Single-token v1: for autoregressive decode the typical 10// inference pattern is one new token at a time, so n_tokens = 1. 11// Multi-token (prefill) is a follow-on brick that adds K^T 12// transpose + multi-row softmax. 13// 14// Algorithm (Vaswani 2017 multi-head + Ainslie 2023 GQA): 15// 16// q_dim = n_heads * head_dim (= hidden_dim typically) 17// kv_dim = n_kv_heads * head_dim 18// group_size = n_heads / n_kv_heads 19// 20// Q = matmul(x, W_q) [1, q_dim] 21// K = matmul(x, W_k) [1, kv_dim] 22// V = matmul(x, W_v) [1, kv_dim] 23// 24// For each query head h in 0..n_heads: 25// kv_head = h / group_size 26// Q_h = Q[h * head_dim : (h+1) * head_dim] 27// K_h = K[kv_head * head_dim : (kv_head+1) * head_dim] 28// V_h = V[kv_head * head_dim : (kv_head+1) * head_dim] 29// score_h = dot(Q_h, K_h) * attn_scale 30// prob_h = softmax([score_h]) (degenerate -> [1.0]) 31// attn_out[h*head_dim : (h+1)*head_dim] = prob_h * V_h 32// 33// out = matmul(attn_out, W_o) [1, hidden_dim] 34// 35// Caller does the residual: x = x + out. 36// 37// genealogy_id: vaswani_2017_multihead + ainslie_2023_gqa 38// lineage_id: substrate_f32_mha_single_token_v1 39 40import "nx_syscalls.nx" 41import "nx_tier.nx" 42import "nx_f32.nx" 43import "nx_f32_matmul.nx" 44import "nx_f32_softmax.nx" 45 46const NX_F32_MHA_OK: nx_int = 0 47const NX_F32_MHA_ERR_BAD_DIM: nx_int = 1 48const NX_F32_MHA_ERR_BAD_GQA: nx_int = 2 49const NX_F32_MHA_N_VERDICTS: nx_int = 3 50 51func nx_f32_mha_verdict_is_valid(v: nx_int) -> nx_int { 52 if v < 0 { return 0 } 53 if v >= NX_F32_MHA_N_VERDICTS { return 0 } 54 return 1 55} 56 57// Single-token MHA forward. Writes the attention sub-block output 58// to `out` (not yet added to x; caller does the residual). 59// 60// Constraints: 61// hidden_dim = n_heads * head_dim (must hold) 62// n_heads % n_kv_heads == 0 (GQA group size is integer) 63// 64// 11 args -- under the 16-arg limit. 65 66func nx_f32_mha_single_token( 67 x: *i64, 68 hidden_dim: nx_int, 69 n_heads: nx_int, 70 n_kv_heads: nx_int, 71 head_dim: nx_int, 72 W_q: *i64, 73 W_k: *i64, 74 W_v: *i64, 75 W_o: *i64, 76 attn_scale: i64, 77 out: *i64) -> nx_int { 78 79 if hidden_dim <= 0 { return NX_F32_MHA_ERR_BAD_DIM } 80 if n_heads <= 0 { return NX_F32_MHA_ERR_BAD_DIM } 81 if n_kv_heads <= 0 { return NX_F32_MHA_ERR_BAD_DIM } 82 if head_dim <= 0 { return NX_F32_MHA_ERR_BAD_DIM } 83 if n_heads * head_dim != hidden_dim { 84 return NX_F32_MHA_ERR_BAD_DIM 85 } 86 if n_heads - (n_heads / n_kv_heads) * n_kv_heads != 0 { 87 return NX_F32_MHA_ERR_BAD_GQA 88 } 89 90 let q_dim: nx_int = n_heads * head_dim 91 let kv_dim: nx_int = n_kv_heads * head_dim 92 let group_size: nx_int = n_heads / n_kv_heads 93 94 // Project x into Q, K, V. 95 let Q: *i64 = sys_mmap(q_dim * 8) as *i64 96 let K: *i64 = sys_mmap(kv_dim * 8) as *i64 97 let V: *i64 = sys_mmap(kv_dim * 8) as *i64 98 nx_f32_matmul(x, W_q, Q, 1, hidden_dim, q_dim) 99 nx_f32_matmul(x, W_k, K, 1, hidden_dim, kv_dim) 100 nx_f32_matmul(x, W_v, V, 1, hidden_dim, kv_dim) 101 102 // Per-head attention. 103 let attn_out: *i64 = sys_mmap(q_dim * 8) as *i64 104 105 var h: nx_int = 0 106 while h < n_heads { 107 let kv_head: nx_int = h / group_size 108 109 // Compute slice base addresses (raw pointer arithmetic). 110 let q_base: i64 = (Q as i64) + h * head_dim * 8 111 let k_base: i64 = (K as i64) + kv_head * head_dim * 8 112 let v_base: i64 = (V as i64) + kv_head * head_dim * 8 113 let Q_h: *i64 = q_base as *i64 114 let K_h: *i64 = k_base as *i64 115 let V_h: *i64 = v_base as *i64 116 117 // score = dot(Q_h, K_h) * attn_scale 118 let raw_score: i64 = nx_f32_dot(Q_h, K_h, head_dim) 119 let score: i64 = nx_f32_mul(raw_score, attn_scale) 120 121 // softmax([score]) = [1.0] (single-element vector) 122 // We run the path for substrate honesty (composes nx_f32_softmax). 123 let scores_slot: *i64 = sys_mmap(8) as *i64 124 let probs_slot: *i64 = sys_mmap(8) as *i64 125 scores_slot[0] = score 126 nx_f32_softmax(scores_slot, 1, probs_slot) 127 let prob_h: i64 = probs_slot[0] 128 129 // attn_out[h*head_dim .. ] = prob_h * V_h 130 let attn_h_base: i64 = (attn_out as i64) + h * head_dim * 8 131 let attn_h: *i64 = attn_h_base as *i64 132 var d: nx_int = 0 133 while d < head_dim { 134 attn_h[d] = nx_f32_mul(prob_h, V_h[d]) 135 d = d + 1 136 } 137 138 h = h + 1 139 } 140 141 // Output projection: [1, q_dim] @ [q_dim, hidden_dim] -> [1, hidden_dim] 142 nx_f32_matmul(attn_out, W_o, out, 1, q_dim, hidden_dim) 143 return NX_F32_MHA_OK 144}