Neo: Saving GPU Memory Crisis with CPU Offloading for Online LLM Inference

发表时间: 2025-05 · arXiv:2411.01142 (MLSys 2025)

原文: https://arxiv.org/abs/2411.01142

Xuanlin Jiang, Yang Zhou, Shiyi Cao, Ion Stoica, Minlan Yu

速读

一句话结论
NEO 提出了一个面向在线大语言模型推理的系统,通过将部分解码阶段的注意力计算和 KV Cache 卸载到宿主机 CPU,在不牺牲延迟和精度的前提下,显著提升了 GPU 的批处理大小和系统吞吐量。

要解决什么问题
在线大语言模型推理的吞吐量高度依赖于请求的批处理大小,但现代 GPU 的显存容量成为了核心卡点。在推理过程中,为了复用历史计算结果,GPU 显存中需要保存随上下文长度线性增长的 KV Cache。这导致显存很快被耗尽,无法继续增加批处理大小,大量 GPU 算力被闲置。现有的卸载方案要么采用逐层交换牺牲了在线推理必需的低延迟,要么像 FastDecode 那样将所有解码阶段的注意力计算全部卸载到 CPU,并采用对称流水线设计。这种做法不仅让 CPU 成为严重的性能瓶颈,而且在面对真实场景中长短不一的动态输入输出时,静态的调度策略会导致流水线产生大量空泡,GPU 经常处于闲置等待状态。

怎么做的
NEO 的核心思路是仅将部分请求的解码阶段注意力计算及其对应的 KV Cache 卸载到宿主机 CPU 上,利用 CPU 较大的内存容量换取 GPU 显存空间的释放,从而扩大全局批处理大小。由于解码阶段的注意力计算是典型的访存密集型操作,而现代 CPU 与 GPU 在内存带宽上的差距远小于算力差距,因此 CPU 完全有能力承担这部分工作。为了绕开传统对称流水线导致的 GPU 闲置问题,NEO 设计了两个关键部件。第一个是非对称 GPU-CPU 流水线。系统在每次迭代时构建两个不对称的子批次:Batch-0 包含所有处于 Prefill(预填充)阶段的请求和部分留在 GPU 上的 Decode(解码)请求,其特点是矩阵乘法等线性层计算量大、注意力计算极少;Batch-1 则包含绝大多数被卸载到 CPU 的 Decode 请求,特点是线性层计算极短、注意力计算耗时长。通过让 Batch-0 的 GPU 线性层计算与 Batch-1 的 CPU 注意力计算在时间上重叠,完美掩盖了 CPU 的计算延迟。第二个是负载感知调度器。面对动态变化的请求长度,调度器在每次迭代时在线决定请求的去向。设单次迭代时间为 $T$,总批处理大小为 $x$,调度器的目标是最大化吞吐效率 $T/x$。单次迭代的核心耗时 $T_{tr}$ 由网络层数 $L$、两个子批次的线性层耗时($T_{l}$)和注意力耗时($T_{ga}$ 为 GPU 注意力,$T_{ca}$ 为 CPU 注意力)决定:

$$T_{tr} \approx L \times (\max\{T_{l_0}, T_{ca_1}\} + \max\{T_{l_1} + T_{ga_0}, T_{ca_0}\})$$


为了不让 CPU 拖慢全局进度,调度器会严格维持 $T_{l_0} \ge T_{ca_1}$ 和 $T_{l_1} + T_{ga_0} \ge T_{ca_0}$ 这两个关键不变量。在每个调度周期,它会分别评估纯 GPU 执行和非对称流水线执行的预期吞吐量,贪心地选择更优方案,从而在极端负载下也能平滑退化为纯 GPU 模式。

效果如何
实验在 T4、A10G 和 8xH100 GPU 硬件环境下进行,测试了 LLaMa-2-7B、LLaMa-3.1-8B 和 70B 模型,使用了 Azure Code 和 OpenAI Summarization 等真实生产数据集。对比基线包括:代表当前主流纯 GPU 路线的 SOTA 引擎 vLLM、代表全量解码注意力卸载路线的 FastDecode+,以及作为 NEO 底座的纯 GPU 引擎 SwiftLLM。在保持与 vLLM 相同在线延迟的前提下,NEO 在 T4、A10G 和 H100 上分别实现了最高 7.5 倍、26% 和 14% 的吞吐量提升。在 T4 这种显存极度受限的低端 GPU 上,卸载 KV Cache 带来的批处理扩大红利最为显著。此外,当宿主机 CPU 的内存带宽翻倍时,NEO 在 A10G 上的吞吐量增益可进一步提升至 79.3%。该方法的局限在于:在请求并发率极低时,由于系统仍会主动寻找卸载机会,调度和显存交换的开销会导致延迟略微高于 vLLM;当工作负载的输出长度过长时,CPU 侧的访存会逐渐成为瓶颈,此时 NEO 的调度器会回退到纯 GPU 模式,表现与基线持平甚至因离线性能预测的误差而出现轻微的性能下降;同时,目前基于 Ray 实现的多 GPU 张量并行版本在工程优化上仍落后于工业级的 vLLM。

主要贡献

在线大语言模型(LLM)推理广泛依赖于请求批处理(batching)来提高吞吐量,以期在昂贵的GPU加速器上实现成本效益。然而,有限的GPU内存极大地限制了实际中能达到的批处理大小,导致大量的GPU计算资源被浪费。

本文提出了NEO,一个用于在线LLM推理的系统。该系统将部分注意力计算(attention compute)和KV Cache状态从GPU卸载(offload)到本地主机CPU,从而有效增加了GPU的批处理大小并提高了推理吞吐量。为此,NEO提出了非对称GPU-CPU流水线(asymmetric GPU-CPU pipelining)负载感知调度(load-aware scheduling),以平衡GPU和CPU的负载,并充分利用它们的计算和内存资源。

在广泛的工作负载(如代码生成、文本摘要)、GPU(T4, A10G, H100)和LLM模型(7B, 8B, 70B)上的评估表明,在保持相同延迟的前提下,NEO在T4、A10G和H100 GPU上相比纯GPU方法分别实现了最高$7.5\times$、$26\%$和$14\%$的吞吐量提升;配合更强大的CPU,NEO在A10G GPU上实现了高达$79.3\%$的吞吐量增益。

背景知识与设计原则

自回归Transformer LLM的基础工作原理:基于Transformer的自回归LLM将token视为语言的基本元素,主要任务是预测给定序列的下一个token。形式上,给定输入$[t_1, t_2, \dots, t_n]$,对于词汇表中的每个token $x$,LLM需要返回该token作为序列下一个token的概率,即$\Pr[t_{n+1}=x | t_1, \dots, t_n]$。为了进行这种预测,Transformer模型首先将token转换为嵌入向量(embedding vectors)。随后,这些向量通过一系列称为Transformer层的基本块,并利用注意力机制进行处理。嵌入向量在通过每个Transformer层后形状保持不变,但变得更加精确和具有上下文感知能力,最后通过全连接层将嵌入向量转换为每个可能token的对应概率。

推理的两个阶段与KV Cache:一个推理请求通常包含预填充(prefilling)和解码(decoding)阶段,并严重依赖GPU内存中的KV Cache来复用先前的计算。预填充阶段在消耗所有输入token后生成初始KV Cache;而解码阶段则重复读取并追加KV Cache,自回归地生成输出token,直到遇到EOS(序列结束)。为了优化此过程,现代LLM推理引擎(如vLLM)利用迭代级调度来适应请求的不同输入/输出长度,利用选择性批处理来通过批处理矩阵乘法提高性能,并利用分页注意力(paged attention)来高效管理GPU内存。

吞吐量瓶颈:LLM推理的吞吐量高度依赖于引擎能够达到的批处理大小,而这本质上受到GPU内存大小的限制(由于庞大的KV Cache)。先前的研究表明,在现代A10、V100和H100 GPU上,随着批处理大小增加到数百,吞吐量几乎呈线性增长;但遗憾的是,这些GPU无法容纳数百甚至几十的批处理大小,导致显著的GPU计算资源未被利用。
图1. 基于Transformer的LLM推理工作流程。

CPU内存作为解决方案的机遇与挑战:增加批处理大小的一种方法是将溢出的KV Cache存储在CPU的主存中,在需要时传输到GPU,不需要时传回。然而,这种重复的KV Cache换入换出会导致系统严重受限于GPU-CPU PCIe带宽。幸运的是,只有解码阶段的注意力操作依赖于KV Cache,因此将这部分计算卸载到CPU可以避免在GPU和CPU之间重复传输KV Cache。此外,该操作在计算中仅占极小比例,且不需要加载模型权重。

解码注意力操作的带宽限制特征:由于算术强度(即每次内存加载的FLOP)较低,解码注意力操作在GPU和CPU上都受限于内存带宽。GPU和CPU之间的内存带宽差距远小于它们的计算差距。例如,A10G GPU具有$600\text{ GB/s}$内存带宽和$125\text{ TFLOPS}$,而现代x86服务器具有约$200\text{ GB/s}$内存带宽和$1.2\text{ TFLOPS}$。因此,尽管GPU和CPU之间看似存在巨大的计算差距,但解码注意力操作的实际性能差距相对较小,因为它们的内存带宽更为接近。

系统设计的核心挑战:构建高效的LLM推理系统面临两个主要挑战。首先,在每个推理迭代中,如何高效地重叠GPU和CPU?GPU具有更多的计算和内存带宽但内存容量有限,而CPU内存大但算力弱。必须仔细重构LLM推理的流水线,将不同的模块适配到正确的硬件,同时不使任何硬件超载,且不能破坏Transformer语义。其次,在动态变化的真实工作负载中,如何跨推理迭代调度推理请求?先前的研究通常假设输入/输出长度固定,并采用静态最优调度策略。但在真实的动态设置中,差异巨大的输入/输出长度会破坏静态调度的最优性,这需要一种自适应的调度策略在每次迭代级别确定最佳的请求分配。

方法细节

NEO系统整体架构:NEO包含一个运行在CPU上的请求调度器,该调度器维护一个预填充等待队列、一个GPU解码运行队列和一个CPU解码运行队列。该调度器在迭代级别做出自适应调度决策,决定传入的请求应在GPU还是CPU上运行。NEO具有两个关键技术:1)非对称流水线(asymmetric pipelining),用于充分利用GPU和CPU的计算资源而不使其超载;2)负载感知调度(load-aware scheduling),用于处理动态变化的工作负载(例如跨请求的不同输入/输出长度)。
图2. NEO的整体架构。“runQ”表示“运行队列”。

简单卸载(Simple offloading)方案的缺陷:为了探索非对称流水线的设计动机,首先考察一个简单的草图设计。该设计提取解码注意力并将其计算和KV Cache卸载到CPU,而将其余部分留在GPU。其余部分包括预填充注意力和token独立的线性操作(主要涉及矩阵乘法)。然而,在这些线性操作期间,CPU始终保持空闲。因此,这种设计未能有效利用CPU的计算和内存资源。
图3. 简单卸载方案草图,将所有请求的KV Cache和解码注意力计算卸载到CPU。“Comm”代表GPU-CPU通信;“TrQKV”代表传输Q,K,V张量到CPU;“TrO”代表将注意力输出传输到GPU。

对称流水线(Symmetric pipelining)方案的缺陷:减少CPU空闲时间的一种直接方法是将单个解码批次均匀分成两个子批次,并重叠它们的线性操作(在GPU上)和注意力操作(在CPU上)。对于预填充批次,对称流水线仅在GPU上运行而不进行卸载。这种看似高效的设计存在三个主要问题。首先,导致GPU内存严重未被充分利用。GPU仅保留模型权重和运行时激活,而原本用于存储KV Cache的内存保持未使用状态。其次,未能实现平衡的GPU-CPU重叠。因为它完全忽略了预填充阶段和KV Cache换出时间(此时CPU空闲);且由于注意力的自回归特性和高内存带宽消耗,GPU上解码子批次的线性阶段通常比CPU上的注意力阶段短得多,导致CPU成为瓶颈。最后,在真实工作负载中,由于不同的输入长度和不可预测的输出长度,几乎不可能确保单个批次能被分成两个完全相同的子批次,这会导致流水线中出现大量空闲期或“气泡”。
预填充阶段图示1
预填充阶段图示2

非对称流水线与部分卸载:为了充分利用GPU内存,NEO执行部分卸载。NEO的KV Cache系统被分为两个独立的组件:位于GPU HBM中的“GPU-cache”和位于CPU主存中的“CPU-cache”。对于系统中已预填充的请求,其KV Cache要么完全驻留在GPU-cache中(称为“GPU-request”),要么完全驻留在CPU-cache中(称为“CPU-request”)。请求被优先存储在GPU-cache中以最大化GPU内存利用率。

预填充集成与层级交换技术:为了实现完全的GPU-CPU重叠,NEO将预填充阶段的计算集成到GPU解码子批次中,使得预填充阶段计算(在GPU中)与CPU注意力计算并行发生。这种选择性批处理大大延长了GPU计算的持续时间,允许更长的整体CPU计算时间,并使更多CPU请求能够被纳入批次。此外,NEO利用层级交换(layer-wise swapping)技术来促进KV值传输与计算的重叠。由于新预填充请求的KV值是逐层计算的,系统可以在每层KV值计算后立即启动PCIe传输,而不是将此过程推迟到整个迭代结束。
图5. 非对称流水线将预填充阶段集成到一个子批次(红箭头),并将大部分解码注意力操作集成到另一个子批次(蓝箭头)。“pr”代表预投影,“po”代表后投影+FFN操作;“attn”代表注意力操作;“Comm”代表GPU-CPU通信。

非对称批次划分策略:为了简化批次划分方案并最小化空闲期,NEO引入了非对称批次划分。系统不再试图创建两个相同的批次,而是将所有预填充请求、GPU解码请求和少量CPU解码请求合并到batch-0,同时将绝大多数CPU解码请求分派给batch-1。这两个批次是互补的:batch-0具有较长的线性阶段和极少的CPU注意力计算,而batch-1包含冗长的注意力阶段和极短的线性阶段。这种安排不仅简化了系统实现,而且促进了两个批次之间的有效重叠,导致推理迭代呈现交替的“长阶段”和“短阶段”。此外,这种有意的不对称性最小化了开始和结束时的非重叠段,并显著减少了GPU内核启动开销。在非对称流水线中,GPU注意力内核每次迭代仅被调用一次。

负载感知调度的四大原则:在复杂的真实世界推理工作负载中,NEO调度器自适应地决定传入请求应放置在GPU还是CPU上。调度器遵循四个原则:1)贪婪(Greedy):在每次迭代开始时,生成纯GPU推理调度和双批次非对称流水线调度,并选择预估吞吐量较高者。2)平衡(Balancing):对于非对称流水线,调度器应最小化流水线“气泡”,使预估的CPU繁忙时间和GPU繁忙时间尽可能接近。3)隐藏CPU(Hiding CPU):对于非对称流水线,不应出现CPU繁忙而GPU空闲的情况。4)最大化GPU(Maximizing GPU):调度器应尽可能多地从预填充等待队列以及GPU和CPU解码运行队列中挑选请求。

迭代时间的数学估计:设每次迭代时间为$T$,批次大小为$x$,NEO的目标是最大化$T/x$。$T$主要由Transformer层时间$T_{tr}$组成(通常占每次迭代时间的95%以上)。$T$的估计公式为:

$$T \approx T_{tr} = L \times (\max\{T_{po_0} + T_{pr_0}, T_{ca_1}\} + \max\{T_{po_1} + T_{pr_1} + T_{ga_0}, T_{ca_0}\})$$


其中$T_{po_x}, T_{pr_x}, T_{ga_x}, T_{ca_x}$分别表示批次$x$的后投影、预投影、GPU注意力和CPU注意力时间。定义$T_{l_x} = T_{po_x} + T_{pr_x}$,代表在一个Transformer层中主要由激活值与模型权重相乘消耗的时间。为了坚持“平衡”和“隐藏CPU”的原则,系统希望保证$T_{l_0} \geq T_{ca_1}$且$T_{l_1} + T_{ga_0} \geq T_{ca_0}$,并最小化这两个不等式左右两边的差距。为了估计这些时间值,NEO对典型的输入/输出长度进行离线分析,并使用线性插值来近似其他长度的值。

负载感知调度器的具体执行步骤
* 1. 初始化:初始化两个空批次调度:batch-0主要在GPU上运行,包含预填充和解码阶段的请求;batch-1主要在CPU上运行,仅包含计算解码注意力的请求。
* 2. 调度GPU解码请求:尝试将GPU解码运行队列中的每个请求放入batch-0。然后,如果GPU内存无法容纳所有新KV Cache则换出请求,如果GPU有充足空间则换入请求(最大化GPU)。
* 3. 调度预填充请求:弹出预填充等待队列并将请求放入batch-0:如果有足够的GPU内存,则将生成的KV Cache保留在GPU上,否则换出生成的KV Cache。重复此操作,直到GPU无法容纳批次中的激活值(最大化GPU)。
* 4. 调度CPU解码请求:扫描CPU解码运行队列,将请求放入batch-0或batch-1,同时维持$T_{ca_0} \leq T_{l_1} + T_{ga_0}$和$T_{ca_1} \leq T_{l_0}$。如果将请求放入任一梯队会违反不等式,则跳过该请求并将其留到下一次迭代(平衡和隐藏CPU)。
* 5. 减少预填充请求:只要上述不等式成立,就从batch-0中移除任何需要换出生成的KV Cache的预填充请求。这是为了避免CPU空闲(平衡)。
* 6. 做出决策:现在双批次调度已制定完成;通过采用batch-0并排除步骤4中添加的所有CPU解码请求来制定纯GPU调度。最后,比较它们的预估$T_{tr}/x$值,并选择值较大的调度(贪婪)。

系统实现与CPU内核优化:NEO基于SwiftLLM(一个用于研究的轻量级LLM推理系统)实现,包含4K行Python代码和1.5K行C++代码。为了实现高效的CPU内核,实现了一个名为Paged-Attention-for-CPU (PACPU) 的自定义C++ Torch扩展库。PACPU调用由ISPC编译器编写的扩展库,ISPC代码可编译为具有不同向量化指令集的目标(如AVX2、AVX512和ARM Neon)。PACPU利用类似于vLLM的分页KV Cache和Paged Attention算法来减轻内存碎片。针对解码注意力在CPU上严重受限于内存带宽的问题,在单核内利用SIMD内存加载/存储指令最小化指令开销,并仔细组织内存访问顺序以确保连续性;在跨核间,使用类似Flash Decoding的并行策略,将请求的计算沿请求维度划分为独立任务,均匀分发给所有线程处理,最后聚合部分输出。

减少内核启动开销与多GPU支持:由于Python的全局解释器锁(GIL)导致多线程执行性能差,数据平面CPU内核无法与控制平面CUDA内核启动调用并行运行。为了减少内核启动开销,将SwiftLLM的大部分Triton-JIT内核替换为CUDA C++编写的内核,有效减少了额外的内核选择和启动开销。此外,重新设计了架构以支持模型分片和张量并行,利用Ray actor持有模型分片,使用PyTorch的预构建通信库(基于NCCL)处理跨GPU通信。每个Ray actor也有其CPU KV Cache分区,负责一部分KV头以避免跨CPU通信。

实验环境

实验结果

在线延迟与负载评估
* 实验内容:在不同请求率(服从泊松过程采样)下评估NEO的在线延迟,计算每个请求的平均每token延迟。
* 实验结果:NEO在所有列出的硬件/模型设置中都能比vLLM承受更高的负载,同时在低速率下提供可比的延迟。在2秒延迟下,NEO在H100上实现$14.3\%$的吞吐量提升,在A10G上提升$6.40\%$;在1秒延迟下,在T4上提升$563\%$(近$6\times$)。在极低请求率下,NEO的延迟行为与vLLM完全相同;随着请求率增加,延迟略高,这是因为NEO的工程优化较少,且主动寻找卸载机会带来了调度和交换的系统开销。
* 分析结论:NEO的吞吐量提升并不以牺牲延迟为代价,特别是在内存极度受限的T4 GPU上,卸载KV Cache到CPU带来了巨大的批处理大小提升。
图6. NEO与vLLM在3种不同设置下的负载-延迟曲线对比。
图7. 在1.6/s请求率下,A10G+LLaMa-3.1-8B+AC设置中NEO和vLLM的延迟分布。

与FastDecode+的对比评估
* 实验内容:在$2\times\text{H100}+\text{LLaMa-3.1-70B}$设置下,对比NEO与FastDecode+的在线延迟和离线吞吐量(固定输入长度为2000,改变输出长度)。
* 实验结果:FastDecode+表现出更高的延迟,且随着输出长度增长,其吞吐量迅速下降至基线的$60\%$以下。而NEO在吞吐量上始终优于基线。
* 分析结论:FastDecode+的高延迟是由其处理不规则工作负载的不灵活性引起的(当CPU队列有请求而预填充队列空时,被迫启动CPU批次阻碍整体性能)。而NEO的调度器可以决定回退到纯GPU模式,从而始终保持优势。
图8. 2xH100+LLaMa-3.1-70B设置下NEO和FastDecode+的比较。图(a)比较AC数据集上的延迟;图(b)显示相对基线的吞吐量。

变化输入/输出长度的评估
* 实验内容:固定输入长度,调整输出长度,在合成工作负载上测试相对吞吐量。
* 实验结果:NEO在H100、A10G和T4上分别实现了最高$14\%$、$26\%$和$750\%$的吞吐量增益。当输出长度较短时,NEO可能表现略差于基线;随着输出长度增加,增益增长到最高点(此时GPU和CPU时间完全平衡);随后逐渐下降并逼近基线水平。
* 分析结论:输出极短时,尝试卸载并换回请求会产生微小开销;极长时,系统启动更大比例的纯GPU批次,性能趋近基线。
图9. 不同设置和不同合成工作负载下的相对吞吐量。

敏感性研究评估
* 实验内容:1)在具有不同CPU核心数和内存带宽的AWS g5实例(g5.2/4/8/16xlarge)上测试CPU容量的影响。2)在单GPU和多GPU设置下对比纯GPU基线(SwiftLLM)与vLLM的吞吐量。
* 实验结果:1)在不同CPU容量下,NEO分别实现了最高$12.2\%$、$13.3\%$、$29.7\%$和$79.3\%$的吞吐量提升。峰值吞吐量增益与CPU内存带宽呈正相关。2)SwiftLLM在单GPU下与vLLM吞吐量相当,但在双GPU下略低($8.8\%$)。
* 分析结论:决定CPU上注意力操作性能的因素是内存带宽,而不是计算能力(核心数)。基线系统在多GPU下的张量并行实现相比生产级的vLLM优化较少。
图10. 敏感性研究。图(a)显示不同AWS EC2 g5实例上的相对吞吐量;图(b)比较SwiftLLM与vLLM的吞吐量。

补充细节

与chunked-prefill的比较:NEO的性能优势来源于更大的GPU批处理大小以及将不可批处理的解码注意力转移到CPU。这与Sarathi-Serve提出的chunked-prefill技术有相似之处。然而,chunked-prefill消耗显著更多的GPU内存带宽(需重复加载先前块的KV Cache),且在内存受限的GPU上效果不佳。相比之下,NEO依赖大容量CPU内存来确保吞吐量增益。这两者可以结合以提供更大的设计空间。

将其他部分卸载到CPU:NEO隐式假设将所有模型权重放在GPU上,仅将注意力计算卸载到CPU是平衡负载的最高效方式。然而,在极端工作负载(如输出token极少)下,NEO会受限于GPU计算而CPU大部分空闲。在这些情况下,将一些密集操作卸载到CPU(如FlexGen等工作所做)可能会缓解GPU压力。

NEO的使用场景与局限:NEO最适合GPU内存受限导致批处理大小受限、GPU算力未充分利用的场景。随着GPU算力持续增长而内存容量相对停滞(如H100算力是A100三倍但内存同为80GB),这种场景将长期存在。当GPU有足够内存达到饱和时,NEO将退化为非卸载模式。此外,NEO专注于使用本地主机CPU,扩展到远程CPU仍受限于高昂的跨机器传输延迟和网络带宽成本。

结论

NEO是一个用于在线LLM推理的CPU卸载系统,旨在增加GPU批处理大小并提高推理吞吐量。它采用非对称流水线和负载感知调度,在不超载的情况下充分利用GPU和CPU资源。在各种工作负载和模型大小下,NEO在保持相同延迟的同时,实现了比纯GPU推理系统高出$14\%$到$7.5\times$(取决于GPU)的吞吐量。未来工作将开源NEO代码库,以鼓励对经济高效的LLM推理进行更多研究。

参考文献引用汇总

在方法细节和背景阐述中,作者引用了多项关键文献来支撑其设计选择和系统比较:
* vLLM(Kwon 等人,2023,Proceedings of ACM SOSP):被引用以说明现代LLM推理引擎广泛依赖批处理、分页注意力(PagedAttention)来高效管理GPU内存,并指出批处理大小本质上受到KV Cache带来的GPU内存大小限制。
* FastDecode(He & Zhai,2024,arXiv):被多次引用。首先用于说明计算卸载可以避免重复交换KV Cache;其次作为对称流水线的反面案例,指出其将单个解码批次均匀拆分会导致CPU严重成为瓶颈;此外,还被用于对比说明跨机器传输延迟会成为远程CPU卸载的瓶颈。
* FlexGen(Sheng 等人,2023,ICML):被引用以说明先前的卸载工作(将模型权重、KV Cache和计算卸载到CPU)通常通过使用巨大的GPU批处理大小和逐层交换来牺牲推理延迟以换取吞吐量,因此不适合在线推理;并且其依赖静态的最优卸载策略。
* Orca(Yu 等人,2022,USENIX OSDI):被引用以说明迭代级调度(iteration-level scheduling)和选择性批处理(selective batching)技术。NEO的非对称流水线中将预填充阶段计算集成到GPU解码子批次中的设计,与Orca中的选择性批处理技术精神相似。
* Sarathi-Serve(Agrawal 等人,2024,USENIX OSDI):在讨论章节中被引用,说明NEO转移不可批处理的解码注意力操作的思想,与Sarathi-Serve中的chunked-prefill技术有相似之处。
* Flash Decoding(Dao 等人,2023):被引用以说明NEO在实现CPU跨核并行策略时,采用了类似于Flash Decoding的策略,即沿请求维度划分计算任务。