Check fused attention kernels#

The problem#

You are checking a FlashAttention-style kernel. It processes the keys in blocks with a running softmax, rounds the weights to bfloat16 before the product with V, and maybe uses exp2 or a reciprocal of the row sum. Its output differs from PyTorch's math path by more than a float32 tolerance, and you cannot tell whether the difference is the kernel's method or a bug.

Why the usual approach falls short#

Each of those methods rounds in a place the definition of attention does not. A tolerance for the definition rejects a correct kernel; a tolerance loosened until it passes accepts real bugs. What you need is the bound for exactly the kernel's methods.

The solution#

Describe the kernel's methods with an AttentionAlgorithm and check its outputs against it. Here a kernel is emulated in PyTorch: blocks of 4 keys, a running maximum and sum, and weights rounded to bfloat16.

Code
import math
import random
import torch
import exactref as er

# 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 value per element: normal, with mean 0 and standard deviation 1
    values = [rng.gauss(0, 1) for _ in range(math.prod(shape))]
    # torch.tensor makes float32, PyTorch's default for Python floats
    return torch.tensor(values).reshape(shape)

# The kernel, for batch 0 and head 0, one query row at a time
def flash_attention(q, k, v, block=4):
    # Attention's default scale, 1/sqrt(E) for queries of E values
    scale = 1 / math.sqrt(q.shape[-1])
    # The outputs: one row per query, and one logsumexp per query
    out, lse = torch.zeros(q.shape), torch.zeros(q.shape[:-1])
    for i in range(q.shape[-2]):
        # The running maximum score and running sum of weights for query i
        running_max, row_sum = torch.tensor(float("-inf")), torch.tensor(0.0)
        # The running weighted sum of value rows
        acc = torch.zeros(v.shape[-1])
        for s in range(0, k.shape[-2], block):
            # Scaled scores of query i against the keys of this block
            scores = (k[0, 0, s:s + block] @ q[0, 0, i]) * scale
            # The maximum score so far, this block included
            new_max = torch.maximum(running_max, scores.max())
            # This block's weights, relative to the new maximum
            weights = torch.exp(scores - new_max)
            # The factor that moves the earlier sums to the new maximum
            rescale = torch.exp(running_max - new_max)
            # Rescale the earlier sum, then add this block's weights
            row_sum = row_sum * rescale + weights.sum()
            # The same for the value rows, with the weights rounded to bfloat16
            # before the product with V
            acc = acc * rescale + weights.bfloat16().float() @ v[0, 0, s:s + block]
            running_max = new_max
        # Divide by the row sum, so the weights add up to 1
        out[0, 0, i] = acc / row_sum
        # The log of the sum of exp(score) over the row
        lse[0, 0, i] = running_max + torch.log(row_sum)
    return out, lse

# Query, key and value: one batch, one head, 8 rows of 16 values each
q, k, v = randn(1, 1, 8, 16), randn(1, 1, 8, 16), randn(1, 1, 8, 16)
# Run the kernel
out, lse = flash_attention(q, k, v)

The definition's bounds reject the output, because of the bfloat16 weights. The algorithm with the kernel's methods accepts both outputs:

Code
# The definition's reference and bounds for the same inputs
definition = er.scaled_dot_product_attention(q, k, v)
print("definition covers output:", definition["output"].covers(out).all())

# The kernel's methods: a running softmax over blocks of 4 keys, bfloat16 weights
flash = er.nn.attention.AttentionAlgorithm(
    online_softmax=True, block_size=4, weights="bfloat16")
# Check both outputs; covers() returns one bool array per name in the dict
checked = flash.covers({"output": out, "logsumexp": lse}, q, k, v)
print("algorithm covers output:", checked.output.all())
print("algorithm covers logsumexp:", checked.logsumexp.all())
Output
definition covers output: False
algorithm covers output: True
algorithm covers logsumexp: True

Leaving out a method the kernel uses fails it, so the check also confirms which methods the kernel really has:

Code
# The same methods without the bfloat16 weights
no_bf16 = er.nn.attention.AttentionAlgorithm(online_softmax=True, block_size=4)
# Check the output alone
print("output covered:", no_bf16.covers({"output": out}, q, k, v).output.all())
Output
output covered: False

For a fused backward that is given the forward's out and logsumexp, use backward_covers(gradients, q, k, v, grad, out, logsumexp); without them, it describes a backward that recomputes the softmax.