scaled_dot_product_attention#

Scaled dot-product attention is the core of transformer models. Each query vector scores every key vector with a dot product, scaled by 1/sqrt(E) for vectors of length E. A softmax turns each query's scores into weights: exponentials normalized to sum to 1. The output for the query is the weighted average of the value vectors. A mask can hide keys, such as future positions in a language model. Implementations differ a lot: PyTorch's math path computes the weights as a full matrix, while fused kernels such as FlashAttention process the keys in blocks with a running softmax and round in extra places.

In ExactRef, er.scaled_dot_product_attention(query, key, value, ...), also exactref.nn.attention.scaled_dot_product_attention, computes softmax(Q K^T * scale + mask) V over the key axis, with the arguments, shapes and special values of torch.nn.functional.scaled_dot_product_attention. It takes real inputs.

Arguments#

Argument

Meaning

Default

query, key, value

real arrays of shapes (..., L, E), (..., S, E) and (..., S, Ev)

required

attn_mask

a bool mask (True keeps a key) or a float mask added to the scaled scores (-inf removes a key)

None

dropout_p

must be 0: dropout is random and has no reference

0.0

is_causal

keep key j for query i only when j <= i

False

scale

the factor on the scores

None, which means 1/sqrt(E)

enable_gqa

allow fewer key and value heads than query heads

False

dtype

the output format, of the same kind as the inputs

the input format

rounding

the rounding of the reference and of the final step into the output format

er.Rounding.NEAREST_EVEN

overflow

what a result past the output format's range becomes: "ieee", "saturate" or "nan" for floats; "wrap", "error" or "saturate" for integers

each format's own rule; integers wrap

Examples#

Code
import exactref as er

# A shorthand for these examples: the values as a float32 array
def f32(values):
    return er.asarray(values, "float32")

# Print a (1, L, Ev) output one row per line
def rows(array):
    for row in array.tolist()[0]:
        print(row)

# A batch of 1: three queries, keys and values, each a vector of length 2
q = f32([[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]])
k = f32([[[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]]])
v = f32([[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]])

# ExactRef's reference for attention on these inputs
a = er.scaled_dot_product_attention(q, k, v)
# a.rounded holds both outputs, by name
print("outputs:", list(a.rounded))
print("output shape:", a.rounded.output.shape)
print("logsumexp shape:", a.rounded.logsumexp.shape)
# The output, each element the exact value rounded once to float32
print("output:")
rows(a.rounded.output)
Output
outputs: ['output', 'logsumexp']
output shape: (1, 3, 2)
logsumexp shape: (1, 3)
output:
[2.728677272796631, 3.728677272796631]
[3.1905195713043213, 4.1905198097229]
[3.0, 4.0]

Masks, causality and scale#

Code
# A shorter name for the operation
sdpa = er.scaled_dot_product_attention

# Query i sees only keys 0 to i
print("causal:")
rows(sdpa(q, k, v, is_causal=True).rounded.output)

# A float mask, added to the scaled scores: -inf hides key 1 from every query
mask = f32([[0.0, float("-inf"), 0.0]] * 3)
print("masked:")
rows(sdpa(q, k, v, attn_mask=mask).rounded.output)

# Scores multiplied by 1 instead of 1/sqrt(2)
print("scale 1:")
rows(sdpa(q, k, v, scale=1.0).rounded.output)
Output
causal:
[1.0, 2.0]
[2.3395230770111084, 3.3395230770111084]
[3.0, 4.0]
masked:
[2.6500840187072754, 3.6500840187072754]
[3.3499159812927246, 4.349915981292725]
[3.0, 4.0]
scale 1:
[2.601430892944336, 3.601430892944336]
[3.24174427986145, 4.241744518280029]
[3.0, 4.0]

With enable_gqa, the query's head axis (the third from the end) is a multiple of the key's and value's, as in grouped-query attention:

Code
# 2 query heads, each with one query
gq = f32([[[[1.0, 0.0]], [[0.0, 1.0]]]])
# 1 key head and 1 value head, for two keys
gk = f32([[[[1.0, 0.0], [0.0, 1.0]]]])
gv = f32([[[[1.0, 2.0], [3.0, 4.0]]]])

# Both query heads use the one key and value head
print("output shape:", sdpa(gq, gk, gv, enable_gqa=True).rounded.output.shape)
Output
output shape: (1, 2, 1, 2)

logsumexp and its format#

logsumexp is in the compute format of the output: float32 for formats narrower than float32, as fused kernels return it.

Code
# Attention on q, k and v converted to bfloat16
b = sdpa(*(er.asarray(t, "bfloat16") for t in (q, k, v)))
print("output dtype:", b.rounded.output.dtype.name)
print("logsumexp dtype:", b.rounded.logsumexp.dtype.name)
Output
output dtype: bfloat16
logsumexp dtype: float32

Tolerance options#

The tolerance budgets exp and log within ulps units in the last place (default 2). Kernels with a less accurate exp need a larger value:

Code
# The largest tolerance over the output's elements, with the default ulps
print("default:", max(a.tolerance().output.scalars()))
# The same when exp and log may each be off by up to 8 ulps
print("ulps=8:", max(a.tolerance(ulps=8).output.scalars()))
Output
default: 5.9604644775390625e-06
ulps=8: 1.4543533325195312e-05

Special values and dropout#

A NaN input makes its row NaN. A row with no kept key gets weight 0 on every key: its output is 0 and its logsumexp -inf. Dropout is random and has no reference:

Code
# One query that holds a NaN
print("NaN row:", sdpa(f32([[[float("nan"), 0.0]]]), k, v).rounded.output.tolist())

# Any dropout_p other than 0 raises
try:
    sdpa(q, k, v, dropout_p=0.1)
except ValueError as error:
    print(f"{type(error).__name__}: {error}")
Output
NaN row: [[[nan, nan]]]
ValueError: dropout is random and has no reference; dropout_p must be 0

Reading complicated tolerances shows a fully masked row.

Checking PyTorch#

Code
import math
import random
import torch

# Seeded, so every run uses the same values
rng = random.Random(0)

# A tensor of the given shape, filled with random values
def randn(*shape):
    # One normally distributed value (mean 0, standard deviation 1) per element
    values = [rng.gauss(0, 1) for _ in range(math.prod(shape))]
    # From Python floats, torch.tensor makes float32, PyTorch's default
    return torch.tensor(values).reshape(shape)

# A batch of 1 with 2 heads, each with 8 vectors of length 16
tq, tk, tv = randn(1, 2, 8, 16), randn(1, 2, 8, 16), randn(1, 2, 8, 16)

# PyTorch's causal attention
out = torch.nn.functional.scaled_dot_product_attention(tq, tk, tv, is_causal=True)
# Could a correct implementation produce PyTorch's output? covers() checks one
# output at a time, so ["output"] picks it.
print("covered:", sdpa(tq, tk, tv, is_causal=True)["output"].covers(out).all())
Output
covered: True

Fused kernels such as FlashAttention round in places this tolerance does not budget; use an attention algorithm.

Padding masks#

Padding masks, such as the ones Hugging Face Transformers builds, hide a key by adding the format's lowest finite value instead of -inf. A query whose keys are all padded still keeps every key, with the same huge value added to each score. In exact arithmetic that shift cancels in the softmax, so the exact output weights the values by the scores. An implementation rounds each score plus its mask value once (the logits step), and at the lowest value that rounding erases the score, so the weights come out equal. The tolerance stays finite and covers that output:

Code
# The lowest finite float32 value
low = torch.finfo(torch.float32).min
# Two queries, two keys and two values, each of length 1
query = torch.tensor([[[1.0], [0.5]]])
key = torch.tensor([[[1.0], [-1.0]]])
value = torch.tensor([[[1.0], [2.0]]])
# Query 0: every key padded. Query 1: key 1 padded.
mask = torch.tensor([[[low, low], [0.0, low]]])

# PyTorch's attention with this padding mask
out = torch.nn.functional.scaled_dot_product_attention(
    query, key, value, attn_mask=mask)
print("PyTorch:")
rows(out)

# ExactRef's reference for the same output
p = sdpa(query, key, value, attn_mask=mask)["output"]
print("rounded:")
rows(p.rounded)
print("tolerance:")
rows(p.tolerance())
# Could a correct implementation produce PyTorch's output?
print("covered:", p.covers(out).all())
Output
PyTorch:
[1.5]
[1.0]
rounded:
[1.119202971458435]
[1.0]
tolerance:
[0.38079822063446045]
[1.3113021850585938e-06]
covered: True

The tolerance is wide because it is measured from rounded, the exact answer. finite_range() shows where correct outputs land, around the even average 1.5:

Code
# The lowest and highest finite output of a correct implementation, per element
lo, hi = p.finite_range()
# Element 0 is query 0's output, whose keys are all padded
print("low:", lo.scalars()[0])
print("high:", hi.scalars()[0])
Output
low: 1.499998927116394
high: 1.5000011920928955

Backward#

er.scaled_dot_product_attention.backward(query, key, value, grad, ...) takes the forward's arguments with grad after value, and returns grad_query, grad_key and grad_value, each rounded once.

Code
# A gradient of 1 for every output element, in the output's shape (1, 3, 2)
grad = f32([[[1.0, 1.0]] * 3])
# The reference backward, and the names of its outputs
print("outputs:", list(sdpa.backward(q, k, v, grad).rounded))
Output
outputs: ['grad_query', 'grad_key', 'grad_value']

Backward checks these gradients against PyTorch's autograd.

Full specification#

er.scaled_dot_product_attention#

scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p: float = 0.0, is_causal: bool = False, scale=None, enable_gqa: bool = False, *, dtype=None, rounding: Rounding = Rounding.NEAREST_EVEN, overflow: str | None = None) -> Result#

softmax(Q K^T * scale + mask) V over the key axis.

Shapes follow torch.nn.functional.scaled_dot_product_attention: query (..., L, E), key (..., S, E), value (..., S, Ev); with enable_gqa the head axis (-3) of query is a multiple of key/value's. A bool mask keeps keys where True; a float mask is added to the scaled scores (-inf removes a key). is_causal keeps key j for query i when j <= i. scale defaults to 1/sqrt(E); with E = 0 the scores are empty sums and stay 0, as in PyTorch, so the logits are the mask alone.

Every output is the exact value correctly rounded to dtype. Non-finite values propagate through the formula as IEEE arithmetic does on exact values, as in PyTorch's math path. A masked key (False in a bool mask, -inf in a float mask, or above the diagonal with is_causal) gets a -inf bias. A key whose logit is -inf gets weight 0; a logit of +inf or NaN makes the row NaN; a row with no finite logit gets weight 0 on every key. In the weighted sum of value rows, a non-finite value in a key with positive weight gives that infinity (NaN for both signs or NaN), and one in a key with weight 0 gives NaN (0 * inf). An exact zero is +0, or -0 when rounding toward -inf. Dropout is random and has no reference, so dropout_p must be 0.

The second output, logsumexp (shape (..., Hq, L), as the fused kernels return it), is log(sum_j exp(l_j)) over each row's logits, correctly rounded to the compute format of dtype (float32 for formats narrower than float32, else dtype itself). It follows IEEE arithmetic on the exact values: NaN if a logit is NaN, +inf if one is +inf, and -inf, log(0), for a row with no kept key.

tolerance() budgets the required boundaries: the score and value accumulations, the scale factor, one rounding of each logit, the shift by the row maximum, exp within ulps (default 2), the row sum, the final division, and the rounding to dtype; for logsumexp, log of the row sum within ulps and one rounding of the row maximum plus that logarithm. The three sums (scores, row sum, value sums) are normal accumulations by default (index order or 2 or 4 SIMD-style lanes, as matmul describes); standard="any_order" covers any order. normal() is the whole model with normal accumulation in index order and the scale, exp and log correctly rounded. Other roundings are not budgeted: a separate mask rounding, materialized probabilities, casting the weights to the input dtype before the value product (FlashAttention does this for bf16 and fp16), online-softmax rescaling, or a reciprocal of the row sum; named algorithms (AttentionAlgorithm) budget all of these except materialized probabilities (P normalized and rounded before the product with V, as PyTorch's math path does), which no tolerance budgets. Pre-scaling Q and K changes the inputs: a consumer that wants it passes the pre-scaled values with scale=1. An implementation that adds such roundings can exceed the tolerance.

Every input, the mask included, may also be an earlier result, as Result describes.

er.scaled_dot_product_attention.backward#

backward(query, key, value, grad, attn_mask=None, dropout_p: float = 0.0, is_causal: bool = False, scale=None, enable_gqa: bool = False, *, dtype=None, rounding: Rounding = Rounding.NEAREST_EVEN, overflow: str | None = None) -> Result#

Gradients of scaled_dot_product_attention for the output gradient grad.

The arguments are the forward's, with grad after value; grad has the output's shape and the inputs' dtype. Returns grad_query, grad_key and grad_value with the inputs' shapes: with P the attention weights, dP = grad V^T, D_i = sum_j P_ij dP_ij and dS = P * (dP - D),

grad_value = P^T grad, grad_query = scale * dS K, grad_key = scale * dS^T Q,

each summed over the copies that broadcasting or enable_gqa made of its input. Every output is the exact value correctly rounded to dtype.

Non-finite values propagate as IEEE arithmetic does on exact values through these formulas, as in PyTorch's math path: masked keys have P = 0, so an infinite dP there makes D NaN; a NaN row of the forward makes every gradient it touches NaN. A row with no kept key has P = 0 and contributes 0. An exact zero is +0, or -0 when rounding toward -inf; an empty sum is +0. Dropout must be 0, and a float mask gets no gradient.

Exactness: grad_query of a row is decided exactly (a Lindemann-Weierstrass test). grad_key, grad_value, and grad_query summed over broadcast copies of a query row, are found exact when every term is rational; otherwise refinement stops at a precision cap and raises ReferenceUnavailableError rather than return an uncertain value (the one accepted case where an exact result can reach the cap).

tolerance() budgets the required boundaries of the gradient computation: the forward's scores, scale factor, logits, shift, exp within ulps (default 2) and row sum, then the probabilities P (one rounding each), the accumulations of dP, D, grad_value, dS K and dS^T Q, the rounding of dP - D, of dS and of the scale product, and the rounding to dtype. The accumulations are normal accumulations by default (index order or 2 or 4 SIMD-style lanes, as matmul describes); standard="any_order" covers any order. normal() is that program with normal accumulation in index order and the scale and exp correctly rounded. Not budgeted: computing P with a reciprocal of the row sum, rounding to the input dtype in between, or computing D from the forward output; named algorithms (AttentionAlgorithm) budget these. Pre-scaling Q and K (PyTorch's math path scales each by the square root of the scale) changes the inputs: a consumer that wants it passes the pre-scaled values with scale=1. PyTorch's math path multiplies the gradients of the pre-scaled Q and K by the square root of the scale at the end; with pre-scaled inputs and scale=1, the reference is the gradient before that rounded product. Every input and grad may also be an earlier result, as Result describes.

See also#