《从 0 实现 SGLang》第 1 篇 · LLM 推理引擎到底在做什么

千行代码一步步搭出一个现代 LLM 推理引擎吃透大模型推理的每一项关键技术。本阶段目标 — 最简推理实现用最朴素的方式把端到端推理跑通先搭起整体框架再逐个模块替换为完整实现。整个阶段共 5 篇短文序号主题1序章LLM 推理引擎到底在做什么本篇2核心数据结构Req与SamplingParams3一个可运行的 Decoder Layer一次性版本4Naive KV Cache5单 GPU Engine v0 与 greedy 生成本篇你将学到一句话能说清LLM 推理引擎在做什么理解 prefill 与 decode 为何要分两步KV cache 为何是核心看到推理引擎 4 类进程的全景图作为后续阶段展开的起点获得一张系列阅读路线表可按自己关心的方向选择性阅读。本篇不涉及实现只界定问题与目标。1. LLM 推理引擎到底在做什么一句话定义给定一段 prompt逐 token 生成续写直到遇到停止条件。把这一过程包装成能以请求形式被外部高效调用的服务系统便是 LLM 推理引擎。与训练框架的区别很多人第一次接触 LLM 推理时会以为加载模型然后调用model.generate()即可。这在演示场景下尚可成立但在生产场景下远远不够。维度训练框架HF transformers / Megatron推理引擎SGLang / vLLM输入固定长度的 batch任意长度、随时到达的请求流forward 次数每个 batch 一次每个请求 N 次N 输出长度显存关注点激活值 梯度KV cache梯度根本不存在时间预算多分钟一个 step 可以接受单 token 的 P99 延迟以毫秒计并发模型数据并行模型并行上面这些 请求级调度training-vs-inference它要解决的真问题吞吐量每秒处理多少请求、延迟首 token 多快出、token 间隔多稳、显存利用率KV cache 不爆、公平性长请求不能饿死短请求。怎么解决把一次 forward 跑模型拆成 prefill 与 decode 两个阶段、引入 KV cache、把请求打包成 batch、用 scheduler 决定每一步谁能上、用分页显存避免碎片、用 CUDA Graph 消除 launch 开销——这正是本系列后续 8 个阶段要逐一构建的东西。本篇不展开任何一个机制只把问题边界画清楚。2. Prefill 与 Decode为什么必须分两步一个朴素假设初次接触推理时常会产生这样的疑问既然 transformer 是一次 forward为何不能在 prompt 后拼接 N 个待生成的 PAD 一并送入一次性输出 N 个 token答案藏在 causal mask 里。Transformer 的自注意力中第 k 个位置只能看到前 k-1 个位置这是因果性的硬约束——你必须先生成第 k 个 token才能为第 k1 个 token 提供 K/V。换句话说生成阶段的 token 之间存在数据依赖串行无法绕开。因此自然分为两个阶段Prefill把整个 prompt 一次性送入模型。所有 prompt token 之间没有未来依赖问题它们是已知输入可以并行算 attention。并行度 prompt 长度通常几十到几千算力密集。Decode每一步只生成 1 个新 token。新 token 需要对全部历史 KV 算 attention但自己只有一个 query。并行度 batch 内同时在 decode 的请求数通常 16~256显存带宽密集。工程上的连锁后果两个阶段在硬件瓶颈、kernel 实现、调度策略上几乎完全不同这是这套引擎后续所有复杂度的根源kernel 选择不同prefill 用flash_attn_varlen_func变长拼一起算decode 用flash_attn_with_kvcache带 cache 的特殊路径。后文会展开。调度策略不同prefill 的约束在于避免显存溢出token budgetdecode 的约束在于避免 batch 过小带宽利用率。后文会展开。优化重点不同prefill 阶段 7PU 一般已经满载性能上限是算力decode 阶段 7PU 经常空转性能上限是 CPU launch 开销。这也是后续引入 CUDA Graph 和 overlap scheduling 的原因。拓展PD 分离既然 prefill 和 decode 在硬件瓶颈、kernel 选型、调度策略上都不同能不能把它们放到不同机器上这就是PD 分离Prefill-Decode DisaggregationSplitwise / DistServe / MoonCake 等近期工作的核心思想prefill 池算力密集用 H100/H200 这类高 FLOPs 卡算完 prompt 把 KV cache 通过 RDMA / NVLink 传给 decode 池decode 池带宽密集用显存大、带宽高但 FLOPs 不必那么强的卡只做 token 逐个生成好处主要有四个干扰彻底消除prefill 的算力大锤不再卡住 decode 的毫秒级 token 流TTFT首 token 延迟和 ITLtoken 间间隔可以分别优化硬件可以错配prefill 节点配高算力卡、decode 节点配大显存卡每一类各用所长整体 $ / token 更优独立扩缩容长输出场景下 decode 工作量远大于 prefill可以单独扩 decode pool 而不动 prefillSLO 解耦两个池的指标互不干扰prefill 池专注吞吐、decode 池专注尾延迟代价节点间要高速传 KV cache依赖 RDMA / NVLink switch调度系统要跨机 routing 容灾复杂度明显上升小规模部署反而不划算。本系列只构建单机引擎PD 分离不展开——但理解为什么会出现这种架构对后面读真实产线系统很有帮助。下面这张图把两阶段的负载差异画出来可作为后续讨论的原点。3. KV Cache核心的核心它存在的根本原因回看 decode 的每一步算 attention(Q_new, K_all, V_all)。其中 Q_new 是当前这个新 token 的 queryK_all、V_all 是从 prompt 第 1 个 token 到上一个生成 token 的全部 key 和 value。关键观察K_all 和 V_all 的前 n-1 项在上一步已经算过了没有任何理由再算一次。把它们存下来、增量追加新一项即可——这就是 KV cache。再深入一层为什么位置 i 的 K_i、V_i 不会随时间变化因为它们都是 hidden state 的线性投影——K_i W_k · h_i、V_i W_v · h_i。而每一层 transformer 的 h_i 又来自 attention(Q_i, K_{0…i}, V_{0…i})——causal mask 让它只依赖位置 ≤ i跟未来生成的 token完全无关。换句话说一旦位置 i 的 K_i、V_i 写入 cache后续无论生成多少新 token 都不会改它。这就是 KV cache 安全可重用的根本依据。复杂度差异有多大设隐藏维度为 d生成第 k 个 token 时需要的计算量操作不加 cache加 cacheQ 投影O(k · d²)O(1 · d²)K/V 投影O(k · d²)O(1 · d²) 仅新 tokenattention 矩阵 QKᵀO(k² · d)O(k · d)输出投影O(k · d²)O(1 · d²)单步总计O(k² · d)O(k · d)生成 N 个 token 总计O(N³ · d)O(N² · d)上面这张表里每一项的复杂度都能从矩阵形状直接读出来——(m×p) × (p×n)的乘法成本是m · n · p从 N³ 降至 N²。当 N1000 时差距达到 1000 倍。这并非一项普通优化而是让 LLM 推理在工程上可行的前提。下文的示例代码将直观呈现这一差距。它带来的新问题伏笔KV cache 解决了一个根本瓶颈但同时引入了两个新的工程问题这两个问题贯穿后面几个阶段显存占用爆炸Llama-3 8B 在 4096 上下文下单条请求的 KV cache ≈ 2 GB。100 个并发请求就吃掉 200 GB——而一块 H100 只有 80 GB 显存。这是后续引入 Paged KV Cache的动机。跨请求的浪费很多生产场景里所有请求共享同一份 system prompt。这部分 KV 完全可以复用而不是每条请求都从头算一遍。这是后续引入 Radix Cache的动机。本篇不展开这两个机制只埋下伏笔——为后续的分页与 prefix tree 留好动机。 示例不加 cache 与加 cache 的复杂度对比 用一个 2 层、hidden4096、vocab100 的简化 GPT 跑两种生成方式 A) naive — 每生成 1 个 token 都重新跑整段序列 B) cached — prefill 一次后decode 每步只算 1 个新 token 观察 wall-time 差距随生成长度增长是否符合理论 naive ∈ O(N³)cached ∈ O(N²) importtimeimporttorchimporttorch.nnasnnimporttorch.nn.functionalasF torch.manual_seed(42)# 有 GPU 就跑 GPU没有就退回 CPU示例模型的参数很小非推理开销占比大实测倍数可能远小于理论值devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)HIDDEN_DIM,NUM_HEADS,VOCAB_SIZE4096,32,100# hidden4096 接近真实 LLM 尺度HEAD_DIMHIDDEN_DIM//NUM_HEADS# 128classDemoAttn(nn.Module):def__init__(self):super().__init__()self.qkvnn.Linear(HIDDEN_DIM,3*HIDDEN_DIM,biasFalse)self.out_projnn.Linear(HIDDEN_DIM,HIDDEN_DIM,biasFalse)defforward(self,x,kv_cacheNone):# x 形状: [batch, seq, hidden]kv_cacheNone 表 prefill否则 decode (seq1)batch,seq,_x.shape q,k,vself.qkv(x).chunk(3,dim-1)# 拆头[batch, seq, hidden] → [batch, heads, seq, head_dim]qq.view(batch,seq,NUM_HEADS,HEAD_DIM).transpose(1,2)kk.view(batch,seq,NUM_HEADS,HEAD_DIM).transpose(1,2)vv.view(batch,seq,NUM_HEADS,HEAD_DIM).transpose(1,2)# 关键decode 时把新 K/V 追加到历史 cache 上——这就是 cache 复用的本质ifkv_cacheisnotNone:cached_k,cached_vkv_cache ktorch.cat([cached_k,k],dim2)vtorch.cat([cached_v,v],dim2)new_kv_cache(k,v)attn_outF.scaled_dot_product_attention(q,k,v,is_causal(kv_cacheisNone))attn_outattn_out.transpose(1,2).reshape(batch,seq,HIDDEN_DIM)returnself.out_proj(attn_out),new_kv_cacheclassDemoGPT(nn.Module):def__init__(self):super().__init__()self.embeddingnn.Embedding(VOCAB_SIZE,HIDDEN_DIM)self.layersnn.ModuleList([DemoAttn()for_inrange(2)])self.lm_headnn.Linear(HIDDEN_DIM,VOCAB_SIZE,biasFalse)defforward(self,token_ids,kv_cachesNone):xself.embedding(token_ids)new_kv_caches[]fori,layerinenumerate(self.layers):x,layer_cachelayer(x,Noneifkv_cachesisNoneelsekv_caches[i])new_kv_caches.append(layer_cache)returnself.lm_head(x),new_kv_cachesdefgenerate_naive(model,prompt,max_new_tokens):# 每一步把整段 token_ids 重新过一遍模型——完全不复用历史 K/Vseqpromptfor_inrange(max_new_tokens):logits,_model(seq)seqtorch.cat([seq,logits[:,-1:].argmax(-1)],dim1)returnseqdefgenerate_cached(model,prompt,max_new_tokens):# prefill 一次性写入 cache之后 decode 每步只把新 token 喂进去logits,kv_cachesmodel(prompt)last_tokenlogits[:,-1:].argmax(-1)output_pieces[prompt,last_token]for_inrange(max_new_tokens-1):logits,kv_cachesmodel(last_token,kv_cacheskv_caches)last_tokenlogits[:,-1:].argmax(-1)output_pieces.append(last_token)returntorch.cat(output_pieces,dim1)defsync():# CUDA op 是异步的timing 前要同步CPU 上是空操作ifdevice.typecuda:torch.cuda.synchronize()modelDemoGPT().to(device).eval()prompttorch.randint(0,VOCAB_SIZE,(1,32),devicedevice)MAX_NEW_TOKENS64withtorch.no_grad():# warmup避开第一次 forward 的 kernel 编译 / cache 分配开销generate_naive(model,prompt,4)generate_cached(model,prompt,4)sync()t0time.perf_counter();generate_naive(model,prompt,MAX_NEW_TOKENS);sync();t1time.perf_counter()t2time.perf_counter();generate_cached(model,prompt,MAX_NEW_TOKENS);sync();t3time.perf_counter()print(fdevice :{device})print(fprompt_len32, max_new_tokens{MAX_NEW_TOKENS})print(fnaive :{t1-t0:.3f}s)print(fcached :{t3-t2:.3f}s ({(t1-t0)/(t3-t2):.1f}x faster))/opt/miniconda3/lib/python3.12/site-packages/torch/cuda/__init__.py:61: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. import pynvml # type: ignore[import] device : cpu prompt_len32, max_new_tokens64 naive : 4.900 s cached : 3.152 s (1.6x faster)4. 推理引擎的核心模块到此为止单条请求内部发生的事情已经讲完。但工程上面对的从来都不是单条请求而是并发的请求流。单进程为何不够理论上可以但实际上不可行原因有三Python GILGlobal Interpreter LockCPython 同一时刻只允许一个线程执行 Python 字节码。tokenizer / Detokenizer / HTTP 处理都是 CPU 密集的纯 Python 代码一旦它们抢到 GILScheduler 主循环就只能等——GPU 因此空转。GPU 调度需要独占主循环Scheduler 必须在每个 step 之后立刻决定下一个 batch 怎么组——任何阻塞都会让 GPU 空转。故障隔离tokenizer 抛异常不应该让整个引擎挂掉。因此自然分成 4 类进程我们要构建的推理引擎把整个系统拆成 4 个相互独立的 Python 进程靠 ZMQ 消息总线相互连接API Server接 HTTP 请求、对接 OpenAI 兼容接口、维护 SSE 流式连接SSE Server-Sent Events单向的 HTTP 推送让客户端能边收 token 边显示Tokenizer Worker把文本切成 token idCPU 密集独立进程不阻塞 GPU 调度Scheduler拥有 GPU。维护请求队列、KV cache 池、Engine 实例每一步决定哪些请求进入 batchDetokenizer Worker把生成的 token id 增量解码回文本要处理多字节 UTF-8 的边界问题独立成进程便于隔离。后续阶段都在拆这张图本系列的大部分阶段本质上都是在展开其中某一个框后续大部分阶段都发生在Scheduler 框内KV cache、batch、调度、CUDA Graph、Radix、TP直到系列末段才把API Server / Tokenizer / Detokenizer三个框展开讲它们怎么和 Scheduler 通信、SSE 怎么实现、/v1/chat/completions怎么对齐 OpenAI 协议。下方这张图给出整体地图先把握其结构内部细节后续逐一展开。5. 系列阶段一览整个系列分 9 个阶段1~9每个阶段以一个可运行的里程碑为终点阶段主题里程碑1最简推理实现单卡 greedy 输出 100 token2Batch 与 Attention 后端单卡同时跑 16~32 条请求3模型层与权重加载直接 load HF 模型Llama / Qwen4Scheduler 与连续 Batching完整引擎接近 SGLang 基线吞吐5CUDA Graph 与 Overlapdecode GPU 利用率拉到 90%6Radix Cache跨请求 prefix 复用吞吐 2~10x7Tensor Parallelism4 卡跑 Qwen3-32B8多进程架构与 OpenAI Server一行命令启动 OpenAI 兼容服务9进阶MoE、kernel、benchmark对标 SGLang 的能力清单每个阶段末尾建议停下来通读自己写完的代码再进入下一阶段。小结用一句话回到开篇所立的定义LLM 推理引擎就是把prompt → 续写这件事包装成一个能并发服务、显存利用率高、延迟可控的系统。它的核心结构有三层算法层prefill / decode 两阶段 KV cache把生成 N 个 token 的复杂度从 O(N³) 压到 O(N²)调度层batch、scheduler、paged cache、prefix tree让多请求高效共享 GPU架构层4 类进程 ZMQ 消息总线让 CPU、网络、调度互不阻塞。接下来的每个阶段都是在把这三层中的某一块换成完整实现。下一篇预告2 核心数据结构Req与SamplingParams——下一篇将引入用于表示一条请求的Req结构看它如何用cached_len/device_len等几个长度字段同时支撑 prefill 和 decode 两阶段并理解 greedy 快路径的判定条件。