Implement Multi-Head Attention

Math Foundations

Single head attention projects queries, keys, and values into a single subspace.

Multi-Head Attention (Vaswani et al., 2017) projects queries, keys, and values $H$ times with different learnable linear projections to dimension $d_k = d_{\text{model}} / H$:

$$\text{head}_i = \text{Attention}\left( Q W_i^Q, K W_i^K, V W_i^V \right)$$

$$\text{MultiHead}(Q, K, V) = \text{Concat}\left(\text{head}_1, \dots, \text{head}_H\right) W^O$$

Input X ──► Linear Projections ──► Split into H Heads ──► Parallel Scaled Dot-Product ──► Concat Heads ──► Output Linear W_O

Efficient Tensor Reshaping Strategy

To execute Multi-Head Attention efficiently on GPUs without slow Python for loops across heads:

Project inputs to full model dimension $d_{\text{model}}$, then transpose dimensions to move num_heads into the batch dimension:

Linear Projection: [B, T, d_model]
Reshape:           [B, T, H, d_h]    (where d_h = d_model / H)
Transpose:         [B, H, T, d_h]    <── Parallel GPU Matrix Multiplication!

PyTorch Implementation from Scratch

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


class MultiHeadAttentionFromScratch(nn.Module):

  def __init__(self, d_model, num_heads):
    super().__init__()
    assert (
        d_model % num_heads == 0
    ), "d_model must be evenly divisible by num_heads"

    self.d_model = d_model
    self.num_heads = num_heads
    self.d_h = d_model // num_heads

    # Linear projections for Query, Key, Value, and Output
    self.w_q = nn.Linear(d_model, d_model, bias=False)
    self.w_k = nn.Linear(d_model, d_model, bias=False)
    self.w_v = nn.Linear(d_model, d_model, bias=False)
    self.w_o = nn.Linear(d_model, d_model, bias=False)

  def forward(self, x, mask=None):
    B, T, _ = x.shape

    # 1. Linear Projections and Reshape to [B, H, T, d_h]
    q = self.w_q(x).view(B, T, self.num_heads, self.d_h).transpose(1, 2)
    k = self.w_k(x).view(B, T, self.num_heads, self.d_h).transpose(1, 2)
    v = self.w_v(x).view(B, T, self.num_heads, self.d_h).transpose(1, 2)

    # 2. Scaled Dot-Product Attention across all heads in parallel
    # Q @ K^T shape: [B, H, T, T]
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_h)

    if mask is not None:
      scores = scores.masked_fill(mask == 0, float("-inf"))

    attn_weights = F.softmax(scores, dim=-1)

    # 3. Multiply Values and Reshape Back: [B, H, T, d_h] -> [B, T, d_model]
    context = torch.matmul(attn_weights, v)  # [B, H, T, d_h]
    context = context.transpose(1, 2).contiguous().view(B, T, self.d_model)

    # 4. Final Output Linear Projection
    output = self.w_o(context)
    return output

Say this out loud

Multi Head Attention projects Query, Key, and Value tensors into multiple parallel attention heads. Transposing dimensions to batch shape allows GPUs to execute scaled dot product attention across all heads in parallel. Concatenating head outputs and passing through a final linear projection combines diverse representation subspaces.

Followups to expect

1. What is FlashAttention? An exact attention algorithm that reorganizes memory accesses to compute self-attention in GPU SRAM memory blocks, cutting memory usage from $O(T^2)$ to $O(T)$ and speeding up processing.
2. What is Grouped Query Attention (GQA)? An optimization sharing key and value heads across multiple query heads, reducing KV cache memory footprint during LLM inference.