@_avichawla: 一个棘手的LLM面试题:你在vLLM上部署推理模型,长序列时GPU内存总是不够用。于是你加入KV缓存压缩,驱逐了90%的缓存token。显存占用依旧,GPU仍然内存不足。为什么?
摘要
解释了为什么在vLLM上部署推理模型时,驱逐90%的KV缓存token无法释放GPU内存,原因是分页注意力碎片化。同时介绍了NVIDIA的TriAttention解决方案,可实现2.5倍加速和10.7倍内存缩减。
查看缓存全文
缓存时间: 2026/06/27 13:55
一个棘手的LLM面试题:
你在vLLM上部署一个推理模型,但在长追踪序列中GPU显存持续耗尽。
于是你添加了KV缓存压缩,并驱逐了90%的缓存token。
VRAM使用量保持不变,GPU仍然耗尽显存。
为什么?
(答案如下)
驱逐90%的KV缓存几乎无法释放它所使用的内存。
这听起来违反直觉,但它直接源于生产服务器当前存储缓存的方式。
KV缓存随着模型生成的每个token而增长。每个token都会在每一层追加其key和value向量,并且在生成过程中不会释放任何内容。
这是推理模型的主要内存开销。
如果一条32K token的思维链缓存了约32K token的KV向量,那么一个采用4位权重的Qwen3-32B模型在24GB GPU上大约在24K token时就会耗尽显存。
一个显而易见的解决方案是保留重要token并丢弃其余部分,因为注意力足够稀疏,可以这样做。
但这并不能解决内存问题。
原因在于分页注意力,它是vLLM和大多数生产服务器背后的内存管理器。
在底层,它将GPU内存划分为固定的物理块,每个块容纳大约16个token的KV信息。
只有当块内的每个槽位都为空时,该块才会返回给分配器。
由于驱逐逻辑根据重要性选择token,而这类token分散在各个块中……
……因此,尽管进行了驱逐,几乎每个块都会保留至少一些幸存token。
例如,如果该逻辑在1000个块中驱逐了16k个token中的14k个,那么很可能每个块仍然有一个token。
这意味着分配器几乎什么也没释放。
将新token放置到那些已释放的槽位中并不理想,因为它破坏了缓存的布局。
假设第16,001个token到来,它被放置在第40个token曾经占用的槽位中。缓存现在读取位置38,然后是16,001,然后是41,因此缓存不再按token顺序排列。
注意力仍然可以从中计算出正确答案,但前提是每个槽位现在都附带一个单独的记录,标明它实际持有的位置。
这会引入另一种记账开销,而按顺序布局本身就避免了这种开销。
因此,缓存在逻辑上小了90%,但在物理上仍然是相同大小。许多压缩结果忽略了这一点,因为它们是在预先分配的连续张量上测量的,而不是在分页服务器上。
还有另一个问题。
驱逐方法通过查看注意力分数本身来决定保留哪些token(正如预期的那样)。
但生产环境中使用的快速注意力内核(如FlashAttention)从不会保存这些分数。
它们以小片段计算注意力,并在进行时丢弃完整的分数网格,这也是它们速度快的原因。
因此,驱逐方法所需的精确信号在内存中不可用。变通方法是回退到 eager attention 并构建完整的矩阵,这放弃了FlashAttention本应提供的速度。
NVIDIA发布了一种名为TriAttention的方法来解决这两个问题。
它永远不需要注意力分数。相反,它根据模型在应用RoPE之前的key和query向量的几何形状来对token进行评分,这些向量位于稳定的簇中。
对于内存问题,它每解码128个token就执行一次压缩过程。
幸存token向前滑动以填补驱逐造成的空洞,因此整个块变空并返回给分配器,同时缓存保持token顺序。
在长推理轨迹上,该方法匹配全注意力精度,同时解码速度快2.5倍,KV内存使用量少10.7倍。
KV缓存压缩是一个重大的基础设施问题。决定其是否有效的数字是被释放块的数量,而不是被驱逐token的数量。
你可以在以下位置找到NVIDIA的文章:https://research.nvidia.com/labs/eai/blogs/kv-cache-compression-and-its-infra-problems/
我写了一篇关于KV缓存工作原理的原理解析。它详细介绍了为什么模型要存储key和value,为什么缓存会随着每个token增长,以及有KV缓存和无KV缓存情况下LLM生成速度的比较。
请阅读下文。
KV缓存压缩及其基础设施问题
来源:https://research.nvidia.com/labs/eai/blogs/kv-cache-compression-and-its-infra-problems/ 考虑在本地GPU上使用一个强大的推理模型运行OpenClaw。代理接到了一个简单的任务——阅读几份文档,编写一份周报——却在半途崩溃了。不是因为任务困难,而是因为GPU显存耗尽了。
这篇博客从基础设施的角度审视KV缓存压缩:不仅关注算法能做什么,还关注它们是否真的能在生产中运行。整个领域都在致力于压缩KV缓存,以便长推理能够适应GPU显存;其中大部分在实验室中有效,但在实际部署中效果差得多。这一差距归结为两个论文几乎从不讨论的基础设施问题。本文概述了该领域的主要思想,追溯了它们与这些问题的冲突之处,并展示了RoPE之下隐藏的几何属性如何被用来解决这两个问题。
KV缓存压缩的难点不在于选择保留哪些token——而是与生产基础设施的两次冲突。现有方法需要历史注意力分数,但FlashAttention内核从不将这些分数写入GPU内存。分页注意力仅在某个块完全不包含KV信息时才能回收内存——然而在token驱逐后,幸存者使几乎所有块都部分占用,因此“释放的”内存实际上永远不会回来。
目录
- 背景:KV缓存与长推理中的内存耗尽 (https://research.nvidia.com/labs/eai/blogs/kv-cache-compression-and-its-infra-problems/#section-1)
- 现有方法——以及两个基础设施问题 (https://research.nvidia.com/labs/eai/blogs/kv-cache-compression-and-its-infra-problems/#section-2)
- 两个基础设施问题的系统解决方案 (https://research.nvidia.com/labs/eai/blogs/kv-cache-compression-and-its-infra-problems/#section-3)
- 视频生成中的KV缓存基础设施 (https://research.nvidia.com/labs/eai/blogs/kv-cache-compression-and-its-infra-problems/#section-4)
- 总结思考 (https://research.nvidia.com/labs/eai/blogs/kv-cache-compression-and-its-infra-problems/#section-5)
1. 背景:KV缓存与长推理中的内存耗尽
当Transformer生成一个token时,它会计算一个query、一个key和一个value向量;新的query会关注所有之前的key和value,以读取模型自身的历史。KV缓存使这一点变得可行:每个token的K和V只计算一次并保存,之后的每一步都会重用它们,而不是重新计算——每一步的工作量从O(n²)降到了O(n)。但它永远不会缩小:每个生成的token都会在所有层和注意力头上追加新的K和V行。
算术结果很糟糕。Qwen3-32B采用4位量化权重——一个人们实际部署的配置——在24 GB GPU上大约在生成24,000个token后崩溃并出现内存不足错误,而推理模型在困难问题上产生的32K token轨迹远高于这个数字(图1)。
KV缓存增长消耗GPU显存(示意图) 一个水平条代表GPU显存,示意且未按比例绘制。约三分之一的固定部分容纳模型权重加上运行时;剩余三分之二随着KV缓存的增长而填充,直到达到条右端的虚线内存不足线。在条下方,一个LLM框将输出token发射到一个历史框中,该框直接对齐于KV缓存部分下方;历史框逐个填充token方块,每新增一个token,上方的KV缓存部分同步延伸。 GPU存储 (示意图 — 未按比例绘制) 权重 + 运行时 (固定) KV缓存 每生成一个token就增长 → OOM LLM 生成中… 输出token 历史 — 每个token的K/V都留在这里 每个生成的token都将其key和value追加到历史中——这个历史就是KV缓存。在生成过程中没有释放任何东西,因此长输出会将条推向OOM墙。 图1. 示意图,未按比例绘制:模型生成的每个token都将其key和value追加到运行历史(即KV缓存)中,因此缓存内存随输出长度增长,而权重保持不变。让生成持续足够长的时间,仅缓存就会导致GPU内存耗尽。 压缩之所以可能,是因为注意力是稀疏的。一小部分token收集了绝大多数的注意力权重;一个永远不被关注的token原则上可以从缓存中删除,而不会改变模型的输出。问题是:哪些token可以安全删除,以及何时删除?
2. 现有方法——以及两个基础设施问题
该领域于2023年从一个观察起步:注意力高度非均匀。模型将不成比例的注意力权重分配给最开始的几个token——“注意力陷阱”,并非语义重要但始终存在(StreamingLLM[2])——而大约20%的token收集了总权重的80%,这些“重击者”倾向于随时间保持重要性(H2O[3];Scissorhands[4])。标准配方随之而来(图2):始终保留陷阱,始终保留最近token的滑动窗口,并从中间保留一定预算的重击者。在长文档任务(如问答、摘要、代码检索)上,该配方效果很好。
StreamingLLM:注意力陷阱加滑动窗口 一行token。前两个是注意力陷阱,始终保留。一个窗口轮廓覆盖了最近的四个token。随着生成进行,新token出现在右侧,窗口向前滑动,掉出窗口的token被驱逐到虚线幽灵槽中。陷阱永不移动,永不驱逐。 StreamingLLM — 保留陷阱,保留滑动窗口,驱逐其余部分 ← 最旧 token 序列在生成过程中增长 最新 → 陷阱 KEEP ✓ 陷阱 KEEP ✓ 滑动窗口 — 保留 滑动窗口 — 保留 滑动窗口 — 保留 滑动窗口 — 保留 滑动窗口 — 保留 滑动窗口 — 保留 滑动窗口 — 保留 注意力陷阱 — 始终保留 在滑动窗口中 — 保留 掉出窗口 — 驱逐 图2. StreamingLLM仅保留两种token:开头的一小部分注意力陷阱,以及最近token的滑动窗口。随着生成进行(动画效果),窗口向前滑动,所有掉出窗口的内容都被驱逐——内存保持恒定,但窗口外的信息永久丢失。 该配方留下了一个关键问题:一个方法如何知道哪些中间token是重击者?定义最大一类压缩方法的答案是读取模型自身的历史注意力分数。每个解码步骤无论如何都会计算对整个历史的注意力,而这些分数是模型实际使用哪些缓存token的运行记录。H2O[3]是典型例子:它为每个缓存token维护一个累积和,即该token在所有解码步骤中收到的注意力总和,并在每一步后驱逐得分最低的token,以将缓存保持在固定预算内(图3)。
H2O:累积历史注意力评分 顶部的查询框关注八个缓存token。每步解码,每个token的分数条都会增加它刚收到的注意力。几个步骤后,累积分数高的token被保留;分数低的token被驱逐到虚线幽灵槽中。 H2O — 每个解码步骤都会增加每个token的累积注意力分数 新token 步骤t 步骤t+1 步骤t+2 陷阱 KEEP ✓ 驱逐 ★ KEEP ✓ 驱逐 驱逐 ★ KEEP ✓ 驱逐 最近 KEEP ✓ 累积注意力分数(每步增长) 最低分数 — 每步驱逐 图3. H2O的“历史注意力”评分。每个解码步骤,最新token关注整个缓存(闪烁箭头),每个缓存token的累积分数会增加它刚收到的注意力。随着解码进行,得分最低的token每步被驱逐一个,将缓存保持在固定预算内。这些每步分数正是FlashAttention从不写入内存的内容——下面的基础设施问题1。 SnapKV[6]是同一思想最具影响力的改进。它不连续累积分数,而是一次性评分(图4):提示词最近的W个token形成一个“观察窗口”,它们对整个历史付出的注意力决定了——在预填充结束时,一劳永逸地——哪些token留下。这消除了每步记账,并避免了累积评分对存在时间更久的token的偏向。该类别还有基于相同成分的许多进一步变体:Scissorhands[4]和TOVA[5]替换了解码时的评分信号;PyramidKV[7]和Ada-KV[8]将一次性预算按层和头划分。
但观察窗口不能做得太大。因为RoPE——将每个token的位置编码到其query向量中的旋转——根据其位置定向每个query,只有最近的查询(经验上约25个)反映了模型实际关注的位置。因此,任何在推理过程某个阶段没有引起注意的信息都会被驱逐——即使后续推理依赖于它。
观察窗口快照评分 最近的W个token构成一个观察窗口。它们对历史付出的注意力识别出重击者token。高得分token在预算B内保留;低得分token被驱逐。 观察窗口快照评分(SnapKV风格) ← token历史(最旧左 → 最新右)→ 陷阱 KEEP ✓ — 驱逐 — 驱逐 ★ KEEP ✓ — 驱逐 — 驱逐 ★ KEEP ✓ — 驱逐 — 驱逐 W KEEP ✓ W KEEP ✓ W KEEP ✓ 观察窗口(W个最近token) 注意力分数↑ 图4. 基于快照的方法(SnapKV风格)通过测量每个缓存token从最近查询窗口收到的注意力来分配其重要性分数;不在前B预算内的token被驱逐。 整个类别——无论是像H2O那样连续评分,还是像SnapKV那样一次性评分——都依赖于观察注意力分数;在上述方法中,只有StreamingLLM的固定陷阱加最近规则不需要。这个共同需求,而不是任何特定启发式的质量,以两种方式与生产基础设施发生冲突。
基础设施问题1:FlashAttention不暴露注意力分数。 生产推理运行在FlashAttention上,它将注意力计算分块到SRAM(GPU小而快的片上内存)中,并且从不将N×N分数矩阵具体化到HBM(GPU的主内存)中。这就是它速度快的原因——也意味着驱逐方法想要观察的分数从未写入压缩代码可以读取的任何地方。 这在方法需要完整注意力历史时打击最严重。H2O风格的累积评分需要每个解码步骤的注意力分数——如果不观察每一步计算出的分数,就无法更新运行总和。参考H2O实现通过回退到 eager attention,具体化完整的
相似文章
内存
解释了为什么由于KV缓存随上下文长度和并发用户数扩展,LLM推理越来越受内存带宽限制,以及像vLLM和PagedAttention这样的系统如何提高内存利用率。
@tom_doerr: 在单个4GB GPU上运行70B大语言模型 https://github.com/lyogavin/airllm
AirLLM是一个开源工具,优化推理内存使用,无需量化即可在单个4GB GPU上运行70B大语言模型,并支持在8GB显存上运行405B模型。
本地LLM CPU用户……你们做任何事情要花多长时间?
关于在CPU上本地运行大语言模型性能的讨论,特别是大上下文尺寸的情况,以及显存限制带来的挑战。
@techNmak: 你的LLM推理正在消耗50%的计算资源在已经完成的工作上。如果你正在运行RAG或多轮对话,……
LMCache是一个开源库,它使KV缓存持久化并可在请求之间共享,消除了RAG和多轮对话工作负载中的重复计算,实现了高达15倍的吞吐量提升和3-10倍的首令牌时间减少。
@KL_Div:随着生成长度增加,LLM 占用的 GPU 内存持续攀升。能否在几乎不牺牲精度的前提下,让 GPU 内存占用保持恒定?
IceCache 通过“动态连续索引”(DCI)技术,在超长生成任务中将 GPU 内存占用压到恒定,且精度损失极小。