Abstract
Softmax Attention has a quadratic time complexity, which becomes prohibitive to run at long contexts, even with highly optimized GPU kernels. For example, FlashAttention (an exact, GPU-optimized implementation of Softmax Attention) cannot complete a single forward-backward pass of a multi-head attention layer once the context exceeds ~4 million tokens on an NVIDIA GH200 (96 GB). We introduce RACE Attention, a kernel-inspired alternative to Softmax Attention that is linear in sequence length and embedding dimension. RACE Attention replaces the exponential kernel with a sharpened angular (cosine) similarity, and approximates attention outputs via randomized projections and soft Locality-Sensitive Hashing (LSH). Across language modeling, masked language modeling, and text classification, RACE Attention matches the accuracy of strong baselines while reducing runtime and memory. In a controlled scale test, it processes up to 12 million tokens during a single forward-backward pass on an NVIDIA GH200 GPU and 75 million tokens on an Intel Xeon Gold 5220R CPU, well beyond the practical limits of the current state-of-the-art attention implementations. RACE Attention thus offers a practical, theoretically grounded mechanism for outrageously long context windows on today's hardware. We hope that it gets adopted in practice.
Abstract (translated)
Softmax注意力机制的时间复杂度为二次方，这在处理长上下文时变得非常不可行，即使使用了高度优化的GPU内核也是如此。例如，FlashAttention（一种精确且经过GPU优化的Softmax Attention实现）在NVIDIA GH200 GPU上一旦上下文超过约4百万令牌，就无法完成多头注意力层的一次前向-反向传递。 我们引入了一种名为RACE Attention的新方法，这是一种受内核启发的替代方案，它的时间复杂度与序列长度和嵌入维度呈线性关系。RACE Attention通过使用锐化的余弦相似度（即角相似度）来替换指数核，并且通过随机投影和软局部敏感哈希（LSH）近似注意力输出。 在语言建模、遮蔽语言建模以及文本分类等多种任务中，RACE Attention能够在减少运行时间和内存消耗的同时与强大的基准模型达到相同的准确性。在一个受控的规模测试中，在NVIDIA GH200 GPU上处理多达1.2亿个令牌，在Intel Xeon Gold 5220R CPU上处理多达7.5亿个令牌的一次前向-反向传递，远远超过了当前最先进的注意力机制实现的实际限制。 因此，RACE Attention为当今硬件上的极长上下文窗口提供了理论上合理且实用的解决方案。我们希望它能在实践中被采用。
URL
https://arxiv.org/abs/2510.04008
