Scaled Dot-Product Attention

The operation the last eight years of AI are built on. Write it from scratch.

Implement attention(q, k, v, mask=None):

attention(Q, K, V) = softmax(Q Kᵀ / √d_k) V

• q, k, v have shape (seq_len, d_k)
• mask, when given, is a boolean array of shape (seq_len, seq_len) where True means allowed and False means the position must not be attended to
• Return the output of shape (seq_len, d_k)

Three things this problem is actually testing:

1. The √d_k. Without it, dot products grow with dimension, softmax saturates, and gradients vanish.
2. The mask. Masked positions must contribute exactly zero attention weight, not "a small amount".
3. Stability. Your softmax must not overflow.