Compute Scaled Dot-Product Attention #182

Implement the core attention calculation. Take the Q and K matrices (post-RoPE) and compute the scaled dot-product for each head: for each head, calculate the matrix of dot products $A = Q \cdot K^T$ (of shape [seq_len, seq_len]). Apply the scale factor $1/\sqrt{d_{\text{head}}}$ to the dot products to stabilize gradients (in inference this is just a constant factor). Because Q and K are quantized 8-bit values in practice, perform this computation in a higher precision (e.g., int32 or float) to preserve accuracy. Next, apply a softmax to each row of $A$ to obtain attention weights (probabilities) for each token’s attention over all other tokens. (In implementation, you might convert the dot products to float32 for the softmax step for simplicity.) Ensure the softmax is numerically stable (you may need to subtract the max value per row before exponentiating). This will produce an attention weight matrix for each head of size [seq_len, seq_len].