Scaled Dot-Product Attention: Q · K^T → softmax → · V

import torch, math

def scaled_dot_product_attention(Q, K, V, mask=None):
    # Q, K, V: (..., seq_len, d_k or d_v)
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)   # (..., n, n)
    if mask is not None:
        scores = scores.masked_fill(mask, float('-inf'))
    weights = torch.softmax(scores, dim=-1)              # row-wise softmax
    return weights @ V                                    # (..., n, d_v)

# Equivalent to torch.nn.functional.scaled_dot_product_attention,
# minus the optimized kernels.

import torch.nn.functional as F

Q = torch.randn(2, 4, 16, 64)
K = torch.randn(2, 4, 16, 64)
V = torch.randn(2, 4, 16, 64)

ref = F.scaled_dot_product_attention(Q, K, V)
mine = scaled_dot_product_attention(Q, K, V)

print(torch.allclose(ref, mine, atol=1e-5))   # True