cs.LGJun 5, 2026

Attention at the Theoretical Minimum: A Mathematics of Arrays Framework for Memory-Optimal Transformer Kernels

Authors: Lenore MullinGaetan Hains

Abstract

The attention mechanism is the dominant computational bottleneck in modern transformer-based AI. Its standard implementation incurs quadratic memory traffic in the sequence length~nn, and DRAM accesses cost 100--1000×\times more energy than arithmetic operations on contemporary hardware, so any analysis focused solely on FLOP counts fundamentally mischaracterises the bottleneck. We present a Mathematics of Arrays (MoA) reformulation of scaled dot-product attention and its numerically stable softmax, deriving a Denotational Normal Form (DNF) that eliminates all intermediate arrays -- including the implicit transposed-key buffer and every softmax temporary -- by algebraic construction rather than empirical tuning. The DNF achieves O(ndk+ndv)O(n_{dk} + n{_{dv}}) data movement versus O(n2+ndk+ndv)O(n^2 + n_{dk} + n_{dv}) for the standard implementation, where nn is the sequence length, dkdk is the key dimensionality and dvdv the value dimensionality, and is verified numerically against PyTorch at full double-precision floating-point on concrete inputs. Unlike hardware-specific accelerators or empirical tiling schemes such as FlashAttention, MoA simultaneously provides array fusion, shape-transformation correctness, and predictive cost models from a single algebraic framework. Memory minimality is a theorem established before any code is written. A predictive performance model projects 22--100×100\times speedup and 22--50×50\times energy reduction, with the advantage widening at exascale. The derivation establishes a formally verified pipeline from Python specification through (ONF) Operational Normal Form, and dimension-lifted hardware mapping, providing performance-portable AI kernels of direct relevance to DARPA edge-deployment and DOE exascale priorities.

Explore similar work

Jul 21, 2026cs.LG

MoA-Structured Decode Attention DNF Derivation, KV-Cache Accumulation, GQA/MQA, and OpenACC Kernel

We derive four memory-optimal inference artifacts for transformer attention using the Mathematics of Arrays (MoA), each following directly from the forward-pass Denotational Normal Form (DNF) of with the query-row index fixed to the current decode step. The artifacts are: (1)~a single-query decode DNF in which the ψψ-reduction eliminates the KK^\top buffer algebraically, achieving (dk+ndk+ndv+dv)×4B(d_k + nd_k+ nd_v+ d_v)\times4\,{B} Dynamic Random Access Memory (DRAM) traffic result numerically verified to err2×107\|{err}\|_\leq2\times10^{-7}; (2)~a C/OpenACC Graphics Processing Unit (GPU) kernel with Operational Normal Form (ONF) stride arithmetic and hardware-coalesced memory access, verified to err=0\|\mathrm{err}\|_\infty=0 (exact IEEE-754 floating-point arithmetic); (3)~a multi-step KV-cache with O(dk+dv)O(d_k+d_v) per-step append via MoA concatenation #\#; and (4)~Grouped-Query Attention (GQA) and Multi-Query Attention (MQA) derived via ψψ-selection, achieving a proven hqhkv\frac {h_q} { h_{kv} } reduction in KV traffic. All programs are verified against PyTorch scaled_dot_product_attention.
Lenore Mulin, Gaetan Hains
May 22, 2026cs.LG

Approaching I/O-optimality for Approximate Attention

We revisit the I/O complexity of attention in large language models. Given query-key-value matrices Q,K,VRn×dQ,K,V\in\mathbb{R}^{n\times d}, and a machine with fast memory size MM, the goal is to compute the "attention matrix" A=softmax(QK/d)VA=\text{softmax}(Q K ^{\top}/\sqrt{d}) V with the minimal number of data transfers between fast and slow memory. Existing methods in the literature, most notably FlashAttention and its variants, incur an I/O cost that depends quadratically on nn, while a trivial lower bound only requires Ω(nd)Ω(nd) I/O's to read the inputs and write the output. In this work, we present a technique for computing attention where the I/O cost only depends almost-linearly on nn in most parameter regimes. This is achieved by developing I/O-efficient algorithms inspired by the recent approximate attention framework of Alman and Song. We also prove corresponding lower bounds in each parameter regime to show that our algorithms are indeed close to I/O-optimal.
Pál András Papp, Aleksandros Sobczyk, Anastasios Zouzias
Sep 21, 2026cs.LG

FlashBoB: I/O-Efficient Exact Backward-over-Backward for Softmax Attention

Transformer models built on the attention mechanism have become a central building block in modern deep learning, yet softmax attention remains a major bottleneck for long-context workloads. While FlashAttention makes the forward and first backward passes I/O-efficient, it does not support backward-over-backward (BoB), which enables exact differentiation through the backward pass for applications such as second-order optimization, test-time training, gradient-based memory, and meta-learning. Existing BoB implementations either materialize large intermediate tensors or exhaust GPU memory at long sequence lengths. We present FlashBoB, an exact, I/O-efficient algorithm for BoB in softmax attention that keeps computation within on-chip tiles and avoids all N×NN \times N intermediate tensors, where NN is the sequence length. The key insight is a hierarchical affine structure in the softmax double backward: two row-wise scalars determine all outputs through affine transformations. This yields a two-pass schedule with bounded on-chip static random-access memory (SRAM) usage and minimal off-chip high-bandwidth memory (HBM) traffic. FlashBoB achieves Θ(N2d2/M)Θ(N^2 d^2/M) HBM traffic (dd is the head dimension and MM is the memory size) and, within the standard FlashAttention-style score-recomputation model, matches the inherited large-cache lower bound for exact forward attention. Empirically, it scales exact attention BoB to N=262KN=262\text{K} on a single A100 80GB GPU, where prior PyTorch exact baselines fail by N=16KN=16\text{K}, and is up to 6.3×6.3\times faster than FlashBack. These results make exact second-order attention practical at long-context sequence lengths where prior implementations cannot run efficiently.
Anthony Givans, Michael Crawshaw, Mingrui Liu