土法炼钢兴趣小组的算法知识备份

【大模型基础设施工程】11：推理引擎基础

2026-04-22 · architecture · ai-infra · #llm #infra #inference #prefill #decode #kv-cache #gqa #mla #continuous-batching #ttft #flash-decoding

文章导航

分类入口
architectureai-infra
标签入口
#llm#infra#inference#prefill#decode#kv-cache#gqa#mla#continuous-batching#ttft#flash-decoding

目录

训练把模型”炼”出来,推理把模型”用”出来。前者是离线任务、按天计量、看吞吐;后者是在线服务、按毫秒计量、看尾延迟。训练工程师可以容忍一次 checkpoint 恢复花半小时,推理工程师被 P99 延迟超了 50ms 就要被叫到会议室。

本篇是系列的推理篇开篇,目的是把”大模型推理”这件事从第一性原理讲清楚:为什么要有 KV Cache、为什么要区分 Prefill 和 Decode、为什么 vLLM 的 Continuous Batching 能把吞吐翻几倍、为什么 DeepSeek / Kimi 都要做 Prefill / Decode 分离。之后的第 12–16 篇会分别展开 PagedAttention、主流引擎对比、量化、推测解码、长上下文。

一、推理 vs 训练的工程差异

1.1 两种 workload 的本质区别

训练是一个离线、同步、吞吐导向的批处理任务。一次 step 处理一个固定 shape 的 batch,forward + backward + optimizer step 三段式,所有样本同生同死。你不关心单样本何时返回,只关心每秒能吃多少 token、每轮 epoch 多长、loss 收敛多快。

推理是在线、异步、延迟导向的服务。请求一个一个来,长度不等、到达时间不等、产出 token 数不等。你不能让后到的请求等前面的请求跑完,不能让短请求陪长请求一起排队。你需要在亚秒级给出首 token,在几十毫秒级给出每个后续 token,还要把 GPU 利用率拉起来。

维度 训练 推理
时效 离线,可以跑几周 在线,毫秒级 SLO
输入 固定 shape 的 batch 变长、动态到达
方向 Forward + Backward 只 Forward
显存 参数 + 梯度 + optimizer state + 激活 参数 + KV Cache
瓶颈 算力 / 通信 访存 / KV 空间
核心指标 Tokens/sec、MFU TTFT、TPOT、并发数
计量 卡时 QPS、1K tokens 价格
失败处理 Checkpoint 回滚 请求重试 / 降级

1.2 为什么推理有自己的一套系统

一台训推混用的方案看上去很美,实际几乎没有人这样做。原因是两边的优化方向几乎相反:

所以主流路线是:训练用 Megatron / DeepSpeed / TorchTitan 等;推理用 vLLM / SGLang / TensorRT-LLM / TGI 等专用引擎。模型从训练态导出(通常是 HF safetensors),推理态再做量化、图优化、权重切分。

1.3 推理的 SLO 模型

一个在线推理服务的 SLO 典型长这样:

这套 SLO 直接决定了引擎怎么调参:batch 不能太大(大 batch 会拉长 TTFT 和 TPOT 尾延迟),也不能太小(小 batch 浪费算力)。所有”继续 batching vs 立刻 decode”的取舍都围绕这条线。

1.4 模型从训练态到推理态

模型权重的生命周期不止是”加载一下那么简单”:

  1. 训练结束 → HuggingFace safetensors / Megatron 分布式 checkpoint。
  2. 格式转换:把分布式 checkpoint 合并为单文件,重命名 key,切分到推理侧的 TP 拓扑。
  3. 量化:GPTQ / AWQ / FP8 校准,生成量化权重 + scale。
  4. 图编译(TRT-LLM 专属):用 trtllm-build 生成 engine,针对 batch size 桶和 seq 长度做 kernel 自动调优。
  5. 部署打包:权重 + tokenizer + 配置上传对象存储(OSS/S3),节点拉取。

这条链路一旦稳定下来,就是每天跑几十次的”模型发布流水线”,和 DevOps 世界的 CI/CD 平行存在。

二、Transformer 推理流程回顾

2.1 自回归生成

把大模型当黑盒,推理接口就是:

输入:  prompt = [t1, t2, ..., tn]
输出:  [tn+1, tn+2, ..., tn+m]   (直到 EOS 或 max_new_tokens)

每一步生成新 token 的伪代码是:

tokens = prompt
while len(tokens) - len(prompt) < max_new_tokens:
    logits = model(tokens)          # shape: [seq, vocab]
    next_id = sample(logits[-1])    # 只取最后一个位置
    tokens.append(next_id)
    if next_id == EOS: break

这份代码能跑,但极慢。问题在 model(tokens) 这一行:每生成一个 token 就把前缀重新 forward 了一遍,计算量随长度平方增长。

2.2 KV Cache 的由来

Transformer 的每一层 attention 计算是:

Q = x @ Wq        # [seq, d]
K = x @ Wk        # [seq, d]
V = x @ Wv        # [seq, d]
out = softmax(Q @ K^T / sqrt(d)) @ V

自回归推理时,生成第 t 个 token 只需要新的 Q_t、但需要与所有历史位置的 K、V 做 attention。历史位置的 K、V 在上一步已经算过了,下一步不会变(因为是 causal mask)。

这就是 KV Cache 的核心:把每一层每个位置的 K、V 缓存下来,下一步只算新位置的 K、V,拼到缓存上,再做 attention。

# 带 KV cache 的 forward
def attn_step(x_new, kv_cache_k, kv_cache_v):
    q = x_new @ Wq
    k_new = x_new @ Wk
    v_new = x_new @ Wv
    k = concat(kv_cache_k, k_new, dim=seq)
    v = concat(kv_cache_v, v_new, dim=seq)
    out = softmax(q @ k.T / sqrt(d)) @ v
    return out, k, v     # 新的 k/v 回写 cache

复杂度从每步 O(n²) 降到每步 O(n),总体从 O(n³) 降到 O(n²)。代价是显存:要存下所有层所有位置的 K 和 V。

三、Prefill 和 Decode 两阶段

3.1 阶段划分

KV Cache 的引入,天然把推理分成两个阶段:

一个典型 4K prompt 生成 512 tokens 的请求,它的时间线是这样的:

Prefill (4096 tokens 一次性):  ~200ms   → 给出 TTFT
Decode (512 步,每步 ~30ms):   ~15s     → 每步给出一个 token

3.2 Arithmetic intensity 的巨大差异

这两个阶段的计算访存比完全不同,决定了它们的瓶颈也不同。

先复习 roofline:一个 kernel 的上限性能由 arithmetic intensity I(计算 FLOPs / 访存 bytes)决定。I 高就是 compute-bound,吃算力(Tensor Core);I 低就是 memory-bound,吃带宽(HBM)。

A100 SXM4 的 peak:FP16 算力 312 TFLOPS,HBM 带宽 2.0 TB/s,拐点 I ≈ 156 FLOPs/byte。

对于一个 Linear 层 y = x @ W,x 形状 [B, d_in],W 形状 [d_in, d_out]:

看到没,算术强度几乎就是 batch size。

Prefill:输入是整个 prompt,有效 batch 是 bs · seq。一个 4K prompt + bs=4 就是 16K 的”宽度”,I 远远超过拐点,compute-bound,吃满 Tensor Core。

Decode:每一步输入只有 1 个 token,有效 batch 就是 bs。bs=4、bs=16 的 I 都远低于 156,memory-bound,瓶颈在读权重和读 KV。

这条结论非常重要,后面几乎所有推理优化都围绕它展开:

3.3 Prefill / Decode 时序图

Prefill / Decode 时序图

这张图说明了一切:单请求在 Decode 阶段把 GPU”空转”得很厉害。把多个请求的 Decode 合并成一个大 batch,就是 continuous batching 的动机。

四、KV Cache 深入

4.1 维度与大小估算

标准 Multi-Head Attention(MHA)的 KV Cache 形状:

kv_cache: [num_layers, 2, batch_size, num_heads, seq_len, head_dim]
                       ^
                       K / V

单请求单层单 token 的 KV 字节数(FP16):

per_layer_per_token = 2 (K+V) * num_heads * head_dim * 2 bytes

对几个代表性模型算一下每 token 的 KV 开销:

模型 层数 heads × head_dim(KV) 单 token 单层 总 KV/token
LLaMA-2 7B (MHA) 32 32 × 128 2 × 32 × 128 × 2 = 16 KB 512 KB
LLaMA-2 70B (GQA, 8 kv_heads) 80 8 × 128 4 KB 320 KB
LLaMA-3 70B (GQA, 8) 80 8 × 128 4 KB 320 KB
DeepSeek-V2 (MLA) 60 latent 512 ~1 KB 级别 ~70 KB

拿 LLaMA-2 70B 算 4K 上下文:

KV per seq = 320 KB/tok × 4096 tok ≈ 1.3 GB(单序列!)

考虑一些实现把 K/V 分开对齐、加 padding,实际逼近 2 GB/序列。这个数字解释了为什么长上下文 + 多并发推理的显存瓶颈不在参数,而在 KV Cache:H100 80GB 减去权重后剩 ~40GB,只能同时服务 ~20 个 4K 上下文,或者 ~5 个 16K 上下文。这就是 KV 是推理显存头号杀手 的由来。

4.2 GQA / MQA:砍 KV head 数

MQA(Multi-Query Attention,2019):所有 Q head 共享同一组 K、V。KV 大小除以 num_heads,极致压缩,但精度损失大,训练难。

GQA(Grouped-Query Attention,2023):折中方案。Q 分 num_heads 组,KV 分 num_kv_heads 组(通常 8),每组 KV 被多个 Q 共享。LLaMA-2 70B / LLaMA-3 / Mistral / Qwen2 / DeepSeek-V1 全部用 GQA。

标准 MHA:   Q[H], K[H], V[H]
MQA:        Q[H], K[1], V[1]
GQA(G=8):  Q[H], K[8], V[8]  —— H=64 时 K/V 内存是 MHA 的 1/8

GQA 带来 4-8× 的 KV 节省,对长上下文 / 高并发场景是决定性的。代价是微小的 PPL 上升,且要在训练阶段就选好组数(不能事后转换)。

4.3 MLA:DeepSeek 的 KV 压缩方案

DeepSeek-V2 / V3 引入了 Multi-head Latent Attention(MLA)。核心思想是把 K、V 投影到一个低维 latent 空间缓存,算 attention 前再”解压”回来。

简化表述:

# 传统 MHA:缓存 K, V ∈ R^{H·d_h}   一共 2 · H · d_h
# MLA:    缓存 c_kv ∈ R^{d_c}      一共 d_c(通常 d_c << 2·H·d_h)

c_kv = x @ W_kv_down      # 压到 latent
# 推理时:
K = c_kv @ W_k_up
V = c_kv @ W_v_up

实际实现里用了”吸收”技巧:W_k_up 与 W_q 可以预先合并,使 decode 时不需要真的把 K 解压到高维,而是把 query 投到 latent 空间,直接和 c_kv 做 attention。这样计算量也没增加多少。

DeepSeek-V3 的 MLA 把 KV Cache 从 GQA 的 ~4KB/层/token 进一步压到约 1-1.5KB/层/token,是长上下文成本下降的关键之一。DeepSeek V3 能在低成本上跑 128K 上下文,MLA 功不可没。

4.4 KV Cache 结构图

KV Cache 结构图

4.5 KV Cache 的量化与驱逐

除了结构上减小 KV,还有两类正交手段:

4.6 KV Cache 的显存预算推导

在一台 H100 80GB 上部署 LLaMA-3-70B(FP16),实操预算大致如下:

权重:        70B × 2 bytes   = 140 GB   → 4 卡 TP=4,单卡占 35 GB
框架保留:    激活 / 临时 buffer   ≈ 5 GB
可用 KV:     80 - 35 - 5       ≈ 40 GB / 卡   →  4 卡合计 ~160 GB KV
单 token KV: GQA(8 kv_heads)   = 320 KB / tok
可服务的"活跃 token 总量":
  160 GB / 320 KB ≈ 500K tokens

500K 活跃 token 听起来不少,换算到并发就缩水得很快:

再乘上 PagedAttention 的碎片率(5-10%)和系统保留,真实数字还要再打折。这个预算直接决定了集群规模和商业化价格模型。

4.7 KV 的数据生命周期

一个请求内 KV 的生命周期:

[admit] → [prefill 产生 KV] → [decode N 步,每步 append 1 行] → [done / timeout → 回收]

完成后 KV 的处理策略分三种:

多轮对话里,活跃会话的 KV 往往整会话保留(避免每轮从头 prefill),这是”会话粘性”必须做的工程决策。

五、Batching 的三代演化

单请求推理下,GPU 永远吃不饱。要拉吞吐必须 batch。但 LLM 的变长 + 分阶段让 batching 远比视觉模型复杂。

5.1 第一代:Static Batching

把同时到达的 N 个请求 pad 到最长,一起 forward。伪代码:

batch = collect_requests_for(timeout=50ms, max=8)
outputs = model.generate(pad(batch))

问题很明显:

Triton、TorchServe 在 LLM 前时代都是这种模式。

5.2 第二代:Dynamic Batching

Triton Inference Server 提供的模式:请求到达队列,调度器在 max_queue_delay 内窥视,尽量凑够 preferred_batch_size,再一起发到模型。

对 CV 模型和 BERT 有效,对 LLM 依然不够:LLM 每一步输入/输出都在变,不是”一次 forward 出完整结果”,batch 的粒度不应该是”请求”,而应该是”每一步”。

5.3 第三代:Continuous Batching / In-flight Batching

这是 vLLM(Orca 论文思想 + PagedAttention 实现)和 TensorRT-LLM 的核心贡献。

关键观察:Decode 阶段每一步都是独立的 [bs, 1, d] forward,不同请求的 decode step 完全可以拼到一个 batch。请求的”寿命”由它们产出多少 token 决定,但每一步 batch 的组成可以动态变化。

新的调度循环(简化版):

running = []                          # 正在 decode 的请求
waiting = []                          # 排队的 prefill 请求

while True:
    # 1. 有新请求且算力允许:做 prefill,加入 running
    if waiting and can_admit():
        req = waiting.pop(0)
        prefill(req)                  # 算出它的首 token + KV cache
        running.append(req)

    # 2. 把 running 里所有请求的"下一步 decode"拼成大 batch
    batch = [r.last_token for r in running]
    logits = model.decode_step(batch, kv_caches=[r.kv for r in running])

    # 3. 每个请求采样各自的 next token
    for r, lg in zip(running, logits):
        r.append(sample(lg))
        if r.done(): running.remove(r); yield r.output

好处:

坏处 / 挑战:

现代引擎(vLLM、SGLang、TRT-LLM、TGI、Mindie 等)默认都开 continuous batching。

5.4 Chunked Prefill

一个棘手问题:某个请求的 4K prompt prefill 要 200ms,这 200ms 里所有 decode 请求都在干等,TPOT 尾延迟会爆炸。

解法:把长 prefill 切成小块,每次只做比如 512 个 token 的 prefill,和一批 decode 拼成一个”混合 batch”一起 forward。vLLM 从 0.4 开始的 --enable-chunked-prefill、SGLang 的默认调度、DeepSeek-V3 的分离式 serve 方案都使用了这个思想。

5.5 Continuous Batching 时序对比

Continuous Batching 时序对比

六、推理关键指标

6.1 延迟指标

6.2 吞吐指标

6.3 SLO 与 P50/P95/P99

在线服务从不看平均值。P95/P99 才是真实用户体验。

P50 TTFT: 150ms      —— 一般用户
P95 TTFT: 400ms      —— 刚好满足
P99 TTFT: 1200ms     —— 1% 用户等了 1.2 秒

推理引擎的尾延迟常常来自:

生产上必须 per-model 维度监控 TTFT/TPOT/Queue 三条线的 P95/P99,超限直接告警或触发扩容。

6.4 输入 vs 输出 token 的瓶颈

不同业务的输入输出配比,决定了到底卡在 Prefill 还是 Decode。

场景 输入 token 输出 token 瓶颈
短对话 200 200 均衡偏 Decode
角色扮演 / 情感陪伴 500 1000 Decode
Code 补全 2K 50 Prefill
文档摘要 / 长文翻译 10K 500 Prefill 主
RAG 问答 5K 200 Prefill 主
Agent 多轮工具调用 累计 20K 300 × N 轮 Prefill 持续主导
推理类(o1 / R1)长思考 500 10K+ Decode

不同配比下最佳配置天差地别:RAG 场景要把更多算力分给 Prefill(更大 TP、更多 Prefill-only 节点);陪伴类聊天要压 Decode 显存(MLA、KV 量化、超长 continuous batching)。

6.5 压测与 benchmark 工具

生产上看的几套工具链:

一个标准 benchmark 报告至少包含:

模型 · 硬件 · 引擎版本 · 量化 · 并行配置 · 请求到达分布 · 输入/输出长度分布
-----------------------------------------------------------------------
QPS    TTFT P50/P95/P99    TPOT P50/P95/P99    Output tok/s    GPU util    KV util

没有到达分布(poisson 或真实 trace)和长度分布的 benchmark 基本可以忽略 —— 用固定长度 + 同步压力的数字严重偏离真实场景。

七、Decode kernel 层优化

Decode 是 memory-bound 的关键阶段,这个阶段每一毫秒都要精打细算。

7.1 FlashAttention 系列在 decode 中的变体

标准 FlashAttention(v1/v2/v3)优化的是 prefill 那种大 seq × seq 的 attention,通过 tiling 把 softmax 融到寄存器里,避免物化 QK^T。

Decode 场景有两个特点:

针对性优化是 FlashDecoding(2023)和 FlashDecoding++(2024):把 KV 维度切成多个 chunk,每个 SM 处理一段,分别算出局部 softmax,最后用 log-sum-exp 合并。让长 context 下的 decode 也能打满 SM。

7.2 Paged Attention kernel

PagedAttention(vLLM)不只是调度机制,还是一个 kernel:它在 CUDA 里实现了按”页表”访问非连续 KV 块的 attention。代价是访存比连续布局略低一些,收益是 KV 内存碎片几乎消失、跨请求共享 KV 变得平凡。下一篇有详解。

7.3 CUDA Graph:消灭 launch 开销

Decode 每步 forward 只要几十毫秒,但一个 70B 模型一次 forward 里有几百上千个 kernel,每个 kernel launch 要 5-20μs,累计就是毫秒级开销,占 TPOT 的相当比例。

CUDA Graph 把一段 kernel 序列录制成图,后续只用一次提交就能执行整张图,launch 开销降到几乎为 0。对固定 shape 的 decode step 尤其合适:batch 大小和 seq 长度虽然每步不同,但只要做成”按 batch 分桶”的几张 graph 就够用。

vLLM、TRT-LLM 都默认启用 CUDA Graph。注意事项:

7.4 权重加载与 kernel fusion

Decode 受限于读权重,那就少读一次。

TRT-LLM 的 plugin 系统、vLLM 的自定义 kernel、SGLang 对 FlashInfer 的依赖都在做这件事。

7.5 Speculative 执行与多流

高级优化里还有两条线:

7.6 一次 Decode step 的性能拆解(70B / H100 · 参考数值)

单 token 单请求(bs=1)total ≈ 30 ms,其中:
  读权重 70B × 2B / (3TB/s HBM)       ≈ 23 ms   <—— 主导
  读 KV   320KB × 4K tok × 1 / HBM    ≈  0.4 ms
  计算 FLOPs 2 × 70B × 1  / 989 TFLOPS ≈  0.14 ms
  kernel launch × 1000                ≈  5-10 ms  <—— CUDA Graph 能砍掉
  同步 / all-reduce(TP=4)             ≈  1-2 ms

看清楚了吗:绝大部分时间在读权重,其次是 kernel launch。所以:

八、Prefill / Decode 分离

8.1 为什么要分离

前面讲过 Prefill 和 Decode 的特性截然不同:

维度 Prefill Decode
计算特性 compute-bound memory-bound
偏好硬件 算力强(H100、B200) 带宽高(HBM3e、大 NVLink 域)
最佳 TP 中等(2-4) 小(1-2)或极大(靠 NVLink 域)
对延迟敏感 TTFT TPOT
对 batch 敏感 拼不太动 越大越好

把两者混在一个 pod 里,必然互相妥协。分离后,每个 pod 可以按自己的负载独立选 TP、选型号、选副本数。

8.2 典型架构

Client
   │
   ▼
Router ──► Prefill Pool(少量大 GPU · 高 TP · 侧重算力)
              │
              │ 通过 RDMA / NVLink 把 KV Cache 传输过去
              ▼
           Decode Pool(多节点 · 小 TP · 侧重带宽 · 长期持有 KV)
              │
              ▼
           Streaming 输出给客户端

关键难点在”KV Cache 传输”:prefill 完一个 4K prompt,KV 可能是几 GB 级别,要通过 RDMA(IB / RoCE)快速打给 Decode 节点,传输时间必须远小于 decode 一整轮的时间。这是一个纯粹的网络+内存工程问题。

8.3 代表性系统

分离方案目前仍在快速演进。对多数团队,先用合一部署 + continuous batching + chunked prefill 能满足大多数负载;当单副本 P95 打不住,或 MoE / 超长上下文场景下,再考虑上分离方案。

8.4 PD 分离的工程挑战

8.5 何时”不”分离

PD 分离不是银弹。以下场景保持合一部署更划算:

九、KV Cache 复用

9.1 Prefix Caching

真实业务里同一个 system prompt 被成千上万次复用:“你是一个专业助手…(2000 字设定)…请回答用户问题”。没有 prefix caching 的引擎,每次都要从头 prefill 这 2000 token,纯浪费。

Prefix caching 的思想:把最近处理过的 prefix 的 KV Cache 保留在显存(或下沉到 CPU / NVMe),新请求来时查前缀命中,命中的段直接复用 KV,只 prefill 新增部分。

命中一段长度 L 的前缀,直接省掉 L 个 token 的 prefill 时间 —— RAG / 长 system prompt / 多轮对话场景下这是数量级的加速。

9.2 Radix Tree:SGLang 的贡献

SGLang 把 prefix 组织成 Radix Tree(前缀树):每个节点是一段 token,边是一段 KV。新请求来时自底向上匹配,命中多深就复用多深。

相比”哈希存整个 prefix”的平面方案,Radix Tree 天然支持多分叉复用:比如一个 system prompt 下衍生出 100 个不同用户的对话,共享前缀只存一份,每个对话的私有部分挂在前缀节点下。这对 Agent 场景(同一套工具定义 prompt + 千差万别的后续)特别友好。

vLLM 在 0.5+ 也支持了 prefix caching(非 radix,基于 block hash),TRT-LLM、TGI 都有各自实现。这已经是现代推理引擎的标配。

9.3 分布式 KV 存储

单节点显存装不下所有活跃 prefix,于是出现多级:GPU HBM → CPU DRAM → NVMe → 分布式 KV 服务。Mooncake 的 KVCache Store、阿里 PAI-EAS 的 KV 分层、字节的 ByteCache 都是这条路线。命中远端 KV 也比重新 prefill 快(因为长 prompt 的 prefill 本身就是几百 ms)。

9.4 共享前缀的 attention 正确性

有个坑:prefix cache 命中后,不同请求共享同一段 KV,attention 计算在”共享段”上要正确处理 causal mask 和位置编码。RoPE 这类相对位置编码天然兼容(同样的 token 在同样的位置),ALiBi 也类似。但如果位置编码依赖 batch 内全局坐标,就会出错。主流引擎(RoPE 为主)已经默认兼容。

9.5 命中率观测

生产上 prefix cache 是”看起来有用,实际看命中率”。必须暴露指标:

prefix_cache_hit_tokens_total / prefix_cache_miss_tokens_total
prefix_cache_hit_rate(tokens 维度,不是 requests 维度)
prefix_cache_size_bytes / prefix_cache_evictions_total

RAG 类业务命中率往往不高(每次 context 都不同);Agent / 系统 prompt 驱动的业务命中率能到 70%+,TTFT 直接折半。

十、采样策略

生成 token 的时候,最后一层 logits 要转成一个具体的 id。策略不只影响质量,也影响工程 —— 某些采样(比如 beam search)会把 batch 规模膨胀。

10.1 基础采样

一组常用默认:temperature=0.7, top_p=0.9, presence_penalty=0.1。

10.2 结构化生成

不少业务要求输出是严格 JSON / SQL / XML。主流方案是 约束解码:在每一步采样前,按文法规则屏蔽掉非法 token 的 logits,只在合法集合里采。代表实现:Outlines、xgrammar、lm-format-enforcer,以及 SGLang、TRT-LLM、vLLM 的内置约束解码。

10.3 推测解码(下一篇详讲)

简单提一下:用一个小”草稿模型”一次吐 k 个 candidate token,大模型并行验证(一个 forward 验证所有 k 个,因为 causal 下可以同时算)。接受的那些直接采纳,不接受的从分歧点继续。在大模型上能带来 2-3× 的 decode 加速。第 15 篇专题讨论。

10.4 采样在 batch 下的工程细节

当上百个请求拼到同一个 decode step,每个请求的采样参数可能都不同(temperature、top_p、seed 都是请求级)。主流引擎会做”向量化采样”:

logits: [B, V]  →  per-request divide temperature
                →  per-request top-k/top-p 掩码(gather + scatter)
                →  per-request categorical sample(使用不同 seed 生成的 uniform noise)

这些操作 vLLM 放在一个叫 Sampler 的融合 kernel 里完成。对高并发 decode 来说,采样 kernel 本身也能占 1-3ms,不容忽视。

十一、多模态推理

视觉语言模型(VLM)、音频语言模型(ALM)、端到端多模态模型(Qwen2-VL / GPT-4o / Gemini / GLM-4V / MiniCPM-V / InternVL)的推理链路大致是:

图像 ──► ViT / Vision Encoder ──► 一串 image tokens ──┐
文本 ──► Tokenizer ───────────────► 一串 text tokens ──┴──► LLM ──► 文本输出

工程上的独特挑战:

11.1 视觉 encoder 与 LLM 的分离

大厂的工程实践里,Vision Encoder 通常和 LLM 分开部署:

分开的好处是 Vision 和 LLM 可以独立弹性:长文本场景 Vision 少用,短图多问答场景 LLM decode 少用。

11.2 视频推理的特殊性

视频输入一秒就是几十帧,token 数会爆炸。工程上:

主流引擎对多模态支持:vLLM、SGLang 均已原生支持主流 VLM;TRT-LLM 需要针对性插件;国内推理平台(火山方舟、阿里 PAI-EAS、腾讯 TI、百度千帆)都把 VLM 作为一等公民。

十二、动手:最简 Decode 循环

把所有概念落到一段可跑的 PyTorch 代码里。这里用 HuggingFace 的一个小模型演示手写 decode + KV cache。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = "Qwen/Qwen2.5-0.5B"   # 任意小 decoder-only 模型都行
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(
    MODEL, torch_dtype=torch.float16
).cuda().eval()

@torch.inference_mode()
def generate(prompt: str, max_new=64, temperature=0.7, top_p=0.9):
    # === Prefill ===
    input_ids = tok(prompt, return_tensors="pt").input_ids.cuda()
    out = model(input_ids, use_cache=True)
    logits = out.logits[:, -1, :]         # 只要最后一位
    past_kv = out.past_key_values         # [(K0,V0), (K1,V1), ...] per layer

    generated = [input_ids]

    # === Decode loop ===
    for _ in range(max_new):
        # 采样
        if temperature == 0:
            next_id = logits.argmax(-1, keepdim=True)
        else:
            probs = torch.softmax(logits / temperature, dim=-1)
            # top-p
            sorted_p, sorted_i = torch.sort(probs, descending=True)
            cum = sorted_p.cumsum(-1)
            mask = cum - sorted_p > top_p
            sorted_p[mask] = 0
            sorted_p /= sorted_p.sum(-1, keepdim=True)
            pick = torch.multinomial(sorted_p, 1)
            next_id = sorted_i.gather(-1, pick)

        generated.append(next_id)
        if next_id.item() == tok.eos_token_id:
            break

        # 只把新 token 送进模型,KV cache 继续增长
        out = model(next_id, past_key_values=past_kv, use_cache=True)
        logits = out.logits[:, -1, :]
        past_kv = out.past_key_values

    return tok.decode(torch.cat(generated, dim=1)[0], skip_special_tokens=True)

print(generate("解释一下什么是 KV Cache:", max_new=128))

这段代码做了两件关键的事:

  1. 第一次 model(input_ids) 是 Prefill,建立所有层的 KV Cache,返回 past_key_values。
  2. 之后每一步只喂 1 个新 token,配合上一次的 past_key_values 做 Decode,HF 内部会把新 K/V append 上去。

把 past_key_values=past_kv 这一行改成 None,强制每步重算整个前缀,你会立刻感受到几十倍的速度差 —— 这就是 KV Cache 的价值。

想看 KV Cache 的真实大小,加一行:

k0, v0 = past_kv[0]
print(k0.shape, k0.dtype)     # e.g. [1, num_kv_heads, seq, head_dim] fp16
total = sum(k.numel()*k.element_size() + v.numel()*v.element_size()
            for k, v in past_kv)
print(f"KV bytes: {total/1024/1024:.2f} MB")

对 Qwen2.5-0.5B + 4096 上下文,这个值大约几十 MB;换成 LLaMA-2 70B,同样上下文就是上面算过的 ~1.3 GB。

十三、引擎选型思路预告

13.1 一个推理引擎的内部结构

站在这里回头看,一个现代推理引擎的骨架大致是:

         ┌──────────────────────────────────────────┐
  HTTP   │  API Server(OpenAI 兼容 / gRPC)          │
  ───►  │  · 鉴权、限流、日志                          │
         │  · 把请求放入 RequestQueue                 │
         └────────────────┬─────────────────────────┘
                          ▼
         ┌──────────────────────────────────────────┐
         │  Scheduler(continuous batching 核心)     │
         │  · admit / preempt / schedule             │
         │  · prefix cache 查表                       │
         │  · 混合 batch(prefill chunks + decodes)  │
         └────────────────┬─────────────────────────┘
                          ▼
         ┌──────────────────────────────────────────┐
         │  Model Executor(可多 worker / TP / PP)   │
         │  · Paged KV Manager(block table)         │
         │  · CUDA Graph 池                          │
         │  · Flash / FlashDecoding / 量化 kernels    │
         │  · Sampler(fused)                        │
         └──────────────────────────────────────────┘

每一层后面都会有专文。但凡一个 production-grade 引擎,少不了这四层。各家差别在于:调度策略谁更聪明、KV 管理谁更省、kernel 谁更快、API 谁更兼容。

13.2 最终决策

下一篇会详聊 PagedAttention 与 Continuous Batching 的实现细节,第 13 篇做 vLLM / SGLang / TRT-LLM / TGI / LMDeploy 的横向对比。这里先给一张决策概览:

场景 推荐起点
通用 HF 模型 / 快速上线 vLLM(生态最广、默认 continuous batching + paged attention + prefix cache)
极致 TTFT / 复杂 prompt 编排 / Agent SGLang(RadixAttention、Python DSL)
NVIDIA 生产级 / 极致单卡吞吐 / 量化 TensorRT-LLM(闭源友好、FP8 / INT4 最强)
HuggingFace 官方 serve TGI
国产卡 / 华为昇腾 MindIE;寒武纪 MagicMind;海光 DTK;摩尔线程 / 沐曦自研
国产云托管 火山方舟、阿里 PAI-EAS、腾讯 TI-ONE、百度千帆

十四、小结与常见误区

14.1 核心要点回顾

  1. LLM 推理分 Prefill + Decode 两阶段,前者 compute-bound,后者 memory-bound。
  2. KV Cache 是推理第一优化对象;显存主要被它吃。GQA / MLA / KV 量化 / 驱逐是四条正交路线。
  3. Continuous Batching 是吞吐提升的关键;Chunked Prefill 解决 prefill 阻塞 decode。
  4. Prefill/Decode 分离 是大规模 serve 的终极形态,DeepSeek、Kimi、Azure 都在用。
  5. Prefix Caching / Radix Tree 让多用户共享 system prompt 成本骤降。
  6. FlashDecoding + CUDA Graph + kernel fusion 把 decode 每毫秒榨干。
  7. 核心指标 TTFT / TPOT / Throughput,必须按 P50/P95/P99 监控。

14.2 常见误区

14.3 给推理工程师的 checklist

上线一个模型之前,逐项确认:

十五、延伸阅读与下一步

推理是 LLM 基础设施里发展最快、开源最活跃的一块。本篇只是从工程第一性原理搭了个骨架,后续几篇会逐层填肉:

推荐在读完本篇后亲自操作一遍:

  1. 在单张 24GB 卡上用 vLLM 起一个 Qwen2.5-7B,用 benchmark_serving.py 跑 ShareGPT,观察 GPU KV cache usage 和 TTFT/TPOT 随并发变化的曲线。
  2. 关掉 --enable-prefix-caching,重新跑同一组负载,对比 TTFT 差异。
  3. 切到 --enforce-eager(关掉 CUDA Graph),对比 TPOT 尾延迟。
  4. 切换 --max-num-seqs 从 16 到 256,观察吞吐与 P99 的 trade-off。

手上有 4 张 H100 的读者,可以复现 DistServe 的 PD 分离实验,直观感受分离带来的 P95 TTFT 改善。没有硬件的读者,可以基于本篇第十二节的代码,把 KV cache enabled/disabled 切换,观察单序列速度差 —— 这是所有推理优化的”起点”。

还想深挖的方向有:

推理工程的魅力在于:模型一旦定型,剩下的每 1% 提升都来自工程。把上面的拼图一块一块拼齐,你就能把一个”能跑”的模型变成”能赚钱”的服务。

上一篇:【大模型基础设施工程】10:Checkpoint 与故障容忍 下一篇:【大模型基础设施工程】12:PagedAttention 与 Continuous Batching

参考资料

  1. Kwon et al. Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM), SOSP 2023.
  2. Yu et al. Orca: A Distributed Serving System for Transformer-Based Generative Models, OSDI 2022.
  3. Dao. FlashAttention-2 / FlashAttention-3.
  4. Hong et al. FlashDecoding++, 2024.
  5. Zhong et al. DistServe: Disaggregating Prefill and Decoding for Goodput-optimized LLM Serving, OSDI 2024.
  6. Patel et al. Splitwise: Efficient Generative LLM Inference Using Phase Splitting, Microsoft, ISCA 2024.
  7. Moonshot AI. Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving, 2024.
  8. DeepSeek-AI. DeepSeek-V2 / V3 Technical Report(MLA 章节).
  9. Ainslie et al. GQA: Training Generalized Multi-Query Transformer Models, 2023.
  10. Zheng et al. SGLang: Efficient Execution of Structured Language Model Programs, 2024.
  11. Xiao et al. StreamingLLM / H2O / SnapKV 系列 KV 驱逐论文.
  12. NVIDIA TensorRT-LLM、HuggingFace TGI、vLLM、SGLang、LMDeploy、MindIE 官方文档.

同主题继续阅读

把当前热点继续串成多页阅读，而不是停在单篇消费。

2026-04-22 · architecture / ai-infra

大模型基础设施工程

面向中国工程团队的大模型基础设施系列。从 GPU/CUDA/互联、训练框架与 3D 并行、vLLM/SGLang 推理引擎、量化与推测解码、RAG/Agent 到服务化、网关、可观测性与安全合规，覆盖 LLMOps 全链路。

By .