用局部敏感哈希采样优化大模型生成，提速5倍且保持高精度。
MagicPIG: LSH Sampling for Efficient LLM Generation
- 基于局部敏感哈希进行关键键值采样，替代传统TopK方法。
- 在96k上下文下实现54ms延迟，推理吞吐提升最高5倍。
- 适合需要长序列生成的场景，如文档摘要与对话系统。
长上下文大语言模型（LLMs）备受关注，但其存储的键值缓存（KV cache）成为瓶颈。现有动态稀疏或基于TopK的注意力近似方法假设注意力具有稀疏性，但实际中该假设不总成立。本文指出，仅选取注意力分数最高的键值会导致下游任务性能下降。相反，具有理论保障的采样方法能更准确估计注意力输出。为此，提出MagicPIG，一种基于局部敏感哈希（LSH）的异构系统。它将LSH哈希表存储于内存，并在CPU上执行注意力计算，显著降低计算负担，支持更长上下文和更大批量。MagicPIG在多种GPU硬件上可实现最高5倍的解码吞吐提升，在单张RTX 4090上对Llama-3.1-8B-Instruct模型处理96k令牌上下文时达到54ms解码延迟。代码已开源。
原文摘要 · Abstract (English)
Large language models (LLMs) with long context windows have gained significant attention. However, the KV cache, stored to avoid re-computation, becomes a bottleneck. Various dynamic sparse or TopK-based attention approximation methods have been proposed to leverage the common insight that attention is sparse. In this paper, we first show that TopK attention itself suffers from quality degradation in certain downstream tasks because attention is not always as sparse as expected. Rather than selecting the keys and values with the highest attention scores, sampling with theoretical guarantees can provide a better estimation for attention output. To make the sampling-based approximation practical in LLM generation, we propose MagicPIG, a heterogeneous system based on Locality Sensitive Hashing (LSH). MagicPIG significantly reduces the workload of attention computation while preserving high accuracy for diverse tasks. MagicPIG stores the LSH hash tables and runs the attention computation on the CPU, which allows it to serve longer contexts and larger batch sizes with high approximation accuracy. MagicPIG can improve decoding throughput by up to $5\times$ across various GPU hardware and achieve 54ms decoding latency on a single RTX 4090 for Llama-3.1-8B-Instruct model with a context of 96k tokens. The code is available at https://github.com/Infini-AI-Lab/MagicPIG.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。