Comparatively, the original transformer attention mechanism is as such (from https://nlp.seas.harvard.edu/2018/04/03/attention.html 1):
def attention(query, key, value, mask=None, dropout=None):
    """Compute 'Scaled Dot Product Attention"""
    d_k = query.size(-1) # no. of hidden dimension of query tensor.
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

Usage:
>>> l, d = 4, 5
>>> q, k, v = torch.rand(l, d), torch.rand(l, d), torch.rand(l, d)
>>> attention(q, k, v)
(tensor([[0.5511, 0.6426, 0.6486, 0.4216, 0.3083],
         [0.5560, 0.6441, 0.6665, 0.4169, 0.3312],
         [0.5590, 0.6426, 0.6647, 0.4173, 0.3283],
         [0.5678, 0.6450, 0.6291, 0.4012, 0.2998]]),
 tensor([[0.2579, 0.2811, 0.2288, 0.2322],
         [0.2724, 0.2851, 0.2585, 0.1840],
         [0.2993, 0.2884, 0.2382, 0.1740],
         [0.2769, 0.2465, 0.2166, 0.2600]]))