Multi-Head Self-Attention Forward Pass

Implement a full multi-head self-attention layer for one sequence X of shape (T, d):

1. project: Q = X @ Wq, K = X @ Wk, V = X @ Wv (each d × d)
2. split the last dimension into h heads of size d_h = d / h
3. per head: $\text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_h}} + M\right)V$, where M is 0 or $-\infty$ from an optional causal mask
4. concatenate heads and apply Wo

Implement mha(X, Wq, Wk, Wv, Wo, h, causal=False) returning the (T, d) output.