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

📄 arXiv: 2409.04992v1 📥 PDF

作者: Xiurui Pan, Endian Li, Qiao Li, Shengwen Liang, Yizhou Shan, Ke Zhou, Yingwei Luo, Xiaolin Wang, Jie Zhang

分类: cs.AR, cs.CL

发布日期: 2024-09-08

💡 一句话要点

InstInfer：面向长文本LLM推理，提出基于计算存储的Attention卸载方案，提升成本效益。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型推理 计算存储 KV缓存 Attention机制 长文本处理

📋 核心要点

  1. 现有LLM推理方案在长文本和高并发场景下，KV缓存占用大量GPU显存，导致资源受限设备性能瓶颈。
  2. InstInfer将Attention计算和KV缓存卸载到计算存储设备（CSD），利用CSD内部高带宽，减少PCIe传输开销。
  3. 实验表明，InstInfer相比于SSD方案FlexGen，在13B模型长序列推理中，吞吐量提升高达11.1倍。

📝 摘要（中文）

大型语言模型（LLM）的广泛应用标志着生成式人工智能的一个重要里程碑。然而，离线LLM推理中不断增长的上下文长度和批处理大小增加了键值（KV）缓存的内存需求，给GPU VRAM带来了巨大的负担，尤其是在资源受限的场景（例如，边缘计算和个人设备）中。一些经济高效的解决方案利用主机内存或SSD来降低离线推理场景的存储成本并提高吞吐量。然而，由于有限的PCIe带宽，它们会因密集的KV缓存访问而遭受显著的性能损失。为了解决这些问题，我们提出了一种新颖的LLM推理系统InstInfer，它将性能最关键的计算（即解码阶段的注意力）和数据（即KV缓存）部分卸载到计算存储驱动器（CSD），从而最大限度地减少了巨大的KV传输开销。InstInfer设计了一个专用的闪存感知型存储内注意力引擎，具有KV缓存管理机制，以利用CSD的高内部带宽，而不是受PCIe带宽的限制。GPU和CSD之间优化的P2P传输进一步降低了数据迁移开销。实验结果表明，对于使用NVIDIA A6000 GPU的13B模型，与现有的基于SSD的解决方案（如FlexGen）相比，InstInfer将长序列推理的吞吐量提高了高达11.1倍。

🔬 方法详解

问题定义：论文旨在解决长文本LLM推理中，KV缓存占用大量GPU显存，导致推理速度受限的问题。现有方法如FlexGen等，虽然利用SSD或主机内存扩展存储，但受限于PCIe带宽，频繁的KV缓存访问成为性能瓶颈。

核心思路：论文的核心思路是将Attention计算和KV缓存卸载到计算存储设备（CSD）上。CSD具有内置计算能力和高速存储，可以利用其内部高带宽进行Attention计算，避免了通过PCIe总线频繁传输KV缓存带来的性能瓶颈。

技术框架：InstInfer的整体架构包含GPU和CSD两部分。GPU负责模型的前后处理以及部分计算，CSD负责Attention计算和KV缓存管理。系统通过优化的P2P传输机制在GPU和CSD之间传输必要的数据。主要模块包括：闪存感知型存储内注意力引擎、KV缓存管理机制和优化的P2P传输模块。

关键创新：InstInfer的关键创新在于将Attention计算卸载到CSD上，并设计了专用的闪存感知型存储内注意力引擎。该引擎充分利用CSD内部的高带宽，避免了PCIe带宽的限制，从而显著提高了推理速度。此外，优化的P2P传输机制也减少了数据迁移的开销。

关键设计：InstInfer的关键设计包括：(1) 闪存感知型存储内注意力引擎，针对闪存的特性进行了优化，例如减少写入次数，提高读取速度。(2) KV缓存管理机制，根据访问频率和重要性，对KV缓存进行分层管理，提高缓存命中率。(3) 优化的P2P传输机制，减少GPU和CSD之间的数据传输延迟。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，对于使用NVIDIA A6000 GPU的13B模型，InstInfer相比于现有的基于SSD的解决方案（如FlexGen），将长序列推理的吞吐量提高了高达11.1倍。这表明InstInfer能够显著提高长文本LLM推理的性能，尤其是在资源受限的场景下。

🎯 应用场景

InstInfer适用于资源受限的边缘计算设备和个人设备，可以支持在这些设备上进行长文本LLM推理。该研究成果有助于降低LLM推理的成本，并扩展LLM的应用范围，例如智能助手、文本生成、机器翻译等。

📄 摘要（原文）

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.