arXiv:2409.04992cs.ARcs.CL2024-09被引 38

将大模型推理的注意力计算和缓存移至存储设备,大幅降低显存压力。

InstInfer: In-Storage Attention Offloading for Cost-Effective Long-Context LLM Inference

  • 把注意力计算和KV缓存放到计算型固态硬盘中,避开显卡带宽瓶颈。
  • 在13B模型上,长序列推理吞吐提升最高达11.1倍。
  • 适合资源受限场景如边缘设备、个人电脑等低显存环境使用。

大型语言模型(LLMs)的普及标志着生成式AI的重要进展。然而,离线推理中上下文长度和批处理规模的增加导致键值(KV)缓存内存需求激增,给GPU显存带来巨大压力,尤其在资源受限场景(如边缘计算和个人设备)下更为严重。现有成本优化方案利用主机内存或SSD降低存储开销并提升吞吐量,但受限于有限的PCIe带宽,频繁访问KV缓存仍造成显著性能损耗。为此,我们提出InstInfer,一种新型大模型推理系统,将解码阶段最耗性能的注意力计算与KV缓存数据卸载至计算型存储设备(CSDs),从而最小化庞大的数据传输开销。InstInfer设计了专用的闪存感知存储内注意力引擎及缓存管理机制,充分利用CSDs内部高带宽,突破PCIe带宽限制。同时优化了GPU与CSD间点对点传输,进一步减少数据迁移开销。实验结果表明,在NVIDIA A6000 GPU上,对于13B模型,InstInfer相比基于SSD的FlexGen等方案,长序列推理吞吐最高提升11.1倍。

原文摘要 · Abstract (English)

The widespread of Large Language Models (LLMs) marks a significant milestone in generative AI. Nevertheless, the increasing context length and batch size in offline LLM inference escalate the memory requirement of the key-value (KV) cache, which imposes a huge burden on the GPU VRAM, especially for resource-constraint scenarios (e.g., edge computing and personal devices). Several cost-effective solutions leverage host memory or SSDs to reduce storage costs for offline inference scenarios and improve the throughput. Nevertheless, they suffer from significant performance penalties imposed by intensive KV cache accesses due to limited PCIe bandwidth. To address these issues, we propose InstInfer, a novel LLM inference system that offloads the most performance-critical computation (i.e., attention in decoding phase) and data (i.e., KV cache) parts to Computational Storage Drives (CSDs), which minimize the enormous KV transfer overheads. InstInfer designs a dedicated flash-aware in-storage attention engine with KV cache management mechanisms to exploit the high internal bandwidths of CSDs instead of being limited by the PCIe bandwidth. The optimized P2P transmission between GPU and CSDs further reduces data migration overheads. Experimental results demonstrate that for a 13B model using an NVIDIA A6000 GPU, InstInfer improves throughput for long-sequence inference by up to 11.1$\times$, compared to existing SSD-based solutions such as FlexGen.

大模型推理计算存储显存优化长序列

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。