How does Scaled Dot-Product Attention work in PyTorch from scratch?

Understanding Transformer architectures requires getting comfortable with attention calculations. While PyTorch offers high-level layers like nn.MultiheadAttention, implementing Scaled Dot-Product Attention directly helps clarify matrix dimensions and numerical stability tricks like softmax scaling.

Mathematical Equation to Code

The standard attention equation scales the matrix product of Queries ($Q$) and Keys ($K$) by the square root of key dimension ($d_k$) before applying softmax and multiplying by Values ($V$).

import torch
import torch.nn as nn
import math

def scaled_dot_product_attention(query, key, value, mask=None):
    # Obtain key dimensionality for scaling factor calculation
    d_k = query.size(-1)
    
    # Compute Q * K^T / sqrt(d_k) for similarity scores
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    
    # Apply mask if provided (e.g., causal masking in decoder)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    # Softmax normalizes scores across the last dimension into probabilities
    attention_weights = torch.softmax(scores, dim=-1)
    
    # Multiply weights by Values matrix to yield context vector
    output = torch.matmul(attention_weights, value)
    return output, attention_weights

# Dummy batch test
q = torch.randn(2, 8, 64)  # [batch_size, sequence_length, d_k]
k = torch.randn(2, 8, 64)
v = torch.randn(2, 8, 64)
out, weights = scaled_dot_product_attention(q, k, v)
print("Output shape:", out.shape)

My Core Questions:

• Why is the scaling factor $\sqrt{d_k}$ necessary to prevent vanishing gradients during backward passes in deep networks?
• How does PyTorch handle fused FlashAttention under the hood to bypass raw tensor allocations for large sequences?