Attention is a weighted average where the weights are learned. Everything else is bookkeeping.
The obvious implementation
Write the loop version first. It is unusably slow and completely clear, which is the correct order to learn this in.
The batched implementation
Then collapse it into matrix multiplications and confirm the outputs match to within floating-point tolerance.
Causal masking
The one detail that silently ruins training if you get it wrong.