Hydragen: Inferenza ad Alto Rendimento per LLM con Prefissi Condivisi

Hydragen: High-Throughput LLM Inference with Shared Prefixes

February 7, 2024

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

cs.AI

Abstract

I modelli linguistici di grandi dimensioni (LLM) basati su Transformer sono ora utilizzati da centinaia di milioni di utenti. L'inferenza degli LLM viene comunemente eseguita su batch di sequenze che condividono un prefisso, come esempi few-shot o il prompt di un sistema di chatbot. Il decoding in questo contesto di grandi batch può essere limitato dall'operazione di attenzione, che legge grandi cache di chiavi-valori (KV) dalla memoria e calcola prodotti matrice-vettore inefficienti per ogni sequenza nel batch. In questo lavoro, introduciamo Hydragen, un'implementazione esatta e consapevole dell'hardware dell'attenzione con prefissi condivisi. Hydragen calcola l'attenzione sul prefisso condiviso e sui suffissi unici separatamente. Questa scomposizione consente un'attenzione efficiente sul prefisso raggruppando le query tra le sequenze, riducendo le letture ridondanti dalla memoria e permettendo l'uso di moltiplicazioni matriciali ottimizzate per l'hardware. Il nostro metodo può migliorare il throughput end-to-end degli LLM fino a 32 volte rispetto a baseline competitive, con un'accelerazione che cresce con la dimensione del batch e la lunghezza del prefisso condiviso. Hydragen consente inoltre l'uso di contesti condivisi molto lunghi: con un batch di grandi dimensioni, aumentare la lunghezza del prefisso da 1K a 16K token riduce il throughput di Hydragen di meno del 15%, mentre il throughput delle baseline diminuisce di oltre il 90%. Hydragen si generalizza oltre la semplice scomposizione prefisso-suffisso e può essere applicato a schemi di condivisione dei prompt basati su alberi, permettendoci di ridurre ulteriormente il tempo di inferenza su problemi di programmazione competitiva del 55%.

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%.