Technical Summary: ThriftAttention

Problem Statement

Efficient inference for Large Language Models (LLMs) in long-context workloads is currently bottlenecked by the quadratic computational cost and memory traffic of the attention mechanism. While NVIDIA's Blackwell architecture introduces native FP4 Tensor Cores offering 4× arithmetic throughput over FP16, utilizing these for attention via uniform quantization (e.g., SageAttention3) results in significant quality degradation, particularly as sequence lengths increase. Conversely, existing sparsity methods that attempt to mitigate this by dropping KV blocks entirely often require aggressive sparsity ratios (dropping >75% of blocks) to match FP4 latency, leading to irrecoverable performance losses. There is a fundamental tension between inference efficiency and output quality in long-context settings.

Methodology: ThriftAttention

The authors propose ThriftAttention, a training-free, selective mixed-precision attention mechanism designed to deliver near-FP16 quality at FP4 inference efficiency. The core insight driving this method is that quantization error is not uniformly distributed; rather, the impact of error is highly non-uniform and concentrates in a small number of query-key block pairs where attention scores are highest and most functionally relevant to the final output distribution.

The algorithm proceeds in two stages:

1. Heuristic Block Selection:

   • The query ($Q$) and key/value ($K, V$) tensors are partitioned into blocks.
   • A lightweight heuristic computes an importance score for each query-key block pair $(i, j)$ using the dot product of the token means: $\hat{S}_{ij} = \bar{q}_i \cdot \bar{k}_j$, where $\bar{q}_i$ and $\bar{k}_j$ are the mean vectors of the respective blocks.
   • The top-$k$ highest-scoring block pairs are identified for high-precision computation.

2. Mixed-Precision Computation & Merging:

   • FP16 Path: The selected top-$k$ block pairs are computed in full FP16 precision.
   • FP4 Path: The remaining blocks are computed using FP4 quantization (utilizing the NVFP4 microscaling format supported on Blackwell GPUs).
   • Online Merging: Both paths are computed and merged via an online softmax procedure (similar to FlashAttention-2) to produce a single output. The FP4 path utilizes a two-level quantization scheme for the probability blocks, adapted from SageAttention3, to maintain numerical stability during accumulation.

The implementation is realized as a single fused CUDA kernel. To manage register pressure, the kernel processes non-selected blocks via the FP4 path first, then utilizes a separate FP16 helper routine for the promoted blocks, loading FP16 query tiles only when necessary. Warps that contain no top-$k$ blocks bypass the FP16 path entirely to avoid unnecessary memory loads.

Key Contributions

• Selective Mixed-Precision Attention: The introduction of ThriftAttention, which computes the most critical block interactions in FP16 while retaining the remainder in FP4. The authors note this is the first work to utilize sub-byte formats in this specific mixed-precision manner for attention computation.
• Training-Free Approach: The method requires no model retraining or fine-tuning, relying solely on a heuristic selection mechanism during inference.
• Comprehensive Evaluation: The approach is evaluated across diverse model families (Llama, Qwen, Ministral) and long-context benchmarks (LongBench-v1, HELMET, RULER, PG-19), demonstrating efficacy across varying architectures and context lengths.

Results

• Quality Recovery: At a 5% FP16 block budget (computing only 5% of query-key blocks in FP16), ThriftAttention recovers an average of 89.1% of the performance gap between FP4 and FP16 attention. This recovery increases to 91.8% and 92.4% at 10% and 25% budgets, respectively.
• Efficiency Gains:
  • Prefill: Achieves up to a 1.7× kernel speedup over FlashAttention-2 (FP16) and consistent end-to-end prefill improvements (approx. 1.2× at 131k context).
  • Decode: Achieves a 3×–5.5× speedup over FlashAttention-2, translating to nearly a 2× end-to-end generation speedup over full FP16 attention at 131k context lengths.
• Context Length Scaling: The advantage of ThriftAttention grows with sequence length. While uniform FP4 attention degrades significantly as context length increases (e.g., FP4 retention drops from 50% at 8k to 32% at 131k for Llama3.1-8B), ThriftAttention's relative improvement over FP4 increases, reaching 2.2× at 131k.
• Comparison to Sparsity: When matched for compute (FP16-equivalent FLOPs), ThriftAttention significantly outperforms inference-time sparse attention baselines (e.g., Quest, Sparse Top-k). The authors attribute this to the fact that sparse methods completely discard blocks, incurring tail errors, whereas ThriftAttention retains all interactions in FP4, degrading them smoothly via quantization rather than deletion.

Significance and Claims

The paper claims that ThriftAttention offers a practical path toward long-context inference that approaches FP4 latency while preserving near-FP16 quality. By identifying that quantization error is concentrated in specific high-importance interactions, the method mitigates the systematic quality degradation observed in uniform low-bit attention at long contexts.

The authors position this work as a solution to the "fundamental tension" between efficiency and quality in Blackwell-era inference. They emphasize that preserving the full attention support in low precision (FP4), rather than aggressively sparsifying and computing only a subset in high precision, yields superior results at matched compute budgets. The approach is specifically targeted at consumer and data-center Blackwell GPUs to leverage native FP4 Tensor Cores, though the authors note current limitations regarding KV-cache memory footprint (increased by 28% due to dual caching) and the current focus on inference rather than training.