nx_f32_mha.nx

nx_f32_mha.nx -- bits-up f32 multi-head attention with GQA. L7 / L8 composition brick. Closes the single-head v1 limitation by decomposing hidden_dim into n_heads * head_dim and computing per-head attention in parallel, then concatenating + projecting through W_o. Supports Grouped-Query Attention (Llama-3, Qwen2, Mistral) via separate n_kv_heads. Single-token v1: for autoregressive decode the typical inference pattern is one new token at a time, so n_tokens = 1. Multi-token (prefill) is a follow-on brick that adds K^T transpose + multi-row softmax. Algorithm (Vaswani 2017 multi-head + Ainslie 2023 GQA): q_dim = n_heads * head_dim (= hidden_dim typically) kv_dim = n_kv_heads * head_dim group_size = n_heads / n_kv_heads Q = matmul(x, W_q) [1, q_dim] K = matmul(x, W_k) [1, kv_dim] V = matmul(x, W_v) [1, kv_dim] For each query head h in 0..n_heads: kv_head = h / group_size Q_h = Q[h * head_dim : (h+1) * head_dim] K_h = K[kv_head * head_dim : (kv_head+1) * head_dim] V_h = V[kv_head * head_dim : (kv_head+1) * head_dim] score_h = dot(Q_h, K_h) * attn_scale prob_h = softmax([score_h]) (degenerate -> [1.0]) attn_out[h*head_dim : (h+1)*head_dim] = prob_h * V_h out = matmul(attn_out, W_o) [1, hidden_dim] Caller does the residual: x = x + out. genealogy_id: vaswani_2017_multihead + ainslie_2023_gqa lineage_id: substrate_f32_mha_single_token_v1