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

1 Introduction

• Online softmax rescales accumulated contributions when a later tile raises the running maximum, while FP32 atomic gradient additions vary with CTA arrival order.Deterministic modes impose an order and cost up to 38% of throughput in FlashAttention-3 and 25% in FlashAttention-4 [28] [44].Source: §1 · p. 1
• FoldAttention fixes a finite reference per query row before scanning, making each weight final and allowing partials to add across disjoint key ranges.The reference need only keep weights within the exponent range of FP32 accumulation and BF16 tensor-core operands.Source: §1 · p. 2, Figure 1 · p. 2
• Final weights gate key and value reads, while a shared integer grid lets backward CTAs add partial gradients in any order.A per-call depth T can cut keys below 2^-T while retaining their mass; shared prefixes are read once for every request that holds them.Source: §1 · p. 2
• At depth 16, decode is 1.36–2.30× faster than the fastest BF16 kernel across seven generations; deterministic backward is up to 1.84× faster than deterministic FlashAttention-3/4 and 1.05× faster than the fastest nondeterministic kernel.The decode output error is within 1.5% of the lowest BF16 error in six cases; the backward gradients are bit-identical across runs, batches, and packings.Source: §1 · pp. 2–3

2 Background √

• Online softmax [21] tracks a running maximum, numerator, and denominator, rescaling accumulated terms when a tile raises the maximum; split-KV partials are merged by comparing log-normalizers and rescaling.The split-KV partial state is (o, ℓ), with ℓ = log L + m [7].Source: §2 · p. 3
• FlashInfer, LeanAttention, and Hydragen use rescaled partial-state merges for split-KV, cascades, and shared prefixes [14] [30] [41].Stable merging compares the partial log-normalizers and rescales one side.Source: §2 · p. 3
• At D = 128, a BF16 decode step reads 512 bytes per cached token, and the fastest BF16 kernel reaches 93–99% of measured H100 bandwidth at 16K keys and beyond.FP8 caches halve traffic but raise output FP32-relative error by 27–133× in the reported benchmarks.Source: §2 · p. 3
• FP32 atomic additions make dQ bits depend on CTA finishing order, while deterministic modes serialize the additions [28] [44].Each CTA holds a key block and adds a partial into every row of dQ it touches.Source: §2 · p. 3

3 The additive contract

• For a finite Zi, disjoint key sets’ partials add, and their numerator-to-denominator quotient equals softmax attention for any Zi.The proposition also bounds overflow and flush-to-zero error under stated limits on key count, values, and score range.Source: §3 · p. 3
• Because each term depends only on its score and Zi, additive partials compose across split-KV, shared-prefix cascades, and different warp completion orders without comparing maxima or rescaling.The decode kernel applies the identity to the inputs selected by its gates.Source: §3 · p. 3
• Any Zi within about 80 binades of the row’s largest score keeps weights and sums in range for n ≤ 2^22 keys and values below 2^16, losing less than 2^-24 of the denominator to flushing.Softmax output is unchanged by an over- or underestimated reference because it scales every weight by the same factor.Source: §3 · p. 4
• The estimated Zi places the true log-sum-exp −0.4 to 17.2 binades above it across seven generations and −0.9 to 28.2 in Qwen3-8B at 64K–128K context.Figure 2 compares the true log-sum-exp and largest score with estimated Zi across rows; its caption notes FP16 overflow at the right edge and BF16/FP32 range to +127.Source: §3 · p. 4, Figure 2 · p. 4
• Backward CTAs round bounded partial gradients onto one shared power-of-two grid, then sum integers in any order without overflow when C < 2^(b+1).A preprocess uses order-free integer maxima to derive the shared bound and grid; the bound uses a row-stochastic attention matrix and Cauchy–Schwarz.Source: §3 · pp. 4–5

Kernel organization

• A front kernel appends the token and estimates each row’s reference before a split-KV kernel uses that fixed reference for read decisions.The cache is paged, and one CTA works on one tile.Source: §4 · p. 5

Cache format

• Keys use two INT8 planes with a per-key BF16 scale, while values are BF16 or two E4M3 planes under a per-layer power-of-two scale.The orthonormal Hadamard rotation is applied to keys and queries; plane A and its scale cost 130 bytes per key at D = 128.Source: §4 · p. 5

Weight-gated reads

• Final-weight verdicts gate reads: only keys needing refinement fetch plane B, and only live keys fetch values; empty tiles skip value work.Verdicts are ORed across the GQA group before bytes are requested, and refinement thresholds depend on request length.Source: §4 · pp. 5–6

Value product

• Split pairs (A, L) add in fixed slot order, and shared-prefix partials can be decoded once for stacked query rows then added to each suffix.Speculative drafts and verification trees are represented as extra query columns masked to their ancestors.Source: §4 · p. 6

Depth as a dial

• At T = 16, error is within 1.5% of the lowest BF16 error on six of seven generations; T = 14 reads fewer bytes within the BF16 error range.Cut keys contribute their weight to L but not their value to A; T is set per call, so one cache serves every depth.Source: §4 · p. 6

Weight-gated reads

• At D = 128, dense decode reads 130 + 128 rK + 256 bytes per key; finite-depth BF16 value traffic is 256 λ bytes per key.Across the ablation cells, these measure 389–423, 261–280, and 196–334 bytes per key; every configuration stores both planes, totaling 514 bytes per key.Source: §4 · p. 6

Kernel structure

• The backward keeps FlashAttention-3’s Hopper structure—a TMA producer warp, two ping-ponged MMA warpgroups, and one key block per CTA—while changing partial-gradient addition.Training pairs FlashAttention-4’s forward with this FoldAttention backward.Source: §5 · p. 6

Order-free reduction

• Figure 4 contrasts ordered FP32 adds, arrival-ordered atomics, and an integer fold that rounds each CTA partial onto one grid and adds in any order.Its second panel shows σ riding the softmax exponential so dQ partials leave the tensor core on the grid.Source: Figure 4 · p. 6

dQ reduction

• For dQ, the grid exponent rides the softmax exponential so partials leave the tensor core on-grid, are rounded to integers, and add asynchronously without overflow.The grid leaves 30 bits for the bound; a postprocess scales the integers back.Source: §5 · pp. 6–7

dK and dV

• For dK and dV, a planner chooses ownership by request shape; per-head CTAs round partials onto integer grids when each head has its own CTA.When sufficient parallel work exists, one CTA processes a key block across all G heads; wider groups split into subgroups whose last CTA adds FP32 partials in subgroup order.Source: §5 · p. 7

Scheduling and guarantees

• Bit-identical gradients across runs, batches, and packings result from reductions that need not wait on an order; reversing the work-list order leaves every bit unchanged.The persistent kernel claims tiles from a cost-sorted work list, and request gradients match whether run alone or packed with others.Source: §5 · p. 7

6.1 Decode on model generations

• At depth 16, decode is 1.36–2.30× faster than the fastest BF16 kernel across seven model generations, with error within 1.5% of the lowest BF16 error in six cases.Dense decode is 1.14–1.30× faster, with error within 1.2% of the lowest BF16 error and below it in six cases.Source: §6.1 · p. 7

6.2 Ablation

• Gating the second key plane refines 2–29% of keys at D = 128, reduces traffic by 18–24%, and yields 1.18–1.27× speedup.The table compares decode speed and bytes read per key as mechanisms are added, at 16K context.Source: §6.2 · p. 8, Table 2 · p. 8

6.3 Context and group-size sweep

• Across 84 cells, depth 16 reaches 3.09× speedup while remaining within the BF16 error range in 82 cells.Dense decode reaches 1.35× and is slower only once; depth 14 reaches 3.50× and stays within the BF16 error range in 67 cells.Source: §6.3 · p. 8

6.4 Serving and quality

• On Qwen3-8B, depth 16 makes a whole decode step 1.35–1.46× faster, while reported likelihood and long-context accuracy match those under BF16 kernels.With BF16 values, finite depth requires a block model fitted on the host; depth-16 steps recover that fitting cost within 38–65 generated tokens.Source: §6.4 · pp. 8–9

6.5 Cascades and drafts

• The Qwen3-8B evaluation reports KL divergence, teacher-forced likelihood difference, RULER accuracy, LongBench scores, and greedy-generation divergence across contexts.Contexts past 32K use YaRN; the caption notes that cuDNN and 8-bit value cells at those contexts come from a second run.Source: Table 3 · p. 10

6.6 Backward

• On eight GQA and MHA shapes at 16K tokens or more, FoldAttention is 1.29× faster than deterministic FlashAttention-3, 1.23× faster than -4, and 1.05× faster than the fastest nondeterministic kernel.Its gradients are bit-identical across reruns, batching, and packing; a whole-model 1B Llama training step is within 0.5% of nondeterministic kernels.Source: §6.6 · p. 10, Table 4 · p. 10

• Hydragen, vLLM’s cascade path, PAT, and FastTree merge shared-prefix online-softmax state into each request, while FoldAttention adds the levels’ pairs.SGLang uses online-softmax state merging for draft trees.Source: §7 · p. 11
• Prior methods fix or freeze references to reduce rescaling, or skip blocks below a pseudo- or running maximum while discarding their mass; FoldAttention uses final weights for per-key reads and additive merges.FlashDecoding++ uses a profiled constant with a fallback, FlashAttention-4 skips small rescales, and VFA freezes a maximum from block summaries.Source: §7 · p. 11
• Quantized caches compress KV storage, while Quest, SparQ, and CoSA select pages or skip value reads using approximate scores; the cited methods primarily evaluate task accuracy rather than BF16-kernel agreement.Among the methods emulated in the paper, none reads fewer bytes than a BF16 kernel at its output error.Source: §7 · p. 11
• Pre-rounding summands to a common grid makes floating-point sums reproducible, while FoldAttention pre-rounds attention-backward partials using cheaply computed bounds so determinism needs no ordering.DASH instead reschedules ordered backward as a DAG; batch-invariant kernels fix reduction topologies and split sizes.Source: §7 · p. 11

• The kernels target Hopper, decode head dimensions 64 and 128, backward dimensions 64, 96, and 128, and MHA and GQA caches; latent caches such as MLA are not supported, and Blackwell porting is future work.Source: §8 · p. 11

A Proofs

• Disjoint key sets split the sums defining F_i, and multiplying numerator and denominator by 2^Z_i shows the quotient is independent of Z_i.Source: A Proofs · p. 16
• With n ≤ 2^22 keys and values below 2^16, the proof bounds overflow and flush-to-zero loss under its stated finite-precision conditions.On the measured generations, m_i − Z_i lies in [−8.2, 16.7], and the reported flush bound is below 2^-95.Source: A Proofs · p. 16
• Each CTA computes an FP32 gradient partial g_c with the power-of-two scale α applied inside the exponential; the bound covers partials because each is a subset of the bounded terms.The grid uses b = 30 with 32-bit accumulators for dQ, and b = 61 with 64-bit accumulators for dK and, under the stated condition, dV.Source: A Proofs · p. 16

B Algorithms

• Algorithm 1 processes a CTA’s 64-key tiles, gates key and value reads using score verdicts, accumulates weighted values, and writes split partials for combination.Finite depth adds cut-key mass to a virtual row; 8-bit values use a gated second value plane.Source: B Algorithms · p. 17
• Algorithm 2 assigns a key block to a backward work-list tile, bulk-adds c dQ into an integer accumulator, and postprocesses integer dQ, dK, and dV sums.The algorithm also writes FP32 partials when a key block’s group is split into subgroups.Source: B Algorithms · pp. 17–18

C Experimental details

• Each library’s fastest valid configuration is selected across the supported cache, GQA, split-count, and backend options.Configurations that fail are dropped; the remaining ones are timed over nine rounds.Source: C Experimental details · p. 19
• Final timings use CUDA-graph replay, L2 eviction before each sample, rotated kernel order, and medians of paired ratios across at least 20 rounds.The first of at least 21 rounds is dropped.Source: C Experimental details · p. 19
• Decode outputs and backward gradients are scored against FP32 references, and kernels exceeding 20 times the best same-precision error are excluded from ratios.Backward reference gradients are computed one head at a time without TF32.Source: C Experimental details · p. 19
• DASH runs separately from FlashAttention-3, with process times joined through FoldAttention; DASH is measured on dense batches only.A one-line masking-bound patch is applied; the cross-process join agrees within 2.5% at every shape.Source: C Experimental details · p. 19
• Shared-prefix baseline results omit cases failing accuracy or execution checks, and PAT is timed eagerly because its default-stream launch cannot be captured by CUDA graphs.PAT has no reported gpt-oss-20b cells; FastTree fails accuracy on three single-prefix cascades.Source: C Experimental details · p. 19

D.1 Latency against error on every generation

• Figure 6 compares decode latency with FP32-relative error across seven generations; on gpt-oss-20b layer 21, no FoldAttention variant meets the most accurate BF16 kernel’s error.Source: Figure 6 · p. 20

D.2 Context sweeps

• Dense FoldAttention is 0.95× at D = 64, G = 8 and 1K context, and 1.03–1.31× elsewhere on uniform-batch context sweeps.The figures extend sweeps across measured group sizes and include 8-bit values; dense 8-bit decode reaches 1.59–1.86× at 32K.Source: §D.2 · pp. 21–23, Figure 7 · p. 21, Figure 8 · p. 22

D.3 Reference quality

• At depths 12–14, the exact-log-sum-exp oracle uses 0.94–0.99 of the estimate’s bytes at matched error, while dense decode differs by at most 3% in bytes and 1.1% in error.The reference estimate’s measured deviation and the real Qwen3-8B ranges are summarized in Figure 2.Source: §D.3 · p. 24

D.4 Tile layouts

• At D = 64, groups of up to four rows use 128-key tiles so copies, round trips, and barriers serve twice as many keys; the wider D = 128 tile gains nothing.Source: §D.4 · p. 24

D.5 Cache construction

• BF16 finite-depth decoding requires a block model fitted in 6.6–11.3 ms per layer per prefill, whereas finite depth with 8-bit values adds no prefill cost.Writing the prompt into the dense cache takes 0.18–0.25 ms per layer versus 0.05–0.09 ms for a BF16 page append.Source: §D.5 · p. 24

D.6 Sparse and quantized decode at matched error

• Figure 10 compares exact-arithmetic byte counts, including metadata, for FoldAttention, Quest, and Faster Flash Decoding on each generation’s final state.FoldAttention is swept over depth, while Quest and Faster Flash Decoding vary their budgets.Source: §D.6 · pp. 24–25, Figure 10 · p. 25

D.7 Cascades and drafts

• On cascades, dense FoldAttention is 0.90–1.23× (1.07×) over the fastest kernel, and on trees it is 0.79–1.13× (0.98×); draft-verification speedups vary with draft count.For chains of two drafts, it is 1.07–1.24× faster; for sixteen drafts, it is 0.43–0.74×. Shared draft bytes do not prevent work from growing with draft count.Source: §D.7 · pp. 25–26, Table 5 · p. 26

E Additional backward results

• Figures 11–12 cover masked and unmasked MHA and eight-query-head-per-KV-head GQA backward throughput, while Tables 6–7 report shape timings and gradient-error comparisons.Table 6 marks shapes below 16K tokens per batch as excluded from the headline geometric mean.Source: E Additional backward results · p. 27, Figure 11 · p. 27, Figure 12 · p. 28
• FoldAttention’s second and third Llama training runs match its first at every step, while the figure compares validation loss and training-loss differences across libraries.Source: Figure 13 · p. 29

F Determinism

• Table 8 evaluates the largest gradient-element difference when a request is rerun, batched, or packed, and reports repeatability across backward benchmark shapes.Source: F Determinism · p. 30, Table 8 · p. 30
• A fixed number of keys per split makes decode invariant to split count, but changing the fixed split from 4 to 8 changes combine grouping and output bits.Table 9 also reports repeat-run and alone-versus-ragged-batch comparisons; cuDNN 9.26 paged decode is not repeatable on the D = 128 batch.Source: F Determinism · p. 30, Table 9 · p. 30