FoldAttention: Declared-Reference Softmax for Fast Decode and Deterministic Backward
On the map: Efficient LLM inference and serving → KV cache optimization
Abstract (arXiv)
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 Zi before scanning the KV cache. Each weight 2^sij-Zi 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× faster than the fastest BF16 baseline, and up to 3.09× 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× faster at matched error. We validate on Qwen3-8B that a whole decode step is up to 1.46× 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× faster than deterministic FlashAttention-3/4 and 1.05× faster than the fastest nondeterministic kernel.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Placed in KV cache optimization, one of 12 directions in Efficient LLM inference and serving. Builds on 38 of its 45 references that are in the map. First posted on arXiv 27 Sep 2026.
Builds on
38 of 45 references.
Closest ancestors · 5
Same direction, newest first
- 2025PAT: Accelerating LLM Decoding via Prefix-Aware Attention with Resource Efficient Multi-Tile Kernel27 Nov
- 2025FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving2 Jan
- 2024Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference16 Jun
- 2023SparQ Attention: Bandwidth-Efficient LLM Inference8 Dec
- 2023Efficient Memory Management for Large Language Model Serving with PagedAttention12 Sep
Other references · 33
- 2017Attention Is All You Needcited by 59,932
- 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
- 2024The Llama 3 Herd of Modelscited by 14,192
- 2025Qwen3 Technical Reportcited by 8,195
- 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,581
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
- 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,551
- 2023Efficient Streaming Language Models with Attention Sinkscited by 2,115
- 2022Fast Inference from Transformers via Speculative Decodingcited by 1,615
- 2023LongBench: A Bilingual, Multitask Benchmark for Long Context Understandingcited by 1,378
- 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,311
- 2025gpt-oss-120b & gpt-oss-20b Model Cardcited by 1,258
21 more
- 2023SGLang: Efficient Execution of Structured Language Model Programscited by 1,107
- 2024ChatGLM: A Family of Large Language Models from GLM-130B to GLM-4 All Toolscited by 1,096
- 2024KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cachecited by 604
- 2024QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMscited by 598
- 2023YaRN: Efficient Context Window Extension of Large Language Modelscited by 584
- 2024KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantizationcited by 571
- 2024FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precisioncited by 507
- 2023SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verificationcited by 359
- 2022FP8 Formats for Deep Learningcited by 187
- 2018Online normalizer calculation for softmaxcited by 154
- 2024Hydragen: High-Throughput LLM Inference with Shared Prefixescited by 71
- 2023FlashDecoding++: Faster Large Language Model Inference on GPUscited by 70
- 2026FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scalingcited by 58
- 2024Lean Attention: Hardware-Aware Scalable Attention Mechanism for the Decode-Phase of Transformerscited by 16
- 2025Deterministic Inference across Tensor Parallel Sizes That Eliminates Training-Inference Mismatchcited by 13
- 2025BLASST: Dynamic BLocked Attention Sparsity via Softmax Thresholdingcited by 12
- 2026Diagnosing Training Inference Mismatch in LLM Reinforcement Learning via a Zero-Mismatch Referencecited by 5
- 2026Faster Than Flash: Exploiting Attention Sparsity for Efficient Long-Context Decodingcited by 4
- 2026DASH: Deterministic Attention Scheduling for High-throughput Reproducible LLM Trainingcited by 3
- 2026VFA: Relieving Vector Operations in Flash Attention with Global Maximum Pre-computationcited by 2
- 2026CoSA: Accelerating Long-Context Inference via Proxy-Kernel Co-Designed Sparse Attentioncited by 1
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2026Stochastic Sparse Attention for Memory-Bound Inference
- 2026Where Activation Sparsity and KV-Cache Sparsity Cross in LLM Decoding
- 2026CORE: COverage CAlibration and Evicted-Mass REdistribution for KV Cache
- 2026Quality Recovery for Quantized KV Caches via Low-Rank Attention Adaptation
- 2026ResKV: Reconstructing Omitted Attention Contributions for Fixed-Budget KV Cache Compression
- 2025Mustafar: Promoting Unstructured Sparsity for KV Cache Pruning in LLM Inference
- 2026The risk of KV cache compression
- 2026EntmaxKV: Support-Aware Decoding for Entmax Attention