Hydragen: Inferencia de LLM de Alto Rendimiento con Prefijos Compartidos

Resumen

Los grandes modelos de lenguaje (LLMs) basados en Transformers ahora están desplegados para cientos de millones de usuarios. La inferencia de LLMs comúnmente se realiza en lotes de secuencias que comparten un prefijo, como ejemplos de pocos disparos o un mensaje de sistema de un chatbot. La decodificación en este entorno de grandes lotes puede verse limitada por la operación de atención, que lee grandes cachés de clave-valor (KV) desde la memoria y calcula productos matriz-vector ineficientes para cada secuencia en el lote. En este trabajo, presentamos Hydragen, una implementación exacta y consciente del hardware de la atención con prefijos compartidos. Hydragen calcula la atención sobre el prefijo compartido y los sufijos únicos por separado. Esta descomposición permite una atención eficiente sobre el prefijo al agrupar consultas entre secuencias, reduciendo lecturas redundantes de memoria y permitiendo el uso de multiplicaciones matriciales amigables con el hardware. Nuestro método puede mejorar el rendimiento de extremo a extremo de los LLMs hasta 32 veces en comparación con líneas base competitivas, con una aceleración que crece con el tamaño del lote y la longitud del prefijo compartido. Hydragen también permite el uso de contextos compartidos muy largos: con un tamaño de lote alto, aumentar la longitud del prefijo de 1K a 16K tokens disminuye el rendimiento de Hydragen en menos del 15%, mientras que el rendimiento de las líneas base cae más del 90%. Hydragen se generaliza más allá de la simple descomposición prefijo-sufijo y puede aplicarse a patrones de compartición de mensajes basados en árboles, permitiéndonos reducir aún más el tiempo de inferencia en problemas de programación competitiva en un 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%.

Hydragen: Inferencia de LLM de Alto Rendimiento con Prefijos Compartidos

Hydragen: High-Throughput LLM Inference with Shared Prefixes

Resumen

Support