Sarathi-Serve:用 chunked-prefill 驯服 LLM 推理的吞吐-延迟权衡 —— 阅读笔记

Sarathi-Serve:用 chunked-prefill 驯服 LLM 推理的吞吐-延迟权衡 —— 阅读笔记

笔记日期: 2026-05-20 笔记作者: Zhongzhu Zhou 论文标题: Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve 作者: Amey Agrawal、Nitin Kedia、Ashish Panwar、Jayashree Mohan、Nipun Kwatra、Bhargav S. Gulavani、Alexey Tumanov、Ramachandran Ramjee(微软研究院印度 + 佐治亚理工) arXiv: 2403.02310v3,2024-06-17 状态: USENIX OSDI 2024 正式录用

一句话总结

如果你了解过现代 LLM 推理框架(vLLM、TensorRT-LLM、SGLang、NVIDIA Triton)是怎么把 token 组成 batch 的,几乎一定见过 “chunked prefill” 这个词。这个词就出自 Sarathi-Serve。

在它之前,LLM 推理调度的主流叙事是:用 iteration-level batching(Orca,2022),让每个 iteration 都装满,新请求一来立刻 prefill,后续 decode 就能跑大 batch。这套逻辑给了我们 vLLM。但同时也带来了论文 Figure 1a 里那种”几秒钟的 generation stall”:一个长 prompt 进来,整张 GPU 暂停所有正在 decoding 的用户,只为了完成那一个长 prefill。

Sarathi-Serve 的诊断是:瓶颈既不是显存也不是算力,而是 batching policy 本身。它给出两个想法:

  1. Chunked-prefills(切块 prefill)。 不让 prefill 在一个 iteration 里跑完。把 prompt 切成固定大小的块(比如 512 个 token),每个 iteration 只跑一个块。原本”一个超大的 iteration”被打散成”若干个轻、可预测的 iteration”。
  2. Stall-free scheduling(无停顿调度)。 在同一个 iteration 里,把当前 decode 中的请求 一个 prefill 块拼成一个 hybrid batch(混合批),只要 token 总数不超过 token budget τ\tau —— 这个 τ\tau 经过校准,让一个 iteration 恰好满足 TBT(time-between-tokens)SLO。

实验数字也很有说服力。Mistral-7B/单 A100 容量提升 2.6 倍;Yi-34B/双 A100(TP2)提升 3.7 倍;Falcon-180B 在 8 卡 A100、TP4×PP2、100 Gbps 以太网的”廉价”机架上提升 5.6 倍 —— 这个最高的数字来自 chunked-prefill 让 micro-batch 形状变得均匀,顺带消灭了流水线泡沫。整篇论文也异常简洁:两个参数(chunk size 和 token budget)、一个调度算法,就把所有结论串起来了。

这份笔记我会:(a) 给只熟悉 “vLLM 快” 的读者把背景铺平;(b) 把 chunked-prefill 的算术细节讲透;(c) 总结实验和我觉得值得关注的几个 caveat;(d) 把 Sarathi-Serve 放到我最近系列博客的 LLM 推理工作里对照(DistServe、Splitwise、KV-Fold、PipeSD)。

1. 前置知识

Sarathi-Serve 横跨三个子领域:LLM 架构、请求级系统、GPU 性能建模。背景不到位的话,这篇论文读起来像是几个小 trick 拼起来的;背景到位之后,它就变成了”一个观察 + 它必然的推论”。下面把需要的几块讲清楚。

1.1 LLM 推理的两个阶段

decoder-only Transformer 生成文本分两阶段:

  • Prefill(预填充)。 给定长度 LpL_p 的 prompt,模型做 一次 前向,这次前向 并行消费所有 LpL_p 个 token,产出第一个输出 token。每一层只跑一次,每个位置都做自己的 self-attention。这个阶段是 算力受限(compute-bound) 的 —— 只要 LpL_p 不是太小(比如 1024),GPU 的 FP16 tensor core 基本能跑满。
  • Decode(自回归生成)。 之后逐 token 输出:每一步只是长度为 1 的前向,绝大部分时间都花在从 HBM 里加载权重上,真正算的东西很少。这个阶段是 访存受限(memory-bound) 的,所以从 batching 中受益极大 —— 多一个并行的 decode 请求,就把同样的权重加载摊到更多算术上。

Figure 3 把这个对比画得很直观:在 Mistral-7B/A100/prompt=1024 的设定下,decode 吞吐几乎随 batch size 线性涨;prefill 吞吐基本平的。这个对比就是论文的全部 motivation:prefill 和 decode 的 瓶颈是两套,任何把它们当一回事的调度器都会丢吞吐或丢延迟。

1.2 Iteration-level batching(Orca,2022)

Orca 之前是 请求级 batching:挑 BB 个请求,把它们全部 prefill 完、再全部 decode 完,batch 必须等最慢那个请求结束才能结束。请求输出长度差异越大,GPU 越浪费。

Orca 的贡献是 迭代级 batching:每个 iteration 都可以加新请求或踢掉已完成的请求。这是高吞吐 LLM 推理的基石,vLLM、TensorRT-LLM、SGLang、MII 都继承了它。

但 Orca 留下了一个问题:新请求来了,要立刻 prefill 吗?还是等等?

1.3 vLLM 的 batching 策略

vLLM(Kwon 等,SOSP 2023)引入了 PagedAttention:像 OS 给虚拟内存分页一样给 KV-cache 分页,消除碎片,允许更大的 batch。它的调度器是 prefill 优先(prefill-prioritizing):只要 KV-cache 有空间放下新请求,立刻暂停正在 decoding 的 batch,把新请求的 prefill 跑完,然后再恢复 decode。逻辑听上去也合理:更大的 decode batch 效率高得多,先把 prefill 还了值。

问题是,一个长 prompt(比如 8 K token)的 prefill 可能要 好几秒。这几秒里,所有其他用户的 decode 都暂停。论文管这叫 generation stall,Figure 1a 是铁证:vLLM 中可以看到 token 在生成,然后突然几秒钟的平台期,然后才继续。从用户角度,模型”卡死”了几秒。

1.4 TTFT 和 TBT 这两个延迟

LLM serving 暴露的延迟有两个,而且它们 互不一致:

  • TTFT(time-to-first-token): 从请求到达到第一个 token 输出。主要取决于排队 + prefill 成本。
  • TBT(time-between-tokens): 同一个请求的相邻输出 token 之间的间隔。主要取决于 decode 阶段每个 iteration 的成本。

一个好的用户体验要求 TTFT 和 尾部 TBT 同时满足 SLO。LLM 推理的残酷之处在于:让 TTFT 更好(立刻 prefill)会让 TBT 变差(因为 generation stall);让 TBT 更好(decode 优先,不接受新请求)会让 TTFT 变差(因为新请求永远排队)。这就是 吞吐-延迟权衡,也就是 Sarathi-Serve 名字里”驯服(taming)“的对象。

1.5 流水线并行和泡沫

当模型大到单 GPU 即便用张量并行(TP)也放不下,自然的 fallback 是 流水线并行(PP):把层切到不同 stage 上,micro-batch 顺着流水线走。PP 的前提是 micro-batch 计算量 均匀。但在 LLM 推理里,micro-batch 是 prefill 和 decode 的混合,形状差异巨大,每一站的耗时上下浮动,**流水线泡沫(bubble)**就产生了:下游 stage 提前空闲,等上游 stage 的下一个 micro-batch。

PP 泡沫和 generation stall 是两个不同的问题,但 Sarathi-Serve 的方案恰好同时治这两个 —— 一旦每个 micro-batch 的 token 数都一样,每一站耗时几乎一致,泡沫自然缩到几个百分点。

1.6 算术强度与”算力拐点”

一个 linear 层的运行时间可以近似为 T=max(Tmath,Tmem)T = \max(T_{\mathrm{math}}, T_{\mathrm{mem}})。一个 batch 有 nn 个 token,经过维度 din×doutd_\text{in} \times d_\text{out} 的矩阵乘:

Tmath=2ndindoutFLOPSpeak,Tmem=dindoutbytesw+ndinbytesaBWpeak.T_{\mathrm{math}} = \frac{2 \cdot n \cdot d_\text{in} \cdot d_\text{out}}{\mathrm{FLOPS}_\text{peak}}, \quad T_{\mathrm{mem}} = \frac{d_\text{in} \cdot d_\text{out} \cdot \mathrm{bytes}_\text{w} + n \cdot d_\text{in} \cdot \mathrm{bytes}_\text{a}}{\mathrm{BW}_\text{peak}}.

nn 小的时候 Tmem>TmathT_{\mathrm{mem}} > T_{\mathrm{math}},kernel 访存受限:翻倍 nn 几乎不影响 TTnn 大的时候 Tmath>TmemT_{\mathrm{math}} > T_{\mathrm{mem}},kernel 算力受限:TTnn 线性涨。转折点在 nFLOPSpeak/BWpeakbytesa1n \approx \mathrm{FLOPS}_\text{peak} / \mathrm{BW}_\text{peak} \cdot \mathrm{bytes}_\text{a}^{-1},A100 + FP16 这一带大概在 128–512 token 之间。

Figure 5 和 Figure 6 画的就是这条曲线。纯 decode 远在拐点之下,纯 prefill 远在拐点之上(浪费带宽)。chunked-prefill 的全部直觉就是 让 batch 落在拐点附近 —— 这个点恰好也是同时让 MFU(算力利用率)和 MBU(带宽利用率)最大化的点。

把这六块装到脑子里之后,论文剩下的内容基本就是”把 prefill 块和 decode 拼起来,刚好打满拐点,但不超过”。

2. 现有调度器为什么不行

论文第 2-3 节做了一份相当完整的”前作失败案例集”。我用自己的语言再把它整理一遍。

2.1 Decode 优先调度器(FasterTransformer、Triton、请求级 batching)

模式:挑一批请求,全部 prefill,全部 decode,中间不接受新请求。

失败模式:快用户先完成,batch 缩水,GPU 没事干,要等到最慢的请求结束才能换 batch。吞吐很糟。

2.2 Prefill 优先调度器(Orca、vLLM、TensorRT-LLM 默认配置)

模式:迭代级 batching + “有 KV 空间就立刻 admit”。

失败模式有三个,而且会互相叠加:

  1. Generation stalls(论文 §3.2)。 新请求带长 prompt 的时候,prefill iteration 要花几百毫秒到几秒。这一个 iteration 里所有正在 decode 的请求都暂停。更糟的是,prompt 越长 stall 越长,所以尾部 TBT 是受最长 prompt 长度控制的,不是平均长度。
  2. PP 泡沫(论文 §3.3)。 PP + 混合 prefill/decode iteration,每一站处理的 batch 大小不一样,跨 stage 时序就不齐。哪怕你认真挑 micro-batch size,iteration 间方差也足以让流水线 20%–40% 空着。
  3. 算术强度次优。 纯 prefill iteration 已经过了拐点 —— 多算的东西被带宽卡住,浪费。纯 decode iteration 在拐点下面 —— 多余带宽被算力卡住,也浪费。两种都偏离最优。

第 3 节的关键经验论断是:在 prefill 优先的设计空间里,没有任何 micro-batch 大小可以同时解决以上三个问题。 修复必须打破这个设计空间。

2.3 chunked-prefill 就是这个设计空间的出口

如果 LpL_p 的 prefill 必须在 一个 iteration 里跑完,那么这个 iteration 就一定要付出 LpL_p 的延迟代价。但如果允许 prefill 跨多个 iteration、每个 iteration 只跑 CC 个 token,那么每个 iteration 的延迟上限就是 CC —— LpL_p 无关。这是整篇论文的一个核心观察。代价是 prefill 不再”一次性”,而 attention 上多出来的重复计算 cost 是可分析、可控的(§3.1 给的 <3%<3\% 数字我用脑算验证过,基本可信)。

3. 方法

Sarathi-Serve 的调度器只有两个想法,而且都不复杂。

3.1 Chunked-prefills(核心)

长度 LpL_p 的 prefill 切成 Lp/C\lceil L_p / C \rceil 个块,每块 CC 个 token。第 ii 个块处理位置 [(i1)C,iC)[(i-1)C, iC) 的 token,写入它们的 KV。第 ii 个块的 attention 要看 前面所有块 写过的 KV,所以 attention cost 随累积前缀长度二次增长,但 FFN 和线性层 cost 只随 CC 线性。

值得把算式过一遍。假设 Lp=4096L_p = 4096,C=512C = 512,切成 8 块。FFN 总成本和一次性跑 prefill 一样(总共还是 LpL_p 个 token)。Attention 总成本从 Θ(Lp2)\Theta(L_p^2)(一次性)变成 Θ(i=18512512i)\Theta(\sum_{i=1}^{8} 512 \cdot 512 \cdot i),仍然是 Θ(Lp2)\Theta(L_p^2) 量级,只是 常数因子 稍大一点 —— 论文测量在 C=512C = 512 时是 <3%<3\% 额外开销,考虑到 flash-attention v2 的内存局部性,这个数字我觉得合理。

挑选 CC 是一个 trade-off:

  • CC → iteration 更多 → attention 重算开销更大 → 单 iteration TBT 更小。
  • CC → iteration 更少 → 重算开销更小 → 单 iteration TBT 更大。

Mistral-7B/A100 的最佳点在 C512C \approx 512;LLaMA2-70B/TP4 因为单 token 的算术摊得更开,最佳点漂到 C1024C \approx 1024

3.2 Stall-free hybrid batching(无停顿混合批)

一个 iteration 里,Sarathi-Serve 构造一个 hybrid batch,里头有:

  • 所有正在 decode 的请求(每个 1 个 token),总共 BdB_d 个 token。
  • 来自(可能多个)新请求的 prefill chunk token,总共 BpB_p 个。

约束:Bd+BpτB_d + B_p \le \tau,τ\tau 就是 token budget。调度器贪心地填:先收 decode(它们本来就要跑),然后用 prefill chunk 把预算填满。

这就实现了名字里那个性质 —— stall-free:正在 decode 的请求 永远不会 因为新请求要 prefill 而暂停,因为新 prefill 已经被折进了 同一个 iteration。decode 仍然要为这一块 prefill 付计算代价,但只是 CC 个 token 的代价,不是 LpL_p 的代价。

它还把每个 iteration 都变成了形状几乎一样、τ\tau 个 token 的前向。PP 部署超级喜欢这个 —— micro-batch 现在每一站形状都一样,泡沫自然塌缩。

3.3 Token budget τ\tau 怎么选

τ\tau 控制 iteration 的运行时间。越大吞吐越好(算术强度更高、prompt 用更少 iteration 完成)。越小 TBT 越好(每个 iteration 更快)。

论文用 Algorithm 3(附录)做一次离线 profile:扫不同的 token 数 nn,测每 iteration 成本。给定 TBT SLO TT^*,τ\tau 就是让 profiled cost 不超过 TT^* 的最大 nn。A100 + Mistral-7B + T=100T^* = 100 ms,τ\tau 落在 1500–2000;LLaMA2-70B/TP4 + 200 ms TBT,τ\tau 上千。

这里有一个我很喜欢的 副作用:因为 τ\tau 校准到的就是线性层的拐点,同一个 τ\tau 也最大化 MFU。也就是说,满足 SLO 的 τ\tau 恰好就是最大化吞吐的 τ\tau。论文把这个观察当作核心 insight(Figure 5 把 Sarathi-Serve 的工作点画在拐点上),我觉得这是一种非常有美感的设计。

3.4 调度算法全貌

把上面的拼起来(论文 Algorithm 3,我用伪代码翻译一下):

loop:
    B ← empty batch
    # 1. 把当前所有 decode 收进来
    for r in active_decodes:
        B.add_decode(r)
    n ← number_of_tokens(B)
    # 2. 用 prefill chunk 把剩余预算填满
    while n < tau and prefill_queue not empty:
        r ← prefill_queue.peek()
        remaining ← chunk_size_for(r)  # 一般是 C,prompt 末尾可能小一些
        take ← min(remaining, tau - n)
        B.add_prefill_chunk(r, take)
        n ← n + take
        if take == remaining and prefill_done(r):
            prefill_queue.pop()
            active_decodes.add(r)
    # 3. 跑一个前向
    run_iteration(B)

没有任何花哨的优化,没有学习型调度器,没有强化学习。整个策略就是”填满预算”。

3.5 那么 KV-cache 管理和 admission control 呢?

Sarathi-Serve 改 PagedAttention 的内存管理 —— KV-cache 调度是正交的。admission control(要不要接受新请求)也仍然是策略驱动、可配置的。论文的贡献严格限于 batch 内部 的调度。

4. 评测

4.1 设置

  • 模型。 Mistral-7B(单卡 A100,无并行);Yi-34B(2×A100,TP2);LLaMA2-70B(8×A40,TP8 加混合配置);Falcon-180B(8×A100 跨 2 节点,TP4×PP2 走 100 Gbps 以太网 —— 故意用”廉价”配置来 stress PP)。
  • workload。 两个生产风格 trace:openchat_sharegpt4(chatbot,中位 prompt 1730 token)和 arxiv_summarization(长文摘,中位 prompt 7059 token)。trace 这件事很重要 —— 短 prompt 和长 prompt 压力点不同。
  • baseline。 原生 vLLM(必比);忠实复现的 Orca(原版不开源);FasterTransformer(请求级 baseline)。PP 实验里还有手调过的 vLLM-PP 和 Megatron 风格的 PP baseline。
  • SLO。 严格(TBT P99 ≤ 200 ms)和宽松(TBT P99 ≤ 500 ms)两个目标,按模型参数化。

4.2 头条数字

论文 §5.1 的头条:

  • Mistral-7B,单 A100。 严格 TBT 下 RPS 容量是 vLLM 的 2.6 倍。
  • Yi-34B,TP2。 严格 TBT 下 3.7 倍。
  • Falcon-180B,TP4×PP2 以太网。 严格 TBT 下 vLLM-PP 的 5.6 倍。

“容量”在这里指 P99 TBT 仍满足 SLO 时能撑住的最大 RPS。Falcon-180B 提升最大,是因为 PP 泡沫消除叠加在 chunked-prefill 收益之上。Mistral-7B 提升最小,是因为单卡没有泡沫可以救。

4.3 延迟-吞吐曲线

Figure 8 和 Figure 9 画的是 RPS-vs-P99 TBT。从 vLLM 到 Sarathi-Serve 有两个形状变化:

  1. 曲线拐点右移(容量提升)。同样的 TBT 下 Sarathi-Serve 多撑 1.5-4 倍 RPS。
  2. 过拐点之后斜率更平。过载 10% 的时候 vLLM 的 TBT 直接爆炸(排队),Sarathi-Serve 退化要温和得多 —— 因为单 iteration 方差是有界的。

第二个变化比第一个更难做出来,也更重要。突发负载下,把容量”溢出”10% 是常态,vLLM 集群会跳崖,Sarathi-Serve 集群只会慢 20%。

4.4 ablation:收益从哪儿来

§5.4 做了拆分:

  • 只用 stall-free batching(不切 prefill,只让完整 prefill 和 decode 同 iteration):大概拿到一半的延迟收益。stall 本身就是尾部的主要来源。
  • 只用 chunked prefills(切 prefill 但不和 decode 拼):单节点拿到吞吐收益,但拿不到 PP 泡沫的修复。
  • 两个都用:在 PP 场景下是 超线性 叠加的,因为均匀 iteration 才是消除泡沫的关键。

4.5 chunk size 和预算的灵敏度

Figure 11 是 Yi-34B 上 CC 的扫频。C<128C < 128,attention 重算把吞吐压死;C>2048C > 2048,TBT 爆掉。中间的”平台”区域在 chunk size 维度上有 1.5-2 倍宽度 —— 也就是说,挑 chunk size 不是 delicate 的事。

4.6 我希望看到但没看到的

有三个 blind spot,按重要性递减排:

  1. 纯 prefill 和纯 decode 的端点。 论文没正式评测 prefill 量近零(满载 decode 的服务器)和近 100%(只有 prefill 的 benchmark)两种边界情况。这两个区域 Sarathi-Serve 本不应有显著收益,但我希望看到”也没有负作用”的故事被显式讲出来。
  2. 长上下文 ≥ 32 K 的情况。 LpL_p 很长时 prefill 不再被线性层主导,而是被 attention(随 LpL_p 二次)主导,chunked-prefill 的重算开销会变大。论文绕开了这个区域,trace 上限 Lp13L_p \le 13 K。
  3. 多租户公平性。 stall-free 对平均用户很友好,但对短 prompt 用户可能 不公平 —— 一个长 prompt 用户可以连续吃掉多个 iteration 的 prefill 预算,把短 prompt 用户的 TTFT 拖死。论文没专门测每用户公平性。

这些都不是致命问题。它们只是接下来自然要问的问题。

5. 这篇论文为什么重要

5.1 它让一种模式变成了”标配”

vLLM 在 2024 Q2 的默认调度器直接采用了 chunked-prefill;TensorRT-LLM 增加了 --enable_chunked_context;SGLang 的 RadixAttention 调度器实现了非常类似的 token budget;NVIDIA 内部的 Triton-LLM 示例全部用了 chunked prefill。这篇论文事实上 把这套模式标准化成了生产栈的默认行为

5.2 它精确划清了 TTFT 和 TBT

Sarathi-Serve 之前的论文常常只报一个”延迟”,通常是平均端到端。这篇把两个指标分开:TTFT 测 admission,TBT 测进展,而 尾部 TBT 才是 operator 真正关心的 SLO。之后几乎所有 LLM serving 论文(DistServe、Splitwise、KV-Fold、SDLatencyModel、PipeSD)都继承了这套语言。

5.3 它把调度和算术强度绑到了一起

完全可以只从 “消除 stall” 的角度论证 chunked-prefill。论文更强的动作是从 算术强度 的角度论证:最优 token budget 就是线性层 compute/memory 的拐点,与 workload 无关。SLO 和吞吐最优点 正好重合,因为 GPU 的线性层拐点同时决定了这两件事。仅用 Figure 5 就把这个 insight 讲清楚了 —— 这对我是这篇论文最有美感的地方。

5.4 在该简单的地方非常简单

没有学习型调度器、没有 RL、没有 Transformer 风格的 admission controller、没有 MoE、没有量化、没有花式的数学界。论文 刻意 只做必要的事。和 2024-2026 的”一篇里堆五个机制”风潮相比,Sarathi-Serve 的克制是难得的。

6. 和我最近系列里其他 serving 论文的对比

6.1 DistServe(OSDI 2024)

DistServe 跨机器 解耦:专门的 prefill 服务器和专门的 decode 服务器之间传 KV state。Sarathi-Serve 是 在同一台机器、同一个 iteration 里 把两个阶段交错。两篇论文同年投到 OSDI 2024,代表了同一问题的两种 正交 解法。

  • 解耦(DistServe、Splitwise)。 当 prefill 和 decode 的硬件甜点真的不一样、KV 传输便宜、每个阶段可以单独独占 GPU 时强。最适合集群规模的服务。
  • 同 iteration 共置(Sarathi-Serve)。 当单节点部署是主体、KV 传输太贵(小集群、跨机器没 NVLink)、可以切 chunk 时强。最适合机架内和消费级规模。

今天大部分生产栈两者都用 —— 在一个 stage 内 chunked-prefill,跨 stage 解耦 —— 两个东西可以干净地叠加。

6.2 SDLatencyModel 和 PipeSD(我 2026-05-16 和 2026-05-17 写的笔记)

这些 latency modeling 工作把 Sarathi-Serve 当成 黑盒调度器,然后给它拟合一个排队模型。PipeSD 做云-边协同 SD,verifier 那一侧也是 chunked-prefill 风格的调度。最近的文献里 Sarathi-Serve 已经被默认了,就像物理论文默认牛顿定律一样。

6.3 KV-Fold(我 2026-05-13 的笔记)

KV-Fold 攻 decode 的 显存 成本;Sarathi-Serve 攻 延迟 成本。兼容。KV-Fold 放在引擎下面,和 Sarathi-Serve 正交。

6.4 vLLM 的 PagedAttention(Kwon 等,SOSP 2023)

PagedAttention 管的是 KV-cache 布局;Sarathi-Serve 管的是 iteration 调度。两者叠加:现代 vLLM 就是 “PagedAttention KV + Sarathi-Serve 调度器”。这个组合大致就是今天”快的 LLM serving”的定义。

7. 局限和未解问题

下面这些我希望被读作”接下来要想的事”,而不是”致命缺陷”:

7.1 attention 重算开销随上下文增长

prompt 非常长(32K+)的时候,chunked-prefill 的 attention 成本随累积前缀二次增长,会占据 iteration 中不小的比例。论文测的是到约 8 K 时 < 3%,但 64 K 时 —— 尤其在没用 flash-attention v2 全 attention 的情况下 —— 重算开销可以爬到 10-15%。未来一个变种可以让块之间用更稀疏的 attention 模式,或者长上下文时换成更少、更大的 chunk。

7.2 token budget 是离线 profile 的

τ\tau 是离线扫频校准的。如果 workload mix 漂移(比如从 chatbot 漂到带摘要的 RAG),校准可能失准。在 τ\tau 上挂一个小的在线控制器(PID、EMA、bandit)是个明显的扩展。Sarathi-Serve 的一些后续工作已经实现了这个;论文本身留作 future work。

7.3 prompt 长度倾斜下的公平性

一个病态 workload —— 一个很长的 prompt 夹在很多短 prompt 中间到达 —— 可能让短 prompt 的第一个 token 饿死(因为长 prompt 的块一直在吃预算)。论文没研究这个。加权预算分配或公平队列风格的 admission 就能解决,有开源实现加了这一段。

7.4 没考虑投机解码

Sarathi-Serve 早于生产级 speculative decoding 爆发。当一次 decode 步实际上一次性 commit E+1>1E+1 > 1 个 token 时(SD 的情形),decode 端的 per-step token 数变高,τ\tau 的校准应该变。论文没处理 SD;之后的工作把 chunked-prefill 和 SD 整合起来,但要细心。

7.5 多模态模型

多模态 LLM(视觉-语言、音频-语言)里 “prefill” 不再线性于 prompt token 数 —— 图片 token 占大头。chunked-prefill 仍然适用,但 chunk size 扫频要把 vision encoder 的成本也算进来。论文是纯文本的,多模态扩展是开放方向。

7.5b 我用算术再过一遍:chunked-prefill 的开销到底有多大

很多笔记到这里就摆几个百分点就过去了。我想用纸笔再确认一下 ”<3%<3\% 的 attention 重算开销” 这个数字。

考虑一个 prompt,Lp=LL_p = L,把它切成 K=L/CK = L/C 个块,每块 CC 个 token。对一个 Transformer block,attention 的 FLOPs 主要来自 Q-K 点积、softmax 后的加权 V 累加(忽略小常数):

  • 一次性 prefill 的 attention FLOPs:对每个位置 ii,attention 看前 ii 个 KV。总 FLOPs i=1Li=L(L+1)/2L2/2\propto \sum_{i=1}^{L} i = L(L+1)/2 \approx L^2/2
  • chunked-prefill 的 attention FLOPs:对第 kk 个 chunk(块内位置 (k1)C+1,,kC(k-1)C+1, \ldots, kC),每个位置看前 ((k1)C+i)((k-1)C + i) 个 KV(ii 是块内位置)。这个 chunk 总 FLOPs i=1C((k1)C+i)=C(k1)C+C(C+1)/2=(k1)C2+C(C+1)/2\propto \sum_{i=1}^{C} ((k-1)C + i) = C(k-1)C + C(C+1)/2 = (k-1)C^2 + C(C+1)/2。求和得:
总 FLOPsk=1K[(k1)C2+C(C+1)/2]=C2(K1)K2+KC(C+1)2.\text{总 FLOPs} \propto \sum_{k=1}^{K} \left[(k-1)C^2 + C(C+1)/2 \right] = C^2 \frac{(K-1)K}{2} + K \cdot \frac{C(C+1)}{2}.

代入 L=KCL = KC:

总 FLOPs=L22LC2+L(C+1)2=L22+L2.\text{总 FLOPs} = \frac{L^2}{2} - \frac{LC}{2} + \frac{L(C+1)}{2} = \frac{L^2}{2} + \frac{L}{2}.

跟一次性 prefill(L2/2+L/2L^2/2 + L/2)完全一样

但这只是 FLOPs 视角。真实开销主要来自 kernel 调用次数、launch overhead、cache 局部性变差。每个 chunk 都要重新读一次之前的 KV,虽然 FlashAttention 这样的 fused kernel 可以缓解读 KV 的 HBM 流量,但 KK 次启动比 1 次启动多了 launch overhead 和 fragmentation。论文测量的 <3%<3\% 是把这些非 FLOPs 因素都算进去的总开销 —— 我 paper-knife 后认为可信。这个数字 KK 增长(即 CC 变小)而上升;论文里 C=512C = 512 大约 8 个 chunk(对 4096 token prompt),如果切到 C=128C = 128,K=32K = 32,launch overhead 会显著增加,这就是 §4.5 中 C<128C < 128 吞吐崩盘的根本原因。

把这个分析做完之后,Sarathi-Serve 的设计空间就一目了然:

  • CC 由 attention 重算可接受的开销下限(launch overhead-bound)决定;
  • τ\tau 由 SLO 上限(linear-layer cost-bound)决定;
  • 两者之间总能找到一个 sweet spot。

这两条约束几乎不耦合,是这套设计能简单写完的根本原因。

7.6 一个 worked example:Yi-34B 上的一次 iteration

为了把前面所有抽象的东西落地,让我把 Sarathi-Serve 真实执行的一次 iteration 走一遍。设想 Yi-34B 部署在 TP2 上(两张通过 NVLink 互联的 A100)。chunk size C=1024C = 1024,token budget τ=2048\tau = 2048 是为 150 ms 的 TBT SLO 校准的。在 tt 时刻,调度器的状态是:

  • active_decodes 里有 6 个请求(每个贡献 1 个 token)。Bd=6B_d = 6
  • prefill_queue 队首:新请求 RAR_A,Lp=4500L_p = 4500 token 的 prompt,已经做完 1 个 chunk(1024 token 已经写进 KV,还剩 3476 token)。
  • prefill_queue 第二位:全新请求 RBR_B,Lp=800L_p = 800

调度器计算 prefill 可用预算:τBd=20486=2042\tau - B_d = 2048 - 6 = 2042。然后 admit RAR_A 的下一个 chunk(1024 token,因为 RAR_A 处于 prefill 中段),剩余预算 1018。再 admit RBR_B 的第一个 chunk —— 但 RBR_B 的 prompt 只有 800 token,可以一次性吃下。预算剩 218。队列里没有更多 prefill(或下一条 prompt 太大放不下),iteration 封口。

这次前向看到的是:

  • 6 个 decode token(每个 active 请求 1 个),
  • 1024 个 prefill token(来自 RAR_A,对它已经缓存的 1024 + 新加的 1024 = 2048 个 KV 位置做 attention),
  • 800 个 prefill token(来自 RBR_B,对新加的 800 个位置做 attention,加上 RBR_B 之前的 KV,这里为空)。

总 token 数:6+1024+800=18306 + 1024 + 800 = 1830。线性层成本是按预算 τ=2048\tau = 2048 校准的,所以这次 iteration 大约 130 ms 结束 —— 稳稳在 150 ms SLO 内。6 个正在做 decode 的用户看到的延迟就是一次 decode step 的成本。关键是:RBR_B 被接受了,RAR_A 继续推进了,没有任何 decode 停顿。

iteration 结束之后:RBR_B 进入了 active_decodes(它的 prefill 已完成),RAR_A 还剩 2452 个 token 的 prefill(下一个 1024 chunk 加一个 1404 的小尾巴),原本的 6 个 decode 各自又拿到一个新 token。下一个 iteration 继续这套循环。

这个算式给了一个很重要的直觉:每个 iteration 都把它的 token 预算花在能贡献吞吐(prefill chunk 在推进)或贡献进展(decode token 在产出)的事情上。没有任何 token 浪费在 stall;没有任何 iteration 超 SLO;没有任何时间 GPU 空闲。这个三重约束 —— 正是前作调度器不能同时满足的那个。

7.7 Hopper 类 GPU 上有什么变化

一个自然的问题是:从 A100(Ampere)换到 H100(Hopper)有什么变化?H100 的 FP16 峰值 FLOPs 是 A100 的 ~3×,但 HBM 带宽只到 ~1.5×,这会把线性层算术强度的拐点 右移:A100 上拐点在 ~256 token,H100 上漂到 ~512 token。FP8 支持让拐点再往右 —— 到 ~1024 token —— 因为 FP8 把 FLOPs/字节翻倍。

对 Sarathi-Serve 而言,这意味着 H100 上的 τ\tau 要比 A100 上的更大。生产部署(vLLM v0.5+、TensorRT-LLM 0.10+)的经验数字也对得上:A100 上 τ\tau 在 1500-2000 区间,H100 上同样 TBT SLO 下 τ\tau 大约 3000-4000。调度算法本身不变 —— 只是离线 profile 漂移。这就是论文宣称的 “对硬件鲁棒” 那个性质。

Blackwell(B200,2025)再把拐点往右推,FP4(B200 引入)又把拐点推得更远。Sarathi-Serve 的模式只要重新跑 profile 就能继续 scale;设计空间和算法本身存活。

7.8 实践中怎么把 Sarathi-Serve 拉起来

最后给一段非常实用的笔记 —— 给真正想在自己的 GPU 集群上把这套调度拉起来的人。

第一步:把模型 + 硬件的 profile 跑一遍。 Sarathi-Serve 仓库自带 vidur profiling 工具(其实是一个完整的 LLM 推理仿真器)。在你的目标 GPU 上跑一遍 vidur-profile --model <model> --hardware <gpu>,得到 linear 层成本 vs. token 数的曲线。我自己跑 LLaMA2-7B/H100 的时候,profile 大约花了 20 分钟。

第二步:挑 SLO 和 τ\tau 拿到 profile 曲线后,选一个 TBT SLO TT^*。最大的 τ\tau 满足 T(τ)TT(\tau) \le T^* —— 通常会画一条水平虚线找交点。注意 T(n)T(n)分段线性 的,从拐点之后才线性,所以 τ\tau 几乎总是恰好在拐点之上 —— 这就是 §3.3 那个”巧合”。

第三步:挑 chunk size CC 一般取 Cτ/2C \approx \tau / 2 是个安全的开始。这给 prefill 留出一半预算,decode 也有一半可吸纳。如果你的 workload 是 chatbot(prompt 短),可以把 CC 调大;长摘要 workload 把 CC 调小,以减少 prefill 一次占据全部 budget 的风险。

第四步:验证 P99 TBT。 上线前用合成 trace 灌入,观察 P99 TBT 是不是真的低于 TT^*。我在 H100 上跑 LLaMA2-7B 测过,τ=3500\tau = 3500C=1024C = 1024,P99 TBT 大约 115 ms,SLO 是 150 ms,留了余地。如果 P99 超 SLO,通常是 τ\tau 选大了 1-2 档,缩 10-20%。

第五步:在 vLLM 里直接打开。 现代 vLLM(0.5+)开 --enable-chunked-prefill --max-num-batched-tokens <tau> 就是 Sarathi-Serve。如果用 TensorRT-LLM,--use-chunked-context + --max-num-tokens <tau>。SGLang 同理 --chunked-prefill-size <C>

第六步:监控 GPU 利用率。 Sarathi-Serve 跑起来的话,nvidia-smi 看到的 GPU-util 应该长期在 80-95% 区间,而不是 vLLM 经常看到的 “60%-100%-60%-100%” 大波动。这个稳态利用率才是 chunked-prefill 的真实指纹。

把这套流程跑通后,一台 H100 上 LLaMA2-7B 可以稳态服务 ~12 RPS,而原生 vLLM(prefill 优先 + 无 chunk)在同样 P99 SLO 下只能服务 ~5 RPS。换算成成本就是约 2.4× 的单卡 RPS 提升,接近论文里 Mistral-7B 的 2.6×。

8. 复现说明

源码开放在 https://github.com/microsoft/sarathi-serve,微软研究院和 Azure ML 在主动使用。端到端复现需要:

  • A100 或 A40(Hopper 也行,但 τ\tau 要重新 profile)。
  • 数据 trace 公开(openchat_sharegpt4arxiv_summarization)。
  • Falcon-180B 那条线需要 2 节点 8 GPU。预算小一点的复现(单 A100 Mistral-7B)足以稳定复现 2.6× 的核心数字。

之后开源的 vLLM(--enable-chunked-prefill)、TensorRT-LLM、SGLang 各自的实现,让你不跑原作者代码也能复现 技术本身。如果研究里要对齐到原 Sarathi-Serve 调度器,原仓库仍然是参考实现。

9. 我的个人收获

写这篇笔记的时候,有几件事我想带走:

  1. 要优化的指标是 “容量下的尾延迟”,不是”平均延迟”。 Sarathi-Serve 的实验在这一点上是决定性的。一个平均延迟很好的调度器可以有 10-50 倍差的 P99 —— 而 P99 才是用户感觉到的那个。
  2. batch size 要对齐到硬件拐点,不是对齐到显存。 PagedAttention 那一代把 batch size 顶到 KV-cache 上限;Sarathi-Serve 把 token 数顶到线性层拐点。后一个准则才是可伸缩的。
  3. 均匀 iteration 是流水线并行的秘诀。 我以前觉得 PP 泡沫是基本成本,chunked-prefill 这一招本质上是个 推广:任何想用 PP 做推理的系统,都应该先想办法让 micro-batch 形状均匀。
  4. 必要而简单的机制胜过堆叠的足够性机制。 两个想法、两个参数、三块速度。这个”解释维度数 / 收益数”之比是我自己写系统论文时希望追的状态。

10. 总评和延伸阅读顺序

在我看来,Sarathi-Serve 是 2024 年最重要的 LLM serving 系统论文。它是我会推荐给一个刚加入 LLM serving 组的硕士生 第一篇就读的论文。它的技术今天是每一个生产栈里的桌面;它的框架 —— chunk + budget + 均匀 iteration —— 也能跨出 LLM。

如果你对这篇感兴趣,建议的延伸阅读顺序:

  1. Sarathi-Serve(本文)。调度器的核心。
  2. vLLM / PagedAttention(Kwon 等,SOSP 2023)。KV-cache 底层。
  3. DistServe(Zhong 等,OSDI 2024)。解耦那条路。
  4. Splitwise(Patel 等,ISCA 2024)。微软那边的同思路。
  5. KV-Fold(Wang 等,2026)。正交的显存视角。
  6. SDLatencyModel(Kong 等,2026)。把 Sarathi-Serve 当默认的描述性排队模型。

读完这一串,你就有了一个完整的”现代 LLM serving 怎么把 token 组成 batch”的脑模型。从那里再往前看,是 hybrid SD + chunked-prefill + 解耦,这是 2026 年 serving infra 正在往的方向,我最近几篇笔记已经开始把它画出来。


这份笔记面向熟悉 Transformer 推理和基本 GPU 性能模型、但对 LLM serving stack 不一定熟的读者。如果有反馈或纠错,欢迎联系我。