PaperlayerView exact revision on arXiv

Source-linked AI summary

ThriftAttention: Selective Mixed Precision for Long-Context FP4 Attention

Joe Sharratt

arXiv:2605.23081v1cs.LG

TL;DR

Low-bit attention can degrade quality in long-context workloads despite its efficiency benefits. ThriftAttention selects important query-key blocks for FP16 while computing the rest in FP4. With 5% of blocks in FP16, it recovers an average 89.1% of the FP4-to-FP16 performance gap.

  • Problem

    Low-bit attention variants degrade in long-context settings, a concern for efficient inference where attention is a key bottleneck.

    Source: §1 · pp. 1–2

  • Method

    ThriftAttention uses a lightweight heuristic to select important query-key block pairs for FP16, computes the remaining blocks in FP4, and merges both paths via online softmax.

    Source: §1 · p. 2

  • Results

    89.1% average recovery of the FP4→FP16 performance gap across model families and evaluation benchmarks is achieved with a 5% FP16 block budget.

    Source: §1 · p. 2, §4.2 · p. 6

  • Takeaways & Limitations

    ThriftAttention’s relative performance advantage over FP4 grows with context length, as FP4 quality degrades relative to FP16.

    Source: §4.3 · p. 7

    ThriftAttention increases KV-cache memory footprint by 28% because it stores FP16 and FP4 caches.

    Source: §5 · p. 9

Abstract

from arXiv · show

Efficient attention algorithms are critical to mitigate the quadratic cost of attention in long-context workloads. Prior work utilises block-scaled quantisation techniques on Blackwell GPUs to move attention computation to 4-bit precision to accelerate inference. However, these techniques result in significant quality degradation in long-context settings. We show that the output impact of quantisation error is highly non-uniform and increases with the importance of each query-key interaction, concentrating functionally relevant error in a small number of attention blocks that contain the most important tokens. We propose ThriftAttention, a low-bit attention variant that delivers near-FP16 long-context quality at FP4 inference efficiency. This approach proceeds in two stages. First, a heuristic rapidly selects a small number of important query-key block pairs for FP16 precision. Second, the selected blocks are computed in FP16 and the remaining blocks in FP4, with both paths merged via online softmax into a single output. We demonstrate across long-context benchmarks and model families that by computing only 5% of query-key blocks in FP16, ThriftAttention recovers on average 89.1% of the FP4-to-FP16 performance gap. We show ThriftAttention's advantage grows with sequence length, mitigating the systematic FP4 quality degradation observed at longer contexts. The code is available at https://github.com/joesharratt1229/ThriftAttention.

1 Introduction

ThriftAttention uses selective FP16 computation within FP4 attention to recover long-context quality while retaining low-bit inference efficiency. The introduction motivates this approach by locating relevant quantisation error in a small number of important query-key interactions.

Motivation

  • On Blackwell, native FP4 Tensor Cores provide 4× the arithmetic throughput of equivalent FP16 instructions while reducing KV-cache traffic similarly, addressing attention’s quadratic cost and memory burden.Attention’s quadratic cost and KV-cache memory traffic dominate execution time in long-context workloads.Source: §1 · p. 1
  • Prior FP4 methods accept quality degradation, while sparsity methods [Tang et al., 2024, Zhang et al., 2025e, 2023] must drop at least 75% of KV blocks to match FP4 generation latency, making omitted-block errors irrecoverable.Source: §1 · p. 1
  • Functionally relevant quantisation error concentrates in a small number of query-key interactions with high-magnitude attention scores that matter most to the final output distribution.Source: §1 · p. 2

Method

  • A training-free heuristic ranks query-key blocks by token-mean similarity, assigns the top-k to FP16 and the rest to FP4, then merges both paths with online softmax.The approach computes the most important block interactions in FP16 and the remainder in FP4; the authors describe it as the first sub-byte mixed-precision approach for attention.Source: §1 · p. 2

Results

  • At 5%, 10%, and 25% FP16 block budgets, ThriftAttention recovers 89.1%, 91.8%, and 92.4% of the FP4-to-FP16 gap, with up to 2× lower long-context generation latency.Its advantage grows with context length, where uniform FP4 attention degrades most severely.Source: §1 · p. 2

Evaluation

  • Evaluation spans LongBench-v1, HELMET, RULER, and PG-19 across Llama, Qwen, and Ministral families, testing whether a small FP16 budget can recover most FP16 quality with low-bit efficiency.Source: §1 · p. 2

2 Related Work

Related work spans I/O-efficient, quantised, and sparse attention; ThriftAttention instead assigns full precision to the most important blocks rather than quantising uniformly.

I/O-Efficient Attention

Quantised Attention

  • SageAttention [Zhang et al., 2025d,a] uses INT8/FP8 quantisation with outlier smoothing, while SageAttention3 [Zhang et al., 2025c] extends attention quantisation to FP4 on Blackwell using two-level microscaling.Prior work also quantises linear layers, compresses KV caches, or combines quantised matmuls with sparsity.Source: §2 · p. 3

Sparse Attention

  • SpargeAttn [Zhang et al., 2025e] is the closest prior work, combining block-sparsity prediction with quantised attention, but uses sparsity as its primary acceleration and not sub-8-bit formats.Other sparse approaches select or evict KV blocks, learn sparse structure, or use linear attention, which is less effective for peaked attention distributions.Source: §2 · p. 3

Positioning

  • ThriftAttention assigns full precision to the most important blocks; for each block, its error is upper-bounded by FP4 quantisation noise rather than skipped or approximated attention-score magnitudes in sparse methods.This contrasts its precision allocation with uniform quantisation and sparse approximation.Source: §2 · p. 3

3 Method

ThriftAttention motivates selective precision by the non-uniform impact of FP4 attention errors, then routes selected query-key blocks through FP16 and the rest through FP4. Its fused kernel is structured to limit register pressure and avoid unnecessary FP16 work.

3.1 Motivation for ThriftAttention

  • FP4 attention error concentrates in a few high-impact blocks, motivating their selective promotion to FP16 while the remaining blocks stay in FP4.A key’s output-error contribution depends on its quantisation error, attention weight, and value deviation from the output; large attention weights amplify error. The concentrated blocks are typically near-diagonal or non-initial attention sinks.Source: §3.1 · pp. 3–4, Figure 3 · p. 4

3.2 ThriftAttention Algorithm

  • ThriftAttention partitions queries and keys/values into blocks, scores block pairs using token-wise block means, and routes each query block’s top-k key blocks to FP16.Q, K, and V are quantised in NVFP4 format; non-selected blocks use FP4, and the FP16 and FP4 output updates are merged online with softmax. Causal attention applies the usual mask and limits selection to visible key blocks.Source: §3.2 · pp. 4–5

3.3 Implementation and Optimisation on Hardware

  • A single fused CUDA kernel processes non-selected blocks through FP4 before invoking an FP16 helper for promoted blocks.FP16 query fragments are scoped to the selected-block phase to reduce register pressure; shared memory is aliased for K/V tiles, FP4 KV tiles are double-buffered, and CTAs with no selected blocks skip the FP16 path and its loads and register allocation.Source: §3.3 · pp. 5–6

4 Experiments

Experiments assess ThriftAttention’s efficiency and quality across long-context benchmarks, model families, sequence lengths, and matched-compute sparse baselines, including kernel speed and per-token negative log-likelihood.

4.1 Efficiency and Effectiveness

  • At 131k context, ThriftAttention reaches roughly 1.2× end-to-end prefill speedup relative to FP16 attention and near 2× generation speedup over Qwen3-8B using full-FP16 attention.Prefill kernels reach up to 1.7× speedup over FlashAttention-2; decode kernels are 3×–5.5× faster than FlashAttention-2 with minimal overhead versus full FP4 attention.Source: §4.1 · p. 6

4.2 ThriftAttention Benchmark Evaluation

  • A 5% FP16 block budget recovers an average 89.1% of the FP4→FP16 performance gap across models and long-context benchmarks.Recovery rises to 91.8% at 10% and 92.4% at 25%, with diminishing returns from additional FP16 blocks.Source: §4.2 · p. 6, Table 1 · p. 7
  • LongBench v1 accuracy saturates at a 5% FP16 budget, RULER improves steadily from 5% to 25%, and HELMET scales unevenly between them.The benchmark results show that required FP16 budgets differ across tasks.Source: §4.2 · pp. 6–7, Table 1 · p. 7

4.3 Sequence Length Experiments

  • At 5% FP16 budget, ThriftAttention’s relative performance over FP4 rises from 2.00× at 8k to 2.2× at 131k on HELMET.Llama3.1-8B FP4 retention falls from 50% at 8k to 32% at 131k, with comparable declines for two Qwen variants and Llama3.2-3B.Source: §4.3 · p. 7, Table 2 · p. 8

4.4 Negative Log-Likelihood Analysis

  • At 64k and 128k, FP4 ∆NLL reaches 0.10 for tokens at sequence ends, while ThriftAttention keeps ∆NLL at ≤0.02 across context lengths.Across 300 PG-19 documents, the relative ∆NLL reduction grows from approximately 2× at 8k to approximately 5× at the end of a 128k sequence. The widening advantage is attributed to important attention interactions occupying a smaller fraction of all blocks as sequence length grows, so a fixed fractional top-k budget captures more attention mass.Source: §4.4 · pp. 7–8, Figure 5 · p. 9

4.5 Comparison to Sparse Attention Baselines at Matched Compute

  • At equivalent inference efficiency, ThriftAttention outperforms sparse-attention approaches by retaining all interactions in FP4 rather than deleting missed blocks.The comparison matches ThriftAttention’s 5% FP16 budget to sparse methods at a 71.3% sparsity ratio; sparse methods can incur large tail errors when blocks are removed entirely.Source: §4.5 · pp. 8–9, Table 3 · p. 8

5 Limitations and Future Work

ThriftAttention’s current kernel targets consumer Blackwell GPUs and inference, with higher KV-cache memory use; future work could extend the kernel to data-center Blackwell and explore selective FP16 attention for training.

5 Limitations and Future Work

  • Extending the consumer-Blackwell kernel to data-center Blackwell could exploit increased asynchrony and move FP4 performance closer to the 4× theoretical throughput advantage over FP16.Source: §5 · p. 9
  • Storing both FP16 and FP4 caches increases KV-cache memory footprint by 28%.Source: §5 · p. 9
  • ThriftAttention currently accelerates inference only; selectively promoting sensitive interactions to FP16 in forward and backward attention could help address stability issues in sub-byte training.Source: §5 · p. 9

6 Conclusion

ThriftAttention selectively promotes a small number of attention blocks to FP16 during long-context FP4 inference, preventing systematic performance degradation. Across model families and benchmarks, it approaches FP4 latency while preserving near-FP16 quality.

  • ThriftAttention selectively promotes a small number of attention blocks to FP16 for long-context FP4 inference.The conclusion presents selective mixed precision as the method's central contribution.Source: §6 · p. 9
  • Across model families and evaluation benchmarks, selective FP16 promotion prevents systematic performance degradation in long-context settings.Source: §6 · p. 9
  • Selective precision offers a practical path to long-context inference that approaches FP4 latency while preserving near-FP16 quality.Source: §6 · p. 9

A Heuristic Ablation

The heuristic ablation compares ThriftAttention’s block selection with random and diagonal-only selection, and reports that ThriftAttention performs better on a subset of HELMET tasks.

A Heuristic Ablation

  • The ablation compares the ThriftAttention heuristic against randomly promoted blocks and diagonal-only selection.Source: A Heuristic Ablation · p. 14
  • ThriftAttention outperforms random and diagonal block selection on the evaluated HELMET tasks.The comparison uses a 5% FP16 budget at sequence length 65,536.Source: Table 4 · p. 14
  • The evaluated subset comprises json kv, kilt popqa 3, and long narrative qa.Source: Table 4 · p. 14

B Experiment Design

The experiments evaluate ThriftAttention across long-context benchmarks and sequence lengths, using fixed block sizes and FP16 budgets alongside specified inference and NLL protocols. The reported downstream evaluations used a single Blackwell GPU and approximately 600 GPU-hours.

Code and environment

  • Experiments ran on one 96GB NVIDIA RTX PRO 6000 Blackwell GPU with CUDA 12.8 and PyTorch 2.8.0, using approximately 600 GPU-hours for downstream benchmarks and 5 for NLL analysis.The paper reports no larger internal cluster or additional undisclosed compute for its results.Source: B Experiment Design · p. 14

Context, and decoding

  • The maximum context was 131,072 tokens, clamped by model limits; generation used greedy argmax decoding.Qwen3 used base RoPE through 32,768 tokens and YaRN beyond that, while Llama and Ministral used checkpoint RoPE settings.Source: B Experiment Design · p. 14

ThriftAttention settings

  • ThriftAttention used 64-token query and key blocks with target FP16 budgets of 5%, 10%, and 25%.For n = ⌊L/64⌋ key blocks, integer top-k was rounded and clamped to [1, n].Source: B Experiment Design · p. 14

Benchmarks

  • Evaluation covered LongBench v1, RULER, and HELMET, with benchmark-specific task sets and context lengths up to 131,072 tokens.LongBench used the English subset and overall.weighted avg.; RULER used 100 samples per task-length cell across 13 tasks, and HELMET used 50 across its listed task categories.Source: B Experiment Design · p. 14

NLL and sequence-length analysis

  • NLL was measured on Qwen3-8B at seven sequence lengths from 2,048 to 131,072 using teacher-forced next-token cross-entropy on 300 packed PG-19 sequences.The sequences used seed 42 with EOS inserted between documents; sequence-length analysis used the same HELMET settings and greedy decoding.Source: B Experiment Design · p. 14

C Short-context benchmark results

This section evaluates ThriftAttention on short-context reasoning and knowledge benchmarks, with Table 5 covering accuracy across four models at varying top-k values.

C Short-context benchmark results

  • The evaluation covers BBH, MMLU-Pro, and GSM8K; Table 5 reports accuracy across four models at varying top-k values, with Recov. denoting percentage recovery from FP4 to FP16 performance.The supplied passages do not include Table 5's cell values, so they do not establish comparative results.Source: C Short-context benchmark results · p. 15, Table 5 · p. 15

Related papers

Loading 2605.23081v1…