Fast and Exact Relative-Position Attention via Structure-Aware Kernels

Delyan Boychev; Ivo Petrov; Dimitar I. Dimitrov

cs.LGOct 31, 2026

Fast and Exact Relative-Position Attention via Structure-Aware Kernels

Authors: Delyan Boychev, Ivo Petrov, Dimitar I. Dimitrov

Organizations: INSAIT, Sofia University “St. Kliment Ohridski”.

Abstract

Relative-position attention mechanisms offer compelling alternatives to RoPE across various transformer architectures, delivering superior performance in domains such as audio. These mechanisms typically modify attention scores by adding content-position or bias-position terms. Unlike RoPE, which encodes position via per-token rotations directly on the queries and keys, relative-position terms generally require pairwise offset-dependent computation or lookups, complicating their integration into fused attention kernels and introducing severe memory, bandwidth, and compute overheads as sequences scale. To address this bottleneck, we propose an exact attention algorithm that incorporates additive relative-position terms into FlashAttention-style tiled computation. Our method (1) uses offset-aware tiling to avoid redundant indexing and memory accesses, (2) computes only the necessary key/query-relation products and reuses them across token pairs in both forward and backward passes, and (3) uses offset-based gradient sharding and aggregation, eliminating atomic contention. We provide a unified, efficient Triton implementation across diverse relative-position attention mechanisms and evaluate it against naive and specialized implementations. For 8K-token sequences, our method speeds up the forward (backward) pass by 1.8-6.4 $\times$ (2.0-21 $\times$ ) and reduces peak memory up to 2.9-128 $\times$ (1.4-54 $\times$ ) over the strongest baseline for each mechanism, enabling long-context end-to-end transcription that is 1.8 $\times$ faster and uses 2.4 $\times$ less memory. Our method thus provides a drop-in, hardware-efficient solution for scaling relative-position attention transformers.

Figures & tables

Figure 1: Overview of our method . (a) Each score adds the relation terms of its row $r=\phi(i,j)$ , which stays constant on tiles far from the diagonal. (b) Our method first computes each query group $g$ only on its reachable rows $\Omega^{Q}_{g}$ . (c) Each shared entry $Z^{Q}_{i,r}$ is loaded once and reused across key tiles. (d) One thread block per query tile sums gradients on chip, avoiding atomic contention.

	Shaw	TXL	DeBERTa	T5
General (dense bias)
PyTorch SDPA	F+B	F+B	–	F+B
FlexAttention ( Dong et al., 2025 )	F+B	F+B	F+B †	F+B
FlashAttention-4 ( Zadouri et al., 2026 , FA4;)	F	F	F	F
FlashInfer ( Ye et al., 2025 )	F	F	F	F
Specialized

Table 1: The systems we compare against and their support for each relative-position mechanism. “F” means forward, “B” backward, “†” the backward pass fails our FP32 check from 4K tokens, “–” unsupported.

Figure 2: Attention throughput on an H200 (FP16). PyTorch, general, and decomposition bars show the fastest system of their kind; labels give the speedup of our method over the fastest system without our decomposition.

	Shaw		TXL		DeBERTa		T5
	F	B	F	B	F	B	F	B
vs. our decomposition (H200)
Speedup	1.3	3.5	1.1	1.7	2.2	4.3	1.5	228
Memory saving	0.99	2.9	1.0	3.0	0.51	2.2	1.0	0.75
vs. production systems: speedup per GPU
H200 (9.0)	6.4	21	1.8	2.0	3.7	7.1	4.6	13

Table 2: Improvements of our method at 8K tokens. Top: against implementations with our decomposition (H200). Bottom: speedup (

\times

) over the fastest production system. “F”/“B”: forward/backward.

		Shaw		TXL		DeBERTa		T5
	MiB	128	8K	128	8K	128	8K	128	8K
Forward	PyTorch eager	192	14K	288	18K	216	12K	144	7.5K
	Leanest baseline	98	14K	168	9.0K	48	48	24	3.0K
	+ our decomp.	120	221	120	3.1K	216	408	24	24
	Ours	120	222	120	3.1K	216	804	24	24
Backward	PyTorch eager	168	19K	264	12K	288	9.0K	168	9.0K

Table 3: Peak memory (MiB) on the H200 (FP16) at 128 and 8K.

	Segments			Whole talks
Attention	WER	RTFx	GiB	WER	RTFx	GiB
NeMo (native)	2.80	2726	4.7	2.88	492	53
NeMo (SDPA)	2.82	3761	4.7	3.01	613	27
Ours	2.80	4560	4.7	2.89	1103	11

Table 4: Parakeet-TDT-0.6B-v3 transcription of TED-LIUM 3’s test set in NeMo, H200, BF16, as segments and as whole talks.

	Step (ms)		Peak (GiB)		WER (%)
$W$	SDPA	Ours	SDPA	Ours	SDPA	Ours
30 s	281	257 (1.09 $\times$ )	44	43	3.73	3.82
2 min	388	325 (1.19 $\times$ )	49	49	3.20	3.04
5 min	506	390 (1.30 $\times$ )	60	60	3.35	3.24
9 min	740	546 (1.36 $\times$ )	79	79	3.73	4.34
15 min	oom	876	oom	99	oom	4.44

Table 5: Fine-tuning Canary-1B-v2 on windows of up to

W

. oom denotes when training ran out of memory.

	Accuracy (%)		Step (ms)		Peak (GiB)
$L$	HF	Ours	HF	Ours	HF	Ours
512	66.6	66.7	116	105 (1.10 $\times$ )	8.2	6.1
2K	67.3	67.6	942	430 (2.19 $\times$ )	15	7.1
4K	67.7	67.6	2185	915 (2.39 $\times$ )	27	7.1

Table 6: Fine-tuning DeBERTa-v3-base with content truncated to

L

tokens. HF denotes Transformers’ performance.

w/o	Latency ( $\times$ )		Peak memory ( $\times$ )
(128–8K)	Fwd	Bwd	Fwd	Bwd
Atomic-free gradients	–	2.9–4.4	–	1.8–2.6
Projections	2.0–2.9	3.5–5.4	0.19–0.43	0.20–0.43
Reachable ranges	1.0–1.1	1.0–1.1	1.1–1.5	1.1–1.5
Constant-range reuse	1.1–1.3	1.2–1.3	1.0	1.0
Bucket copies	1.0–1.1	1.1	0.70–1.0	0.81–1.0

Table 7: Ablations on an H200 (FP16): latency and peak memory without each method, relative to our method , as the range over the mechanisms it applies to. “–”: not applicable.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning
Tiled Attention
$N,d$	Sequence length and head dimension.
$B,H$	Batch size and number of heads.
$i,j$	Query and key positions in $\{0,\ldots,N-1\}$ .
$Q,K,V,O$	Query, key, value, and output matrices in $\mathbb{R}^{N\times d}$ , with rows $q_{i}^{\top},k_{j}^{\top},v_{j}^{\top},o_{i}^{\top}$ .
$S,P$	Complete softmax logits and row-normalized probabilities in $\mathbb{R}^{N\times N}$ .

Appendix

Table 8: Notation.

Kernel	Lines	Parallel over	Writes
Forward ( Algorithm 1 )
Projection	4	tokens, relation rows	$Z^{Q},Z^{K}$
Attention	5 – 21	queries	$O,L$
Backward ( Algorithm 2 )
Preprocessing	1	queries	$D$
Query-parallel pass	2 – 12	queries	$dQ,d\tilde{Z}^{Q}$

Appendix

Table 9: Kernels of our method in launch order, parallelized over heads.

		Shaw		TXL		DeBERTa		T5
w/o	Pass	Lat.	Mem.	Lat.	Mem.	Lat.	Mem.	Lat.	Mem.
Atomic-free gradients	Bwd	3.92	1.75	2.94	2.60	3.99	2.16	4.39	1.97
Projections	Fwd	2.47	0.43	1.99	0.20	2.91	0.19	–	–
	Bwd	4.98	0.43	3.48	0.20	5.41	0.22	–	–
Reachable ranges	Fwd	–	–	1.12	1.48	1.01	1.08	–	–
	Bwd	–	–	1.09	1.50	1.01	1.08	–	–

Appendix

Table 10: Ablations without each method, relative to our method : (a) latency and peak memory on an H200 (FP16), as the geometric mean over 128 to 8K tokens, and (b) HBM traffic (bytes read and written) on an H100 at 4K tokens. Retained projections span 4 layers. Without projections, each tile computes its own relation products, timed at its fastest tile configuration. –: not applicable.

		Forward (TFLOP/s $\uparrow$ )							Backward (TFLOP/s $\uparrow$ )
	System	128	256	512	1K	2K	4K	8K	128	256	512	1K	2K	4K	8K
Shaw	PyTorch eager	27	37	31	26	18	12	7.1	45	64	57	41	27	15	8.0
	PyTorch compile	72	89	84	82	62	57	44	64	89	19	4.6	1.7	0.74	0.35
	PyTorch SDPA	67	89	68	49	30	17	8.9	67	74	60	41	26	15	7.8
	FlexAttention	54	66	51	37	25	15	8.3	61	78	64	44	27	15	7.9
	FA4	45	53	45	33	23	14	8.1	–	–	–	–	–	–	–

Appendix

Table 11: Throughput of every system on the H200 (FP16) at every length. Engines running our decomposition use their faster variant per mechanism and pass. Bold : the fastest system. †: fails the FP32 check. –: not supported.

		Forward (MiB $\downarrow$ )							Backward (MiB $\downarrow$ )
	System	128	256	512	1K	2K	4K	8K	128	256	512	1K	2K	4K	8K
Shaw	PyTorch eager	192	336	624	1.2K	2.3K	5.1K	14K	168	312	600	1.1K	2.3K	4.7K	19K
	PyTorch compile	122	224	440	920	2.0K	5.0K	14K	144	193	289	512	1.3K	3.6K	11K
	PyTorch SDPA	98	200	418	904	2.0K	5.1K	14K	169	264	456	840	1.6K	4.7K	19K
	FlexAttention	98	200	418	904	2.0K	5.1K	14K	216	360	648	1.2K	2.3K	4.7K	19K
	FA4	98	200	418	904	2.0K	5.1K	14K	–	–	–	–	–	–	–

Appendix

Table 12: Peak memory of every system on the H200 (FP16) at every length. Engines running our decomposition use their faster variant per mechanism and pass. Bold : the leanest system. †: fails the FP32 check. –: not supported.

		Forward (TFLOP/s $\uparrow$ )							Backward (TFLOP/s $\uparrow$ )
	System	128	256	512	1K	2K	4K	8K	128	256	512	1K	2K	4K	8K
Shaw	PyTorch eager	23	23	18	15	12	8.6	5.8	32	36	27	22	16	11	6.6
	PyTorch compile	51	49	37	31	27	23	17	49	57	17	5.1	2.0	0.88	0.41
	PyTorch SDPA	53	57	42	30	21	13	7.5	51	61	45	33	22	13	7.4
	FlexAttention	47	46	33	25	18	12	7.1	39	44	32	24	18	11	6.8
	FA4	44	38	26	20	15	10	6.5	–	–	–	–	–	–	–

Appendix

Table 13: Throughput of every system on the RTX PRO 6000 Blackwell (FP16) at every length. Engines running our decomposition use their faster variant per mechanism and pass. Bold : the fastest system. †: fails the FP32 check. –: not supported.

		Forward (TFLOP/s $\uparrow$ )							Backward (TFLOP/s $\uparrow$ )
	System	128	256	512	1K	2K	4K	8K	128	256	512	1K	2K	4K	8K
Shaw	PyTorch eager	7.5	9.0	7.5	6.2	5.3	4.1	3.0	11	13	11	8.0	5.4	3.3	1.8
	PyTorch compile	12	16	14	12	11	8.9	6.7	16	21	11	4.0	1.7	0.75	0.35
	PyTorch SDPA	17	20	16	13	10	6.9	4.3	14	16	12	9.1	6.1	3.6	2.0
	FlexAttention	14	17	14	11	8.8	6.2	4.0	13	16	12	8.6	5.5	3.3	1.8
	FlashInfer	13	15	11	8.6	7.1	5.3	3.6	–	–	–	–	–	–	–

Appendix

Table 14: Throughput of every system on the RTX A6000 (FP16) at every length. Engines running our decomposition use their faster variant per mechanism and pass. Bold : the fastest system. †: fails the FP32 check. –: not supported.

Efficient Attention

Efficient attention reduces the compute, memory traffic, and cache costs of Transformer attention while preserving useful long-range interactions.

See all 148 papers

Aug 9, 2026

Position Encoding in Transformers: From Absolute and Relative Methods to Rotary Position Embeddings and Long-Context Scaling

Self-attention models content-dependent interactions between tokens but does not by itself encode token order. Position encoding addresses this limitation by introducing absolute coordinates, relative distances, or position-dependent rotations into Transformer representations and attention scores. This technical survey develops a unified account of sinusoidal and learned absolute position embeddings, Shaw-style relative position representations, Transformer-XL, T5 relative position bias, ALiBi, and Rotary Position Embeddings (RoPE). We derive how RoPE converts absolute position indices into relative phase differences in Query-Key inner products and compare these methods in terms of where position is injected, computational cost, compatibility with KV caching, and length extrapolation. We then examine long-context extensions, including Position Interpolation, RoPE scaling laws, NTK-aware scaling, Dynamic NTK, NTK-by-parts, YaRN, LongRoPE, and LongRoPE2, with emphasis on frequency allocation, attention rescaling, training length, and target context length. We also summarize implementation considerations, evaluation protocols, and position-encoding choices in representative large language models. A central conclusion is that the ability to compute positional features beyond the training length does not imply reliable long-context generalization; context extension must be evaluated through short-context retention, position-wise perplexity, retrieval, reasoning, and long-context code tasks.

Jiguo Li

Rotary Positional Embeddings Long-Context Modeling Positional Encoding LLM Context Management Relative Positional Encoding

Ask PDF

Oct 3, 2026

WavePrune: One period is often enough for RoPE

Rotary Position Embedding (RoPE) encodes token positions by rotating each two-dimensional channel of the query and key vectors at a channel-specific frequency, making the attention logits invariant to a common shift of positions. However, this rotation is periodic, and it leads to position aliasing where relative positions separated by a full rotation period become hard to tell apart. To address this, we propose WavePrune, which restricts each channel to its first rotation period. We show that it removes the distractions in attention maps created by position aliasing and improves overall long-context performance. Specifically, WavePrune raises the HELMET score on four of five models we test without any extra tuning (e.g., 35.7 -> 40.0 on Qwen3-8B). When pretraining models from scratch, WavePrune also achieves lower validation loss at extrapolated lengths than pretraining without it. Because WavePrune restricts each channel to a sliding window, it induces a fine-grained sparsity that our hardware-aligned CUDA kernels exploit for 1.15x prefill and 1.24x decoding speedups over FlashAttention-2 at 32K context. Together, these results show that RoPE's periodic structure, widely regarded as essential, is largely redundant beyond the first rotation period.

Guancheng Du, Luotian Huang, Shaowen Wang +2

Tsinghua University

Rotary Positional Embeddings Positional Encoding Relative Positional Encoding Long-Context Language Modeling LLM Inference Acceleration

Ask PDF

Jun 4, 2026

Where does Absolute Position come from in decoder-only Transformers?

RoPE-trained transformers distinguish absolute position in their attention patterns, even though RoPE encodes only relative offsets in the inner product. We trace this leakage to two architectural components, The causal mask is responsible for the first: its per-query softmax denominator depends on the absolute query position by construction. The residual stream supplies the second. Under causal attention the activation at position

0

attends only to itself and runs as a closed dynamical system from the embedding of the token at that position; downstream attention reads this trajectory through sink-reading heads. Both components appear in all three architectures we study, in architecturally specific balance: NTK scaling suppresses the residual-stream component, sliding-window attention allows it to accumulate with depth, and standard RoPE sits between. Replacing the BOS embedding before the forward pass removes

40\%

of the residual-stream component at early queries. Attention sinks are token-anchored stabilizers that pass forward a deterministic fingerprint of the token at position

0

, constant across inputs when that token is the auto-prepended BOS and varying with it otherwise.

Valeria Ruscio, Umberto Nanni, Fabrizio Silvestri

Intuition Machines · Sapienza University of Rome

Rotary Positional Embeddings Transformer Interpretability Decoder-Only Language Models Attention Sinks Position Bias

Ask PDF

Abstract

Efficient Attention

Explore similar work

Position Encoding in Transformers: From Absolute and Relative Methods to Rotary Position Embeddings and Long-Context Scaling

WavePrune: One period is often enough for RoPE

Where does Absolute Position come from in decoder-only Transformers?