Superlinear Multi-Step Attention

TL;DR

Superlinear attention introduces a trainable multi-step mechanism that reduces long-context attention from quadratic to subquadratic scaling while preserving random context access. By framing attention as iterative span-search and span-attention rounds across an accumulation-driven representation, it achieves a baseline $O(L^{3/2})$ for $N=2$ and generalizes toward $O(L^{1+1/N})$ with more steps. The authors detail architectural components, a bucketed GPU kernel for irregular spans, and feasibility experiments showing practical throughput gains for very long contexts and learnability on long-context retrieval tasks. The work lays groundwork for scalable long-context processing in transformers, combining architectural design, scaling analysis, and system-level feasibility with clear directions for broader empirical evaluation and higher-step extensions.

Abstract

In this paper, we propose \textbf{Superlinear attention}, a fully trainable multi-step attention architecture that achieves subquadratic complexity for long sequences while preserving \textbf{random context access} (a.k.a.\ structural non-exclusion): no eligible token position is structurally excluded from being selected for attention. Superlinear attention reformulates standard causal self-attention as a multi-step search problem with $N$ steps, yielding an overall complexity of $O(L^{1+\frac{1}{N}})$. To illustrate the architecture, we present a baseline $N=2$ implementation, which is algorithmically analogous to standard jump search. In this $O(L^{3/2})$ instantiation, the first step performs $O(L^{3/2})$ span-search to select relevant spans of the sequence, and the second step applies $O(L^{3/2})$ span-attention (standard attention restricted to the selected spans). In an upscaled $O(L^{1.54})$ configuration for robustness, we achieve an average decoding throughput of 114 tokens/sec at 1M context length and 80 tokens/sec at 10M context in our implementation on a modified 30B hybrid MoE model on a single B200 GPU. With limited training, we also obtain strong performance on the NIAH (Needle In A Haystack) task up to 256K context length, demonstrating that the routed span selection is learnable end-to-end. This paper emphasizes architectural formulation, scaling analysis, and systems feasibility, and presents initial validation; comprehensive quality evaluations across diverse long-context tasks are left to future work.

Figures (6)

• Figure 1: Left: Overview of the proposed Superlinear attention architecture. The architecture consists of four main components: accumulation, span-search, span-attention, and combination. Each component plays a crucial role in achieving subquadratic complexity while maintaining random context access. In the example, we search over 4 spans, and top-$2$ spans are selected for attention and combined for the final output. The schematic in this figure only shows the 1D attention mask for a single query token for simplicity; Figure \ref{['fig:span_search_and_attention']} shows the full 2D attention map for the span-search and span-attention steps. Right: overall model architecture with the Superlinear attention layer integrated into a hybrid transformer architecture using Nemotron-3-Nano nvidia2025nemotron3nanoopennvidia2025nvidianemotron3efficient as the base model. All the standard attention layers in the base model are replaced by Superlinear attention layers.
• Figure 2: Left: Stride pattern for the span-search step with $N=2$, where the search complexity is $O(L^{3/2})$. Right: Illustration of the selected spans for the span-attention step, which also has a complexity of $O(L^{3/2})$. The span tokens are shown in green, while the diagonal tokens are shown in gray. Each row corresponds to the attention map for a single query token. Figure \ref{['fig:superlinear_overview']} shows how each row fits into the Superlinear attention architecture. In actual implementations, the diagonal tokens are expanded into sliding windows to capture local dependencies, and multiple spans are used and combined to improve expressiveness and make the span-search trainable via backpropagation. Combined, the span-search and span-attention steps achieve an overall complexity of $O(L^{3/2})$ while ensuring random context access.
• Figure 3: Illustration of random context access (structural non-exclusion) induced by the $N=2$ span construction. In this example, query position is $i=30$, the backward extent is $b=2$ and the forward extent is $f=0$, the span length is therefore $(b+f) \cdot \lceil i^{1/2} \rceil = 2 \cdot \lceil\sqrt{30}\rceil = 12$. $j$ is the key position index, and $s$ is the stride index for the span-search step. Then the anchor tokens for search are located at $j=i - (s+1)^2 + 1$, or $30, 27, 22, 15, 6$ for $s=0, 1, 2, 3, 4$ respectively. As shown, all token positions from $0$ to $30$ fall inside at least one candidate span (centered at some anchor), meaning each position is reachable if the routing selects an appropriate anchor.
• Figure 4: Comparison of kernel designs for irregular span-attention patterns. (a) The original scattered span-attention pattern. (b) The sorting approach reorders queries to create a contiguous diagonal band, but incurs $O(L \log L)$ overhead from sorting and permutation. (c) The proposed bucketed approach groups queries by their key-block footprint (color-coded) into buckets. (d) Distribution of queries per bucket (key-block range). The kernel uses dynamic work-stealing to handle varying bucket sizes, ensuring high GPU occupancy. Queries in the same bucket are processed together, ensuring efficient memory access without the need for global sorting. This design is critical for realizing the theoretical speedup of the Superlinear attention mechanism.
• Figure 5: Throughput comparison between Superlinear attention (our baseline $N=2$ implementation) and FlashAttention-2. (a) Prefill throughput (tokens/sec) under chunked prefilling (32K chunk size) versus context length. (b) Decode throughput (tokens/sec) versus context length. Superlinear attention scales more gracefully with context and enables efficient multi-million-token inference.