【随笔】KV缓存卸载：逐层预取能藏住多少传输时间

昨天的随笔讨论了Prefill与Decode位于不同实例时怎样交接KV状态。今天把范围缩小到一台机器：模型权重已经能放进GPU，短提示也能正常生成，长提示却在生成过程中耗尽显存。把KV缓存移到CPU内存能腾出多少空间？下一枚Token又要为此等待多久？

一、先弄清哪部分内存随上下文增长

模型权重用于每次前向计算，同一个模型不会因为多生成一个Token就多出一份权重。KV缓存保存各注意力层已经计算过的键和值。随着请求继续生成，全注意力层保留的历史越来越长。两个请求即便使用同一套权重，也各自需要与其Token历史匹配的缓存。

本文讨论Transformers 4.57.1提供的逐层缓存卸载。GPU执行当前层时，预取下一层需要的缓存；当前层完成注意力计算后，把更新过的缓存送回主存。预取允许计算与部分传输重叠，临时缓冲也要占用显存，不能简单把实际峰值当成一层缓存的大小。Transformers缓存策略

栗子看管CPU侧各层缓存，小橘在GPU侧执行当前层。下一层要提前送来，当前层更新后的状态也要写回，两个方向都需要预算。

二、为什么"省下1GiB"还不能回答速度问题

设一条请求有32层全注意力，每层8个KV头，头维度128，元素占2字节，历史长度8192。每层历史KV的大小为：

text 复制代码

2份张量 × 8个KV头 × 128维 × 2字节 × 8192个Token = 32 MiB
32层 × 32 MiB = 1 GiB

若按这个简化布局，每步取回所有32层的历史KV，单向读取量便是1GiB，尚未加上写回。实际卸载层数、滑动窗口和实现的复制粒度都会改变搬运量。这里算的是每次访问的数据，不是一次分配之后就结束的成本。

设每层有2ms计算可与下一层预取重叠。32MiB在16GiB/s有效读带宽下需要1.953ms，在8GiB/s下需要3.906ms。前一种配置有机会在计算结束之前备好下一层，后一种配置在每次交接处留下约1.906ms等待。这个比较使用明确的带宽与计算时间假设。

三、一个可运行的预取预算程序

下面把读方向单独拿出来计算。第0层必须先到位，此后31次切换各有一次隐藏读延迟的机会。它适合判断该测量哪段传输。

python 复制代码

def budget(layers, layer_mib, gib_per_s, compute_ms):
    read_ms = 1000 * layer_mib / (1024 * gib_per_s)
    exposed_ms = max(0, read_ms - compute_ms)
    total_ms = read_ms + layers * compute_ms + (layers - 1) * exposed_ms
    return read_ms, exposed_ms, total_ms

for bw in (16, 8):
    read, wait, total = budget(32, 32, bw, 2)
    print(f"{bw} GiB/s: read={read:.3f} wait={wait:.3f} total={total:.3f} ms")

按模型得到：

text 复制代码

16 GiB/s: read=1.953 wait=0.000 total=65.953 ms
8 GiB/s: read=3.906 wait=1.906 total=127.000 ms

总时长包含首次读取、32次计算和各次切换中露出来的等待时间。真实生成还涉及写回、同步、内存分配和其他算子。尤其当读写共享传输资源时，单测读取带宽容易高估可用的重叠时间。

四、怎样在同一份生成代码里比较两种策略

在模型权重已加载到GPU、输入张量也已经放到对应设备后，可以用同一份输入调用两次generate。这里让两次调用分别创建自己的缓存：

python 复制代码

options = dict(max_new_tokens=32, do_sample=False, use_cache=True)
resident = model.generate(**inputs, **options, cache_implementation="dynamic")
offloaded = model.generate(**inputs, **options, cache_implementation="offloaded")
print(resident.shape, offloaded.shape)
print(resident.equal(offloaded))

先用两种方式都能容纳的短输入做这项对照，确认模型、输入和停止条件一致。再逐渐拉长提示，记录在哪个长度开始出现显存不足。若常驻策略已经耗尽显存，单独运行卸载策略验证是否能完成同一个请求，不把失败请求与成功请求的耗时混在一个均值里。

固定最大容量时，文档还提供offloaded_static。使用滑动窗口的层在达到窗口上限后，缓存不会像全注意力层那样继续增长。决定是否卸载这些层时，可以算一算腾出的字节是否值得每层发生传输。

五、验收时记录哪些数

问题	应比较的量	怎样解释
能否完成目标请求	相同输入、输出上限、批量下的成功率与GPU峰值	先满足容量目标
避免OOM后是否可用	首Token时间、生成阶段逐Token延迟及高分位	识别预取赶不上计算的情况
瓶颈去了哪里	CPU内存峰值、双向传输量、GPU等待区间	区分主存容量、传输和计算瓶颈
是否仍得到预期结果	相同解码配置下的Token和停止原因	发现状态错配或回归

不要只看GPU利用率的单个平均值。一段运行可能在计算当前层时很忙，又在切换到下一层时反复等待。把传输和计算放在同一条时间线上，才知道能否靠预取隐藏开销。

六、🧠 思维导图

#mermaid-svg-IZoeTSDfhnVTvFqs{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-IZoeTSDfhnVTvFqs .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-IZoeTSDfhnVTvFqs .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-IZoeTSDfhnVTvFqs .error-icon{fill:#552222;}#mermaid-svg-IZoeTSDfhnVTvFqs .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-IZoeTSDfhnVTvFqs .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-IZoeTSDfhnVTvFqs .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-IZoeTSDfhnVTvFqs .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-IZoeTSDfhnVTvFqs .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-IZoeTSDfhnVTvFqs .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-IZoeTSDfhnVTvFqs .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-IZoeTSDfhnVTvFqs .marker{fill:#333333;stroke:#333333;}#mermaid-svg-IZoeTSDfhnVTvFqs .marker.cross{stroke:#333333;}#mermaid-svg-IZoeTSDfhnVTvFqs svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-IZoeTSDfhnVTvFqs p{margin:0;}#mermaid-svg-IZoeTSDfhnVTvFqs .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-IZoeTSDfhnVTvFqs .cluster-label text{fill:#333;}#mermaid-svg-IZoeTSDfhnVTvFqs .cluster-label span{color:#333;}#mermaid-svg-IZoeTSDfhnVTvFqs .cluster-label span p{background-color:transparent;}#mermaid-svg-IZoeTSDfhnVTvFqs .label text,#mermaid-svg-IZoeTSDfhnVTvFqs span{fill:#333;color:#333;}#mermaid-svg-IZoeTSDfhnVTvFqs .node rect,#mermaid-svg-IZoeTSDfhnVTvFqs .node circle,#mermaid-svg-IZoeTSDfhnVTvFqs .node ellipse,#mermaid-svg-IZoeTSDfhnVTvFqs .node polygon,#mermaid-svg-IZoeTSDfhnVTvFqs .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-IZoeTSDfhnVTvFqs .rough-node .label text,#mermaid-svg-IZoeTSDfhnVTvFqs .node .label text,#mermaid-svg-IZoeTSDfhnVTvFqs .image-shape .label,#mermaid-svg-IZoeTSDfhnVTvFqs .icon-shape .label{text-anchor:middle;}#mermaid-svg-IZoeTSDfhnVTvFqs .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-IZoeTSDfhnVTvFqs .rough-node .label,#mermaid-svg-IZoeTSDfhnVTvFqs .node .label,#mermaid-svg-IZoeTSDfhnVTvFqs .image-shape .label,#mermaid-svg-IZoeTSDfhnVTvFqs .icon-shape .label{text-align:center;}#mermaid-svg-IZoeTSDfhnVTvFqs .node.clickable{cursor:pointer;}#mermaid-svg-IZoeTSDfhnVTvFqs .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-IZoeTSDfhnVTvFqs .arrowheadPath{fill:#333333;}#mermaid-svg-IZoeTSDfhnVTvFqs .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-IZoeTSDfhnVTvFqs .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-IZoeTSDfhnVTvFqs .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-IZoeTSDfhnVTvFqs .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-IZoeTSDfhnVTvFqs .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-IZoeTSDfhnVTvFqs .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-IZoeTSDfhnVTvFqs .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-IZoeTSDfhnVTvFqs .cluster text{fill:#333;}#mermaid-svg-IZoeTSDfhnVTvFqs .cluster span{color:#333;}#mermaid-svg-IZoeTSDfhnVTvFqs div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-IZoeTSDfhnVTvFqs .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-IZoeTSDfhnVTvFqs rect.text{fill:none;stroke-width:0;}#mermaid-svg-IZoeTSDfhnVTvFqs .icon-shape,#mermaid-svg-IZoeTSDfhnVTvFqs .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-IZoeTSDfhnVTvFqs .icon-shape p,#mermaid-svg-IZoeTSDfhnVTvFqs .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-IZoeTSDfhnVTvFqs .icon-shape .label rect,#mermaid-svg-IZoeTSDfhnVTvFqs .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-IZoeTSDfhnVTvFqs .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-IZoeTSDfhnVTvFqs .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-IZoeTSDfhnVTvFqs :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} KV缓存卸载
容量
逐层历史状态
GPU与主存
执行
预取下一层
更新后写回
验收
峰值与延迟
结果与传输

七、总结

总结要点

卸载把一部分KV驻留空间交给CPU内存，代价是生成过程中持续发生数据传输。是否值得启用，取决于它能否让目标长度的请求完成，以及完成后的逐Token延迟能否被业务接受。

从单层KV大小出发，算出预取时长，再和当前层可供重叠的计算时间比较，就能找到一次有针对性的性能实验。实际验收需把写回也放到时间线上。

下一篇讨论KV缓存量化，看看缩短每个元素的表示以后，容量、转换开销和输出质量怎样一起验收。

👉 如果你觉得这篇文章对你有所帮助，欢迎点赞、收藏、分享！😊