Algorithms#

An operation's own tolerance covers every implementation that rounds only where the operation requires. A named algorithm describes one specific way of computing it: a fixed order of additions, an accumulator format, or the extra roundings of a fused kernel. Use one when the definition's tolerance does not describe your implementation, or to compare ways of computing the same thing. The reference does not change: an algorithm's bounds measure the distance from the operation's rounded.

matmul algorithms#

The matmul algorithms are accumulation orders, each in a format, with or without fused multiply-add. They live in er.linalg.products:

Constructor

How the products of one output element are added

matmul_sequential(format, *, fma=True)

one at a time, in index order

matmul_pairwise(format, *, fma=True)

as a balanced binary tree

matmul_blocked(format, block_size, *, fma=True)

in order within blocks of block_size, then the block sums in order

matmul_exact_accumulation(format)

exactly, then rounded once to format

matmul_slots16(format)

rounded, then summed in 16 slots by magnitude

With fma=True the exact products are added (a fused multiply-add); with fma=False each product is first rounded to format. matmul_slots16 adds each slot's values in increasing magnitude with alternating signs, then the slots from the smallest up. Each constructor also takes a rounding for its accumulator: to nearest, NEAREST_EVEN (the default) or NEAREST_AWAY. A directed rounding raises ValueError: past its largest value such an accumulator saturates, which can move the result by up to the size of the terms, and no tolerance bounds that usefully.

Each algorithm has:

  • run(A, B): what that implementation computes, emulated exactly.

  • tolerance(A, B), finite_range(A, B), may_be(kind, A, B) and covers(output, A, B): its bounds, as a result's methods describe them, with dtype, rounding, standard and flush_to_zero options.

  • error_model(k): its error formula for k products, which depends only on shapes and formats, for comparing algorithms.

Adding 300 ones in bfloat16 shows how much the order can matter:

Code
import exactref as er

# The module that holds the matmul algorithms
products = er.linalg.products
# A 1x300 row of bfloat16 ones, and the same ones as a 300x1 column
ones = er.asarray([[1.0] * 300], "bfloat16")
column = ones.reshape((300, 1))

# The reference: the exact sum of the 300 ones, rounded once to bfloat16
print("reference:", er.matmul(ones, column).rounded.tolist())

# Products rounded to bfloat16, then added one at a time in bfloat16
sequential_algorithm = products.matmul_sequential("bfloat16", fma=False)
# run() gives what that implementation computes, emulated exactly
print("sequential:", sequential_algorithm.run(ones, column).tolist())

# The same, but added as a balanced tree
pairwise_algorithm = products.matmul_pairwise("bfloat16", fma=False)
print("pairwise:", pairwise_algorithm.run(ones, column).tolist())
Output
reference: [[300.0]]
sequential: [[256.0]]
pairwise: [[300.0]]

Added one at a time in bfloat16, the sum gets stuck at 256, where bfloat16 values are 2 apart and adding 1 rounds back down. Added as a tree, every partial sum fits, and the answer is exact.

The error models compare algorithms without any data:

Code
# A bound on the error of a pairwise float32 sum of 64 rounded products
print("pairwise:", products.matmul_pairwise("float32", fma=False).error_model(64))
Output
pairwise: 4.77e-07 * sum over k of |A[i,k] * B[k,j]| + 1.34e-43

The order-only versions of these methods, for any list of terms, are in exactref.accumulation: exact, sequential, pairwise and blocked, each with the number of roundings a term passes through and the standard worst-case bound built on it (Higham, Accuracy and Stability of Numerical Algorithms, chapters 3 and 4).

Attention algorithms#

Fused attention kernels round in places the definition does not. An er.nn.attention.AttentionAlgorithm describes those places by method, and gives the bounds of any kernel that rounds there:

er.nn.attention.AttentionAlgorithm(
    *, online_softmax=False, block_size=None, exp2=False, weights=None,
    accumulate=None, bias=None, reciprocal=False,
)

Method

What the kernel does

online_softmax=True

processes the keys in blocks with a running maximum and sum, rescaling earlier partial sums

block_size=n

with online_softmax: at most ceil(S / n) blocks of at most n keys, in any order (without it, any blocking)

exp2=True

computes the weights as exp2 of logits in base-2 units

weights="bfloat16"

rounds the weights to this format before the product with V (in the backward, also P and dS)

accumulate="bfloat16"

does every accumulation in this format instead of the compute format

bias="bfloat16"

rounds float-mask values to this format before adding them, as a kernel computing ALiBi biases in lower precision does

reciprocal=True

normalizes by a rounded reciprocal of the row sum instead of dividing

PyTorch's CPU softmax and its CPU flash kernel normalize by a reciprocal; that kernel's logsumexp for bfloat16 and float16 inputs also needs ulps=8. Any format can stand in for "bfloat16" above.

The methods combine. With none, the bounds are the definition's. With online_softmax, both standards give the any-order bound.

An attention algorithm has tolerance(), finite_range(), may_be() and covers(), taking scaled_dot_product_attention's arguments, and their backward_ forms for the gradients. It has no run(), because kernels differ in details these methods leave open, and no error model, because attention's error grows with the size of the logits, which no formula in shapes and formats alone can bound.

Code
import math
import random
import torch

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

# Like torch.randn, but drawing from rng
def randn(*shape):
    # Normally distributed, mean 0 and standard deviation 1, one per element
    values = [rng.gauss(0, 1) for _ in range(math.prod(shape))]
    return torch.tensor(values).reshape(shape)

# Query, key and value of shape (batch, heads, positions, features)
q, k, v = randn(1, 2, 64, 16), randn(1, 2, 64, 16), randn(1, 2, 64, 16)

# No methods: the definition's bounds
plain = er.nn.attention.AttentionAlgorithm()
# Keys in blocks of at most 16 with a running maximum and sum; bfloat16 weights
flash = er.nn.attention.AttentionAlgorithm(
    online_softmax=True, block_size=16, weights="bfloat16")

# Each algorithm's largest tolerance on the attention output, to 3 significant digits
for name, alg in (("plain", plain), ("flash", flash)):
    print(f"{name}: {max(alg.tolerance(q, k, v).output.scalars()):.3g}")
Output
plain: 1.31e-05
flash: 0.00545

Reading complicated tolerances checks an emulated kernel with covers().

For the backward, backward_tolerance(q, k, v, grad, out, logsumexp) describes a fused backward kernel that is given the forward's out and logsumexp: it recomputes P from logsumexp and takes D from out. Without them, backward_tolerance(q, k, v, grad) describes a backward that recomputes the softmax, as PyTorch's math path does. out and logsumexp may be the forward result's own outputs, fwd["output"] and fwd["logsumexp"].

Scaling Q and K before the product is not a method: it changes the inputs. To check such a kernel, pass the pre-scaled values with scale=1.