Papers
Topics
Authors
Recent
Search
2000 character limit reached

ThriftAttention: Mixed-Precision Long-Context Inference

Updated 4 July 2026
  • ThriftAttention is a selective mixed-precision attention mechanism that partitions the attention matrix into fixed blocks to address quality collapse in long-context inference.
  • It uses a heuristic based on mean query–key scores to promote a small percentage of blocks (as low as 5%) to FP16 while computing remaining blocks in FP4, balancing precision and efficiency.
  • Empirical results show it recovers up to 89.1% of the FP4-to-FP16 performance gap and accelerates decoding by up to 5.5×, making it a practical choice for long-context applications.

ThriftAttention is a training-free, selective mixed-precision attention mechanism for long-context inference that targets the quality collapse of uniform FP4 attention while preserving most of its efficiency benefits. It partitions the attention matrix into fixed query–key blocks, rapidly identifies a small subset of important blocks, computes those blocks in FP16, computes the remaining blocks in FP4, and merges both paths with a single online softmax. In the reported long-context setting, computing only 5% of query–key blocks in FP16 recovers on average 89.1% of the FP4-to-FP16 performance gap, and the advantage increases with sequence length (Sharratt, 21 May 2026).

1. Problem setting and motivation

ThriftAttention is situated in the long-context regime, where standard scaled dot-product attention

Attention(Q,K,V)=softmax ⁣(QKdk)V\text{Attention}(Q,K,V)=\mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V

becomes a dominant bottleneck because forming QKQK^\top costs O(N2dk)O(N^2 d_k) and the associated memory traffic scales with long prefixes. For large-context LLM inference, this makes attention a major wall-clock and bandwidth bottleneck, particularly during decoding, where each new token attends to the full prefix (Sharratt, 21 May 2026).

The immediate systems motivation is NVIDIA Blackwell support for native FP4 tensor cores. FP4 attention is attractive because FP4 matmul instructions provide roughly 4×4\times the arithmetic throughput of FP16 tensor cores on the same chip, and storing QQ, KK, and VV in FP4 can reduce KV-cache memory traffic by about 4×4\times. Prior work such as SageAttention3 showed that attention kernels can be moved almost entirely into FP4 using microscaling, but long-context quality degrades substantially under naive uniform FP4 quantization (Sharratt, 21 May 2026).

The central observation behind ThriftAttention is that this degradation is not spatially uniform across the attention matrix. The output impact of quantization error is highly non-uniform and increases with the importance of each query–key interaction. Functionally relevant error concentrates in a small number of attention blocks containing the most important tokens, especially near-diagonal blocks and “attention sinks.” This motivates selective precision promotion rather than uniform precision assignment (Sharratt, 21 May 2026).

2. Analytical basis: non-uniform error concentration

For a single query vector qq attending over keys kjk_j with values QKQK^\top0, the attention output is

QKQK^\top1

If FP4 quantization perturbs the scores as QKQK^\top2, then to first order

QKQK^\top3

Hence

QKQK^\top4

The decisive term is the multiplicative factor QKQK^\top5: even if quantization noise QKQK^\top6 is roughly uniform across positions, the contribution of that noise to output error is weighted by attention mass. Low-probability positions are naturally damped, whereas high-probability positions amplify their own quantization error (Sharratt, 21 May 2026).

This leads to a block-level formulation. ThriftAttention partitions the sequence into blocks and treats a query-block–key-block pair as the basic unit of mixed-precision allocation. Because important interactions occupy a small subset of blocks, promoting only those blocks to FP16 can recover most of the lost quality while leaving the majority of interactions in FP4. The paper’s empirical heatmaps of

QKQK^\top7

show that the largest discrepancies are concentrated in a small number of blocks per query, not distributed uniformly over the full QKQK^\top8 matrix (Sharratt, 21 May 2026).

A common misconception is that ThriftAttention is a sparsity method in the ordinary sense. It is not. It never drops any query–key block. Unimportant blocks are still computed, but in FP4 rather than FP16. This distinction is central to its accuracy profile: the approximation error is bounded by quantization noise rather than by deleting interactions outright (Sharratt, 21 May 2026).

3. Block heuristic, quantization scheme, and online softmax

ThriftAttention groups tokens into fixed-size blocks along the sequence axis. With QKQK^\top9, the matrices are partitioned as

O(N2dk)O(N^2 d_k)0

with O(N2dk)O(N^2 d_k)1 and O(N2dk)O(N^2 d_k)2. The conceptual attention score block is

O(N2dk)O(N^2 d_k)3

Importance selection is heuristic rather than learned. For each query block O(N2dk)O(N^2 d_k)4, ThriftAttention computes the mean query

O(N2dk)O(N^2 d_k)5

and for each key block O(N2dk)O(N^2 d_k)6, the mean key

O(N2dk)O(N^2 d_k)7

The block importance score is

O(N2dk)O(N^2 d_k)8

For each query block, the top-O(N2dk)O(N^2 d_k)9 key blocks under 4×4\times0 are designated important and computed in FP16; all others are computed in FP4 (Sharratt, 21 May 2026).

The FP4 path uses NVFP4 microscaling. For 4×4\times1, ThriftAttention stores an FP4 data tensor 4×4\times2 and an FP8 scale tensor 4×4\times3, with one FP8 scale shared by each group of 16 channels in a row. Quantization is applied independently to 4×4\times4, 4×4\times5, and 4×4\times6: 4×4\times7 Selected blocks use FP16 score and value products; non-selected blocks use FP4 matmuls with the corresponding FP8 scales (Sharratt, 21 May 2026).

The two precision paths are merged by an online softmax in the style of FlashAttention-2. For each query block 4×4\times8, the kernel maintains a running rowwise maximum 4×4\times9, a running denominator QQ0, and a partial output QQ1 while scanning key blocks. FP16 and FP4 contributions update the same running softmax state, so the final output corresponds to a single softmax over all blocks rather than separate normalizations. This is the mechanism that preserves exact softmax structure while changing numerical precision locally (Sharratt, 21 May 2026).

For causal attention, the same scheme applies with the standard causal restriction: only causally visible key blocks are eligible for scoring and selection. The method is therefore a drop-in replacement for standard causal scaled dot-product attention rather than a different attention operator (Sharratt, 21 May 2026).

4. Empirical results

The reported experiments use Llama3.2-3B, Llama3.1-8B, Qwen3-4B, Qwen3-8B, and Ministral-3-8B-Base-2512 on a single NVIDIA RTX PRO 6000 Blackwell GPU with 96GB RAM, CUDA 12.8, and PyTorch 2.8.0. Generation uses greedy decoding, maximum context length is 131k tokens, and block sizes are fixed at QQ2 (Sharratt, 21 May 2026).

Evaluation spans LongBench-v1, HELMET, RULER, and PG-19. The principal quantitative result is that pure FP4 attention performs far below FP16 on long-context benchmarks, whereas ThriftAttention with 5% FP16 blocks recovers on average 89.1% of the FP4-to-FP16 gap; 10% and 25% FP16 budgets recover 91.8% and 92.4%, respectively. For Qwen3-8B on RULER, the scores are FP16 QQ3, FP4 QQ4, ThriftAttention 5% QQ5, and ThriftAttention 10% QQ6. On LongBench-v1 the corresponding scores are QQ7, QQ8, QQ9, and KK0. On HELMET they are KK1, KK2, KK3, and KK4, with the 10% setting slightly surpassing FP16 for that model/benchmark pair (Sharratt, 21 May 2026).

The long-context trend is especially important. On HELMET, pure FP4 retention degrades with sequence length: for Llama3.1-8B, the FP4 score is KK5 versus FP16 KK6 at 8k, but KK7 versus KK8 at 131k. PG-19 NLL analysis shows that the per-token KK9 between FP16 and FP4 grows from about VV0 at shorter contexts to about VV1 at 64k–128k tokens. ThriftAttention at 5% keeps VV2 across lengths, reducing the error by roughly VV3–VV4 relative to pure FP4 (Sharratt, 21 May 2026).

Efficiency remains close to the FP4 regime. On RTX PRO 6000 Blackwell, ThriftAttention’s attention kernel is up to VV5 faster than FlashAttention-2 in prefill, with end-to-end prefill latency improvement around VV6 versus FP16 attention at 131k context. In decode, ThriftAttention kernels achieve VV7–VV8 speedup versus FlashAttention-2, and end-to-end generation speedup is close to VV9 for Qwen3-8B at 131k context. Relative to pure FP4 attention, the overhead is small (Sharratt, 21 May 2026).

The heuristic itself is materially better than naive alternatives. On a HELMET subset at 65,536 tokens with a 5% FP16 budget, the ThriftAttention heuristic scores 4×4\times0, compared with 4×4\times1 for random block selection and 4×4\times2 for diagonal-only promotion (Sharratt, 21 May 2026).

5. Position within efficient attention and KV-cache research

ThriftAttention belongs to the quantized-attention branch of efficient inference, but it is best understood in relation to several adjacent families. “Top-Theta Attention” replaces dynamic top-4×4\times3 selection with calibrated per-layer, per-head, per-row thresholds and combines sparsification with Softmax Denominator Compensation and V-Mean Compensation; its emphasis is threshold-based pruning and reduced V-cache row access, not mixed-precision preservation of dense attention (Berestizshevsky et al., 12 Feb 2025). ADORE maintains a fixed-size KV cache, uses a GRU-based controller to approximate ideal top-4×4\times4 sparse attention, and rebuilds released tokens on demand; its objective is token retention and cache management rather than numeric precision allocation over a dense block grid (Zhang et al., 2024).

Other works operate on orthogonal axes of the same systems problem. GTA reduces attention FLOPs by sharing attention maps across grouped heads and compressing values into a latent space with a nonlinear decoder, thereby shrinking the KV cache and accelerating both prefill and decode (Sun et al., 15 Jun 2025). StiefAttention performs post-training low-rank KV-cache compression by learning orthonormal projection bases on the Stiefel manifold that minimize decoder-layer output reconstruction error (Benfenati et al., 29 Jan 2026). Transactional Attention introduces a sponsorship term 4×4\times5 to protect dormant tokens such as credentials from KV eviction and explicitly does not modify 4×4\times6 or softmax (Basu, 13 Apr 2026).

These contrasts clarify what ThriftAttention does not do. It does not sparsify away interactions, does not compress KV states via projection, and does not solve dormant-token eviction. It preserves the dense softmax computation structurally and intervenes only in the precision assigned to different blocks. A plausible implication is that it is most naturally combined with methods from those other families, because its control variable—precision—is orthogonal to token-selection, retention, and low-rank compression policies.

A broader interpretability perspective also helps contextualize the method. In hybrid SSM–Transformer models, retrieval has been shown to depend exclusively on self-attention layers, and sparsifying attention to just 15% of heads can maintain near-perfect retrieval while preserving 84% MMLU performance (Michalak et al., 21 Oct 2025). This suggests that preserving the quality of a small set of high-impact interactions can be disproportionately important, which is consistent with ThriftAttention’s block-level precision allocation, although the two interventions operate at different granularities.

6. Limitations, scope, and future directions

ThriftAttention is currently hardware-targeted. The implementation is optimized for consumer Blackwell, and the paper explicitly notes that data-center Blackwell with SM100 features may permit further speedups that are not explored. It also incurs a KV-cache memory overhead of about 28% relative to pure FP4 attention because it stores both FP4 and FP16 KV representations, or otherwise requires access to FP16 activations for promoted blocks (Sharratt, 21 May 2026).

The method is inference-only in its present form. The paper does not develop or evaluate a training-time analogue, and it notes that existing FP4 training methods often keep attention at higher precision because of stability concerns. The selection heuristic is also fixed and heuristic: 4×4\times7 is effective, but it is not claimed to be optimal, and tasks with unusual attention structure may require richer importance signals (Sharratt, 21 May 2026).

A second misconception is that a larger FP16 budget should always be expected to deliver proportionally better quality. The reported results show diminishing returns: 5% already recovers most of the FP4-to-FP16 gap on average, and moving from 10% to 25% yields only modest additional recovery. The underlying reason is precisely the concentration phenomenon that motivates the method: most impactful blocks are already captured in the top few percent (Sharratt, 21 May 2026).

Future directions identified in the paper include richer block-importance predictors, extension to data-center Blackwell kernels, and selective promotion during training so that FP4 attention could be stabilized in both forward and backward passes. More generally, ThriftAttention defines a design pattern for long-context inference: treat attention precision as a locally allocable resource, preserve dense softmax structure, and spend FP16 only where quantization error is functionally amplified (Sharratt, 21 May 2026).

Topic to Video (Beta)

No one has generated a video about this topic yet.

Whiteboard

No one has generated a whiteboard explanation for this topic yet.

Follow Topic

Get notified by email when new papers are published related to ThriftAttention.