Simran Arora

Sequence models face stark tradeoffs between recall quality and memory efficiency. The ability to use information over long sequences (i.e. recall) is critical for sequence modeling tasks ranging from information extraction to reasoning. Prior work has shown that in theory, *linear* attention models with sufficient recurrent state sizes can expand the Pareto frontier of the recall-memory tradeoff space beyond alternative architectures such as softmax attention and state space models. However, it is difficult to scale the linear attention state size due to hardware bottlenecks. I/O-aware algorithms store the linear attention states in thread registers, however state sizes beyond $\approx 3$ megabytes exhaust register memory and trigger expensive register spills. In this work, we introduce Cylon, a hardware-aware modification to linear attention architectures which partitions the recurrent state across the registers of multiple GPU processors and asynchronously combines the partitions. When applying Cylon to popular architectures, such as Hedgehog, Mamba-2, and Gated Deltanet, we unlock up to $6\times$ higher forward pass throughput and $10\times$ higher backward pass throughput compared to prior linear attention algorithms on Hopper and Blackwell GPUs. Moreover, Cylon preserves recall quality compared to non-partitioned baselines. Finally, Cylon makes large states available to model designers by unlocking sizes ($\geq 131$ MB) which are not achievable by the existing linear attention kernels.