Hydragen: Высокопроизводительный вывод LLM с общими префиксами
Hydragen: High-Throughput LLM Inference with Shared Prefixes
February 7, 2024
Авторы: Jordan Juravsky, Bradley Brown, Ryan Ehrlich, Daniel Y. Fu, Christopher Ré, Azalia Mirhoseini
cs.AI
Аннотация
Трансформерные большие языковые модели (LLM) теперь используются сотнями миллионов пользователей. Вывод LLM обычно выполняется на пакетах последовательностей, имеющих общий префикс, таких как примеры few-shot или системный промпт чат-бота. Декодирование в таких условиях с большими пакетами может быть ограничено операцией внимания, которая считывает большие кэши ключей и значений (KV) из памяти и вычисляет неэффективные матрично-векторные произведения для каждой последовательности в пакете. В данной работе мы представляем Hydragen, аппаратно-ориентированную точную реализацию внимания с общими префиксами. Hydragen вычисляет внимание для общего префикса и уникальных суффиксов отдельно. Такая декомпозиция позволяет эффективно вычислять внимание для префикса, объединяя запросы из разных последовательностей в пакеты, что уменьшает избыточные чтения из памяти и позволяет использовать аппаратно-оптимизированные матричные умножения. Наш метод может увеличить общую пропускную способность LLM до 32 раз по сравнению с конкурирующими базовыми методами, причем ускорение растет с увеличением размера пакета и длины общего префикса. Hydragen также позволяет использовать очень длинные общие контексты: при большом размере пакета увеличение длины префикса с 1K до 16K токенов снижает пропускную способность Hydragen менее чем на 15%, в то время как пропускная способность базовых методов падает более чем на 90%. Hydragen обобщается за пределы простой декомпозиции на префикс и суффикс и может применяться к древовидным схемам разделения промптов, что позволяет дополнительно сократить время вывода на задачах конкурсного программирования на 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%.