ThriftAttention: Selective Mixed Precision for Long-Context FP4 Attention

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.

📋 한 줄 요약

**[FP4 Attention / Long Context]** ThriftAttention이 중요 query-key 5%만 FP16·나머지 FP4로 계산하고 online softmax로 병합 — Blackwell FP4 quality degradation을 long context에서 평균 89.1%·FP4→FP16 gap 복구, 시퀀스 길이 증가에 따라 이점 증가.

🎯 핵심 기여도

Long-context 워크로드의 quadratic attention 비용 완화를 위한 efficient attention 알고리즘의 critical 성을 동기.
선행 연구가 Blackwell GPU에서 block-scaled quantization으로 attention을 4-bit 정밀도로 옮겨 inference 가속하지만, long-context 환경에서 significant quality degradation을 일으킨다는 한계 진단.
Quantization error의 output impact가 highly non-uniform이며 each query-key interaction의 importance와 함께 커진다 — functionally relevant error가 most important token을 포함한 소수 attention block에 집중됨을 발견.
ThriftAttention 제안 — low-bit attention variant로 FP4 inference 효율로 near-FP16 long-context 품질 달성. 2단계: (1) heuristic이 small number의 important query-key block pair를 FP16 정밀도로 빠르게 선택, (2) 선택된 block은 FP16·나머지는 FP4로 계산해 online softmax로 단일 출력에 병합.

💡 핵심 아이디어

Long-context attention의 FP4 quality degradation은 모든 block을 균일 처리하는 것이 비효율이며, 소수의 functionally important block에 quantization error가 집중된다는 비균일성을 활용해 단 5%만 FP16, 나머지는 FP4로 처리하고 online softmax로 병합하면 거의 모든 quality gap을 복구할 수 있다.

🔬 기술적 접근법

**방법론**: ThriftAttention — selective mixed precision FP4 attention.
**핵심 기법**: (1) Quantization error의 비균일 분포 관찰 — importance가 큰 block에 functionally relevant error 집중, (2) Heuristic으로 important query-key block pair 빠르게 선택, (3) 선택 block은 FP16·나머지는 FP4 계산, (4) Online softmax로 두 path를 single output으로 병합, (5) Blackwell GPU의 block-scaled quantization과 호환.

📊 주요 결과

Long-context 벤치마크·model family에서 query-key block의 단 5%만 FP16 계산.
FP4→FP16 performance gap을 평균 89.1% 복구.
ThriftAttention의 이점이 시퀀스 길이와 함께 증가 — 긴 context에서 관찰되는 systematic FP4 quality degradation 완화.
Code 공개: github.com/joesharratt1229/ThriftAttention.

💭 의의 및 한계

**의의**: Blackwell FP4 inference의 실용성 끌어올림, quantization error의 importance-aware 분포 관찰을 architecture-level 알고리즘으로 전환, long context inference 비용 절감과 품질 보존 동시. **한계**: Heuristic-based block selection의 robustness는 task별 차이 가능, 5% 선택 비율 hyperparameter, FP4 hardware(Blackwell) 종속.

🚀 실용적 활용

Long-context LLM inference 비용 절감.
Blackwell GPU 기반 production deployment.
Mixed-precision attention 일반 디자인 패턴.