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

Sriman Achanta

cs.LGSep 27, 2026

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

Authors: Sriman Achanta

Organizations: Virginia Commonwealth University

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.30 $\times$ faster than the fastest BF16 baseline, and up to 3.09 $\times$ 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.30 $\times$ faster at matched error. We validate on Qwen3-8B that a whole decode step is up to 1.46 $\times$ 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.84 $\times$ faster than deterministic FlashAttention-3/4 and 1.05 $\times$ faster than the fastest nondeterministic kernel.

Figures & tables

Model, layer, batch $\times$ context	Fastest BF16	BF16 err.	Fold dense	Fold $T{=}16$	Fold $T{=}14$
Qwen3-30B L24, 8 $\times$ 16K	FA-3 99	1.66–3.00	77 / 1.63 (1.30 $\times$ )	54 / 1.68 † (1.84 $\times$ )	45 / 2.33 † (2.22 $\times$ )
Qwen3-30B L24, 32 $\times$ 8–16K	FlashInfer 264	1.71–2.33	206 / 1.67 (1.28 $\times$ )	126 / 1.70 (2.10 $\times$ )	103 / 2.19 † (2.54 $\times$ )
Qwen3-30B L24, 16 $\times$ 16–32K	FlashInfer 262	1.71–3.13	207 / 1.65 (1.27 $\times$ )	114 / 1.73 † (2.30 $\times$ )	99 / 2.39 † (2.65 $\times$ )
gpt-oss-20b L9, 32 $\times$ 8–16K	cuDNN 269	1.68–2.83	212 / 1.67 (1.27 $\times$ )	163 / 1.70 † (1.65 $\times$ )	134 / 2.06 † (2.00 $\times$ )
gpt-oss-20b L21, 32 $\times$ 8–16K	cuDNN 267	1.68–2.72	216 / 1.70 † (1.24 $\times$ )	171 / 1.97 † (1.57 $\times$ )	142 / 2.46 † (1.89 $\times$ )
GLM-4-9B L8, 32 $\times$ 8–16K	FlashInfer 139	1.70–3.05	119 / 1.66 (1.17 $\times$ )	91 / 1.67 (1.53 $\times$ )	83 / 1.82 † (1.69 $\times$ )

Table 1: Decode on eight-step generations: latency in

\mu

s / FP32-relative

\ell_{2}

error in

10^{-3}

(speedup over the fastest BF16 baseline).

\dagger

: error above the lower end of the BF16 baselines’ range.

	$D{=}128$ , $G{=}8$	$D{=}128$ , $G{=}4$	$D{=}128$ , $G{=}16$	$D{=}64$ , $G{=}8$	Error
	Qwen	Qwen	GLM	gpt-oss
Both planes	0.98 $\times$ (514 B)	0.99 $\times$ (514 B)	0.98 $\times$ (514 B)	1.00 $\times$ (258 B)	1.64–1.66
Gated plane B (dense)	1.26 $\times$ (390 B)	1.27 $\times$ (389 B)	1.18 $\times$ (423 B)	1.24 $\times$ (204 B)	1.66–1.67
+ 8-bit values	1.76 $\times$ (262 B)	1.79 $\times$ (261 B)	1.41 $\times$ (275 B)	1.56 $\times$ (134 B)	1.68–1.70
+ depth $T{=}16$ , BF16 values	2.15 $\times$ (217 B)	2.37 $\times$ (196 B)	1.53 $\times$ (308 B)	1.65 $\times$ (141 B)	1.70–1.72

Table 2: Decode kernel speed over the fastest BF16 baseline’s decode and bytes read per key, adding one mechanism at a time (ragged batches of 256 KV heads, 16K context). Errors are FP32-relative

\ell_{2}

10^{-3}

; the most accurate BF16 baseline’s is 1.67–1.69 in these cells, and XQA’s 2.31–2.83.

FA-3	0	0	0	0	87.1	0 / 0	68.0	49.3	–
	KL, 8K / 16K / 32K			$\Delta$ NLL	RULER	KL	RULER	LongBench	Diverge
Arm				32K	8–32K	64K / 128K	64–128K	16 tasks
FlashInfer	7.9	8.0	8.6	$-2.9$	86.9	9.1 / 9.8	66.8	49.3	55
cuDNN	8.0	8.1	8.9	$-1.1$	87.1	8.8 / 10.2	69.9	–	61
FlashInfer FP8	73.9	109.0	131.8	$+75.1$	87.3	124.8 / 174.6	66.0	49.2	26
Fold dense	8.1	8.1	8.5	$+0.2$	87.6	8.9 / 10.5	67.6	49.3	90

Table 3: Qwen3-8B decoding for real in each arm. KL is FlashAttention-3’s next-token distribution against the arm’s, in

10^{-4}

nats per token, and

\Delta

NLL the arm’s teacher-forced negative log-likelihood minus FlashAttention-3’s, in

10^{-4}

nats, on 32 WikiText-103 documents of 1024 tokens. RULER is the mean exact match over four tasks. Contexts past 32K use YaRN [ 27 ] ; their cuDNN and 8-bit

V

cells come from a second run, in which FlashAttention-3 scores 70.3 on RULER. “Diverge” is the median position at which greedy generation first leaves FlashAttention-3’s.

Set	FA-3	FA-3 det.	FA-4	FA-4 det.	cuDNN	DASH	Fastest det.	Fastest nondet.
Model shapes, backward (8)	1.07	1.29	1.06	1.23	1.23	1.32	1.23	1.05
Packed varlen, backward (5)	1.08	1.30	1.07	1.29	–	–	1.28	1.06
MHA grid, backward (24)	0.99	1.18	0.98	1.11	1.17	1.12	1.07	0.97
GQA grid ( $G{=}8$ ), backward (24)	1.06	1.32	1.04	1.25	1.23	1.28	1.23	1.04
Model shapes, train (8)	0.99	1.15	1.01	1.14	1.26	1.17	1.12	0.99

Table 4: Backward speedup of FoldAttention over each baseline (geometric mean). “Fastest” takes the best kernel of its kind at each shape. Training time combines FA-4’s forward with each backward.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Set	FlashInfer cascade	vLLM cascade	PAT	FastTree	Fastest
Cascades (16)	3.93–8.88 (6.18)	1.08–1.76 (1.34)	1.13–1.48 (1.32)	1.23–1.95 (1.61)	1.04–1.42 (1.22)
Prefix trees (8)	3.38–8.88 (5.27)	1.48–2.88 (1.80)	0.90–1.29 (1.07)	1.02–1.71 (1.24)	0.90–1.33 (1.09)

Appendix

Table 5: Cascade speedup of dense FoldAttention over each kernel: the range of the kernel’s time over ours, with the geometric mean in parentheses. “Fastest” takes the fastest kernel of any kind in each cell. PAT has no gpt-oss-20b cells and FastTree no result on three cascades (Appendix C ).

$B$	$H/H_{KV}$	$S$	$D$	mask	ours ( $\mu$ s)	FA-3	FA-3 det.	FA-4	FA-4 det.	cuDNN	DASH
2	32/8	8192	128	causal	4942	1.03	1.29	1.04	1.19	1.21	1.37
4	32/8	4096	128	causal	2599	1.05	1.35	1.05	1.28	1.20	1.37
8	32/8	2048	128	causal	1548	1.10	1.38	1.08	1.34	1.18	1.39
16	32/8	1024	128	causal	1056	1.21	1.43	1.16	1.39	1.18	1.41
2	32/32	8192	128	causal	4865	1.00	1.25	0.99	1.14	1.19	1.30
2	32/8	8192	128	full	9026	1.03	1.17	1.00	1.07	1.16	1.16

Appendix

Table 6: Backward time of FoldAttention and each baseline’s time over it.

H/H_{KV}

gives query and KV heads. Shapes marked

\dagger

are below 16K tokens per batch and are excluded from the headline geometric mean.

Step, layer	$\ell_{2}$ Fold / FA-3	p99 Fold / FA-3	0 to $-13$	$-14$ to $-19$	share $\leq-20$	median $\leq-20$ , Fold / FA-3
500, 0	2.99 / 2.99	1.45 / 1.52	1.00–1.03	1.07–1.12	2.7%	0.58 / 0.46
500, 7	5.01 / 5.01	12.5 / 19.0	1.00	1.00	6.3%	1.12 / 2.20
500, 15	5.01 / 5.01	12.0 / 14.8	1.00	1.00	5.7%	1.30 / 2.03
1999, 0	3.36 / 3.36	1.23 / 1.25	1.00–1.04	1.08–1.21	2.1%	0.40 / 0.31
1999, 7	5.67 / 5.67	3.17 / 3.27	1.00	1.00	0.9%	2.26 / 2.67
1999, 15	4.43 / 4.44	1.69 / 1.54	1.00	1.00–2.70	2.2%	0.46 / 0.0047

Appendix

Table 7: Per-element gradient error against FP64 on attention operands captured from the training run of Section 6.4 (a Llama-3.2-1B-shaped model,

H/H_{KV}=32/8

D=64

, 8192 tokens), for FoldAttention and FlashAttention-3.

\ell_{2}

is the

dQ

error in

10^{-3}

; p99 is the 99th percentile of per-element relative error. The binade columns give FoldAttention’s median relative error over FlashAttention-3’s among

dQ

elements that many binades below the largest of their request and KV head; the last columns cover elements

2^{-20}

or more below it.

dK

and

dV

match FlashAttention-3’s

\ell_{2}

error to three digits in every layer.

Kernel	Ten runs	Alone vs. batch of 4	Alone vs. packed	Repeatable	Worst error
FA-3	$3.9\times 10^{-3}$	$2.0\times 10^{-3}$	$2.0\times 10^{-3}$	1/73	2.53
FA-4	$2.0\times 10^{-3}$	$4.9\times 10^{-4}$	$2.4\times 10^{-4}$	3/141	2.53
cuDNN	$4.9\times 10^{-4}$	$4.9\times 10^{-4}$	–	11/136	3.02
FA-3 det.	0	0	0	73/73	2.53
FA-4 det.	0	0	0	141/141	2.53
DASH	0	0	–	68/68	2.52

Appendix

Table 8: Largest absolute difference in any

dQ,dK,dV

element of one request against its first result, over three dense shapes and two packed ones (0 means every bit matched). “Repeatable” counts backward benchmark shapes with identical bits over three calls; error is the worst FP32-relative

\ell_{2}

over

dQ,dK,dV

10^{-3}

	Qwen3-30B, $D=128$		gpt-oss-20b, $D=64$
Kernel	repeat	alone vs. batch	repeat	alone vs. batch
FA-3	0	$4.9\times 10^{-4}$	0	$1.6\times 10^{-2}$
FA-4	0	0	0	0
FlashInfer	0	$2.0\times 10^{-3}$	0	$1.6\times 10^{-2}$
XQA	0	$2.0\times 10^{-3}$	0	$3.1\times 10^{-2}$
cuDNN	$2.4\times 10^{-4}$	$9.8\times 10^{-4}$	0	$1.6\times 10^{-2}$

Appendix

Table 9: Largest absolute output difference for the first request of a ragged batch of 32 at 16K context, run twice, and run alone instead of in the batch. FoldAttention’s default split is chosen from the batch (8 splits in the batch, 64 alone); a fixed split of 4 makes it batch invariant. Changing the fixed split from 4 to 8 changes the combine’s grouping and the bits. cuDNN 9.26’s paged decode is not repeatable on the

D=128

batch.

Explore similar work

FlashBoB: I/O-Efficient Exact Backward-over-Backward for Softmax Attention
Sep 21, 2026Anthony Givans, Michael Crawshaw, Mingrui LiuBlock Sparse Flash Attention Gumbel-Softmax Relaxation
Stochastic Sparse Attention for Memory-Bound Inference
May 3, 2026Kyle Lee, Corentin Delacour, Kevin Callahan-Coray +5Dynamic Sparse Attention Autoregressive Decoding
Department of Electrical and Computer Engineering, University of California, Santa Barbara, CA, USA · Department of Electrical and Computer Engineering, University of Michigan, Ann Arbor, MI, USA · Department of Electrical and Computer Engineering, Carnegie Mellon University, Pittsburgh, PA, USA
Kascade: A Practical Sparse Attention Method for Long-Context LLM Inference
Dec 18, 2025Dhruv Deshmukh, Saurabh Goyal, Nipun Kwatra +1Dynamic Sparse Attention LLM Inference Optimization
Microsoft Research India.

Sep 21, 2026, cs.LGJ/K move · Enter open · S save

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

Anthony Givans, Michael Crawshaw, Mingrui Liu

Block Sparse Flash Attention Gumbel-Softmax Relaxation Backward Passes Activation Functions Gradient-Based Optimization

Open paper AskPDF arXiv

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 \times N

intermediate tensors, where

N

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

Θ(N^2 d^2/M)

HBM traffic (

d

is the head dimension and

M

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=262\text{K}

on a single A100 80GB GPU, where prior PyTorch exact baselines fail by

N=16\text{K}

, and is up to

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.

Abstract

Explore similar work

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

Stochastic Sparse Attention for Memory-Bound Inference

Kascade: A Practical Sparse Attention Method for Long-Context LLM Inference