LLM 推理系统：SGLang 深度解析

2026-06-03

一个 LLM 推理请求从用户输入到最终输出的完整链路：请求生命周期、Continuous Batching、KV Cache 管理、Scheduler 设计、Quantization、Speculative Decoding 等。本文正文为 markdown，关键机制配有可交互动画（点按钮演示）。

4.1 请求生命周期 (Request Lifecycle)

一个 LLM 推理请求从用户输入到最终输出，需要经历多个阶段。在 SGLang 中，这个过程被精心设计为一个高效的 pipeline：TokenizerManager → Scheduler → TpModelWorker → DetokenizerManager。

下面动画展示请求流经 Tokenize → Schedule → Prefill → Decode → Detokenize 五个阶段：

各阶段详解

阶段	操作	瓶颈	SGLang 优化
`Tokenize`	将用户输入文本转换为 token ID 序列	CPU 密集	异步 tokenization，与 GPU 计算重叠
`Schedule`	将请求分配到 batch slot，管理优先级	调度延迟	Continuous Batching + 优先级队列
`Prefill`	并行处理所有 prompt tokens，填充 KV Cache	Compute bound	Chunked Prefill，Radix Cache 复用
`Decode`	自回归逐 token 生成	Memory bandwidth	PageAttention，高效 KV Cache 访问
`Detokenize`	将生成的 token IDs 转换回文本	CPU 密集	增量 detokenization + streaming

Prefill vs Decode 的本质区别

Prefill 阶段是 compute-bound（计算量大，GPU 计算单元是瓶颈）：所有 prompt tokens 可以并行计算 attention，矩阵运算密集。

Decode 阶段是 memory-bandwidth-bound（需要频繁读取 KV Cache，显存带宽是瓶颈）：每步只生成一个 token，但需要读取完整的 KV Cache。

// Prefill: 并行处理 N 个 tokens
// FLOPs ≈ 2 * N * d_model^2 * num_layers (compute bound)
attention_output = softmax(Q @ K.T / sqrt(d)) @ V  // Q: [N, d], K: [N, d]

// Decode: 每次只处理 1 个 token
// Memory reads ≈ 2 * seq_len * d_model * num_layers (memory bound)
attention_output = softmax(q @ K_cache.T / sqrt(d)) @ V_cache  // q: [1, d], K: [seq, d]

性能指标: TTFT 与 TPOT

指标	全称	含义	决定因素
TTFT	Time To First Token	用户发出请求到第一个 token 出现的延迟	Prefill 速度 + 排队时间
TPOT	Time Per Output Token	后续每个 token 的生成间隔 (用户感知的 “打字速度”)	Decode 速度

用户体验: 1000 / TPOT = tokens/sec。TPOT=50ms → 20 tok/s (流畅), TPOT=200ms → 5 tok/s (卡顿)。

Prefill 速度推导: 从 GPU FLOPS 到 tokens/sec

Prefill 是 compute-bound, 速度由 GPU 算力 / 每 token 计算量 决定:

# 核心公式:
prefill_speed = GPU_FLOPS × MFU / FLOPs_per_token

# Step 1: 每 token 的 FLOPs ≈ 2 × model_params
# 为什么 2×? 矩阵乘 y=x@W (W有P参数) = P次乘法 + P次加法 = 2P FLOPs
# 模型 forward = 过一遍所有权重矩阵, 所以 FLOPs/token ≈ 2 × total_params

# Step 2: GPU 实际算力 = 峰值 × MFU (Model FLOPs Utilization)
# Prefill 是大矩阵乘, MFU 较高: 50-70%
# H100 BF16 峰值 989 TFLOPS, 实际 ~500-700 TFLOPS

# Step 3: 相除
# 7B on H100 BF16, MFU=60%:
#   = 989 × 0.6 TFLOPS / (2 × 7 GFLOP/token)
#   = 593 TFLOPS / 14 GFLOP = ~42K tokens/sec

速记公式 (H100 BF16, MFU≈60%)：

# prefill_speed ≈ 300,000 / B  (tokens/sec, 单卡)
#
# 7B   → ~42K tok/s
# 13B  → ~23K tok/s
# 70B  → ~4.3K tok/s  (单卡; TP=8 → ~30K tok/s)
# 405B → ~740 tok/s   (单卡; TP=8 → ~5K tok/s)

Decode 速度推导: Memory-Bandwidth Bound

Decode 时每个权重只被 1 个 token 用一次, 瓶颈在 读数据的速度。每生成一个 token 需要从 HBM 读取:

# 完整公式: 每 token 需要读的总数据量
bytes_per_token = model_weights + kv_cache_read
#                 ↑ 固定 (每次都要读完整模型)   ↑ 随 seq_len 增长!

# TPOT = bytes_per_token / memory_bandwidth
# Decode throughput = memory_bandwidth / bytes_per_token

简化情况: 短序列 + GQA (KV cache 可忽略)：

# LLaMA-70B GQA, seq_len=2048 on 2×A100-80GB (TP=2):
model_weights = 140  # GB (BF16), TP=2 → 每卡读 70 GB
kv_cache = 80 * 8 * 128 * 2048 * 2 * 2 / 1e9  # = 0.67 GB (GQA 只有 8 KV heads)
# KV cache 只占 0.5% → 可忽略

# TP=2: 两卡并行, 每卡独立读自己的 70 GB 权重
# A100 带宽 2.0 TB/s:
#   理论 TPOT = 70 GB / 2.0 TB/s = 35 ms
#   + TP AllReduce (80层×2次) + kernel launch overhead ≈ +15ms
# 实测: ~53.8 ms (Cursor blog, bs=1)
# → MBU = 35/53.8 ≈ 65% (合理: AllReduce + launch 吃掉了 35% 的时间)

KV Cache 何时不可忽略？

# 关键: KV cache 大小 = 2 × layers × kv_heads × head_dim × seq_len × 2
# 它随 seq_len 线性增长!

# 小模型 + MHA + 长序列: KV cache 可能比模型还大
# LLaMA-7B (MHA, 32 KV heads), seq_len=32768:
model_weights = 14    # GB
kv_cache = 32 * 32 * 128 * 32768 * 2 * 2 / 1e9  # = 16.8 GB (!)
# KV cache 比模型还大! TPOT 翻倍:
# (14 + 16.8) / 3350 = 9.2 ms (vs 无 KV 的 4.2ms)

# 大模型 + GQA + 短序列: KV cache 微不足道
# LLaMA-70B (GQA, 8 KV heads), seq_len=4096:
# KV = 80×8×128×4096×2×2 / 1e9 = 1.3 GB vs 140 GB model → ~1%

场景	Model Size	KV Cache	KV 占比	TPOT 影响
70B + GQA + 4K seq	140 GB	1.3 GB	1%	可忽略
7B + MHA + 4K seq	14 GB	2.1 GB	13%	轻微
7B + MHA + 32K seq	14 GB	16.8 GB	55%	TPOT 翻倍!
7B + MHA + 128K seq	14 GB	67 GB	83%	完全主导
70B + GQA + 128K seq	140 GB	41 GB	23%	不可忽略

结论: GQA/MLA 不仅省显存, 更重要的是省带宽 — 减少每次 decode 需要读取的 KV cache 量, 直接降低 TPOT。

实测数据验证 (理论 vs 实际)：

模型	GPU	理论 TPOT	实测 TPOT	MBU	来源
LLaMA-7B (BF16)	1×A100 80GB	7.2 ms	9.3 ms (bs=4)	~77%	HuggingFace TGI
LLaMA-70B (BF16)	2×A100 80GB	36 ms	53.8 ms (bs=1)	~67%	Cursor blog
LLaMA-70B (BF16)	H100 (推算)	21 ms	~35 ms	~60%	Databricks (H100 比 A100 快 36%)

MBU (Model Bandwidth Utilization) = 理论 TPOT / 实测 TPOT。实际 60-80% 是正常的, 因为还有 kernel launch overhead、AllReduce (TP)、attention 计算等非权重读取的开销。

为什么 Decode 要 Batching？

# 权重从 HBM 读一次, 被 batch 中的 B 个 token 共享:
# Arithmetic Intensity = 2B (B个token每个做2次运算/读一次权重)
# 当 B 增大, 从 memory-bound 逐渐变成 compute-bound
#
# 临界 batch size (compute 与 memory 平衡):
#   B_critical = FLOPS / (Bandwidth × 2) = 989e12 / (3.35e12 × 2) ≈ 148
#
# B < 148: memory-bound, 增大 batch → 吞吐线性增长, TPOT 不变
# B > 148: compute-bound, 增大 batch → TPOT 开始增长
#
# 实测验证 (Cursor blog, LLaMA-70B on 2×A100):
#   bs=1:  18.6 tok/s  ← memory-bound
#   bs=8:  17.2 tok/s  ← 仍然 memory-bound (每请求 TPOT 不变)
#   bs=32: 12.4 tok/s  ← 开始 saturate
#
# LLM-Viewer 论文 (arXiv 2402.16363) 证实:
# "在 decode 阶段, 所有计算层的 arithmetic intensity ≈ 1 OPs/byte,
#  firmly in memory-bound region"

KV Cache 大小估算

从模型 config 推算 KV cache per token:

# 公式:
kv_per_token = 2 × n_layers × n_kv_heads × head_dim × dtype_bytes
#              ↑       ↑          ↑           ↑          ↑
#             K+V     层数     KV heads数   每head维度   BF16=2

# 速算: kv_per_token = 512 × n_layers × n_kv_heads  (bytes, BF16)

模型	参数量	Layers	KV Heads	Attention	KV/token (BF16)
LLaMA-7B	7B	32	32 (MHA)	MHA	512 KB
LLaMA-13B	13B	40	40 (MHA)	MHA	800 KB
LLaMA-70B	70B	80	8 (GQA)	GQA	320 KB
LLaMA-405B	405B	126	8 (GQA)	GQA	504 KB
DeepSeek-V3	671B	61	MLA	MLA (512d)	122 KB

关键观察: 不能从参数量线性推断！

# 7B  → 512 KB/token  (73 KB per B of params)
# 70B → 320 KB/token  (4.6 KB per B of params) ← 差 16×!
#
# 原因: 大模型用 GQA (8 KV heads 共享给 64 Q heads)
# KV cache 不随 Q heads 增长, 只随 KV heads 增长
# MLA 更激进: 把 KV 压缩到低维 latent (512 vs 7168)

Serving 容量估算：

# 70B on 2×H100 (TP=2), 模型占 70GB/卡, 剩余 10GB 给 KV cache
available_kv = 10 * 1024  # 10240 MB
kv_per_token = 0.32       # MB (320 KB, LLaMA-70B GQA)

# 如果平均 seq_len = 4096:
max_concurrent = available_kv / (kv_per_token * 4096)
# = 10240 / 1310 ≈ 7-8 个并发请求 (单卡很紧张!)
# 这就是为什么 PagedAttention 和内存管理如此重要

GPU 显存到底分给了谁

KV cache 不是孤零零一项 — 推理时一张卡的显存被切给好几个”住户”。vLLM 那个 gpu_memory_utilization=0.9 留出来的 10%, 到底在给谁? 把每一项列清楚, 才能算明白能塞多少并发, 以及什么情况下会 OOM。

项	作用	大小 (70B, TP=8, FP16, 单卡)	随并发 / seq 增长?
Weights	模型参数 — Linear/Embedding/LayerNorm	140 GB / 8 = 17.5 GB	✗ 固定
KV cache	每个历史 token 的 K, V tensor, 解码时必读	0 — 几十 GB	✓ 唯一线性涨的项
Activations (单层峰值)	当前 forward 层的中间结果 (QKV / attn out / MLP)	~100 MB (decode) — ~1 GB (长 prefill)	✓ 仅 prefill 时
cuBLAS / cuDNN workspace	GEMM kernel 的 split-K / reduce 临时缓冲	每 stream 几百 MB	✗ 启动时申请
NCCL / RCCL 通信 buffer	TP all-reduce / all-gather 的收发 staging	100 MB - 1 GB	✗ 启动时申请
CUDA graph 内存	vLLM 默认按 batch size 各 capture 一份 graph	数百 MB - 1+ GB	✗ 启动时申请
PyTorch caching allocator 碎片	历史分配未归还、未合并的 block 余量	0.5 - 2 GB	✗ 随时间累积
CUDA driver / runtime baseline	进程一启动 `nvidia-smi` 就能看见的底噪	~500 MB	✗ 固定

70B 单卡的非-KV-cache 总开销 ≈ 17.5 (weights) + 2-3 (workspaces/graph/碎片) + 0.5 (driver) ≈ 20 GB. 80 GB 卡剩下 ~60 GB 给 KV cache + prefill spike — 这就是为什么 gpu_memory_utilization=0.9 要留 8 GB 余量。

各项的关键 caveat

Weights — TP 切完每卡只有 1/TP, 但 embedding 通常不切 (词表小, 切完反而要 all-gather 更贵). 70B vocab 32k × hidden 8192 × 2 bytes ≈ 0.5 GB, 不切问题不大。

Activations — 训练里是大头 (上一篇里 70B 训练光 activation 就 145 GB), 推理里只占 100 MB 级别。原因: 推理不存反向, 一层算完就丢; FlashAttention 让 O(s²) 的 attention score 不物化。唯一会爆的两种场景:

超长 prefill: s=128k 的 prompt 单层就 10·128k·8192/8 ≈ 1.3 GB
大 prefill batch: 同时 prefill 多条请求, batch 翻倍 activation 也翻倍

vLLM 的 --max-num-batched-tokens 就是在限这个 prefill 峰值的上限。

cuBLAS workspace — 跑 torch.cuda.memory_summary() 经常能看到几百 MB 的 Allocated by CUDA libs. 主要是 cuBLAS 的 GEMM split-K reduction buffer — 大矩阵乘把 K 维切成多个部分各算各的, 最后做归约, 中间结果需要 staging. 每个 CUDA stream 独立申请, 多 stream 时翻倍。

NCCL buffer — TP=8 时每卡为 all-reduce 准备的 staging ≈ hidden × bs × seq × 2 bytes × 2 (recv+send). NCCL 默认按需扩, 大 batch 时膨胀更明显。可以用 NCCL_BUFFSIZE=8388608 等环境变量限制上限。

CUDA graph — vLLM 为常见 batch size (1, 2, 4, 8, 16, 32, ...) 各 capture 一份 graph 避开 kernel launch overhead. 每份 graph 自己 ~几十 MB workspace, 累计可观。--enforce-eager 关掉能省 1-2 GB 但 decode 性能跌 10-20%。

Allocator 碎片 — PyTorch caching allocator 切大 block 给小请求后, free 时块不一定能合并。Long-running serving 进程通常会比刚启动多 1-2 GB “看不见”的占用。PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 可以显著缓解。

实际显存账: 70B 在 8 × 80GB H100 上

每卡 80 GB 怎么分:
  Weights (TP=8, FP16):                17.5 GB
  CUDA driver / runtime:                0.5 GB
  cuBLAS + cuDNN workspaces:            0.6 GB
  NCCL TP buffers:                      0.8 GB
  CUDA graphs (vLLM default):           0.6 GB
  PyTorch allocator 碎片余量:           1.0 GB
  ─────────────────────────────────────────────
  非 KV cache 总开销:                   ~21 GB
  剩余 KV cache + prefill spike:        ~59 GB

gpu_memory_utilization=0.9 配置:
  vLLM 总预算 = 80 × 0.9 = 72 GB
  扣掉 weights 17.5 → 划给 KV cache pool: 54.5 GB
  (driver/workspaces 等从 72 之外吃, 留 8 GB 余量就是给它们)

70B GQA, kv/token = 320 KB:
  KV cache pool: 54.5 GB / 320 KB ≈ 170k tokens 总容量
  平均 seq=4k 时 ≈ 40 并发请求 / 卡

什么时候要调小 `gpu_memory_utilization`?

症状	真正原因	不是因为
启动就 OOM	weights + workspaces 已超 90%	activation 涨
Prefill burst 时 OOM	长 prompt / 大 prefill batch 让 activation 临时膨胀到 GB 级	KV cache 涨 (KV cache 是渐增, 不会瞬间爆)
Long-running 之后 OOM	allocator 碎片累积	weights / KV cache

调小 gpu_memory_utilization (例如 0.9 → 0.85) 主要在帮 workspace + prefill 峰值 留缓冲, 跟常态 activation 关系不大。

一句话直觉

训练显存 = (weights + grads + optim) (能 ZeRO 切) + activation (随 batch×s² 长, 必须 checkpointing + TP + SP 压)

推理显存 = weights (静态) + KV cache (随并发线性涨, 唯一可调的池子) + workspace (启动时申请) + activation spike (prefill 才出现)

gpu_memory_utilization 真正在留位置给的是 workspace + spike, 不是常态 activation。这个数字越小 → KV cache 池越小 → 能塞的并发越少; 越大 → 越容易在 prefill burst 时 OOM。

activation 到底”长”在 Transformer 哪里

上面表格里 activation 一栏含糊一句 “中间结果”, 但具体是哪些 tensor、什么形状, 决定了 prefill spike 会有多大、FlashAttention 救了什么。走一遍一个 pre-LN transformer block 的前向, 把每个 activation 标出来:

x_in                                  shape: (b, s, h)         ← 上一层来的, 也是残差用
│
├─ x1 = LN1(x_in)                     (b, s, h)                ← LN 输出
├─ Q = x1 @ W_Q                       (b, s, h)
├─ K = x1 @ W_K                       (b, s, h)
├─ V = x1 @ W_V                       (b, s, h)                ← Q, K, V 三份
│
├─ split heads → (b, n_h, s, d_k)
├─ scores = Q @ K^T / √d_k            (b, n_h, s, s)           ← ★ O(s²), 最大单项
├─ attn_w = softmax(scores)           (b, n_h, s, s)           ← ★ 又一份 O(s²)
├─ ctx = attn_w @ V                   (b, n_h, s, d_k)
├─ merge → (b, s, h)
├─ attn_out = ctx @ W_O               (b, s, h)
├─ x_attn = x_in + attn_out           (b, s, h)                ← 残差和
│
├─ x2 = LN2(x_attn)                   (b, s, h)
├─ up = x2 @ W_up                     (b, s, 4h)               ← ★ MLP 中间, 4× hidden
├─ act = GeLU(up)                     (b, s, 4h)               ← ★ 又一份 4×
├─ down = act @ W_down                (b, s, h)
└─ x_out = x_attn + down              (b, s, h)                ← 下一层输入

每个 ★ 是体积大头, 其它都是 (b, s, h) 级别。按数量级把 activation 分成三档:

量级	形状	哪些 tensor	单个大小 (70B, b=1, s=2048)
O(b·s·h)	`(b, s, h)`	LN io, Q/K/V, attn_out, 残差和, LN2 io	每个 ≈ 34 MB
O(b·s·4h)	`(b, s, 4h)`	MLP up output, GeLU output	每个 ≈ 134 MB (大 4×)
O(b·n_h·s²)	`(b, n_h, s, s)`	Attention scores, softmax output	每个 ≈ 1.0 GB (s² 翻车)

s=8192 时 attention score 一份就 17 GB / 层, 完全压死一切 — 所以 long context 训练几乎必上 FlashAttention。

FlashAttention 抹掉哪一档

FlashAttention 的核心: scores 和 softmax 输出永远不物化到 HBM, 只在 SRAM 里分块算完直接喂给 @ V。所以加 FA 之后:

★ O(s²) 那两块: 从 activation 列表里消失 (反向也用同样的分块重算, 见 Flash Attention 的反向)
其它项不变

Megatron 的 Act per layer = 34·s·b·h + 5·a·s²·b 公式里, 5·a·s²·b 这项被 FA 直接砍掉。这就是为什么 Megatron 的 selective activation recomputation 默认只重算 attention score — 单点最大, 重算代价最低 (~10% 额外 FLOPs 换 ~80% activation 节省)。

边界处的两个”非 layer” activation

Tensor	形状	大小 (70B, b=1, s=2048)	备注
Input embedding 输出	`(b, s, h)`	34 MB	一份, 训练时也要存 (embedding 反向需要)
Output logits	`(b, s, vocab)`	LLaMA-2 vocab=32k: 130 MB; LLaMA-3 vocab=128k: 524 MB	比单层 activation 还大, 单点容易 OOM
Loss softmax 中间	`(b, s, vocab)`	同上	常用 fused CE kernel (Liger/Apex) 干掉

Output logits 那一份是单点大户 — 这就是为什么训练 long-context + 大词表模型时, fused cross-entropy 几乎必上 (思路跟 FlashAttention 一模一样: 不物化中间)。

训练 vs 推理: 哪些 activation 活着, 哪些被回收

	训练 (反向需要)	推理 (无反向)
LN io / Q / K / V / MLP 中间	全部存到反向结束	一层算完就丢, 内存可被下一层复用
Attention scores (O(s²))	不用 FA: 存; 用 FA: 不存	不用 FA: 临时; 用 FA: 不存
Output logits	存到 loss backward	用完就丢 (生成一个 token 即可)
同时活着的 activation 总量	所有层 × 上面所有项 (除非 checkpointing)	大约只有当前层的中间

这就是为什么训练里 activation 是 ~150 GB / 卡量级 (70B), 推理里是 ~100 MB / 卡量级 — 不是因为推理”省”了什么, 而是 inference 根本不需要在反向开始前把它们都留着。

一句话直觉

“短” activation (O(b·s·h)): 像 hidden state 在流水线上一节节传, 每节一份, 不大但量多。

“宽” activation (O(b·s·4h)): MLP 中间膨胀到 4×, 是每层最大的非-O(s²) 项, 总 activation 里贡献最多。

“方” activation (O(b·n_h·s²)): attention 的注意力矩阵, s 一长就爆 — FlashAttention 是专门为干掉它生的。

4.2 Continuous Batching vs Static Batching

传统的 Static Batching 会等待一个 batch 内所有请求都生成完毕后才开始处理下一个 batch，造成严重的 GPU 资源浪费。Continuous Batching 通过动态地添加和移除请求，让 GPU 始终保持高利用率。

下面动画对比两种 batching 方式的 GPU 利用率（红色为浪费的 slot）：

Continuous Batching 核心实现

class ContinuousBatchScheduler:
    def __init__(self, max_batch_size, max_seq_len):
        self.running_batch = []        # 正在处理的请求
        self.waiting_queue = []        # 等待队列
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len
        self.kv_pool = KVCachePool()  # KV Cache 内存池

    def schedule_step(self):
        # 1. 移除已完成的请求 (遇到 EOS 或达到 max_len)
        finished = [r for r in self.running_batch if r.is_finished()]
        for req in finished:
            self.running_batch.remove(req)
            self.kv_pool.free(req.kv_blocks)  # 释放 KV Cache
            yield_response(req)

        # 2. 从等待队列中取出新请求填充空位
        available_slots = self.max_batch_size - len(self.running_batch)
        while available_slots > 0 and self.waiting_queue:
            new_req = self.waiting_queue.pop(0)
            if self.kv_pool.can_allocate(new_req.prompt_len):
                new_req.kv_blocks = self.kv_pool.allocate(new_req.prompt_len)
                self.running_batch.append(new_req)
                available_slots -= 1
            else:
                self.waiting_queue.insert(0, new_req)
                break  # 内存不足，停止调度

        # 3. 执行一个 decode step
        if self.running_batch:
            model_forward(self.running_batch)

关键设计要点：

Iteration-level scheduling：每个 decode iteration 后都可能有请求加入或退出
无需 padding：不同长度的请求通过 PageAttention 高效处理
内存感知调度：只有当 KV Cache pool 有足够空间时才接纳新请求
优先级支持：短请求优先完成，降低平均 latency

性能对比

指标	Static Batching	Continuous Batching	提升
GPU 利用率	30-50%	85-95%	~2x
Throughput (tokens/s)	基准	2-3x	显著
平均 Latency	高 (等待最长请求)	低 (及时释放)	50-70%↓
尾部 Latency (P99)	极高	可控	大幅改善

4.3 KV Cache 管理

PageAttention: KV Cache 的虚拟内存

灵感来自操作系统的分页内存管理。PageAttention 将 KV Cache 分割成固定大小的 pages (blocks)，通过 page table 实现非连续的物理内存映射，解决了传统方法中的内存碎片问题。

class KVCachePool:
    """类似 OS 虚拟内存的 KV Cache 管理器"""
    def __init__(self, num_blocks, block_size, num_layers, num_heads, head_dim):
        # block_size: 每个 block 存储的 token 数 (通常 16 或 32)
        self.block_size = block_size
        self.num_blocks = num_blocks

        # 物理内存池: [num_blocks, 2, num_layers, block_size, num_heads, head_dim]
        # 2 代表 K 和 V
        self.kv_pool = torch.zeros(
            num_blocks, 2, num_layers, block_size, num_heads, head_dim,
            dtype=torch.float16, device="cuda"
        )
        self.free_blocks = list(range(num_blocks))  # 空闲 block 列表
        self.page_tables = {}  # request_id → [block_indices]

    def allocate(self, request_id, num_tokens):
        # 计算需要的 block 数量 (向上取整)
        num_needed = ceil_div(num_tokens, self.block_size)
        if len(self.free_blocks) < num_needed:
            raise MemoryError("KV Cache pool exhausted")

        blocks = [self.free_blocks.pop() for _ in range(num_needed)]
        self.page_tables[request_id] = blocks
        return blocks

    def free(self, request_id):
        blocks = self.page_tables.pop(request_id)
        self.free_blocks.extend(blocks)

    def append_token(self, request_id, layer, k, v):
        # 追加新 token 的 KV，可能需要分配新 block
        blocks = self.page_tables[request_id]
        current_len = self.token_counts[request_id]
        block_idx = current_len // self.block_size
        offset = current_len % self.block_size

        if offset == 0 and block_idx == len(blocks):
            # 需要新 block
            new_block = self.free_blocks.pop()
            blocks.append(new_block)

        physical_block = blocks[block_idx]
        self.kv_pool[physical_block, 0, layer, offset] = k
        self.kv_pool[physical_block, 1, layer, offset] = v

Radix Cache: 前缀共享树

SGLang 的核心创新之一。当多个请求共享相同的 system prompt 或前缀时，Radix Cache 通过树结构共享 KV Cache，避免重复计算。

下面是一个可交互的 Radix Tree：点击按钮添加请求，观察相同前缀如何被共享、新请求如何匹配最长前缀：

Radix Cache 实现原理

class RadixCacheNode:
    def __init__(self):
        self.children = {}       # token_id → RadixCacheNode
        self.kv_indices = []     # 对应的 KV cache block indices
        self.ref_count = 0       # 引用计数 (用于 eviction)
        self.last_access = 0    # LRU 时间戳

class RadixCache:
    def __init__(self, kv_pool):
        self.root = RadixCacheNode()
        self.kv_pool = kv_pool

    def match_prefix(self, token_ids):
        """找到最长匹配前缀，返回 (匹配长度, KV blocks)"""
        node = self.root
        matched_len = 0
        kv_blocks = []

        for token_id in token_ids:
            if token_id in node.children:
                node = node.children[token_id]
                node.ref_count += 1
                node.last_access = current_time()
                kv_blocks.extend(node.kv_indices)
                matched_len += 1
            else:
                break

        return matched_len, kv_blocks

    def insert(self, token_ids, kv_blocks):
        """插入新的 KV cache 到树中"""
        node = self.root
        for i, token_id in enumerate(token_ids):
            if token_id not in node.children:
                new_node = RadixCacheNode()
                new_node.kv_indices = [kv_blocks[i]]
                node.children[token_id] = new_node
            node = node.children[token_id]
            node.ref_count += 1

    def evict_lru(self, num_blocks_needed):
        """当内存不足时，按 LRU 策略驱逐"""
        # 收集所有叶子节点，按 last_access 排序
        leaves = self._collect_leaves(self.root)
        leaves.sort(key=lambda n: n.last_access)

        freed = 0
        for leaf in leaves:
            if leaf.ref_count == 0 and freed < num_blocks_needed:
                self.kv_pool.free_blocks(leaf.kv_indices)
                freed += len(leaf.kv_indices)
                self._remove_node(leaf)
        return freed

Radix Cache 的优势：

前缀共享：相同 system prompt 的请求共享 KV Cache，节省显存
自动匹配：新请求自动找到最长匹配前缀，跳过已缓存的 prefill 计算
LRU 驱逐：内存不足时按 LRU 策略回收不常用的缓存
Multi-turn 加速：多轮对话中前几轮的 KV 被缓存复用

4.4 Scheduler 设计

Prefill 与 Decode 的调度矛盾

Prefill 和 Decode 有根本性的计算特性差异，如何在同一 GPU 上高效调度是核心挑战：

特性	Prefill	Decode
计算特征	Compute-bound (大矩阵乘)	Memory-bound (读 KV Cache)
延迟敏感度	影响 TTFT	影响 TPOT
Token 数量	数百~数千	1 per step
GPU 利用模式	高计算密度	高带宽需求

SGLang 的调度策略（Overlap Scheduling: CPU 调度与 GPU 计算重叠）：

class OverlapScheduler:
    """Prefill-Decode Overlap Scheduling"""

    def schedule_batch(self):
        # 策略: Chunked Prefill + Decode 混合
        batch = MixedBatch()

        # 1. 优先加入正在 decode 的请求 (保证 TPOT)
        for req in self.running_decode:
            batch.add_decode_request(req)

        # 2. 用剩余算力做 prefill (chunked)
        remaining_budget = self.compute_budget - batch.decode_flops
        for req in self.waiting_prefill:
            chunk_size = min(req.remaining_prefill, remaining_budget)
            if chunk_size > 0:
                batch.add_prefill_chunk(req, chunk_size)
                remaining_budget -= chunk_size

        # 3. 拥塞控制: 如果等待队列过长，限制新请求
        if len(self.waiting_prefill) > self.congestion_threshold:
            self.apply_backpressure()

        return batch

Chunked Prefill: 为什么需要 & 如何计算 chunk size

问题: 长 prefill 阻塞 decode。 Prefill 和 decode 共享 GPU 的 SM (计算单元)。一个大 prefill kernel 提交到 GPU 后独占所有 SM 直到执行完, 后续的 decode kernel 必须排队:

# 无 chunked prefill: 32K tokens 的 prefill 占用 GPU ~1 秒
# GPU timeline:
# |<--- prefill 32K (占满所有 SM, ~1s) --->|<- decode ->|
#                                           ↑ decode 等了 1 秒
# 用户正在打字 → 突然停顿 1 秒 → 体验极差

注意: 不能用多 stream 解决 — 如果一个 kernel 已经占满所有 SM, 另一个 stream 的 kernel 虽然 submit 了也拿不到 SM。Chunked prefill 是在 kernel 边界 做交替, 不是真正的并行。

Chunk Size 推导：

# 约束: 每个 chunk 的执行时间 ≤ TPOT 预算
# chunk_time = chunk_size / prefill_speed ≤ TPOT_budget
# → max_chunk_size = prefill_speed × TPOT_budget

prefill_speed = 30000   # tokens/sec (7B on H100, from 300K/B formula)
tpot_budget = 0.050     # 50ms, decode latency SLO
max_chunk_size = prefill_speed * tpot_budget  # = 1500 tokens

# 各请求切分:
# 100 tokens:   1 chunk (< 1500, 不需要切)
# 500 tokens:   1 chunk
# 2000 tokens:  2 chunks (1500 + 500)
# 8000 tokens:  6 chunks, 总 prefill 时间 = 6×50ms = 300ms
# 32000 tokens: 22 chunks, 总 prefill 时间 ≈ 1.1s
#               但 decode 每 51ms 仍然出 1 个 token!

# 实际时间线:
# [chunk 1500tok][decode 1tok][chunk 1500tok][decode 1tok]...
#     50ms         ~1ms          50ms         ~1ms

调度策略: SJF (Shortest Job First) + 内存感知：

# SJF: 短请求优先 prefill → 最小化平均等待时间
# 内存感知: 如果 KV cache 不够放这个请求, 跳过它

def schedule_prefills(waiting_queue, available_kv_memory):
    # 按 prompt length 排序 (SJF)
    sorted_queue = sorted(waiting_queue, key=lambda r: r.prompt_len)

    for req in sorted_queue:
        kv_needed = req.prompt_len * kv_per_token
        if kv_needed <= available_kv_memory:
            schedule(req)
            available_kv_memory -= kv_needed
        else:
            # 放不下 → 跳过 (不阻塞后面更小的请求)
            continue

Preemption (抢占): 内存不够时怎么办

当 KV cache 内存不足以容纳新请求时, 需要 “踢掉” 某些正在运行的请求来腾出空间。

# 场景: 7B 模型, H100, KV cache 预算 50 GB
# 当前状态:
running_decode = 20  # 正在 decode 的请求, 平均 seq_len=2000
kv_used = 20 * 2000 * 0.5  # MB = 20 GB (0.5 MB/token for 7B MHA)

# 新请求涌入:
new_requests = [100, 500, 2000, 8000, 32000]  # prompt lengths
kv_needed = sum(new_requests) * 0.5  # MB = 21.3 GB

# 当前 20 + 21.3 = 41.3 GB < 50 GB → 暂时够用

# ⚠ 但 decode 会增长! 每个 decode 请求的 seq_len 持续增加:
# 如果每个 decode 最终生成到 4096 tokens:
future_kv = 20 * 4096 * 0.5  # = 40 GB
# 加上 32000 的新请求 (16 GB) → 56 GB > 50 GB → 爆了!

Preemption 策略：

策略	做法	恢复速度	额外开销
Swap to CPU	KV cache 搬到 CPU RAM	快 (PCIe copy back)	需要 CPU 内存 + PCIe 带宽
Recompute	丢弃 KV cache, 恢复时重新 prefill	慢 (重算)	无额外内存, 但浪费算力
Kill	直接终止请求	N/A	用户体验差

谁应该被 preempt？

# 优先 preempt 32000 token 的请求:
# 1. KV cache 最大 (16 GB), 踢掉它释放最多内存
# 2. 还在 chunked prefill, 没产出任何 output → 用户感知代价低
# 3. 短请求已在 decode, preempt 它们 = 用户看到输出中断 → 体验更差
#
# SGLang/vLLM 策略:
#   - LIFO preemption (后来的先被踢)
#   - 优先保护已经产出 output 的请求
#   - 大请求设置低优先级

拥塞控制机制，类似 TCP 的拥塞控制，当系统负载过高时：

监控等待队列长度和 KV Cache 利用率
超过阈值时减少新接纳的 prefill 请求
极端情况下触发 preemption (暂停低优先级请求)
负载降低后逐步恢复接纳速率

4.5 Quantization 概述

Quantization 通过降低权重和/或激活的数值精度来减少模型的内存占用和计算开销。对于 LLM 推理来说，这是在有限 GPU 显存上部署大模型的关键技术。

下面拖动滑块切换精度级别，对比 FP32 / FP16 / FP8 / INT4 的显存占用：

常见量化策略

方法	权重精度	激活精度	特点
W16A16	FP16	FP16	基准，无损
W8A8 (FP8)	FP8 E4M3	FP8 E5M2	H100 原生支持，几乎无损
W4A16 (GPTQ/AWQ)	INT4	FP16	权重量化，激活保持高精度
W4A8	INT4	INT8	极致压缩，需要校准

FP8 量化 (H100 Tensor Core 原生支持)：

# FP8 E4M3: 1 sign + 4 exponent + 3 mantissa (权重)
# FP8 E5M2: 1 sign + 5 exponent + 2 mantissa (梯度/激活)
# 范围: E4M3 → [-448, 448], E5M2 → [-57344, 57344]

def quantize_to_fp8(tensor, dtype="e4m3"):
    # Per-tensor 或 per-channel scaling
    scale = tensor.abs().max() / fp8_max(dtype)
    quantized = (tensor / scale).to(torch.float8_e4m3fn)
    return quantized, scale

def fp8_matmul(x_fp8, w_fp8, x_scale, w_scale):
    # H100 可以直接做 FP8 矩阵乘, 吞吐是 FP16 的 2x
    result = torch._scaled_mm(x_fp8, w_fp8.T,
                              scale_a=x_scale, scale_b=w_scale,
                              out_dtype=torch.bfloat16)
    return result

Mixed Precision 策略。并非所有层都适合量化。通常的策略：

Attention QKV/O 投影：适合 FP8 量化
FFN 层：适合 INT4 权重量化 (参数量大)
LayerNorm：保持 FP32 (对精度敏感)
第一层和最后一层：通常保持高精度

4.6 Speculative Decoding

Speculative Decoding 的核心思想：用一个小而快的 draft model 先生成多个候选 token，然后用大的 target model 并行验证这些 token。因为验证 N 个 token 的开销与生成 1 个 token 几乎相同 (都是一次 forward pass)，所以当 acceptance rate 足够高时可以显著加速。

下面动画演示 draft model 生成 5 个候选，target model 并行验证，接受 3 个 + 1 个修正 token：

验证算法

def speculative_decode(target_model, draft_model, prompt, gamma=5):
    """
    gamma: draft 步数 (每次投机生成的 token 数)
    """
    generated = []

    while not is_finished(generated):
        # Step 1: Draft model 快速生成 gamma 个 token
        draft_tokens = []
        draft_probs = []
        for _ in range(gamma):
            p_draft = draft_model.forward(prompt + generated + draft_tokens)
            token = sample(p_draft)
            draft_tokens.append(token)
            draft_probs.append(p_draft)

        # Step 2: Target model 并行验证所有 draft tokens
        # 一次 forward pass 验证所有位置!
        target_probs = target_model.forward_batch(
            prompt + generated + draft_tokens  # 并行处理 gamma 个位置
        )

        # Step 3: 逐个验证，使用 rejection sampling
        accepted = 0
        for i in range(gamma):
            # 接受概率: min(1, p_target / p_draft)
            accept_prob = min(1.0, target_probs[i][draft_tokens[i]] /
                                     draft_probs[i][draft_tokens[i]])
            if random() < accept_prob:
                generated.append(draft_tokens[i])
                accepted += 1
            else:
                # 拒绝: 从修正分布中采样
                adjusted_dist = normalize(
                    max(0, target_probs[i] - draft_probs[i])
                )
                generated.append(sample(adjusted_dist))
                break  # 后续 draft tokens 作废
        else:
            # 所有 draft tokens 都被接受，bonus: 再采样一个
            bonus_token = sample(target_probs[gamma])
            generated.append(bonus_token)

    return generated

加速比分析

设 acceptance rate 为，draft 步数为，draft model 速度为 target 的倍：

# 期望接受的 token 数: E[accepted] = (1 - α^(γ+1)) / (1 - α)
# 每轮开销: γ * cost_draft + 1 * cost_target
# 加速比 ≈ E[accepted] / (γ * c + 1)

# 例: α=0.8, γ=5, c=0.1 (draft 是 target 的 10 倍快)
# E[accepted] = (1 - 0.8^6) / (1 - 0.8) = 3.36 tokens
# 开销 = 5 * 0.1 + 1 = 1.5 (相对 target 单步)
# 加速比 = 3.36 / 1.5 = 2.24x

期望接受 token 数的闭式公式：

4.7 Constraint Decoding

结构化输出生成

在实际应用中，我们经常需要 LLM 生成符合特定格式的输出（如 JSON、SQL、代码等）。Constraint Decoding 通过在解码过程中动态屏蔽不合法的 token，保证输出始终符合给定的语法规则。

核心方法：

CFG (Context-Free Grammar) 引导：定义语法规则，解码时只允许合法的下一个 token
JSON Schema 约束：根据 JSON Schema 动态计算合法 token mask
Regex 约束：基于正则表达式状态机过滤 token
FSM (Finite State Machine)：将约束编译为有限状态机，高效查找合法 token

class JSONConstraintDecoder:
    """基于 JSON Schema 的约束解码器"""

    def __init__(self, schema, tokenizer):
        self.schema = schema
        self.tokenizer = tokenizer
        # 预编译 schema 为 FSM 状态转移表
        self.fsm = self.compile_schema_to_fsm(schema)
        # 预计算每个状态对应的合法 token mask
        self.token_masks = self.precompute_masks()

    def get_valid_tokens(self, current_state):
        """返回当前状态下合法的 token IDs"""
        return self.token_masks[current_state]

    def apply_constraint(self, logits, state):
        """将非法 token 的 logits 设为 -inf"""
        valid_mask = self.get_valid_tokens(state)
        logits[~valid_mask] = float('-inf')
        return logits

    def decode_step(self, model, input_ids, state):
        logits = model.forward(input_ids)
        constrained_logits = self.apply_constraint(logits, state)
        token = sample(constrained_logits)
        new_state = self.fsm.transition(state, token)
        return token, new_state

SGLang 中的优化：SGLang 将 FSM 状态转移与 Radix Cache 结合——相同约束规则下的前缀可以共享 FSM 状态和 token mask 计算结果，避免重复编译。此外，通过 jump-forward decoding，当 FSM 确定接下来只有一条合法路径时（如 JSON key 的引号），可以跳过模型推理直接输出。

4.8 Model Loading Pipeline

从磁盘到 GPU 的完整流程

加载一个大型模型涉及多个步骤：权重文件解析 → 权重映射 → Tensor Parallel 切分 → GPU 内存分配 → 量化转换。

class ModelLoader:
    def load_model(self, model_path, tp_size, dtype="auto"):
        # 1. 解析模型配置
        config = load_config(model_path / "config.json")

        # 2. 创建模型骨架 (空权重)
        with torch.device("meta"):
            model = LlamaForCausalLM(config)

        # 3. 确定权重映射 (safetensors index)
        weight_map = load_weight_map(model_path / "model.safetensors.index.json")

        # 4. Tensor Parallel 切分计划
        tp_plan = self.compute_tp_plan(config, tp_size)
        # 例: {
        #   "q_proj.weight": ShardDim(0),   # 按行切分
        #   "k_proj.weight": ShardDim(0),   # 按行切分
        #   "o_proj.weight": ShardDim(1),   # 按列切分
        #   "gate_proj.weight": ShardDim(0),
        #   "embed_tokens.weight": Replicate(),  # 复制到所有 rank
        # }

        # 5. 逐 shard 加载 (流式，避免 CPU 内存爆炸)
        for shard_file in weight_map.shard_files:
            tensors = load_safetensors(shard_file)
            for name, tensor in tensors.items():
                # 按 TP plan 切分
                if name in tp_plan:
                    tensor = self.shard_tensor(
                        tensor, tp_plan[name], tp_rank, tp_size
                    )
                # 量化转换 (如 FP16 → FP8)
                if dtype != "auto":
                    tensor = self.quantize(tensor, dtype)
                # 传输到 GPU
                model.load_weight(name, tensor.to("cuda"))

        return model

    def shard_tensor(self, tensor, plan, rank, world_size):
        """按 TP plan 切分 tensor"""
        if isinstance(plan, ShardDim):
            chunks = tensor.chunk(world_size, dim=plan.dim)
            return chunks[rank].contiguous()
        elif isinstance(plan, Replicate):
            return tensor  # 全量复制

关键优化点：

流式加载：逐个 shard 文件加载，避免 CPU 内存峰值
异步 H2D 传输：CPU → GPU 传输与下一个 shard 的 CPU 加载重叠
Safetensors 格式：零拷贝内存映射，按需加载特定 tensor
预计算切分：TP plan 提前确定，避免运行时分配

练习题

练习 1: 计算 KV Cache 内存

给定模型参数 (Llama-70B, num_layers=80, num_kv_heads=8 GQA, head_dim=128, seq_len=4096, batch_size=32, FP16)，计算 KV Cache 内存。问题：(a) 单个请求的 KV Cache 大小？ (b) 整个 batch 的 KV Cache 大小？ (c) 若 GPU 80GB、模型权重 140GB (TP=2, 每卡 70GB)，还剩多少给 KV Cache？最大 batch？

完整解答：

# (a) 单个请求的 KV Cache
single = 2 * 80 * 8 * 128 * 4096 * 2 bytes
  K 或 V 单层: 8 * 128 * 4096 * 2 = 8 MB
  K 和 V 单层: 16 MB
  所有层: 16 MB * 80 = 1,280 MB ≈ 1.25 GB

# (b) 整个 batch
batch_total = 1.25 GB * 32 = 40 GB

# (c) TP=2, 每卡 70GB 权重, 80GB 显存
available = 80 - 70 = 10 GB per GPU
per_request_per_gpu = 1.25 / 2 = 0.625 GB  # KV heads 也分到 2 卡
max_batch = 10 / 0.625 = 16 个请求
# 预留激活内存后, 真实可用约 7-8 GB → realistic_max_batch ≈ 12 个请求

练习 2: Scheduler 策略设计

场景：H100 (80GB) TP=1，7B 模型 (14GB)，KV Cache pool 剩余 50GB，Running batch 20 个 decode 请求 (平均 seq_len=2000)，Waiting queue 5 个 prefill 请求 (prompt_len 100/500/2000/8000/32000)，SLO: TTFT < 2s, TPOT < 50ms。(a) prefill 顺序？(b) chunk size？(c) 是否 preempt？

完整解答：

# (a) SJF + 内存感知
order = [100, 500, 2000, 8000, 32000]  # 按 prompt length 升序

# (b) Chunked Prefill
# 7B prefill 速度 ≈ 30K tok/s on H100, TPOT 50ms
# max_chunk_size = 30000 * 0.05 = 1500 tokens
# 100/500: 1 chunk; 2000: 2 chunks; 8000: 6 chunks; 32000: 22 chunks
chunk_size = 1500

# (c) Preemption 分析
# 当前 decode: 20 * 2000 * 0.5MB ≈ 20 GB
# 新 prefill 总计: (100+500+2000+8000+32000) * 0.5MB ≈ 21.3 GB
# 总需求: 41.3 GB < 50 GB (够用)
# 但 decode 增长到 4096 时: 20 * 4096 * 0.5MB = 40 GB → 超预算!
# 策略: 暂不 preempt, 但给 32000 请求设置低优先级

练习 3: Speculative Decoding 加速比分析

Target model 70B 单次 forward 50ms，Draft model 7B 单次 8ms，，。(a) 期望每轮接受 token 数 (b) 每轮总延迟 (c) 加速比 (d) 时？

完整解答：

# (a) E[tokens] = (1 - α^(γ+1)) / (1 - α)
E_tokens = (1 - 0.75**6) / (1 - 0.75) = 0.822 / 0.25 = 3.29 tokens

# (b) 每轮延迟
round_latency = 5 * 8ms + 1 * 50ms = 90ms

# (c) 加速比
speedup = E_tokens * target_latency / round_latency = 3.29 * 50 / 90 = 1.83x

# (d) α = 0.5
E_tokens_05 = (1 - 0.5**6) / (1 - 0.5) = 1.97 tokens
speedup_05 = 1.97 * 50 / 90 = 1.09x  # 几乎没收益!
# 盈亏平衡点: E_tokens = 1.8 → α ≈ 0.47
# 实际需 α > 0.6 才有显著收益

练习 4: Radix Cache 命中率计算

chatbot 共享 system prompt (1000 tokens)。60% 请求共享前缀 1200，25% 共享 1150，15% 仅共享 1000；平均 user query 300 tokens。(a) 平均 cache hit token 数 (b) 平均新计算 prefill token 数 (c) 节省多少？

完整解答：

# (a) 平均 cache hit tokens
avg_hit = 0.60 * 1200 + 0.25 * 1150 + 0.15 * 1000 = 1157.5 tokens

# (b) 平均新计算
avg_new_compute = 0.60 * 300 + 0.25 * 300 + 0.15 * 300 = 300 tokens

# (c) 节省比例
avg_total = 0.60 * 1500 + 0.25 * 1450 + 0.15 * 1300 = 1457.5 tokens
savings = 1 - 300 / 1457.5 = 79.4%
# Radix Cache 节省约 79% 的 prefill 计算, TTFT 接近 5x 快
# System prompt (1.0 * 1000 = 1000) 被所有请求共享, 最值得缓存

编程练习 A: 实现简单的 Continuous Batching Scheduler

实现支持动态插入和完成请求的 Continuous Batching Scheduler，每个 step 移除 is_done 的请求并从 waiting queue 填充空 slot。

import random

class Request:
    def __init__(self, req_id, prompt_tokens, max_len):
        self.req_id = req_id
        self.prompt_tokens = prompt_tokens
        self.generated_tokens = 0
        self.max_len = max_len
        self.is_done = False

    def step(self):
        self.generated_tokens += 1
        if self.generated_tokens >= self.max_len:
            self.is_done = True

class ContinuousBatchingScheduler:
    def __init__(self, max_batch_size):
        self.max_batch_size = max_batch_size
        self.running = []
        self.waiting = []
        self.completed = []

    def add_request(self, request):
        self.waiting.append(request)

    def step(self):
        # (1) 每个 running request 生成一个 token
        for req in self.running:
            req.step()

        # (2) 移除已完成的请求
        done = [r for r in self.running if r.is_done]
        self.completed.extend(done)
        self.running = [r for r in self.running if not r.is_done]

        # (3) 从 waiting 填充空位
        free_slots = self.max_batch_size - len(self.running)
        while free_slots > 0 and self.waiting:
            req = self.waiting.pop(0)
            self.running.append(req)
            free_slots -= 1

    def get_utilization(self):
        return len(self.running) / self.max_batch_size

# ======== 测试 ========
random.seed(42)
scheduler = ContinuousBatchingScheduler(max_batch_size=4)

requests = []
for i in range(10):
    req = Request(
        req_id=i,
        prompt_tokens=random.randint(50, 200),
        max_len=random.randint(3, 8)  # 短生成方便观察
    )
    requests.append(req)

# 前 5 个立即到达, 后 5 个在 step 5 到达
for req in requests[:5]:
    scheduler.add_request(req)

print("Step | Running | Waiting | Completed | Utilization")
print("-" * 55)

for step in range(20):
    if step == 5:
        for req in requests[5:]:
            scheduler.add_request(req)
        print(f"  -- [Step {step}] 5 new requests arrived --")

    scheduler.step()
    util = scheduler.get_utilization()
    print(f"  {step:2d}  |    {len(scheduler.running)}    |    {len(scheduler.waiting)}    |"
          f"     {len(scheduler.completed)}     |   {util:.0%}")

    if not scheduler.running and not scheduler.waiting:
        print("  -- All requests completed! --")
        break

编程练习 B: 实现 Speculative Decoding 的 Verification 逻辑

给定 draft_probs 和 target_probs，实现逐 token 验证。Accept 条件 rand() < min(1, target_p[x] / draft_p[x])，reject 时从 max(0, target_p - draft_p) 归一化后采样。理论 acceptance rate per token = 。

import numpy as np

def speculative_verify(draft_tokens, draft_probs, target_probs, vocab_size):
    """
    Speculative Decoding verification 算法.
    Returns: accepted_tokens, bonus_token
    """
    accepted_tokens = []

    for i, token in enumerate(draft_tokens):
        q_x = draft_probs[i][token]    # draft prob of chosen token
        p_x = target_probs[i][token]   # target prob of chosen token

        # Acceptance probability
        accept_prob = min(1.0, p_x / max(q_x, 1e-10))

        if np.random.random() < accept_prob:
            accepted_tokens.append(token)
        else:
            # Reject: sample from corrected distribution
            corrected = np.maximum(target_probs[i] - draft_probs[i], 0)
            corrected_sum = corrected.sum()
            if corrected_sum > 0:
                corrected /= corrected_sum
            else:
                corrected = target_probs[i]  # fallback
            bonus_token = np.random.choice(vocab_size, p=corrected)
            return accepted_tokens, bonus_token

    # 所有 draft tokens 都被 accepted, 从 target 的下一个位置采样
    bonus_token = np.random.choice(vocab_size, p=target_probs[-1])
    return accepted_tokens, bonus_token

def theoretical_acceptance_rate(p, q):
    """理论 acceptance rate = sum_x min(p(x), q(x))"""
    return np.minimum(p, q).sum()

# ======== 测试 ========
np.random.seed(42)
vocab_size = 100
K = 5  # draft 生成 5 个 token

def make_distribution(vocab_size, temperature):
    logits = np.random.randn(vocab_size)
    probs = np.exp(logits / temperature)
    return probs / probs.sum()

num_trials = 5000
total_accepted = 0
theoretical_rates = []

for trial in range(num_trials):
    draft_probs = []
    target_probs = []
    draft_tokens = []

    for k in range(K):
        q = make_distribution(vocab_size, temperature=1.5)  # draft: flatter
        p = make_distribution(vocab_size, temperature=0.8)  # target: sharper
        draft_probs.append(q)
        target_probs.append(p)
        draft_tokens.append(np.random.choice(vocab_size, p=q))

    accepted, bonus = speculative_verify(draft_tokens, draft_probs, target_probs, vocab_size)
    total_accepted += len(accepted)

    for k in range(K):
        theoretical_rates.append(theoretical_acceptance_rate(target_probs[k], draft_probs[k]))

avg_accepted = total_accepted / num_trials
avg_theoretical_rate = np.mean(theoretical_rates)

print(f"Draft length K = {K}")
print(f"Empirical average accepted tokens: {avg_accepted:.2f}")
print(f"Theoretical per-token acceptance rate: {avg_theoretical_rate:.3f}")
print(f"Effective speedup ≈ (avg_accepted + 1) / 1 = {avg_accepted + 1:.2f}x")

System Design 面试题 (Anthropic/OpenAI 风格)

以下题目模拟 AI Infra 面试中的系统设计和性能分析问题。要求: 给出数值推导过程, 不只是定性回答。

面试题 1: Capacity Planning — 405B 模型的 serving 集群

部署 405B 模型 (FP8, 405 GB)，SLO: P99 TPOT ≤ 100ms / TTFT ≤ 2s，QPS 200，平均 input 2000 / output 500，GPU H100 80GB。

完整解答：

# (1) 模型放置
min_gpus = 405 / 80  # = 5.06 → 最少 6 张, 推荐 TP=8 (对齐节点)

# (2) TPOT
tpot_theory = 50.6 / 3350 * 1000  # = 15.1 ms
tpot_actual = tpot_theory / 0.65    # ≈ 23 ms << 100ms SLO ✓

# (3) Prefill 速度
prefill_speed = 8 * 1979e12 * 0.55 / (810e9)  # ≈ 10,750 tok/s
ttft_2000 = 2000 / 10750  # = 186 ms → P99 约 500ms-1s, 满足 2s ✓

# (4) Decode 吞吐
# 同时 decode 请求数 = 200 × (500/43) ≈ 2326 concurrent
# 每 TP group batch 64 → 需要 2326/64 ≈ 37 个 TP groups
num_replicas = 37

# (5) 总 GPU 和成本
total_gpus = 37 * 8  # = 296 H100s
annual_cost = 296 * 2 * 24 * 365  # = $5.19M / year
# 加 20-30% buffer → ~380 GPUs, ~$6.7M/year

面试题 2: Decode Batching 的天花板在哪?

LLaMA-70B (BF16) on 1×H100，推导 batch size 的三个上限。

完整解答：

# 上限 1: Memory (KV cache 装不下)
# 单卡放不下 → TP=2: 每卡 70 GB model, 剩 10 GB for KV
kv_per_seq = 0.32 * 4096  # = 1.28 GB per request
max_batch_memory = 10 / 1.28  # ≈ 7-8 requests ← 最紧!

# 上限 2: Compute (变成 compute-bound)
# B_critical = FLOPS / (Bandwidth × 2) = 989e12 / (3.35e12 × 2) ≈ 148

# 上限 3: Latency SLO
# 100ms = 2×70e9×B / 989e12 → B = 706 (远大于内存限制)

# 结论: max batch ≈ 7 (被 KV cache memory 限死)
# TPOT vs batch:
# B=1:   TPOT=35ms,  吞吐=28 tok/s   (memory-bound)
# B=4:   TPOT=35ms,  吞吐=114 tok/s
# B=148: TPOT=35ms,  吞吐=4229 tok/s (理论极限, 但 memory 早不够)

面试题 3: Prefill-Decode Disaggregation 的 Trade-off

Mooncake/DistServe 把 Prefill 和 Decode 放到不同 GPU 集群，分析 KV cache 传输代价。

完整解答：

# (1) KV cache 传输延迟 (70B GQA, IB 400 Gbps = 50 GB/s)
kv_size = 2000 * 320 / 1024  # = 625 MB
transfer_time = 625 / (50 * 1024) * 1000  # = 12.2 ms (+ RDMA ~14ms)

# (2) 对 TTFT 的影响
# co-located TTFT = 2000 / 4300 = 465 ms (70B TP=8)
# Disaggregated = 465 + 14 = 479 ms (+3%, 可接受)

# (3) 什么时候值得?
# ✓ decode 多, prefill 抢 SM; prefill 长度变化大; decode 用更高带宽 GPU
# ✗ KV cache 大 (MHA + 长序列); QPS 低; 短请求为主

面试题 4: 为什么 Speculative Decoding 在高 batch size 下收益递减?

完整解答：

# 核心: spec decoding 增加 compute, 减少 decode steps
# 但 verify 阶段 (K tokens 并行) 是 compute → 高 batch 下变瓶颈

# (1) Verify 计算量: [B×K, h] × [h, 4h] → compute 增大 K 倍
# (2) effective B' = B × K > B_critical → verify 变 compute-bound
#     例: B=64, K=5 → B'=320 > 148
#     临界 B_crossover = B_critical / K = 148 / 5 ≈ 30

# (3) 自适应策略
def should_use_spec_decoding(current_batch_size, draft_k, b_critical):
    if current_batch_size * draft_k > b_critical:
        return False
    if current_batch_size < b_critical * 0.3:
        return True
    return True, max(1, b_critical // current_batch_size)

面试题 5: 设计一个 Multi-turn Chat 的 KV Cache 策略

70B GQA (KV per token 320 KB)，4×H100 (TP=4, 30 GB/卡 KV)，平均 8 轮对话，每轮 user 200 + assistant 500 tokens。

完整解答：

# (1) 单 session KV cache
tokens_8_rounds = 8 * (200 + 500)  # = 5600 tokens
kv_per_session = 5600 * 320 / 1024 / 1024  # = 1.71 GB
# TP=4 → 每卡 0.43 GB/session → max sessions = 30/0.43 ≈ 70

# (2) Eviction: LRU + session 价值加权
#   expected_cost = cost_to_evict × probability_of_return
#   优先 evict 短 session + idle 久的

# (3) Re-prefill 代价
# 5600 tokens at 4300 tok/s = 1.3s TTFT (不可接受)
# 优化: swap to CPU = 1.71 GB / 64 GB/s = 27ms (远好于 re-prefill)

# (4) Prefix caching (system prompt 1000 tokens)
# 每 session 省 312 MB; 70 session 总省 5.25 GB/卡 → 多放 12 session
# 节省 17.5% 内存, 新 session TTFT 省 233ms

面试题 6: 量化对 Serving 性能的影响分析

对比 BF16 / FP8 / INT4 部署 LLaMA-70B。

完整解答：

# 模型大小:
# BF16: 140 GB → TP=2;  FP8: 70 GB → TP=1;  INT4: 35 GB → TP=1

# TPOT (batch=1, H100 3.35 TB/s):
# BF16 TP=2: 70/3.35 = 20.9ms → ~32ms (MBU 65%)
# FP8 TP=1:  70/3.35 = 20.9ms → ~32ms  ← 和 BF16 相同!
# INT4 TP=1: 35/3.35 = 10.4ms → ~16ms
# ⚠ decode 是 memory-bound, FP8 的 2× compute 优势用不上
#   FP8 真正优势: TP=1 (省 GPU) + 更大 KV cache 空间

# 最大 batch (seq_len=4096, KV=320KB/token, 1.31 GB/request):
# BF16 TP=2: 10/1.31 ≈ 7;  FP8 TP=1: 7;  INT4 TP=1: 45/1.31 ≈ 34 (4.8×)

# Cost-per-quality:
# FP8: $0.0046/1K tokens ← quality-sensitive 场景的甜蜜点
# INT4: $0.0005/1K tokens (quality≈95%) ← 高吞吐场景 / draft model

面试题 7: 突发流量下的降级策略

正常 100 QPS 突然飙到 500 QPS (5× spike) 持续 2 分钟，无自动扩容。设计优雅降级方案。

完整解答：

# 立即生效 (< 1s):
# 1. 减小 max_output_tokens (2048 → 512): KV cache 释放更快
# 2. 增大 chunked prefill chunk_size: 优先消化现有 decode
# 3. Admission control: 长请求返回 429, 短请求仍接纳

# 短期措施 (10-30s):
# 4. 动态量化: 部分 replica BF16 → FP8 (释放 KV cache 内存)
# 5. 减少 speculative decoding 的 K (K=5 → K=2 或关闭)
# 6. 低价值请求 early stopping

# 请求分级:
# Tier 1 (保护): 付费用户, multi-turn 中间轮
# Tier 2 (降级): 新 session, 长 prompt, 非付费
# Tier 3 (丢弃): 重试请求, 超过 5s 排队的请求

# Spike 结束后恢复:
# 1. 不立即恢复 max_output_tokens → 等队列 drain
# 2. 逐步增加接纳率 (类似 TCP slow start)
# 3. 被 preempt 的 session 优先 re-prefill
# 4. P99 TPOT 恢复 SLO 内后切回正常模式

本文是 ML Systems 系列 Chapter 4。正文 markdown 渲染，5 个交互动画通过自定义 {% anim %} 标签以隔离 iframe 嵌入，源自 Arkive 教程。

LLM 推理系统：SGLang 深度解析

4.1 请求生命周期 (Request Lifecycle)

各阶段详解

Prefill vs Decode 的本质区别

性能指标: TTFT 与 TPOT

Prefill 速度推导: 从 GPU FLOPS 到 tokens/sec

Decode 速度推导: Memory-Bandwidth Bound

KV Cache 大小估算

GPU 显存到底分给了谁

各项的关键 caveat

实际显存账: 70B 在 8 × 80GB H100 上

什么时候要调小 gpu_memory_utilization?

一句话直觉

activation 到底”长”在 Transformer 哪里

FlashAttention 抹掉哪一档

边界处的两个”非 layer” activation

训练 vs 推理: 哪些 activation 活着, 哪些被回收

一句话直觉

4.2 Continuous Batching vs Static Batching

Continuous Batching 核心实现

性能对比

4.3 KV Cache 管理

PageAttention: KV Cache 的虚拟内存

Radix Cache: 前缀共享树

Radix Cache 实现原理

4.4 Scheduler 设计

Prefill 与 Decode 的调度矛盾

Chunked Prefill: 为什么需要 & 如何计算 chunk size

Preemption (抢占): 内存不够时怎么办

4.5 Quantization 概述

常见量化策略

4.6 Speculative Decoding

验证算法

加速比分析

4.7 Constraint Decoding

结构化输出生成

4.8 Model Loading Pipeline

从磁盘到 GPU 的完整流程

练习题

System Design 面试题 (Anthropic/OpenAI 风格)

什么时候要调小 `gpu_memory_utilization`?