结构感知内核实现精确相对位置注意力,长序列提速最高6.4倍
Fast and Exact Relative-Position Attention via Structure-Aware Kernels
把相对位置注意力塞进 FlashAttention 式内核,8K 序列反向最快快 21 倍,还省了 128 倍内存,做音频或长序列 transformer 的可以看看。
论文提出一种将加性相对位置项纳入 FlashAttention 式分块计算的精确注意力算法。方法通过 offset 感知分块、只计算必要的 key/query 关系积并跨 token 对复用、基于 offset 的梯度分片聚合来消除原子竞争。统一 Triton 实现覆盖多种相对位置注意力机制,在 8K token 序列上前向提速 1.8-6.4 倍、反向提速 2.0-21 倍,峰值内存最多降低 128 倍。端到端长上下文语音转录因此快 1.8 倍、省 2.4 倍内存。
Fast and Exact Relative-Position Attention via Structure-Aware Kernels
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.