Chapter 3 · Bonus material

Same Attention, Different Speed

Stacked heads, one big weight split, einsum, PyTorch's fused kernel — nine implementations of multi-head attention return the same-shaped tensor, but disagree on the clock.

Adapted from the original notebook by Sebastian Raschka · the Python shown is the book's; the demos run a faithful JavaScript port in your browser

Bonus · Efficient attention

every implementation returnstorch.Size([8, 1024, 768])batch 8 · 1024 tokens · 768 dims · 12 heads

same output — what changes is speed

01 / 05

There are many ways to write multi-head attention. All of them return the exact same output — shape torch.Size([8, 1024, 768]). What differs is how fast they get there.

The real code

Both classes come from the main chapter. The stacked wrapper is the literal reading of "multi-head": run several single-head attentions and glue their outputs together.

ch03.ipynb · §3.6.1

class MultiHeadAttentionWrapper(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
        super().__init__()
        self.heads = nn.ModuleList(
            [CausalAttention(d_in, d_out, context_length, dropout, qkv_bias)
             for _ in range(num_heads)]
        )
 
    def forward(self, x):
        # each head has its own W_q/W_k/W_v; run them in a loop, concat the results
        return torch.cat([head(x) for head in self.heads], dim=-1)

The weight-split version does the same math with one big Q/K/V projection, reshaped into heads so every head runs in a single batched matmul:

ch03.ipynb · §3.6.2

def forward(self, x):
    b, num_tokens, d_in = x.shape
 
    keys = self.W_key(x)        # one big projection each: (b, num_tokens, d_out)
    queries = self.W_query(x)
    values = self.W_value(x)
 
    # split d_out into heads: (b, num_tokens, num_heads, head_dim)
    keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
    values = values.view(b, num_tokens, self.num_heads, self.head_dim)
    queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)
 
    # (b, num_heads, num_tokens, head_dim) — heads become a batch dimension
    keys = keys.transpose(1, 2)
    queries = queries.transpose(1, 2)
    values = values.transpose(1, 2)
 
    attn_scores = queries @ keys.transpose(2, 3)  # one batched matmul, all heads
    # ... causal mask + softmax (elided) ...
    attn_weights = torch.softmax(attn_scores / keys.shape[-1] ** 0.5, dim=-1)
    context_vec = (attn_weights @ values).transpose(1, 2)
    context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
    return self.out_proj(context_vec)      # combine heads

The full MultiHeadAttention class

ch03.ipynb · §3.6.2 (full)

class MultiHeadAttention(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
        super().__init__()
        assert d_out % num_heads == 0, "d_out must be divisible by num_heads"
 
        self.d_out = d_out
        self.num_heads = num_heads
        self.head_dim = d_out // num_heads    # match desired output dim
 
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.out_proj = nn.Linear(d_out, d_out)   # combine head outputs
        self.dropout = nn.Dropout(dropout)
        self.register_buffer(
            "mask",
            torch.triu(torch.ones(context_length, context_length), diagonal=1),
        )
 
    def forward(self, x):
        b, num_tokens, d_in = x.shape
 
        keys = self.W_key(x)          # (b, num_tokens, d_out)
        queries = self.W_query(x)
        values = self.W_value(x)
 
        # unroll last dim into heads: (b, num_tokens, num_heads, head_dim)
        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
        values = values.view(b, num_tokens, self.num_heads, self.head_dim)
        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)
 
        # (b, num_heads, num_tokens, head_dim)
        keys = keys.transpose(1, 2)
        queries = queries.transpose(1, 2)
        values = values.transpose(1, 2)
 
        attn_scores = queries @ keys.transpose(2, 3)   # dot product per head
        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]
        attn_scores.masked_fill_(mask_bool, -torch.inf)
 
        attn_weights = torch.softmax(attn_scores / keys.shape[-1] ** 0.5, dim=-1)
        attn_weights = self.dropout(attn_weights)
 
        context_vec = (attn_weights @ values).transpose(1, 2)   # (b, num_tokens, num_heads, head_dim)
        context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
        context_vec = self.out_proj(context_vec)    # optional projection
        return context_vec

Three things to remember

The output shape is fixed; the implementation is a choice. Every correct multi-head attention returns a tensor of the same shape — with its own weights, the actual numbers differ, but the differences that matter here are entirely about how the work is scheduled onto hardware.
Fewer, bigger operations win. The weight-split version replaces a loop of small per-head matmuls with one large batched matmul, which is why it beats the stacked wrapper on real hardware.
The fused kernel wins. PyTorch's scaled_dot_product_attention takes the top two spots on these CPU bars — the from-scratch versions are for understanding what it computes. Its FlashAttention kernel is what pulls ahead on the A100 (1.1 ms vs 1.8); on this CPU the default kernel is actually the fastest of the nine.

Adapted from Build a Large Language Model (From Scratch) by Sebastian Raschka — original notebook (Apache 2.0). The two classes shown are its; the benchmark numbers are its published M3 CPU / A100 GPU measurements, shown as attributed data rather than measured live in your browser.