Hydragen: Hochdurchsatz-LLM-Inferenz mit gemeinsamen Präfixen

Hydragen: High-Throughput LLM Inference with Shared Prefixes

February 7, 2024

Autoren: Jordan Juravsky, Bradley Brown, Ryan Ehrlich, Daniel Y. Fu, Christopher Ré, Azalia Mirhoseini

cs.AI

Zusammenfassung

Transformer-basierte große Sprachmodelle (LLMs) werden mittlerweile von Hunderten Millionen Nutzern eingesetzt. Die Inferenz von LLMs wird üblicherweise in Batches von Sequenzen durchgeführt, die ein gemeinsames Präfix teilen, wie beispielsweise Few-Shot-Beispiele oder ein System-Prompt für Chatbots. Das Decodieren in diesem Large-Batch-Setting kann durch den Attention-Mechanismus zum Engpass werden, der große Key-Value (KV)-Caches aus dem Speicher liest und ineffiziente Matrix-Vektor-Produkte für jede Sequenz im Batch berechnet. In dieser Arbeit stellen wir Hydragen vor, eine hardwarebewusste, exakte Implementierung von Attention mit gemeinsamen Präfixen. Hydragen berechnet die Attention für das gemeinsame Präfix und die einzigartigen Suffixe separat. Diese Zerlegung ermöglicht eine effiziente Präfix-Attention, indem Anfragen über Sequenzen hinweg gebündelt werden, redundante Speicherzugriffe reduziert werden und hardwarefreundliche Matrixmultiplikationen genutzt werden können. Unsere Methode kann den end-to-end-Durchsatz von LLMs um bis zu 32x im Vergleich zu konkurrierenden Baselines steigern, wobei die Beschleunigung mit der Batch-Größe und der Länge des gemeinsamen Präfixes zunimmt. Hydragen ermöglicht auch die Verwendung sehr langer gemeinsamer Kontexte: Bei einer hohen Batch-Größe verringert sich der Durchsatz von Hydragen bei einer Erhöhung der Präfixlänge von 1K auf 16K Token um weniger als 15%, während der Durchsatz der Baselines um über 90% sinkt. Hydragen geht über eine einfache Präfix-Suffix-Zerlegung hinaus und kann auf baumbasierte Prompt-Sharing-Muster angewendet werden, wodurch wir die Inferenzzeit bei Wettbewerbsprogrammierungsproblemen um weitere 55% reduzieren können.

English

Transformer-based large language models (LLMs) are now deployed to hundreds of millions of users. LLM inference is commonly performed on batches of sequences that share a prefix, such as few-shot examples or a chatbot system prompt. Decoding in this large-batch setting can be bottlenecked by the attention operation, which reads large key-value (KV) caches from memory and computes inefficient matrix-vector products for every sequence in the batch. In this work, we introduce Hydragen, a hardware-aware exact implementation of attention with shared prefixes. Hydragen computes attention over the shared prefix and unique suffixes separately. This decomposition enables efficient prefix attention by batching queries together across sequences, reducing redundant memory reads and enabling the use of hardware-friendly matrix multiplications. Our method can improve end-to-end LLM throughput by up to 32x against competitive baselines, with speedup growing with the batch size and shared prefix length. Hydragen also enables the use of very long shared contexts: with a high batch size, increasing the prefix length from 1K to 16K tokens decreases Hydragen throughput by less than 15%, while the throughput of baselines drops by over 90%. Hydragen generalizes beyond simple prefix-suffix decomposition and can be applied to tree-based prompt sharing patterns, allowing us to further reduce inference time on competitive programming problems by 55%.