How do you implement scaled dot-product attention from scratch?

In this article

1. Before you start
2. Why interviewers ask this
3. Scaled dot-product 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 a step of the PyTorch and machine learning (ML) coding round guide, and the code is plain PyTorch tensor operations, with no nn.Module needed.

Why interviewers ask this

Attention is the core operation of a transformer, so a coding round that asks you to write it tests whether you understand the model or only call it. Candidates report being asked to implement self-attention, an attention head or masked attention in ML coding rounds at AI labs and large tech companies; these are reports, not a public statement from any company.

A weak answer writes the formula from memory and gets the shapes or the scale wrong. A strong answer states the shape of every tensor, explains why the scale is the square root and not d_k itself, and puts the mask before the softmax. A Staff-level answer adds the cost: the score matrix grows with the square of the sequence length, and says how production code avoids storing it.

Scaled dot-product attention from first principles

Scaled dot-product attention

Scaled dot-product attention is a function that returns, for each query, a weighted average of value vectors, where the weights come from a softmax over that query's dot products with every key, divided by the square root of the key size d_k so the softmax stays well behaved.

Attention weights

Attention weights are the softmax output for one query: one non-negative number per key, summing to one, that says how much of each value goes into that query's output.

In plain words

Think of a reader with a question looking at a shelf of index cards. Each card has a short title, which is its key, and a longer note, which is its value. The reader compares the question with every title, decides how relevant each card is, and writes a blended summary of the notes, leaning toward the relevant ones. The analogy breaks in one place. The reader picks a few cards, while attention blends all of them with smooth weights.

In one precise sentence, the output is the softmax of Q times the transpose of K, divided by the square root of d_k, then multiplied by V (Vaswani et al., 2017, section 3.2.1).

For a single sequence the shapes are: Q is (T_q, d_k), K is (T_k, d_k), V is (T_k, d_v). The scores are (T_q, T_k) and the output is (T_q, d_v). A batch adds a leading dimension to each.

Where do Q, K and V come from in self-attention?

Q, K and V are learned linear projections of the token embeddings. Self-attention is attention in which Q, K and V all come from the same sequence. If x holds one embedding per token, with shape (T, d_model), then q = x @ W_q, k = x @ W_k and v = x @ W_v. W_q and W_k have shape (d_model, d_k), and W_v has shape (d_model, d_v). The three matrices are weights that training learns.

Cross-attention takes Q from one sequence, such as the decoder's tokens, and K and V from another, such as the encoder's output; the causal mask and cross-attention page covers it. A single head of self-attention is then three projections and one call to the function defined below:

q, k, v = x @ w_q, x @ w_k, x @ w_v   # each projection is a matrix product
mask = torch.triu(torch.ones(x.size(0), x.size(0), dtype=torch.bool), diagonal=1)
out, weights = sdpa(q, k, v, mask)    # out: (T, d_v)

Step 1: Score every query against every key

The scores are the dot products of each query with each key, computed for all pairs at once as Q times the transpose of K. Input: Q of shape (T_q, d_k) and K of shape (T_k, d_k). Output: a (T_q, T_k) matrix where entry (i, j) says how well query i matches key j.

In code this is q @ k.transpose(-2, -1). Transposing only the last two dimensions leaves any batch or head dimension in front untouched, which is why -2, -1 is written instead of a plain .T. You now have raw match scores, but they are not yet safe to feed to a softmax, as the next step shows.

Step 2: Divide the scores by the square root of d_k

Dividing by the square root of d_k keeps the scores at a size where the softmax still has useful gradients. If the entries of a query and a key are independent, with mean 0 and variance 1, their dot product is a sum of d_k terms and has variance d_k, so its standard deviation grows as the square root of d_k. Dividing by that square root brings the variance back to 1. Vaswani et al. (2017) use d_k of 64, a scale of 0.125.

Input: the raw scores. Output: scaled scores of the same shape. The scale is the square root, not d_k: dividing by d_k would shrink the scores to a variance of 1 over d_k and make every weight nearly equal. The worked example measures the effect.

Step 3: Apply the mask before the softmax, not after

A mask blocks some query-key pairs by setting their scores to minus infinity before the softmax, so their weights come out as exactly zero. A causal mask is a mask that blocks every key after the current position, so a decoder cannot read tokens it has not produced yet.

Build it as an upper-triangular boolean matrix: torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1), where True means blocked. Then call scores.masked_fill(mask, float("-inf")). Zeroing the weights after the softmax is the common error: the surviving weights no longer sum to one, so the output is silently shrunk. The worked example shows the exact row sums.¹ The causal mask page covers encoder, decoder and cross-attention masks in full.

Step 4: Softmax the rows, then take the weighted sum of V

The softmax runs along the last dimension, so each query gets its own weights over the keys, and the output is those weights times V. Input: masked scores (T_q, T_k) and V (T_k, d_v). Output: (T_q, d_v), one blended value vector per query.

Use torch.softmax(scores, dim=-1), and never dim=0 or dim=1 on batched input, because the softmax must normalise over keys. Library softmax subtracts the row maximum before exponentiating so large scores do not overflow; the stable softmax page shows how. Return the weights as well as the output while you debug, since a row that does not sum to one points straight at a mask bug.

1 · Scores2 · Scale3 · Mask4 · Softmax and mixReplay

Step-by-step description

1. Q times K transposed gives one score per pair.
2. Divide by the square root of d_k.
3. Blocked scores become minus infinity.
4. Weights multiply V to give the output.

• The scores are the dot products of queries and keys.
• The scores are divided by the square root of d_k.
• Blocked positions are set to minus infinity before the softmax.
• The softmax weights multiply V to give the output.

The whole function, with True meaning blocked:

import math
import torch

def sdpa(q, k, v, mask=None):
    # q: (..., T_q, d_k)  k: (..., T_k, d_k)  v: (..., T_k, d_v)
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / math.sqrt(d_k)   # steps 1 and 2
    if mask is not None:                                 # step 3
        scores = scores.masked_fill(mask, float("-inf"))
    weights = torch.softmax(scores, dim=-1)              # step 4
    return weights @ v, weights

PyTorch also ships this as torch.nn.functional.scaled_dot_product_attention; read its documentation for how it reads a boolean mask. In that function a boolean True in attn_mask means the position takes part in attention, the reverse of the convention above.

Worked example

Use three tokens with d_k of 4 and d_v of 2. The values are illustrative. Q has rows (1, 0, 1, 0), (0, 1, 0, 1) and (1, 1, 0, 0); K has rows (1, 0, 1, 0), (0, 1, 0, 1) and (1, 1, 1, 1); V has rows (1, 0), (0, 1) and (1, 1).¹

The raw scores for query 1 are 2, 0 and 2. The scale is 1 over the square root of 4, which is 0.5, so the scaled scores are 1, 0 and 1, and the softmax gives weights 0.4223, 0.1554 and 0.4223. The output for query 1 is 0.4223 times (1, 0) plus 0.1554 times (0, 1) plus 0.4223 times (1, 1), which is (0.8446, 0.5777).¹ Every row of weights sums to 1.¹

With the causal mask, query 1 sees only itself, so its weights are 1, 0 and 0. Query 2 gets 0.2689 and 0.7311, and query 3 keeps the full row 0.2741, 0.2741 and 0.4519.¹ If the same mask is applied after the softmax, query 1 keeps only 0.4223 of its weight and query 2 keeps 0.5777, so the rows sum to less than one.¹

Text description of the diagram

• Masked before the softmax, every row of weights sums to one.
• Masked after the softmax, query 1 keeps 0.4223 and query 2 keeps 0.5777 of their weight.

To see why the scale matters, the script draws random vectors with mean 0 and variance 1 (standard normals) and measures them over 20,000 pairs per size.¹

The largest-weight columns average over 2,000 random queries, each attending to 8 random keys.¹ At d_k of 64 the unscaled standard deviation is 7.93, so one key tends to take almost all the weight, and the softmax gradient for the other keys becomes tiny. The scaled version stays near 0.36 at every size.¹

Text description of the diagram

• Unscaled, the mean largest softmax weight is 0.525 at d_k of 4 and 0.935 at d_k of 256.
• Scaled, the mean largest weight stays near 0.36 at every size.

Trade-offs and edge cases

The main cost is memory. The score matrix has T_q times T_k entries per head, so it grows with the square of the sequence length: 1,048,576 entries at length 1,024 and 67,108,864 at length 8,192, which is 4,194,304 bytes and 268,435,456 bytes per head in 32-bit floats.¹ Dao et al. (2022) describe FlashAttention, which computes the same exact result in blocks without storing the full matrix.

A second edge case is a fully masked row, for example a padded query whose keys are all blocked. Every score is minus infinity, so the softmax divides zero by zero and returns not-a-number (NaN) values that then spread through later layers. Guard it by replacing those rows with zeros or by keeping at least one unmasked key.

Finally, this function has no notion of order: swapping two tokens swaps their rows and nothing else (without a mask), which is why models add positional encodings. To run several attention heads in parallel, see multi-head attention in PyTorch without undefined.

How to answer it in an interview

The spoken structure: name the shapes, say the formula, justify the square root, place the mask before the softmax, then give the memory cost. Say the shapes out loud as you write each line; interviewers often check those first.

Common mistakes

• "I scale by d_k." (wrong) The scale is the square root of d_k. (right) "I divide by the square root of d_k so the score variance stays near one."
• "I apply the causal mask after the softmax." (wrong) Rows then sum to less than one. (right) "I fill blocked scores with minus infinity before the softmax."
• "I softmax over dimension 0." (wrong) That normalises across queries or the batch. (right) "I softmax over the last dimension, the keys."
• "I transpose K with .T." (wrong) On a batched tensor this is the wrong permutation or an error. (right) "I use k.transpose(-2, -1) so batch and head dimensions stay in front."

Follow-up questions

Why do you divide by the square root of d_k?

Because a dot product of two d_k-sized vectors with zero-mean, unit-variance entries has variance d_k, and the square root of that is the standard deviation. Dividing by it returns the scores to variance 1, which keeps the softmax away from saturation (Vaswani et al., 2017).

How do you add a causal mask?

Build an upper-triangular boolean matrix with torch.triu(..., diagonal=1) and fill those score positions with minus infinity before the softmax. The causal mask page shows the padding and cross-attention variants.

How does multi-head attention build on this function?

It projects the input into several smaller queries, keys and values, runs this function on each head, concatenates the results and applies one output projection. Vaswani et al. (2017) use 8 heads. The next step is the multi-head attention walkthrough, and the generative pre-trained transformer (GPT) build puts it in a decoder.

What is the difference between self-attention and scaled dot-product attention?

Scaled dot-product attention is the function; self-attention is using it with Q, K and V all projected from the same sequence. Cross-attention takes Q from one sequence and K and V from another.

What if d_k is 1 or very large?

With d_k of 1 the scale is 1, so scaling changes nothing. At very large d_k the scale matters most: in the worked example the unscaled largest weight reaches 0.935 at d_k of 256, against 0.358 when scaled.¹

What are multi-query and grouped-query attention?

Multi-query attention (MQA) shares one key and value head across all query heads, which shrinks the cache used in decoding (Shazeer, 2019). Grouped-query attention (GQA) shares one key and value head per group of query heads, a middle ground (Ainslie et al., 2023).

Sources

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