KV Cache 让自回归推理保留过去 token 的 K/V,避免每一步重新跑完整前缀。随着上下文和并发增加,这些状态既占据设备内存,也在每一步 Decode 被反复读取。把缓存装得下、让它读得快、让请求复用已有状态,是三个不同的系统问题。

1. 一条请求怎样建立和读取缓存

在 decoder-only Transformer 中,因果 mask 保证历史 token 在同一次生成中看不到未来 token。固定权重、位置编码和可见历史时,每层已经算出的历史 K/V 可以直接复用。Vaswani 等,NeurIPS 2017

  • Prefill 批量计算 prompt 的 Q/K/V,按层写入 K/V。最后一个有效 prompt 位置的 logits 通常用于采样第一个输出 token。
  • Decode 将刚采出的 token 作为下一次前向的输入。在每层计算它的 Q/K/V,将新 K/V 写入该位置,用当前 Q 读取历史和当前位置,再预测下一个 token。历史 Q 在这一步无须保存。

对长度为 \(T\) 的当前序列,使用缓存后,每步每层只为新增位置做投影,密集 Attention 仍需扫描 \(T\) 个可见位置。处理固定长度 prompt 后,再运行 \(G\) 次单 token 前向,Attention 扫描长度之和随 \(G T + G^2\) 增长。缓存避免了旧位置的重新投影与旧 query 的重复计算,却保留了当前 query 对历史的依赖。融合扫描如何调度,见《Attention 技术演进》。

2. 先算容量,再算 Decode 每步读多少

假设各层 head 数、维度和精度相同,K/V head 维度均为 \(d\)。层数为 \(N\),KV head 数为 \(H_{kv}\),每元素 \(b\) 字节,第 \(i\) 个请求缓存长度为 \(T_i\)。按无共享的有效载荷计算:

\[B_{\rm KV}=2N H_{kv}db\sum_i T_i.\]

前面的 2 分别对应 K 和 V。按页分配时还要加上尾块空槽、页表、量化元数据和布局 padding。对混合全局层、窗口层或不同精度的模型,应按层累加。做 tensor parallel 时,还要按实际 head 分片或复制规则算单卡容量。MQA 的一个 KV head 如果在多个 rank 复制,简单除以卡数会低估内存。

沿用上一篇的 32 个 Q head、8 个 KV head、BF16,再取 32 层、head 维度 128。每 token 的 KV 有效载荷是 128 KiB,一条 8,192 token 请求占 1 GiB。相同配置如果用 32 个 KV head,则是 4 GiB。理想情况下每个 KV 元素每步从 HBM 读一次,那么一条 8K GQA 请求的一次 Decode 就要读约 1 GiB。若设备给这一条路径的有效带宽为 1 TiB/s,仅这部分读取约需 0.98 ms。这个数字是设定带宽下的算术示例。权重、索引、内核间同步和并发调度都会影响实际延迟。

3. 从连续缓冲到 PagedAttention 和前缀复用

连续缓冲区让索引简单,但请求长度和结束时间各不相同。为每条请求按最大长度预留连续空间,会在请求尚短时占住大量空槽。PagedAttention 把序列划成固定 token 数的逻辑块,用块表定位设备内存池里的物理块。设块大小为 \(P\),位置 \(t\) 的物理块与块内偏移为

\[p=\mathrm{blockTable}_i[\lfloor t/P\rfloor],\qquad r=t\bmod P.\]

逻辑序列可连续增长,物理块按需分配和回收。独占的最后一块最多浪费 \(P-1\) 个 token 槽。块越大,访问越容易合并,索引相对更少;块越小,尾部浪费越低,前缀共享粒度越细。Kwon 等,SOSP 2023

两条序列共享已填满的前缀块,并把新 token 写进各自尾块

图 1,原创示意。两个请求的逻辑前缀一致,引用同一个物理块。各自的尾块可以继续追加,物理编号可以不连续。

前缀复用进一步把计算过的完整块留在池中。块的 KV 取决于可见历史,因此命中条件要包含前缀身份及当前位置的 token。权重版本、适配器或多模态输入不同,也会改变缓存语义。共享块应保持只读;若实现允许共享未填满的块,分叉后继续写入需要写时复制。跨租户复用时,缓存键或缓存池还应体现隔离边界。vLLM 前缀缓存设计 SGLang / RadixAttention,2023

前缀命中可以省掉命中部分的 Prefill,并让多个请求引用同一份物理状态。Decode 仍然按各自的 query 读取这些 token。共享块能否进一步减少 HBM 读量,取决于执行时的调度、片上复用和硬件缓存。

4. 每 token 还能存得更小吗

方法 主要改变 需要核对
MQA / GQA 多个 Q head 共享一个 KV head,容量按 \(H_{kv}\) 缩放 模型结构与训练方式、同组加载复用、TP 分片与复制
MLA 保存压缩 latent 和单独的位置 key latent 维度、RoPE 维度、投影吸收后的内核形状
KIVI 等低比特量化 降低缓存载荷的元素位宽 scale/zero point、残余区、解包与反量化、生成质量
StreamingLLM 等保留策略 减少驻留或参与计算的 token 数 长程检索、任务精度、被移除 token 的访问语义

GQA 在保持多个独立 Q head 的同时,让每组共享 K/V。8 个 KV head 对 32 个 Q head,载荷是同维度 MHA 的四分之一。若算子把共享 K/V 复制成 32 份或让各 head 各自读 HBM,容量节省和实际带宽节省就会出现差距。

MLA 把内容 K/V 投影共享到一个低维向量 \(c_t^{KV}\),并单独存放解耦的 RoPE key。对 DeepSeek-V2 的这一缓存表示,每层每 token 有 \(d_c+d_R\) 个缓存元素,其中 \(d_R\) 是共享位置 key 的维度。内容打分可用矩阵乘的结合律改写,其中 \(W_h^{UK}\) 表示第 \(h\) 个 head 的 key 上投影:

\[q_h^{\mathsf T}W_h^{UK}c_t^{KV} =\big((W_h^{UK})^{\mathsf T}q_h\big)^{\mathsf T}c_t^{KV}.\]

类似地,value 的线性上投影可以挪到加权汇总之后。把内容和位置分开,保留了这种重排的机会。实际内核可能针对 Prefill 和 Decode 选择不同的展开或吸收路径。DeepSeek-V2 技术报告,§2.1

5. 从量化位宽算到内存事务

KIVI 的分组轴来自 K、V 的统计差异。K 在某些通道持续出现离群值,按通道分开定标可以减少离群通道对其他通道的影响。V 按 token 定标,便于单独保留每个位置的动态范围。K 的量化组跨多个 token,因此持续追加时先放进高精度残余区,凑足组再量化。每组有 \(g\) 个元素、位宽 \(w\)、scale 与 zero point 合计 \(s\) 字节时,平均载荷为 \(w/8+s/g\) 字节/元素。例如 2-bit、32 元素一组、4 B 参数,就是 0.375 B/元素,约为 BF16 的 5.33 倍压缩,尚未计入高精度残余区和布局 padding。KIVI

KVQuant 进一步在 RoPE 之前量化 key,并研究非均匀量化和离群值的稀疏高精度表示。预旋转 key 读取后仍要保持原位置语义,离群值也要从单独的表示恢复。这个方向用额外恢复逻辑换质量。FP8、NVFP4 等运行时格式则要看算子能否直接消费对应布局和 scale。只在主机或磁盘的冷层用 FP4 时,节省的是冷层容量和传输字节,Decode 的设备读量取决于热层格式。KVQuant TensorRT-LLM 的冷热层压缩

Prefill 路径适合批量生成 scale 和打包。Decode 中每步写一个新 token,却读很多已量化 token。NPU 可在页加载后由 Vector 逐 tile 解包,将恢复后的 tile 直接送给矩阵路径。若解包占满 Vector 或产生高精度 HBM 临时缓冲,带宽节省会被转换和写回抵消。对比时把 KV dtype、scale 字节、实际 HBM 流量、Vector 利用率和长上下文任务质量一起记录。

6. 淘汰、选择与页回收是三件事

工作 什么时候决定 保留什么
StreamingLLM 随 Decode 窗口滚动 最初的 attention sink 与最近 token
H₂O 生成时更新得分 近期 token 与累计注意力权重高的历史位置
SnapKV Prefill 末尾观察窗口投票 每个 head 的高分历史位置与观察窗口
PyramidKV 按层分配预算,再选择位置 浅层更宽的历史范围,深层更集中

这几种策略用有限容量交换历史可见性。举例来说,prompt 末尾问的是合同生效日期,SnapKV 的观察窗口可能偏向日期相关段落。后续生成转向违约条款时,先前压缩保留的历史未必包含新问题需要的信息。H₂O 可以在线更新,但在融合内核中暴露每个 token 的累计权重会增加统计开销。对 GQA 而言,若组内 Q head 选择的历史不同,联合保留会抬高页数,单独保留可能破坏共享和复用。

算子少访问若干 token,也不会自动归还设备页。页内还有存活 token 时,页管理器通常不能整页释放。要把选择转换成容量收益,得把存活 token 压实并更新索引,或者直接按页做决策。前者增加搬运和写放大,后者损失选择粒度。DSA 的逐 query Top-K 主要减少本步主 Attention 的读量,未来的 query 仍可能选到当前跳过的位置,因此原始历史通常还得保留。

7. 同一份 KV 可以在层间、请求间或设备间复用

层间共享由模型决定。YOCO 的 cross-decoder 多层读取由 self-decoder 生成的一套全局 K/V。缓存大小应拆成 self-decoder 的局部窗或递推状态,以及共享全局缓存。容量收益与每层是否重读 HBM 要分开测。它与第 3 节的 Prefix Caching 可以叠加,因为前缀缓存关心的是能否跨请求重用已经算出的状态。YOCO

请求间复用靠身份匹配。SGLang 的 RadixAttention 用前缀树管理复用边界,vLLM 则按块哈希组织。对于共享的 6K prompt 和独占的 2K 后缀,两个请求可以只保留一份 6K 的物理缓存。若两条 Decode 的 query 独立执行,单步合计读取仍可能接近两份 8K;能否减少这部分 HBM 流量要看调度是否把共享部分放进同一片上工作集。SGLang vLLM 前缀缓存

页间接寻址可以发生在不同层次。PagedAttention 把块表交给 Attention kernel。vAttention 保留连续虚拟地址,在后端按需映射物理内存,使内核可继续使用连续张量接口。FlashInfer 是支持分页、可组合稀疏布局及不规则请求调度的内核引擎,页池生命周期仍由调用方管理。对 NPU,若驱动的虚拟内存能力与页粒度不合适,可以用显式块表。评估时分别看页分配/回收、映射或索引开销,以及算子侧的实际 DMA 效率。vAttention FlashInfer

8. 层级存储决定的是从哪里取、何时取

SRAM 用作算子 tile 和短期复用工作集,HBM 适合活跃历史,主机 DDR 可容纳暂时不跑的请求或热门前缀的副本。CXL 内存只有在平台具备合适的可访问路径时才可纳入扩展层,需要测实际链路带宽、时延与拷贝方式。SSD 更适合可提前预测、能做批量 I/O 的冷状态。FlexGen 展示了在 GPU、CPU 和磁盘之间联合调度权重、激活和 KV 的吞吐路线;交互式小 batch 的 ITL 预算与离线吞吐场景不同。FlexGen

对刚到达的请求,可比较加载已缓存前缀与重新 Prefill 的关键路径:

\[T_{\rm fetch}\approx t_{\rm queue}+t_{\rm map}+B_{\rm cold}/\beta_{\rm link}+t_{\rm convert}.\]

考虑 128 KiB/token 的 BF16 GQA,一段 4K 前缀是 512 MiB。若有效链路带宽为 25 GiB/s,载荷需要 20 ms。若重算花 30 ms,读取路径只剩约 10 ms 的排队、映射和转换余量。反过来,如果对同一份冷页每步 Decode 都这样搬一次,20 ms 会反复落到 ITL 上。因此运行时通常把可预见的热历史重新放回 HBM,或用稀疏选择把每步必须搬运的工作集限制在更小范围。

KV Cache 在计算单元、设备与远端之间的流向

图 2,原创。热路径按 tile 读页,跨节点路径转移可重新映射的页载荷。

9. NPU 上的分页读取和 Prefill/Decode 分离

以昇腾的增量推理接口为例,页池的 token 粒度、物理 stride、head 布局与块表必须同时匹配算子约束。CANN 8.5 的 aclnnIncreFlashAttentionV3 文档列出分页和 KV 反量化能力。具体能组合哪些特性,以目标产品和对应接口参数为准。昇腾 CANN 接口文档

在内核里,地址生成先通过块表定位 K/V。搬运单元按连续区间把页内切片送入片上缓冲,低比特路径可能还要由 Vector 做解包和缩放,再供矩阵路径消费。页与计算 tile 不匹配时,tile 可能跨多页,增加小粒度搬运;页过大则可能增加尾块浪费。可以固定上下文总 token 数,分别测连续页和打散页,再扫页大小、历史长度和 batch,区分 HBM 带宽不足、搬运粒度过碎和 Vector 反量化跟不上这三类现象。

分离式推理把 Prefill 和 Decode 放在不同 worker。Prefill 完成后,Decode 需要取得各层 KV 与有效长度,并在自己的页池中建立块表。生产端的物理块号只在自己的内存池有意义。两端若采用不同的并行策略或数据布局,还要计入切分、重排及同步成本。DistServe,OSDI 2024 P/D-Serve,2024

DistServe 以 TTFT 和每 token 延迟约束下的 goodput 为目标选择两阶段资源。Mooncake 进一步把缓存看成跨节点共享资源,利用 CPU、DRAM、SSD 和高速网络组成缓存池。分布式目录记录缓存身份和存放位置,调度器可优先把请求分配给已有前缀的节点。缓存命中也要付出实际传输成本。接收端需在 KV 与长度元数据都就绪后才开放 Decode,失败重试和请求取消后需回收未完成的页引用。Mooncake

若完整 KV 大小为 \(B_{\rm KV}\),跨 worker 路径的有效带宽为 \(\beta_{\rm link}\),一次完整迁移仅传输载荷就至少需要 \(B_{\rm KV}/\beta_{\rm link}\)。用上面的 1 GiB 和假设的 25 GiB/s 计算,是 40 ms。按层流水或前缀命中可以减少暴露在关键路径上的时间。评估时分别记录跨 worker 字节数、排队和迁移时间、TTFT、首个 Decode 间隔及后续 ITL 尾分位,才能区分计算隔离的收益和搬运引入的成本。

10. 沿七条路线推导缓存的生命周期

以下用 32 层、8 个 KV head、128 维、BF16 的 GQA 作为基线。每层每 token 4 KiB,整模型每 token 128 KiB。每项优化先回答保存什么,再算写入、读取与回收。

10.1 架构压缩如何改变总容量

请求长度为 \(T_i\) 时,无共享的全局缓存为 \(128\ {\rm KiB}\sum_iT_i\)。如果 16 个请求各有 8K token,就是 16 GiB 有效载荷。把 KV head 从 8 改为 2,可将同维度载荷降至 4 GiB;模型质量和权重需要对应的架构支持。

MLA 用 latent 与位置 key 代替逐 head 的 K/V。若每层缓存 576 个 BF16 元素,32 层每 token 是 36 KiB,同样 16 条 8K 请求为 4.5 GiB。这是固定层数、固定示例维度的容量对照。每层主 Attention 的 FLOPs、投影权重和并行复制量还要另外计算。容量较小的模型不一定有较小的 Decode 延迟。

10.2 量化误差怎样进入 Attention

先看一个量化组。设 scale 为 \(s\),零点为 \(z\),整数表示为 \(q=\operatorname{clip}(\operatorname{round}(x/s)+z)\),恢复值为 \(\hat x=s(q-z)\)。K/V 恢复后可记成 \(\hat K=K+E_K,\hat V=V+E_V\)。

分数扰动是 \(\Delta S=QE_K^{\mathsf T}/\sqrt d\)。对一行 Softmax,其 Jacobian 为 \(J=\operatorname{diag}(p)-pp^{\mathsf T}\),一阶近似下输出误差为

\[\Delta o\approx(J\Delta s)^{\mathsf T}V+p^{\mathsf T}E_V.\]

K 的误差先改变权重,再作用到 V;V 的误差直接进入加权求和。按通道或按 token 选择量化轴,本质上是在约束不同位置的误差如何传播。scale 的粒度越细,局部动态范围越容易拟合,元数据和解码成本也越高。KIVI、KVQuant 的量化轴和离群值处理正是在这些约束下取舍。KIVI KVQuant

设量化主体长度 \(T-R\),保留 \(R\) 个近期高精度 token,量化后每元素含元数据占 \(\bar b\),则每层缓存约为

\[B=2H_{kv}d\big((T-R)\bar b+Rb\big).\]

如果 \(T=8192,R=128,b=2,\bar b=0.375\),相对 BF16 载荷的比例约为 20.0%,压缩约 5 倍。这个算例还没有计入页对齐。Prefill 可以批量打包,Decode 要在残余区满时完成转换,并保证读者不会读到只更新了一半的量化页。

10.3 淘汰后怎样恢复正确的索引

假设保留的原位置为 \([0,1,101,500,8191]\)。压实后的物理槽是 \([0,1,2,3,4]\),位置编码却仍对应原位置。若把物理槽号当位置号重新应用 RoPE,就改变了模型看到的相对距离。实现需要保存逻辑位置,或者保留已按原位置变换好的 key。

H₂O 更新历史累计得分,SnapKV 在观察窗口上对候选投票,PyramidKV 进一步按层分预算。三者都需要明确“得分对应哪个原位置”。若每页有 16 个槽,淘汰 15 个而留下 1 个,页池仍占完整一页。压实时的最低数据搬运量至少包含存活 token 的读和写,还要更新映射。频繁压实可能吃掉容量节省带来的收益,因此应把决策周期、页粒度和生成长度一起测。H₂O SnapKV PyramidKV

10.4 前缀共享的一笔完整账

两个请求有相同的 6K 前缀,各自再延伸 2K。按本节基线,独立保存需要 \(2\times8K\times128\) KiB,即 2 GiB;共享前缀后需要 \((6K+2K+2K)\times128\) KiB,即 1.25 GiB。节省 0.75 GiB。

Prefill 命中前缀后,新后缀仍要读取该前缀。后缀长度为 \(C\)、共享前缀长度为 \(P\),后缀产生的有效因果配对数为 \(CP+C(C+1)/2\)。跳过旧前缀的重新计算并不消除新 query 对旧 key 的访问。

生命周期上,共享页的引用计数随请求进入、分叉、取消和结束变化。只有最后一个引用释放后才可回收。未填满尾页如果需要继续写入,可在分叉时复制已有有效槽。前缀键要纳入会影响 KV 的模型、适配器和输入身份,目录中的命中必须对应真实可复用的状态。vLLM 前缀缓存

10.5 页大小怎样影响碎片和读取

页大小 \(P\)、序列长度 \(T\) 时,占用槽数是 \(P\lceil T/P\rceil\),尾部浪费为

\[W=P\lceil T/P\rceil-T,\qquad 0\le W<P.\]

取 \(T=8193\),16-token 页浪费 15 槽,128-token 页浪费 127 槽。按整模型每 token 128 KiB 计,两者分别约为 1.875 MiB 和 15.875 MiB。大页会减少块表项和部分地址管理开销,但请求短、并发多时,尾部浪费会累加。

vAttention 将动态物理分配放在连续虚拟地址背后,需要底层映射机制;PagedAttention 在内核显式索引页表;FlashInfer 负责消费相应布局并组织请求的计算。选择时要先确认分配器、驱动和算子的接口能否一致。单看 kernel latency 会漏掉页分配、映射和请求回收成本。vAttention FlashInfer

10.6 卸载的临界点如何计算

假设某段缓存大小为 \(B\),链路有效带宽为 \(\beta\),固定提交和映射开销为 \(t_0\)。读取时间约为 \(t_0+B/\beta\)。只有重算时间大于这一值,且请求延迟预算允许时,读回才有机会带来收益。

例如 512 MiB 缓存、25 GiB/s 有效链路,纯传输 20 ms。可提前 30 ms 发起搬运时,传输可能被前序计算遮住;依赖上一 token 的动态选择直到最后才确定时,就很难获得同样的提前量。SSD、DDR、CXL 或远端内存的实测带宽和排队时延应分别带入。多个请求一起迁移时,共享链路带宽还会改变每条请求的到达时间。

10.7 分布式缓存怎样进入一次请求

一次分离式请求可以沿以下过程检查。调度器先查前缀所在位置,选择 Prefill worker;Prefill 按层生成 KV;传输端发送缓存身份、布局、有效长度和载荷;Decode worker 为这些逻辑块分配自己的物理页;每层所需数据就绪后才能执行对应计算。取消请求时,双方都要释放尚未完成的传输和页引用。

设每层传输时间 \(t_x\),Prefill 每层计算 \(t_p\),共有 \(L\) 层。若每层产生后即可发送,理想流水完成时间近似为

\[T_{\rm ready}\approx t_p+t_x+(L-1)\max(t_p,t_x).\]

全部计算完才传输则约为 \(L(t_p+t_x)\)。前式要求层间依赖、发送缓冲、接收页和链路能够支撑重叠,还未计入排队与同步。传输追不上生产速度时,待发送 KV 会暂时同时占据生产端和接收端资源,因此容量规划必须计入迁移中的副本。DistServe Mooncake


本站由 Zane Jiang 使用 Stellar 1.33.1 主题创建,一款很棒的 Hexo 主题!

总访问 次 || 本页访问 次
总访客 人 || 本页访客 人