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

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 \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.

Anthony Givans, Michael Crawshaw, Mingrui Liu

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

Ask PDF

May 3, 2026cs.LG

Stochastic Sparse Attention for Memory-Bound Inference

Autoregressive decoding becomes bandwidth-limited at long contexts, as generating each token requires reading all

n_k

key and value vectors from KV cache. We present Stochastic Additive No-mulT Attention (SANTA), a method that sparsifies value-cache access by sampling

S \ll n_k

indices from the post-softmax distribution and aggregates only those value rows. This yields an unbiased estimator of the post-softmax value aggregation while replacing value-stage multiply-accumulates with gather-and-add. We introduce stratified and systematic sampling to design variance-reduced, GPU-friendly variants. Evaluated on Llama-3.1-8B-Instruct at 32k-token contexts, S

^2

ANTA matches baseline accuracy while achieving up to

1.5\times

decode-step attention-kernel speedup over FlashInfer and FlashDecoding on an NVIDIA RTX 6000 Ada. In batched long-context generation, these kernel gains translate to up to

1.25\times

end-to-end decode-latency speedup. Finally, we propose Bernoulli

qK^\mathsf{T}

sampling as a complementary technique to sparsify the score stage, reducing key-feature access through stochastic ternary queries. Both methods are complementary to upstream quantization, low-rank projection, KV-cache compression, and KV-cache selection methods. Together, they point toward sparse, multiplier-free, and energy-efficient inference. We open-source our kernels at: https://github.com/OPUSLab/SANTA.git

Kyle Lee, Corentin Delacour, Kevin Callahan-Coray +5

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

Dynamic Sparse Attention Autoregressive Decoding Key-Value Cache Key-Value Cache Compression Gumbel-Softmax Relaxation Stochastic

Ask PDF

Dec 18, 2025cs.LG

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

Attention is the dominant source of latency during long-context LLM inference, an increasingly popular workload with reasoning models and RAG. We propose Kascade, a training-free sparse attention method that leverages known observations such as 1) post-softmax attention is intrinsically sparse, and 2) the identity of high-weight keys is stable across nearby layers. Kascade computes exact Top-k indices in a small set of anchor layers, then reuses those indices in intermediate reuse layers. The anchor layers are selected algorithmically, via a dynamic-programming objective that maximizes cross-layer similarity over a development set, allowing easy deployment across models. The method incorporates efficient implementation constraints (e.g. tile-level operations), across both prefill and decode attention. The Top-k selection and reuse in Kascade is head-aware and we show in our experiments that this is critical for high accuracy. Kascade achieves up to 4.1x speedup in decode attention and 2.2x speedup in prefill attention over FlashAttention-3 baseline on H100 GPUs while closely matching dense attention accuracy on long-context benchmarks such as LongBench and AIME-24.

Dhruv Deshmukh, Saurabh Goyal, Nipun Kwatra +1

Microsoft Research India.

Dynamic Sparse Attention LLM Inference Optimization Cross-Layer Interactions

Ask PDF