S108 · Q/K/V Math — Scaled Dot-Product Attention Derived

# qkv_math.py — scaled dot-product attention, from scratch AND with the PyTorch primitive.
# Also verifies the sqrt(d_k) scaling by showing what happens without it.
import math
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
 
 
def attention_scratch(
    Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
    mask: torch.Tensor | None = None,
    scale: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Scaled dot-product attention — every step explicit."""
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1)                # [B, T, T]
    if scale:
        scores = scores / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask, float("-inf"))
    weights = F.softmax(scores, dim=-1)             # [B, T, T]
    output = weights @ V                             # [B, T, d_v]
    return output, weights
 
 
def attention_builtin(
    Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
    is_causal: bool = False,
) -> torch.Tensor:
    """The optimised PyTorch primitive — uses FlashAttention when available."""
    return F.scaled_dot_product_attention(Q, K, V, is_causal=is_causal)
 
 
def verify_equivalence() -> None:
    torch.manual_seed(0)
    B, T, D = 2, 8, 16
    X = torch.randn(B, T, D)
    W_q = torch.randn(D, D) * 0.1
    W_k = torch.randn(D, D) * 0.1
    W_v = torch.randn(D, D) * 0.1
    Q, K, V = X @ W_q, X @ W_k, X @ W_v
 
    out_scratch, weights = attention_scratch(Q, K, V)
    out_builtin = attention_builtin(Q, K, V)
 
    diff = (out_scratch - out_builtin).abs().max().item()
    print(f"Max diff (scratch vs builtin): {diff:.2e}")
    print(f"  → should be ~1e-6 (float precision noise)")
    print(f"\nAttention weights row 0: {weights[0, 0].tolist()}")
    print(f"  → sums to: {weights[0, 0].sum().item():.4f} (should be 1.0)")
 
 
def demonstrate_scaling_matters() -> None:
    """Show the softmax saturation without sqrt(d_k) scaling."""
    torch.manual_seed(0)
    B, T = 1, 8
    for D in [16, 64, 256, 1024]:
        X = torch.randn(B, T, D)
        Q = K = V = X
        _, w_scaled = attention_scratch(Q, K, V, scale=True)
        _, w_unscaled = attention_scratch(Q, K, V, scale=False)
        max_scaled = w_scaled[0, 0].max().item()
        max_unscaled = w_unscaled[0, 0].max().item()
        print(f"d_k={D:4d}  max_attn (scaled)={max_scaled:.3f}  (unscaled)={max_unscaled:.3f}")
 
 
def benchmark() -> None:
    torch.manual_seed(0)
    B, T, D = 8, 512, 128
    Q = torch.randn(B, T, D)
    K = torch.randn(B, T, D)
    V = torch.randn(B, T, D)
 
    # Warm up
    _ = attention_scratch(Q, K, V)
    _ = attention_builtin(Q, K, V)
 
    t0 = time.time()
    for _ in range(20):
        _ = attention_scratch(Q, K, V)
    t_scratch = time.time() - t0
 
    t0 = time.time()
    for _ in range(20):
        _ = attention_builtin(Q, K, V)
    t_builtin = time.time() - t0
 
    print(f"\nBenchmark (B={B}, T={T}, D={D}, 20 iters):")
    print(f"  scratch: {t_scratch*1000:.1f} ms")
    print(f"  builtin: {t_builtin*1000:.1f} ms  ({t_scratch/t_builtin:.2f}× faster)")
 
 
def visualise_attention() -> None:
    """Print an attention matrix that reveals the structure."""
    torch.manual_seed(42)
    T, D = 6, 8
    # Construct Q, K such that token i has high similarity to token i-1 (like a shift op).
    X = torch.randn(T, D)
    Q = X.clone()
    K = torch.roll(X, shifts=1, dims=0)  # K_i = X_{i-1}, so Q_i · K_i is high for i>0
    V = X.clone()
 
    _, weights = attention_scratch(Q.unsqueeze(0), K.unsqueeze(0), V.unsqueeze(0))
    w = weights[0].numpy()
    print("\nSynthetic 'shift by 1' attention pattern:")
    print("     " + "  ".join(f"k{i:d}" for i in range(T)))
    for i in range(T):
        row = "  ".join(f"{v:.2f}" for v in w[i])
        print(f"q{i:d}  {row}")
    print("(Notice heavy weight on the diagonal-shifted-by-1 position.)")
 
 
if __name__ == "__main__":
    verify_equivalence()
    print("\n--- Softmax saturation without scaling ---")
    demonstrate_scaling_matters()
    benchmark()
    visualise_attention()