Hydragen : Inférence à haut débit pour les modèles de langage avec préfixes partagés

Hydragen: High-Throughput LLM Inference with Shared Prefixes

February 7, 2024

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

cs.AI

Résumé

Les grands modèles de langage (LLM) basés sur l'architecture Transformer sont désormais déployés auprès de centaines de millions d'utilisateurs. L'inférence des LLM est généralement effectuée sur des lots de séquences partageant un préfixe commun, comme des exemples few-shot ou un prompt de système de chatbot. Le décodage dans ce contexte de grands lots peut être limité par l'opération d'attention, qui lit de larges caches clé-valeur (KV) en mémoire et calcule des produits matrice-vecteur inefficaces pour chaque séquence du lot. Dans ce travail, nous introduisons Hydragen, une implémentation exacte et optimisée pour le matériel de l'attention avec des préfixes partagés. Hydragen calcule l'attention sur le préfixe partagé et les suffixes uniques séparément. Cette décomposition permet une attention efficace sur le préfixe en regroupant les requêtes entre les séquences, réduisant les lectures redondantes en mémoire et permettant l'utilisation de multiplications matricielles adaptées au matériel. Notre méthode peut améliorer le débit end-to-end des LLM jusqu'à 32 fois par rapport à des bases de référence compétitives, avec une accélération qui augmente avec la taille du lot et la longueur du préfixe partagé. Hydragen permet également l'utilisation de contextes partagés très longs : avec un lot de grande taille, augmenter la longueur du préfixe de 1K à 16K tokens réduit le débit de Hydragen de moins de 15 %, tandis que le débit des bases de référence chute de plus de 90 %. Hydragen se généralise au-delà de la simple décomposition préfixe-suffixe et peut être appliqué à des motifs de partage de prompts basés sur des arbres, nous permettant de réduire encore le temps d'inférence sur des problèmes de programmation compétitive de 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%.