Add the wrapper for multi-head scaled dot product attention.