How do you write multi-head attention in PyTorch from scratch?

In this article

1. Before you start
2. Why interviewers ask this
3. Multi-head attention from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page is one step of the PyTorch and machine learning (ML) coding round guide. If you want the whole model around this layer, the build generative pre-trained transformer (GPT) from scratch walkthrough uses it inside a full network.

Why interviewers ask this

Writing this layer on a blank editor is a common test in machine-learning coding rounds, because one short function exposes whether you understand tensor shapes, masking and the cost of attention. Candidates report variants at AI labs and large tech companies: add a causal mask, support cross-attention, or extend to a key-value cache with grouped-query attention in the same hour.

A weak answer calls nn.MultiheadAttention or pastes the formula without shapes. A strong answer writes the code, states the shape after every line and explains why the scores are divided by a square root. A Staff-level answer also knows what the mask convention is, what happens to a fully blocked row, and why sharing key and value heads shrinks the memory used in generation. Ask whether F.scaled_dot_product_attention is allowed; the PyTorch documentation lists it as a built-in, and many interviewers want the explicit version instead.

Multi-head attention from first principles

Head dimension

The head dimension is the number of features each attention head works on, usually equal to the model dimension divided by the number of heads, so the heads together use the model's width (some models decouple the two).

Padding mask

A padding mask is a boolean tensor that marks the filler positions added to shorter sequences in a batch, so attention blocks those positions and real tokens never take information from padding.

In plain words

Picture a group of people who each read the same report, but each looks for one thing: one tracks names, one tracks dates, one tracks causes. Afterwards they pool their notes into one summary. Each reader sees only part of the detail but together they cover more than one reader skimming everything. The analogy breaks in one place: nobody assigns the topics. The model learns what each head looks for during training.

In one precise sentence, scaled dot-product attention is the function that scores each query against every key, scales the scores, turns them into weights with a softmax and returns the weighted sum of the values, and multi-head attention runs it on several projected slices and merges the results. Vaswani et al. (2017) introduced both.

Here Q, K and V are the query, key and value matrices and $d_k$ is the head dimension.

Step 1: Project the input into queries, keys and values

Three linear layers turn the input x of shape (B, T, D) (batch, tokens, model dimension) into queries, keys and values of the same shape. Each token's vector is multiplied by three different learned matrices, so the same token can ask one thing as a query and advertise another as a key.

Input: x of shape (B, T, D). Output: three tensors of shape (B, T, D). A fourth linear layer, the output projection, is kept for the last step. In the base model of Vaswani et al. (2017) the model dimension is 512, so each of the four matrices holds 262,144 weights, and four together hold 1,048,576 without biases.¹ That number does not change with the number of heads. You now have full-width queries, keys and values; the next step cuts them into heads.

Step 2: Split the heads with a reshape and a transpose

The projections are split into heads, so each head works on its own slice of the features. Splitting the heads takes two operations: reshape (B, T, D) into (B, T, h, d_k), where h is the number of heads and d_k the head dimension, then swap the token and head axes to get (B, h, T, d_k). In PyTorch that is x.view(B, T, h, d_k).transpose(1, 2).

The order matters. A direct reshape of (B, T, D) to (B, h, T, d_k) is always legal once D is divisible by h, so PyTorch raises no error, but it cuts the flattened data into chunks that mix the features of neighbouring tokens. That is the silent bug from the lead. The check that catches it is a comparison: write a slow version that loops over the heads, giving each head its own column slice of the weight matrices, and confirm that the batched code matches it.

In the toy run the correct split matches the loop exactly, while the direct reshape differs from it by up to 1.0495.¹ A second test, changing token 3 and confirming that outputs for tokens 0 to 2 do not move, checks the causal mask only; the direct reshape passes it too, because its chunks still run in token order.¹ Input: (B, T, D). Output: (B, h, T, d_k). Every head now has its own T by d_k matrix to work on.

Step 3: Score, scale, mask and softmax

For each head, multiply the queries by the transposed keys to get (B, h, T, T) scores, divide by the square root of the head dimension, set blocked positions to minus infinity and take a numerically stable softmax over the last axis. Each row then sums to one.

The division matters. If query and key entries are independent with unit variance (illustrative), a dot product over 64 features has a standard deviation of 8 and scaling brings it to 1.¹ Softmax of the scores 2, 1, 0 gives weights 0.6652, 0.2447 and 0.09; softmax of the same scores multiplied by 8 gives 1.0, 0.000335 and about one in ten million, so the gradient almost vanishes.¹ Vaswani et al. (2017) scale for this reason. A causal mask blocks every position above the diagonal, so a token cannot read later tokens. Output: weights of shape (B, h, T, T).

Step 4: Weigh the values, merge the heads and project out

Multiply the weights by the values to get (B, h, T, d_k), then undo the split: transpose back to (B, T, h, d_k) and reshape to (B, T, D). Finally apply the output projection, which lets the heads' notes mix before the next layer reads them.

Use reshape here, not view, because the transpose leaves the tensor non-contiguous and view raises an error on it. Input: weights and values. Output: (B, T, D), the same shape as the input, so layers can be stacked. The full layer fits in about thirty lines:

import math
import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.h, self.dk = n_heads, d_model // n_heads
        self.wq = nn.Linear(d_model, d_model)
        self.wk = nn.Linear(d_model, d_model)
        self.wv = nn.Linear(d_model, d_model)
        self.wo = nn.Linear(d_model, d_model)

    def split(self, x):                       # (B, T, D) -> (B, h, T, dk)
        B, T, _ = x.shape
        return x.view(B, T, self.h, self.dk).transpose(1, 2)

    def forward(self, x_q, x_kv=None, blocked=None):
        x_kv = x_q if x_kv is None else x_kv  # cross-attention: pass the encoder output
        q = self.split(self.wq(x_q))          # (B, h, Tq, dk)
        k = self.split(self.wk(x_kv))         # (B, h, Tk, dk)
        v = self.split(self.wv(x_kv))         # (B, h, Tk, dk)
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.dk)  # (B, h, Tq, Tk)
        if blocked is not None:               # True = not allowed to attend
            scores = scores.masked_fill(blocked, float("-inf"))
        w = scores.softmax(dim=-1)
        out = (w @ v).transpose(1, 2)         # (B, Tq, h, dk)
        B, Tq = out.shape[0], out.shape[1]
        return self.wo(out.reshape(B, Tq, self.h * self.dk))

1 · Project2 · Split heads3 · Score4 · MergeReplay

Step-by-step description

1. Three linear layers give queries, keys and values.
2. Reshape then transpose gives one matrix per head.
3. Scale, mask and softmax the scores in each head.
4. Weigh values, merge heads, project out.

• The input is projected to queries, keys and values.
• The projections are split into heads.
• Each head scores, masks and applies a softmax.
• The heads are merged and projected back to the input shape.

The mask here is a boolean tensor where true means blocked. Pass torch.ones(T, T, dtype=torch.bool).triu(1) for a causal mask, or pad[:, None, None, :] for a padding mask built from a (B, T) tensor that is true at pad tokens. Library functions differ on whether true means allowed or blocked, so check the documentation of the one you call.

Worked example

The toy run uses a batch of 2, 4 tokens, a model dimension of 8 and 2 heads, so the head dimension is 4. The numbers come from a NumPy mirror of the code above, seeded and run on random inputs (illustrative).¹ Each shape change in the table is one of the steps above.

The row with a causal mask shows the effect: scores 1.0, 2.0, 0.0, 0.0 with the last two blocked give weights 0.2689, 0.7311, 0 and 0.¹ The loop comparison is the bug check: the batched code and the loop agree exactly, while the direct-reshape variant does not.¹ The change-token-3 test is the mask check: it shows that no later token affects earlier outputs, and the weight on future tokens is 0.0 in every batch and head.¹

A padding mask gives zero weight to every padded key.¹ A grouped-query layer with 4 query heads and 2 key-value heads, decoded one token at a time with a cache, matches a full recompute to within 1e-15, and a two-token chunk with the offset mask matches exactly, while the unoffset mask is off by 1.5308.¹ The last table row counts two operations per multiply-add; it is the same for 1, 8 or 16 heads because each head's d_k shrinks as the head count grows.¹

Text description of the diagram

• Scores 1, 2, 0 and 0 with the last two blocked give weights 0.2689, 0.7311, 0 and 0.
• Blocked positions get exactly zero weight, and the open ones share all of it.
• Blocked positions get exactly zero weight, and the open ones share all of it.

Trade-offs and edge cases

More heads leave the cost of the two matrix products unchanged, but the score and weight tensor (B, h, T, T) grows linearly with the head count, and so do the softmax and mask work and the memory: at 1,024 tokens and batch 1 it holds 1,048,576 entries for 1 head and 16,777,216 for 16.¹ Each head also sees fewer features, so very many heads can each become too narrow to be useful. The paper's base model used 8 heads and a head dimension of 64 (Vaswani et al., 2017).

Cross-attention. Cross-attention is attention whose queries come from one sequence and whose keys and values come from another, such as a decoder reading an encoder's output. In the code above it is the same module with a second argument, and the score matrix becomes (Tq, Tk).

Text description of the diagram

• In self-attention queries, keys and values come from one sequence.
• In cross-attention queries come from the decoder and keys and values from the encoder output.
• The cross-attention score matrix has one row per query and one column per memory token.

Fully blocked rows. If every position in a row is minus infinity, the softmax returns not-a-number values that spread through the network. This happens with left padding and a causal mask. Keep at least one position open, or mask with a large finite negative number instead of minus infinity. Replacing those rows after the softmax is not enough, because the gradients through them are still not-a-number.

Cache and grouped heads. The key-value (KV) cache is a store of the keys and values of tokens already processed, so each generation step computes the query, key and value of the new token only and reuses the stored keys and values of earlier tokens. KV cache is called by that name because it holds exactly those two tensors per layer.

Multi-query attention (multi-query attention (MQA), Shazeer, 2019) shares one key-value head across all query heads, and grouped-query attention (grouped-query attention (GQA), Ainslie et al., 2023) shares one per group. Take an illustrative model with 32 layers, a head dimension of 128, 4,096 cached tokens and 16-bit numbers: with 32 key-value heads the cache is 2.0 gibibytes (GiB), with 8 groups 0.5 GiB and with one head 0.0625 GiB.¹ The saving is memory and memory traffic, not the score product.

Text description of the diagram

• With 32 key-value heads the cache is 2.0 gibibytes.
• With 8 groups the cache is 0.5 gibibytes and with one head it is 0.0625.

In code, GQA projects keys and values to g heads, appends the new ones to the cache along the token axis and repeats each group with repeat_interleave(h // g, dim=1) before the score product. This extension reuses the class above:

class GQA(MultiHeadAttention):
    def __init__(self, d_model, n_heads, n_kv_heads):
        super().__init__(d_model, n_heads)
        self.g = n_kv_heads
        self.wk = nn.Linear(d_model, self.g * self.dk)   # g heads, not h
        self.wv = nn.Linear(d_model, self.g * self.dk)

    def forward(self, x, cache=None):            # x holds the new tokens only
        B, Tq, _ = x.shape
        q = self.split(self.wq(x))               # (B, h, Tq, dk)
        k = self.wk(x).view(B, Tq, self.g, self.dk).transpose(1, 2)  # (B, g, Tq, dk)
        v = self.wv(x).view(B, Tq, self.g, self.dk).transpose(1, 2)
        if cache is not None:                    # append on the token axis
            k, v = torch.cat([cache[0], k], dim=2), torch.cat([cache[1], v], dim=2)
        Tk = k.shape[2]
        kk = k.repeat_interleave(self.h // self.g, dim=1)        # (B, h, Tk, dk)
        vv = v.repeat_interleave(self.h // self.g, dim=1)
        s = q @ kk.transpose(-2, -1) / math.sqrt(self.dk)        # (B, h, Tq, Tk)
        if Tq > 1:                               # one decode step needs no mask
            s = s.masked_fill(torch.ones(Tq, Tk, dtype=torch.bool).triu(Tk - Tq + 1), float("-inf"))
        out = (s.softmax(-1) @ vv).transpose(1, 2).reshape(B, Tq, -1)
        return self.wo(out), (k, v)              # the cache keeps g heads

With a cache, Tq and Tk differ, so the causal mask must be offset: triu(diagonal=Tk - Tq + 1) lets each new query see all cached tokens and itself. A single decode step has Tq equal to 1 and needs no mask. The executed mirror of this layer matches a full recompute, and the plain triu(1) mask does not.¹ Generation loops such as beam search are where this cache pays off.

How to answer it in an interview

The spoken structure: say the shape after each line, name the mask convention, then offer the cache and grouped-query extension. Practising the same discipline on another from-scratch task, such as 2D convolution, builds the habit of stating shapes aloud.

Common mistakes

• "I reshape straight to batch, heads, tokens, head dimension." (wrong) That mixes tokens across heads without an error. (right) "I reshape to batch, tokens, heads, head dimension, then transpose."
• "I scale by the square root of the model dimension." (wrong) The scale is set by the size of the dot product. (right) "I divide by the square root of the head dimension."
• "Each head works on the full model dimension." (wrong) Each head gets the model dimension divided by the head count. (right) "The heads split the width, so compute stays similar to one wide head."
• "A causal mask and a padding mask are the same thing." (wrong) One blocks the future, the other blocks pad tokens. (right) "I combine both masks with a logical or."
• "I call view after the transpose." (wrong) It fails on a non-contiguous tensor. (right) "I call reshape, or contiguous first."
• "I use the usual triu(1) mask with a KV cache." (wrong) With a cache the query and key counts differ, so the triangle sits in the wrong place. (right) "I offset the mask by the key count minus the query count, plus one, and skip it for one-token steps."

Follow-up questions

How do you implement multi-head attention in PyTorch from scratch, and make it support cross-attention?

Write the module above and let the keys and values come from a second input. Queries are projected from the decoder sequence, keys and values from the encoder output, and the score matrix has one row per query and one column per memory token.

What are the computationally efficient variants of attention?

Multi-query and grouped-query attention cut the cache and memory traffic by sharing key-value heads. FlashAttention computes the same exact result with less memory movement by working in tiles, so it changes speed, not the maths. Two families cut the quadratic score cost itself and do change the result: sliding-window or local attention, which lets each token see only a window of neighbours (Beltagy et al., 2020), and kernel or linear attention, which replaces the softmax with a feature map so cost grows linearly with length (Katharopoulos et al., 2020).

How do you add a causal mask and a padding mask?

Build one boolean tensor for each, with true meaning blocked, and combine them with a logical or before masked_fill. The causal one is a triangle of shape (T, T); the padding one has shape (B, 1, 1, T) and broadcasts.

How do you extend it to grouped-query attention and a KV cache?

Project keys and values to g heads instead of h, repeat each group h // g times before the score product, and append each new token's keys and values to a stored tensor during generation. The cache then holds g heads per layer instead of h.

Why divide the scores by the square root of the head dimension?

Dividing by the square root of $d_k$ keeps the score spread near 1 so the softmax gradient survives; Step 3 shows the numbers.

Do more heads mean more compute?

The matrix products cost the same, because each head's width shrinks as the head count grows; the score tensor and the softmax grow linearly with the head count, so memory rises with heads. What the extra heads buy is more distinct attention patterns.

Sources

1. Attention Is All You Need (arXiv 1706.03762) — arXiv (Vaswani et al.), 12 Jun 2017
2. torch.nn.functional.scaled_dot_product_attention — PyTorch documentation, 26 Oct 2026
3. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (arXiv 2305.13245) — arXiv (Ainslie et al.), 22 May 2023
4. Fast Transformer Decoding: One Write-Head is All You Need (arXiv 1911.02150) — arXiv (Shazeer), 6 Nov 2019
5. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (arXiv 2205.14135) — arXiv (Dao et al.), 27 May 2022
6. Longformer: The Long-Document Transformer (arXiv 2004.05150) — arXiv (Beltagy et al.), 10 Apr 2020
7. Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention (arXiv 2006.16236) — arXiv (Katharopoulos et al.), 29 Jun 2020