跳转至

第十三章　推理优化：量化、蒸馏、KV Cache 与高效服务

推理优化要围绕目标场景：离线吞吐、在线首 token、持续解码、长上下文、边缘设备各有不同瓶颈。先测 profile，再选技术。量化减少权重带宽，KV Cache 优化减少重复计算，FlashAttention 减少内存访问，连续批处理提高设备利用率，彼此解决的问题不同。

推理优化栈

图 13-1　从模型表示、算子、缓存到调度的分层优化；每层都需质量与性能回归。

13.1 Prefill 与 Decode

先建立直觉。 Prefill 对整段输入并行计算并建立 KV Cache，通常计算密集；Decode 每步只处理新 token，却反复读取大量权重和缓存，通常更受内存带宽与调度影响。 读这一节时，先不要急着记缩写，而要不断追问：输入是什么，经过了哪些可观察变换，输出又怎样被验证。

把过程拆开看：

  1. prefill 产生首 token 与缓存。
  2. decode 读取历史缓存。
  3. 逐步采样并追加缓存。
  4. 遇 EOS/长度/停止串结束。

最小例子。 长提示短回答主要压力在 prefill 和 TTFT；短提示长回答主要压力在 decode、TPOT 和缓存占用。 例子越小，越容易手算、打印中间量并判断实现是否偏离定义。

容易踩坑。 用单一 tokens/s 描述两阶段、把输入 token 和输出 token 同价、忽略并发下缓存竞争，都会误判。 排错时应保存失败输入和版本，先定位最早出现异常的步骤，再决定是否调整模型或超参数。

动手任务：对四种输入/输出长度组合记录 TTFT、TPOT 与总时间，画二维热图。 完成后不要只保存最终输出，还要保存代码、配置、中间张量或检索结果，以及你对结果的解释。

Prefill 对整段输入并行计算，矩阵乘规模大，通常更偏计算密集；Decode 每步只生成一个 token，却要读取大量权重和 KV Cache，常偏内存带宽与调度。TTFT 主要受排队、tokenization 与 prefill 影响；TPOT 主要受 decode 影响。

优化前按输入/输出长度分桶测 TTFT、TPOT、吞吐、显存和质量。平均请求会掩盖长上下文的极端成本。

13.2 量化基础

从问题出发。 量化用较少比特表示权重或激活。核心是选择尺度、零点、粒度和校准数据，让舍入与截断误差可控。 先把名词放到一边，沿着输入、状态变化和输出走一遍；只有能指出证据落在哪一步，概念才算真正掌握。

沿数据流逐步检查：

  1. 统计目标张量范围。
  2. 按 tensor/channel/group 求尺度。
  3. 映射到整数并裁剪。
  4. 反量化或用低比特 kernel 计算。

用小数据走一遍。 对称 int8 可令 q=round(x/scale)，scale=max|x|/127；离群值会放大 scale，使多数小值分辨率下降。 先在纸面预测结果，再让程序打印中间状态；预测与运行不一致的地方，正是需要继续追查的知识缺口。

这里最容易出现的误解。 文件变小不等于推理变快、校准数据与真实分布不同、敏感层全部量化、只看平均精度，都会带来问题。 出现异常时先冻结样本、配置和环境，从最靠近输入的环节开始验证；不要同时更换模型、数据和超参数。

小实验：手写对称 int8 量化，比较不同 group size 的重构误差。 提交物应包含预期、运行证据和失败记录；只有别人能按同样步骤复现，结果才具有学习价值。

线性量化把浮点 \(x\) 映射到整数 \(q\)：

$$ q=\mathrm{clip}(\mathrm{round}(x/s)+z,q_{min},q_{max}), \qquad \hat{x}=s(q-z). $$ \(s\) 是尺度，\(z\) 是零点。对称量化通常 \(z=0\)；非对称量化更适合偏移分布。Per-tensor 使用一个尺度，简单但受离群值影响；Per-channel 或 group-wise 粒度更细，精度更好但元数据与 kernel 更复杂。

import torch

def symmetric_quantize(x: torch.Tensor, bits=8):
    qmax = 2 ** (bits - 1) - 1
    scale = x.abs().amax().clamp_min(1e-8) / qmax
    q = torch.round(x / scale).clamp(-qmax, qmax).to(torch.int8)
    return q, scale

def dequantize(q, scale):
    return q.float() * scale

这是教学版 per-tensor 权重量化，不会带来真实加速，因为计算前又反量化为 FP32。性能收益需要硬件支持的低比特 kernel、打包格式和融合算子。

13.3 BNB、GPTQ、AWQ 与 GGUF

先看它解决什么。 bitsandbytes 常用于易用的运行时量化，GPTQ/AWQ 属于权重量化路线，GGUF 是 llama.cpp 生态常用文件格式与量化集合。算法、格式和执行后端不可混为一谈。 理解的标准不是会复述定义，而是能画出数据流、预测中间结果，并设计一个让错误暴露出来的检查。

可以把实现分成以下环节：

  1. 确认目标硬件与 kernel。
  2. 选择权重/激活位宽。
  3. 准备校准数据。
  4. 比较质量、速度、内存和兼容。

一个可以手算的例子。 同一个 4-bit 文件在没有优化 kernel 的硬件上可能只省内存不加速，甚至因反量化开销更慢。 这里故意不用大模型或大数据，因为小输入能把每个轴、分数和状态变化完整暴露出来。

排错时先看这些地方。 只按模型大小选择方法、把 QLoRA 训练格式直接当部署格式、不同量化命名跨工具等价，都会踩坑。 修复不能止于‘这次跑通’，还要把失败样本变成自动测试，防止同类问题在下一版本重新出现。

现在动手：为 GPU 服务与 CPU 边缘部署分别写量化选择表。 把关键断言写进测试，并在 README 说明怎样运行、怎样判断正确、当前实现还不支持什么。

BitsAndBytes 提供 8-bit/4-bit 线性层与 QLoRA 支持，适合快速实验。GPTQ 使用校准数据做训练后权重量化，逐步减少重构误差；AWQ 关注激活显著通道，保护重要权重。两者最终速度取决于后端 kernel 和硬件。

GGUF 是 llama.cpp 生态常用的模型文件格式，支持多种量化类型，适合 CPU/边缘推理。格式兼容不代表不同量化方案质量相同。选择时以目标硬件实测 tokens/s、首 token、内存和任务质量为准。

不要把 QLoRA 与部署量化混为一谈：QLoRA 的目标是低显存微调；部署时可选择不同的合并与量化流程。

13.4 知识蒸馏

抓住这一节的主线。 知识蒸馏让学生模型学习教师的软分布、中间特征或生成数据。温度提高后，非最大类别概率携带教师对相似性的判断。 遇到新术语，先找它在系统中的位置：它读取什么、保存什么、改变什么，以及失败时会留下什么信号。

真正动手时按这个顺序走：

  1. 教师产生 logits 或答案。
  2. 学生同时学习硬标签与软目标。
  3. 调温度和损失权重。
  4. 在目标与回归集评估。

先做最小实验。 KL 蒸馏常乘 T² 补偿温度造成的梯度尺度变化；具体实现要核对 reduction 与 mask。 把随机性固定并保留中间量，这个例子就能成为后续优化时的回归基线。

别被表面现象带偏。 教师错误被批量放大、合成数据缺少多样性、tokenizer 不同直接逐 token 对齐、只看学生模仿率，都会限制效果。 先区分定义错误、实现错误、数据问题和性能瓶颈。四类问题需要的证据不同，不能靠盲目调参混在一起处理。

验证任务：在小分类任务比较无蒸馏与 logits 蒸馏，并分析错误样本。 实验结束后写三句话：观察到了什么、这些证据支持什么结论、还有哪种解释尚未排除。

蒸馏用教师模型指导学生模型。分类中可组合硬标签交叉熵与温度软化后的 KL：

$$ L=\lambda L_{CE}+(1-\lambda)T^2\,KL(p_T^{teacher}|p_T^{student}). $$ 语言模型还可蒸馏 logits、隐藏状态、注意力或教师生成数据。若只能调用教师 API，常采用数据蒸馏：生成高质量指令/推理样本，过滤后对学生 SFT。学生会继承教师偏差，生成数据还可能缺少真实分布多样性。

import torch.nn.functional as F

def distill_loss(student_logits, teacher_logits, labels, temperature=2.0,
                 alpha=0.5):
    hard = F.cross_entropy(student_logits, labels)
    soft_target = F.softmax(teacher_logits / temperature, dim=-1)
    soft_logprob = F.log_softmax(student_logits / temperature, dim=-1)
    soft = F.kl_div(soft_logprob, soft_target, reduction="batchmean")
    return alpha * hard + (1 - alpha) * (temperature ** 2) * soft

13.5 剪枝

先把概念落到可观察对象上。 剪枝删除对输出影响较小的权重、通道、头或层。只有硬件和 kernel 支持相应稀疏结构时，参数减少才可能转化为真实加速。 这一节的重点是建立因果链。先说明问题，再看计算或流程，最后用可重复实验验证结论。

把抽象概念还原成操作：

  1. 选择重要性指标。
  2. 按非结构化或结构化模式剪枝。
  3. 校准或再训练恢复。
  4. 在目标后端实测。

把它缩小到能逐项检查。 2:4 稀疏要求每四个权重中保留两个，规则性比任意零值更容易被专用硬件利用。 如果这个小例子还不能解释清楚，扩大数据只会让错误更难发现。

需要特别守住的边界。 只报告稀疏率、剪掉后仍用稠密 kernel、注意力头重要性在任务间变化、无恢复训练，都会夸大收益。 保留完整 trace 比猜原因更重要。找到第一次偏离预期的位置，通常比分析最终错误输出更高效。

本节练习：对小 MLP 做幅值剪枝与整列剪枝，比较参数、精度和真实耗时。 除代码外，请保留一份结果说明，标明环境、随机种子、输入规模和你主动检查过的边界。

非结构化剪枝把单个权重置零，只有稀疏 kernel 支持时才加速；结构化剪枝删除头、通道、FFN 单元或层，更容易获得真实速度，但影响能力更大。剪枝后通常需要恢复训练。报告参数稀疏率之外，还要报告真实硬件延迟。

13.6 FlashAttention

先建立直觉。 FlashAttention 通过分块与在线 Softmax 减少 HBM 读写，不显式保存完整注意力矩阵；它计算的是精确注意力，主要改变执行次序和内存复杂度。 读这一节时，先不要急着记缩写，而要不断追问：输入是什么，经过了哪些可观察变换，输出又怎样被验证。

把过程拆开看：

  1. Q/K/V 分块载入片上存储。
  2. 局部计算分数。
  3. 维护在线最大值与归一化和。
  4. 累积输出再写回。

最小例子。 T² 分数矩阵不落到 HBM，可显著降低长序列激活内存；算术量级仍大体是二次。 例子越小，越容易手算、打印中间量并判断实现是否偏离定义。

容易踩坑。 认为 FlashAttention 是稀疏近似、任何形状都同样加速、mask/dtype/kernel 不支持仍强开，都会出错。 排错时应保存失败输入和版本，先定位最早出现异常的步骤，再决定是否调整模型或超参数。

动手任务：在支持环境比较标准 SDPA 与 Flash 后端的峰值显存和时间，并验证输出误差。 完成后不要只保存最终输出，还要保存代码、配置、中间张量或检索结果，以及你对结果的解释。

标准注意力会把大 \(T\times T\) 中间矩阵频繁写入/读取 HBM。FlashAttention 用分块和在线 Softmax，在片上 SRAM 中计算块，减少 HBM I/O，同时保持精确注意力结果。它主要改善训练和 prefill 的注意力内存访问，并不消除 \(O(T^2)\) 的算术复杂度。

是否启用取决于 GPU、dtype、头维、mask 和软件版本。数值差异可能来自运算顺序，需设置合理容差测试。

13.7 KV Cache 与 PagedAttention

从问题出发。 KV Cache 保存每层历史 token 的 K/V，PagedAttention 把逻辑连续序列映射到固定大小物理块，减少预留和碎片并支持动态调度。 先把名词放到一边，沿着输入、状态变化和输出走一遍；只有能指出证据落在哪一步，概念才算真正掌握。

沿数据流逐步检查：

  1. 新请求分配缓存块。
  2. prefill 写入多 token K/V。
  3. decode 追加新位置。
  4. 请求结束释放或复用块。

用小数据走一遍。 不同请求长度差异大时，预先分配最大连续空间浪费严重；分页可按增长需要追加块。 先在纸面预测结果，再让程序打印中间状态；预测与运行不一致的地方，正是需要继续追查的知识缺口。

这里最容易出现的误解。 缓存键未含模型/adapter、prefix cache 跨租户泄漏、块太小元数据多、释放不完整，都会产生问题。 出现异常时先冻结样本、配置和环境，从最靠近输入的环节开始验证；不要同时更换模型、数据和超参数。

小实验：模拟可变长度请求的连续分配与分页分配，比较浪费率。 提交物应包含预期、运行证据和失败记录；只有别人能按同样步骤复现，结果才具有学习价值。

自回归解码时，历史 token 的 K/V 不变，缓存后新一步只计算新 token 的 Q/K/V。缓存随层数、并发、上下文与 K/V 头数增长。GQA/MQA、KV 量化、滑动窗口和前缀缓存都可降低或复用成本。

PagedAttention 把缓存划分为块并维护逻辑到物理映射，类似虚拟内存思想，减少碎片并支持灵活共享。Prefix caching 对重复系统提示或长公共前缀有效；缓存键要包含模型、token 序列和相关配置。

13.8 连续批处理与投机解码

先看它解决什么。 连续批处理动态插入和移除请求；投机解码让小模型提议多个 token，再由大模型一次验证，接受连续正确前缀。二者分别优化调度和解码串行性。 理解的标准不是会复述定义，而是能画出数据流、预测中间结果，并设计一个让错误暴露出来的检查。

可以把实现分成以下环节：

  1. 调度器按 token 预算选请求。
  2. draft 模型生成候选。
  3. target 模型并行验证。
  4. 按接受规则提交并继续。

一个可以手算的例子。 draft 越快且与 target 越一致，投机收益越大；频繁拒绝会浪费验证和 draft 计算。 这里故意不用大模型或大数据，因为小输入能把每个轴、分数和状态变化完整暴露出来。

排错时先看这些地方。 draft tokenizer 不一致、采样接受规则实现错误、只报理想接受率、不算额外显存，都会高估。 修复不能止于‘这次跑通’，还要把失败样本变成自动测试，防止同类问题在下一版本重新出现。

现在动手：用确定性 toy 分布实现投机解码，验证输出分布与直接采样一致。 把关键断言写进测试，并在 README 说明怎样运行、怎样判断正确、当前实现还不支持什么。

连续批处理按调度步动态组合不同请求，减少等待最慢序列的浪费。调度器需要在吞吐、公平、TTFT 和长请求饥饿之间取舍。

投机解码由小 draft 模型提出多个 token，大模型一次验证，接受一致前缀；输出分布可保持目标模型一致，但收益依赖接受率、draft 成本和硬件。若两个模型 tokenizer 或分布差异大，收益有限。

13.9 压测与回归

抓住这一节的主线。 优化必须用代表性负载做回归，质量、延迟、吞吐、显存、能耗和错误率一起报告。先 profile 再改，避免把瓶颈从 GPU 转移到排队或网络。 遇到新术语，先找它在系统中的位置：它读取什么、保存什么、改变什么，以及失败时会留下什么信号。

真正动手时按这个顺序走：

  1. 冻结模型与数据版本。
  2. 建立未优化基线。
  3. 一次引入一项优化。
  4. 做质量和性能双门禁。

先做最小实验。 量化后 P50 更快但 P99 因某些长请求 kernel 回退变慢，不能只报平均吞吐。 把随机性固定并保留中间量，这个例子就能成为后续优化时的回归基线。

别被表面现象带偏。 压测客户端不足、缓存命中率不真实、输出长度未固定、质量评测太小，都会给出错误结论。 先区分定义错误、实现错误、数据问题和性能瓶颈。四类问题需要的证据不同，不能靠盲目调参混在一起处理。

验证任务：设计一张优化实验表，至少记录模型、后端、dtype、长度分布、并发、硬件和置信区间。 实验结束后写三句话：观察到了什么、这些证据支持什么结论、还有哪种解释尚未排除。

建立性能矩阵：模型精度/量化格式 × 输入长度 × 输出长度 × 并发。记录质量基准、显存、TTFT、TPOT、吞吐、功耗和错误。预热后多轮测试，并报告软件、driver、硬件与采样参数。

量化质量不能只测困惑度；代码、数学、工具调用和长上下文对误差敏感程度不同。上线采用灰度和可回滚模型版本。

13.10 选择指南

推理瓶颈与优化手段

图 13-2　先判断瓶颈来自 prefill、decode、KV Cache 还是模型存储，再选择技术。 先把概念落到可观察对象上。 选择优化路线先识别约束：放不下、算不动、带宽不足、并发不够还是长尾过高。每种技术只对特定瓶颈有效。 这一节的重点是建立因果链。先说明问题，再看计算或流程，最后用可重复实验验证结论。

把抽象概念还原成操作：

  1. profile 定位瓶颈。
  2. 选择最小可逆改动。
  3. 在目标硬件验证 kernel。
  4. 保留质量回退和快速回滚。

把它缩小到能逐项检查。 显存不足可先考虑权重量化和 KV 管理；decode 带宽不足可考虑量化/并行；长上下文激活则关注 FlashAttention。 如果这个小例子还不能解释清楚，扩大数据只会让错误更难发现。

需要特别守住的边界。 同时上量化、蒸馏、剪枝后无法归因、追逐理论压缩率、忽略运维兼容，都会增加风险。 保留完整 trace 比猜原因更重要。找到第一次偏离预期的位置，通常比分析最终错误输出更高效。

本节练习：为 7B 单卡、70B 多卡和 CPU 边缘三个场景给出分步优化计划。 除代码外，请保留一份结果说明，标明环境、随机种子、输入规模和你主动检查过的边界。

  • 显存不足：先评估权重量化、GQA 模型、KV 量化和上下文上限。
  • TTFT 高：检查队列、输入长度、prefill batching 和 tokenizer。
  • TPOT 高：检查量化 kernel、张量并行通信、KV 带宽和 batch。
  • CPU/边缘：测试 GGUF/llama.cpp 的多种量化与线程设置。
  • 模型太大且质量可牺牲：考虑蒸馏或结构化剪枝。

练习：对随机矩阵比较 8/4-bit 量化误差；按公式估算不同 GQA 配置的 KV Cache；设计可重复的并发压测并画出吞吐—P99 曲线。

延伸阅读：FlashAttention、vLLM 文档、BitsAndBytes 官方说明、Transformers 量化选择。

本章配套代码

下面的脚本与正文使用相同符号。建议先在代码中打印形状和中间量，再运行断言；若依赖尚未安装，至少先阅读入口函数、输入输出和测试部分。

本章端到端实验：把知识变成可复现证据

本实验不是把本章代码重新抄一遍，而是把概念、实现、测试和解释串成一个小型工程。请新建独立目录，保存 README.md、环境文件、源代码、测试、运行日志和结果图。README 至少说明任务、输入输出、运行命令、预期现象、已知限制和复现条件。

实验步骤

  1. 对四种输入/输出长度组合记录 TTFT、TPOT 与总时间，画二维热图。
  2. 手写对称 int8 量化，比较不同 group size 的重构误差。
  3. 为 GPU 服务与 CPU 边缘部署分别写量化选择表。
  4. 在小分类任务比较无蒸馏与 logits 蒸馏，并分析错误样本。
  5. 对小 MLP 做幅值剪枝与整列剪枝，比较参数、精度和真实耗时。
  6. 在支持环境比较标准 SDPA 与 Flash 后端的峰值显存和时间，并验证输出误差。

每完成一步，先写下预期，再运行代码。若结果与预期不一致，不要覆盖旧日志；建立 failures.md，记录现象、假设、证据、修复和回归测试。这样得到的不是一次性 Demo，而是一份能证明你真正理解本章内容的实验档案。

验收标准

  • 全新环境能够按照 README 从头运行，依赖和随机种子已记录。
  • 关键函数至少有正常、边界和错误输入三类测试；涉及数值计算时检查有限值与合理误差。
  • 结果包含一个基线和至少一个受控改动，能够说明变化来自哪里。
  • 日志保留输入规模、耗时、内存或显存、软件版本和失败样本。
  • 结论区分“实验直接证明的事实”“根据事实做出的推断”和“仍未验证的猜想”。

本章自测

  1. 不看正文，用自己的话解释“Prefill 对整段输入并行计算并建立 KV Cache，通常计算密集”，并给出一个可以证伪的测试。
  2. 不看正文，用自己的话解释“量化用较少比特表示权重或激活”，并给出一个可以证伪的测试。
  3. 不看正文，用自己的话解释“bitsandbytes 常用于易用的运行时量化，GPTQ/AWQ 属于权重量化路线，GGUF 是 llama.cpp 生态常用文件格式与量化集合”，并给出一个可以证伪的测试。
  4. 不看正文，用自己的话解释“知识蒸馏让学生模型学习教师的软分布、中间特征或生成数据”，并给出一个可以证伪的测试。
  5. 不看正文，用自己的话解释“剪枝删除对输出影响较小的权重、通道、头或层”，并给出一个可以证伪的测试。
  6. 不看正文，用自己的话解释“FlashAttention 通过分块与在线 Softmax 减少 HBM 读写，不显式保存完整注意力矩阵”，并给出一个可以证伪的测试。

回答时先画图或写形状，再给结论。若只能说出术语而不能给出最小例子、边界条件和验证方法，说明这一节仍需要回到代码中重做。