CS 与 LLM 学习知识库

推理系统

30分钟阅读

一张 H100 每秒能做约 989 万亿次 bf16 乘加，一个 13B 模型生成一个 token 只需要约 260 亿次运算。按算力算，它应该每秒吐出几万个 token；实际单个请求只有一百多 token/s，GPU 的算力利用率不到 1%。时间去哪了？答案几乎全在”搬数据”上。这篇从这个落差出发，一路推出现代推理引擎（以 vLLM 为例）里的每个零件。

前置：Transformer 架构、资源核算与 Scaling Laws（算术强度）、GPU 与 Kernel（HBM 与带宽）

一个 token 一个 token 地生成

训练时整句话都已知，模型一次前向就能对所有位置同时算 loss，序列维度可以并行，矩阵又大又方，GPU 吃得很饱。生成不一样：第 t+1 个 token 依赖第 t 个 token 的采样结果，只能一步一步来，这叫自回归。

最朴素的写法是每生成一个 token，把”prompt + 已生成的全部 token”整段重新喂给模型，只取最后一个位置的输出。一次前向里 attention 的开销是 O(T²)，生成 T 个 token 总共就是 O(T³)。而这里绝大部分计算是重复的：Transformer 是因果的，后面追加 token 不会改变前面任何位置在任何一层的激活（双向模型就不行，加一个词全部都变）。

既然前面位置的激活不变，attention 里它们贡献的 K、V 向量也不变。把每一层每个位置的 K、V 存下来，下一步只给新 token 算 Q、K、V，用新 Q 去和”历史 K + 新 K”做点积即可。这块缓存就是 KV cache。下面用 4 层单头 attention 对比两种写法，输出完全一致，速度差距随长度拉大：

import torch, time
torch.manual_seed(0)
D, L = 64, 4                         # 4 层单头 attention，够看清 KV cache 的账
Wq, Wk, Wv = (torch.randn(L, D, D) / D**0.5 for _ in range(3))
 
def attn(x, l, cache=None):
    q, k, v = x @ Wq[l], x @ Wk[l], x @ Wv[l]
    if cache is not None:            # 只算新 token 的 K/V，历史的直接从 cache 拿
        cache[l][0].append(k); cache[l][1].append(v)
        k, v = torch.cat(cache[l][0]), torch.cat(cache[l][1])
    T, S = q.shape[0], k.shape[0]
    # 新 token 排在最后，第 i 个新 token 能看到前 S-T+i 个位置
    mask = torch.ones(T, S).tril(S - T).bool()
    s = (q @ k.T / D**0.5).masked_fill(~mask, float("-inf"))
    return x + s.softmax(-1) @ v
 
def forward(x, cache=None):
    for l in range(L): x = attn(x, l, cache)
    return torch.nn.functional.normalize(x, dim=-1) * D**0.5  # 输出直接当下一个 token 的 embedding
 
prompt = torch.randn(16, D)
def gen_naive(n):
    seq = prompt.clone()
    for _ in range(n):               # 每步把整段历史重算一遍
        seq = torch.cat([seq, forward(seq)[-1:]])
    return seq[16:]
def gen_cached(n):
    cache = [([], []) for _ in range(L)]
    out = forward(prompt, cache)[-1:]       # prefill：整个 prompt 一次过
    res = [out]
    for _ in range(n - 1):                  # decode：每步只喂 1 个 token
        out = forward(out, cache); res.append(out)
    return torch.cat(res)
 
for n in (64, 256, 512):
    t = time.perf_counter(); a = gen_naive(n); t1 = time.perf_counter() - t
    t = time.perf_counter(); b = gen_cached(n); t2 = time.perf_counter() - t
    print(f"n={n:3d}  naive {t1*1000:7.1f} ms  cached {t2*1000:6.1f} ms  max diff {(a-b).abs().max():.1e}")
n= 64  naive    21.8 ms  cached   10.1 ms  max diff 1.9e-06
n=256  naive   177.0 ms  cached   63.1 ms  max diff 2.1e-06
n=512  naive   619.4 ms  cached  211.6 ms  max diff 2.5e-06

差值在 1e-6 量级，只是浮点累加顺序不同。cached 版本还不够快，因为每步 torch.cat 都把整段 cache 复制一遍，真实引擎会预先分配显存、原地写入。这个”怎么放 KV cache”的问题后面会变成 PagedAttention。

代码里已经出现了推理的两个阶段。prefill：把整个 prompt 一次送进模型，填满 KV cache，产出第一个 token，和训练前向一样可以沿序列并行。decode：之后每步只输入 1 个新 token，读整个 KV cache，追加一条。用户感受到的两个延迟指标正好对应它们：TTFT（time to first token）主要是 prefill 的时间；TPOT（time per output token）是 decode 每步的时间。服务方还关心吞吐：所有请求合计每秒产出多少 token。

停一下，检查理解单选

为什么 GPT 这类 decoder 能用 KV cache，而 BERT 这类双向 encoder 逐词追加输入时不能？

瓶颈不是算力，是带宽

有了 KV cache，计算量从 O(T³) 降下来了，可 decode 依然慢。要看清原因，得用 资源核算与 Scaling Laws 里的算术强度：一个操作做的 FLOPs 除以它读写的字节数。H100 的算力 989 TFLOP/s 除以 HBM 带宽 3.35 TB/s 约等于 295，算术强度低于 295 的操作，GPU 大部分时间在等数据，叫 memory-bound。

先看 MLP。输入是 B 个序列、每个 T 个 token，权重矩阵 D×F。读权重 2DF 字节（bf16），计算 2·B·T·D·F FLOPs，D、F 远大于 B·T 时强度约等于 B·T。prefill 时 T 是整个 prompt 长度，几百上千，轻松 compute-bound。decode 时 T = 1，强度只剩 B，也就是同时在跑的请求数。每一步都要把全部权重从 HBM 读一遍，却只给 B 个 token 用。

再看 attention。decode 时每个序列用自己的 1 个 query 去和自己的 S 个历史 K、V 做点积，FLOPs 约 4·S·D，读的字节约 4·S·D（K 和 V 各 2 字节/数）。强度是 S·T/(S+T)，T = 1 时等于 S/(S+1)，小于 1，而且和 B 无关：MLP 的权重是所有序列共用的，batch 越大分摊越多；KV cache 是每个序列私有的，batch 加大，字节和 FLOPs 同比例增长。

既然 decode 是 memory-bound，时间就约等于”这一步要读多少字节 ÷ 带宽”。每一步要读全部参数加上所有序列的 KV cache。每个 token 的 KV cache 大小是：

L 是层数，K_kv 是 KV 头数，H 是每头维度。Llama 2 13B（40 层、40 个 KV 头、H=128）每个 token 800 KiB，一条 1024 token 的序列约 0.84 GB。用这个公式算一张 H100 上的延迟和吞吐：

def kv_bytes_per_token(L, kv_heads, head_dim, bytes_per=2):
    return L * kv_heads * head_dim * 2 * bytes_per   # 2 = K 和 V
bw = 3.35e12                          # H100 HBM 带宽 B/s
params = 2 * 13.0e9                   # Llama 2 13B, bf16
for name, K in (("MHA K=40", 40), ("GQA K=8 ", 8)):
    kv = kv_bytes_per_token(40, K, 128) * 1024        # 每序列 1024 token
    for B in (1, 64):
        total = B * kv + params
        lat = total / bw
        print(f"{name} B={B:3d}  显存 {total/1e9:6.2f} GB  {lat*1000:5.1f} ms/step  {B/lat:6.0f} tok/s")
MHA K=40 B=  1  显存  26.84 GB    8.0 ms/step     125 tok/s
MHA K=40 B= 64  显存  79.69 GB   23.8 ms/step    2691 tok/s
GQA K=8  B=  1  显存  26.17 GB    7.8 ms/step     128 tok/s
GQA K=8  B= 64  显存  36.74 GB   11.0 ms/step    5836 tok/s

开头那个”125 token/s”就是这么来的：B = 1 时每一步都在读 26 GB 参数。把 batch 加到 64，参数的读取被 64 个请求分摊，吞吐涨到 2691，但每一步要多读 64 份 KV cache，单请求延迟从 8 ms 涨到 24 ms，显存也快满了。调 batch 时延迟和吞吐是此消彼长的：吞吐 = B·带宽 / (B·KV + 参数)，B 趋于无穷时上限是 带宽 / 每序列 KV ≈ 4000 token/s。

后两行把 KV 头数从 40 降到 8（这就是 GQA，见 注意力变体与 MoE；示意里没扣掉 W_K、W_V 变小省下的参数）。batch 不变时延迟和吞吐同时变好，因为要搬的字节少了。这给出了推理优化的总纲：decode 的时间 ≈ 字节数 ÷ 带宽，任何减少字节的手段都直接变成速度。减少 KV cache 的结构改动（GQA、MLA、滑动窗口、跨层共享）属于模型架构，放在 注意力变体与 MoE；推理侧能做的是下面三件事：把每个数存得更小（量化），让一次读取产出不止一个 token（投机采样），以及把显存管理好、塞进更大的 batch（PagedAttention 与 continuous batching）。

停一下，检查理解多选

decode 阶段（每步每个序列 1 个新 token），下面哪些说法正确？

量化：每个数少存几位

既然时间正比于字节，最直接的办法是让每个数占更少字节。量化把 bf16 的权重（或激活、KV cache）映射到 int8、fp8、int4。最基本的线性量化：选一个 scale 和 zero point，

例如 s = 0.1、z = 4 时 5.2342 量化成 56，反量化回 5.2。scale 通常按通道或按小块单独取，块越小，离群值影响越小。

格式字节/数典型用途
bf162推理默认
fp8（E4M3，最大约 448）1H100 起有硬件支持，权重、激活、KV cache 都可用
int81权重量化，激活量化要处理离群值
int40.5只量化权重（weight-only），计算时反量化回 bf16

权重从 bf16 换成 int4，decode 时读参数的字节变成 1/4，memory-bound 下几乎直接变成速度。注意 weight-only 量化不减少 FLOPs，乘法前还要反量化，所以它主要帮 decode（读得少），对 compute-bound 的 prefill 帮助有限；想让 prefill 也变快，需要权重和激活都低精度、用 fp8/int8 的 tensor core 直接算。

QAT 在训练时模拟量化误差，效果好但要重训；PTQ 训练完再用少量校准数据定 scale，便宜，是开源模型的主流。两个常见 PTQ 方法：GPTQ 逐层量化，每量化一部分权重，就用二阶信息（Hessian）调整还没量化的权重来抵消误差；AWQ 观察到少数激活通道数值特别大，和它们相乘的权重更重要，于是量化前先按激活大小放大这些权重通道（等价地缩小对应激活），让它们在量化格点上更精确。KV cache 同样可以量化，fp8 KV cache 能让同一张卡放下两倍的 token。

投机采样：一次读权重，产出几个 token

量化是有损的。有没有一种加速，输出分布和原模型完全一样？

关键观察来自前面的算术强度：decode 时 GPU 大部分算力闲着。让目标模型一次处理 K 个 token（像一次小 prefill），和处理 1 个 token 的耗时差不多，因为两者都要把权重读一遍。所以”检查 K 个 token”几乎和”生成 1 个 token”一样贵。投机采样（speculative decoding）利用这一点：先让一个小的 draft 模型 p 便宜地连猜 K 个 token，再让目标模型 q 一次前向并行算出这 K 个位置上的分布，逐个决定接受还是拒绝。

规则是：对 draft 采出的 token x，以概率 min(1, q(x)/p(x)) 接受；一旦拒绝，从残差分布 max(q − p, 0)（归一化后）重新采一个 token，本轮结束。如果 K 个全被接受，目标模型那次前向还顺带给出了第 K+1 个位置的分布，再白送一个。所以每轮至少产出 1 个、最多 K+1 个 token。

为什么无损？对任意 token x，得到它有两条路：draft 采到 x 且被接受，概率 p(x)·min(1, q/p) = min(p(x), q(x))；或者 draft 采到别的被拒，再从残差采到 x，残差恰好补上 q(x) − min(p(x), q(x))。两条相加正好是 q(x)。用模拟验证：

import random
random.seed(0)
q = {"A": 0.2, "B": 0.5, "C": 0.3}   # 目标模型
p = {"A": 0.6, "B": 0.2, "C": 0.2}   # draft 模型：猜得不太准
def sample(d): return random.choices(list(d), weights=d.values())[0]
def spec_step():
    x = sample(p)
    if random.random() < min(1, q[x] / p[x]): return x, True
    res = {t: max(q[t] - p[t], 0) for t in q}   # 拒绝后只从"q 比 p 多出来"的部分采样
    return sample(res), False
N = 200_000; cnt = {t: 0 for t in q}; acc = 0
for _ in range(N):
    x, ok = spec_step(); cnt[x] += 1; acc += ok
print({t: round(c / N, 3) for t, c in cnt.items()}, "accept", round(acc / N, 3))
print("理论接受率 sum min(p,q) =", sum(min(p[t], q[t]) for t in q))
{'A': 0.2, 'B': 0.501, 'C': 0.299} accept 0.598
理论接受率 sum min(p,q) = 0.6000000000000001

draft 明显偏爱 A，输出分布却和 q 一致。单步接受率 α = Σ min(p, q)，两个分布越像越高。假设各位置独立，一轮期望产出 1 + α + α² + … + α^K = (1 − α^(K+1))/(1 − α) 个 token。α = 0.8、K = 4 时约 3.36 个；K 加到 8 也只有 4.33 个，因为后面的 draft token 要前面全部被接受才有意义，而 draft 的成本随 K 线性增长。Chen 等人在 Chinchilla 70B 上报告 2–2.5 倍加速。

两个限制要知道。第一，收益来自 decode 时闲置的算力；batch 已经很大、decode 接近 compute-bound 时，验证 K 个 token 不再免费，加速会缩水。第二，draft 要和目标模型共用 tokenizer、分布要像，通常从目标模型蒸馏，或像 EAGLE 那样在目标模型上加小预测头。

停一下，检查理解单选

draft 分布 p = [0.5, 0.5]，目标分布 q = [0.2, 0.8]（两个 token A、B）。draft 采到 A，接受概率是多少？如果被拒绝，输出什么？

PagedAttention：像操作系统一样管 KV cache

前面的账显示，吞吐靠大 batch，而 batch 能开多大取决于显存里能放多少 KV cache。早期推理系统给每个请求按”最大长度”预留一整段连续显存：请求的输出长度事先不知道，只能按上限留。结果一个最终只生成 10 个 token 的请求占着 2048 个位置（内部碎片），各请求区段之间的空隙又凑不成一整段（外部碎片）。vLLM 论文测得这类系统里真正存了 token 的只占 KV 区域的 20%–40%。

这和操作系统早年的内存管理是同一个问题，解法也一样：分页（见 虚拟内存）。PagedAttention 把 KV cache 切成固定大小的块（block，vLLM 默认每块 16 个 token），物理块可以散落在显存任何位置。每个请求有一张 block table，把逻辑块号映射到物理块号，attention kernel 按表去取 K、V。新 token 填满当前块才分配下一块，浪费只发生在每个请求的最后一块里，论文报告浪费降到 4% 以下，同样的显存能放下几倍的请求。

分页还顺带带来共享。同一 system prompt 的多个请求、并行采样同一 prompt 的多个回答，前缀的块完全相同，可以让多张 block table 指向同一批物理块，用引用计数管理，要写入一个共享块时再复制（copy-on-write，和 fork 一样）。vLLM 的 prefix caching 进一步把块按”内容哈希”索引，新请求来时先查已有的前缀块，命中就跳过那部分 prefill，直接降低 TTFT。

停一下，检查理解多选

关于 PagedAttention，下面哪些说法正确？

Continuous batching：每一步都重新组 batch

显存有了，还要把请求真正组成大 batch。线上请求到达时间不同、输出长度不同。传统的静态 batching 凑齐一批一起跑，必须等最长的那个生成完，整批才能退出，先结束的请求空占位置，新来的请求在门外等。

Orca（OSDI 2022）的办法是 iteration-level scheduling：调度粒度从”一个 batch 跑完”缩到”一步 decode”。每一步结束后，生成完的请求立刻移出、释放 KV 块，排队的新请求立刻加入。batch 每步都在变，所以叫 continuous batching。

长度各不相同的请求怎么在一步里一起算？MLP、投影这些对每个 token 独立的操作，不关心 token 属于哪个请求，直接把所有请求的 token 拼成一个 [总 token 数, D] 的矩阵做一次 matmul，权重只读一次；只有 attention 需要按请求分开，各读各的 KV 块。这样完全不需要 padding。

还剩一个冲突：一个 30K token 的长 prompt 进来，它的 prefill 会独占好几百毫秒，这一步里所有正在 decode 的请求都被卡住，用户看到输出突然停顿。chunked prefill 把长 prompt 切成几段，每一步只算一段，和 decode 请求共享这一步的 token 预算。prefill 是 compute-bound、decode 是 memory-bound，混在同一步里反而互补：decode 带来的权重读取顺便喂给 prefill 的计算。规模再大时还可以把两者放到不同 GPU 上（PD 分离），代价是要传输 KV cache。

停一下，检查理解单选

开启 chunked prefill 后，一个 30K token 的新请求到来，而本步 token 预算是 8K。会发生什么？

拼起来：vLLM 的架构与接入一个模型

把以上零件装进一个系统，就是 vLLM（V1 架构）的主干：

请求 -> API Server（OpenAI 兼容接口、分词）
     -> EngineCore（独立进程，循环：schedule -> execute -> update）
          |- Scheduler + KVCacheManager（决定本步每个请求算几个 token、分配 KV 块）
          |- Executor -> Worker（每张 GPU 一个）-> ModelRunner（拼输入、跑模型、采样）
     -> Detokenizer -> 流式返回

调度器把 prefill 和 decode 统一成一个问题：本步给每个请求分配多少个新 token。每步有 token 上限（max_num_batched_tokens）和请求数上限（max_num_seqs）。先给正在 decode 的请求各分 1 个（开投机采样则多几个），剩余预算分给排队的新请求，prompt 太长就只算一段，这就是 chunked prefill。然后给每个请求申请本步需要的 KV 块；块不够时抢占优先级最低的请求，释放它的块，之后重新计算。

显存预算在启动时定好：先加载权重，再用一次最坏情况的 dummy 输入跑前向，测出激活峰值，剩下的 总显存 × gpu_memory_utilization − 权重 − 激活峰值 全部切成 KV 块。比如 80 GB 卡跑 Llama 3 8B（权重约 16 GB，每 token KV 128 KiB），假设激活峰值 4 GB：80 × 0.9 − 16 − 4 = 52 GB，约能放 40 万个 token 的 KV。

decode 步很小，kernel launch 的 CPU 开销会变得显眼，所以 ModelRunner 用 CUDA Graph 把整步录下来重放。attention kernel 本身用 FlashAttention 一类的实现，按 block table 读分页的 K、V。

接入一个新模型，本质是把 HuggingFace 的建模代码改写成符合上面这套运行时的样子，分四步：

  1. config：vLLM 读模型目录的 config.json，按 architectures 字段（如 Qwen2ForCausalLM）在注册表里找实现类，超参从 hf_config 取。
  2. 模型类：构造函数签名统一为 __init__(self, vllm_config, prefix="")，prefix 是模块在 state dict 里的完整名字，量化配置和 attention 层绑定 KV cache 都靠它。forward 收到的 input_ids、positions 是压平的一维张量，没有 batch 维、没有 padding，每个 token 属于哪个请求由调度器给 attention 层的元数据决定。所以要删掉 HF 代码里的 attention_mask、past_key_values 和 dropout、loss 等训练逻辑。线性层要换成支持张量并行（见 并行训练）的版本：QKVParallelLinear 按头切 Q/K/V、MergedColumnParallelLinear 合并 gate 和 up、RowParallelLinear 做输出投影并在末尾 all-reduce。直接搬 nn.Linear 的模型开 TP=4 时每张卡仍放完整权重，等于白开。
  3. load_weights：checkpoint 里是分开的 q_proj、k_proj、v_proj，vLLM 里是合并的 qkv_proj。用一张映射表把每个张量路由到合并参数的对应分片，每个并行层自带的 weight_loader 负责按当前卡号切出自己那片。函数返回已加载参数名的集合，加载器据此检查有没有漏加载的权重，这是最常见的 bug 来源。
  4. 注册：加进内置模型注册表，或者通过插件机制在外部注册，不改 vLLM 源码。
停一下，检查理解单选

在 vLLM 的模型类里，forward 收到的 input_ids 为什么是一维的、没有 batch 维？

记住这三件事

  1. decode 每步都要把参数和全部 KV cache 从显存读一遍，时间约等于字节数 ÷ 带宽；减少字节的手段（GQA、量化、fp8 KV cache）都直接变成速度。
  2. batch 越大吞吐越高、单请求延迟越差；投机采样用闲置的算力一次验证多个 token，且输出分布与目标模型完全一致。
  3. 推理引擎的核心是显存管理和调度：PagedAttention 消除碎片、共享前缀，continuous batching 与 chunked prefill 让 batch 每步都满且不卡顿。

延伸阅读

反向链接