跳到主要内容

给 KV cache 减肥——不动架构的内存优化

这一章讲三件事: KV cache 为什么是动态内存的大头、量化它省在哪;Hugging Face 和 vLLM 两个引擎里同一件事为什么做法不同、记的账也不同;以及为什么「哪个技巧更好」是个伪问题——硬件说了算。 这是注意力章(原书第 8 章)的上半场:完全不碰架构;下半场(下一章)才动刀。

1. 这一章讲什么​

前七章一直在动权重(删掉它们、蒸馏它们、微调它们)。本章换目标:推理时随上下文增长的那块内存——KV cache。给它量化,模型一个参数都不改,却能处理更长的提示、并发(同时服务更多路)请求1。

它在全书链条里的位置: 第 03 章算过那笔「1152 个槽」的账;本章把那本账当成可以直接压缩的对象。它与下一章的「删注意力层」不是二选一,是可以叠加的两笔独立收益(下一章末尾给证据)。

2. 顶层全景​

  推理显存的两本账
    静态:模型权重,加载后不变(第 07 章的 NF4 管这本)
    动态:KV cache,每生成一个 token 涨一点(本章管这本)

  Hugging Face 引擎              vLLM 引擎
  ───────────────────           ───────────────────
  cache 随用随长,长进空闲显存    启动时预分配固定比例显存
  压缩粒度:4-bit(HQQ)          压缩粒度:FP8(硬件理由)
  度量:兆字节(MB)              度量:能容多少 KV token

  图说:同一个「量化 cache」,在两个引擎里连成功的标准都不一样[^2]。

3. 核心原理​

3.1 压缩-还原循环​

KV cache 默认按 16-bit 浮点(FP16 一类)存。量化的做法:注意力模块每产出一对 K、V,写进显存(显卡上那块工作内存)之前压成 8-bit(FP8,1 字节——8 个最小开关位)或 4-bit(0.5 字节);下一个新词的 Query 要跟它们相乘时,再还原回 16-bit。每个 token 走一遍「压缩→写入→读出→还原」的循环2。

省是省了,两笔代价:精度损失(第 07 章见过量化的一般代价),以及——取决于推理库和 GPU 架构——可能的提速反噬。这个「取决于」是本章的主角3。

3.2 Hugging Face 实验:两本账各自独立​

实验设计本身就是结论的一半:同一个 Llama-3.2-3B,三种加载方式——全精度基线 / 4-bit 权重(cache 不动)/ 4-bit KV cache(权重不动)——分别量静态显存、动态增量、吞吐。要证明的正是:权重量化与 cache 量化是作用在不同内存区上的两个独立杠杆4。

工具上,transformers 的量化 cache 需要专用后端,主流是 quanto 与 HQQ;本书选 HQQ——性能更好,而且对 T4 这种老架构(Turing 世代)没有兼容问题(quanto 已并入 Optimum 库)5。一个坑顺手带走:代码里 cache_implementation="quantized" 这个参数不传的话,cache 配置会被静默忽略——你以为在量化,其实在裸奔6。

结果(T4):4-bit 权重把静态显存从 6.1 GB 压到 2.2 GB,近 3 倍;4-bit cache 把动态增量压到约三分之一7。一个反直觉细节:开了 cache 量化,静态显存反而涨了一点——那是 HQQ 内核的一次性预留,L4 上就没有这个现象,证明涨的不是方法本身8。

质量对照,这一节最重要的发现: 三个配置各生成 500 token,基线和 cache 量化的输出连贯;4-bit 权重版从第一个 token 起陷入复读循环,500 个 token 没爬出来9。原因:3B 这个尺寸的模型参数冗余少,对权重量化极敏感;BitsAndBytes 均匀压缩所有权重,不保护少数关键值(AWQ、GPTQ 那类方法有校准步骤,专门先找出这些「离群」权重再压缩)。而 cache 量化不碰权重,没有这个问题10。

3.3 vLLM:换引擎,换做法,换账本​

vLLM 的内存管理哲学完全不同:PagedAttention,分页(把内存切成小格来管理)式地工作,启动时就把可用显存的固定比例(实验里 85%)预分配掉。这个设计直接改写了「省内存」的意义——权重省下的每一字节,自动变成 KV 池的新容量11。所以 vLLM 的成绩单不记兆字节,记「引擎能容多少 KV token」:能装越多,能同时服务的长上下文请求就越多12。

连压缩格式都换了:vLLM 不用 4-bit cache,用 FP8。理由是纯硬件的——字节是 GPU 最小可寻址单位,4-bit 张量每次读取都要做额外的拼装计算,罚延迟;FP8 恰好一字节,不用拼13。

3.4 主走查:同一道题,两代硬件给出相反答案​

走查对象:同一个 Llama-3.2-3B,同一段长提示,三种配置(FP16 基线 / NF4 权重 / FP8 cache),两块 GPU。 全部数字为原书实测:

  T4(老架构,无原生 FP8)
    FP16 基线:      容量 51,248 token · 吞吐 10.8 token/秒
    FP8 cache:      容量 108,192(+111%)· 吞吐 8.6(罚速 -20%)
    NF4 权重:       容量 +68.4%       · 吞吐 31.64(近 3 倍)

  L4(新架构,原生 FP8)
    FP16 基线:      吞吐 35.6 token/秒
    FP8 cache:      容量 226,768(+105%)· 吞吐 35.8(纹丝不动)
    NF4 权重:       容量 +31.8%       · 吞吐 82.5(2.3 倍)

  图说:同一个 FP8,在 T4 上要软件转换(所以罚速),
  在 L4 上是硬件原生(所以白赚)[^15]。

为什么 NF4 压权重能换来近三倍吞吐?原理在第 04 章已经打过照面:自回归(写一个词、看一遍前文、再写下一个)解码是内存带宽受限的活——每生成一个词,整个模型都要从显存流过一遍;模型占内存少,每步搬运的字节就少,吞吐直接上涨14。

质量那一栏也要看:FP8 配置在两块卡上都生成连贯;NF4 权重版在两块卡上都复读——和 Hugging Face 实验互相印证:压权重对 3B 模型的质量伤害是真实的,容量和速度的表格只讲了一半故事15。

所以本章的结论不是排名,是决策表:要上下文容量,量化 cache——新卡免费,老卡接受罚速;要吞吐,压权重——但要选带校准保护的方法并测质量;两者作用于不同内存区,可以同时要16。

4. 作者的判断与证据​

给了证据的: 全部数字出自原书在 T4/L4 上的可复现 notebook(脚本自带解析器(一小段代码),从 vLLM 启动日志(运行过程的文字记录)里抓 KV token 容量);质量对照的生成文本原书原样贴出。

作者的判断: 「硬件的选择不是小事——同一种量化在不同 GPU 架构上表现可以完全不同」是本节的题眼17;「没有哪项技术全面更优,选择取决于目标与手头的硬件」是本节的结论18。这两句都是经验判断,但书内两代 GPU 的对照数据直接支撑。

工程坑(书里明说): vLLM 在 Jupyter/Colab 里启动会因 stdout 缺少 fileno() 报错——书里的解法是把测试写成独立脚本落盘再运行,结果用一行 JSON 取回19。

5. 边界与局限​

  • 实验模型只有 3B。 「NF4 复读」在小模型上严重;更大的模型参数冗余多,结论未必如此——书里没有测更大模型的这一项。
  • 吞吐测量环境嘈杂。 共享 GPU 上的数字本身不可比,作者自比的是方向;本书第 04 章同样的提醒在此适用。
  • 4-bit cache 只在 Hugging Face 路线出现,vLLM 路线没有等价配置;两边数字不能横向互比(度量都不同)。
  • 没测对能力基准的影响。 本章只看生成文本的观感,没有跑基准评分——cache 量化对事实能力的定量影响是空白。

6. 可带走的​

  1. 显存有两本账:静态(权重)与动态(cache);两个独立杠杆,可分别量化、可叠加;
  2. KV 量化的本质是「写前压缩、读前还原」,每个 token 一遍循环;
  3. Hugging Face 路线记 MB,vLLM 路线记 KV token 数——引擎不同,成功的定义都不同;
  4. vLLM 预分配显存:权重省下的内存自动变成上下文容量,这是它「省内存」的意义;
  5. vLLM 用 FP8 不用 4-bit:字节是最小寻址单位,拼装 4-bit 罚延迟;
  6. 同一技巧换硬件可能反转:FP8 在 T4 罚速 20%,在 L4 免费;
  7. 压权重对 3B 模型有真实质量风险(复读循环),选带校准保护的方法(AWQ/GPTQ 类);
  8. cache_implementation="quantized" 忘传,配置静默失效——量化没生效不会报错。

7. 原文地图​

主题原书章原文位置
本章两个主题、注意力是瓶颈8 Attention optimizationtext/30-ch08-8-attention-optimization.txt:41(搜「main bottleneck」) · text/30-ch08-8-attention-optimization.txt:16(搜「quantizing the KV cache」)
压缩-还原循环8.1 KV cache fundamentals and Hugging Face implementationtext/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:43(搜「dequantizes them」) · text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:56(搜「cycle repeats」)
两引擎两卡、独立杠杆8.1 KV cache fundamentals and Hugging Face implementationtext/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:19(搜「not a minor detail」) · text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:80(搜「independent levers」)
后端选择与静默忽略的坑8.1 KV cache fundamentals and Hugging Face implementationtext/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:124(搜「Turing」) · text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:240(搜「silently ignored」)
HF 结果与反直觉细节8.1 KV cache fundamentals and Hugging Face implementationtext/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:316(搜「6.1 GB to 2.2 GB」) · text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:322(搜「3x reduction」) · text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:347(搜「HQQ kernel」)
质量对照与复读8.1 KV cache fundamentals and Hugging Face implementationtext/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:389(搜「repetition loop」) · text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:389(搜「outliers」) · text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:353(搜「native quantization capabilities」)
vLLM 内存哲学8.2 Scaling KV cache for production with vLLMtext/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:7(搜「PagedAttention」) · text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:13(搜「KV tokens」) · text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:59(搜「smallest addressable unit」)
主走查:T4/L4 结果8.2 Scaling KV cache for production with vLLMtext/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:236(搜「108,192」) · text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:242(搜「31.64」) · text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:255(搜「82.5」) · text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:261(搜「35.8 vs 35.6」)
内存带宽受限原理8.2 Scaling KV cache for production with vLLMtext/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:242(搜「memory-bandwidth-bound」)
结论与质量提醒8.2 Scaling KV cache for production with vLLMtext/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:267(搜「universally better」) · text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:273(搜「repetition loop」)

Footnotes​

  1. 出处:「8 Attention optimization」第 41 段(text/30-ch08-8-attention-optimization.txt:41,搜「main bottleneck」)与第 13 段(text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:13,搜「grows with the context」)。 ↩

  2. 出处:「8.1 KV cache fundamentals and Hugging Face implementation」第 43 段(text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:43,搜「dequantizes them」)。 ↩

  3. 出处:「8.1 KV cache fundamentals and Hugging Face implementation」第 62 段(text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:62,搜「performance penalty」)。 ↩

  4. 出处:「8.1 KV cache fundamentals and Hugging Face implementation」第 80-97 段(text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:80,搜「independent levers」)。 ↩

  5. 出处:「8.1 KV cache fundamentals and Hugging Face implementation」第 124 段(text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:124,搜「Turing」)与第 116 行(text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:116,搜「Optimum」)。 ↩

  6. 出处:「8.1 KV cache fundamentals and Hugging Face implementation」第 240 段(text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:240,搜「silently ignored」)。 ↩

  7. 出处:「8.1 KV cache fundamentals and Hugging Face implementation」第 316 段(text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:316,搜「6.1 GB to 2.2 GB」)与第 322 段(text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:322,搜「3x reduction」)。 ↩

  8. 出处:「8.1 KV cache fundamentals and Hugging Face implementation」第 347 段(text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:347,搜「HQQ kernel」)。 ↩

  9. 出处:「8.1 KV cache fundamentals and Hugging Face implementation」第 373-383 行的三个生成对照(text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:389,搜「never exits」)。 ↩

  10. 出处:「8.1 KV cache fundamentals and Hugging Face implementation」第 389 段(text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:389,搜「outliers」)。 ↩

  11. 出处:「8.2 Scaling KV cache for production with vLLM」第 7 段(text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:7,搜「PagedAttention」)与第 205 段(text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:205,搜「available for the KV Cache」)。 ↩

  12. 出处:「8.2 Scaling KV cache for production with vLLM」第 13 段(text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:13,搜「KV tokens」)与第 217 段(text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:217,搜「difference in architecture」)。 ↩

  13. 出处:「8.2 Scaling KV cache for production with vLLM」第 59 段(text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:59,搜「smallest addressable unit」)。 ↩

  14. 出处:「8.2 Scaling KV cache for production with vLLM」第 242 段(text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:242,搜「memory-bandwidth-bound」)。 ↩

  15. 出处:「8.2 Scaling KV cache for production with vLLM」第 273 段(text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:273,搜「repetition loop」)。 ↩

  16. 出处:「8.2 Scaling KV cache for production with vLLM」第 267 段(text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:267,搜「universally better」);叠加性证据在第 08.5 节(text/35-ch08-05-8-5-from-paper-to-practice.txt:128,搜「cumulative」)。 ↩

  17. 出处:「8.1 KV cache fundamentals and Hugging Face implementation」第 19 段(text/31-ch08-01-8-1-kv-cache-fundamentals-and-hugging-face-imple.txt:19,搜「not a minor detail」)。 ↩

  18. 出处:「8.2 Scaling KV cache for production with vLLM」第 267 段(text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:267,搜「No technique is universally better」)。 ↩

  19. 出处:「8.2 Scaling KV cache for production with vLLM」第 39 段(text/32-ch08-02-8-2-scaling-kv-cache-for-production-with-vllm.txt:39,搜「fileno()」)。 ↩