Hydragen:面向共享前缀场景的高吞吐量大语言模型推理
Hydragen: High-Throughput LLM Inference with Shared Prefixes
AI总结:
针对共享前缀场景下LLM推理注意力运算效率低的问题,提出Hydragen方法,通过分解共享前缀与独有后缀的注意力计算提升效率,大幅提高推理吞吐量,还支持长共享上下文与树状提示共享模式。
摘要:
基于Transformer的大语言模型(LLM)目前已服务数亿用户。LLM推理通常在批量的共享前缀序列上执行,例如少样本示例或聊天机器人系统提示词。在这种大批量场景下,解码过程的瓶颈可能出现在注意力运算环节:该运算需要从内存中读取大型键值(KV)缓存,且要为批次中的每条序列计算低效的矩阵-向量乘积。本研究提出Hydragen,一种面向硬件的、支持共享前缀注意力的精确实现方案。Hydragen分别计算共享前缀部分和独有后缀部分的注意力。这种分解方式通过跨序列批量处理查询向量,实现了高效的前缀注意力计算,减少了冗余内存读取,且可采用适配硬件的矩阵乘法运算。与同类基线方法相比,该方法可将CodeLlama-13b的端到端吞吐量提升最高达32倍,且加速效果随批次规模和共享前缀长度增加而提升。Hydragen还支持极长的共享上下文:在大批次规模下,将前缀长度从1K token增加到16K token时,Hydragen的吞吐量下降不足15%,而基线方法的吞吐量下降超过90%。Hydragen的适用范围不止于简单的前后缀分解,还可应用于基于树的提示共享模式,借此可将竞赛编程问题的推理时间进一步缩短55%。
英文摘要:
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 CodeLlama-13b 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 large 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%.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。