Implement scaled dot-product self-attention

import numpy as np

def attention(Q, K, V, causal=False):
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)
    if causal:
        mask = np.triu(np.ones(scores.shape, dtype=bool), k=1)
        scores = np.where(mask, -np.inf, scores)
    scores -= scores.max(axis=-1, keepdims=True)  # numerical stability
    weights = np.exp(scores)
    weights /= weights.sum(axis=-1, keepdims=True)
    return weights @ V

Say out loud why each line exists: the √d_k scaling keeps dot-product variance ~1 so softmax doesn't saturate; the max-subtraction prevents overflow; masked positions get −inf so they contribute exactly 0 probability after softmax.

Follow-ups: Extend to batched multi-head (see the next question). Time and memory complexity? O(n²·d) time, O(n²) memory for the score matrix — this is the lead-in to the FlashAttention question below. What's the difference between padding masks and causal masks, and how do they combine? Why multiple heads instead of one big head? (Multiple learned similarity subspaces.)