InstInfer: In-Storage Attention Offloading for Cost-Effective Long-Context LLM Inference
InstInfer: In-Storage Attention Offloading for Cost-Effective Long-Context LLM Inference
发表时间: 2024-09 · arXiv:2409.04992 (preprint)
原文: https://arxiv.org/abs/2409.04992
Xiurui Pan, Endian Li, Qiao Li, Shengwen Liang, Yizhou Shan, Ke Zhou, Yingwei Luo, Xiaolin Wang, and Jie Zhang
Peking University, Xiamen University, Institute of Computing Technology (Chinese Academy of Sciences), Huawei Cloud, Wuhan National Laboratory for Optoelectronics (Huazhong University of Science and Technology)
速读
一句话结论 本文提出了一个名为 InstInfer 的软硬协同大语言模型推理系统,通过将解码阶段的注意力计算和键值缓存(KV cache)卸载到计算型存储驱动器(CSD)内部,绕过了外部 PCIe 总线的带宽瓶颈,在 13B 参数模型长上下文推理中实现了比现有固态硬盘卸载方案高 11.1 倍的吞吐量。
要解决什么问题 在离线大语言模型推理中,随着上下文长度和批处理大小的增加,自回归解码阶段产生的 KV cache 体积会急剧膨胀,轻易超出单张 GPU 的显存容量。为了降低硬件成本,现有的主流方案(如 DeepSpeed-MII 和 FlexGen)选择将 KV cache 卸载到主机内存或廉价的固态硬盘(SSD)中。然而,解码阶段是一个典型的访存密集型任务,需要频繁读取庞大的 KV cache。当数据被存放在 SSD 时,GPU 与存储设备之间狭窄的外部 PCIe 带宽(通常仅有 3 到 6 GB/s)成为了致命的传输瓶颈。由于 SSD 内部多个闪存通道的聚合带宽实际上远高于外部 PCIe 带宽,传统的卸载方案白白浪费了存储介质内部的并发读取潜力,导致长序列推理时因极高的 KV cache 访问延迟而出现严重的性能断崖。
怎么做的 核心思路是改变粗粒度的任务划分,仅将访存密集且计算强度极低的解码阶段注意力计算以及庞大的 KV cache 卸载到带有计算能力的存储驱动器(CSD)内部,而将计算密集的预填充阶段和其他线性层保留在 GPU 上。这样,注意力计算可以直接利用 CSD 内部极高的闪存通道聚合带宽,GPU 和 CSD 之间只需通过点对点直接内存访问(P2P DMA)传输极小的数据量,彻底绕开了 PCIe 带宽卡点。为了弥补 CSD 算力较弱的问题并适应闪存按页访问的物理特性,系统设计了软硬协同的 SparF 注意力机制。该机制首先选取查询向量中绝对值最大的 $r$ 个隐藏层特征 $i \gets \mathrm{argtopk}(|\pmb{q}|, r)$,据此从闪存中读取对应的键(K)缓存来近似计算注意力分数 $\hat{s}$;接着选出分数最大的 $k$ 个词元 $j \gets \mathrm{argtopk}(\hat{s} + m, k)$,再读取完整的键值(KV)缓存计算最终输出。针对闪存按页读取导致的读放大问题,SparF 采用了双步加载策略:在粗粒度阶段,直接跳过不包含任何强相关词元或特征的整个闪存页;在细粒度阶段,数据到达闪存控制器(NFC)内部的缓冲区后,由硬件过滤单元剔除页内剩余的无关数据,从而大幅降低了实际传输和计算的数据量。此外,为了支持上述算法中对词元和隐藏层特征的两种随机索引需求,系统在闪存转换层(FTL)中设计了双重地址映射机制,将 K 缓存分别按词元和隐藏层特征存储两份,并将连续的词元打包成与闪存页大小匹配的数据组,跨通道交错存储以最大化并发带宽。
效果如何 实验在配备单张 48GB 显存的 NVIDIA A6000 GPU 和基于 Zynq7045 FPGA 模拟的 CSD 硬件上进行,测试模型为 OPT-13B,输入和输出序列长度均设定为 1024 个词元。对比基线包括代表主机内存卸载路线的 DeepSpeed、代表常规 SSD 卸载路线的 FlexGen,以及结合了 SparQ 稀疏算法的 FlexGen-SparQ。量化结果显示,在批处理大小为 64 时,基础版 InstInfer 的吞吐量是 FlexGen 的 6.85 倍;在批处理大小扩大到 256 且开启 1/8 稀疏度的 SparF 机制后,InstInfer 的吞吐量达到了 FlexGen 的 11.1 倍,并将解码阶段的 KV cache 访问延迟开销降低了 94.0%。当扩展到双 CSD 阵列时,该方法依然能比双盘 FlexGen-SparQ 路线提升 3.11 倍吞吐量。不过作者也指出,由于 CSD 的内部总带宽(11.2 GB/s)仍低于主机内存到 GPU 的 PCIe 带宽(32 GB/s),在批处理较小、KV cache 尚未超出主机内存容量的场景下,其绝对吞吐量略低于纯内存卸载的 DeepSpeed;同时,为了支持双向索引,该方法需要牺牲额外的廉价闪存空间来双份存储 K 缓存。
主要贡献
大型语言模型(LLMs)的离线推理服务中,随着上下文长度和批处理大小(batch size)的不断增加,键值(KV)缓存的内存需求急剧上升,这给GPU显存带来了巨大的负担,特别是在资源受限的场景下(如边缘计算和个人设备)。为了降低存储成本并提高吞吐量,现有的高性价比解决方案(如FlexGen、DeepSpeed-MII)倾向于将KV缓存卸载(offload)到主机内存或SSD中。然而,由于PCIe带宽的限制,密集访问KV缓存会导致严重的性能下降。
为了解决上述存储成本和带宽瓶颈问题,本文提出了InstInfer,这是一个新颖的LLM推理系统,它将对性能最关键的计算部分(即解码阶段的注意力机制)和数据(即KV缓存)卸载到计算存储驱动器(Computational Storage Drives, CSDs)中,从而最小化巨大的KV数据传输开销。本文的核心创新点如下:
1. 开创性的基于CSD的长上下文LLM异构推理系统:详细分析表明,解码阶段的注意力计算由于算术强度极低且需要通过受限的PCIe带宽访问庞大的KV缓存,成为了最关键的性能瓶颈。InstInfer将KV缓存和解码阶段的注意力计算卸载到CSD,利用闪存芯片的高聚合带宽,将数据迁移开销有效降低了高达$94.0\%$,并通过优化的点对点(P2P)DMA机制进一步缓解了预填充阶段的开销。
2. 软硬件协同设计的存内注意力引擎:为了弥补GPU和CSD之间的带宽和算力差距,提出了一种带宽高效的SparF算法。该算法不仅降低了计算强度,还在保持精度的同时最小化了解码阶段所需的KV缓存量。考虑到闪存按页访问的粒度特性,InstInfer采用双步加载策略来管理序列中的稀疏性:首先在页级别加载,随后在Token级别过滤。此外,在硬件内核中实现了存内SparF注意力引擎,采用细粒度并行设计来隐藏闪存芯片的长访问延迟。
3. 面向KV缓存的高效检索FTL设计:由于SparF算法在Token和隐藏层嵌入(hidden embeddings)中均识别出稀疏模式,导致在闪存芯片中对KV缓存产生随机访问。InstInfer通过引入分别为Token索引和隐藏层嵌入索引量身定制的双地址映射机制来解决这一挑战。同时,将KV缓存张量精心组织成与闪存页大小对齐的组,并以步幅(stridden)方式分布在多个闪存块和芯片中,从而充分利用固有的高并行性。
背景知识与关键观察
LLM推理基础
LLM架构与自回归推理:主流的大型语言模型主要采用仅解码器(decoder-only)的Transformer架构【58, Llama: Open and efficient foundation language models 2023 arXiv】、【60, Visionllm: Large language model is also an open-ended decoder for vision-centric tasks 2024 NeurIPS】、【65, Harnessing the power of llms in practice: A survey on chatgpt and beyond 2024 ACM TKDD】、【71, Opt: Open pre-trained transformer language models 2022 arXiv】。该架构由多个堆叠的解码器块组成,每个块包含一个自注意力模块和前馈神经网络(FFN)模块。给定输入序列$X = [x_1, ..., x_s]$,模型通过参数矩阵将其线性变换为$Q$、$K$、$V$三个嵌入矩阵,然后执行注意力机制:$Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$。多头注意力(MHA)【59, Attention is all you need 2017 NeurIPS】进一步将矩阵划分为更小的矩阵以捕捉不同上下文。推理过程分为预填充(prefilling)和解码(decoding)两个阶段。预填充阶段并行处理所有输入Prompt的Token以生成第一个预测Token;解码阶段则基于已有序列逐个预测新Token,直到生成结束符(EOS)或达到上下文限制。
KV缓存
重计算减少与稀疏注意力:在解码阶段,注意力操作需要重复计算之前所有Token的$QKV$矩阵,导致每次迭代的计算复杂度为$O(s^2)$。KV缓存【30, Efficient memory management for large language model serving with pagedattention 2023 SOSP】通过在GPU显存中缓存已生成Token的$KV$矩阵来避免冗余计算,将复杂度降低至$O(s)$。然而,随着上下文长度增加,KV缓存会消耗大量存储空间并产生极高的I/O需求。为了进一步降低内存访问需求,稀疏注意力成为常用方法【7, An attentive survey of attention models 2021 ACM TIST】、【8, Scatterbrain: Unifying sparse and low-rank attention 2021 NeurIPS】。例如,SparQ Attention【50, Sparq attention: Bandwidth-efficient llm inference 2023 arXiv】基于查询向量($Q$)中最大的$r$个隐藏层嵌入值来近似注意力分数,然后选出最重要的前$k$个Token计算最终输出。尽管SparQ减少了传输带宽,但它需要同时按Token和隐藏层嵌入维度索引$K$缓存,导致KV缓存的内存占用增加$1.5\times$,限制了其在内存受限场景中的应用。
SSD与存内计算
SSD基础与计算存储驱动器:现代基于NAND闪存的SSD包含闪存裸片(dies)、SSD控制器和DRAM模块。闪存按页(通常为4KB至16KB)进行读写,并在写入前需按块(block)擦除【3, Design tradeoffs for SSD performance 2008 USENIX ATC】。SSD控制器运行闪存转换层(FTL)以管理逻辑到物理地址的映射。现代SSD具有8至16个闪存通道,聚合带宽可达数十GB/s,远超SSD的外部PCIe带宽($3 \sim 6 GB/s$)【51, Samsung 980pro nvme ssd URL】。计算存储驱动器(CSD)通过在SSD内部集成ARM核、NPU或FPGA等计算引擎【28, Summarizer: trading communication with computing near storage 2017 MICRO】、【34, Cognitive SSD: A deep learning engine for In-Storage data retrieval 2019 USENIX ATC】、【38, Genstore: A high-performance in-storage processing system for genome sequence analysis 2022 ASPLOS】,使SSD具备直接处理数据的能力。为了充分利用高闪存通道带宽,计算引擎最好放置在靠近闪存裸片或NFC的位置。
传统KV缓存卸载的局限性
KV缓存分析与卸载导致的性能下降:当前LLM推理系统的上下文长度和批处理大小不断增加,导致KV缓存体积急剧膨胀。例如,对于批处理大小为128、长度为2K的序列,OPT-13B模型的KV缓存高达200GB,远超其24GB的模型权重。由于PCIe带宽显著低于GPU显存带宽,将KV缓存卸载到主机内存或SSD会导致性能显著下降。实验评估了DeepSpeed【19, Deepspeedfastgen: High-throughput text generation for llms via mii and deepspeedinference 2024 arXiv】和FlexGen【53, Flexgen: High-throughput generative inference of large language models with a single gpu 2023 ICML】在长上下文场景下的表现。结果显示,随着批处理大小增加,当KV缓存超出GPU显存并被迫卸载到内存或SSD时,两者均出现严重的吞吐量下降。DeepSpeed在发生从主机内存到SSD的内核交换时,性能下降了$97.01\%$。FlexGen的解码延迟分解表明,在较大的批处理大小下,KV缓存访问开销飙升至高达$98.94\%$,成为绝对的性能瓶颈。
CSD的卸载机遇
重构任务解耦方案:与内存和NVMe SSD相比,将KV缓存卸载到CSD内的闪存芯片可以直接利用更高的闪存通道带宽。然而,由于CSD的计算性能远低于GPU,简单的预填充-解码分离架构【48, Splitwise: Efficient generative llm inference using phase splitting 2023 ISCA】、【75, DistServe: Disaggregating prefill and decoding for goodput-optimized large language model serving 2024 OSDI】并不适用。通过对NVIDIA A6000 GPU和基于Zynq7045 FPGA的CSD进行Roofline模型分析【68, Llm inference unveiled: Survey and roofline model insights 2024 arXiv】发现,预填充阶段的所有操作均应保留在GPU上;在解码阶段,$QKV$投影、$O$投影和FFN虽然在GPU上受限于内存带宽,但其计算强度接近CSD的算力上限,且仅依赖权重矩阵进行GeMM计算【20, Flashdecoding++: Faster large language model inference 2024 MLSys】。相反,解码阶段的注意力操作(Logit和Attend)涉及极低强度的GeMV计算,且需要直接访问KV缓存。因此,将解码阶段的注意力操作卸载到CSD,保留其他操作在GPU上,能够显著降低KV缓存传输开销并最小化CSD的计算负担。
方法细节
InstInfer架构概述
系统硬件组件与整体协同:InstInfer是首个面向长上下文和大批量离线LLM推理的配备通用GPU的存内注意力卸载系统。系统主要由三个硬件组件构成:1)InstCSD,负责执行解码阶段的注意力计算并存储庞大的KV缓存;2)InstGPU,负责执行其他推理计算并在预填充阶段生成KV缓存;3)InstHost,运行软件栈,负责调度推理任务并协调GPU与InstCSD之间的数据传输。为了平衡CSD上庞大的带宽需求与较低的存储容量要求,InstInfer提出了SparF注意力机制,这是一种针对闪存优化的SparQ算法增强版。由于闪存按页访问,SparF在与闪存页大小对应的组级别组织Token,通过双步机制(粗粒度组级别+细粒度Token级别)识别和获取KV缓存。基于此,InstCSD上设计了硬件加速器,并提出了两种地址映射机制以实现高效检索。在解码阶段,GPU与InstCSD之间仅传输$qkv$向量和注意力输出,且通过P2PDMA绕过主机内存进行传输,传输过程与计算重叠以隐藏延迟。
计算注意力输出
闪存感知的稀疏注意力(SparF Attention):解码阶段的注意力算子在CSD上仍然严重受限于内存,因为它们极度依赖闪存中的KV缓存。传统的稀疏算法会产生大量随机访问,导致严重的写放大【22, Write amplification analysis in flash-based solid state drives 2009 SYSTOR】和带宽浪费。为此,InstInfer提出了SparF Attention,如Algorithm 1所示。SparF Attention首先通过选择$q$向量中最大的前$r$个隐藏层嵌入值(步骤1)来近似完整的注意力分数$\hat{s}$,并根据这些索引从闪存中加载对应的$K$缓存。随后,基于近似分数$\hat{s}$,选择前$k$个最大的Token,并从闪存中加载这些Token完整的$K, V$缓存。为了适配闪存页大小,KV缓存加载分为两步(步骤2-3,8-9):首先在闪存控制器(NFC)层面过滤掉包含弱元素的整个闪存页,不进行读取;当粗粒度的稀疏KV缓存被取到NFC后,NFC执行细粒度过滤,丢弃剩余的弱单元。这种双步加载方案不仅减少了闪存通道的数据传输量,还减轻了计算引擎的负担。
Input: q, \bar{v} \in \mathbb{R}^{d_h}, K, V \in \mathbb{R}^{S \times d_h}, r, k, l, m, n \in \mathbb{N}
Output: out \in \mathbb{R}^{d_h}
1: i \leftarrow [1 \text{ if } i \in \text{argtopk}(|q|, r) \text{ else } 0]_{i=1}^S
2: load K_{[:, i_1]}^\top \text{ if } i_{[m\lfloor \frac{i_1}{m} \rfloor : m\lceil \frac{i_1}{m} \rceil]} = 0
3: K_{[:, i]}^\top \leftarrow \text{filter } K_{[i_2, :]}^\top \text{ if } i_2 = 0
4: \hat{s} \leftarrow \text{softmax} \left( q_{[i]} \cdot K_{[:, i]}^\top / \sqrt{d_h \frac{||q_{[i]}||_1}{||q||_1}} \right)
5: m \leftarrow [1 \text{ if } i > S \text{ else } 0]_{i=1}^S
6: j \leftarrow [1 \text{ if } j \in \text{argtopk}(\hat{s} + m, k) \text{ else } 0]_{j=1}^S
7: \alpha \leftarrow \text{sum}(\hat{s}_{[j]})
8: load K_{[j_1, :]}^\top, V_{[j_1, :]} \text{ if } j_{[n\lfloor \frac{j_1}{n} \rfloor : n\lceil \frac{j_1}{n} \rceil]} = 0
9: K_{[j, :]}^\top, V_{[j, :]} \leftarrow \text{filter } K_{[j_2, :]}^\top, V_{[j_2, :]} \text{ if } j_2 = 0
10: s \leftarrow \text{softmax} \left( q \cdot K_{[j, :]}^\top / \sqrt{d_h} \right)
11: out \leftarrow \alpha s \cdot V_{[j, :]} + (1 - \alpha) \bar{v}
基于硬件的注意力引擎:在InstCSD上设计了与SSD控制器集成的硬件注意力引擎。引擎主要包括注意力内核、argtopk单元以及集成在每个NFC中的过滤器。$q$向量首先进入argtopk单元过滤出前$r$个最大的隐藏层嵌入索引,随后这些索引发送到NFC以获取$K_{[:, i]}$缓存。数据页在通过NFC时,NFC利用细粒度索引信息过滤掉所有稀疏单元。过滤后的$q_{[i]}$和$K_{[:, i]}$被送入Attention Kernel 1计算近似注意力分数,随后再次送入argtopk单元过滤出前$k$个最大的Token索引。基于这些索引,以页为粒度从闪存加载$K_{[j, :]}$和$V_{[j, :]}$缓存并通过NFC过滤。稀疏的$q, K_{[j, :]}$张量被送入Attention Kernel 2,而$V_{[j, :]}$并行加载以隐藏延迟。两个注意力内核结构相同,均包含多个GeMV单元和Softmax单元,根据实时负载进行调度。
管理与传输KV缓存
Token索引映射:主流大模型(如OPT、LLaMA)的注意力头隐藏层大小通常为128(FP16格式),最小读取粒度为256B。面对4KB的闪存页,传统FTL的随机读取会导致高达$16\times$的性能下降。InstInfer将连续16个Token的$K$或$V$缓存分为一组,存储在同一个页内,并跨通道步进(stridden)存储。在稀疏模式分析中,只有当组内所有Token都未达到前$k$阈值时,该组才会被忽略。这种设计确保了KV缓存始终以页粒度读取,并最大化利用通道带宽。结合双步加载机制,第一步加载包含稀疏Token的完整页,第二步由NFC内的过滤器剔除组内的稀疏Token,仅将强单元加载到片上缓冲区。测试表明,双步加载在第一步保持了约一半的稀疏性,第二步达到完全稀疏。
隐藏层嵌入索引映射:由于每个隐藏层嵌入索引需要连续访问多个Token,InstInfer将多个Token对应的$K$缓存连续存储在单个页内,每个Token仅占用一个隐藏层嵌入。对于4KB页,每页可存储2K个Token。同样采用双步加载机制,将2-8个隐藏层嵌入分组到一个闪存页中。最小存储粒度为256-1K个Token,组大小可根据运行时输入长度和模型最大上下文长度动态调整。
批量写入请求:在预填充阶段后,解码阶段会持续生成新Token的KV向量并传输至CSD。由于闪存按页写入,逐个生成的Token会先缓存在CSD内的DRAM组缓冲区中,满后在后台刷新到闪存。为减轻写放大,必须确保以块(block)粒度进行写入。在读取时,不同注意力头在CSD上的并行度较小,而同一注意力头内的不同Token组需要最大化并行读取以利用通道带宽。在写入时,GPU并行生成所有注意力头的新$k, v$向量,因此可以将不同注意力头的组批量写入同一个闪存块中。InstInfer优先将不同注意力头的组放入同一块中,同时将同一注意力头的组分布在不同块(跨通道)中,从而在读取时充分利用所有通道带宽。
系统集成与扩展
GPU-CSD协同:InstInfer采用流水线协作模式。在预填充阶段,GPU处理所有计算,并将生成的KV缓存通过层级流水线(layer-wise pipeline)传输至CSD,即第$i$层生成的KV缓存与第$i+1$层的计算同时进行传输。在解码阶段,CSD仅接收$q, k, v$向量,计算完成后将注意力输出返回给GPU进行后续的$O$投影和FFN推理。与传统卸载系统相比,PCIe总线上的数据传输量减少了$s/2$($s$为序列长度)。GPU与CSD之间采用P2P DMA直接传输数据,绕过主机内存,避免了传统GPUDirect Storage【44, Gpudirect rdma URL】依赖主机文件系统的开销。InstCSD内的FTL独立管理所有数据映射和地址转换,元数据存储在其内部DRAM中。
扩展至CSD阵列:InstInfer可以无缝扩展至多个CSD。由于多头注意力机制中各个头独立计算,且InstCSD专门处理注意力模块,因此可以将不同的注意力头分配给多个CSD。对于配备$n$个CSD和$n_{head}$个注意力头的配置(通常$n_{head} \gg n$),每个CSD处理$n_{head}/n$个注意力头。最终,各CSD的输出传回GPU并拼接成最终结果。
实验环境
- 数据集:ShareGPT【52】、WikiText-2【41】、SQuAD【55】、TriviaQA【25】。输入和输出序列长度均设置为1024,以充分测试长上下文场景。
- 模型架构:OPT-13B,采用FP16格式。
-
硬件配置:
- GPU:NVIDIA A6000(48GB VRAM)。
- CPU:2.2GHz Intel Xeon 5320(96GB DDR4内存)。
- SSD:Samsung 980pro 2TB。
- 连接:PCIe Gen4x16。
- 实际部署CSD(Daisyplus OpenSSD):Xilinx ZU17EG UltraScale+ MPSoC(四核ARM,2GB DRAM),PCIe Gen3x4。
- 软件定义CSD(NVMeVirt迁移):Xilinx Zynq7045 FPGA,扩展至8通道,1.4GB/s带宽,以对齐现代SSD。
-
软件配置:基于FlexGen【53】构建,重构了数据路径以支持P2P DMA。驱动程序适配自SmartIO【40】,修改了NVMe命令。
实验结果
- 精度评估:在OPT-13B和LLaMA-2-7b模型上测试了SparF Attention、SparQ Attention、H2O【73】和Local attention的精度。结果表明,在压缩率低至1/8时,SparF Attention的精度几乎与原始SparQ Attention一致,且优于H2O和Local attention。这是因为SparF通过双步加载有效识别了KV缓存中的稀疏性。后续实验默认采用1/8压缩率(图11)。
- 单SSD(CSD)吞吐量评估:DeepSpeed在小批量(4-16)时因利用主机内存带宽而表现较好,但在批量达到32时因内存耗尽发生内核交换,吞吐量下降$32.6\times$。FlexGen最高支持批量64,但受限于PCIe带宽,吞吐量较低,且在批量128时因预填充阶段的中间KV缓存超出显存而发生OOM。InstI通过层级传输大幅降低了显存需求,支持更大的批处理大小。InstI-SparF进一步减少了KV缓存需求,在批量256时,其吞吐量比基线FlexGen提高了$11.1\times$(图12)。
- 双SSD(CSD)吞吐量评估:传统卸载方案由于依赖主机文件系统,即使增加SSD数量,性能提升也微乎其微。而InstI通过P2P DMA和内部闪存通道扩展,在配备2个CSD、批量为256时,InstI(密集)比FlexGen(批量32)的最大吞吐量高出$10.5\times$;InstI-SparF比FlexGen-SparQ高出$3.11\times$(图13)。
- 解码延迟分解:对于密集推理,与FlexGen(批量64)相比,InstI和InstI-2将KV缓存访问占比从$98.9\%$分别降至$80.7\%$和$76.4\%$;对于稀疏推理,从$92.4\%$降至$82.3\%$和$74.0\%$。端到端延迟方面,密集的InstI和InstI-SparF分别减少了$88.1\%$和$94.0\%$的KV缓存访问开销(图14、图15)。对SparF引擎的分析表明,额外的Logit-0计算过程有效帮助识别了序列中的稀疏性(图16)。
- 可扩展性与敏感性测试:部署20个CSD时,密集和稀疏推理的吞吐量分别提升了$8.99\times$和$7.29\times$,展现出极佳的注意力头级别并行扩展性。在不同压缩率下,得益于双步加载机制,InstInfer能有效应对细粒度随机访问带来的挑战,从更高的压缩率中获益(图17)。
补充细节
基于PIM的Transformer加速:多项研究【10, Unleashing the potential of pim: Accelerating large batched inference of transformerbased generative models 2023 IEEE CAL】、【18, Neupims: Npu-pim heterogeneous acceleration for batched llm inferencing 2024 ASPLOS】、【63, Pim gpt a hybrid process in memory accelerator for autoregressive transformers 2024 npj UC】、【76, Transpim: A memorybased acceleration via software-hardware co-design for transformer 2022 HPCA】利用存内处理(PIM)技术将计算单元集成在内存单元中,以加速受内存限制的注意力计算。然而,这些研究均基于模拟器,且PIM设备制造成本高昂,难以在资源受限场景中广泛部署,且其容量无法与闪存芯片相比。InstInfer则在真实硬件上部署,采用经济的FPGA和SSD,提供了更具成本效益和可扩展性的解决方案。
KV缓存管理优化:vLLM【30】借鉴虚拟内存机制,以块粒度管理GPU显存和主机内存中的KV缓存以减少碎片。LMDeploy【23】和CachedAttention【14, Cost-Efficient large language model serving for multiturn conversations with CachedAttention 2024 USENIX ATC】专注于管理主机内存和SSD上的KV缓存,以减少多轮对话中的重计算。这些工作主要优化在线推理中的预填充阶段,不适用于输出序列较长的离线推理。其他方案【17, Fastdecode: High-throughput gpu-efficient llm serving using heterogeneous pipelines 2024 arXiv】、【31, InfiniGen: Efficient generative inference of large language models with dynamic KV cache management 2024 OSDI】、【48】、【49, Mooncake: Kimi’s kvcache-centric architecture for llm serving 2024 arXiv】、【75】利用解耦资源池加速推理,但不适合资源受限场景。InstInfer利用高性价比的CSD,有效解决了解码阶段的KV缓存瓶颈。
结论
本文提出了InstInfer,这是一个新颖的基于CSD的LLM离线推理系统,旨在以高性价比的方式解决KV缓存带来的巨大存储和带宽挑战。通过将关键的解码阶段注意力计算和KV缓存卸载到具有闪存感知设计的CSD中,InstInfer充分利用了闪存芯片的高通道带宽,规避了外部PCIe带宽的限制。评估结果表明,在资源受限场景下的长上下文推理中,InstInfer的性能比现有的SSD卸载系统高出$11.1\times$。未来的工作可以进一步探索在更大规模的CSD集群上部署该系统,以及针对更多样化的LLM架构进行适配。
💬 评论讨论
欢迎在这里分享您的想法和见解!