FoldAttention: Declared-Reference Softmax for Fast Decode and Deterministic Backward

arXiv:2609.33410v1 Announce Type: new Abstract: Autoregressive decode repeatedly streams a growing KV cache, making attention a major cost at long context. Existing high-performance kernels use online softmax, which discovers a row's normalization reference as it scans keys. Earlier contributions therefore remain provisional and may require rescaling. We argue that the reference need not be discovered: softmax is invariant to a common shift, so the reference only has to keep the weights in range. We present FoldAttention, an additive formulation of softmax attention that fixes a finite reference Z_i before scanning the KV cache. Each weight 2^{s_{ij}-Z_i} is then final when computed, so contributions add across disjoint key ranges and their quotient equals softmax attention in real arithmetic. We use this property to develop two techniques for Hopper decode: (1) final weights gate key and value reads before the bytes are fetched, and a per-call depth T cuts keys below 2^{-T} while keeping their mass, and (2) additive partials compose split KV and shared-prefix cascades without rescaling. On H100 at T=16, FoldAttention decodes seven real-model generations 1.36-2.30imes faster than the fastest BF16 baseline, and up to 3.09imes faster across MHA and GQA shapes, at an error within 1.5% of the lowest BF16 error on six of the seven; reading every key, it is 1.14-1.30imes faster at matched error. We validate on Qwen3-8B that a whole decode step is up to 1.46imes faster while likelihood and long-context accuracy match those under BF16 kernels. The same principle makes the backward deterministic: CTAs round bounded partial gradients onto an integer grid declared before the reduction and add them in any order. FoldAttention thereby removes the determinism tax: its deterministic backward is up to 1.84imes faster than deterministic FlashAttention-3/4 and 1.05imes faster than the fastest nondeterministic kernel.

Source: arXiv cs.LG | 2026-09-29