FlexiCache: Leveraging Temporal Stability of Attention Heads for Efficient KV Cache Management
FlexiCache: Leveraging Temporal Stability of Attention Heads for Efficient KV Cache Management
作者/机构: Nazmul Takbir 1 Hamidreza Alikhani 1 Nikil Dutt 1 Sangeetha Abdu Jyothi 1
A1 主要贡献
大型语言模型(LLM)服务日益受到键值(KV)缓存大小的限制,其大小随上下文长度和生成长度的增加而增长。先前的工作表明,注意力计算由一小部分关键的token主导,但现有系统难以在不降低准确性的情况下有效利用这一点,尤其是在长文本生成场景中。本文提出了一个关键观察:这些关键token的时间稳定性在不同的KV头之间存在显著差异。一些头(稳定头)持续关注相同的token,而另一些头(不稳定头)的关注点则频繁变化。
基于这一发现,本文引入了 FlexiCache,一个分层KV缓存管理系统,它利用KV头的时间稳定性来减少GPU内存使用和计算开销,同时保持模型准确性。FlexiCache 将KV头分为稳定或不稳定两类:它将不稳定头的所有KV缓存页保留在GPU内存中,而对于稳定头,仅在GPU上保留Top-K的页,并将其余页卸载到主机内存中。通过利用时间稳定性,FlexiCache对稳定头进行周期性重排序,以获取新晋升的Top页。
核心贡献如下:
* 发现KV头的模型内在时间稳定性:本文识别出KV头在选择其Top-K KV页进行注意力计算时,表现出一种模型固有的、跨任务一致的时间稳定性。
* 开发FlexiCache系统:设计并实现了一个利用KV头时间稳定性的LLM服务系统,该系统具有分层KV缓存管理机制。
* 显著的性能提升:在vLLM基础上实现的FlexiCache,能够:
* 将长上下文请求的GPU内存占用减少高达70%。
* 将离线服务吞吐量提高1.38-1.55倍。
* 将在线服务的token延迟降低1.6-2.1倍。
* 在长上下文、长生成场景中保持模型的高准确性。
- 即将开源:FlexiCache将很快开源。
A3 背景知识与关键观察
基于Transformer架构【【22,Attention is all you need,2017,NIPS】】的LLM推理分为两个阶段:prefill(预填充)和decode(解码)。在prefill阶段,所有输入token被并行处理,模型为每个token生成查询(Q)、键(K)和值(V)向量,其中K和V向量存储在KV缓存中供后续使用。在decode阶段,模型一次生成一个token,并在每一步追加新的K/V对。decode阶段是内存带宽受限的,提高解码吞吐量的常用策略是增加批处理大小。然而,批处理大小最终受限于可用于容纳所有活动请求的KV缓存的GPU内存。此外,GPU内存碎片化进一步限制了有效的批处理大小。PagedAttention【【14,Efficient memory management for large language model serving with pagedattention,2023,SOSP'23】】通过将KV存储在通过块表解引用的固定大小页面中来解决此问题,从而最大限度地减少了碎片并实现了更大的批处理大小。
由于KV缓存随上下文长度线性增长,大的上下文长度不仅限制了批处理大小,还增加了自注意力的计算和I/O开销,因为每个解码步骤都需要获取整个KV缓存。先前的工作表明,在每个解码步骤中,一小部分KV缓存页主导了注意力。Quest【【21,Quest: query-aware sparsity for efficient long-context llm inference,2024,ICML'24】】引入了一种方法,在每个解码步骤中为每个注意力头独立地识别这些子集。虽然在每一步选择新的top-K页可以保持语言建模的准确性,但这会产生大量的计算开销,并要求完整的KV缓存驻留在GPU上。为了减轻这一成本,LServe【【28,LServe: Efficient long-sequence LLM serving with unified sparse attention,2025c,MLSys'25】】只对一半的头进行重排序,并周期性地(每四个解码步骤)更新它们。然而,这种部分重排序策略可能导致准确性下降,特别是在长生成场景中,因为未被重排序的头可能会永久丢弃后来变得相关的KV条目。
2.1 Top-K选择中的时间稳定性
我们做出了一个关键观察,即许多KV头在其top-K页面的选择上表现出时间稳定性:在连续的解码步骤中,相同的页面倾向于保持在最受关注的页面之列。此外,在给定的模型中,稳定性的程度在不同的头之间系统性地变化——一些保持高度稳定,而另一些则持续不稳定。在接下来的分析中,我们(i)形式化了这种时间稳定性的量化度量,(ii)为每个模型识别出最不稳定的头,并证明它们的不稳定模式是模型固有的,并且在不同任务中是一致的,以及(iii)利用这些见解来设计一种高效的KV缓存管理机制。
图 1. KV头的时间稳定性模式。以 Llama3.1-8B-Instruct 第4层为例。一些头在不同偏移量上保持高RCO,而其他头则持续显示低值。
2.2 量化KV头的时间稳定性
数据收集。我们首先从LongBench【【2,LongBench: A bilingual, multitask benchmark for long context understanding,2024,ACL】】和L-Eval【【1,L-eval: Instituting standardized evaluation for long context language models,2024,ACL】】中的几个长上下文、长生成任务样本中,收集每个解码步骤的top-K页面索引集。具体来说,对于一个给定的样本,我们将数据存储在一个形状为[D, L, H, K]的张量中,其中D表示解码步骤数,L表示层索引,H表示层内的KV头索引,K表示选择的top-K页面数量。换句话说,对于每个样本和每个层-头对,我们记录在每个解码步骤中被关注的top-K页面的索引。
随机校正重叠(RCO)度量。对于头h和解码步骤s,令$S_h^{(s)}$表示特定样本的top-K页面索引集。我们计算$S_h^{(s)}$与同一头在固定窗口W步内的后续top-K选择之间的重叠程度。为了将真正的时间持久性与随机机会导致的重叠分离开,我们计算步骤s和步骤$t = s + \Delta (1 \leq \Delta < W)$之间页面选择的随机校正重叠(RCO)如下:
$$\mathrm{RCO}_{h}(s, t)=\max \left(0, \frac{|S_{h}^{(s)} \cap S_{h}^{(t)}| / K-K / N_{t}}{1-K / N_{t}}\right),$$其中$N_t$是步骤t可用的候选页面数。该公式减去了超几何随机抽取模型下的预期重叠($K/N_t$),并由最大值$1 - K/N_t$进行归一化,该值代表了随机机会重叠与完美重叠之间的差距。直观地,RCO量化了观察到的重叠比随机机会好多少。RCO为0表示随机水平的重叠,1表示完美重叠。
时间稳定性分数。我们将窗口内所有对的RCO值进行聚合,并为每个KV头定义一个时间稳定性分数如下:
$$TS_{h}(s)=\frac{1}{W-1} \sum_{t=s+1}^{s+W-1} \operatorname{RCO}_{h}(s, t).$$较低的稳定性值表示一个头在其注意力焦点在解码步骤窗口内的变化更大。图1展示了Llama-3.1-8B-Instruct第4层的八个KV头在L-Eval的SPACE(酒店评论摘要)任务上的这种行为。对于参考解码步骤的每个偏移量,我们绘制了在所有窗口和样本中平均的RCO均值。一些头持续表现出低RCO(不稳定),而其他头则保持高RCO,并随偏移量逐渐衰减(稳定)。
表 1. 不稳定头的跨任务重叠。针对 Llama-3.1-8B-Instruct。值为交集大小除以不稳定头集大小($|A_i \cap A_j|/64$)。GovReport 来自 LongBench【【2,LongBench: A bilingual, multitask benchmark for long context understanding,2024,ACL】】;其余来自 L-Eval【【1,L-eval: Instituting standardized evaluation for long context language models,2024,ACL】】。
2.3 识别模型固有的稳定性模式
头分类。给定每个KV头的每个窗口的时间稳定性分数,我们通过选择最不稳定的25%的头作为不稳定头,其余75%作为稳定头,将头分为稳定和不稳定两类。具体来说,在所有样本内和样本间的所有解码窗口中,我们按其稳定性分数对KV头进行排序,并将排名后四分之一中最频繁出现的头指定为不稳定头,其余的则被认为是稳定头。
跨任务一致性。我们对每个评估任务独立执行此过程,为给定模型每个任务生成一个不稳定头集。这些集中的每一个都具有相同的基数(总头数的25%),从而可以直接跨任务进行比较。为了评估头的不稳定性是否是模型的固有属性,我们测量了同一模型在不同任务之间的不稳定头集的成对重叠。Llama3.1-8B-Instruct的跨任务重叠矩阵如表1所示。重叠度始终很高,跨任务的平均重叠度为0.83,有几对接近0.90,这表明不稳定头的身份在不同任务中基本保持不变。因此,头的不稳定性是模型的固有特性:某些头在解码步骤中天生表现出不稳定的KV页选择行为。此外,由于不稳定头在不同任务中是一致的,每个模型只需进行一次离线分析即可识别它们,以供后续在线服务使用。我们在其他模型(Mistral-7B, Mistral-Small-24B, 和 Qwen2.5-32B)上也观察到类似的高跨任务重叠度;详见附录A.1。
图 2. FlexiCache 系统架构。在工作节点(worker),top-K选择器为每个头识别最相关的KV页,并根据头的稳定性以不同频率更新它们。稀疏解码核仅关注这些选定的页。GPU内存存储不稳定头的完整KV缓存和稳定头的top-K页,其余页存储在主机内存中。块分配器管理这种分层KV布局,而KV传输模块和调度器则将主机-GPU的KV传输与计算流水线化。
A2 方法细节
我们提出了FlexiCache,一个LLM服务系统,它利用KV头固有的稀疏性和时间稳定性来高效管理KV缓存。FlexiCache的主要目标是:(i)G1:减少注意力计算开销,(ii)G2:优化GPU内存使用,(iii)G3:最小化I/O传输,以及(iv)G4:保持高质量的语言建模能力。FlexiCache通过优先处理并仅在GPU内存中保留最关键的KV缓存页,即使在长上下文和长生成场景下也能同时实现这些目标。
3.1 FlexiCache系统概述
多管齐下的设计。FlexiCache采用多管齐下的设计来实现其目标。利用关键KV页选择的时间稳定性,它首先将头分为稳定和不稳定两类。为了减少计算负载(G1),FlexiCache采用稀疏注意力,在解码阶段所有头仅使用top-K的KV页进行注意力计算。为了提高GPU内存效率(G2),它仅将稳定头的top-K KV页保留在GPU内存中,同时将不稳定头的所有KV页保留在GPU上。稳定头的其余KV页则存储在主机内存中。为了最小化主机-GPU数据传输(G3),FlexiCache控制重排序的频率:不稳定头每一步都进行重排序,但由于其所有KV页都在GPU上,因此无需传输;稳定头仅周期性地重排序,每次重排序后,只从主机内存中获取新晋升的top-K页。最后,通过共同利用KV头的稳定性和注意力稀疏性,FlexiCache确保top-K页面选择作用于整个KV缓存(因为没有页面被永久丢弃),从而维持了高语言建模性能(G4)。
系统架构。FlexiCache在vLLM之上实现,其主要组件如图2所示。稳定性感知重排序器(stability-aware re-ranker)为每个头独立识别top-K KV页,并根据其稳定性以不同频率进行。稀疏解码模块(sparse decode module)仅使用每个头的top-K KV页进行注意力计算。KV传输模块(KV transfer module)管理GPU和主机内存之间的数据移动:它将稳定头的所有KV页卸载到主机内存,并在每次重排序后周期性地仅重新加载新晋升的top-K页。FlexiCache控制器协调工作节点(worker)上的请求执行,并包括两个关键组件。FlexiCache调度器(Scheduler)有效地将计算和KV重载在批处理请求中流水线化,以最大化GPU利用率。FlexiCache块分配器(Block Allocator)扩展了vLLM的分配器,具有分层内存感知能力,为稳定头的top-K KV页和不稳定头的整个KV缓存分配足够的GPU KV块,同时分配主机KV块来存储稳定头的完整KV缓存。
3.2 KV页评分机制
评分方法。FlexiCache使用的页面重要性评分机制与Quest【【21,Quest: query-aware sparsity for efficient long-context llm inference,2024,ICML'24】】和LServe【【28,LServe: Efficient long-sequence LLM serving with unified sparse attention,2025c,MLSys'25】】中的类似。对于每个KV页面,我们维护两个向量——该页面内每个维度的最小和最大键值。给定当前查询向量q,页面p的重要性被估计为q与该页面中任何键之间可能的最大点积:
$$s_p = \sum_i \max(q_i \cdot k_{p,i}^{\min}, q_i \cdot k_{p,i}^{\max})$$关键差异。然而,FlexiCache在两个关键方面有所不同。首先,页面评分对稳定头和不稳定头以不同的频率进行:不稳定头每一步都进行评分,而稳定头由于其时间稳定性,仅周期性地(每16步)进行评分。如果一个层完全由稳定头组成,该层的评分通常可以完全跳过。其次,FlexiCache通过将min-max向量存储在专用的MinMax缓存中,将其与KV缓存页解耦,即使相应的KV页被卸载到主机内存,该缓存也保留在GPU上。这使得系统能够评估一个已卸载页面的重要性,并在必要时将其重新加载到GPU内存。MinMax缓存维护自己的块管理器,并分配大块,每个块包含多达128个KV缓存块的min-max向量,以最小化分配开销。这些优化共同显著降低了页面评分和排序的开销。
图 3. 分层KV缓存布局。示例为一个在双层、双头模型上运行的、包含四个逻辑KV页的请求。
3.3 分层KV缓存管理
基于稳定性的分层布局。KV缓存的分层管理建立在KV头在选择其top-K KV页时的时间稳定性(§ 2.2)之上。我们将最不稳定的25%的头分类为不稳定头,其余的头为稳定头。对于稳定头,FlexiCache仅将top-K最重要KV页保留在GPU内存中,同时将其余的卸载到主机内存。相反,对于不稳定头,它将整个KV缓存保留在GPU上。
图示示例。图3用一个包含两层、每层两个KV头的模型的简单示例来说明这种分层布局。该请求包含四个逻辑页,其中三个是满的,top-K大小为2。对于稳定头,只有两个top-K页驻留在GPU内存中,而主机内存存储所有三个完整的KV页。最后一个部分填充的页,由于正在被追加,所以总是在top-K中,一旦它变满,将被卸载到主机内存。对于不稳定头,所有四个KV页都保留在GPU内存中,并且永远不会被卸载。
图 4. FlexiCache流水线。KV卸载与同一请求的计算重叠,而KV重载与批处理中其他请求的计算重叠。
开销最小化策略。为了使分层KV缓存管理变得实用,FlexiCache通过以下方式最小化在主机和GPU内存之间移动KV页的开销:
最小化传输大小。在主机到GPU的方向上,我们只传输新晋升的页面——即那些在之前的top-K集合中但不在当前集合中的页面——在对稳定头进行周期性top-K页面重排序之后。这减少了传输的频率和大小,因为top-K集合本身只占总KV缓存的一小部分,而时间稳定性进一步减少了top-K集合之间的差异。为了最小化GPU到主机的传输,稳定头的每个KV页只被卸载一次。在prefill阶段之后会发生一次大的卸载,随后在解码过程中随着新填充页面的完成,会进行小的增量卸载。因此,主机内存始终保存着稳定头的完整KV缓存,而GPU只保留top-K页面。
将传输与计算重叠。为防止KV缓存移动拖慢计算流水线,GPU和主机内存之间的传输与计算被精心重叠,如图4所示。在prefill阶段之后,初始的GPU→主机卸载在一个低优先级的后台流上异步开始,而主流程继续解码。一旦完成,块管理器会释放不再属于top-K的页面的GPU块。这显著减少了每个请求的GPU内存使用,并实现了更大的批处理大小,因为现在更多的请求可以容纳在GPU内存中。后续的增量卸载,在页面变满时触发,也以类似的方式重叠。
处理主机到GPU的传输。重叠主机→GPU传输更具挑战性,因为晋升的KV页会修改GPU上的缓存。对同一请求的并发解码可能导致不一致的读取,特别是在启用物理块回收(第3.4节)时。为防止这种情况,FlexiCache调度器仅暂时暂停正在重新加载的请求,而其他请求继续解码。性能影响微乎其微:重新加载的大小很小(第3.3节),并且通常在一个解码步骤内完成。由于不同的请求在不同的时间暂停,任何时候只有一小部分批次(对于重排频率为16的情况,大约为1/16)处于空闲状态,从而保持了整体的高吞吐量。
减少碎片化开销。Paged attention将每个请求的KV缓存分散在内存池中。一个简单的GPU↔主机传输(例如,使用PyTorch的高级索引)会首先在源端收集分散的块,复制一个连续的缓冲区,然后在目标端再次分散它们。CPU端的收集/分散过程主导了此过程的延迟。
自定义CUDA内核。为消除这种开销,我们实现了一个自定义CUDA内核,它通过统一虚拟寻址(UVA)直接从固定的主机KV缓存中读取或写入,完全避免了CPU端的收集/分散操作。然而,主机映射的读写受限于PCIe带宽,因此它们在等待I/O时会占用GPU流式多处理器(SMs)。为了减少与主计算内核的竞争,我们采用了三种技术:(1)在低优先级的后台流上启动传输;(2)限制网格大小以保持SM占用率低;以及(3)将大传输分割成多个较小的块,启动一系列短内核。这些策略减少了每次启动的停滞时间,并为CUDA调度器提供了在低优先级传输内核之间交错执行高优先级计算内核的机会。
3.4 高效的块表管理
问题背景。Paged attention维护一个块表,将每个请求的逻辑索引映射到GPU内存中的物理块。在标准的密集注意力中,所有头和层的KV缓存均匀增长,允许整个模型使用一个共享的块表。FlexiCache打破了这一假设,因为每个头-层对根据其自己的top-K选择独立地晋升和驱逐KV页。因此,它们的内存布局会发生分歧,一个单一的共享块表无法准确地表示它们。因此,FlexiCache为每个头-层对执行独立的块分配、释放和映射。为了保持这种设计的效率,我们引入了三个关键优化:
脏追踪(Dirty Tracking)。采用每个头-层块映射方案将块表张量从(B, N)扩展到(B, L, H, N),其中B、L、H和N分别表示批处理大小、层数、头数和每个序列的最大块数。通常,这个块表在主机内存中构建和更新,然后在每个解码步骤复制到GPU内存。由此导致的L × H倍的表大小增加使得这种频繁的PCIe传输成为一个显著的性能瓶颈。为了减轻这种开销,FlexiCache跟踪自上次GPU同步以来CPU端块表的哪些区域已被修改,并且只传输这些“脏”段。由于脏区域通常是零散的,我们采用一个自定义的CUDA内核,类似于第3.3节中的方法,以高效地传输更新的段,同时避免主机端的收集开销。
重排序期间的物理块重用。在对稳定头进行周期性top-K重排序期间,简单地释放被驱逐页面的物理块并为晋升的页面分配新块会导致过度的块流失。频繁调用基于CPU的块分配器使得这个过程成本高昂。为了消除这种开销,FlexiCache转而回收物理块。具体来说,我们实现了一个融合的CUDA内核,它(i)识别被驱逐和晋升的逻辑块集合,(ii)将被驱逐集合的物理块重新分配给晋升集合,同时原地更新GPU端的块表,以及(iii)为随后的主机到GPU的KV缓存传输生成一个从主机KV块到GPU KV块的映射。
通过空块实现统一性。从GPU内存中驱逐非关键块会产生一个不规则的块表结构,这会通过使诸如追加新块之类的简单操作难以向量化而降低性能。此外,当布局不规则时,按层、头和逻辑索引在表中定位一个槽位变得更加复杂。为了解决这个问题,我们使用一个特殊的空块在块表中强制实现一个统一、密集的逻辑结构。当一个逻辑块被释放时,它的物理块被返回到池中,但它的表条目被保留并重定向到空块。由于这些条目被排除在注意力计算之外,它们永远不会被解引用,使得该方案既安全又高效。
A4 实验
4.1 实验设置
实现。我们在vLLM【【14,Efficient memory management for large language model serving with pagedattention,2023,SOSP'23】】的基础上实现了FlexiCache。我们修改了vLLM基于Triton的FlashDecoding【【7,Flash-decoding for long-context inference,2023,PyTorch Blog】】内核,使其仅对每个头的top-K页面执行解码注意力。我们用CUDA实现了专门的KV传输内核,并将它们编译为PyTorch扩展,以便与vLLM无缝集成。
测试平台。所有实验都在一台配备94GB HBM内存的NVIDIA H100 GPU上进行。主机系统配备了2个AMD EPYC 9554 64核处理器和1.1TB的DDR5系统内存。除非另有说明,我们为主机KV缓存分配了180GB的内存。GPU通过PCIe 5.0连接,峰值双向带宽为64 GB/s。我们的软件环境包括PyTorch 2.7.0、CUDA 12.8和cuDNN 9.5。
模型。我们在两个系列的模型上评估FlexiCache:Llama-3.1-8B-Instruct【【10,The llama 3 herd of models,2024,arXiv】】和Mistral-7B-Instruct-v0.2【【12,From clip to dino: Visual encoders shout in multi-modal large language models,2023,arXiv】】。
评估指标。对于模型准确性,我们报告FlexiCache与同一模型的密集注意力基准得分之比,以提供一个跨基准和模型的统一性能保持度量,尽管使用了稀疏注意力。对于系统性能,我们使用总token吞吐量(用于离线推理)和每输出token时间(TPOT)(用于在线服务)。虽然FlexiCache主要优化解码阶段,我们也在在线设置中分析了其对首token时间(TTFT)的影响。
基线。我们将FlexiCache与LServe【【28,LServe: Efficient long-sequence LLM serving with unified sparse attention,2025c,MLSys'25】】在稀疏注意力下的性能保持度进行比较,因为LServe是概念上最相关的系统。两个框架都支持paged attention,并减少了注意力计算的有效token预算以及驻留在GPU内存中的KV缓存的token数量。关键区别在于,LServe永久丢弃了较不重要token的KV缓存,而FlexiCache将其卸载到主机内存以备潜在重用。对于系统性能,我们使用vLLM基于Triton的Flash-Decoding后端作为我们的基线,这是FlexiCache实现的相同运行时,以通过隔离FlexiCache缓存管理的影响来确保公平比较。
4.2 准确性
长上下文。我们在LongBench【【2,LongBench: A bilingual, multitask benchmark for long context understanding,2024,ACL】】的16个不同任务中评估了FlexiCache相对于密集注意力基线的准确性保持率,如表2所示。FlexiCache与密集注意力分数的平均比率表明,FlexiCache在两种架构上都有效地保持了模型性能。在此实验中,top-K页预算为64页(相当于页面大小为16时的1024个token),64个最不稳定的头被分类为不稳定头,稳定头的重排序频率设置为16。为防止任务特定的泄露,我们对所有其他LongBench任务使用了GovReport的KV头稳定性配置文件,而对于GovReport本身,我们使用了从L-Eval任务【【1,L-eval: Instituting standardized evaluation for long context language models,2024,ACL】】派生的配置文件。强大的性能保持率进一步突显了跨任务稳定性配置的一致性,使得一次离线分析步骤即可在推理时重复使用。
长上下文与长生成。虽然LongBench涵盖了各种任务,但在评估稀疏注意力方面有一个关键限制:大多数任务生成短输出。除少数任务(如GovReport和MultiNews)外,平均生成长度低于50个token,有些任务平均甚至少于10个。因此,永久丢弃较不重要token的KV缓存的稀疏解码方法在LongBench上仍然可以表现良好,因为这些token在生成后期很少被需要。为了解决这个限制,我们进一步在L-Eval【【1,L-eval: Instituting standardized evaluation for long context language models,2024,ACL】】上评估FlexiCache,该基准测试包括几个具有长提示和长生成的任务(如附录中的表7所示)。与基于永久驱逐的稀疏解码不同,FlexiCache从不丢弃token:它将完整的KV缓存保留在主机内存中,并动态地将新变重要的页面提升到GPU内存,从而在长生成过程中保持模型性能。
表 2. LongBench【【2,LongBench: A bilingual, multitask benchmark for long context understanding,2024,ACL】】上的准确性保持率
与LServe的比较。表3比较了FlexiCache和LServe【【28,LServe: Efficient long-sequence LLM serving with unified sparse attention,2025c,MLSys'25】】在L-Eval任务上的表现。对于FlexiCache,我们逐步引入了重排序、稳定性感知重排序和更大的token预算。最终配置(token预算为2048,重排序频率为16,64个不稳定头)在两种模型上都实现了对密集注意力准确性的强力保持。另外两个模型(附录A.3中的Mistral-Small-24B-Instruct-2501和Qwen2.5-32B-Instruct)也显示出相同的趋势。对于LServe,我们遵循默认配置,流式头的token预算为384,检索头的token预算为4096。由于LServe是建立在QServe【【17,Qserve: W4a8kv4 quantization and system co-design for efficient llm serving,2024,arXiv】】之上,并且设计上采用了量化的KV缓存,因此不能直接与非量化的密集基线进行比较。为了隔离稀疏性的影响,我们将LServe的量化-稀疏配置与其仅量化基线进行比较,并报告添加稀疏性导致的相对下降。
表 3. L-Eval【【1,L-eval: Instituting standardized evaluation for long context language models,2024,ACL】】上的准确性保持率
4.3 端到端效率
我们在长上下文和长生成场景下,对FlexiCache在离线和在线服务设置中进行了评估。工作负载我们从L-Eval基准测试中随机抽样提示,而不是使用合成token,因为重排序期间晋升的KV缓存大小取决于工作负载特性。使用真实提示确保了晋升缓存行为的真实性。对于Llama-3.1-8B-Instruct,具有192个稳定头,fp16 KV缓存和2048个token预算,稳定头的top-K KV缓存为192 MB。晋升的KV大小在不同任务之间变化为44-67 MB(23-35%)(如附录表7所示)。这表明时间稳定性使晋升集相对较小但不可忽略,强调了随着token重要性的演变进行重排序的必要性。所有实验均使用64个不稳定头,稳定头重排序频率为16,token预算为1024或2048。
图 5. 端到端吞吐量。FlexiCache在Llama-3.1-8B和Mistral-7B上的token吞吐量始终优于vLLM,且增益随输出长度增长而增加。在输出长度为500时,请求吞吐量也观察到类似的改进。
离线吞吐量。图5显示了FlexiCache与vLLM基线的端到端吞吐量比较。对于Llama-3.1-8B和Mistral-7B,FlexiCache始终实现更高的总token吞吐量(提示+输出token/秒)。由于FlexiCache主要优化解码阶段,随着输出长度的增加和解码在总延迟中占主导地位,改进变得更加明显。工作负载包括从L-Eval随机抽样的500个请求,输入长度在10k-30k之间,输出长度在50-1500之间。1024-token预算通常获得更高的吞吐量,由于请求抢占的随机性,存在轻微波动。在不同输出长度上平均,FlexiCache将Llama的token吞吐量提高了1.38-1.55倍,Mistral提高了1.44-1.46倍。另一方面,如图5(c)所示,对于输出长度为500的情况,FlexiCache也比vLLM实现了高达1.46倍(1024)和1.41倍(2048)的端到端请求吞吐量加速。结果在Mistral-24B和Qwen-32B上也成立(见附录A.4)。
在线服务。图6比较了FlexiCache和vLLM基线在端到端在线服务性能方面的表现,指标为平均每输出token时间(TPOT)和首token时间(TTFT)。工作负载由500个从L-Eval随机抽样的请求组成,输入长度在10k-30k之间,输出长度在20-2000之间,请求根据泊松过程以不同的请求率到达。
图 6. 在线服务。FlexiCache在所有到达率下都降低了平均TPOT。通过降低每个请求的GPU内存使用量,它推迟了由高负载下队列积压引起的TTFT恶化,从而在保持较低尾部TTFT的同时,实现了更高的持续请求率。
图 7. 解码速度与批处理大小。每个输入有10k个token。随着批处理大小的增加,FlexiCache实现了更大的加速。
如图6(a)所示,FlexiCache在平均TPOT上始终优于基线,性能增益在更高的请求率下更为明显。在0.4 req/s时,FlexiCache-1024的平均TPOT为34.6,而基线为71.5,提升了2.1倍;FlexiCache-2048则达到44.9,提升了1.6倍。图6(b)显示,在0.35 req/s之前,各系统的TTFT相似,但基线在此点之后崩溃,因为GPU内存被完全占用,迫使新请求在队列中等待,导致TTFT急剧飙升。相比之下,FlexiCache通过减少解码期间每个请求的GPU内存使用,防止了内存饱和和队列积压,从而即使在更高负载下也能保持稳定的TTFT。如图6(c)所示,即使在中等负载(0.35 req/s)下,FlexiCache的TTFT尾部也更短,反映了排队请求更少。总体而言,FlexiCache比基线能承受更高的请求率,并在任何给定的请求率下提供更低的TPOT。
4.4 微基准测试
我们使用Llama3.1-8B-Instruct模型进行微基准测试实验,并将FlexiCache配置为重排序频率为16,除非另有说明。
图 8. 稳定性感知重排序的优势。每个输入有10k个token;随着批处理大小的增加,性能增益也增加。
仅解码速度提升。图7比较了密集解码和FlexiCache的稀疏解码(token预算为1024和2048)在不同批处理大小下,每个输入有10k提示token时的逐层解码延迟。对于密集解码,延迟反映了解码内核的执行时间,而对于FlexiCache,它包括了top-K选择和稀疏解码内核的时间。随着批处理大小的增加,FlexiCache实现了更大的加速(对于FlexiCache-1024,在批处理大小为40时高达4倍)。在较小的批处理大小下,内核启动开销占主导地位,而在较大的批处理大小下,不断增长的总KV缓存大小放大了稀疏解码的好处。
稳定性感知重排序的影响。图8显示,在批处理大小为40,每个输入包含10k token的情况下,稳定性感知重排序比朴素重排序快2.44倍。在朴素重排序中,所有头在每一步都进行重排序。相比之下,稳定性感知重排序在每一步只更新25%的不稳定头,并周期性地重排序稳定头,从而使平均MinMax缓存更小,并减少了延迟。
GPU内存节省。在GPU上,FlexiCache仅为25%的头存储完整的KV缓存,其余的头则存储top-K的KV缓存。图9显示了在不同序列长度下,单个请求在token预算为1024和2048时的相对GPU内存节省。随着序列长度的增长,内存节省渐近地接近75%,在20k token及以上时达到约70%。
图 9. GPU内存节省。在序列长度大于20k且token预算为1024时,实现了超过70%的内存节省。
A7 补充细节
5 讨论与未来工作
与其他服务优化的集成。FlexiCache优化了解码阶段,当专注于解码内核时,速度提升高达4倍(图7)。我们预计,当FlexiCache与优化工作流程其他部分的补充技术(如分离prefill和decode【【29,{DistServe}: Disaggregating prefill and decoding for goodput-optimized large language model serving,2024,OSDI'24】】、内核融合和推测性解码【【15,Fast inference from transformers via speculative decoding,2023,ICML】】)相结合时,将获得更大的端到端增益。
超越GPU-CPU层级。FlexiCache目前假设一个两级的GPU-CPU内存层级。这种设计可以扩展到涉及NVMe或分布式内存池的多层层级。在多节点环境中高效的KV缓存传输仍然是一个开放的挑战,而基于页面关键性和头稳定性的KV缓存管理有潜力进一步改进现有解决方案【【18,Cachegen: Kv cache compression and streaming for fast large language model serving,2024,SIGCOMM'24】;【5,Lmcache: An efficient KV cache layer for enterprise-scale LLM inference,2025,White paper】】。
联合内存管理与集群调度。在大型生产服务系统中,跨GPU集群的高效调度和负载均衡仍然是活跃的挑战。巨大的KV缓存大小使得快速的负载均衡和任务迁移变得困难【【20,Llumnix: Dynamic scheduling for large language model serving,2024,OSDI'24】】。头的时间稳定性为利用头级稳定性信号联合优化请求批处理、KV放置和预取时机开辟了道路。
6 相关工作
稀疏注意力。先前的工作表明,在每个解码步骤中,一小部分关键token主导了注意力输出。不同的启发式方法被用来识别这个子集。这些方法主要分为两类。一类是永久丢弃较不重要token的KV条目,如StreamingLLM【【24,Efficient streaming language models with attention sinks,2023,arXiv】】和SnapKV【【16,Snapkv: Llm knows what you are looking for before generation,2024,NeurIPS】】。这些方法减少了GPU内存、计算和I/O,但在长生成中存在准确性损失的风险,因为被丢弃的token后来可能变得重要。第二类是保留完整的KV缓存,但在每一步只关注最相关的token。Quest【【21,Quest: query-aware sparsity for efficient long-context llm inference,2024,ICML'24】】是这种方法的代表,它通过为每个KV缓存页增加每通道的min-max元数据,并根据当前查询对其进行评分,只获取Top-K页。这减少了计算和I/O,同时基本保持了准确性,但内存使用没有改变。一种混合方法是LServe【【28,LServe: Efficient long-sequence LLM serving with unified sparse attention,2025c,MLSys'25】】,它将大约一半的头转换为流式掩码(丢弃这些头的长程KV),并对其余的头应用动态的、查询感知的选择。这比纯粹的驱逐方法更好地平衡了效率和准确性,但由于一半的头永久驱逐KV,在长生成中仍然存在准确性损失的风险。
分层KV缓存管理。当一批请求的KV缓存超过GPU内存时,一个常见的解决方案是将其部分卸载到主机内存甚至磁盘,并在稍后重新加载。像FlexGen【【19,Flexgen: ´ high-throughput generative inference of large language models with a single gpu,2023,ICML'23】】和MoE-Lightning【【4,Moe-lightning: High-throughput moe inference on memory-constrained gpus,2025,ASPLOS'25】】这样的系统采用这种策略,通过精心将数据传输与计算流水线化——例如,在当前批次或层执行时预取下一批次或层的KV缓存或权重。然而,与FlexiCache不同,这些设计将同一层内的KV缓存视为统一的;没有一部分被认为比另一部分对解码更关键。它们的一些优化,如逐层加载,与FlexiCache是正交的,可以用来进一步增强它。LMCache【【5,Lmcache: An efficient KV cache layer for enterprise-scale LLM inference,2025,White paper】】跨内存层级卸载KV缓存,以在请求之间重用公共前缀KV,并支持prefill-decode分离。这减少了冗余的prefill开销并降低了TTFT,但它不针对卸载已经处于解码阶段的请求的KV缓存。NEO【【13,Neo: Saving gpu memory crisis with cpu offloading for online llm inference,2024b,arXiv】】在在线服务中使用CPU卸载KV缓存以提高相比于纯GPU系统的性能。然而,其卸载粒度是子批次级别而非头级别。
A5 结论
我们观察到LLM中的KV头在其选择关键KV页面的过程中,在解码步骤之间表现出不同程度的时间稳定性。基于这一洞察,我们将头分为稳定和不稳定两类,并提出了FlexiCache,一种稳定性感知的层级化KV缓存管理策略。FlexiCache根据KV页的关键性动态地将其分布在GPU和主机内存中,并通过周期性的、查询感知的机制来提升先前非关键的页面。这种设计减少了GPU的KV占用,而不会永久丢弃可能重要的token——这对于长上下文、长生成工作负载是至关重要的特性。
A6 附录
A.1 不稳定头的跨任务重叠
表4、5和6报告了三种模型——Mistral-7B-Instruct-v0.2、Mistral-Small-24B-Instruct-2501和Qwen2.5-32B-Instruct——不稳定头的成对跨任务重叠,平均重叠分别为0.83、0.76和0.81。这些结果表明,某些KV头在解码步骤中选择top-KV页时的任务无关不稳定性,在不同家族和大小的模型中都具有良好的泛化性,与正文中为Llama报告的观察结果一致。对于每个模型,我们将25%的头分类为不稳定头,分别对应Mistral-7B-Instruct-v0.2的64个头、Mistral-Small-24B-Instruct-2501的80个头和Qwen2.5-32B-Instruct的128个头。
表 4. 不稳定头的跨任务重叠。针对Mistral-7B-Instruct-v0.2。值为交集大小除以不稳定头集大小($|A_i \cap A_j|/64$)。GovReport来自LongBench【【2,LongBench: A bilingual, multitask benchmark for long context understanding,2024,ACL】】;其余来自L-Eval【【1,L-eval: Instituting standardized evaluation for long context language models,2024,ACL】】。
表 5. 不稳定头的跨任务重叠。针对Mistral-Small-24B-Instruct-2501。值为交集大小除以不稳定头集大小($|A_i \cap A_j|/80$)。GovReport来自LongBench【【2,LongBench: A bilingual, multitask benchmark for long context understanding,2024,ACL】】;其余来自L-Eval【【1,L-eval: Instituting standardized evaluation for long context language models,2024,ACL】】。
A.2 L-Eval统计数据
表7显示了L-Eval中不同任务的平均提示长度(以token为单位)和平均生成长度。它还报告了在稳定头周期性重排序期间,每个请求的平均晋升KV缓存大小(以MB为单位)。在此实验中,我们使用了2048的top-K大小和Llama-3.1-8B-Instruct模型,该模型有192个稳定头。稳定头的top-K KV缓存总大小为192 MB,观察到的传输大小范围为23%至35%。
表 6. 不稳定头的跨任务重叠。针对Qwen2.5-32B-Instruct。值为交集大小除以不稳定头集大小($|A_i \cap A_j|/128$)。GovReport来自LongBench【【2,LongBench: A bilingual, multitask benchmark for long context understanding,2024,ACL】】;其余来自L-Eval【【1,L-eval: Instituting standardized evaluation for long context language models,2024,ACL】】。
这个范围虽然由于时间稳定性而相对较小,但仍然不可忽略——这表明被卸载的KV缓存在后续的解码步骤中可能再次变得重要。平均提示长度和晋升的KV缓存大小之间,以及平均生成长度和晋升的KV缓存大小之间都存在正相关关系:当提示和生成长度都很大时,最初被认为不那么关键的KV缓存更有可能在以后变得重要。这进一步强调了像FlexiCache这样的系统的必要性,它能够在减少GPU内存使用的同时实现稀疏注意力,而无需永久丢弃KV缓存的任何部分。
表 7. L-Eval任务统计:平均提示和生成长度,以及晋升的KV大小(Llama-3.1-8B,token预算为2048,192个稳定头)
A.3 准确性保持率
表8报告了FlexiCache在Mistral-Small-24B-Instruct-2501和Qwen2.5-32B-Instruct上,跨不同L-Eval任务实现的准确性保持率。结果与正文中其他模型的报告趋势相似:随着周期性重排序的引入、不稳定头的包含以及更大的token预算,准确性保持率得到改善。最终的配置——64个不稳定头,稳定头每16步重排序一次,token预算为2048——显示出近乎无损的性能。
表 8. L-Eval结果
A.4 端到端吞吐量
图10显示了在FlexiCache上执行Mistral-Small-24B-Instruct-2501和Qwen2.5-32B-Instruct时,相较于vLLM基线的端到端吞吐量增益。工作负载由100个来自L-Eval的随机输入组成,每个输入的长度范围从10k到30k个token,输出长度为500。FlexiCache实现了1.37倍的改进,表明性能增益也扩展到了更大的模型。
图 10. 端到端吞吐量。FlexiCache在Mistral-Small-24B和Qwen2.5-32B模型上实现了显著的吞吐量提升。
💬 评论讨论
欢迎在这里分享您的想法和见解!