滑动窗口注意力与环形缓存：decode阶段显存优化深度解析

滑动窗口注意力在decode阶段使用环形缓存是一个看起来很细节、其实直接影响显存占用和生成速度的设计。它的核心思路是不要让KV cache无限增长而是用一块固定大小的内存只保留最近若干个token的键值新token进来时把最旧的token覆盖掉。这样在长序列生成时显存占用不会随生成长度线性涨decode单步的计算量也被限制在窗口大小内。这篇文章适合做大模型推理、在本地跑LLM、或者看推理框架源码时卡住的人。我会先解释为什么缓存会膨胀再拆环形缓存怎么工作最后给出一套排查和验证方法。1. 先搞懂滑动窗口注意力到底省了什么1.1 全注意力、窗口注意力与KV cache的增长问题在标准的Transformer decoder里每一步生成token时当前token都要和前面所有token计算注意力。从数学上看第 t 步的注意力覆盖范围是 1 到 t所以随着序列变长计算量和缓存量都在增长。这里有一个常被忽略的点注意力计算并不需要重新算一遍历史token的隐藏状态只需要把之前每一步算好的Key和Value缓存下来。这就是KV cache。KV cache的作用是避免重复计算思路看起来没问题可一旦序列很长KV cache本身就会变成显存杀手。全注意力的KV cache大小是随序列长度线性增长的。生成1000个token存1000份生成10000个token存10000份。如果模型层数多、头数多、维度大这个线性增长会非常快。很多人在本地跑长文本生成写到后面直接OOM最常见的原因不是模型参数占了多少而是KV cache越攒越多。滑动窗口注意力对这个问题的解法很直接当前token不再看全部历史只看最近 w 个token。w 就是窗口大小。既然只看最近w个那么历史中更早的Key和Value就没有必要继续保存。整个KV cache可以被固定在一段大小不超过w的内存区域里。这就是“省内存”的根本原理不是压缩而是主动丢弃。窗口注意力不是把信息压缩进一个更小的表示而是直接规定超出窗口的信息不再参与计算。1.2 窗口之外的历史信息为什么可以丢弃很多人第一次看到窗口注意力会产生疑问把前面的信息全丢光生成质量不会崩吗这个问题要分两层看。第一层是语言本身的局部性。大多数文本里的强相关性确实集中在邻近区域。比如一个词的意思主要受它前后几句话影响一段代码里的变量通常也在附近被引用。滑窗假设认为对绝大多数token来说跨越几千个位置去attend收益有限。这也是Longformer、Mistral这类模型敢采用局部注意力的底气。第二层是模型结构上的补偿。很多使用窗口注意力的模型并不是“只有窗口”而是会混合一些全局token。比如在序列开头放几个特殊token让这些全局token参与所有位置的注意力相当于给模型留了一条跨距离通信的通道。这样既保住了大部分局部能力又不会让缓存无限增长。但必须承认滑窗不是万能的。如果任务是跨越大段文本的精确信息提取比如从一篇长文档开头找一个细节然后要求结尾处复述纯滑窗模型很可能丢掉这个信息。这是滑动窗口注意力的边界不是bug。在用这类模型做长文档生成时要提前判断任务是否强依赖远距离信息。1.3 用一个小例子说明不同注意力的内存占用差异直接算一笔账可能比讲道理更清楚。假设窗口大小 w512模型只有一个注意力层每个token的Key和Value缓存占用为1个单位。再假设我们要生成到第10000个token。全注意力需要缓存10000个token的K和V占用10000个单位。滑动窗口注意力只需要缓存最近512个token占用512个单位。差距不是线性倍率而是随生成长度持续扩大。生成长度到100万时全注意力需要100万单位的缓存滑窗注意力仍然只需要512个单位。如果进入真实模型每层的缓存都要乘以层数。一个32层的模型窗口设为512那KV cache总量就是2 * 32 * num_heads * head_dim * 512 * 精度字节数。这个公式后面会细算。总之滑动窗口注意力本身就是为了让长序列生成时的缓存可控而环形缓存就是实现这种“固定大小缓存”最顺手的工程结构。2. decode阶段为什么不能照搬prefill的缓存方式2.1 prefill是一次性计算decode是逐token增量更新Transformer推理通常分成两个阶段prefill和decode。prefill阶段处理的是输入提示。这个阶段可以一次性并行计算所有输入token的Key和Value因为输入是完整给定的不存在“一个token依赖另一个token生成结果”的问题。计算完这整段输入后得到一份完整的KV cache然后进入decode。decode阶段不一样。每步只生成一个token把这个token追加到序列末尾然后下一步要用这个新token作为输入。这是一个典型的自回归过程每个新token都要与之前的KV cache做注意力生成自己的Key和Value再写进缓存。关键差异在于prefill是一次性批量写decode是每步只写一条。如果照着prefill的思路来管理decode缓存很容易写出这样的逻辑每生成一个token就把它追加到缓存尾部。代码很直观效果就是缓存数组越来越长。短序列没问题生成几百个token也没问题但一旦连续生成长文本内存就会一路走高。decode阶段真正需要的缓存管理能力是固定空间、持续覆盖、快速读取。这三个需求环形缓存刚好都能满足。2.2 如果不在decode复用缓存每步都要重新算前面的注意力有人可能觉得反正我不存缓存每步重新计算全部历史不就行了理论上可以实际上代价极大。假设你在decode第5000步如果没有任何缓存你就要重新计算前4999个token的Key和Value然后才能做一个token的注意力。再下一步又要把前5000个token全部重算一遍。这个计算量是平方级别的。别说是本地CPU就是A100也扛不住长时间这样跑。KV cache就是为了避免这种重复计算而存在的。用空间换时间每步只算新token的K和V历史K和V直接从缓存里读。这也是为什么KV cache的访问效率和更新效率会直接影响整体生成速度。如果不复用缓存滑动窗口注意力本身的意义就消失了。既然每步都重新编码全部历史那窗口限制只影响注意力范围不影响计算量。你会得到一个既不快、又不省内存的尴尬实现。所以一个合理的decode实现必然要考虑缓存是否被正确复用。2.3 缓存满了之后面临的问题覆盖还是全量重算滑动窗口注意力把KV cache上限定为w。这时有个现实问题当序列长度超过w缓存已经满了下一步写入新token时旧token怎么办最简单的做法是把缓存数组整体往左挪一位丢掉最左边腾出最后边写入新token。这个思路没问题但每次移动都是O(w)的数组拷贝。如果w是512、1024拷贝还算能忍如果w是4096或更大每生成一个token都要拷贝整段缓存速度下降会很明显。另一种做法是不移动数据只移动指针。用一个固定大小的数组按顺序循环写入。写入指针到达数组末尾后取模回到头部覆盖最旧的数据。这就是环形缓存。对比一下两种方式操作数组移动环形缓存写入新token先整体左移再写末尾直接写到当前位置指针取模时间复杂度O(w)O(1)内存分配可能触发重新分配初始化时一次性分配缓存是否固定否是decode阶段每步都要做一次写入这个操作会被执行成千上万次。哪怕单次差距只有几百纳秒累积起来也非常可观。更重要的是环形缓存让显存占用变成有上界的常量训练和推理框架可以提前分配好内存避免反复申请和释放。这也是为什么在实际推理框架里这是一个非常常见的实现选择。3. 环形缓存的工作原理与实现细节3.1 用数组头尾指针模拟固定长度缓存环形缓存本质是一个固定大小的数组配合两个指针读指针和写指针。在KV cache场景下我们只需要一个写指针因为读取时是读取整个窗口内的所有有效数据不是单点读取。初始化时分配一个长度为w的数组写指针指向0。每次写入一个token的Key和Value数据放到写指针指向的位置然后写指针加1。当写指针到达数组末尾时加1之后通过取模回到0。这样一来数组被反复循环使用。伪代码如下class RingBuffer: def __init__(self, capacity): self.capacity capacity self.keys [None] * capacity self.values [None] * capacity self.write_pos 0 self.current_size 0 def append(self, key, value): # 覆盖最旧位置 self.keys[self.write_pos] key self.values[self.write_pos] value if self.current_size self.capacity: self.current_size 1 self.write_pos (self.write_pos 1) % self.capacity这里有一个非常关键的细节current_size在前w步内是递增的之后保持为w因为新数据开始覆盖旧数据。实际使用时当前有效的数据就是数组里最近写入的w个位置可能不是从0到w-1连续排布。3.2 写入、覆盖、读取为什么环形结构天然适合“先进先出”环形缓存本质上实现了一种FIFO行为先写入的token在缓存空位满了之后会被下一个新token覆盖。它的顺序不是“数组从左到右”而是“从最旧有效位置到最新写入位置”。读取时需要区分两种状态缓存未满有效区域从数组0开始到写指针前一格。缓存已满有效区域从写指针当前位置开始到写指针前一格中间可能跨越数组末尾。这两种状态的处理方式不同但都可以通过一步简单的指针计算拿到顺序。实际工程中多数推理框架不会真的按顺序从缓存里逐条取数据而是直接用矩阵乘法和mask来读取整块缓存。不过理解顺序问题仍然是排查注意力错误的基础。环形结构天然适合“先进先出”的原因很简单覆盖位置是固定的。新token永远写到当前指针处指针永远按顺序前进。不用比较时间戳不用记录淘汰策略因为滑动窗口注意力里唯一需要淘汰的就是“最旧token”。这就把淘汰策略简化为一个指针移动。3.3 与KV cache结合每步生成一个token就写入一个位置覆盖最旧的位置把环形缓存放到decode流程里看整体流程是这样的prefill阶段输入提示一次性计算得到第一批K和V。将这些K和V写入环形缓存。如果输入长度不足w缓存未满如果输入长度超过w写入过程中就会开始覆盖。decode阶段每步生成一个新token计算它的K和V写入环形缓存当前指针指向的位置。计算当前token与缓存中所有有效K、V的注意力得到输出。这里有一个容易踩坑的点初始输入超过窗口大小怎么办如果模型本身支持滑窗prefill时通常不会把超出w的部分全部存入KV cache。很多实现会对输入做分块处理或者在输入序列过长时只保留最后w个token的K和V作为初始缓存。否则一开始就存入大量超过窗口的缓存后续decode又用不到白白浪费显存。还有一种做法是即使窗口外的K和V被丢弃位置编码仍然保留原来的绝对位置。这取决于模型如何构建mask和位置信息。涉及到具体的推理框架要去看它实际是怎么处理边界条件的。3.4 真正的工程细节位置索引、mask与批量隔离先看位置索引。环形缓存里数据存放位置和token的真实序号并不等价。比如窗口大小为8第10个token写入后数组里可能存在第3到第10个token但它们在数组里的位置可能是乱序的。更准确地说写入指针write_pos记录的是下一个写入位置。有效数据的顺序要以真实token顺序为准不能直接按数组下标排列。所以很多实现会额外维护一份位置索引或者通过数学计算来还原真实顺序。再看mask。注意力的mask也要跟着真实token顺序走。窗口注意力要求当前token只能attend到它之前、且距离不超过w的token。如果缓存数组是循环的但mask按固定下标生成就会导致注意力范围错误表现为生成的文本突然混乱、重复。批量decode场景更要注意隔离。假设一个batch里有多个序列每个序列都有自己独立的write_pos和缓存内容。不能用一个全局指针管理多个序列。要么为每个序列单独分配一段缓存要么用张量索引来区分。许多高性能框架会预分配一整块二维或三维缓存然后用每一行的起始位置和写位置来管理。代码层面看起来很简单的环形缓存真正落到多batch、多层的KV cache里复杂度和日志排查量都会上升。正因为如此理解它背后的数据结构比死记硬背某个框架的API更有价值。4. 参数、资源边界与常见坑点4.1 窗口大小w怎么选太小影响效果太大浪费显存窗口大小直接决定了模型能看到多远的上下文。太小生成内容容易失去前后连贯性太大KV cache占用的显存又会涨上去滑动窗口注意力的优势被削弱。常见模型里窗口大小通常在512到4096这个区间。例如部分长文档模型使用4096窗口很多滑窗模型使用1024或2048。具体选多少要看任务类型和可用显存。如果只是本地玩一玩建议先用模型默认的窗口大小。不要一上来就手动调大窗口因为很多模型的位置编码和mask是围绕默认窗口设计的。调大窗口可能不仅没有效果反而让计算变慢。如果是自己设计一个小模型做实验可以从256或512开始。先看输出质量能不能接受再看显存占用最后决定要不要扩大。不要只盯显存也不要只盯效果要两个指标一起看。4.2 显存占用如何估算一个可执行的公式使用滑动窗口注意力时KV cache的显存占用可以按这个公式估算KV cache大小 2 * 层数 * 注意力头数 * head维度 * 窗口大小 * 每个元素字节数乘2是因为Key和Value分别占一份。窗口大小就是有效缓存长度。每个元素字节数取决于数据类型FP32是4字节FP16/BF16是2字节INT8量化后是1字节。举个例子一个12层、12头、head维度64的模型窗口大小1024使用FP16推理2 * 12 * 12 * 64 * 1024 * 2 36,864,000 字节约 36 MB这个模型很小。如果换成70亿参数级别的模型层数32、头数32、head维度128窗口4096FP162 * 32 * 32 * 128 * 4096 * 2 1,073,741,824 字节约 1 GB这只是KV cache还不算模型权重和中间激活。所以窗口大小每翻一倍KV cache显存就翻一倍。理解这个公式之后再去看推理框架的日志就能知道显存到底花在哪儿了。4.3 常见问题排查输出变差、OOM、速度没有提升这一类问题在实践里非常常见。先看现象再按顺序排查不要急着改参数。现象一生成到一定长度后输出质量突然下降甚至开始重复。排查顺序先看窗口大小是不是设得太小。如果生成的文本长度本身就超过了窗口模型看不到前面关键信息质量下降是必然的。看位置信息是否保留。如果滑动窗口只保存K和V却没有保存token的绝对位置模型就会对距离产生错觉。看mask是否正确构造。很多环形缓存实现里如果mask只按数组下标生成没有考虑token真实顺序注意力就会计算出错。这里最容易误判的是以为模型“笨了”其实是缓存内容或者mask错了。现象二生成一段时间后OOM。排查顺序确认是否真的限制了缓存长度。如果代码里只是不断往列表尾部追加环形缓存根本没生效OOM只是时间问题。看显存曲线。如果显存随生成长度线性上涨基本可以推断KV cache没有固定。看是否有额外的内存碎片。频繁重新分配小块张量也可能导致显存碎片化但通常先检查KV cache策略。现象三环形缓存已经用了速度却没有明显提升。排查顺序看是否每步都对整个窗口做矩阵运算。如果注意力实现里的Q和K维度是w速度自然受窗口大小影响。看是否有额外拷贝。有些框架为了读取方便会在每步把环形缓存重新排列成连续数组这个拷贝开销可能抵消缓存收益。看是否还有其他瓶颈比如采样器、日志输出、文件写入、批量排队。不要一慢就怪缓存。4.4 批量decode场景下的环形缓存当多个序列同时生成时环形缓存的管理难度会上升。首先每个序列的写入指针都是独立的。因为不同序列长度不同、生成进度不同不能共享一个指针。很多框架会把KV cache预分配成形状为(batch_size, num_layers, num_heads, max_length, head_dim)的张量然后为每个batch位置维护自己的当前长度和写指针。其次不同序列的窗口内有效token数量可能不一样。短序列的缓存可能还没写满长序列的缓存已经开始覆盖。如果整个batch走同一个mask逻辑就必须区分哪些位置是有效缓存、哪些是填充。这里的处理方式通常是用一个mask矩阵把无效位置置为负无穷。第三显存分配策略也要调整。如果batch内所有序列都分配相同的最大窗口大小而窗口又都很大显存压力会成倍增加。这时可以考虑动态分配或者按序列长度分组。不过这会增加调度复杂度适合在推理框架层做不太适合自己写一个简单demo时硬搞。批量场景下的环形缓存真正要盯住的指标是batch内每个序列的显存占用是否可控以及是否存在跨序列的指针串扰。后者一旦发生经常表现为不同序列的输出互相混入排查起来非常痛苦。5. 从“为什么”到“怎么验证”一套实测路线5.1 先确认是否真的使用了滑动窗口很多人以为自己在用滑动窗口注意力实际看代码时发现KV cache其实还在无限增长。所以第一件事是确认模型和推理框架到底做了什么。可以检查这几处模型配置里是否有关键字段指向窗口大小比如sliding_window、window_size或者attention_window。推理日志里KV cache的显存占用是否稳定。如果显存曲线一直往上走说明没有真正生效。源码里寻找缓存写入逻辑。是直接append到列表还是预分配数组并使用取模索引。在不清楚的情况下不要只看参数名要看具体用法。5.2 用一个小实验对比全量缓存 vs 环形缓存这里不建议直接上大模型跑很长的生成成本高且不好定位。更稳妥的做法是写一个最小实验模拟两种缓存策略。可以构造一个很小的单层注意力模块输入一段随机token序列分别用“全量KV cache”和“环形缓存”两种方式做decode。然后用相同输入跑几轮记录生成到不同长度时的显存或内存占用。这种实验不需要完整的模型因为问题核心在缓存结构不在模型效果。你会发现全量缓存的内存占用随序列长度线性上涨环形缓存则稳定在一个固定值附近。在真实模型上也可以做类似验证但更建议先看两条曲线生成token数量与显存占用。如果显存是一条接近水平的直线说明环形缓存生效如果是一条斜线说明还没固定。5.3 需要重点看的指标验证这个设计是否合理不需要看太多花哨指标我一般只看三个。第一是显存/内存占用。这是环形缓存最直接的目标。看它是否随生成长度保持稳定。第二是每token生成速度。环形缓存本身就是为了避免数组移动和重新分配所以速度应该比“每步移动整段缓存”的方案更快。但如果你用高性能框架它可能已经把移动优化得很好差距不会特别大。这时要关注的是“长时间生成后速度是否稳定”而不是单纯的初次速度。第三是输出质量。生成到远超窗口长度时输出是否还会崩溃。如果窗口是512生成到2000多个token时输出依然连贯说明实现基本正确。如果生成到接近窗口长度时文本开始混乱优先怀疑mask或位置信息有问题。5.4 什么时候不需要环形缓存不是所有decode场景都需要环形缓存。以下情况可以不用生成的序列长度远小于窗口大小。这时候缓存根本不会写满用普通的追加数组就够了。模型本身没有采用滑动窗口注意力而是标准全注意力。强制加一个环形缓存反而会破坏模型能力。短prompt一次性生成。例如问答场景生成结果通常只有几十到几百个token缓存不会膨胀到不可控简单方案更容易维护。理解环形缓存的意义不是让你在每种场景都强行引入它而是当你面对长序列生成、显存受限、批量推理、框架选型时心里有一个判断标准什么时候该固定缓存什么时候该直接追加什么时候该考虑覆盖策略。我自己更倾向的做法是先跑通一个最小样例确认滑动窗口注意力的mask和位置信息正确再切换成环形缓存。不要一上来就同时调整缓存结构和参数。很多问题看起来复杂最后发现只是顺序错了应该先在内存里把数据结构和mask搞清楚再优化显存。