Splitting into multiple heads -- multihead self attention
In principle, the pseudocode is correct, but it is not how it is implemented. The projection and dot-product attention can be done efficiently using matrix multiplication only for all heads simultaneously.
Instead of looping over the heads, you can use only one dense layer for queries, one for keys and one for values. E.g., for the keys, you would do a dense layer of dimension $h \cdot d_k$ and then reshape it.
Assuming you have batch size $b$, sequence length $l$ and model dimension $d_m$:

The input of the dense layers is of shape $b \times l \times d_m$.

The output of the dense layer is has shape $b \times l \times hd_k$ (or $d_v$ respectively).

Then you can reshape the queries and keys to have shape $b \times l \times h \times d_k$.


Now, if you permute the dimensions, such that the queries have shape $b \times h \times l \times d_k$ and the keys have shape $b \times h \times d_k \times l$, you can do the batch matrix multiplication in the last two dimensions and you end with attention energies of shape $b \times h \times l \times l$. Then, if you do softmax in the last dimension, you get the attention distribution for each head and each query.
By doing the batch matrix multiplication with projected values that now with transposed dimensions to $b \times h \times l \times d_v$, you get the weighted average. So finally, the &quot;concatenation&quot; of all heads is in fact just another tensor reshaping.