ThriftAttention: Selektive gemischte Präzision für FP4-Attention mit langem Kontext

Zusammenfassung

Effiziente Aufmerksamkeitsalgorithmen sind entscheidend, um die quadratischen Kosten der Aufmerksamkeit in langen Kontexten zu reduzieren. Frühere Arbeiten nutzen blockweise skalierte Quantisierungstechniken auf Blackwell-GPUs, um die Aufmerksamkeitsberechnung auf 4-Bit-Präzision zu verlagern und die Inferenz zu beschleunigen. Diese Techniken führen jedoch in Langkontext-Szenarien zu erheblichen Qualitätseinbußen. Wir zeigen, dass die Auswirkung von Quantisierungsfehlern auf die Ausgabe stark nicht-uniform ist und mit der Bedeutung jeder Query-Key-Interaktion zunimmt, wobei funktional relevante Fehler in einer kleinen Anzahl von Aufmerksamkeitsblöcken konzentriert sind, die die wichtigsten Token enthalten. Wir schlagen ThriftAttention vor, eine Low-Bit-Aufmerksamkeitsvariante, die nahezu FP16-Qualität für lange Kontexte bei FP4-Inferenz-Effizienz liefert. Dieser Ansatz erfolgt in zwei Schritten. Zunächst wählt eine Heuristik schnell eine kleine Anzahl wichtiger Query-Key-Blockpaare für FP16-Präzision aus. Anschließend werden die ausgewählten Blöcke in FP16 und die restlichen Blöcke in FP4 berechnet, wobei beide Pfade über Online-Softmax zu einer einzigen Ausgabe zusammengeführt werden. Wir demonstrieren über Langkontext-Benchmarks und Modellfamilien hinweg, dass ThriftAttention durch die Berechnung von nur 5% der Query-Key-Blöcke in FP16 im Durchschnitt 89,1% des FP4-zu-FP16-Leistungsabstands wiederherstellt. Wir zeigen, dass der Vorteil von ThriftAttention mit der Sequenzlänge wächst und die systematische FP4-Qualitätsverschlechterung bei längeren Kontexten abmildert. Der Code ist verfügbar unter https://github.com/joesharratt1229/ThriftAttention.

English

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.