PagedAttention

PagedAttention ist ein Attention-Algorithmus und Speicherverwalter für das Bereitstellen von großen Sprachmodellen (LLMs). Er speichert den KV-Cache jeder Anfrage in kleinen Blöcken fester Größe, die an beliebiger Stelle im GPU-Speicher abgelegt werden. Die Blöcke werden über eine Tabelle pro Anfrage nachverfolgt – ähnlich wie ein Betriebssystem virtuelle Speicherseiten abbildet. Woosuk Kwon, Zhuohan Li und Kollegen an der UC Berkeley und Partneruniversitäten stellten PagedAttention gemeinsam mit der Bereitstellungs-Engine vLLM in der Studie Efficient Memory Management for Large Language Model Serving with PagedAttention vor, die auf der SOSP 2023 präsentiert wurde. Da keine zusammenhängenden Cache-Bereiche mehr zugewiesen werden müssen, entsteht keine entsprechende Fragmentierung. So kann eine GPU deutlich mehr Sequenzen gleichzeitig speichern. Die Studie berichtet bei gleicher Latenz von einem 2- bis 4-mal höheren Durchsatz als bei FasterTransformer und Orca.

Das gelöste Speicherproblem#

Bei der Textgenerierung behält ein Transformer die Schlüssel- und Wertvektoren aller vorherigen Tokens, damit sie nicht erneut berechnet werden müssen. Der Cache wächst pro Schicht mit jedem Token um einen Eintrag. Eine einzelne Sequenz des LLaMA-Modells mit 13 Milliarden Parametern kann bis zu 1,7 GB belegen. Beim Eingang einer Anfrage ist die Ausgabelänge unbekannt. Frühere Bereitstellungssysteme reservierten deshalb pro Anfrage einen zusammenhängenden Speicherbereich, der auf die maximal mögliche Länge ausgelegt war. Der größte Teil dieses Bereichs blieb leer, und die Lücken zwischen den Bereichen waren zu klein, um sie wiederzuverwenden. Das vLLM-Team maß bei solchen Systemen einen ungenutzten Anteil von 60 bis 80 % des KV-Cache-Speichers. Dadurch war die mögliche Batchgröße begrenzt, lange bevor die Rechenkapazität ausgeschöpft war.

So funktioniert Paging#

PagedAttention teilt den Cache in Blöcke auf, die jeweils Schlüssel und Werte für eine feste Anzahl von Tokens enthalten. In vLLM sind es standardmäßig 16. Eine Anfrage erhält erst dann einen neuen Block, wenn ihr letzter Block voll ist. Der neue Block kann sich an jeder Stelle befinden, an der freier Speicher verfügbar ist. Die Blocktabelle hält die Reihenfolge fest. Der Kernel des Attention-Mechanismus liest anhand dieser Tabelle die Schlüssel und Werte blockweise ein. Der Cache einer Sequenz muss daher nicht zusammenhängend gespeichert werden. Der Speicherverlust beschränkt sich auf den nicht gefüllten Rest des letzten Blocks jeder Sequenz. Laut den Messungen des vLLM-Teams liegt er unter 4 %.

Dieses PyTorch-Beispiel speichert die Schlüssel einer Anfrage verteilt auf Blöcken eines gemeinsam genutzten Pools und liest sie mithilfe der Blocktabelle wieder ein:

import torch

block_size, head_dim = 16, 64
key_pool = torch.zeros(64, block_size, head_dim)  # physical blocks shared by every request
free_blocks = [41, 7, 19, 2, 33]  # free blocks are scattered across the pool

# A 40-token request needs ceil(40 / 16) = 3 blocks, wherever they happen to be free
keys = torch.randn(40, head_dim)
num_blocks = -(-len(keys) // block_size)
block_table = [free_blocks.pop(0) for _ in range(num_blocks)]
for pos, key in enumerate(keys):
    key_pool[block_table[pos // block_size], pos % block_size] = key

# The attention kernel reads the request's keys back through its block table
gathered = key_pool[block_table].reshape(-1, head_dim)[: len(keys)]
print(f"block table {block_table}, keys intact: {torch.equal(gathered, keys)}")
print(f"unused slots: {num_blocks * block_size - len(keys)} paged vs {2048 - len(keys)} with a 2,048-token reservation")

Ausgegeben werden die Blocktabelle [41, 7, 19] und keys intact: True. Anschließend zeigt die Ausgabe 8 ungenutzte Plätze für die Anfrage mit Paging gegenüber 2.008 bei einem für 2.048 Tokens reservierten Bereich.

Gemeinsam genutzte Blöcke für Sequenzen#

Mithilfe von Blocktabellen können auch mehrere Sequenzen auf dieselben physischen Blöcke verweisen. Paralleles Sampling, Beam Search und Anfragen, die mit demselben System-Prompt beginnen, verwenden eine einzige Kopie des gemeinsamen Präfixes. Für jeden Block wird eine Referenzanzahl geführt. Muss eine Sequenz in einen Block schreiben, den andere noch verwenden, kopiert die Engine diesen Block zuerst. Dieses Copy-on-Write-Verfahren ist Betriebssystemen entlehnt.

Die Autoren von vLLM berichten, dass die gemeinsame Nutzung den Speicherbedarf von parallelem Sampling und Beam Search um bis zu 55 % senkt. Das ermöglicht einen bis zu 2,2-mal höheren Durchsatz. Durch das Hashing vollständiger Blöcke lässt sich das Verfahren auf mehrere Anfragen ausweiten. So implementiert vLLM Prompt-Caching.

Verwandte Verfahren#

PagedAttention legt fest, wo gecachte Schlüssel und Werte gespeichert werden. Das Verfahren lässt sich mit Methoden kombinieren, die die zu speichernde Datenmenge oder die Lesegeschwindigkeit verändern:

• Flash Attention: Teilt die Attention-Berechnung in Kacheln auf, sodass die vollständige Bewertungsmatrix nie in den GPU-Hauptspeicher geschrieben wird. Aktuelle Engines führen Kernel nach dem FlashAttention-Prinzip aus, die Daten aus paginierten Blöcken lesen.
• Kontinuierliches Batching: Nimmt eine neue Anfrage auf, sobald eine andere abgeschlossen ist. Durch Paging werden die Blöcke der abgeschlossenen Sequenz sofort wieder dem freien Pool hinzugefügt, sodass für die neue Anfrage Platz vorhanden ist.
• Attention mit gruppierten Abfragen: Teilt Schlüssel- und Wert-Heads zwischen Abfrage-Heads, wodurch der Inhalt jedes Blocks kleiner wird.
• FP8-Quantisierung des Caches: Speichert Schlüssel und Werte mit 8-Bit-Präzision. Dadurch halbiert sich der Cache-Speicherbedarf ungefähr im Vergleich zu 16-Bit-Formaten.

Die LLM-Schnittstelle von Ultralytics leitet Argumente von Anfragen unverändert weiter. Ein Chat-Completions-Aufruf mit n=3, der über base_url an einen lokalen vLLM-Server gesendet wird, gibt daher drei Stichproben zurück, die dieselben gecachten Prompt-Blöcke verwenden. Nur die voneinander abweichenden Fortsetzungen beanspruchen neuen Speicher.