arXiv preprint ·
Fast and Exact Relative-Position Attention via Structure-Aware Kernels
- Published
- Source
- arXiv preprint, posted before peer review
- Original
- Read it at arXiv
- Topics
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.
The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.
Read it on ReadSearch
Each opens after you create an account, and brings you back to this paper.
More in Machine Learning
Robustness Starts with the Objective: Designing and Certifying Robust LLM-Generated Heuristics
Token-as-a-Utility: A Model-Agnostic Architecture for Metering and Provisioning AI Inference Services
Weights That Stay: On-Chip Megakernels for Serial Heads in Speculative Decoding
Serving Agentic Robot Fleets: Workload Characterization, KV-Cache Thrashing, and Admission Control
Running a 20B Video Diffusion Transformer on Consumer Hardware: A Deployment and Measurement Report for MiniMax H3 on Apple Silicon
Get the few papers worth your time in Machine Learning every morning.
ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.