PyTorch Implementation of Attention

Overview

The theory of 02 Scaled Dot Product Attention fits in about twenty lines of PyTorch. This note walks through three progressively more complete implementations: a bare-bones single-head attention, a full multi-head layer, and the modern fused F.scaled_dot_product_attention that production models actually use.

• Goal: build intuition by reading + running real code
• Pattern: naive math version → correct multi-head version → production fused version
• Prerequisite: 01 Query, Key & Value, 02 Scaled Dot Product Attention, 04 Multi Head Attention

---

1. Bare-Bones Scaled Dot-Product Attention

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None, causal=False):
    """
    Q, K, V: (batch, seq_len, head_dim)
    Returns: (batch, seq_len, head_dim)
    """
    d_k = Q.size(-1)

    # Step 1: raw scores (batch, seq_len, seq_len)
    scores = Q @ K.transpose(-2, -1)

    # Step 2: scale to prevent softmax saturation
    scores = scores / (d_k ** 0.5)

    # Causal mask: token i attends only to j <= i
    if causal:
        seq_len = Q.size(1)
        mask = torch.triu(
            torch.full((seq_len, seq_len), float("-inf")),
            diagonal=1,
        ).to(Q.device)

    # Step 3: normalize to weights
    weights = F.softmax(scores + mask, dim=-1) if mask is not None \
              else F.softmax(scores, dim=-1)

    # Step 4: weighted sum of values
    output = weights @ V
    return output

# Toy usage — same numbers as the worked example in 02
seq, d = 3, 2
Q = torch.tensor([[[1., 0.], [0., 1.], [1., 1.]]])
K = torch.tensor([[[1., 0.], [0., 1.], [1., 0.]]])
V = torch.tensor([[[1., 2.], [3., 4.], [5., 6.]]])

out = scaled_dot_product_attention(Q, K, V, causal=True)
print(out)
# tensor([[[1.00, 2.00],      <- "The" sees only itself
#          [2.34, 3.34],      <- "cat" blends "The" + itself
#          [3.00, 4.00]]])    <- "sat" blends everything

Verify against the hand-computed example

Example A (non-causal) from [02 Scaled Dot Product Attention](/02-llm-modeling/00-fundamentals/01-attention/02-scaled-dot-product-attention/):
  inputs: Q=[[1,0],[0,1],[1,1]], K=[[1,0],[0,1],[1,0]], V=[[1,2],[3,4],[5,6]]
  expected outputs: all rows ≈ [3.00, 4.00]

Run this file's function with causal=False → same result ✅
(Softmax rounding differences of ±0.01 are expected.)

---

2. Complete Multi-Head Attention Module

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    """Scaled dot-product multi-head attention.

    Shapes:
      x:      (batch, seq_len, d_model)
      output: (batch, seq_len, d_model)
    """

    def __init__(self, d_model=512, n_heads=8, dropout=0.0):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"

        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads          # head dimension (64 for 512/8)

        # Single combined projection for Q, K, V (common optimization)
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.out_proj = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        batch, seq_len, _ = x.shape

        # 1. Project and split into heads: (b, seq, 3*d_model)
        qkv = self.qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)

        # 2. Reshape to (b, n_heads, seq, d_k)
        q = q.view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        k = k.view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        v = v.view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # 3. Attention scores: (b, n_heads, seq, seq)
        scores = q @ k.transpose(-2, -1) / (self.d_k ** 0.5)

        # 4. Apply mask (e.g., causal) if given
        if mask is not None:
            scores = scores + mask

        # 5. Softmax + dropout (training regularization)
        weights = F.softmax(scores, dim=-1)
        weights = self.dropout(weights)

        # 6. Weighted sum of values: (b, n_heads, seq, d_k)
        context = weights @ v

        # 7. Concatenate heads: (b, seq, n_heads * d_k) -> (b, seq, d_model)
        context = context.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model)

        # 8. Final projection (mixes information across heads)
        return self.out_proj(context)


# Causal mask helper (matches the mask in [03 Types Of Attention](/02-llm-modeling/00-fundamentals/01-attention/03-types-of-attention/))
def causal_mask(seq_len):
    """(seq_len, seq_len) mask: 0 allowed, -inf masked."""
    return torch.triu(torch.full((seq_len, seq_len), float("-inf")), diagonal=1)


# Smoke test
torch.manual_seed(0)
x = torch.randn(2, 10, 512)            # (batch=2, seq=10, d_model=512)
mha = MultiHeadAttention(d_model=512, n_heads=8)
out = mha(x, mask=causal_mask(10))
assert out.shape == (2, 10, 512), f"bad shape: {out.shape}"
print("Multi-head attention output shape:", out.shape)

What changed vs. version 1

Version 1 (bare-bones):      single head, explicit steps, hand-shaped Q/K/V
Version 2 (module):
  - learns W_Q/W_K/W_V as one Linear (3·d_model)  → chunk into 3
  - reshapes into (batch, heads, seq, head_dim)   → parallel heads
  - adds dropout on attention weights
  - output projection W_O mixes head outputs
  - reusable as a PyTorch module with mask support

---

3. Production: F.scaled_dot_product_attention

PyTorch ≥ 2.0 ships a fused, flash-attention-backed implementation. It is what you should use in real code:

import torch
import torch.nn.functional as F

def fused_attention(q, k, v, is_causal=True, attn_mask=None):
    """
    q, k, v: (batch, n_heads, seq_len, head_dim)
    Returns: (batch, n_heads, seq_len, head_dim)
    """
    return F.scaled_dot_product_attention(
        q, k, v,
        attn_mask=attn_mask,     # optional custom mask (floats: 0/-inf)
        is_causal=is_causal,     # built-in causal optimization
        dropout_p=0.0,
        enable_gqa=True,         # PyTorch 2.4+: Grouped Query Attention
    )

# Usage inside a real model
batch, heads, seq, d_k = 2, 32, 4096, 128
q = torch.randn(batch, heads, seq, d_k)
k = torch.randn(batch, heads, seq, d_k)
v = torch.randn(batch, heads, seq, d_k)

out = fused_attention(q, k, v, is_causal=True)
print(out.shape)   # (2, 32, 4096, 128)

# Why this matters:
#  - uses FlashAttention kernels when available (SDPA dispatch)
#  - never materializes the full (seq, seq) matrix in HBM
#  - supports GQA without custom code — see [Gqa](/02-llm-modeling/01-architecture/01-core-designs/multi-query-attention-(mqa)-&-grouped-query-attention-(gqa)|gqa/)

⚠️ Note: versions 1–2 are for learning. Production models use fused kernels (FlashAttention CUDA), KV-cache optimized decode loops, and enable_gqa — never the naive loops. See 06 Complexity & Cost for why.

---

Putting It All Together: A Tiny GPT-Style Block

class SelfAttentionBlock(nn.Module):
    """One decoder layer: pre-norm → causal multi-head attention → MLP."""

    def __init__(self, d_model=512, n_heads=8, d_ff=2048):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.attn = MultiHeadAttention(d_model, n_heads)
        self.norm2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
        )

    def forward(self, x):
        # Pre-LN + residual connection (modern decoder style)
        x = x + self.attn(self.norm1(x), mask=causal_mask(x.size(1)))
        x = x + self.mlp(self.norm2(x))
        return x


x = torch.randn(2, 16, 512)
block = SelfAttentionBlock(d_model=512, n_heads=8, d_ff=2048)
print("Block output:", block(x).shape)   # (2, 16, 512)

---

Key Takeaways

🐍 20 lines capture the entire attention idea
🔨 Module version adds learning: Linear projections + head split + W_O
⚡ Fused SDPA is what real models use — no manual loops
🎯 Verification: run the toy example and compare against the hand math
🧱 Composable: attention blocks stack into GPT-style decoders
🚀 Efficiency matters: naive code is correct but unusable at 8K+ context

---

Related Notes

• 00 Attention Mechanisms — chapter overview
• 02 Scaled Dot Product Attention — the math this code implements
• 01 Query, Key & Value — the Q/K/V role intuition
• 04 Multi Head Attention — the module structure above
• Flash Attention — the kernel behind fused SDPA
• 06 Complexity & Cost — why we can't just run the naive version