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

Long-context FP4 attention improves efficiency but can substantially degrade quality, while sparse alternatives omit blocks and incur irrecoverable error. ThriftAttention selectively promotes important query-key blocks to FP16, merges them with FP4 computation, and recovers most of the FP4-to-FP16 gap with a small FP16 budget. Its advantage increases with sequence length, while the current implementation remains limited to inference and increases KV-cache memory footprint.

  • Problem

    FP4 attention improves inference efficiency but causes significant long-context quality degradation, while sparse methods can require aggressive block omission with irrecoverable error.

  • Method

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

  • Results

    Across long-context benchmarks and model families, a 5% FP16 block budget recovers 89.1% of the FP4→FP16 performance gap on average.

  • Takeaways & Limitations

    ThriftAttention approaches FP4 latency while preserving near-FP16 quality, with relative improvement over FP4 increasing as context length grows.

  • Takeaways & Limitations

    The current kernel targets consumer Blackwell GPUs, increases KV-cache memory footprint by 28%, and is designed for inference acceleration only.

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

Long-context attention faces a quality–efficiency tension: FP4 improves inference efficiency but degrades quality, while sparse methods must omit many blocks. ThriftAttention selectively computes important interactions in FP16 and the remainder in FP4 to preserve quality efficiently.

  • Motivation: Blackwell FP4 Tensor Cores offer 4× FP16 arithmetic throughput while reducing KV-cache memory traffic, but low-bit attention introduces a quality–efficiency tension.Attention’s quadratic cost and KV-cache traffic dominate long-context execution time.
  • Motivation: Sparse methods must drop at least 75% of KV blocks to match FP4 latency during generation, making omitted-block error irrecoverable.This aggressive sparsity can substantially degrade inference-only performance.
  • Approach: ThriftAttention concentrates FP16 computation on query-key blocks with high-magnitude attention scores and strong influence on the final output distribution.The approach is training-free and targets degradation in long-context low-bit attention.
  • Approach: A lightweight heuristic scores query-key block pairs using token-mean query and key vectors, selecting the top-k blocks for FP16 and assigning the rest to FP4.Both precision paths are merged into one output through online softmax.
  • Results: At a 5% FP16 block budget, ThriftAttention recovers 89.1% of the FP4→FP16 performance gap on average across long-context evaluations.Recovery reaches 91.8% and 92.4% at 10% and 25% budgets, respectively.
  • Results: The method is evaluated across LongBench-v1, HELMET, RULER, and PG-19 using Llama, Qwen, and Ministral model families.These evaluations assess whether a small FP16 budget preserves FP16 quality while retaining low-bit efficiency.

2 Related Work

Prior work accelerates attention through I/O-efficient kernels, quantisation, KV-cache compression, and sparsity, but ThriftAttention positions selective precision as a distinct alternative.

  • I/O-Efficient Attention: FlashAttention and later variants reduce GPU memory I/O through tiling, improved parallelism, and hardware-specific optimisations.This line of work targets efficient attention execution rather than selective numeric precision.
  • Quantised Attention: Quantised attention methods extend post-training quantisation from linear layers to attention, including FP4 attention on Blackwell using two-level microscaling.Related approaches also target KV-cache compression or combine quantised matrix multiplications with sparsity.
  • Sparse Attention: Sparse attention methods select or predict important KV blocks, while SpargeAttn combines block-sparsity prediction with quantised INT8/FP8 attention.SpargeAttn is described as the closest prior work but does not use sub-8-bit numeric formats.
  • Positioning: ThriftAttention assigns full precision to important blocks instead of uniformly quantising all blocks or skipping them.For each block, its error is bounded by FP4 quantisation noise rather than by the magnitude of skipped or approximated attention scores.

3 Method

ThriftAttention identifies query-key blocks likely to matter most, computes those blocks in FP16, and processes the remainder in FP4 before merging both paths online. Its design targets the non-uniform impact of FP4 quantisation error while preserving efficient attention execution.

  • Error analysis: Each key’s output-error contribution combines score quantisation error, attention weight, and value deviation from the output.The attention weight makes error non-uniform across keys.
  • Error analysis: High pre-softmax scores produce larger attention weights and amplify their quantisation errors, whereas low scores dampen their effects.This motivates prioritising interactions with large attention scores.
  • Error analysis: Quantisation error concentrates in a small number of query-key blocks, typically near-diagonal blocks and non-initial attention sinks.The pattern is illustrated across representative layers and heads in Qwen3-8B at sequence length 4096.
  • Block selection: ThriftAttention partitions Q, K, and V into query-key blocks and computes token-wise means for block-importance scoring.The algorithm represents Q, K, and V as block sequences before selecting key blocks.
  • Mixed-precision computation: For each query block, the top-k scored key blocks use FP16 while remaining pairs use FP4, with their outputs accumulated through online softmax.The selected set is routed to the FP16 path and the rest to the FP4 path before online merging.
  • Hardware implementation: A fused kernel processes non-selected blocks through FP4 and selected blocks through a separate FP16 helper routine.The implementation aliases shared memory across paths and lets warps without selected blocks bypass the FP16 path.

4 Experiments

Experiments evaluate ThriftAttention across long-context benchmarks, model families, sequence lengths, and matched-compute sparse-attention baselines. The results show that small FP16 budgets recover most FP4 quality, with larger benefits at longer contexts and competitive efficiency.

  • Efficiency: ThriftAttention achieves up to 1.7× kernel speedup over FlashAttention-2 for prefill and nearly 2× end-to-end generation speedup over full-FP16 attention at 131k context.Decode kernels are 3×–5.5× faster than FlashAttention-2, with minimal overhead compared with full FP4 attention.
  • Benchmark evaluation: ThriftAttention recovers 89.1% of the FP4–FP16 performance gap with 5% FP16 blocks, rising to 91.8% and 92.4% at 10% and 25%.The evaluation spans five models and long-context benchmarks; additional FP16 blocks provide diminishing returns after the most important blocks are promoted.
  • Benchmark evaluation: Benchmark response to FP16 budget varies: LongBench-v1 saturates at 5%, RULER improves steadily through 25%, and HELMET scales unevenly.These results show that the quality returned by additional FP16 blocks depends on the benchmark.
  • Sequence-length scaling: At 5% FP16 budget, relative improvement over FP4 increases from 2.00× at 8k to 2.2× at 131k on HELMET.FP4 retention for Llama3.1-8B falls from 50% at 8k to 32% at 131k, while ThriftAttention becomes relatively more effective.
  • Negative log-likelihood: Across PG-19 context lengths, ThriftAttention reduces ∆NLL to ≤0.02, while FP4 ∆NLL reaches 0.10 for end-of-sequence tokens at 64k and 128k.The relative ∆NLL reduction grows from approximately 2× at 8k to approximately 5× at the end of a 128k sequence.
  • Matched-compute comparison: At equivalent inference efficiency, ThriftAttention outperforms inference-time sparsity approaches by retaining all interactions in FP4 instead of deleting missed blocks.The matched-compute comparison uses a 5% FP16 budget for ThriftAttention and a 71.3% sparsity ratio for sparse baselines.

5 Limitations and Future Work

The current implementation and scope impose hardware, memory, and application boundaries. Future work includes data-center Blackwell support and selective mixed-precision attention for training.

  • Hardware and memory: The current kernel targets consumer Blackwell GPUs, while data-center Blackwell support could exploit increased asynchrony for higher FP4 throughput.The authors also report a 28% KV-cache memory increase from storing FP16 and FP4 caches.
  • Application scope: ThriftAttention is currently designed for inference acceleration only, leaving selective FP16 promotion during attention training as future work.The proposed training direction is motivated by potential stability issues in sub-byte attention training.

6 Conclusion

ThriftAttention selectively promotes a small number of attention blocks to FP16 for long-context FP4 inference, preventing systematic quality degradation while approaching FP4 latency and near-FP16 quality.

  • ThriftAttention selectively promotes only a small number of attention blocks to FP16 while computing the remainder in FP4.
  • The approach prevents systematic performance degradation in long-context settings across model families and evaluation benchmarks.
  • Selective precision offers a practical path toward long-context inference that approaches FP4 latency while preserving near-FP16 quality.

A Heuristic Ablation

The heuristic ablation compares ThriftAttention’s block selection against random and diagonal-only alternatives, showing that the proposed heuristic performs better on selected long-context tasks.

  • ThriftAttention’s heuristic is compared against randomly promoted blocks and diagonal-only selection.
  • At sequence length 65,536 with a 5% FP16 budget, ThriftAttention outperforms random and diagonal block selection on three HELMET tasks.The evaluated tasks are json kv, kilt popqa 3, and long narrative qa.

B Experiment Design

Experiments use a single Blackwell GPU and evaluate multiple model checkpoints under long-context settings up to 131,072 tokens, with specified precision budgets and benchmark protocols.

  • Experiments use CUDA 12.8, PyTorch 2.8.0, and a single NVIDIA RTX PRO 6000 Blackwell GPU with 96GB memory.Reported downstream evaluations required approximately 600 GPU-hours, while NLL analysis required approximately 5 GPU-hours.
  • The evaluated checkpoints include Llama3.2, Llama3.1, Qwen3, and Ministral 3 models ranging from 3B to 8B parameters.
  • The maximum context is 131,072 tokens, with model-specific RoPE settings and greedy argmax decoding.
  • All experiments use 64-token query and key blocks with target FP16 budgets of 5%, 10%, and 25%.For n = ⌊L/64⌋ key blocks, the integer top-k is rounded and clamped to [1, n].
  • LongBench v1 reports overall weighted average scores, while RULER evaluates 13 tasks across lengths from 4,096 to 131,072 tokens.
  • NLL analysis evaluates Qwen3-8B on PG19 using teacher-forced next-token cross-entropy across sequence lengths from 2,048 to 131,072 tokens.The analysis uses 300 packed sequences and does not use generation.

C Short-context benchmark results

Short-context evaluation covers BBH, MMLU-Pro, and GSM8K, with downstream accuracy reported across four models and varying top-k settings.

  • Short-context reasoning and knowledge evaluation uses BBH, MMLU-Pro, and GSM8K.
  • Table 5 reports ThriftAttention downstream accuracy across four models and varying top-k values on these three benchmarks.The table’s Recov. metric denotes percentage recovery from FP4 to FP16 performance.
Loading 2605.23081v1…