07｜KV Cache、FlashAttention与PagedAttention的内部过程
本章目录
本章预备知识
缓存保留可复用的旧计算；FlashAttention减少中间数据搬运；分页缓存改变存储管理。三者都可能改善效率，但不是同一个算法，更不能都理解为压缩语义。
本章各题列出所需基础。可先完成相应基础课，也可通过自检确认掌握程度。
进入深读前的小检查：对一层3个位置、每行4个数的K和V，先数24个数。再乘层数和每数的字节数，最后讨论是不是还保留全部历史。
术语带虚线时可点击速查。公式中的字母先看定义；不要用猜缩写替代推导。
这是容易被一句话带过、也最值得逐步推导的部分。先回答“哪些计算真的不会变”，再回答“怎么少搬数据”和“怎么分配内存”。三类优化不能混淆。
沿着对象与状态理解本章
这里有三个不同的存储问题：KV缓存跨轮保留历史；FlashAttention安排单次注意力的中间读写；PagedAttention管理物理缓存块。动画分别给出具体数据，避免将三种节约误认为都减少了相同计算。
下方图解的固定输入在“输入与适用条件”中声明。它用于解释本章机制，不与本章其他独立例子的数值暗中接续。
FlashAttention：一边读一边累计，不攒整张表
不把完整注意力表存下来，怎样仍算出同一输出？
想象分批称重汇总：每批到来时更新“最大刻度、总重量、内容加权和”，不用把所有中间单据再写一遍显存。最大刻度变化时，旧累计必须跟着换刻度。这个在线归一化过程解释了省读写如何与同一数学结果兼容。
输入与适用条件：独立单查询例：读取顺序C/A/B，分数[0,2,1]，对应V为[[5,5],[10,0],[0,20]]。为展示最大值变化而设，不接用F09候选顺序。
1 / 5 · 普通写法先想象整张分数表
窄屏可左右滑动图形，或点击“放大图形”。文字解释紧随图后。
在长序列中，把所有查询对所有Key的分数完整存下来会产生很大的中间表。FlashAttention改变分块与读写组织，目标不是随意丢弃低分位置。
本步状态由图形与说明共同给出。默认静止。用“下一步”看完解释再继续；播放到最后一步自动停止。键盘聚焦播放器后，可用左右键逐步查看、空格播放/暂停。图中颜色辅助区分，含义同时用文字注明。
为什么需要
大中间矩阵造成显存读写与容量压力。
怎样实现
分块计算并以在线softmax维护归一化统计，减少外部显存访问。
可能的代价与误解
实现复杂；数值顺序变化有舍入差异；不适用的硬件或形状收益可能很小。
如何判断是否做对
与完整softmax加权结果做容差对照，再在目标硬件实测而非把教学计数当加速倍数。
Q025｜为什么历史K/V可以缓存，而Q通常不缓存？请给出逐层解释。
核心回答
因果Decoder中，旧位置的隐藏状态只依赖它及其之前的Token；追加新Token不会改变这些旧状态。因此每层旧位置投影得到的K/V也不变。新位置只需要自己的Q，去查询该层历史K/V；旧Q产生的旧输出已经算完，通常无需再次使用。[CACHE][TR]
不是一句“缓存历史”就足够
用归纳法看。第0层的旧Token embedding和位置编码不变。假设第l层输入在旧位置不变，由于因果mask禁止看新增位置，旧位置的attention输入集合不变，attention输出也不变；FFN和归一化作用于原有相同输入，残差结果不变。因此第l+1层旧状态不变，归纳成立。
生成新Token时，它在第l层先得到该层归一化状态，计算q_new,k_new,v_new，把新的K/V附加到该层缓存，执行q_new @ K_allᵀ、softmax和加权求和，再通过Wo、残差、FFN进入下一层。直到全部层算完，才得到预测再下一个Token的logits。缓存不是一份跨层共享的文本向量，而是每一层独立的K/V数组。
若只保存最顶层K/V，下一Token在底层找不到所需历史表示。若保存旧Q，没有K/V，旧Q也不能让新位置恢复历史内容。常见缓存布局为K/V各[B,H_kv,T,Dh]，内核可能采用不同内存排布，但语义轴不应混淆。
用因果依赖逐层证明旧K/V为什么不变
输入与条件。两层模型先处理Token A、B，再追加C。所有层都是固定权重的因果层，归一化与FFN不跨位置混合。
| 步骤与对象 | 内部动作 | 中间结果 / 新状态 |
|---|---|---|
| 1第0层输入 | A、B的Embedding和位置未变 | 旧输入相同。 |
| 2第0层可见集合 | A只看A，B只看A/B | 追加C不在旧查询可见范围内。 |
| 3第0层结果 | 旧输入和可见集合都不变 | A、B的第0层输出不变。 |
| 4第1层输入 | 来自第0层旧结果 | 因此第1层A、B的K/V也不变。 |
| 5新C的过程 | 在每层只算C的Q/K/V，并查该层A/B缓存 | 不是只查顶层。 |
| 6反例 | 把B改成D后仍复用B及其后缓存 | 因果依赖前提已破坏，需要从改动处重算。 |
把过程连起来。缓存成立是依赖关系结论，不是因为历史发生过就自动永久有效。
边界与反例。双向模型、跨位置统计或动态位置变换需另行证明；本例限标准固定因果层。
依据与适用范围：[CACHE]。示例输入和场景由本书明确设定。
追问：编辑Prompt第100个Token后，哪些缓存还能用？
在相同权重、模板、位置与运行约束下，第0至99个Token对应的缓存仍可能复用；第100个及其后继依赖已改变，需要重算。这里说的是模型实际Token前缀，不是人眼看到的相同字符。更换系统指令、工具schema、adapter或位置规则可能在更早处改变前缀。不能只把编辑的那个Token缓存改掉，而保留其后所有上下文化状态。
追问：为什么训练通常不直接复用跨优化步的KV？
权重更新后，原K/V不再是新参数对原前缀的计算结果；而且训练需要相应计算图和梯度。将旧缓存用于新权重前向会破坏目标一致性。推理时no_grad和固定参数是常见缓存前提，不应把推理优化直接套到普通训练循环。
追问：哪些模型不能按这个证明直接缓存？
双向编码器中旧位置可依赖后面Token，追加输入会改变旧状态；某些跨Token归一化、动态位置变换或修改历史的架构也需另行分析。滑动窗口、循环状态和混合模型有各自状态机制。本章证明适用于固定前缀、固定参数的标准因果层，不能由此断言所有模型都有同样缓存。
Q026｜KV缓存到底节约多少计算，又占多少显存？
核心回答
它避免对整个历史反复执行投影、FFN和旧查询的attention；新查询仍需读取历史K/V。一次decode的attention仍随已缓存长度增长，整个生成过程也不是无条件线性。显存由层数、batch、KV头数、长度、头维度和dtype决定。
把复杂度拆开
设模型宽度d、前缀P、生成M个Token。忽略常数，单层对长度T全量前向的线性投影/FFN约O(Td²)，attention约O(T²d)。没有缓存，每生成一步都重做越来越长的前缀。
使用缓存后，Prefill仍处理P个位置；每个decode只做新Token的投影/FFN约O(d²)，以及新Q对T个K/V的attention约O(Td)。M步attention合计近似O((PM+M²)d)，另有O(Md²)的层内计算。相比反复全量前向大幅减少，但不能表述为“KV让attention从平方永远变成线性”。讨论复杂度必须说明是单步还是整段、固定P还是P与M一起增长。
缓存字节数的常用估算是：
KV bytes = 2 × L × B × T × H_kv × Dh × bytes_per_element
前面的2来自K和V。教学例：L=32，B=1，T=32768，H_kv=8，Dh=128，FP16每元素2字节，结果是4GiB；每多一个Token约增加128KiB。若改为32个KV头而其他不变，就是16GiB。这个估算不包括权重、页面碎片、workspace和服务元数据。实际设备能接多少并发不能只拿总显存除以这个数字。
把一条长上下文的K/V字节数从单位算出来
输入与条件。L=32层，B=1，T=32768，Hkv=8，Dh=128，FP16每数2字节。
| 步骤与对象 | 内部动作 | 中间结果 / 新状态 |
|---|---|---|
| 1一份K | 32×1×32768×8×128个数 | 1,073,741,824个元素。 |
| 2换字节 | K元素数×2 | 2,147,483,648字节，即2GiB。 |
| 3加V | K和V大小相同 | 合计4GiB。 |
| 4增一个Token | 2×32×1×8×128×2 | 131,072字节，即128KiB。 |
| 5改MHA对照 | KV头从8增到32，其余不变 | 缓存16GiB，但这不是整个模型显存。 |
把过程连起来。公式前面的2来自K/V两份，最后的2来自FP16字节数，二者不能混在一起忘记。
边界与反例。按GiB=2³⁰计算；不含权重、页表碎片、工作区和服务开销。
依据与适用范围：[CACHE]。示例输入和场景由本书明确设定。
追问：为什么长上下文decode可能受显存带宽限制？
新查询计算量相对有限，却需要逐层读取大量历史K/V；若计算强度低，搬运速度限制总时间。增加batch可以提高某些权重读取的复用，但也增加KV总量与服务延迟。GPU峰值算力高不意味着单请求长上下文一定更快，需看实际内核和内存访问。
追问：KV量化与权重量化是同一种优化吗？
量化对象不同。权重量化主要减少模型参数存储和读取；KV量化主要减少随上下文与并发增长的动态状态。K的误差改变attention分数，V的误差改变聚合结果，两者都可能影响长程信息利用。应评测不同长度和任务，而不只看短答案质量。
追问：Prompt Cache与KV Cache什么关系？
KV是计算状态形式；Prompt Cache是服务对重复前缀复用计算的一种产品/系统能力，底层可能利用KV等状态。缓存命中需要满足服务规则，并非会话持久化就自然命中。API里的cached tokens计费与具体KV物理驻留不是一一可见映射。[PCACHE]
Q027｜FlashAttention为什么省内存？不存完整注意力矩阵怎样算softmax？
核心回答
FlashAttention把attention计算按块组织，利用更快的片上存储，避免将完整T×T分数和概率矩阵反复写回高带宽显存。它用在线softmax维护每个查询的归一化统计，计算的是同一类精确attention，不是删除低权重项的近似检索。[FLASH]
推导在线softmax合并
对一个固定查询，一批分数为s_j、值向量为v_j。维护三个量：当前最大值m，指数和l，以及未归一化加权和u。
m = max(s_j)
l = Σ exp(s_j−m)
u = Σ exp(s_j−m) v_j
最终输出 o = u/l
旧块统计为m_old,l_old,u_old，新块分数最大值为m_b。令m_new=max(m_old,m_b)，则：
α = exp(m_old−m_new)
l_new = α l_old + Σnew exp(s_j−m_new)
u_new = α u_old + Σnew exp(s_j−m_new) v_j
关键是当最大值变化时，旧统计必须同步缩放。否则不同块使用不同归一化基准，结果就不等于全量softmax。对每个查询遍历K/V块，最后除以l即可。矩阵内核会并行处理查询块，并融合mask、缩放等操作；训练反向可根据保存的少量统计重新计算所需块，而非保存全部概率矩阵。
举例：旧块只有分数1和值2，则m=1,l=1,u=2。新块分数3、值10，合并后m=3，l=exp(−2)+1，u=2exp(−2)+10，输出约9.0464。若忘记给旧u乘exp(−2)，输出会明显错误。
两块分数如何不存完整概率表也得到同一结果
输入与条件。固定一个查询，第一块只含score=1、value=2；第二块score=3、value=10。维护最大值m、指数和l、未归一化加权和u。
| 步骤与对象 | 内部动作 | 中间结果 / 新状态 |
|---|---|---|
| 1读第1块 | m=1，l=exp(0)=1，u=2 | 临时输出2。 |
| 2读第2块 | 新最大值m=3 | 旧统计需乘exp(1−3)=exp(−2)。 |
| 3合并归一化分母 | l=exp(−2)+1 | 约1.135335。 |
| 4合并加权分子 | u=2exp(−2)+10 | 约10.270671。 |
| 5最后相除 | u/l | 约9.046377，与对[1,3]做全量softmax再加权相同。 |
| 6错误对照 | 只更新m却不给旧u、l缩放 | 不同块不在同一指数基准，结果错误。 |
把过程连起来。在线统计保留的是计算所需的充分数值状态，不是对文本做摘要。FlashAttention用分块与融合减少中间读写。
边界与反例。这只实现一个查询的在线softmax核心，不是GPU FlashAttention内核，也不证明任何加速倍数。
依据与适用范围：[FLASH]。示例输入和场景由本书明确设定。
追问：既然是精确算法，结果为什么可能不逐bit相同？
实数数学上等价，浮点加法次序、累加精度和内核融合会造成微小误差。“exact attention”主要是区别于稀疏/低秩近似，不是保证不同硬件逐bit一致。验证应设置合理数值容差，并区分随机性与数值舍入。
追问：FlashAttention能代替KV Cache吗？
不能。前者优化一次attention如何计算与搬数据，后者避免下一步重算历史状态。两者通常互补。没有KV的快速attention仍会重复计算；有KV但低效内核仍可能搬运过多数据。
追问：它是否把时间复杂度从O(T²)变成O(T)？
标准密集attention的配对数量仍是平方量级，只是减少内存流量和中间存储。不同硬件/长度下加速比不同。把论文的某个实验加速倍数直接搬到自己的模型与GPU是不严谨的，应分Prefill/Decode做实际基准。
Q028｜PagedAttention的“分页”具体分页什么？为什么能提升并发？
核心回答
分页把每条序列逻辑上连续的KV分为固定大小块，映射到物理KV块池。序列增长时按需分配，不必预留一大段连续显存；多个请求和共享前缀可以更灵活地使用内存。它主要改变KV的存储管理与访问映射，不改变模型看到的Token顺序。[PAGED]
一个块大小为4的例子
请求A已有10个Token，需要3个逻辑块。块表可能是A:[7,2,19]，分别指向物理块7、2、19。逻辑位置5位于逻辑块1、块内偏移1，实际访问物理块2的偏移1。A再追加第11、12个Token可填满块19；第13个Token才需要新块。
请求B可以使用物理块3、15，完全不必与A连续。请求结束时释放引用，块回到池中。最后一个块可能有内部碎片，但通常比每请求按最大长度预留大段空间更可控。内核借助块表按逻辑序列访问K/V，必须保持位置、长度和mask正确。
若A、B有完全相同前缀，可在满足实现规则时共享前缀块并增加引用计数。对共享的未满末块写入新Token时可能需要copy-on-write；否则A的追加会污染B。前缀哈希也必须包含足以区分模型、Token前缀、适配器等有效身份的信息，不能只按人类输入文本共享。
页表把逻辑位置5映射到哪个物理格子
输入与条件。每块容纳4个Token。请求A已有10个位置0—9，逻辑块映射为 [7,2,19]。
| 步骤与对象 | 内部动作 | 中间结果 / 新状态 |
|---|---|---|
| 1定位逻辑块 | 位置5整除4 | 逻辑块1。 |
| 2定位偏移 | 5模4 | 块内偏移1。 |
| 3查页表 | A的逻辑块1对应物理块2 | 访问物理块2、偏移1，而不是物理位置5。 |
| 4追加两个Token | 新位置10、11 | 填满物理块19的偏移2、3。 |
| 5再追加位置12 | 进入第4个逻辑块 | 这时才分配新块。 |
| 6共享反例 | A/B共享未满末块后各写不同Token | 需按实现复制共享块或分离写入，不能互相覆盖。 |
把过程连起来。页表改变存放位置，不改变模型逻辑上的Token顺序与因果位置。
边界与反例。物理块号是假想编号；块大小、共享粒度、内核访问策略依实现而定。
依据与适用范围：[PAGED]。示例输入和场景由本书明确设定。
追问：为什么块越小不一定越好？
小块减少尾部碎片，却增加块表、寻址和管理成本；大块减少元数据但浪费更多尾部空间。最佳大小取决于长度分布、内核和设备。分页不是免费的操作系统魔法，而是具体数据布局的取舍。
追问：内存不够时，服务可以怎样处理？
可以暂停/抢占部分请求、把状态移到其他存储、重新计算前缀，或拒绝新请求。每种策略影响尾延迟和吞吐。单纯提高最大并发可能触发反复换出/重算，导致整体速度变差。应观察KV占用、抢占次数、排队时间和成功完成率。
追问：PagedAttention是否等于Prefix Caching？
分页提供适合共享与管理的基础设施，前缀缓存决定何时识别、保留和复用相同前缀。二者相关但不同；可以分页却不跨请求缓存，也可以用其他内存管理方式实现前缀复用。不要把一个机制的收益都归给另一个名字。