Tiled attention cuts costs of relative-position terms
Fast and Exact Relative-Position Attention via Structure-Aware Kernels · arXiv:2611.01182
In plain words
Transformer attention can use learned information about how far apart two words are, but adding that information can require a large table or repeated calculations. This paper introduces an exact way to calculate it in small blocks, reusing work and organizing gradient sums.
In tests of four attention types on an H200 GPU, it was faster and used less peak memory than the fastest production implementation for each type at 8K tokens. Those measurements cover an attention layer, not a whole model.
The research
Long sequences make it costly to preserve learned information about token distance in transformer attention.
The results change the expectation that these terms must bring a large memory or runtime penalty in an attention layer, though the measured advantage depends on the relation map and does not describe a whole model.
An exact tiled algorithm reuses relative-position calculations and organizes gradient sums, reducing attention-layer runtime and peak memory in long-sequence tests.
The algorithm projects queries and keys only onto the relative-position entries their attention tiles can reach, reuses those entries, and sums gradients by offset before reducing them to relation rows.
The authors compared its attention-layer speed and memory with production implementations across four mechanisms.
At 8K tokens, it was faster than the fastest production implementation for each mechanism and used less peak memory than the leanest. An independent FP32 reference check found its outputs and gradients met the paper’s numerical criteria.
The performance results come from attention-block tests on an H200 GPU, not every full model workload. Identity maps still require quadratic projection work and storage, and some other relation maps use a slower, nondeterministic atomic fallback.
What the evidence shows
- The cost of token distance
The paper presents an exact tiled algorithm for additive relative-position attention, where learned terms based on token distance adjust attention scores. Existing implementations can materialize a dense pairwise table, repeat calculations or lookups, and contend when accumulating gradients.
The proposed approach targets those costs while keeping the relative-position terms in the attention calculation.
- Faster attention blocks
Comparing each mechanism with its fastest production implementation on the same H200 attention layer at 8K tokens shows that the speed advantage holds across all four tested mechanisms, though its size varies.
Transformer-XL gains less because its identity map keeps projection work quadratic in sequence length.
- Lower peak memory
Comparing peak device memory beyond the inputs with each mechanism’s leanest production implementation shows that the method uses less memory in both passes across the four tested mechanisms. The size of the savings varies by mechanism and pass, and these are measurements of one attention layer.
- Beyond the shared component
Because both conditions use the same relation decomposition, the remaining speed difference in these 8K-token attention-block tests indicates that the decomposition alone does not account for the complete method’s measured performance.
This comparison is limited to programmable engines and does not establish that all other implementation details were held constant.
- Accuracy and remaining limits
The method’s outputs and gradients passed the authors’ numerical check against an independent FP32 reference, using stated elementwise or relative L2 error limits. That check supports the interpretation that the speedups accompany numerically acceptable results under the paper’s criteria.
The conclusion remains bounded: identity maps retain quadratic projection costs, and some maps require a slower, nondeterministic fallback.
Deep
- Original paper authors
- Delyan Boychev, Ivo Petrov, Dimitar I. Dimitrov
- Submitted to arXiv
- License
- arXiv non-exclusive