Hydragen: Hoogdoorvoer LLM-inferentie met gedeelde prefixen

Samenvatting

Transformer-gebaseerde grote taalmodellen (LLM's) worden nu ingezet voor honderden miljoenen gebruikers. LLM-inferentie wordt meestal uitgevoerd op batches van sequenties die een prefix delen, zoals few-shot voorbeelden of een systeemprompt voor een chatbot. Decodering in deze grote-batch setting kan worden beperkt door de aandachtoperatie, die grote sleutel-waarde (KV) caches uit het geheugen leest en inefficiënte matrix-vectorproducten berekent voor elke sequentie in de batch. In dit werk introduceren we Hydragen, een hardwarebewuste exacte implementatie van aandacht met gedeelde prefixes. Hydragen berekent aandacht over de gedeelde prefix en unieke suffixen afzonderlijk. Deze decompositie maakt efficiënte prefix-aandacht mogelijk door queries over sequenties te batchen, wat redundante geheugenleesoperaties vermindert en het gebruik van hardwarevriendelijke matrixvermenigvuldigingen mogelijk maakt. Onze methode kan de end-to-end LLM-doorvoer met tot 32x verbeteren ten opzichte van competitieve baselines, waarbij de snelheidswinst toeneemt met de batchgrootte en de lengte van de gedeelde prefix. Hydragen maakt ook het gebruik van zeer lange gedeelde contexten mogelijk: bij een hoge batchgrootte vermindert het verhogen van de prefixlengte van 1K naar 16K tokens de doorvoer van Hydragen met minder dan 15%, terwijl de doorvoer van baselines met meer dan 90% daalt. Hydragen generaliseert verder dan eenvoudige prefix-suffix decompositie en kan worden toegepast op boomgebaseerde prompt-delingpatronen, waardoor we de inferentietijd op competitieve programmeerproblemen met 55% kunnen verminderen.

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

Hydragen: Hoogdoorvoer LLM-inferentie met gedeelde prefixen

Hydragen: High-Throughput LLM Inference with Shared Prefixes

Samenvatting

Support