Attention benchmark

The baseline

Materialised attention is three lines and one allocation per head:

s = (q @ k.transpose(-2, -1)) * sm_scale
p = torch.softmax(s, dim=-1)
o = p @ v

That is the baseline for a claim about removing the score matrix, because it is the code the kernel replaces. torch.nn.functional.scaled_dot_product_attention is a different comparison: it is already fused, so beating it is a stronger claim than beating the three lines above, and a table that does not say which baseline it used is not a table.

The sweep

Sequence length and head dimension are the two axes that move the answer. Sweep N_CTX geometrically, sweep HEAD_DIM over the sizes your model uses, and separate causal from non-causal with one dtype per table. The baseline’s traffic grows with N while the fused kernel’s does not, so the ratio is expected to grow with sequence length. That shape is the result; a single cell is an anecdote.

Report

Time both paths in one session, interleaved, with warmup discarded, and print <ms> and <speedup> per shape with the band behind each. Record the baseline’s peak allocation beside its time. At the longest lengths the baseline may fail to allocate, and a row saying so at <N_CTX> is a result: it is the IO argument arriving as a measurement.

Things that break

Beating an already-fused baseline and calling the win FlashAttention, when the claim was about materialising scores. Timing a kernel whose output nobody consumes, so the launch is measured rather than the work. Comparing across dtypes, or changing baseline between rows, which produces a ratio that describes the table. And reporting <speedup> with no band, when the band is what tells a reader whether the gap exceeds the noise. The harness follows Triton 3.9 do_bench, which returns milliseconds; confirm the return units against your installed version before a figure reaches the table.