q, k, v = x @ Wq, x @ Wk, x @ Wv             # [B, T, d] each
scores = q @ k.transpose(-1, -2) / sqrt(d)  # [B, T, T]
weights = scores.softmax(-1)                # every row sums to 1
out = weights @ v                           # [B, T, d]