深入解析LLM推理服务中Continuous Batching的工作原理,说明Prefill/Decode阶段成本差异,GPU内存带宽瓶颈,以及批量调度如何提升利用率。附step_cost_ms估算模型。
语言模型并不是一口气写出答案的。它运行一次完整的 forward pass 来生成一个 token,将该 token 追加到自己的输入中,然后再运行一次。一条 300 token 的回复需要穿过数十亿参数进行 300 次顺序 passes——而在此之前,包括模型自己在内,没人知道它会是 300,直到序列结束 token 真正出现的那一刻。
因此,你的服务系统必须调度一项无法测量、无法预测、无法提前取消的工作。LLM 服务中的一切困难都源于这一事实。
在 decode 阶段,GPU 不缺算术运算,缺的是内存带宽。为单个用户生成一个 token,它必须把模型中的每一个权重从 HBM 中流出,然后用这些权重做一点点微不足道的数学运算。逐个服务请求几乎浪费了整张卡。
把 32 个序列打包成一批,你只需要流一次这些权重,然后为 32 个 token 复用它们:
def step_cost_ms(batch, prefill_tokens):
return (22.0 # 固定开销:kernel 启动 + 流权重
+ 0.35 * batch # decode:多一个序列几乎不花钱
+ 60.0 * prefill_tokens / 1000)
step_cost_ms(batch=1, prefill_tokens=0) # 22.4 ms -> 1 token
step_cost_ms(batch=32, prefill_tokens=0) # 33.2 ms -> 32 tokens
32 倍的输出,只花了 1.5 倍的时间。吞吐量图上真正追踪的是 batch size,而不是时钟频率。
请求级(静态)批处理对你处理其他工作负载的方式如出一辙:收集 N 个请求,将它们 padding 到一个 tensor 中,运行生成循环直到全部完成,一起返回结果,再取下一批 N 个。
缺陷是结构性的。循环只有在最慢的成员完成时才退出,因此一个在 12 个 token 后发出 EOS 的请求,它所在的 tensor 行会在最长成员需要的数百步中持续被占用。它的槽位在每一次 pass 中都被计算,但什么也没贡献。
这使得静态批处理的 makespan 有一个简洁的闭式——每批的开销等于其最长成员:
def static_makespan(reqs, cap, chunk):
t = 0
for g in [reqs[i:i + cap] for i in range(0, len(reqs), cap)]:
start = max(t, max(r["arrive"] for r in g)) # 等待批次填满
max_pf = max(prefill_steps(r["prompt"], chunk) for r in g)
max_out = max(r["out"] for r in g) # 那个拖后腿的
t = start + max_pf + max_out - 1
return t
来量化这个浪费。一次 forward pass 中一个批槽位被占用称为一个 slot-step,只有当它真正输出了一个 token 时才称为有用:
useful = sum(step["active"] for step in log)
total = cap * len(log)
print(useful / total) # 0.22 静态 / 0.69 连续
在一个现实的的长尾工作负载上、8 个槽位,静态批处理的结果约为 22%。超过四分之三的昂贵 GPU 花费在计算 padding 上。批越大情况越糟,因为更大的组更可能包含一个其他所有人都得等的"巨无霸"成员。
修复方法是把批次不再当作一个工作单元。批次只需要作为一次 forward pass 的单元即可;没有任何要求下一次 pass 必须包含相同的序列。
这在 OSDI 2022 的 Orca 论文中以迭代级调度(iteration-level scheduling)为名发表。vLLM、TGI、TensorRT-LLM 和 SGLang 都实现了它,通常称之为 continuous batching 或 in-flight batching。核心思想是三个阶段:
def run_continuous(reqs, cap, chunk):
slots, t, qi, done = [None] * cap, 0, 0, 0
while done < len(reqs):
# 1. ADMIT - 从队列头部填满每一个空闲槽位
for s in range(cap):
if slots[s] is None and qi < len(reqs) and reqs[qi]["arrive"] <= t:
slots[s] = new_state(reqs[qi], chunk); qi += 1
# 2. STEP - 对驻留序列恰好执行一次 forward pass
for r in filter(None, slots):
r.advance() # prefill chunk 或吐出一个 token
# 3. RETIRE - EOS 立即释放槽位,而非等到本批结束
for s in range(cap):
if slots[s] and slots[s].finished:
slots[s] = None; done += 1
t += 1
一个序列在第 t 步结束,就在第 t+1 步被替换。这就是全部的诀窍。在默认跟踪条件下——64 个请求、8 个槽位、长尾输出长度——它需要 1 122 次 forward pass 而不是 3 470 次,利用率从 22% 提升到 69%,吞吐量乘以 3.11 倍。同样的权重、同样的硬件模型、同样的请求。
Prefill 一次读取整个 prompt 并并行计算其 key/value tensors:数千个 token 的矩阵乘法,是计算密集型且昂贵的。Decode 每个序列每步生成一个 token:极少量的数学运算,以流权重为主,是带宽密集型且每多一个序列几乎不花钱。
简单混合两者就会得到经典病症——一个用户粘贴了一篇 4 000 token 的文档,它的 prefill 落进了一个 decode 迭代中,其他所有用户的流在那一步明显冻结。再看成本模型:step_cost_ms(batch=16, prefill_tokens=4000) 是 268 ms,而同样批次不带 prefill 只有 27.6 ms。
Chunked prefill 将 prompt 切片成固定大小的块,这样就没有单独一步会被某一个用户的 prompt 独占:
def prefill_steps(prompt, chunk):
return max(1, math.ceil(prompt / chunk))
# vLLM: enable_chunked_prefill=True, max_num_batched_tokens=512
由此产生的一个有用细节:最后一个 prefill 块已经产出了第一个输出 token,这就是为什么 TTFT 本质上等于排队时间加上 prefill 时间。
很容易认为 batch size 受算术运算限制。其实不是——一个 decode 步骤几乎注意不到多出 10 个序列。真正用尽的是内存。每个驻留序列都持有一个 KV-cache,每一层、每一个注意力头每看到一个 token 就增长一个条目,这些都存在存放权重的同一块 HBM 中:
def kv_bytes_per_token(layers, kv_heads, head_dim, dtype=2):
return 2 * layers * kv_heads * head_dim * dtype # K 和 V
per_token = kv_bytes_per_token(32, 8, 128) # Llama-3-8B bf16 -> 128 KiB/token
max_tokens = (0.92 * 80e9 - 16e9) / per_token # ~ 450,000 tokens
cap = int(max_tokens // 2048) # ~ 220 个并发 2k-token 对话
这个数字就是你的最大批大小。这也是为什么 PagedAttention 和 continuous batching 是一个故事的两半:调度器决定谁跑,paging 决定能容纳多少。
而且由于输出长度未知,一个在入队时能容纳的批可能在 20 步后溢出。服务器此时必须 preempt——将受害者的 KV 换出到主机内存,或者丢弃它并在重试时重新计算 prefill(通常比一次 PCIe 往返更便宜)。Preemption 是一个正确性机制,而非调优选项,但 preempt 计数器持续上升意味着你过度承诺了。
Continuous batching 不会让 GPU 更快。它只是阻止了输出长度方差对 GPU 的浪费。把槽位上限设为 1,或者让每个请求长度完全相同,加速比就会精确塌陷到 1.00——因为在那两种情况下静态批处理根本不浪费什么。演示中的 skew sweep 直白地展示了这一点:skew 0 时 ×1.00,skew 0.3 时 ×2.31,skew 0.6 时 ×3.11。
这里也有真实的代价。更多序列驻留意味着每次 forward pass 都略微变慢,因此单个请求在舰队整体加速三倍的同时可能在空闲机器上更晚完成。吞吐量和单请求延迟是不同的目标,调度器在优化前者。请汇报 goodput:即满足了你 TTFT 和 token 间延迟目标的每秒请求数。
把 max_num_seqs 调高直到 p95 TTFT 打破你的 SLO,然后停手。
下面的页面把整个过程作为一个真实的离散事件模拟器运行——一条带种子的一次到达和未知输出长度的跟踪,通过两种调度器回放,并有一个实时的 slot×step 占用时间线,展示静态批处理下的 padding 和连续批处理下的紧凑排列。60 条断言用独立的闭式和暴力实现验证了该引擎:https://dev48v.infy.uk/ai/days/day60-continuous-batching.html