Hydragen: Inferência de LLM de Alta Taxa de Transferência com Prefixos Compartilhados

Hydragen: High-Throughput LLM Inference with Shared Prefixes

February 7, 2024

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

cs.AI

Resumo

Modelos de linguagem de grande escala (LLMs) baseados em Transformers agora são implantados para centenas de milhões de usuários. A inferência de LLMs é comumente realizada em lotes de sequências que compartilham um prefixo, como exemplos de few-shot ou um prompt de sistema de chatbot. A decodificação nesse cenário de grandes lotes pode ser limitada pela operação de atenção, que lê grandes caches de chave-valor (KV) da memória e calcula produtos matriciais ineficientes para cada sequência no lote. Neste trabalho, introduzimos o Hydragen, uma implementação exata e consciente do hardware da atenção com prefixos compartilhados. O Hydragen calcula a atenção sobre o prefixo compartilhado e os sufixos únicos separadamente. Essa decomposição permite uma atenção eficiente no prefixo ao agrupar consultas entre sequências, reduzindo leituras redundantes de memória e possibilitando o uso de multiplicações matriciais amigáveis ao hardware. Nosso método pode melhorar a taxa de transferência de ponta a ponta de LLMs em até 32x em relação a baselines competitivas, com o ganho de velocidade aumentando com o tamanho do lote e o comprimento do prefixo compartilhado. O Hydragen também permite o uso de contextos compartilhados muito longos: com um lote grande, aumentar o comprimento do prefixo de 1K para 16K tokens diminui a taxa de transferência do Hydragen em menos de 15%, enquanto a taxa de transferência das baselines cai mais de 90%. O Hydragen generaliza além da simples decomposição prefixo-sufixo e pode ser aplicado a padrões de compartilhamento de prompt baseados em árvore, permitindo-nos reduzir ainda mais o tempo de inferência em problemas de programação competitiva em 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%.