Attention

AttentionKernels computes O = softmax(Scale · Q Kᵀ + Mask) V for every batch entry and head in one kernel. The score matrix is never written to memory: a workgroup keeps a block of 64 query rows and walks over the keys one tile at a time, keeping for each row a running maximum, a running sum of exp(s − max) and the accumulated exp(s − max) · V, and rescales them whenever a tile raises the maximum (Dao et al., FlashAttention, 2022, for the tiling; Milakov and Gimelshein, Online normalizer calculation for softmax, 2018, for the normaliser).

using var attention = new AttentionKernels(device);
var p = new AttentionProblem(batch, heads, qLen, kLen, headDim)
{
    Q = GpuAttentionLayout.PositionsFirst(heads, qLen, headDim),   // [batch, qLen, heads, headDim], as a projection leaves it
    K = GpuAttentionLayout.PositionsFirst(heads, kLen, headDim),
    V = GpuAttentionLayout.PositionsFirst(heads, kLen, headDim),
    O = GpuAttentionLayout.PositionsFirst(heads, qLen, headDim),
    Mask = GpuAttentionLayout.ForMask([batch, 1, 1, kLen], batch, heads, qLen, kLen),   // optional additive mask
};
var encoder = device.BeginCommands();
attention.Encode(encoder, p, q, k, v, o, mask: maskBuffer, mode: GpuKernelMode.Deterministic);
device.Submit(encoder);

What it covers

Special values

• A key hidden by -∞ in the mask, by its batch entry's key length or by the causal rule has weight 0.
• A query row with no visible key gives NaN in every output (what 0 / 0 gives, and what a plain softmax over an all--∞ row gives), or zeros with MaskedRows = GpuAttentionMaskedRows.Zero (a row that attends to nothing contributes nothing). KLen = 0 is such a row.
• A row whose score is NaN or +∞ gives NaN. V is not inspected.
• In a browser, WGSL does not promise infinities and NaN behave, so use a large finite negative mask value there (-1e30); the tests do.

Variants and ordering

Keys are folded into each query row one at a time, in increasing key order, with the online normaliser, from key and value tiles held in workgroup memory. Two kernels share that structure (AttentionKernels.Variant, chosen once in the constructor):

• rows: one thread owns a query row (64 rows per workgroup). Needs nothing from the device.
• pairs: two neighbouring threads own a row, each the half of the head dimension, and add their halves of a dot product after one subgroup shuffle (32 rows per workgroup). This halves the registers per thread. Chosen automatically on Apple GPUs, which spill past about 128 registers per thread (a 64-wide head needs 128 for the query and the accumulator); on the RTX 4070 it took longer for large problems and less time for small ones, so Auto picks rows there. GpuAttentionVariant forces either; allowSubgroups: false rules out pairs.

The compiler keeps loops as loops and a thread-local array indexed by a loop variable can live in slow memory (Apple's compiler does that), so the SPIR-V and Metal builds write every access to the query and accumulator registers as a separate statement with a literal index (macros in Attention.hlsl). The WebGPU build for the software adapter (AttentionLoops.hlsl) keeps plain loops instead: its compiler does not finish compiling the unrolled form. A browser on a real GPU gets the unrolled form like the other targets (see In a browser below).

In a browser

Chrome's WGSL compiler did not keep a loop-indexed private array in registers: with plain loops the fused kernel was 3 to 18 times slower than GEMM + softmax + GEMM on an RTX 4070 (Chrome 155, wall time per call from a batch of calls closed by a one-element readback; indicative). Two changes closed the gap, and neither moves a bit on the other targets:

• A WebGPU device on a hardware adapter uses the unrolled statements (WGSL is generated from Attention.hlsl); a software adapter (SwiftShader) keeps AttentionLoops.hlsl.
• Key and value tiles are read as float4 (VEC4 build). WGSL checks every dynamic index of a workgroup array, so reading four floats per access also removes three of four checks, and the unrolled form's index arithmetic is the same either way. The sums keep the element-wise order (each group of four adds its products in increasing element order), so the results equal the scalar form's bit for bit. Apple GPUs ran the float4 form 17 to 25% slower (rows variant, M2) and keep the scalar one.

The 0.14 ms rows are the floor of one submit and readback in the browser, not kernel time. The same kernel in Vulkan on that card takes 1.09 ms for the Whisper encoder (GPU timestamps), so the browser form is within 1.8x of native (wall time, which includes the readback). AttentionBrowserTiming reruns the table in a browser on a hardware adapter.

Long calls

A single submission that keeps an Apple GPU busy for a long time (the one-thread-per-row kernel on a batch of 8 Whisper encoder inputs was the example) is stopped by the system and leaves the output half written. Encode therefore cuts a call whose estimated work exceeds the device's SubmissionWorkBudget (submissionWorkBudget in the constructor overrides it; 0 never cuts) into dispatches of whole workgroups and tells the encoder the work after each (ReportWork), so the encoder can end the submission between them. Every workgroup computes its own rows from the same inputs in the same order, so a cut call equals the whole call bit for bit; the tests compare them on every device. The estimate is 4 · batch · heads · QLen · KLen · HeadDim operations, times 1.5 for the matrix kernel (it multiplies Q Kᵀ twice) and times 4 for the scalar kernels (several times slower per operation, measured on the M2).

Accuracy at long key lengths

Adding one product per key into a single running sum lets the rounding error grow with the number of keys: at 4096 keys with moderately spread scores the plain sequential form was several times less accurate than GEMM + softmax + GEMM (whose sums are shorter or tree-shaped). For head dimensions 32 and 64 the scalar kernels therefore sum blockwise: the products of one tile of keys go into a second accumulator that starts at zero, and each finished tile is added to the main accumulator after one rescale by exp(old max - max). Every sum is over one tile (32 keys) or over the tiles, in a fixed order, so determinism is unchanged. The price is PD more registers per thread (a batch of 8 BERT-base inputs ran about 30% slower on the RTX 4070), so only workgroups that walk 640 keys or more use it (the cut ranges of a split call are usually shorter), and head dimensions 128 and 256 keep the sequential form (they spill already). Rounding the float32 scores themselves is a floor both paths share: with scores spread over a standard deviation of about 12 or more, both the fused and the unfused path sit near 1e-5 whatever the summation.

Fast and Deterministic run the same kernel: every sum has an order the code fixes, there are no atomics and no subgroup arithmetic, and the one shuffle adds two numbers (which does not depend on operand order). Two runs on one device and driver give identical bits. rows and pairs add a dot product in a different order, so they may differ from each other, and devices from each other, in the last bits.

When a problem has too few query blocks to fill a GPU, the keys are cut into at most 16 contiguous ranges of a multiple of 32 keys (at least 128 keys each, enough ranges for about 256 workgroups); each range stores its partial (maximum, sum, accumulator) in a scratch buffer and a second pass merges them in range order. The cut depends only on the shape. ScratchLength(problem) says how many floats are needed (0 without a cut); pass a buffer as scratch, or leave it null and the family keeps one of its own and releases it on Dispose.

Matrix hardware (cooperative-matrix)

For half and bfloat16 tensors with a head dimension up to 128, Auto picks a kernel that runs Q Kᵀ and P V as 16 x 16 x 16 cooperative matrices (VK_KHR_cooperative_matrix on Vulkan, simdgroup matrices on Metal), when the device lists such matrices of the element type with a float32 accumulator, runs 32-wide subgroups (pinned where the device can run several widths) with full subgroups, and offers the kernel's shader features. Everything else (float32, head dimensions 129 to 256, devices without matrix hardware such as lavapipe and the RADV iGPU, WebGPU) keeps rows or pairs; GpuAttentionVariant.Cooperative forces the matrix kernel (and throws where it cannot run) and allowMatrixHardware: false rules it out. Variant names the kernel used for half tensors, FloatVariant the one for float32.

A workgroup owns 64 query rows as four subgroups of 16 rows and visits the keys in tiles of 32 (16 for 128-wide heads), staged through workgroup memory (at most 12.8 KB, so it also fits a Metal limit that shader validation counts twice). The elements of a cooperative matrix have no specified layout, so an accumulator cannot be scaled per row without a round trip through workgroup memory; the kernel avoids the rescaling of the online normaliser instead. A first pass over the keys computes only the scores and finds the exact maximum of every row; a second pass recomputes them, forms p = exp(score − max) and accumulates P V in matrix registers. The price is a second Q Kᵀ, 1.5x the matrix work. The probabilities are rounded to the element type before the product with V, and the normaliser sums the rounded values, so numerator and denominator see the same p. Scores and probabilities pass between the matrix units and the lanes through workgroup memory, where lane l of a subgroup owns row l & 15 and every other column of each 16-wide block; the two lanes of a row meet once per pass in a shuffle that maximises or adds two numbers (operand order does not matter). The key-range split and the merge pass are those of the scalar kernels.

This kernel serves Fast and Deterministic alike: every order is fixed by the code (tiles in increasing key order, the hardware's matrix product depends only on its operands), there are no atomics and no subgroup arithmetic, and two runs on one device and driver give identical bits (checked per device by the tests). Devices differ from each other in the last bits, and so do the matrix and scalar kernels.

Apple simdgroup kernel (AttentionSimd.hlsl, Metal only)

On Apple GPUs every element type (float32 with head dimension up to 64, half, bfloat16; head dimensions 32, 64 and 128) runs a single-pass kernel on 8 x 8 simdgroup matrices, in place of pairs/rows (those stay as the fallback and for GpuAttentionVariant overrides). Variant and FloatVariant name the kernel. The scalar kernels were 10-20x slower than GEMM + softmax + GEMM here; this one is at or below the unfused chain at every reference shape (table below).

• Scores transposed. Sᵀ = K Qᵀ, so a row of the softmax is a column of the matrix and each lane owns the same row across blocks; no per-row scaling of an accumulator through memory is needed unless the maximum changes.
• Reference maximum. Each row keeps a reference m raised only when a block's maximum exceeds it by more than 8 (REF_SLACK): exp(score − m) is then at most e⁸, far from overflow, and the accumulators are rescaled by exp(old − new) through workgroup memory only on a raise (the decision is made per workgroup from flags, so the barriers stay uniform). The first block sets m with no rescale. Whether a rescale happens depends only on the data, never on timing.
• Query-direct variant. For 16 or more queries the 4 simdgroups of a workgroup (64 rows) reload their Q fragments from the tensor at each step instead of keeping them in registers (the last 10-15 %); the last 16-row block slides back so no row is written twice. Decoder steps and very short queries keep Q in registers, one simdgroup per 16 rows.
• K and V fragments load straight from the buffers when the key block is full and the head dimension is a multiple of 16; otherwise through a staged tile.

Indicative M2 wall time, ms, validation off, 12 heads of 64 (Whisper encoder) or BERT/MiniLM heads: unfused f32 GEMM + softmax + GEMM against the fused kernel before (pairs) and after.

Accuracy (scaled, double reference): float32 1.5e-6 on the model shapes and 4.5e-6 on scores that keep rising or are wide (the same as the scalar kernels there), half 5.0e-4, bfloat16 4.4e-3; fused against unfused float32 at most 9.7e-7. Two runs are bit-identical. As with the other kernels, Fast and Deterministic run the same code. The 8x8 product is a fixed function of its operands and each output tile is summed by one simdgroup in key order, so it counts as a fixed order; the scores and the O rescale use no atomics or subgroup arithmetic.

Tolerances

Scaled error (max |gpu − reference| / max |reference| per output, reference in double on the CPU, inputs rounded to half first for half rows): float32 ≤ 1e-5, half with float32 accumulation ≤ 2e-3, bfloat16 with float32 accumulation ≤ 8e-3 (the probabilities and the output are rounded to bfloat16's eight bits of precision, which costs about 4e-3). The tests assert these on the shapes of the reference models (MiniLM, BERT-base, Whisper-small encoder, decoder self- and cross-attention) and on tile edges, remainders, size 1, non-power-of-two sizes, masks with -∞, fully masked rows, key lengths of 0, grouped heads and odd head dimensions; they also compare the fused kernel with GEMM + add + softmax + GEMM.