32. Cross-Attention Between Encoder and Decoder Sequences

Problem

Every self-attention layer covered so far has queries, keys, and values all drawn from the same sequence. Encoder-decoder architectures (the original Transformer, T5, and most seq2seq speech/text pipelines where one sequence gets converted into another) add a second kind of attention layer: the decoder's queries attend over the encoder's output sequence instead of over themselves. This has two consequences a self-attention implementation doesn't need to handle: the query sequence length and the key/value sequence length can differ, and there's no causal mask (a decoder position may freely see the entire encoder output, not just a causal prefix of it) — but the encoder output is commonly padded to a fixed length across a batch, so a key-padding mask is needed instead.

Implement multi-head cross-attention: queries come from one sequence length T_dec, keys/values come from a different sequence length T_enc, and an optional per-batch boolean mask marks which encoder positions are padding and must receive zero attention weight regardless of their score.

Source: src/32_cross_attention.py

def cross_attention(
    q: torch.Tensor,                              # (B, H, T_dec, D) decoder queries
    k: torch.Tensor,                              # (B, H, T_enc, D) encoder keys
    v: torch.Tensor,                              # (B, H, T_enc, D) encoder values
    key_padding_mask: torch.Tensor | None = None,  # (B, T_enc) bool, True = PAD (ignore)
) -> torch.Tensor:  # (B, H, T_dec, D)
    ...

>>> q = torch.randn(1, 2, 4, 8)   # 4 decoder positions
>>> k = v = torch.randn(1, 2, 7, 8)  # 7 encoder positions
>>> cross_attention(q, k, v).shape
torch.Size([1, 2, 4, 8])

Step-by-Step Approach

1. Compute raw scaled dot-product scores between the decoder queries and encoder keys: Q @ K^T / sqrt(D), shape (B, H, T_dec, T_enc) — rectangular, not square, since T_dec and T_enc can differ.
2. If a key_padding_mask is given (shape (B, T_enc), True = padding), reshape it to (B, 1, 1, T_enc) so it broadcasts across both the head dimension and every decoder query position.
3. Apply the mask with masked_fill, setting scores at padded key positions to -inf — this is the only masking cross-attention needs; there is no causal component.
4. Explicitly do NOT build a causal mask: every decoder query is allowed to attend to any (non-padded) encoder position, since the encoder already saw its entire input sequence when producing those keys/values.
5. Softmax over the last dimension (T_enc) and matmul with V to get the (B, H, T_dec, D) output.
6. Sanity-check the shape contract carefully: the output's sequence length always matches the QUERY side (T_dec), never the key/value side (T_enc) — a common place to introduce a silent shape bug.

The key insight is that cross-attention is mechanically identical to self-attention (same scaled-dot-product-softmax-matmul recipe) once you stop assuming Q and K/V share both a source and a sequence length; the only real new considerations are a rectangular score matrix and mask semantics that are about padding, not causality.

Reference solution

import math

import torch


def cross_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    key_padding_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    B, H, T_dec, D = q.shape
    T_enc = k.shape[2]
    scale = 1.0 / math.sqrt(D)

    # unlike self-attention, T_dec and T_enc need not match -- scores is
    # (B, H, T_dec, T_enc), not necessarily square
    scores = torch.matmul(q, k.transpose(-1, -2)) * scale  # (B, H, T_dec, T_enc)

    if key_padding_mask is not None:
        # key_padding_mask: (B, T_enc), True = pad. Reshape to (B, 1, 1, T_enc)
        # so it broadcasts across heads and every decoder query position.
        mask = key_padding_mask.view(B, 1, 1, T_enc)
        scores = scores.masked_fill(mask, float("-inf"))

    # no causal mask at all: every decoder position may see the whole
    # (non-padded) encoder output, since encoding already saw the full input
    attn = torch.softmax(scores, dim=-1)
    return torch.matmul(attn, v)  # (B, H, T_dec, D)


TEST_CASES = [
    {"name": "no padding, T_dec != T_enc", "B": 1, "H": 2, "T_dec": 4, "T_enc": 7, "D": 8, "n_pad": 0},
    {"name": "encoder sequence has trailing padding", "B": 2, "H": 2, "T_dec": 3, "T_enc": 6, "D": 4, "n_pad": 2},
    {"name": "T_dec == T_enc but still no causal masking (unlike self-attention)", "B": 1, "H": 1, "T_dec": 5, "T_enc": 5, "D": 4, "n_pad": 0},
]


def main():
    import torch.nn.functional as F

    torch.manual_seed(0)

    for i, case in enumerate(TEST_CASES):
        B, H, T_dec, T_enc, D, n_pad = case["B"], case["H"], case["T_dec"], case["T_enc"], case["D"], case["n_pad"]
        print(f"Test {i}: {case['name']} (B={B}, H={H}, T_dec={T_dec}, T_enc={T_enc}, D={D}, n_pad={n_pad})")
        q = torch.randn(B, H, T_dec, D)
        k = torch.randn(B, H, T_enc, D)
        v = torch.randn(B, H, T_enc, D)

        key_padding_mask = None
        additive_mask = None
        if n_pad > 0:
            key_padding_mask = torch.zeros(B, T_enc, dtype=torch.bool)
            key_padding_mask[:, T_enc - n_pad :] = True
            additive_mask = torch.zeros(B, 1, 1, T_enc).masked_fill(
                key_padding_mask.view(B, 1, 1, T_enc), float("-inf")
            )

        out = cross_attention(q, k, v, key_padding_mask)
        print(f"  out.shape={tuple(out.shape)} (q was T_dec={T_dec}, k/v were T_enc={T_enc})")
        assert out.shape == (B, H, T_dec, D)

        expected = F.scaled_dot_product_attention(q, k, v, attn_mask=additive_mask)
        torch.testing.assert_close(out, expected, atol=1e-4, rtol=1e-4)
        print("PASSED")


if __name__ == "__main__":
    main()

Key Functions & Tricks

• torch.matmul(q, k.transpose(-1,-2)) — produces a (B, H, T_dec, T_enc) score matrix; rectangular whenever the two sequence lengths differ, unlike self-attention's always-square scores.
• key_padding_mask.view(B, 1, 1, T_enc) — reshapes a per-batch, per-key-position boolean mask so it broadcasts across heads and across every decoder query position uniformly.
• masked_fill(mask, float('-inf')) — the same additive-masking mechanism as causal attention, just applied for a different reason (padding, not future-position leakage).
• no causal torch.triu here — cross-attention deliberately omits the causal mask construction that self-attention decoder layers require.
• F.scaled_dot_product_attention(q, k, v, attn_mask=...) — used as the test oracle, taking an explicit additive mask built independently from the same key_padding_mask.

How to Recognize This Pattern