推理优化

推理优化

本节摘要:两阶段定义 LLM 推理。Prefill 并行处理你的提示——算力受限;Decode 一次一个 token 生成——内存受限。每个优化都针对其一或两者。你在 4×A100 上部署 Llama 3 70B,单用户约 50 token/秒挺快,100 用户同时来,掉到每用户 3 token/秒——2.5 万美元/月的 GPU 账单服务得比人打字还慢。模型没变,变的是调度。本节带你实现 KV 缓存、连续批处理、PagedAttention(vLLM 的虚拟内存式 KV 管理)、投机解码(小模型起草大模型校验)、前缀缓存。这些技术是把 2.5 万/月账单降到 5 千/月、服务同样流量的分水岭。

学习目标

阅读完本节,你应当能够:

1. 实现 KV 缓存,消除自回归 token 生成中的冗余计算。
2. 解释 LLM 推理的 prefill 与 decode 阶段,以及各自不同瓶颈(算力受限 vs 内存受限)。
3. 实现连续批处理与 PagedAttention 概念,在并发请求下最大化 GPU 利用率。
4. 对比推理优化技术(KV 缓存、投机解码、Flash Attention)及其吞吐/延迟权衡。

一、问题与直觉

你在 4×A100 上部署 Llama 3 70B。单用户约 50 token/秒,挺快。然后 100 用户同时打端点,吞吐掉到每用户 3 token/秒。2.5 万美元/月的 GPU 账单,服务得比人打字还慢。

模型在 1 用户与 100 用户间没变——同样的权重、架构、数学。变的是怎么调度工作。朴素推理浪费 90%+ 可用 GPU 算力:一个等第 47 个 token 的用户占着一整个批槽位,而 GPU 内存总线在矩阵乘之间闲着。同时,一个新用户的 2000-token 提示本可填那段死时间做有用计算。

这不是扩容问题,是调度问题。本节的技术——KV 缓存、连续批处理、PagedAttention、投机解码、前缀缓存——是把 2.5 万/月账单降到 5 千/月、服务同样流量的分水岭。vLLM 在 4×A100-80GB 上服务 Llama 3 70B,低并发约 50 token/秒/用户,通过连续批处理与 PagedAttention 在 100 并发下维持 15~25 token/秒/用户。无这些优化,同样硬件在那个并发下只有 5 token/秒。同样 GPU同样模型,4 倍吞吐。

Prefill vs Decode

每个 LLM 推理请求有两个截然不同的阶段。Prefill 处理整个输入提示——所有 token 已知,注意力可跨全序列并行算,这是大矩阵乘,GPU 核心忙,瓶颈是算力(A100 在 BF16 下 312 TFLOPS,70B 模型 4096-token 提示 prefill 单 A100 约 400ms)。Decode 一次生成一个输出 token——每个新 token 注意所有先前 token,但每次前向只产一个 token,权重矩阵与 prefill 同大,但你用单向量乘而非矩阵,GPU 核心微秒级完成然后等下一批权重从内存到,瓶颈是内存带宽(A100 有 2TB/秒,70B FP16 是 140GB,读完整模型一次 70ms——这是单步 decode 的下限)。

ops:byte 比(算术强度)抓住这个权衡:每从内存读一字节做多少次运算。Prefill 批 4096 token 时,约每权重 4096 次乘加,比高——算力受限。Decode 批 1 时,约每权重 1 次运算,比低——内存受限。根本洞见:decode 内存受限是因为你为产一个 token 要读整个模型。下面每个优化要么减少读取、要么增加每次读取处理的 token 批、要么完全避免读取。

KV 缓存

注意力里,每个 token 的查询注意所有先前 token 的键和值向量。不缓存的话,生成第 N 个 token 要为所有 N-1 个先前 token 重算键值投影——第 1 个 token 在生成第 2 时投影一次,第 3 时再投影一次……到第 1000 时已投影 999 次。KV 缓存存住所有先前 token 的键值投影,生成第 N 时只算第 N 的键值,再与缓存中 1 到 N-1 的 K/V 拼接。

KV 缓存显存公式:2 × 层数 × KV 头数 × 头维 × 序列长 × 每参数字节数。Llama 3 70B(80 层、GQA 8 个 KV 头、头维 128、BF16):每 token 2×80×8×128×2 = 327,680 字节 = 320KB;4096 token = 1.28GB;128K token = 40GB。一条 128K 上下文对话消耗 40GB KV 缓存——半张 A100。100 并发各 4K 即需 128GB——这就是 KV 缓存管理是推理优化核心挑战的原因。

连续批处理

静态批处理等到 N 个请求的批到齐,一起处理,等所有完成才接新请求——若一个要 500 token、另一个要 10,短的在完成后闲 490 步。连续批处理(迭代级批处理)在任意请求完成时立即插入新请求到批,每个 decode 步重评批。10 token 完成的请求立即被等待的请求替换。吞吐提升取决于输出长度变化——均匀长度时连续批与静态批持平,变长时(常见情况)连续批可高 2~5 倍吞吐,因为 GPU 槽位永不被空占。

PagedAttention

每请求的 KV 缓存是一块连续内存。请求来去时内存碎片化——就像操作系统的 RAM 碎片。4K-token 请求需 1.28GB 连续,即使你有 2GB 空闲总量也未必有 1.28GB 连续,你要么浪费内存要么拒绝请求。PagedAttention(vLLM)把操作系统式虚拟内存用于 KV 缓存:不每请求分配一大块连续,而是分配固定大小「页」(典型每页 16 token),页可在物理 GPU 内存任意位置,页表把每请求的逻辑序列位置映射到物理页位置。PagedAttention 还为共享前缀启用写时复制:若 50 个请求共享同一系统提示,该系统提示的 KV 缓存页只存一次,被 50 个请求引用,只在请求分叉(不同用户消息)时才得自己的页——对有共享系统提示的应用大幅省内存。vLLM 报告通过 PagedAttention 内存浪费近零(4% vs 朴素分配的 6080%)。

投机解码

Decode 慢是因为串行——生成一个、喂回、生成下一个。但若你能便宜猜下 5 个 token,然后一次性全校验呢?投机解码用一个又小又快的草稿模型生成 K 个候选 token,大目标模型然后单次前向处理全部 K 个候选(看起来像 prefill——并行、算力受限、高效)。若目标模型同意草稿的预测,你在一个目标前向时间里收下全部 K 个 token;若在第 j 位不同意,收下 1 到 j-1、丢弃其余。加速取决于接受率——草稿预测匹配目标的频率。Llama 3 8B 给 Llama 3 70B 起草,自然语言上接受率 7085% 典型,转化为 23 倍 decode 加速。

三种方法:草稿-目标(Leviathan 等,独立小模型,接受率 7085%,代价是草稿模型内存)、EAGLE(Li 等,目标模型隐藏状态上的轻量头,接受率 7590%,约 1% 额外参数)、N-gram 查表(token n-gram 表,接受率 40~60%,代价可忽略)。EAGLE 在目标模型自身表示上训一个小的自回头(而非独立模型),以最小额外内存获更高接受率。N-gram 投机维护当前上下文或预建语料的 n-gram 续接表,若草稿匹配之前同对话出现过的(重复模式、代码、结构化输出),零神经网络开销触发。投机解码是数学精确的——输出分布与目标模型分布相同,不是近似;校验步骤确保每个被接受的 token 恰有目标模型会赋的概率。

前缀缓存

许多请求共享同一前缀:聊天机器人系统提示、RAG 上下文块、少样本示例集。无前缀缓存,每请求从头重算这些共享 token 的 KV 缓存。前缀缓存存住常见前缀的 KV 缓存,跨请求复用。新请求带已知前缀到达时,系统复制(或引用)缓存的 KV 项,只算唯一后缀的 KV。对 2000-token 系统提示(所有请求共享),前缀缓存消除每请求约 400ms prefill——100 请求/秒即每秒省 40 秒 GPU 计算,多于一张 GPU 的活。SGLang 的 RadixAttention 用基数树(trie)按 token 内容索引前缀实现前缀缓存,支持部分前缀匹配(与缓存项共享 2000 前缀里的 1500,就复用那 1500 只重算 500)。

推理引擎

三大引擎主导生产 LLM 服务:vLLM(PagedAttention、连续批处理,通用服务、兼容性最广)、SGLang(RadixAttention 前缀缓存、结构化生成,多轮聊天/约束解码,常比 vLLM 快 2~5 倍靠前缀复用)、TensorRT-LLM(NVIDIA 核融合、FP8,单 GPU 吞吐最高但仅 NVIDIA 卡、配置多)。

ops:byte 框架

你无法优化你不度量的。ops:byte 比告诉你算力受限还是内存受限,决定哪些优化要紧。计算屋顶:GPU 峰值 FLOPS;内存屋顶:峰值带宽 × ops:byte 比。ops:byte 低(decode、小批)时撞内存带宽屋顶,加算力(更高时钟、更多核)没用,需减内存读取(量化、KV 压缩)或增大批把读取摊到更多有用工作上。ops:byte 高(prefill、大批)时撞计算屋顶,内存带宽优化没用,需更快 GPU、核融合、降精度挤更多 FLOPS。A100 上的交叉点约 ops:byte=156(312 TFLOPS / 2TB/秒),低于 156 内存受限,高于算力受限。

二、从零实现

Step 1:KV 缓存

构建多头 KV 缓存,每层每头存键值投影,展示内存增长模式。

class KVCache:
    def __init__(self, num_layers, num_heads, head_dim, max_seq_len, dtype=np.float16):
        self.k_cache = np.zeros((num_layers, num_heads, max_seq_len, head_dim), dtype)
        self.v_cache = np.zeros((num_layers, num_heads, max_seq_len, head_dim), dtype)
        self.seq_len = 0
    def update(self, layer_idx, new_keys, new_values):
        n = new_keys.shape[1]; end = self.seq_len + n
        self.k_cache[layer_idx, :, self.seq_len:end, :] = new_keys
        self.v_cache[layer_idx, :, self.seq_len:end, :] = new_values
        return self.k_cache[layer_idx, :, :end, :], self.v_cache[layer_idx, :, :end, :]
​

Step 2:连续批处理模拟器

模拟静态与连续批处理的调度差异,统计平均延迟、p50/p99、吞吐。

Step 3:前缀缓存(字典树)

基于 trie 的前缀缓存,存共享前缀的 KV 项,支持查找、插入、命中率统计。

Step 4:投机解码模拟器

模拟草稿-目标投机解码,可配接受率。核心:草稿生成 K token、目标单次前向校验全部、按接受率收下前缀、拒绝处从修正分布重采。

def speculative_decode(draft, target, context, num_spec=5):
    draft_tokens = draft.generate(context, num_spec)
    target_probs = target.get_probs(context, draft_tokens)
    accepted = 0
    for i, tok in enumerate(draft_tokens):
        r = np.random.random()
        if r < min(1.0, target_probs[i][tok] / draft.get_probs(...)[tok]):
            accepted += 1; context.append(tok)
        else:
            context.append(np.random.choice(vocab, p=target_probs[i])); break
    if accepted == num_spec:  # 全对,送一个奖励 token
        context.append(np.random.choice(vocab, p=target.get_bonus_probs(context)))
    return {"speedup": sequential_cost/speculative_cost, "avg_accepted": accepted}
​

Step 5:KV 缓存显存分析器

算真实模型配置的 KV 缓存显存需求,回答「这个模型在 N 张卡上能服务多少并发用户」。

三、框架对比

vLLM

from vllm import LLM, SamplingParams
llm = LLM(model="meta-llama/Llama-3-70B-Instruct", tensor_parallel_size=4,
          enable_prefix_caching=True, gpu_memory_utilization=0.9)
outputs = llm.generate(["解释推理优化。"], SamplingParams(temperature=0.7, max_tokens=256))
​

SGLang(前缀缓存 + 结构化输出)

import sglang as sgl
@sgl.function
def classify(s, text):
    s += sgl.system("你是分类器,只输出 JSON。")
    s += sgl.user(f"分类:{text}")
    s += sgl.assistant(sgl.gen("result", regex=r'\{"label":"(正|负|中)"\}'))
​

TensorRT-LLM

把模型编译成优化 NVIDIA GPU 核,融合操作、用 FP8、集成 Triton。

四、可复用产物

本节产出 outputs/skill-inference-optimization.md——一个诊断与优化 LLM 推理服务的技能,含 prefill/decode 瓶颈分析、KV 缓存预算、批处理策略、引擎选型。

五、练习

1. (Easy) 改 KV 缓存分析器对比 FP16/FP8/INT4 KV 量化,Llama 3 70B 4K 上下文在 4×A100-80GB 上各能支持多少并发(INT4 应约 4 倍)。

2. (Medium) 扩展连续批处理模拟器追踪 GPU 利用率(每步批槽位填充率),对 50 个帕累托分布输出长度的请求绘利用率曲线,连续批应保持 >80%。

3. (Medium) 实现 GQA 版 KV 缓存,Llama 3 70B 用 64 查询头但只 8 KV 头,算省的显存(8 倍 KV 缓存缩减)。

4. (Hard) 构建 LRU 驱逐的前缀缓存,max_entries=500,生 1000 请求其中 60% 共享 5 个常见前缀,测命中率(应保持 >55%)。

5. (Hard) 扩展投机解码模拟器实现树状投机(EAGLE-2 风格),每 3 层 2 分支=8 叶候选,对比线性投机的每轮接受 token 数。

本节要点回顾

1. Prefill 算力受限,Decode 内存受限:prefill 并行处理全提示满吞吐,decode 一次一 token 干等内存读权重。
2. ops:byte 比定瓶颈:低则内存受限(量化、KV 压缩),高则算力受限(核融合、FP8),A100 交叉点约 156。
3. KV 缓存省重复计算:存住先前 token 的 K/V,生成新 token 只算新 Q,Llama 3 70B 128K 上下文 KV 缓存 40GB。
4. 连续批处理填空槽:请求完成立即插新请求,变长输出下吞吐高 2~5 倍。
5. PagedAttention 虚拟内存化 KV:固定页替代连续块,消除碎片,写时复制共享前缀,浪费从 60~80% 降到 ~4%。
6. 投机解码数学精确:小模型起草、大模型单次前向校验,输出分布与目标相同,2~3 倍加速。
7. 接受率定加速:8B 给 70B 起草,自然语言接受率 7085%,EAGLE 用目标自身表示更高(7590%)。
8. 前缀缓存省重复 prefill:共享系统提示的 KV 复用,SGLang 的 RadixAttention 用 trie 部分前缀匹配。
9. 三大引擎:vLLM(通用)、SGLang(前缀缓存+结构化)、TensorRT-LLM(NVIDIA 单卡最高)。
10. TTFT 由 prefill 主导:首 token 延迟,长提示时关键指标。

下一节,完整 LLM 流水线:把分词→预训练→SFT→对齐→评估→量化→服务串成一条可复现的端到端运行。