210. Backprop: Attention head

Hand-derive the gradient of L = sum(Attention(Q, K, V)) w.r.t. Q (queries), with K and V held fixed.

Forward (single head, no mask):

• Scores: S = Q @ K.T / sqrt(d_k)
• Probs: P = softmax(S) along the last axis
• Output: y = P @ V

Q, K have shape (N_q, d_k) and (N_k, d_k); V has shape (N_k, d_v); y has shape (N_q, d_v).

Implement:

• attention_forward(Q, K, V) -> y
• attention_backward(Q, K, V) -> dL/dQ of shape (N_q, d_k)

Chain (with L = sum(y) so dL/dy = ones(N_q, d_v)):

1. dL/dP = ones @ V.T — each row is the column sums of V.
2. Per-row softmax backward: dL/dS[i, :] = P[i, :] * (dL/dP[i, :] - sum_j P[i, j] * dL/dP[i, j]).
3. dL/dQ = (dL/dS) @ K / sqrt(d_k).