Deep Learning Fundamentals

How do you implement causal attention in PyTorch?

By Interview Clozo · Published 27 October 2026 · Last verified 26 October 2026 · 15 min read

A causal attention implementation in PyTorch is scaled dot-product attention with one change: set every score above the diagonal to minus infinity before the softmax, so position i attends only to positions up to i. One ordering mistake, masking after the softmax, leaves rows that no longer sum to one, and Step 2 shows the numbers.¹

In this article

Before you start
Why interviewers ask this
Causal attention from first principles
Worked example
Trade-offs and edge cases
How to answer it in an interview
Common mistakes
Follow-up questions
Sources

Before you start

Start with how neural networks work if attention is new to you; this page assumes only a little maths.

Why interviewers ask this

Causal attention sits inside every decoder-only language model, so a question on it tests whether a candidate understands both the idea and the code. A weak answer says "we mask the future" and cannot say where in the computation the mask goes. A strong answer says the mask is applied to the scores before the softmax, explains why the value is minus infinity and not zero, and knows which PyTorch arguments to use.

The consequence of getting it wrong is concrete: a model that sees future tokens during training looks excellent on the training loss and fails when it must generate text one token at a time. A Staff-level answer adds the memory picture, the padding-mask and cross-attention distinctions, and the traps in the PyTorch interface.

Causal attention from first principles

Causal attention: Causal attention is a form of self-attention in which each position may read only itself and earlier positions, enforced by a mask on the attention scores, so a model predicting the next token never sees the tokens after it. It is the attention used in decoder-only language models.
Causal mask: A causal mask is a square grid of allowed and blocked pairs of positions, lower-triangular so that row i allows columns 0 to i. It is applied by replacing blocked scores with minus infinity before the softmax, so their weights are exactly zero.

In plain words

Imagine reading a book with a card covering everything to the right of your current word. At each word you may look at anything on the left, never at what the card hides. Causal attention is that card, applied to every word at the same time. The analogy breaks in one place: the model does not read word by word during training. It processes the whole sequence in one pass, and the mask gives each position its own private card.

An auto-regressive model is one that produces each output from the outputs before it, and the mask is what keeps training faithful to that rule. Vaswani et al. (2017) put it as preventing "positions from attending to subsequent positions", so that predictions for position i depend only on known outputs at positions before i. In one precise sentence, causal attention is attention whose grid of attention weights is forced to be lower-triangular.

Step 1: Scaled dot-product scores compare every query with every key

Scaled dot-product attention is the operation that scores each query vector against each key vector with a dot product, divides by the square root of the head size, and later uses the scores to mix value vectors. The input is a batch of sequences projected into queries, keys and values by learned weight matrices. The output of this step is one score for every pair of positions: a square grid of T by T numbers per head for a sequence of T tokens.

PyTorch's reference code uses a default scale of one over the square root of the last dimension, which is 0.125 for a head size of 64.¹ At this point every position can still score every other position, including the ones in its future. The next step fixes that.

Step 2: The causal mask sets illegal scores to minus infinity before the softmax

The input is the score grid; the output is the same grid with every entry above the diagonal replaced by minus infinity. Vaswani et al. do this "inside of scaled dot-product attention by masking out (setting to −∞) all values in the input of the softmax which correspond to illegal connections." In PyTorch the allowed pattern is torch.ones(L, S, dtype=torch.bool).tril(diagonal=0); blocked entries get float("-inf").

The value must be minus infinity because the exponential of minus infinity is exactly zero. A score of zero would still give an exponential of one, so a blocked position would keep a real weight. Take the third token with illustrative scores of 1.0, 2.0 and 0.5 for the allowed positions and 3.0 for the blocked one: after masking, the weights are about 0.23, 0.63 and 0.14, and the blocked position gets a weight of exactly zero.¹ Masking after the softmax is wrong: the unmasked softmax gives the blocked position about 0.63 of the weight, and zeroing it afterwards leaves a row summing to only about 0.37.¹ That follows from the softmax algebra.

Step 3: Softmax turns each masked row into weights that sum to one

The input is the masked grid and the output is a grid of weights in which each row is a probability distribution. The softmax is applied along the key dimension, so each query row is normalised separately. Because blocked entries are exactly zero after the exponential, softmax over a masked row equals softmax over only its first i plus one scores.

One useful consequence: row 0 has only one allowed entry, so the first token always puts all of its weight on itself, whatever its score.¹ Another is that a row's weights never depend on tokens to its right, which is the property the mask exists to enforce. If you apply dropout to attention weights, do it here, after the softmax; PyTorch's function applies dropout according to dropout_p whenever it is called, so pass dropout_p of zero when evaluating. See activation functions for the softmax itself.

Step 4: The weights mix the value vectors into the output

The input is the weight grid and the value vectors; the output is one new vector per position. Each output is a weighted sum of the value vectors at the positions that row is allowed to see, so position i is a blend of positions 0 to i. In code this is one matrix multiplication of the weight grid with the values.

After this, the heads are put back side by side and a final linear layer mixes them, usually inside a block with a feed-forward network and residual connections. The whole four-step path is the figure below: scores, mask, softmax, weighted sum. Illegal scores are set to minus infinity before the softmax, each row of weights sums to one after it, and the weights then mix the value vectors.

Masking happens between the scores and the softmax, so masked positions get exactly zero weight.CC BY 4.0 — source: How do you implement causal attention in PyTorch?

Step-by-step description

Every query is scored against every key.
Scores above the diagonal become minus infinity.
Each row of weights sums to one.
Weights blend the value vectors.

Illegal scores are set to minus infinity before the softmax.
Each row of weights sums to one after the softmax.
The weights then mix the value vectors.

Knowing it is half the job. Saying it is the other half.

Interview Clozo is an AI interview copilot for real interview rounds: it hears the question and puts a strong answer on your screen, and you say it in your own words. See the features or how it works.

Start Free Trial →

Worked example

The hand-written module below is the whole algorithm. It uses one fused linear layer that produces query, key and value together, as the PyTorch attention tutorial does, then follows the four steps in order.

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

class CausalSelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        assert embed_dim % num_heads == 0
        self.h = num_heads
        self.hd = embed_dim // num_heads
        self.qkv = nn.Linear(embed_dim, 3 * embed_dim)  # fused projection
        self.proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x, fused=False):
        B, T, E = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q, k, v = (t.view(B, T, self.h, self.hd).transpose(1, 2) for t in (q, k, v))
        if fused:
            y = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=True)
        else:
            scores = q @ k.transpose(-2, -1) / math.sqrt(self.hd)      # step 1
            allowed = torch.ones(T, T, dtype=torch.bool, device=x.device).tril()
            scores = scores.masked_fill(~allowed, float("-inf"))      # step 2
            y = scores.softmax(dim=-1) @ v                            # steps 3 and 4
        return self.proj(y.transpose(1, 2).reshape(B, T, E))

The branch fused=True calls F.scaled_dot_product_attention with is_causal=True, which gives the causal mask without building it. It is expected to match the hand-written branch up to floating-point rounding; this page did not run that comparison, so check it with torch.allclose on your own inputs. Train with dropout_p above zero if you want attention dropout; keep it at zero here.

Cross-attention changes only where the inputs come from and drops the causal mask. This variant takes queries from x and keys and values from a separate context, so the query length L and the key length S can differ.

class CrossAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.h, self.hd = num_heads, embed_dim // num_heads
        self.q_proj = nn.Linear(embed_dim, embed_dim)        # queries from x
        self.kv_proj = nn.Linear(embed_dim, 2 * embed_dim)   # keys, values from context
        self.proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x, context, pad_mask=None):  # x: (B, L, E), context: (B, S, E)
        B, L, E = x.shape
        S = context.shape[1]
        q = self.q_proj(x).view(B, L, self.h, self.hd).transpose(1, 2)   # (B, h, L, hd)
        k, v = self.kv_proj(context).chunk(2, dim=-1)
        k, v = (t.view(B, S, self.h, self.hd).transpose(1, 2) for t in (k, v))  # (B, h, S, hd)
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=pad_mask)  # no is_causal
        return self.proj(y.transpose(1, 2).reshape(B, L, E))

The optional pad_mask is a boolean tensor of shape (B, 1, 1, S) in which True marks a real encoder token that takes part, so padded encoder positions get zero weight; the call passes no is_causal.

A causal mask over T tokens allows T(T+1)/2 pairs, so about half of the score grid is blocked, while the score grid itself grows with the square of the length.
Quantity	Value (computed or illustrative)
Allowed pairs in a 4-token mask	10
Blocked pairs in a 4-token mask	6
Allowed pairs at 2,048 tokens	2,098,176 of 4,194,304 (0.5002)
Default scale for head size 64	0.125
Fused projection width for embedding size 512	1,536
Third-token weights, illustrative scores	0.23, 0.63, 0.14
Row sum if masked after the softmax, illustrative	0.37
Score matrix, 8 heads, 16-bit floats, 512 tokens	4.19 MB
Score matrix, 8 heads, 16-bit floats, 1,024 tokens	16.78 MB
Score matrix, 8 heads, 16-bit floats, 2,048 tokens	67.11 MB
Score matrix, 8 heads, 16-bit floats, 4,096 tokens	268.44 MB

The tutorial uses 8 heads of 64 dimensions, an embedding size of 512, so the fused layer maps 512 numbers to 1,536.¹ The score matrix for 8 heads in 16-bit floats (see 16-bit floating point (FP16) vs brain floating point 16 (BF16)) is about 67 MB per layer at 2,048 tokens, from 8 heads times 2,048 squared entries times 2 bytes, for one sequence.¹ Doubling the length quadruples the score matrix.

Doubling the sequence length quadruples the score matrix, which is why fused kernels matter.CC BY 4.0 — source: How do you implement causal attention in PyTorch?

Text description of the diagram

The score matrix for 8 heads in 16-bit floats is about 67 MB per layer at 2,048 tokens.
Doubling the length quadruples the score matrix.

Trade-offs and edge cases

A hand-written mask is clear but stores the full score matrix. Dao et al. (2022) note that the time and memory of standard attention are quadratic in sequence length, and their FlashAttention computes exact attention with memory linear in sequence length. It is not an approximation: it changes memory access, not the maths, and the arithmetic stays quadratic. The repository README reports 10 times memory savings at sequence length 2K and 20 times at 4K. According to its documentation, PyTorch's function dispatches to FlashAttention-2, a memory-efficient kernel or a plain C++ implementation of the maths.

The three common routes build the same mask and differ in cost and in how they fail.
Route	What to know
Hand-written mask	The mask is built explicitly with `tril` and filled with minus infinity, which is clear but stores the full score matrix.
`is_causal=True`	No mask is built, and PyTorch raises an error if `attn_mask` is also set.
`nn.MultiheadAttention`	The `is_causal` argument is only a hint that `attn_mask` is causal, so you still pass the mask.

In the square grid below, each row allows only the keys up to its own position, and the third row carries the illustrative weights 0.23, 0.63, 0.14 and 0.¹

Each row of a causal grid allows only the keys up to its own position, and blocked keys get exactly zero weight.CC BY 4.0 — source: How do you implement causal attention in PyTorch?

Text description of the diagram

Each row of the causal grid allows only the keys up to its own position.
The third row carries the illustrative weights 0.23, 0.63, 0.14 and 0.

Edge cases follow. Scaled dot-product attention (SDPA) raises an error if both attn_mask and is_causal are set, according to the SDPA documentation. The is_causal argument on nn.MultiheadAttention is a hint that attn_mask is the causal mask (MultiheadAttention documentation), and a reported issue against PyTorch 2.0.0 says it was ignored when need_weights=True; its current status is unverified, so pass an explicit attn_mask and set need_weights=False to use the optimised path.

Boolean masks have opposite meanings: in SDPA True means a position takes part, while in the key_padding_mask of nn.MultiheadAttention it means the position is ignored. A causal mask also does not handle padding; that is a separate mask. If padding plus the causal mask blocks every key of a row, for example a padded query row under left padding, the softmax over all minus infinity gives not-a-number values that spread through later layers, so keep at least one allowed key per row or zero those rows out.

When the query length L and key length S differ, as in decoding with cached keys, is_causal=True aligns the diagonal at the upper left, which is wrong for a cache: a single new query would see only the first key. Use causal_lower_right(L, S) from torch.nn.attention.bias. For one new token at a time, no mask is needed, because the newest query may see every cached key. In the grid below, with two new queries and four cached keys, upper-left alignment lets the first new query see only key 0, while lower-right alignment lets the last new query see every cached key.

For decoding with a cache, the diagonal must sit at the lower right, or new queries lose most of the cache.CC BY 4.0 — source: How do you implement causal attention in PyTorch?

Text description of the diagram

Upper-left alignment lets the first new query see only key 0.
Lower-right alignment lets the last new query see every cached key.

Normalisation layers around attention are covered in batch norm vs layer norm.