探秘Transformer系列之(25)--- KV Cache优化之处理长文本序列

探秘Transformer系列之(25)--- KV Cache优化之处理长文本序列

目录

  • [探秘Transformer系列之(25)--- KV Cache优化之处理长文本序列](#探秘Transformer系列之(25)--- KV Cache优化之处理长文本序列)
    • [0x00 概述](#0x00 概述)
    • [0x01 优化依据](#0x01 优化依据)
      • [1.1 稀疏性](#1.1 稀疏性)
      • [1.2 重要性](#1.2 重要性)
      • [1.3 小结](#1.3 小结)
    • [0x02 稀疏化](#0x02 稀疏化)
      • [1.1 分类](#1.1 分类)
      • [1.2 静态稀疏化](#1.2 静态稀疏化)
        • [1.2.1 Window Attention](#1.2.1 Window Attention)
        • [1.2.2 StreamingLLM](#1.2.2 StreamingLLM)
      • [1.3 动态稀疏化](#1.3 动态稀疏化)
      • [1.3 针对prefill的稀疏化](#1.3 针对prefill的稀疏化)
      • [1.4 针对层特点的稀疏化](#1.4 针对层特点的稀疏化)
        • [1.4.1 PyramidKV](#1.4.1 PyramidKV)
        • [1.4.2 PyramidInfer](#1.4.2 PyramidInfer)
        • [1.4.3 ZigZagKV](#1.4.3 ZigZagKV)
        • [1.4.4 Layer-Condensed KV Cache](#1.4.4 Layer-Condensed KV Cache)
      • [1.5 其它方案](#1.5 其它方案)
        • [1.5.1 针对头特点的稀疏化](#1.5.1 针对头特点的稀疏化)
        • [1.5.2 投机采样稀疏化](#1.5.2 投机采样稀疏化)
        • [1.5.3 聚类稀疏化](#1.5.3 聚类稀疏化)
    • [0x02 复用](#0x02 复用)
      • [2.1 KV Cache合并](#2.1 KV Cache合并)
        • [2.1.1 层间合并](#2.1.1 层间合并)
        • [2.3 层内合并](#2.3 层内合并)
      • [2.2 前缀复用](#2.2 前缀复用)
        • [2.2.1 需求](#2.2.1 需求)
        • [2.2.1 Prompt Cache](#2.2.1 Prompt Cache)
        • [2.2.2 Hydragen](#2.2.2 Hydragen)
        • [2.2.3 ChunkAttention](#2.2.3 ChunkAttention)
        • [2.2.4 AttentionStore](#2.2.4 AttentionStore)
        • [2.2.5 Radix Attention](#2.2.5 Radix Attention)
    • [0x03 依据检索的方案](#0x03 依据检索的方案)
      *
    • [0xFF 参考](#0xFF 参考)

0x00 概述

即便做了KV Cache,在长序列场景下,也只是减少了K和V本身结果计算量。但是和后面的部分依然计算量很大,并且K和V存储的压力也非常大。这其实严重阻碍了我们使用更长的序列做输入以及生成更长的序列。要解决这个问题,我们就需要对KV Cache进行优化。

KV Cache优化的核心在于减少内存消耗的原则。这一目标可以通过进一步压缩键值对(Key-Value Pairs)中的"K"(键)或"V"(值)来实现。这些压缩技术的选择直接影响模型的效率,特别是在内存使用和处理速度方面。压缩键值对就是让K和V变短,也就是说我们要舍弃或者压缩一些K和V。通过丢弃这些 token ,我们实际上将相应的注意力分数设置为零,并将注意力矩阵近似为一个更稀疏的矩阵。

减少序列长度是KV-Cache压缩领域研究最多的点。目前大致分为几个大方向:

  • 抛弃不重要的token(稀疏化)。通过attention score或者attention的系数特征来判别不重要的token,然后把这些token对应的KV都抛弃掉。这部分研究跟LLM的量化一样,都在activation outlier提出来之后取得了比较大的进展。虽然这种方法可以提高效率,但存在丢弃关键词元的风险,特别是对于需要深入理解远距离上下文的任务。
  • 复用token。通过复用来降低序列长度。比如复用system prompt等,这是一种工程优化,从算法层面上看,是一种无损优化,无需训练侧介入。
  • 基于检索(retrieval/取回)的方法。比如把KVCache offload到cpu,以页或者聚类的形式组织。q和每个页或者聚类进行相似度计算来决定使用哪些页或者聚类。每次只把得分最高的top-k个页或者聚类加载到显存计算。

注:全部文章列表在这里,估计最终在35篇左右,后续每发一篇文章,会修改此文章列表。
cnblogs 探秘Transformer系列之文章列表


0x01 优化依据

减少序列长度的理论根源主要是以下两点:

  • 大型语言模型(LLMs)注意力机制中的内在稀疏性。
  • token的重要性程度不同。

1.1 稀疏性

尽管transformers存在巨大的计算需求,但LLMs的注意力机制内部确实存在固有的稀疏性。同理,KV Cache中做为注意力机制的历史记录,也必然是很稀疏的。当然,稀疏性的程度可能因特定的下游任务而异。

下图(a) 展示了在使用CNN/DailyMail数据集进行摘要任务时,不同模型之间注意力稀疏性的多样化水平。这种变化在模型的各个层面都有体现,包括整个模型、单个层以及模型的特定部分。在(b)中,累积分布函数(CDF)描绘了注意力分数与总上下文占比之间的关系。可以看出来,即使KV缓存大小减少了50%,关键token的平均累积注意力得分仍然保持在一个相当高的水平,大约在90%到95%之间。进而也可以产生同样的token。

论文"H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models"的作者也发现,KV Cache其实是很稀疏的,只要5%就可以达到相同效果,即产生同样的token。

因此,我们可以在内存消耗和模型准确性之间做一定的妥协和权衡。

1.2 重要性

导致稀疏性的一个原因可能是因为token的重要性不同。

LLM 最核心的模块是 Attention:通过 Q 、 K 两个矩阵进行矩阵乘,计算\(O(N^2)\)大小的注意力得分矩阵,并使用注意力得分矩阵对 V 矩阵求加权平均值,得到注意力层的输出矩阵。其中注意力得分矩阵每一行表示 Q 中每个 Token 和 K 中每个 Token 的相关性得分,注意力层输出矩阵的每一行表示每个 Token 与其他 Token 的加权平均值。即然是加权平均,就说明每个 Token 的重要性并不相同,有的 Token 权重更大,而有的 Token 权重更小。

实际上,在文本生成过程中,只有一小部分 Token 获得了最多的注意力。这强调了特定关键token的重要性以及它们在理解上下文和促进文本生成中的核心作用。为了发现这些token,不同研究人员做了不同的工作。

  • 有些工作已经显示了对初始token的偏向性。这种偏向性源于在解码迭代过程中累积的注意力分数倾向于初始token的累积效应。即无论它们与语言建模任务的相关性如何,初始tokens都被分配了惊人的大量注意力得分(很大的原因并非是 token 决定的,而是其起始位置决定的),人们将tokens称为"attention sinks"。当文本长度超过额定长度时,头部的 token 就会被遗弃掉,这就会在 softmax 阶段产生很大的问题,进而影响后续的推理结果。
  • 也有一些工作发现稀疏的KV缓存可以有强时间局部性。稀疏性分布满足幂律分布,少部分固定的token一直主导着decoding。即在计算attention score时,只有一小部分token对模型的价值贡献最大,而且这些token在整段生成中都是重要的。人们称这些token为"Heavy-Hitters"(H2)。比如,论文"Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time"发现,单个 token 的注意力分数呈现强幂律分布,这意味着某些 token 对 Attention 机制的重要性远高于其他 token。作者据此提出了"重要性的持久性"(Persistence of Importance)假设:在生成过程中,只有具有显著重要性的 token 会对后续生成步骤产生持续影响。论文进一步指出,大多数 Transformer 层之间的注意力模式高度相似,重叠比率超过 90%。因此,通过识别"高权重分数 token",可以有效减少 KV Cache 的内存占用。此外,还可通过预先分配 KV Cache 的空间,仅存储必要的 token,从而将计算复杂度简化为常数级别。

当然,也有另一种可能是:因为模型经过训练,不需要关注整个序列(例如,Mistral AI的Mistral-7B),因此会对某些token不再进行关注或者很少关注。

1.3 小结

我们已经知道,注意力模块倾向于不成比例地持续将更多注意力集中在序列中的少数几个词元上。那么我们就可以不去存储模型几乎不关注或很少关注的词元的键和值。

通过丢弃 KVCache 中一些权重小的 Token,让它们不再参与注意力计算,我们将相应的注意力分数设为零,并用一个更稀疏的注意力矩阵来近似注意力矩阵。这样就可以在保证模型效果的前提下压缩 KVCache 长度,从而在一定量的显存下保存更多的 Token,提高长文本的推理效率。另外,不将序列中所有词元的键和值都存储起来的一个可能的原因是:我们可以在每次迭代中重新计算缺失的键和值,即用时间换空间。

因此,压缩 KVCache 长度的问题就可以转化为通过算法设计找出一个序列中相关性更高的 Token。

0x02 稀疏化

减少序列长度的主要手段就是稀疏化。KV Cache 稀疏化(KVCache Sparsification)是通过减少不必要的存储和计算来提高效率的一种优化方法。它的核心思想是选取需要关注的关键 token,并丢弃无关的 token,这是一种算法层面的优化,直接降低了计算,以减少 KV Cache 的体积和计算复杂度。

1.1 分类

KV Cache 稀疏化的各种方案本质上就是选取需要关注的 token 方法的不同。但由于舍弃了部分 Token 的注意力计算,所以理论上是有损的,可能需要训练侧保证模型效果。稀疏化属于目前比较主流的KVCache研究方向,不同论文对其的具体分类也不尽相同。两篇综述论文将稀疏化称为Token Dropping和KV Cache Selection。下图就是两篇综述中对稀疏化的不同称呼和细分。

  • Token丢弃是一种识别不重要token并将其丢弃的技术。然而,这些方法的关键挑战在于如何确定哪些token是不重要的。
  • KV Cache选择机制旨在降低KV缓存的内存利用率,最大限度地减少推理延迟,并提高大型语言模型的整体吞吐量。KV Cache选择机制以运行过程进行分类:静态KV Cache选择机制仅在预填充阶段执行token过滤,所选token在后续解码步骤中保持不变;动态KV Cache选择机制在解码阶段不断更新KV缓存,实现自适应缓存管理。

本文对稀疏化主要分为以下:

  • 静态稀疏化。在计算预测下一个 Token 时,只维护一个窗口大小的历史 Token 信息。然后通过调整注意力窗口来调整序列长度。静态稀疏化的窗口内的 Token 是固定的。
  • 动态稀疏化。通过调整注意力窗口来调整序列长度。动态稀疏化窗口内的 Token 是不固定的。
  • 针对prefill阶段的稀疏化。主要是优化prompt的序列长度,取出其中重要的token。
  • 针对层特点的稀疏化。LLM模型的注意力层有自己的特点,这也影响了KV Cache稀疏化策略。因此,在KV Cache管理中,即需要依据注意力权重的贡献进行区分处理或者累积处理,也需要动态调整每层参与注意力计算的键token数量。
  • 其它方案。比如结合投机采样进行的稀疏化;针对头特点的稀疏化等。

我们也用论文"A Survey on Large Language Model Acceleration based on KV Cache Management"来进行比对学习。该论文中与稀疏化对应的部分是KV Cache Selection、KV Cache Budget Allocation和KV Cache Merging。因为与笔者的分类方式不是正交的,所以我们在图上做出标注。

我们接下来通过案例进行学习。

1.2 静态稀疏化

静态稀疏化可以通过注意力机制来实现,即稀疏注意力机制或者线性注意力。稀疏注意力机制的核心思想就是在推理中选择合适的 Token 来进行相应的计算,这种方案在序列比较长时尤其有帮助,可以大幅降低 Attention 部分的 KV Cache 大小和计算量。线性注意力机制(如Linear Transformer、RWKV和Mamba等)通过将标准注意力机制替换为与序列长度线性相关的机制来减少内存需求。然而,这种方法可能会降低模型的表达能力,导致在需要复杂、长距离词元依赖关系的任务中性能下降。

比如,Mistral-7B使用被称为滑动窗口注意力(SWA)或局部注意力的注意力机制变体。局部注意力保证我们不需要关注整个序列,只通过关注最后(4096个)相邻的词元来构建词元表示,这样在KV缓存中永远不会存储超过窗口大小(例如4096)的张量对。再比如,

论文"Efficient Streaming Language Models with Attention Sinks"中分析了四种 Attention 实现,具体如下图所示。其中T 为待预测的第 T 个 Token,L 表示 pretrain 时的最大序列长度(一般为 2k、4k 等),由于主要针对的是长序列场景,因此 T 远大于 L。

第一种是原始注意力实现,后三种是稀疏化实现,这些稀疏化方法都是静态的,即 Token 之间的相关性是一种固定的范式,每个 Token 的相关的 Token 都是固定距离的。

  • Dense Attention是vanilla Attention 实现。每个 Token 都能看到该 Token 及之前的 Token,其计算复杂度为 \(O(T^2)\),KVCache 存储复杂度为\(O(T)\)。由于复杂度比较高,T 在预训练的时候会比较小。在推理的时候,当文本长度超过了预训练时的长度,模型效果就会大幅降级,所以表现出来的 PPL 值也比较大。
  • 窗口注意力(Window Attention)缓存最近的 n 个tokens的KV。这种方法相当于只在最近的tokens的KV状态上维护一个固定大小的滑动窗口。其中灰色的虚线框表示超过窗口长度后从 cache 中淘汰的 Token。计算复杂度为 O(T*L),计算复杂度低,cache 比较小,在推理中效率高。但是模型效果很差(PPL 很大),因为虽然但一旦开始tokens的键和值被删除,性能就会急剧下降。
  • 带重计算的滑动窗口(Sliding Window with Re-computation)。 这种 Attention 与 Window Attention 类似,区别是 Sliding Window 不缓存窗口的 Tokens,而是为每个新token重建来自 \(L\)个最近tokens的\(KV\)状态。这种方案的模型效果对比Window Attention 高,PPL 值远低于 Window Attention。主要原因是重计算把窗口中的 Tokens 作为初始 Tokens,这样既保留初始 Tokens 又保证计算只在一个窗口内,大大降低了KVCache 存储复杂度。由于重计算的存在,其精度可以保证,但是性能损失比较大。虽然它在长文本上表现良好,但由于在上下文重计算中的二次注意力,其的O(TL\^2 ) 复杂度使其相当缓慢,使得这种方法不适用于实际的流式应用。
  • StreamingLLM 保留了用于稳定注意力计算的attention sink(几个初始tokens),并结合了最近的tokens。它高效并且在扩展文本上提供稳定的性能。黄色的框表示初始的几个 Token,也就是 Attention Sink,其计算复杂度为 O(T*L),模型精度很不错,和带重计算的滑动窗口 Attention 相当。

我们接下来分析窗口注意力和StreamingLLM。

1.2.1 Window Attention

Window Attention采用滑动窗口机制来解决长文本推理挑战,其中落在窗口外的token被永久驱逐并变得无法访问。

在原始注意力中,每一个token,都要和它之前所有的token做Attention。而通过我们平时语言习惯中可以知道,一段话中每个字之间的相关性差别很大。对于当前token,一般来说越相近的字相关性越强,距离越远的token,能提供的信息量往往越低,因此似乎没有必要浪费资源和这些远距离的token做Attention。因此,滑动窗口(Sliding Window Attention)机制就是:每一个token只和包含其本身在内的前 L(L 表示窗口长度)个token做Attention,即只和当前token距离相近的token做注意力计算。这样就可以把cache的容控制在 L 内。因为每个 Token 只和邻近的 Token 做 Attention 计算,所以计算复杂度为 O(TL) ,KVCache 存储复杂度为O(L)。这种方法极大的降低了 KVCache 的存储开销,从线性复杂度降低到常数复杂度。

典型代表即 Longformer,具体如下图所示。

  • (a)展示了传统的"全注意力"机制,其中每个新生成的token都会关注序列中所有先前的token。

  • (b)描述了"滑动窗口注意力"。该方法维护一个固定大小的最近token滑动窗口,给定固定的窗口大小w,每个token都会关注其两侧 \\frac{1}{2}w \\(个token。滑动窗口允许我们拥有一个固定大小的缓存。一旦序列增长超出滑动窗口token数,我们可以在缓存中循环并开始覆盖。由于w可以设置的比较小,所以可以和输入序列长度n呈线性关系,其计算复杂性为\\)O(n×w),更加适合浅层捕获局部信息。而且,因为缓存中的位置无关紧要,所有与位置相关的信息都通过位置嵌入编码,所以很容易实现并且效果很好。然而,这种方法限制了模型从过去捕获全面语义信息的能力,导致文本生成质量较低且准确性下降。

    人们可能会疑惑,虽然距离越远的token涵盖的信息量可能越少,但不意味着它们对当前token一点用处都没有,在滑动窗口中直接杜绝了它们的参与,这真的合理吗?其实,Silding Window Attention并非完全不利用窗口外的token信息,而是随着模型层数的增加,间接性地利用起窗口外的tokens。这里可以类比为CNN技术中的感受野。如果堆叠多个层,从layer0开始,每往上走一层,对应token的感受野就往前拓宽W,虽然在每一层它"最远"只能看到前置序列中部分token,但是只要模型够深,它一定能够在某一层看到所有的前置tokens。该层窗口注意力可以访问所有输入位置,产生一个大的感受野,可以构建跨整个输入信息的表示。在具有l层的Transformer中,顶层的感受野大小为\(l \times w\)(假设对于所有层w是固定的)。根据应用程序的不同,可能为每个层使用不同的w 值更好,以在效率和模型表示能力之间取得平衡。具体如下图所示。

  • (c)展示了一种称为"膨胀窗口注意力"的变体,其准确性与窗口注意力类似,也存在类似的限制。其特点如下。

    • 为了进一步增加感受野而不增加计算量,滑动窗口可以"扩张"。这类似于dilated conv(空洞卷积或扩张卷积),其中窗口具有大小为扩张值d的间隔,这样可以将感受位置的间隔放大,让token在当前位置适当捕捉到得更远处的信息,关注更大的视野,考虑了更加全面的上下文信息。由于考虑了更加全面的上下文信息,空洞滑窗机制比普通的滑窗机制表现更佳。
    • 每个位置的Q关注的K还是W个,但是W个不是连续的而是跳跃的。定义跳跃的间隔是d, 那么Q关注的K的范围就是w * d。计算方面依然保持滑动窗口在计算性能方面的优点。
    • 假设对于所有层都是固定的d和w,那么感受野大小为l × d × w ,即使对于较小的d值,也可以触及成千上万个token。在多头注意力中,每个注意力头计算不同的注意力分数。每个头部具有不同的扩张配置的设置可以通过允许一些没有扩张的头部关注局部上下文,而其他具有扩张的头部专注于更长的上下文,从而提高性能。
  • (d)是Global+sliding window(融合全局信息的滑窗机制),滑动窗口和扩张注意力不足以学习特定任务的表示,需要从其它方面借鉴。因此引入了全局注意力对这两者进行补充。其特点如下:

    • 全局注意力就和普通的transformer是一样的,都是关注全部的K。依据具体任务的不同,在特定选择的位置引入关于全局的计算,使得模型在计算过程中能够接收到全局的信息。
    • 具有全局注意力的token会关注整个序列中的所有token,而序列中的所有token也会关注它。由于这些token的数量相对较小且独立于n,因此结合局部和全局注意力的复杂度仍然为O(n)。虽然指定全局注意力是与任务相关的,但这是一种向模型的注意力添加归纳偏差的简单方式,比起使用复杂架构将信息跨越较小的输入块组合的现有任务特定方法要简单得多。

1.2.2 StreamingLLM

StreamingLLM 利用"注意力沉积"效应,用早期Token 的KV结合近期上下文来优化长序列处理,实现了对无限长度输入的支持,同时生成无限长度的输出。

StreamingLLM 来自论文《Efficient Streaming Language Models with Attention Sinks》,其目标是让模型能够处理无限长度的输入(这里的"无限长度输入"与"无限长度上下文"有所不同------前者无需记住所有输入内容)。

Window Attention方案虽然压缩了KVCache的长度被压缩,但是模型效果却不好,主要原因是最初始的 Token 被丢弃了。这种 Token 的重要性其实非常高,丢弃会严重影响模型效果,精度下降比较大。而StreamingLLM恰恰弥补了这个缺憾。StreamingLLM发现注意力中的一个关键现象,即初始序列token中保留的键值对保持了关键的模型性能。这种注意力下沉效应通过早期位置的不对称注意力权重积累来体现,而与语义意义无关。该方法利用了这一特性,将注意力汇位置与最近的上下文相结合,以实现高效处理。

StreamingLLM策略在 Window Attention 的基础上,通过仅保留最初的位置词元(sink tokens)和最后相邻的词元(局部注意力)来构建了一个滑动窗口缓存。因此,StreamingLLM KV Cache 的长度是固定的,既有固定部分(通常1到4个 token ),也有滑动部分。这样既保留 Window Attention的特性,使得 KVCache 存储复杂度为O(L),计算复杂度为 O(TL) ,同时又保证了模型效果不因丢失初始 Tokens 而大幅下降,PPL 和 Sliding Window Attention with recomputation 相似。

Attention sink

StreamingLLM作者发现,window attention只要加上开头的几个token比如就4个,再操作一下位置编码,模型输出的PPL就不会变好,输出到20k长度都很平稳。论文将这种现象叫做Attention Sink(注意力池)。就好像文本的注意力被"沉溺"到了开头的位置。即初始词元聚集了大量注意力。而且intial tokens与生成token的绝对距离距离和语义信息都不重要,重要的是intial tokens是第一个或者前面几个token,就是这些token作为锚点,所以其权重特别大。

上图给出了lama-2-7B中256个句子的平均注意力逻辑的可视化,每个句子的长度为16。观察结果包括:

  • 前两层(第0层和第1层)的注意力图呈现出"局部"模式,最近的token受到了更多的关注。
  • 除了前两层之外,之后的注意力主要集中在first token。

如果删除这些初始token的KV,将导致注意力计算中SoftMax函数的相当一部分分母也被删除。这种变化导致注意力得分的分布发生了显著变化,偏离了正常推理环境中的预期。

为何会出现这种现象?Evan Miller 在"Attention Is Off By One" 做了精彩的解读。

The problem with using softmax is that it forces each attention head to make an annotation, even if it has no information to add to the output vector. Using softmax to choose among discrete alternatives is great; using it for optional annotation (i.e. as input into addition) is, like, not cool, man. The problem here is exacerbated with multi-head attention, as a specialized head is more likely to want to "pass" than a general-purpose one. These attention heads are needlessly noisy, a deafening democracy where abstention is disallowed.

在Attention机制中,Softmax的输出代表了key/query的匹配程度的概率,如果softmax在某个位置的值非常大,那么在反向传播时,这个位置的权重就会被大幅度地更新。然而,有时候attention机制并不能确定哪个位置更值得关注,但是Softmax又需要所有位置的值的总和为1,因此模型必须"表态"给某些位置较大的权重,因此,模型倾向于将不必要的注意力值转嫁给特定的token,这些token就是initial tokens。

因此,Evan Miller改进了一下Softmax,提出了"SoftMax-off-by-One":把softmax的分母加了个1,这样所有位置值可以加和不为1,Attention就有了可以不对任何位置"表态"的权利。

思路

根据attention sinks特性,论文作者参考之前多轮对话场景的解决方法,给出StreamingLLM的解决方案,即在window attention基础上,新引入了initial tokens的KV。

  • 锚点加窗口 。StreamingLLM中的KV cache可以分为两部分,既有一个固定部分(通常为1到4个词元),又有一个滑动部分。
    • 固定部分。保留整个序列的初始 Tokens(attention sink部分),图中画的是4个initial tokens;其作用是识别并保存模型固有的attention sink,锚定其推理的初始token来进行注意力计算并稳定模型性能。
    • 滑动部分。滑动窗口的KV(Rolling KV部分)cache,保留了最近的3个tokens,窗口值固定为3。其目的是:想让大模型无限输出,达到"Streaming"效果:不用KV cache会计算太慢,但是用了KV cache,显存占用随长度增长太多。因此需要丢弃一些KV cache。
  • 计算。每个 token 只和窗口内的 Tokens 以及序列的初始 Tokens 进行 Attention 计算。而且训练期间,用"SoftMax-off-by-One"替代常规softmax,解决attention sink问题。

在实现中,会用"Rolling Buffer Cache"技术来丢弃KV cache中的某些行,同时也会控制位置编码。下图给出了Rolling Buffer Cache的运作流程,当我们使用滑动窗口后,KV Cache就不需要保存所有tokens的KV信息了,你可以将其视为一个固定容量(W)的cache,随着token index增加,我们来"滚动更新" KV Cache。

问题

StreamingLLM虽然记住了Attention sink以及最近的窗口里面的token对应的KV cache,然而仅仅在注意力机制中依赖这k个token并不能提供必要的准确性。因为目前是按照位置把中间的token都丢掉了,万一中间的token就是重要的怎么办。这就是静态 Token 稀疏化的问题:Token 候选集过于固定。这种 Token 候选集的设计源于作者通过观察某些数据集中 Token 之间的相关性发现的规律。这种规律不能保证具备普适性,会导致关键token注意力中近期上下文的丢失和窗口注意力中关键上下文的缺失,会损害模型在长序列任务上的性能。

1.3 动态稀疏化

动态 Token 稀疏化本质是为当前处理 Token 维护一个相关性高的历史 Token 集合,但与静态 Token 稀疏化不同,这个集合的构造不再由 Token 距离或者固定 Token 决定,而是设计一种算法去筛选历史的 Token。静态 Token 稀疏化是动态 Token 稀疏化的子集,所以理论上动态 Token 稀疏化的效果会比静态 Token 稀疏化更好。

然而,在推理过程中,尤其是当输入序列包含未知或未见过的token时,动态确定哪些token作为关键token是一个相当大的挑战,针对这挑战主要有两种方法:

  • 评分函数。利用评分函数识别关键token。比如我们为每个token引入了一个评分函数fθ,以从总共n个token中识别出k个关键token。
  • 驱逐策略。保留系统认为重要的KV,舍弃那些不重要的。其中重要性的判断非常重要,大部分的系统都是对于每一层按照attention score的大小,保留top-B的结果。

其实这两者是一枚硬币的两面:通过动态的评价策略来判断需要保留和废弃的KV值,通过运行时键/值驱逐来减少KV缓存大小。

需要注意的是,大部分工作都假设了注意力模式在不同迭代之间的持久性,即,如果某个token在当前迭代中被认为是不重要的(即具有较低的注意力权重),那么在生成未来token时,它也很可能保持不重要。基于这一假设,当KV缓存大小超过其预算时,他们会在每次迭代中从KV缓存中驱逐具有低注意力权重的token。被驱逐token的键和值在从内存中删除后,将永久排除在后续迭代之外。

我们通过几个案例来学习下。

1.3.1 H2O

论文"H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models"提出的H2O是比较经典的方案。H2O观察到,注意力计算主要由一组被称为重拳("pivotal tokens"或"heavy hitters")的高影响力token来驱动。因此,H2O将缓存优化重新表述为动态子模型问题,利用累积注意力得分来指导token保留决策。

动机

稀疏的KV缓存并不一定具有100%的时间局部性。比如,LLM在decode第6个token时不一定需要第1至第5个token,而是有可能需要其中的任意数量的token。想象一种最极端的情况,decode第6个token时,需要第1、2、4个token,decode第7个token时需要第3、5、6个token。因为在同一时刻虽然不会使用所有token,但是每个token都有可能在任意时刻被使用,所以这种情况下必须在内存中保存所有历史token的KV缓存。论文将这类情况称为Dynamic Sparsity,见下图左一。

一种直观的想法是,也许稀疏的KV缓存可以有类似密集KV缓存那样的强时间局部性:在两个decoding步骤之间仅有一小部分KV缓存不同,并且在某个decoding步骤被稀疏掉的token永远不会在接下来再被使用,这样就实现了与密集KV缓存类似的时间局部性同时减少了KV缓存内存占用。论文将这类情况称为本地稀疏(Static Sparsity (Local))和等步长稀疏(Static Sparsity (Strided)),见下图左二左三。

H2O认为,本地稀疏、等步长稀疏等强时间局部性的稀疏都太naïve了,光看最近的K个元素或者等步长最近的K个元素无法代表所有的历史信息。在满足强时间局部性的稀疏策略中,应该存在更优的策略。H2O利用累积的token重要性在decoding过程中找到了这一策略,见图左四。

H2O作者认为,最为理想的drop token方式是:预知未来的token对当前token的注意力分数,然后drop掉那些分数低的token。因为未来的token还不知道呢,所以只能预测。

因为attention分数比较大的token对计算结果的影响比较大,所以作者拿前面的attention分数相加的结果来预测之后的attention分数。作者在论文后续附录中证明了在submodular情形下,H2O假设的贪心地构造KV缓存的策略是near-optimal的。即随着规模的增长存在边际效应。因此,H2O的设计问题已经收敛到了一个很小的问题:dynamic submodular problem(动态次模问题)。即如何设计贪心算法,满足每次驱逐的token都是submodular最优的。即每次丢掉的token对生成的影响最小。

方案

H2O的设计目标就是确定一个KV缓存维护/驱逐策略。H2O设置最大缓存 token 数量(预算),每次达到缓存预算时,该策略可以通过贪心法去计算token的打分,每次丢弃累积注意力分数最低的 token ,因此只保留在迭代中始终实现高注意力分数的 token。这样使基于稀疏KV缓存的decoding效果与基于密集KV缓存的decoding效果相近。

H2O作者作者发现:在给定步骤中具有影响力的词元("pivotal tokens"或"heavy hitters")在未来步骤中仍然具有影响力。换句话说,我们可以确信被丢弃的低影响力词元在未来生成步骤中仍然可以相对被忽略,因此可以安全丢弃。基于此观察,作者提出了一种名为"Heavy-Hitters Oracle"(H2O)的贪心的历史 Token驱逐策略:对于某个token A ,A以及A之前的token都会对A有注意力,把这些注意力加起来,可以得到A的"重要性分数"。将每个token的重要性分数都计算出来,然后丢弃分数最低的那些token即可。H2O选择了注意力分数作为评分函数,记作\(f_θ(acc; attn)\)。这样就可以识别出在提示和token生成阶段持续获得较高注意力分数的token作为最关键或关键的token。此处其实有提高空间,因为是使用attention score的绝对值来反映token的影响力,求和隐含了"越久的token越重要"的假设,越久的token求和项越多,最终的值有可能越大。

算法具体流程如下图所示:

  • 初始化相关性历史 Token 集合 ,设置集合最大容量。

  • 迭代更新历史 Token 集合。

    • 当历史 Token 集合大小等于最大容量时,将当前 Token 加入历史 Token 集合中;
    • 当历史 Token 集合大小大于等于最大容量时,计算当前 Token 和历史 Token 集合中所有 Token 的相关性得分(实现起来是对于注意力矩阵按列求和再取topk),将最低分的 Token 从集合中剔除,将当前 Token 加入到集合中。即每生成一步会将当前 Token 相关性最低的历史 Token 永久剔除,被剔除的历史 Token 将不会参与后续任何一个 Token 的 Attention 计算。用自然语言描述,就是KV缓存中被稀疏掉的部分对应的attention score设置为0,其它部分正常计算。

    这样可以保证前后两个KV缓存的大小相等,最多有一个位置不同。即:要么从前一个时刻的KV缓存中删掉一个,然后加上本次生成的token的KV缓存;要么保留前一个时刻的KV缓存,不加本次生成的token的KV缓存。这样的KV缓存始终满足强时间局部性。

图上的反斜杠"\"是一种集合运算符号:相对差集。"U \ A"表示在集合U中,但不在集合A中的所有元素,比如相对差集{1,2,3} \ {2,3,4} 为{1} ,而相对差集{2,3,4} \ {1,2,3} 为{4} 。

与H2O类似的方案是Scissorhands。H2O算法一次仅丢弃一个词元,而Scissorhands则根据目标压缩比例(例如,减少30%的KV缓存)丢弃尽可能多的词元。Scissorhands简单地保留最近的词元和在历史窗口内具有最高注意力分数的词元。而H2O则丢弃累计注意力分数最低的词元,只保留那些在迭代过程中始终具有高注意力分数的词元。

优化点

现有通过驱逐token来减少KV Cache大小的研究本质上存在一些挑战。类似H2O的方法假设注意力模式在不同迭代之间不会改变,换句话说,H2O其实是用上一层KV缓存的稀疏性来指导下一层的KV缓存的稀疏性。但实际上情况并非如此,因为注意力模式在迭代之间具有动态性,在当前迭代中被认为不重要的token在后续迭代中可能会变得重要。在前面生成过程中被剔除的 Token 可能与当前的 Token 有着较高的相关性,但因为 Token 在前面的生成过程中已经被剔除,导致无法把 Token 召回和当前 Token 完成 Attention 计算。

因此,H2O的评估窗口较窄,导致其在KV缓存预算范围内表现出较高的效率,但随着序列长度超出KV缓存预算,它开始难以应对注意力模式的动态性,未能很好捕捉这些动态变化。

另外,H2O方法将保留的键/值token数量设置为输入序列长度的一个固定百分比。在这种情况下,无论生成了多少个token,KV Cache预算始终保持恒定。然而,这种固定KV Cache预算存在一些局限性。每一层所需的键/值token数量是不同的,而每个查询token也需要不同数量的键/值token来有效地表示基线模型的注意力模式。此外,即使是相邻的查询token,所需的键token数量也会发生显著变化。固定的KV Cache预算忽视了查询token之间的这种差异性,如果保留的键/值token数量不足以准确表示基线模型的注意力权重,会导致KV Cache管理效率低下。因此,需要根据每个查询token的需求,动态调整所加载和计算的键/值token数量,同时考虑不同层和查询token之间的差异性,以实现更高效的KV Cache管理。

1.3.2 keyformer

在生成推理中,大约90%的注意力权重集中在特定的一小部分token子集上,这些token被称为"关键"token。这些token对于大型语言模型(LLMs)理解上下文至关重要,但可能不在窗口注意力的滑动窗口内。Keyformer引入了一种混合注意力方法,如下图(d)所示,该方法在生成下一个token时结合了最近的token和前面的关键token。此外,Keyformer揭示了token移除会扭曲底层softmax概率分布。考虑到softmax分布在token显著性评估中的关键作用,Keyformer 结合了正则化技术来减轻这些分布扰动。

动机

下图比较了完整注意力下的得分函数分布与KV缓存减少后的分布情况。如方程2所示,当KV缓存减少时,分数较低的token会被丢弃,因为isoftmax函数的性质,这会破坏得分函数的分布,进而导致上下文信息的丢失、准确性的降低以及可能生成的文本质量下降。因此需要增强模型提取近期关键token的能力。

方案
评分函数

Keyformer通过对未归一化的对数几率(logits)应用正则化来调整由于丢弃tokens而导致的评分函数变化,从而识别关键tokens。具体而言,Keyformer提出了一个新颖的评分函数,记作fθ(Keyformer),以解决基于累积注意力的评分函数(fθ(acc attn))的局限性。这个新的评分函数将Gumbel噪声分布整合到未归一化的logits中。然而,它在形成基础概率分布时未能考虑被丢弃的token。为了纠正这一点,Keyformer引入了一个温度参数,记作τ。

具体剖析如下:

  • Keyformer使用对数几率正则化技术。通过引入额外的分布(ζ)来正则化减少的对数几率,使模型保持鲁棒性和适应性。即使在推理过程中存在未知上下文的情况下,它也能帮助识别关键token。Keyformer将这种噪声添加到从QKT派生的未归一化对数几率中。此外,添加的分布类型对最终的概率分布产生显著影响。对应下图标号1。在这里,yi是经过调整的对数几率(logits),xi是未归一化的对数几率,ζi是用于正则化的额外分布。
  • 计算概率时,用温度参数 τ 考虑被丢弃的token。"温度"参数(τ)在调节概率分布的平滑性方面起着关键作用。这个参数控制概率中的随机程度。在token从KV缓存中被移除时,τ 尤为重要,因为它们不能在没有重新计算其键和值的情况下被重新引入。较高的τ值(τ → ∞)产生均匀概率,将所有token赋予相同的分数。相反,较低的τ值(τ → 0)产生更尖锐的分布,根据未归一化的logits优先考虑特定的token。另外,随着更多的token被丢弃,我们需要一个更均匀或随机的概率分布。因此通过逐步增加τ和∆τ来系统地增加我们评分函数fθ中的随机性。
  • 选择的分布受到Gumbel分布的启发。Gumbel分布特别适合我们的关键token识别任务,因为它描述了样本集中最大值的分布,并且偏向于初始token。这使得它成为长序列中建模关键token的理想选择。标号3展示了应用于未归一化logits的标准Gumbel概率密度函数(pdf),而标号4显示了添加Gumbel噪声后logits的pdf。
算法

尽管当前token的相关性在识别关键token方面很重要,但模型在大多数生成的token中应该保持一致的行为。为了基于这种一致的行为来识别关键token,Keyformer在提示阶段和token生成阶段都应用累积评分函数(fθ)。如果没有累积,token只能依赖于当前token与先前token的相关性。

在提示处理阶段,Keyformer计算提示长度\(S_n\)内所有n个token的键和值,以预测第一个token。给定KV缓存预算,Keyformer保留一个最近的w个token的窗口,从n-w个token的窗口中丢弃n-k个token,并且基于Keyformer评分函数来识别出k-w个token。关键token(k-w)和最近token(w)的组合形成了缩减的KV缓存。

在token生成阶段,Keyformer使用缩减的KV缓存进行操作。第一个生成的token仅关注KV缓存中的k个token,如下图图中的解码步骤2所示。最近窗口w向右移动一个token,而评分函数fθ则累积了来自前一个解码步骤的评分函数。在token生成阶段的每个解码步骤中,模型会从大小为k+1-w的窗口中识别出k-w个关键token。因此,最近窗口中添加了一个token并删除了一个token。

此外,温度参数τ通过增加∆τ来调整评分函数概率分布中移除token的数量。

Keyformer的详细算法参见下图。

优化点

与H2O类似的技术如下:

  • SubGen是一种为 KV cache开发的高效压缩技术。经验证据表明,attention模块中的关键嵌入存在显著的聚类趋势。基于这一关键见解,SubGen设计了一种具有亚线性复杂度的新颖缓存方法,对关键token采用在线聚类并对值进行在线 (2)采样。

  • LESS(Low-rank Embedding Sidekick with Sparse policy)会学习原始注意力输出和稀疏策略近似的注意力输出之间的残差,通过将稀疏策略丢弃的信息累积到恒定大小的低等级缓存或状态中来实现此目的,从而允许查询仍然访问信息以恢复注意力图中先前省略的区域。

上述基于永久驱逐的方法面临两个重大限制。

  • 首先,对token的不可逆驱逐可能会损害模型在长序列任务上的性能,特别是在needle-in-a-haystack场景中,这些方法被证明难以适应多回合对话环境。
  • 其次,在解码阶段选择KV缓存会引入计算开销,对解码延迟产生不利影响,并影响端到端加速。

为了应对这些挑战,几项研究侧重于开发解码阶段KV缓存选择策略,而无需永久驱逐。比如:

  • InfLLM:块级存储全KV Cache,GPU保留关键Token,次要部分卸载至CPU。
  • SqueezedAttention:离线K均值聚类生成语义键簇,推理时加载相关键。
  • SparQ:基于查询向量显著元素选择性检索缓存键的隐藏维度,近似计算注意力。

这些方法通常采用多层缓存系统(例如CPU-GPU分层缓存),并利用先进的数据结构和系统级增强功能来优化检索效率,从而在减少GPU KV缓存占用的情况下实现高效推理。为了加速关键token的检索,这几项研究提出了基于索引的方法,以块或集群粒度组织和访问KV缓存,采用多级缓存(如CPU-GPU层次存储)避免永久移除,优化检索效率,从而实现高效的查询和提取操作。

1.3.3 小结

论文"A Survey on Large Language Model Acceleration based on KV Cache Management"给出了静态稀疏化和动态稀疏化方法的一些总结。

1.3 针对prefill的稀疏化

现有的KV缓存压缩方法主要关注生成阶段的优化,忽略了输入阶段KV缓存的压缩。而实践中,大模型应用如对话(特别是RAG)和文档处理的特点是:输入很长,而输出相对较短。输入可能就已经把显存撑爆。或者即便显存还可以,每一步计算也会因为需要交互的token太多而非常慢。因此输入阶段的KV缓存是内存和计算瓶颈所在。比如我们输入的prompt有16K,那么存储这16K文本对应的所有KV,无论对于显存还是计算都有极大的压力。因此我们可以在prefill阶段对KV Cache进行处理,drop一些不重要的token,即压缩prompt生成的KV。而解码阶段就照常进行,这样不需要在解码阶段进行额外的计算,又可以加速。

1.3.1 FastGen

FastGen 通过识别五种注意力模式(局部、标点导向、稀疏分布等),再在预填充期间识别重要的token。针对每个token,FastGen 对不同注意力的特点进行针对性的分析(profiling),从而来辨别注意模块的内在结构。然后,相应地在已识别的不同结构上进行了不同的压缩策略,以自适应的方式构建KV缓存。FastGen不设置缓存预算,而是设置注意力矩阵的最大近似误差,因此专注于模型精度保持。

洞见

以前许多压缩KV缓存的研究都没有利用LLM中复杂的注意力结构。而LLM中的注意力结构有如下特点:

  • 注意模块中有丰富的结构。
  • 不同的注意头通常具有不同的功能,遵循不同模式。
  • 并非所有注意模块都需要关注所有表征。

上述特点表明需要根据每个单独的注意头定制压缩策略。下图左侧是四种常见的注意力结构,右侧是同一层的三个注意头的注意图。其中,特殊词元(绿色)+ 标点词元(橙色)+ 局部注意力(蓝色)。灰色为其它词元。

算法

除了传统的全KV缓存策略外,FastGen考虑了四种基本的KV缓存压缩策略。四种KV缓存压缩策略是:

  • 特殊token(Special Tokens)。在KV缓存中只保留特殊的token,如句子开头token <s>、指令token[INST]等。

  • 标点符号(Punctuation)。在KV缓存中只保留标点符号,如"."、":"、"?"。

  • 局部性(Locality)。此策略驱逐长期上下文,保留最后相邻词元。有些注意力模块主要关注局部上下文,也就是我们常说的 Locality。在此策略中,一旦上下文token和当前token之间的相对距离超过阈值,上下文token的KV缓存将被逐出。阈值由预定义的局部上下文长度预算与输入序列长度的比率所确定。

  • 频率(Frequency)。监控每个token的累计注意力分数,然后将这些分数视为token使用频率,并只将最频繁的token保存在KV缓存中。

FastGen可以分为两步:

  • 首先,在预填充阶段结束时,对模型的注意力层进行剖析(profiling),识别各种注意头的行为和结构模式,在这个profiling的指导下,自适应地为每个头选择最合适的满足误差目标的压缩策略,构建KV Cache。与其他方法一样,FastGen假设已识别的注意力模式将在未来的生成步骤中保持不变。如果误差目标过于严格,无法达到,则FastGen会退回到常规的KV缓存。
  • 然后,在每个生成步骤,不是不加区分地为每个新生成token添加新KV向量,而是根据每个token选择的压缩策略来管理其KV缓存。比如驱逐注意头上强调局部上下文的长程文本,丢弃以特殊token为中心的注意头上的非特殊token,并且仅对广泛关注所有token的注意头使用标准KV缓存。

1.3.2 SnapKV

SnapKV 的核心思想是根据用户问题识别重要的token。LLMs在生成过程中对输入tokens的关注模式具有很强的一致性和稳定性,而这种模式可以在输入序列末尾的一个小窗口内就观察到。因此,SnapKV只保留那些被注意力heads持续关注的输入tokens对应的KV,从而大幅压缩KV缓存的尺寸。同时,SnapKV还引入了池化等聚类技术,既能捕捉局部的完整信息,又能过滤掉大部分无关tokens。

洞见

论文作者做了实验,把输入按照128个token为单位来划分window,取靠近尾部的20个window来做计算。对于prefix里面的任意一个token,假设叫做\(t_i\)。我们取20个window里面的某一个window,假设叫\(w_i\)。\(w_i\)的每一个元素和\(t_i\)都会有交互,因为\(t_i\)在\(w_i\)的元素前面。那么,\(w_i\)的每个元素和\(t_i\)都会计算出一个attention score。对于窗口\(w_i\),\(t_i\)会有128个attention score。取这128个attention score的均值作为\(t_i\)的重要性值。这样,对于窗口\(w_i\),prefix的每个token都可以计算出重要性值,然后根据平均score选出生成阶段prefix中最重要的topK个token。

然后,作者比较窗口选择出来的重要token和生成阶段选择出来的重要token的重叠率。通过比较20个窗口选择的重要token和生成阶段实际选择出来的重要token,作者发现最后一个窗口选择出来的和生成阶段选择的是高度重合的。因此这就说明:输入里面哪些重要可以使用输入prompt本身就可以决定。即论文所说的:

  • Pattern can be identified before generation.
  • Pattern is consistent during generation.

另外,作者还观察到几个现象。

  • 无论生成上下文的长度如何,提示中的特定 Token 在生成过程中始终展现出更高的注意力权重,而且比较稳定。
  • 提示中问题的位置并不会显著改变观察到的注意力模式的一致性。也就是,无论问题的位置如何,都可以获得相关特征的注意力。
  • 注意力模式高度依赖上下文,表明这些模式与用户提出的具体指令有很强的关联。

因此,一个上下文感知的 KV Cache 压缩方法可能会带来更好的性能。

方案

SnapKV简化了FastGen的方法,只专注于根据token的重要性得分检索token。在所有提示token中,只有一部分携带了响应生成的关键信息,这些token在生成阶段保持了其重要性。SnapKV采用末端定位的观察窗口(end-positioned observation window)来检测这些重要的上下文token。然后,将它们对应的键值对与观察窗口中的token连接起来

SnapKV在 Prefill 阶段不是保留所有输入 Token 的 KV Cache,而是采用稀疏化的方式。

  • 在prefill阶段,针对每个 Attention Head 将 Prompt 分为 Prefix 和 Window 两部分,即下图中橙色和绿色部分,可设为representative token 和 window token。对应下图标号1。可以从位于 Prompt 末尾的 "Observation" 窗口获得每个 Attention Head 在生成过程中始终关注特定的注意力特征。

  • 通过 Window 中 Token 与 Prefix 中 Token 的 Attention Score 来选择稀疏化的 Token。即假设window token之前的某个token为A,window token会对A有一个注意力分数,将所有window token对A的注意力分数加起来,就是A的"重要性分数"。重要性分数最高的那些token就是representative token。对应下图标号2。

  • 依据注意力分数选出topk的token,此外,在选取topk之前还要对attention矩阵进行一维的pooling操作,文中解释这样可以帮助LLM保留更完整的信息,因为只使用窗口选择的内容,会造成生成的信息不全。对token附近token做pooling操作可以补齐上下文,类似高斯模糊。对应下图标号3。

  • 最后,将选出token的 KV Cache 和 最后一个窗口中 所有Token 的 KV Cache 一起作为 Prompt 的 KV Cache。对应下 图标号3。

SnapKV只在 Prefill 阶段筛选出关键的 KV Cache,然后在 Decoding 阶段使用这些稀疏化的 KV Cache,以加速长上下文场景的生成速度。此外,Decoding 阶段也不会再更新 Prompt 的 KV Cache。

具体算法如下。

1.4 针对层特点的稀疏化

LLM的分层架构导致跨层的不同信息提取模式,每层的KV缓存对模型性能的贡献不同。

LLM模型的注意力层有自己的特点,这也影响了KV Cache稀疏化策略。比如,为了估计需要保留多少KV Cache中的键/值,InfiGen作者对每个查询token的注意力权重按降序排序,并累计键token的权重,直到累计权重达到0.9。下图展示了在不同层(Layer 0 和 Layer 18)中,为了达到总注意力权重的0.9所需的键token数量分布情况。直方图的横轴表示键token数量,纵轴表示查询token的数量。

  • 第0层的注意力模式更加分散,不同查询token所需的键token数量差异较大,需要参考更多的键token才能捕获有效的注意力。
  • 第18层绝大部分查询token仅需16至几十个键token即可达到累计权重0.9,只有极少数token需要更多的键token,说明高层的注意力模式更集中,依赖较少的键token即可获得有效注意力。

从这些观察中我们可以知道,每个层对于总注意力其实都有着不同程度的贡献,所以不能直接驱逐在外。而且,不同层的注意力模式差异显著。因此,跨层的统一KV缓存压缩可能不是最优的。

在KV Cache管理中,即需要针对LLM层间信息提取异质性,依据注意力权重的贡献进行区分处理或者累积处理,也需要动态调整每层参与注意力计算的键token数量。通过根据每个组件对预测准确性的重要性智能分配内存资源,从而优化内存利用率,同时最大限度地减少准确性下降。

注:论文"A Survey on Large Language Model Acceleration based on KV Cache Management"是将基于层特点的稀疏化和基于头特点的稀疏化都归结到"KV Cache Budget Allocation"类别中,KV缓存预算分配通过根据每个组件对预测准确性的重要性来智能分配内存资源来解决这一挑战。比如将更多的KV缓存分配给信息更加分散的较低层(前几层),同时减少较高层信息比较集中的KV缓存等。论文给出了具体总结如下图所示。

我们具体看看一些案例。

1.4.1 PyramidKV

论文"PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling"通过研究不同layer间的注意力机制,探索了跨层共享kvcache的效果,发现较低层在输入序列中表现出均匀的注意力分布,而上层则表现出对特定token的集中注意力。因此,PyramidKV采用了金字塔形的内存分配策略,将更多的KV缓存分配给信息更加分散的较低层(前几层),每个KV包含的信息较少,同时减少较高层的KV缓存,使得高层中信息变得集中在较少的关键token中,同时在每一层选择具有高关注值的token。此外,在解码阶段,PyramidInfer通过由注意力值驱动的频繁更新动态维护一组重要token。

洞察

作者研究了Llama模型进行多文档问答的逐层注意力图,发现了注意力层中存在金字塔形信息汇聚模式(Pyramidal Information Funneling):

  • 在模型的低层(例如第0层)中,注意力得分呈现近似均匀分布,信息比较分散。这表明模型在较低层时从所有可用内容中全局聚合信息,而不会优先关注特定的段落。
  • 当编码信息进行到中间层(6-18)时,逐渐转变为聚焦在段落内部的注意力模式(Localized Attention)。在这个阶段,注意力主要集中在同一文档内的Token上,表明模型在单个段落内进行了段落内部的信息聚合。即模型已经将需要的信息集中到了部分token中。
  • 这种趋势在上层(24-30)继续并加强,而且可以观察到"Attention Sink"和"Massive Activation"现象。

Massive Activation出自论文"Massive activations in large language models",即在模型的隐藏状态中存在极少数激活值(activations)远大于其他激活值的情况,这些被激活称为"massive activations"。

方案

作者选择通过基于注意力模式动态分配缓存预算来提高压缩效率。针对层的不同特性使用不同的KV Cache机制。

  • 在前几层使用全部的kvcache。因为这些层在聚合全局信息,所以信息更加分散。
  • 越深的层使用越少的kvcache。在这些层中,注意力机制极大地集中在少数几个关键Token上,因此只需要保留这些关键Token就能让输出保持一致并且减少显存占用。而其他的token由于值很低,最后被采样的概率也很低,因此可以考虑丢弃这部分token的cache。

KV Cache最终是一个金字塔形状,即论文中提到的PyrimadKV。一旦为每一层确定了KV缓存预算,PyramidKV在每一个Transformer层中选择根据注意力选择要缓存的KV。最后的部分(\(\alpha\)个)Token的KV缓存,即Instruction Token,则会在所有Transformer层中保留。

为了量化每个token在生成过程中的重要性,作者测量每个token从指令token中获得的关注程度,并使用此测量值为KV缓存选择重要token。具体实现流程如下:

  • 因为最前面的token和最近邻的token都包含重要信息,对生成效果影响显著,不能舍弃,所以用超参数控制这些token的选择,比如始终保留初始64个token(sink token)和最近的128个token(instruction token)。

  • 确定第一层和最后一层要保留的token总数,比如第一层是全部,最后一层是1/10(超参数控制)的token。那么中间layer需要保存的token数目从第一层到最后一层线性递减。具体计算token是用 instruction token和该token计算注意力分数来确定的。

  • 确定了token数目k之后,就可以进行topk的选择。由于qk矩阵每一行的权重和value值加权平均后得到最终的hidden_status,因此筛选规则就是qk矩阵每一行的权重值。

  • 每一层layer在decoding时,只加载对应数目的kvcache进行计算。

1.4.2 PyramidInfer

PyramidInfer也采用金字塔形的预算分配策略,同时在每一层选择具有高注意力值的token。此外,在解码阶段,PyramidInfer依据注意力值来更新动态维护重要token。

动机

之前很多方案忽视了层与层之间的相互依赖以及预计算过程中巨大的内存消耗。PyramidInfer 的作者发现影响未来生成的关键 KV 的数量逐层减少,并且可以通过注意力权重的一致性来提取这些对将来有关键作用的 KV。基于这些发现,作者提出了 PyramidInfer,通过逐层保留关键上下文来压缩 KV Cache。

下图为 PyramidInfer 与 StreamingLLM 和 H2O 的区别,PyramidInfer 中 KV Cache 会逐层递减,越往后越稀疏。

方案

PyramidInfer 的执行过程如下图所示:

  • 在 Prefill 阶段,PyramidInfer 只保留每层的关键上下文(Pivotal Context, PvC)来压缩 KV Cache。
  • 在 Decoding 阶段,PyramidInfer 根据新的最近的 Token 来更新 PvC。

1.4.3 ZigZagKV

ZigZagKV 其实是在SnapKV上继续做优化,也就是说,它也是输入prompt的KV Cache压缩算法,但是ZigZagKV依据层特点进行了稀疏化,因此放在此处。

动机

为了更好的改进,ZigZagKV 定义如下几个指标:

  • MBA(M inimum B udget size required to maintain 90% of the totalAttention score),其实就是维持90%的attention score需要的最小KV数。
  • LMBA(L ayer M inimum B udget size to maintain Attention score),由于每个attention head都可以计算出一个MBA值,我们把一层里面所有的attention head的MBA的均值叫做LMBA。LMBA越高也就意味着需要更多的token来维持精度,也就意味着需要分配更多的显存来存储KV Cache。
  • LMBO(L ayer-wise M inimum B udget for Output):此参数的作用是以输出的相似性这个维度分析一下需要保留多少的KV才能让输出类似。

具体参见下图,其中,\(a_i\)是A的第i个元素。n是当前KV Cache里面的token数量。

通过实验,作者发现:

  • 不同的模型的对应层的LMBA差距比较大;同一个模型不同层的LMBA差距也很大。
  • 同一个模型不同层在保持输出基本一致的限制下需要的KV数量差距很大。

基于上面的两个观察,作者得出的结论就是给模型的每一层分配同样大小显存来存储KV Cache是不合理的

空间分配

为每一层分配多大的空间存储KV Cache是合理的呢?

一个直观的想法就是每一层按照LMBA计算得到的比例来分配,大层多份,小的少分。具体参见下图标号1,其中B是单层的Budget。这么分其实有一个问题,就是如果某一层的LMBA特别大的时候可能导致某些层分配到的空间特别小,也就意味着这些层基本无法保存任何KV。一个层基本没有KV Cache,效果自然保证不了。所以,我们给每个层保留一个最小大小,分配公式改为下图标号2。这就是ZigZagKV的主要改进点。SanpKV给每一层均匀分配空间。而ZigZagKV自适应地分配空间。

解决了每一层KV Cache的空间分配问题,还有一个问题就是要保存那些K和V。ZigZagKV里面token的选择跟SnapKV非常类似。但是,SnapKV给所有的layer的head是平均分配的,ZigZagKV没有给每个attention head分配同样多的显存,而是根据每个head可以保留绝大部分信息需要的token数来进行分配。

1.4.4 Layer-Condensed KV Cache

LCKV建议只计算和缓存一小部分层的键和值,甚至只计算顶层,然后让底层的查询与保存的键和值配对进行推理。这种方法不仅大大提高了推理速度,降低了内存消耗,而且与现有的内存节省技术正交,能够直接集成以进一步优化。虽然这种机制使下一个注意力计算依赖于前一个注意力的顶层密钥和值,这与Transformer的并行训练相矛盾,但LCKV引入了一种近似训练方法来支持并行训练。

具体而言,Transformer 对模型输入进行层层编码的过程,其实也就是对token信息一步步加工的过程,因此Layer-Condensed KV Cache作者认为i最后一层的信息是最重要最全面的,可以忽略前面几层处理得到的信息,只保留最后一层的KV cache。当生成后续 Token 时,其对应的 KV Cache 都从最后一层取。这样模型不再需要计算以及保存前面这些层token的key和value,也不需要前面几层的\(W^K W^V\) 矩阵,极大的节省了计算和内存开销。如下图所示。

1.5 其它方案

1.5.1 针对头特点的稀疏化

针对头特点的稀疏化进行更细粒度的逐头预算分配,它能够在每一层内的单个注意力头之间进行精确的内存分配,提供更灵活和有针对性的优化机会。

  • AdaKV:基于注意力模式差异为各头分配缓存,通过优化L1损失界限最大化保留信息。
  • RazorAttention:描述了两类不同的检索头。"回声(echo)头"专注于先前出现的相同token,"归纳(induction)头"关注与当前token重复的之前token。该框架实现了差异化缓存策略,为检索头维护完整的缓存条目,同时将远程注意力压缩为非检索头的合并补偿注意力。
  • DuoAttention:DuoAttention引入了一种参数化方法来区分两类注意力机制:检索头,对全面的长上下文处理至关重要;流式头,主要处理最近的注意力和注意力吸收器。这种分类是通过学习参数实现的,这些参数可以自动识别需要全注意力的检索头。

1.5.2 投机采样稀疏化

投机采样稀疏化的思路是:通过draft模型生成候选token,然后选取其中重要的token。

论文"TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding"提出了 TriForce,这是一种层次化投机采样系统,主要用于长序列生成。TriForce想解决的是长序列生成背景下的加速问题。在长序列生成设定下,除了模型本身的权重,模型推理的 KV cache 也会占据大量的显存,并在上下文足够长的时候同样成为制约推理速度的关键因素。

因此,TriForce引入了一个只使用部分 KV cache 的 target model 来减小全量的 KV cache 对于推理速度的制约,同时采用多级投机推理策略进一步降低 draft model 的推理开销。TriForce的核心思路如下图所示,其主要分为三个过程,涉及两级投机采样:

  • Speculation Phase 1(此阶段可以执行多次,直到获得满足条件的多个候选 Token):
    • 使用比较小的 LLaMA-68M 模型作为草稿模型,同时使用 StreamingLLM 的方案快速生成多个候选。
    • 使用 LLaMA2-7B-128K 模型作为 Target 模型,采用稀疏 key、value cache 方案进行并行验证。验证后的 Token 作为新的草稿。因为采用的稀疏 key、value cache,因此相应的结果可能是有损的。
  • Speculation Phase 2(此阶段执行一次):使用 LLaMA2-7B-128K 模型作为 Target 模型,采用全量 key、value cahe 方案进行并行验证。此阶段可以保证生成的结果无损。

在缓存的使用上,TriForce采用基于检索的方案:

  • 将 key、value 划分为不同的 Chunk。
  • 使用 query 与每个 Chunk 的中心(使用均值作为中心)计算距离。
  • 选择 topk 的 Chunk 作为稀疏 key、value cache 进行 Attention 计算。

1.5.3 聚类稀疏化

为了加速关键注意力的检索,一些研究工作提出了基于索引的方法,以块或集群(cluster )粒度组织和访问KV缓存,实现了高效的查询和提取操作。InfLLM在块中维护完整的KV缓存,同时通过分层存储策略促进长序列处理。该框架采用CPU-GPU内存编排,在GPU内存中保留基本注意力和当前计算单元,同时将访问频率较低的单元卸载到CPU内存中。为了进一步提高top-k块检索精度,Quest框架提出了一种基于KV缓存块中最小和最大键值的精细块表示方法。

SqueezedAttention在离线阶段采用K-means聚类来对语义相似的密钥进行分组,每组由一个质心表示。在推理过程中,它将输入查询与这些质心进行比较,以识别和加载上下文中语义相关的键。同样,RetrievalAttention使用近似最近邻搜索(ANNS)技术索引KV缓存注意力。此外,EM-LLM动态地将传入的注意力分段为偶发事件。此外,它实现了一种混合检索机制,将语义相似性匹配与时间上下文相结合,以有效地访问相关的KV缓存段。同样,ClusterKV将注意力分组到语义簇中,并在推理过程中选择性地召回它们,从而为LLM实现了高精度和高效率。

我们用ClusterKV为例进行学习。

在如何舍弃K和V的方向的研究主要又细分为两个方向:

  • Non-recallable KV Cache:直接把一些系统判定为不重要的K和V舍弃掉。如何判定不重要是里面关键点。比如,保留attention sink以及最近的窗口里面的token对应的KV cache,把其它的都扔掉,以此来使得计算和显存都可控。

  • Recallable KV Cache:实际上,KV Cache里面的某一个token的重要性是动态变化的。一个token也许在前面不重要,但是在后边就变重要了。为了得到更好的结果,我们不能把某个KV直接抛弃掉,而应该根据具体的情况在一些时候抛掉KV,一些时候又把KV重载载入显存继续使用。

下图给出了几种方法的对比,图上的星号表示重要的token。我们可以看出,重要的token其实是动态变化的。

ClusterKV的一个重要观察就是,对于一个q重要的token对应的K的cos距离都比较小。基于这个观察,作者使用K-means聚类的方法把类似的K和对应的V进行聚类(使用K的相似性作为分类依据),这样每个类都有一个中心向量,计算重要性的时候就使用这个中心向量来计算。如果q和一个中心向量的cos距离(1-cos值)比较小,那么q和这个类里面所有的向量的cos距离都应该比较小,换句话说距离比较小的类 都会比较大。所以我们可以把这一个整个类都加载进来做后续计算,并且保证漏掉的重要token尽量少。

聚类方法如下:

  1. 先把attention sink部分保留,论文简单地直接把前16个token对应的都保留;
  2. 对于prefill阶段的结果进行聚类。作者实验发现,对于32K的context,400个类是最佳权衡;
  3. decoding阶段,每m步做一次聚类。论文中默认情况下每320个token生成完之后新作一次聚类。

在选择类别是,先设置固定大小B,然后按照类别的重要性从最重要的开始取,一直取到所有类里面的元素加起来超过B为止。

0x02 复用

复用也是减少KV序列长度的一个重要手段。在本节,我们主要介绍两种方案:KV Cache合并(Merging)和前缀复用。

2.1 KV Cache合并

面对是否删除不重要的键值对以节省内存,还是保持原样但进行压缩而不删除的问题,尚无定论。删除操作或许能立即缓解内存压力,但可能会潜在地影响模型的性能。相反,更先进的压缩技术则力求在减少内存占用的同时,保持信息的完整性。

KV cache合并(KV cache merging )探索通过压缩或整合KV缓存,尽可能保留被丢弃token的信息,而不会显著降低模型精度,这可以视为对前面提到的token丢弃策略的扩展。KV缓存合并技术不是采用统一的压缩策略,而是利用层内和层间的固有冗余来动态优化内存利用率。这些方法旨在减小KV缓存的大小,同时保留准确注意力计算所需的关键信息,从而在资源受限的环境中实现高效推理。现有的KV缓存合并策略可分为两种主要方法:层内合并,侧重于在各个层内合并KV缓存,以减少每层的内存使用;跨层合并,针对跨层冗余,以消除不必要的重复。这两种方法都提供了互补的优势,提供了平衡内存节省和模型性能下降的灵活性。KV缓存合并的总结如表4所示。

2.1.1 层间合并

对于层间合并,我们以MiniCache为例进行学习。

MiniCache 会合并中高层相似KV对。具体而言,MiniCache并将将相邻相似层中每个token的KV对分解为向量共享方向向量和幅度,仅存储共享的方向向量、注意力幅度和不可缩放的注意力,以最大限度地提高内存效率。

动机

作者发现,KV Cache 在 LLM 中的深层部分的相邻层之间表现出了高度相似性,可以基于这些相似性对 KV Cache 进行压缩与合并。此外,作者还引入了 Token 保留策略,对高度不同的 KV Cache 不进行合并。并且这种方法可以与其他的 KV Cache 量化方案正交使用。

思路

下图给出了MiniCache和以前方法之间的主要区别。MiniCache的特点是沿着LLM的深度(depth)维度研究KV缓存的层间冗余,这是之前基于层内方法忽略的一个方面。这里,T是指预解码的最后一个时间戳,T+1是指解码的第一个时间戳。

下图是MiniCache的压缩策略和保留策略。

  • (a)描述了跨层压缩过程。模型从层l和l-1中获取KV缓存,并通过方程式(3)将它们合并到共享状态中。此外,模型也选择不可合并的注意力进行保留,然后将合并的缓存、保留注意力和数量存储在C中。
  • (b)说明了层l和l-1的恢复过程,其中包括方程式(2)中的数量重新缩放和保留注意力恢复。

算法

下图为MiniCache详细的执行流程。

  • 获取 KV Cache:在 Prefill 阶段,逐层生成 KV Cache。
  • 跨层合并:当到达合并开始层 S 时,将当前层 L 的 KV Cache 与前一层 L-1 的 KV Cache 进行合并,以减少冗余。
  • 缓存:将合并后的 KV Cache 存储起来,以便将来使用。
  • 删除:在 Decoding 阶段,删除不必要的或冗余的 KV Cache,以优化内存使用。
  • 加载和生成:获取所需的 KV Cache,用于生成输出。
  • 恢复:对获取的 KV Cache 应用误差抑制机制,包括 rescaling 和 retention recovery,以最小化合并和压缩过程中引入的误差。
  • 更新:在恢复阶段后,使用最终的 KV Cache 更新共享的 KV Cache。

2.3 层内合并

我们使用DMC来学习层内合并。

在普通 Transformer 中,在每个时间步 𝑡,𝑘𝑡 和 𝑣𝑡 都会被追加到 KV Cache中。

在 DMC 中,KV Cache更新过程有所不同,如下图算法 1 所示。在每个时间步长,DMC 会决定是将当前的key-value对直接添加(append)到缓存中,还是对它们与缓存中顶部的项进行加权平均(weighted average)。

具体而言,DMC 使用一个决策变量 𝛼𝑡∈ {0, 1} (只能取0和1) 和一个重要性变量 ω𝑡∈[0,1]。我们用 𝑘𝑡 和 𝑞𝑡 中的第一个神经元来提取这两个分数。DMC 根据 𝛼𝑡 来决定是将 KV 表示 𝑘𝑡 和 𝑣𝑡 追加到缓存中,还是将其与缓存的最后一个元素累加。

DMC 中的内存以亚线性方式增长,虽然比不上Linear Attention Transformer推理时的内存恒定,但明显好于Transformer。

2.2 前缀复用

LLM服务已经从无状态发展为有状态系统,利用推理请求中固有的依赖关系可以优化KV Cache。这些利用依赖关系的技术可以分为两种类型:跨请求和请求内。

  • 跨请求技术利用跨请求的依赖关系。一个显著的技术是前缀复用,也称之为上下文缓存,它为共享相同提示前缀的请求重用KV缓存,从而加快预填充阶段。能够在请求之间重复使用KV缓存将显著提高时延(尤其是首词元时延)和吞吐量(通过大幅减少具有共享前缀的并发请求的内存占用)。
  • 另一方面,请求内技术利用单个请求内的依赖关系。一个著名的例子是分离推理,将一个请求分成两个子请求以获得更好的调度,以及序列并行性,将一个请求分成多个子请求以分发负载。

本篇介绍前缀复用,下一篇介绍分离推理。

2.2.1 需求

当前大部分 LLM 应用在构造输入时一般遵循 system prompt + user prompt 范式,我们也经常会面临具有长system prompt的场景以及多轮对话的场景。

  • 长system prompt的场景。在LM程序的批量执行过程中,如果有多个大型语言模型(LLM)调用共享了相同的前缀,system prompt在不同的请求中是相同的,并且此部分位于提示的开头,因此其键和值在请求之间也是相同的。system prompt可能在较长的时间跨进程共享,从几秒钟到甚至几小时或几天。这些预定义的长系统提示在命中率低的情况下会造成内存浪费,而且在大规模部署中频繁刷新时是静态和不灵活的。

  • 多轮对话场景。在多轮对话场景下,下一轮的 prompt 实际上可以视作上一轮的 prompt 与 completion 的组合。每一轮对话需要依赖所有历史轮次对话的上下文,历史轮次中的KV Cache在后续每一轮中都要被重新计算。

这两种情况下,理论上可以多次重用KV缓存中的数据。如果能把system prompt和历史轮次中的KV Cache保存下来,留给后续的请求复用,那么可以避免出现多个冗余的 system prompt KVCache,将会极大地降低首Token的耗时,能提高服务的极限吞吐,并降低新一轮请求的首次 Token 计算时间(Time To First Token,TTFT)。比如,对于system prompt,如果这些缓存以自注意力机制中的键值对命名。具有相同提示符前缀的请求就可以共享同一KV缓存,从而有效减少计算冗余和内存消耗。在多轮对话的应用中,如果复用,则可以消除历史轮次中生成对话的recompute。

能够跨请求重用KV Cache将实现显著的延迟(特别是第一个 token 延迟)和吞吐量(通过大大减少具有共享前缀的并发请求的内存占用)改进。

然而,这种优化非常复杂,面临很多技术难点,因为它需要额外的存储和更复杂的内存管理。比如每个推理进程都是独立应用,这些副本是并行的。需要如何跨越多个调用和实例来高效地重用KV缓存(共享提示词前缀)。然而,现有的系统却缺乏有效的机制来支持这种重用,正因如此,目前大多数领先的推理系统都采取了一种相对保守的做法:并未分析整个大批次中的前缀共享特性,也未将共享相同前缀的 Request 一并调度。而是以 Request 为粒度进行任务调度,在完成一个请求的处理后就会丢弃其KV缓存,这不仅可能导致可共享的 KV Cache 过早被驱逐,这一宝贵资源无法在多次请求间共享,还延长了可共享 KV Cache 的生命周期,从而拖慢了执行速度,导致了大量的重复计算和内存浪费。

下图是KV缓存的四种典型模型:小样本学习示例(few-shot learning examples)、自一致性问题(questions in self-consistency) 、多轮聊天中的聊天记录(chat history in multi-turn chat)以及思维树中的搜索记录(search history in tree-of-thought )。蓝色框表示可共享提示部分,绿色框表示不可共享部分,黄色框标记不可共享的模型输出。现有的系统很难自动处理所有模式。

接下来将介绍 Prefix Caching 在一些大型模型推理系统中的实现。

2.2.1 Prompt Cache

论文"Prompt Cache: Modular Attention Reuse For Low-latency Inference"提出了Prompt Cache,这是一种通过在不同的 LLM 提示中重用注意力来加速[LLM)推理的方法。Prompt Cache 采用一种称为schema(议程)的方法来明确定义可重用的文本段,并将其称为提示模块。该schema可确保注意力重用期间的位置准确性。然后,Prompt Cache 在加载一个schema时填充其缓存。在推理时,如果这些"缓存"段出现在输入提示词中,系统使用内存中预先计算的键值注意力,只计算未缓存的文本。Prompt Cache在本质上也是以空间换时间的技术。

下图是完全自回归生成(a)、KV Cache(b)和 Prompt Cache(c)之间的区别,其中每种方法均显示三个步骤(1 到 3步),每个框表示一个token。蓝色框代表提示。

  • LLM 接收提示并预测下一个 token (A) 。然后将生成的 token (A) 附加到提示以预测下一个 token (B)。此过程将持续到满足停止条件。
  • KV Cache 仅对提示计算一次注意力,并在后续步中重复使用KV Cache。
  • Prompt Cache 可以跨服务重用 KV 状态,从而省略提示词的注意力计算。在加载模式时,Prompt Cache 填充缓存,并在基于该模式的提示词中复用缓存状态

随着缓存段的大小增加,Prompt Cache的性能优势会变得更加明显,因为注意力的计算开销与输入序列大小成二次方关系,而 Prompt Cache 的空间和计算复杂度与大小成线性关系。

洞见

许多输入提示都有重叠的文本段,例如系统消息、提示模板和为上下文提供的文档。通过预先计算并将这些频繁出现的文本段注意力存储在推理服务器上,可以在这些文本端出现在用户提示中时有效地重用它们。

然而,跨提示重用注意力时会出现两个挑战。

  • 注意力是依赖位置编码的,文本段的注意力只有在该段出现在相同位置时才能复用,怎么保证复用的注意力和当前 prompt 的位置相同呢?
  • 系统如何高效地识别可能已缓存的文本段?

为了解决这两个问题,Prompt Cache 结合了两个想法。

  • 使用提示标记语言 (Prompt Markup Language, PML) 来格式化提示词的结构。PML 将可重用的文本段标记为模块,即提示模块。它不仅解决了上述第二个问题,而且为解决第一个问题打开了思路,因为可以为每个提示模块分配唯一的独特的位置 ID。
  • 作者发现,LLM 可以处理具有不连续位置 ID 的注意力。这意味着可以提取不同部分的注意力并将它们连接起来形成子集。这一点使用户能够根据需要选择提示模块,甚至在运行时更新某些提示模块。
方案

一个schema是定义提示模块并描述其相对位置和层次结构的文档。LLM 用户用 PML 编写提示,目的是可以根据提示模块重复使用注意力。重要的是,用户必须从议程中派生出提示,该议程也是用 PML 编写的。

当 Prompt Cache 收到提示时,它首先处理其schema并计算其提示模块的注意力。它将这些状态重用于提示中的提示模块以及从同一议程派生的其他提示。

上图显示了基于示例schema的示例提示,可以解释 Prompt Cache 的重用机制:

  • PML 在schema和提示中都明确了可重用的提示模块。提示可以包含在排除的模块位置上新的文本段、参数以及末尾。
  • Prompt Cache预先计算议程中所有模块的注意力(1),并将其缓存以供将来重用。
  • 当向 Prompt Cache 提供提示时,Prompt Cache 会对其进行解析,以确保与声明的schema一致。Prompt Cache会验证导入模块的有效性。然后,Prompt Cache 从缓存中检索导入提示模块缓存的注意力(2),计算出参数(3)和新文本段(4)的注意力,并将它们连接起来,生成整个提示的注意力。

2.2.2 Hydragen

Hydragen是一种硬件感知的精确实现注意力机制的方法,适用于共享前缀的情况。Hydragen分别计算共享前缀和唯一后缀的注意力。利用单个共享的KV缓存来处理多个请求中的公共前缀,然后在跨序列将查询拼接在一起,实现了高效的前缀注意力,减少了冗余的内存读取。

下图给出了Hydragen对应的结构描述:

  • 左:LLM 推理场景示例,其中聊天机器人模型处理许多共享大共享前缀(系统提示)。
  • 右:Hydragen 的概述,其中整体注意力被分解为共享前缀上的注意力以及在独立后缀上的注意力。

在计算注意力时,共享公共前缀的序列具有部分重叠的键和值。我们的目标是将这种部分重叠的计算分为两个单独的操作:关注共享前缀,其中存在总键值重叠;关注唯一后缀,其中没有重叠。

下图给出了推导过程。Hydragen希望避免直接计算SDP (Q, K, V ),转而使用子计算SDP (Q, K_1, V_1) \\(和\\)SDP (Q, K_2, V_2) 的结果来计算这个值。划分注意力的挑战在于softmax运算,因为softmax分母是通过将序列中所有指数化的注意力得分相加来计算的。为了组合两个子计算,Hydragen使用了一种分母重缩放技巧,其灵感来自FlashAttention的blocked softmax计算,通过该技巧,Hydragen可以通过计算全序列softmax分母并相应地重新缩放注意力输出来得到最终结果SDP (Q, K, V )。

下图给出了代码示例。

2.2.3 ChunkAttention

ChunkAttention 将KV Cache分块组织为前缀树,支持运行时检测和共享多个请求之间的公共前缀,通过两阶段分区算法提升数据局部性。 具体而言,ChunkAttention通过将整体的键/值张量分解为较小的块,然后将它们结构化到辅助前缀树中。而且,ChunkAttention在基于前缀树的KV Cache基础之上,又实现了两阶段分区算法,这样在存在共享系统提示的情况下,可以增强自注意力计算时的数据局部性。

前缀树

ChunkAttention认为KV缓存应该是前缀感知的,因此,ChunkAttention将所有解码序列的KV缓存组织成前缀树。KV缓存仅存储当前正在解码的序列的键/值张量,因此,KV缓存是前缀感知的,并且可以在运行时动态检测和删除冗余,无需人为干预。

前缀树中的每条路径定义了一个序列。服务器中可以同时存在多个树(森林),例如,应用程序开发人员设计不同的系统提示。树中的父子关系定义了每个块所覆盖的序列子集。根节点覆盖所有序列,叶子节点只覆盖一个序列。前缀树的一个关键属性是,在前缀树中每个块所覆盖的序列在序列索引维度上是连续的。

下图显示了存储在前缀树中的KV缓存的结构。\(S_0\)、\(S_1\)、\(S_2\)提示中的指令和示例是通用的,可以共享。但是问题各不相同,不可分享。每个节点是一个存储三个关键要素的块(chunk)C,三个要素是:

  • 由序列\(S_i,...,S_j\)共享的c个上下文标记token,以便进行前缀树操作。
  • 对应着c个token的,大小为b×h×c×d的键张量的切片。
  • 相应的值张量切片。

在推理过程中,存在三种可能的情况:i)新序列加入,ii)完成的序列离开,iii)所有序列一起解码一个token。每种情况都可以转化为前缀树操作。当新序列加入时,搜索并更新前缀树以插入新路径。当完成的序列离开时,更新前缀树以删除其路径。在每次解码迭代中,我们将新token追加到叶子块中,或者当叶子块已满时增加新块。

通过共享公共前缀,ChunkAttention将可以同时处理的序列数量增加至约\(1/(1-r)\)。共享比率r由共享token数\(n_s/(n_p + n_c)\)定义,其中\(n_c\)是完成token数。在内存受限的推理场景中,这有助于增加批处理大小,从而提高吞吐量。

两阶段划分(TPP)

在预填充期间,ChunkAttention执行前缀查找,以避免对匹配的提示前缀重复计算KV投影和位置嵌入。对于不匹配的后缀token,仍然计算KV投影和位置嵌入,并将键/值张量划分并插入前缀树中。然后,我们在整个键/值张量上应用现有的高度优化的自注意力内核,例如FlashAttention。

ChunkAttention中的分区策略是基于online softmax构建的,并受到FlashAttention的启发。然而,FlashAttention在处理非连续内存或可变序列长度时不够灵活,更适合于模型训练而不是推理。在解码过程中,当查询token计数始终为1时,几乎没有收益。ChunkAttention擅长处理解码期间可变的序列长度,并将几个序列的注意力操作批处理以减少内存操作。因此,ChunkAttention和FlashAttention是互补的。

在迭代解码过程中,自注意力计算被分为块优先和序列优先两个阶段。这两个阶段关注查询张量的不同切片、KV缓存块,并使用不同的并行化策略。具有匹配提示前缀的序列的查询张量被批处理在一起,与键/值张量一起执行注意力。

该过程如下图所示。这里隐含了头维度的处理。

块优先阶段只处理被多个序列共享的块。由于GPU具有比头数(Llama 7B为32)更多的流式多处理器(A100为108),而通过头进行划分会低效利用硬件资源,因此我们会对键/值进行额外的划分。具体算法如下图所示:遍历前缀树中的共享块,在遍历中执行部分注意力内核partial_attn,并将部分注意力结果保存到内存中。

在序列优先阶段,我们从块优先阶段加载共享块的部分注意力结果,并继续处理与一个特定序列相关的块。我们通过切片Q的第i行来划分序列,每个由序列优先内核处理的q都是一个向量,如下图算法所示。

2.2.4 AttentionStore

AttentionStore允许在同一对话的后续轮次中重用KV缓存,从而有效减少了多轮对话中的重复计算开销。为了提高AttentionStore的效率,其作者计了重叠KV缓存访问、分层放置KV缓存和位置编码解耦的KV缓存截断方案。

动机

当前LLM服务引擎在执行多轮对话时存在显著的效率问题,主要原因是多轮对话中KV缓存的重复计算。

在一次对话中,LLM引擎会将中间数据,即键值(KV)对,存储在GPU上的高带宽内存(HBM)中。然而,当对话结束且会话变得不活跃时,为了给其他活跃会话腾出空间,LLM引擎通常会丢弃与该会话相关的KV缓存。当同一个会话再次变得活跃(即用户发送下一个消息)时,LLM引擎需要重新计算整个KV缓存。这就导致了同样的KV缓存被重复计算,浪费了宝贵的GPU计算资源。随着对话轮次的增加,重复计算的开销也会线性增加。

如下图(a)所示:在对话的第一个轮次,LLM服务引擎会生成q1和a1的KV缓存。然而,在完成第一个轮次后,为了回收HBM(高带宽内存)空间,LLM服务引擎会丢弃这些KV缓存。在第二轮次中,LLM服务引擎需要重新生成q1和a1的KV缓存。在第三轮次中,它需要重新生成q1、a1、q2和a2的KV缓存。随着会话的扩展,历史token不断积累,重复计算的开销显著增加。

方案

AttentionStore引入了一个分层的KV缓存系统,利用成本效益高的存储介质来保存所有请求的KV缓存,从而显著提高了多轮对话的效率。

在传统的注意力机制中,每次对话结束后,KV缓存都会被丢弃。AttentionStore将KV缓存保存在一个KV缓存系统中。当对话会话不活跃时,相关的KV缓存会被保存下来。如果同一会话在未来重新激活,其KV缓存将从KV缓存系统中获取并重新用于推理。这样,AttentionStore只需对新一轮对话中的新输入token进行部分预填充,而无需重新预填充所有历史token。例如,在执行第三轮对话推理时,q1、a1、q2和a2的KV缓存被复用,只需预填充q3,如上图(b)所示。

另外,为了应对技术挑战,AttentionStore也做了很多设计,通过这些设计,AttentionStore有效解决了多轮对话中的KV缓存管理问题,显著提高了多轮对话的效率和性能。

  • 为了减少从慢速存储介质访问KV缓存的开销,AttentionStore采用了分层预加载和异步保存方案,将KV缓存的访问与GPU计算重叠进行。
  • 同时,为确保需要访问的KV缓存始终处于最快速的层级,AttentionStore使用了调度感知的获取和驱逐方案,根据推理任务调度器提供的提示,智能地将KV缓存放置在不同的层级中。
  • 此外,AttentionStore通过解耦位置编码和有效截断KV缓存,避免了因上下文窗口溢出而导致的KV缓存失效。这一创新使得保存的KV缓存始终保持有效。

2.2.5 Radix Attention

SGLang论文提出的RadixAttention技术是实现KV缓存重用的典型方案之一,也是高效内存管理的代表。

与在生成请求完成后丢弃 KV 缓存的现有系统不同,RadixAttention会利用基数树实现了缓存的快速匹配、插入和替换。而且,在RadixAttention中,无论是Prefix还是Generate阶段产生的KV Cache,都会被缓存。这可以最大程度增加KV Cache被新请求复用的几率。作为对比,ChunkAttention只聚焦在prefix本身。

思路

SGLang的Prefix Caching方案关键点如下:

  • 跟以前那些"用完就扔"的系统不一样,在每次处理完请求后,RadixAttention 算法不会在完成生成请求后丢弃KV缓存(提示符和生成的结果对应的KV Cache),而是将其保留在GPU内存中。
  • RadixAttention 算法还将序列 token 与 KV Cache 之间的映射关系保存在 RadixTree 数据结构中。相较于传统的字典树(或称前缀树),Radix树的边不仅可以标记为单个元素,还可以标记为长度可变的元素序列。在必要的时候,一个已经在Tree中的大node可以动态地分裂成小node,以满足动态shared prefix的需求。这一特性极大地加速了数据处理速度。
  • 收到新请求时,调度器使用radix tree进行前缀匹配。如果有缓存命中,则调度器重用缓存的KV张量来满足请求。
  • 由于GPU内存有限,缓存的KV张量不能无限期保留。因此,RadixAttention算法包括一个驱逐策略(例如最近最少使用(LRU)驱逐策略)。最优缓存重用可能与像先来先服务这样的调度不兼容。因此,RadixAttention配备了一个修改后的调度器,优先考虑与缓存前缀匹配的请求(缓存感知调度)。

方案

鉴于GPU内存资源宝贵且易于被迅速消耗的特点,不能永远保留缓存的KV张量。为了让缓存更加高效,RadixAttention还特别设计了两种策略:

  • "最近最少使用"(LRU)的淘汰策略。该策略优先移除那些最长时间未被访问的叶子节点,给新的内容腾地方。通过这一策略,我们不仅有效释放了被淘汰叶子节点所占用的内存,还充分利用了这些节点的共同祖先节点,实现内存的最大化复用。随着系统运行,当这些祖先节点也变成孤立的叶子节点并符合淘汰条件时,它们同样会被有序地逐出,从而保持整个缓存系统的健康与高效。
  • 缓存感知的调度策略。我们将缓存命中率定义为已缓存提示token数除以提示token数(number of cached prompt tokens / number of prompt tokens)。当等待队列中有许多请求时,执行请求的顺序会显著影响缓存命中率。例如,如果请求调度程序频繁地在不同的、不相关的请求之间进行切换,则会导致高速缓存抖动和较低的命中率。因此,RadixAttention设计了一个具有缓存感知的调度算法来提高缓存命中率。这种算法会按匹配的前缀长度对请求进行排序,并优先处理具有较长匹配前缀的请求,而不是使用先到先服务的调度。

图中(1)~(9)表示Radix Tree的动态变化。

整个树结构表示对请求的共享和分支操作。每条边都带有一个标签,代表子字符串或token序列。节点用颜色编码以反映不同的状态: 绿色表示新添加的节点,蓝色表示在当前时间点内访问的缓存节点,红色表示已被驱逐的节点。新到达的请求会先从 Radix Tree 中匹配最长前缀,寻找最远的一个节点。当请求仍有 Tokens 没匹配到已有的前缀,则会新增 Token 子节点。每次新增 Token 时,将新增 Token 节点token为绿色,并将其前缀的节点token为蓝色。这种token主要是更新节点的访问时间,当缓存空间不足时,会将最近没被访问的节点清除。具体操作流程如下:

  • 在步骤 (1) 中,基数树最初为空。
  • 在步骤 (2) 中,system prompt(系统提示)为"You are a helpful assistant",服务器处理传入的用户消息"Hello",并使用 LLM 输出"hi"进行回复。整个对话(系统提示词 "You are a helpful assistant.",用户消息"Hello",以及 LLM 回复"hi")被合并为一个大node,保存到Radix Tree的新节点a。
  • 在步骤 (3) 中,同一个用户输入了一个新的prompt,服务器在基数树中找到历史对话前缀(即第一轮对话),并重新利用其 KV 缓存。新回合(Solve this problem...)作为新节点"b"被追加到树中。
  • 在步骤 (4) 中,开始新的聊天会话。服务器将 (3) 中的大节点"b"拆分为两个节点"c"和"d",以允许两个聊天会话共享system prompt。
  • 在步骤 (5) 中,第二个聊天会话继续,用户请求 "Write a story"。但是,由于内存限制,必须驱逐 (4) 中的节点"c"。新回合在 (4) 中被缓存到新的节点"d"。
  • 在步骤 (6) 中,服务器收到一批小样本学习查询,这些查询共享相同的小样本示例集。服务器对其进行处理并将其插入树中。因为新查询与现有节点不共享任何前缀,所以根节点被拆分,生成节点"e",以启用共享。
  • 在步骤 (7) 中,服务器收到一批其他的小样本学习查询。这些查询和(6)共享相同的小样本示例集,因此我们拆分 (6) 中的节点,插入新节点"f","g"和"h"。
  • 在步骤 (8) 中,服务器从第一个聊天会话收到一条新消息。它驱逐了第二个聊天会话(节点"g"和"h")中的所有节点,因为它们是最近最少使用的。
  • 在步骤 (9) 中,服务器收到一个请求,以对 (8) 中节点"j"中的问题进行更多答案采样。为了给这些请求腾出空间,我们驱逐了 (8) 中的节点"i"、"k"和"l"。

0x03 依据检索的方案

依据检索的方案是一种更为激进的策略:如果GPU内存不足,就将KV Cache存储于更丰富但速度较慢的外部设备,甚至可能是另一种存储介质上。这种方法将KV Cache管理转变为一项检索挑战,即按需检索相关键值对并重新集成到模型中。虽然这能减少主设备上的内存使用,但会增加检索和集成过程的复杂性。

3.1 InfiniGen

现代LLM服务系统(DeepSpeed Inference & FlexGen)支持将数据卸载到CPU内存中,以便在有限的硬件预算内高效地为LLMs提供服务。这些基于卸载的推理系统甚至开始支持将KV Cache卸载到CPU内存中,从而允许用户在GPU内存容量之外生成更长的上下文。然而,因为CPU和GPU之间的PCIe带宽较低,将庞大的KV Cache从CPU内存传输到GPU成为LLM推理中新的性能瓶颈。

论文"InfiniGen: Efficient generative inference of large language models with dynamic KV cache management"通过预测机制,将KV Cache的主要部分卸载到CPU上,仅保留关键组件。在有新请求时,InfiniGen利用有限保留的信息(利用从上一层查询选择的重要KV条目的索引来检索当前层中的KV缓存条目)来推测性地将一小部分KV Cache重新加载到GPU上,从而在不影响性能的前提下,节省高速计算设备的内存空间,同时避免大量数据交换。

InfiniGen是为了解决Offload中由于PCIe带宽不足带宽的推理延迟增加的问题。只是最后的解决方案是从算法的角度出发的。主要是节省decode的时间。

问题

下图展示了不同方式来执行Transformer块的高层时序图,从中可以看出性能瓶颈。

  • (a) 表示KV Cache完全位于GPU内存中的情况。在这种情况下,KV Cache的加载延迟(Load Cache)仅涉及从GPU内存读取操作,由于GPU内存的高带宽,这个延迟几乎可以忽略不计。然而,最大批量大小或序列长度受限于GPU内存容量,而GPU内存容量通常比CPU内存要小。
  • (b)表示将KV Cache卸载到CPU内存中(KV Cache在CPU上)。尽管基于卸载的推理系统缓解了批量大小和序列长度的限制(可以支持更大的批量大小或更长的序列长度),但由于PCIe带宽的限制,将数百GB的KV Cache传输到GPU进行注意力计算会显著增加Transformer块的整体执行时间。
  • (c)表示传统的预取技术(预取KV Cache)。在这种情况下,只有部分加载延迟能够被前一个Transformer块的计算所隐藏。虽然通过量化压缩KV Cache可能会在基于卸载的系统中减少数据传输的开销, 但这并不能解决根本问题,因为量化并未解决KV Cache问题的根本原因------即KV条目随着序列长度的线性增长。因此,必须进行智能的KV Cache管理,以减轻性能开销,同时保留其优势。
  • (d)表示通过识别计算注意力得分的关键键和值,减少需要加载的KV Cache量。在注意力计算中,某些tokens的键和值比其他token更为重要。如果选择适当的token,跳过一些不太关键token的注意力计算并不会显著降低模型的准确性。
洞察

InfiniGen的核心洞察在于,通过对当前层输入以及后续层部分查询权重和键缓存进行最小限度的回顾,可以推测出对后续Transformer注意力层计算至关重要的少数关键token。这一方法使我们能够仅预取必要的KV Cache条目(而非全部),从而缓解了基于卸载的LLM服务系统中从主机内存获取数据时的开销。

大型语言模型中的异常值

大型语言模型在Transformer块的输入张量中存在异常值。这些异常值指的是与其他元素相比,具有明显更大数值的元素。在LLMs中,这些异常值会在层与层之间出现在少数固定的通道上(即二维矩阵中的列)。模型输入的任何一行与权重矩阵(如 \(W^K\))的一列之间的点积可能具有同样大的幅度,这又导致了查询和键矩阵中的异常值通道。

注意力分数就高度依赖于查询和键矩阵中这少数具有较大幅度的几列。这些大幅度列对注意力模式有很大影响,是因为查询和键之间的点积高度受这些少数几列的影响。

InfiniGen作者观察到,通过倾斜查询(Query)和键(Key)矩阵,使得一小部分通道远大于其他通道,并仅利用这些通道来计算注意力得分矩阵,可以有效地预测哪些token是重要的。本质上,我们将Q和K矩阵与正交矩阵A相乘,使其与Q拉伸最明显的方向对齐,从而生成相应的倾斜矩阵Q和K。

为了找到这样的正交矩阵A,我们采用了奇异值分解(Singular Value Decomposition,简称SVD),这是线性代数中广泛使用的矩阵因式分解技术。对于一个大小为m×n的实数矩阵Q,其SVD分解可以表示为如下形式:

\[\mathbf{Q} = \mathbf{U} \mathbf{Σ} \mathbf{V}^{\top} \]

其中,U和V分别是大小为m×m和n×n的正交矩阵,Σ是一个m×n的对角矩阵,其对角线上的非零值(σ1,σ2,...,σk)表示奇异值,其中k = min(m,n)。

例如,下图展示了当m和n都为2时,\(V^T\)的列向量v1和v2如何变换为Q的列向量q1和q2。在(a)中,正交单位向量v1和v2首先被拉伸到椭圆上的点,这些点的半轴长度对应于Σ对角线上的元素。然后,这些向量通过矩阵U进行旋转和/或反射,变成q1和q2。另一方面,(b)展示了正交矩阵A如何进行旋转,使得得到的q1远大于q2。具体来说,A将向量v1和v2旋转到e1和e2,它们分别映射到椭圆的半轴上。这样,向量就通过Σ矩阵被拉伸到最大和最小值。这个过程强调了q1相对于q2的幅度,使得我们能够有效地仅使用q1来预测注意力得分,而忽略q2。

设计

InfiniGen建立在两个关键设计原则之上。

  • 首先,在计算Layer i-1 层的时候,同时计算出第 i 层需要哪些token的kv cache。即推测并选择对产生下一个输出token至关重要的KV Cache条目,同时丢弃非关键条目。
  • 其次,它利用CPU内存容量,在CPU上维护KV Cache池,而不是在GPU上,以确保在缓解长内容生成中GPU内存容量有限问题的同时,能够为大窗口大小的所有输出和层识别关键的KV Cache值。
总体方案

下图展示了InfiniGen的操作流程。

离线阶段

在离线阶段,InfiniGen修改权重矩阵以生成偏斜的查询和键矩阵。为了实现这一点,InfiniGen首先使用样本输入执行模型的前向传递一次。在此过程中,InfiniGen从每一层收集查询矩阵,并对每个查询矩阵执行奇异值分解(SVD)。使用查询矩阵的分解矩阵获得每一层的偏斜矩阵(Ai)。然后,将该偏斜矩阵与相应层中的每个查询和键权重矩阵相乘,得到偏斜的查询和键矩阵。

需要注意的是:

  • 偏斜不会改变原始功能。即使在偏斜之后,注意力层也会产生相同的计算结果。

  • 偏斜是一个一次性的离线过程,并且不会在运行时产生任何开销,因为我们修改的是运行时不变的权重矩阵。

  • 由于我们利用了列模式,这种模式源自模型的固有属性而非输入,因此在偏斜之后,无论我们为不同的输入计算查询和键,其值都会表现出高度的偏斜性,从而提高了我们预取模块的有效性。

上图标号1 展示了InfiniGen如何创建这些部分矩阵。

Prefill阶段

在预填充阶段,InfiniGen从查询权重矩阵和键缓存中选择几个重要的列,以推测注意力模式,并生成在解码阶段使用的部分查询权重和键缓存矩阵。

InfiniGen的预取模块建立在以下关键观察之上,大型语言模型(LLMs)中连续注意力层的注意力输入高度相似。Transformer模型中连续层的注意力输入之所以具有相似性,主要是因为层归一化(LayerNorm)和残差连接(Residual Connections)的作用。层归一化减少了输入数据的尺度差异,使得各层的输入在数值上更加接近。而残差连接使得每一层的输出都直接加回到输入上,这进一步增强了层与层之间输入的相似性。因此,尽管每一层都会处理输入并产生新的输出,但由于层归一化和残差连接的存在,这些输出在作为下一层输入时,仍然保留了与上一层输入较高的相似性。

此处对应图上标号2。

解码阶段

标号3 展示了InfiniGen如何计算推测的注意力分数。

解码阶段:在解码阶段,InfiniGen推测下一层的注意力模式,并确定要预取的关键键和值。

在层i-1中,我们使用在预填充阶段确定的层i的部分查询权重矩阵和键缓存,以及层i-1的注意力输入。在将部分查询和部分键缓存相乘后,InfiniGen仅预取注意力分数大于最高注意力分数减去α的token的键和值。由于多个注意力头是并行计算的,我们确保同一层中的每个头通过平均各个头之间最高分数和阈值之间的token数来预取相同数量的token。

通过减少需要加载和计算的键值(KV)缓存量,InfiniGen有效地降低了加载延迟(即从CPU到GPU的数据传输时间),同时保持了与具有完整KV Cache的原始模型相似的输出质量。此外,由于InfiniGen不需要从CPU内存中加载固定数量的token,因此它仅利用必要的PCI互连带宽。InfiniGen从第1层开始进行推测和预取,因为在第0层的计算过程中会出现对于利用输入相似性至关重要的异常值。

KV Cache资源池管理

我们将键值(KV)缓存管理为一个池,将其卸载到CPU内存中,并仅将必要数量的数据预取到GPU中。虽然CPU内存的成本更低且容量大于GPU内存,但其容量仍然有限。因此,在某些部署场景中,限制KV Cache池的大小并移除那些由查询token不常选择的、重要性较低的KV条目可能至关重要。我们扩展了设计,以纳入用户定义的内存大小限制。在运行时,当CPU内存的大小达到用户定义的限制时,KV Cache池管理器会选择一个受害者KV条目进行逐出。随后,管理器会用新生成的键和值覆盖选定的受害者条目,并更新GPU中相应的部分键缓存。需要注意的是,KV条目的顺序可以是任意的,只要相同token的键和值在KV Cache池中保持相同的相对位置即可。

选择受害者(即要被逐出的KV条目)的策略非常重要,因为它直接影响模型的准确性。我们考虑了基于计数器的策略以及两种广泛使用的软件缓存逐出策略:先进先出(FIFO)和最近最少使用(LRU)。

0xFF 参考

Z. Zheng, X. Ji, T. Fang, F. Zhou, C. Liu, and G. Peng, "Batchllm: Optimizing large batched llm inference with global prefix sharing and throughput-oriented token batching," 2024. [Online]. Available: https://arxiv.org/abs/2412.03594

L. Zheng, L. Yin, Z. Xie, C. Sun, J. Huang, C. H. Yu, S. Cao, C. Kozyrakis, I. Stoica, J. E. Gonzalez, C. Barrett, and Y. Sheng, "Sglang: Efficient execution of structured language model programs," 2024. [Online]. Available: https: <//arxiv.org/abs/2312.07104>

L. Ye, Z. Tao, Y. Huang, and Y. Li, "Chunkattention: Efficient selfattention with prefix-aware kv cache and two-phase partition," 2024. [Online]. Available: https://arxiv.org/abs/2402.15220

C. Hu, H. Huang, J. Hu, J. Xu, X. Chen, T. Xie, C. Wang, S. Wang, Y. Bao, N. Sun, and Y. Shan, "Memserve: Context caching for disaggregated llm serving with elastic memory pool," 2024. [Online]. Available: https://arxiv.org/abs/2406.17565

GPU加速LLM论文夸赞|InfiniGen通过动态KV Cache管理高效实现LLM的生成式推理:动机 GPUer

llm 推理 latency 分析

ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition 赵尚春

AttentionStore: Cost-effective Attention Reuse across Multi-turn Conversations in* *Large Language Model Serving

Modular Attention Reuse for Low-Latency Inference(@yale.edu)

论文: 《Efficiently Programming Large Language Models using SGLang》

ZigZagKV,KVCache80%的压缩率可以达到近乎无损的效果

ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition(@microsoft.com)

Efficient Prompt Caching via Embedding Similarity(@UC Berkeley)

Modular Attention Reuse for Low-Latency Inference(@yale.edu)

33 倍 LLM 推理性能提升:Character.AI 的最佳实践

[CLA\] Reducing Transformer Key-Value Cache Size with Cross-Layer Attention](https://arxiv.org/pdf/2405.12981) \[[GQA\] GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://arxiv.org/pdf/2305.13245) [\[KV Cache优化\]🔥MQA/GQA/YOCO/CLA/MLKV笔记: 层内和层间KV Cache共享](https://zhuanlan.zhihu.com/p/697311739) [DefTruth](https://www.zhihu.com/people/qyjdef) [\[LLM 推理服务优化\] DistServe速读------Prefill \& Decode解耦、模型并行策略\&GPU资源分配解耦](https://zhuanlan.zhihu.com/p/698341542) [阿杰](https://www.zhihu.com/people/zhou-zhou-69-22) [\[LLM\]KV cache详解 图示,显存,计算量分析,代码](https://zhuanlan.zhihu.com/p/646577898) [莫笑傅立叶](https://www.zhihu.com/people/zwwwk) [\[LLM性能优化\]聊聊长文本推理性能优化方向](https://zhuanlan.zhihu.com/p/698308542) [阿杰](https://www.zhihu.com/people/zhou-zhou-69-22) \[[MLKV\] MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding](https://arxiv.org/pdf/2406.09297) \[[MQA\] Fast Transformer Decoding: One Write-Head is All You Need](https://arxiv.org/pdf/1911.02150) [\[Prefill优化\]\[万字\]🔥原理\&图解vLLM Automatic Prefix Cache(RadixAttention): 首Token时延优化](https://zhuanlan.zhihu.com/p/693556044) [DefTruth](https://www.zhihu.com/people/qyjdef) \[[YOCO\] You Only Cache Once: Decoder-Decoder Architectures for Language Models](https://arxiv.org/pdf/2405.05254) [A Survey on Efficient Inference for Large Language Models](https://arxiv.org/abs/2404.14294) [akaihaoshuai](https://www.zhihu.com/people/wan-wan-mei-xiang-dao-71-71) [Cascade Inference: Memory Bandwidth Efficient Shared Prefix Batch Decoding](https://flashinfer.ai/2024/02/02/cascade-inference.html) [Confused with four attention mechanism and their performance mentioned by paper · Issue #33 · mit-han-lab/streaming-ll](https://github.com/mit-han-lab/streaming-llm/issues/33) [Continuous Batching:一种提升 LLM 部署吞吐量的利器](https://zhuanlan.zhihu.com/p/657586838) 幻方AI [DejaVu:KV Cache Streaming For LLM](http://mp.weixin.qq.com/s?__biz=MzA5MTQxMTM0Nw==&mid=2247484113&idx=1&sn=7fce34f26756903637051938c1c02715&chksm=907d842da70a0d3b9ff13f6be7098b816512d875900b20174f495737b211eb0647df85828890&scene=21#wechat_redirect) [Efficient Streaming Language Models with Attention Sinks](https://arxiv.org/abs/2309.17453) [GPT模型与K/V Cache导读](https://zhuanlan.zhihu.com/p/633653755) [皓月争曦](https://www.zhihu.com/people/yang-min-64-15) [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://arxiv.org/pdf/2305.13245v1.pdf) [H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models](https://arxiv.org/abs/2306.14048) [Hydragen: High-Throughput LLM Inference with Shared Prefixes(@Stanford University)](https://arxiv.org/pdf/2402.05099v1) [Inference without Interference:混合下游工作负载的分离式LLM推理服务](https://mp.weixin.qq.com/s?__biz=Mzg5MzkxMDA5Ng==&mid=2247490528&idx=1&sn=929ac6336d52fcc92d8767dcf860b0b1&chksm=c026f6dff7517fc96b31ceecefd0cd6e212fdf9b06dce6bcd1a5716c525007f89ee4dcaa5435&mpshare=1&scene=1&srcid=08127gtQOc8VjuHDqqjnJhuq&sharer_shareinfo=15587159e1e71e9f291462fe5e0c2abc&sharer_shareinfo_first=15587159e1e71e9f291462fe5e0c2abc#rd) 竹言见智 [kimi chat大模型的200万长度无损上下文可能是如何做到的?](https://mp.weixin.qq.com/s/ow7n7W2Or9VmmtxRxaj9Hg) \[GiantPandaCV\](javascript:void(0)😉 [KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache](https://arxiv.org/pdf/2402.02750) [LaViT:这也行,微软提出直接用上一层的注意力权重生成当前层的注意力权重 \| CVPR 2024](https://mp.weixin.qq.com/s?__biz=MzkyMDE2OTA3Mw==&mid=2247531933&idx=2&sn=27fa53abb35fbec186cd3801983779f8&chksm=c060731f8427a4c6d32de787447d8604fecf0739f6dcacc599153285a7c4683cf9dfb590d239&mpshare=1&scene=1&srcid=0821UThfJG2X8cWUw2eUOuwb&sharer_shareinfo=41b3507fb5f2282a9516d6c002e13797&sharer_shareinfo_first=41b3507fb5f2282a9516d6c002e13797#rd) VincentLee [LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding](https://arxiv.org/abs/2404.16710) [LLaMa 原理+源码------拆解 (KV-Cache, Rotary Positional Embedding, RMS Norm, Grouped Query Attention, SwiGLU)](https://blog.csdn.net/weixin_54338498/article/details/135269411) [Yuezero_](https://blog.csdn.net/weixin_54338498) [LLM Inference Series: 3. KV caching unveiled](https://medium.com/%40plienhar/llm-inference-series-3-kv-caching-unveiled-048152e461c8) [LLM 推理优化探微 (2) :Transformer 模型 KV 缓存技术详解](https://zhuanlan.zhihu.com/p/685547693) Baihai IDP [LLM 推理优化探微 (3) :如何有效控制 KV 缓存的内存占用,优化推理速度?](https://zhuanlan.zhihu.com/p/686107883) Baihai IDP [LLM推理入门指南②:深入解析KV缓存](https://zhuanlan.zhihu.com/p/689870031) OneFlow [LLM推理加速02 KV Cache](https://zhuanlan.zhihu.com/p/692183150) [沉思的斯多葛九狗](https://www.zhihu.com/people/yuan-xiao-jiu-56) [LLM推理加速3:推理优化总结Mooncake/AttentionStore/vllm0.5/cache优化 etc](https://zhuanlan.zhihu.com/p/706144160) [LLM推理加速调研](https://zhuanlan.zhihu.com/p/699776257) [akaihaoshuai](https://www.zhihu.com/people/wan-wan-mei-xiang-dao-71-71) [LLM推理算法简述](https://zhuanlan.zhihu.com/p/685794495) [陆淳](https://www.zhihu.com/people/zhang-da-pu-6) [MiniCache 和 PyramidInfer 等 6 种优化 LLM KV Cache 的最新工作](https://mp.weixin.qq.com/s?__biz=Mzk0ODU3MjcxNA==&mid=2247486886&idx=1&sn=a72b7be7e1f258f444781c034b0d905c&chksm=c364cae3f41343f59ebb5851fd71567b240af08e40979fa50448414d14c9ce1a7f419b9d255b&scene=178&cur_album_id=3209901770408230914#rd) AI闲谈 [MLKV:跨层 KV Cache 共享,降低内存占用](http://mp.weixin.qq.com/s?__biz=Mzk0ODU3MjcxNA==&mid=2247487052&idx=1&sn=6f5ba90078e463ee5a8fe69df8701256&chksm=c364c909f413401fc442f16fdb29d3b4bbdfe2838710c377e6ed1478757da4bd7c6ed6c17afe&scene=21#wechat_redirect) [MLSys 2024 Keyformer:选择关键token,降低50%KV 缓存](https://mp.weixin.qq.com/s?__biz=MzA5MTQxMTM0Nw==&mid=2247484353&idx=1&sn=d0fd89049075e6d8c83738fe5dfb0b80&chksm=9177c54b162507324bb6ce3a5b4e732b5f5093db9fdc0ad5a343259a91abdc71e8e65c03093f&mpshare=1&scene=1&srcid=0822vl21jwhFjSE5LfrOSBC1&sharer_shareinfo=cc67f8ea061a7b70db2e1d6de07a8edd&sharer_shareinfo_first=cc67f8ea061a7b70db2e1d6de07a8edd#rd) MLSys2024 \[MLSys2024\](javascript:void(0)😉 [MODEL TELLS YOU WHAT TO DISCARD: ADAPTIVE KV CACHE COMPRESSION FOR LLMS](https://arxiv.org/pdf/2310.01801) [Mooncake (1): 在月之暗面做月饼,Kimi 以 KVCache 为中心的分离式推理架构](https://zhuanlan.zhihu.com/p/705754254) [ZHANG Mingxing](https://www.zhihu.com/people/james0zan) [Mooncake (2):Kimi "泼天的流量"怎么接,分离架构下基于预测的调度策略](https://zhuanlan.zhihu.com/p/706204757) [ZHANG Mingxing](https://www.zhihu.com/people/james0zan) [Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving](https://zhuanlan.zhihu.com/p/706109023) [手抓饼熊](https://www.zhihu.com/people/tongsanpang) [Orca: How to Serve Large-scale Transformer Models](https://zhuanlan.zhihu.com/p/630281367) [WXB506](https://www.zhihu.com/people/WXB506) [OSDI22 ORCA:LLM Serving系统优化经典之作,动态batch,吞吐提升36倍](https://mp.weixin.qq.com/s?__biz=MzA5MTQxMTM0Nw==&mid=2247484528&idx=1&sn=0a4a0dfe650f36929c23a331ca333c3e&chksm=907d828ca70a0b9a4da4a66993b612687cf0dafb8267f47f57788a14325b51b1b13965211e53&mpshare=1&scene=1&srcid=0812L4PeQePrB3qOqN0Rac2M&sharer_shareinfo=df32ee39f343b6a7e66d1969d3635356&sharer_shareinfo_first=df32ee39f343b6a7e66d1969d3635356#rd) MLSys2024 [OSDI24最新力作:InfiniGen大模型推理高效KV缓存管理框架,吞吐提升3倍,准确性提升32%](https://mp.weixin.qq.com/s/Jus8GR5XCbhKNWGBR_Z45w) MLSys2024 [SnapKV: KV Cache 稀疏化,零微调加速长序列 LLM 推理](https://mp.weixin.qq.com/s?__biz=Mzk0ODU3MjcxNA==&mid=2247486640&idx=1&sn=2ff9f43414c6ccc8297f172afcfc8b07&chksm=c364cbf5f41342e364f066b5ecf509457ce280f138863dff681db2e29f32953d3ccf358b1ef4&scene=178&cur_album_id=3209901770408230914#rd) AI闲谈 [SnapKV: LLM Knows What You are Looking for Before Generation](2404.14469%5D) [Token 稀疏化 + Channel 稀疏化:16 倍 LLM 推理加速](https://mp.weixin.qq.com/s?__biz=Mzk0ODU3MjcxNA==&mid=2247487617&idx=1&sn=13c07766f50a514e21c04aee55a99086&chksm=c26e4be65eac4814fcf1ec50fb471a1d571b3b2da91b6519eceb9f5867a4a77a0eeab02ab52f&mpshare=1&scene=1&srcid=0822VwhkTpwNz5fooxdqqas0&sharer_shareinfo=efc873c3d8b5e45ba76ca5a464b98f84&sharer_shareinfo_first=efc873c3d8b5e45ba76ca5a464b98f84#rd) AI闲谈 [Transformers KV Caching Explained](https://medium.com/%40joaolages/kv-caching-explained-276520203249) [TriForce:KV Cache 稀疏化+投机采样,2.3x LLM 无损加速](http://mp.weixin.qq.com/s?__biz=Mzk0ODU3MjcxNA==&mid=2247486626&idx=1&sn=9eb019aa7025592b5ba29b0d43f76a42&chksm=c364cbe7f41342f1c2597c8537cd14d09a4db9eb759e3e61c20990fc3f3ef6f02d0d9853c345&scene=21#wechat_redirect) [vAttention:用于在没有Paged Attention的情况下Serving LLM](https://zhuanlan.zhihu.com/p/702606462) [BBuf](https://www.zhihu.com/people/zhang-xiao-yu-45-67-74) [YOCO:长序列优化,1M 长文本10x LLM 推理加速](https://mp.weixin.qq.com/s?__biz=Mzk0ODU3MjcxNA==&mid=2247486818&idx=1&sn=0f4e071b279f93a8c7789948b9eb4268&chksm=c364ca27f4134331359a117ebf47288f46fb7b6410c28b0c7c0a4b87637ef118eba87efea63f&scene=178&cur_album_id=3209901770408230914#rd) AI闲谈 [You Only Cache Once: Decoder-Decoder Architectures for Language Models](https://arxiv.org/abs/2405.05254) [ZHANG Mingxing:Mooncake (1): 在月之暗面做月饼,Kimi 以 KVCache 为中心的分离式推理架构](https://zhuanlan.zhihu.com/p/705754254) [【手撕LLM- KV Cache】为什么没有Q-Cache ??](https://mp.weixin.qq.com/s?__biz=MzkzNzU4MTU5Nw==&mid=2247485108&idx=1&sn=08d3b4a987d11299f4a4dff6b63f4bdc&chksm=c33c46ba14746340a104a2259071b38b1f82f08c2fc0f459807e60fa256a909bb33dc2560a5b&mpshare=1&scene=1&srcid=0820LYSbynXSLvsFcmtEFjL1&sharer_shareinfo=2224627bd558275d6b1933a2149f2f1a&sharer_shareinfo_first=2224627bd558275d6b1933a2149f2f1a#rd) 小冬瓜AIGC [一文读懂KVCache](https://zhuanlan.zhihu.com/p/686183300) [游凯超](https://www.zhihu.com/people/youkaichao) [为什么加速LLM推断有KV Cache而没有Q Cache?](https://www.zhihu.com/question/653658936/answer/3532237418) [DefTruth](https://www.zhihu.com/people/qyjdef) [全面解析 LLM 推理优化:技术、应用与挑战](https://mp.weixin.qq.com/s?__biz=Mzk0ODU3MjcxNA==&mid=2247486732&idx=1&sn=9887fdc9b6d1151aaf8c2b443d3c595d&chksm=c364ca49f413435f5f93e68195a38708cb195454272d044d7df368d78586958bef75f5075c8e&scene=178&cur_album_id=3209901770408230914#rd) AI闲谈 [关于 Mooncake 的碎碎念](https://zhuanlan.zhihu.com/p/705910725) [许欣然](https://www.zhihu.com/people/ultimategeeker) [刀刀宁:笔记:Llama.cpp 代码浅析(一):并行机制与KVCache](https://zhuanlan.zhihu.com/p/670515231) [刀刀宁:笔记:Llama.cpp 代码浅析(二):数据结构与采样方法](https://zhuanlan.zhihu.com/p/671761052) [分析 GPT 模型自回归生成过程中的 KV Cache 优化技术](https://zhuanlan.zhihu.com/p/668931470) [sudit](https://www.zhihu.com/people/sudit) [图解Mixtral 8 \* 7b推理优化原理与源码实现](https://zhuanlan.zhihu.com/p/691066049) [猛猿](https://www.zhihu.com/people/lemonround) [图解大模型计算加速系列:分离式推理架构1,从DistServe谈起](https://zhuanlan.zhihu.com/p/706761664) [猛猿](https://www.zhihu.com/people/lemonround) [图解大模型计算加速系列:分离式推理架构2,模糊分离与合并边界的chunked-prefills](https://zhuanlan.zhihu.com/p/710165390) [猛猿](https://www.zhihu.com/people/lemonround) [图解大模型计算加速系列:分离式推理架构2,模糊分离与合并边界的chunked-prefills](https://zhuanlan.zhihu.com/p/710165390) 猛猿 [大模型并行推理的太祖长拳:解读Jeff Dean署名MLSys 23杰出论文](https://zhuanlan.zhihu.com/p/660715870) [方佳瑞](https://www.zhihu.com/people/feifeibear) [大模型推理优化--Prefill阶段seq_q切分](https://zhuanlan.zhihu.com/p/658443665) kimbaol [大模型推理优化实践:KV cache复用与投机采样](https://mp.weixin.qq.com/s/M6bisR_rTHM-vyeOD9ILXA) 米基 \[阿里技术\](javascript:void(0)😉 [大模型推理优化技术-KV Cache](https://zhuanlan.zhihu.com/p/700197845) \[吃果冻不吐果冻皮\]( /people/liguodong-iot) [大模型推理加速:KV Cache Sparsity(稀疏化)方法](https://zhuanlan.zhihu.com/p/701580870) [歪门正道](https://www.zhihu.com/people/wu-guan-yu-72) [大模型推理加速:看图学KV Cache](https://zhuanlan.zhihu.com/p/662498827?utm_id=0) [看图学](https://www.zhihu.com/people/aaron-73-65) [大模型推理性能优化之KV Cache解读](https://zhuanlan.zhihu.com/p/630832593) [Young](https://www.zhihu.com/people/xiao-zhao-23-57) [大模型百倍推理加 速之KV Cache稀疏篇](https://zhuanlan.zhihu.com/p/704710823) [zhang](https://www.zhihu.com/people/zhang-20-3) [大白话解说Continous Batching](https://zhuanlan.zhihu.com/p/680123256) [Gnuey Iup](https://www.zhihu.com/people/ning-pei-yang) [大语言模型推理性能优化综述](https://zhuanlan.zhihu.com/p/656485997) [Young](https://www.zhihu.com/people/xiao-zhao-23-57) [微软 MInference:百万 Token 序列,10x 加速](http://mp.weixin.qq.com/s?__biz=Mzk0ODU3MjcxNA==&mid=2247487205&idx=1&sn=b320a056476ff1628108f68d5f952fc5&chksm=c364c9a0f41340b6c952950e2880c5de23ea8ea229cb0e97677739a0a8ab662c354e66926232&scene=21#wechat_redirect) [思辨|H2O: 从缓存视角解决LLM长上下文问题](https://mp.weixin.qq.com/s?__biz=MzUyMDk1OTc1OA==&mid=2247489215&idx=1&sn=ed5dba0759b24c38e4984667b6abf08a&chksm=f899f01270a179c659dc7ba1373e09f747844726a794f9af559735e97abdedb7704a9f5e57c2&mpshare=1&scene=1&srcid=0823e3q6xPRmqEcrB1GfBeQZ&sharer_shareinfo=645ae299948537831034d0ce9b3b2382&sharer_shareinfo_first=645ae299948537831034d0ce9b3b2382#rd) DSTS Center \[数据空间技术与系统\](javascript:void(0)😉 [打造高性能大模型推理平台之Prefill、Decode分离系列(一):微软新作SplitWise,通过将PD分离提高GPU的利用率](https://zhuanlan.zhihu.com/p/701772045) 哆啦不是梦 [打造高性能大模型推理平台之Prefill、Decode分离系列(一):微软新作SplitWise,通过将PD分离提高GPU的利用率](https://zhuanlan.zhihu.com/p/701772045) [哆啦不是梦](https://www.zhihu.com/people/18-28-80-57-94) [月之暗明的月饼:以KVCache为中心,分离式推理架构,LLM吞吐能力飙升75%](http://mp.weixin.qq.com/s?__biz=MzA5MTQxMTM0Nw==&mid=2247484450&idx=1&sn=e60991ee6522cdccf30c78dc1e85427b&chksm=907d82dea70a0bc8ee7c2abf29b8e0aa7af46e2671f62ac2bba607f5de3616cc8f5f21a10cfe&scene=21#wechat_redirect) [极限探索:LLM推理速度的深度解析](https://mp.weixin.qq.com/s/IjRBcFNlZcTgurjEZD1sTg) MLSys2024 [歪门正道:魔改KV Cache,让大模型无限输出!](https://zhuanlan.zhihu.com/p/696355658) [符尧博士-全栈Transformer推理优化第二季:部署长上下文模型-翻译](https://zhuanlan.zhihu.com/p/697244539) [强化学徒](https://www.zhihu.com/people/heda-he-28) [聊聊大模型推理中的分离式推理](https://zhuanlan.zhihu.com/p/706469785) [刀刀宁](https://www.zhihu.com/people/zzningxp) [聊聊大模型推理服务中的优化问题](https://zhuanlan.zhihu.com/p/677650022) [刀刀宁](https://www.zhihu.com/people/zzningxp) [阿杰:vLLM \& PagedAttention 论文深度解读(一)------ LLM 服务现状与优化思路](https://zhuanlan.zhihu.com/p/656939628) [高效生成:Splitwise论文解读](http://mp.weixin.qq.com/s?__biz=MzA5MTQxMTM0Nw==&mid=2247483889&idx=1&sn=2704c97da231eb5f0ae7403f443704db&chksm=907d870da70a0e1b9fad34e00fe88f4ad39014ac2ffbdf0cde19e065be642b9766963a8353f6&scene=21#wechat_redirect) [https://yaofu.notion.site/Full-Stack-Transformer-Inference-Optimization-Season-2-Deploying-Long-Context-Models-ee25d3a77ba14f73b8ae19147f77d5e2#d03b8a2a583246cea5a1fdeb8f583586](https://yaofu.notion.site/Full-Stack-Transformer-Inference-Optimization-Season-2-Deploying-Long-Context-Models-ee25d3a77ba14f73b8ae19147f77d5e2%23d03b8a2a583246cea5a1fdeb8f583586) [Longformer](https://arxiv.org/pdf/2004.05150v2.pdf) [Efficient Streaming Language Models with Attention Sinks](https://link.zhihu.com/?target=http%3A//arxiv.org/abs/2309.17453) [Massive activations in large language models. arXiv preprint arXiv:2402.17762, 2024. Mingjie Sun, Xinlei Chen, J Zico Kolter, and Zhuang Liu.]() [PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference](2405.12532) [MiniCache: KV Cache Compression in Depth Dimension for Large Language Models](2405.14366) [TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding](2404.11912) [ZigZagKV,KVCache80%的压缩率可以达到近乎无损的效果](https://mp.weixin.qq.com/s?__biz=Mzg5NDU2MDcyNQ==&mid=2247488362&idx=1&sn=ee5d1de21813e8db8f83a9025a284fa8&scene=21#wechat_redirect) 杜凌霄 [SnapKV: LLM Knows What You are Looking for Before Generation](https://link.zhihu.com/?target=http%3A//arxiv.org/abs/2404.14469) [FastGen:Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs]() \[**Prompt Cache** \]: [Modular Attention Reuse for Low-Latency Inference(@yale.edu)](https://arxiv.org/pdf/2311.04934) \[**Cache Similarity** \]: [Efficient Prompt Caching via Embedding Similarity(@UC Berkeley)](https://arxiv.org/pdf/2402.01173) \[**Shared Prefixes** \]: [Hydragen: High-Throughput LLM Inference with Shared Prefixes(@Stanford University)](https://arxiv.org/pdf/2402.05099v1) \[**ChunkAttention** \]: [ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition(@microsoft.com)](https://arxiv.org/pdf/2402.15220) [ClusterKV,接近无损的KV Cache压缩方法](https://mp.weixin.qq.com/s?__biz=Mzg5NDU2MDcyNQ==&mid=2247488234&idx=1&sn=18586ba3f0c4b4cee7a4ba88b8bd1b6d&chksm=c1b58c4961502f75d72c8a01fc0c76fded3ba924bfa65134dfe2e316d0332a123e95a3b55977&mpshare=1&scene=1&srcid=1208y4hg4sk6gzwZsqGGrfyU&sharer_shareinfo=6cc9038bb19cc28e8875cb49753992c3&sharer_shareinfo_first=6cc9038bb19cc28e8875cb49753992c3#rd) 杜凌霄 [图解大模型计算加速系列:分离式推理架构2,模糊分离与合并边界的chunked-prefills](https://mp.weixin.qq.com/s/_nm2Fwz2FlkcLuXnWDB9PA) 猛猿 [《基于KV Cache管理的LLM加速研究综述》精炼版](https://mp.weixin.qq.com/s/E-kvVIQUAkgMfQL8Hm7xEA) 常华Andy [\[KV Cache优化\]🔥MQA/GQA/YOCO/CLA/MLKV笔记: 层内和层间KV Cache共享](https://zhuanlan.zhihu.com/p/697311739) [DefTruth](https://www.zhihu.com/people/qyjdef) Venkat Raman. Essential math \& concepts for llm inference, 2024. URL u/essential-math-concepts-for-llm-inference#heading-insights-from-model-lat ency-amp-understanding-hardware-utilization-on-modern-gpus. Yao Fu. Challenges in deploying long-context transformers: A theoretical peak performance analysis. arXiv preprint arXiv:2405.08944, 2024 [GQA,MLA之外的另一种KV Cache压缩方式:动态内存压缩(DMC)](https://zhuanlan.zhihu.com/p/702408621) BBuf [TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding](https://zhuanlan.zhihu.com/p/693478388) [灰瞳六分](https://www.zhihu.com/people/detached_sextant) [TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding](https://arxiv.org/abs/2404.11912)