Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve

model 2403.0231
llm-inferenceschedulingchunked-prefillthroughput-latencyroofline
← model Survey Source ↗

Sarathi-Serve — L2 深度蒸馏 #

1. TL;DR #

在线 LLM serving 中,prefill(计算受限)与 decode(内存受限)交织导致吞吐-延迟二选一。Sarathi-Serve 用 chunked-prefill + stall-free batching:把长 prefill 切成受 token budget 约束的小块,piggyback 到 decode 批次而不打断 decode。Mistral-7B 提升 2.6×、Falcon-180B 提升 6.9× serving capacity。

2. Q1 / Q2 / Q3 #

Q1 — 痛点 (痛点) #

现有 LLM serving 调度器在吞吐与延迟之间存在硬性权衡,根源是每个请求分两阶段:

两类现有调度器各有致命缺陷:

Q2 — 方法 (方法) #

利用 decode 批次的算术强度 slack(memory-bound 时算力空闲),把 prefill 拆块塞进 decode 批次:

  1. Chunked-prefills:长 prompt 的 prefill 分成多个 chunk,跨多个 iteration 计算(来自前作 Sarathi)。
  2. Stall-free batching(Algorithm 3):先按 SLO 算出每个 batch 的 token budget τ;每个 iteration 先装满所有正在跑的 decode,再放至多一个正在进行的 prefill chunk,最后在剩余预算内接纳新请求。decode 永不被 prefill 打断。
  3. Token budget 决策:budget 越小 TBT 越低,但过度切块引入两种 overhead — (a) GPU 利用率降低,(b) attention 反复读取前序 chunk 的 KV-cache。还须避开 tile-quantization 效应(batch token 数须为 tile size 的倍数)。

    4. 核心技术壁垒:并非"切块"这个动作,而是把 token budget 校准到 decode iteration 的内存受限拐点之下——既榨干 decode 的算力 slack、又不越过 compute-bound 阈值触发 TBT 爆炸(naive full-prefill hybrid 会把 TBT 抬高 28.3×)。同时预算须落在 tile-size 倍数上,否则单个多余 token 就带来一整块 tile 的浪费(257 vs 256 → +32% prefill 时间)。这个"贴着两个物理边界走"的联合校准是最难复制的洞见。

    Q3 — 结果 (结果) #

    • Serving capacity(满足 SLO 前提下的最大 QPS)全面超越 Orca 与 vLLM。
    • 严格 SLO 下最高 3.2× vs Orca(Mistral-7B, openchat);宽松 SLO 下最高 2.7× vs vLLM(Yi-34B, openchat)。
    • Falcon-180B 端到端 6.9×,主要得益于均匀 iteration 计算减少 pipeline bubble。
    • 消融显示 chunked-prefill 与 hybrid-batching 二者互补,单用任一都会在某一维度退化。

    3. 架构 / 方法图 #

    Sarathi-Serve 的核心是调度策略的对比,而非新的模型结构。下图给出四种调度策略在同一时间线上的行为对比(A/B 处于 decode,C/D 随后到达):

    Figure 4: 四种调度策略的时间线对比

    Paper's Figure 4. vLLM 把尽可能多的 prefill 排在恢复 decode 之前,产生 generation stall;Orca 支持 hybrid batch 但含长 prompt 的批次执行时间仍高,无法消除 stall;FasterTransformer 跑完 decode 才排新 prefill,无 stall 但批量小、吞吐低;Sarathi-Serve 把 C 的 prefill 切成 $p_0,p_1$ 两块,与 A/B 的 decode 合批,既不停 decode 也不停 prefill。

    宏观的两难关系由 Figure 2 概括:

    Figure 2: 吞吐-延迟权衡与调度策略

    Paper's Figure 2. 优先 prefill → 高吞吐但牺牲 TBT 尾延迟;优先 decode → 反之。Sarathi-Serve 通过 stall-free batching 同时拿到高吞吐与低 TBT。读者应注意 Sarathi-Serve 位于两条曲线的"帕累托前沿之外"。

    调度算法的控制流(Algorithm 3)可用如下状态流表示:

    flowchart TB A[新 iteration: n_t ← 0] --> B[遍历 B 中已完成 prefill 的请求
    装入所有 decode token] B --> C{存在未完成 prefill?} C -->|是| D["取一个 chunk c = get_next_chunk_size(R, τ, n_t)
    n_t += c"] C -->|否| E[接纳新请求循环] D --> E E --> F{"can_allocate ∧ n_t < τ ?"} F -->|是, c>0| G[加入新请求 chunk
    n_t += c] G --> F F -->|否 / c=0| H[process_hybrid_batch] H --> I[filter_finished_requests] I --> A

    顺序:先塞满 decode(保证 decode 不被延迟)→ 再放一个 ongoing prefill chunk → 最后在 τ 剩余预算内接纳新请求。

    4. 作者证明 #

    本文以实证为主,唯一形式化模型是 §3.2 的 roofline。给出符号表与物理意义,并做最小检查。

    符号表 #

    符号含义
    $T$算子总执行时间
    $T_{\text{math}}$数学运算耗时
    $T_{\text{mem}}$从 HBM 取数耗时
    $\tau$每 batch 的 token budget
    $N$一个 prompt 被切成的 chunk 数

    方程物理意义 #

    roofline 执行时间模型:

    $$T = \max(T_{\text{math}}, T_{\text{mem}})$$

    物理意义:matmul kernel 把访存与计算 overlap,取二者较慢者(用 $\max$ 而非求和)。当 $T_{\text{math}} < T_{\text{mem}}$ 为 memory-bound(decode);当 $T_{\text{math}} = T_{\text{mem}}$ 时算力与带宽利用率同时最大化,对应最优算术强度 = 设备 FLOPS/带宽比。

    chunked-prefill 的 KV-cache 重读代价:若 prompt 切成 $N$ 块,则第 $k$ 块的 KV-cache 被加载 $N-k$ 次(第 1 块读 $N-1$ 次,以此类推),计算量不变但 HBM 读增加。

    最小检查(6 项) #

    1. 量纲一致性:$T$、$T_{\text{math}}$、$T_{\text{mem}}$ 均为时间量纲,$\max$ 保持量纲一致 ✓。
    2. 极限行为:batch → 单 decode token 时 $T_{\text{math}} \ll T_{\text{mem}}$,$T \approx T_{\text{mem}}$,与"decode memory-bound"一致 ✓。
    3. 拐点自洽:理论上 A100 约 200 token 转 compute-bound,实测因高 TP 固定开销推迟到 ~500-600 token(footnote 2),模型给出的是理想上界 ✓。
    4. KV 重读求和:总额外读取 $\sum_{k=1}^{N-1}(N-k) = N(N-1)/2$,随 $N$ 二次增长,解释为何"过度切块"有惩罚 ✓。
    5. 单调性:token budget ↓ → 每 iteration prefill token ↓ → TBT ↓,但 chunk 数 ↑ → overhead ↑,存在权衡 ✓。
    6. 实测锚点:Fig 5 显示 1 个 decode token 的 linear 代价 ≈ 128 prefill token,印证 $T_{\text{mem}}$ 主导下"塞进上百 prefill token 几乎免费" ✓。

      7. 5. 实验与数据 #

      Prefill vs Decode 的 batching 行为(动机基础) #

      Figure 3: prefill/decode 吞吐 vs batch size

      Paper's Figure 3 (Mistral-7B, A100, prompt=1024). decode 吞吐随 batch size 近线性增长,prefill 单请求已近饱和。这是 Takeaway-1,也是整个 hybrid-batching 论证的地基:decode 有 batching slack,prefill 没有。

      算术强度与线性算子占比 #

      Figure 5: prefill/decode 各算子耗时

      Paper's Figure 5. linear 层主导 prefill 与 decode 运行时;由于 decode 算术强度极低,1 个 decode token 的 linear 代价 ≈ 128 个 prefill token —— 这正是"可免费 piggyback prefill"的量化依据。

      Figure 7: 线性层执行时间 vs token 数

      Paper's Figure 7 (LLaMA2-70B, 不同 TP). token 数少时执行时间被 HBM 取权重主导,在 128-512 区间几乎持平(尤其高 TP);越过临界阈值后随 token 数线性上升。这条曲线定义了 token budget 应落在的"平台区"。

      Chunk 合批的增量代价 #

      Figure 8: coalescing prefill 到 decode 的增量代价

      Paper's Figure 8. Decode+Full-Prefill(Orca)对 decode 延迟冲击巨大;Decode+Chunked-Prefill(Sarathi-Serve)在固定 token budget 下冲击小得多,且 decode batch 越大、context 越长,相对影响越小。

      Capacity 主结果 #

      Figure 9: Mistral-7B / Yi-34B capacity

      Paper's Figure 9. 严格(SLO-S)与宽松(SLO-R)两档下,Sarathi-Serve 的可持续 QPS 一致高于 Orca/vLLM。关键现象:Orca/vLLM 往往在达到最大吞吐前就已违反 P99 TBT SLO。

      Figure 10: LLaMA2-70B / Falcon-180B capacity

      Paper's Figure 10. 带 pipeline parallelism 的大模型上,均匀 iteration 计算额外压缩了 pipeline bubble,Falcon-180B 端到端提升达 6.9×。

      组件消融(互补性) #

      Table 4(Yi-34B, TP2, 128 请求, budget 1024)显示两个组件单用都会在某维度退化,合用最佳:

      Scheduleropenchat P50 TTFTopenchat P99 TBTarxiv P50 TTFTarxiv P99 TBT
      hybrid-batching-only0.561.084.181.76
      chunked-prefills-only1.040.344.860.38
      Sarathi-Serve (combined)0.850.293.030.35

      hybrid-only 因长 prefill 仍造 stall → TBT 高;chunked-only 因 chunk 略低效 → TTFT 高;合用两维度都改善。

      6. 论证链 #

      #论证步骤依据(paper 内部)
      1prefill 计算受限、decode 内存受限,batching 只对 decode 有效§2.2 Takeaway-1,Fig 3
      2因此现有交织调度必然在吞吐与 TBT 间二选一(prefill-prioritizing 造 stall;decode-prioritizing 吞吐低)§3.1 Takeaway-2,Fig 4
      3decode 批次处于 memory-bound regime,存在算力 slack 可塞入更多 token§3.2 Takeaway-3,Fig 5/6/7
      4直接塞整段 prefill 会因越过 compute-bound 阈值把 TBT 抬到 28.3×,故须 chunk 并限制 token budget τ§4.2,Fig 8
      5在 τ 约束下先装 decode 再装 prefill chunk(stall-free),decode 永不被延迟§4.2 Algorithm 3
      6结果:在 SLO 内 capacity 全面超越 baseline,大模型因 bubble 缩减额外获益§5.1,Fig 9/10,Table 4

      7. 实现 cross-reference #

      实现建立在 vLLM 开源代码之上(PyTorch + xFormers 提供 matmul/attention kernel,NCCL 负责 PP/TP 通信),扩展了调度策略、chunked-prefill、pipeline parallelism 与 telemetry(§5 Implementation 段)。论文正文未给出行级代码路径,故:

      • 调度算法:见 Algorithm 3(§4.2,L1 line ~350),compute_token_buget、get_next_chunk_size 为关键函数,论文未公开其源码行号 → [实现未公开]。
      • 参考开源实现:项目对应 sarathi-serve / 已上游到 vLLM chunked-prefill 路径,论文未在正文给出 file:line 锚点 → [实现未公开]。

      核心技术壁垒(§7 展开):最难复制的是 token budget 的联合校准。它必须同时满足三个约束——(1) 落在 decode 内存受限"平台区"以内(Fig 7),越界即触发 TBT 线性爆炸(naive full-prefill 达 28.3×);(2) 是 GPU tile size 的整数倍,否则 tile-quantization 让单个多余 token(257 vs 256)带来 +32% prefill 时间;(3) 满足应用的 P99 TBT SLO。论文用一次性 profiling + 取 SLO 内最大可打包 token 数来定 τ,但把它调到"贴着物理边界又不越界"依赖对具体模型/硬件的细致标定,不是照搬公式就能得到。

      关键实现细节(易漏):

      1. 调度顺序不可换:必须先装满所有 decode,再放至多一个 ongoing prefill chunk,最后才接纳新请求;顺序保证 decode 从不被 prefill 挤占(Algorithm 3 lines 6-20)。
      2. KV-cache 重读代价随 N 二次增长:$N$ 块切分总额外 HBM 读为 $N(N-1)/2$;但因 attention 只占总时间小头,过度切块的主要惩罚其实来自 linear 层利用率下降与更频繁的 allreduce 固定开销(Fig 11,§5.2)。

        3. Appendix: 模型架构图 #

        N/A — 本文为在线 serving 调度系统(Sarathi-Serve),而非模型发布(model release)。所评测的 Mistral-7B / Yi-34B / LLaMA2-70B / Falcon-180B 均为已有开源模型(均采用 decoder-only transformer,GQA/GQA-SW attention,见 Table 1),本文未提出新的模型结构,故不附代码驱动的模型架构图。

        评测模型的关键配置(§5,Table 1):

        ModelAttentionGPU 配置总显存(每卡)
        Mistral-7BGQA-SW1 A10080GB (80GB)
        Yi-34BGQA2 A100 (TP2)160GB (80GB)
        LLaMA2-70BGQA8 A40 (TP4-PP2)384GB (48GB)
        Falcon-180BGQA4 A100 × 2 节点 (TP4-PP2)640GB (80GB)